diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 6ff6fbc902..4e0b6e98a4 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3031,6 +3031,7 @@ dependencies = [ "codex-utils-process", "codex-utils-pty", "codex-utils-redacted-string", + "codex-utils-rustls-provider", "codex-utils-stream-parser", "codex-utils-string", "codex-web-search-extension", @@ -3057,6 +3058,7 @@ dependencies = [ "predicates", "pretty_assertions", "rand 0.9.3", + "rcgen", "regex-lite", "rmcp", "serde", @@ -3070,6 +3072,7 @@ dependencies = [ "test-log", "thiserror 2.0.18", "tokio", + "tokio-rustls", "tokio-tungstenite", "tokio-util", "toml 0.9.11+spec-1.1.0", @@ -3106,6 +3109,7 @@ dependencies = [ "codex-home", "codex-image-generation-extension", "codex-login", + "codex-model-provider", "codex-model-provider-info", "codex-models-manager", "codex-protocol", diff --git a/codex-rs/agent-identity/src/lib.rs b/codex-rs/agent-identity/src/lib.rs index 14a9e35190..8323e1d9e7 100644 --- a/codex-rs/agent-identity/src/lib.rs +++ b/codex-rs/agent-identity/src/lib.rs @@ -54,22 +54,23 @@ pub enum ChatGptEnvironment { impl ChatGptEnvironment { pub fn from_chatgpt_base_url(chatgpt_base_url: &str) -> Result { - match chatgpt_base_url.trim_end_matches('/') { - "https://chatgpt.com" - | "https://chatgpt.com/backend-api" - | "https://chatgpt.com/codex" - | "https://chatgpt.com/backend-api/codex" - | "https://chat.openai.com" - | "https://chat.openai.com/backend-api" - | "https://chat.openai.com/codex" - | "https://chat.openai.com/backend-api/codex" => Ok(Self::Production), - "https://chatgpt-staging.com" - | "https://chatgpt-staging.com/backend-api" - | "https://chatgpt-staging.com/codex" - | "https://chatgpt-staging.com/backend-api/codex" => Ok(Self::Staging), - _ => anyhow::bail!( - "Agent Identity only supports production and staging ChatGPT environments" - ), + let url = codex_http_client::validate_endpoint_url(chatgpt_base_url, "chatgpt_base_url") + .map_err(anyhow::Error::msg)?; + let host = url.host_str().unwrap_or_default(); + anyhow::ensure!( + codex_http_client::is_allowed_chatgpt_host(host) + && url.port_or_known_default() == Some(443) + && url.query().is_none() + && matches!( + url.path().trim_end_matches('/'), + "" | "/backend-api" | "/codex" | "/backend-api/codex" + ), + "Agent Identity only supports production and staging ChatGPT environments" + ); + if host == "chatgpt-staging.com" || host.ends_with(".chatgpt-staging.com") { + Ok(Self::Staging) + } else { + Ok(Self::Production) } } @@ -893,12 +894,34 @@ J1bwkqKZTB5dHolX9A58e/xXnfZ5P8f3Z83+Izap3FwqQulk7b1WO1MQcHuVg2NN ChatGptEnvironment::Staging.agent_identity_authapi_base_url(), "https://auth.api.openai.org/api/accounts" ); + assert_eq!( + ChatGptEnvironment::from_chatgpt_base_url("https://us.chatgpt.com/backend-api")?, + ChatGptEnvironment::Production + ); + assert_eq!( + ChatGptEnvironment::from_chatgpt_base_url( + "https://us.chatgpt-staging.com/backend-api" + )?, + ChatGptEnvironment::Staging + ); Ok(()) } #[test] fn chatgpt_environment_rejects_custom_urls() { - assert!(ChatGptEnvironment::from_chatgpt_base_url("http://localhost:8080").is_err(),); + for base in [ + "http://localhost:8080", + "http://us.chatgpt.com", + "https://chatgpt.com.evil.example", + "https://us.chatgpt.com:8443", + "https://u@us.chatgpt.com", + "https://chatgpt.com/custom", + ] { + assert!( + ChatGptEnvironment::from_chatgpt_base_url(base).is_err(), + "{base}" + ); + } } #[test] diff --git a/codex-rs/app-server/src/in_process.rs b/codex-rs/app-server/src/in_process.rs index 8073046e59..4691c05134 100644 --- a/codex-rs/app-server/src/in_process.rs +++ b/codex-rs/app-server/src/in_process.rs @@ -393,6 +393,7 @@ pub async fn start(mut args: InProcessStartArgs) -> IoResult, + thread_manager: Arc, startup_cwd: Option, state: Mutex, } @@ -35,12 +35,12 @@ impl ModelCatalog { config_manager: ConfigManager, startup_cwd: Option, config: Arc, - models_manager: SharedModelsManager, - auth_manager: Arc, + thread_manager: Arc, ) -> Self { + let models_manager = thread_manager.get_models_manager(); Self { config_manager, - auth_manager, + thread_manager, startup_cwd, state: Mutex::new(CatalogState { needs_reload: false, @@ -72,14 +72,6 @@ impl ModelCatalog { } } - pub(crate) async fn set_api_key_model_discovery_enabled(&self, enabled: bool) { - self.state - .lock() - .await - .models_manager - .set_api_key_model_discovery_enabled(enabled); - } - #[expect( clippy::await_holding_invalid_type, reason = "catalog reloads are serialized with account invalidation; model fetching holds no lock" @@ -100,11 +92,16 @@ impl ModelCatalog { if config.model_provider_id != state.config.model_provider_id || config.model_provider != state.config.model_provider || config.model_catalog != state.config.model_catalog + || config.features.enabled(Feature::RuntimeAgentInstances) + != state + .config + .features + .enabled(Feature::RuntimeAgentInstances) + || config.runtime_auth_mode != state.config.runtime_auth_mode { - state.models_manager = - codex_core::build_models_manager(&config, Arc::clone(&self.auth_manager)); + state.models_manager = self.thread_manager.models_manager_for_config(&config); } - state.models_manager.set_api_key_model_discovery_enabled( + self.thread_manager.set_api_key_model_discovery_enabled( config.features.enabled(Feature::ApiKeyModelDiscovery), ); state.config = Arc::new(config); diff --git a/codex-rs/app-server/src/models_refresh_worker_tests.rs b/codex-rs/app-server/src/models_refresh_worker_tests.rs index c0fd170e6f..7355b343db 100644 --- a/codex-rs/app-server/src/models_refresh_worker_tests.rs +++ b/codex-rs/app-server/src/models_refresh_worker_tests.rs @@ -11,6 +11,7 @@ use codex_core::config::ConfigBuilder; use codex_core::config::ConfigOverrides; use codex_feedback::CodexFeedback; use codex_http_client::HttpClientFactory; +use codex_models_manager::manager::ModelCatalogSettings; use codex_models_manager::manager::ModelsEndpointClient; use codex_models_manager::manager::ModelsEndpointFuture; use codex_models_manager::manager::ModelsEndpointResponse; @@ -28,6 +29,38 @@ use tracing_subscriber::layer::SubscriberExt; use super::*; +fn test_thread_manager( + config: &codex_core::config::Config, + models_manager: Option, + auth_manager: Arc, +) -> Arc { + let factory = match models_manager { + Some(models) => codex_core::test_support::fixed_model_catalog(models), + None => Arc::new(ModelCatalogSettings::disk(config.codex_home.to_path_buf())), + }; + let startup = codex_core::build_startup_models( + config, + auth_manager, + uuid::Uuid::new_v4().to_string(), + factory, + ); + Arc::new(codex_core::ThreadManager::new( + config, + startup, + codex_core::CodexAppsToolsCache::default(), + codex_protocol::protocol::SessionSource::Exec, + Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + codex_extension_api::empty_extension_registry(), + Arc::new(codex_core::test_support::EmptyUserInstructionsProvider), + /*analytics_events_client*/ None, + codex_core::passthrough_image_store(), + codex_core::ThreadStoreBuilder::from_config(config, /*state_db*/ None).build(), + /*agent_graph_store*/ None, + /*attestation_provider*/ None, + /*external_time_provider*/ None, + )) +} + #[derive(Debug)] struct TestModelsEndpoint { fetch_count: AtomicUsize, @@ -148,9 +181,8 @@ async fn refreshes_immediately_periodically_and_stops_when_dropped() { let catalog = Arc::new(ModelCatalog::new( config_manager, Some(config.cwd.to_path_buf()), - config, - Arc::clone(&models_manager), - auth_manager, + Arc::clone(&config), + test_thread_manager(&config, Some(Arc::clone(&models_manager)), auth_manager), )); let worker = spawn_with_interval(&catalog, Duration::from_millis(/*millis*/ 10)); @@ -232,9 +264,8 @@ async fn background_refresh_blocks_invalid_requirements_then_recovers() -> anyho let catalog = Arc::new(ModelCatalog::new( config_manager, Some(config.cwd.to_path_buf()), - config, - models_manager, - auth_manager, + Arc::clone(&config), + test_thread_manager(&config, Some(models_manager), auth_manager), )); let feedback = CodexFeedback::new(); // This test uses Tokio's current-thread runtime, so spawned refreshes share the subscriber. @@ -319,9 +350,11 @@ async fn catalog_reload_recovers_after_cloud_failure() -> anyhow::Result<()> { /*enable_codex_api_key_env*/ false, ) .await?; - let models = codex_core::build_models_manager(&config, Arc::clone(&auth)); let catalog = Arc::new(ModelCatalog::new( - manager, /*startup_cwd*/ None, config, models, auth, + manager, + /*startup_cwd*/ None, + Arc::clone(&config), + test_thread_manager(&config, /*models_manager*/ None, auth), )); failed.store(true, Ordering::SeqCst); catalog.invalidate().await; @@ -370,8 +403,12 @@ async fn catalog_reload_recovers_original_project_after_startup_fallback() -> an /*enable_codex_api_key_env*/ false, ) .await?; - let models = codex_core::build_models_manager(&config, Arc::clone(&auth)); - let catalog = ModelCatalog::new(manager, Some(project_path.clone()), config, models, auth); + let catalog = ModelCatalog::new( + manager, + Some(project_path.clone()), + Arc::clone(&config), + test_thread_manager(&config, /*models_manager*/ None, auth), + ); let mut model = codex_models_manager::bundled_models_response()? .models .remove(0); diff --git a/codex-rs/app-server/src/request_processors/account_processor/workspace_routing_tests.rs b/codex-rs/app-server/src/request_processors/account_processor/workspace_routing_tests.rs index dafe4e8e98..187084e21d 100644 --- a/codex-rs/app-server/src/request_processors/account_processor/workspace_routing_tests.rs +++ b/codex-rs/app-server/src/request_processors/account_processor/workspace_routing_tests.rs @@ -102,8 +102,7 @@ async fn background_routing_skips_config_until_chatgpt_sign_in( manager.clone(), Some(config.cwd.to_path_buf()), Arc::clone(&config), - threads.get_models_manager(), - Arc::clone(&auth_manager), + Arc::clone(&threads), )); let processor = AccountRequestProcessor::new( Arc::clone(&auth_manager), diff --git a/codex-rs/app-server/src/request_processors/config_processor.rs b/codex-rs/app-server/src/request_processors/config_processor.rs index 1ef91ec491..893598efba 100644 --- a/codex-rs/app-server/src/request_processors/config_processor.rs +++ b/codex-rs/app-server/src/request_processors/config_processor.rs @@ -357,16 +357,9 @@ impl ConfigRequestProcessor { .map_err(|_| internal_error("failed to update feature enablement"))?; let config = self.load_latest_config(/*fallback_cwd*/ None).await?; - self.thread_manager - .get_models_manager() - .set_api_key_model_discovery_enabled( - config.features.enabled(Feature::ApiKeyModelDiscovery), - ); - self.model_catalog - .set_api_key_model_discovery_enabled( - config.features.enabled(Feature::ApiKeyModelDiscovery), - ) - .await; + self.thread_manager.set_api_key_model_discovery_enabled( + config.features.enabled(Feature::ApiKeyModelDiscovery), + ); if should_start_background_rollout_migration && config.features.enabled(feature) { self.thread_manager.start_background_rollout_migration(); } diff --git a/codex-rs/cli/src/doctor.rs b/codex-rs/cli/src/doctor.rs index 03c4781946..d939632a6e 100644 --- a/codex-rs/cli/src/doctor.rs +++ b/codex-rs/cli/src/doctor.rs @@ -41,6 +41,7 @@ use codex_core::config::ConfigOverrides; use codex_core::config::LoaderOverrides; use codex_core::config::find_codex_home; use codex_features::FEATURES; +use codex_features::Feature; use codex_http_client::ClientRouteClass; use codex_http_client::HttpClientFactory; use codex_http_client::OutboundProxyPolicy; @@ -2343,14 +2344,38 @@ async fn websocket_reachability_check( provider.websocket_connect_timeout().as_millis() )); - let runtime_provider = create_model_provider(provider.clone(), auth_manager); - let auth = runtime_provider.auth().await; - details.push(format!( - "auth mode: {}", - auth.as_ref().map(auth_mode_name).unwrap_or("none") - )); - - let (api_provider, api_auth) = match runtime_provider.api_provider_and_auth().await { + let setup = async { + let runtime_auth = if config.features.enabled(Feature::RuntimeAgentInstances) + && let Some(credentials) = &auth_manager + { + let installation_id = codex_core::resolve_installation_id(&config.codex_home).await?; + codex_core::build_runtime_auth_manager( + config, + Arc::clone(credentials), + &installation_id, + ) + .map(|runtime| (runtime, config.runtime_auth_mode)) + } else { + None + }; + let runtime_provider = match auth_manager { + Some(credentials) => codex_model_provider::ProviderAuthContext { + credentials, + runtime_auth, + } + .model_provider(provider.clone()), + None => create_model_provider(provider.clone(), /*auth_manager*/ None), + }; + codex_model_provider::resolve_request( + runtime_provider.as_ref(), + /*scope*/ None, + Some(&config.workspace_routing_context()), + config.http_client_factory(), + ) + .await + } + .await; + let setup = match setup { Ok(setup) => setup, Err(err) => { return websocket_probe_warning( @@ -2360,6 +2385,11 @@ async fn websocket_reachability_check( ); } }; + details.push(format!( + "auth mode: {}", + setup.auth.as_ref().map(auth_mode_name).unwrap_or("none") + )); + let api_provider = setup.api_provider; match api_provider.websocket_url_for_path("responses") { Ok(url) => { details.push(format!("endpoint: {url}")); @@ -2383,8 +2413,8 @@ async fn websocket_reachability_check( OPENAI_BETA_HEADER, HeaderValue::from_static(RESPONSES_WEBSOCKETS_V2_BETA_HEADER_VALUE), ); - let client = ResponsesWebsocketClient::new(api_provider, api_auth); - let http_client_factory = config.http_client_factory(); + let client = ResponsesWebsocketClient::new(api_provider, setup.api_auth); + let http_client_factory = setup.http_client_factory; match tokio::time::timeout( provider.websocket_connect_timeout(), client.probe_handshake( diff --git a/codex-rs/cli/src/main.rs b/codex-rs/cli/src/main.rs index 1c3fe17c1c..fff4dd4656 100644 --- a/codex-rs/cli/src/main.rs +++ b/codex-rs/cli/src/main.rs @@ -2189,7 +2189,20 @@ async fn run_debug_models_command( .await?; let auth_manager = AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ true).await?; - let models_manager = build_models_manager(&config, auth_manager); + let runtime_auth = if config + .features + .enabled(codex_features::Feature::RuntimeAgentInstances) + { + let installation_id = codex_core::resolve_installation_id(&config.codex_home).await?; + codex_core::build_runtime_auth_manager( + &config, + Arc::clone(&auth_manager), + &installation_id, + ) + } else { + None + }; + let models_manager = build_models_manager(&config, auth_manager, runtime_auth); models_manager .raw_model_catalog( RefreshStrategy::OnlineIfUncached, diff --git a/codex-rs/core-api/Cargo.toml b/codex-rs/core-api/Cargo.toml index b82d5ee480..ed1fbff4da 100644 --- a/codex-rs/core-api/Cargo.toml +++ b/codex-rs/core-api/Cargo.toml @@ -26,6 +26,7 @@ codex-image-generation-extension = { workspace = true } codex-exec-server = { workspace = true } codex-features = { workspace = true } codex-login = { workspace = true } +codex-model-provider = { workspace = true } codex-model-provider-info = { workspace = true } codex-models-manager = { workspace = true } codex-protocol = { workspace = true } diff --git a/codex-rs/core-api/src/lib.rs b/codex-rs/core-api/src/lib.rs index 4255f0a3e5..6a10fec243 100644 --- a/codex-rs/core-api/src/lib.rs +++ b/codex-rs/core-api/src/lib.rs @@ -53,6 +53,7 @@ pub use codex_core::PromptCacheKey; pub use codex_core::RecoverTurnRequest; pub use codex_core::StartIfIdleSubmission; pub use codex_core::StartThreadOptions; +pub use codex_core::StartupModels; pub use codex_core::StateDbHandle; pub use codex_core::SteerSubmission; pub use codex_core::SuspendTurnOutcome; @@ -66,6 +67,8 @@ pub use codex_core::TurnStartOptions; pub use codex_core::WaitForEnvironmentToolConfig; pub use codex_core::WithTurnExtensionData; pub use codex_core::build_models_manager; +pub use codex_core::build_runtime_auth_manager; +pub use codex_core::build_startup_models; pub use codex_core::config::Config; pub use codex_core::config::Constrained; pub use codex_core::config::ExtraConfig; @@ -117,8 +120,11 @@ pub use codex_login::ExternalAuthFuture; pub use codex_login::ExternalAuthRefreshContext; pub use codex_login::ExternalAuthRefreshReason; pub use codex_login::default_client::set_default_originator; +pub use codex_model_provider::ModelCatalogFactory; +pub use codex_model_provider::SharedModelProvider; pub use codex_model_provider_info::OPENAI_PROVIDER_ID; pub use codex_model_provider_info::built_in_model_providers; +pub use codex_models_manager::manager::ModelCatalogSettings; pub use codex_models_manager::manager::RefreshStrategy; pub use codex_models_manager::manager::SharedModelsManager; pub use codex_protocol::ThreadId; @@ -138,6 +144,7 @@ pub use codex_protocol::error::Result as CodexResult; pub use codex_protocol::models::PermissionProfile; pub use codex_protocol::models::PermissionProfileSnapshot; pub use codex_protocol::openai_models::ModelPreset; +pub use codex_protocol::openai_models::ModelsResponse; pub use codex_protocol::protocol::AskForApproval; pub use codex_protocol::protocol::EnvironmentConfigState; pub use codex_protocol::protocol::EventMsg; diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index 30d92564c2..e822a6d790 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -157,6 +157,7 @@ codex-protocol-test-support = { workspace = true } codex-test-binary-support = { workspace = true } codex-utils-cargo-bin = { workspace = true } codex-utils-redacted-string = { workspace = true } +codex-utils-rustls-provider = { workspace = true } codex-web-search-extension = { workspace = true } codex-windows-sandbox-test-support = { workspace = true } core_test_support = { workspace = true } @@ -166,6 +167,7 @@ maplit = { workspace = true } opentelemetry = { workspace = true } predicates = { workspace = true } pretty_assertions = { workspace = true } +rcgen = { workspace = true } test-case = "3.3.1" opentelemetry_sdk = { workspace = true, features = [ "experimental_metrics_custom_reader", @@ -174,6 +176,7 @@ opentelemetry_sdk = { workspace = true, features = [ serial_test = { workspace = true } tempfile = { workspace = true } test-log = { workspace = true } +tokio-rustls = { workspace = true } tracing-opentelemetry = { workspace = true } tracing-subscriber = { workspace = true } tracing-test = { workspace = true, features = ["no-env-filter"] } diff --git a/codex-rs/core/src/agent/control/spawn_publication_tests.rs b/codex-rs/core/src/agent/control/spawn_publication_tests.rs index ce679c37df..785c6392d7 100644 --- a/codex-rs/core/src/agent/control/spawn_publication_tests.rs +++ b/codex-rs/core/src/agent/control/spawn_publication_tests.rs @@ -113,8 +113,12 @@ async fn publication_harness( }); harness.manager = ThreadManager::new( &harness.config, - harness.manager.auth_manager(), - harness.manager.get_models_manager(), + crate::build_startup_models( + &harness.config, + harness.manager.auth_manager(), + "spawn-publication-cleanup-test".to_string(), + crate::test_support::fixed_model_catalog(harness.manager.get_models_manager()), + ), crate::CodexAppsToolsCache::default(), SessionSource::Exec, harness.manager.environment_manager(), @@ -124,7 +128,6 @@ async fn publication_harness( harness.manager.image_store(), crate::ThreadStoreBuilder::from_config(&harness.config, harness.state_db.clone()).build(), Some(graph.clone()), - "spawn-publication-cleanup-test".to_string(), /*attestation_provider*/ None, /*external_time_provider*/ None, ) diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index ce05ae9e52..93a0771986 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -49,6 +49,7 @@ use codex_history::ResumedHistory; use codex_history::RolloutItem; use codex_login::AuthManager; use codex_login::CodexAuth; +use codex_models_manager::manager::ModelCatalogSettings; use codex_models_manager::model_info::model_info_from_slug; use codex_protocol::AgentPath; use codex_protocol::ResponseItemId; @@ -3785,8 +3786,14 @@ async fn spawn_agent_full_fork_legacy_compaction_rebuilds_child_instructions_onc let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("dummy")); harness.manager = ThreadManager::new( &parent_config, - auth_manager.clone(), - crate::thread_manager::build_models_manager(&parent_config, auth_manager), + crate::build_startup_models( + &parent_config, + auth_manager.clone(), + uuid::Uuid::new_v4().to_string(), + Arc::new(ModelCatalogSettings::disk( + parent_config.codex_home.to_path_buf(), + )), + ), crate::CodexAppsToolsCache::default(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), @@ -3797,7 +3804,6 @@ async fn spawn_agent_full_fork_legacy_compaction_rebuilds_child_instructions_onc crate::ThreadStoreBuilder::from_config(&parent_config, harness.state_db.clone()) .build(), crate::thread_manager::local_agent_graph_store_from_state_db(harness.state_db.as_ref()), - uuid::Uuid::new_v4().to_string(), /*attestation_provider*/ None, /*external_time_provider*/ None, ); @@ -4941,8 +4947,12 @@ async fn resume_thread_subagent_restores_stored_metadata() { let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("dummy")); let manager = ThreadManager::new( &config, - auth_manager.clone(), - crate::thread_manager::build_models_manager(&config, auth_manager), + crate::build_startup_models( + &config, + auth_manager.clone(), + uuid::Uuid::new_v4().to_string(), + Arc::new(ModelCatalogSettings::disk(config.codex_home.to_path_buf())), + ), crate::CodexAppsToolsCache::default(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), @@ -4952,7 +4962,6 @@ async fn resume_thread_subagent_restores_stored_metadata() { crate::thread_manager::passthrough_image_store(), thread_store.clone(), /*agent_graph_store*/ None, - uuid::Uuid::new_v4().to_string(), /*attestation_provider*/ None, /*external_time_provider*/ None, ); diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index c1f1027388..e86bb5297f 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -54,7 +54,6 @@ use codex_api::RealtimeSessionConfig as ApiRealtimeSessionConfig; use codex_api::Reasoning; use codex_api::ReasoningContext; use codex_api::RequestTelemetry; -use codex_api::ReqwestTransport; use codex_api::ResponseCreateWsRequest; use codex_api::ResponsesApiRequest; use codex_api::ResponsesClient as ApiResponsesClient; @@ -71,22 +70,17 @@ use codex_api::auth_header_telemetry; use codex_api::build_session_headers; use codex_api::create_text_param_for_request; use codex_api::response_create_client_metadata; -use codex_features::RuntimeAuthMode; -use codex_http_client::ClientRouteClass; use codex_http_client::HttpClientFactory; use codex_login::AuthManager; use codex_login::CodexAuth; use codex_login::RefreshTokenError; use codex_login::UnauthorizedRecovery; -use codex_login::default_client::ClientRedirectPolicy; use codex_login::default_client::add_originator_header; -use codex_login::default_client::create_client_for_route; use codex_otel::SessionTelemetry; use codex_otel::WEBSOCKET_CONTINUATION_COUNT_METRIC; use codex_otel::current_span_w3c_trace_context; use codex_protocol::ResponseItemId; use codex_protocol::auth::AuthMode; -use codex_runtime_agent_instance::RuntimeAuthManager; use codex_protocol::ThreadId; use codex_protocol::config_types::ReasoningSummary as ReasoningSummaryConfig; @@ -147,14 +141,13 @@ use codex_login::auth_env_telemetry::collect_auth_env_telemetry; use codex_model_provider::AgentIdentitySessionFallback; use codex_model_provider::ProviderAuthScope; use codex_model_provider::ProviderUnauthorizedRecovery; -use codex_model_provider::ResolvedProviderAuth; +use codex_model_provider::ResolvedProviderRequest as CurrentClientSetup; use codex_model_provider::ResponsesConnectionKey; use codex_model_provider::SharedModelProvider; use codex_model_provider::WorkspaceRoutingContext; use codex_model_provider::create_model_provider; #[cfg(test)] use codex_model_provider_info::DEFAULT_WEBSOCKET_CONNECT_TIMEOUT_MS; -use codex_model_provider_info::ModelProviderInfo; use codex_model_provider_info::WireApi; use codex_protocol::error::CodexErr; use codex_protocol::error::CodexErrorDetails; @@ -235,44 +228,6 @@ enum ClientRouting { ConfiguredProvider, } -/// Resolved API client setup for a single request attempt. -/// -/// Keeping this as a single bundle ensures prewarm and normal request paths -/// share the same auth/provider setup flow. -struct CurrentClientSetup { - auth: Option, - auth_owner_generation: Option, - auth_revision: Option, - api_provider: ApiProvider, - http_client_factory: HttpClientFactory, - redirect_policy: ClientRedirectPolicy, - api_auth: SharedAuthProvider, - agent_identity_telemetry: Option, - runtime_id: Option, -} - -impl CurrentClientSetup { - fn build_api_transport(&self, endpoint: &str) -> Result { - let api_provider = &self.api_provider; - let redirect_policy = if api_provider - .headers - .contains_key(codex_model_provider::ACCOUNT_ROUTING_HEADER) - { - ClientRedirectPolicy::Reject - } else { - self.redirect_policy - }; - let client = create_client_for_route( - &self.http_client_factory, - &api_provider.url_for_path(endpoint), - ClientRouteClass::Api, - redirect_policy, - ) - .map_err(std::io::Error::from)?; - Ok(ReqwestTransport::from_http_client(client)) - } -} - #[derive(Clone, Copy)] struct RequestRouteTelemetry { endpoint: &'static str, @@ -309,7 +264,6 @@ pub struct ModelClient { executed_tool_calls: Option, // Resolved once when the session is created, like other session feature flags. api_key_cyber_access_programs: cyber_access_program::ApiKeyCyberAccessPrograms, - runtime_auth: Option<(Arc, RuntimeAuthMode)>, } /// A turn-scoped streaming session created from a [`ModelClient`]. @@ -581,10 +535,9 @@ impl ModelClient { /// HTTP client factory must come from the effective session configuration so every transport /// observes the resolved outbound proxy policy. pub fn new( - auth_manager: Option>, + model_provider: SharedModelProvider, agent_identity_policy: AgentIdentityAuthPolicy, thread_id: ThreadId, - provider_info: ModelProviderInfo, session_source: SessionSource, originator: String, model_verbosity: Option, @@ -600,7 +553,6 @@ impl ModelClient { workspace_routing: WorkspaceRoutingContext, request_contributors: Vec>, ) -> Self { - let model_provider = create_model_provider(provider_info, auth_manager); let codex_api_key_env_enabled = model_provider .auth_manager() .as_ref() @@ -647,7 +599,6 @@ impl ModelClient { codex_responses_headers: None, event_sender: None, http_client_factory, - runtime_auth: None, restored_history: false, request_contributors, executed_tool_calls: None, @@ -672,24 +623,6 @@ impl ModelClient { self } - /// Use the harness-owned runtime for this first-party provider. The owner must share this - /// client's credential manager. Provider-specific overrides need their own runtime owner. - pub fn with_runtime_auth( - mut self, - manager: Arc, - mode: RuntimeAuthMode, - ) -> Result { - let info = self.state.provider.info(); - codex_model_provider::validate_runtime_destination(info.base_url.as_deref())?; - codex_model_provider::validate_runtime_provider( - info, - self.auth_manager().as_ref(), - &manager, - )?; - self.runtime_auth = Some((manager, mode)); - Ok(self) - } - pub(crate) fn with_session_context( mut self, prompt_cache_key: PromptCacheKey, @@ -834,7 +767,7 @@ impl ModelClient { if let Some(api_provider) = api_provider_override { client_setup.api_provider = api_provider; } - let transport = client_setup.build_api_transport(REALTIME_CALLS_ENDPOINT)?; + let transport = client_setup.http_transport(REALTIME_CALLS_ENDPOINT)?; let response = ApiRealtimeCallClient::new(transport, client_setup.api_provider, client_setup.api_auth) .create_with_session_and_headers(sdp, session_config, extra_headers) @@ -882,7 +815,7 @@ impl ModelClient { let client_setup = self .current_client_setup(ClientRouting::ConfiguredProvider) .await?; - let transport = client_setup.build_api_transport(MEMORIES_SUMMARIZE_ENDPOINT)?; + let transport = client_setup.http_transport(MEMORIES_SUMMARIZE_ENDPOINT)?; let request_telemetry = Self::build_request_telemetry( session_telemetry, AuthRequestTelemetryContext::new( @@ -1206,61 +1139,6 @@ impl ModelClient { /// This centralizes setup used by both prewarm and normal request paths so they stay in /// lockstep when auth/provider resolution changes. async fn current_client_setup(&self, routing: ClientRouting) -> Result { - if let Some((manager, mode)) = &self.runtime_auth { - let selected = codex_model_provider::select_runtime_auth( - manager, - &self.state.provider.info().build_header_map()?, - ) - .await - .map_err(|error| self.state.provider.map_api_error(error.into()))?; - let auth = selected.auth().clone(); - let auth_revision = Some(selected.change_state().generation); - let auth_owner_generation = Some(selected.change_state().owner_generation); - let api_provider = self.state.provider.api_provider_for_auth(Some(&auth))?; - let (api_provider, redirect_policy) = match routing { - ClientRouting::Workspace => { - let resolved = self - .state - .provider - .route_responses_api_provider( - api_provider, - Some(&auth), - &self.state.workspace_routing, - ) - .await?; - (resolved.provider, resolved.redirect_policy) - } - ClientRouting::ConfiguredProvider => (api_provider, ClientRedirectPolicy::Default), - }; - let runtime = selected.runtime(); - codex_model_provider::validate_runtime_selectors( - &api_provider.headers, - runtime.actor().info(), - )?; - let runtime_id = Some(runtime.id().to_string()); - let http_client_factory = selected.http_client_factory().clone(); - let resolved_auth = self - .state - .provider - .compose_auth(ResolvedProviderAuth { - auth: selected.clone().into_auth_provider(*mode), - agent_identity_telemetry: None, - runtime: Some(selected), - }) - .await?; - tracing::debug!(runtime_auth_mode = ?mode, "selected runtime authentication"); - return Ok(CurrentClientSetup { - auth: Some(auth), - auth_owner_generation, - auth_revision, - api_provider, - http_client_factory, - redirect_policy, - api_auth: resolved_auth.auth, - agent_identity_telemetry: None, - runtime_id, - }); - } self.provider_client_setup(routing, &self.state.provider) .await } @@ -1284,46 +1162,19 @@ impl ModelClient { routing: ClientRouting, provider: &SharedModelProvider, ) -> Result { - let routing = &routing; - codex_model_provider::prepare_provider_request( + let routing = match routing { + ClientRouting::Workspace => Some(&self.state.workspace_routing), + ClientRouting::ConfiguredProvider => None, + }; + codex_model_provider::resolve_request( provider.as_ref(), - provider.info().auth.is_none(), - |state| async move { - let auth = provider.auth().await; - let (api_provider, redirect_policy) = match routing { - ClientRouting::Workspace => { - let resolved = provider - .responses_api_provider(&self.state.workspace_routing) - .await?; - (resolved.provider, resolved.redirect_policy) - } - ClientRouting::ConfiguredProvider => ( - provider.api_provider().await?, - ClientRedirectPolicy::Default, - ), - }; - let resolved_auth = provider - .api_auth_for_scope(ProviderAuthScope { - agent_identity_policy: self.agent_identity_policy, - session_source: self.state.session_source.clone(), - agent_identity_session_fallback: self - .state - .agent_identity_session_fallback - .clone(), - }) - .await?; - Ok(CurrentClientSetup { - auth, - auth_owner_generation: state.map(|state| state.owner_generation), - auth_revision: state.map(|state| state.generation), - api_provider, - http_client_factory: self.http_client_factory.clone(), - redirect_policy, - api_auth: resolved_auth.auth, - agent_identity_telemetry: resolved_auth.agent_identity_telemetry, - runtime_id: None, - }) - }, + Some(ProviderAuthScope { + agent_identity_policy: self.agent_identity_policy, + session_source: self.state.session_source.clone(), + agent_identity_session_fallback: self.state.agent_identity_session_fallback.clone(), + }), + routing, + self.http_client_factory.clone(), ) .await } @@ -1953,7 +1804,7 @@ impl ResponsesHttpExecutor<'_> { .client .responses_headers(client_setup.auth.as_ref(), &model_info.slug); tracing::Span::current().record("api.path", "/responses"); - let transport = client_setup.build_api_transport("/responses")?; + let transport = client_setup.http_transport("/responses")?; let request_auth_context = AuthRequestTelemetryContext::new( client_setup.auth.as_ref().map(CodexAuth::auth_mode), client_setup.api_auth.as_ref(), diff --git a/codex-rs/core/src/client_runtime_auth_tests.rs b/codex-rs/core/src/client_runtime_auth_tests.rs index 432db1606b..e1f5b32da5 100644 --- a/codex-rs/core/src/client_runtime_auth_tests.rs +++ b/codex-rs/core/src/client_runtime_auth_tests.rs @@ -70,16 +70,19 @@ async fn selected_runtime_transport_stays_revoked(mode: RuntimeAuthMode) { )); let mut client = test_model_client(SessionSource::Exec); client.http_client_factory = factory; - Arc::get_mut(&mut client.state).unwrap().provider = create_model_provider( - ModelProviderInfo::create_openai_provider(Some(format!("{}/v1", server.uri()))), - Some(manager), - ); - let client = client.with_runtime_auth(runtime, mode).unwrap(); - let setup = client + Arc::get_mut(&mut client.state).unwrap().provider = codex_model_provider::ProviderAuthContext { + credentials: manager, + runtime_auth: Some((runtime, mode)), + } + .model_provider(ModelProviderInfo::create_openai_provider( + /*base_url*/ None, + )); + let mut setup = client .current_client_setup(ClientRouting::Workspace) .await .unwrap(); - let transport = setup.build_api_transport("/responses").unwrap(); + setup.api_provider.base_url = format!("{}/v1", server.uri()); + let transport = setup.http_transport("/responses").unwrap(); let request = setup .api_provider .build_request(http::Method::POST, "/responses"); @@ -130,15 +133,10 @@ async fn selected_runtime_transport_stays_revoked(mode: RuntimeAuthMode) { server.verify().await; } -#[test_case::test_case(RuntimeAuthMode::DualHeader, false; "dual_header")] -#[test_case::test_case(RuntimeAuthMode::AssertionOnly, false; "assertion_only")] -#[test_case::test_case(RuntimeAuthMode::DualHeader, true; "runtime_provider_dual_header")] -#[test_case::test_case(RuntimeAuthMode::AssertionOnly, true; "runtime_provider_assertion_only")] +#[test_case::test_case(RuntimeAuthMode::DualHeader; "dual_header")] +#[test_case::test_case(RuntimeAuthMode::AssertionOnly; "assertion_only")] #[tokio::test] -async fn realtime_keeps_provider_auth_until_runtime_integration( - mode: RuntimeAuthMode, - runtime_provider: bool, -) { +async fn realtime_keeps_provider_auth_until_runtime_integration(mode: RuntimeAuthMode) { let (mut client, _) = model_client_with_counting_attestation(/*include_attestation*/ true); let registration = MockServer::start().await; let manager = client.auth_manager().unwrap(); @@ -158,16 +156,13 @@ async fn realtime_keeps_provider_auth_until_runtime_integration( }, format!("{}/register", registration.uri()).parse().unwrap(), )); - if runtime_provider { - let info = client.state.provider.info().clone(); - Arc::get_mut(&mut client.state).unwrap().provider = - codex_model_provider::ProviderAuthContext { - credentials: manager, - runtime_auth: Some((Arc::clone(&runtime), mode)), - } - .model_provider(info); + let info = client.state.provider.info().clone(); + Arc::get_mut(&mut client.state).unwrap().provider = codex_model_provider::ProviderAuthContext { + credentials: manager, + runtime_auth: Some((runtime, mode)), } - let client = client.with_runtime_auth(runtime, mode).unwrap(); + .model_provider(info); + assert!(client.state.provider.runtime_auth().is_some()); let sideband = client .realtime_sideband_headers(http::HeaderMap::new()) .await diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index b6a539f56c..d8ecd9082b 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -125,10 +125,9 @@ fn test_model_client_with_thread_id( ) -> ModelClient { let provider = create_oss_provider_with_base_url("https://example.com/v1", WireApi::Responses); ModelClient::new( - /*auth_manager*/ None, + codex_model_provider::create_model_provider(provider, /*auth_manager*/ None), AgentIdentityAuthPolicy::JwtOnly, thread_id, - provider, session_source, "test_originator".to_string(), /*model_verbosity*/ None, @@ -227,13 +226,14 @@ async fn workspace_routed_http_rejects_redirects_without_a_routing_header( format!("{}/register", registration.uri()).parse().unwrap(), )); let mut client = test_model_client(SessionSource::Exec); - Arc::get_mut(&mut client.state).unwrap().provider = create_model_provider( - ModelProviderInfo::create_openai_provider(/*base_url*/ None), - Some(manager), - ); - if let Some(mode) = runtime_mode { - client = client.with_runtime_auth(runtime, mode).unwrap(); - } + Arc::get_mut(&mut client.state).unwrap().provider = + codex_model_provider::ProviderAuthContext { + credentials: manager, + runtime_auth: runtime_mode.map(|mode| (runtime, mode)), + } + .model_provider(ModelProviderInfo::create_openai_provider( + /*base_url*/ None, + )); let mut setup = client .current_client_setup(super::ClientRouting::Workspace) .await @@ -246,7 +246,7 @@ async fn workspace_routed_http_rejects_redirects_without_a_routing_header( } // Exercise the resolved route's redirect policy against loopback HTTP servers. setup.api_provider.base_url = origin.uri(); - let transport = setup.build_api_transport("/responses").unwrap(); + let transport = setup.http_transport("/responses").unwrap(); let request = setup .api_provider .build_request(http::Method::POST, "/responses") @@ -2147,10 +2147,9 @@ fn model_client_with_counting_attestation( ) }; let model_client = ModelClient::new( - auth_manager, + codex_model_provider::create_model_provider(provider, auth_manager), AgentIdentityAuthPolicy::JwtOnly, ThreadId::new(), - provider, SessionSource::Exec, "test_originator".to_string(), /*model_verbosity*/ None, diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index d4227af4b1..31c65b3454 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -114,6 +114,7 @@ pub(crate) async fn run_codex_thread_interactive( instructions, installation_id: parent_session.installation_id.clone(), auth_manager, + runtime_auth: parent_session.services.runtime_auth.clone(), models_manager, git_root_discovery: Arc::clone(&parent_session.services.git_root_discovery), environment_manager: parent_session diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index 2f1a78651c..6990af063c 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -76,6 +76,7 @@ use codex_features::Features; use codex_features::FeaturesToml; use codex_features::MultiAgentV2ConfigToml; use codex_features::NetworkProxyConfigToml; +use codex_features::RuntimeAuthMode; use codex_features::SleepToolMode; use codex_features::TokenBudgetConfigToml; use codex_git_utils::resolve_root_git_project_for_trust; @@ -1166,6 +1167,8 @@ pub struct Config { pub current_time_reminder: Option, /// How the sleep tool is selected when its feature gate is enabled. pub sleep_tool_mode: SleepToolMode, + /// Authentication mode used when runtime agent instances are enabled. + pub runtime_auth_mode: RuntimeAuthMode, /// Centralized feature flags; source of truth for feature gating. pub features: ManagedFeatures, @@ -4033,6 +4036,15 @@ impl Config { FeatureToml::Config(config) => config.mode, }) .unwrap_or_default(); + let runtime_auth_mode = cfg + .features + .as_ref() + .and_then(|features| features.runtime_agent_instances.as_ref()) + .and_then(|feature| match feature { + FeatureToml::Enabled(_) => None, + FeatureToml::Config(config) => config.mode, + }) + .unwrap_or_default(); let terminal_resize_reflow = resolve_terminal_resize_reflow_config(&cfg); let agent_roles = @@ -4754,6 +4766,7 @@ impl Config { rollout_budget, current_time_reminder, sleep_tool_mode, + runtime_auth_mode, features, runtime_feature_defaults: BTreeMap::new(), suppress_unstable_features_warning: cfg diff --git a/codex-rs/core/src/guardian/test_host.rs b/codex-rs/core/src/guardian/test_host.rs index 3f98917afb..a30b96b1a8 100644 --- a/codex-rs/core/src/guardian/test_host.rs +++ b/codex-rs/core/src/guardian/test_host.rs @@ -27,8 +27,12 @@ pub(crate) fn install(session: &Session, config: &Config) { ); let manager = Arc::new(crate::ThreadManager::new( config, - Arc::clone(&session.services.auth_manager), - Arc::clone(&session.services.models_manager), + crate::build_startup_models( + config, + Arc::clone(&session.services.auth_manager), + session.installation_id.clone(), + crate::test_support::fixed_model_catalog(Arc::clone(&session.services.models_manager)), + ), crate::CodexAppsToolsCache::default(), SessionSource::Exec, session.services.turn_environments.environment_manager(), @@ -40,7 +44,6 @@ pub(crate) fn install(session: &Session, config: &Config) { crate::passthrough_image_store(), Arc::clone(&session.services.thread_store), /*agent_graph_store*/ None, - session.installation_id.clone(), /*attestation_provider*/ None, /*external_time_provider*/ None, )); diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 906e74b345..26940e692d 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -183,10 +183,12 @@ pub use thread_manager::InternalSessionParent; pub use thread_manager::NewThread; pub use thread_manager::PromptCacheKey; pub use thread_manager::StartThreadOptions; +pub use thread_manager::StartupModels; pub use thread_manager::ThreadManager; pub use thread_manager::ThreadRegistration; pub use thread_manager::ThreadShutdownReport; pub use thread_manager::build_models_manager; +pub use thread_manager::build_startup_models; pub use thread_manager::local_agent_graph_store_from_state_db; pub use thread_manager::passthrough_image_store; pub use thread_store_builder::ThreadStoreBuilder; @@ -281,3 +283,6 @@ pub mod otel_init; // Captured environment bindings can be passed back to ThreadManager by internal reviewers. pub use environment_selection::TurnEnvironmentSnapshot; pub use environment_selection::validate_environment_ids_and_cwds; + +mod runtime_auth; +pub use runtime_auth::build_runtime_auth_manager; diff --git a/codex-rs/core/src/prompt_debug.rs b/codex-rs/core/src/prompt_debug.rs index c54f109aab..9a8b6d8278 100644 --- a/codex-rs/core/src/prompt_debug.rs +++ b/codex-rs/core/src/prompt_debug.rs @@ -5,6 +5,7 @@ use codex_exec_server::ExecServerRuntimeOptions; use codex_extension_api::ExtensionRegistry; use codex_extension_api::UserInstructionsProvider; use codex_login::AuthManager; +use codex_models_manager::manager::ModelCatalogSettings; use codex_protocol::error::CodexErr; use codex_protocol::error::Result as CodexResult; use codex_protocol::models::ResponseItem; @@ -44,10 +45,15 @@ pub async fn build_prompt_input( let thread_store = ThreadStoreBuilder::from_config(&config, state_db.clone()).build(); let installation_id = resolve_installation_id(&config.codex_home).await?; - let thread_manager = ThreadManager::new( + let startup = crate::build_startup_models( &config, Arc::clone(&auth_manager), - crate::thread_manager::build_models_manager(&config, Arc::clone(&auth_manager)), + installation_id, + Arc::new(ModelCatalogSettings::disk(config.codex_home.to_path_buf())), + ); + let thread_manager = ThreadManager::new( + &config, + startup, crate::CodexAppsToolsCache::default(), SessionSource::Exec, Arc::new( @@ -65,7 +71,6 @@ pub async fn build_prompt_input( crate::thread_manager::passthrough_image_store(), thread_store, crate::local_agent_graph_store_from_state_db(state_db.as_ref()), - installation_id, /*attestation_provider*/ None, /*external_time_provider*/ None, ); diff --git a/codex-rs/core/src/runtime_auth.rs b/codex-rs/core/src/runtime_auth.rs new file mode 100644 index 0000000000..554591aca8 --- /dev/null +++ b/codex-rs/core/src/runtime_auth.rs @@ -0,0 +1,53 @@ +//! Construct and share the harness-owned runtime credential context. + +use crate::config::Config; +use codex_features::Feature; +use codex_login::AuthManager; +use codex_runtime_agent_instance::RuntimeAuthManager; +use codex_runtime_agent_instance::RuntimeRegistrationMetadata; +use codex_runtime_agent_instance::RuntimeRole; +use std::sync::Arc; +use std::sync::OnceLock; + +// A later thread may enable runtime auth even when harness startup left it disabled. +pub(crate) type RuntimeAuthSlot = Arc>>; + +/// Share this owner between thread sessions and first-party API clients. +/// Construction does not register a runtime or generate its key. +pub fn build_runtime_auth_manager( + config: &Config, + credentials: Arc, + installation_id: &str, +) -> Option> { + runtime_for_config( + &RuntimeAuthSlot::default(), + config, + credentials, + installation_id, + ) +} + +pub(crate) fn runtime_for_config( + slot: &RuntimeAuthSlot, + config: &Config, + credentials: Arc, + installation_id: &str, +) -> Option> { + config + .features + .enabled(Feature::RuntimeAgentInstances) + .then(|| { + Arc::clone(slot.get_or_init(|| { + Arc::new(RuntimeAuthManager::for_chatgpt( + credentials, + installation_id.to_owned(), + config.chatgpt_base_url.clone(), + RuntimeRegistrationMetadata { + harness: Some("codex".to_string()), + harness_version: Some(env!("CARGO_PKG_VERSION").to_string()), + role: RuntimeRole::Orchestrator, + }, + )) + })) + }) +} diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index e67f317163..1575de9244 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -218,6 +218,10 @@ impl Session { published: false, }; let auth = self.services.auth_manager.auth().await; + // Capture after claiming the refresh so a new invalidation cannot be consumed + // with the environment snapshot taken before its permissions changed. + let environments = self.services.turn_environments.snapshot().await; + let environment_selections = environments.all_selections(); let desired = self.latest_mcp_desired_state(auth, environments).await; let selected_capability_roots = self .resolve_selected_capability_roots_for_step(&desired.environments) diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 58cd57bd69..4113bfbea4 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -2,6 +2,7 @@ pub(crate) mod startup; use crate::PromptCacheKey; use crate::WithTurnExtensionData; +use crate::runtime_auth::RuntimeAuthSlot; use std::borrow::Cow; use std::collections::HashMap; use std::collections::HashSet; @@ -235,7 +236,6 @@ use crate::thread_rollout_truncation::initial_history_has_prior_user_turns; use codex_config::CONFIG_TOML_FILE; use codex_config::ConfigLayerSource; use codex_config::types::McpServerConfig; -use codex_model_provider::create_model_provider; use codex_model_provider_info::ModelProviderInfo; use codex_protocol::error::CodexErr; use codex_protocol::error::CodexErrorDetails; @@ -509,6 +509,7 @@ pub(crate) struct SessionSpawnArgs { pub(crate) instructions: SessionInstructions, pub(crate) installation_id: String, pub(crate) auth_manager: Arc, + pub(crate) runtime_auth: RuntimeAuthSlot, pub(crate) models_manager: SharedModelsManager, pub(crate) git_root_discovery: Arc, pub(crate) environment_manager: Arc, @@ -630,6 +631,7 @@ impl Session { instructions, installation_id, auth_manager, + runtime_auth, models_manager, git_root_discovery, environment_manager, @@ -758,12 +760,19 @@ impl Session { }; let mut config = Arc::new(config); + let configured_runtime = crate::runtime_auth::runtime_for_config( + &runtime_auth, + config.as_ref(), + Arc::clone(&auth_manager), + &installation_id, + ); let refresh_strategy = if session_source.is_non_root_agent() { codex_models_manager::manager::RefreshStrategy::Offline } else { codex_models_manager::manager::RefreshStrategy::OnlineIfUncached }; - if config.model.is_none() + if configured_runtime.is_some() + || config.model.is_none() || !matches!( refresh_strategy, codex_models_manager::manager::RefreshStrategy::Offline @@ -951,12 +960,15 @@ impl Session { let enable_full_access_form_input = !root_capabilities_restricted && matches!(thread_source.as_ref(), Some(ThreadSource::User)) && client_mcp_extensions.contains(OPENAI_STANDARD_FORM_INPUT_EXTENSION_ID); + let provider_auth = codex_model_provider::ProviderAuthContext { + credentials: Arc::clone(&auth_manager), + runtime_auth: configured_runtime.map(|runtime| (runtime, config.runtime_auth_mode)), + }; + // Replace any inherited context before startup hooks see this thread's effective auth. + thread_extension_init.insert(provider_auth.clone()); let session_configuration = SessionConfiguration { turn_extension_init, - provider: create_model_provider( - config.model_provider.clone(), - Some(Arc::clone(&auth_manager)), - ), + provider: provider_auth.model_provider(config.model_provider.clone()), step_settings: Arc::new(StepSettings { collaboration_mode, reasoning_summary: config.model_reasoning_summary, @@ -1014,6 +1026,7 @@ impl Session { instructions, installation_id, auth_manager.clone(), + runtime_auth, models_manager.clone(), git_root_discovery, model_info, diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index ce336c5a9d..dc4ae7ba02 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -20,6 +20,7 @@ use crate::mcp_tool_call::McpToolApprovalContext; use crate::responses_metadata::CodexResponsesMetadata; use crate::responses_metadata::CodexResponsesRequestKind; use crate::responses_metadata::CompactionTurnMetadata; +use crate::runtime_auth::RuntimeAuthSlot; use crate::shell_snapshot::ShellSnapshot; use crate::shell_snapshot::SnapshotCredentialBrokerState; use crate::state::ActiveTurn; @@ -833,6 +834,7 @@ impl Session { mut instructions: SessionInstructions, installation_id: String, auth_manager: Arc, + runtime_auth: RuntimeAuthSlot, models_manager: SharedModelsManager, git_root_discovery: Arc, model_info: CatalogModel, @@ -1794,6 +1796,48 @@ impl Session { ); let reasoning_effort_override_enabled = config.features.enabled(Feature::ReasoningEffortOverride) && !title_request; + let model_client = ModelClient::new( + Arc::clone(&session_configuration.provider), + if config.features.enabled(Feature::UseAgentIdentity) { + AgentIdentityAuthPolicy::ChatGptAuth + } else { + AgentIdentityAuthPolicy::JwtOnly + }, + thread_id, + session_configuration.session_source.clone(), + session_configuration.originator.clone(), + config.model_verbosity, + config.features.enabled(Feature::ContentItemKinds), + reasoning_effort_override_enabled, + config.features.enabled(Feature::EnableRequestCompression), + config.features.enabled(Feature::RuntimeMetrics), + Self::build_model_client_beta_features_header(config.as_ref()), + /*concurrent_reasoning_summaries_enabled*/ + config + .features + .enabled(Feature::ConcurrentReasoningSummaries), + config.features.enabled(Feature::OutputTokenReplay), + attestation_provider.clone(), + config.http_client_factory(), + workspace_routing.as_ref().clone(), + extensions.model_request_contributors().to_vec(), + ) + .with_executed_tool_calls(executed_tool_calls.clone()) + .with_restored_history(matches!( + &initial_history, + InitialHistory::Resumed(_) | InitialHistory::Forked(_) + )) + .with_session_context( + prompt_cache_key, + crate::guardian::prompt_cache_key_override_for_review_session( + &session_configuration.session_source, + session_configuration.parent_thread_id, + ) + .or(fork_cache_key), + tx_event.clone(), + codex_responses_headers, + crate::cyber_access_program::ApiKeyCyberAccessPrograms::from_config(&config), + ); // Once tracing starts, startup cleanup must own the session before the next await. let rollout_thread_trace = if matches!( session_configuration.session_source, @@ -1832,6 +1876,7 @@ impl Session { show_raw_agent_reasoning: config.show_raw_agent_reasoning, exec_policy, auth_manager: Arc::clone(&auth_manager), + runtime_auth, openai_file_upload_client_pool: RouteAwareClientPool::new_without_request_logging( config.http_client_factory(), ClientRouteClass::Api, @@ -1866,49 +1911,7 @@ impl Session { thread_store: Arc::clone(&thread_store), attestation_provider: attestation_provider.clone(), time_provider, - model_client: ModelClient::new( - Some(Arc::clone(&auth_manager)), - if config.features.enabled(Feature::UseAgentIdentity) { - AgentIdentityAuthPolicy::ChatGptAuth - } else { - AgentIdentityAuthPolicy::JwtOnly - }, - thread_id, - session_configuration.provider.info().clone(), - session_configuration.session_source.clone(), - session_configuration.originator.clone(), - config.model_verbosity, - config.features.enabled(Feature::ContentItemKinds), - reasoning_effort_override_enabled, - config.features.enabled(Feature::EnableRequestCompression), - config.features.enabled(Feature::RuntimeMetrics), - Self::build_model_client_beta_features_header(config.as_ref()), - /*concurrent_reasoning_summaries_enabled*/ - config - .features - .enabled(Feature::ConcurrentReasoningSummaries), - config.features.enabled(Feature::OutputTokenReplay), - attestation_provider, - config.http_client_factory(), - workspace_routing.as_ref().clone(), - extensions.model_request_contributors().to_vec(), - ) - .with_executed_tool_calls(executed_tool_calls.clone()) - .with_restored_history(matches!( - &initial_history, - InitialHistory::Resumed(_) | InitialHistory::Forked(_) - )) - .with_session_context( - prompt_cache_key, - crate::guardian::prompt_cache_key_override_for_review_session( - &session_configuration.session_source, - session_configuration.parent_thread_id, - ) - .or(fork_cache_key), - tx_event.clone(), - codex_responses_headers, - crate::cyber_access_program::ApiKeyCyberAccessPrograms::from_config(&config), - ), + model_client, executed_tool_calls: executed_tool_calls.clone(), code_mode_service: crate::tools::code_mode::CodeModeService::new( thread_id, diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index c40209fcbf..1ab26f39fb 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -895,10 +895,12 @@ fn test_model_client_session() -> crate::client::ModelClientSession { let thread_id = ThreadId::try_from("00000000-0000-4000-8000-000000000001") .expect("test thread id should be valid"); crate::client::ModelClient::new( - /*auth_manager*/ None, + codex_model_provider::create_model_provider( + ModelProviderInfo::create_openai_provider(/*base_url*/ None), + /*auth_manager*/ None, + ), AgentIdentityAuthPolicy::JwtOnly, thread_id, - ModelProviderInfo::create_openai_provider(/* base_url */ /*base_url*/ None), codex_protocol::protocol::SessionSource::Exec, "test_originator".to_string(), /*model_verbosity*/ None, @@ -7077,6 +7079,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { show_raw_agent_reasoning: config.show_raw_agent_reasoning, exec_policy, auth_manager: auth_manager.clone(), + runtime_auth: Default::default(), openai_file_upload_client_pool: RouteAwareClientPool::new_without_request_logging( config.http_client_factory(), ClientRouteClass::Api, @@ -7116,10 +7119,12 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { attestation_provider: None, time_provider: Arc::new(crate::current_time::SystemTimeProvider), model_client: ModelClient::new( - Some(auth_manager.clone()), + codex_model_provider::create_model_provider( + session_configuration.provider.info().clone(), + Some(auth_manager.clone()), + ), AgentIdentityAuthPolicy::JwtOnly, thread_id, - session_configuration.provider.info().clone(), session_configuration.session_source.clone(), session_configuration.originator.clone(), config.model_verbosity, @@ -7350,6 +7355,7 @@ async fn make_session_with_config_and_rx( SessionInstructions::default(), "11111111-1111-4111-8111-111111111111".to_string(), auth_manager, + /*runtime_auth*/ Default::default(), models_manager, Arc::default(), model_info, @@ -7489,6 +7495,7 @@ async fn make_session_with_history_source_and_agent_control_and_rx( SessionInstructions::default(), "11111111-1111-4111-8111-111111111111".to_string(), auth_manager, + /*runtime_auth*/ Default::default(), models_manager, Arc::default(), model_info, @@ -9407,6 +9414,7 @@ where show_raw_agent_reasoning: config.show_raw_agent_reasoning, exec_policy, auth_manager: Arc::clone(&auth_manager), + runtime_auth: Default::default(), openai_file_upload_client_pool: RouteAwareClientPool::new_without_request_logging( config.http_client_factory(), ClientRouteClass::Api, @@ -9446,10 +9454,12 @@ where attestation_provider: None, time_provider: Arc::new(crate::current_time::SystemTimeProvider), model_client: ModelClient::new( - Some(Arc::clone(&auth_manager)), + codex_model_provider::create_model_provider( + session_configuration.provider.info().clone(), + Some(Arc::clone(&auth_manager)), + ), AgentIdentityAuthPolicy::JwtOnly, thread_id, - session_configuration.provider.info().clone(), session_configuration.session_source.clone(), session_configuration.originator.clone(), config.model_verbosity, diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index 7ac50acc55..24df4e9d0a 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -1338,6 +1338,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { instructions: Default::default(), installation_id: "11111111-1111-4111-8111-111111111111".to_string(), auth_manager, + runtime_auth: Default::default(), models_manager, git_root_discovery: Arc::default(), environment_manager: Arc::new(EnvironmentManager::default_for_tests()), diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index 3671386bd4..2bb25c7557 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -15,6 +15,7 @@ use crate::environment_selection::ThreadEnvironments; use crate::exec_policy::ExecPolicyManager; use crate::mcp::McpManager; use crate::mcp_tool_exposure::McpHandlerCache; +use crate::runtime_auth::RuntimeAuthSlot; use crate::tools::ExecutedToolCalls; use crate::tools::code_mode::CodeModeService; use crate::tools::handlers::ToolSearchHandlerCache; @@ -62,6 +63,7 @@ pub(crate) struct SessionServices { pub(crate) show_raw_agent_reasoning: bool, pub(crate) exec_policy: Arc, pub(crate) auth_manager: Arc, + pub(crate) runtime_auth: RuntimeAuthSlot, /// Upload-only clients shared across turns without logging signed blob URLs. pub(crate) openai_file_upload_client_pool: RouteAwareClientPool, pub(crate) models_manager: SharedModelsManager, diff --git a/codex-rs/core/src/test_support.rs b/codex-rs/core/src/test_support.rs index 2c33f0c78e..f6ba23047a 100644 --- a/codex-rs/core/src/test_support.rs +++ b/codex-rs/core/src/test_support.rs @@ -41,6 +41,30 @@ use crate::responses_metadata::subagent_metadata_kind; use crate::thread_manager; use crate::unified_exec; +/// Supply fixed model metadata to tests without implementing a provider. +pub fn fixed_model_catalog( + manager: SharedModelsManager, +) -> Arc { + struct FixedCatalog(SharedModelsManager); + + impl codex_model_provider::ModelCatalogFactory for FixedCatalog { + fn build( + &self, + _provider: codex_model_provider::SharedModelProvider, + _config_model_catalog: Option, + _api_key_model_discovery_default: bool, + ) -> SharedModelsManager { + Arc::clone(&self.0) + } + + fn set_api_key_model_discovery_enabled(&self, enabled: bool) { + self.0.set_api_key_model_discovery_enabled(enabled); + } + } + + Arc::new(FixedCatalog(manager)) +} + static TEST_MODEL_PRESETS: Lazy> = Lazy::new(|| { let mut response = bundled_models_response() .unwrap_or_else(|err| panic!("bundled models.json should parse: {err}")); diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 4baf6fca8b..c11fda446e 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -1,4 +1,5 @@ mod managed; +mod model_catalog; mod shared_instructions; mod shutdown; @@ -72,7 +73,6 @@ use codex_login::AuthManager; use codex_login::CodexAuth; use codex_login::default_client::CODEX_INTERNAL_ORIGINATOR_OVERRIDE_ENV_VAR; use codex_login::default_client::originator; -use codex_model_provider::create_model_provider; use codex_model_provider_info::ModelProviderInfo; use codex_model_provider_info::OPENAI_PROVIDER_ID; use codex_models_manager::manager::RefreshStrategy; @@ -429,6 +429,7 @@ struct ThreadSpawnRequest { startup: Option>, options: StartThreadOptions, auth_manager: Arc, + inherited_runtime_auth: Option, agent_control: AgentControlInit, parent_thread_id: Option, parent_originator: Option, @@ -452,6 +453,7 @@ impl ThreadSpawnRequest { startup: None, options, auth_manager, + inherited_runtime_auth: None, agent_control: agent_control.into(), parent_thread_id: None, parent_originator: None, @@ -521,8 +523,8 @@ pub(crate) struct ThreadManagerState { thread_id_generator: ThreadIdGenerator, agent_control_factory: Option, auth_manager: Arc, - models_manager: SharedModelsManager, - models_manager_provider: ModelProviderInfo, + runtime_auth: crate::runtime_auth::RuntimeAuthSlot, + model_catalog: model_catalog::ThreadModelCatalog, git_root_discovery: Arc, environment_manager: Arc, starting_mcp_runtimes: std::sync::Mutex>>, @@ -544,20 +546,9 @@ pub(crate) struct ThreadManagerState { ops_log: Option, } -pub fn build_models_manager( - config: &Config, - auth_manager: Arc, -) -> SharedModelsManager { - let provider = create_model_provider(config.model_provider.clone(), Some(auth_manager)); - let manager = provider.models_manager( - config.codex_home.to_path_buf(), - config.model_catalog.clone(), - ); - manager.set_api_key_model_discovery_enabled( - config.features.enabled(Feature::ApiKeyModelDiscovery), - ); - manager -} +pub use model_catalog::StartupModels; +pub use model_catalog::build_models_manager; +pub use model_catalog::build_startup_models; /// Constructs the default image store that preserves inline images. pub fn passthrough_image_store() -> Arc { @@ -577,8 +568,7 @@ impl ThreadManager { #[allow(clippy::too_many_arguments)] pub fn new( config: &Config, - auth_manager: Arc, - models_manager: SharedModelsManager, + startup: StartupModels, codex_apps_tools_cache: CodexAppsToolsCache, session_source: SessionSource, environment_manager: Arc, @@ -588,10 +578,15 @@ impl ThreadManager { image_store: Arc, thread_store: Arc, agent_graph_store: Option>, - installation_id: String, attestation_provider: Option>, external_time_provider: Option>, ) -> Self { + let StartupModels { + credentials: auth_manager, + runtime_auth, + installation_id, + catalog: model_catalog, + } = startup; let codex_home = config.codex_home.clone(); let restriction_product = session_source.restriction_product(); let (thread_created_tx, _) = broadcast::channel(THREAD_CREATED_CHANNEL_CAPACITY); @@ -629,8 +624,7 @@ impl ThreadManager { thread_created_tx, thread_id_generator: default_thread_id_generator(), agent_control_factory: None, - models_manager, - models_manager_provider: config.model_provider.clone(), + model_catalog, git_root_discovery: Arc::default(), environment_manager, starting_mcp_runtimes: std::sync::Mutex::new(Vec::new()), @@ -646,6 +640,7 @@ impl ThreadManager { attestation_provider, external_time_provider, auth_manager, + runtime_auth, session_source, installation_id, analytics_events_client, @@ -834,14 +829,24 @@ impl ThreadManager { let agent_graph_store = local_agent_graph_store_from_state_db(state_db.as_ref()); Self { state: Arc::new(ThreadManagerState { + runtime_auth: Default::default(), threads: Arc::new(RwLock::new(HashMap::new())), shared_thread_instructions: Default::default(), thread_created_tx, thread_id_generator: default_thread_id_generator(), agent_control_factory: None, - models_manager_provider: provider.clone(), - models_manager: create_model_provider(provider, Some(auth_manager.clone())) - .models_manager(codex_home, /*config_model_catalog*/ None), + model_catalog: model_catalog::ThreadModelCatalog::new( + provider, + /*catalog*/ None, + Arc::new(codex_models_manager::manager::ModelCatalogSettings::disk( + codex_home, + )), + /*api_key_model_discovery_default*/ false, + codex_model_provider::ProviderAuthContext { + credentials: Arc::clone(&auth_manager), + runtime_auth: None, + }, + ), git_root_discovery: Arc::default(), environment_manager, starting_mcp_runtimes: std::sync::Mutex::new(Vec::new()), @@ -982,7 +987,31 @@ impl ThreadManager { } pub fn get_models_manager(&self) -> SharedModelsManager { - self.state.models_manager.clone() + self.state.model_catalog.manager.clone() + } + + /// Select the app-server catalog using the same owner as thread requests. + pub fn models_manager_for_config(&self, config: &Config) -> SharedModelsManager { + let credentials = self.auth_manager(); + let runtime = crate::runtime_auth::runtime_for_config( + &self.state.runtime_auth, + config, + Arc::clone(&credentials), + &self.state.installation_id, + ); + self.state.model_catalog.for_config( + config, + config.model_catalog.as_ref(), + credentials, + runtime, + /*same_owner*/ true, + ) + } + + pub fn set_api_key_model_discovery_enabled(&self, enabled: bool) { + self.state + .model_catalog + .set_api_key_model_discovery_enabled(enabled); } pub async fn list_models( @@ -991,13 +1020,14 @@ impl ThreadManager { http_client_factory: codex_http_client::HttpClientFactory, ) -> Vec { self.state - .models_manager + .model_catalog + .manager .list_models(refresh_strategy, http_client_factory) .await } pub fn list_collaboration_modes(&self) -> Vec { - self.state.models_manager.list_collaboration_modes() + self.state.model_catalog.manager.list_collaboration_modes() } pub async fn list_thread_ids(&self) -> Vec { @@ -2081,8 +2111,9 @@ impl ThreadManagerState { (options.client_mcp_extensions, options.prompt_cache_key) = self.inherited_child_settings(parent_thread_id).await; let forked = matches!(options.initial_history, InitialHistory::Forked(_)); - let mut request = - ThreadSpawnRequest::new(options, Arc::clone(&self.auth_manager), agent_control); + let (auth_manager, runtime_auth) = self.inherited_child_auth(parent_thread_id).await; + let mut request = ThreadSpawnRequest::new(options, auth_manager, agent_control); + request.inherited_runtime_auth = Some(runtime_auth); request.parent_thread_id = parent_thread_id; request.inherited_exec_policy = inherited_exec_policy; if forked { @@ -2103,8 +2134,9 @@ impl ThreadManagerState { options.reserved_thread_id.ok_or_else(|| { CodexErr::InvalidRequest("prepared child startup requires a reserved thread ID".into()) })?; - let mut request = - ThreadSpawnRequest::new(options, Arc::clone(&self.auth_manager), agent_control); + let (auth_manager, runtime_auth) = self.inherited_child_auth(Some(parent_thread_id)).await; + let mut request = ThreadSpawnRequest::new(options, auth_manager, agent_control); + request.inherited_runtime_auth = Some(runtime_auth); request.ready_notification = ThreadReadyNotification::Defer; request.parent_thread_id = Some(parent_thread_id); let history_initialization = match &request.options.initial_history { @@ -2175,6 +2207,23 @@ impl ThreadManagerState { Box::pin(self.spawn_thread(request)).await } + async fn inherited_child_auth( + &self, + parent_thread_id: Option, + ) -> (Arc, crate::runtime_auth::RuntimeAuthSlot) { + let threads = self.threads.read().await; + match parent_thread_id.and_then(|id| threads.get(&id)) { + Some(parent) => ( + Arc::clone(&parent.session.services.auth_manager), + Arc::clone(&parent.session.services.runtime_auth), + ), + None => ( + Arc::clone(&self.auth_manager), + Arc::clone(&self.runtime_auth), + ), + } + } + async fn inherited_child_settings( &self, parent_thread_id: Option, @@ -2226,6 +2275,7 @@ impl ThreadManagerState { startup, options, auth_manager, + inherited_runtime_auth, agent_control, parent_thread_id, parent_originator, @@ -2527,14 +2577,36 @@ impl ThreadManagerState { } else { codex_sandboxing::WindowsSandboxProxySettingsMode::Reconcile }; - let models_manager = if config.model_provider == self.models_manager_provider { - Arc::clone(&self.models_manager) + // Resume may supply a different credential owner. Keep its runtime separate from the + // harness default, and let sessions using that parent's credentials inherit its owner. + let runtime_auth = if let Some(runtime_auth) = inherited_runtime_auth { + runtime_auth + } else if Arc::ptr_eq(&auth_manager, &self.auth_manager) { + Arc::clone(&self.runtime_auth) } else { - build_models_manager(&config, Arc::clone(&auth_manager)) + let threads = self.threads.read().await; + parent_thread_id + .or(forked_from_thread_id) + .and_then(|id| threads.get(&id)) + .filter(|parent| Arc::ptr_eq(&parent.session.services.auth_manager, &auth_manager)) + .map(|parent| Arc::clone(&parent.session.services.runtime_auth)) + .unwrap_or_default() }; - let models_manager = models_manager + let runtime = crate::runtime_auth::runtime_for_config( + &runtime_auth, + &config, + 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), + ); + let models_manager = manager .for_originator(&startup_context.originator) - .unwrap_or(models_manager); + .unwrap_or(manager); let spawn_result = Session::spawn(SessionSpawnArgs { attached_thread, startup, @@ -2544,6 +2616,7 @@ impl ThreadManagerState { instructions: startup_context.instructions, installation_id: self.installation_id.clone(), auth_manager, + runtime_auth, models_manager, git_root_discovery: Arc::clone(&self.git_root_discovery), environment_manager: Arc::clone(&self.environment_manager), diff --git a/codex-rs/core/src/thread_manager/model_catalog.rs b/codex-rs/core/src/thread_manager/model_catalog.rs new file mode 100644 index 0000000000..0bb0645c08 --- /dev/null +++ b/codex-rs/core/src/thread_manager/model_catalog.rs @@ -0,0 +1,166 @@ +//! Keep each catalog bound to its provider and auth owner while preserving host cache policy. + +use crate::config::Config; +use crate::runtime_auth::RuntimeAuthSlot; +use crate::runtime_auth::runtime_for_config; +use codex_features::Feature; +use codex_features::RuntimeAuthMode; +use codex_login::AuthManager; +use codex_model_provider::ModelCatalogFactory; +use codex_model_provider::ProviderAuthContext; +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_runtime_agent_instance::RuntimeAuthManager; +use std::sync::Arc; + +/// The startup catalog and the credential/runtime owner used to construct it. +pub struct StartupModels { + pub(super) credentials: Arc, + pub(super) runtime_auth: RuntimeAuthSlot, + pub(super) installation_id: String, + pub(super) catalog: ThreadModelCatalog, +} + +impl StartupModels { + /// Share the startup catalog with host discovery and prefetch callers. + pub fn models_manager(&self) -> SharedModelsManager { + Arc::clone(&self.catalog.manager) + } +} + +/// Construct the catalog and runtime owner together using the host's catalog factory. +/// Use `ModelCatalogSettings` for standard provider catalogs with a disk, custom, or no cache. +/// Runtime registration remains lazy; this does not register an instance or generate a key. +pub fn build_startup_models( + config: &Config, + credentials: Arc, + installation_id: String, + factory: Arc, +) -> StartupModels { + let runtime_auth = RuntimeAuthSlot::default(); + let runtime = runtime_for_config( + &runtime_auth, + config, + Arc::clone(&credentials), + &installation_id, + ); + let catalog = ThreadModelCatalog::new( + config.model_provider.clone(), + config.model_catalog.clone(), + factory, + config.features.enabled(Feature::ApiKeyModelDiscovery), + ProviderAuthContext { + credentials: Arc::clone(&credentials), + runtime_auth: runtime.map(|runtime| (runtime, config.runtime_auth_mode)), + }, + ); + StartupModels { + credentials, + runtime_auth, + installation_id, + catalog, + } +} + +pub(super) struct ThreadModelCatalog { + pub(super) manager: SharedModelsManager, + provider: ModelProviderInfo, + catalog: Option, + factory: Arc, + api_key_model_discovery_default: bool, + runtime_mode: Option, +} + +pub fn build_models_manager( + config: &Config, + credentials: Arc, + runtime: Option>, +) -> SharedModelsManager { + let auth = ProviderAuthContext { + credentials, + runtime_auth: runtime + .filter(|_| config.features.enabled(Feature::RuntimeAgentInstances)) + .map(|runtime| (runtime, config.runtime_auth_mode)), + }; + ModelCatalogSettings::disk(config.codex_home.to_path_buf()).build( + auth.model_provider(config.model_provider.clone()), + config.model_catalog.clone(), + config.features.enabled(Feature::ApiKeyModelDiscovery), + ) +} + +impl ThreadModelCatalog { + pub(super) fn new( + provider: ModelProviderInfo, + catalog: Option, + factory: Arc, + api_key_model_discovery_default: bool, + auth: ProviderAuthContext, + ) -> Self { + let runtime_mode = auth.runtime_auth.as_ref().map(|(_, mode)| *mode); + let manager = factory.build( + auth.model_provider(provider.clone()), + catalog.clone(), + api_key_model_discovery_default, + ); + Self { + manager, + provider, + catalog, + factory, + api_key_model_discovery_default, + runtime_mode, + } + } + + pub(super) fn for_thread( + &self, + config: &Config, + credentials: Arc, + runtime: Option>, + same_owner: bool, + ) -> SharedModelsManager { + let models = config.model_catalog.as_ref().or_else(|| { + self.catalog + .as_ref() + .filter(|_| self.provider == config.model_provider) + }); + self.for_config(config, models, credentials, runtime, same_owner) + } + + pub(super) fn for_config( + &self, + config: &Config, + models: Option<&ModelsResponse>, + credentials: Arc, + runtime: Option>, + same_owner: bool, + ) -> SharedModelsManager { + let runtime = runtime.map(|runtime| (runtime, config.runtime_auth_mode)); + let api_key_model_discovery_default = + config.features.enabled(Feature::ApiKeyModelDiscovery); + if same_owner + && self.runtime_mode == runtime.as_ref().map(|(_, mode)| *mode) + && self.provider == config.model_provider + && self.catalog.as_ref() == models + && self.api_key_model_discovery_default == api_key_model_discovery_default + { + return Arc::clone(&self.manager); + } + let auth = ProviderAuthContext { + credentials, + runtime_auth: runtime, + }; + self.factory.build( + auth.model_provider(config.model_provider.clone()), + models.cloned(), + api_key_model_discovery_default, + ) + } + + 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 c4601f0e59..ffb2463c00 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -25,6 +25,7 @@ use codex_extension_api::SessionIsolation; use codex_extension_api::empty_extension_registry; use codex_history::InitialHistory; use codex_history::ResumedHistory; +use codex_models_manager::manager::ModelCatalogSettings; use codex_models_manager::manager::RefreshStrategy; use codex_protocol::ResponseItemId; use codex_protocol::approvals::ExecPolicyAmendment; @@ -678,6 +679,128 @@ async fn thread_id_generator_does_not_replace_resumed_thread_id() { .expect("shut down resumed thread"); } +#[test_case::test_case(false; "local_child")] +#[test_case::test_case(true; "host_child")] +#[tokio::test] +async fn child_session_inherits_resumed_parent_auth(host_child: bool) { + let home = tempdir().expect("temporary home"); + let mut config = test_config().await; + config.codex_home = home.path().join("codex-home").abs(); + config.cwd = config.codex_home.abs(); + std::fs::create_dir_all(&config.codex_home).expect("create codex home"); + let manager = ThreadManager::with_models_provider_and_home_for_tests( + CodexAuth::from_api_key("key-a"), + config.model_provider.clone(), + config.codex_home.to_path_buf(), + Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + ); + let original = manager + .start_thread(StartThreadOptions { + history_mode: Some(ThreadHistoryMode::Legacy), + ..StartThreadOptions::new(config.clone()) + }) + .await + .expect("start original thread"); + original.thread.ensure_rollout_materialized().await; + original + .thread + .flush_rollout() + .await + .expect("flush original history"); + original + .thread + .shutdown_and_wait() + .await + .expect("stop original thread"); + manager.remove_thread(&original.thread_id).await; + let replacement = AuthManager::from_auth_for_testing_with_home( + CodexAuth::from_api_key("key-b"), + config.codex_home.to_path_buf(), + ); + let parent = manager + .resume_thread_with_history( + config.clone(), + InitialHistory::Resumed(ResumedHistory { + history_revision: None, + conversation_id: original.thread_id, + history: Arc::new(Vec::new()), + rollout_path: original.thread.rollout_path(), + }), + Arc::clone(&replacement), + /*parent_trace*/ None, + ClientMcpExtensions::default(), + ) + .await + .expect("resume with replacement credentials"); + let options = StartThreadOptions { + reserved_thread_id: host_child.then(|| manager.reserve_thread_id()), + session_source: Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id: parent.thread_id, + depth: 1, + agent_path: None, + agent_nickname: None, + agent_role: None, + })), + ..StartThreadOptions::new(config) + }; + let control = parent + .thread + .session + .services + .local_agent_runtime + .control(parent.thread.session.session_id()); + let child = if host_child { + manager + .state + .spawn_host_child_thread( + options, + control.into(), + parent.thread_id, + crate::agent::host::PreparedChildInheritance { + instructions: parent.thread.session.inherited_instructions().await, + exec_policy: None, + originator: parent.thread.session.originator().await, + attachment_source: None, + }, + ) + .await + } else { + manager + .state + .spawn_child_thread( + options, + control, + Some(parent.thread_id), + /*inherited_exec_policy*/ None, + ) + .await + } + .expect("spawn child under the resumed parent"); + + assert!( + Arc::ptr_eq(&child.thread.session.services.auth_manager, &replacement), + "child must use the resumed parent's credentials", + ); + assert!( + Arc::ptr_eq( + &child.thread.session.services.runtime_auth, + &parent.thread.session.services.runtime_auth, + ), + "child must share the resumed parent's runtime owner", + ); + assert!( + !Arc::ptr_eq( + &child.thread.session.services.runtime_auth, + &manager.state.runtime_auth + ), + "replacement-auth descendants must stay separate from the harness runtime", + ); + let report = manager + .shutdown_all_threads_bounded(Duration::from_secs(10)) + .await; + assert_eq!(report.completed.len(), 2); +} + #[tokio::test] async fn child_session_inherits_client_mcp_extensions() { let temp_dir = tempdir().expect("tempdir"); @@ -1250,8 +1373,12 @@ async fn mcp_invalidation_refreshes_threads_that_are_still_starting() { let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("dummy")); let manager = Arc::new(ThreadManager::new( &config, - Arc::clone(&auth_manager), - build_models_manager(&config, auth_manager), + crate::build_startup_models( + &config, + Arc::clone(&auth_manager), + TEST_INSTALLATION_ID.to_string(), + Arc::new(ModelCatalogSettings::disk(config.codex_home.to_path_buf())), + ), crate::CodexAppsToolsCache::default(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), @@ -1261,7 +1388,6 @@ async fn mcp_invalidation_refreshes_threads_that_are_still_starting() { passthrough_image_store(), ThreadStoreBuilder::from_config(&config, /*state_db*/ None).build(), /*agent_graph_store*/ None, - TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, /*external_time_provider*/ None, )); @@ -1918,8 +2044,12 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); let manager = ThreadManager::new( &config, - auth_manager.clone(), - build_models_manager(&config, auth_manager), + crate::build_startup_models( + &config, + auth_manager.clone(), + TEST_INSTALLATION_ID.to_string(), + Arc::new(ModelCatalogSettings::disk(config.codex_home.to_path_buf())), + ), crate::CodexAppsToolsCache::default(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), @@ -1929,7 +2059,6 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() passthrough_image_store(), ThreadStoreBuilder::from_config(&config, /*state_db*/ None).build(), /*agent_graph_store*/ None, - TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, /*external_time_provider*/ None, ); @@ -2365,8 +2494,12 @@ async fn resume_and_fork_do_not_restore_thread_environments_from_rollout() { AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); let manager = ThreadManager::new( &config, - auth_manager.clone(), - build_models_manager(&config, auth_manager.clone()), + crate::build_startup_models( + &config, + auth_manager.clone(), + TEST_INSTALLATION_ID.to_string(), + Arc::new(ModelCatalogSettings::disk(config.codex_home.to_path_buf())), + ), crate::CodexAppsToolsCache::default(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), @@ -2376,7 +2509,6 @@ async fn resume_and_fork_do_not_restore_thread_environments_from_rollout() { passthrough_image_store(), ThreadStoreBuilder::from_config(&config, /*state_db*/ None).build(), /*agent_graph_store*/ None, - TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, /*external_time_provider*/ None, ); @@ -2525,8 +2657,12 @@ async fn explicit_installation_id_skips_codex_home_file() { let thread_store = ThreadStoreBuilder::from_config(&config, state_db.clone()).build(); let manager = ThreadManager::new( &config, - auth_manager.clone(), - build_models_manager(&config, auth_manager), + crate::build_startup_models( + &config, + auth_manager.clone(), + installation_id.clone(), + Arc::new(ModelCatalogSettings::disk(config.codex_home.to_path_buf())), + ), crate::CodexAppsToolsCache::default(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), @@ -2536,7 +2672,6 @@ async fn explicit_installation_id_skips_codex_home_file() { passthrough_image_store(), thread_store, local_agent_graph_store_from_state_db(state_db.as_ref()), - installation_id.clone(), /*attestation_provider*/ None, /*external_time_provider*/ None, ); @@ -2569,8 +2704,12 @@ async fn resume_active_thread_from_rollout_returns_running_thread() { AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); let manager = ThreadManager::new( &config, - auth_manager.clone(), - build_models_manager(&config, auth_manager.clone()), + crate::build_startup_models( + &config, + auth_manager.clone(), + TEST_INSTALLATION_ID.to_string(), + Arc::new(ModelCatalogSettings::disk(config.codex_home.to_path_buf())), + ), crate::CodexAppsToolsCache::default(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), @@ -2580,7 +2719,6 @@ async fn resume_active_thread_from_rollout_returns_running_thread() { passthrough_image_store(), ThreadStoreBuilder::from_config(&config, /*state_db*/ None).build(), /*agent_graph_store*/ None, - TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, /*external_time_provider*/ None, ); @@ -2635,8 +2773,12 @@ async fn resume_stopped_thread_from_rollout_spawns_new_thread() { AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); let manager = ThreadManager::new( &config, - auth_manager.clone(), - build_models_manager(&config, auth_manager.clone()), + crate::build_startup_models( + &config, + auth_manager.clone(), + TEST_INSTALLATION_ID.to_string(), + Arc::new(ModelCatalogSettings::disk(config.codex_home.to_path_buf())), + ), crate::CodexAppsToolsCache::default(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), @@ -2646,7 +2788,6 @@ async fn resume_stopped_thread_from_rollout_spawns_new_thread() { passthrough_image_store(), ThreadStoreBuilder::from_config(&config, /*state_db*/ None).build(), /*agent_graph_store*/ None, - TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, /*external_time_provider*/ None, ); @@ -2708,8 +2849,12 @@ async fn resume_stopped_thread_from_rollout_preserves_thread_source() { let thread_store = ThreadStoreBuilder::from_config(&config, state_db.clone()).build(); let manager = ThreadManager::new( &config, - auth_manager.clone(), - build_models_manager(&config, auth_manager.clone()), + crate::build_startup_models( + &config, + auth_manager.clone(), + TEST_INSTALLATION_ID.to_string(), + Arc::new(ModelCatalogSettings::disk(config.codex_home.to_path_buf())), + ), crate::CodexAppsToolsCache::default(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), @@ -2719,7 +2864,6 @@ async fn resume_stopped_thread_from_rollout_preserves_thread_source() { passthrough_image_store(), thread_store, local_agent_graph_store_from_state_db(state_db.as_ref()), - TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, /*external_time_provider*/ None, ); @@ -2795,8 +2939,12 @@ async fn subtree_listing_uses_injected_graph_store_without_state_db() { AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); let manager = ThreadManager::new( &config, - auth_manager.clone(), - build_models_manager(&config, auth_manager), + crate::build_startup_models( + &config, + auth_manager.clone(), + TEST_INSTALLATION_ID.to_string(), + Arc::new(ModelCatalogSettings::disk(config.codex_home.to_path_buf())), + ), crate::CodexAppsToolsCache::default(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), @@ -2806,7 +2954,6 @@ async fn subtree_listing_uses_injected_graph_store_without_state_db() { passthrough_image_store(), ThreadStoreBuilder::from_config(&config, /*state_db*/ None).build(), Some(agent_graph_store), - TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, /*external_time_provider*/ None, ); @@ -2843,8 +2990,12 @@ async fn rollout_path_resume_and_fork_read_history_through_thread_store() { .expect("configured in-memory store"); let manager = ThreadManager::new( &config, - auth_manager.clone(), - build_models_manager(&config, auth_manager.clone()), + crate::build_startup_models( + &config, + auth_manager.clone(), + TEST_INSTALLATION_ID.to_string(), + Arc::new(ModelCatalogSettings::disk(config.codex_home.to_path_buf())), + ), crate::CodexAppsToolsCache::default(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), @@ -2854,7 +3005,6 @@ async fn rollout_path_resume_and_fork_read_history_through_thread_store() { passthrough_image_store(), thread_store.clone(), local_agent_graph_store_from_state_db(state_db.as_ref()), - TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, /*external_time_provider*/ None, ); @@ -2953,8 +3103,12 @@ async fn metadata_update_without_result_reads_only_when_the_caller_needs_the_thr .expect("configured in-memory store"); let manager = ThreadManager::new( &config, - auth_manager.clone(), - build_models_manager(&config, auth_manager), + crate::build_startup_models( + &config, + auth_manager.clone(), + TEST_INSTALLATION_ID.to_string(), + Arc::new(ModelCatalogSettings::disk(config.codex_home.to_path_buf())), + ), crate::CodexAppsToolsCache::default(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), @@ -2964,7 +3118,6 @@ async fn metadata_update_without_result_reads_only_when_the_caller_needs_the_thr passthrough_image_store(), thread_store.clone(), /*agent_graph_store*/ None, - TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, /*external_time_provider*/ None, ); @@ -3080,8 +3233,12 @@ async fn new_uses_active_provider_for_model_refresh() { AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); let manager = ThreadManager::new( &config, - auth_manager.clone(), - build_models_manager(&config, auth_manager), + crate::build_startup_models( + &config, + auth_manager.clone(), + TEST_INSTALLATION_ID.to_string(), + Arc::new(ModelCatalogSettings::disk(config.codex_home.to_path_buf())), + ), crate::CodexAppsToolsCache::default(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), @@ -3091,7 +3248,6 @@ async fn new_uses_active_provider_for_model_refresh() { passthrough_image_store(), ThreadStoreBuilder::from_config(&config, /*state_db*/ None).build(), /*agent_graph_store*/ None, - TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, /*external_time_provider*/ None, ); @@ -3157,7 +3313,7 @@ async fn new_uses_active_provider_for_model_refresh() { } #[tokio::test] -async fn injected_models_manager_controls_refresh_policy() { +async fn startup_catalog_cache_policy_controls_refresh() { let server = MockServer::start().await; let _ = mount_models_once(&server, ModelsResponse { models: vec![] }).await; let _ = mount_models_once(&server, ModelsResponse { models: vec![] }).await; @@ -3172,15 +3328,14 @@ async fn injected_models_manager_controls_refresh_policy() { let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); - let provider = create_model_provider( - config.model_provider.clone(), - Some(Arc::clone(&auth_manager)), - ); - let models_manager = provider.models_manager_without_cache(config.model_catalog.clone()); let manager = ThreadManager::new( &config, - auth_manager, - models_manager, + crate::build_startup_models( + &config, + auth_manager, + TEST_INSTALLATION_ID.to_string(), + Arc::new(ModelCatalogSettings::default()), + ), crate::CodexAppsToolsCache::default(), SessionSource::Custom("test-embedder".to_string()), Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), @@ -3190,7 +3345,6 @@ async fn injected_models_manager_controls_refresh_policy() { passthrough_image_store(), ThreadStoreBuilder::from_config(&config, /*state_db*/ None).build(), /*agent_graph_store*/ None, - TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, /*external_time_provider*/ None, ); @@ -3460,8 +3614,12 @@ async fn interrupted_fork_snapshot_does_not_synthesize_turn_id_for_legacy_histor let state_db = init_state_db(&config).await; let manager = ThreadManager::new( &config, - auth_manager.clone(), - build_models_manager(&config, auth_manager.clone()), + crate::build_startup_models( + &config, + auth_manager.clone(), + TEST_INSTALLATION_ID.to_string(), + Arc::new(ModelCatalogSettings::disk(config.codex_home.to_path_buf())), + ), crate::CodexAppsToolsCache::default(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), @@ -3471,7 +3629,6 @@ async fn interrupted_fork_snapshot_does_not_synthesize_turn_id_for_legacy_histor passthrough_image_store(), ThreadStoreBuilder::from_config(&config, state_db.clone()).build(), local_agent_graph_store_from_state_db(state_db.as_ref()), - TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, /*external_time_provider*/ None, ); @@ -3572,8 +3729,12 @@ async fn interrupted_fork_snapshot_preserves_explicit_turn_id() { let state_db = init_state_db(&config).await; let manager = ThreadManager::new( &config, - auth_manager.clone(), - build_models_manager(&config, auth_manager.clone()), + crate::build_startup_models( + &config, + auth_manager.clone(), + TEST_INSTALLATION_ID.to_string(), + Arc::new(ModelCatalogSettings::disk(config.codex_home.to_path_buf())), + ), crate::CodexAppsToolsCache::default(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), @@ -3583,7 +3744,6 @@ async fn interrupted_fork_snapshot_preserves_explicit_turn_id() { passthrough_image_store(), ThreadStoreBuilder::from_config(&config, state_db.clone()).build(), local_agent_graph_store_from_state_db(state_db.as_ref()), - TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, /*external_time_provider*/ None, ); @@ -3677,8 +3837,12 @@ async fn interrupted_fork_snapshot_uses_persisted_mid_turn_history_without_live_ let state_db = init_state_db(&config).await; let manager = ThreadManager::new( &config, - auth_manager.clone(), - build_models_manager(&config, auth_manager.clone()), + crate::build_startup_models( + &config, + auth_manager.clone(), + TEST_INSTALLATION_ID.to_string(), + Arc::new(ModelCatalogSettings::disk(config.codex_home.to_path_buf())), + ), crate::CodexAppsToolsCache::default(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), @@ -3688,7 +3852,6 @@ async fn interrupted_fork_snapshot_uses_persisted_mid_turn_history_without_live_ passthrough_image_store(), ThreadStoreBuilder::from_config(&config, state_db.clone()).build(), local_agent_graph_store_from_state_db(state_db.as_ref()), - TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, /*external_time_provider*/ None, ); diff --git a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs index 4ee569c712..254f603029 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -37,6 +37,7 @@ use codex_login::AuthManager; use codex_login::CodexAuth; use codex_model_provider::create_model_provider; use codex_model_provider_info::built_in_model_providers; +use codex_models_manager::manager::ModelCatalogSettings; use codex_models_manager::manager::StaticModelsManager; use codex_otel::MULTI_AGENT_SPAWN_FAILURE_METRIC; use codex_otel::MetricsClient; @@ -4319,8 +4320,12 @@ async fn tool_handlers_cascade_close_and_resume_and_keep_explicitly_closed_subtr let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("dummy")); let manager = ThreadManager::new( &config, - auth_manager.clone(), - crate::thread_manager::build_models_manager(&config, auth_manager), + crate::build_startup_models( + &config, + auth_manager.clone(), + "11111111-1111-4111-8111-111111111111".to_string(), + Arc::new(ModelCatalogSettings::disk(config.codex_home.to_path_buf())), + ), crate::CodexAppsToolsCache::default(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), @@ -4330,7 +4335,6 @@ async fn tool_handlers_cascade_close_and_resume_and_keep_explicitly_closed_subtr crate::thread_manager::passthrough_image_store(), ThreadStoreBuilder::from_config(&config, state_db.clone()).build(), local_agent_graph_store_from_state_db(state_db.as_ref()), - "11111111-1111-4111-8111-111111111111".to_string(), /*attestation_provider*/ None, /*external_time_provider*/ None, ); diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index 25a09467ec..ea01655918 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -44,6 +44,7 @@ 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::manager::ModelCatalogSettings; use codex_models_manager::manager::SharedModelsManager; use codex_protocol::config_types::CollaborationMode; use codex_protocol::config_types::ModeKind; @@ -424,6 +425,7 @@ pub struct TestCodexBuilder { code_mode_host_program: Option, history_mode: Option, models_manager: Option, + model_catalog_settings: Option, thread_store: Option>, agent_graph_store: Option>, image_store: Arc, @@ -479,6 +481,11 @@ impl TestCodexBuilder { self } + pub fn with_model_catalog_settings(mut self, settings: ModelCatalogSettings) -> Self { + self.model_catalog_settings = Some(settings); + self + } + pub fn with_models_manager(mut self, models_manager: SharedModelsManager) -> Self { self.models_manager = Some(models_manager); self @@ -863,14 +870,23 @@ impl TestCodexBuilder { )) }); let auth_manager = self.auth.manager_for_home(config.codex_home.as_path()); - let models_manager = self - .models_manager - .clone() - .unwrap_or_else(|| codex_core::build_models_manager(&config, auth_manager.clone())); let code_mode_host_program = self .code_mode_host_program .take() .or_else(|| codex_utils_cargo_bin::cargo_bin("codex-code-mode-host").ok()); + let factory = + match &self.models_manager { + Some(models) => codex_core::test_support::fixed_model_catalog(Arc::clone(models)), + None => Arc::new(self.model_catalog_settings.clone().unwrap_or_else(|| { + ModelCatalogSettings::disk(config.codex_home.to_path_buf()) + })), + }; + let startup = codex_core::build_startup_models( + &config, + Arc::clone(&auth_manager), + installation_id, + factory, + ); let thread_manager = Arc::new_cyclic(|manager| { let mut extensions = self.extensions.to_builder(); codex_app_server::install_agent_message_board(&mut extensions, manager.clone()); @@ -881,8 +897,7 @@ impl TestCodexBuilder { } let thread_manager = ThreadManager::new( &config, - auth_manager.clone(), - models_manager, + startup, self.apps_tools_cache.clone(), SessionSource::Exec, Arc::clone(&environment_manager), @@ -894,7 +909,6 @@ impl TestCodexBuilder { self.agent_graph_store.clone().or_else(|| { codex_core::local_agent_graph_store_from_state_db(state_db.as_ref()) }), - installation_id, /*attestation_provider*/ None, /*external_time_provider*/ self.external_time_provider.clone(), ); @@ -1573,6 +1587,7 @@ pub fn test_codex() -> TestCodexBuilder { // These fixtures exercise legacy-only resume/fork helpers; store-default tests opt out. history_mode: Some(ThreadHistoryMode::Legacy), models_manager: None, + model_catalog_settings: None, thread_store: None, agent_graph_store: None, image_store: codex_core::passthrough_image_store(), diff --git a/codex-rs/core/tests/responses_headers.rs b/codex-rs/core/tests/responses_headers.rs index ef0102393a..bcb28d6f1c 100644 --- a/codex-rs/core/tests/responses_headers.rs +++ b/codex-rs/core/tests/responses_headers.rs @@ -129,10 +129,9 @@ async fn responses_stream_includes_subagent_header_on_review() { ); let client = ModelClient::new( - /*auth_manager*/ None, + codex_model_provider::create_model_provider(provider.clone(), /*auth_manager*/ None), AgentIdentityAuthPolicy::JwtOnly, thread_id, - provider.clone(), session_source.clone(), "test_originator".to_string(), config.model_verbosity, @@ -279,10 +278,9 @@ async fn responses_stream_includes_subagent_header_on_other() { ); let client = ModelClient::new( - /*auth_manager*/ None, + codex_model_provider::create_model_provider(provider.clone(), /*auth_manager*/ None), AgentIdentityAuthPolicy::JwtOnly, thread_id, - provider.clone(), session_source.clone(), "test_originator".to_string(), config.model_verbosity, @@ -414,10 +412,9 @@ async fn responses_respects_model_info_overrides_from_config() { ); let client = ModelClient::new( - /*auth_manager*/ None, + codex_model_provider::create_model_provider(provider.clone(), /*auth_manager*/ None), AgentIdentityAuthPolicy::JwtOnly, thread_id, - provider.clone(), session_source.clone(), "test_originator".to_string(), config.model_verbosity, diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index 6967284e64..8455498749 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -26,6 +26,7 @@ use codex_model_provider_info::ModelProviderInfo; use codex_model_provider_info::WireApi; use codex_model_provider_info::built_in_model_providers; use codex_models_manager::bundled_models_response; +use codex_models_manager::manager::ModelCatalogSettings; use codex_otel::SessionTelemetry; use codex_otel::TelemetryAuthMode; use codex_protocol::ResponseItemId; @@ -1823,12 +1824,14 @@ async fn send_request_with_provider(provider: ModelProviderInfo) { SessionSource::Exec, ); let client = ModelClient::new( - Some(AuthManager::from_auth_for_testing(CodexAuth::from_api_key( - "unused-api-key", - ))), + codex_model_provider::create_model_provider( + provider, + Some(AuthManager::from_auth_for_testing(CodexAuth::from_api_key( + "unused-api-key", + ))), + ), AgentIdentityAuthPolicy::JwtOnly, thread_id, - provider, SessionSource::Exec, "test_originator".to_string(), config.model_verbosity, @@ -2074,8 +2077,12 @@ async fn prefers_apikey_when_config_prefers_apikey_even_with_chatgpt_tokens() { .expect("resolve installation id"); let thread_manager = ThreadManager::new( &config, - auth_manager.clone(), - codex_core::build_models_manager(&config, auth_manager), + codex_core::build_startup_models( + &config, + auth_manager.clone(), + installation_id, + Arc::new(ModelCatalogSettings::disk(config.codex_home.to_path_buf())), + ), codex_core::CodexAppsToolsCache::default(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), @@ -2085,7 +2092,6 @@ async fn prefers_apikey_when_config_prefers_apikey_even_with_chatgpt_tokens() { codex_core::passthrough_image_store(), ThreadStoreBuilder::from_config(&config, /*state_db*/ None).build(), /*agent_graph_store*/ None, - installation_id, /*attestation_provider*/ None, /*external_time_provider*/ None, ); @@ -3329,10 +3335,9 @@ async fn azure_responses_request_does_not_store_and_preserves_prefixed_item_ids( ); let client = ModelClient::new( - /*auth_manager*/ None, + codex_model_provider::create_model_provider(provider.clone(), /*auth_manager*/ None), AgentIdentityAuthPolicy::JwtOnly, thread_id, - provider.clone(), SessionSource::Exec, "test_originator".to_string(), config.model_verbosity, diff --git a/codex-rs/core/tests/suite/client_websockets.rs b/codex-rs/core/tests/suite/client_websockets.rs index eea77b19ee..cf0b7e1448 100755 --- a/codex-rs/core/tests/suite/client_websockets.rs +++ b/codex-rs/core/tests/suite/client_websockets.rs @@ -2827,6 +2827,28 @@ async fn websocket_harness_with_provider_options_and_auth( concurrent_reasoning_summaries_enabled: bool, enabled_features: &[Feature], auth: Option, +) -> WebsocketTestHarness { + websocket_harness_with_provider_factory( + provider, + runtime_metrics_enabled, + concurrent_reasoning_summaries_enabled, + enabled_features, + auth, + |provider, auth| codex_model_provider::create_model_provider(provider, auth), + ) + .await +} + +async fn websocket_harness_with_provider_factory( + provider: ModelProviderInfo, + runtime_metrics_enabled: bool, + concurrent_reasoning_summaries_enabled: bool, + enabled_features: &[Feature], + auth: Option, + provider_factory: impl FnOnce( + ModelProviderInfo, + Option>, + ) -> codex_model_provider::SharedModelProvider, ) -> WebsocketTestHarness { let codex_home = TempDir::new().unwrap(); let mut config = load_default_config_for_test(&codex_home).await; @@ -2898,10 +2920,9 @@ async fn websocket_harness_with_provider_options_and_auth( ) }); let client = ModelClient::new( - client_auth_manager, + provider_factory(provider.clone(), client_auth_manager), AgentIdentityAuthPolicy::JwtOnly, thread_id, - provider.clone(), SessionSource::Exec, "test_originator".to_string(), config.model_verbosity, diff --git a/codex-rs/core/tests/suite/client_websockets_runtime_auth.rs b/codex-rs/core/tests/suite/client_websockets_runtime_auth.rs index 5d8408b629..53041d393e 100644 --- a/codex-rs/core/tests/suite/client_websockets_runtime_auth.rs +++ b/codex-rs/core/tests/suite/client_websockets_runtime_auth.rs @@ -58,12 +58,33 @@ async fn runtime_harness( ) -> (WebsocketTestHarness, Arc) { provider.requires_openai_auth = true; let initial = CodexAuth::from_api_key("key-a"); - let mut harness = websocket_harness_with_provider_options_and_auth( + let registration_url = format!("{}/register", registration.uri()).parse().unwrap(); + let harness = websocket_harness_with_provider_factory( provider, /*runtime_metrics_enabled*/ false, /*concurrent_reasoning_summaries_enabled*/ false, &[], Some(initial.clone()), + move |provider, credentials| { + let credentials = credentials.expect("runtime fixture requires credentials"); + let runtime = Arc::new(RuntimeAuthManager::new( + Arc::clone(&credentials), + AgentHostAuthorization::InstallationId { + installation_id: Uuid::nil(), + }, + RuntimeRegistrationMetadata { + harness: Some("codex".into()), + harness_version: None, + role: RuntimeRole::OrchestratorAndExecutor, + }, + registration_url, + )); + codex_model_provider::test_support::runtime_model_provider( + provider, + credentials, + (runtime, mode), + ) + }, ) .await; let credentials = Arc::new(Credentials(RwLock::new(initial))); @@ -72,19 +93,6 @@ async fn runtime_harness( .set_external_auth(credentials.clone()) .await .unwrap(); - let runtime = Arc::new(RuntimeAuthManager::new( - harness.auth_manager.clone(), - AgentHostAuthorization::InstallationId { - installation_id: Uuid::nil(), - }, - RuntimeRegistrationMetadata { - harness: Some("codex".into()), - harness_version: None, - role: RuntimeRole::OrchestratorAndExecutor, - }, - format!("{}/register", registration.uri()).parse().unwrap(), - )); - harness.client = harness.client.with_runtime_auth(runtime, mode).unwrap(); (harness, credentials) } diff --git a/codex-rs/core/tests/suite/injected_models_cache.rs b/codex-rs/core/tests/suite/injected_models_cache.rs index e206437c90..e287ec0dcb 100644 --- a/codex-rs/core/tests/suite/injected_models_cache.rs +++ b/codex-rs/core/tests/suite/injected_models_cache.rs @@ -1,5 +1,6 @@ //! Integration coverage for model catalogs supplied through the public cache interface. +use codex_core::StartThreadOptions; use codex_core::TurnInputRequest; use std::sync::Arc; use std::sync::Mutex; @@ -12,12 +13,12 @@ use chrono::Utc; use codex_http_client::HttpClientFactory; use codex_login::AuthManager; use codex_login::CodexAuth; -use codex_model_provider::create_model_provider; use codex_model_provider_info::built_in_model_providers; 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::ModelsEndpointClient; use codex_models_manager::manager::ModelsEndpointFuture; use codex_models_manager::manager::ModelsEndpointResponse; @@ -44,6 +45,7 @@ use core_test_support::test_codex::test_codex; use core_test_support::wait_for_event; use pretty_assertions::assert_eq; use tempfile::TempDir; +use test_case::test_case; use tokio::sync::oneshot; use tokio::time::Duration; use tokio::time::timeout; @@ -202,7 +204,15 @@ async fn run_agent_with_model(models_manager: SharedModelsManager, model_slug: & .await; let mut builder = test_codex() .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) - .with_models_manager(models_manager); + .with_models_manager(models_manager) + .with_config(|config| { + config + .features + .enable(codex_features::Feature::RuntimeAgentInstances) + .expect("test config should allow enabling runtime auth"); + // Model inference uses a third-party fixture; a custom ChatGPT base must not block startup. + config.chatgpt_base_url = "https://custom.example.test/backend-api".to_string(); + }); let test = builder.build(&server).await?; let available_models = test .thread_manager @@ -245,6 +255,7 @@ async fn run_agent_with_model(models_manager: SharedModelsManager, model_slug: & response_mock.single_request().body_json()["model"], model_slug ); + assert!(!test.codex_home_path().join("models_cache.json").exists()); Ok(()) } @@ -295,6 +306,124 @@ async fn injected_cache_error_falls_back_for_agent_model_selection() -> Result<( Ok(()) } +#[test_case(false; "without_cache")] +#[test_case(true; "with_host_cache")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn threads_follow_provider_changes_and_preserve_host_cache_policy( + use_cache: bool, +) -> Result<()> { + let host = responses::start_mock_server().await; + let inference = responses::start_mock_server().await; + let model = "host-catalog-model"; + wiremock::Mock::given(wiremock::matchers::method("GET")) + .and(wiremock::matchers::path("/host-models")) + .respond_with( + wiremock::ResponseTemplate::new(200).set_body_json(ModelsResponse { + models: vec![remote_model(model)], + }), + ) + .expect(1..) + .mount(&host) + .await; + let auth = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("host-catalog-key")); + let mut info = + built_in_model_providers(Some(format!("{}/v1", inference.uri())))["openai"].clone(); + info.model_catalog_url = Some(format!("{}/host-models", host.uri()).into()); + let cache = Arc::new(TestModelsCache { + entry: None, + load_error: false, + stored_entries: Mutex::default(), + store_gate: Mutex::default(), + }); + let settings = if use_cache { + ModelCatalogSettings::with_cache(cache.clone()) + } else { + ModelCatalogSettings::default() + }; + settings.set_api_key_model_discovery_enabled(/*enabled*/ true); + let test = test_codex() + .with_auth_manager(auth) + .with_model_catalog_settings(settings) + .with_config(move |config| { + config.model_provider = info; + config + .features + .enable(codex_features::Feature::RuntimeAgentInstances) + .expect("runtime identity feature should be available"); + config.model = None; + }) + .build_with_auto_env(&inference) + .await?; + assert_eq!(test.session_configured.model, model); + let mut disabled = test.config.clone(); + disabled + .features + .disable(codex_features::Feature::RuntimeAgentInstances)?; + let sibling = test + .thread_manager + .start_thread(StartThreadOptions::new(disabled)) + .await?; + assert_eq!(sibling.session_configured.model, model); + let second_host = responses::start_mock_server().await; + let second_inference = responses::start_mock_server().await; + let second_model = "second-provider-model"; + wiremock::Mock::given(wiremock::matchers::method("GET")) + .and(wiremock::matchers::path("/second-models")) + .and(wiremock::matchers::header( + "authorization", + "Bearer host-catalog-key", + )) + .respond_with( + wiremock::ResponseTemplate::new(200).set_body_json(ModelsResponse { + models: vec![remote_model(second_model)], + }), + ) + .expect(1..) + .mount(&second_host) + .await; + let mut changed = test.config.clone(); + changed.model_provider.base_url = Some(format!("{}/v1", second_inference.uri())); + changed.model_provider.model_catalog_url = + Some(format!("{}/second-models", second_host.uri()).into()); + let changed_thread = test + .thread_manager + .start_thread(StartThreadOptions::new(changed)) + .await?; + assert_eq!(changed_thread.session_configured.model, second_model); + if use_cache { + assert_eq!( + cache + .stored_entries + .lock() + .expect("stored entries lock should not be poisoned") + .last() + .expect("provider catalog should be cached") + .models, + vec![remote_model(second_model)], + ); + } + for (server, path) in [(&host, "/host-models"), (&second_host, "/second-models")] { + for request in server + .received_requests() + .await + .expect("mock requests should be recorded") + { + assert_eq!(request.url.path(), path); + assert_eq!(request.headers["authorization"], "Bearer host-catalog-key"); + } + } + assert_eq!( + !cache + .stored_entries + .lock() + .expect("stored entries lock should not be poisoned") + .is_empty(), + use_cache + ); + assert!(!test.codex_home_path().join("models_cache.json").exists()); + 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; @@ -313,8 +442,6 @@ async fn account_switch_during_cache_store_preserves_new_catalog_for_next_turn() stored_entries: Mutex::new(Vec::new()), store_gate: Mutex::new(None), }); - let manager = create_model_provider(provider_info.clone(), Some(auth.clone())) - .models_manager_with_cache(/*config_model_catalog*/ None, cache.clone()); let mut old_model = remote_model("gpt-5.4"); old_model.default_reasoning_summary = ReasoningSummary::Concise; let initial = responses::mount_models_once( @@ -326,15 +453,16 @@ async fn account_switch_during_cache_store_preserves_new_catalog_for_next_turn() .await; let test = test_codex() .with_home(home.clone()) - .with_auth(initial_auth.clone()) + .with_auth_manager(Arc::clone(&auth)) .with_model("gpt-5.4") - .with_models_manager(manager.clone()) + .with_model_catalog_settings(ModelCatalogSettings::with_cache(cache.clone())) .with_config(move |config| { config.model_provider = provider_info; config.model_reasoning_summary = None; }) .build_with_auto_env(&server) .await?; + let manager = test.thread_manager.get_models_manager(); assert_eq!(initial.requests().len(), 1); // Pause after the old catalog has been stored, before its refresh can publish it. @@ -379,7 +507,6 @@ async fn account_switch_during_cache_store_preserves_new_catalog_for_next_turn() }))?, )?; auth.reload().await; - test.thread_manager.auth_manager().reload().await; let mut new_model = remote_model("gpt-5.4"); new_model.default_reasoning_summary = ReasoningSummary::Detailed; diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index d78151c35b..70d5f144a1 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/codex-rs/core/tests/suite/mod.rs @@ -208,6 +208,7 @@ mod rmcp_client; mod rollout_budget; mod rollout_compression; mod rollout_list_find; +mod runtime_auth; mod safety_buffering; mod safety_check_downgrade; mod scenarios; diff --git a/codex-rs/core/tests/suite/runtime_auth.rs b/codex-rs/core/tests/suite/runtime_auth.rs new file mode 100644 index 0000000000..65a831ca7c --- /dev/null +++ b/codex-rs/core/tests/suite/runtime_auth.rs @@ -0,0 +1,650 @@ +//! Exercise harness-owned auth through real sessions, subagents, and delegated review. + +use anyhow::Context; +use anyhow::Result; +use base64::Engine; +use codex_core::CodexThread; +use codex_core::StartThreadOptions; +use codex_core::TurnInputRequest; +use codex_features::Feature; +use codex_features::RuntimeAuthMode; +use codex_history::InitialHistory; +use codex_history::ResumedHistory; +use codex_login::CodexAuth; +use codex_model_provider::ProviderAuthContext; +use codex_models_manager::manager::RefreshStrategy; +use codex_protocol::mcp::ClientMcpExtensions; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::InternalSessionSource; +use codex_protocol::protocol::Op; +use codex_protocol::protocol::ReviewRequest; +use codex_protocol::protocol::ReviewTarget; +use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::SubAgentSource; +use codex_protocol::user_input::UserInput; +use core_test_support::responses; +use core_test_support::skip_if_no_network; +use core_test_support::test_codex::test_codex; +use core_test_support::wait_for_event; +use pretty_assertions::assert_eq; +use serde_json::Value; +use serde_json::json; +use std::collections::BTreeSet; +use std::sync::Arc; +use std::time::Duration; +use tokio::io::AsyncReadExt; +use tokio::io::AsyncWriteExt; +use tokio::net::TcpListener; +use tokio::net::TcpStream; +use tokio::process::Command; +use tokio::task::JoinSet; +use tokio_rustls::TlsAcceptor; +use tokio_rustls::rustls; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::body_string_contains; +use wiremock::matchers::method; +use wiremock::matchers::path; + +const PROBE_SERVER: &str = "CODEX_RUNTIME_AUTH_TEST_SERVER"; +const TEST_NAME: &str = "suite::runtime_auth::harness_shares_runtime_and_honors_thread_auth_config"; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn harness_shares_runtime_and_honors_thread_auth_config() -> Result<()> { + skip_if_no_network!(Ok(())); + if let Ok(server_uri) = std::env::var(PROBE_SERVER) { + return exercise_harness(server_uri).await; + } + + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/agent/register")) + .respond_with(|request: &wiremock::Request| { + let body: Value = request.body_json().expect("registration body should parse"); + let replacement = body["actor_auth"]["credential"] == "key-b"; + ResponseTemplate::new(200).set_body_json(json!({ + "runtime_agent_instance_id": if replacement { "runtime-replacement" } else { "runtime-shared" }, + "agent_host_id": "host-a", + "actor": {"kind": "platform", "user_id": if replacement { "user-b" } else { "user-a" }, "organization_id": "org-a", "project_id": null}, + })) + }) + .expect(2) + .mount(&server) + .await; + // Both runtime modes retain the normal same-origin redirect behavior. + for label in ["parent-dual", "child-assertion"] { + Mock::given(method("POST")) + .and(path("/v1/responses")) + .and(body_string_contains(label)) + .respond_with( + ResponseTemplate::new(307) + .insert_header("Location", "https://api.openai.com/redirected/responses"), + ) + .up_to_n_times(1) + .expect(1) + .with_priority(1) + .mount(&server) + .await; + } + let responses = responses::mount_response_sequence( + &server, + (0..14) + .map(|i| { + let response = + responses::sse_response(responses::sse(vec![responses::ev_completed( + &format!("resp-{i}"), + )])); + if (1..=4).contains(&i) { + response.insert_header("X-Models-Etag", format!("catalog-refreshed-{i}")) + } else { + response + } + }) + .collect(), + ) + .await; + Mock::given(method("GET")) + .and(path("/catalog-isolation")) + .respond_with(ResponseTemplate::new(200).set_body_json(catalog_fixture())) + .expect(2) + .mount(&server) + .await; + + for mode in ["dual", "assertion", "disabled"] { + Mock::given(method("GET")) + .and(path(format!("/catalog-reload/{mode}"))) + .respond_with( + ResponseTemplate::new(200) + .insert_header("ETag", "catalog-initial") + .set_body_json(catalog_fixture()), + ) + .expect(2) + .mount(&server) + .await; + } + + Mock::given(method("GET")) + .and(path("/resumed-off-models")) + .respond_with(ResponseTemplate::new(200).set_body_json(catalog_fixture())) + .mount(&server) + .await; + + // The production destination check stays intact. A test-only CA and CONNECT proxy route + // the two real host names to wiremock, only inside this isolated test subprocess. + codex_utils_rustls_provider::ensure_rustls_crypto_provider(); + let rcgen::CertifiedKey { cert, signing_key } = rcgen::generate_simple_self_signed(vec![ + "api.openai.com".into(), + "auth.openai.com".into(), + ])?; + let ca_dir = tempfile::tempdir()?; + let ca_path = ca_dir.path().join("ca.pem"); + std::fs::write(&ca_path, cert.pem())?; + let acceptor = TlsAcceptor::from(Arc::new( + rustls::ServerConfig::builder() + .with_no_client_auth() + .with_single_cert(vec![cert.der().clone()], signing_key.into())?, + )); + let proxy = TcpListener::bind("127.0.0.1:0").await?; + let proxy_url = format!("http://{}", proxy.local_addr()?); + let target = *server.address(); + let proxy_task = tokio::spawn(async move { + let mut connections = JoinSet::new(); + loop { + let (mut stream, _) = proxy.accept().await?; + let acceptor = acceptor.clone(); + connections.spawn(async move { + let mut headers = Vec::new(); + while !headers.ends_with(b"\r\n\r\n") && headers.len() < 8192 { + headers.push(stream.read_u8().await?); + } + let headers = std::str::from_utf8(&headers)?; + anyhow::ensure!( + matches!( + headers.lines().next(), + Some( + "CONNECT api.openai.com:443 HTTP/1.1" + | "CONNECT auth.openai.com:443 HTTP/1.1" + ) + ), + "unexpected proxy destination" + ); + stream + .write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n") + .await?; + let mut stream = acceptor.accept(stream).await?; + let mut target = TcpStream::connect(target).await?; + tokio::io::copy_bidirectional(&mut stream, &mut target).await?; + Ok::<_, anyhow::Error>(()) + }); + } + #[allow(unreachable_code)] + Ok::<_, anyhow::Error>(()) + }); + let mut command = Command::new(std::env::current_exe()?); + command + .args(["--exact", TEST_NAME, "--nocapture"]) + .kill_on_drop(true); + for &key in codex_network_proxy::PROXY_ENV_KEYS { + command.env_remove(key); + } + command + .env(PROBE_SERVER, server.uri()) + .env("HTTPS_PROXY", proxy_url) + .env("NO_PROXY", codex_network_proxy::DEFAULT_NO_PROXY_VALUE) + .env("CODEX_CA_CERTIFICATE", ca_path); + let output = tokio::time::timeout(Duration::from_secs(90), command.output()).await; + proxy_task.abort(); + let output = output??; + assert!( + output.status.success(), + "harness subprocess failed\n{}\n{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + + let requests = responses.requests(); + let expected = [ + (Some("startup-disabled"), "Bearer key-a", None), + ( + Some("discovery-off"), + "Bearer key-a", + Some("runtime-shared"), + ), + (Some("catalog-dual"), "Bearer key-a", Some("runtime-shared")), + ( + Some("catalog-assertion"), + "AgentAssertion ", + Some("runtime-shared"), + ), + (Some("catalog-disabled"), "Bearer key-a", None), + (Some("parent-dual"), "Bearer key-a", Some("runtime-shared")), + // Explicit first-party inference with a third-party catalog. + (None, "Bearer key-a", Some("runtime-shared")), + ( + Some("child-assertion"), + "AgentAssertion ", + Some("runtime-shared"), + ), + (None, "AgentAssertion ", Some("runtime-shared")), + // The inline review delegate inherits the parent's dual-header mode. + (None, "Bearer key-a", Some("runtime-shared")), + (Some("thread-disabled"), "Bearer key-a", None), + (Some("third-party"), "Bearer key-a", None), + ( + Some("replacement-root"), + "Bearer key-b", + Some("runtime-replacement"), + ), + ( + Some("replacement-child"), + "Bearer key-b", + Some("runtime-replacement"), + ), + ]; + assert_eq!(requests.len(), expected.len()); + let mut proof_ids = BTreeSet::new(); + for (request, (label, expected_auth, expected_runtime)) in requests.iter().zip(expected) { + if let Some(label) = label { + assert_eq!( + request + .message_input_texts("user") + .last() + .map(String::as_str), + Some(label) + ); + } + let auth = request.header("authorization").unwrap(); + let proof = if expected_auth == "AgentAssertion " { + assert!(request.header("x-openai-agent-assertion").is_none()); + Some( + auth.strip_prefix(expected_auth) + .context("assertion-only request")? + .to_owned(), + ) + } else { + assert_eq!(auth, expected_auth); + request.header("x-openai-agent-assertion") + }; + assert_eq!(proof.is_some(), expected_runtime.is_some(), "{label:?}"); + if let Some(proof) = proof { + let payload: Value = serde_json::from_slice( + &base64::engine::general_purpose::URL_SAFE_NO_PAD + .decode(proof.split('.').nth(2).unwrap())?, + )?; + assert_eq!(payload["runtime_agent_instance"], expected_runtime.unwrap()); + assert!(proof_ids.insert(payload["id"].as_str().unwrap().to_owned())); + } + } + let resumed_catalog: Vec<_> = server + .received_requests() + .await + .unwrap() + .into_iter() + .filter(|request| request.url.path() == "/resumed-off-models") + .collect(); + assert!( + !resumed_catalog.is_empty(), + "resumed thread must fetch its catalog" + ); + for request in resumed_catalog { + assert_eq!(request.headers["authorization"], "Bearer key-b"); + assert!(!request.headers.contains_key("openai-agent-assertion")); + } + let catalogs: Vec<_> = server + .received_requests() + .await + .unwrap() + .into_iter() + .filter(|request| request.url.path() == "/catalog-isolation") + .collect(); + assert_eq!(catalogs.len(), 2); + for request in catalogs { + assert!( + request.headers.contains_key("x-openai-agent-assertion") + || request.headers["authorization"] + .to_str()? + .starts_with("AgentAssertion ") + ); + } + let requests = server.received_requests().await.unwrap(); + let reloaded_catalogs: Vec<_> = requests + .iter() + .filter(|request| request.url.path().starts_with("/catalog-reload/")) + .collect(); + assert_eq!(reloaded_catalogs.len(), 6); + for request in reloaded_catalogs { + assert_eq!(request.headers["originator"], "codex_work_cca"); + let authorization = request.headers["authorization"].to_str()?; + if request.url.path().ends_with("/disabled") { + assert_eq!(authorization, "Bearer key-a"); + assert!(!request.headers.contains_key("x-openai-agent-assertion")); + continue; + } + let proof = if request.url.path().ends_with("/dual") { + assert_eq!(authorization, "Bearer key-a"); + request.headers["x-openai-agent-assertion"].to_str()? + } else { + assert!(!request.headers.contains_key("x-openai-agent-assertion")); + authorization + .strip_prefix("AgentAssertion ") + .context("catalog assertion")? + }; + let payload: Value = serde_json::from_slice( + &base64::engine::general_purpose::URL_SAFE_NO_PAD + .decode(proof.split('.').nth(2).unwrap())?, + )?; + assert_eq!(payload["runtime_agent_instance"], "runtime-shared"); + assert!(proof_ids.insert(payload["id"].as_str().unwrap().to_owned())); + } + server.verify().await; + Ok(()) +} + +async fn exercise_harness(server_uri: String) -> Result<()> { + let server = MockServer::start().await; + let test = test_codex() + .with_auth(CodexAuth::from_api_key("key-a")) + .with_config(|config| { + config.model_provider.base_url = Some("https://api.openai.com/v1".into()); + config.chatgpt_base_url = "https://us.chatgpt.com/backend-api".to_string(); + config.model_provider.env_http_headers = None; + config + .features + .disable(Feature::RuntimeAgentInstances) + .expect("test config should allow disabling runtime auth"); + config.model_catalog = Some( + codex_models_manager::bundled_models_response() + .expect("bundled models should parse"), + ); + }) + .build_with_auto_env(&server) + .await?; + turn(&test.codex, "startup-disabled").await?; + + // Enable after startup, then share one lazy registration with a child and an inline delegate. + let mut enabled = test.config.clone(); + enabled.features.enable(Feature::RuntimeAgentInstances)?; + // Initial discovery and response-triggered ETag refresh use each thread's effective auth. + for (mode, runtime_enabled) in [ + (RuntimeAuthMode::DualHeader, true), + (RuntimeAuthMode::AssertionOnly, true), + (RuntimeAuthMode::DualHeader, false), + ] { + let mut catalog_config = enabled.clone(); + catalog_config.model_catalog = None; + catalog_config + .features + .enable(Feature::ApiKeyModelDiscovery)?; + catalog_config.runtime_auth_mode = mode; + if !runtime_enabled { + catalog_config + .features + .disable(Feature::RuntimeAgentInstances)?; + } + catalog_config.model = None; + let label = if !runtime_enabled { + "disabled" + } else if mode == RuntimeAuthMode::AssertionOnly { + "assertion" + } else { + "dual" + }; + catalog_config.model_provider.model_catalog_url = + Some(format!("https://api.openai.com/catalog-reload/{label}").into()); + let catalog_thread = test + .thread_manager + .start_thread(StartThreadOptions { + metrics_service_name: Some("codex_work_cca".into()), + ..StartThreadOptions::new(catalog_config) + }) + .await?; + assert_eq!( + catalog_thread.session_configured.model, + "catalog-isolation-model" + ); + if label == "dual" { + // The first new ETag must not fetch while discovery is disabled. A second, + // distinct ETag on the same thread must fetch once discovery is enabled again. + test.thread_manager + .set_api_key_model_discovery_enabled(/*enabled*/ false); + turn(&catalog_thread.thread, "discovery-off").await?; + test.thread_manager + .set_api_key_model_discovery_enabled(/*enabled*/ true); + } + turn(&catalog_thread.thread, &format!("catalog-{label}")).await?; + } + let parent = test + .thread_manager + .start_thread(StartThreadOptions::new(enabled.clone())) + .await?; + turn(&parent.thread, "parent-dual").await?; + exercise_provider_destinations(&parent.thread, &enabled).await?; + let mut child_config = enabled.clone(); + child_config.runtime_auth_mode = RuntimeAuthMode::AssertionOnly; + parent.thread.ensure_rollout_materialized().await; + parent.thread.flush_rollout().await?; + let prepared = test + .thread_store + .prepare_fork(codex_thread_store::PrepareForkParams { + thread_id: parent.thread_id, + boundary: codex_thread_store::ForkBoundary::Latest, + }) + .await?; + let child = test + .thread_manager + .fork_prepared_thread( + StartThreadOptions { + session_source: Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id: parent.thread_id, + depth: 1, + agent_path: None, + agent_nickname: None, + agent_role: None, + })), + ..StartThreadOptions::new(child_config) + }, + prepared, + ) + .await?; + turn(&child.thread, "child-assertion").await?; + exercise_provider_destinations(&child.thread, &enabled).await?; + parent + .thread + .submit(Op::Review { + review_request: ReviewRequest { + target: ReviewTarget::Custom { + instructions: "Review this change".into(), + }, + user_facing_hint: None, + }, + }) + .await?; + wait_for_event(&parent.thread, |e| matches!(e, EventMsg::TurnComplete(_))).await; + + let disabled = test + .thread_manager + .start_thread(StartThreadOptions::new(test.config.clone())) + .await?; + turn(&disabled.thread, "thread-disabled").await?; + // A third-party OpenAI-compatible endpoint still uses its existing provider authentication. + enabled.model_provider.base_url = Some(format!("{server_uri}/v1")); + enabled.chatgpt_base_url = "https://custom.example.test/backend-api".to_string(); + let third_party = test + .thread_manager + .start_thread(StartThreadOptions::new(enabled)) + .await?; + turn(&third_party.thread, "third-party").await?; + + // Resume a stopped thread under a separately supplied credential context. An internal child + // must inherit the resumed parent's runtime, not the ThreadManager's original default. + let rollout = parent + .thread + .rollout_path() + .context("parent rollout should exist")?; + parent.thread.submit(Op::Shutdown).await?; + wait_for_event(&parent.thread, |e| matches!(e, EventMsg::ShutdownComplete)).await; + test.thread_manager.remove_thread(&parent.thread_id).await; + let context = test + .thread_store + .load_latest_model_context(codex_thread_store::LoadThreadHistoryParams { + thread_id: parent.thread_id, + include_archived: false, + }) + .await?; + let replacement = codex_core::test_support::auth_manager_from_auth_with_home( + CodexAuth::from_api_key("key-b"), + test.codex_home_path().to_path_buf(), + ); + let history = InitialHistory::Resumed(ResumedHistory { + history_revision: None, + conversation_id: context.thread_id, + history: Arc::new(context.items), + rollout_path: Some(rollout), + }); + let mut disabled = test.config.clone(); + disabled.model_catalog = None; + disabled.model = None; + disabled.model_provider.model_catalog_url = + Some("https://api.openai.com/resumed-off-models".into()); + let resumed_off = test + .thread_manager + .resume_thread_with_history( + disabled, + history.clone(), + Arc::clone(&replacement), + /*parent_trace*/ None, + ClientMcpExtensions::default(), + ) + .await?; + assert_eq!( + resumed_off.session_configured.model, + "catalog-isolation-model" + ); + resumed_off.thread.submit(Op::Shutdown).await?; + wait_for_event(&resumed_off.thread, |e| { + matches!(e, EventMsg::ShutdownComplete) + }) + .await; + test.thread_manager + .remove_thread(&resumed_off.thread_id) + .await; + + let mut resumed_config = test.config.clone(); + resumed_config + .features + .enable(Feature::RuntimeAgentInstances)?; + let resumed = test + .thread_manager + .resume_thread_with_history( + resumed_config.clone(), + history, + replacement, + /*parent_trace*/ None, + ClientMcpExtensions::default(), + ) + .await?; + turn(&resumed.thread, "replacement-root").await?; + let internal = test + .thread_manager + .spawn_internal_session( + resumed.thread_id, + StartThreadOptions { + session_source: Some(SessionSource::Internal( + InternalSessionSource::MemoryConsolidation, + )), + ..StartThreadOptions::new(resumed_config) + }, + ) + .await?; + turn(&internal.thread, "replacement-child").await?; + Ok(()) +} + +fn catalog_fixture() -> codex_protocol::openai_models::ModelsResponse { + let mut models = codex_models_manager::bundled_models_response() + .expect("bundled catalog should provide a valid model fixture"); + models.models.truncate(1); + models.models[0].slug = "catalog-isolation-model".into(); + models +} + +async fn exercise_provider_destinations( + thread: &CodexThread, + config: &codex_core::config::Config, +) -> Result<()> { + let external = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/catalog-isolation")) + .respond_with(ResponseTemplate::new(200).set_body_json(catalog_fixture())) + .expect(1) + .mount(&external) + .await; + let external_responses = responses::mount_sse_once( + &external, + responses::sse(vec![responses::ev_completed("external-response")]), + ) + .await; + let auth = thread + .thread_extension_data() + .get::() + .context("thread should publish its provider auth context")?; + for (inference, catalog) in [ + ("https://api.openai.com/v1".to_string(), external.uri()), + (external.uri(), "https://api.openai.com".to_string()), + ] { + let mut info = config.model_provider.clone(); + info.base_url = Some(inference); + info.model_catalog_url = Some(format!("{catalog}/catalog-isolation").into()); + let provider = auth.model_provider(info); + let setup = codex_model_provider::resolve_request( + provider.as_ref(), + /*scope*/ None, + /*routing*/ None, + config.http_client_factory(), + ) + .await?; + let client = codex_api::ResponsesClient::new( + setup.http_transport("responses")?, + setup.api_provider, + setup.api_auth, + ); + let _stream = client + .stream( + json!({"model": "test", "input": []}), + http::HeaderMap::new(), + codex_api::Compression::None, + /*turn_state*/ None, + ) + .await?; + let models = provider.models_manager_without_cache(/*config_model_catalog*/ None); + models.set_api_key_model_discovery_enabled(/*enabled*/ true); + models + .raw_model_catalog(RefreshStrategy::Online, config.http_client_factory()) + .await; + assert_eq!(models.get_remote_models().await, catalog_fixture().models); + } + assert_eq!(external_responses.requests().len(), 1); + let requests = external + .received_requests() + .await + .context("external destination should record requests")?; + assert_eq!(requests.len(), 2); + for request in requests { + assert_eq!(request.headers["authorization"], "Bearer key-a"); + assert!(!request.headers.contains_key("x-openai-agent-assertion")); + } + external.verify().await; + Ok(()) +} + +async fn turn(thread: &Arc, label: &str) -> Result<()> { + thread + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: label.into(), + text_elements: Vec::new(), + }])) + .await?; + wait_for_event(thread, |e| matches!(e, EventMsg::TurnComplete(_))).await; + Ok(()) +} diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/extension.rs b/codex-rs/ext/guardian-v2/src/async_scorer/extension.rs index 58818f9937..034333defe 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/extension.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/extension.rs @@ -72,12 +72,8 @@ impl ThreadLifecycleContributor for GuardianV2Extension { .constrain_guardian_policy(&mut policy, &model.slug); } let scoring_enabled = policy.scoring_enabled(); - let sampler_config = super::startup::sampler_config( - &input, - Arc::clone(&self.auth_manager), - self.thread_manager.upgrade(), - ) - .await; + let sampler_config = + super::startup::sampler_config(&input, self.thread_manager.upgrade()).await; if scoring_enabled && guardian_config.transcript.include_images { input diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs index bcb74a512e..ea73bd791d 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs @@ -265,8 +265,13 @@ async fn installed_extension_warms_connections_without_blocking_thread_start() - async fn installed_extension_uses_http_after_warm_socket_auth_expires() -> Result<()> { skip_if_no_network!(Ok(())); + let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("original")); + auth_manager + .set_external_auth(Arc::new(RefreshableAuth(std::sync::Mutex::new("original")))) + .await?; let thread_server = responses::start_mock_server().await; let test = test_codex() + .with_auth_manager(Arc::clone(&auth_manager)) .with_config(|config| config.approvals_reviewer = ApprovalsReviewer::AutoReview) .build_with_auto_env(&thread_server) .await?; @@ -281,10 +286,6 @@ async fn installed_extension_uses_http_after_warm_socket_auth_expires() -> Resul let server = responses::start_websocket_server(connections).await; let http = responses::start_mock_server().await; let http_mock = responses::mount_sse_once(&http, responses::sse(events)).await; - let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("original")); - auth_manager - .set_external_auth(Arc::new(RefreshableAuth(std::sync::Mutex::new("original")))) - .await?; let mut config = test.config.clone(); config.model_provider = ModelProviderInfo::create_openai_provider(Some( proxy_websocket_servers_with_http( diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool.rs index 23165a419f..b89fb46c7c 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool.rs @@ -25,6 +25,9 @@ use codex_api::SharedAuthProvider; use codex_api::TransportError; use codex_api::build_session_headers; use codex_http_client::ClientRouteClass; +use codex_http_client::HttpClientFactory; +use codex_http_client::NetworkPolicy; +use codex_login::AuthChangeState; use codex_login::CodexAuth; use codex_login::default_client::ClientRedirectPolicy; use codex_login::default_client::add_originator_header; @@ -33,7 +36,6 @@ use codex_login::default_client::default_headers; use codex_model_provider::ACCOUNT_ROUTING_HEADER; use codex_model_provider::AgentIdentitySessionFallback; use codex_model_provider::ProviderAuthScope; -use codex_model_provider::ResolvedResponsesProvider; use codex_model_provider::ResponsesConnectionKey; use codex_protocol::ThreadId; use codex_websocket_client::WebSocketConnector; @@ -76,6 +78,7 @@ pub(super) struct ConnectionPool { struct HttpTransport { url: String, + network_policy: NetworkPolicy, redirect_policy: ClientRedirectPolicy, transport: OnceCell, } @@ -86,17 +89,21 @@ pub(super) struct PooledConnection { // The bridge routes by thread ID, so each socket needs its own identity. pub(super) thread_id: String, pub(super) expires_at: Instant, - auth_changes: Option>, + auth_changes: Option>, key: ResponsesConnectionKey, + runtime_id: Option, _permit: OwnedSemaphorePermit, } struct ClientSetup { provider: Provider, + http_client_factory: HttpClientFactory, + runtime_id: Option, + request_kind: RequestMode, auth: SharedAuthProvider, redirect_policy: ClientRedirectPolicy, key: ResponsesConnectionKey, - auth_changes: Option>, + auth_changes: Option>, } enum Connection { @@ -218,6 +225,9 @@ impl ConnectionPool { let ClientSetup { mut provider, auth, + http_client_factory, + runtime_id, + request_kind, redirect_policy, key, auth_changes, @@ -232,6 +242,7 @@ impl ConnectionPool { match idle { Some(connection) if connection.key == key + && connection.runtime_id == runtime_id && connection .auth_changes .as_ref() @@ -264,7 +275,6 @@ impl ConnectionPool { self.replenish(); // Sampling owns the retry budget across both transports. provider.retry.max_attempts = 0; - let request_kind = self.responses_request_kind().await; let url = provider.url_for_path("/responses"); let redirect_policy = if provider.headers.contains_key(ACCOUNT_ROUTING_HEADER) { ClientRedirectPolicy::Reject @@ -279,11 +289,13 @@ impl ConnectionPool { if let Some(transport) = cached.as_ref() && transport.url == url && transport.redirect_policy == redirect_policy + && transport.network_policy == *http_client_factory.network_policy() { Arc::clone(transport) } else { let transport = Arc::new(HttpTransport { url: url.clone(), + network_policy: http_client_factory.network_policy().clone(), redirect_policy, transport: OnceCell::new(), }); @@ -295,7 +307,7 @@ impl ConnectionPool { .transport .get_or_try_init(|| async { let client = create_client_for_route_async( - self.config.http_client_factory.clone(), + http_client_factory.clone(), url, ClientRouteClass::Api, redirect_policy, @@ -353,45 +365,41 @@ impl ConnectionPool { } async fn client_setup(&self) -> Result { + let setup = codex_model_provider::resolve_request( + self.config.provider.as_ref(), + Some(ProviderAuthScope { + agent_identity_policy: self.config.agent_identity_policy, + session_source: self.config.session_source.clone(), + agent_identity_session_fallback: AgentIdentitySessionFallback::default(), + }), + Some(&self.config.workspace_routing), + self.config.http_client_factory.clone(), + ) + .await + .map_err(LunaSamplerError::Provider)?; let mut auth_changes = self .config .provider .auth_manager() - .map(|manager| manager.auth_change_receiver()); - let revision = auth_changes.as_mut().map(|auth| *auth.borrow_and_update()); - let ResolvedResponsesProvider { - provider, - redirect_policy, - } = self - .config - .provider - .responses_api_provider(&self.config.workspace_routing) - .await - .map_err(LunaSamplerError::Provider)?; - let auth = self - .config - .provider - .api_auth_for_scope(ProviderAuthScope { - agent_identity_policy: self.config.agent_identity_policy, - session_source: self.config.session_source.clone(), - agent_identity_session_fallback: AgentIdentitySessionFallback::default(), - }) - .await - .map_err(LunaSamplerError::Provider)? - .auth; + .map(|manager| manager.auth_change_state_receiver()); if auth_changes - .as_ref() - .is_some_and(|auth| auth.has_changed().unwrap_or(true)) + .as_mut() + .map(|auth| auth.borrow_and_update().generation) + != setup.auth_revision { return Err(LunaSamplerError::Api(ApiError::Stream( "authentication changed while resolving routing".into(), ))); } - let key = ResponsesConnectionKey::new(&provider, revision); + let key = ResponsesConnectionKey::new(&setup.api_provider, setup.auth_revision); + let request_kind = self.responses_request_kind(setup.auth.as_ref()); Ok(ClientSetup { - provider, - auth, - redirect_policy, + provider: setup.api_provider, + auth: setup.api_auth, + redirect_policy: setup.redirect_policy, + http_client_factory: setup.http_client_factory, + runtime_id: setup.runtime_id, + request_kind, key, auth_changes, }) @@ -430,15 +438,9 @@ impl ConnectionPool { } Ok(headers) } - async fn responses_request_kind(&self) -> RequestMode { + fn responses_request_kind(&self, auth: Option<&CodexAuth>) -> RequestMode { let provider = self.config.provider.info(); - if self - .config - .provider - .auth() - .await - .as_ref() - .is_some_and(CodexAuth::uses_codex_backend) + if auth.is_some_and(CodexAuth::uses_codex_backend) && provider.is_openai() && provider.requires_openai_auth && provider.env_key.is_none() @@ -459,12 +461,14 @@ impl ConnectionPool { let ClientSetup { provider, auth, + http_client_factory, + runtime_id, + request_kind, redirect_policy: _, key, auth_changes, } = self.client_setup().await?; let thread_id = ThreadId::new().to_string(); - let request_kind = self.responses_request_kind().await; let mut headers = self.headers(&thread_id, request_kind)?; headers.insert( "openai-beta", @@ -488,9 +492,17 @@ impl ConnectionPool { }) }) .await?; + let connector = connector + .with_account_binding(http_client_factory.network_policy()) + .ok_or_else(|| { + ApiError::Transport(TransportError::Build( + "websocket credentials and transport have different account policies" + .into(), + )) + })?; client .connect_with_connector( - connector, + &connector, headers, default_headers(), /*turn_state*/ None, @@ -532,6 +544,7 @@ impl ConnectionPool { expires_at: Instant::now() + MAX_WEBSOCKET_AGE, auth_changes, key, + runtime_id, _permit: permit, }) } diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool_tests.rs index 86bd03dfd2..1121b668ef 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool_tests.rs @@ -513,6 +513,71 @@ impl WorkspaceRoutingResolver for RoutingPolicy { } } +#[tokio::test] +async fn sampling_uses_new_credentials_after_logout_and_login() -> Result<()> { + skip_if_no_network!(Ok(())); + let home = tempfile::tempdir()?; + let auth = AuthManager::from_auth_for_testing_with_home( + CodexAuth::from_api_key("initial-key"), + home.path().to_path_buf(), + ); + let server = responses::start_mock_server().await; + let requests = responses::mount_sse_sequence( + &server, + vec![ + responses::sse(vec![ + responses::ev_output_text_delta("low"), + responses::ev_completed("before-logout"), + ]), + responses::sse(vec![ + responses::ev_output_text_delta("low"), + responses::ev_completed("after-login"), + ]), + ], + ) + .await; + let mut config = sampler_config(format!("{}/v1", server.uri())); + config.provider = create_model_provider(config.provider.info().clone(), Some(auth.clone())); + let sampler = LunaSampler::new(config); + // Keep both requests on HTTP without starting background WebSocket opens. + *sampler.connections.retry_after.lock().unwrap() = Some(Instant::now() + CONNECT_COOLDOWN); + assert_eq!( + sampler + .sample(sample_request("before-logout")) + .await + .expect("sample before logout"), + "low" + ); + + auth.logout().await?; + codex_login::login_with_api_key( + home.path(), + "replacement-key", + codex_login::AuthCredentialsStoreMode::File, + codex_login::AuthKeyringBackendKind::default(), + )?; + auth.reload().await; + + assert_eq!( + sampler + .sample(sample_request("after-login")) + .await + .expect("sample after signing in again"), + "low" + ); + let requests = requests.requests(); + assert_eq!(requests.len(), 2); + assert_eq!( + requests[0].header("authorization").as_deref(), + Some("Bearer initial-key") + ); + assert_eq!( + requests[1].header("authorization").as_deref(), + Some("Bearer replacement-key") + ); + Ok(()) +} + #[tokio::test] async fn http_cache_preserves_no_constraint_redirect_policy_and_configured_headers() -> Result<()> { let base_url = "https://localhost:9/backend-api/codex".to_owned(); diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler_routing_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler_routing_tests.rs index 5a94aed3ed..f352fe388c 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler_routing_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler_routing_tests.rs @@ -184,12 +184,14 @@ impl WorkspaceRoutingResolver for ChangedBootstrap { #[tokio::test] async fn idle_classifier_retains_the_owning_threads_workspace_binding() -> Result<()> { core_test_support::skip_if_no_network!(Ok(())); - let server = responses::start_mock_server().await; - let test = core_test_support::test_codex::test_codex() - .build_with_auto_env(&server) - .await?; let auth = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); + let thread_server = responses::start_mock_server().await; + let test = core_test_support::test_codex::test_codex() + .with_auth_manager(Arc::clone(&auth)) + .build_with_auto_env(&thread_server) + .await?; + let server = responses::start_mock_server().await; let resolver: Arc = Arc::new(ChangedBootstrap(AtomicBool::new(/*v*/ true))); auth.set_workspace_routing_resolver(Arc::downgrade(&resolver)); @@ -213,8 +215,7 @@ async fn idle_classifier_retains_the_owning_threads_workspace_binding() -> Resul .thread_store .get::() .expect("Core seeds the owning thread's routing context before extensions start"); - let classifier = - crate::async_scorer::startup::sampler_config(&input, auth.clone(), /*manager*/ None).await; + let classifier = crate::async_scorer::startup::sampler_config(&input, /*manager*/ None).await; classifier .provider .responses_api_provider(&owning_context) diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/startup.rs b/codex-rs/ext/guardian-v2/src/async_scorer/startup.rs index c7521b050e..5027026e45 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/startup.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/startup.rs @@ -12,7 +12,6 @@ use codex_features::Feature; use codex_login::AgentIdentityAuthPolicy; use codex_login::AuthManager; use codex_login::CodexAuth; -use codex_model_provider::create_model_provider; use super::decisions::DecisionsError; use super::decisions::DecisionsSampler; @@ -78,7 +77,6 @@ pub(super) fn decisions_api_key( pub(super) async fn sampler_config( input: &ThreadStartInput<'_, Config>, - auth_manager: Arc, manager: Option>, ) -> LunaSamplerConfig { let luna_model = if let Some(thread_manager) = manager { @@ -102,6 +100,11 @@ pub(super) async fn sampler_config( ) }); let luna_compaction_hash = luna_model.and_then(|model| model.comp_hash.clone()); + #[expect(clippy::expect_used, reason = "Required thread startup invariant")] + let auth = input + .thread_store + .get::() + .expect("Core installs provider auth before starting Guardian"); LunaSamplerConfig { api_key_cyber_access_programs: ApiKeyCyberAccessPrograms::from_config(input.config), workspace_routing: input @@ -109,7 +112,7 @@ pub(super) async fn sampler_config( .get_or_init(|| input.config.workspace_routing_context()) .as_ref() .clone(), - provider: create_model_provider(input.config.model_provider.clone(), Some(auth_manager)), + provider: auth.model_provider(input.config.model_provider.clone()), http_client_factory: input.config.http_client_factory(), agent_identity_policy: if input.config.features.enabled(Feature::UseAgentIdentity) { AgentIdentityAuthPolicy::ChatGptAuth diff --git a/codex-rs/ext/history-notes/src/backend.rs b/codex-rs/ext/history-notes/src/backend.rs index fa9db40af6..8218b764c1 100644 --- a/codex-rs/ext/history-notes/src/backend.rs +++ b/codex-rs/ext/history-notes/src/backend.rs @@ -53,11 +53,16 @@ impl HistoryNotesBackend { }), ); - let (provider, auth) = self.provider.api_provider_and_auth().await.map_err(|_| { - format!("{OPERATION_ERROR_PREFIX} Could not prepare the backend request.") - })?; + let setup = codex_model_provider::resolve_request( + self.provider.as_ref(), + /*scope*/ None, + /*routing*/ None, + self.http_client_factory.clone(), + ) + .await + .map_err(|_| format!("{OPERATION_ERROR_PREFIX} Could not prepare the backend request."))?; - let mut request = provider.build_request(Method::POST, path); + let mut request = setup.api_provider.build_request(Method::POST, path); let encoded_truncation_policy = serde_json::to_string(&truncation_policy).map_err(|_| { format!("{OPERATION_ERROR_PREFIX} Could not encode the output truncation policy.") @@ -84,12 +89,13 @@ impl HistoryNotesBackend { } request.body = Some(RequestBody::Json(arguments)); request.timeout = Some(HISTORY_NOTES_BACKEND_TIMEOUT); - let request = auth.apply_auth(request).await.map_err(|_| { + let request = setup.api_auth.apply_auth(request).await.map_err(|_| { format!("{OPERATION_ERROR_PREFIX} Could not apply backend authentication.") })?; let transport = create_transport_for_routes_async( - self.http_client_factory.clone(), + setup.http_client_factory, ClientRouteClass::Api, + setup.redirect_policy, ) .await .map_err(|_| format!("{OPERATION_ERROR_PREFIX} The backend request failed."))?; diff --git a/codex-rs/ext/history-notes/src/extension.rs b/codex-rs/ext/history-notes/src/extension.rs index 8af60180d4..e5590588f1 100644 --- a/codex-rs/ext/history-notes/src/extension.rs +++ b/codex-rs/ext/history-notes/src/extension.rs @@ -18,7 +18,6 @@ use codex_extension_api::ToolCall; use codex_extension_api::ToolContributor; use codex_extension_api::ToolExecutor; use codex_login::AuthManager; -use codex_model_provider::create_model_provider; use codex_protocol::AgentPath; use codex_protocol::openai_models::HistoryToolMessages; use codex_protocol::openai_models::NotesToolMessages; @@ -46,19 +45,24 @@ struct HistoryNotesAgentIdentity { impl HistoryNotesExtension { fn update_config(&self, thread_store: &ExtensionData, config: &Config) { + let auth = thread_store + .get::() + .unwrap_or_else(|| { + Arc::new(codex_model_provider::ProviderAuthContext { + credentials: Arc::clone(&self.auth_manager), + runtime_auth: None, + }) + }); if config .token_budget .as_ref() .is_some_and(|token_budget| token_budget.use_history_notes_extension) && config.model_provider.is_openai() - && self.auth_manager.current_auth_uses_codex_backend() + && auth.credentials.current_auth_uses_codex_backend() { thread_store.insert(HistoryNotesExtensionConfig { backend: HistoryNotesBackend::new( - create_model_provider( - config.model_provider.clone(), - Some(self.auth_manager.clone()), - ), + auth.model_provider(config.model_provider.clone()), config.http_client_factory(), ), }); diff --git a/codex-rs/ext/image-generation/src/backend.rs b/codex-rs/ext/image-generation/src/backend.rs index 694e7b08e3..69d14c7994 100644 --- a/codex-rs/ext/image-generation/src/backend.rs +++ b/codex-rs/ext/image-generation/src/backend.rs @@ -78,18 +78,26 @@ impl CodexImagesBackend { /// Resolves the provider and auth required for the current image API request. async fn client(&self) -> Result, ImageBackendError> { - let (provider, auth) = self - .provider - .api_provider_and_auth() - .await - .map_err(|err| ImageBackendError::from_message(err.to_string()))?; - let transport = create_transport_for_routes_async( + let setup = codex_model_provider::resolve_request( + self.provider.as_ref(), + /*scope*/ None, + /*routing*/ None, self.http_client_factory.clone(), - ClientRouteClass::Api, ) .await .map_err(|err| ImageBackendError::from_message(err.to_string()))?; - Ok(ImagesClient::new(transport, provider, auth)) + let transport = create_transport_for_routes_async( + setup.http_client_factory, + ClientRouteClass::Api, + setup.redirect_policy, + ) + .await + .map_err(|err| ImageBackendError::from_message(err.to_string()))?; + Ok(ImagesClient::new( + transport, + setup.api_provider, + setup.api_auth, + )) } /// Sends a standalone image generation request through the configured Images client. diff --git a/codex-rs/ext/image-generation/src/extension.rs b/codex-rs/ext/image-generation/src/extension.rs index eea41d3f4e..e240e4a807 100644 --- a/codex-rs/ext/image-generation/src/extension.rs +++ b/codex-rs/ext/image-generation/src/extension.rs @@ -13,7 +13,6 @@ use codex_extension_api::ToolContributor; use codex_extension_api::ToolExecutor; use codex_http_client::HttpClientFactory; use codex_login::AuthManager; -use codex_model_provider::create_model_provider; use codex_model_provider_info::ModelProviderInfo; use codex_utils_absolute_path::AbsolutePathBuf; @@ -98,10 +97,18 @@ impl ToolContributor for ImageGenerationExtension { if !config.available { return Vec::new(); } + let auth = thread_store + .get::() + .unwrap_or_else(|| { + Arc::new(codex_model_provider::ProviderAuthContext { + credentials: Arc::clone(&self.auth_manager), + runtime_auth: None, + }) + }); vec![Arc::new(ImageGenerationTool::new( CodexImagesBackend::new( - create_model_provider(config.provider.clone(), Some(self.auth_manager.clone())), + auth.model_provider(config.provider.clone()), config.http_client_factory.clone(), thread_store .get::() diff --git a/codex-rs/ext/web-search/src/extension.rs b/codex-rs/ext/web-search/src/extension.rs index bebdf5ebf3..f5e172437d 100644 --- a/codex-rs/ext/web-search/src/extension.rs +++ b/codex-rs/ext/web-search/src/extension.rs @@ -20,7 +20,6 @@ use codex_extension_api::ToolContributor; use codex_http_client::HttpClientFactory; use codex_login::AuthManager; use codex_model_provider::WorkspaceRoutingContext; -use codex_model_provider::create_model_provider; use codex_model_provider_info::ModelProviderInfo; use codex_protocol::config_types::WebSearchContextSize; use codex_protocol::config_types::WebSearchMode; @@ -137,14 +136,19 @@ impl ToolContributor for WebSearchExtension { if !config.provider_supports_search || *mode == WebSearchMode::Disabled { return Vec::new(); } + let auth = thread_store + .get::() + .unwrap_or_else(|| { + Arc::new(codex_model_provider::ProviderAuthContext { + credentials: Arc::clone(&self.auth_manager), + runtime_auth: None, + }) + }); vec![Arc::new(WebSearchTool { session_id: session_store.level_id().to_string(), http_client_factory: config.http_client_factory.clone(), - provider: create_model_provider( - config.provider.clone(), - Some(self.auth_manager.clone()), - ), + provider: auth.model_provider(config.provider.clone()), workspace_routing, settings: SearchSettings { external_web_access: Some(external_web_access_for_mode(*mode)), diff --git a/codex-rs/ext/web-search/src/tool.rs b/codex-rs/ext/web-search/src/tool.rs index 4b92af27d0..34b00c3ed2 100644 --- a/codex-rs/ext/web-search/src/tool.rs +++ b/codex-rs/ext/web-search/src/tool.rs @@ -29,8 +29,6 @@ use codex_login::default_client::create_transport_for_routes_async; use codex_model_provider::ACCOUNT_ROUTING_HEADER; use codex_model_provider::SharedModelProvider; use codex_model_provider::WorkspaceRoutingContext; -use codex_model_provider::prepare_provider_request; -use codex_model_provider::provider_uses_first_party_auth_path; use codex_protocol::models::WebSearchAction as CoreWebSearchAction; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::WebSearchBeginEvent; @@ -113,45 +111,38 @@ impl WebSearchTool { ) -> Result, FunctionCallError> { let commands = parse_commands(call.function_arguments()?, self.supports_literature_search)?; let command_action = command_action(&commands); - let provider = self.provider.as_ref(); - let first_party = provider_uses_first_party_auth_path(provider.info()); - let client = prepare_provider_request(provider, first_party, |_| async { - let routing_auth = if first_party { - provider.auth().await - } else { - None - }; - let (endpoint, api_auth) = provider.api_provider_and_auth().await?; - let resolved = provider - .route_responses_api_provider( - endpoint, - routing_auth.as_ref(), - &self.workspace_routing, - ) - .await?; - let endpoint = resolved.provider; - let transport = if resolved.redirect_policy == ClientRedirectPolicy::Reject - || endpoint.headers.contains_key(ACCOUNT_ROUTING_HEADER) - { - create_client_for_route_async( - self.http_client_factory.clone(), - endpoint.url_for_path("alpha/search"), - ClientRouteClass::Api, - ClientRedirectPolicy::Reject, - ) - .await - .map(ReqwestTransport::from_http_client) - } else { - create_transport_for_routes_async( - self.http_client_factory.clone(), - ClientRouteClass::Api, - ) - .await - }?; - Ok(SearchClient::new(transport, endpoint, api_auth)) - }) + let setup = codex_model_provider::resolve_request( + self.provider.as_ref(), + /*scope*/ None, + Some(&self.workspace_routing), + self.http_client_factory.clone(), + ) .await .map_err(|err| FunctionCallError::RespondToModel(search_error_message(err)))?; + let transport = if setup.redirect_policy == ClientRedirectPolicy::Reject + || setup + .api_provider + .headers + .contains_key(ACCOUNT_ROUTING_HEADER) + { + create_client_for_route_async( + setup.http_client_factory, + setup.api_provider.url_for_path("alpha/search"), + ClientRouteClass::Api, + ClientRedirectPolicy::Reject, + ) + .await + .map(ReqwestTransport::from_http_client) + } else { + create_transport_for_routes_async( + setup.http_client_factory, + ClientRouteClass::Api, + setup.redirect_policy, + ) + .await + } + .map_err(|err| FunctionCallError::RespondToModel(search_error_message(err)))?; + let client = SearchClient::new(transport, setup.api_provider, setup.api_auth); let request = SearchRequest { id: self.session_id.clone(), model: call.model.clone(), diff --git a/codex-rs/login/src/auth/default_client.rs b/codex-rs/login/src/auth/default_client.rs index 7d524a0e7c..42d023e339 100644 --- a/codex-rs/login/src/auth/default_client.rs +++ b/codex-rs/login/src/auth/default_client.rs @@ -383,6 +383,7 @@ async fn create_client_for_route_async_with_builder( pub async fn create_transport_for_routes_async( http_client_factory: HttpClientFactory, route_class: ClientRouteClass, + redirect_policy: ClientRedirectPolicy, ) -> std::io::Result { let permit = ROUTE_AWARE_CLIENT_BUILD_PERMIT .acquire() @@ -395,7 +396,8 @@ pub async fn create_transport_for_routes_async( http_client_factory.outbound_proxy_policy(), OutboundProxyPolicy::ReqwestDefault ); - if !http_client_factory.network_policy().is_managed() + if redirect_policy == ClientRedirectPolicy::Default + && !http_client_factory.network_policy().is_managed() && !http_client_factory.protects_sensitive_headers() && (transport_default_proxy || sandboxed) { @@ -405,7 +407,10 @@ pub async fn create_transport_for_routes_async( let mut pool = RouteAwareClientPool::with_builder( http_client_factory, route_class, - default_http_client_builder(), + match redirect_policy { + ClientRedirectPolicy::Default => default_http_client_builder(), + ClientRedirectPolicy::Reject => default_http_client_builder().without_redirects(), + }, ); if sandboxed { pool = pool.with_legacy_direct_proxy_and_custom_ca_fallback(); diff --git a/codex-rs/memories/write/src/runtime.rs b/codex-rs/memories/write/src/runtime.rs index 7eff8ef926..1b995b9098 100644 --- a/codex-rs/memories/write/src/runtime.rs +++ b/codex-rs/memories/write/src/runtime.rs @@ -21,8 +21,8 @@ use codex_login::auth::AgentIdentityAuthPolicy; use codex_login::auth_env_telemetry::collect_auth_env_telemetry; use codex_login::default_client::originator; use codex_model_provider::ModelProvider; +use codex_model_provider::ProviderAuthContext; use codex_model_provider::SharedModelProvider; -use codex_model_provider::create_model_provider; use codex_otel::SessionTelemetry; use codex_otel::TelemetryAuthMode; use codex_protocol::MemoryVersion; @@ -147,13 +147,19 @@ impl MemoryStartupContext { config: &Config, source: SessionSource, ) -> Self { - let provider = create_model_provider( - config.model_provider.clone(), - Some(Arc::clone(&auth_manager)), - ); + let auth = thread + .thread_extension_data() + .get::() + .unwrap_or_else(|| { + Arc::new(ProviderAuthContext { + credentials: auth_manager, + runtime_auth: None, + }) + }); + let provider = auth.model_provider(config.model_provider.clone()); Self::new_with_provider( thread_manager, - auth_manager, + Arc::clone(&auth.credentials), thread_id, thread, config, @@ -327,10 +333,9 @@ impl MemoryStartupContext { let session_id = SessionId::from(self.thread_id); let session_id_string = session_id.to_string(); let model_client = ModelClient::new( - Some(Arc::clone(&self.auth_manager)), + Arc::clone(&self.provider), AgentIdentityAuthPolicy::JwtOnly, self.thread_id, - config.model_provider.clone(), session_source.clone(), config_snapshot.originator, config.model_verbosity, diff --git a/codex-rs/model-provider/src/auth.rs b/codex-rs/model-provider/src/auth.rs index 6cbdeaef48..8d3e407d2e 100644 --- a/codex-rs/model-provider/src/auth.rs +++ b/codex-rs/model-provider/src/auth.rs @@ -12,6 +12,7 @@ use codex_api::AuthProvider; use codex_api::SharedAuthProvider; use codex_features::RuntimeAuthMode; use codex_http_client::CHATGPT_ACCOUNT_ID_HEADER; +use codex_http_client::HttpClientFactory; use codex_login::AuthHeaders; use codex_login::AuthManager; use codex_login::CodexAuth; @@ -74,6 +75,8 @@ impl AgentIdentitySessionFallback { /// Provider auth resolved for a request, plus metadata describing the effective auth. #[derive(Clone)] pub struct ResolvedProviderAuth { + /// Login selected with these headers; routing must use the same credential. + pub credential: Option, pub auth: SharedAuthProvider, pub agent_identity_telemetry: Option, pub runtime: Option, @@ -82,15 +85,45 @@ pub struct ResolvedProviderAuth { impl ResolvedProviderAuth { pub(crate) fn new(auth: SharedAuthProvider) -> Self { Self { + credential: None, auth, agent_identity_telemetry: None, runtime: None, } } + pub(crate) fn for_runtime(selected: RuntimeAuth, mode: RuntimeAuthMode) -> Self { + Self { + credential: Some(selected.auth().clone()), + auth: selected.clone().into_auth_provider(mode), + agent_identity_telemetry: None, + runtime: Some(selected), + } + } + + /// Preserve the caller's transport configuration and bind it to the selected login. + pub fn bind_http_client_factory( + &self, + factory: HttpClientFactory, + ) -> Result { + match &self.runtime { + Some(runtime) => runtime.bind_http_client_factory(factory), + None => Ok(factory), + } + } + + /// Reject account selectors that disagree with the selected runtime. + pub fn validate_headers(&self, headers: &HeaderMap) -> codex_protocol::error::Result<()> { + if let Some(runtime) = &self.runtime { + crate::validate_runtime_selectors(headers, runtime.runtime().actor().info())?; + } + Ok(()) + } + fn for_agent_identity(auth: AgentIdentityAuth) -> Self { let agent_identity_telemetry = agent_identity_telemetry(&auth); Self { + credential: None, auth: Arc::new(AgentIdentityAuthProvider { auth }), agent_identity_telemetry: Some(agent_identity_telemetry), runtime: None, diff --git a/codex-rs/model-provider/src/catalog_factory.rs b/codex-rs/model-provider/src/catalog_factory.rs new file mode 100644 index 0000000000..9ede9430d5 --- /dev/null +++ b/codex-rs/model-provider/src/catalog_factory.rs @@ -0,0 +1,39 @@ +//! Host catalog construction from the provider and credentials selected for a thread. + +use crate::SharedModelProvider; +use codex_models_manager::manager::ModelCatalogSettings; +use codex_models_manager::manager::SharedModelsManager; +use codex_protocol::openai_models::ModelsResponse; + +/// Builds catalogs for the selected provider while retaining the host's cache policy. +/// Provider-backed implementations must use the supplied provider, which carries the +/// current credential owner and runtime mode. Hosts may also supply independent metadata. +pub trait ModelCatalogFactory: Send + Sync { + fn build( + &self, + provider: SharedModelProvider, + config_model_catalog: Option, + api_key_model_discovery_default: bool, + ) -> SharedModelsManager; + + /// Publish a live discovery setting to existing and future catalogs. + /// Sources that do not perform API-key discovery may ignore it. + fn set_api_key_model_discovery_enabled(&self, _enabled: bool) {} +} + +impl ModelCatalogFactory for ModelCatalogSettings { + fn build( + &self, + provider: SharedModelProvider, + config_model_catalog: Option, + api_key_model_discovery_default: bool, + ) -> SharedModelsManager { + let mut settings = self.clone(); + settings.api_key_model_discovery_default = api_key_model_discovery_default; + provider.models_manager_with_settings(config_model_catalog, settings) + } + + fn set_api_key_model_discovery_enabled(&self, enabled: bool) { + ModelCatalogSettings::set_api_key_model_discovery_enabled(self, enabled); + } +} diff --git a/codex-rs/model-provider/src/lib.rs b/codex-rs/model-provider/src/lib.rs index e8b578e0ec..3b04f298cc 100644 --- a/codex-rs/model-provider/src/lib.rs +++ b/codex-rs/model-provider/src/lib.rs @@ -2,6 +2,7 @@ mod amazon_bedrock; mod auth; mod bearer_auth_provider; mod capabilities; +mod catalog_factory; mod chatgpt_base_url; mod combined_auth; mod gcp; @@ -9,6 +10,7 @@ mod models_endpoint; mod models_identity; mod provider; mod request_preparation; +mod request_setup; mod runtime_auth; mod shared_state; pub mod test_support; @@ -31,6 +33,7 @@ pub use auth::unauthenticated_auth_provider; pub use bearer_auth_provider::BearerAuthProvider; pub use bearer_auth_provider::BearerAuthProvider as CoreAuthProvider; pub use capabilities::ProviderCapabilities; +pub use catalog_factory::ModelCatalogFactory; pub use chatgpt_base_url::ChatgptBaseUrl; pub use chatgpt_base_url::DEFAULT_CHATGPT_BASE_URL; pub use codex_model_provider_info::AMAZON_BEDROCK_PROVIDER_ID; @@ -52,7 +55,8 @@ pub use provider::provider_uses_first_party_auth_path; pub use provider::resolve_api_provider; pub use request_preparation::prepare_provider_request; -pub use runtime_auth::select_runtime_auth; +pub use request_setup::ResolvedProviderRequest; +pub use request_setup::resolve_request; pub use runtime_auth::validate_runtime_destination; pub use runtime_auth::validate_runtime_provider; pub use runtime_auth::validate_runtime_selectors; diff --git a/codex-rs/model-provider/src/models_endpoint.rs b/codex-rs/model-provider/src/models_endpoint.rs index 46fa826aed..1e9ca0843e 100644 --- a/codex-rs/model-provider/src/models_endpoint.rs +++ b/codex-rs/model-provider/src/models_endpoint.rs @@ -39,12 +39,11 @@ use codex_runtime_agent_instance::RuntimeAuthManager; use http::HeaderMap; use tokio::time::timeout; -use crate::auth::ResolvedProviderAuth; use crate::auth::agent_identity_telemetry; -use crate::auth::resolve_provider_auth; use crate::auth::uses_fedramp_api_key; use crate::combined_auth::compose_auth; use crate::provider::resolve_api_provider; +use crate::request_setup::resolve_request_auth; const MODELS_REFRESH_TIMEOUT: Duration = Duration::from_secs(5); const MODELS_ENDPOINT: &str = "/models"; @@ -90,9 +89,15 @@ impl OpenAiModelsEndpoint { } async fn auth(&self) -> Option { - match self.auth_manager.as_ref() { - Some(auth_manager) => auth_manager.auth().await, - None => None, + let auth_manager = self.auth_manager.clone()?; + // Catalog startup has its own deadline. Let rotating credentials finish + // saving even when that caller stops waiting. + match tokio::spawn(async move { auth_manager.auth().await }).await { + Ok(auth) => auth, + Err(error) => { + tracing::warn!(?error, "model catalog authentication task failed"); + None + } } } @@ -121,48 +126,20 @@ impl OpenAiModelsEndpoint { client_version: &str, http_client_factory: HttpClientFactory, ) -> CoreResult { - let provider_info = self.provider_info.clone(); - let auth_manager = self.auth_manager.clone(); - let runtime_auth = self.runtime_auth.clone(); - // A timed-out catalog fetch must not discard a rotated refresh token before - // it is saved. Keep credential refresh and registration owned until they finish. - let (auth, selected_runtime_auth) = tokio::spawn(async move { - let selected_runtime_auth = match runtime_auth { - Some((runtime, mode)) => { - crate::validate_runtime_destination( - provider_info - .model_catalog_url - .as_deref() - .map(String::as_str) - .or(provider_info.base_url.as_deref()), - )?; - crate::validate_runtime_provider( - &provider_info, - auth_manager.as_ref(), - &runtime, - )?; - Some(( - crate::runtime_auth::select_runtime_auth( - &runtime, - &provider_info.build_header_map()?, - ) - .await - .map_err(|error| map_api_error(error.into()))?, - mode, - )) - } - None => None, - }; - let auth = match &selected_runtime_auth { - Some((selected_runtime_auth, _)) => Some(selected_runtime_auth.auth().clone()), - None => match auth_manager { - Some(auth_manager) => auth_manager.auth().await, - None => None, - }, - }; - Ok::<_, CodexErr>((auth, selected_runtime_auth)) - }) - .await??; + let resolved = resolve_request_auth( + self.provider_info.clone(), + self.auth_manager.clone(), + self.runtime_auth.clone(), + /*scope*/ None, + self.provider_info + .model_catalog_url + .as_deref() + .map(String::as_str) + .or(self.provider_info.base_url.as_deref()) + .map(str::to_owned), + ) + .await?; + let auth = resolved.credential.clone(); let metric_auth_mode = if self.has_provider_api_key() || auth.as_ref().is_some_and(CodexAuth::is_api_key_auth) { @@ -192,25 +169,14 @@ impl OpenAiModelsEndpoint { } .to_string(); } - let http_client_factory = match &selected_runtime_auth { - Some((selected_runtime_auth, _)) => selected_runtime_auth - .bind_http_client_factory(http_client_factory) - .map_err(|error| map_api_error(error.into()))?, - None => http_client_factory, - }; - let primary = if let Some((selected_runtime_auth, mode)) = selected_runtime_auth { - crate::validate_runtime_selectors( - &api_provider.headers, - selected_runtime_auth.runtime().actor().info(), - )?; - selected_runtime_auth.into_auth_provider(mode) - } else { - resolve_provider_auth(auth.as_ref(), &self.provider_info)? - }; + let http_client_factory = resolved + .bind_http_client_factory(http_client_factory) + .map_err(|error| map_api_error(error.into()))?; + resolved.validate_headers(&api_provider.headers)?; let resolved = compose_auth( &self.provider_info, self.gateway_auth_manager.as_ref(), - ResolvedProviderAuth::new(primary), + resolved, ) .await?; let api_auth = resolved.auth; diff --git a/codex-rs/model-provider/src/models_endpoint_refresh_tests.rs b/codex-rs/model-provider/src/models_endpoint_refresh_tests.rs index 88cd130ac3..7a18508142 100644 --- a/codex-rs/model-provider/src/models_endpoint_refresh_tests.rs +++ b/codex-rs/model-provider/src/models_endpoint_refresh_tests.rs @@ -6,6 +6,8 @@ use codex_http_client::OutboundProxyPolicy; use codex_login::AuthCredentialsStoreMode; use codex_login::AuthKeyringBackendKind; use codex_login::AuthRouteConfig; +use codex_models_manager::manager::ModelsManager; +use codex_models_manager::manager::OpenAiModelsManager; use codex_runtime_agent_instance::AgentHostAuthorization; use codex_runtime_agent_instance::RuntimeRegistrationMetadata; use codex_runtime_agent_instance::RuntimeRole; @@ -20,6 +22,15 @@ use wiremock::matchers::path; #[tokio::test] async fn catalog_timeout_preserves_refreshed_tokens_with_and_without_runtime_auth() { + check_catalog_timeout_preserves_refreshed_tokens(/*refresh_after_auth_change*/ false).await; +} + +#[tokio::test] +async fn catalog_startup_timeout_preserves_refreshed_tokens() { + check_catalog_timeout_preserves_refreshed_tokens(/*refresh_after_auth_change*/ true).await; +} + +async fn check_catalog_timeout_preserves_refreshed_tokens(refresh_after_auth_change: bool) { for mode in [ None, Some(RuntimeAuthMode::DualHeader), @@ -138,20 +149,40 @@ async fn catalog_timeout_preserves_refreshed_tokens_with_and_without_runtime_aut }); let mut info = ModelProviderInfo::create_openai_provider(Some(server.uri())); info.env_http_headers = None; - let endpoint = OpenAiModelsEndpoint::new( - info, - Some(Arc::clone(&credentials)), - /*gateway_auth_manager*/ None, - ) - .with_runtime_auth(runtime_auth); - let error = endpoint - .list_models("test", factory.clone()) - .await - .unwrap_err(); - assert!(matches!( - error.details(), - codex_protocol::error::CodexErrorDetails::RequestTimeout - )); + let endpoint = Arc::new( + OpenAiModelsEndpoint::new( + info, + Some(Arc::clone(&credentials)), + /*gateway_auth_manager*/ None, + ) + .with_runtime_auth(runtime_auth), + ); + if refresh_after_auth_change { + let manager = OpenAiModelsManager::new_without_cache( + endpoint.clone(), + Some(Arc::clone(&credentials)), + ); + manager.refresh_after_auth_change(factory.clone()).await; + } else { + let error = endpoint + .list_models("test", factory.clone()) + .await + .unwrap_err(); + assert!(matches!( + error.details(), + codex_protocol::error::CodexErrorDetails::RequestTimeout + )); + } + assert_eq!( + server + .received_requests() + .await + .unwrap() + .iter() + .filter(|request| request.url.path() == "/oauth/token") + .count(), + 1, + ); // Wait through the existing refresh lock; this must reuse the completed refresh. let current = timeout(Duration::from_secs(5), credentials.auth()) diff --git a/codex-rs/model-provider/src/models_endpoint_runtime_auth_tests.rs b/codex-rs/model-provider/src/models_endpoint_runtime_auth_tests.rs index 1959823a3d..3dee3965dd 100644 --- a/codex-rs/model-provider/src/models_endpoint_runtime_auth_tests.rs +++ b/codex-rs/model-provider/src/models_endpoint_runtime_auth_tests.rs @@ -129,13 +129,28 @@ async fn model_requests_keep_the_selected_project_while_reusing_one_runtime() { headers.insert("openai-project".into(), project.into()); } info.http_headers = Some(headers); - let auth = context - .model_provider(info.clone()) - .api_auth() - .await - .unwrap(); + let provider = context.model_provider(info.clone()); + let request = crate::resolve_request( + provider.as_ref(), + /*scope*/ None, + /*routing*/ None, + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ) + .await + .unwrap(); + assert_eq!(request.runtime_id.as_deref(), Some("runtime-a")); + assert_eq!( + request + .api_provider + .headers + .get("openai-project") + .map(|value| value.to_str().unwrap()), + project, + ); assert!( - auth.to_auth_headers() + request + .api_auth + .to_auth_headers() .unwrap() .contains_key("authorization") ); diff --git a/codex-rs/model-provider/src/provider.rs b/codex-rs/model-provider/src/provider.rs index 986f153675..dbdb4ef011 100644 --- a/codex-rs/model-provider/src/provider.rs +++ b/codex-rs/model-provider/src/provider.rs @@ -16,6 +16,7 @@ use codex_login::WorkspaceRoutingRequest; use codex_login::default_client::ClientRedirectPolicy; use codex_model_provider_info::ModelProviderInfo; use codex_models_manager::cache::ModelsCache; +use codex_models_manager::manager::ModelCatalogSettings; use codex_models_manager::manager::OpenAiModelsManager; use codex_models_manager::manager::SharedModelsManager; use codex_models_manager::manager::StaticModelsManager; @@ -37,6 +38,7 @@ use crate::auth::uses_fedramp_api_key; use crate::combined_auth::compose_auth; use crate::gcp::GcpModelProvider; use crate::models_endpoint::OpenAiModelsEndpoint; +use crate::request_setup::resolve_request_auth; use crate::workspace_routing::WorkspaceRoutingContext; /// Current app-visible account state for a model provider. @@ -353,6 +355,15 @@ pub trait ModelProvider: fmt::Debug + Send + Sync { } } + /// Resolve the endpoint before headers for providers whose setup may refresh credentials. + fn request_provider_and_auth( + &self, + scope: Option, + ) -> ModelProviderFuture<'_, codex_protocol::error::Result<(Provider, ResolvedProviderAuth)>> + { + Box::pin(async move { Ok((self.api_provider().await?, self.request_auth(scope).await?)) }) + } + /// Creates the model manager implementation appropriate for this provider. fn models_manager( &self, @@ -360,6 +371,15 @@ pub trait ModelProvider: fmt::Debug + Send + Sync { config_model_catalog: Option, ) -> SharedModelsManager; + /// Build a thread catalog with the same cache and live discovery policy as its harness. + fn models_manager_with_settings( + &self, + config_model_catalog: Option, + _settings: ModelCatalogSettings, + ) -> SharedModelsManager { + self.models_manager_without_cache(config_model_catalog) + } + /// Creates a model manager with caching disabled. /// /// Providers that fetch model catalogs should override this method. The default uses an @@ -462,12 +482,6 @@ struct ConfiguredModelProvider { catalog_runtime_auth: Option<(Arc, RuntimeAuthMode)>, } -enum ModelsCacheConfig { - Disk { codex_home: PathBuf }, - Disabled, - Custom(Arc), -} - impl ConfiguredModelProvider { fn new( info: ModelProviderInfo, @@ -513,74 +527,16 @@ impl ConfiguredModelProvider { &self, scope: Option, ) -> codex_protocol::error::Result { - let resolved = if let Some((runtime, mode)) = &self.runtime_auth { - crate::validate_runtime_destination(self.info.base_url.as_deref())?; - crate::validate_runtime_provider(&self.info, self.auth_manager.as_ref(), runtime)?; - let headers = self.info.build_header_map()?; - let selected = crate::runtime_auth::select_runtime_auth(runtime, &headers) - .await - .map_err(|error| self.map_api_error(error.into()))?; - crate::validate_runtime_selectors(&headers, selected.runtime().actor().info())?; - ResolvedProviderAuth { - auth: selected.clone().into_auth_provider(*mode), - agent_identity_telemetry: None, - runtime: Some(selected), - } - } else { - let auth = self.auth().await; - match scope { - Some(scope) if provider_uses_first_party_auth_path(&self.info) => { - resolve_provider_auth_for_scope( - self.auth_manager.clone(), - auth.as_ref(), - &self.info, - scope, - ) - .await? - } - _ => ResolvedProviderAuth::new(resolve_provider_auth(auth.as_ref(), &self.info)?), - } - }; + let resolved = resolve_request_auth( + self.info.clone(), + self.auth_manager.clone(), + self.runtime_auth.clone(), + scope, + self.info.base_url.clone(), + ) + .await?; self.compose_auth(resolved).await } - - fn create_models_manager( - &self, - config_model_catalog: Option, - cache: ModelsCacheConfig, - ) -> SharedModelsManager { - if let Some(model_catalog) = config_model_catalog { - return Arc::new(StaticModelsManager::new( - self.auth_manager.clone(), - model_catalog, - codex_models_manager::ModelOrigin::ExplicitConfig, - )); - } - let endpoint = Arc::new( - OpenAiModelsEndpoint::new( - self.info.clone(), - self.auth_manager.clone(), - self.gateway_auth_manager.clone(), - ) - .with_runtime_auth(self.catalog_runtime_auth.clone()), - ); - let auth_manager = self.auth_manager.clone(); - let manager = match cache { - ModelsCacheConfig::Disk { codex_home } => { - OpenAiModelsManager::new(codex_home, endpoint, auth_manager) - } - ModelsCacheConfig::Disabled => { - OpenAiModelsManager::new_without_cache(endpoint, auth_manager) - } - ModelsCacheConfig::Custom(cache) => { - OpenAiModelsManager::new_with_cache(cache, endpoint, auth_manager) - } - }; - match &self.info.model_catalog_url { - Some(_) => Arc::new(manager.with_provider_catalog()), - None => Arc::new(manager), - } - } } impl ModelProvider for ConfiguredModelProvider { @@ -643,6 +599,18 @@ impl ModelProvider for ConfiguredModelProvider { Box::pin(self.resolve_auth(scope)) } + fn request_provider_and_auth( + &self, + scope: Option, + ) -> ModelProviderFuture<'_, codex_protocol::error::Result<(Provider, ResolvedProviderAuth)>> + { + Box::pin(async move { + let resolved = self.resolve_auth(scope).await?; + let provider = self.api_provider_for_auth(resolved.credential.as_ref())?; + Ok((provider, resolved)) + }) + } + fn api_auth( &self, ) -> ModelProviderFuture<'_, codex_protocol::error::Result> { @@ -733,14 +701,45 @@ impl ModelProvider for ConfiguredModelProvider { codex_home: PathBuf, config_model_catalog: Option, ) -> SharedModelsManager { - self.create_models_manager(config_model_catalog, ModelsCacheConfig::Disk { codex_home }) + self.models_manager_with_settings( + config_model_catalog, + ModelCatalogSettings::disk(codex_home), + ) + } + + fn models_manager_with_settings( + &self, + config_model_catalog: Option, + settings: ModelCatalogSettings, + ) -> SharedModelsManager { + if let Some(model_catalog) = config_model_catalog { + return Arc::new(StaticModelsManager::new( + self.auth_manager.clone(), + model_catalog, + codex_models_manager::ModelOrigin::ExplicitConfig, + )); + } + let endpoint = Arc::new( + OpenAiModelsEndpoint::new( + self.info.clone(), + self.auth_manager.clone(), + self.gateway_auth_manager.clone(), + ) + .with_runtime_auth(self.catalog_runtime_auth.clone()), + ); + let auth_manager = self.auth_manager.clone(); + let manager = OpenAiModelsManager::with_settings(settings, endpoint, auth_manager); + match &self.info.model_catalog_url { + Some(_) => Arc::new(manager.with_provider_catalog()), + None => Arc::new(manager), + } } fn models_manager_without_cache( &self, config_model_catalog: Option, ) -> SharedModelsManager { - self.create_models_manager(config_model_catalog, ModelsCacheConfig::Disabled) + self.models_manager_with_settings(config_model_catalog, ModelCatalogSettings::default()) } fn models_manager_with_cache( @@ -748,7 +747,10 @@ impl ModelProvider for ConfiguredModelProvider { config_model_catalog: Option, cache: Arc, ) -> SharedModelsManager { - self.create_models_manager(config_model_catalog, ModelsCacheConfig::Custom(cache)) + self.models_manager_with_settings( + config_model_catalog, + ModelCatalogSettings::with_cache(cache), + ) } } diff --git a/codex-rs/model-provider/src/request_setup.rs b/codex-rs/model-provider/src/request_setup.rs new file mode 100644 index 0000000000..c7e33a0152 --- /dev/null +++ b/codex-rs/model-provider/src/request_setup.rs @@ -0,0 +1,158 @@ +//! Resolve credentials, routing, and transport policy as one request context. + +use crate::ModelProvider; +use crate::ProviderAuthScope; +use crate::ResolvedProviderAuth; +use crate::WorkspaceRoutingContext; +use crate::auth::resolve_provider_auth; +use crate::auth::resolve_provider_auth_for_scope; +use codex_api::AgentIdentityTelemetry; +use codex_api::Provider; +use codex_api::ReqwestTransport; +use codex_api::SharedAuthProvider; +use codex_features::RuntimeAuthMode; +use codex_http_client::ClientRouteClass; +use codex_http_client::HttpClientFactory; +use codex_login::AuthManager; +use codex_login::CodexAuth; +use codex_login::default_client::ClientRedirectPolicy; +use codex_login::default_client::create_client_for_route; +use codex_model_provider_info::ModelProviderInfo; +use codex_protocol::error::Result; +use codex_runtime_agent_instance::RuntimeAuth; +use codex_runtime_agent_instance::RuntimeAuthManager; +use std::sync::Arc; + +/// Select auth for inference and discovery, letting rotating credentials finish saving +/// if the requesting operation is cancelled. +pub(crate) async fn resolve_request_auth( + info: ModelProviderInfo, + manager: Option>, + runtime: Option<(Arc, RuntimeAuthMode)>, + scope: Option, + destination: Option, +) -> codex_protocol::error::Result { + tokio::spawn(async move { + if let Some((runtime, mode)) = runtime { + crate::validate_runtime_destination(destination.as_deref())?; + crate::validate_runtime_provider(&info, manager.as_ref(), &runtime)?; + let headers = info.build_header_map()?; + let selected = crate::runtime_auth::select_runtime_auth(&runtime, &headers) + .await + .map_err(|error| codex_api::map_api_error(error.into()))?; + let resolved = ResolvedProviderAuth::for_runtime(selected, mode); + resolved.validate_headers(&headers)?; + tracing::debug!(runtime_auth_mode = ?mode, "selected runtime authentication"); + return Ok(resolved); + } + + let credential = match &manager { + Some(manager) => manager.auth().await, + None => None, + }; + let mut resolved = match scope { + Some(scope) if crate::provider_uses_first_party_auth_path(&info) => { + resolve_provider_auth_for_scope(manager, credential.as_ref(), &info, scope).await? + } + _ => ResolvedProviderAuth::new(resolve_provider_auth(credential.as_ref(), &info)?), + }; + resolved.credential = credential; + Ok(resolved) + }) + .await? +} + +pub struct ResolvedProviderRequest { + pub auth: Option, + pub auth_owner_generation: Option, + pub auth_revision: Option, + pub api_provider: Provider, + pub redirect_policy: ClientRedirectPolicy, + pub api_auth: SharedAuthProvider, + pub agent_identity_telemetry: Option, + pub runtime_id: Option, + pub http_client_factory: HttpClientFactory, +} + +impl ResolvedProviderRequest { + pub fn http_transport(&self, endpoint: &str) -> Result { + let redirect_policy = if self + .api_provider + .headers + .contains_key(crate::ACCOUNT_ROUTING_HEADER) + { + ClientRedirectPolicy::Reject + } else { + self.redirect_policy + }; + let client = create_client_for_route( + &self.http_client_factory, + &self.api_provider.url_for_path(endpoint), + ClientRouteClass::Api, + redirect_policy, + ) + .map_err(std::io::Error::from)?; + Ok(ReqwestTransport::from_http_client(client)) + } +} + +/// Resolve one attempt, including any workspace route, without mixing credential revisions. +/// The supplied factory retains the caller's proxy/TLS settings; runtime requests bind its +/// application policy to the same selected account used to sign the request. +pub async fn resolve_request( + provider: &dyn ModelProvider, + scope: Option, + routing: Option<&WorkspaceRoutingContext>, + http_client_factory: HttpClientFactory, +) -> Result { + let check_auth = provider.runtime_auth().is_none() + && crate::provider_uses_first_party_auth_path(provider.info()); + crate::prepare_provider_request(provider, check_auth, |state| { + let scope = scope.clone(); + let http_client_factory = http_client_factory.clone(); + async move { + let (api_provider, resolved_auth) = provider.request_provider_and_auth(scope).await?; + let auth = match &resolved_auth.credential { + Some(auth) => Some(auth.clone()), + None => provider.auth().await, + }; + let (api_provider, redirect_policy) = match routing { + Some(routing) => { + let resolved = provider + .route_responses_api_provider(api_provider, auth.as_ref(), routing) + .await?; + (resolved.provider, resolved.redirect_policy) + } + None => (api_provider, ClientRedirectPolicy::Default), + }; + resolved_auth.validate_headers(&api_provider.headers)?; + let runtime_state = resolved_auth + .runtime + .as_ref() + .map(RuntimeAuth::change_state); + let http_client_factory = resolved_auth + .bind_http_client_factory(http_client_factory) + .map_err(|error| provider.map_api_error(error.into()))?; + let runtime_id = resolved_auth + .runtime + .as_ref() + .map(|selected| selected.runtime().id().to_owned()); + Ok(ResolvedProviderRequest { + auth, + auth_owner_generation: runtime_state + .map(|state| state.owner_generation) + .or(state.map(|state| state.owner_generation)), + auth_revision: runtime_state + .map(|state| state.generation) + .or(state.map(|state| state.generation)), + api_provider, + redirect_policy, + api_auth: resolved_auth.auth, + agent_identity_telemetry: resolved_auth.agent_identity_telemetry, + runtime_id, + http_client_factory, + }) + } + }) + .await +} diff --git a/codex-rs/model-provider/src/runtime_auth.rs b/codex-rs/model-provider/src/runtime_auth.rs index 3a817f428a..4c4b992f15 100644 --- a/codex-rs/model-provider/src/runtime_auth.rs +++ b/codex-rs/model-provider/src/runtime_auth.rs @@ -18,7 +18,7 @@ use http::HeaderMap; use std::sync::Arc; /// Select the runtime for the resource scope carried by the model request. -pub async fn select_runtime_auth( +pub(crate) async fn select_runtime_auth( manager: &RuntimeAuthManager, headers: &HeaderMap, ) -> Result { diff --git a/codex-rs/models-manager/src/manager.rs b/codex-rs/models-manager/src/manager.rs index 0ce3c8fda7..3837f1b206 100644 --- a/codex-rs/models-manager/src/manager.rs +++ b/codex-rs/models-manager/src/manager.rs @@ -28,8 +28,6 @@ use std::future::Future; use std::path::PathBuf; use std::pin::Pin; use std::sync::Arc; -use std::sync::atomic::AtomicBool; -use std::sync::atomic::Ordering; use std::time::Duration; use tokio::sync::Mutex; use tokio::sync::MutexGuard; @@ -165,7 +163,7 @@ pub trait ModelsManager: fmt::Debug + Send + Sync { None } - /// Supply startup API-key discovery policy; live changes require a new session. + /// Update API-key discovery for this catalog and its derived catalogs. /// Static catalogs ignore this setting. fn set_api_key_model_discovery_enabled(&self, _enabled: bool) {} @@ -344,16 +342,54 @@ impl CatalogSource { } } +/// Cache policy and live feature state shared by catalogs with different request credentials. +#[derive(Clone, Debug, Default)] +pub struct ModelCatalogSettings { + cache: Option>, + pub api_key_model_discovery_default: bool, + // Before the host publishes a live change, each thread uses its configured default. + live_discovery: Arc>>, +} + +impl ModelCatalogSettings { + pub fn disk(codex_home: PathBuf) -> Self { + Self::with_cache(Arc::new(FileModelsCache::new( + codex_home.join(MODEL_CACHE_FILE), + DEFAULT_MODEL_CACHE_TTL, + ))) + } + + pub fn with_cache(cache: Arc) -> Self { + Self { + cache: Some(cache), + ..Self::default() + } + } + + pub fn set_api_key_model_discovery_enabled(&self, enabled: bool) { + *self + .live_discovery + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(enabled); + } + + fn api_key_model_discovery_enabled(&self) -> bool { + self.live_discovery + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .unwrap_or(self.api_key_model_discovery_default) + } +} + /// OpenAI-compatible model manager backed by bundled models, cache, and `/models`. #[derive(Debug)] pub struct OpenAiModelsManager { originator_managers: std::sync::Mutex>, remote_models: RwLock, catalog_updates: watch::Sender>>, - cache: Option>, + settings: ModelCatalogSettings, endpoint_client: SharedModelsEndpointClient, catalog_source: CatalogSource, - api_key_model_discovery_enabled: Arc, auth_manager: Option>, } @@ -372,12 +408,8 @@ impl OpenAiModelsManager { endpoint_client: Arc, auth_manager: Option>, ) -> Self { - let cache_path = codex_home.join(MODEL_CACHE_FILE); - Self::new_with_optional_cache( - Some(Arc::new(FileModelsCache::new( - cache_path, - DEFAULT_MODEL_CACHE_TTL, - ))), + Self::with_settings( + ModelCatalogSettings::disk(codex_home), endpoint_client, auth_manager, ) @@ -388,7 +420,11 @@ impl OpenAiModelsManager { endpoint_client: Arc, auth_manager: Option>, ) -> Self { - Self::new_with_optional_cache(/*cache*/ None, endpoint_client, auth_manager) + Self::with_settings( + ModelCatalogSettings::default(), + endpoint_client, + auth_manager, + ) } /// Constructs an OpenAI-compatible model manager with a caller-provided cache. @@ -400,11 +436,15 @@ impl OpenAiModelsManager { endpoint_client: Arc, auth_manager: Option>, ) -> Self { - Self::new_with_optional_cache(Some(cache), endpoint_client, auth_manager) + Self::with_settings( + ModelCatalogSettings::with_cache(cache), + endpoint_client, + auth_manager, + ) } - fn new_with_optional_cache( - cache: Option>, + pub fn with_settings( + settings: ModelCatalogSettings, endpoint_client: Arc, auth_manager: Option>, ) -> Self { @@ -418,8 +458,7 @@ impl OpenAiModelsManager { identity: endpoint_client.identity(), models: load_remote_models_from_file().unwrap_or_default(), })), - cache, - api_key_model_discovery_enabled: Arc::new(AtomicBool::new(false)), + settings, endpoint_client, catalog_source: CatalogSource::Default, auth_manager, @@ -503,6 +542,7 @@ impl ModelsManager for OpenAiModelsManager { state.model_catalog(self.models_for_entry(&state.entry)) }) } + fn for_originator(&self, originator: &str) -> Option { let originator = ModelCatalogOriginator::from_originator(originator)?; let mut managers = self @@ -514,23 +554,22 @@ impl ModelsManager for OpenAiModelsManager { } let endpoint = self.endpoint_client.for_originator(originator)?; // A backend that cannot partition must not overwrite the default catalog. - let cache = self + let mut settings = self.settings.clone(); + settings.cache = settings .cache .as_ref() .and_then(|cache| cache.for_originator(originator)); - let mut manager = Self::new_with_optional_cache(cache, endpoint, self.auth_manager.clone()); + let mut manager = Self::with_settings(settings, endpoint, self.auth_manager.clone()); if matches!(self.catalog_source, CatalogSource::ExplicitProvider(_)) { manager = manager.with_provider_catalog(); } - manager.api_key_model_discovery_enabled = Arc::clone(&self.api_key_model_discovery_enabled); let manager: SharedModelsManager = Arc::new(manager); managers.insert(originator, Arc::clone(&manager)); Some(manager) } fn set_api_key_model_discovery_enabled(&self, enabled: bool) { - self.api_key_model_discovery_enabled - .store(enabled, Ordering::SeqCst); + self.settings.set_api_key_model_discovery_enabled(enabled); } fn refresh_after_auth_change( @@ -657,7 +696,7 @@ impl OpenAiModelsManager { && Some(&identity) == self.endpoint_client.identity().as_ref() && current_etag.as_deref() == Some(etag.as_str()) { - if let Some(cache) = self.cache.as_ref() + if let Some(cache) = self.settings.cache.as_ref() && let Err(err) = cache .refresh_ttl(&crate::client_version_to_whole(), &identity, &etag) .await @@ -761,7 +800,7 @@ impl OpenAiModelsManager { self.catalog_updates.send_replace(None); } if let Some(identity) = request_identity - && let Some(cache) = self.cache.as_ref() + && let Some(cache) = self.settings.cache.as_ref() { let invalidation = async { if let Some(mut entry) = cache.load(&client_version).await? @@ -798,7 +837,7 @@ impl OpenAiModelsManager { identity: Some(identity), models, }; - if let Some(cache) = self.cache.as_ref() + if let Some(cache) = self.settings.cache.as_ref() && let Err(err) = cache.store(&entry).await { error!("failed to write models cache: {err}"); @@ -821,7 +860,7 @@ impl OpenAiModelsManager { self.uses_api_key_auth() && !self.endpoint_client.has_command_auth() && (!self.endpoint_client.supports_api_key_models() - || !self.api_key_model_discovery_enabled.load(Ordering::SeqCst)) + || !self.settings.api_key_model_discovery_enabled()) } fn uses_api_key_auth(&self) -> bool { @@ -874,7 +913,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 { - let Some(cache) = self.cache.as_ref() else { + let Some(cache) = self.settings.cache.as_ref() else { return false; }; let _timer = diff --git a/codex-rs/runtime-agent-instance/src/manager.rs b/codex-rs/runtime-agent-instance/src/manager.rs index aa526b679d..15d7b42c5e 100644 --- a/codex-rs/runtime-agent-instance/src/manager.rs +++ b/codex-rs/runtime-agent-instance/src/manager.rs @@ -9,6 +9,7 @@ use crate::RuntimeAuth; use crate::RuntimeRegistration; use crate::RuntimeRegistrationClient; use crate::RuntimeRegistrationMetadata; +use codex_agent_identity::ChatGptEnvironment; use codex_http_client::HttpClientFactory; use codex_http_client::NetworkPolicyDenied; use codex_login::AuthChangeState; @@ -23,12 +24,52 @@ use url::Url; /// Owned alongside an AuthManager, not created for every thread or request. pub struct RuntimeAuthManager { credentials: Arc, - host: AgentHostAuthorization, + registration: RegistrationTarget, metadata: RuntimeRegistrationMetadata, - registration_url: Url, state: Mutex>, } +enum RegistrationTarget { + Explicit { + host: AgentHostAuthorization, + url: Url, + }, + ChatGpt { + installation_id: String, + base_url: String, + }, +} + +impl RegistrationTarget { + fn host(&self) -> Result { + match self { + Self::Explicit { host, .. } => Ok(host.clone()), + Self::ChatGpt { + installation_id, .. + } => Ok(AgentHostAuthorization::InstallationId { + installation_id: installation_id + .parse() + .map_err(|_| RegistrationError::InvalidRequest)?, + }), + } + } + + fn url(&self) -> Result { + match self { + Self::Explicit { url, .. } => Ok(url.clone()), + Self::ChatGpt { base_url, .. } => { + let url = match ChatGptEnvironment::from_chatgpt_base_url(base_url) + .map_err(|_| RegistrationError::InvalidEndpoint)? + { + ChatGptEnvironment::Production => crate::DEFAULT_REGISTRATION_URL, + ChatGptEnvironment::Staging => "https://auth.api.openai.org/agent/register", + }; + url.parse().map_err(|_| RegistrationError::InvalidEndpoint) + } + } + } +} + struct CachedRuntime { source_generation: u64, selection: PlatformResourceSelection, @@ -61,9 +102,29 @@ impl RuntimeAuthManager { ) -> Self { Self { credentials, - host, + registration: RegistrationTarget::Explicit { + host, + url: registration_url, + }, + metadata, + state: Mutex::new(None), + } + } + + /// Resolve the ChatGPT registration origin only when a first-party request needs a runtime. + pub fn for_chatgpt( + credentials: Arc, + installation_id: String, + base_url: String, + metadata: RuntimeRegistrationMetadata, + ) -> Self { + Self { + credentials, + registration: RegistrationTarget::ChatGpt { + installation_id, + base_url, + }, metadata, - registration_url, state: Mutex::new(None), } } @@ -149,7 +210,7 @@ impl RuntimeAuthManager { state: RuntimeState::Pending { registration: RuntimeRegistration::new( actor_auth.clone(), - self.host.clone(), + self.registration.host()?, self.metadata.clone(), ) .map_err(|_| RuntimeAuthError::KeyGeneration)?, @@ -172,10 +233,8 @@ impl RuntimeAuthManager { registration } }; - let client = RuntimeRegistrationClient::new( - http_client_factory.clone(), - self.registration_url.clone(), - )?; + let client = + RuntimeRegistrationClient::new(http_client_factory.clone(), self.registration.url()?)?; let result = client.register(registration).await; if self.auth_state(&auth)?.source_generation != change_state.source_generation { // A late response must not publish an old identity after a credential switch. diff --git a/codex-rs/thread-manager-sample/src/main.rs b/codex-rs/thread-manager-sample/src/main.rs index b065645a0b..ecee9645d4 100644 --- a/codex-rs/thread-manager-sample/src/main.rs +++ b/codex-rs/thread-manager-sample/src/main.rs @@ -35,6 +35,7 @@ use codex_core_api::GhostSnapshotConfig; use codex_core_api::History; use codex_core_api::MemoriesConfig; use codex_core_api::ModelAvailabilityNuxConfig; +use codex_core_api::ModelCatalogSettings; use codex_core_api::MultiAgentV2Config; use codex_core_api::NewThread; use codex_core_api::Notice; @@ -64,7 +65,6 @@ use codex_core_api::UserInput; use codex_core_api::WebSearchMode; use codex_core_api::arg0_dispatch_or_else; use codex_core_api::bootstrap_auth_config; -use codex_core_api::build_models_manager; use codex_core_api::built_in_model_providers; use codex_core_api::find_codex_home; use codex_core_api::init_state_db; @@ -146,10 +146,15 @@ async fn run_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { install_image_generation_extension(&mut extensions, auth_manager.clone(), |config: &Config| { Some(config.codex_home.clone()) }); - let thread_manager = ThreadManager::new( + let startup = codex_core_api::build_startup_models( &config, Arc::clone(&auth_manager), - build_models_manager(&config, auth_manager), + installation_id, + Arc::new(ModelCatalogSettings::disk(config.codex_home.to_path_buf())), + ); + let thread_manager = ThreadManager::new( + &config, + startup, CodexAppsToolsCache::default(), SessionSource::Exec, environment_manager, @@ -159,7 +164,6 @@ async fn run_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { passthrough_image_store(), Arc::clone(&thread_store), local_agent_graph_store_from_state_db(state_db.as_ref()), - installation_id, /*attestation_provider*/ None, /*external_time_provider*/ None, ); @@ -394,6 +398,7 @@ async fn new_config( rollout_budget: None, current_time_reminder: None, sleep_tool_mode: Default::default(), + runtime_auth_mode: Default::default(), features: Default::default(), prefer_mxc: false, diff --git a/codex-rs/websocket-client/src/lib.rs b/codex-rs/websocket-client/src/lib.rs index 4bd8513fe7..ee97e5a8e2 100644 --- a/codex-rs/websocket-client/src/lib.rs +++ b/codex-rs/websocket-client/src/lib.rs @@ -112,6 +112,20 @@ impl WebSocketConnector { self } + /// Reuse the TLS/proxy configuration with a captured credential's account policy. + pub fn with_account_binding(&self, policy: &codex_http_client::NetworkPolicy) -> Option { + let policy = self + .http_client_factory + .network_policy() + .clone() + .with_account_binding(policy)?; + Some(Self { + http_client_factory: self.http_client_factory.clone().with_network_policy(policy), + tls_config: self.tls_config.clone(), + tcp_nodelay: self.tcp_nodelay, + }) + } + /// Connects a WebSocket after resolving the request destination through the configured proxy /// policy. pub async fn connect(