From d83d95231107ecfd956a67a88800e7b1e8224a1c Mon Sep 17 00:00:00 2001 From: celia-oai Date: Thu, 9 Jul 2026 10:39:29 -0700 Subject: [PATCH] Refresh loaded model provider sessions --- .../request_processors/account_processor.rs | 38 ++-- .../request_processors/thread_processor.rs | 61 +++-- .../thread_processor_tests.rs | 44 ++++ codex-rs/app-server/tests/suite/v2/account.rs | 59 +++++ .../suite/v2/provider_runtime_refresh.rs | 151 +++++++++++++ codex-rs/core/src/client.rs | 52 +++++ codex-rs/core/src/codex_delegate.rs | 7 + codex-rs/core/src/compact.rs | 2 +- .../core/src/compact_remote_v2_attempt.rs | 2 +- codex-rs/core/src/guardian/review_session.rs | 2 +- codex-rs/core/src/lib.rs | 1 + codex-rs/core/src/model_provider_runtime.rs | 126 +++++++++-- codex-rs/core/src/realtime_conversation.rs | 2 +- codex-rs/core/src/responses_retry.rs | 2 +- codex-rs/core/src/session/mod.rs | 6 + codex-rs/core/src/session/review.rs | 9 +- codex-rs/core/src/session/session.rs | 72 +++--- codex-rs/core/src/session/tests.rs | 81 ++++++- .../core/src/session/tests/guardian_tests.rs | 7 + codex-rs/core/src/session/turn.rs | 6 +- codex-rs/core/src/session/turn_context.rs | 210 +++++++++++++++--- codex-rs/core/src/session_startup_prewarm.rs | 6 +- codex-rs/core/src/state/mod.rs | 1 + codex-rs/core/src/state/service.rs | 18 ++ codex-rs/core/src/tasks/mod.rs | 5 - codex-rs/core/src/tasks/review.rs | 2 +- codex-rs/core/src/thread_manager.rs | 80 +++++-- codex-rs/core/src/thread_manager_tests.rs | 138 +++++++++++- codex-rs/login/src/auth/bedrock_api_key.rs | 92 +++++++- .../login/src/auth/bedrock_api_key_tests.rs | 138 +++++++++--- codex-rs/login/src/auth/manager.rs | 122 +++++++++- .../model-provider/src/amazon_bedrock/mod.rs | 43 +--- codex-rs/model-provider/src/provider.rs | 16 +- 33 files changed, 1348 insertions(+), 253 deletions(-) diff --git a/codex-rs/app-server/src/request_processors/account_processor.rs b/codex-rs/app-server/src/request_processors/account_processor.rs index 6975f2ef9c..727b1ef416 100644 --- a/codex-rs/app-server/src/request_processors/account_processor.rs +++ b/codex-rs/app-server/src/request_processors/account_processor.rs @@ -184,12 +184,15 @@ impl AccountRequestProcessor { .set_auth_mode(self.auth_manager.get_api_auth_mode()); } - fn current_account_updated_notification(&self) -> AccountUpdatedNotification { - let auth = self.auth_manager.auth_cached(); + async fn current_account_updated_notification(&self) -> AccountUpdatedNotification { + let config = self.load_latest_config().await; + let provider = + create_model_provider(config.model_provider, Some(self.auth_manager.clone())); + let auth = provider.auth().await; AccountUpdatedNotification { auth_mode: auth .as_ref() - .map(CodexAuth::api_auth_mode) + .map(CodexAuth::auth_mode) .map(auth_mode_to_api), plan_type: auth.as_ref().and_then(CodexAuth::account_plan_type), } @@ -424,7 +427,11 @@ impl AccountRequestProcessor { self.config.auth_keyring_backend_kind(), ) .map_err(|err| internal_error(format!("failed to save Amazon Bedrock auth: {err}")))?; - self.auth_manager.reload().await; + self.auth_manager + .reload_bedrock_api_key_auth() + .map_err(|err| { + internal_error(format!("failed to reload Amazon Bedrock auth: {err}")) + })?; refresh_default_model_provider(&self.config_manager, &self.thread_manager).await?; Ok(LoginAccountResponse::AmazonBedrock {}) } @@ -782,7 +789,7 @@ impl AccountRequestProcessor { self.outgoing .send_server_notification(ServerNotification::AccountUpdated( - self.current_account_updated_notification(), + self.current_account_updated_notification().await, )) .await; } @@ -835,11 +842,9 @@ impl AccountRequestProcessor { } async fn logout_common(&self) -> std::result::Result, JSONRPCErrorError> { - let managed_bedrock_auth = matches!( - self.auth_manager.auth_cached(), - Some(CodexAuth::BedrockApiKey(_)) - ); let config = self.load_latest_config().await; + let managed_bedrock_auth = config.model_provider.is_amazon_bedrock() + && self.auth_manager.bedrock_api_key_auth_cached().is_some(); if config.model_provider.is_amazon_bedrock() && !managed_bedrock_auth { return Err(invalid_request( "cannot log out while Amazon Bedrock is using AWS-managed credentials; manage those credentials through AWS or switch model providers before logging out Codex authentication", @@ -854,16 +859,17 @@ impl AccountRequestProcessor { } } - match self.auth_manager.logout_with_revoke().await { - Ok(_) => {} - Err(err) => { - return Err(internal_error(format!("logout failed: {err}"))); - } - } - if managed_bedrock_auth { + self.auth_manager + .logout_bedrock_api_key_auth() + .map_err(|err| internal_error(format!("logout failed: {err}")))?; clear_user_model_provider_if_bedrock(&self.config_manager).await?; refresh_default_model_provider(&self.config_manager, &self.thread_manager).await?; + } else { + self.auth_manager + .logout_with_revoke() + .await + .map_err(|err| internal_error(format!("logout failed: {err}")))?; } Self::maybe_refresh_plugin_caches_for_current_config( diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index 271573f41e..44a5560024 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -1,6 +1,7 @@ use super::*; use crate::error_code::method_not_found; use codex_app_server_protocol::SelectedCapabilityRoot; +use codex_core::ModelProviderSelection; use codex_extension_api::ExtensionDataInit; use codex_protocol::config_types::MultiAgentMode; use codex_protocol::models::BUILT_IN_PERMISSION_PROFILE_DANGER_FULL_ACCESS; @@ -23,6 +24,19 @@ struct ThreadListFilters { relation_filter: Option, } +fn thread_start_model_provider_selection( + model_provider: Option<&str>, + request_config: Option<&HashMap>, +) -> ModelProviderSelection { + if model_provider.is_some() + || request_config.is_some_and(|config| config.contains_key("model_provider")) + { + ModelProviderSelection::Explicit + } else { + ModelProviderSelection::RuntimeDefault + } +} + fn collect_resume_override_mismatches( request: &ThreadResumeParams, config_snapshot: &ThreadConfigSnapshot, @@ -968,6 +982,8 @@ impl ThreadRequestProcessor { "`permissions` cannot be combined with `sandbox`", )); } + let model_provider_selection = + thread_start_model_provider_selection(model_provider.as_deref(), config.as_ref()); let environment_selections = resolve_turn_environment_selections(self.thread_manager.as_ref(), environments)?; let runtime_workspace_roots = runtime_workspace_roots.map(resolve_runtime_workspace_roots); @@ -1012,6 +1028,7 @@ impl ThreadRequestProcessor { supports_openai_form_elicitation, config, typesafe_overrides, + model_provider_selection, dynamic_tools, selected_capability_roots.unwrap_or_default(), history_mode.map(Into::into), @@ -1089,6 +1106,7 @@ impl ThreadRequestProcessor { supports_openai_form_elicitation: bool, config_overrides: Option>, typesafe_overrides: ConfigOverrides, + model_provider_selection: ModelProviderSelection, dynamic_tools: Option>, selected_capability_roots: Vec, history_mode: Option, @@ -1217,25 +1235,32 @@ impl ThreadRequestProcessor { .. } = listener_task_context .thread_manager - .start_thread_with_options(StartThreadOptions { - config, - allow_provider_model_fallback, - initial_history: match session_start_source - .unwrap_or(codex_app_server_protocol::ThreadStartSource::Startup) - { - codex_app_server_protocol::ThreadStartSource::Startup => InitialHistory::New, - codex_app_server_protocol::ThreadStartSource::Clear => InitialHistory::Cleared, + .start_thread_with_options_and_provider_selection( + StartThreadOptions { + config, + allow_provider_model_fallback, + initial_history: match session_start_source + .unwrap_or(codex_app_server_protocol::ThreadStartSource::Startup) + { + codex_app_server_protocol::ThreadStartSource::Startup => { + InitialHistory::New + } + codex_app_server_protocol::ThreadStartSource::Clear => { + InitialHistory::Cleared + } + }, + history_mode, + session_source: None, + thread_source, + dynamic_tools, + metrics_service_name: service_name, + parent_trace: request_trace, + environments, + thread_extension_init, + supports_openai_form_elicitation, }, - history_mode, - session_source: None, - thread_source, - dynamic_tools, - metrics_service_name: service_name, - parent_trace: request_trace, - environments, - thread_extension_init, - supports_openai_form_elicitation, - }) + model_provider_selection, + ) .instrument(tracing::info_span!( "app_server.thread_start.create_thread", otel.name = "app_server.thread_start.create_thread", diff --git a/codex-rs/app-server/src/request_processors/thread_processor_tests.rs b/codex-rs/app-server/src/request_processors/thread_processor_tests.rs index 8c0d59efed..4175ae701e 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor_tests.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor_tests.rs @@ -1,3 +1,47 @@ +mod model_provider_selection_tests { + use super::super::thread_start_model_provider_selection; + use codex_core::ModelProviderSelection; + use pretty_assertions::assert_eq; + use serde_json::json; + use std::collections::HashMap; + + #[test] + fn omitted_provider_follows_the_runtime_default() { + assert_eq!( + thread_start_model_provider_selection( + /*model_provider*/ None, /*request_config*/ None + ), + ModelProviderSelection::RuntimeDefault + ); + + let request_config = HashMap::from([("model".to_string(), json!("mock-model"))]); + assert_eq!( + thread_start_model_provider_selection( + /*model_provider*/ None, + Some(&request_config), + ), + ModelProviderSelection::RuntimeDefault + ); + } + + #[test] + fn typed_or_config_provider_override_is_explicit() { + assert_eq!( + thread_start_model_provider_selection(Some("provider-a"), /*request_config*/ None,), + ModelProviderSelection::Explicit + ); + + let request_config = HashMap::from([("model_provider".to_string(), json!("provider-a"))]); + assert_eq!( + thread_start_model_provider_selection( + /*model_provider*/ None, + Some(&request_config), + ), + ModelProviderSelection::Explicit + ); + } +} + mod thread_list_cwd_filter_tests { use super::super::normalize_thread_list_cwd_filters; use codex_app_server_protocol::ThreadListCwdFilter; diff --git a/codex-rs/app-server/tests/suite/v2/account.rs b/codex-rs/app-server/tests/suite/v2/account.rs index 09a2bf7d95..f396504e4f 100644 --- a/codex-rs/app-server/tests/suite/v2/account.rs +++ b/codex-rs/app-server/tests/suite/v2/account.rs @@ -1386,6 +1386,65 @@ async fn logout_managed_bedrock_restores_default_account() -> Result<()> { Ok(()) } +#[tokio::test] +async fn managed_bedrock_login_and_logout_preserve_cached_openai_auth() -> Result<()> { + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), CreateConfigTomlParams::default())?; + login_with_api_key( + codex_home.path(), + "sk-test-key", + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + + let mut mcp = TestAppServer::builder() + .with_codex_home(codex_home.path()) + .without_auto_env() + .with_env_overrides(&[("OPENAI_API_KEY", None)]) + .build() + .await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let request_id = mcp + .send_login_account_amazon_bedrock_request("managed-bedrock-api-key", "us-west-2") + .await?; + let response = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + assert_eq!( + to_response::(response)?, + LoginAccountResponse::AmazonBedrock {} + ); + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("account/login/completed"), + ) + .await??; + assert_account_updated(&mut mcp, Some(AuthMode::BedrockApiKey)).await?; + + let request_id = mcp.send_logout_account_request().await?; + let response = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + assert_eq!( + to_response::(response)?, + LogoutAccountResponse {} + ); + assert_account_updated(&mut mcp, Some(AuthMode::ApiKey)).await?; + assert_eq!( + read_account(&mut mcp).await?, + GetAccountResponse { + account: Some(Account::ApiKey {}), + requires_openai_auth: true, + } + ); + Ok(()) +} + #[tokio::test] async fn logout_aws_managed_bedrock_errors_without_changing_auth_or_config() -> Result<()> { let codex_home = TempDir::new()?; diff --git a/codex-rs/app-server/tests/suite/v2/provider_runtime_refresh.rs b/codex-rs/app-server/tests/suite/v2/provider_runtime_refresh.rs index f4c1d9b712..3d96df9560 100644 --- a/codex-rs/app-server/tests/suite/v2/provider_runtime_refresh.rs +++ b/codex-rs/app-server/tests/suite/v2/provider_runtime_refresh.rs @@ -17,8 +17,13 @@ use codex_app_server_protocol::ModelProviderCapabilitiesReadResponse; use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ThreadListParams; use codex_app_server_protocol::ThreadListResponse; +use codex_app_server_protocol::ThreadStartParams; +use codex_app_server_protocol::ThreadStartResponse; +use codex_app_server_protocol::TurnStartParams; +use codex_app_server_protocol::UserInput; use codex_app_server_protocol::WriteStatus; use codex_model_provider_info::AMAZON_BEDROCK_GPT_5_5_MODEL_ID; +use core_test_support::responses; use pretty_assertions::assert_eq; use serde_json::json; use tempfile::TempDir; @@ -107,6 +112,53 @@ fn bedrock_capabilities() -> ModelProviderCapabilitiesReadResponse { } } +async fn start_provider_runtime_thread( + mcp: &mut TestAppServer, + model_provider: Option<&str>, +) -> Result { + let request_id = mcp + .send_thread_start_request_with_auto_env(ThreadStartParams { + model: Some("mock-model".to_string()), + model_provider: model_provider.map(str::to_string), + ..Default::default() + }) + .await?; + let response: JSONRPCResponse = timeout( + DEFAULT_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + to_response(response) +} + +async fn run_provider_runtime_turn( + mcp: &mut TestAppServer, + thread_id: &str, + text: &str, +) -> Result<()> { + let request_id = mcp + .send_turn_start_request(TurnStartParams { + thread_id: thread_id.to_string(), + input: vec![UserInput::Text { + text: text.to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }) + .await?; + timeout( + DEFAULT_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + timeout( + DEFAULT_TIMEOUT, + mcp.read_stream_until_notification_message("turn/completed"), + ) + .await??; + Ok(()) +} + #[tokio::test] async fn config_write_publishes_provider_runtime_before_responding() -> Result<()> { let codex_home = TempDir::new()?; @@ -227,3 +279,102 @@ async fn managed_bedrock_login_and_logout_publish_provider_runtime() -> Result<( ); Ok(()) } + +#[tokio::test] +async fn loaded_threads_adopt_provider_switch_only_when_following_runtime_default() -> Result<()> { + let provider_a_server = responses::start_mock_server().await; + let provider_b_server = responses::start_mock_server().await; + let provider_a_mock = responses::mount_sse_once( + &provider_a_server, + responses::sse(vec![ + responses::ev_response_created("provider-a-response"), + responses::ev_assistant_message("provider-a-message", "provider A"), + responses::ev_completed("provider-a-response"), + ]), + ) + .await; + let provider_b_mock = responses::mount_sse_once( + &provider_b_server, + responses::sse(vec![ + responses::ev_response_created("provider-b-response"), + responses::ev_assistant_message("provider-b-message", "provider B"), + responses::ev_completed("provider-b-response"), + ]), + ) + .await; + + let codex_home = TempDir::new()?; + std::fs::write( + codex_home.path().join("config.toml"), + format!( + r#" +model = "mock-model" +model_provider = "provider_a" + +[model_providers.provider_a] +name = "Provider A" +base_url = "{}/v1" +wire_api = "responses" +request_max_retries = 0 +stream_max_retries = 0 + +[model_providers.provider_b] +name = "Provider B" +base_url = "{}/v1" +wire_api = "responses" +request_max_retries = 0 +stream_max_retries = 0 +"#, + provider_a_server.uri(), + provider_b_server.uri(), + ), + )?; + let mut mcp = TestAppServer::builder() + .with_codex_home(codex_home.path()) + .build() + .await?; + timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; + + let runtime_default_thread = start_provider_runtime_thread(&mut mcp, None).await?; + let explicit_thread = start_provider_runtime_thread(&mut mcp, Some("provider_a")).await?; + assert_eq!(runtime_default_thread.model_provider, "provider_a"); + assert_eq!(explicit_thread.model_provider, "provider_a"); + + let request_id = mcp + .send_config_value_write_request(ConfigValueWriteParams { + key_path: "model_provider".to_string(), + value: json!("provider_b"), + merge_strategy: MergeStrategy::Replace, + file_path: None, + expected_version: None, + }) + .await?; + let response: JSONRPCResponse = timeout( + DEFAULT_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + assert_eq!( + to_response::(response)?.status, + WriteStatus::Ok + ); + + // Both threads were idle when the runtime changed. Selection is applied at the next turn + // boundary, so no in-flight request is expected to migrate between providers. + run_provider_runtime_turn( + &mut mcp, + &runtime_default_thread.thread.id, + "use the refreshed runtime provider", + ) + .await?; + run_provider_runtime_turn( + &mut mcp, + &explicit_thread.thread.id, + "stay on the explicitly selected provider", + ) + .await?; + + assert_eq!(provider_b_mock.requests().len(), 1); + assert_eq!(provider_a_mock.requests().len(), 1); + Ok(()) +} diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 251b7c1dcf..30abb05721 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -426,6 +426,39 @@ impl ModelClient { http_client_factory: HttpClientFactory, ) -> Self { let model_provider = create_model_provider(provider_info, auth_manager); + Self::new_with_provider( + model_provider, + agent_identity_policy, + thread_id, + session_source, + originator, + model_verbosity, + enable_request_compression, + include_timing_metrics, + beta_features_header, + item_ids_enabled, + concurrent_reasoning_summaries_enabled, + attestation_provider, + http_client_factory, + ) + } + + #[allow(clippy::too_many_arguments)] + pub(crate) fn new_with_provider( + model_provider: SharedModelProvider, + agent_identity_policy: AgentIdentityAuthPolicy, + thread_id: ThreadId, + session_source: SessionSource, + originator: String, + model_verbosity: Option, + enable_request_compression: bool, + include_timing_metrics: bool, + beta_features_header: Option, + item_ids_enabled: bool, + concurrent_reasoning_summaries_enabled: bool, + attestation_provider: Option>, + http_client_factory: HttpClientFactory, + ) -> Self { let codex_api_key_env_enabled = model_provider .auth_manager() .as_ref() @@ -458,6 +491,25 @@ impl ModelClient { } } + pub(crate) fn rebind_provider(&self, model_provider: SharedModelProvider) -> Self { + Self::new_with_provider( + model_provider, + self.agent_identity_policy, + self.state.thread_id, + self.state.session_source.clone(), + self.state.originator.clone(), + self.state.model_verbosity, + self.state.enable_request_compression, + self.state.include_timing_metrics, + self.state.beta_features_header.clone(), + self.state.item_ids_enabled, + self.state.concurrent_reasoning_summaries_enabled, + self.state.attestation_provider.clone(), + self.http_client_factory.clone(), + ) + .with_prompt_cache_key_override(self.prompt_cache_key_override.clone()) + } + pub(crate) fn with_prompt_cache_key_override( mut self, prompt_cache_key_override: Option, diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index d2911df407..2dc47ffd6e 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -45,6 +45,7 @@ use crate::mcp_tool_call::build_guardian_mcp_tool_review_request; use crate::mcp_tool_call::is_mcp_tool_approval_question_id; use crate::mcp_tool_call::lookup_mcp_tool_metadata; use crate::mcp_tool_call::mcp_approvals_reviewer; +use crate::model_provider_runtime::build_explicit_model_provider_runtime_with_models_manager; use crate::session::Codex; use crate::session::CodexSpawnArgs; use crate::session::CodexSpawnOk; @@ -91,6 +92,11 @@ pub(crate) async fn run_codex_thread_interactive( instructions: parent_session.user_instructions().await, warnings: Vec::new(), }; + let model_provider_runtime = build_explicit_model_provider_runtime_with_models_manager( + &config, + Arc::clone(&auth_manager), + Arc::clone(&models_manager), + ); let CodexSpawnOk { codex, .. } = Box::pin(Codex::spawn(CodexSpawnArgs { config, allow_provider_model_fallback: false, @@ -98,6 +104,7 @@ pub(crate) async fn run_codex_thread_interactive( installation_id: parent_session.installation_id.clone(), auth_manager, models_manager, + model_provider_runtime, environment_manager: parent_session .services .turn_environments diff --git a/codex-rs/core/src/compact.rs b/codex-rs/core/src/compact.rs index d2a9fd5222..a318d13199 100644 --- a/codex-rs/core/src/compact.rs +++ b/codex-rs/core/src/compact.rs @@ -237,7 +237,7 @@ async fn run_compact_task_inner_impl( let max_retries = turn_context.provider.info().stream_max_retries(); let mut retries = 0; - let mut client_session = sess.services.model_client.new_session(); + let mut client_session = turn_context.model_client.new_session(); // Reuse one client session so turn-scoped state (sticky routing, websocket incremental // request tracking) // survives retries within this compact turn. diff --git a/codex-rs/core/src/compact_remote_v2_attempt.rs b/codex-rs/core/src/compact_remote_v2_attempt.rs index a145d00ae6..ebc9ce3307 100644 --- a/codex-rs/core/src/compact_remote_v2_attempt.rs +++ b/codex-rs/core/src/compact_remote_v2_attempt.rs @@ -97,7 +97,7 @@ pub(super) async fn run_remote_compact_v2_attempt( let mut owned_client_session = None; let client_session = match client_session { Some(client_session) => client_session, - None => owned_client_session.insert(sess.services.model_client.new_session()), + None => owned_client_session.insert(turn_context.model_client.new_session()), }; let compaction_output_result = run_remote_compaction_request_v2( sess, diff --git a/codex-rs/core/src/guardian/review_session.rs b/codex-rs/core/src/guardian/review_session.rs index 05ef58b4dd..5d5d5a05db 100644 --- a/codex-rs/core/src/guardian/review_session.rs +++ b/codex-rs/core/src/guardian/review_session.rs @@ -665,7 +665,7 @@ async fn spawn_guardian_review_session( let codex = Box::pin(run_codex_thread_interactive( spawn_config, parent_session.services.auth_manager.clone(), - parent_session.services.models_manager.clone(), + Arc::clone(&parent_turn.models_manager), Arc::clone(parent_session), Arc::clone(parent_turn), cancel_token.clone(), diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 9487344fa8..bc116a736a 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -116,6 +116,7 @@ mod thread_manager; pub(crate) mod web_search; pub(crate) mod windows_sandbox_read_grants; pub use thread_manager::ForkSnapshot; +pub use thread_manager::ModelProviderSelection; pub use thread_manager::NewThread; pub use thread_manager::StartThreadOptions; pub use thread_manager::ThreadManager; diff --git a/codex-rs/core/src/model_provider_runtime.rs b/codex-rs/core/src/model_provider_runtime.rs index e8c2c64c0c..0b5c171e0f 100644 --- a/codex-rs/core/src/model_provider_runtime.rs +++ b/codex-rs/core/src/model_provider_runtime.rs @@ -2,6 +2,7 @@ use crate::config::Config; use arc_swap::ArcSwap; use codex_login::AuthManager; use codex_model_provider::ProviderCapabilities; +use codex_model_provider::SharedModelProvider; use codex_model_provider::create_model_provider; use codex_model_provider_info::ModelProviderInfo; use codex_models_manager::manager::SharedModelsManager; @@ -17,35 +18,44 @@ struct ModelProviderRuntimeIdentity { provider_info: ModelProviderInfo, codex_home: PathBuf, config_model_catalog: Option, + auth_revision: u64, } impl ModelProviderRuntimeIdentity { - fn from_config(config: &Config) -> Self { + fn from_config(config: &Config, auth_manager: &AuthManager) -> Self { Self { provider_id: config.model_provider_id.clone(), provider_info: config.model_provider.clone(), codex_home: config.codex_home.to_path_buf(), config_model_catalog: config.model_catalog.clone(), + auth_revision: auth_revision(auth_manager), } } - fn matches_config(&self, config: &Config) -> bool { + fn matches_config(&self, config: &Config, auth_manager: &AuthManager) -> bool { self.provider_id == config.model_provider_id && self.provider_info == config.model_provider && self.codex_home == config.codex_home.to_path_buf() && self.config_model_catalog == config.model_catalog + && self.auth_revision == auth_revision(auth_manager) } } +fn auth_revision(auth_manager: &AuthManager) -> u64 { + let receiver = auth_manager.auth_change_receiver(); + *receiver.borrow() +} + /// Immutable process-default model-provider state published as one coherent value. -struct DefaultModelProviderSnapshot { +pub(crate) struct ModelProviderRuntimeSnapshot { generation: u64, identity: ModelProviderRuntimeIdentity, + provider: SharedModelProvider, models_manager: SharedModelsManager, capabilities: ProviderCapabilities, } -impl DefaultModelProviderSnapshot { +impl ModelProviderRuntimeSnapshot { fn new( generation: u64, identity: ModelProviderRuntimeIdentity, @@ -55,34 +65,78 @@ impl DefaultModelProviderSnapshot { identity.provider_info.clone(), Some(Arc::clone(&auth_manager)), ); - let capabilities = provider.capabilities(); let models_manager = provider.models_manager( identity.codex_home.clone(), identity.config_model_catalog.clone(), ); + Self::from_provider(generation, identity, provider, models_manager) + } + + fn from_provider( + generation: u64, + identity: ModelProviderRuntimeIdentity, + provider: SharedModelProvider, + models_manager: SharedModelsManager, + ) -> Self { + let capabilities = provider.capabilities(); Self { generation, identity, + provider, models_manager, capabilities, } } - fn matches_config(&self, config: &Config) -> bool { - self.identity.matches_config(config) + fn matches_config(&self, config: &Config, auth_manager: &AuthManager) -> bool { + self.identity.matches_config(config, auth_manager) } + + pub(crate) fn generation(&self) -> u64 { + self.generation + } + + pub(crate) fn provider_info(&self) -> &ModelProviderInfo { + &self.identity.provider_info + } + + pub(crate) fn provider(&self) -> SharedModelProvider { + Arc::clone(&self.provider) + } + + pub(crate) fn models_manager(&self) -> SharedModelsManager { + Arc::clone(&self.models_manager) + } + + pub(crate) fn apply_to_config(&self, config: &mut Config) { + config.model_provider_id = self.identity.provider_id.clone(); + config.model_provider = self.identity.provider_info.clone(); + config.model_catalog = self.identity.config_model_catalog.clone(); + } +} + +#[derive(Clone)] +pub(crate) enum ModelProviderRuntimeSource { + RuntimeDefault(Arc), + Explicit, +} + +#[derive(Clone)] +pub(crate) struct InitialModelProviderRuntime { + pub(crate) source: ModelProviderRuntimeSource, + pub(crate) snapshot: Arc, } /// Owns the atomically replaceable process-default model-provider snapshot. pub(crate) struct DefaultModelProviderRuntime { - snapshot: ArcSwap, + snapshot: ArcSwap, refresh_lock: Mutex<()>, } impl DefaultModelProviderRuntime { pub(crate) fn new(config: &Config, auth_manager: Arc) -> Self { Self::from_identity( - ModelProviderRuntimeIdentity::from_config(config), + ModelProviderRuntimeIdentity::from_config(config, &auth_manager), auth_manager, ) } @@ -100,6 +154,7 @@ impl DefaultModelProviderRuntime { provider_info, codex_home, config_model_catalog, + auth_revision: auth_revision(&auth_manager), }, auth_manager, ) @@ -110,7 +165,7 @@ impl DefaultModelProviderRuntime { auth_manager: Arc, ) -> Self { Self { - snapshot: ArcSwap::from_pointee(DefaultModelProviderSnapshot::new( + snapshot: ArcSwap::from_pointee(ModelProviderRuntimeSnapshot::new( /*generation*/ 0, identity, auth_manager, @@ -125,13 +180,13 @@ impl DefaultModelProviderRuntime { .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); let current = self.snapshot.load_full(); - if current.matches_config(config) { + if current.matches_config(config, &auth_manager) { return false; } - let next = DefaultModelProviderSnapshot::new( + let next = ModelProviderRuntimeSnapshot::new( current.generation.saturating_add(1), - ModelProviderRuntimeIdentity::from_config(config), + ModelProviderRuntimeIdentity::from_config(config, &auth_manager), auth_manager, ); self.snapshot.store(Arc::new(next)); @@ -142,17 +197,8 @@ impl DefaultModelProviderRuntime { Arc::clone(&self.snapshot.load().models_manager) } - pub(crate) fn models_manager_for_config( - &self, - config: &Config, - auth_manager: Arc, - ) -> SharedModelsManager { - let current = self.snapshot.load(); - if current.matches_config(config) { - return Arc::clone(¤t.models_manager); - } - - build_models_manager(config, auth_manager) + pub(crate) fn snapshot(&self) -> Arc { + self.snapshot.load_full() } pub(crate) fn provider_id(&self) -> String { @@ -168,6 +214,38 @@ impl DefaultModelProviderRuntime { } } +pub(crate) fn build_explicit_model_provider_runtime( + config: &Config, + auth_manager: Arc, +) -> InitialModelProviderRuntime { + InitialModelProviderRuntime { + source: ModelProviderRuntimeSource::Explicit, + snapshot: Arc::new(ModelProviderRuntimeSnapshot::new( + /*generation*/ 0, + ModelProviderRuntimeIdentity::from_config(config, &auth_manager), + auth_manager, + )), + } +} + +pub(crate) fn build_explicit_model_provider_runtime_with_models_manager( + config: &Config, + auth_manager: Arc, + models_manager: SharedModelsManager, +) -> InitialModelProviderRuntime { + let identity = ModelProviderRuntimeIdentity::from_config(config, &auth_manager); + let provider = create_model_provider(identity.provider_info.clone(), Some(auth_manager)); + InitialModelProviderRuntime { + source: ModelProviderRuntimeSource::Explicit, + snapshot: Arc::new(ModelProviderRuntimeSnapshot::from_provider( + /*generation*/ 0, + identity, + provider, + models_manager, + )), + } +} + pub(crate) fn build_models_manager( config: &Config, auth_manager: Arc, diff --git a/codex-rs/core/src/realtime_conversation.rs b/codex-rs/core/src/realtime_conversation.rs index 3915c8190d..b1b2a6d860 100644 --- a/codex-rs/core/src/realtime_conversation.rs +++ b/codex-rs/core/src/realtime_conversation.rs @@ -1012,7 +1012,7 @@ async fn handle_start_inner( codex_response_handoff_prefix, realtime_call_api_provider, session_config, - model_client: sess.services.model_client.clone(), + model_client: sess.services.model_runtime.load().model_client.clone(), sdp, }; let start_output = sess.conversation.start(start).await?; diff --git a/codex-rs/core/src/responses_retry.rs b/codex-rs/core/src/responses_retry.rs index 9a55fc8114..9f8b43d155 100644 --- a/codex-rs/core/src/responses_retry.rs +++ b/codex-rs/core/src/responses_retry.rs @@ -60,7 +60,7 @@ pub(crate) async fn handle_retryable_response_stream_error( // transient reconnect messages. In debug builds, keep full visibility for diagnosis. let report_error = retry_count > 1 || cfg!(debug_assertions) - || !sess.services.model_client.responses_websocket_enabled(); + || !turn_context.model_client.responses_websocket_enabled(); if report_error { // Surface retry information to any UI/front-end so the user understands what is // happening instead of staring at a seemingly frozen screen. diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 9339413b33..5de251aee8 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -302,6 +302,8 @@ use crate::SkillsService; use crate::exec_policy::ExecPolicyUpdateError; use crate::guardian::GuardianReviewSessionManager; use crate::mcp::McpManager; +use crate::model_provider_runtime::InitialModelProviderRuntime; +use crate::model_provider_runtime::ModelProviderRuntimeSource; use crate::network_policy_decision::execpolicy_network_rule_amendment; use crate::rollout::map_session_init_error; use crate::session_startup_prewarm::SessionStartupPrewarmHandle; @@ -311,6 +313,7 @@ use crate::skills::SkillLoadOutcome; use crate::state::AutoCompactWindowIds; use crate::state::AutoCompactWindowSnapshot; use crate::state::PendingRequestPermissions; +use crate::state::SessionModelRuntime; use crate::state::SessionServices; use crate::state::SessionState; #[cfg(test)] @@ -412,6 +415,7 @@ pub(crate) struct CodexSpawnArgs { pub(crate) installation_id: String, pub(crate) auth_manager: Arc, pub(crate) models_manager: SharedModelsManager, + pub(crate) model_provider_runtime: InitialModelProviderRuntime, pub(crate) environment_manager: Arc, pub(crate) skills_service: Arc, pub(crate) plugins_manager: Arc, @@ -504,6 +508,7 @@ impl Codex { installation_id, auth_manager, models_manager, + model_provider_runtime, environment_manager, skills_service, plugins_manager, @@ -686,6 +691,7 @@ impl Codex { installation_id, auth_manager.clone(), models_manager.clone(), + model_provider_runtime, exec_policy, tx_event.clone(), agent_status_tx.clone(), diff --git a/codex-rs/core/src/session/review.rs b/codex-rs/core/src/session/review.rs index 6b68447b92..6029b9c679 100644 --- a/codex-rs/core/src/session/review.rs +++ b/codex-rs/core/src/session/review.rs @@ -13,8 +13,7 @@ pub(super) async fn spawn_review_thread( .review_model .clone() .unwrap_or_else(|| parent_turn_context.model_info.slug.clone()); - let review_model_info = sess - .services + let review_model_info = parent_turn_context .models_manager .get_model_info(&model, &config.to_models_manager_config()) .await; @@ -24,8 +23,7 @@ pub(super) async fn spawn_review_thread( let _ = review_features.disable(Feature::WebSearchCached); let _ = review_features.disable(Feature::Goals); let review_web_search_mode = WebSearchMode::Disabled; - let available_models = sess - .services + let available_models = parent_turn_context .models_manager .list_models( RefreshStrategy::OnlineIfUncached, @@ -117,6 +115,9 @@ pub(super) async fn spawn_review_thread( model_info: model_info.clone(), session_telemetry: session_telemetry_for_context, provider: provider_for_context, + models_manager: Arc::clone(&parent_turn_context.models_manager), + model_client: parent_turn_context.model_client.clone(), + model_provider_generation: parent_turn_context.model_provider_generation, reasoning_effort, reasoning_summary, session_source, diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 7c59ec22b7..0c57a21b66 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -484,6 +484,7 @@ impl Session { installation_id: String, auth_manager: Arc, models_manager: SharedModelsManager, + model_provider_runtime: InitialModelProviderRuntime, exec_policy: Arc, tx_event: Sender, agent_status: watch::Sender, @@ -1048,6 +1049,43 @@ impl Session { }).await; } + let model_client = ModelClient::new_with_provider( + model_provider_runtime.snapshot.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::EnableRequestCompression), + config.features.enabled(Feature::RuntimeMetrics), + Self::build_model_client_beta_features_header(config.as_ref()), + /*item_ids_enabled*/ config.features.enabled(Feature::ItemIds) + || matches!( + session_configuration.history_mode, + ThreadHistoryMode::Paginated + ), + /*concurrent_reasoning_summaries_enabled*/ config + .features + .enabled(Feature::ConcurrentReasoningSummaries), + attestation_provider.clone(), + config.http_client_factory(), + ) + .with_prompt_cache_key_override( + crate::guardian::prompt_cache_key_override_for_review_session( + &session_configuration.session_source, + session_configuration.parent_thread_id, + ), + ); + let model_runtime = SessionModelRuntime { + source: model_provider_runtime.source, + snapshot: model_provider_runtime.snapshot, + model_client: model_client.clone(), + }; + let services = SessionServices { // Initialize the MCP connection manager with an uninitialized // instance. It will be replaced with one created via @@ -1073,6 +1111,7 @@ impl Session { show_raw_agent_reasoning: config.show_raw_agent_reasoning, exec_policy, auth_manager: Arc::clone(&auth_manager), + model_runtime: arc_swap::ArcSwap::from_pointee(model_runtime), session_telemetry, models_manager: Arc::clone(&models_manager), tool_approvals: Mutex::new(ApprovalStore::default()), @@ -1102,38 +1141,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.clone(), - session_configuration.session_source.clone(), - session_configuration.originator.clone(), - config.model_verbosity, - config.features.enabled(Feature::EnableRequestCompression), - config.features.enabled(Feature::RuntimeMetrics), - Self::build_model_client_beta_features_header(config.as_ref()), - /*item_ids_enabled*/ config.features.enabled(Feature::ItemIds) - || matches!( - session_configuration.history_mode, - ThreadHistoryMode::Paginated - ), - /*concurrent_reasoning_summaries_enabled*/ config - .features - .enabled(Feature::ConcurrentReasoningSummaries), - attestation_provider, - config.http_client_factory(), - ) - .with_prompt_cache_key_override( - crate::guardian::prompt_cache_key_override_for_review_session( - &session_configuration.session_source, - session_configuration.parent_thread_id, - ), - ), + model_client, code_mode_service: crate::tools::code_mode::CodeModeService::new(Arc::clone( &code_mode_session_provider, )), diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 696a8adbd8..b3e88d82b8 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -520,6 +520,44 @@ fn test_model_client_session() -> crate::client::ModelClientSession { .new_session() } +fn session_model_runtime_for_tests( + config: &Config, + auth_manager: &Arc, + models_manager: &SharedModelsManager, + thread_id: ThreadId, + session_configuration: &SessionConfiguration, +) -> SessionModelRuntime { + let initial = + crate::model_provider_runtime::build_explicit_model_provider_runtime_with_models_manager( + config, + Arc::clone(auth_manager), + Arc::clone(models_manager), + ); + let model_client = ModelClient::new_with_provider( + initial.snapshot.provider(), + AgentIdentityAuthPolicy::JwtOnly, + thread_id, + session_configuration.session_source.clone(), + session_configuration.originator.clone(), + config.model_verbosity, + config.features.enabled(Feature::EnableRequestCompression), + config.features.enabled(Feature::RuntimeMetrics), + Session::build_model_client_beta_features_header(config), + /*item_ids_enabled*/ config.features.enabled(Feature::ItemIds), + /*concurrent_reasoning_summaries_enabled*/ + config + .features + .enabled(Feature::ConcurrentReasoningSummaries), + /*attestation_provider*/ None, + config.http_client_factory(), + ); + SessionModelRuntime { + source: initial.source, + snapshot: initial.snapshot, + model_client, + } +} + fn developer_input_texts(items: &[ResponseItem]) -> Vec<&str> { items .iter() @@ -5287,6 +5325,12 @@ async fn session_new_fails_when_zsh_fork_enabled_without_packaged_zsh() { /*bundled_skills_enabled*/ true, )); let environment_manager = Arc::new(EnvironmentManager::default_for_tests()); + let model_provider_runtime = + crate::model_provider_runtime::build_explicit_model_provider_runtime_with_models_manager( + config.as_ref(), + Arc::clone(&auth_manager), + Arc::clone(&models_manager), + ); let result = Session::new( session_configuration, Arc::clone(&config), @@ -5294,6 +5338,7 @@ async fn session_new_fails_when_zsh_fork_enabled_without_packaged_zsh() { "11111111-1111-4111-8111-111111111111".to_string(), auth_manager, models_manager, + model_provider_runtime, Arc::new(ExecPolicyManager::default()), tx_event, agent_status_tx, @@ -5473,6 +5518,13 @@ 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(), + model_runtime: arc_swap::ArcSwap::from_pointee(session_model_runtime_for_tests( + config.as_ref(), + &auth_manager, + &models_manager, + thread_id, + &session_configuration, + )), session_telemetry: session_telemetry.clone(), models_manager: Arc::clone(&models_manager), tool_approvals: Mutex::new(ApprovalStore::default()), @@ -5547,12 +5599,13 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { .skills_service .snapshot_for_config(&skills_input, Some(Arc::clone(&skill_fs))) .await; + let model_runtime = services.model_runtime.load_full(); let turn_context = Session::make_turn_context( thread_id, SessionId::from(thread_id), Some(Arc::clone(&auth_manager)), &session_telemetry, - session_configuration.provider.clone(), + model_runtime.as_ref(), &session_configuration, config.multi_agent_version_from_features(), services.user_shell.as_ref(), @@ -5560,7 +5613,6 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { services.main_execve_wrapper_exe.as_ref(), per_turn_config, model_info, - &models_manager, /*network*/ None, resolved_turn_environments, session_configuration.cwd().clone(), @@ -5673,6 +5725,12 @@ async fn make_session_with_config_and_rx( /*bundled_skills_enabled*/ true, )); let environment_manager = Arc::new(EnvironmentManager::default_for_tests()); + let model_provider_runtime = + crate::model_provider_runtime::build_explicit_model_provider_runtime_with_models_manager( + config.as_ref(), + Arc::clone(&auth_manager), + Arc::clone(&models_manager), + ); let session = Session::new( session_configuration, @@ -5681,6 +5739,7 @@ async fn make_session_with_config_and_rx( "11111111-1111-4111-8111-111111111111".to_string(), auth_manager, models_manager, + model_provider_runtime, Arc::new(ExecPolicyManager::default()), tx_event, agent_status_tx, @@ -5781,6 +5840,12 @@ async fn make_session_with_history_source_and_agent_control_and_rx( /*bundled_skills_enabled*/ true, )); let environment_manager = Arc::new(EnvironmentManager::default_for_tests()); + let model_provider_runtime = + crate::model_provider_runtime::build_explicit_model_provider_runtime_with_models_manager( + config.as_ref(), + Arc::clone(&auth_manager), + Arc::clone(&models_manager), + ); let session = Session::new( session_configuration, @@ -5789,6 +5854,7 @@ async fn make_session_with_history_source_and_agent_control_and_rx( "11111111-1111-4111-8111-111111111111".to_string(), auth_manager, models_manager, + model_provider_runtime, Arc::new(ExecPolicyManager::default()), tx_event, agent_status_tx, @@ -7604,6 +7670,13 @@ where show_raw_agent_reasoning: config.show_raw_agent_reasoning, exec_policy, auth_manager: Arc::clone(&auth_manager), + model_runtime: arc_swap::ArcSwap::from_pointee(session_model_runtime_for_tests( + config.as_ref(), + &auth_manager, + &models_manager, + thread_id, + &session_configuration, + )), session_telemetry: session_telemetry.clone(), models_manager: Arc::clone(&models_manager), tool_approvals: Mutex::new(ApprovalStore::default()), @@ -7678,12 +7751,13 @@ where .skills_service .snapshot_for_config(&skills_input, Some(Arc::clone(&skill_fs))) .await; + let model_runtime = services.model_runtime.load_full(); let turn_context = Arc::new(Session::make_turn_context( thread_id, SessionId::from(thread_id), Some(Arc::clone(&auth_manager)), &session_telemetry, - session_configuration.provider.clone(), + model_runtime.as_ref(), &session_configuration, config.multi_agent_version_from_features(), services.user_shell.as_ref(), @@ -7691,7 +7765,6 @@ where services.main_execve_wrapper_exe.as_ref(), per_turn_config, model_info, - &models_manager, /*network*/ None, resolved_turn_environments, session_configuration.cwd().clone(), diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index d13239d810..f55764fb1c 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -715,6 +715,12 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { codex_thread_store::LocalThreadStoreConfig::from_config(&config), /*state_db*/ None, )); + let model_provider_runtime = + crate::model_provider_runtime::build_explicit_model_provider_runtime_with_models_manager( + &config, + Arc::clone(&auth_manager), + Arc::clone(&models_manager), + ); let CodexSpawnOk { codex, .. } = Codex::spawn(CodexSpawnArgs { config, @@ -723,6 +729,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { installation_id: "11111111-1111-4111-8111-111111111111".to_string(), auth_manager, models_manager, + model_provider_runtime, environment_manager: Arc::new(EnvironmentManager::default_for_tests()), skills_service, plugins_manager, diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index bafb368b32..97fc1e8a1f 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -148,7 +148,7 @@ pub(crate) async fn run_turn( cancellation_token: CancellationToken, ) -> CodexResult> { let mut client_session = - prewarmed_client_session.unwrap_or_else(|| sess.services.model_client.new_session()); + prewarmed_client_session.unwrap_or_else(|| turn_context.model_client.new_session()); // TODO(ccunningham): Pre-turn compaction runs before context updates and the // new user message are recorded. Estimate pending incoming items (context // diffs/full reinjection + user input) and trigger compaction preemptively @@ -871,7 +871,7 @@ async fn maybe_run_previous_model_inline_compact( let previous_model = previous_turn_settings.model; let previous_model_turn_context = Arc::new( turn_context - .with_model(previous_model.clone(), &sess.services.models_manager) + .with_model(previous_model.clone(), &turn_context.models_manager) .await, ); @@ -2271,7 +2271,7 @@ async fn try_run_sampling_request( } ResponseEvent::ModelsEtag(etag) => { // Update internal state with latest models etag - sess.services + turn_context .models_manager .refresh_if_new_etag(etag, turn_context.config.http_client_factory()) .await; diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index e200413d3d..9c3e36256e 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -1,10 +1,10 @@ use super::*; use crate::environment_selection::TurnEnvironmentSnapshot; +use crate::model_provider_runtime::ModelProviderRuntimeSnapshot; use crate::shell_snapshot::ShellSnapshotFile; use codex_core_skills::HostSkillsSnapshot; use codex_file_system::FileSystemSandboxContext; use codex_model_provider::SharedModelProvider; -use codex_model_provider::create_model_provider; use codex_protocol::SessionId; use codex_protocol::ThreadId; use codex_protocol::models::AdditionalPermissionProfile; @@ -109,6 +109,9 @@ pub struct TurnContext { pub(crate) model_info: ModelInfo, pub(crate) session_telemetry: SessionTelemetry, pub(crate) provider: SharedModelProvider, + pub(crate) models_manager: SharedModelsManager, + pub(crate) model_client: ModelClient, + pub(crate) model_provider_generation: u64, pub(crate) reasoning_effort: Option, pub(crate) reasoning_summary: ReasoningSummaryConfig, pub(crate) session_source: SessionSource, @@ -269,6 +272,9 @@ impl TurnContext { .clone() .with_model(model.as_str(), model_info.slug.as_str()), provider: self.provider.clone(), + models_manager: Arc::clone(&self.models_manager), + model_client: self.model_client.clone(), + model_provider_generation: self.model_provider_generation, reasoning_effort, reasoning_summary: self.reasoning_summary, session_source: self.session_source.clone(), @@ -421,7 +427,121 @@ fn local_time_context() -> (String, String) { } } +struct PreparedModelProviderAdoption { + current: Arc, + latest: Arc, + model: String, + model_info: ModelInfo, +} + impl Session { + async fn prepare_latest_default_model_provider_adoption( + &self, + ) -> Option { + let current = self.services.model_runtime.load_full(); + let ModelProviderRuntimeSource::RuntimeDefault(runtime) = ¤t.source else { + return None; + }; + let latest = runtime.snapshot(); + if latest.generation() == current.snapshot.generation() { + return None; + } + + let (refresh_strategy, mut config) = { + let state = self.state.lock().await; + let refresh_strategy = if state + .session_configuration + .session_source + .is_non_root_agent() + { + RefreshStrategy::Offline + } else { + RefreshStrategy::OnlineIfUncached + }; + let config = (*state.session_configuration.original_config_do_not_use).clone(); + (refresh_strategy, config) + }; + let models_manager = latest.models_manager(); + let _ = models_manager + .list_models(refresh_strategy, config.http_client_factory()) + .await; + let model = models_manager + .get_default_model( + /*model*/ &None, + /*allow_provider_model_fallback*/ true, + refresh_strategy, + config.http_client_factory(), + ) + .await; + + latest.apply_to_config(&mut config); + config.model = Some(model.clone()); + let model_info = models_manager + .get_model_info(model.as_str(), &config.to_models_manager_config()) + .await; + Some(PreparedModelProviderAdoption { + current, + latest, + model, + model_info, + }) + } + + fn apply_default_model_provider_adoption( + &self, + state: &mut SessionState, + prepared: Option, + ) -> Arc { + let installed = self.services.model_runtime.load_full(); + let Some(prepared) = prepared else { + return installed; + }; + if !Arc::ptr_eq(&installed.snapshot, &prepared.current.snapshot) { + return installed; + } + + let PreparedModelProviderAdoption { + current, + latest, + model, + model_info, + } = prepared; + let supported_reasoning_levels = model_info + .supported_reasoning_levels + .iter() + .map(|preset| preset.effort.clone()) + .collect::>(); + let current_reasoning_effort = state + .session_configuration + .collaboration_mode + .reasoning_effort(); + let reasoning_effort = current_reasoning_effort + .filter(|effort| supported_reasoning_levels.contains(effort)) + .or_else(|| model_info.default_reasoning_level.clone()); + let mut config = (*state.session_configuration.original_config_do_not_use).clone(); + latest.apply_to_config(&mut config); + config.model = Some(model.clone()); + config.model_reasoning_effort = reasoning_effort.clone(); + + state.session_configuration.provider = latest.provider_info().clone(); + state.session_configuration.collaboration_mode = + state.session_configuration.collaboration_mode.with_updates( + Some(model), + Some(reasoning_effort), + /*developer_instructions*/ None, + ); + state.session_configuration.original_config_do_not_use = Arc::new(config); + let next = Arc::new(SessionModelRuntime { + source: current.source.clone(), + snapshot: latest.clone(), + model_client: current.model_client.rebind_provider(latest.provider()), + }); + self.services.model_runtime.store(Arc::clone(&next)); + // A prewarmed websocket is bound to the provider runtime that created it. + let _ = state.take_session_startup_prewarm(); + next + } + /// Don't expand the number of mutated arguments on config. We are in the process of getting rid of it. pub(crate) fn build_per_turn_config( session_configuration: &SessionConfiguration, @@ -482,7 +602,7 @@ impl Session { session_id: SessionId, auth_manager: Option>, session_telemetry: &SessionTelemetry, - provider: ModelProviderInfo, + model_runtime: &SessionModelRuntime, session_configuration: &SessionConfiguration, multi_agent_version: MultiAgentVersion, user_shell: &shell::Shell, @@ -490,7 +610,6 @@ impl Session { main_execve_wrapper_exe: Option<&PathBuf>, per_turn_config: Config, model_info: ModelInfo, - models_manager: &SharedModelsManager, network: Option, environments: TurnEnvironmentSnapshot, cwd: AbsolutePathBuf, @@ -506,8 +625,9 @@ impl Session { model_info.slug.as_str(), ); let session_source = session_configuration.session_source.clone(); - let auth_manager_for_context = auth_manager.clone(); - let provider_for_context = create_model_provider(provider, auth_manager); + let auth_manager_for_context = auth_manager; + let provider_for_context = model_runtime.snapshot.provider(); + let models_manager = model_runtime.models_manager(); let session_telemetry_for_context = session_telemetry; let available_models = models_manager.try_list_models().unwrap_or_default(); let unified_exec_shell_mode = UnifiedExecShellMode::for_session( @@ -549,6 +669,9 @@ impl Session { model_info, session_telemetry: session_telemetry_for_context, provider: provider_for_context, + models_manager, + model_client: model_runtime.model_client.clone(), + model_provider_generation: model_runtime.snapshot.generation(), reasoning_effort, reasoning_summary, session_source, @@ -589,8 +712,11 @@ impl Session { updates: SessionSettingsUpdate, ) -> CodexResult> { let notify_config_contributors = !self.services.extensions.config_contributors().is_empty(); + let prepared_model_provider = self.prepare_latest_default_model_provider_adoption().await; let update_result: CodexResult<_> = { let mut state = self.state.lock().await; + let model_runtime = + self.apply_default_model_provider_adoption(&mut state, prepared_model_provider); match state.session_configuration.clone().apply(&updates) { Ok(next) => { let previous_permission_profile = @@ -611,6 +737,7 @@ impl Session { state.session_configuration = next.clone(); Ok(( next, + model_runtime, permission_profile_changed, previous_config, new_config, @@ -620,22 +747,27 @@ impl Session { } }; - let (session_configuration, permission_profile_changed, previous_config, new_config) = - match update_result { - Ok(update) => update, - Err(err) => { - let message = err.to_string(); - self.send_event_raw(Event { - id: sub_id.clone(), - msg: EventMsg::Error(ErrorEvent { - message: message.clone(), - codex_error_info: Some(CodexErrorInfo::BadRequest), - }), - }) - .await; - return Err(CodexErr::InvalidRequest(message)); - } - }; + let ( + session_configuration, + model_runtime, + permission_profile_changed, + previous_config, + new_config, + ) = match update_result { + Ok(update) => update, + Err(err) => { + let message = err.to_string(); + self.send_event_raw(Event { + id: sub_id.clone(), + msg: EventMsg::Error(ErrorEvent { + message: message.clone(), + codex_error_info: Some(CodexErrorInfo::BadRequest), + }), + }) + .await; + return Err(CodexErr::InvalidRequest(message)); + } + }; self.emit_config_changed_contributors(previous_config.as_ref(), new_config.as_ref()); if permission_profile_changed { @@ -647,6 +779,7 @@ impl Session { .new_turn_from_configuration( sub_id, session_configuration, + model_runtime, updates.final_output_json_schema, ) .await) @@ -656,11 +789,13 @@ impl Session { &self, sub_id: String, session_configuration: SessionConfiguration, + model_runtime: Arc, final_output_json_schema: Option>, ) -> Arc { self.new_turn_context_from_configuration( sub_id, session_configuration, + model_runtime, final_output_json_schema, TurnMultiAgentRuntime::ResolveAndStore, ) @@ -671,10 +806,12 @@ impl Session { &self, sub_id: String, session_configuration: SessionConfiguration, + model_runtime: Arc, ) -> Arc { self.new_turn_context_from_configuration( sub_id, session_configuration, + model_runtime, /*final_output_json_schema*/ None, TurnMultiAgentRuntime::Preview, ) @@ -686,6 +823,7 @@ impl Session { &self, sub_id: String, session_configuration: SessionConfiguration, + model_runtime: Arc, final_output_json_schema: Option>, multi_agent_runtime: TurnMultiAgentRuntime, ) -> Arc { @@ -706,9 +844,8 @@ impl Session { .set_permission_profile(session_configuration.permission_profile()); } - let model_info = self - .services - .models_manager + let model_info = model_runtime + .models_manager() .get_model_info( session_configuration.collaboration_mode.model(), &per_turn_config.to_models_manager_config(), @@ -752,7 +889,7 @@ impl Session { self.session_id(), Some(Arc::clone(&self.services.auth_manager)), &self.services.session_telemetry, - session_configuration.provider.clone(), + model_runtime.as_ref(), &session_configuration, multi_agent_version, self.services.user_shell.as_ref(), @@ -760,7 +897,6 @@ impl Session { self.services.main_execve_wrapper_exe.as_ref(), per_turn_config, model_info, - &self.services.models_manager, self.services .network_proxy .load_full() @@ -820,10 +956,11 @@ impl Session { } pub(crate) async fn new_default_turn_with_sub_id(&self, sub_id: String) -> Arc { - let session_configuration = self.default_turn_configuration().await; + let (session_configuration, model_runtime) = self.default_turn_configuration().await; self.new_turn_from_configuration( sub_id, session_configuration, + model_runtime, /*final_output_json_schema*/ None, ) .await @@ -833,13 +970,20 @@ impl Session { &self, sub_id: String, ) -> Arc { - let session_configuration = self.default_turn_configuration().await; - self.new_startup_prewarm_turn_from_configuration(sub_id, session_configuration) - .await + let (session_configuration, model_runtime) = self.default_turn_configuration().await; + self.new_startup_prewarm_turn_from_configuration( + sub_id, + session_configuration, + model_runtime, + ) + .await } - async fn default_turn_configuration(&self) -> SessionConfiguration { - let state = self.state.lock().await; - state.session_configuration.clone() + async fn default_turn_configuration(&self) -> (SessionConfiguration, Arc) { + let prepared_model_provider = self.prepare_latest_default_model_provider_adoption().await; + let mut state = self.state.lock().await; + let model_runtime = + self.apply_default_model_provider_adoption(&mut state, prepared_model_provider); + (state.session_configuration.clone(), model_runtime) } } diff --git a/codex-rs/core/src/session_startup_prewarm.rs b/codex-rs/core/src/session_startup_prewarm.rs index a5a4c5e437..7641294501 100644 --- a/codex-rs/core/src/session_startup_prewarm.rs +++ b/codex-rs/core/src/session_startup_prewarm.rs @@ -182,10 +182,10 @@ impl SessionStartupPrewarmHandle { impl Session { pub(crate) async fn schedule_startup_prewarm(self: &Arc, base_instructions: String) { - if !self.services.model_client.responses_websocket_enabled() { + let model_client = self.services.model_runtime.load().model_client.clone(); + if !model_client.responses_websocket_enabled() { // Without websocket prewarm, resolve auth once so Agent Identity bootstrap can // register or engage this session's bearer fallback before the first user request. - let model_client = self.services.model_client.clone(); tokio::spawn(async move { if let Err(err) = model_client.prewarm_auth().await { warn!("startup auth prewarm failed: {err:#}"); @@ -303,7 +303,7 @@ async fn schedule_startup_prewarm_inner( window_id, CodexResponsesRequestKind::Prewarm, ); - let mut client_session = session.services.model_client.new_session(); + let mut client_session = startup_turn_context.model_client.new_session(); let websocket_warmup_started_at = Instant::now(); client_session .prewarm_websocket( diff --git a/codex-rs/core/src/state/mod.rs b/codex-rs/core/src/state/mod.rs index 20ff78b54f..4142465247 100644 --- a/codex-rs/core/src/state/mod.rs +++ b/codex-rs/core/src/state/mod.rs @@ -7,6 +7,7 @@ mod turn; pub(crate) use additional_context::AdditionalContextStore; pub(crate) use auto_compact_window::AutoCompactWindowIds; pub(crate) use auto_compact_window::AutoCompactWindowSnapshot; +pub(crate) use service::SessionModelRuntime; pub(crate) use service::SessionServices; pub(crate) use session::SessionState; pub(crate) use turn::ActiveTurn; diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index 43f411e31b..5c3cabed4b 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -16,6 +16,8 @@ use crate::exec_policy::ExecPolicyManager; use crate::guardian::GuardianRejection; use crate::guardian::GuardianRejectionCircuitBreaker; use crate::mcp::McpManager; +use crate::model_provider_runtime::ModelProviderRuntimeSnapshot; +use crate::model_provider_runtime::ModelProviderRuntimeSource; use crate::session::McpRuntimeSnapshot; use crate::tools::code_mode::CodeModeService; use crate::tools::handlers::ToolSearchHandlerCache; @@ -68,6 +70,9 @@ pub(crate) struct SessionServices { pub(crate) show_raw_agent_reasoning: bool, pub(crate) exec_policy: Arc, pub(crate) auth_manager: Arc, + /// Atomically replaceable provider runtime used to build future turns. Each turn clones one + /// immutable runtime so an in-flight turn cannot observe a provider refresh. + pub(crate) model_runtime: ArcSwap, pub(crate) models_manager: SharedModelsManager, pub(crate) session_telemetry: SessionTelemetry, pub(crate) tool_approvals: Mutex, @@ -103,6 +108,19 @@ pub(crate) struct SessionServices { pub(crate) turn_environments: Arc, } +#[derive(Clone)] +pub(crate) struct SessionModelRuntime { + pub(crate) source: ModelProviderRuntimeSource, + pub(crate) snapshot: Arc, + pub(crate) model_client: ModelClient, +} + +impl SessionModelRuntime { + pub(crate) fn models_manager(&self) -> SharedModelsManager { + self.snapshot.models_manager() + } +} + impl SessionServices { /// Installs the manager before validating required servers so startup-time elicitation can /// resolve through the session's manager while validation waits. diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 7a930bb578..a3afbf5e52 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -37,7 +37,6 @@ use crate::state::TaskKind; use codex_analytics::TurnProfileFact; use codex_analytics::TurnTokenUsageFact; use codex_login::AuthManager; -use codex_models_manager::manager::SharedModelsManager; use codex_otel::SessionTelemetry; use codex_otel::TURN_E2E_DURATION_METRIC; use codex_otel::TURN_MEMORY_METRIC; @@ -197,10 +196,6 @@ impl SessionTaskContext { pub(crate) fn auth_manager(&self) -> Arc { Arc::clone(&self.session.services.auth_manager) } - - pub(crate) fn models_manager(&self) -> SharedModelsManager { - Arc::clone(&self.session.services.models_manager) - } } /// Async task that drives a [`Session`] turn. diff --git a/codex-rs/core/src/tasks/review.rs b/codex-rs/core/src/tasks/review.rs index 30f8d47d81..8db29c8022 100644 --- a/codex-rs/core/src/tasks/review.rs +++ b/codex-rs/core/src/tasks/review.rs @@ -125,7 +125,7 @@ async fn start_review_conversation( (run_codex_thread_one_shot( sub_agent_config, session.auth_manager(), - session.models_manager(), + Arc::clone(&ctx.models_manager), input, session.clone_session(), ctx.clone(), diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index d98e8bbcc8..6843e9e798 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -9,6 +9,9 @@ use crate::environment_selection::TurnEnvironmentSnapshot; use crate::environment_selection::default_thread_environment_selections; use crate::mcp::McpManager; use crate::model_provider_runtime::DefaultModelProviderRuntime; +use crate::model_provider_runtime::InitialModelProviderRuntime; +use crate::model_provider_runtime::ModelProviderRuntimeSource; +use crate::model_provider_runtime::build_explicit_model_provider_runtime; use crate::rollout::truncation; use crate::session::Codex; use crate::session::CodexSpawnArgs; @@ -200,6 +203,14 @@ pub struct StartThreadOptions { pub supports_openai_form_elicitation: bool, } +/// Determines whether a thread follows the process-default model provider or keeps its initial +/// provider pinned for the lifetime of the loaded session. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ModelProviderSelection { + RuntimeDefault, + Explicit, +} + fn originator_from_service_name(service_name: Option<&str>) -> Option { let service_name = service_name?.trim(); for originator in ["codex_work_desktop", "codex_work_web", "codex_work_mobile"] { @@ -241,7 +252,7 @@ pub(crate) struct ThreadManagerState { threads: Arc>>>, thread_created_tx: broadcast::Sender, auth_manager: Arc, - default_model_provider_runtime: DefaultModelProviderRuntime, + default_model_provider_runtime: Arc, environment_manager: Arc, skills_service: Arc, plugins_manager: Arc, @@ -334,10 +345,10 @@ impl ThreadManager { state: Arc::new(ThreadManagerState { threads: Arc::new(RwLock::new(HashMap::new())), thread_created_tx, - default_model_provider_runtime: DefaultModelProviderRuntime::new( + default_model_provider_runtime: Arc::new(DefaultModelProviderRuntime::new( config, auth_manager.clone(), - ), + )), environment_manager, skills_service, plugins_manager, @@ -456,13 +467,13 @@ impl ThreadManager { state: Arc::new(ThreadManagerState { threads: Arc::new(RwLock::new(HashMap::new())), thread_created_tx, - default_model_provider_runtime: DefaultModelProviderRuntime::from_parts( + default_model_provider_runtime: Arc::new(DefaultModelProviderRuntime::from_parts( OPENAI_PROVIDER_ID.to_string(), provider, codex_home, /*config_model_catalog*/ None, auth_manager.clone(), - ), + )), environment_manager, skills_service, plugins_manager, @@ -706,14 +717,31 @@ impl ThreadManager { &self, options: StartThreadOptions, ) -> CodexResult { - self.start_thread_with_options_and_fork_source(options, /*forked_from_thread_id*/ None) - .await + self.start_thread_with_options_and_provider_selection( + options, + ModelProviderSelection::Explicit, + ) + .await + } + + pub async fn start_thread_with_options_and_provider_selection( + &self, + options: StartThreadOptions, + model_provider_selection: ModelProviderSelection, + ) -> CodexResult { + self.start_thread_with_options_and_fork_source( + options, + /*forked_from_thread_id*/ None, + model_provider_selection, + ) + .await } async fn start_thread_with_options_and_fork_source( &self, options: StartThreadOptions, forked_from_thread_id: Option, + model_provider_selection: ModelProviderSelection, ) -> CodexResult { let agent_control = self.agent_control_for_config(&options.config); let (resumed_session_source, resumed_thread_source) = options @@ -742,6 +770,7 @@ impl ThreadManager { options.thread_extension_init, options.supports_openai_form_elicitation, /*user_shell_override*/ None, + model_provider_selection, )) .await } @@ -779,8 +808,12 @@ impl ThreadManager { inherited_multi_agent_version, ), ); - self.start_thread_with_options_and_fork_source(options, Some(forked_from_thread_id)) - .await + self.start_thread_with_options_and_fork_source( + options, + Some(forked_from_thread_id), + ModelProviderSelection::Explicit, + ) + .await } pub async fn resume_thread_from_rollout( @@ -839,6 +872,7 @@ impl ThreadManager { /*thread_extension_init*/ ExtensionDataInit::default(), supports_openai_form_elicitation, /*user_shell_override*/ None, + ModelProviderSelection::Explicit, )) .await } @@ -910,6 +944,7 @@ impl ThreadManager { /*thread_extension_init*/ ExtensionDataInit::default(), supports_openai_form_elicitation, /*user_shell_override*/ Some(user_shell_override), + ModelProviderSelection::Explicit, )) .await } @@ -1412,6 +1447,7 @@ impl ThreadManagerState { /*thread_extension_init*/ ExtensionDataInit::default(), /*supports_openai_form_elicitation*/ false, /*user_shell_override*/ None, + ModelProviderSelection::Explicit, )) .await } @@ -1452,6 +1488,7 @@ impl ThreadManagerState { /*thread_extension_init*/ ExtensionDataInit::default(), /*supports_openai_form_elicitation*/ false, /*user_shell_override*/ None, + ModelProviderSelection::Explicit, )) .await } @@ -1494,6 +1531,7 @@ impl ThreadManagerState { thread_extension_init, /*supports_openai_form_elicitation*/ false, /*user_shell_override*/ None, + ModelProviderSelection::Explicit, )) .await } @@ -1537,6 +1575,7 @@ impl ThreadManagerState { thread_extension_init, supports_openai_form_elicitation, user_shell_override, + ModelProviderSelection::Explicit, )) .await } @@ -1544,7 +1583,7 @@ impl ThreadManagerState { #[allow(clippy::too_many_arguments)] pub(crate) async fn spawn_thread_with_source( &self, - config: Config, + mut config: Config, initial_history: InitialHistory, history_mode: Option, allow_provider_model_fallback: bool, @@ -1563,6 +1602,7 @@ impl ThreadManagerState { thread_extension_init: ExtensionDataInit, supports_openai_form_elicitation: bool, user_shell_override: Option, + model_provider_selection: ModelProviderSelection, ) -> CodexResult { let is_resumed_thread = matches!(&initial_history, InitialHistory::Resumed(_)); if let InitialHistory::Resumed(resumed) = &initial_history { @@ -1610,9 +1650,22 @@ impl ThreadManagerState { forked_from_thread_id, ) .await; - let models_manager = self - .default_model_provider_runtime - .models_manager_for_config(&config, Arc::clone(&auth_manager)); + let model_provider_runtime = match model_provider_selection { + ModelProviderSelection::RuntimeDefault => { + let snapshot = self.default_model_provider_runtime.snapshot(); + snapshot.apply_to_config(&mut config); + InitialModelProviderRuntime { + source: ModelProviderRuntimeSource::RuntimeDefault(Arc::clone( + &self.default_model_provider_runtime, + )), + snapshot, + } + } + ModelProviderSelection::Explicit => { + build_explicit_model_provider_runtime(&config, Arc::clone(&auth_manager)) + } + }; + let models_manager = model_provider_runtime.snapshot.models_manager(); let CodexSpawnOk { codex, thread_id, .. } = Box::pin(Codex::spawn(CodexSpawnArgs { @@ -1622,6 +1675,7 @@ impl ThreadManagerState { installation_id: self.installation_id.clone(), auth_manager, models_manager, + model_provider_runtime, environment_manager: Arc::clone(&self.environment_manager), skills_service: Arc::clone(&self.skills_service), plugins_manager: Arc::clone(&self.plugins_manager), diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index cbeebf34d7..2260d7fdaf 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -64,6 +64,34 @@ async fn listed_model_ids(models_manager: &SharedModelsManager) -> Vec { .collect() } +async fn start_thread_with_provider_selection( + manager: &ThreadManager, + config: Config, + model_provider_selection: ModelProviderSelection, +) -> NewThread { + let environments = manager.default_environment_selections(&config.cwd); + manager + .start_thread_with_options_and_provider_selection( + StartThreadOptions { + config, + allow_provider_model_fallback: false, + initial_history: InitialHistory::New, + history_mode: None, + session_source: None, + thread_source: None, + dynamic_tools: Vec::new(), + metrics_service_name: None, + parent_trace: None, + environments, + thread_extension_init: ExtensionDataInit::default(), + supports_openai_form_elicitation: false, + }, + model_provider_selection, + ) + .await + .expect("start thread") +} + struct FakeAgentGraphStore { root_thread_id: ThreadId, descendant_thread_ids: Vec, @@ -1492,10 +1520,12 @@ async fn new_threads_select_models_manager_for_their_effective_config() { vec!["model-b".to_string()] ); - let default_thread = manager - .start_thread(provider_b_config) - .await - .expect("start default-provider thread"); + let default_thread = start_thread_with_provider_selection( + &manager, + provider_b_config, + ModelProviderSelection::RuntimeDefault, + ) + .await; assert!(Arc::ptr_eq( &default_thread.thread.codex.session.services.models_manager, &provider_b_models_manager, @@ -1529,6 +1559,106 @@ async fn new_threads_select_models_manager_for_their_effective_config() { ); } +#[tokio::test] +async fn loaded_default_provider_session_adopts_refresh_between_turns() { + let temp_dir = tempdir().expect("tempdir"); + let mut provider_a_config = test_config().await; + provider_a_config.codex_home = temp_dir.path().join("codex-home").abs(); + provider_a_config.cwd = provider_a_config.codex_home.abs(); + std::fs::create_dir_all(&provider_a_config.codex_home).expect("create codex home"); + provider_a_config.model = None; + provider_a_config.model_provider_id = "provider-a".to_string(); + provider_a_config.model_provider.name = "Provider A".to_string(); + provider_a_config.model_catalog = Some(static_model_catalog("model-a")); + + let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("dummy")); + let manager = ThreadManager::new( + &provider_a_config, + auth_manager, + SessionSource::Exec, + Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + empty_extension_registry(), + Arc::new(crate::test_support::EmptyUserInstructionsProvider), + /*analytics_events_client*/ None, + thread_store_from_config(&provider_a_config, /*state_db*/ None), + /*agent_graph_store*/ None, + TEST_INSTALLATION_ID.to_string(), + /*attestation_provider*/ None, + /*external_time_provider*/ None, + ); + let default_thread = start_thread_with_provider_selection( + &manager, + provider_a_config.clone(), + ModelProviderSelection::RuntimeDefault, + ) + .await; + let explicit_thread = start_thread_with_provider_selection( + &manager, + provider_a_config.clone(), + ModelProviderSelection::Explicit, + ) + .await; + + let captured_a_turn = default_thread.thread.codex.session.new_default_turn().await; + assert_eq!(captured_a_turn.config.model_provider_id, "provider-a"); + assert_eq!(captured_a_turn.model_info.slug, "model-a"); + assert_eq!(captured_a_turn.model_provider_generation, 0); + + let mut provider_b_config = provider_a_config; + provider_b_config.model_provider_id = "provider-b".to_string(); + provider_b_config.model_provider.name = "Provider B".to_string(); + provider_b_config.model_catalog = Some(static_model_catalog("model-b")); + assert!(manager.refresh_default_model_provider(&provider_b_config)); + + // An already-created turn retains the exact provider runtime it captured. + assert_eq!(captured_a_turn.config.model_provider_id, "provider-a"); + assert_eq!(captured_a_turn.model_info.slug, "model-a"); + assert_eq!( + listed_model_ids(&captured_a_turn.models_manager).await, + vec!["model-a".to_string()] + ); + + let adopted_turn = default_thread.thread.codex.session.new_default_turn().await; + assert_eq!(adopted_turn.config.model_provider_id, "provider-b"); + assert_eq!(adopted_turn.model_info.slug, "model-b"); + assert_eq!(adopted_turn.model_provider_generation, 1); + assert_eq!( + listed_model_ids(&adopted_turn.models_manager).await, + vec!["model-b".to_string()] + ); + assert_eq!( + default_thread + .thread + .config_snapshot() + .await + .model_provider_id, + "provider-b" + ); + + let pinned_turn = explicit_thread + .thread + .codex + .session + .new_default_turn() + .await; + assert_eq!(pinned_turn.config.model_provider_id, "provider-a"); + assert_eq!(pinned_turn.model_info.slug, "model-a"); + assert_eq!(pinned_turn.model_provider_generation, 0); + + let mut expected_completed = vec![default_thread.thread_id, explicit_thread.thread_id]; + expected_completed.sort_by_key(std::string::ToString::to_string); + assert_eq!( + manager + .shutdown_all_threads_bounded(Duration::from_secs(10)) + .await, + ThreadShutdownReport { + completed: expected_completed, + submit_failed: Vec::new(), + timed_out: Vec::new(), + } + ); +} + #[test] fn interrupted_fork_snapshot_appends_interrupt_boundary() { let committed_history = diff --git a/codex-rs/login/src/auth/bedrock_api_key.rs b/codex-rs/login/src/auth/bedrock_api_key.rs index 3445878365..45b11e7e03 100644 --- a/codex-rs/login/src/auth/bedrock_api_key.rs +++ b/codex-rs/login/src/auth/bedrock_api_key.rs @@ -1,21 +1,34 @@ +use std::fmt; use std::path::Path; use codex_config::types::AuthCredentialsStoreMode; use serde::Deserialize; use serde::Serialize; +use super::manager::AuthManager; +use super::manager::load_auth_dot_json; +use super::manager::logout; use super::manager::save_auth; use super::storage::AuthDotJson; use super::storage::AuthKeyringBackendKind; use codex_protocol::auth::AuthMode; /// Managed Amazon Bedrock API key persisted in `auth.json`. -#[derive(Deserialize, Serialize, Clone, Debug, PartialEq, Eq)] +#[derive(Deserialize, Serialize, Clone, PartialEq, Eq)] pub struct BedrockApiKeyAuth { pub api_key: String, pub region: String, } +impl fmt::Debug for BedrockApiKeyAuth { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("BedrockApiKeyAuth") + .field("api_key", &"") + .field("region", &self.region) + .finish() + } +} + /// Writes an `auth.json` that contains only the Amazon Bedrock API key auth. pub fn login_with_bedrock_api_key( codex_home: &Path, @@ -44,6 +57,83 @@ pub fn login_with_bedrock_api_key( ) } +pub(super) fn load_stored_bedrock_api_key( + codex_home: &Path, + auth_credentials_store_mode: AuthCredentialsStoreMode, + keyring_backend_kind: AuthKeyringBackendKind, +) -> std::io::Result> { + load_stored_bedrock_api_key_auth( + codex_home, + auth_credentials_store_mode, + keyring_backend_kind, + ) + .map(|auth| auth.and_then(|auth| auth.bedrock_api_key)) +} + +fn load_stored_bedrock_api_key_auth( + codex_home: &Path, + auth_credentials_store_mode: AuthCredentialsStoreMode, + keyring_backend_kind: AuthKeyringBackendKind, +) -> std::io::Result> { + load_auth_dot_json( + codex_home, + auth_credentials_store_mode, + keyring_backend_kind, + ) + .map(|auth| auth.filter(AuthDotJson::is_bedrock_api_key)) +} + +impl AuthManager { + pub fn has_stored_bedrock_api_key_auth(&self) -> std::io::Result { + load_stored_bedrock_api_key_auth( + &self.codex_home, + self.auth_credentials_store_mode, + self.keyring_backend_kind, + ) + .map(|auth| auth.is_some()) + } + + /// Persists managed Bedrock auth without changing the in-memory auth snapshot. + /// + /// The account coordinator publishes the provider-scoped cache update only + /// after the associated config mutation succeeds. + pub fn persist_bedrock_api_key_auth(&self, api_key: &str, region: &str) -> std::io::Result<()> { + login_with_bedrock_api_key( + &self.codex_home, + api_key, + region, + self.auth_credentials_store_mode, + self.keyring_backend_kind, + ) + } + + /// Reloads only managed Bedrock auth, preserving cached general Codex auth. + pub fn reload_bedrock_api_key_auth(&self) -> std::io::Result { + let bedrock_api_key = load_stored_bedrock_api_key( + &self.codex_home, + self.auth_credentials_store_mode, + self.keyring_backend_kind, + )?; + Ok(self.set_cached_bedrock_api_key_auth(bedrock_api_key)) + } + + /// Removes only managed Bedrock auth, preserving cached general Codex auth. + pub fn logout_bedrock_api_key_auth(&self) -> std::io::Result { + let stored_bedrock_auth = self.has_stored_bedrock_api_key_auth()?; + let removed = if stored_bedrock_auth { + logout( + &self.codex_home, + self.auth_credentials_store_mode, + self.keyring_backend_kind, + )? + } else { + false + }; + self.set_cached_bedrock_api_key_auth(/*new_bedrock_api_key*/ None); + Ok(removed) + } +} + #[cfg(test)] #[path = "bedrock_api_key_tests.rs"] mod tests; diff --git a/codex-rs/login/src/auth/bedrock_api_key_tests.rs b/codex-rs/login/src/auth/bedrock_api_key_tests.rs index 182834af66..990e4093c4 100644 --- a/codex-rs/login/src/auth/bedrock_api_key_tests.rs +++ b/codex-rs/login/src/auth/bedrock_api_key_tests.rs @@ -1,13 +1,15 @@ use codex_config::types::AuthCredentialsStoreMode; use codex_protocol::auth::AuthMode; +use codex_protocol::config_types::ForcedLoginMethod; use pretty_assertions::assert_eq; use serial_test::serial; use tempfile::tempdir; use super::*; +use crate::auth::AuthConfig; use crate::auth::AuthKeyringBackendKind; use crate::auth::AuthManager; -use crate::auth::CodexAuth; +use crate::auth::enforce_login_restrictions; use crate::auth::storage::AuthStorageBackend; use crate::auth::storage::FileAuthStorage; @@ -42,20 +44,24 @@ fn bedrock_auth() -> BedrockApiKeyAuth { } } +fn auth_config(codex_home: &std::path::Path, forced_login_method: ForcedLoginMethod) -> AuthConfig { + AuthConfig { + codex_home: codex_home.to_path_buf(), + auth_credentials_store_mode: AuthCredentialsStoreMode::File, + keyring_backend_kind: AuthKeyringBackendKind::default(), + forced_login_method: Some(forced_login_method), + chatgpt_base_url: None, + forced_chatgpt_workspace_id: None, + auth_route_config: None, + } +} + #[tokio::test] #[serial(codex_auth_env)] -async fn login_with_bedrock_api_key_replaces_openai_auth() -> anyhow::Result<()> { +async fn managed_bedrock_persistence_preserves_cached_openai_auth() -> anyhow::Result<()> { let codex_home = tempdir()?; let storage = FileAuthStorage::new(codex_home.path().to_path_buf()); storage.save(&api_key_auth())?; - login_with_bedrock_api_key( - codex_home.path(), - "bedrock-api-key-test", - "us-east-1", - AuthCredentialsStoreMode::File, - AuthKeyringBackendKind::default(), - )?; - let auth_manager = AuthManager::new( codex_home.path().to_path_buf(), /*enable_codex_api_key_env*/ false, @@ -67,6 +73,8 @@ async fn login_with_bedrock_api_key_replaces_openai_auth() -> anyhow::Result<()> ) .await; + auth_manager.persist_bedrock_api_key_auth("bedrock-api-key-test", "us-east-1")?; + let loaded = storage.load()?.expect("auth should be stored"); let expected = AuthDotJson { auth_mode: Some(AuthMode::BedrockApiKey), @@ -78,25 +86,33 @@ async fn login_with_bedrock_api_key_replaces_openai_auth() -> anyhow::Result<()> bedrock_api_key: Some(bedrock_auth()), }; assert_eq!(loaded, expected); - assert_eq!(auth_manager.auth_mode(), Some(AuthMode::BedrockApiKey)); + assert_eq!(auth_manager.auth_mode(), Some(AuthMode::ApiKey)); + assert_eq!(auth_manager.bedrock_api_key_auth_cached(), None); + assert!(auth_manager.reload_bedrock_api_key_auth()?); assert_eq!( - auth_manager.auth_cached().and_then(|auth| match auth { - CodexAuth::BedrockApiKey(auth) => Some(auth), - CodexAuth::ApiKey(_) - | CodexAuth::Chatgpt(_) - | CodexAuth::ChatgptAuthTokens(_) - | CodexAuth::Headers(_) - | CodexAuth::AgentIdentity(_) - | CodexAuth::PersonalAccessToken(_) => None, - }), + auth_manager + .auth_cached() + .and_then(|auth| auth.api_key().map(str::to_string)), + Some("sk-test-key".to_string()) + ); + assert_eq!( + auth_manager.bedrock_api_key_auth_cached(), Some(bedrock_auth()) ); Ok(()) } +#[test] +fn bedrock_auth_debug_redacts_api_key() { + assert_eq!( + format!("{:?}", bedrock_auth()), + "BedrockApiKeyAuth { api_key: \"\", region: \"us-east-1\" }" + ); +} + #[tokio::test] #[serial(codex_auth_env)] -async fn logout_removes_bedrock_auth() -> anyhow::Result<()> { +async fn auth_manager_logout_removes_bedrock_api_key_auth() -> anyhow::Result<()> { let codex_home = tempdir()?; let storage = FileAuthStorage::new(codex_home.path().to_path_buf()); login_with_bedrock_api_key( @@ -121,12 +137,13 @@ async fn logout_removes_bedrock_auth() -> anyhow::Result<()> { assert_eq!(storage.load()?, None); assert_eq!(auth_manager.auth_cached(), None); + assert_eq!(auth_manager.bedrock_api_key_auth_cached(), None); Ok(()) } #[tokio::test] #[serial(codex_auth_env)] -async fn bedrock_only_auth_storage_creates_primary_auth() -> anyhow::Result<()> { +async fn bedrock_only_auth_storage_does_not_create_codex_auth() -> anyhow::Result<()> { let codex_home = tempdir()?; let storage = FileAuthStorage::new(codex_home.path().to_path_buf()); storage.save(&bedrock_only_auth())?; @@ -142,19 +159,44 @@ async fn bedrock_only_auth_storage_creates_primary_auth() -> anyhow::Result<()> ) .await; - assert_eq!(auth_manager.auth_mode(), Some(AuthMode::BedrockApiKey)); + assert_eq!(auth_manager.auth_mode(), None); assert_eq!( - auth_manager.auth_cached().and_then(|auth| match auth { - CodexAuth::BedrockApiKey(auth) => Some(auth), - CodexAuth::ApiKey(_) - | CodexAuth::Chatgpt(_) - | CodexAuth::ChatgptAuthTokens(_) - | CodexAuth::Headers(_) - | CodexAuth::AgentIdentity(_) - | CodexAuth::PersonalAccessToken(_) => None, - }), + auth_manager.bedrock_api_key_auth_cached(), Some(bedrock_auth()) ); + assert_eq!(auth_manager.auth_cached(), None); + Ok(()) +} + +#[tokio::test] +#[serial(codex_auth_env)] +async fn bedrock_logout_preserves_cached_openai_auth() -> anyhow::Result<()> { + let codex_home = tempdir()?; + let storage = FileAuthStorage::new(codex_home.path().to_path_buf()); + storage.save(&api_key_auth())?; + let auth_manager = AuthManager::new( + codex_home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, + /*chatgpt_base_url*/ None, + AuthKeyringBackendKind::default(), + /*auth_route_config*/ None, + ) + .await; + auth_manager.persist_bedrock_api_key_auth("bedrock-api-key-test", "us-east-1")?; + auth_manager.reload_bedrock_api_key_auth()?; + + assert!(auth_manager.logout_bedrock_api_key_auth()?); + + assert_eq!(storage.load()?, None); + assert_eq!(auth_manager.bedrock_api_key_auth_cached(), None); + assert_eq!( + auth_manager + .auth_cached() + .and_then(|auth| auth.api_key().map(str::to_string)), + Some("sk-test-key".to_string()) + ); Ok(()) } @@ -180,3 +222,35 @@ async fn login_with_api_key_clears_bedrock_api_key() -> anyhow::Result<()> { assert_eq!(storage.load()?, Some(api_key_auth())); Ok(()) } + +#[tokio::test] +async fn forced_chatgpt_login_removes_bedrock_auth() -> anyhow::Result<()> { + let codex_home = tempdir()?; + let storage = FileAuthStorage::new(codex_home.path().to_path_buf()); + storage.save(&bedrock_only_auth())?; + + let error = + enforce_login_restrictions(&auth_config(codex_home.path(), ForcedLoginMethod::Chatgpt)) + .await + .expect_err("Bedrock auth should violate forced ChatGPT login"); + + assert_eq!( + error.to_string(), + "ChatGPT login is required, but an API key is currently being used. Logging out." + ); + assert_eq!(storage.load()?, None); + Ok(()) +} + +#[tokio::test] +async fn forced_api_login_allows_bedrock_auth() -> anyhow::Result<()> { + let codex_home = tempdir()?; + let storage = FileAuthStorage::new(codex_home.path().to_path_buf()); + let auth = bedrock_only_auth(); + storage.save(&auth)?; + + enforce_login_restrictions(&auth_config(codex_home.path(), ForcedLoginMethod::Api)).await?; + + assert_eq!(storage.load()?, Some(auth)); + Ok(()) +} diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 9509efc348..c056127773 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -36,6 +36,7 @@ use super::agent_identity::record_needs_task_registration; use super::agent_identity::register_managed_chatgpt_agent_identity; use super::agent_identity::require_agent_identity_authapi_base_url; use super::agent_identity::verified_record_from_jwt; +use super::bedrock_api_key::load_stored_bedrock_api_key; use super::external_bearer::BearerTokenRefresher; use super::revoke::revoke_auth_tokens; use crate::auth::AuthHeaders; @@ -1063,7 +1064,7 @@ async fn enforce_login_restrictions_with_agent_identity_authapi_base_url( config: &AuthConfig, agent_identity_authapi_base_url: Option<&str>, ) -> std::io::Result<()> { - let Some(auth) = load_auth( + let auth = load_auth( &config.codex_home, /*enable_codex_api_key_env*/ true, config.auth_credentials_store_mode, @@ -1073,13 +1074,23 @@ async fn enforce_login_restrictions_with_agent_identity_authapi_base_url( agent_identity_authapi_base_url, config.auth_route_config.as_ref(), ) - .await? - else { + .await?; + let stored_bedrock_api_key = load_stored_bedrock_api_key( + &config.codex_home, + config.auth_credentials_store_mode, + config.keyring_backend_kind, + )?; + let auth_mode = auth.as_ref().map(CodexAuth::auth_mode).or_else(|| { + stored_bedrock_api_key + .as_ref() + .map(|_| AuthMode::BedrockApiKey) + }); + let Some(auth_mode) = auth_mode else { return Ok(()); }; if let Some(required_method) = config.forced_login_method { - let method_violation = match (required_method, auth.auth_mode()) { + let method_violation = match (required_method, auth_mode) { (ForcedLoginMethod::Api, AuthMode::ApiKey) | (ForcedLoginMethod::Api, AuthMode::BedrockApiKey) => None, (ForcedLoginMethod::Chatgpt, AuthMode::Chatgpt) @@ -1113,7 +1124,10 @@ async fn enforce_login_restrictions_with_agent_identity_authapi_base_url( } if let Some(expected_account_ids) = config.forced_chatgpt_workspace_id.as_deref() { - let chatgpt_account_id = match &auth { + let Some(auth) = auth.as_ref() else { + return Ok(()); + }; + let chatgpt_account_id = match auth { CodexAuth::ApiKey(_) | CodexAuth::Headers(_) | CodexAuth::BedrockApiKey(_) => { return Ok(()); } @@ -1234,7 +1248,10 @@ async fn load_auth( AuthCredentialsStoreMode::Ephemeral, AuthKeyringBackendKind::default(), ); - if let Some(auth_dot_json) = ephemeral_storage.load()? { + if let Some(auth_dot_json) = ephemeral_storage + .load()? + .filter(|auth| !auth.is_bedrock_api_key()) + { let auth = CodexAuth::from_auth_dot_json( codex_home, auth_dot_json, @@ -1283,7 +1300,8 @@ async fn load_auth( keyring_backend_kind, ); let auth_dot_json = match storage.load()? { - Some(auth) => auth, + Some(auth) if !auth.is_bedrock_api_key() => auth, + Some(_) => return Ok(None), None => return Ok(None), }; @@ -1458,6 +1476,11 @@ fn refresh_token_endpoint() -> String { } impl AuthDotJson { + /// Returns whether this stored payload resolves to managed Amazon Bedrock API key auth. + pub fn is_bedrock_api_key(&self) -> bool { + self.resolved_mode() == AuthMode::BedrockApiKey + } + fn from_external_access_token( access_token: &str, chatgpt_account_id: &str, @@ -1520,6 +1543,7 @@ impl AuthDotJson { #[derive(Clone)] struct CachedAuth { auth: Option, + bedrock_api_key: Option, /// Permanent refresh failure cached for the current auth snapshot so /// later refresh attempts for the same credentials fail fast without network. permanent_refresh_failure: Option, @@ -1538,6 +1562,7 @@ impl Debug for CachedAuth { "auth_mode", &self.auth.as_ref().map(CodexAuth::api_auth_mode), ) + .field("has_bedrock_api_key", &self.bedrock_api_key.is_some()) .field( "permanent_refresh_failure", &self @@ -1765,12 +1790,12 @@ impl UnauthorizedRecovery { /// `reload()` is called explicitly. This matches the design goal of avoiding /// different parts of the program seeing inconsistent auth data mid‑run. pub struct AuthManager { - codex_home: PathBuf, + pub(super) codex_home: PathBuf, inner: RwLock, auth_change_tx: watch::Sender, enable_codex_api_key_env: bool, - auth_credentials_store_mode: AuthCredentialsStoreMode, - keyring_backend_kind: AuthKeyringBackendKind, + pub(super) auth_credentials_store_mode: AuthCredentialsStoreMode, + pub(super) keyring_backend_kind: AuthKeyringBackendKind, forced_chatgpt_workspace_id: RwLock>>, chatgpt_base_url: Option, agent_identity_authapi_base_url: Option, @@ -1849,6 +1874,13 @@ impl AuthManager { ) -> Self { let agent_identity_authapi_base_url = agent_identity_authapi_base_url(chatgpt_base_url.as_deref()).ok(); + let bedrock_api_key = load_stored_bedrock_api_key( + &codex_home, + auth_credentials_store_mode, + keyring_backend_kind, + ) + .ok() + .flatten(); let managed_auth = load_auth( &codex_home, enable_codex_api_key_env, @@ -1867,6 +1899,7 @@ impl AuthManager { codex_home, inner: RwLock::new(CachedAuth { auth: managed_auth, + bedrock_api_key, permanent_refresh_failure: None, }), auth_change_tx, @@ -1888,6 +1921,7 @@ impl AuthManager { pub fn from_auth_for_testing(auth: CodexAuth) -> Arc { let cached = CachedAuth { auth: Some(auth), + bedrock_api_key: None, permanent_refresh_failure: None, }; let (auth_change_tx, _auth_change_rx) = watch::channel(0); @@ -1914,6 +1948,7 @@ impl AuthManager { pub fn from_auth_for_testing_with_home(auth: CodexAuth, codex_home: PathBuf) -> Arc { let cached = CachedAuth { auth: Some(auth), + bedrock_api_key: None, permanent_refresh_failure: None, }; let (auth_change_tx, _auth_change_rx) = watch::channel(0); @@ -1943,6 +1978,7 @@ impl AuthManager { ) -> Arc { let cached = CachedAuth { auth: Some(auth), + bedrock_api_key: None, permanent_refresh_failure: None, }; let (auth_change_tx, _auth_change_rx) = watch::channel(0); @@ -1974,6 +2010,7 @@ impl AuthManager { codex_home: PathBuf::from("non-existent"), inner: RwLock::new(CachedAuth { auth: None, + bedrock_api_key: None, permanent_refresh_failure: None, }), auth_change_tx, @@ -2001,6 +2038,14 @@ impl AuthManager { .and_then(|cached| cached.auth.clone()) } + /// Current managed Amazon Bedrock API key without applying Codex auth precedence. + pub fn bedrock_api_key_auth_cached(&self) -> Option { + self.inner + .read() + .ok() + .and_then(|cached| cached.bedrock_api_key.clone()) + } + /// Subscribes to cached auth changes that can affect request recovery. pub fn auth_change_receiver(&self) -> watch::Receiver { self.auth_change_tx.subscribe() @@ -2103,7 +2148,8 @@ impl AuthManager { pub async fn reload(&self) -> bool { tracing::info!("Reloading auth"); let new_auth = self.load_auth().await; - self.set_cached_auth(new_auth) + let new_bedrock_api_key = self.load_bedrock_api_key_from_storage(); + self.set_cached_auth_state(new_auth, new_bedrock_api_key) } async fn reload_if_account_id_matches( @@ -2219,6 +2265,60 @@ impl AuthManager { .flatten() } + fn load_bedrock_api_key_from_storage(&self) -> Option { + load_stored_bedrock_api_key( + &self.codex_home, + self.auth_credentials_store_mode, + self.keyring_backend_kind, + ) + .ok() + .flatten() + } + + fn set_cached_auth_state( + &self, + new_auth: Option, + new_bedrock_api_key: Option, + ) -> bool { + if let Ok(mut guard) = self.inner.write() { + let previous = guard.auth.as_ref(); + let auth_changed = !Self::auths_equal(previous, new_auth.as_ref()); + let auth_changed_for_refresh = + !Self::auths_equal_for_refresh(previous, new_auth.as_ref()); + let bedrock_auth_changed = guard.bedrock_api_key != new_bedrock_api_key; + if auth_changed_for_refresh { + guard.permanent_refresh_failure = None; + } + let changed = auth_changed || bedrock_auth_changed; + tracing::info!("Reloaded auth, changed: {changed}"); + guard.auth = new_auth; + guard.bedrock_api_key = new_bedrock_api_key; + if auth_changed_for_refresh || bedrock_auth_changed { + self.auth_change_tx.send_modify(|revision| *revision += 1); + } + changed + } else { + false + } + } + + pub(super) fn set_cached_bedrock_api_key_auth( + &self, + new_bedrock_api_key: Option, + ) -> bool { + if let Ok(mut guard) = self.inner.write() { + let changed = guard.bedrock_api_key != new_bedrock_api_key; + tracing::info!("Reloaded managed Bedrock auth, changed: {changed}"); + guard.bedrock_api_key = new_bedrock_api_key; + if changed { + self.auth_change_tx.send_modify(|revision| *revision += 1); + } + changed + } else { + false + } + } + fn set_cached_auth(&self, new_auth: Option) -> bool { if let Ok(mut guard) = self.inner.write() { let previous = guard.auth.as_ref(); diff --git a/codex-rs/model-provider/src/amazon_bedrock/mod.rs b/codex-rs/model-provider/src/amazon_bedrock/mod.rs index 2a7230e081..a0e6cf3bc6 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/mod.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/mod.rs @@ -39,13 +39,13 @@ use mantle::runtime_base_url; pub(crate) struct AmazonBedrockModelProvider { pub(crate) info: ModelProviderInfo, pub(crate) aws: ModelProviderAwsAuthInfo, - auth_manager: Option>, + managed_auth: Option, } impl AmazonBedrockModelProvider { pub(crate) fn new( provider_info: ModelProviderInfo, - auth_manager: Option>, + managed_auth: Option, ) -> Self { let aws = provider_info .aws @@ -57,23 +57,12 @@ impl AmazonBedrockModelProvider { Self { info: provider_info, aws, - auth_manager, + managed_auth, } } fn managed_auth(&self) -> Option { - self.auth_manager - .as_ref() - .and_then(|auth_manager| auth_manager.auth_cached()) - .and_then(|auth| match auth { - CodexAuth::BedrockApiKey(auth) => Some(auth), - CodexAuth::ApiKey(_) - | CodexAuth::Chatgpt(_) - | CodexAuth::ChatgptAuthTokens(_) - | CodexAuth::Headers(_) - | CodexAuth::AgentIdentity(_) - | CodexAuth::PersonalAccessToken(_) => None, - }) + self.managed_auth.clone() } async fn auth(&self) -> Option { @@ -127,8 +116,7 @@ impl ModelProvider for AmazonBedrockModelProvider { } fn auth_manager(&self) -> Option> { - self.managed_auth() - .and_then(|_| self.auth_manager.as_ref().cloned()) + None } fn auth(&self) -> ModelProviderFuture<'_, Option> { @@ -208,22 +196,15 @@ mod tests { api_key: "managed-bedrock-api-key".to_string(), region: "us-east-1".to_string(), }; - let auth_manager = - AuthManager::from_auth_for_testing(CodexAuth::BedrockApiKey(managed_auth.clone())); let provider = AmazonBedrockModelProvider::new( ModelProviderInfo::create_amazon_bedrock_provider(Some(ModelProviderAwsAuthInfo { profile: Some("aws-profile-that-should-not-be-loaded".to_string()), region: Some("us-west-2".to_string()), })), - Some(auth_manager.clone()), + Some(managed_auth.clone()), ); - assert!(Arc::ptr_eq( - &provider - .auth_manager() - .expect("managed Bedrock auth manager should be exposed"), - &auth_manager, - )); + assert!(provider.auth_manager().is_none()); assert_eq!( provider.auth().await, Some(CodexAuth::BedrockApiKey(managed_auth)) @@ -256,12 +237,10 @@ mod tests { } #[tokio::test] - async fn openai_auth_is_not_exposed_to_bedrock() { + async fn bedrock_without_managed_auth_uses_aws_account() { let provider = AmazonBedrockModelProvider::new( ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None), - Some(AuthManager::from_auth_for_testing(CodexAuth::from_api_key( - "openai-api-key", - ))), + /*managed_auth*/ None, ); assert!(provider.auth_manager().is_none()); @@ -281,7 +260,7 @@ mod tests { fn capabilities_disable_unsupported_hosted_tools() { let provider = AmazonBedrockModelProvider::new( ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None), - /*auth_manager*/ None, + /*managed_auth*/ None, ); assert_eq!( @@ -298,7 +277,7 @@ mod tests { fn approval_review_preferred_model_uses_bedrock_gpt_5_4() { let provider = AmazonBedrockModelProvider::new( ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None), - /*auth_manager*/ None, + /*managed_auth*/ None, ); assert_eq!( diff --git a/codex-rs/model-provider/src/provider.rs b/codex-rs/model-provider/src/provider.rs index 6bc532ee50..230d175ba9 100644 --- a/codex-rs/model-provider/src/provider.rs +++ b/codex-rs/model-provider/src/provider.rs @@ -220,7 +220,10 @@ pub fn create_model_provider( auth_manager: Option>, ) -> SharedModelProvider { if provider_info.is_amazon_bedrock() { - Arc::new(AmazonBedrockModelProvider::new(provider_info, auth_manager)) + let managed_auth = auth_manager + .as_ref() + .and_then(|auth_manager| auth_manager.bedrock_api_key_auth_cached()); + Arc::new(AmazonBedrockModelProvider::new(provider_info, managed_auth)) } else { Arc::new(ConfiguredModelProvider::new(provider_info, auth_manager)) } @@ -529,17 +532,6 @@ mod tests { assert!(provider.auth_manager().is_none()); } - #[tokio::test] - async fn create_model_provider_uses_managed_auth_for_amazon_bedrock_provider() { - let auth = bedrock_api_key_auth(); - let provider = create_model_provider( - ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None), - Some(AuthManager::from_auth_for_testing(auth.clone())), - ); - - assert_eq!(provider.auth().await, Some(auth)); - } - #[test] fn openai_provider_returns_unauthenticated_openai_account_state() { let provider = create_model_provider(