diff --git a/codex-rs/core-plugins/src/marketplace_staging_tests.rs b/codex-rs/core-plugins/src/marketplace_staging_tests.rs index e5ac5edc65..b53fd100ab 100644 --- a/codex-rs/core-plugins/src/marketplace_staging_tests.rs +++ b/codex-rs/core-plugins/src/marketplace_staging_tests.rs @@ -94,7 +94,10 @@ fn cleanup_reclaims_unpublished_and_aged_legacy_clones() -> Result<()> { staging_root.join(format!("{}file", MarketplaceStagingDir::directory_prefix())); fs::write(&unrelated_file, "keep")?; - let now = SystemTime::now(); + // Base synthetic ages on the filesystem clock, which can differ from the wall clock. + let now = fs::metadata(&legacy_add)? + .modified()? + .max(fs::metadata(&legacy_upgrade)?.modified()?); for (age, expected_legacy_dirs) in [ (Duration::ZERO, [true, true]), ( diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index b2d22649a6..10ab88e044 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -2404,7 +2404,10 @@ impl ThreadManagerState { instructions, installation_id: self.installation_id.clone(), auth_manager, - models_manager: Arc::clone(&self.models_manager), + models_manager: self + .models_manager + .for_originator(&originator) + .unwrap_or_else(|| Arc::clone(&self.models_manager)), git_root_discovery: Arc::clone(&self.git_root_discovery), environment_manager: Arc::clone(&self.environment_manager), skills_service: Arc::clone(&self.skills_service), diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index 9578c92459..98d4547f2f 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -2908,6 +2908,57 @@ async fn new_uses_active_provider_for_model_refresh() { ) .await; assert_eq!(models_mock.requests().len(), 1); + + // Thread startup must select the originator catalog before resolving model metadata. + for originator in [ + "codex_work_cca", + "codex_work_web", + "codex_work_desktop", + "codex_work_mobile", + "chatgpt_cca", + ] { + let mut model = codex_models_manager::model_info::model_info_from_slug("gpt-5.5"); + model.description = Some(originator.to_string()); + wiremock::Mock::given(wiremock::matchers::method("GET")) + .and(wiremock::matchers::path("/models")) + .and(wiremock::matchers::header("originator", originator)) + .respond_with( + wiremock::ResponseTemplate::new(200).set_body_json(ModelsResponse { + models: vec![model], + }), + ) + .expect(1) + .mount(&server) + .await; + let started = manager + .start_thread(StartThreadOptions { + metrics_service_name: Some(originator.to_string()), + environments: Some(Vec::new()), + ..StartThreadOptions::new(config.clone()) + }) + .await + .expect("start scoped thread"); + let model = started + .thread + .session + .services + .models_manager + .get_model_info("gpt-5.5", &config.to_models_manager_config()) + .await; + assert_eq!(model.description.as_deref(), Some(originator)); + } + let legacy = manager + .start_thread(StartThreadOptions { + metrics_service_name: Some("codex_cli_rs".to_string()), + environments: Some(Vec::new()), + ..StartThreadOptions::new(config) + }) + .await + .expect("start legacy thread"); + assert!(Arc::ptr_eq( + &legacy.thread.session.services.models_manager, + &manager.get_models_manager() + )); } #[tokio::test] diff --git a/codex-rs/core/tests/suite/scenarios_model_instructions.rs b/codex-rs/core/tests/suite/scenarios_model_instructions.rs index 2e16542430..15dce4189f 100644 --- a/codex-rs/core/tests/suite/scenarios_model_instructions.rs +++ b/codex-rs/core/tests/suite/scenarios_model_instructions.rs @@ -1,11 +1,127 @@ -//! Instruction refresh and compaction with an unchanged model and tool catalog. +//! Model-visible instructions and tools from catalog routing, refresh, and compaction. use super::*; use codex_models_manager::manager::RefreshStrategy; use codex_models_manager::model_info::model_info_from_slug; +use codex_protocol::openai_models::ApplyPatchToolType; use codex_protocol::openai_models::ModelVisibility; use codex_protocol::openai_models::ModelsResponse; use core_test_support::responses; +use pretty_assertions::assert_eq; + +/// Originator routing supplies the actual turn's instructions and tools without changing legacy callers. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn originator_catalogs_control_turn_instructions_and_tools() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = wiremock::MockServer::start().await; + for (originator, instructions, patch_tool) in [ + ("provider-default", "Default catalog instructions.", None), + ( + "codex_work_cca", + "Work catalog instructions.", + Some(ApplyPatchToolType::Freeform), + ), + ("chatgpt_cca", "Chat catalog instructions.", None), + ] { + let mut model = model_info_from_slug("gpt-5.5"); + model.visibility = ModelVisibility::List; + model.use_responses_lite = false; + model.tool_mode = Some(ToolMode::Direct); + model.apply_patch_tool_type = patch_tool; + model + .model_messages + .get_or_insert_default() + .instructions_template = Some(instructions.into()); + wiremock::Mock::given(wiremock::matchers::method("GET")) + .and(wiremock::matchers::path("/v1/models")) + .and(wiremock::matchers::header("originator", originator)) + .respond_with( + wiremock::ResponseTemplate::new(200).set_body_json(ModelsResponse { + models: vec![model], + }), + ) + .expect(1) + .mount(&server) + .await; + } + let mock = mount_sse_sequence( + &server, + (1..=5) + .map(|index| sse(vec![ev_completed(&format!("response-{index}"))])) + .collect(), + ) + .await; + let mut test = test_codex() + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_config(|config| { + configure_scenario_catalog(config); + config.model_catalog = None; + config + .model_provider + .http_headers + .get_or_insert_default() + .insert("originator".into(), "provider-default".into()); + }) + .build_with_auto_env(&server) + .await?; + test.submit_turn("Describe the task.").await?; + for originator in ["codex_work_cca", "chatgpt_cca", "unknown", "codex_work_cca"] { + let started = test + .thread_manager + .start_thread(StartThreadOptions { + metrics_service_name: Some(originator.into()), + ..test.start_thread_options().await + }) + .await?; + test.codex = started.thread; + test.session_configured = started.session_configured; + test.submit_turn("Describe the task.").await?; + } + let requests = mock.requests(); + assert_eq!(requests.len(), 5); + for (request, expected) in requests.iter().zip([ + "Default catalog instructions.", + "Work catalog instructions.", + "Chat catalog instructions.", + "Default catalog instructions.", + "Work catalog instructions.", + ]) { + assert_eq!(request.instructions_text(), expected); + } + insta::assert_snapshot!( + "originator_catalog_instructions_and_tools", + context_snapshot::format_request_history_snapshot( + "Default, Work, Chat, unsupported, and repeated Work threads use their own catalog instructions and tools. Unsupported callers preserve the default request; repeated Work reuses its catalog.", + &requests, + &ContextSnapshotOptions::default() + .rewrite_known_segments() + .include_request_settings(), + ) + ); + let tools = requests + .iter() + .map(|request| { + request.body_json()["tools"] + .as_array() + .expect("Responses tool catalog") + .clone() + }) + .collect::>(); + assert_eq!(tools[0], tools[3]); + assert_eq!(tools[1], tools[4]); + assert_eq!( + tools + .iter() + .map(|declarations| declarations + .iter() + .any(|tool| tool["name"] == "apply_patch")) + .collect::>(), + vec![false, true, false, false, true], + ); + server.verify().await; + Ok(()) +} #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn catalog_refresh_replaces_instructions_and_compaction_installs_latest_base() -> Result<()> { diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__model_instructions__originator_catalog_instructions_and_tools.snap b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__model_instructions__originator_catalog_instructions_and_tools.snap new file mode 100644 index 0000000000..1f977eeabe --- /dev/null +++ b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__model_instructions__originator_catalog_instructions_and_tools.snap @@ -0,0 +1,95 @@ +--- +source: core/tests/suite/scenarios_model_instructions.rs +expression: "context_snapshot::format_request_history_snapshot(\"Default, Work, Chat, unsupported, and repeated Work threads use their own catalog instructions and tools. Unsupported callers preserve the default request; repeated Work reuses its catalog.\",\n&requests,\n&ContextSnapshotOptions::default().rewrite_known_segments().include_request_settings(),)" +--- +Scenario: Default, Work, Chat, unsupported, and repeated Work threads use their own catalog instructions and tools. Unsupported callers preserve the default request; repeated Work reuses its catalog. + +## Window 1 +Settings: + include: ["reasoning.encrypted_content"] + model: "gpt-5.5" + parallel_tool_calls: true + prompt_cache_key: "" + reasoning: {"summary":"auto"} + store: false + stream: true + tool_choice: "auto" + tools (6; hash=d558f97965d6db2b): + - function/exec_command: Runs a command in a PTY, returning output or a session ID for ongoing interaction.; args=[cmd, justification, login, max_output_tokens, prefix_rule, sandbox_permissions, shell, tty, workdir, yield_time_ms]; hash=c04108e51a8548c5 + - function/write_stdin: Writes characters to an existing unified exec session and returns recent output.; args=[chars, max_output_tokens, session_id, yield_time_ms]; hash=7fb62d8a3e2df4e2 + - function/request_user_input: Request user input for one to three short questions and wait for the response. This tool is only ...; args=[questions]; hash=64a729b25c424fc6 + - function/view_image: View a local image file from the filesystem when visual inspection is needed. Use this for images...; args=[path]; hash=1908de9c0be7118a + - namespace/multi_agent_v1: Tools for spawning and managing sub-agents.; hash=6ac4ec60e68e50c0 + - function/close_agent + - function/resume_agent + - function/send_input + - function/spawn_agent + - function/wait_agent + - web_search; hash=8586237e354ee5ed +-- request 1 (turn) -- +00:message/developer: + +01:message/developer: + +02:message/user: + +03:message/user: + Describe the task. + +## Window 2 (after request 1: input diverged at item 00) +Settings: relative to window 1 + prompt_cache_key: "" + tools (7; hash=7cc079a7a7aa5604): + + custom/apply_patch: The `apply_patch` tool can be used to edit files. This is a FREEFORM tool, so do not wrap the pat...; hash=dea9f12b0111a4ba +-- request 2 (turn) -- +00:message/developer: + +01:message/developer: + +02:message/user: + +03:message/user: + Describe the task. + +## Window 3 (after request 2: input diverged at item 00) +Settings: relative to window 2 + prompt_cache_key: "" + tools (6; hash=d558f97965d6db2b): + - removed custom/apply_patch +-- request 3 (turn) -- +00:message/developer: + +01:message/developer: + +02:message/user: + +03:message/user: + Describe the task. + +## Window 4 (after request 3: input diverged at item 00) +Settings: relative to window 3 + prompt_cache_key: "" +-- request 4 (turn) -- +00:message/developer: + +01:message/developer: + +02:message/user: + +03:message/user: + Describe the task. + +## Window 5 (after request 4: input diverged at item 00) +Settings: relative to window 4 + prompt_cache_key: "" + tools (7; hash=7cc079a7a7aa5604): + + custom/apply_patch: The `apply_patch` tool can be used to edit files. This is a FREEFORM tool, so do not wrap the pat...; hash=dea9f12b0111a4ba +-- request 5 (turn) -- +00:message/developer: + +01:message/developer: + +02:message/user: + +03:message/user: + Describe the task. diff --git a/codex-rs/model-provider/src/models_endpoint.rs b/codex-rs/model-provider/src/models_endpoint.rs index 9734a38abf..be9f44af6e 100644 --- a/codex-rs/model-provider/src/models_endpoint.rs +++ b/codex-rs/model-provider/src/models_endpoint.rs @@ -25,6 +25,7 @@ use codex_login::default_client::ClientRedirectPolicy; use codex_login::default_client::create_client_for_route_async; use codex_model_provider_info::CHATGPT_CODEX_BASE_URL; use codex_model_provider_info::ModelProviderInfo; +use codex_models_manager::ModelCatalogOriginator; use codex_models_manager::manager::ModelsEndpointClient; use codex_models_manager::manager::ModelsEndpointFuture; use codex_models_manager::manager::ModelsEndpointResponse; @@ -211,6 +212,25 @@ impl OpenAiModelsEndpoint { } impl ModelsEndpointClient for OpenAiModelsEndpoint { + fn for_originator( + &self, + originator: ModelCatalogOriginator, + ) -> Option> { + let mut provider_info = self.provider_info.clone(); + let headers = provider_info.http_headers.get_or_insert_default(); + headers.retain(|name, _| !name.eq_ignore_ascii_case("originator")); + headers.insert("originator".to_string(), originator.as_str().into()); + if let Some(headers) = &mut provider_info.env_http_headers { + headers.retain(|name, _| !name.eq_ignore_ascii_case("originator")); + } + Some(Arc::new(Self { + provider_info, + auth_manager: self.auth_manager.clone(), + gateway_auth_manager: self.gateway_auth_manager.clone(), + transport_builder: Arc::clone(&self.transport_builder), + })) + } + fn supports_api_key_models(&self) -> bool { self.provider_info.model_catalog_url.is_some() || (self.provider_info.is_openai() && self.provider_info.base_url.is_none()) @@ -396,6 +416,10 @@ impl RequestTelemetry for ModelsRequestTelemetry { #[path = "models_endpoint_timeout_tests.rs"] mod timeout_tests; +#[cfg(test)] +#[path = "models_originator_tests.rs"] +mod originator_tests; + #[cfg(test)] mod tests { use std::num::NonZeroU64; diff --git a/codex-rs/model-provider/src/models_originator_tests.rs b/codex-rs/model-provider/src/models_originator_tests.rs new file mode 100644 index 0000000000..392d3c4ebe --- /dev/null +++ b/codex-rs/model-provider/src/models_originator_tests.rs @@ -0,0 +1,189 @@ +//! End-to-end catalog partitioning through provider headers, manager state, and file storage. + +use super::*; +use codex_http_client::OutboundProxyPolicy; +use codex_models_manager::manager::ModelsManager; +use codex_models_manager::manager::OpenAiModelsManager; +use codex_models_manager::manager::RefreshStrategy; +use codex_models_manager::model_info::model_info_from_slug; +use codex_protocol::openai_models::ModelsResponse; +use pretty_assertions::assert_eq; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::header; +use wiremock::matchers::method; +use wiremock::matchers::path; + +/// Runtime discovery updates reach existing and newly created scopes, including cached catalogs. +#[tokio::test] +async fn originator_managers_share_api_key_discovery_policy() { + let server = MockServer::start().await; + let home = tempfile::tempdir().unwrap(); + let mut provider = ModelProviderInfo::create_openai_provider(Some(server.uri())); + provider.model_catalog_url = Some(format!("{}/models", server.uri()).into()); + provider.http_headers = Some(std::collections::HashMap::from([( + codex_model_provider_info::RESIDENCY_HEADER_NAME.to_string(), + "us".into(), + )])); + let mut model = model_info_from_slug("gpt-5.5"); + model.visibility = codex_protocol::openai_models::ModelVisibility::List; + model.used_fallback_model_metadata = false; + let catalog = ModelsResponse { + models: vec![model], + }; + for originator in ["codex_work_cca", "chatgpt_cca"] { + Mock::given(method("GET")) + .and(path("/models")) + .and(header("originator", originator)) + .respond_with(ResponseTemplate::new(200).set_body_json(&catalog)) + .expect(1) + .mount(&server) + .await; + } + let auth = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("test-key")); + let endpoint = Arc::new(OpenAiModelsEndpoint::new( + provider, + Some(Arc::clone(&auth)), + /*gateway_auth_manager*/ None, + )); + let manager = OpenAiModelsManager::new(home.path().to_path_buf(), endpoint, Some(auth)); + let work = manager.for_originator("codex_work_cca").unwrap(); + let bundled = codex_models_manager::bundled_models_response().unwrap(); + let client = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault); + assert_eq!( + work.raw_model_catalog(RefreshStrategy::Online, client.clone()) + .await, + bundled + ); + assert!(server.received_requests().await.unwrap().is_empty()); + + manager.set_api_key_model_discovery_enabled(/*enabled*/ true); + let chat = manager.for_originator("chatgpt_cca").unwrap(); + for child in [&work, &chat] { + assert_eq!( + child + .raw_model_catalog(RefreshStrategy::OnlineIfUncached, client.clone()) + .await, + catalog + ); + } + + manager.set_api_key_model_discovery_enabled(/*enabled*/ false); + for child in [&work, &chat] { + assert_eq!( + child + .raw_model_catalog(RefreshStrategy::Online, client.clone()) + .await, + bundled + ); + assert_eq!(child.try_get_remote_models().unwrap(), bundled.models); + } + + manager.set_api_key_model_discovery_enabled(/*enabled*/ true); + for child in [&work, &chat] { + assert_eq!( + child + .raw_model_catalog(RefreshStrategy::OnlineIfUncached, client.clone()) + .await, + catalog + ); + } + server.verify().await; +} + +/// Work, Chat, and legacy callers retain their own catalogs across concurrent access and restarts. +#[tokio::test] +async fn originator_catalogs_are_isolated_and_reused_from_disk() { + let server = MockServer::start().await; + let home = tempfile::tempdir().unwrap(); + let mut provider = ModelProviderInfo::create_openai_provider(Some(server.uri())); + // The managed-residency test toggles the process-wide US policy. Keep this + // provider's cache identity stable whether that policy is enabled or disabled. + provider.http_headers = Some(std::collections::HashMap::from([ + ("Originator".to_string(), "provider-default".into()), + ( + codex_model_provider_info::RESIDENCY_HEADER_NAME.to_string(), + "us".into(), + ), + ])); + let mut catalogs = Vec::new(); + for originator in [ + "provider-default", + "codex_work_cca", + "codex_work_web", + "codex_work_desktop", + "codex_work_mobile", + "chatgpt_cca", + ] { + let mut model = model_info_from_slug("gpt-5.5"); + model.visibility = codex_protocol::openai_models::ModelVisibility::List; + model.used_fallback_model_metadata = false; + model.description = Some(originator.to_string()); + let catalog = ModelsResponse { + models: vec![model], + }; + Mock::given(method("GET")) + .and(path("/models")) + .and(header("originator", originator)) + .respond_with(ResponseTemplate::new(200).set_body_json(&catalog)) + .expect(1) + .mount(&server) + .await; + catalogs.push(catalog); + } + let client = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault); + for _ in 0..2 { + let auth = + AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); + let endpoint = Arc::new(OpenAiModelsEndpoint::new( + provider.clone(), + Some(Arc::clone(&auth)), + /*gateway_auth_manager*/ None, + )); + let manager = OpenAiModelsManager::new(home.path().to_path_buf(), endpoint, Some(auth)); + let work = manager.for_originator("codex_work_cca").unwrap(); + let chat = manager.for_originator("chatgpt_cca").unwrap(); + assert!(Arc::ptr_eq( + &work, + &manager.for_originator("codex_work_cca").unwrap() + )); + assert!(Arc::ptr_eq( + &chat, + &manager.for_originator("chatgpt_cca").unwrap() + )); + let (default_catalog, work_catalog, chat_catalog) = tokio::join!( + manager.raw_model_catalog(RefreshStrategy::OnlineIfUncached, client.clone()), + work.raw_model_catalog(RefreshStrategy::OnlineIfUncached, client.clone()), + chat.raw_model_catalog(RefreshStrategy::OnlineIfUncached, client.clone()), + ); + let mut actual = vec![default_catalog, work_catalog]; + for originator in ["codex_work_web", "codex_work_desktop", "codex_work_mobile"] { + actual.push( + manager + .for_originator(originator) + .unwrap() + .raw_model_catalog(RefreshStrategy::OnlineIfUncached, client.clone()) + .await, + ); + } + actual.push(chat_catalog); + assert_eq!(actual, catalogs); + for originator in [ + "", + "unknown", + "codex_cli_rs", + "codex_vscode", + "chatgpt_cca_local_dev", + ] { + assert!(manager.for_originator(originator).is_none()); + } + } + assert!(home.path().join("models_cache.json").exists()); + assert!( + home.path() + .join("models_cache.codex_work_cca.json") + .exists() + ); + assert!(home.path().join("models_cache.chatgpt_cca.json").exists()); +} diff --git a/codex-rs/models-manager/src/cache.rs b/codex-rs/models-manager/src/cache.rs index 7a0890760a..41c96d3ea8 100644 --- a/codex-rs/models-manager/src/cache.rs +++ b/codex-rs/models-manager/src/cache.rs @@ -1,3 +1,4 @@ +use crate::ModelCatalogOriginator; use chrono::DateTime; use chrono::Utc; use codex_protocol::openai_models::ModelInfo; @@ -9,6 +10,7 @@ use std::io; use std::io::ErrorKind; use std::path::PathBuf; use std::pin::Pin; +use std::sync::Arc; use std::time::Duration; use tokio::fs; use tracing::info; @@ -23,6 +25,12 @@ use tracing::info; /// Cache failures are non-fatal. The models manager logs them and falls back to the configured /// models endpoint. pub trait ModelsCache: fmt::Debug + Send + Sync { + /// Return a separate cache partition, or `None` to disable storage for this originator. + /// Implementations must isolate loads, writes, and TTL renewal from the default catalog. + fn for_originator(&self, _originator: ModelCatalogOriginator) -> Option> { + None + } + /// Loads a fresh entry for `client_version`. /// /// Returns `Ok(None)` for a normal miss, including an absent, stale, or version-mismatched @@ -141,6 +149,14 @@ impl FileModelsCache { } impl ModelsCache for FileModelsCache { + fn for_originator(&self, originator: ModelCatalogOriginator) -> Option> { + Some(Arc::new(Self::new( + self.cache_path + .with_extension(format!("{}.json", originator.as_str())), + self.cache_ttl, + ))) + } + fn load<'a>( &'a self, client_version: &'a str, diff --git a/codex-rs/models-manager/src/lib.rs b/codex-rs/models-manager/src/lib.rs index cd7b59a053..f3da87e511 100644 --- a/codex-rs/models-manager/src/lib.rs +++ b/codex-rs/models-manager/src/lib.rs @@ -4,10 +4,12 @@ pub(crate) mod config; pub mod manager; pub mod model_info; pub mod model_presets; +mod originator; pub mod test_support; pub use codex_protocol::auth::AuthMode; pub use config::ModelsManagerConfig; +pub use originator::ModelCatalogOriginator; /// Load the bundled model catalog shipped with `codex-models-manager`. pub fn bundled_models_response() diff --git a/codex-rs/models-manager/src/manager.rs b/codex-rs/models-manager/src/manager.rs index a78f51c698..6fba03dab8 100644 --- a/codex-rs/models-manager/src/manager.rs +++ b/codex-rs/models-manager/src/manager.rs @@ -1,4 +1,5 @@ use super::cache::FileModelsCache; +use crate::ModelCatalogOriginator; use crate::cache::ModelsCache; use crate::cache::ModelsCacheEntry; use crate::collaboration_mode_presets::builtin_collaboration_mode_presets; @@ -17,6 +18,7 @@ use codex_protocol::openai_models::ModelPreset; use codex_protocol::openai_models::ModelVisibility; use codex_protocol::openai_models::ModelsResponse; use std::borrow::Cow; +use std::collections::HashMap; use std::fmt; use std::future::Future; use std::path::PathBuf; @@ -42,6 +44,14 @@ const DEFAULT_MODEL_CACHE_TTL: Duration = Duration::from_secs(300); /// manager owns refresh policy, cache behavior, and catalog merging; it calls /// this endpoint only when it decides a remote refresh should happen. pub trait ModelsEndpointClient: fmt::Debug + Send + Sync { + /// Return an endpoint whose requests and identity use this catalog originator. + fn for_originator( + &self, + _originator: ModelCatalogOriginator, + ) -> Option> { + None + } + /// Opaque identity of the current provider and credentials, without resolving auth. /// /// Must distinguish provider routing and authenticated principals within the @@ -117,6 +127,11 @@ type SharedModelsEndpointClient = Arc; /// Coordinates model discovery plus cached metadata on disk. pub trait ModelsManager: fmt::Debug + Send + Sync { + /// Select an isolated catalog for a supported thread originator. `None` retains this manager. + fn for_originator(&self, _originator: &str) -> Option { + None + } + /// Supply startup API-key discovery policy; live changes require a new session. /// Static catalogs ignore this setting. fn set_api_key_model_discovery_enabled(&self, _enabled: bool) {} @@ -288,11 +303,12 @@ impl CatalogSource { /// OpenAI-compatible model manager backed by bundled models, cache, and `/models`. #[derive(Debug)] pub struct OpenAiModelsManager { + originator_managers: std::sync::Mutex>, remote_models: RwLock, cache: Option>, endpoint_client: SharedModelsEndpointClient, catalog_source: CatalogSource, - api_key_model_discovery_enabled: AtomicBool, + api_key_model_discovery_enabled: Arc, auth_manager: Option>, } @@ -347,6 +363,7 @@ impl OpenAiModelsManager { auth_manager: Option>, ) -> Self { Self { + originator_managers: std::sync::Mutex::new(HashMap::new()), remote_models: RwLock::new(ModelsCacheEntry { fetched_at: Utc::now(), etag: None, @@ -355,7 +372,7 @@ impl OpenAiModelsManager { models: load_remote_models_from_file().unwrap_or_default(), }), cache, - api_key_model_discovery_enabled: AtomicBool::new(false), + api_key_model_discovery_enabled: Arc::new(AtomicBool::new(false)), endpoint_client, catalog_source: CatalogSource::Default, auth_manager, @@ -382,6 +399,31 @@ impl StaticModelsManager { } impl ModelsManager for OpenAiModelsManager { + fn for_originator(&self, originator: &str) -> Option { + let originator = ModelCatalogOriginator::from_originator(originator)?; + let mut managers = self + .originator_managers + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(manager) = managers.get(&originator) { + return Some(Arc::clone(manager)); + } + let endpoint = self.endpoint_client.for_originator(originator)?; + // A backend that cannot partition must not overwrite the default catalog. + let cache = self + .cache + .as_ref() + .and_then(|cache| cache.for_originator(originator)); + let mut manager = Self::new_with_optional_cache(cache, endpoint, self.auth_manager.clone()); + if matches!(self.catalog_source, CatalogSource::ExplicitProvider(_)) { + manager = manager.with_provider_catalog(); + } + manager.api_key_model_discovery_enabled = Arc::clone(&self.api_key_model_discovery_enabled); + let manager: SharedModelsManager = Arc::new(manager); + managers.insert(originator, Arc::clone(&manager)); + Some(manager) + } + fn set_api_key_model_discovery_enabled(&self, enabled: bool) { self.api_key_model_discovery_enabled .store(enabled, Ordering::SeqCst); diff --git a/codex-rs/models-manager/src/originator.rs b/codex-rs/models-manager/src/originator.rs new file mode 100644 index 0000000000..2880fc6e7d --- /dev/null +++ b/codex-rs/models-manager/src/originator.rs @@ -0,0 +1,33 @@ +//! Bounded catalog partitions for callers migrated to CCA. Other callers keep the default catalog. + +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +pub enum ModelCatalogOriginator { + Work, + WorkWeb, + WorkDesktop, + WorkMobile, + Chat, +} + +impl ModelCatalogOriginator { + pub fn from_originator(originator: &str) -> Option { + match originator { + "codex_work_cca" => Some(Self::Work), + "codex_work_web" => Some(Self::WorkWeb), + "codex_work_desktop" => Some(Self::WorkDesktop), + "codex_work_mobile" => Some(Self::WorkMobile), + "chatgpt_cca" => Some(Self::Chat), + _ => None, + } + } + + pub fn as_str(self) -> &'static str { + match self { + Self::Work => "codex_work_cca", + Self::WorkWeb => "codex_work_web", + Self::WorkDesktop => "codex_work_desktop", + Self::WorkMobile => "codex_work_mobile", + Self::Chat => "chatgpt_cca", + } + } +}