refactor: load agent identity runtime eagerly (#19763)

## Summary

AgentIdentity auth previously registered the process task lazily behind
a `OnceCell`. That meant the auth object could be constructed before its
runtime task binding was known.

This PR makes AgentIdentity auth load the runtime task at auth load time
and stores the resulting process task id directly on the auth object.
The model-provider call path can then read a concrete task id instead of
handling a missing lazy value.

## Stack

1. [refactor: make auth loading
async](https://github.com/openai/codex/pull/19762) (merged)
2. **This PR:** [refactor: load AgentIdentity runtime
eagerly](https://github.com/openai/codex/pull/19763)
3. [fix: configure AgentIdentity AuthAPI base
URL](https://github.com/openai/codex/pull/19904)
4. [feat: verify AgentIdentity JWTs with
JWKS](https://github.com/openai/codex/pull/19764)

## Important call sites

| Area | Change |
| --- | --- |
| `AgentIdentityAuth::load` | Registers the process task during auth
loading and stores `process_task_id`. |
| `CodexAuth::from_agent_identity_jwt` | Awaits AgentIdentity auth
loading. |
| model-provider auth | Reads a concrete `process_task_id` instead of an
optional lazy value. |
| AgentIdentity auth tests | Mock task registration now covers eager
runtime allocation. |

## Design decisions

AgentIdentity auth now treats task registration as part of constructing
a usable auth object. That matches how callers use the value: once auth
is present, the model-provider path expects the task-scoped assertion
data to be ready.

## Testing

Tests: targeted Rust auth test compilation, formatter, scoped Clippy
fix, and Bazel lock check.
This commit is contained in:
efrazer-oai authored and GitHub committed 2026-04-27 21:09:26 -07:00
1 parent 2307aa8d98
commit c08177f7d0
6 files changed
+43 -188

No files matched your search

+2 -2
View File
@@ -148,7 +148,7 @@ pub async fn list_cached_accessible_connectors_from_mcp_tools(
let auth = auth_manager.auth().await; let auth = auth_manager.auth().await;
if !config if !config
.features .features
.apps_enabled_for_auth(auth.as_ref().is_some_and(CodexAuth::is_chatgpt_auth)) .apps_enabled_for_auth(auth.as_ref().is_some_and(CodexAuth::uses_codex_backend))
{ {
return Some(Vec::new()); return Some(Vec::new());
} }
@@ -220,7 +220,7 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_environment_manager(
let auth = auth_manager.auth().await; let auth = auth_manager.auth().await;
if !config if !config
.features .features
.apps_enabled_for_auth(auth.as_ref().is_some_and(CodexAuth::is_chatgpt_auth)) .apps_enabled_for_auth(auth.as_ref().is_some_and(CodexAuth::uses_codex_backend))
{ {
return Ok(AccessibleConnectorsStatus { return Ok(AccessibleConnectorsStatus {
connectors: Vec::new(), connectors: Vec::new(),
+20 -40
View File
@@ -1,9 +1,6 @@
use std::sync::Arc;
use codex_agent_identity::AgentIdentityKey; use codex_agent_identity::AgentIdentityKey;
use codex_agent_identity::register_agent_task; use codex_agent_identity::register_agent_task;
use codex_protocol::account::PlanType as AccountPlanType; use codex_protocol::account::PlanType as AccountPlanType;
use tokio::sync::OnceCell;
use crate::default_client::build_reqwest_client; use crate::default_client::build_reqwest_client;
@@ -11,50 +8,33 @@ use super::storage::AgentIdentityAuthRecord;
const AGENT_IDENTITY_AUTHAPI_BASE_URL: &str = "https://auth.openai.com/api/accounts"; const AGENT_IDENTITY_AUTHAPI_BASE_URL: &str = "https://auth.openai.com/api/accounts";
#[derive(Debug)] #[derive(Clone, Debug)]
pub struct AgentIdentityAuth { pub struct AgentIdentityAuth {
record: AgentIdentityAuthRecord, record: AgentIdentityAuthRecord,
process_task_id: Arc<OnceCell<String>>, process_task_id: String,
}
impl Clone for AgentIdentityAuth {
fn clone(&self) -> Self {
Self {
record: self.record.clone(),
process_task_id: Arc::clone(&self.process_task_id),
}
}
} }
impl AgentIdentityAuth { impl AgentIdentityAuth {
pub fn new(record: AgentIdentityAuthRecord) -> Self { pub async fn load(record: AgentIdentityAuthRecord) -> std::io::Result<Self> {
Self { let process_task_id = register_agent_task(
&build_reqwest_client(),
AGENT_IDENTITY_AUTHAPI_BASE_URL,
key(&record),
)
.await
.map_err(std::io::Error::other)?;
Ok(Self {
record, record,
process_task_id: Arc::new(OnceCell::new()), process_task_id,
} })
} }
pub fn record(&self) -> &AgentIdentityAuthRecord { pub fn record(&self) -> &AgentIdentityAuthRecord {
&self.record &self.record
} }
pub fn process_task_id(&self) -> Option<&str> { pub fn process_task_id(&self) -> &str {
self.process_task_id.get().map(String::as_str) &self.process_task_id
}
pub async fn ensure_runtime(&self) -> std::io::Result<()> {
self.process_task_id
.get_or_try_init(|| async {
register_agent_task(
&build_reqwest_client(),
AGENT_IDENTITY_AUTHAPI_BASE_URL,
self.key(),
)
.await
.map_err(std::io::Error::other)
})
.await
.map(|_| ())
} }
pub fn account_id(&self) -> &str { pub fn account_id(&self) -> &str {
@@ -76,11 +56,11 @@ impl AgentIdentityAuth {
pub fn is_fedramp_account(&self) -> bool { pub fn is_fedramp_account(&self) -> bool {
self.record.chatgpt_account_is_fedramp self.record.chatgpt_account_is_fedramp
} }
}
fn key(&self) -> AgentIdentityKey<'_> { fn key(record: &AgentIdentityAuthRecord) -> AgentIdentityKey<'_> {
AgentIdentityKey { AgentIdentityKey {
agent_runtime_id: &self.record.agent_runtime_id, agent_runtime_id: &record.agent_runtime_id,
private_key_pkcs8_base64: &self.record.agent_private_key, private_key_pkcs8_base64: &record.agent_private_key,
}
} }
} }
+3 -28
View File
@@ -625,33 +625,6 @@ impl Drop for EnvVarGuard {
} }
} }
#[tokio::test]
#[serial(codex_auth_env)]
async fn load_auth_reads_agent_identity_from_env() {
let codex_home = tempdir().unwrap();
let expected_record = agent_identity_record("account-123");
let agent_identity = fake_agent_identity_jwt(&expected_record).expect("fake agent identity");
let _agent_guard = EnvVarGuard::set(CODEX_AGENT_IDENTITY_ENV_VAR, &agent_identity);
let auth = super::load_auth(
codex_home.path(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
)
.await
.expect("env auth should load")
.expect("env auth should be present");
let CodexAuth::AgentIdentity(agent_identity) = auth else {
panic!("env auth should load as agent identity");
};
assert_eq!(agent_identity.record(), &expected_record);
assert!(
!get_auth_file(codex_home.path()).exists(),
"env auth should not write auth.json"
);
}
#[tokio::test] #[tokio::test]
#[serial(codex_auth_env)] #[serial(codex_auth_env)]
async fn load_auth_keeps_codex_api_key_env_precedence() { async fn load_auth_keeps_codex_api_key_env_precedence() {
@@ -805,9 +778,11 @@ async fn enforce_login_restrictions_blocks_env_api_key_when_chatgpt_required() {
} }
fn agent_identity_record(account_id: &str) -> AgentIdentityAuthRecord { fn agent_identity_record(account_id: &str) -> AgentIdentityAuthRecord {
let key_material =
codex_agent_identity::generate_agent_key_material().expect("generate agent key material");
AgentIdentityAuthRecord { AgentIdentityAuthRecord {
agent_runtime_id: "agent-runtime-id".to_string(), agent_runtime_id: "agent-runtime-id".to_string(),
agent_private_key: "private-key".to_string(), agent_private_key: key_material.private_key_pkcs8_base64,
account_id: account_id.to_string(), account_id: account_id.to_string(),
chatgpt_user_id: "user-id".to_string(), chatgpt_user_id: "user-id".to_string(),
email: "user@example.com".to_string(), email: "user@example.com".to_string(),
+7 -20
View File
@@ -212,7 +212,7 @@ impl CodexAuth {
"agent identity auth is missing an agent identity token.", "agent identity auth is missing an agent identity token.",
)); ));
}; };
return Self::from_agent_identity_jwt(&agent_identity); return Self::from_agent_identity_jwt(&agent_identity).await;
} }
let storage_mode = auth_dot_json.storage_mode(auth_credentials_store_mode); let storage_mode = auth_dot_json.storage_mode(auth_credentials_store_mode);
@@ -246,9 +246,9 @@ impl CodexAuth {
.await .await
} }
pub fn from_agent_identity_jwt(jwt: &str) -> std::io::Result<Self> { pub async fn from_agent_identity_jwt(jwt: &str) -> std::io::Result<Self> {
let record = AgentIdentityAuthRecord::from_agent_identity_jwt(jwt)?; let record = AgentIdentityAuthRecord::from_agent_identity_jwt(jwt)?;
Ok(Self::AgentIdentity(AgentIdentityAuth::new(record))) Ok(Self::AgentIdentity(AgentIdentityAuth::load(record).await?))
} }
pub fn auth_mode(&self) -> AuthMode { pub fn auth_mode(&self) -> AuthMode {
@@ -322,16 +322,6 @@ impl CodexAuth {
} }
} }
pub async fn initialize_runtime(
&self,
_chatgpt_base_url: Option<String>,
) -> std::io::Result<()> {
match self {
Self::AgentIdentity(auth) => auth.ensure_runtime().await,
Self::ApiKey(_) | Self::Chatgpt(_) | Self::ChatgptAuthTokens(_) => Ok(()),
}
}
/// Returns `None` if Codex backend auth does not expose an account id. /// Returns `None` if Codex backend auth does not expose an account id.
pub fn get_account_id(&self) -> Option<String> { pub fn get_account_id(&self) -> Option<String> {
match self { match self {
@@ -749,7 +739,9 @@ async fn load_auth(
} }
if let Some(agent_identity) = read_codex_agent_identity_from_env() { if let Some(agent_identity) = read_codex_agent_identity_from_env() {
return CodexAuth::from_agent_identity_jwt(&agent_identity).map(Some); return CodexAuth::from_agent_identity_jwt(&agent_identity)
.await
.map(Some);
} }
// Fall back to the configured persistent store (file/keyring/auto) for managed auth. // Fall back to the configured persistent store (file/keyring/auto) for managed auth.
@@ -1400,12 +1392,7 @@ impl AuthManager {
tracing::error!("Failed to refresh token: {}", err); tracing::error!("Failed to refresh token: {}", err);
return Some(auth); return Some(auth);
} }
let auth = self.auth_cached()?; self.auth_cached()
if let Err(err) = auth.initialize_runtime(self.chatgpt_base_url.clone()).await {
tracing::error!("Failed to initialize auth runtime: {err}");
return None;
}
Some(auth)
} }
/// Force a reload of the auth information from auth.json. Returns /// Force a reload of the auth information from auth.json. Returns
+11 -17
View File
@@ -21,23 +21,17 @@ struct AgentIdentityAuthProvider {
impl AuthProvider for AgentIdentityAuthProvider { impl AuthProvider for AgentIdentityAuthProvider {
fn add_auth_headers(&self, headers: &mut HeaderMap) { fn add_auth_headers(&self, headers: &mut HeaderMap) {
let record = self.auth.record(); let record = self.auth.record();
let header_value = self let header_value = authorization_header_for_agent_task(
.auth AgentIdentityKey {
.process_task_id() agent_runtime_id: &record.agent_runtime_id,
.ok_or_else(|| std::io::Error::other("agent identity process task is not initialized")) private_key_pkcs8_base64: &record.agent_private_key,
.and_then(|task_id| { },
authorization_header_for_agent_task( AgentTaskAuthorizationTarget {
AgentIdentityKey { agent_runtime_id: &record.agent_runtime_id,
agent_runtime_id: &record.agent_runtime_id, task_id: self.auth.process_task_id(),
private_key_pkcs8_base64: &record.agent_private_key, },
}, )
AgentTaskAuthorizationTarget { .map_err(std::io::Error::other);
agent_runtime_id: &record.agent_runtime_id,
task_id,
},
)
.map_err(std::io::Error::other)
});
if let Ok(header_value) = header_value if let Ok(header_value) = header_value
&& let Ok(header) = HeaderValue::from_str(&header_value) && let Ok(header) = HeaderValue::from_str(&header_value)
@@ -9,9 +9,6 @@ use codex_login::ExternalAuth;
use codex_login::ExternalAuthRefreshContext; use codex_login::ExternalAuthRefreshContext;
use codex_login::ExternalAuthTokens; use codex_login::ExternalAuthTokens;
use codex_login::TokenData; use codex_login::TokenData;
use codex_login::auth::AgentIdentityAuth;
use codex_login::auth::AgentIdentityAuthRecord;
use codex_protocol::account::PlanType;
use codex_protocol::openai_models::ModelsResponse; use codex_protocol::openai_models::ModelsResponse;
use pretty_assertions::assert_eq; use pretty_assertions::assert_eq;
use serde_json::json; use serde_json::json;
@@ -237,18 +234,6 @@ c2ln",
.expect("auth should be present") .expect("auth should be present")
} }
fn agent_identity_auth_for_tests() -> CodexAuth {
CodexAuth::AgentIdentity(AgentIdentityAuth::new(AgentIdentityAuthRecord {
agent_runtime_id: "agent-runtime-id".to_string(),
agent_private_key: "agent-private-key".to_string(),
account_id: "account-id".to_string(),
chatgpt_user_id: "chatgpt-user-id".to_string(),
email: "agent@example.com".to_string(),
plan_type: PlanType::Pro,
chatgpt_account_is_fedramp: false,
}))
}
#[tokio::test] #[tokio::test]
async fn get_model_info_tracks_fallback_usage() { async fn get_model_info_tracks_fallback_usage() {
let codex_home = tempdir().expect("temp dir"); let codex_home = tempdir().expect("temp dir");
@@ -713,43 +698,6 @@ async fn refresh_available_models_fetches_with_chatgpt_auth_tokens() {
); );
} }
#[tokio::test]
async fn refresh_available_models_fetches_with_agent_identity() {
let dynamic_slug = "dynamic-model-only-for-test-agent-identity";
let codex_home = tempdir().expect("temp dir");
let endpoint = TestModelsEndpoint::new(vec![vec![remote_model(
dynamic_slug,
"Agent Identity",
/*priority*/ 1,
)]]);
let manager = openai_manager_for_tests_with_auth(
codex_home.path().to_path_buf(),
endpoint.clone(),
Some(AuthManager::from_auth_for_testing(
agent_identity_auth_for_tests(),
)),
);
manager
.refresh_available_models(RefreshStrategy::Online)
.await
.expect("refresh should fetch with agent identity");
assert!(
manager
.get_remote_models()
.await
.iter()
.any(|candidate| candidate.slug == dynamic_slug),
"remote refresh should include models fetched with agent identity"
);
assert_eq!(
endpoint.fetch_count(),
1,
"endpoint should fetch models with agent identity"
);
}
#[test] #[test]
fn build_available_models_picks_default_after_hiding_hidden_models() { fn build_available_models_picks_default_after_hiding_hidden_models() {
let manager = static_manager_for_tests(ModelsResponse { models: Vec::new() }); let manager = static_manager_for_tests(ModelsResponse { models: Vec::new() });
@@ -768,35 +716,6 @@ fn build_available_models_picks_default_after_hiding_hidden_models() {
assert_eq!(available, vec![expected_hidden, expected_visible]); assert_eq!(available, vec![expected_hidden, expected_visible]);
} }
#[tokio::test]
async fn static_manager_treats_agent_identity_as_backend_auth_for_filtering() {
let chatgpt_only_model = {
let mut model = remote_model("chatgpt-only", "ChatGPT Only", /*priority*/ 0);
model.supported_in_api = false;
model
};
let api_model = remote_model("api-model", "API Model", /*priority*/ 1);
let manager = StaticModelsManager::new(
Some(AuthManager::from_auth_for_testing(
agent_identity_auth_for_tests(),
)),
ModelsResponse {
models: vec![chatgpt_only_model, api_model],
},
CollaborationModesConfig::default(),
);
let agent_identity_models = manager.list_models(RefreshStrategy::Online).await;
assert_eq!(
agent_identity_models
.iter()
.map(|model| model.model.as_str())
.collect::<Vec<_>>(),
vec!["chatgpt-only", "api-model"]
);
}
#[tokio::test] #[tokio::test]
async fn static_manager_reads_latest_auth_mode() { async fn static_manager_reads_latest_auth_mode() {
let auth_manager = let auth_manager =