Persist provider and auth identity with model catalog caches (#43897)

## What changed

Add a SHA-256 identity derived from provider routing, headers, and authentication scope. Return it with each model catalog response and persist it in `ModelsCacheEntry`. ChatGPT credentials with stable account and user metadata retain the same identity across token refreshes; opaque API credentials contribute to the digest.

## Testing

Add identity tests covering account, user, email, plan, auth mode, provider routing, headers, and API credential changes, plus stability across ChatGPT token refreshes. Update cache tests to include the persisted identity.

GitOrigin-RevId: 3f51c6cabcb01bc03505150a768a61dfe5d6569f
This commit is contained in:
Ahmed Ibrahim
2026-09-08 19:48:37 +00:00
committed by copyberry
parent 94e4b3d0bd
commit f31bd3adff
11 changed files with 278 additions and 18 deletions
+2
View File
@@ -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",
@@ -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<String> {
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<ModelInfo>, Option<String>)>> {
) -> ModelsEndpointFuture<'a, CoreResult<ModelsEndpointResponse>> {
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"),
})
})
}
}
@@ -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<String> {
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<ModelInfo>, Option<String>)>> {
) -> ModelsEndpointFuture<'a, CoreResult<ModelsEndpointResponse>> {
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()),
+2
View File
@@ -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 }
+1
View File
@@ -2,6 +2,7 @@ mod amazon_bedrock;
mod auth;
mod bearer_auth_provider;
mod models_endpoint;
mod models_identity;
mod provider;
mod shared_state;
+19 -5
View File
@@ -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<ModelInfo>, Option<String>)> {
) -> CoreResult<ModelsEndpointResponse> {
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<String> {
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<ModelInfo>, Option<String>)>> {
) -> ModelsEndpointFuture<'a, CoreResult<ModelsEndpointResponse>> {
Box::pin(OpenAiModelsEndpoint::list_models(
self,
client_version,
@@ -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<String> {
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;
@@ -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());
}
+3
View File
@@ -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<String>,
/// Opaque provider and auth identity associated with this catalog.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub identity: Option<String>,
/// Models returned by the catalog endpoint.
#[serde(
deserialize_with = "codex_protocol::openai_models::deserialize_model_infos_with_legacy_base"
+23 -2
View File
@@ -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<String>;
/// 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<ModelInfo>, Option<String>)>>;
) -> ModelsEndpointFuture<'a, CoreResult<ModelsEndpointResponse>>;
}
/// Catalog and validation metadata captured using the credentials for one request.
#[derive(Debug)]
pub struct ModelsEndpointResponse {
pub models: Vec<ModelInfo>,
pub etag: Option<String>,
/// 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<Box<dyn Future<Output = T> + 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 {
+30 -6
View File
@@ -214,7 +214,7 @@ impl TestModelsEndpoint {
.expect("observed proxy policy lock should not be poisoned")
}
async fn list_models(&self) -> CoreResult<(Vec<ModelInfo>, Option<String>)> {
async fn list_models(&self) -> CoreResult<ModelsEndpointResponse> {
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<String> {
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<ModelInfo>, Option<String>)>> {
) -> ModelsEndpointFuture<'a, CoreResult<ModelsEndpointResponse>> {
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<ModelInfo>, Option<String>)> {
async fn list_models(&self) -> CoreResult<ModelsEndpointResponse> {
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<String> {
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<ModelInfo>, Option<String>)>> {
) -> ModelsEndpointFuture<'a, CoreResult<ModelsEndpointResponse>> {
Box::pin(TestAuthAwareModelsEndpoint::list_models(self))
}
}