diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index 91da0c1645..41c3dd4048 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -107,6 +107,7 @@ mod message_processor; mod models; mod models_refresh_worker; mod outgoing_message; +mod provider_runtime_refresh; mod request_processors; mod request_serialization; mod server_request_error; diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 62eff72660..6b61a79856 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -298,9 +298,8 @@ impl MessageProcessor { )), ) }); - let models_manager = thread_manager.get_models_manager(); let models_refresh_worker = - crate::models_refresh_worker::spawn(&models_manager, config.http_client_factory()); + crate::models_refresh_worker::spawn(&thread_manager, config.http_client_factory()); thread_manager .plugins_manager() .set_analytics_events_client(analytics_events_client.clone()); diff --git a/codex-rs/app-server/src/models_refresh_worker.rs b/codex-rs/app-server/src/models_refresh_worker.rs index ab785867a6..75313d123d 100644 --- a/codex-rs/app-server/src/models_refresh_worker.rs +++ b/codex-rs/app-server/src/models_refresh_worker.rs @@ -1,6 +1,7 @@ use std::sync::Arc; use std::time::Duration; +use codex_core::ThreadManager; use codex_http_client::HttpClientFactory; use codex_models_manager::manager::RefreshStrategy; use codex_models_manager::manager::SharedModelsManager; @@ -28,18 +29,42 @@ impl Drop for ModelsRefreshWorker { } pub(crate) fn spawn( - models_manager: &SharedModelsManager, + thread_manager: &Arc, http_client_factory: HttpClientFactory, ) -> ModelsRefreshWorker { - spawn_with_interval(models_manager, http_client_factory, MODELS_REFRESH_INTERVAL) + let thread_manager = Arc::downgrade(thread_manager); + spawn_with_source( + Arc::new(move || { + thread_manager + .upgrade() + .map(|thread_manager| thread_manager.get_models_manager()) + }), + http_client_factory, + MODELS_REFRESH_INTERVAL, + ) } +#[cfg(test)] fn spawn_with_interval( models_manager: &SharedModelsManager, http_client_factory: HttpClientFactory, refresh_interval: Duration, ) -> ModelsRefreshWorker { let models_manager = Arc::downgrade(models_manager); + spawn_with_source( + Arc::new(move || models_manager.upgrade()), + http_client_factory, + refresh_interval, + ) +} + +type ModelsManagerSource = Arc Option + Send + Sync>; + +fn spawn_with_source( + models_manager_source: ModelsManagerSource, + http_client_factory: HttpClientFactory, + refresh_interval: Duration, +) -> ModelsRefreshWorker { let shutdown = CancellationToken::new(); let worker_shutdown = shutdown.clone(); let task = tokio::spawn(async move { @@ -47,7 +72,7 @@ fn spawn_with_interval( if worker_shutdown.is_cancelled() { break; } - let Some(models_manager) = models_manager.upgrade() else { + let Some(models_manager) = models_manager_source() else { break; }; models_manager diff --git a/codex-rs/app-server/src/models_refresh_worker_tests.rs b/codex-rs/app-server/src/models_refresh_worker_tests.rs index 989bfd8608..9c926532ba 100644 --- a/codex-rs/app-server/src/models_refresh_worker_tests.rs +++ b/codex-rs/app-server/src/models_refresh_worker_tests.rs @@ -1,4 +1,5 @@ use std::sync::Arc; +use std::sync::RwLock; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::time::Duration; @@ -95,3 +96,45 @@ async fn refreshes_immediately_periodically_and_stops_when_dropped() { assert_eq!(endpoint.fetch_count.load(Ordering::SeqCst), 2); } + +#[tokio::test] +async fn refreshes_the_latest_models_manager_each_iteration() { + let codex_home = tempdir().expect("temp dir"); + let first_endpoint = TestModelsEndpoint::new(); + let first_models_manager: SharedModelsManager = Arc::new(OpenAiModelsManager::new( + codex_home.path().to_path_buf(), + first_endpoint.clone(), + /*auth_manager*/ None, + )); + let second_endpoint = TestModelsEndpoint::new(); + let second_models_manager: SharedModelsManager = Arc::new(OpenAiModelsManager::new( + codex_home.path().to_path_buf(), + second_endpoint.clone(), + /*auth_manager*/ None, + )); + let current_models_manager = Arc::new(RwLock::new(first_models_manager)); + let models_manager_source = { + let current_models_manager = Arc::clone(¤t_models_manager); + Arc::new(move || { + Some( + current_models_manager + .read() + .expect("models manager lock") + .clone(), + ) + }) + }; + let worker = spawn_with_source( + models_manager_source, + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + Duration::from_millis(25), + ); + + first_endpoint.wait_for_fetch_count(/*expected*/ 1).await; + *current_models_manager.write().expect("models manager lock") = second_models_manager; + second_endpoint.wait_for_fetch_count(/*expected*/ 1).await; + drop(worker); + + assert_eq!(first_endpoint.fetch_count.load(Ordering::SeqCst), 1); + assert_eq!(second_endpoint.fetch_count.load(Ordering::SeqCst), 1); +} diff --git a/codex-rs/app-server/src/provider_runtime_refresh.rs b/codex-rs/app-server/src/provider_runtime_refresh.rs new file mode 100644 index 0000000000..c6bccda9f0 --- /dev/null +++ b/codex-rs/app-server/src/provider_runtime_refresh.rs @@ -0,0 +1,19 @@ +use crate::config_manager::ConfigManager; +use crate::error_code::internal_error; +use codex_app_server_protocol::JSONRPCErrorError; +use codex_core::ThreadManager; + +pub(crate) async fn refresh_default_model_provider( + config_manager: &ConfigManager, + thread_manager: &ThreadManager, +) -> Result { + let config = config_manager + .load_latest_config(/*fallback_cwd*/ None) + .await + .map_err(|err| { + internal_error(format!( + "failed to reload config for model provider refresh: {err}" + )) + })?; + Ok(thread_manager.refresh_default_model_provider(&config)) +} 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 a598542247..6975f2ef9c 100644 --- a/codex-rs/app-server/src/request_processors/account_processor.rs +++ b/codex-rs/app-server/src/request_processors/account_processor.rs @@ -3,6 +3,7 @@ use super::bedrock_auth::set_user_model_provider_to_bedrock; use super::*; use crate::auth_mode::auth_mode_to_api; use crate::external_auth::ExternalAuthBridge; +use crate::provider_runtime_refresh::refresh_default_model_provider; use chrono::DateTime; use codex_model_provider::is_supported_amazon_bedrock_region; @@ -424,6 +425,7 @@ impl AccountRequestProcessor { ) .map_err(|err| internal_error(format!("failed to save Amazon Bedrock auth: {err}")))?; self.auth_manager.reload().await; + refresh_default_model_provider(&self.config_manager, &self.thread_manager).await?; Ok(LoginAccountResponse::AmazonBedrock {}) } .await; @@ -861,6 +863,7 @@ impl AccountRequestProcessor { if managed_bedrock_auth { clear_user_model_provider_if_bedrock(&self.config_manager).await?; + refresh_default_model_provider(&self.config_manager, &self.thread_manager).await?; } Self::maybe_refresh_plugin_caches_for_current_config( diff --git a/codex-rs/app-server/src/request_processors/config_processor.rs b/codex-rs/app-server/src/request_processors/config_processor.rs index 0eeb1aaca9..23185d898d 100644 --- a/codex-rs/app-server/src/request_processors/config_processor.rs +++ b/codex-rs/app-server/src/request_processors/config_processor.rs @@ -6,6 +6,7 @@ use crate::error_code::internal_error; use crate::error_code::invalid_request; use crate::outgoing_message::ConnectionRequestId; use crate::outgoing_message::OutgoingMessageSender; +use crate::provider_runtime_refresh::refresh_default_model_provider; use codex_analytics::AnalyticsEventsClient; use codex_app_server_protocol::ClientResponsePayload; use codex_app_server_protocol::ComputerUseRequirements; @@ -41,7 +42,6 @@ use codex_config::SandboxModeRequirement as CoreSandboxModeRequirement; use codex_core::ThreadManager; use codex_features::canonical_feature_for_key; use codex_features::feature_for_key; -use codex_model_provider::create_model_provider; use codex_plugin::PluginId; use codex_protocol::config_types::WebSearchMode; use serde_json::json; @@ -159,9 +159,7 @@ impl ConfigRequestProcessor { pub(crate) async fn model_provider_capabilities_read( &self, ) -> Result { - let config = self.load_latest_config(/*fallback_cwd*/ None).await?; - let provider = create_model_provider(config.model_provider, /*auth_manager*/ None); - let capabilities = provider.capabilities(); + let capabilities = self.thread_manager.default_model_provider_capabilities(); Ok(ModelProviderCapabilitiesReadResponse { namespace_tools: capabilities.namespace_tools, image_generation: capabilities.image_generation, @@ -169,9 +167,11 @@ impl ConfigRequestProcessor { }) } - pub(crate) async fn handle_config_mutation(&self) { + pub(crate) async fn handle_config_mutation(&self) -> Result<(), JSONRPCErrorError> { self.thread_manager.plugins_manager().clear_cache(); self.thread_manager.skills_service().clear_cache(); + refresh_default_model_provider(&self.config_manager, &self.thread_manager).await?; + Ok(()) } async fn handle_config_mutation_result( @@ -179,7 +179,7 @@ impl ConfigRequestProcessor { result: std::result::Result, ) -> Result { let response = result?; - self.handle_config_mutation().await; + self.handle_config_mutation().await?; Ok(response) } diff --git a/codex-rs/app-server/src/request_processors/external_agent_config_processor.rs b/codex-rs/app-server/src/request_processors/external_agent_config_processor.rs index c71abcb884..a59fd38980 100644 --- a/codex-rs/app-server/src/request_processors/external_agent_config_processor.rs +++ b/codex-rs/app-server/src/request_processors/external_agent_config_processor.rs @@ -226,8 +226,13 @@ impl ExternalAgentConfigRequestProcessor { let (pending_session_imports, session_validation_result) = self.validate_pending_session_imports(¶ms); let import_outcome = self.import_external_agent_config(params).await; - if needs_runtime_refresh { - self.config_processor.handle_config_mutation().await; + if needs_runtime_refresh + && let Err(err) = self.config_processor.handle_config_mutation().await + { + tracing::warn!( + error = %err.message, + "failed to refresh runtime state after external agent config import" + ); } self.outgoing .send_response( 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 fc6a3281b4..271573f41e 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -3785,7 +3785,7 @@ impl ThreadRequestProcessor { } } None if relation_filter.is_some() => None, - None => Some(vec![self.config.model_provider_id.clone()]), + None => Some(vec![self.thread_manager.default_model_provider_id()]), }; let (allowed_sources_vec, source_kind_filter) = if relation_filter.is_some() && source_kinds.is_none() { diff --git a/codex-rs/app-server/tests/suite/v2/mod.rs b/codex-rs/app-server/tests/suite/v2/mod.rs index 3d55079dc6..6e2159fc30 100644 --- a/codex-rs/app-server/tests/suite/v2/mod.rs +++ b/codex-rs/app-server/tests/suite/v2/mod.rs @@ -46,6 +46,7 @@ mod plugin_read; mod plugin_share; mod plugin_uninstall; mod process_exec; +mod provider_runtime_refresh; mod rate_limit_reset_credits; mod rate_limits; mod realtime_conversation; 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 new file mode 100644 index 0000000000..f4c1d9b712 --- /dev/null +++ b/codex-rs/app-server/tests/suite/v2/provider_runtime_refresh.rs @@ -0,0 +1,229 @@ +use std::time::Duration; + +use anyhow::Result; +use app_test_support::TestAppServer; +use app_test_support::create_fake_rollout; +use app_test_support::to_response; +use codex_app_server_protocol::ConfigValueWriteParams; +use codex_app_server_protocol::ConfigWriteResponse; +use codex_app_server_protocol::JSONRPCResponse; +use codex_app_server_protocol::LoginAccountResponse; +use codex_app_server_protocol::LogoutAccountResponse; +use codex_app_server_protocol::MergeStrategy; +use codex_app_server_protocol::ModelListParams; +use codex_app_server_protocol::ModelListResponse; +use codex_app_server_protocol::ModelProviderCapabilitiesReadParams; +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::WriteStatus; +use codex_model_provider_info::AMAZON_BEDROCK_GPT_5_5_MODEL_ID; +use pretty_assertions::assert_eq; +use serde_json::json; +use tempfile::TempDir; +use tokio::time::timeout; + +const DEFAULT_TIMEOUT: Duration = Duration::from_secs(60); + +fn write_mock_provider_config(codex_home: &TempDir) -> Result<()> { + std::fs::write( + codex_home.path().join("config.toml"), + r#" +model = "mock-model" +model_provider = "mock_provider" + +[model_providers.mock_provider] +name = "Mock provider" +base_url = "http://127.0.0.1:0/v1" +wire_api = "responses" +request_max_retries = 0 +stream_max_retries = 0 +"#, + )?; + Ok(()) +} + +async fn read_capabilities( + mcp: &mut TestAppServer, +) -> Result { + let request_id = mcp + .send_model_provider_capabilities_read_request(ModelProviderCapabilitiesReadParams {}) + .await?; + let response: JSONRPCResponse = timeout( + DEFAULT_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + to_response(response) +} + +async fn list_models(mcp: &mut TestAppServer) -> Result { + let request_id = mcp + .send_list_models_request(ModelListParams { + limit: Some(100), + cursor: None, + include_hidden: Some(true), + }) + .await?; + let response: JSONRPCResponse = timeout( + DEFAULT_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + to_response(response) +} + +async fn list_threads_for_default_provider(mcp: &mut TestAppServer) -> Result { + let request_id = mcp + .send_thread_list_request(ThreadListParams { + cursor: None, + limit: Some(100), + sort_key: None, + sort_direction: None, + model_providers: None, + source_kinds: None, + archived: None, + cwd: None, + use_state_db_only: false, + search_term: None, + parent_thread_id: None, + ancestor_thread_id: None, + }) + .await?; + let response: JSONRPCResponse = timeout( + DEFAULT_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + to_response(response) +} + +fn bedrock_capabilities() -> ModelProviderCapabilitiesReadResponse { + ModelProviderCapabilitiesReadResponse { + namespace_tools: true, + image_generation: false, + web_search: false, + } +} + +#[tokio::test] +async fn config_write_publishes_provider_runtime_before_responding() -> Result<()> { + let codex_home = TempDir::new()?; + write_mock_provider_config(&codex_home)?; + let mock_thread_id = create_fake_rollout( + codex_home.path(), + "2026-01-01T00-00-00", + "2026-01-01T00:00:00Z", + "mock provider thread", + Some("mock_provider"), + /*git_info*/ None, + )?; + let bedrock_thread_id = create_fake_rollout( + codex_home.path(), + "2026-01-01T00-00-01", + "2026-01-01T00:00:01Z", + "Bedrock provider thread", + Some("amazon-bedrock"), + /*git_info*/ None, + )?; + let mut mcp = TestAppServer::builder() + .with_codex_home(codex_home.path()) + .build() + .await?; + timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; + + assert_eq!( + list_threads_for_default_provider(&mut mcp) + .await? + .data + .into_iter() + .map(|thread| thread.id) + .collect::>(), + vec![mock_thread_id] + ); + + let request_id = mcp + .send_config_value_write_request(ConfigValueWriteParams { + key_path: "model_provider".to_string(), + value: json!("amazon-bedrock"), + 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 + ); + + assert_eq!(read_capabilities(&mut mcp).await?, bedrock_capabilities()); + assert!( + list_models(&mut mcp) + .await? + .data + .iter() + .any(|model| model.model == AMAZON_BEDROCK_GPT_5_5_MODEL_ID) + ); + assert_eq!( + list_threads_for_default_provider(&mut mcp) + .await? + .data + .into_iter() + .map(|thread| thread.id) + .collect::>(), + vec![bedrock_thread_id] + ); + Ok(()) +} + +#[tokio::test] +async fn managed_bedrock_login_and_logout_publish_provider_runtime() -> Result<()> { + let codex_home = TempDir::new()?; + write_mock_provider_config(&codex_home)?; + let mut mcp = TestAppServer::builder() + .with_codex_home(codex_home.path()) + .with_env_overrides(&[("OPENAI_API_KEY", None)]) + .build() + .await?; + timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; + + let login_request_id = mcp + .send_login_account_amazon_bedrock_request("managed-bedrock-api-key", "us-west-2") + .await?; + let login_response: JSONRPCResponse = timeout( + DEFAULT_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(login_request_id)), + ) + .await??; + assert_eq!( + to_response::(login_response)?, + LoginAccountResponse::AmazonBedrock {} + ); + assert_eq!(read_capabilities(&mut mcp).await?, bedrock_capabilities()); + + let logout_request_id = mcp.send_logout_account_request().await?; + let logout_response: JSONRPCResponse = timeout( + DEFAULT_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(logout_request_id)), + ) + .await??; + assert_eq!( + to_response::(logout_response)?, + LogoutAccountResponse {} + ); + assert_eq!( + read_capabilities(&mut mcp).await?, + ModelProviderCapabilitiesReadResponse { + namespace_tools: true, + image_generation: true, + web_search: true, + } + ); + Ok(()) +} diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index cfc26c5711..9487344fa8 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -69,6 +69,7 @@ pub use codex_mcp::SandboxState; mod mcp_openai_file; mod mcp_tool_call; pub(crate) mod mention_syntax; +mod model_provider_runtime; pub(crate) mod utils; pub use mention_syntax::PLUGIN_TEXT_MENTION_SIGIL; pub use mention_syntax::TOOL_MENTION_SIGIL; diff --git a/codex-rs/core/src/model_provider_runtime.rs b/codex-rs/core/src/model_provider_runtime.rs new file mode 100644 index 0000000000..e8c2c64c0c --- /dev/null +++ b/codex-rs/core/src/model_provider_runtime.rs @@ -0,0 +1,180 @@ +use crate::config::Config; +use arc_swap::ArcSwap; +use codex_login::AuthManager; +use codex_model_provider::ProviderCapabilities; +use codex_model_provider::create_model_provider; +use codex_model_provider_info::ModelProviderInfo; +use codex_models_manager::manager::SharedModelsManager; +use codex_protocol::openai_models::ModelsResponse; +use std::path::PathBuf; +use std::sync::Arc; +use std::sync::Mutex; + +/// The config inputs that determine which model manager is valid for a runtime. +#[derive(Debug, Clone, PartialEq)] +struct ModelProviderRuntimeIdentity { + provider_id: String, + provider_info: ModelProviderInfo, + codex_home: PathBuf, + config_model_catalog: Option, +} + +impl ModelProviderRuntimeIdentity { + fn from_config(config: &Config) -> 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(), + } + } + + fn matches_config(&self, config: &Config) -> 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 + } +} + +/// Immutable process-default model-provider state published as one coherent value. +struct DefaultModelProviderSnapshot { + generation: u64, + identity: ModelProviderRuntimeIdentity, + models_manager: SharedModelsManager, + capabilities: ProviderCapabilities, +} + +impl DefaultModelProviderSnapshot { + fn new( + generation: u64, + identity: ModelProviderRuntimeIdentity, + auth_manager: Arc, + ) -> Self { + let provider = create_model_provider( + 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 { + generation, + identity, + models_manager, + capabilities, + } + } + + fn matches_config(&self, config: &Config) -> bool { + self.identity.matches_config(config) + } +} + +/// Owns the atomically replaceable process-default model-provider snapshot. +pub(crate) struct DefaultModelProviderRuntime { + snapshot: ArcSwap, + refresh_lock: Mutex<()>, +} + +impl DefaultModelProviderRuntime { + pub(crate) fn new(config: &Config, auth_manager: Arc) -> Self { + Self::from_identity( + ModelProviderRuntimeIdentity::from_config(config), + auth_manager, + ) + } + + pub(crate) fn from_parts( + provider_id: String, + provider_info: ModelProviderInfo, + codex_home: PathBuf, + config_model_catalog: Option, + auth_manager: Arc, + ) -> Self { + Self::from_identity( + ModelProviderRuntimeIdentity { + provider_id, + provider_info, + codex_home, + config_model_catalog, + }, + auth_manager, + ) + } + + fn from_identity( + identity: ModelProviderRuntimeIdentity, + auth_manager: Arc, + ) -> Self { + Self { + snapshot: ArcSwap::from_pointee(DefaultModelProviderSnapshot::new( + /*generation*/ 0, + identity, + auth_manager, + )), + refresh_lock: Mutex::new(()), + } + } + + pub(crate) fn refresh(&self, config: &Config, auth_manager: Arc) -> bool { + let _refresh_guard = self + .refresh_lock + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let current = self.snapshot.load_full(); + if current.matches_config(config) { + return false; + } + + let next = DefaultModelProviderSnapshot::new( + current.generation.saturating_add(1), + ModelProviderRuntimeIdentity::from_config(config), + auth_manager, + ); + self.snapshot.store(Arc::new(next)); + true + } + + pub(crate) fn models_manager(&self) -> SharedModelsManager { + 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 provider_id(&self) -> String { + self.snapshot.load().identity.provider_id.clone() + } + + pub(crate) fn capabilities(&self) -> ProviderCapabilities { + self.snapshot.load().capabilities + } + + pub(crate) fn generation(&self) -> u64 { + self.snapshot.load().generation + } +} + +pub(crate) fn build_models_manager( + config: &Config, + auth_manager: Arc, +) -> SharedModelsManager { + let provider = create_model_provider(config.model_provider.clone(), Some(auth_manager)); + provider.models_manager( + config.codex_home.to_path_buf(), + config.model_catalog.clone(), + ) +} diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index de37484f8b..d98e8bbcc8 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -8,6 +8,7 @@ use crate::current_time::TimeProvider; 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::rollout::truncation; use crate::session::Codex; use crate::session::CodexSpawnArgs; @@ -36,7 +37,7 @@ use codex_login::AuthManager; use codex_login::CodexAuth; use codex_login::default_client::CODEX_INTERNAL_ORIGINATOR_OVERRIDE_ENV_VAR; use codex_login::default_client::originator; -use codex_model_provider::create_model_provider; +use codex_model_provider::ProviderCapabilities; use codex_model_provider_info::ModelProviderInfo; use codex_model_provider_info::OPENAI_PROVIDER_ID; use codex_models_manager::manager::RefreshStrategy; @@ -240,7 +241,7 @@ pub(crate) struct ThreadManagerState { threads: Arc>>>, thread_created_tx: broadcast::Sender, auth_manager: Arc, - models_manager: SharedModelsManager, + default_model_provider_runtime: DefaultModelProviderRuntime, environment_manager: Arc, skills_service: Arc, plugins_manager: Arc, @@ -263,11 +264,7 @@ pub fn build_models_manager( config: &Config, auth_manager: Arc, ) -> SharedModelsManager { - let provider = create_model_provider(config.model_provider.clone(), Some(auth_manager)); - provider.models_manager( - config.codex_home.to_path_buf(), - config.model_catalog.clone(), - ) + crate::model_provider_runtime::build_models_manager(config, auth_manager) } pub fn thread_store_from_config( @@ -337,7 +334,10 @@ impl ThreadManager { state: Arc::new(ThreadManagerState { threads: Arc::new(RwLock::new(HashMap::new())), thread_created_tx, - models_manager: build_models_manager(config, auth_manager.clone()), + default_model_provider_runtime: DefaultModelProviderRuntime::new( + config, + auth_manager.clone(), + ), environment_manager, skills_service, plugins_manager, @@ -456,8 +456,13 @@ impl ThreadManager { state: Arc::new(ThreadManagerState { threads: Arc::new(RwLock::new(HashMap::new())), thread_created_tx, - models_manager: create_model_provider(provider, Some(auth_manager.clone())) - .models_manager(codex_home, /*config_model_catalog*/ None), + default_model_provider_runtime: DefaultModelProviderRuntime::from_parts( + OPENAI_PROVIDER_ID.to_string(), + provider, + codex_home, + /*config_model_catalog*/ None, + auth_manager.clone(), + ), environment_manager, skills_service, plugins_manager, @@ -539,7 +544,7 @@ impl ThreadManager { } pub fn get_models_manager(&self) -> SharedModelsManager { - self.state.models_manager.clone() + self.state.default_model_provider_runtime.models_manager() } pub async fn list_models( @@ -547,14 +552,35 @@ impl ThreadManager { refresh_strategy: RefreshStrategy, http_client_factory: codex_http_client::HttpClientFactory, ) -> Vec { - self.state - .models_manager + self.get_models_manager() .list_models(refresh_strategy, http_client_factory) .await } pub fn list_collaboration_modes(&self) -> Vec { - self.state.models_manager.list_collaboration_modes() + self.get_models_manager().list_collaboration_modes() + } + + /// Publishes a new process-default model-provider runtime when its config identity changed. + pub fn refresh_default_model_provider(&self, config: &Config) -> bool { + self.state + .default_model_provider_runtime + .refresh(config, Arc::clone(&self.state.auth_manager)) + } + + /// Returns the effective provider ID for new threads that follow the process default. + pub fn default_model_provider_id(&self) -> String { + self.state.default_model_provider_runtime.provider_id() + } + + /// Returns the provider-owned capabilities for the current process default. + pub fn default_model_provider_capabilities(&self) -> ProviderCapabilities { + self.state.default_model_provider_runtime.capabilities() + } + + /// Returns the generation of the current process-default provider snapshot. + pub fn default_model_provider_generation(&self) -> u64 { + self.state.default_model_provider_runtime.generation() } pub async fn list_thread_ids(&self) -> Vec { @@ -1584,6 +1610,9 @@ 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 CodexSpawnOk { codex, thread_id, .. } = Box::pin(Codex::spawn(CodexSpawnArgs { @@ -1592,7 +1621,7 @@ impl ThreadManagerState { user_instructions, installation_id: self.installation_id.clone(), auth_manager, - models_manager: Arc::clone(&self.models_manager), + models_manager, 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 d29d6d04e5..cbeebf34d7 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -9,12 +9,17 @@ use crate::session::tests::make_session_and_context; use crate::tasks::InterruptedTurnHistoryMarker; use crate::tasks::interrupted_turn_history_marker; use codex_extension_api::empty_extension_registry; +use codex_model_provider::ProviderCapabilities; +use codex_model_provider_info::AMAZON_BEDROCK_PROVIDER_ID; use codex_models_manager::manager::RefreshStrategy; +use codex_models_manager::manager::SharedModelsManager; +use codex_models_manager::model_info::model_info_from_slug; use codex_protocol::capabilities::CapabilityRootLocation; use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_protocol::models::ContentItem; use codex_protocol::models::ReasoningItemReasoningSummary; use codex_protocol::models::ResponseItem; +use codex_protocol::openai_models::ModelVisibility; use codex_protocol::openai_models::ModelsResponse; use codex_protocol::protocol::AgentMessageEvent; use codex_protocol::protocol::InitialHistory; @@ -37,6 +42,28 @@ use wiremock::MockServer; const TEST_INSTALLATION_ID: &str = "11111111-1111-4111-8111-111111111111"; +fn static_model_catalog(slug: &str) -> ModelsResponse { + let mut model = model_info_from_slug(slug); + model.visibility = ModelVisibility::List; + model.priority = 0; + model.used_fallback_model_metadata = false; + ModelsResponse { + models: vec![model], + } +} + +async fn listed_model_ids(models_manager: &SharedModelsManager) -> Vec { + models_manager + .list_models( + RefreshStrategy::Offline, + crate::test_support::default_http_client_factory(), + ) + .await + .into_iter() + .map(|model| model.id) + .collect() +} + struct FakeAgentGraphStore { root_thread_id: ThreadId, descendant_thread_ids: Vec, @@ -1369,6 +1396,139 @@ async fn new_uses_active_provider_for_model_refresh() { assert_eq!(models_mock.requests().len(), 1); } +#[tokio::test] +async fn refresh_default_model_provider_publishes_one_coherent_snapshot() { + let mut config = test_config().await; + let manager = ThreadManager::with_models_provider_for_tests( + CodexAuth::from_api_key("dummy"), + config.model_provider.clone(), + ); + let initial_models_manager = manager.get_models_manager(); + + assert_eq!(manager.default_model_provider_id(), OPENAI_PROVIDER_ID); + assert_eq!(manager.default_model_provider_generation(), 0); + assert_eq!( + manager.default_model_provider_capabilities(), + ProviderCapabilities::default() + ); + + config.model_provider_id = AMAZON_BEDROCK_PROVIDER_ID.to_string(); + config.model_provider = ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None); + config.model_catalog = None; + + assert!(manager.refresh_default_model_provider(&config)); + let bedrock_models_manager = manager.get_models_manager(); + assert!(!Arc::ptr_eq( + &initial_models_manager, + &bedrock_models_manager + )); + assert_eq!( + ( + manager.default_model_provider_id(), + manager.default_model_provider_generation(), + manager.default_model_provider_capabilities(), + ), + ( + AMAZON_BEDROCK_PROVIDER_ID.to_string(), + 1, + ProviderCapabilities { + namespace_tools: true, + image_generation: false, + web_search: false, + }, + ) + ); + + assert!(!manager.refresh_default_model_provider(&config)); + assert_eq!(manager.default_model_provider_generation(), 1); + assert!(Arc::ptr_eq( + &bedrock_models_manager, + &manager.get_models_manager() + )); +} + +#[tokio::test] +async fn new_threads_select_models_manager_for_their_effective_config() { + 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 provider_a_models_manager = manager.get_models_manager(); + assert_eq!( + listed_model_ids(&provider_a_models_manager).await, + vec!["model-a".to_string()] + ); + + let mut provider_b_config = provider_a_config.clone(); + 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)); + + let provider_b_models_manager = manager.get_models_manager(); + assert_eq!( + listed_model_ids(&provider_b_models_manager).await, + vec!["model-b".to_string()] + ); + + let default_thread = manager + .start_thread(provider_b_config) + .await + .expect("start default-provider thread"); + assert!(Arc::ptr_eq( + &default_thread.thread.codex.session.services.models_manager, + &provider_b_models_manager, + )); + + let explicit_thread = manager + .start_thread(provider_a_config) + .await + .expect("start explicitly configured provider thread"); + let explicit_models_manager = &explicit_thread.thread.codex.session.services.models_manager; + assert!(!Arc::ptr_eq( + explicit_models_manager, + &provider_b_models_manager + )); + assert_eq!( + listed_model_ids(explicit_models_manager).await, + vec!["model-a".to_string()] + ); + + 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 =