From 07df3bf4493b5a6ef409b2b7b36d7c9702d8812f Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 11 Oct 2026 14:49:05 +0000 Subject: [PATCH 1/9] fix(bedrock): derive the SigV4 key and canonical request as AWS does The signing key chain passed every HMAC its key and data swapped and dropped the AWS4 prefix on the secret, so AWS rejected every request signed with an access key. The canonical request also left path segments unencoded (the ':' in Bedrock model ids), kept the query unsorted and did not collapse header whitespace. Bedrock, Claude on AWS and Polly share this signer. The new tests pin signatures computed by botocore's SigV4Auth. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01A3WxDBWBA4ituJjXsHaEAa --- aimux-providers/src/bedrock/sigv4.rs | 218 +++++++++++++++++++++++++-- 1 file changed, 208 insertions(+), 10 deletions(-) diff --git a/aimux-providers/src/bedrock/sigv4.rs b/aimux-providers/src/bedrock/sigv4.rs index e915d95ca..b5fad16c3 100644 --- a/aimux-providers/src/bedrock/sigv4.rs +++ b/aimux-providers/src/bedrock/sigv4.rs @@ -55,6 +55,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 +89,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 +119,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 +158,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 +185,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); @@ -220,4 +319,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=" + ); + } } From 9503d91a64e26aebd83098a45d273612ebd6b709 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 12 Oct 2026 09:31:06 +0000 Subject: [PATCH 2/9] fix(providers): keep AWS secret keys and Codex tokens out of Debug output AwsCredentials and CodexTokens derived Debug, so a {:?} in logs or an error message printed the secret access key, session token, and the Codex access and refresh tokens. Both now have hand-written Debug impls that print [REDACTED] for the secrets. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01A3WxDBWBA4ituJjXsHaEAa --- aimux-providers/src/bedrock/sigv4.rs | 31 +++++++++++++++++++++++++++- aimux-providers/src/codex.rs | 31 +++++++++++++++++++++++++++- 2 files changed, 60 insertions(+), 2 deletions(-) diff --git a/aimux-providers/src/bedrock/sigv4.rs b/aimux-providers/src/bedrock/sigv4.rs index b5fad16c3..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 { @@ -282,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 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")); + } } From 13fe482105c3b0d5430b90b52d01f8399d5fc7a7 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 13 Oct 2026 02:58:15 +0000 Subject: [PATCH 3/9] fix(core): ignore non-finite Retry-After values in retry_after_hint Rust's f64 parser accepts `inf`, `Infinity` and `NaN`, and `as i64` saturates infinity, so `retry-after-ms: Infinity` produced an i64::MAX hint. The core retry loop dropped it through its window check, but the public hint (FFI `aimux_error_retry_ms`, the Node and Python `retry_ms`) reported it. A non-finite value, including seconds that overflow once converted to milliseconds, is now no hint, as upstream never uses it either. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01A3WxDBWBA4ituJjXsHaEAa --- aimux-core/src/error.rs | 42 ++++++++++++++++++++++++++++++++++++++--- 1 file changed, 39 insertions(+), 3 deletions(-) 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] From d38a9d3e68d441671b79829f924e165cff517cf5 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 13 Oct 2026 02:59:04 +0000 Subject: [PATCH 4/9] fix(core): redact the Prodia and Azure subscription key headers is_sensitive_key matched only cookie, key, authorization, api-key and *token names, so X-Prodia-Key and Azure's Ocp-Apim-Subscription-Key were written in clear text into recordings (and logs, which share the policy) whenever recording was on. Both names are now sensitive; a recording test sends them through the HTTP layer and checks that the stored exchange holds [REDACTED]. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01A3WxDBWBA4ituJjXsHaEAa --- aimux-core/src/recording.rs | 8 +++ .../tests/http_recording_test.rs | 53 +++++++++++++++++++ 2 files changed, 61 insertions(+) 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-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() { From 0fce438f622399c0a1227897f833119e2def95d6 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 13 Oct 2026 03:00:24 +0000 Subject: [PATCH 5/9] fix(provider-utils): strip URL userinfo from logs, errors and recordings sanitized_request_url dropped only the query, so a base URL written as https://user:password@host kept its credentials in tracing fields, ApiCallError.url and recorded exchanges. It now clears the username and password too, also for a URL that does not parse. The download guard's rejection message reuses it, so a blocked ftp://user:pass@ URL no longer echoes the password. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01A3WxDBWBA4ituJjXsHaEAa --- aimux-provider-utils/src/download_guard.rs | 15 ++++++-- aimux-provider-utils/src/http.rs | 43 +++++++++++++++++++++- 2 files changed, 54 insertions(+), 4 deletions(-) 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" + ); } } From 3d578954cdf08594d43dc265487c3efb36ee9144 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 13 Oct 2026 03:03:34 +0000 Subject: [PATCH 6/9] fix(flutter): give the NativeFinalizer a pointer-sized drop entry point The Dart NativeFinalizer called aimux_drop_handle(uint64_t) through a void (*)(void *) pointer. On 32-bit ARM (armeabi-v7a, which the release ships) a uint64_t argument does not travel like a pointer, so the drop received a garbage high word and a GC'd Model or ProviderHandle leaked its handle. aimux-ffi now exports aimux_drop_handle_ptr(void *), which takes the handle as a pointer-sized value and drops it, and the Flutter finalizer uses it. Tests: the Rust unit test and the export smoke test release handles through it, and a Flutter test checks it releases the handle its address names. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01A3WxDBWBA4ituJjXsHaEAa --- aimux-ffi/aimux-ffi.h | 10 ++++++++ aimux-ffi/src/lib.rs | 31 ++++++++++++++++++++++++ aimux-ffi/tests/exports_smoke_test.rs | 22 +++++++++-------- bindings/flutter/lib/aimux.dart | 22 ++++++++++------- bindings/flutter/test/aimux_test.dart | 35 +++++++++++++++++++++++++++ docs/api/c.md | 1 + 6 files changed, 102 insertions(+), 19 deletions(-) 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/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/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) | From 478522dbbd0e8e345d87db633c801c237abcf89b Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 13 Oct 2026 03:08:37 +0000 Subject: [PATCH 7/9] fix(python): release the GIL around every blocking native call Fourteen of the binding's block_on calls (generate_text, generate_object, consume_stream_text, generate_text_as_openai, ProviderHandle.list_models, get_model_specs and the eight multimodal calls) held the GIL while they waited on the network, so every other Python thread stopped for the whole request. They now run inside allow_threads, as decide() and the stream iterator already did, and so do the recording flushes, which wait on the writer thread. Results and errors are unchanged. test_gil.py serves the mock API from a thread of the test process: it can only answer while the call releases the GIL. Without the fix both calls fail with their 10 s total timeout. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01A3WxDBWBA4ituJjXsHaEAa --- bindings/python/src/lib.rs | 65 +++++++++++++++------- bindings/python/src/multimodal.rs | 64 ++++++++++++++------- bindings/python/tests/test_e2e.py | 6 +- bindings/python/tests/test_gil.py | 92 +++++++++++++++++++++++++++++++ 4 files changed, 184 insertions(+), 43 deletions(-) create mode 100644 bindings/python/tests/test_gil.py 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 From d5a10bd0d788042f5068d6bb2b7565d5489a38d0 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 13 Oct 2026 03:12:21 +0000 Subject: [PATCH 8/9] fix(web): accept env:VAR only for the provider's registered key variable A request's `api_key: "env:VAR"` read any environment variable of the console process. Together with a caller-supplied base_url, a client of a non-loopback console could send any server variable to a host of its choice. resolve_api_key now accepts only the variable the provider's key is registered under (the native providers' standard variable, or the registry / overlay entry's) and refuses every other name before reading it. This covers the calls, cache-probe and replay paths, which all resolve keys through it. aimux-providers gains provider_api_key_env(), the variable provider() reads a key from when none is passed, so the console does not keep its own copy of the registry. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01A3WxDBWBA4ituJjXsHaEAa --- aimux-providers/src/lib.rs | 2 +- aimux-providers/src/provider.rs | 71 +++++++++++++++++++++++ rfc/0029-web-console.md | 2 +- tools/aimux-web/README.md | 3 +- tools/aimux-web/src/model_builder.rs | 40 +++++++------ tools/aimux-web/src/wire.rs | 86 ++++++++++++++++++++++++++-- 6 files changed, 178 insertions(+), 26 deletions(-) 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/rfc/0029-web-console.md b/rfc/0029-web-console.md index 456451a2b..2b0fb41e3 100644 --- a/rfc/0029-web-console.md +++ b/rfc/0029-web-console.md @@ -222,7 +222,7 @@ 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。 diff --git a/tools/aimux-web/README.md b/tools/aimux-web/README.md index 06541c334..b7e79b3f2 100644 --- a/tools/aimux-web/README.md +++ b/tools/aimux-web/README.md @@ -64,7 +64,8 @@ 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 管理自动禁用,只能走 diff --git a/tools/aimux-web/src/model_builder.rs b/tools/aimux-web/src/model_builder.rs index 3823e10ff..c0ad80023 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,19 @@ pub fn build_model( } } +/// 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 { diff --git a/tools/aimux-web/src/wire.rs b/tools/aimux-web/src/wire.rs index 27c646919..d5bd75228 100644 --- a/tools/aimux-web/src/wire.rs +++ b/tools/aimux-web/src/wire.rs @@ -274,8 +274,9 @@ pub fn to_generate_options( /// /// Priority: /// 1. **explicit request spec** — `Some("env:VAR")` reads the environment -/// variable; a plaintext literal is accepted only when the console is -/// loopback-bound (`allow_plaintext`, one-shot Playground use); +/// variable, which must be the provider's registered key variable +/// ([`registered_key_env`]); a plaintext literal is accepted only when the +/// console is loopback-bound (`allow_plaintext`, one-shot Playground use); /// 2. **Settings key store** — a key saved for this provider in the console /// Settings page (`store.get(provider)`); /// 3. `None` — the provider reads its registered env var. @@ -300,6 +301,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 +336,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::*; @@ -338,14 +366,31 @@ mod tests { ); } + /// 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 + ) + .unwrap(), Some("k-123".to_string()) ); // SAFETY: Cleanup of the unique test key set above. @@ -387,16 +432,47 @@ mod tests { let err = resolve_api_key(Some("sk-literal"), "openai", &store, 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 + ) + .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).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).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).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 resolve_key_rejects_plaintext_and_missing_env() { let store = KeyStore::from_path(None); From bf0d19efc723661136c73dafb93e51794a780095 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 13 Oct 2026 03:16:36 +0000 Subject: [PATCH 9/9] fix(web): send no stored or env key to a caller-supplied base_url MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit /api/calls and /api/cache-probe resolved the key from the request spec, then the Settings key store, then the provider's env var, and passed the request's base_url to the provider unchanged. A client could point base_url at its own host and receive the console's key. Replay had the same hole through an imported recording's base URL. When the base URL differs from the provider's own default, only a key carried in the request (plaintext, loopback binds only) is used: the Settings key and the env var are not read, env: references are refused, and a call without a key goes out with an empty key instead of the provider's env fallback. The test sends calls with a custom base_url to a local server and checks that neither the stored key nor the env key arrives. The alignment doc's §0.6 note now records the web fixes as done. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01A3WxDBWBA4ituJjXsHaEAa --- docs/aisdk-architecture-alignment.md | 2 +- rfc/0029-web-console.md | 2 + tools/aimux-web/README.md | 4 + tools/aimux-web/src/api/cache_probe.rs | 1 + tools/aimux-web/src/api/calls.rs | 143 ++++++++++++++++++++++++- tools/aimux-web/src/api/replay.rs | 26 ++++- tools/aimux-web/src/model_builder.rs | 52 +++++++++ tools/aimux-web/src/wire.rs | 88 ++++++++++++--- 8 files changed, 299 insertions(+), 19 deletions(-) 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/rfc/0029-web-console.md b/rfc/0029-web-console.md index 2b0fb41e3..defbb364e 100644 --- a/rfc/0029-web-console.md +++ b/rfc/0029-web-console.md @@ -226,6 +226,8 @@ POST /api/calls 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 b7e79b3f2..6104b90f8 100644 --- a/tools/aimux-web/README.md +++ b/tools/aimux-web/README.md @@ -71,6 +71,10 @@ cargo run -p aimux-web -- --port 8787 --no-open 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 c0ad80023..9fa6f4640 100644 --- a/tools/aimux-web/src/model_builder.rs +++ b/tools/aimux-web/src/model_builder.rs @@ -54,6 +54,23 @@ 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 { @@ -78,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 @@ -357,11 +383,11 @@ 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 ); } @@ -388,7 +414,8 @@ mod tests { Some("env:AIMUX_WEB_TEST_KEY"), "web-test-priority", &store, - true + true, + false ) .unwrap(), Some("k-123".to_string()) @@ -402,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 ); } @@ -425,11 +452,11 @@ 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"); @@ -440,6 +467,7 @@ mod tests { Some("env:AIMUX_WEB_TEST_KEY_LAN"), "web-test-lan", &store, + false, false ) .unwrap(), @@ -454,13 +482,14 @@ mod tests { 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).unwrap_err(); + 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).is_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).unwrap_err(); + 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}" @@ -473,14 +502,45 @@ mod tests { ); } + #[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]