From 44ccf394573cabb95e1f751a59f9d76f6f79faf1 Mon Sep 17 00:00:00 2001 From: evanz-oai Date: Sat, 7 Nov 2026 02:57:57 +0000 Subject: [PATCH] Preserve model catalogs when starting child agents (#60049) ## Why Child startup can lose remote model capabilities such as `tool_search` when the selected catalog is uninitialized or disk data is older than the parent's in-memory snapshot. ## What changed - Initialize the final provider, auth, and originator catalog before resolving a child's startup model, then read it from memory. - Preserve loaded catalogs, including empty responses and updates that failed to persist, without overwriting concurrent refreshes. - Reuse catalogs across retry and WebSocket transport overrides while retaining routing and header distinctions. ## Testing Add regression coverage for child tool search with warm, cold, and disabled caches; failed cache writes; concurrent refreshes; empty catalogs; and Guardian reviewer catalog reuse. GitOrigin-RevId: 01d995a65b623cc7d6f26e2dc9b6e8cf380de782 --- codex-rs/core/src/session/mod.rs | 9 +- .../core/src/session/step_activation_tests.rs | 4 + codex-rs/core/src/thread_manager.rs | 58 +++-- .../core/src/thread_manager/model_catalog.rs | 33 ++- codex-rs/core/src/thread_manager_tests.rs | 79 +++++- .../core/tests/suite/injected_models_cache.rs | 68 ++++- codex-rs/core/tests/suite/items.rs | 3 +- .../core/tests/suite/models_cache_auth.rs | 13 +- codex-rs/core/tests/suite/models_cache_ttl.rs | 84 +++++- codex-rs/core/tests/suite/remote_models.rs | 246 ++++++++++++++++++ codex-rs/core/tests/suite/wake_reservation.rs | 4 + .../src/cache_identity_tests.rs | 8 +- codex-rs/models-manager/src/catalog.rs | 4 + codex-rs/models-manager/src/manager.rs | 125 +++++++-- codex-rs/models-manager/src/manager_tests.rs | 24 ++ 15 files changed, 682 insertions(+), 80 deletions(-) diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 570fda86dc..c52dd4e9c7 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -605,15 +605,14 @@ impl Session { session_source: &SessionSource, allow_provider_model_fallback: bool, ) -> CodexResult<(String, CatalogModel)> { + // Child catalogs are initialized by their selector. Reading only memory here + // preserves the loaded snapshot shared with the parent. let refresh_strategy = if session_source.is_non_root_agent() { - RefreshStrategy::Offline + RefreshStrategy::InMemory } else { RefreshStrategy::OnlineIfUncached }; - if config.features.enabled(Feature::RuntimeAgentInstances) - || config.model.is_none() - || refresh_strategy != RefreshStrategy::Offline - { + if !session_source.is_non_root_agent() { let _ = models_manager .list_models(refresh_strategy, config.http_client_factory()) .await; diff --git a/codex-rs/core/src/session/step_activation_tests.rs b/codex-rs/core/src/session/step_activation_tests.rs index 6b325d8a53..4ad17f2300 100644 --- a/codex-rs/core/src/session/step_activation_tests.rs +++ b/codex-rs/core/src/session/step_activation_tests.rs @@ -360,6 +360,10 @@ impl GatedModelsManager { } impl ModelsManager for GatedModelsManager { + fn ensure_initialized(&self, factory: HttpClientFactory) -> ModelsManagerFuture<'_, ()> { + self.inner.ensure_initialized(factory) + } + fn model_catalog_with_origins( &self, strategy: RefreshStrategy, diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 7362fb3bc7..8cbc7bc9b2 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -1116,12 +1116,16 @@ impl ThreadManager { /*parent_originator*/ None, ) .await; - let manager = self.state.models_for_thread( - &options.config, - &self.state.auth_manager, - &self.state.runtime_auth, - &originator, - ); + let manager = self + .state + .models_for_thread( + &options.config, + &self.state.auth_manager, + &self.state.runtime_auth, + &originator, + &source, + ) + .await; let (_, info) = Session::resolve_startup_model( manager.as_ref(), &options.config, @@ -1919,12 +1923,13 @@ impl ThreadManager { } impl ThreadManagerState { - fn models_for_thread( + async fn models_for_thread( &self, config: &Config, auth_manager: &Arc, runtime_auth: &crate::runtime_auth::RuntimeAuthSlot, originator: &str, + session_source: &SessionSource, ) -> SharedModelsManager { let runtime = crate::runtime_auth::runtime_for_config( runtime_auth, @@ -1932,13 +1937,16 @@ impl ThreadManagerState { Arc::clone(auth_manager), &self.installation_id, ); - let manager = self.model_catalog.for_thread( - config, - Arc::clone(auth_manager), - runtime, - Arc::ptr_eq(auth_manager, &self.auth_manager), - ); - manager.for_originator(originator).unwrap_or(manager) + self.model_catalog + .for_thread( + config, + Arc::clone(auth_manager), + runtime, + Arc::ptr_eq(auth_manager, &self.auth_manager), + originator, + session_source, + ) + .await } pub(crate) fn shared_thread_instructions_provider( @@ -2775,12 +2783,22 @@ impl ThreadManagerState { } else { Arc::new(mcp_manager.for_plugins_manager(Arc::clone(&plugins_manager))) }; - let models_manager = self.models_for_thread( - &config, - &auth_manager, - &runtime_auth, - &startup_context.originator, - ); + let models_manager = tokio::select! { + biased; + _ = agent_control.runtime().shutdown.cancelled() => { + if let Some(startup_guard) = startup_guard { + startup_guard.cleanup().await; + } + return Err(CodexErr::TurnAborted); + } + manager = self.models_for_thread( + &config, + &auth_manager, + &runtime_auth, + &startup_context.originator, + &session_source, + ) => manager, + }; let spawn_result = Session::spawn(SessionSpawnArgs { attached_thread, startup, diff --git a/codex-rs/core/src/thread_manager/model_catalog.rs b/codex-rs/core/src/thread_manager/model_catalog.rs index 9af414c718..3e4a4357fa 100644 --- a/codex-rs/core/src/thread_manager/model_catalog.rs +++ b/codex-rs/core/src/thread_manager/model_catalog.rs @@ -12,6 +12,7 @@ use codex_model_provider_info::ModelProviderInfo; use codex_models_manager::manager::ModelCatalogSettings; use codex_models_manager::manager::SharedModelsManager; use codex_protocol::openai_models::ModelsResponse; +use codex_protocol::protocol::SessionSource; use codex_runtime_agent_instance::RuntimeAuthManager; use std::sync::Arc; @@ -115,23 +116,35 @@ impl ThreadModelCatalog { } } - pub(super) fn for_thread( + #[allow(clippy::too_many_arguments)] + pub(super) async fn for_thread( &self, config: &Config, credentials: Arc, runtime: Option>, same_owner: bool, + originator: &str, + session_source: &SessionSource, ) -> SharedModelsManager { let models = config.model_catalog.as_ref().or_else(|| { self.catalog .as_ref() - .filter(|_| self.provider == config.model_provider) + .filter(|_| self.matches_catalog_provider(&config.model_provider)) }); let auth = ProviderAuthContext { credentials, runtime_auth: runtime.map(|runtime| (runtime, config.runtime_auth_mode)), }; - self.for_config(config, models, auth, same_owner) + let manager = self.for_config(config, models, auth, same_owner); + let manager = manager.for_originator(originator).unwrap_or(manager); + if session_source.is_non_root_agent() { + // The final provider/auth/originator catalog may be new even when the base + // manager was reused. Preserve loaded snapshots, including unsaved updates. + manager + .ensure_initialized(config.http_client_factory()) + .await; + } + manager } pub(super) fn for_config( @@ -145,7 +158,7 @@ impl ThreadModelCatalog { config.features.enabled(Feature::ApiKeyModelDiscovery); if same_owner && self.runtime_mode == auth.runtime_auth.as_ref().map(|(_, mode)| *mode) - && self.provider == config.model_provider + && self.matches_catalog_provider(&config.model_provider) && self.catalog.as_ref() == models && self.api_key_model_discovery_default == api_key_model_discovery_default { @@ -158,6 +171,18 @@ impl ThreadModelCatalog { ) } + fn matches_catalog_provider(&self, provider: &ModelProviderInfo) -> bool { + // Guardian overrides inference transport and retries without changing the catalog. + // Keep all other provider fields, including routing and headers, in the comparison. + self.provider + == ModelProviderInfo { + request_max_retries: self.provider.request_max_retries, + stream_max_retries: self.provider.stream_max_retries, + supports_websockets: self.provider.supports_websockets, + ..provider.clone() + } + } + pub(super) fn set_api_key_model_discovery_enabled(&self, enabled: bool) { self.factory.set_api_key_model_discovery_enabled(enabled); } diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index e46b32918c..c6ca9736e4 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -2248,6 +2248,8 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() let lifecycle_observed = Arc::new(std::sync::Mutex::new(Vec::new())); let mcp_observed = Arc::new(std::sync::Mutex::new(Vec::new())); let mcp_loaded = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + // This fixture counts explicit MCP projections, excluding background websocket prewarming. + config.model_provider.supports_websockets = false; let recorder = Arc::new(InitialDataRecorder { lifecycle_observed: Arc::clone(&lifecycle_observed), mcp_observed: Arc::clone(&mcp_observed), @@ -3432,8 +3434,10 @@ async fn metadata_update_without_result_reads_only_when_the_caller_needs_the_thr ); } +#[test_case::test_case(true; "disk_cache")] +#[test_case::test_case(false; "without_cache")] #[tokio::test] -async fn new_uses_active_provider_for_model_refresh() { +async fn new_uses_active_provider_for_model_refresh(with_cache: bool) { let server = MockServer::start().await; let models_mock = mount_models_once(&server, ModelsResponse { models: vec![] }).await; @@ -3453,7 +3457,11 @@ async fn new_uses_active_provider_for_model_refresh() { &config, auth_manager.clone(), TEST_INSTALLATION_ID.to_string(), - Arc::new(ModelCatalogSettings::disk(config.codex_home.to_path_buf())), + Arc::new(if with_cache { + ModelCatalogSettings::disk(config.codex_home.to_path_buf()) + } else { + ModelCatalogSettings::default() + }), ), crate::CodexAppsToolsCache::default(), SessionSource::Exec, @@ -3476,6 +3484,16 @@ async fn new_uses_active_provider_for_model_refresh() { .await; assert_eq!(models_mock.requests().len(), 1); + let mut thread_config = config.clone(); + thread_config.model = Some("remote-search-model".to_string()); + thread_config.model_provider.request_max_retries = Some(1); + thread_config.model_provider.stream_max_retries = Some(1); + thread_config + .model_provider + .http_headers + .get_or_insert_default() + .insert("X-OpenAI-Product-Sku".to_string(), "test-product".into()); + // Thread startup must select the originator catalog before resolving model metadata. for originator in [ "codex_work_cca", @@ -3484,24 +3502,32 @@ async fn new_uses_active_provider_for_model_refresh() { "codex_work_mobile", "chatgpt_cca", ] { - let mut model = codex_models_manager::model_info::model_info_from_slug("gpt-5.5"); + let mut model = + codex_models_manager::model_info::model_info_from_slug("remote-search-model"); model.description = Some(originator.to_string()); + model.supports_search_tool = true; + model.tool_mode = Some(codex_protocol::openai_models::ToolMode::CodeModeOnly); + model.use_responses_lite = true; wiremock::Mock::given(wiremock::matchers::method("GET")) .and(wiremock::matchers::path("/models")) .and(wiremock::matchers::header("originator", originator)) + .and(wiremock::matchers::header( + "X-OpenAI-Product-Sku", + "test-product", + )) .respond_with( wiremock::ResponseTemplate::new(200).set_body_json(ModelsResponse { models: vec![model], }), ) - .expect(1) + .expect(if with_cache { 1..=1 } else { 1..=2 }) .mount(&server) .await; let started = manager .start_thread(StartThreadOptions { metrics_service_name: Some(originator.to_string()), environments: Some(Vec::new()), - ..StartThreadOptions::new(config.clone()) + ..StartThreadOptions::new(thread_config.clone()) }) .await .expect("start scoped thread"); @@ -3510,9 +3536,50 @@ async fn new_uses_active_provider_for_model_refresh() { .session .services .models_manager - .get_model_info("gpt-5.5", &config.to_models_manager_config()) + .get_model_info( + "remote-search-model", + &thread_config.to_models_manager_config(), + ) .await; assert_eq!(model.description.as_deref(), Some(originator)); + assert!(model.supports_search_tool); + + let child = manager + .state + .spawn_child_thread( + StartThreadOptions { + session_source: Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id: started.thread_id, + depth: 1, + agent_path: None, + agent_nickname: None, + agent_role: None, + })), + environments: Some(Vec::new()), + ..StartThreadOptions::new(thread_config.clone()) + }, + started + .thread + .session + .services + .local_agent_runtime + .control(started.thread.session.session_id()), + Some(started.thread_id), + /*inherited_exec_policy*/ None, + ) + .await + .expect("spawn child with the parent's provider headers"); + let (turn, _) = child + .thread + .session + .new_turn_with_sub_id( + "child-turn".to_string(), + SessionSettingsUpdate::default(), + Default::default(), + ) + .await + .expect("prepare child's first turn"); + assert_eq!(turn.model_info().as_ref(), &model); } let legacy = manager .start_thread(StartThreadOptions { diff --git a/codex-rs/core/tests/suite/injected_models_cache.rs b/codex-rs/core/tests/suite/injected_models_cache.rs index e287ec0dcb..5a6787c9d8 100644 --- a/codex-rs/core/tests/suite/injected_models_cache.rs +++ b/codex-rs/core/tests/suite/injected_models_cache.rs @@ -129,14 +129,14 @@ impl ModelsCache for TestModelsCache { #[derive(Debug)] struct TestModelsEndpoint { - models: Vec, + models: Mutex>, fetch_count: AtomicUsize, } impl TestModelsEndpoint { fn new(models: Vec) -> Arc { Arc::new(Self { - models, + models: Mutex::new(models), fetch_count: AtomicUsize::new(0), }) } @@ -163,7 +163,7 @@ impl ModelsEndpointClient for TestModelsEndpoint { Box::pin(async move { self.fetch_count.fetch_add(1, Ordering::SeqCst); Ok(ModelsEndpointResponse { - models: self.models.clone(), + models: self.models.lock().expect("model response lock").clone(), etag: None, identity: self.identity().expect("test endpoint identity"), }) @@ -424,6 +424,57 @@ async fn threads_follow_provider_changes_and_preserve_host_cache_policy( Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn initialization_preserves_a_concurrent_refresh() -> Result<()> { + let (store_entered, entered) = oneshot::channel(); + let (release_store, release) = oneshot::channel(); + let cache = Arc::new(TestModelsCache { + entry: None, + load_error: false, + stored_entries: Mutex::new(Vec::new()), + store_gate: Mutex::new(Some(StoreGate { + entered: store_entered, + release, + })), + }); + let mut model = remote_model("gpt-5.4"); + model.default_reasoning_summary = ReasoningSummary::Concise; + let endpoint = TestModelsEndpoint::new(vec![model.clone()]); + let manager = models_manager(cache, endpoint.clone()); + let initialization = tokio::spawn({ + let manager = manager.clone(); + async move { + manager + .ensure_initialized(codex_core::test_support::default_http_client_factory()) + .await; + } + }); + timeout(Duration::from_secs(/*secs*/ 10), entered).await??; + + // A normal refresh must proceed while the initial load is waiting on storage. + model.default_reasoning_summary = ReasoningSummary::Detailed; + *endpoint.models.lock().expect("model response lock") = vec![model.clone()]; + timeout( + Duration::from_secs(/*secs*/ 10), + manager.list_models( + RefreshStrategy::Online, + codex_core::test_support::default_http_client_factory(), + ), + ) + .await + .expect("refresh must not wait for an initialization cache write"); + assert_eq!(manager.get_remote_models().await, vec![model.clone()]); + + release_store.send(()).expect("release initial cache write"); + timeout(Duration::from_secs(/*secs*/ 10), initialization).await??; + assert_eq!(manager.get_remote_models().await, vec![model.clone()]); + assert_eq!( + manager.refreshed_model_catalog().as_deref(), + Some(std::slice::from_ref(&model)) + ); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn account_switch_during_cache_store_preserves_new_catalog_for_next_turn() -> Result<()> { let server = wiremock::MockServer::start().await; @@ -517,12 +568,15 @@ async fn account_switch_during_cache_store_preserves_new_catalog_for_next_turn() }, ) .await; - manager - .list_models( + timeout( + Duration::from_secs(/*secs*/ 10), + manager.list_models( RefreshStrategy::Online, codex_core::test_support::default_http_client_factory(), - ) - .await; + ), + ) + .await + .expect("new account catalog must load while the old account's cache write is paused"); assert_eq!(manager.get_remote_models().await, vec![new_model.clone()]); release_store.send(()).unwrap(); timeout(Duration::from_secs(10), old_refresh).await??; diff --git a/codex-rs/core/tests/suite/items.rs b/codex-rs/core/tests/suite/items.rs index 69ae12fe6a..034f443be4 100644 --- a/codex-rs/core/tests/suite/items.rs +++ b/codex-rs/core/tests/suite/items.rs @@ -309,10 +309,11 @@ async fn web_search_item_is_emitted() -> anyhow::Result<()> { let server = start_mock_server().await; - let TestCodex { codex, .. } = test_codex() + let test = test_codex() .with_history_mode(ThreadHistoryMode::Paginated) .build(&server) .await?; + let codex = test.codex.clone(); let web_search_added = ev_web_search_call_added_partial("web-search-1", "in_progress"); let web_search_done = ev_web_search_call_done("web-search-1", "completed", "weather seattle"); diff --git a/codex-rs/core/tests/suite/models_cache_auth.rs b/codex-rs/core/tests/suite/models_cache_auth.rs index 9c8bd121f8..96e0a83e22 100644 --- a/codex-rs/core/tests/suite/models_cache_auth.rs +++ b/codex-rs/core/tests/suite/models_cache_auth.rs @@ -25,6 +25,8 @@ use codex_protocol::openai_models::ToolMode; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::InterAgentCommunication; use codex_protocol::protocol::Op; +use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::ThreadSettingsOverrides; use codex_protocol::user_input::UserInput; use core_test_support::responses; @@ -385,8 +387,10 @@ async fn auth_and_provider_switches_do_not_reuse_chatgpt_catalog() -> Result<()> Ok(()) } +#[test_case(SessionSource::VSCode; "root")] +#[test_case(SessionSource::SubAgent(SubAgentSource::Review); "child")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn host_model_info_uses_the_thread_catalog() -> Result<()> { +async fn host_model_info_uses_the_thread_catalog(source: SessionSource) -> Result<()> { let slug = "provider-default"; let server = MockServer::start().await; let test = test_codex() @@ -403,7 +407,11 @@ async fn host_model_info_uses_the_thread_catalog() -> Result<()> { .model_catalog .take() .expect("runtime catalog"); - let mut model = catalog.models.into_iter().find(|m| m.slug == slug).unwrap(); + let mut model = catalog + .models + .into_iter() + .find(|m| m.slug == slug) + .expect("runtime catalog contains the selected model"); model.comp_hash = Some("thread-revision".into()); let mut next_model = model.clone(); next_model.slug = "after-update".into(); @@ -422,6 +430,7 @@ async fn host_model_info_uses_the_thread_catalog() -> Result<()> { options.config.model = None; options.config.model_context_window = Some(32_000); options.metrics_service_name = Some("codex_work_cca".into()); + options.session_source = Some(source); let selected = test.thread_manager.resolve_startup_model(&options).await?; assert_eq!(selected.slug, slug); assert_eq!(selected.comp_hash.as_deref(), Some("thread-revision")); diff --git a/codex-rs/core/tests/suite/models_cache_ttl.rs b/codex-rs/core/tests/suite/models_cache_ttl.rs index d37d67b115..8ec562a976 100644 --- a/codex-rs/core/tests/suite/models_cache_ttl.rs +++ b/codex-rs/core/tests/suite/models_cache_ttl.rs @@ -1,3 +1,4 @@ +use codex_core::StartThreadOptions; use codex_core::TurnInputRequest; use core_test_support::test_codex::local_requests; use std::path::Path; @@ -10,6 +11,7 @@ use chrono::Utc; use codex_core::config::Constrained; use codex_login::CodexAuth; use codex_models_manager::client_version_to_whole; +use codex_models_manager::manager::ModelCatalogSettings; use codex_models_manager::manager::RefreshStrategy; use codex_protocol::config_types::ApprovalsReviewer; use codex_protocol::config_types::CollaborationMode; @@ -28,6 +30,8 @@ use codex_protocol::openai_models::TruncationPolicyConfig; use codex_protocol::openai_models::default_input_modalities; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::InternalSessionSource; +use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::ThreadSettingsOverrides; use codex_protocol::user_input::UserInput; use core_test_support::responses; @@ -46,7 +50,11 @@ use pretty_assertions::assert_eq; use serde::Deserialize; use serde::Serialize; use serde_json::json; +use wiremock::Mock; use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::method; +use wiremock::matchers::path_regex; const ETAG: &str = "\"models-etag-ttl\""; const CACHE_FILE: &str = "models_cache.json"; @@ -55,11 +63,19 @@ const VERSIONED_MODEL: &str = "codex-test-versioned"; const MISSING_VERSION_MODEL: &str = "codex-test-missing-version"; const DIFFERENT_VERSION_MODEL: &str = "codex-test-different-version"; +#[test_case::test_case(true; "expired_cache")] +#[test_case::test_case(false; "disabled_cache")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn guardian_reused_reviewer_avoids_stale_catalog_lookup() -> Result<()> { +async fn guardian_reused_reviewer_avoids_stale_catalog_lookup(with_cache: bool) -> Result<()> { skip_if_no_network!(Ok(())); let server = MockServer::start().await; + let websocket_fallback = Mock::given(method("GET")) + .and(path_regex(".*/responses$")) + .respond_with(ResponseTemplate::new(426)) + .expect(1) + .mount_as_scoped(&server) + .await; let bundled = codex_models_manager::bundled_models_response()?; let catalog = ModelsResponse { models: bundled @@ -77,23 +93,30 @@ async fn guardian_reused_reviewer_avoids_stale_catalog_lookup() -> Result<()> { .with_model("gpt-5.5"); builder = builder.with_config(|config| { config.permissions.approval_policy = Constrained::allow_any(AskForApproval::OnRequest); - config.approvals_reviewer = ApprovalsReviewer::AutoReview; + config.approvals_reviewer = ApprovalsReviewer::User; config.model_provider.request_max_retries = Some(0); + config.model_provider.supports_websockets = true; }); + if !with_cache { + builder = builder.with_model_catalog_settings(ModelCatalogSettings::default()); + } let test = builder.build_with_auto_env(&server).await?; - test.thread_manager - .get_models_manager() - .list_models( - RefreshStrategy::OnlineIfUncached, - codex_core::test_support::default_http_client_factory(), - ) - .await; assert_eq!(initial_models.requests().len(), 1); - rewrite_cache_timestamp( - &test.config.codex_home.join(CACHE_FILE), - Utc::now() - chrono::Duration::hours(1), + // Complete the parent's fallback before Guardian can prewarm its own session. + tokio::time::timeout( + std::time::Duration::from_secs(5), + websocket_fallback.wait_until_satisfied(), ) .await?; + if with_cache { + rewrite_cache_timestamp( + &test.config.codex_home.join(CACHE_FILE), + Utc::now() - chrono::Duration::hours(1), + ) + .await?; + } else { + assert!(!test.config.codex_home.join(CACHE_FILE).exists()); + } let repeat_models = responses::mount_models_once_with_etag(&server, catalog, ETAG).await; let deny = r#"{"risk_level":"high","user_authorization":"low","outcome":"deny"}"#; let mut sequence = Vec::new(); @@ -121,7 +144,22 @@ async fn guardian_reused_reviewer_avoids_stale_catalog_lookup() -> Result<()> { ev_completed("parent-done"), ])); let response_mock = mount_sse_sequence(&server, sequence).await; - test.submit_text_turn("Review both shell requests.").await?; + test.codex + .start_or_steer_turn( + TurnInputRequest::user_input(vec![UserInput::Text { + text: "Review both shell requests.".into(), + text_elements: Vec::new(), + }]) + .with_thread_settings(ThreadSettingsOverrides { + approvals_reviewer: Some(ApprovalsReviewer::AutoReview), + ..Default::default() + }), + ) + .await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; let requests = response_mock.requests(); assert_eq!( @@ -151,6 +189,26 @@ async fn guardian_reused_reviewer_avoids_stale_catalog_lookup() -> Result<()> { Some(guardian_thread.as_str()), "the second approval must reuse the Guardian reviewer" ); + + // Parallel reviewers inherit HTTP fallback as well as Guardian's retry overrides. + let fork_reviewer = || async { + let mut config = test.config.clone(); + config.model_provider.request_max_retries = Some(1); + config.model_provider.stream_max_retries = Some(1); + config.model_provider.supports_websockets = false; + test.thread_manager + .fork_internal_session( + test.session_configured.session_id.into(), + StartThreadOptions { + session_source: Some(SessionSource::Internal(InternalSessionSource::Guardian)), + ..StartThreadOptions::new(config) + }, + Vec::new(), + ) + .await + }; + let (first, second) = tokio::try_join!(fork_reviewer(), fork_reviewer())?; + assert_ne!(first.thread_id, second.thread_id); assert_eq!(repeat_models.requests().len(), 0); Ok(()) } diff --git a/codex-rs/core/tests/suite/remote_models.rs b/codex-rs/core/tests/suite/remote_models.rs index 46a523f0ae..adede0808c 100644 --- a/codex-rs/core/tests/suite/remote_models.rs +++ b/codex-rs/core/tests/suite/remote_models.rs @@ -1,11 +1,17 @@ #![cfg(not(target_os = "windows"))] use anyhow::Result; +use codex_core::StartThreadOptions; use codex_core::TurnInputRequest; use codex_features::Feature; use codex_login::CodexAuth; use codex_model_provider_info::ModelProviderInfo; use codex_model_provider_info::built_in_model_providers; use codex_models_manager::bundled_models_response; +use codex_models_manager::cache::ModelsCache; +use codex_models_manager::cache::ModelsCacheEntry; +use codex_models_manager::cache::ModelsCacheError; +use codex_models_manager::cache::ModelsCacheFuture; +use codex_models_manager::manager::ModelCatalogSettings; use codex_models_manager::manager::RefreshStrategy; use codex_models_manager::manager::SharedModelsManager; use codex_models_manager::model_info::BASE_INSTRUCTIONS; @@ -32,6 +38,7 @@ use core_test_support::load_default_config_for_test; use core_test_support::responses::ev_assistant_message; use core_test_support::responses::ev_completed; use core_test_support::responses::ev_function_call; +use core_test_support::responses::ev_function_call_with_namespace; use core_test_support::responses::ev_response_created; use core_test_support::responses::mount_models_once; use core_test_support::responses::mount_models_once_with_delay; @@ -49,6 +56,10 @@ use core_test_support::wait_for_event; use core_test_support::wait_for_event_match; use pretty_assertions::assert_eq; use serde_json::json; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; use tempfile::TempDir; use test_case::test_case; use tokio::time::Duration; @@ -64,6 +75,241 @@ use wiremock::matchers::path; const REMOTE_MODEL_SLUG: &str = "codex-test"; +enum ChildCatalogCache { + Warm, + Cold, + Disabled, +} + +#[test_case(ChildCatalogCache::Warm; "warm_cache")] +#[test_case(ChildCatalogCache::Cold; "cold_cache")] +#[test_case(ChildCatalogCache::Disabled; "without_cache")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn child_loads_remote_tool_search_with_thread_provider_headers( + cache: ChildCatalogCache, +) -> Result<()> { + skip_if_no_network!(Ok(())); + let server = MockServer::start().await; + let cache_home = TempDir::new()?; + let catalog_settings = match cache { + ChildCatalogCache::Warm | ChildCatalogCache::Cold => { + ModelCatalogSettings::disk(cache_home.path().to_path_buf()) + } + ChildCatalogCache::Disabled => ModelCatalogSettings::default(), + }; + let mut model = test_remote_model( + REMOTE_MODEL_SLUG, + ModelVisibility::List, + /*priority*/ 0, + ); + model.supports_search_tool = true; + let _models = Mock::given(method("GET")) + .and(path("/v1/models")) + .respond_with(ResponseTemplate::new(200).set_body_json(ModelsResponse { + models: vec![model], + })) + .expect(match cache { + ChildCatalogCache::Warm => 2..=2, + ChildCatalogCache::Cold | ChildCatalogCache::Disabled => 2..=3, + }) + .mount_as_scoped(&server) + .await; + let mut test = test_codex() + .with_model(REMOTE_MODEL_SLUG) + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_model_catalog_settings(catalog_settings) + .with_config(|config| { + config + .features + .enable(Feature::Collab) + .expect("enable collaboration"); + config + .features + .disable(Feature::MultiAgentV2) + .expect("disable multi-agent v2"); + config.agent_max_depth = 2; + }) + .build_with_auto_env(&server) + .await?; + test.codex.shutdown_and_wait().await?; + test.config + .model_provider + .http_headers + .get_or_insert_default() + .insert("X-OpenAI-Product-Sku".to_string(), "test-product".into()); + let root = test + .thread_manager + .start_thread(StartThreadOptions { + metrics_service_name: Some("codex_work_cca".to_string()), + ..test.start_thread_options().await + }) + .await?; + test.codex = root.thread; + test.session_configured = root.session_configured; + if matches!(cache, ChildCatalogCache::Cold) { + tokio::fs::remove_file(cache_home.path().join("models_cache.codex_work_cca.json")).await?; + } + assert_child_retains_tool_search(&mut test, &server).await +} + +#[derive(Debug, Default)] +struct FailingWritesCache { + entry: Mutex>, + reject_writes: AtomicBool, +} + +impl ModelsCache for FailingWritesCache { + fn load<'a>( + &'a self, + _client_version: &'a str, + ) -> ModelsCacheFuture<'a, Result, ModelsCacheError>> { + Box::pin(async { Ok(self.entry.lock().expect("cache lock").clone()) }) + } + + fn store<'a>( + &'a self, + entry: &'a ModelsCacheEntry, + ) -> ModelsCacheFuture<'a, Result<(), ModelsCacheError>> { + Box::pin(async move { + if self.reject_writes.load(Ordering::SeqCst) { + return Err(ModelsCacheError::new("injected cache write failure")); + } + *self.entry.lock().expect("cache lock") = Some(entry.clone()); + Ok(()) + }) + } + + fn refresh_ttl<'a>( + &'a self, + _client_version: &'a str, + _identity: &'a str, + _etag: &'a str, + ) -> ModelsCacheFuture<'a, Result<(), ModelsCacheError>> { + Box::pin(async { Err(ModelsCacheError::new("unexpected TTL refresh")) }) + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn child_preserves_parent_catalog_after_cache_write_failure() -> Result<()> { + skip_if_no_network!(Ok(())); + let server = MockServer::start().await; + let cache = Arc::new(FailingWritesCache::default()); + let mut model = test_remote_model( + REMOTE_MODEL_SLUG, + ModelVisibility::List, + /*priority*/ 0, + ); + model.supports_search_tool = false; + let initial = mount_models_once( + &server, + ModelsResponse { + models: vec![model.clone()], + }, + ) + .await; + let mut test = test_codex() + .with_model(REMOTE_MODEL_SLUG) + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_model_catalog_settings(ModelCatalogSettings::with_cache(cache.clone())) + .with_config(|config| { + config + .features + .enable(Feature::Collab) + .expect("enable collaboration"); + config + .features + .disable(Feature::MultiAgentV2) + .expect("disable multi-agent v2"); + config.agent_max_depth = 2; + }) + .build_with_auto_env(&server) + .await?; + assert_eq!(initial.requests().len(), 1); + + model.supports_search_tool = true; + let updated = mount_models_once( + &server, + ModelsResponse { + models: vec![model], + }, + ) + .await; + cache.reject_writes.store(/*val*/ true, Ordering::SeqCst); + let manager = test.thread_manager.get_models_manager(); + manager + .list_models(RefreshStrategy::Online, test.config.http_client_factory()) + .await; + let model_config = test.config.to_models_manager_config(); + let before = manager + .get_model_info(REMOTE_MODEL_SLUG, &model_config) + .await; + assert!(before.supports_search_tool); + assert!( + !cache + .entry + .lock() + .expect("cache lock") + .as_ref() + .expect("older cached catalog") + .models[0] + .supports_search_tool + ); + + assert_child_retains_tool_search(&mut test, &server).await?; + assert_eq!( + manager + .get_model_info(REMOTE_MODEL_SLUG, &model_config) + .await, + before + ); + assert_eq!(updated.requests().len(), 1); + Ok(()) +} + +async fn assert_child_retains_tool_search(test: &mut TestCodex, server: &MockServer) -> Result<()> { + let responses = mount_sse_sequence( + server, + vec![ + sse(vec![ + ev_response_created("parent-spawn"), + ev_function_call_with_namespace( + "spawn-child", + "multi_agent_v1", + "spawn_agent", + &json!({"message": "Say done."}).to_string(), + ), + ev_completed("parent-spawn"), + ]), + sse(vec![ev_response_created("done-1"), ev_completed("done-1")]), + sse(vec![ev_response_created("done-2"), ev_completed("done-2")]), + ], + ) + .await; + let mut created = test.thread_manager.subscribe_thread_created(); + test.submit_turn("Spawn a child.").await?; + let child_id = timeout(Duration::from_secs(10), created.recv()).await??; + let child = test.thread_manager.get_thread(child_id).await?; + wait_for_event(&child, |event| matches!(event, EventMsg::TurnComplete(_))).await; + let requests = responses.requests(); + let child_request = requests + .iter() + .find(|request| request.header("thread-id") == Some(child_id.to_string())) + .expect("child's first Responses request"); + for (thread, request) in [("root", &requests[0]), ("child", child_request)] { + assert_eq!(request.body_json()["model"], REMOTE_MODEL_SLUG); + assert!( + request.body_json()["tools"] + .as_array() + .expect("request tools array") + .iter() + .any(|tool| { tool["type"] == "tool_search" }), + "{thread} must retain deferred tool search" + ); + } + child.shutdown_and_wait().await?; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn disabled_update_plan_preserves_custom_catalog_instructions() -> Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/core/tests/suite/wake_reservation.rs b/codex-rs/core/tests/suite/wake_reservation.rs index 9bfcbfa98a..f24b30c550 100644 --- a/codex-rs/core/tests/suite/wake_reservation.rs +++ b/codex-rs/core/tests/suite/wake_reservation.rs @@ -49,6 +49,10 @@ struct GatedModelsManager { } impl ModelsManager for GatedModelsManager { + fn ensure_initialized(&self, factory: HttpClientFactory) -> ModelsManagerFuture<'_, ()> { + self.inner.ensure_initialized(factory) + } + fn model_catalog_with_origins( &self, strategy: RefreshStrategy, diff --git a/codex-rs/models-manager/src/cache_identity_tests.rs b/codex-rs/models-manager/src/cache_identity_tests.rs index b4fbe252c0..aeb28506eb 100644 --- a/codex-rs/models-manager/src/cache_identity_tests.rs +++ b/codex-rs/models-manager/src/cache_identity_tests.rs @@ -23,12 +23,12 @@ async fn mismatched_and_legacy_cache_entries_fetch_the_current_catalog() { CodexAuth::create_dummy_chatgpt_auth_for_testing(), )), ); + manager + .ensure_initialized(DEFAULT_HTTP_CLIENT_FACTORY) + .await; assert_eq!( manager - .raw_model_catalog( - RefreshStrategy::OnlineIfUncached, - DEFAULT_HTTP_CLIENT_FACTORY - ) + .raw_model_catalog(RefreshStrategy::InMemory, DEFAULT_HTTP_CLIENT_FACTORY) .await .models, vec![current.clone()] diff --git a/codex-rs/models-manager/src/catalog.rs b/codex-rs/models-manager/src/catalog.rs index 831ed78df6..6137875a63 100644 --- a/codex-rs/models-manager/src/catalog.rs +++ b/codex-rs/models-manager/src/catalog.rs @@ -29,6 +29,8 @@ impl<'a> SelectedCatalog<'a> { pub(crate) struct CatalogState { pub entry: ModelsCacheEntry, pub provider_model_ids: HashSet, + // Distinguishes a successfully loaded empty catalog from the initial bundled seed. + pub initialized: bool, } impl CatalogState { @@ -36,6 +38,7 @@ impl CatalogState { Self { entry, provider_model_ids: HashSet::new(), + initialized: false, } } @@ -49,6 +52,7 @@ impl CatalogState { Self { entry, provider_model_ids, + initialized: true, } } diff --git a/codex-rs/models-manager/src/manager.rs b/codex-rs/models-manager/src/manager.rs index 3837f1b206..a8aafa125d 100644 --- a/codex-rs/models-manager/src/manager.rs +++ b/codex-rs/models-manager/src/manager.rs @@ -129,6 +129,13 @@ impl fmt::Display for RefreshStrategy { } } +/// Initialization preserves a loaded snapshot; explicit refreshes follow their requested policy. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum CatalogLoadMode { + Initialize, + Refresh(RefreshStrategy), +} + type SharedModelsEndpointClient = Arc; /// A successful in-process provider response, before bundled fallback is merged in. @@ -140,6 +147,14 @@ pub struct RefreshedModelCatalog { /// Coordinates model discovery plus cached metadata on disk. pub trait ModelsManager: fmt::Debug + Send + Sync { + /// Initialize a newly selected catalog without replacing an already-loaded snapshot. + /// Stateful managers must check readiness for their current provider and credentials; + /// a successful empty catalog is initialized too. Static catalogs need no work. + fn ensure_initialized( + &self, + http_client_factory: HttpClientFactory, + ) -> ModelsManagerFuture<'_, ()>; + /// Observe successful provider refreshes and invalidations, never cache-only reads. fn model_catalog_updates(&self) -> Option>>> { None @@ -490,6 +505,25 @@ impl StaticModelsManager { } impl ModelsManager for OpenAiModelsManager { + fn ensure_initialized( + &self, + http_client_factory: HttpClientFactory, + ) -> ModelsManagerFuture<'_, ()> { + Box::pin(async move { + // Resolve lazy credentials before checking the selected catalog's identity. + self.should_refresh_models().await; + if self.catalog_is_initialized(&*self.remote_models.read().await) { + return; + } + if let Err(err) = self + .load_catalog(CatalogLoadMode::Initialize, &http_client_factory) + .await + { + error!("failed to initialize model catalog: {err}"); + } + }) + } + fn model_catalog_updates(&self) -> Option>>> { Some(self.catalog_updates.subscribe()) } @@ -672,6 +706,12 @@ impl ModelsManager for OpenAiModelsManager { } impl OpenAiModelsManager { + fn catalog_is_initialized(&self, catalog: &CatalogState) -> bool { + catalog.initialized + && catalog.entry.identity.is_some() + && catalog.entry.identity == self.endpoint_client.identity() + } + /// Selects the catalog available under the current provider and auth identity. fn models_for_entry<'a>(&self, entry: &'a ModelsCacheEntry) -> SelectedCatalog<'a> { if !self.api_key_discovery_disabled() @@ -719,6 +759,18 @@ impl OpenAiModelsManager { &self, refresh_strategy: RefreshStrategy, http_client_factory: &HttpClientFactory, + ) -> CoreResult<()> { + self.load_catalog( + CatalogLoadMode::Refresh(refresh_strategy), + http_client_factory, + ) + .await + } + + async fn load_catalog( + &self, + mode: CatalogLoadMode, + http_client_factory: &HttpClientFactory, ) -> CoreResult<()> { let _refresh_guard = self.catalog_source.refresh_lock().await; // API-key discovery must be enabled and supported before reusing a remote catalog. @@ -727,34 +779,49 @@ impl OpenAiModelsManager { if self.api_key_discovery_disabled() { return Ok(()); } - if !self.should_refresh_models().await { + let should_refresh = self.should_refresh_models().await; + if mode == CatalogLoadMode::Initialize + && self.catalog_is_initialized(&*self.remote_models.read().await) + { + return Ok(()); + } + if !should_refresh { if matches!( - refresh_strategy, - RefreshStrategy::Offline | RefreshStrategy::OnlineIfUncached + mode, + CatalogLoadMode::Initialize + | CatalogLoadMode::Refresh( + RefreshStrategy::Offline | RefreshStrategy::OnlineIfUncached + ) ) { - self.try_load_cache().await; + self.try_load_cache(mode).await; } return Ok(()); } - match refresh_strategy { - RefreshStrategy::InMemory => Ok(()), - RefreshStrategy::Offline => { - self.try_load_cache().await; + match mode { + CatalogLoadMode::Refresh(RefreshStrategy::InMemory) => Ok(()), + CatalogLoadMode::Refresh(RefreshStrategy::Offline) => { + self.try_load_cache(mode).await; Ok(()) } - RefreshStrategy::OnlineIfUncached => { - if self.try_load_cache().await { + CatalogLoadMode::Initialize + | CatalogLoadMode::Refresh(RefreshStrategy::OnlineIfUncached) => { + if self.try_load_cache(mode).await { return Ok(()); } - self.fetch_and_update_models(http_client_factory).await + self.fetch_and_update_models(http_client_factory, mode) + .await + } + CatalogLoadMode::Refresh(RefreshStrategy::Online) => { + self.fetch_and_update_models(http_client_factory, mode) + .await } - RefreshStrategy::Online => self.fetch_and_update_models(http_client_factory).await, } } async fn fetch_and_update_models( &self, http_client_factory: &HttpClientFactory, + mode: CatalogLoadMode, ) -> CoreResult<()> { let client_version = crate::client_version_to_whole(); let request_identity = self.endpoint_client.identity(); @@ -787,6 +854,7 @@ impl OpenAiModelsManager { let CatalogState { entry: current, provider_model_ids, + initialized, } = &mut *current; if request_identity != self.endpoint_client.identity() { return Err(error); @@ -797,6 +865,7 @@ impl OpenAiModelsManager { provider_model_ids.clear(); current.etag = None; current.identity.clone_from(&request_identity); + *initialized = true; self.catalog_updates.send_replace(None); } if let Some(identity) = request_identity @@ -823,7 +892,10 @@ impl OpenAiModelsManager { return Err(error); } }; - if Some(&identity) != self.endpoint_client.identity().as_ref() { + if Some(&identity) != self.endpoint_client.identity().as_ref() + || (mode == CatalogLoadMode::Initialize + && self.catalog_is_initialized(&*self.remote_models.read().await)) + { return Ok(()); } let refreshed = Arc::new(RefreshedModelCatalog { @@ -843,7 +915,7 @@ impl OpenAiModelsManager { error!("failed to write models cache: {err}"); } let mut current = self.remote_models.write().await; - if self.apply_remote_models(&mut current, entry) { + if self.apply_remote_models(&mut current, entry, mode) { self.catalog_updates.send_replace(Some(refreshed)); } Ok(()) @@ -878,8 +950,17 @@ impl OpenAiModelsManager { } /// Publish only while the request identity still matches, including after async storage. - fn apply_remote_models(&self, current: &mut CatalogState, entry: ModelsCacheEntry) -> bool { - if entry.identity != self.endpoint_client.identity() { + fn apply_remote_models( + &self, + current: &mut CatalogState, + entry: ModelsCacheEntry, + mode: CatalogLoadMode, + ) -> bool { + // A refresh may finish while initialization awaits cache or network I/O. + // Check under the write lock so initialization cannot replace its loaded snapshot. + if entry.identity != self.endpoint_client.identity() + || (mode == CatalogLoadMode::Initialize && self.catalog_is_initialized(current)) + { return false; } let mut catalog = CatalogState::provider(entry); @@ -912,7 +993,7 @@ impl OpenAiModelsManager { } /// Attempt to satisfy the refresh from the cache when it matches the provider and TTL. - async fn try_load_cache(&self) -> bool { + async fn try_load_cache(&self, mode: CatalogLoadMode) -> bool { let Some(cache) = self.settings.cache.as_ref() else { return false; }; @@ -961,11 +1042,19 @@ impl OpenAiModelsManager { return false; } let mut current = self.remote_models.write().await; - self.apply_remote_models(&mut current, cache_entry) + (mode == CatalogLoadMode::Initialize && self.catalog_is_initialized(¤t)) + || self.apply_remote_models(&mut current, cache_entry, mode) } } impl ModelsManager for StaticModelsManager { + fn ensure_initialized( + &self, + _http_client_factory: HttpClientFactory, + ) -> ModelsManagerFuture<'_, ()> { + Box::pin(std::future::ready(())) + } + fn model_catalog_with_origins( &self, _refresh_strategy: RefreshStrategy, diff --git a/codex-rs/models-manager/src/manager_tests.rs b/codex-rs/models-manager/src/manager_tests.rs index 72732f4a3d..fd24d1b2d9 100644 --- a/codex-rs/models-manager/src/manager_tests.rs +++ b/codex-rs/models-manager/src/manager_tests.rs @@ -566,6 +566,9 @@ async fn injected_cache_write_error_does_not_fail_remote_refresh() { .await; assert_eq!(catalog.models, remote_models); + manager + .ensure_initialized(DEFAULT_HTTP_CLIENT_FACTORY) + .await; assert_eq!( manager .model_catalog_with_origins(RefreshStrategy::InMemory, DEFAULT_HTTP_CLIENT_FACTORY) @@ -581,6 +584,27 @@ async fn injected_cache_write_error_does_not_fail_remote_refresh() { assert_eq!(endpoint.fetch_count(), 1); } +#[tokio::test] +async fn initialization_retains_loaded_catalogs_including_empty_responses() { + for models in [ + Vec::new(), + vec![remote_model("remote", "Remote", /*priority*/ 0)], + ] { + let endpoint = TestModelsEndpoint::new(vec![models.clone()]); + let manager = + OpenAiModelsManager::new_without_cache(endpoint.clone(), /*auth_manager*/ None) + .with_provider_catalog(); + manager + .ensure_initialized(DEFAULT_HTTP_CLIENT_FACTORY) + .await; + manager + .ensure_initialized(DEFAULT_HTTP_CLIENT_FACTORY) + .await; + assert_eq!(manager.get_remote_models().await, models); + assert_eq!(endpoint.fetch_count(), 1); + } +} + #[tokio::test] async fn injected_cache_ttl_refresh_preserves_cached_payload() { let cached_models = vec![remote_model("cached", "Cached", /*priority*/ 0)];