diff --git a/src/providers/codex/auth/manager.rs b/src/providers/codex/auth/manager.rs index 2226e14..ad4d691 100644 --- a/src/providers/codex/auth/manager.rs +++ b/src/providers/codex/auth/manager.rs @@ -9,6 +9,14 @@ use crate::auth::AuthStorage; pub struct CodexAuthManager> { pub store: CodexTokenStore, cached: Arc>>, + // Serializes token refreshes (single-flight). The `cached` mutex only + // guards cache reads/writes; without this lock, N concurrent requests + // hitting the expiry margin each POST /oauth/token with the SAME + // (single-use, rotating) refresh token — the first rotates it, the rest + // get 401 and used to clear_auth(), destroying the winner's fresh tokens. + // Observed in production as minutes-long all-requests-401 windows during + // agent fan-outs at token-expiry boundaries. + refresh_flight: Arc>, } impl> CodexAuthManager { @@ -16,6 +24,7 @@ impl> CodexAuthManager { Self { store, cached: Arc::new(Mutex::new(None)), + refresh_flight: Arc::new(Mutex::new(())), } } @@ -71,6 +80,29 @@ impl> CodexAuthManager { } fn refresh_now(&self, current: &StoredAuth) -> Result { + // Single-flight: hold the flight lock for the whole refresh. Racing + // callers block here, then discover the winner's tokens on re-check + // below and return without touching the token endpoint. + let _flight = self + .refresh_flight + .lock() + .map_err(|e| anyhow::anyhow!("{e}"))?; + + // Re-check under the lock: a concurrent flight (or another process + // sharing the store) may have refreshed while we waited. The store is + // the persisted truth; prefer it over both `current` and the cache. + let current = match self.store.load_auth()? { + Some(latest) => { + if latest.expires > Self::now_ms() + REFRESH_MARGIN_MS { + let mut guard = self.cached.lock().map_err(|e| anyhow::anyhow!("{e}"))?; + *guard = Some(latest.clone()); + return Ok(latest); + } + latest + } + None => current.clone(), + }; + if current.refresh.is_empty() { anyhow::bail!("No refresh token stored; re-authenticate"); } @@ -90,6 +122,17 @@ impl> CodexAuthManager { let status = resp.status().as_u16(); if status == 401 || status == 403 { + // Before destroying auth state: if the store's refresh token has + // rotated since we read `current`, a concurrent writer (e.g. + // another process sharing the Keychain entry) beat us — its + // tokens are good; return them instead of clobbering the store. + if let Ok(Some(latest)) = self.store.load_auth() { + if latest.refresh != current.refresh && latest.expires > Self::now_ms() { + let mut guard = self.cached.lock().map_err(|e| anyhow::anyhow!("{e}"))?; + *guard = Some(latest.clone()); + return Ok(latest); + } + } { let mut guard = self.cached.lock().map_err(|e| anyhow::anyhow!("{e}"))?; *guard = None; @@ -197,4 +240,60 @@ mod tests { .contains("Not authenticated") ); } + + #[test] + fn refresh_recheck_returns_concurrently_rotated_tokens_without_network() { + // Cache holds an EXPIRED token; the store already holds a FRESH one + // (as after a concurrent flight or another process refreshed). The + // re-check under the flight lock must return the store's tokens and + // never reach the token endpoint (no HTTP mock exists here — reaching + // the network would fail the test with a refresh error). + let store = test_store(); + let fresh = StoredAuth { + access: "rotated_access".into(), + refresh: "rotated_refresh".into(), + expires: 9_999_999_999_999, + account_id: Some("acct_1".into()), + }; + store.save_auth(fresh.clone()).unwrap(); + let manager = CodexAuthManager::new(store); + manager.set_cached(StoredAuth { + access: "stale_access".into(), + refresh: "stale_refresh".into(), + expires: 0, // expired -> get_auth enters refresh_now + account_id: Some("acct_1".into()), + }); + let result = manager.get_auth().unwrap(); + assert_eq!(result.access, "rotated_access"); + assert_eq!(result.refresh, "rotated_refresh"); + } + + #[test] + fn concurrent_get_auth_single_flights_to_rotated_tokens() { + use std::sync::Arc as StdArc; + let store = test_store(); + store + .save_auth(StoredAuth { + access: "rotated_access".into(), + refresh: "rotated_refresh".into(), + expires: 9_999_999_999_999, + account_id: None, + }) + .unwrap(); + let manager = StdArc::new(CodexAuthManager::new(store)); + manager.set_cached(StoredAuth { + access: "stale".into(), + refresh: "stale".into(), + expires: 0, + account_id: None, + }); + let mut handles = Vec::new(); + for _ in 0..8 { + let m = StdArc::clone(&manager); + handles.push(std::thread::spawn(move || m.get_auth().unwrap().access)); + } + for h in handles { + assert_eq!(h.join().unwrap(), "rotated_access"); + } + } } diff --git a/src/providers/codex/mod.rs b/src/providers/codex/mod.rs index f775f8b..5fa52cc 100644 --- a/src/providers/codex/mod.rs +++ b/src/providers/codex/mod.rs @@ -37,7 +37,7 @@ use self::translate::model_allowlist::{ assert_allowed_model, resolve_model_request, uses_responses_lite, }; use self::translate::reducer::finish_metadata_from_upstream; -use self::translate::request::{TranslateOptions, translate_request}; +use self::translate::request::{TranslateOptions, has_hosted_web_search, translate_request}; const MAX_RETRYABLE_LIVE_STREAM_RETRIES: u32 = 10; use self::translate::stream::translate_stream_bytes_with_traffic; @@ -109,7 +109,8 @@ impl Provider for CodexProvider { session_id: ctx.session_id.clone(), service_tier: resolved.service_tier.clone(), model: resolved.model.clone(), - use_responses_lite: uses_responses_lite(&resolved.model), + use_responses_lite: uses_responses_lite(&resolved.model) + && !has_hosted_web_search(&body), }, ) { Ok(t) => t, @@ -256,7 +257,8 @@ impl Provider for CodexProvider { session_id: None, service_tier: resolved.service_tier.clone(), model: resolved.model.clone(), - use_responses_lite: uses_responses_lite(&resolved.model), + use_responses_lite: uses_responses_lite(&resolved.model) + && !has_hosted_web_search(&body), }, ) { Ok(t) => t, diff --git a/src/providers/codex/translate/request.rs b/src/providers/codex/translate/request.rs index dcf6125..bd91d9c 100644 --- a/src/providers/codex/translate/request.rs +++ b/src/providers/codex/translate/request.rs @@ -293,6 +293,20 @@ pub fn normalize_strict_json_schema(schema: &Value) -> Value { } } +/// Hosted tools (web_search) are rejected by the Responses Lite lane, which +/// only supports function and custom tools. Requests carrying them must use +/// the full Responses API. +pub fn has_hosted_web_search(req: &MessagesRequest) -> bool { + req.extra + .get("tools") + .and_then(|v| v.as_array()) + .is_some_and(|tools| { + tools + .iter() + .any(|tool| tool.get("type").and_then(|v| v.as_str()) == Some("web_search_20250305")) + }) +} + pub fn translate_request( req: &MessagesRequest, opts: TranslateOptions, @@ -367,6 +381,21 @@ pub fn translate_request( out.tools = Some(tools); } + // Never force a web_search tool_choice the request didn't register — + // upstream 502s instead of ignoring it. + if matches!( + out.tool_choice, + Some(ResponsesToolChoice::WebSearch { .. }) + ) { + let has_web_search = out + .tools + .as_ref() + .is_some_and(|t| t.iter().any(|tool| matches!(tool, ResponsesTool::WebSearch(_)))); + if !has_web_search { + out.tool_choice = Some(ResponsesToolChoice::Auto); + } + } + if let Some(sid) = opts.session_id { out.prompt_cache_key = Some(sid); } @@ -463,13 +492,16 @@ fn read_tools(req: &MessagesRequest) -> Result>, anyho .collect() }); } - let has_filters = - filters.allowed_domains.is_some() || filters.blocked_domains.is_some(); + // The Codex backend zeroes out results when domain filters are + // present (allowed_domains=[github.com] -> 0 results for queries + // that hit the domain unfiltered), so don't forward them; the + // model filters from titles/URLs instead. + let _ = &filters; out.push(ResponsesTool::WebSearch(ResponsesWebSearchTool { kind: "web_search".to_string(), - external_web_access: false, + external_web_access: true, search_content_types: vec!["text".to_string(), "image".to_string()], - filters: if has_filters { Some(filters) } else { None }, + filters: None, })); } else { let name = tool @@ -885,6 +917,90 @@ mod tests { )); } + #[test] + fn has_hosted_web_search_detects_web_search_tool() { + let with: MessagesRequest = serde_json::from_value(json!({ + "model": "gpt-5.6-sol", + "messages": [{"role":"user", "content":"find it"}], + "tools": [ + {"name":"Bash", "input_schema":{}}, + {"type":"web_search_20250305", "name":"web_search"} + ] + })) + .unwrap(); + assert!(has_hosted_web_search(&with)); + + let without: MessagesRequest = serde_json::from_value(json!({ + "model": "gpt-5.6-sol", + "messages": [{"role":"user", "content":"run it"}], + "tools": [{"name":"Bash", "input_schema":{}}] + })) + .unwrap(); + assert!(!has_hosted_web_search(&without)); + } + + #[test] + fn responses_lite_downgrades_unregistered_web_search_tool_choice() { + // On the lite lane tools travel in the AdditionalTools developer + // prefix, so a top-level web_search tool_choice would reference a + // tool upstream doesn't know about and 502. + let req: MessagesRequest = serde_json::from_value(json!({ + "model": "gpt-5.6-sol", + "messages": [{"role":"user", "content":"find it"}], + "tools": [{ + "type":"web_search_20250305", + "name":"web_search" + }], + "tool_choice": {"type":"tool", "name":"web_search"} + })) + .unwrap(); + let out = translate_request( + &req, + TranslateOptions { + session_id: None, + service_tier: None, + model: "gpt-5.6-sol".to_string(), + use_responses_lite: true, + }, + ) + .unwrap(); + assert!(out.tools.is_none()); + assert!(matches!(out.tool_choice, Some(ResponsesToolChoice::Auto))); + } + + #[test] + fn full_lane_keeps_web_search_tool_choice_registered() { + let req: MessagesRequest = serde_json::from_value(json!({ + "model": "gpt-5.6-sol", + "messages": [{"role":"user", "content":"find it"}], + "tools": [{ + "type":"web_search_20250305", + "name":"web_search" + }], + "tool_choice": {"type":"tool", "name":"web_search"} + })) + .unwrap(); + let out = translate_request( + &req, + TranslateOptions { + session_id: None, + service_tier: None, + model: "gpt-5.6-sol".to_string(), + use_responses_lite: false, + }, + ) + .unwrap(); + assert!( + out.tools + .as_ref() + .is_some_and(|t| t.iter().any(|tool| matches!(tool, ResponsesTool::WebSearch(_)))) + ); + assert!(matches!( + out.tool_choice, + Some(ResponsesToolChoice::WebSearch { .. }) + )); + } + #[test] fn translate_read_tool_adds_codex_offset_guidance() { let req: MessagesRequest = serde_json::from_value(json!({