mirror of
https://github.com/openai/codex.git
synced 2026-11-13 14:46:14 +00:00
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:
1 parent
94e4b3d0bd
commit
f31bd3adff
11 files changed
+278
-18
No files matched your search
@@ -2,6 +2,7 @@ mod amazon_bedrock;
|
||||
mod auth;
|
||||
mod bearer_auth_provider;
|
||||
mod models_endpoint;
|
||||
mod models_identity;
|
||||
mod provider;
|
||||
mod shared_state;
|
||||
|
||||
|
||||
@@ -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());
|
||||
}
|
||||
Reference in new issue
Block a user