Skip to content

Commit c08177f

Browse files
authored
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](#19762) (merged) 2. **This PR:** [refactor: load AgentIdentity runtime eagerly](#19763) 3. [fix: configure AgentIdentity AuthAPI base URL](#19904) 4. [feat: verify AgentIdentity JWTs with JWKS](#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.
1 parent 2307aa8 commit c08177f

6 files changed

Lines changed: 43 additions & 188 deletions

File tree

codex-rs/core/src/connectors.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -148,7 +148,7 @@ pub async fn list_cached_accessible_connectors_from_mcp_tools(
148148
let auth = auth_manager.auth().await;
149149
if !config
150150
.features
151-
.apps_enabled_for_auth(auth.as_ref().is_some_and(CodexAuth::is_chatgpt_auth))
151+
.apps_enabled_for_auth(auth.as_ref().is_some_and(CodexAuth::uses_codex_backend))
152152
{
153153
return Some(Vec::new());
154154
}
@@ -220,7 +220,7 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_environment_manager(
220220
let auth = auth_manager.auth().await;
221221
if !config
222222
.features
223-
.apps_enabled_for_auth(auth.as_ref().is_some_and(CodexAuth::is_chatgpt_auth))
223+
.apps_enabled_for_auth(auth.as_ref().is_some_and(CodexAuth::uses_codex_backend))
224224
{
225225
return Ok(AccessibleConnectorsStatus {
226226
connectors: Vec::new(),
Lines changed: 20 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -1,60 +1,40 @@
1-
use std::sync::Arc;
2-
31
use codex_agent_identity::AgentIdentityKey;
42
use codex_agent_identity::register_agent_task;
53
use codex_protocol::account::PlanType as AccountPlanType;
6-
use tokio::sync::OnceCell;
74

85
use crate::default_client::build_reqwest_client;
96

107
use super::storage::AgentIdentityAuthRecord;
118

129
const AGENT_IDENTITY_AUTHAPI_BASE_URL: &str = "https://auth.openai.com/api/accounts";
1310

14-
#[derive(Debug)]
11+
#[derive(Clone, Debug)]
1512
pub struct AgentIdentityAuth {
1613
record: AgentIdentityAuthRecord,
17-
process_task_id: Arc<OnceCell<String>>,
18-
}
19-
20-
impl Clone for AgentIdentityAuth {
21-
fn clone(&self) -> Self {
22-
Self {
23-
record: self.record.clone(),
24-
process_task_id: Arc::clone(&self.process_task_id),
25-
}
26-
}
14+
process_task_id: String,
2715
}
2816

2917
impl AgentIdentityAuth {
30-
pub fn new(record: AgentIdentityAuthRecord) -> Self {
31-
Self {
18+
pub async fn load(record: AgentIdentityAuthRecord) -> std::io::Result<Self> {
19+
let process_task_id = register_agent_task(
20+
&build_reqwest_client(),
21+
AGENT_IDENTITY_AUTHAPI_BASE_URL,
22+
key(&record),
23+
)
24+
.await
25+
.map_err(std::io::Error::other)?;
26+
Ok(Self {
3227
record,
33-
process_task_id: Arc::new(OnceCell::new()),
34-
}
28+
process_task_id,
29+
})
3530
}
3631

3732
pub fn record(&self) -> &AgentIdentityAuthRecord {
3833
&self.record
3934
}
4035

41-
pub fn process_task_id(&self) -> Option<&str> {
42-
self.process_task_id.get().map(String::as_str)
43-
}
44-
45-
pub async fn ensure_runtime(&self) -> std::io::Result<()> {
46-
self.process_task_id
47-
.get_or_try_init(|| async {
48-
register_agent_task(
49-
&build_reqwest_client(),
50-
AGENT_IDENTITY_AUTHAPI_BASE_URL,
51-
self.key(),
52-
)
53-
.await
54-
.map_err(std::io::Error::other)
55-
})
56-
.await
57-
.map(|_| ())
36+
pub fn process_task_id(&self) -> &str {
37+
&self.process_task_id
5838
}
5939

6040
pub fn account_id(&self) -> &str {
@@ -76,11 +56,11 @@ impl AgentIdentityAuth {
7656
pub fn is_fedramp_account(&self) -> bool {
7757
self.record.chatgpt_account_is_fedramp
7858
}
59+
}
7960

80-
fn key(&self) -> AgentIdentityKey<'_> {
81-
AgentIdentityKey {
82-
agent_runtime_id: &self.record.agent_runtime_id,
83-
private_key_pkcs8_base64: &self.record.agent_private_key,
84-
}
61+
fn key(record: &AgentIdentityAuthRecord) -> AgentIdentityKey<'_> {
62+
AgentIdentityKey {
63+
agent_runtime_id: &record.agent_runtime_id,
64+
private_key_pkcs8_base64: &record.agent_private_key,
8565
}
8666
}

codex-rs/login/src/auth/auth_tests.rs

Lines changed: 3 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -625,33 +625,6 @@ impl Drop for EnvVarGuard {
625625
}
626626
}
627627

628-
#[tokio::test]
629-
#[serial(codex_auth_env)]
630-
async fn load_auth_reads_agent_identity_from_env() {
631-
let codex_home = tempdir().unwrap();
632-
let expected_record = agent_identity_record("account-123");
633-
let agent_identity = fake_agent_identity_jwt(&expected_record).expect("fake agent identity");
634-
let _agent_guard = EnvVarGuard::set(CODEX_AGENT_IDENTITY_ENV_VAR, &agent_identity);
635-
636-
let auth = super::load_auth(
637-
codex_home.path(),
638-
/*enable_codex_api_key_env*/ false,
639-
AuthCredentialsStoreMode::File,
640-
)
641-
.await
642-
.expect("env auth should load")
643-
.expect("env auth should be present");
644-
645-
let CodexAuth::AgentIdentity(agent_identity) = auth else {
646-
panic!("env auth should load as agent identity");
647-
};
648-
assert_eq!(agent_identity.record(), &expected_record);
649-
assert!(
650-
!get_auth_file(codex_home.path()).exists(),
651-
"env auth should not write auth.json"
652-
);
653-
}
654-
655628
#[tokio::test]
656629
#[serial(codex_auth_env)]
657630
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() {
805778
}
806779

807780
fn agent_identity_record(account_id: &str) -> AgentIdentityAuthRecord {
781+
let key_material =
782+
codex_agent_identity::generate_agent_key_material().expect("generate agent key material");
808783
AgentIdentityAuthRecord {
809784
agent_runtime_id: "agent-runtime-id".to_string(),
810-
agent_private_key: "private-key".to_string(),
785+
agent_private_key: key_material.private_key_pkcs8_base64,
811786
account_id: account_id.to_string(),
812787
chatgpt_user_id: "user-id".to_string(),
813788
email: "user@example.com".to_string(),

codex-rs/login/src/auth/manager.rs

Lines changed: 7 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -212,7 +212,7 @@ impl CodexAuth {
212212
"agent identity auth is missing an agent identity token.",
213213
));
214214
};
215-
return Self::from_agent_identity_jwt(&agent_identity);
215+
return Self::from_agent_identity_jwt(&agent_identity).await;
216216
}
217217

218218
let storage_mode = auth_dot_json.storage_mode(auth_credentials_store_mode);
@@ -246,9 +246,9 @@ impl CodexAuth {
246246
.await
247247
}
248248

249-
pub fn from_agent_identity_jwt(jwt: &str) -> std::io::Result<Self> {
249+
pub async fn from_agent_identity_jwt(jwt: &str) -> std::io::Result<Self> {
250250
let record = AgentIdentityAuthRecord::from_agent_identity_jwt(jwt)?;
251-
Ok(Self::AgentIdentity(AgentIdentityAuth::new(record)))
251+
Ok(Self::AgentIdentity(AgentIdentityAuth::load(record).await?))
252252
}
253253

254254
pub fn auth_mode(&self) -> AuthMode {
@@ -322,16 +322,6 @@ impl CodexAuth {
322322
}
323323
}
324324

325-
pub async fn initialize_runtime(
326-
&self,
327-
_chatgpt_base_url: Option<String>,
328-
) -> std::io::Result<()> {
329-
match self {
330-
Self::AgentIdentity(auth) => auth.ensure_runtime().await,
331-
Self::ApiKey(_) | Self::Chatgpt(_) | Self::ChatgptAuthTokens(_) => Ok(()),
332-
}
333-
}
334-
335325
/// Returns `None` if Codex backend auth does not expose an account id.
336326
pub fn get_account_id(&self) -> Option<String> {
337327
match self {
@@ -749,7 +739,9 @@ async fn load_auth(
749739
}
750740

751741
if let Some(agent_identity) = read_codex_agent_identity_from_env() {
752-
return CodexAuth::from_agent_identity_jwt(&agent_identity).map(Some);
742+
return CodexAuth::from_agent_identity_jwt(&agent_identity)
743+
.await
744+
.map(Some);
753745
}
754746

755747
// Fall back to the configured persistent store (file/keyring/auto) for managed auth.
@@ -1400,12 +1392,7 @@ impl AuthManager {
14001392
tracing::error!("Failed to refresh token: {}", err);
14011393
return Some(auth);
14021394
}
1403-
let auth = self.auth_cached()?;
1404-
if let Err(err) = auth.initialize_runtime(self.chatgpt_base_url.clone()).await {
1405-
tracing::error!("Failed to initialize auth runtime: {err}");
1406-
return None;
1407-
}
1408-
Some(auth)
1395+
self.auth_cached()
14091396
}
14101397

14111398
/// Force a reload of the auth information from auth.json. Returns

codex-rs/model-provider/src/auth.rs

Lines changed: 11 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -21,23 +21,17 @@ struct AgentIdentityAuthProvider {
2121
impl AuthProvider for AgentIdentityAuthProvider {
2222
fn add_auth_headers(&self, headers: &mut HeaderMap) {
2323
let record = self.auth.record();
24-
let header_value = self
25-
.auth
26-
.process_task_id()
27-
.ok_or_else(|| std::io::Error::other("agent identity process task is not initialized"))
28-
.and_then(|task_id| {
29-
authorization_header_for_agent_task(
30-
AgentIdentityKey {
31-
agent_runtime_id: &record.agent_runtime_id,
32-
private_key_pkcs8_base64: &record.agent_private_key,
33-
},
34-
AgentTaskAuthorizationTarget {
35-
agent_runtime_id: &record.agent_runtime_id,
36-
task_id,
37-
},
38-
)
39-
.map_err(std::io::Error::other)
40-
});
24+
let header_value = authorization_header_for_agent_task(
25+
AgentIdentityKey {
26+
agent_runtime_id: &record.agent_runtime_id,
27+
private_key_pkcs8_base64: &record.agent_private_key,
28+
},
29+
AgentTaskAuthorizationTarget {
30+
agent_runtime_id: &record.agent_runtime_id,
31+
task_id: self.auth.process_task_id(),
32+
},
33+
)
34+
.map_err(std::io::Error::other);
4135

4236
if let Ok(header_value) = header_value
4337
&& let Ok(header) = HeaderValue::from_str(&header_value)

codex-rs/models-manager/src/manager_tests.rs

Lines changed: 0 additions & 81 deletions
Original file line numberDiff line numberDiff line change
@@ -9,9 +9,6 @@ use codex_login::ExternalAuth;
99
use codex_login::ExternalAuthRefreshContext;
1010
use codex_login::ExternalAuthTokens;
1111
use codex_login::TokenData;
12-
use codex_login::auth::AgentIdentityAuth;
13-
use codex_login::auth::AgentIdentityAuthRecord;
14-
use codex_protocol::account::PlanType;
1512
use codex_protocol::openai_models::ModelsResponse;
1613
use pretty_assertions::assert_eq;
1714
use serde_json::json;
@@ -237,18 +234,6 @@ c2ln",
237234
.expect("auth should be present")
238235
}
239236

240-
fn agent_identity_auth_for_tests() -> CodexAuth {
241-
CodexAuth::AgentIdentity(AgentIdentityAuth::new(AgentIdentityAuthRecord {
242-
agent_runtime_id: "agent-runtime-id".to_string(),
243-
agent_private_key: "agent-private-key".to_string(),
244-
account_id: "account-id".to_string(),
245-
chatgpt_user_id: "chatgpt-user-id".to_string(),
246-
email: "agent@example.com".to_string(),
247-
plan_type: PlanType::Pro,
248-
chatgpt_account_is_fedramp: false,
249-
}))
250-
}
251-
252237
#[tokio::test]
253238
async fn get_model_info_tracks_fallback_usage() {
254239
let codex_home = tempdir().expect("temp dir");
@@ -713,43 +698,6 @@ async fn refresh_available_models_fetches_with_chatgpt_auth_tokens() {
713698
);
714699
}
715700

716-
#[tokio::test]
717-
async fn refresh_available_models_fetches_with_agent_identity() {
718-
let dynamic_slug = "dynamic-model-only-for-test-agent-identity";
719-
let codex_home = tempdir().expect("temp dir");
720-
let endpoint = TestModelsEndpoint::new(vec![vec![remote_model(
721-
dynamic_slug,
722-
"Agent Identity",
723-
/*priority*/ 1,
724-
)]]);
725-
let manager = openai_manager_for_tests_with_auth(
726-
codex_home.path().to_path_buf(),
727-
endpoint.clone(),
728-
Some(AuthManager::from_auth_for_testing(
729-
agent_identity_auth_for_tests(),
730-
)),
731-
);
732-
733-
manager
734-
.refresh_available_models(RefreshStrategy::Online)
735-
.await
736-
.expect("refresh should fetch with agent identity");
737-
738-
assert!(
739-
manager
740-
.get_remote_models()
741-
.await
742-
.iter()
743-
.any(|candidate| candidate.slug == dynamic_slug),
744-
"remote refresh should include models fetched with agent identity"
745-
);
746-
assert_eq!(
747-
endpoint.fetch_count(),
748-
1,
749-
"endpoint should fetch models with agent identity"
750-
);
751-
}
752-
753701
#[test]
754702
fn build_available_models_picks_default_after_hiding_hidden_models() {
755703
let manager = static_manager_for_tests(ModelsResponse { models: Vec::new() });
@@ -768,35 +716,6 @@ fn build_available_models_picks_default_after_hiding_hidden_models() {
768716
assert_eq!(available, vec![expected_hidden, expected_visible]);
769717
}
770718

771-
#[tokio::test]
772-
async fn static_manager_treats_agent_identity_as_backend_auth_for_filtering() {
773-
let chatgpt_only_model = {
774-
let mut model = remote_model("chatgpt-only", "ChatGPT Only", /*priority*/ 0);
775-
model.supported_in_api = false;
776-
model
777-
};
778-
let api_model = remote_model("api-model", "API Model", /*priority*/ 1);
779-
let manager = StaticModelsManager::new(
780-
Some(AuthManager::from_auth_for_testing(
781-
agent_identity_auth_for_tests(),
782-
)),
783-
ModelsResponse {
784-
models: vec![chatgpt_only_model, api_model],
785-
},
786-
CollaborationModesConfig::default(),
787-
);
788-
789-
let agent_identity_models = manager.list_models(RefreshStrategy::Online).await;
790-
791-
assert_eq!(
792-
agent_identity_models
793-
.iter()
794-
.map(|model| model.model.as_str())
795-
.collect::<Vec<_>>(),
796-
vec!["chatgpt-only", "api-model"]
797-
);
798-
}
799-
800719
#[tokio::test]
801720
async fn static_manager_reads_latest_auth_mode() {
802721
let auth_manager =

0 commit comments

Comments
 (0)