mirror of
https://github.com/openai/codex.git
synced 2026-11-13 14:46:14 +00:00
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:
1 parent
2307aa8d98
commit
c08177f7d0
6 files changed
+43
-188
No files matched your search
@@ -1,9 +1,6 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use codex_agent_identity::AgentIdentityKey;
|
||||
use codex_agent_identity::register_agent_task;
|
||||
use codex_protocol::account::PlanType as AccountPlanType;
|
||||
use tokio::sync::OnceCell;
|
||||
|
||||
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";
|
||||
|
||||
#[derive(Debug)]
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct AgentIdentityAuth {
|
||||
record: AgentIdentityAuthRecord,
|
||||
process_task_id: Arc<OnceCell<String>>,
|
||||
}
|
||||
|
||||
impl Clone for AgentIdentityAuth {
|
||||
fn clone(&self) -> Self {
|
||||
Self {
|
||||
record: self.record.clone(),
|
||||
process_task_id: Arc::clone(&self.process_task_id),
|
||||
}
|
||||
}
|
||||
process_task_id: String,
|
||||
}
|
||||
|
||||
impl AgentIdentityAuth {
|
||||
pub fn new(record: AgentIdentityAuthRecord) -> Self {
|
||||
Self {
|
||||
pub async fn load(record: AgentIdentityAuthRecord) -> std::io::Result<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,
|
||||
process_task_id: Arc::new(OnceCell::new()),
|
||||
}
|
||||
process_task_id,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn record(&self) -> &AgentIdentityAuthRecord {
|
||||
&self.record
|
||||
}
|
||||
|
||||
pub fn process_task_id(&self) -> Option<&str> {
|
||||
self.process_task_id.get().map(String::as_str)
|
||||
}
|
||||
|
||||
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 process_task_id(&self) -> &str {
|
||||
&self.process_task_id
|
||||
}
|
||||
|
||||
pub fn account_id(&self) -> &str {
|
||||
@@ -76,11 +56,11 @@ impl AgentIdentityAuth {
|
||||
pub fn is_fedramp_account(&self) -> bool {
|
||||
self.record.chatgpt_account_is_fedramp
|
||||
}
|
||||
}
|
||||
|
||||
fn key(&self) -> AgentIdentityKey<'_> {
|
||||
AgentIdentityKey {
|
||||
agent_runtime_id: &self.record.agent_runtime_id,
|
||||
private_key_pkcs8_base64: &self.record.agent_private_key,
|
||||
}
|
||||
fn key(record: &AgentIdentityAuthRecord) -> AgentIdentityKey<'_> {
|
||||
AgentIdentityKey {
|
||||
agent_runtime_id: &record.agent_runtime_id,
|
||||
private_key_pkcs8_base64: &record.agent_private_key,
|
||||
}
|
||||
}
|
||||
@@ -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]
|
||||
#[serial(codex_auth_env)]
|
||||
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 {
|
||||
let key_material =
|
||||
codex_agent_identity::generate_agent_key_material().expect("generate agent key material");
|
||||
AgentIdentityAuthRecord {
|
||||
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(),
|
||||
chatgpt_user_id: "user-id".to_string(),
|
||||
email: "user@example.com".to_string(),
|
||||
|
||||
@@ -212,7 +212,7 @@ impl CodexAuth {
|
||||
"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);
|
||||
@@ -246,9 +246,9 @@ impl CodexAuth {
|
||||
.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)?;
|
||||
Ok(Self::AgentIdentity(AgentIdentityAuth::new(record)))
|
||||
Ok(Self::AgentIdentity(AgentIdentityAuth::load(record).await?))
|
||||
}
|
||||
|
||||
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.
|
||||
pub fn get_account_id(&self) -> Option<String> {
|
||||
match self {
|
||||
@@ -749,7 +739,9 @@ async fn load_auth(
|
||||
}
|
||||
|
||||
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.
|
||||
@@ -1400,12 +1392,7 @@ impl AuthManager {
|
||||
tracing::error!("Failed to refresh token: {}", err);
|
||||
return Some(auth);
|
||||
}
|
||||
let auth = 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)
|
||||
self.auth_cached()
|
||||
}
|
||||
|
||||
/// Force a reload of the auth information from auth.json. Returns
|
||||
|
||||
Reference in new issue
Block a user