diff --git a/aimux-core/src/error.rs b/aimux-core/src/error.rs index e5796b5af..789ca4846 100644 --- a/aimux-core/src/error.rs +++ b/aimux-core/src/error.rs @@ -328,18 +328,23 @@ impl AiMuxError { return None; }; let headers = detail.response_headers.as_ref()?; + // Rust's float parser accepts `inf`, `Infinity` and `NaN`. A + // non-finite value is no hint: `as i64` would saturate it to + // `i64::MAX`, and upstream never uses it either (it fails the + // `ms < 60000 || ms < backoff` window). headers .get("retry-after-ms") .and_then(|value| value.parse::().ok()) - .filter(|value| !value.is_nan()) + .filter(|value| value.is_finite()) .map(|value| value as i64) .or_else(|| { headers.get("retry-after").and_then(|value| { value .parse::() .ok() - .filter(|value| !value.is_nan()) - .map(|seconds| (seconds * 1_000.0) as i64) + .map(|seconds| seconds * 1_000.0) + .filter(|milliseconds| milliseconds.is_finite()) + .map(|milliseconds| milliseconds as i64) .or_else(|| { httpdate::parse_http_date(value).ok().map(|date| { match date.duration_since(SystemTime::now()) { @@ -473,6 +478,37 @@ mod tests { assert_eq!(serde_json::to_string(&error).unwrap(), json); } + fn hint_from(name: &str, value: &str) -> Option { + AiMuxError::ApiCall(Box::new(ApiCallError { + response_headers: Some(HashMap::from([(name.into(), value.into())])), + ..api_error("rate limited") + })) + .retry_after_hint() + } + + /// `str::parse::` accepts `inf`, `Infinity` and `NaN`; none of them + /// is a delay, so the hint is absent instead of `i64::MAX`. + #[test] + fn non_finite_retry_after_is_no_hint() { + for value in ["Infinity", "inf", "-Infinity", "NaN"] { + assert_eq!( + hint_from("retry-after-ms", value), + None, + "retry-after-ms: {value}" + ); + assert_eq!( + hint_from("retry-after", value), + None, + "retry-after: {value}" + ); + } + // Finite seconds that overflow once converted to milliseconds. + assert_eq!(hint_from("retry-after", "1e308"), None); + // Finite values keep working. + assert_eq!(hint_from("retry-after-ms", "250"), Some(250)); + assert_eq!(hint_from("retry-after", "1.5"), Some(1_500)); + } + /// A 429 is retryable and carries its hint wherever it lives — the /// classification is the `status_code` field, not a variant. #[test] diff --git a/aimux-core/src/recording.rs b/aimux-core/src/recording.rs index 65e5fce7d..79d7ddf25 100644 --- a/aimux-core/src/recording.rs +++ b/aimux-core/src/recording.rs @@ -572,6 +572,9 @@ pub fn new_call_id() -> String { /// 精确匹配常见 bearer-token 字段和 AWS Bedrock sigv4 的 /// `x-amz-security-token`(name 已 lowercase 归一化)。其余 needle 维持 contains,以覆盖 /// `x-goog-api-key`/`proxy-authorization` 等变体。 +/// +/// Provider auth headers whose names match none of these are listed by name: +/// Prodia's `X-Prodia-Key` and the Azure `Ocp-Apim-Subscription-Key`. #[must_use] pub fn is_sensitive_key(name: &str) -> bool { let n = name.to_ascii_lowercase(); @@ -583,6 +586,8 @@ pub fn is_sensitive_key(name: &str) -> bool { n == "cookie" || n == "set-cookie" || n == "key" + || squashed == "xprodiakey" + || squashed.ends_with("subscriptionkey") || squashed.ends_with("token") || n.contains("authorization") || squashed.contains("apikey") @@ -1925,6 +1930,9 @@ mod tests { "Authorization", "x-api-key", "apiKey", + "X-Prodia-Key", + "Ocp-Apim-Subscription-Key", + "ocp_apim_subscription_key", ] { assert!(is_sensitive_key(key), "{key} must be redacted"); } diff --git a/aimux-ffi/aimux-ffi.h b/aimux-ffi/aimux-ffi.h index c51524221..cc897dccd 100644 --- a/aimux-ffi/aimux-ffi.h +++ b/aimux-ffi/aimux-ffi.h @@ -453,6 +453,16 @@ aimux_error_t *aimux_apply_tool_call_repair_to_result(const char *result_json, */ void aimux_drop_handle(uint64_t handle); +/** + * aimux_drop_handle with the handle passed as a pointer-sized value, for a + * finalizer that can only call `void (*)(void *)` (Dart NativeFinalizer). + * The pointer is never dereferenced: its address is the handle. Safe to call + * with NULL or an unknown handle (no-op). + * + * @param handle Handle from aimux_*_new, cast to a pointer. + */ +void aimux_drop_handle_ptr(void *handle); + /** * Free a string written by any aimux function (`*out_json`, error getters). * (Handles are released with aimux_drop_handle, errors with diff --git a/aimux-ffi/src/lib.rs b/aimux-ffi/src/lib.rs index e3e596c81..de46346a2 100644 --- a/aimux-ffi/src/lib.rs +++ b/aimux-ffi/src/lib.rs @@ -2134,6 +2134,19 @@ pub extern "C" fn aimux_drop_handle(handle: u64) { } } +/// [`aimux_drop_handle`] with the handle passed as a pointer-sized value, for +/// callers that can only register a `void (*)(void *)` callback, such as a +/// Dart `NativeFinalizer`. The pointer is never dereferenced: its address is +/// the handle. Calling `aimux_drop_handle` through a `void (*)(void *)` +/// pointer instead passes a garbage high word on 32-bit targets, where a +/// `uint64_t` argument does not travel like a pointer. +/// +/// Safe to call with NULL or an unknown handle (no-op). +#[unsafe(no_mangle)] +pub extern "C" fn aimux_drop_handle_ptr(handle: *mut c_void) { + aimux_drop_handle(handle.addr() as u64); +} + /// Free a C string previously returned by any aimux function that hands out /// `*mut c_char` — result JSON, error getters. /// @@ -4463,6 +4476,24 @@ mod tests { ); } + #[test] + fn drop_handle_ptr_releases_the_handle_its_address_names() { + let handle = mock_handle("mock", "finalized", "F"); + let other = mock_handle("mock", "kept", "K"); + assert!(get_handle(handle).is_some()); + + aimux_drop_handle_ptr(std::ptr::without_provenance_mut(handle as usize)); + assert!(get_handle(handle).is_none()); + assert!( + get_handle(other).is_some(), + "only the named handle is released" + ); + + // NULL is handle 0: a no-op, like aimux_drop_handle(0). + aimux_drop_handle_ptr(std::ptr::null_mut()); + aimux_drop_handle(other); + } + #[test] fn composites_reject_any_dead_member_handle() { // A dead handle inside the array is a caller bug: reported, never diff --git a/aimux-ffi/tests/exports_smoke_test.rs b/aimux-ffi/tests/exports_smoke_test.rs index 8d7ce9148..ead8ab414 100644 --- a/aimux-ffi/tests/exports_smoke_test.rs +++ b/aimux-ffi/tests/exports_smoke_test.rs @@ -30,7 +30,7 @@ //! //! ## Coverage //! -//! All 129 `#[unsafe(no_mangle)]` exports in `src/lib.rs` are exercised (the +//! All 130 `#[unsafe(no_mangle)]` exports in `src/lib.rs` are exercised (the //! constructor and utility classes in full; the session class one //! representative call per export). [`header_and_exports_agree`] pins the //! count against the two headers. Decision exports are exercised by @@ -51,13 +51,13 @@ use aimux_ffi::{ aimux_azure_new_with_base, aimux_bedrock_new, aimux_bedrock_new_with_base, aimux_codex_refresh, aimux_cohere_embedding_new, aimux_cohere_embedding_new_with_base, aimux_cohere_new, aimux_cohere_new_with_base, aimux_cohere_reranking_new, aimux_cohere_reranking_new_with_base, - aimux_consume_stream_text, aimux_drop_handle, aimux_embed, aimux_error_available_tools, - aimux_error_code, aimux_error_free, aimux_error_message, aimux_error_model_id, - aimux_error_model_type, aimux_error_original_error, aimux_error_provider_code, - aimux_error_provider_id, aimux_error_provider_message, aimux_error_response_body, - aimux_error_retry_ms, aimux_error_retryable, aimux_error_status, aimux_error_t, - aimux_error_tool_input, aimux_error_tool_name, aimux_file_upload, aimux_free_string, - aimux_generate_object, aimux_generate_text, aimux_generate_text_as_openai, + aimux_consume_stream_text, aimux_drop_handle, aimux_drop_handle_ptr, aimux_embed, + aimux_error_available_tools, aimux_error_code, aimux_error_free, aimux_error_message, + aimux_error_model_id, aimux_error_model_type, aimux_error_original_error, + aimux_error_provider_code, aimux_error_provider_id, aimux_error_provider_message, + aimux_error_response_body, aimux_error_retry_ms, aimux_error_retryable, aimux_error_status, + aimux_error_t, aimux_error_tool_input, aimux_error_tool_name, aimux_file_upload, + aimux_free_string, aimux_generate_object, aimux_generate_text, aimux_generate_text_as_openai, aimux_generate_text_result_as_openai, aimux_get_model_specs, aimux_google_embedding_new, aimux_google_embedding_new_with_base, aimux_google_image_new, aimux_google_image_new_with_base, aimux_google_video_new, aimux_google_video_new_with_base, aimux_image_generate, @@ -158,7 +158,7 @@ fn header_and_exports_agree() { exports.sort(); assert_eq!( exports.len(), - 129, + 130, "export count changed; update the headers" ); @@ -550,9 +550,11 @@ fn constructor_exports_build_and_release_handles() { assert_ne!(*handles.last().unwrap(), 0, "abort_signal_new"); // Release every handle created above (aimux_drop_handle is safe with 0). - for h in [traced, audited, router, moa] { + // The pointer-sized entry point (a finalizer's callback) releases one. + for h in [traced, audited, router] { aimux_drop_handle(h); } + aimux_drop_handle_ptr(std::ptr::without_provenance_mut(moa as usize)); // mock_replay_new: empty input must fail cleanly (no recordings) — the // success path needs a real Recording, covered by replay tests elsewhere. diff --git a/aimux-provider-utils/src/download_guard.rs b/aimux-provider-utils/src/download_guard.rs index c009a608e..b57b28261 100644 --- a/aimux-provider-utils/src/download_guard.rs +++ b/aimux-provider-utils/src/download_guard.rs @@ -35,10 +35,9 @@ pub(crate) fn hop_trusted_origin<'a>( trusted_origin.filter(|origin| same_origin(current_url, origin)) } +/// The URL as an error message may show it: no query, no userinfo. fn without_query(parsed: &url::Url) -> String { - let mut redacted = parsed.clone(); - redacted.set_query(None); - redacted.to_string() + crate::http::sanitized_request_url(parsed.as_str()) } /// Parse a host as `url::Url::host_str` yields it into an IP literal. The url @@ -349,6 +348,16 @@ mod tests { } } + #[test] + fn rejected_url_error_omits_userinfo_and_query() { + let error = validate_download_url("ftp://user:secret@example.com/file?token=t0k") + .expect_err("ftp must be blocked"); + let message = error.to_string(); + assert!(message.contains("ftp://example.com/file"), "{message}"); + assert!(!message.contains("secret"), "{message}"); + assert!(!message.contains("t0k"), "{message}"); + } + #[test] fn resolved_address_validation_fails_closed_and_preserves_all_answers() { let empty = validate_resolved_download_addresses("empty.example", Vec::new()) diff --git a/aimux-provider-utils/src/http.rs b/aimux-provider-utils/src/http.rs index 9d8723414..e9e304776 100644 --- a/aimux-provider-utils/src/http.rs +++ b/aimux-provider-utils/src/http.rs @@ -999,12 +999,53 @@ fn request_host(url: &str) -> String { .unwrap_or_else(|| "unknown".into()) } +/// The request URL as logs, errors and recordings show it: without the query +/// (API keys travel there) and without the `user:password@` userinfo. pub(crate) fn sanitized_request_url(url: &str) -> String { match url::Url::parse(url) { Ok(mut parsed) => { parsed.set_query(None); + // Both fail only for URLs that cannot carry userinfo at all. + let _ = parsed.set_username(""); + let _ = parsed.set_password(None); parsed.to_string() } - Err(_) => url.split('?').next().unwrap_or(url).to_owned(), + Err(_) => { + let without_query = url.split('?').next().unwrap_or(url); + // Unparseable: drop anything up to the last '@' of the authority. + if let Some((scheme, rest)) = without_query.split_once("://") { + let authority = &rest[..rest.find('/').unwrap_or(rest.len())]; + if let Some(at) = authority.rfind('@') { + return format!("{scheme}://{}", &rest[at + 1..]); + } + } + without_query.to_owned() + } + } +} + +#[cfg(test)] +mod tests { + use super::sanitized_request_url; + + #[test] + fn sanitized_request_url_drops_userinfo_and_query() { + assert_eq!( + sanitized_request_url("https://user:secret@api.example.com/v1/chat?key=sk-1"), + "https://api.example.com/v1/chat" + ); + assert_eq!( + sanitized_request_url("https://token@api.example.com/v1"), + "https://api.example.com/v1" + ); + // Unparseable (space in the host): userinfo is still dropped. + assert_eq!( + sanitized_request_url("https://user:secret@bad host/v1?key=sk-1"), + "https://bad host/v1" + ); + assert_eq!( + sanitized_request_url("https://api.example.com/v1/chat?key=sk-1"), + "https://api.example.com/v1/chat" + ); } } diff --git a/aimux-provider-utils/tests/http_recording_test.rs b/aimux-provider-utils/tests/http_recording_test.rs index 58dac8160..ba11416d4 100644 --- a/aimux-provider-utils/tests/http_recording_test.rs +++ b/aimux-provider-utils/tests/http_recording_test.rs @@ -169,6 +169,59 @@ async fn one_helper_call_records_one_finalized_exchange() { let _ = std::fs::remove_dir_all(directory); } +/// Provider auth headers whose names carry no `token` / `api-key` (Prodia, +/// Azure Speech) are stored as `[REDACTED]`, like `authorization`. +#[tokio::test] +#[serial] +async fn provider_auth_headers_are_recorded_redacted() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/call")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"ok": true}))) + .expect(1) + .mount(&server) + .await; + let (recorder, directory) = recorder("provider-auth-headers"); + recording::init_recording(Some(recorder.clone())); + + let mut request = request( + format!("{}/call", server.uri()), + "call-auth", + recorder.clone(), + ); + request + .headers + .push(("X-Prodia-Key".into(), "prodia-secret".into())); + request + .headers + .push(("Ocp-Apim-Subscription-Key".into(), "azure-secret".into())); + post_json_to_api( + request, + json!({"prompt": "hi"}), + create_json_response_handler::(), + failed(), + ) + .await + .unwrap(); + finish_call(recorder.as_ref(), "call-auth"); + + let recording = read_recording(&directory); + let headers = &recording.exchanges[0].request.headers; + for name in ["X-Prodia-Key", "Ocp-Apim-Subscription-Key"] { + let (_, value) = headers + .iter() + .find(|(header, _)| header.eq_ignore_ascii_case(name)) + .unwrap_or_else(|| panic!("{name} missing from the recording")); + assert_eq!(value, "[REDACTED]", "{name}"); + } + let stored = std::fs::read_to_string(directory.join("recordings.jsonl")).unwrap(); + assert!(!stored.contains("prodia-secret"), "{stored}"); + assert!(!stored.contains("azure-secret"), "{stored}"); + + recording::init_recording(None); + let _ = std::fs::remove_dir_all(directory); +} + #[tokio::test] #[serial] async fn cancelled_exchange_is_recorded_as_finalized_failure() { diff --git a/aimux-providers/src/bedrock/sigv4.rs b/aimux-providers/src/bedrock/sigv4.rs index e915d95ca..681f91f3b 100644 --- a/aimux-providers/src/bedrock/sigv4.rs +++ b/aimux-providers/src/bedrock/sigv4.rs @@ -18,7 +18,11 @@ use hmac::{Hmac, Mac}; use sha2::{Digest, Sha256}; /// A set of AWS credentials for SigV4 signing. -#[derive(Debug, Clone)] +/// +/// `Debug` is hand-written so the secret key and session token never reach +/// logs or error output; the access key id and region stay visible because +/// they identify the credentials without granting access. +#[derive(Clone)] pub struct AwsCredentials { pub access_key_id: String, pub secret_access_key: String, @@ -27,6 +31,20 @@ pub struct AwsCredentials { pub region: String, } +impl std::fmt::Debug for AwsCredentials { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("AwsCredentials") + .field("access_key_id", &self.access_key_id) + .field("secret_access_key", &"[REDACTED]") + .field( + "session_token", + &self.session_token.as_ref().map(|_| "[REDACTED]"), + ) + .field("region", &self.region) + .finish() + } +} + /// The result of signing a request. #[derive(Debug, Clone)] pub struct SignedRequest { @@ -55,6 +73,27 @@ pub fn sign_request( url: &str, body: &str, extra_headers: &[(String, String)], +) -> SignedRequest { + sign_request_at( + credentials, + service, + method, + url, + body, + extra_headers, + chrono::Utc::now(), + ) +} + +/// [`sign_request`] with an explicit signing time. +fn sign_request_at( + credentials: &AwsCredentials, + service: &str, + method: &str, + url: &str, + body: &str, + extra_headers: &[(String, String)], + now: chrono::DateTime, ) -> SignedRequest { let parsed = url::Url::parse(url).unwrap_or_else(|_| url::Url::parse("http://localhost").unwrap()); @@ -68,11 +107,10 @@ pub fn sign_request( Some(port) => format!("{}:{}", parsed.host_str().unwrap_or(""), port), None => parsed.host_str().unwrap_or("").to_string(), }; - let path = parsed.path(); - let query = parsed.query().unwrap_or(""); + let path = canonical_uri(parsed.path()); + let query = canonical_query(parsed.query().unwrap_or("")); // Timestamp:yyyyMMddTHHmmssZ and date:yyyyMMdd - let now = chrono::Utc::now(); let amz_date = now.format("%Y%m%dT%H%M%SZ").to_string(); let date_stamp = now.format("%Y%m%d").to_string(); @@ -99,7 +137,7 @@ pub fn sign_request( && lower != "x-amz-content-sha256" && lower != "x-amz-security-token" { - canonical_headers.push((lower, v.trim().to_string())); + canonical_headers.push((lower, canonical_header_value(v))); } } @@ -138,16 +176,17 @@ pub fn sign_request( let string_to_sign = format!("AWS4-HMAC-SHA256\n{amz_date}\n{credential_scope}\n{canonical_request_hash}"); - // Signing key: derived through chained HMAC + // Signing key: derived through chained HMAC, each step keyed by the + // previous one and starting from "AWS4" + secret. let k_date = hmac_sha256( + format!("AWS4{}", credentials.secret_access_key).as_bytes(), date_stamp.as_bytes(), - credentials.secret_access_key.as_bytes(), ); - let k_region = hmac_sha256(credentials.region.as_bytes(), &k_date); - let k_service = hmac_sha256(service.as_bytes(), &k_region); - let k_signing = hmac_sha256(b"aws4_request", &k_service); + let k_region = hmac_sha256(&k_date, credentials.region.as_bytes()); + let k_service = hmac_sha256(&k_region, service.as_bytes()); + let k_signing = hmac_sha256(&k_service, b"aws4_request"); - let signature = hex::encode(hmac_sha256(string_to_sign.as_bytes(), &k_signing)); + let signature = hex::encode(hmac_sha256(&k_signing, string_to_sign.as_bytes())); let authorization = format!( "AWS4-HMAC-SHA256 Credential={}/{}, SignedHeaders={}, Signature={}", @@ -164,6 +203,84 @@ pub fn sign_request( SignedRequest { headers } } +/// SigV4 `UriEncode`: every byte except the unreserved set +/// (`A-Z a-z 0-9 - _ . ~`) becomes `%XX` with uppercase hex. +fn uri_encode(input: &[u8]) -> String { + let mut out = String::with_capacity(input.len()); + for &byte in input { + if byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_' | b'.' | b'~') { + out.push(byte as char); + } else { + out.push_str(&format!("%{byte:02X}")); + } + } + out +} + +/// Canonical URI for every service but S3: each segment of the path as sent +/// is URI-encoded again, so `:` signs as `%3A` and an already-escaped `%3A` +/// signs as `%253A`. +fn canonical_uri(path: &str) -> String { + if path.is_empty() { + return "/".to_string(); + } + path.split('/') + .map(|segment| uri_encode(segment.as_bytes())) + .collect::>() + .join("/") +} + +/// Canonical query string: each name and value percent-decoded, then +/// URI-encoded, and the pairs sorted by name and then value. +fn canonical_query(query: &str) -> String { + let mut pairs: Vec<(String, String)> = query + .split('&') + .filter(|pair| !pair.is_empty()) + .map(|pair| { + let (name, value) = pair.split_once('=').unwrap_or((pair, "")); + ( + uri_encode(&percent_decode(name)), + uri_encode(&percent_decode(value)), + ) + }) + .collect(); + pairs.sort(); + pairs + .iter() + .map(|(name, value)| format!("{name}={value}")) + .collect::>() + .join("&") +} + +fn percent_decode(input: &str) -> Vec { + let bytes = input.as_bytes(); + let mut out = Vec::with_capacity(bytes.len()); + let mut i = 0; + while i < bytes.len() { + let escaped = (bytes[i] == b'%') + .then(|| bytes.get(i + 1..i + 3)) + .flatten() + .and_then(|hex| std::str::from_utf8(hex).ok()) + .and_then(|hex| u8::from_str_radix(hex, 16).ok()); + match escaped { + Some(byte) => { + out.push(byte); + i += 3; + } + None => { + out.push(bytes[i]); + i += 1; + } + } + } + out +} + +/// Canonical header value: trimmed, with runs of spaces collapsed to one. +fn canonical_header_value(value: &str) -> String { + value.split_whitespace().collect::>().join(" ") +} + fn hmac_sha256(key: &[u8], data: &[u8]) -> Vec { let mut mac = HmacSha256::new_from_slice(key).expect("HMAC accepts any key length"); mac.update(data); @@ -183,6 +300,17 @@ mod tests { } } + #[test] + fn debug_redacts_secret_key_and_session_token() { + let mut credentials = creds(); + credentials.session_token = Some("FwoGZXIvYXdzEXAMPLETOKEN".to_string()); + let printed = format!("{credentials:?}"); + assert!(!printed.contains("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY")); + assert!(!printed.contains("FwoGZXIvYXdzEXAMPLETOKEN")); + assert!(printed.contains("[REDACTED]")); + assert!(printed.contains("us-east-1")); + } + fn host_header(signed: &SignedRequest) -> String { signed .headers @@ -220,4 +348,103 @@ mod tests { "bedrock-runtime.us-east-1.amazonaws.com" ); } + + fn fixed_time() -> chrono::DateTime { + use chrono::TimeZone; + chrono::Utc.with_ymd_and_hms(2024, 1, 2, 3, 4, 5).unwrap() + } + + fn signature(signed: &SignedRequest) -> String { + signed + .headers + .iter() + .find(|(k, _)| k == "Authorization") + .and_then(|(_, v)| v.rsplit_once("Signature=")) + .expect("signed request carries a signature") + .1 + .to_string() + } + + /// Expected values come from botocore 1.43's `SigV4Auth` (AWS's own + /// Python signer) for the same credentials, headers, body and fixed + /// time, not from this implementation. + #[test] + fn signatures_match_botocore() { + let signed = sign_request_at( + &creds(), + "bedrock", + "POST", + "https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-3-5-sonnet-20240620-v1:0/converse", + r#"{"messages":[]}"#, + &[("X-Custom".to_string(), " value ".to_string())], + fixed_time(), + ); + assert_eq!( + signature(&signed), + "538444fb4b17f3cdd777bbd5f1cd4e3e1f53d6ac97167da96c81ec8d14d6468a" + ); + + let session = AwsCredentials { + session_token: Some("SESSION".to_string()), + ..creds() + }; + let signed = sign_request_at( + &session, + "bedrock", + "POST", + "http://127.0.0.1:54321/model/x/invoke?a=1", + "{}", + &[], + fixed_time(), + ); + assert_eq!( + signature(&signed), + "dd6a2384cb9c8f0db8ea8ee383a61bf06dd85da679cec54d5c944b1b0cb8d9f6" + ); + + let signed = sign_request_at( + &creds(), + "bedrock", + "POST", + "https://bedrock-runtime.us-east-1.amazonaws.com/model/us.anthropic.claude-3-5-sonnet-20240620-v1%3A0/converse-stream", + "{}", + &[ + ("content-type".to_string(), "application/json".to_string()), + ("X-Multi".to_string(), "a b c".to_string()), + ], + fixed_time(), + ); + assert_eq!( + signature(&signed), + "a6abc7f26ad04c0922b03e05c9184a20a8fcf4b68b7dfce358e9a6436c8a9489" + ); + + let signed = sign_request_at( + &creds(), + "bedrock", + "POST", + "https://polly.us-east-1.amazonaws.com/v1/speech?b=2&a=x%20y&a=1", + "{}", + &[], + fixed_time(), + ); + assert_eq!( + signature(&signed), + "70c6e5794c4f1456f33f2b41cebc483ee0d19c98b95660411f4b51ccac464006" + ); + } + + #[test] + fn canonical_uri_and_query_follow_sigv4() { + assert_eq!(canonical_uri(""), "/"); + assert_eq!( + canonical_uri("/model/a.b-v1:0/converse"), + "/model/a.b-v1%3A0/converse" + ); + assert_eq!(canonical_uri("/a+b/c%2Fd"), "/a%2Bb/c%252Fd"); + assert_eq!( + canonical_query("b=2&a=x%20y&a=1&e=&z"), + "a=1&a=x%20y&b=2&e=&z=" + ); + } } diff --git a/aimux-providers/src/codex.rs b/aimux-providers/src/codex.rs index ccf99d3e1..7a46ac47e 100644 --- a/aimux-providers/src/codex.rs +++ b/aimux-providers/src/codex.rs @@ -478,7 +478,10 @@ impl LanguageModel for CodexModel { } /// Tokens returned by the Codex OAuth token endpoint (RFC-0018 §2 V3). -#[derive(Debug, Clone, Serialize, Deserialize)] +/// +/// `Debug` is hand-written so both tokens are redacted from logs and error +/// output. +#[derive(Clone, Serialize, Deserialize)] pub struct CodexTokens { /// Access token to pass as `CodexConfig::subscription(..)`. pub access_token: String, @@ -489,6 +492,19 @@ pub struct CodexTokens { pub expires_in_secs: Option, } +impl std::fmt::Debug for CodexTokens { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CodexTokens") + .field("access_token", &"[REDACTED]") + .field( + "refresh_token", + &self.refresh_token.as_ref().map(|_| "[REDACTED]"), + ) + .field("expires_in_secs", &self.expires_in_secs) + .finish() + } +} + /// Stateless OAuth token-refresh helper for the Codex subscription channel /// (RFC-0018 §3.2). Performs one `POST auth.openai.com/oauth/token` call and /// returns the refreshed tokens — **no persistence, no orchestration**; the @@ -579,4 +595,17 @@ mod tests { assert!(snap.base_url.is_some()); assert_eq!(snap.api_key_source, "explicit"); } + + #[test] + fn tokens_debug_redacts_secrets() { + let tokens = CodexTokens { + access_token: "at-secret-value".to_string(), + refresh_token: Some("rt-secret-value".to_string()), + expires_in_secs: Some(3600), + }; + let printed = format!("{tokens:?}"); + assert!(!printed.contains("at-secret-value")); + assert!(!printed.contains("rt-secret-value")); + assert!(printed.contains("3600")); + } } diff --git a/aimux-providers/src/lib.rs b/aimux-providers/src/lib.rs index 90e926f31..61f8bb1f4 100644 --- a/aimux-providers/src/lib.rs +++ b/aimux-providers/src/lib.rs @@ -39,7 +39,7 @@ pub mod provider_name; pub mod replay; pub use provider::{ ExternalProviderEntry, ProviderOptions, ProviderProfile, is_external_provider, - load_providers_from_json, provider, provider_from_env, provider_handle, + load_providers_from_json, provider, provider_api_key_env, provider_from_env, provider_handle, provider_registry_entry, register_provider, }; pub use replay::{rebuild_decision_provider, rebuild_provider}; diff --git a/aimux-providers/src/provider.rs b/aimux-providers/src/provider.rs index ea74bc148..0aa3ce26f 100644 --- a/aimux-providers/src/provider.rs +++ b/aimux-providers/src/provider.rs @@ -594,6 +594,37 @@ pub fn provider_from_env( provider(name, None, model_id, options) } +/// The environment variable [`provider`] / [`provider_handle`] read `name`'s +/// API key from when no key is passed: the runtime overlay's `api_key` +/// `env:` reference or its `env_var`, `{NAME}_API_KEY` for the local +/// OpenAI-compatible servers (ollama, llamacpp, …), or the registry entry's +/// `env_var`. `None` for an unknown name or an entry without one. +/// +/// Hosts that let a caller pick a key by variable name (the aimux-web +/// console's `env:VAR`) use it to allow only the provider's own variable. +#[must_use] +pub fn provider_api_key_env(name: &str) -> Option { + if let Some(ext) = overlays().read().unwrap().get(name) { + return ext + .api_key + .as_deref() + .and_then(|key| key.strip_prefix("env:")) + .or(ext.env_var.as_deref()) + .filter(|var| !var.is_empty()) + .map(str::to_string); + } + match name { + "ollama" | "llamacpp" | "localai" | "vllm" | "sglang" | "laya" => { + Some(format!("{}_API_KEY", name.to_uppercase())) + } + "openai" => Some("OPENAI_API_KEY".to_string()), + _ => registry() + .iter() + .find(|entry| entry.name == name) + .map(|entry| entry.env_var.clone()), + } +} + /// Public lookup of a registered provider's runtime profile — used by tests /// that assert registry wiring (e.g. `max_tokens_key`) without constructing a /// model. Returns `None` for unknown provider names. @@ -610,6 +641,46 @@ mod tests { use super::*; use crate::provider_name::ProviderName; + #[test] + fn provider_api_key_env_names_the_variable_the_provider_reads() { + assert_eq!( + provider_api_key_env("groq").as_deref(), + Some("GROQ_API_KEY") + ); + assert_eq!( + provider_api_key_env("openai").as_deref(), + Some("OPENAI_API_KEY") + ); + assert_eq!( + provider_api_key_env("ollama").as_deref(), + Some("OLLAMA_API_KEY") + ); + assert_eq!(provider_api_key_env("no-such-provider"), None); + + let name = "key-env-overlay-test"; + let entry = |api_key: Option<&str>, env_var: Option<&str>| ExternalProviderEntry { + name: name.into(), + display: None, + base_url: "https://overlay.example/v1".into(), + env_var: env_var.map(str::to_string), + api_key: api_key.map(str::to_string), + protocol: "openai_compat".into(), + profile: ProviderProfile::default(), + headers: None, + organization: None, + project: None, + max_retries: None, + body_overrides: None, + comment: None, + }; + register_provider(entry(None, Some("OVERLAY_KEY"))).unwrap(); + assert_eq!(provider_api_key_env(name).as_deref(), Some("OVERLAY_KEY")); + // An `env:` api_key reference wins, as in key resolution. + register_provider(entry(Some("env:OVERLAY_REF"), Some("OVERLAY_KEY"))).unwrap(); + assert_eq!(provider_api_key_env(name).as_deref(), Some("OVERLAY_REF")); + clear_overlay(name); + } + #[test] fn provider_builds_groq_model() { let model = match provider("groq", Some("sk-test".into()), "llama-3.3-70b", None) { diff --git a/bindings/flutter/lib/aimux.dart b/bindings/flutter/lib/aimux.dart index cc0ee5c5a..f2e728517 100644 --- a/bindings/flutter/lib/aimux.dart +++ b/bindings/flutter/lib/aimux.dart @@ -265,14 +265,18 @@ final class _AimuxFFI { 'aimux_stream_text_as_openai'); late final dropHandle = _lib.lookupFunction<_DropHandleC, _DropHandleDart>('aimux_drop_handle'); - // NativeFinalizer callback pointer: aimux_drop_handle reinterpreted as - // Void Function(Pointer). The handle value is attached as the token - // (its address == the uint64 handle); on 64-bit the ABI is identical, so - // the finalizer releases the handle when a Model/ProviderHandle is GC'd - // without close(). NativeFinalizerFunction is already a NativeFunction<…>, - // so it is looked up directly (no extra NativeFunction wrapper). + // NativeFinalizer callback pointer: aimux_drop_handle_ptr, declared + // `void aimux_drop_handle_ptr(void *)` — exactly NativeFinalizerFunction + // (Void Function(Pointer)). The handle value is attached as the token + // (its address == the handle), and the finalizer releases the handle when + // a Model/ProviderHandle is GC'd without close(). aimux_drop_handle takes a + // uint64_t and must not be called through this signature: on 32-bit ARM + // (armeabi-v7a) a uint64_t argument does not travel like a pointer, so it + // would receive a garbage high word and leak the handle. + // NativeFinalizerFunction is already a NativeFunction<…>, so it is looked + // up directly (no extra NativeFunction wrapper). late final dropHandlePtr = - _lib.lookup('aimux_drop_handle'); + _lib.lookup('aimux_drop_handle_ptr'); late final initLogging = _lib.lookupFunction<_StrC, _StrDart>('aimux_init_logging'); late final initRecording = @@ -309,8 +313,8 @@ final _AimuxFFI _ffi = _AimuxFFI(); /// Process-wide native finalizer that releases a handle when a [Model] or /// [ProviderHandle] is garbage-collected without [Model.close] / /// [ProviderHandle.close] (T9). The token attached is the handle value -/// encoded as a `Pointer` (its address == the uint64 handle), so the -/// `aimux_drop_handle` native callback receives it as a uint64. +/// encoded as a `Pointer` (its address == the handle), which the +/// `aimux_drop_handle_ptr` native callback takes as a pointer-sized value. final NativeFinalizer _handleFinalizer = NativeFinalizer(_ffi.dropHandlePtr); // ───────────────────────────────────────────────────────────────────────────── diff --git a/bindings/flutter/test/aimux_test.dart b/bindings/flutter/test/aimux_test.dart index aaa09384d..f647b385d 100644 --- a/bindings/flutter/test/aimux_test.dart +++ b/bindings/flutter/test/aimux_test.dart @@ -97,6 +97,41 @@ void main() { } }); + // The NativeFinalizer callback: aimux_drop_handle_ptr takes the handle + // as a pointer-sized value (NativeFinalizerFunction), so it is safe on + // 32-bit targets where aimux_drop_handle(uint64_t) is not. + test('aimux_drop_handle_ptr releases the handle its address names', () { + final dropHandlePtr = lib.lookupFunction), + void Function(Pointer)>('aimux_drop_handle_ptr'); + expect(() => lib.lookup('aimux_drop_handle_ptr'), + returnsNormally); + final outHandle = calloc(); + final outJson = calloc>(); + try { + final e = withUtf8('sk-test-fake-key', + (k) => withUtf8('gpt-4o-mini', (m) => openaiNew(k, m, outHandle))); + expect(e, nullptr); + final handle = outHandle.value; + expect(handle, isNot(0)); + dropHandlePtr(Pointer.fromAddress(handle)); + // The handle is gone: a call on it is the dead-handle failure (203). + expect( + () => withUtf8('"hi"', (p) { + final e = generateText(handle, p, nullptr, outJson); + expect(errorCode(e), 203); + expectAimuxError(e, 'generate_text'); + }), + throwsA(isA()), + ); + expect(outJson.value, nullptr); + // NULL is handle 0: a no-op. + dropHandlePtr(nullptr); + } finally { + calloc.free(outHandle); + calloc.free(outJson); + } + }); + // Bad wire JSON is only reachable when checkJson pre-validation was // bypassed (raw C call here); it is a C ABI failure → StateError. test('bad wire JSON at the raw C ABI → StateError', () { diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs index 527e468a4..d716da7fb 100644 --- a/bindings/python/src/lib.rs +++ b/bindings/python/src/lib.rs @@ -141,12 +141,19 @@ impl Model { /// opts_json: optional JSON-serialized GenerateTextOptions /// Returns JSON-serialized GenerateTextResult. #[pyo3(signature = (prompt_json, opts_json=None))] - fn generate_text(&self, prompt_json: &str, opts_json: Option<&str>) -> PyResult { + fn generate_text( + &self, + py: Python<'_>, + prompt_json: &str, + opts_json: Option<&str>, + ) -> PyResult { let prompt = parse_prompt(prompt_json)?; let opts = parse_opts(opts_json)?; let rt = runtime(); - let result = rt.block_on(async move { generate_text(&*self.inner, prompt, opts).await }); + let result = py.allow_threads(|| { + rt.block_on(async move { generate_text(&*self.inner, prompt, opts).await }) + }); match result { Ok(r) => serialize_result(&r), @@ -162,12 +169,19 @@ impl Model { /// `response_format: { "Json": { ... } }` via opts_json for schema /// control; the function applies JSON repair before parsing. #[pyo3(signature = (prompt_json, opts_json=None))] - fn generate_object(&self, prompt_json: &str, opts_json: Option<&str>) -> PyResult { + fn generate_object( + &self, + py: Python<'_>, + prompt_json: &str, + opts_json: Option<&str>, + ) -> PyResult { let prompt = parse_prompt(prompt_json)?; let opts = parse_opts(opts_json)?; let rt = runtime(); - let result = rt.block_on(async move { generate_object(&*self.inner, prompt, opts).await }); + let result = py.allow_threads(|| { + rt.block_on(async move { generate_object(&*self.inner, prompt, opts).await }) + }); match result { Ok(r) => serialize_result(&r), @@ -182,14 +196,21 @@ impl Model { /// opts_json: optional JSON-serialized GenerateTextOptions /// Returns JSON-serialized StreamTextResultAggregated. #[pyo3(signature = (prompt_json, opts_json=None))] - fn consume_stream_text(&self, prompt_json: &str, opts_json: Option<&str>) -> PyResult { + fn consume_stream_text( + &self, + py: Python<'_>, + prompt_json: &str, + opts_json: Option<&str>, + ) -> PyResult { let prompt = parse_prompt(prompt_json)?; let opts = parse_opts(opts_json)?; let rt = runtime(); - let result = rt.block_on(async move { - let stream_result = stream_text(&*self.inner, prompt, opts).await?; - stream_result.consume().await + let result = py.allow_threads(|| { + rt.block_on(async move { + let stream_result = stream_text(&*self.inner, prompt, opts).await?; + stream_result.consume().await + }) }); match result { @@ -277,6 +298,7 @@ impl Model { #[pyo3(signature = (prompt_json, opts_json=None))] fn generate_text_as_openai( &self, + py: Python<'_>, prompt_json: &str, opts_json: Option<&str>, ) -> PyResult { @@ -284,8 +306,9 @@ impl Model { let opts = parse_opts(opts_json)?; let rt = runtime(); - let result = - rt.block_on(async move { generate_text_as_openai(&*self.inner, prompt, opts).await }); + let result = py.allow_threads(|| { + rt.block_on(async move { generate_text_as_openai(&*self.inner, prompt, opts).await }) + }); match result { Ok(r) => serialize_result(&r), @@ -736,10 +759,10 @@ impl ProviderHandle { /// List models available on this provider (runtime discovery + anya2a spec). /// Returns a JSON array of RuntimeModel. - fn list_models(&self) -> PyResult { + fn list_models(&self, py: Python<'_>) -> PyResult { let rt = runtime(); - let models = rt - .block_on(async { self.inner.list_models().await }) + let models = py + .allow_threads(|| rt.block_on(async { self.inner.list_models().await })) .map_err(|e| to_py_err(&e))?; serialize_result(&models) } @@ -786,9 +809,11 @@ fn create_provider( /// string (serialized `Catalogue`). `source_url` defaults to the anya2a /// `dist/all.json`. #[pyfunction] -fn get_model_specs(source_url: Option<&str>) -> PyResult { - let catalogue = runtime() - .block_on(async { aimux_providers::get_model_specs(source_url).await }) +fn get_model_specs(py: Python<'_>, source_url: Option<&str>) -> PyResult { + let catalogue = py + .allow_threads(|| { + runtime().block_on(async { aimux_providers::get_model_specs(source_url).await }) + }) .map_err(|e| to_py_err(&e))?; serialize_result(&catalogue) } @@ -926,9 +951,9 @@ fn recording_stop() { /// 刷盘全局 recorder(阻塞至 JSONL 落盘;ring 模式 no-op)。 #[pyfunction] -fn recording_flush() { +fn recording_flush(py: Python<'_>) { if let Some(rec) = aimux_core::recording::recorder() { - rec.flush(); + py.allow_threads(|| rec.flush()); } } @@ -938,11 +963,11 @@ fn recording_flush() { /// exception type, not an `AimuxError`. Returns normally when nothing is /// recording. The legacy `recording_flush` stays and never reports. #[pyfunction] -fn recording_try_flush() -> PyResult<()> { +fn recording_try_flush(py: Python<'_>) -> PyResult<()> { let Some(rec) = aimux_core::recording::recorder() else { return Ok(()); }; - rec.try_flush() + py.allow_threads(|| rec.try_flush()) .map_err(|e| AiMuxBindingError::from(e).to_py_err()) } diff --git a/bindings/python/src/multimodal.rs b/bindings/python/src/multimodal.rs index 3c7289617..8f838f0aa 100644 --- a/bindings/python/src/multimodal.rs +++ b/bindings/python/src/multimodal.rs @@ -93,7 +93,12 @@ impl EmbeddingModel { /// `opts_json` is optional JSON-serialized EmbeddingCallOptions. /// Returns JSON-serialized EmbeddingResult. #[pyo3(signature = (values_json, opts_json=None))] - pub fn embed(&self, values_json: &str, opts_json: Option<&str>) -> PyResult { + pub fn embed( + &self, + py: Python<'_>, + values_json: &str, + opts_json: Option<&str>, + ) -> PyResult { let mut opts: EmbeddingCallOptions = match opts_json { Some(s) if !s.trim().is_empty() && s.trim() != "null" => wire_json("opts_json", s)?, _ => EmbeddingCallOptions::new(""), @@ -102,8 +107,10 @@ impl EmbeddingModel { let values: Vec = wire_json("values_json", values_json)?; opts.values = values; - let result = crate::runtime().block_on(async move { - aimux_core::embedding_model::embed(self.inner.as_ref(), opts).await + let result = py.allow_threads(|| { + crate::runtime().block_on(async move { + aimux_core::embedding_model::embed(self.inner.as_ref(), opts).await + }) }); match result { @@ -126,11 +133,13 @@ pub struct SpeechModel { impl SpeechModel { /// Generate speech audio. `opts_json` is JSON-serialized SpeechCallOptions. /// Returns JSON-serialized SpeechResult (audio as base64 in JSON). - pub fn generate(&self, opts_json: &str) -> PyResult { + pub fn generate(&self, py: Python<'_>, opts_json: &str) -> PyResult { let opts: SpeechCallOptions = wire_json("opts_json", opts_json)?; - let result = crate::runtime().block_on(async move { - aimux_core::speech_model::generate_speech(self.inner.as_ref(), opts).await + let result = py.allow_threads(|| { + crate::runtime().block_on(async move { + aimux_core::speech_model::generate_speech(self.inner.as_ref(), opts).await + }) }); match result { @@ -153,11 +162,13 @@ pub struct ImageModel { impl ImageModel { /// Generate images. `opts_json` is JSON-serialized ImageCallOptions. /// Returns JSON-serialized ImageResult (images as base64 in JSON). - pub fn generate(&self, opts_json: &str) -> PyResult { + pub fn generate(&self, py: Python<'_>, opts_json: &str) -> PyResult { let opts: ImageCallOptions = wire_json("opts_json", opts_json)?; - let result = crate::runtime().block_on(async move { - aimux_core::image_model::generate_image(self.inner.as_ref(), opts).await + let result = py.allow_threads(|| { + crate::runtime().block_on(async move { + aimux_core::image_model::generate_image(self.inner.as_ref(), opts).await + }) }); match result { @@ -184,6 +195,7 @@ impl TranscriptionModel { #[pyo3(signature = (audio_base64, media_type, opts_json=None))] pub fn generate( &self, + py: Python<'_>, audio_base64: &str, media_type: &str, opts_json: Option<&str>, @@ -210,8 +222,10 @@ impl TranscriptionModel { } } - let result = crate::runtime().block_on(async move { - aimux_core::transcription_model::transcribe(self.inner.as_ref(), opts).await + let result = py.allow_threads(|| { + crate::runtime().block_on(async move { + aimux_core::transcription_model::transcribe(self.inner.as_ref(), opts).await + }) }); match result { @@ -238,6 +252,7 @@ impl RerankingModel { #[pyo3(signature = (query, docs_json, opts_json=None))] pub fn rerank( &self, + py: Python<'_>, query: &str, docs_json: &str, opts_json: Option<&str>, @@ -264,8 +279,10 @@ impl RerankingModel { } } - let result = crate::runtime().block_on(async move { - aimux_core::reranking_model::rerank(self.inner.as_ref(), opts).await + let result = py.allow_threads(|| { + crate::runtime().block_on(async move { + aimux_core::reranking_model::rerank(self.inner.as_ref(), opts).await + }) }); match result { @@ -288,11 +305,13 @@ pub struct VideoModel { impl VideoModel { /// Generate video. `opts_json` is JSON-serialized VideoCallOptions. /// Returns JSON-serialized VideoResult (typically contains a URL). - pub fn generate(&self, opts_json: &str) -> PyResult { + pub fn generate(&self, py: Python<'_>, opts_json: &str) -> PyResult { let opts: VideoCallOptions = wire_json("opts_json", opts_json)?; - let result = crate::runtime().block_on(async move { - aimux_core::video_model::generate_video(self.inner.as_ref(), opts).await + let result = py.allow_threads(|| { + crate::runtime().block_on(async move { + aimux_core::video_model::generate_video(self.inner.as_ref(), opts).await + }) }); match result { @@ -316,7 +335,7 @@ impl SearchModel { /// Search. `query` is the search query, `opts_json` is optional JSON options. /// Returns JSON-serialized SearchResult. #[pyo3(signature = (query, opts_json=None))] - pub fn search(&self, query: &str, opts_json: Option<&str>) -> PyResult { + pub fn search(&self, py: Python<'_>, query: &str, opts_json: Option<&str>) -> PyResult { let mut opts = SearchCallOptions::new(query.to_string()); if let Some(s) = opts_json { if !s.trim().is_empty() && s.trim() != "null" { @@ -332,8 +351,10 @@ impl SearchModel { } } - let result = crate::runtime().block_on(async move { - aimux_core::search_model::search(self.inner.as_ref(), opts).await + let result = py.allow_threads(|| { + crate::runtime().block_on(async move { + aimux_core::search_model::search(self.inner.as_ref(), opts).await + }) }); match result { @@ -361,6 +382,7 @@ impl Files { #[pyo3(signature = (data_base64, media_type, opts_json=None))] pub fn upload_file( &self, + py: Python<'_>, data_base64: &str, media_type: &str, opts_json: Option<&str>, @@ -380,7 +402,9 @@ impl Files { } } - let result = crate::runtime().block_on(async move { self.inner.upload_file(&opts).await }); + let result = py.allow_threads(|| { + crate::runtime().block_on(async move { self.inner.upload_file(&opts).await }) + }); match result { Ok(r) => serialize_result(&r), diff --git a/bindings/python/tests/test_e2e.py b/bindings/python/tests/test_e2e.py index 6a6157868..a55779448 100644 --- a/bindings/python/tests/test_e2e.py +++ b/bindings/python/tests/test_e2e.py @@ -3,9 +3,9 @@ Verifies the FULL chain: Python → PyO3 → aimux-core → HTTP mock → response parsing → typed result. -The mock server runs in a separate process (multiprocessing) because -PyO3's block_on blocks the calling thread — a threading-based server -would deadlock. +The mock server runs in a separate process (multiprocessing). Native calls +release the GIL, so a server thread in this process works too +(test_gil.py relies on it). """ import json diff --git a/bindings/python/tests/test_gil.py b/bindings/python/tests/test_gil.py new file mode 100644 index 000000000..acb53e1e6 --- /dev/null +++ b/bindings/python/tests/test_gil.py @@ -0,0 +1,92 @@ +"""Blocking native calls release the GIL. + +The mock server here runs in a thread of the test process, so it can only +answer while the calling thread waits inside a native call if that call +released the GIL. Each call carries a total timeout: if the GIL were held, +the server could not answer and the call would fail with a timeout instead +of hanging the suite. +""" + +import json +import threading +import time +from http.server import BaseHTTPRequestHandler, HTTPServer + +import pytest + +from aimux import generate_text, openai, openai_embedding + +CHAT = json.dumps({ + "id": "chatcmpl-gil", + "model": "gpt-4o", + "choices": [{ + "message": {"role": "assistant", "content": "served by a Python thread"}, + "finish_reason": "stop", + }], + "usage": {"prompt_tokens": 1, "completion_tokens": 5, "total_tokens": 6}, +}) + +EMBEDDING = json.dumps({ + "data": [{"embedding": [0.5, 0.25], "index": 0}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 1, "total_tokens": 1}, +}) + +LIMITS = {"max_retries": 0, "timeout": {"total_ms": 10000}} + + +class _ThreadServer: + """HTTP server on a daemon thread of this process, answering slowly.""" + + def __init__(self, body): + progress = self.progress = [] + + class Handler(BaseHTTPRequestHandler): + def do_POST(self): + self.rfile.read(int(self.headers.get("Content-Length", 0))) + # Keep the native call waiting while this thread keeps running. + for _ in range(10): + progress.append(time.monotonic()) + time.sleep(0.02) + payload = body.encode() + self.send_response(200) + self.send_header("content-type", "application/json") + self.send_header("content-length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + def log_message(self, *args): + pass + + self.server = HTTPServer(("127.0.0.1", 0), Handler) + self.url = "http://127.0.0.1:%d" % self.server.server_address[1] + self.thread = threading.Thread(target=self.server.serve_forever, daemon=True) + + def __enter__(self): + self.thread.start() + return self + + def __exit__(self, *args): + self.server.shutdown() + self.server.server_close() + self.thread.join(timeout=5) + + +def test_generate_text_lets_other_python_threads_run(): + with _ThreadServer(CHAT) as server: + model = openai("test-key", "gpt-4o", server.url) + started = time.monotonic() + result = generate_text(model, "hi", LIMITS) + assert result["text"] == "served by a Python thread" + # The server thread ran while generate_text was blocked. + assert len(server.progress) == 10 + assert all(t >= started for t in server.progress) + + +def test_embed_lets_other_python_threads_run(): + with _ThreadServer(EMBEDDING) as server: + embedder = openai_embedding("test-key", "text-embedding-3-small", base_url=server.url) + opts = dict(LIMITS, values=[]) + result = json.loads(embedder.embed(json.dumps(["hello"]), json.dumps(opts))) + assert result["embeddings"][0] == pytest.approx([0.5, 0.25]) + assert len(server.progress) == 10 diff --git a/docs/aisdk-architecture-alignment.md b/docs/aisdk-architecture-alignment.md index 0d0231ac6..8ce469f47 100644 --- a/docs/aisdk-architecture-alignment.md +++ b/docs/aisdk-architecture-alignment.md @@ -94,7 +94,7 @@ FFI / Node / Python / 5 个 C-ABI 绑定 / CLI / Web 由 descriptor + manifest ### 0.6 独立于本文的事项 -- **安全**:aimux-web 绑定非 loopback 地址时,请求可用 `env:任意变量` + 自带 `base_url` 把服务进程的密钥发往任意主机(附录 tl:TOOLS-1,已核实)。已单独立任务修复,不等待本文。 +- **安全**:aimux-web 绑定非 loopback 地址时,请求可用 `env:任意变量` + 自带 `base_url` 把服务进程的密钥发往任意主机(附录 tl:TOOLS-1,已核实)。已在 master 修复,不等待本文:`env:` 只接受该 provider 注册的 key 变量;请求自带非默认 `base_url` 时只用请求内直接给的 key,不用 Settings 已存 key,也不读环境变量(`tools/aimux-web/src/wire.rs` `resolve_api_key`)。 - **真实 bug**:附录 §5 列出 53 条,标注了"重构后自然消除 / 需独立修复"。其中 SSE 1 MiB 单事件上限(as:STREAM-1)、Bedrock converse-stream 非真流式与吞 exception(as:STREAM-2)、Bedrock/Vertex embedding 只嵌入第一个值(ce:BUG-1)、Python 阻塞调用持有 GIL(bd:BIND-4)不依赖重构,建议先行修复。 ### 0.7 与 ROADMAP / RFC-0036 既有承诺的关系 diff --git a/docs/api/c.md b/docs/api/c.md index 2bd71cfa1..954b885fd 100644 --- a/docs/api/c.md +++ b/docs/api/c.md @@ -439,6 +439,7 @@ A pull timeout is a **state**, not an error. | Function | Description | |------|------| | `void aimux_drop_handle(handle)` | Free a handle (`0` is a no-op) | +| `void aimux_drop_handle_ptr(void *handle)` | Same, with the handle cast to a pointer: the `void (*)(void *)` callback a finalizer (Dart `NativeFinalizer`) needs. Never call `aimux_drop_handle` through that signature | | `void aimux_free_string(ptr)` | Free a result / getter string (`NULL` is safe) | | `void aimux_error_free(err)` | Release a returned error (`NULL` is safe); never `aimux_drop_handle` / `aimux_free_string` | | `aimux_error_code` / `aimux_error_message` / `aimux_error_*` | Read an error (see aimux-error.h) | diff --git a/rfc/0029-web-console.md b/rfc/0029-web-console.md index 456451a2b..defbb364e 100644 --- a/rfc/0029-web-console.md +++ b/rfc/0029-web-console.md @@ -222,10 +222,12 @@ POST /api/calls - **API 面**(见 §5.1):`GET` 只返回 `{provider, status, hint}`(hint = 末 4 位,不足 4 位给长度),**明文永不回传**;`PUT`/`DELETE` 非回环绑定时返回 403。 - **回环门控**:后端启动时已知 bind host;非 `127.0.0.1`/`localhost`/`::1` 时 Settings 写接口 403(错误信息引导 `env:` 引用),`GET` 附 `plaintext_entry: false` 供前端隐藏输入框,请求内明文 spec 亦维持拒绝。 - **调用优先级**(`wire.rs` `resolve_api_key`,三个调用方 calls/cache-probe/replay 均带上 provider 名): - 1. **请求内显式 spec** —— `env:VAR` 引用(任意绑定可用);明文仅限回环绑定(Playground 快速用一次、不保存的场景); + 1. **请求内显式 spec** —— `env:VAR` 引用(任意绑定可用;`VAR` 只能是该 provider 注册的 key 变量,如 openai 的 `OPENAI_API_KEY`,其他变量名一律拒绝、不读取);明文仅限回环绑定(Playground 快速用一次、不保存的场景); 2. **Settings 已存该 provider 的 key**(KeyStore,含启动时从磁盘加载的); 3. `None` —— provider 读自己注册的默认 env var。 + 请求自带 `base_url`(Replay 时为录制里的地址)且不是该 provider 的默认地址时,只用请求内直接给的 key(明文,仅回环):第 2、3 级不适用,`env:` 引用被拒绝,不给 key 则不带 key 发出。否则任何客户端都能把服务端的 key 发往自己指定的主机。 + **威胁模型**:信任边界 = 本机回环客户端。回环下浏览器 ↔ 后端信道视为可信,故允许网页内明文录入与请求内明文直传;非回环(如 `--host 0.0.0.0` LAN 共享)下任何局域网客户端都能 POST,故全部明文入口关闭。明文 key 永不进日志、永不回传前端、永不进 Recording(录制层脱敏复用 RFC-0023 规则)。 --- diff --git a/tools/aimux-web/README.md b/tools/aimux-web/README.md index 06541c334..6104b90f8 100644 --- a/tools/aimux-web/README.md +++ b/tools/aimux-web/README.md @@ -64,12 +64,17 @@ cargo run -p aimux-web -- --port 8787 --no-open "记住(写入磁盘)" 会持久化到配置目录的 `keys.json`(权限 `0600`)并在下次 启动自动加载。列表只显示掩码(末 4 位),明文永不回传前端。 2. **shell 环境变量(可选)**:`export OPENAI_API_KEY=sk-…`(或对应 provider - 的注册 env var);Playground 等页的 key 字段也可填 `env:VAR` 引用。 + 的注册 env var);Playground 等页的 key 字段也可填 `env:VAR` 引用,`VAR` + 只能是所选 provider 注册的 key 变量(如 `env:OPENAI_API_KEY`)。 调用优先级:请求内显式 `env:` 引用 > Settings 已存 key > provider 注册的默认 env var。非回环绑定(`--host 0.0.0.0` 等)时网页内 key 管理自动禁用,只能走 `env:` 引用。 +自填 `base_url`(不是该 provider 的默认地址)时只用请求里直接给的 key(仅回环 +绑定可填明文):Settings 已存 key 和环境变量都不会发往该地址,`env:` 引用被拒绝; +key 留空则不带 key 调用。Replay 一条指向非默认地址的录制时同样如此。 + 打开后使用 Playground / Agent 页跑一次调用,Traces 页即可看到完整录制。 ## 开发模式(前端热更新) diff --git a/tools/aimux-web/src/api/cache_probe.rs b/tools/aimux-web/src/api/cache_probe.rs index 8f5152024..95ab747ad 100644 --- a/tools/aimux-web/src/api/cache_probe.rs +++ b/tools/aimux-web/src/api/cache_probe.rs @@ -81,6 +81,7 @@ async fn run_inner(state: AppState, req: ProbeRequest) -> Result, Json(req): Json Result { // api_key: explicit spec (env: ref, or plaintext on loopback) > Settings - // key store > provider's registered env var (RFC-0029 §5.5). + // key store > provider's registered env var (RFC-0029 §5.5). A custom + // base_url gets only a key carried in the request. let key = wire::resolve_api_key( req.api_key.as_deref(), &req.provider, &state.keys, state.loopback, + model_builder::is_custom_base_url(&req.provider, req.base_url.as_deref()), )?; let model: Arc = if req.mock { @@ -135,3 +137,142 @@ async fn stream_response( Err(e) => err_response(e), } } + +#[cfg(test)] +mod tests { + use std::sync::{Arc, Mutex}; + + use axum::http::StatusCode; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::TcpListener; + use tower::ServiceExt; + + use crate::api::router; + use crate::settings::KeyStore; + use crate::state::AppState; + + const CHAT: &str = r#"{"id":"chatcmpl-1","model":"m","choices":[{"index":0,"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}"#; + + /// A one-endpoint OpenAI-compatible server that keeps every raw request. + async fn capture_server() -> (String, Arc>>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("http://{}/v1", listener.local_addr().unwrap()); + let seen = Arc::new(Mutex::new(Vec::new())); + let sink = seen.clone(); + tokio::spawn(async move { + while let Ok((mut socket, _)) = listener.accept().await { + let mut raw = Vec::new(); + let mut buf = [0u8; 4096]; + loop { + let n = socket.read(&mut buf).await.unwrap_or(0); + raw.extend_from_slice(&buf[..n]); + let text = String::from_utf8_lossy(&raw).to_string(); + let complete = text.split_once("\r\n\r\n").is_some_and(|(head, body)| { + let length = head + .lines() + .find_map(|l| { + l.to_ascii_lowercase() + .strip_prefix("content-length:") + .map(|v| v.trim().parse::().unwrap_or(0)) + }) + .unwrap_or(0); + body.len() >= length + }); + if n == 0 || complete { + break; + } + } + sink.lock() + .unwrap() + .push(String::from_utf8_lossy(&raw).to_string()); + let reply = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{CHAT}", + CHAT.len() + ); + let _ = socket.write_all(reply.as_bytes()).await; + } + }); + (url, seen) + } + + async fn call(app: axum::Router, body: serde_json::Value) -> (StatusCode, String) { + let resp = app + .oneshot( + axum::http::Request::post("/api/calls") + .header("content-type", "application/json") + .body(axum::body::Body::from(body.to_string())) + .unwrap(), + ) + .await + .unwrap(); + let status = resp.status(); + let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX) + .await + .unwrap(); + (status, String::from_utf8_lossy(&bytes).to_string()) + } + + fn request(provider: &str, base_url: &str, api_key: Option<&str>) -> serde_json::Value { + serde_json::json!({ + "provider": provider, + "model": "m", + "api_key": api_key, + "base_url": base_url, + "stream": false, + "options": {"max_retries": 0}, + "messages": [{"role": "user", "content": [{"type": "text", "text": "hi"}]}], + }) + } + + /// A caller-supplied base_url never receives the Settings key or the + /// provider's environment key: with no key in the request, the call goes + /// out without one. + #[tokio::test] + async fn custom_base_url_gets_no_stored_or_env_key() { + // A registry-style provider whose key variable this test owns. + aimux_providers::load_providers_from_json( + r#"{"providers":[{"name":"web-test-custom-base","base_url":"https://example.invalid/v1","env_var":"AIMUX_WEB_TEST_CUSTOM_BASE_KEY"}]}"#, + ) + .unwrap(); + // SAFETY: This test owns a unique process-wide env key and removes it below. + unsafe { std::env::set_var("AIMUX_WEB_TEST_CUSTOM_BASE_KEY", "env-secret-7") }; + + let keys = KeyStore::from_path(None); + keys.set("openai", "sk-stored-secret-1", false).unwrap(); + keys.set("web-test-custom-base", "sk-stored-secret-2", false) + .unwrap(); + let app = router(AppState::with_bind_host_and_store("127.0.0.1", keys)); + let (url, seen) = capture_server().await; + + for provider in ["openai", "web-test-custom-base"] { + let (status, body) = call(app.clone(), request(provider, &url, None)).await; + assert_eq!(status, StatusCode::OK, "{provider}: {body}"); + } + // An env: reference is refused outright for a custom base_url. + let (status, body) = call( + app.clone(), + request( + "web-test-custom-base", + &url, + Some("env:AIMUX_WEB_TEST_CUSTOM_BASE_KEY"), + ), + ) + .await; + assert_eq!(status, StatusCode::BAD_REQUEST, "{body}"); + // A key carried in the request (loopback) is what the endpoint gets. + let (status, body) = call(app, request("openai", &url, Some("sk-for-this-endpoint"))).await; + assert_eq!(status, StatusCode::OK, "{body}"); + + // SAFETY: Cleanup of the unique test key set above. + unsafe { std::env::remove_var("AIMUX_WEB_TEST_CUSTOM_BASE_KEY") }; + + let seen = seen.lock().unwrap(); + assert_eq!(seen.len(), 3, "the refused call must not reach the server"); + for raw in seen.iter() { + for secret in ["sk-stored-secret-1", "sk-stored-secret-2", "env-secret-7"] { + assert!(!raw.contains(secret), "{secret} leaked: {raw}"); + } + } + assert!(seen[2].contains("sk-for-this-endpoint"), "{}", seen[2]); + } +} diff --git a/tools/aimux-web/src/api/replay.rs b/tools/aimux-web/src/api/replay.rs index 776685f07..1d5844b60 100644 --- a/tools/aimux-web/src/api/replay.rs +++ b/tools/aimux-web/src/api/replay.rs @@ -12,10 +12,11 @@ use aimux_core::error::AiMuxError; use aimux_core::language_model::LanguageModel; use aimux_core::language_model_message::convert_to_language_model_prompt; use aimux_core::message::ModelPrompt; -use aimux_core::recording::Recording; +use aimux_core::recording::{ProviderRecord, Recording}; use aimux_core::replay::{MockReplayModel, ReplayOverrides, replay_with_model}; use crate::api::err_response; +use crate::model_builder; use crate::state::AppState; use crate::wire::{self, WireMessage}; @@ -63,17 +64,25 @@ async fn run_inner(state: AppState, req: ReplayRequest) -> Result { - wire::resolve_api_key(Some(s), &provider_name, &state.keys, state.loopback)? + wire::resolve_api_key(Some(s), &provider_name, &state.keys, state.loopback, custom)? } - "none" => wire::resolve_api_key(None, &provider_name, &state.keys, state.loopback)?, + "none" => wire::resolve_api_key(None, &provider_name, &state.keys, state.loopback, custom)?, _ => wire::resolve_api_key( req.api_key.as_deref(), &provider_name, &state.keys, state.loopback, + custom, ) .map_err(|_| { AiMuxError::InvalidArgument( @@ -116,6 +125,17 @@ async fn run_inner(state: AppState, req: ReplayRequest) -> Result Option { + p.provider_options + .as_ref() + .and_then(|options| options.get("base_url")) + .and_then(serde_json::Value::as_str) + .or(p.base_url.as_deref()) + .map(str::to_string) +} + /// `POST /api/mock/load` — load NDJSON recordings into a `MockReplayModel` /// and switch the console to offline mock mode (no real API calls). pub async fn mock_load( diff --git a/tools/aimux-web/src/model_builder.rs b/tools/aimux-web/src/model_builder.rs index 3823e10ff..9fa6f4640 100644 --- a/tools/aimux-web/src/model_builder.rs +++ b/tools/aimux-web/src/model_builder.rs @@ -23,8 +23,9 @@ pub fn build_model( base_url: Option<&str>, ) -> Result, AiMuxError> { macro_rules! native { - ($provider_mod:ident, $config:ident, $provider_type:ident, $env:literal) => {{ - let key = native_key(provider, $env, api_key.clone())?; + ($provider_mod:ident, $config:ident, $provider_type:ident) => {{ + let env = native_key_env(provider).unwrap_or_default(); + let key = native_key(provider, env, api_key.clone())?; let mut cfg = aimux_providers::$provider_mod::$config::new(key); if let Some(url) = base_url { cfg = cfg.with_base_url(url); @@ -35,22 +36,12 @@ pub fn build_model( } match provider { - "openai" => native!(openai, OpenAIConfig, OpenAIProvider, "OPENAI_API_KEY"), - "anthropic" => native!( - anthropic, - AnthropicConfig, - AnthropicProvider, - "ANTHROPIC_API_KEY" - ), - "google" => native!( - google, - GoogleConfig, - GoogleProvider, - "GOOGLE_GENERATIVE_AI_API_KEY" - ), - "mistral" => native!(mistral, MistralConfig, MistralProvider, "MISTRAL_API_KEY"), - "xai" => native!(xai, XAIConfig, XAIProvider, "XAI_API_KEY"), - "cohere" => native!(cohere, CohereConfig, CohereProvider, "COHERE_API_KEY"), + "openai" => native!(openai, OpenAIConfig, OpenAIProvider), + "anthropic" => native!(anthropic, AnthropicConfig, AnthropicProvider), + "google" => native!(google, GoogleConfig, GoogleProvider), + "mistral" => native!(mistral, MistralConfig, MistralProvider), + "xai" => native!(xai, XAIConfig, XAIProvider), + "cohere" => native!(cohere, CohereConfig, CohereProvider), _ => { let mut options = ProviderOptions::default(); if let Some(url) = base_url { @@ -63,6 +54,36 @@ pub fn build_model( } } +/// Whether `base_url` points somewhere other than `provider`'s own default +/// endpoint. Stored and environment keys are never sent to such a URL: it +/// came from the caller (a request field, or an imported recording). +/// +/// The default is what [`build_model`] uses when no base URL is given; when +/// it cannot be determined (an unknown provider, a templated registry URL) +/// every given URL counts as custom. +pub fn is_custom_base_url(provider: &str, base_url: Option<&str>) -> bool { + let Some(url) = base_url else { + return false; + }; + let default = build_model(provider, Some(String::new()), "default-base-url", None) + .ok() + .and_then(|model| model.config_snapshot().base_url); + default.is_none_or(|default| default.trim_end_matches('/') != url.trim_end_matches('/')) +} + +/// The env var a native provider (see [`build_model`]) reads its key from. +pub fn native_key_env(provider: &str) -> Option<&'static str> { + Some(match provider { + "openai" => "OPENAI_API_KEY", + "anthropic" => "ANTHROPIC_API_KEY", + "google" => "GOOGLE_GENERATIVE_AI_API_KEY", + "mistral" => "MISTRAL_API_KEY", + "xai" => "XAI_API_KEY", + "cohere" => "COHERE_API_KEY", + _ => return None, + }) +} + /// Resolve the key for a native provider: explicit key, else the standard env var. fn native_key(provider: &str, env: &str, api_key: Option) -> Result { match api_key { @@ -74,3 +95,38 @@ fn native_key(provider: &str, env: &str, api_key: Option) -> Result, provider: &str, store: &crate::settings::KeyStore, allow_plaintext: bool, + custom_base_url: bool, ) -> Result, AiMuxError> { + let blank = spec.is_none_or(|s| s.trim().is_empty()); + if custom_base_url && blank { + return Ok(Some(String::new())); + } + if custom_base_url && spec.is_some_and(|s| s.starts_with("env:")) { + return Err(AiMuxError::InvalidArgument( + "env: key references are not sent to a custom base_url — pass that endpoint's \ + own key in the request (loopback binds only) or leave api_key empty" + .into(), + )); + } + if custom_base_url && !allow_plaintext { + return Err(AiMuxError::InvalidArgument( + "a custom base_url only takes a key carried in the request, and plaintext keys \ + are accepted only while the console is bound to a loopback address — leave \ + api_key empty to call that endpoint without a key" + .into(), + )); + } match spec { None => Ok(store.get(provider)), // An empty/whitespace-only field is "unset", not an empty key — fall @@ -300,6 +327,24 @@ pub fn resolve_api_key( "empty env: reference (use env:VAR_NAME)".into(), )); } + // Only the provider's own key variable: any other name would let a + // caller read an arbitrary server environment variable (and send + // it wherever the request points). + match registered_key_env(provider) { + Some(allowed) if allowed == var => {} + Some(allowed) => { + return Err(AiMuxError::InvalidArgument(format!( + "env:{var} is not allowed for provider '{provider}' — only its \ + registered key variable can be referenced (env:{allowed})" + ))); + } + None => { + return Err(AiMuxError::InvalidArgument(format!( + "provider '{provider}' has no registered key variable, so env: \ + references are not accepted for it" + ))); + } + } match std::env::var(var) { Ok(v) => Ok(Some(v)), Err(_) => Err(AiMuxError::InvalidArgument(format!( @@ -317,6 +362,15 @@ pub fn resolve_api_key( } } +/// The environment variable a provider's key is registered under: the +/// native providers' standard variable, or the registry / runtime overlay +/// entry's. The only name an `env:` reference may use for that provider. +pub fn registered_key_env(provider: &str) -> Option { + crate::model_builder::native_key_env(provider) + .map(str::to_string) + .or_else(|| aimux_providers::provider_api_key_env(provider)) +} + #[cfg(test)] mod tests { use super::*; @@ -329,23 +383,41 @@ mod tests { let store = KeyStore::from_path(None); store.set("openai", "sk-from-store", false).unwrap(); assert_eq!( - resolve_api_key(Some(""), "openai", &store, true).unwrap(), + resolve_api_key(Some(""), "openai", &store, true, false).unwrap(), Some("sk-from-store".to_string()) ); assert_eq!( - resolve_api_key(Some(" "), "unknown-provider", &store, true).unwrap(), + resolve_api_key(Some(" "), "unknown-provider", &store, true, false).unwrap(), None ); } + /// Register a console-test provider whose key variable is `env_var`. + fn test_provider(name: &str, env_var: &str) { + aimux_providers::load_providers_from_json(&format!( + r#"{{"providers":[{{"name":"{name}","base_url":"https://example.invalid/v1","env_var":"{env_var}"}}]}}"# + )) + .unwrap(); + } + #[test] fn resolve_key_priority_explicit_env_over_store() { + test_provider("web-test-priority", "AIMUX_WEB_TEST_KEY"); let store = KeyStore::from_path(None); - store.set("openai", "sk-from-store", false).unwrap(); + store + .set("web-test-priority", "sk-from-store", false) + .unwrap(); // SAFETY: This test owns a unique process-wide env key and removes it below. unsafe { std::env::set_var("AIMUX_WEB_TEST_KEY", "k-123") }; assert_eq!( - resolve_api_key(Some("env:AIMUX_WEB_TEST_KEY"), "openai", &store, true).unwrap(), + resolve_api_key( + Some("env:AIMUX_WEB_TEST_KEY"), + "web-test-priority", + &store, + true, + false + ) + .unwrap(), Some("k-123".to_string()) ); // SAFETY: Cleanup of the unique test key set above. @@ -357,18 +429,18 @@ mod tests { let store = KeyStore::from_path(None); // Tier 3: empty store → provider's registered env var. assert_eq!( - resolve_api_key(None, "openai", &store, false).unwrap(), + resolve_api_key(None, "openai", &store, false, false).unwrap(), None ); // Tier 2: stored key used when no explicit spec. store.set("openai", "sk-from-store", false).unwrap(); assert_eq!( - resolve_api_key(None, "openai", &store, false).unwrap(), + resolve_api_key(None, "openai", &store, false, false).unwrap(), Some("sk-from-store".to_string()) ); // Store lookup is per-provider. assert_eq!( - resolve_api_key(None, "anthropic", &store, false).unwrap(), + resolve_api_key(None, "anthropic", &store, false, false).unwrap(), None ); } @@ -380,31 +452,95 @@ mod tests { // Loopback: explicit plaintext wins (one-shot use, not saved). assert_eq!( - resolve_api_key(Some("sk-literal"), "openai", &store, true).unwrap(), + resolve_api_key(Some("sk-literal"), "openai", &store, true, false).unwrap(), Some("sk-literal".to_string()) ); // Non-loopback: plaintext rejected, even over a stored key. - let err = resolve_api_key(Some("sk-literal"), "openai", &store, false).unwrap_err(); + let err = resolve_api_key(Some("sk-literal"), "openai", &store, false, false).unwrap_err(); assert!(err.to_string().contains("loopback"), "err: {err}"); // env: references still work on non-loopback. + test_provider("web-test-lan", "AIMUX_WEB_TEST_KEY_LAN"); // SAFETY: This test owns a unique process-wide env key and removes it below. unsafe { std::env::set_var("AIMUX_WEB_TEST_KEY_LAN", "k-lan") }; assert_eq!( - resolve_api_key(Some("env:AIMUX_WEB_TEST_KEY_LAN"), "openai", &store, false,).unwrap(), + resolve_api_key( + Some("env:AIMUX_WEB_TEST_KEY_LAN"), + "web-test-lan", + &store, + false, + false + ) + .unwrap(), Some("k-lan".to_string()) ); // SAFETY: Cleanup of the unique test key set above. unsafe { std::env::remove_var("AIMUX_WEB_TEST_KEY_LAN") }; } + #[test] + fn env_reference_is_limited_to_the_providers_key_variable() { + let store = KeyStore::from_path(None); + // Any other server variable is refused before it is read. + for spec in ["env:HOME", "env:PATH", "env:ANTHROPIC_API_KEY"] { + let err = resolve_api_key(Some(spec), "openai", &store, true, false).unwrap_err(); + assert!(err.to_string().contains("not allowed"), "{spec}: {err}"); + assert!(err.to_string().contains("env:OPENAI_API_KEY"), "{err}"); + } + assert!(resolve_api_key(Some("env:HOME"), "groq", &store, false, false).is_err()); + // A provider without a registered variable accepts no env: reference. + let err = + resolve_api_key(Some("env:HOME"), "no-such-provider", &store, true, false).unwrap_err(); + assert!( + err.to_string().contains("no registered key variable"), + "{err}" + ); + + assert_eq!(registered_key_env("groq").as_deref(), Some("GROQ_API_KEY")); + assert_eq!( + registered_key_env("google").as_deref(), + Some("GOOGLE_GENERATIVE_AI_API_KEY") + ); + } + + #[test] + fn custom_base_url_only_takes_a_key_from_the_request() { + let store = KeyStore::from_path(None); + store.set("openai", "sk-from-store", false).unwrap(); + // No key in the request: no stored key, and Some("") so the provider + // does not read its env var either. + for spec in [None, Some(""), Some(" ")] { + assert_eq!( + resolve_api_key(spec, "openai", &store, true, true).unwrap(), + Some(String::new()) + ); + } + // env: references stay on the server. + let err = + resolve_api_key(Some("env:OPENAI_API_KEY"), "openai", &store, true, true).unwrap_err(); + assert!(err.to_string().contains("custom base_url"), "{err}"); + // A key in the request: on loopback only, as everywhere else. + assert_eq!( + resolve_api_key(Some("sk-literal"), "openai", &store, true, true).unwrap(), + Some("sk-literal".to_string()) + ); + assert!(resolve_api_key(Some("sk-literal"), "openai", &store, false, true).is_err()); + } + #[test] fn resolve_key_rejects_plaintext_and_missing_env() { let store = KeyStore::from_path(None); - assert!(resolve_api_key(Some("sk-literal"), "openai", &store, false).is_err()); + assert!(resolve_api_key(Some("sk-literal"), "openai", &store, false, false).is_err()); assert!( - resolve_api_key(Some("env:AIMUX_WEB_MISSING_VAR"), "openai", &store, true).is_err() + resolve_api_key( + Some("env:AIMUX_WEB_MISSING_VAR"), + "openai", + &store, + true, + false + ) + .is_err() ); - assert!(resolve_api_key(Some("env:"), "openai", &store, true).is_err()); + assert!(resolve_api_key(Some("env:"), "openai", &store, true, false).is_err()); } #[test]