diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index ea48d2a4fd..20a55a69ac 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3972,6 +3972,7 @@ dependencies = [ name = "codex-model-provider" version = "0.0.0" dependencies = [ + "base64 0.22.1", "codex-agent-identity", "codex-api", "codex-aws-auth", @@ -3987,6 +3988,7 @@ dependencies = [ "http 1.4.0", "pretty_assertions", "serde_json", + "sha2 0.10.9", "tokio", "tracing", "wiremock", 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..0a7e132407 100644 --- a/codex-rs/app-server/src/models_refresh_worker_tests.rs +++ b/codex-rs/app-server/src/models_refresh_worker_tests.rs @@ -7,11 +7,11 @@ use codex_http_client::HttpClientFactory; use codex_http_client::OutboundProxyPolicy; use codex_models_manager::manager::ModelsEndpointClient; use codex_models_manager::manager::ModelsEndpointFuture; +use codex_models_manager::manager::ModelsEndpointResponse; use codex_models_manager::manager::OpenAiModelsManager; use codex_models_manager::manager::SharedModelsManager; use codex_protocol::error::CodexErr; use codex_protocol::error::Result as CoreResult; -use codex_protocol::openai_models::ModelInfo; use pretty_assertions::assert_eq; use tempfile::tempdir; use tokio::sync::Notify; @@ -46,6 +46,10 @@ impl TestModelsEndpoint { } impl ModelsEndpointClient for TestModelsEndpoint { + fn identity(&self) -> Option { + Some("test-provider".to_string()) + } + fn has_command_auth(&self) -> bool { true } @@ -58,7 +62,7 @@ impl ModelsEndpointClient for TestModelsEndpoint { &'a self, _client_version: &'a str, _http_client_factory: HttpClientFactory, - ) -> ModelsEndpointFuture<'a, CoreResult<(Vec, Option)>> { + ) -> ModelsEndpointFuture<'a, CoreResult> { Box::pin(async move { let fetch_index = self.fetch_count.fetch_add(1, Ordering::SeqCst); self.fetched.notify_one(); @@ -68,7 +72,11 @@ impl ModelsEndpointClient for TestModelsEndpoint { if fetch_index == 1 { self.release_second_fetch.notified().await; } - Ok((Vec::new(), None)) + Ok(ModelsEndpointResponse { + models: Vec::new(), + etag: None, + identity: self.identity().expect("test endpoint identity"), + }) }) } } diff --git a/codex-rs/core/tests/suite/injected_models_cache.rs b/codex-rs/core/tests/suite/injected_models_cache.rs index e909f88149..3d0d114c68 100644 --- a/codex-rs/core/tests/suite/injected_models_cache.rs +++ b/codex-rs/core/tests/suite/injected_models_cache.rs @@ -15,6 +15,7 @@ use codex_models_manager::cache::ModelsCacheError; use codex_models_manager::cache::ModelsCacheFuture; use codex_models_manager::manager::ModelsEndpointClient; use codex_models_manager::manager::ModelsEndpointFuture; +use codex_models_manager::manager::ModelsEndpointResponse; use codex_models_manager::manager::OpenAiModelsManager; use codex_models_manager::manager::RefreshStrategy; use codex_models_manager::manager::SharedModelsManager; @@ -105,6 +106,10 @@ impl TestModelsEndpoint { } impl ModelsEndpointClient for TestModelsEndpoint { + fn identity(&self) -> Option { + Some("test-provider".to_string()) + } + fn has_command_auth(&self) -> bool { false } @@ -117,10 +122,14 @@ impl ModelsEndpointClient for TestModelsEndpoint { &'a self, _client_version: &'a str, _http_client_factory: HttpClientFactory, - ) -> ModelsEndpointFuture<'a, CoreResult<(Vec, Option)>> { + ) -> ModelsEndpointFuture<'a, CoreResult> { Box::pin(async move { self.fetch_count.fetch_add(1, Ordering::SeqCst); - Ok((self.models.clone(), None)) + Ok(ModelsEndpointResponse { + models: self.models.clone(), + etag: None, + identity: self.identity().expect("test endpoint identity"), + }) }) } } @@ -210,6 +219,7 @@ async fn injected_cache_hit_drives_agent_model_selection() -> Result<()> { let model_slug = "injected-cache-model"; let cache = Arc::new(TestModelsCache { entry: Some(ModelsCacheEntry { + identity: Some("test-provider".to_string()), fetched_at: Utc::now(), etag: None, client_version: Some(codex_models_manager::client_version_to_whole()), diff --git a/codex-rs/model-provider/Cargo.toml b/codex-rs/model-provider/Cargo.toml index e72a55a3a7..c10bb3f0bb 100644 --- a/codex-rs/model-provider/Cargo.toml +++ b/codex-rs/model-provider/Cargo.toml @@ -25,10 +25,12 @@ codex-otel = { workspace = true } codex-protocol = { workspace = true } codex-response-debug-context = { workspace = true } http = { workspace = true } +sha2 = { workspace = true } tokio = { workspace = true, features = ["process", "sync", "time"] } tracing = { workspace = true, features = ["log"] } [dev-dependencies] +base64 = { workspace = true } codex-utils-redacted-string = { workspace = true } pretty_assertions = { workspace = true } serde_json = { workspace = true } diff --git a/codex-rs/model-provider/src/lib.rs b/codex-rs/model-provider/src/lib.rs index 98ed5f5e7b..f248f963fe 100644 --- a/codex-rs/model-provider/src/lib.rs +++ b/codex-rs/model-provider/src/lib.rs @@ -2,6 +2,7 @@ mod amazon_bedrock; mod auth; mod bearer_auth_provider; mod models_endpoint; +mod models_identity; mod provider; mod shared_state; diff --git a/codex-rs/model-provider/src/models_endpoint.rs b/codex-rs/model-provider/src/models_endpoint.rs index 4ccb5b7823..76bf039f2a 100644 --- a/codex-rs/model-provider/src/models_endpoint.rs +++ b/codex-rs/model-provider/src/models_endpoint.rs @@ -23,10 +23,10 @@ use codex_login::default_client::create_client_for_route_async; use codex_model_provider_info::ModelProviderInfo; use codex_models_manager::manager::ModelsEndpointClient; use codex_models_manager::manager::ModelsEndpointFuture; +use codex_models_manager::manager::ModelsEndpointResponse; use codex_otel::TelemetryAuthMode; use codex_protocol::error::CodexErr; use codex_protocol::error::Result as CoreResult; -use codex_protocol::openai_models::ModelInfo; use codex_response_debug_context::extract_response_debug_context; use codex_response_debug_context::telemetry_transport_error_message; use http::HeaderMap; @@ -77,10 +77,11 @@ impl OpenAiModelsEndpoint { &self, client_version: &str, http_client_factory: HttpClientFactory, - ) -> CoreResult<(Vec, Option)> { + ) -> CoreResult { let _timer = codex_otel::start_global_timer("codex.remote_models.fetch_update.duration_ms", &[]); let auth = self.auth().await; + let identity = crate::models_identity::identity(&self.provider_info, auth.as_ref())?; let auth_mode = auth.as_ref().map(CodexAuth::auth_mode); let mut api_provider = self.provider_info.to_api_provider(auth_mode)?; enforce_managed_residency(&mut api_provider); @@ -100,7 +101,7 @@ impl OpenAiModelsEndpoint { agent_identity_telemetry, auth_env: self.auth_env(), }); - timeout(MODELS_REFRESH_TIMEOUT, async { + let (models, etag) = timeout(MODELS_REFRESH_TIMEOUT, async { let transport = self .transport_builder .build(http_client_factory, request_url.clone()) @@ -113,7 +114,12 @@ impl OpenAiModelsEndpoint { .map_err(map_api_error) }) .await - .map_err(|_| CodexErr::Timeout)? + .map_err(|_| CodexErr::Timeout)??; + Ok(ModelsEndpointResponse { + models, + etag, + identity, + }) } fn auth_env(&self) -> AuthEnvTelemetry { @@ -126,6 +132,14 @@ impl OpenAiModelsEndpoint { } impl ModelsEndpointClient for OpenAiModelsEndpoint { + fn identity(&self) -> Option { + let auth = self + .auth_manager + .as_ref() + .and_then(|manager| manager.auth_cached()); + crate::models_identity::identity(&self.provider_info, auth.as_ref()).ok() + } + fn has_command_auth(&self) -> bool { self.provider_info.has_command_auth() } @@ -138,7 +152,7 @@ impl ModelsEndpointClient for OpenAiModelsEndpoint { &'a self, client_version: &'a str, http_client_factory: HttpClientFactory, - ) -> ModelsEndpointFuture<'a, CoreResult<(Vec, Option)>> { + ) -> ModelsEndpointFuture<'a, CoreResult> { Box::pin(OpenAiModelsEndpoint::list_models( self, client_version, diff --git a/codex-rs/model-provider/src/models_identity.rs b/codex-rs/model-provider/src/models_identity.rs new file mode 100644 index 0000000000..2baa42f2ce --- /dev/null +++ b/codex-rs/model-provider/src/models_identity.rs @@ -0,0 +1,86 @@ +//! Model catalog cache identity follows provider routing and effective auth. +//! Only a digest is persisted; access tokens for ChatGPT are excluded so token +//! rotation does not discard a catalog for the same account, user, and plan. + +use codex_login::CodexAuth; +use codex_model_provider_info::ModelProviderInfo; +use codex_protocol::error::Result as CoreResult; +use sha2::Digest; +use sha2::Sha256; + +use crate::auth::resolve_provider_auth; +use crate::provider::enforce_managed_residency; + +pub(crate) fn identity( + provider_info: &ModelProviderInfo, + auth: Option<&CodexAuth>, +) -> CoreResult { + let mut provider = provider_info.to_api_provider(auth.map(CodexAuth::auth_mode))?; + enforce_managed_residency(&mut provider); + let mut digest = Sha256::new(); + // Length-prefix every field to avoid ambiguity between adjacent values. + let mut field = |value: &[u8]| { + digest.update((value.len() as u64).to_le_bytes()); + digest.update(value); + }; + field(b"models-cache-v1"); + field(provider.name.as_bytes()); + field(provider.base_url.as_bytes()); + let mut query: Vec<_> = provider.query_params.iter().flatten().collect(); + query.sort(); + field(&(query.len() as u64).to_le_bytes()); + for (name, value) in query { + field(name.as_bytes()); + field(value.as_bytes()); + } + field(&[u8::from(provider_info.requires_openai_auth)]); + field(&[u8::from(provider_info.has_command_auth())]); + field(format!("{:?}", auth.map(CodexAuth::api_auth_mode)).as_bytes()); + if let Some(auth) = auth { + field(format!("{:?}", auth.get_account_id()).as_bytes()); + field(format!("{:?}", auth.get_chatgpt_user_id()).as_bytes()); + field(format!("{:?}", auth.get_account_email()).as_bytes()); + field(format!("{:?}", auth.account_plan_type()).as_bytes()); + field( + format!( + "{:?}", + auth.get_token_data() + .ok() + .and_then(|tokens| tokens.id_token.get_chatgpt_plan_type_raw()) + ) + .as_bytes(), + ); + field(&[u8::from(auth.is_fedramp_account())]); + } + let has_stable_account = auth.is_some_and(|auth| { + matches!( + auth, + CodexAuth::Chatgpt(_) | CodexAuth::ChatgptAuthTokens(_) | CodexAuth::AgentIdentity(_) + ) && auth.get_account_id().is_some() + && (auth.get_chatgpt_user_id().is_some() || auth.get_account_email().is_some()) + }); + let explicit_bearer = + provider_info.api_key()?.is_some() || provider_info.experimental_bearer_token.is_some(); + // Command auth also returns opaque API keys. A changed key may belong to a + // different account, so only credentials with stable owner metadata can be reused. + if !has_stable_account || explicit_bearer { + let headers = resolve_provider_auth(auth, provider_info)?.to_auth_headers(); + provider.headers.extend(headers); + } + let mut headers: Vec<_> = provider.headers.iter().collect(); + headers.sort_by(|(a, av), (b, bv)| { + a.as_str() + .cmp(b.as_str()) + .then(av.as_bytes().cmp(bv.as_bytes())) + }); + field(&(headers.len() as u64).to_le_bytes()); + for (name, value) in headers { + field(name.as_str().as_bytes()); + field(value.as_bytes()); + } + Ok(format!("{:x}", digest.finalize())) +} + +#[cfg(test)] +#[path = "models_identity_tests.rs"] +mod tests; diff --git a/codex-rs/model-provider/src/models_identity_tests.rs b/codex-rs/model-provider/src/models_identity_tests.rs new file mode 100644 index 0000000000..4732265de1 --- /dev/null +++ b/codex-rs/model-provider/src/models_identity_tests.rs @@ -0,0 +1,89 @@ +//! Provider and authentication dimensions that partition model catalogs. + +use super::*; +use base64::Engine; +use base64::engine::general_purpose::URL_SAFE_NO_PAD; +use pretty_assertions::assert_eq; + +fn chatgpt_auth(email: &str, user: &str, account: &str, plan: &str, signature: &str) -> CodexAuth { + let claims = serde_json::json!({ + "email": email, + "https://api.openai.com/auth": {"chatgpt_user_id": user} + }); + let payload = URL_SAFE_NO_PAD.encode(serde_json::to_vec(&claims).unwrap()); + CodexAuth::from_external_chatgpt_tokens( + &format!("e30.{payload}.{signature}"), + account, + Some(plan), + ) + .unwrap() +} + +#[test] +fn cache_identity_tracks_account_email_user_plan_and_auth_mode_but_not_token_refresh() { + let mut provider = ModelProviderInfo::create_openai_provider(/*base_url*/ None); + // The endpoint residency test temporarily sets the process default to US. + provider.http_headers = Some(std::collections::HashMap::from([( + codex_login::default_client::RESIDENCY_HEADER_NAME.to_string(), + "us".into(), + )])); + let initial = chatgpt_auth("one@example.com", "user", "account", "team", "first"); + let key = identity(&provider, Some(&initial)).unwrap(); + let refreshed = chatgpt_auth("one@example.com", "user", "account", "team", "second"); + assert_eq!(identity(&provider, Some(&refreshed)).unwrap(), key); + for auth in [ + chatgpt_auth("two@example.com", "user", "account", "team", "first"), + chatgpt_auth("one@example.com", "other", "account", "team", "first"), + chatgpt_auth("one@example.com", "user", "other", "team", "first"), + chatgpt_auth("one@example.com", "user", "account", "plus", "first"), + CodexAuth::from_api_key("api-key"), + ] { + assert_ne!(identity(&provider, Some(&auth)).unwrap(), key); + } + assert_ne!(identity(&provider, /*auth*/ None).unwrap(), key); +} + +#[test] +fn cache_identity_tracks_provider_routing_and_effective_api_credentials() { + let auth = CodexAuth::from_api_key("first-key"); + let provider = ModelProviderInfo::create_openai_provider(Some("https://one.example/v1".into())); + let key = identity(&provider, Some(&auth)).unwrap(); + let other_key = CodexAuth::from_api_key("second-key"); + assert_ne!(identity(&provider, Some(&other_key)).unwrap(), key); + let mut other_provider = provider.clone(); + other_provider.base_url = Some("https://two.example/v1".into()); + assert_ne!(identity(&other_provider, Some(&auth)).unwrap(), key); + other_provider = provider.clone(); + other_provider.experimental_bearer_token = Some("provider-key".into()); + assert_ne!(identity(&other_provider, Some(&auth)).unwrap(), key); + other_provider = provider; + other_provider.http_headers = Some(std::collections::HashMap::from([( + "openai-project".into(), + "another-project".into(), + )])); + assert_ne!(identity(&other_provider, Some(&auth)).unwrap(), key); +} + +#[test] +fn cache_identity_distinguishes_auth_requirements_and_unknown_plans() { + let mut provider = + ModelProviderInfo::create_openai_provider(Some("https://example.com/v1".into())); + let first = chatgpt_auth( + "one@example.com", + "user", + "account", + "future-plan-one", + "first", + ); + let second = chatgpt_auth( + "one@example.com", + "user", + "account", + "future-plan-two", + "first", + ); + let key = identity(&provider, Some(&first)).unwrap(); + assert_ne!(key, identity(&provider, Some(&second)).unwrap()); + provider.requires_openai_auth = false; + assert_ne!(key, identity(&provider, Some(&first)).unwrap()); +} diff --git a/codex-rs/models-manager/src/cache.rs b/codex-rs/models-manager/src/cache.rs index 15e8cf4dd8..240ec7e11b 100644 --- a/codex-rs/models-manager/src/cache.rs +++ b/codex-rs/models-manager/src/cache.rs @@ -71,6 +71,9 @@ pub struct ModelsCacheEntry { /// The models manager rejects entries whose value is absent or differs from its current version. #[serde(default, skip_serializing_if = "Option::is_none")] pub client_version: Option, + /// Opaque provider and auth identity associated with this catalog. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub identity: Option, /// Models returned by the catalog endpoint. #[serde( deserialize_with = "codex_protocol::openai_models::deserialize_model_infos_with_legacy_base" diff --git a/codex-rs/models-manager/src/manager.rs b/codex-rs/models-manager/src/manager.rs index 45835c37cd..9b37abe10e 100644 --- a/codex-rs/models-manager/src/manager.rs +++ b/codex-rs/models-manager/src/manager.rs @@ -35,6 +35,12 @@ 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 { + /// Opaque identity of the current provider and credentials, without resolving auth. + /// + /// Must change when account, email, plan, provider, or API credentials change. + /// Return `None` when identity is unavailable. + fn identity(&self) -> Option; + /// Returns whether this provider can authenticate command-scoped requests. fn has_command_auth(&self) -> bool; @@ -46,7 +52,17 @@ pub trait ModelsEndpointClient: fmt::Debug + Send + Sync { &'a self, client_version: &'a str, http_client_factory: HttpClientFactory, - ) -> ModelsEndpointFuture<'a, CoreResult<(Vec, Option)>>; + ) -> ModelsEndpointFuture<'a, CoreResult>; +} + +/// Catalog and validation metadata captured using the credentials for one request. +#[derive(Debug)] +pub struct ModelsEndpointResponse { + pub models: Vec, + pub etag: Option, + /// Identity of the credentials actually used to fetch this catalog. + /// Successful responses must always identify their provider and authentication scope. + pub identity: String, } pub type ModelsEndpointFuture<'a, T> = Pin + Send + 'a>>; @@ -414,7 +430,11 @@ impl OpenAiModelsManager { http_client_factory: &HttpClientFactory, ) -> CoreResult<()> { let client_version = crate::client_version_to_whole(); - let (models, etag) = self + let ModelsEndpointResponse { + models, + etag, + identity, + } = self .endpoint_client .list_models(&client_version, http_client_factory.clone()) .await?; @@ -425,6 +445,7 @@ impl OpenAiModelsManager { fetched_at: Utc::now(), etag, client_version: Some(client_version), + identity: Some(identity), models, }; if let Err(err) = cache.store(&entry).await { diff --git a/codex-rs/models-manager/src/manager_tests.rs b/codex-rs/models-manager/src/manager_tests.rs index 324ca42ceb..78ecf95336 100644 --- a/codex-rs/models-manager/src/manager_tests.rs +++ b/codex-rs/models-manager/src/manager_tests.rs @@ -214,7 +214,7 @@ impl TestModelsEndpoint { .expect("observed proxy policy lock should not be poisoned") } - async fn list_models(&self) -> CoreResult<(Vec, Option)> { + async fn list_models(&self) -> CoreResult { self.fetch_count.fetch_add(1, Ordering::SeqCst); let models = self .responses @@ -222,7 +222,11 @@ impl TestModelsEndpoint { .expect("responses lock should not be poisoned") .pop_front() .unwrap_or_default(); - Ok((models, None)) + Ok(ModelsEndpointResponse { + models, + etag: None, + identity: self.identity().expect("test endpoint identity"), + }) } } @@ -259,6 +263,10 @@ impl ExternalAuth for TestUnresolvedExternalApiKeyAuth { } impl ModelsEndpointClient for TestModelsEndpoint { + fn identity(&self) -> Option { + Some("test-provider".to_string()) + } + fn has_command_auth(&self) -> bool { self.has_command_auth } @@ -271,7 +279,7 @@ impl ModelsEndpointClient for TestModelsEndpoint { &'a self, _client_version: &'a str, http_client_factory: HttpClientFactory, - ) -> ModelsEndpointFuture<'a, CoreResult<(Vec, Option)>> { + ) -> ModelsEndpointFuture<'a, CoreResult> { Box::pin(async move { *self .observed_proxy_policy @@ -332,6 +340,7 @@ async fn file_cache_implements_models_cache_contract() { ); let client_version = crate::client_version_to_whole(); let entry = ModelsCacheEntry { + identity: Some("test-provider".to_string()), fetched_at: Utc::now(), etag: Some("file-etag".to_string()), client_version: Some(client_version.clone()), @@ -363,6 +372,7 @@ async fn file_cache_refresh_ttl_renews_expired_entry_without_serving_it_stale() let client_version = crate::client_version_to_whole(); let expired_at = Utc::now() - chrono::Duration::hours(1); let entry = ModelsCacheEntry { + identity: Some("test-provider".to_string()), fetched_at: expired_at, etag: Some("expired-etag".to_string()), client_version: Some(client_version.clone()), @@ -433,6 +443,7 @@ async fn manager_without_cache_fetches_on_every_refresh() { async fn injected_cache_hit_avoids_remote_fetch() { let cached_models = vec![remote_model("cached", "Cached", /*priority*/ 0)]; let cache = TestModelsCache::with_entry(ModelsCacheEntry { + identity: Some("test-provider".to_string()), fetched_at: Utc::now(), etag: Some("cached-etag".to_string()), client_version: Some(crate::client_version_to_whole()), @@ -486,6 +497,7 @@ async fn injected_cache_read_error_falls_back_and_persists_remote_models() { assert_eq!( stored_entries, vec![ModelsCacheEntry { + identity: Some("test-provider".to_string()), fetched_at: stored_entries[0].fetched_at, etag: None, client_version: Some(crate::client_version_to_whole()), @@ -523,6 +535,7 @@ async fn injected_cache_ttl_refresh_preserves_cached_payload() { let cached_models = vec![remote_model("cached", "Cached", /*priority*/ 0)]; let cached_at = Utc::now() - chrono::Duration::minutes(1); let cache = TestModelsCache::with_entry(ModelsCacheEntry { + identity: Some("test-provider".to_string()), fetched_at: cached_at, etag: Some("cached-etag".to_string()), client_version: Some(crate::client_version_to_whole()), @@ -1268,7 +1281,7 @@ impl TestAuthAwareModelsEndpoint { } } - async fn list_models(&self) -> CoreResult<(Vec, Option)> { + async fn list_models(&self) -> CoreResult { self.fetch_count.fetch_add(1, Ordering::SeqCst); let models = self .responses @@ -1276,11 +1289,22 @@ impl TestAuthAwareModelsEndpoint { .expect("responses lock should not be poisoned") .pop_front() .unwrap_or_default(); - Ok((models, None)) + Ok(ModelsEndpointResponse { + models, + etag: None, + identity: self.identity().expect("test endpoint identity"), + }) } } impl ModelsEndpointClient for TestAuthAwareModelsEndpoint { + fn identity(&self) -> Option { + self.auth_manager + .as_ref() + .and_then(|manager| manager.auth_mode()) + .map(|mode| format!("{mode:?}")) + } + fn has_command_auth(&self) -> bool { false } @@ -1293,7 +1317,7 @@ impl ModelsEndpointClient for TestAuthAwareModelsEndpoint { &'a self, _client_version: &'a str, _http_client_factory: HttpClientFactory, - ) -> ModelsEndpointFuture<'a, CoreResult<(Vec, Option)>> { + ) -> ModelsEndpointFuture<'a, CoreResult> { Box::pin(TestAuthAwareModelsEndpoint::list_models(self)) } }