agent/codex
1
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: add run task identity primitives (#19047)

Browse Source 
## Stack

This is PR 1 of the simplified HAI single-run-task stack:

- [#19047](https://github.com/openai/codex/pull/19047) Agent Identity
assertion and task-registration primitives, including the shared
run-task helper used by existing Agent Identity JWT auth.
- [#19049](https://github.com/openai/codex/pull/19049)
Disabled-by-default ChatGPT auth opt-in that provisions/reuses persisted
Agent Identity runtime auth and its single run task.
- [#19051](https://github.com/openai/codex/pull/19051) Run-scoped
provider auth that uses one backend-owned task id for first-party
inference and compaction requests.

[#19054](https://github.com/openai/codex/pull/19054) collapsed out of
the active stack because the simplified design no longer needs a
separate background/control-plane task helper.

## Summary

The simplified POC shape is one backend-owned task per Agent Identity
run. This PR makes the first layer match that final shape directly
instead of introducing task targets, caller-owned external task refs, or
intermediate wrappers that later PRs would need to undo.

What changed:

- keeps the `AgentAssertion` wire payload as `agent_runtime_id`,
`task_id`, `timestamp`, and `signature`
- exposes `register_agent_task` as the single task-registration helper
for both existing Agent Identity JWT auth and the ChatGPT-registration
path added later in the stack
- makes task registration send only the signed registration timestamp;
the backend owns the returned opaque task id
- removes the unused target/task-kind/external-task-ref surfaces from
`codex-agent-identity`
- keeps Agent Identity JWT JWKS lookup separate from agent/task
registration URL derivation
- updates Agent Identity JWT auth to register one run task during auth
construction and share that task across cloned auth handles

This PR intentionally does not enable ChatGPT-derived Agent Identity.
That opt-in and config gate are added in the next PR.

## Testing

- `just test -p codex-agent-identity`
This commit is contained in:
Adrian
2026-06-17 11:23:39 -07:00
committed by GitHub
Unverified
parent 1391d786bc
commit d9e0551564
8 changed files with 546 additions and 198 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/agent-identity/src/lib.rs
+322 -82
View File
@@ -1,4 +1,6 @@
use std::collections::BTreeMap;
use std::error::Error as StdError;
use std::fmt;
use std::time::Duration;
use anyhow::Context;
@@ -34,21 +36,66 @@ const AGENT_TASK_REGISTRATION_TIMEOUT: Duration = Duration::from_secs(30);
const AGENT_IDENTITY_JWKS_TIMEOUT: Duration = Duration::from_secs(10);
const AGENT_IDENTITY_JWT_AUDIENCE: &str = "codex-app-server";
const AGENT_IDENTITY_JWT_ISSUER: &str = "https://chatgpt.com/codex-backend/agent-identity";
const AGENT_REGISTRATION_TIMEOUT: Duration = Duration::from_secs(15);
const PROD_AGENT_IDENTITY_AUTHAPI_BASE_URL: &str = "https://auth.openai.com/api/accounts";
const STAGING_AGENT_IDENTITY_AUTHAPI_BASE_URL: &str = "https://auth.api.openai.org/api/accounts";
const AGENT_IDENTITY_KEY_SEED_BYTES: usize = 64;
const AGENT_IDENTITY_KEY_DERIVATION_CONTEXT: &[u8] = b"codex-agent-identity-ed25519-v1";
/// Stored key material for a registered agent identity.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum ChatGptEnvironment {
#[default]
Production,
Staging,
}
impl ChatGptEnvironment {
pub fn from_chatgpt_base_url(chatgpt_base_url: &str) -> Result<Self> {
match chatgpt_base_url.trim_end_matches('/') {
"https://chatgpt.com"
| "https://chatgpt.com/backend-api"
| "https://chatgpt.com/codex"
| "https://chatgpt.com/backend-api/codex"
| "https://chat.openai.com"
| "https://chat.openai.com/backend-api"
| "https://chat.openai.com/codex"
| "https://chat.openai.com/backend-api/codex" => Ok(Self::Production),
"https://chatgpt-staging.com"
| "https://chatgpt-staging.com/backend-api"
| "https://chatgpt-staging.com/codex"
| "https://chatgpt-staging.com/backend-api/codex" => Ok(Self::Staging),
_ => anyhow::bail!(
"Agent Identity only supports production and staging ChatGPT environments"
),
}
}
pub fn chatgpt_base_url(self) -> &'static str {
match self {
Self::Production => "https://chatgpt.com/backend-api",
Self::Staging => "https://chatgpt-staging.com/backend-api",
}
}
pub fn agent_identity_authapi_base_url(self) -> &'static str {
match self {
Self::Production => PROD_AGENT_IDENTITY_AUTHAPI_BASE_URL,
Self::Staging => STAGING_AGENT_IDENTITY_AUTHAPI_BASE_URL,
}
}
}
/// Borrowed durable signing material for a registered agent identity.
///
/// This intentionally does not include a task id. Task ids are scoped to a
/// single Codex run, while the agent runtime id and private key are the
/// reusable identity material used to register and sign that run task.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct AgentIdentityKey<'a> {
pub agent_runtime_id: &'a str,
pub private_key_pkcs8_base64: &'a str,
}
/// Task binding to use when constructing a task-scoped AgentAssertion.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct AgentTaskAuthorizationTarget<'a> {
pub agent_runtime_id: &'a str,
pub task_id: &'a str,
}
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct AgentBillOfMaterials {
pub agent_version: String,
@@ -103,23 +150,92 @@ struct RegisterTaskResponse {
encrypted_task_id_camel: Option<String>,
}
#[derive(Debug, Serialize)]
struct RegisterAgentRequest {
abom: AgentBillOfMaterials,
agent_public_key: String,
capabilities: Vec<String>,
ttl: Option<u64>,
}
#[derive(Debug, Deserialize)]
struct RegisterAgentResponse {
agent_runtime_id: String,
}
/// HTTP status failure returned by Agent Identity registration endpoints.
#[derive(Debug)]
pub struct AgentIdentityRegistrationHttpError {
operation: &'static str,
status: reqwest::StatusCode,
body: String,
}
impl AgentIdentityRegistrationHttpError {
fn new(operation: &'static str, status: reqwest::StatusCode, body: String) -> Self {
Self {
operation,
status,
body,
}
}
/// HTTP status returned by the registration endpoint.
pub fn status(&self) -> reqwest::StatusCode {
self.status
}
}
impl fmt::Display for AgentIdentityRegistrationHttpError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
if self.body.is_empty() {
write!(f, "{} failed with status {}", self.operation, self.status)
} else {
write!(
f,
"{} failed with status {}: {}",
self.operation, self.status, self.body
)
}
}
}
impl StdError for AgentIdentityRegistrationHttpError {}
/// Returns whether an Agent Identity registration error is safe to retry.
pub fn is_retryable_registration_error(error: &anyhow::Error) -> bool {
error.chain().any(is_retryable_registration_cause)
}
fn is_retryable_registration_cause(cause: &(dyn StdError + 'static)) -> bool {
if let Some(error) = cause.downcast_ref::<AgentIdentityRegistrationHttpError>() {
return is_retryable_registration_status(error.status());
}
if let Some(error) = cause.downcast_ref::<reqwest::Error>() {
if let Some(status) = error.status() {
return is_retryable_registration_status(status);
}
return error.is_timeout() || error.is_connect() || error.is_request();
}
false
}
fn is_retryable_registration_status(status: reqwest::StatusCode) -> bool {
status == reqwest::StatusCode::TOO_MANY_REQUESTS || status.is_server_error()
}
pub fn authorization_header_for_agent_task(
key: AgentIdentityKey<'_>,
target: AgentTaskAuthorizationTarget<'_>,
task_id: &str,
) -> Result<String> {
anyhow::ensure!(
key.agent_runtime_id == target.agent_runtime_id,
"agent task runtime {} does not match stored agent identity {}",
target.agent_runtime_id,
key.agent_runtime_id
);
let timestamp = Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true);
let envelope = AgentAssertionEnvelope {
agent_runtime_id: target.agent_runtime_id.to_string(),
task_id: target.task_id.to_string(),
agent_runtime_id: key.agent_runtime_id.to_string(),
task_id: task_id.to_string(),
timestamp: timestamp.clone(),
signature: sign_agent_assertion_payload(key, target.task_id, &timestamp)?,
signature: sign_agent_assertion_payload(key, task_id, &timestamp)?,
};
let serialized_assertion = serialize_agent_assertion(&envelope)?;
Ok(format!("AgentAssertion {serialized_assertion}"))
@@ -127,10 +243,10 @@ pub fn authorization_header_for_agent_task(
pub async fn fetch_agent_identity_jwks(
client: &reqwest::Client,
chatgpt_base_url: &str,
agent_identity_jwt_base_url: &str,
) -> Result<JwkSet> {
let response = client
.get(agent_identity_jwks_url(chatgpt_base_url))
.get(agent_identity_jwks_url(agent_identity_jwt_base_url))
.timeout(AGENT_IDENTITY_JWKS_TIMEOUT)
.send()
.await
@@ -195,7 +311,7 @@ pub fn sign_task_registration_payload(
pub async fn register_agent_task(
client: &reqwest::Client,
chatgpt_base_url: &str,
agent_identity_authapi_base_url: &str,
key: AgentIdentityKey<'_>,
) -> Result<String> {
let timestamp = Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true);
@@ -203,7 +319,7 @@ pub async fn register_agent_task(
signature: sign_task_registration_payload(key, &timestamp)?,
timestamp,
};
let url = agent_task_registration_url(chatgpt_base_url, key.agent_runtime_id);
let url = agent_task_registration_url(agent_identity_authapi_base_url, key.agent_runtime_id);
let response = client
.post(url)
@@ -220,7 +336,12 @@ pub async fn register_agent_task(
} else {
body
};
anyhow::bail!("failed to register agent task with status {status}: {body}");
return Err(AgentIdentityRegistrationHttpError::new(
"agent task registration",
status,
body,
)
.into());
}
let response = response
@@ -231,6 +352,45 @@ pub async fn register_agent_task(
task_id_from_register_task_response(key, response)
}
pub async fn register_agent_identity(
client: &reqwest::Client,
agent_identity_authapi_base_url: &str,
access_token: &str,
is_fedramp_account: bool,
key_material: &GeneratedAgentKeyMaterial,
abom: AgentBillOfMaterials,
capabilities: Vec<String>,
) -> Result<String> {
let url = agent_registration_url(agent_identity_authapi_base_url);
let request = RegisterAgentRequest {
abom,
agent_public_key: key_material.public_key_ssh.clone(),
capabilities,
ttl: None,
};
let mut request_builder = client
.post(&url)
.bearer_auth(access_token)
.json(&request)
.timeout(AGENT_REGISTRATION_TIMEOUT);
if is_fedramp_account {
request_builder = request_builder.header("X-OpenAI-Fedramp", "true");
}
let response = request_builder
.send()
.await
.with_context(|| format!("failed to send agent identity registration request to {url}"))?
.error_for_status()
.with_context(|| format!("agent identity registration failed for {url}"))?
.json::<RegisterAgentResponse>()
.await
.with_context(|| format!("failed to parse agent identity response from {url}"))?;
Ok(response.agent_runtime_id)
}
fn task_id_from_register_task_response(
key: AgentIdentityKey<'_>,
response: RegisterTaskResponse,
@@ -260,10 +420,17 @@ pub fn decrypt_task_id_response(
}
pub fn generate_agent_key_material() -> Result<GeneratedAgentKeyMaterial> {
let mut secret_key_bytes = [0u8; 32];
let mut seed_material = [0u8; AGENT_IDENTITY_KEY_SEED_BYTES];
OsRng
.try_fill_bytes(&mut secret_key_bytes)
.context("failed to generate agent identity private key bytes")?;
.try_fill_bytes(&mut seed_material)
.context("failed to generate agent identity private key seed material")?;
// Ed25519 stores a 32-byte seed, so derive it from all sampled seed material.
let mut digest = Sha512::new();
digest.update(AGENT_IDENTITY_KEY_DERIVATION_CONTEXT);
digest.update(seed_material);
let digest = digest.finalize();
let mut secret_key_bytes = [0u8; 32];
secret_key_bytes.copy_from_slice(&digest[..32]);
let signing_key = SigningKey::from_bytes(&secret_key_bytes);
let private_key_pkcs8 = signing_key
.to_pkcs8_der()
@@ -296,23 +463,22 @@ pub fn curve25519_secret_key_from_private_key_pkcs8_base64(
Ok(curve25519_secret_key_from_signing_key(&signing_key))
}
pub fn agent_registration_url(chatgpt_base_url: &str) -> String {
let trimmed = chatgpt_base_url.trim_end_matches('/');
format!("{trimmed}/v1/agent/register")
pub fn agent_registration_url(agent_identity_authapi_base_url: &str) -> String {
agent_identity_authapi_url(agent_identity_authapi_base_url, "/v1/agent/register")
}
pub fn agent_task_registration_url(chatgpt_base_url: &str, agent_runtime_id: &str) -> String {
let trimmed = chatgpt_base_url.trim_end_matches('/');
format!("{trimmed}/v1/agent/{agent_runtime_id}/task/register")
pub fn agent_task_registration_url(
agent_identity_authapi_base_url: &str,
agent_runtime_id: &str,
) -> String {
agent_identity_authapi_url(
agent_identity_authapi_base_url,
&format!("/v1/agent/{agent_runtime_id}/task/register"),
)
}
pub fn agent_identity_biscuit_url(chatgpt_base_url: &str) -> String {
let trimmed = chatgpt_base_url.trim_end_matches('/');
format!("{trimmed}/authenticate_app_v2")
}
pub fn agent_identity_jwks_url(chatgpt_base_url: &str) -> String {
let trimmed = chatgpt_base_url.trim_end_matches('/');
pub fn agent_identity_jwks_url(agent_identity_jwt_base_url: &str) -> String {
let trimmed = agent_identity_jwt_base_url.trim_end_matches('/');
if trimmed.contains("/backend-api") {
format!("{trimmed}/wham/agent-identities/jwks")
} else {
@@ -320,15 +486,9 @@ pub fn agent_identity_jwks_url(chatgpt_base_url: &str) -> String {
}
}
pub fn agent_identity_request_id() -> Result<String> {
let mut request_id_bytes = [0u8; 16];
OsRng
.try_fill_bytes(&mut request_id_bytes)
.context("failed to generate agent identity request id")?;
Ok(format!(
"codex-agent-identity-{}",
URL_SAFE_NO_PAD.encode(request_id_bytes)
))
fn agent_identity_authapi_url(agent_identity_authapi_base_url: &str, api_path: &str) -> String {
let base_url = agent_identity_authapi_base_url.trim_end_matches('/');
format!("{base_url}{api_path}")
}
pub fn build_abom(session_source: SessionSource) -> AgentBillOfMaterials {
@@ -412,6 +572,24 @@ mod tests {
use super::*;
#[test]
fn register_task_request_uses_single_run_task_shape() {
let request = RegisterTaskRequest {
timestamp: "2026-04-23T00:00:00Z".to_string(),
signature: "signature".to_string(),
};
let serialized = serde_json::to_value(request).expect("serialize request");
assert_eq!(
serialized,
serde_json::json!({
"timestamp": "2026-04-23T00:00:00Z",
"signature": "signature",
})
);
}
#[test]
fn authorization_header_for_agent_task_serializes_signed_agent_assertion() {
let signing_key = SigningKey::from_bytes(&[7u8; 32]);
@@ -422,13 +600,9 @@ mod tests {
agent_runtime_id: "agent-123",
private_key_pkcs8_base64: &BASE64_STANDARD.encode(private_key.as_bytes()),
};
let target = AgentTaskAuthorizationTarget {
agent_runtime_id: "agent-123",
task_id: "task-123",
};
let header =
authorization_header_for_agent_task(key, target).expect("build agent assertion header");
let header = authorization_header_for_agent_task(key, "task-123")
.expect("build agent assertion header");
let token = header
.strip_prefix("AgentAssertion ")
.expect("agent assertion scheme");
@@ -464,31 +638,6 @@ mod tests {
.expect("signature should verify");
}
#[test]
fn authorization_header_for_agent_task_rejects_mismatched_runtime() {
let signing_key = SigningKey::from_bytes(&[7u8; 32]);
let private_key = signing_key
.to_pkcs8_der()
.expect("encode test key material");
let private_key_pkcs8_base64 = BASE64_STANDARD.encode(private_key.as_bytes());
let key = AgentIdentityKey {
agent_runtime_id: "agent-123",
private_key_pkcs8_base64: &private_key_pkcs8_base64,
};
let target = AgentTaskAuthorizationTarget {
agent_runtime_id: "agent-456",
task_id: "task-123",
};
let error = authorization_header_for_agent_task(key, target)
.expect_err("runtime mismatch should fail");
assert_eq!(
error.to_string(),
"agent task runtime agent-456 does not match stored agent identity agent-123"
);
}
#[test]
fn decode_agent_identity_jwt_reads_claims() {
let jwt = jwt_with_payload(serde_json::json!({
@@ -704,7 +853,98 @@ J1bwkqKZTB5dHolX9A58e/xXnfZ5P8f3Z83+Izap3FwqQulk7b1WO1MQcHuVg2NN
}
#[test]
fn agent_identity_jwks_url_uses_backend_api_base_url() {
fn chatgpt_environment_maps_known_urls_to_authapi() -> anyhow::Result<()> {
assert_eq!(
ChatGptEnvironment::from_chatgpt_base_url("https://chatgpt.com/backend-api/codex")?,
ChatGptEnvironment::Production
);
assert_eq!(
ChatGptEnvironment::Production.agent_identity_authapi_base_url(),
"https://auth.openai.com/api/accounts"
);
assert_eq!(
ChatGptEnvironment::from_chatgpt_base_url("https://chatgpt-staging.com/backend-api")?,
ChatGptEnvironment::Staging
);
assert_eq!(
ChatGptEnvironment::Staging.agent_identity_authapi_base_url(),
"https://auth.api.openai.org/api/accounts"
);
Ok(())
}
#[test]
fn chatgpt_environment_rejects_custom_urls() {
assert!(ChatGptEnvironment::from_chatgpt_base_url("http://localhost:8080").is_err(),);
}
#[test]
fn agent_registration_url_appends_to_authapi_base_url() {
assert_eq!(
agent_registration_url("https://auth.openai.com/api/accounts"),
"https://auth.openai.com/api/accounts/v1/agent/register"
);
assert_eq!(
agent_registration_url("http://localhost:8080"),
"http://localhost:8080/v1/agent/register"
);
assert_eq!(
agent_registration_url("http://localhost:8080/backend-api"),
"http://localhost:8080/backend-api/v1/agent/register"
);
}
#[test]
fn agent_task_registration_url_appends_to_authapi_base_url() {
assert_eq!(
agent_task_registration_url("https://auth.openai.com/api/accounts", "agent-runtime-id"),
"https://auth.openai.com/api/accounts/v1/agent/agent-runtime-id/task/register"
);
assert_eq!(
agent_task_registration_url(
"https://auth.openai.com/api/accounts/",
"agent-runtime-id"
),
"https://auth.openai.com/api/accounts/v1/agent/agent-runtime-id/task/register"
);
assert_eq!(
agent_task_registration_url("http://localhost:8080", "agent-runtime-id"),
"http://localhost:8080/v1/agent/agent-runtime-id/task/register"
);
}
#[test]
fn retryable_registration_error_accepts_429_and_5xx() {
let too_many_requests = anyhow::Error::new(AgentIdentityRegistrationHttpError::new(
"agent registration",
reqwest::StatusCode::TOO_MANY_REQUESTS,
"rate limited".to_string(),
));
let unavailable = anyhow::Error::new(AgentIdentityRegistrationHttpError::new(
"agent registration",
reqwest::StatusCode::SERVICE_UNAVAILABLE,
"try later".to_string(),
));
assert!(is_retryable_registration_error(&too_many_requests));
assert!(is_retryable_registration_error(&unavailable));
}
#[test]
fn retryable_registration_error_rejects_hard_failures() {
let forbidden = anyhow::Error::new(AgentIdentityRegistrationHttpError::new(
"agent registration",
reqwest::StatusCode::FORBIDDEN,
"not allowed".to_string(),
));
let malformed = anyhow::anyhow!("failed to sign registration request");
assert!(!is_retryable_registration_error(&forbidden));
assert!(!is_retryable_registration_error(&malformed));
}
#[test]
fn agent_identity_jwks_url_uses_agent_identity_jwt_route() {
assert_eq!(
agent_identity_jwks_url("https://chatgpt.com/backend-api"),
"https://chatgpt.com/backend-api/wham/agent-identities/jwks"
@@ -716,7 +956,7 @@ J1bwkqKZTB5dHolX9A58e/xXnfZ5P8f3Z83+Izap3FwqQulk7b1WO1MQcHuVg2NN
}
#[test]
fn agent_identity_jwks_url_uses_codex_api_base_url() {
fn agent_identity_jwks_url_uses_jwt_issuer_base_url() {
assert_eq!(
agent_identity_jwks_url("http://localhost:8080/api/codex"),
"http://localhost:8080/api/codex/agent-identities/jwks"
codex-rs/cli/src/main.rs
+6 -3
View File
@@ -1731,9 +1731,12 @@ async fn load_exec_server_remote_auth_provider(
let agent_identity_jwt = read_codex_access_token_from_env().ok_or_else(|| {
anyhow::anyhow!("CODEX_ACCESS_TOKEN is required when --use-agent-identity-auth is set")
})?;
let auth =
CodexAuth::from_agent_identity_jwt(&agent_identity_jwt, Some(&config.chatgpt_base_url))
.await?;
let auth = CodexAuth::from_agent_identity_jwt(
&agent_identity_jwt,
Some(&config.chatgpt_base_url),
/*agent_identity_authapi_base_url_override*/ None,
)
.await?;
return Ok(codex_model_provider::auth_provider_from_auth(&auth));
}
codex-rs/exec/src/lib.rs
+1
View File
@@ -475,6 +475,7 @@ pub async fn run_main(cli: Cli, arg0_paths: Arg0DispatchPaths) -> anyhow::Result
forced_login_method: config.forced_login_method,
forced_chatgpt_workspace_id: config.forced_chatgpt_workspace_id.clone(),
chatgpt_base_url: Some(config.chatgpt_base_url.clone()),
agent_identity_authapi_base_url: None,
})
.await
{
codex-rs/login/src/auth/agent_identity.rs
+138 -82
View File
@@ -1,75 +1,84 @@
use std::sync::Arc;
use codex_agent_identity::AgentIdentityKey;
use codex_agent_identity::register_agent_task;
use codex_protocol::account::PlanType as AccountPlanType;
use std::env;
use crate::default_client::build_reqwest_client;
use super::storage::AgentIdentityAuthRecord;
const PROD_AGENT_IDENTITY_AUTHAPI_BASE_URL: &str = "https://auth.openai.com/api/accounts";
const CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL_ENV_VAR: &str = "CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL";
#[derive(Clone, Debug)]
pub struct AgentIdentityAuth {
inner: Arc<AgentIdentityAuthInner>,
}
#[derive(Debug)]
struct AgentIdentityAuthInner {
record: AgentIdentityAuthRecord,
process_task_id: String,
run_task_id: String,
}
impl AgentIdentityAuth {
pub async fn load(record: AgentIdentityAuthRecord) -> std::io::Result<Self> {
let agent_identity_authapi_base_url = agent_identity_authapi_base_url();
let process_task_id = register_agent_task(
pub async fn load(
record: AgentIdentityAuthRecord,
agent_identity_authapi_base_url: &str,
) -> std::io::Result<Self> {
let run_task_id = register_agent_task(
&build_reqwest_client(),
&agent_identity_authapi_base_url,
key(&record),
agent_identity_authapi_base_url,
key_for_record(&record),
)
.await
.map_err(std::io::Error::other)?;
Ok(Self {
record,
process_task_id,
inner: Arc::new(AgentIdentityAuthInner {
record,
run_task_id,
}),
})
}
pub fn record(&self) -> &AgentIdentityAuthRecord {
&self.record
#[cfg(test)]
fn from_initialized_record(record: AgentIdentityAuthRecord, run_task_id: String) -> Self {
Self {
inner: Arc::new(AgentIdentityAuthInner {
record,
run_task_id,
}),
}
}
pub fn process_task_id(&self) -> &str {
&self.process_task_id
pub fn record(&self) -> &AgentIdentityAuthRecord {
&self.inner.record
}
pub fn run_task_id(&self) -> &str {
&self.inner.run_task_id
}
pub fn account_id(&self) -> &str {
&self.record.account_id
&self.inner.record.account_id
}
pub fn chatgpt_user_id(&self) -> &str {
&self.record.chatgpt_user_id
&self.inner.record.chatgpt_user_id
}
pub fn email(&self) -> &str {
&self.record.email
&self.inner.record.email
}
pub fn plan_type(&self) -> AccountPlanType {
self.record.plan_type
self.inner.record.plan_type
}
pub fn is_fedramp_account(&self) -> bool {
self.record.chatgpt_account_is_fedramp
self.inner.record.chatgpt_account_is_fedramp
}
}
fn agent_identity_authapi_base_url() -> String {
env::var(CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL_ENV_VAR)
.ok()
.map(|base_url| base_url.trim().trim_end_matches('/').to_string())
.filter(|base_url| !base_url.is_empty())
.unwrap_or_else(|| PROD_AGENT_IDENTITY_AUTHAPI_BASE_URL.to_string())
}
fn key(record: &AgentIdentityAuthRecord) -> AgentIdentityKey<'_> {
fn key_for_record(record: &AgentIdentityAuthRecord) -> AgentIdentityKey<'_> {
AgentIdentityKey {
agent_runtime_id: &record.agent_runtime_id,
private_key_pkcs8_base64: &record.agent_private_key,
@@ -78,63 +87,110 @@ fn key(record: &AgentIdentityAuthRecord) -> AgentIdentityKey<'_> {
#[cfg(test)]
mod tests {
use std::sync::Arc;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use codex_agent_identity::generate_agent_key_material;
use pretty_assertions::assert_eq;
use serde_json::json;
use wiremock::Mock;
use wiremock::MockServer;
use wiremock::ResponseTemplate;
use wiremock::matchers::method;
use wiremock::matchers::path;
use super::*;
use serial_test::serial;
#[test]
#[serial(codex_auth_env)]
fn agent_identity_authapi_base_url_prefers_env_value() {
let _guard = EnvVarGuard::set(
CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL_ENV_VAR,
"https://authapi.example.test/api/accounts/",
);
assert_eq!(
agent_identity_authapi_base_url(),
"https://authapi.example.test/api/accounts"
);
}
#[test]
#[serial(codex_auth_env)]
fn agent_identity_authapi_base_url_uses_prod_authapi_by_default() {
let _guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL_ENV_VAR);
assert_eq!(
agent_identity_authapi_base_url(),
PROD_AGENT_IDENTITY_AUTHAPI_BASE_URL
);
}
struct EnvVarGuard {
key: &'static str,
original: Option<std::ffi::OsString>,
}
impl EnvVarGuard {
fn set(key: &'static str, value: &str) -> Self {
let original = env::var_os(key);
unsafe {
env::set_var(key, value);
}
Self { key, original }
}
fn remove(key: &'static str) -> Self {
let original = env::var_os(key);
unsafe {
env::remove_var(key);
}
Self { key, original }
fn agent_identity_record(private_key: String) -> AgentIdentityAuthRecord {
AgentIdentityAuthRecord {
agent_runtime_id: "agent-runtime-1".to_string(),
agent_private_key: private_key,
account_id: "account-1".to_string(),
chatgpt_user_id: "user-1".to_string(),
email: "agent@example.com".to_string(),
plan_type: AccountPlanType::Plus,
chatgpt_account_is_fedramp: false,
}
}
impl Drop for EnvVarGuard {
fn drop(&mut self) {
unsafe {
match &self.original {
Some(value) => env::set_var(self.key, value),
None => env::remove_var(self.key),
fn agent_identity_record_with_generated_key() -> AgentIdentityAuthRecord {
let key_material = generate_agent_key_material().expect("generate key material");
agent_identity_record(key_material.private_key_pkcs8_base64)
}
#[tokio::test]
async fn load_registers_run_task() -> anyhow::Result<()> {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/v1/agent/agent-runtime-1/task/register"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"task_id": "task-run-1",
})))
.expect(1)
.mount(&server)
.await;
let auth =
AgentIdentityAuth::load(agent_identity_record_with_generated_key(), &server.uri())
.await?;
assert_eq!(auth.run_task_id(), "task-run-1");
let requests = server
.received_requests()
.await
.expect("failed to fetch task registration request");
let request_body = requests[0]
.body_json::<serde_json::Value>()
.expect("task registration request should be JSON");
let request_body = request_body
.as_object()
.expect("request body should be object");
assert!(request_body.get("timestamp").is_some());
assert!(request_body.get("signature").is_some());
assert_eq!(request_body.len(), 2);
Ok(())
}
#[test]
fn run_task_is_shared_across_clones() {
let auth = AgentIdentityAuth::from_initialized_record(
agent_identity_record_with_generated_key(),
"task-run-1".to_string(),
);
let cloned = auth.clone();
assert!(Arc::ptr_eq(&auth.inner, &cloned.inner));
assert_eq!(cloned.run_task_id(), "task-run-1");
}
#[tokio::test]
async fn failed_run_task_registration_can_be_retried_on_next_call() -> anyhow::Result<()> {
let server = MockServer::start().await;
let request_count = Arc::new(AtomicUsize::new(0));
let response_count = Arc::clone(&request_count);
Mock::given(method("POST"))
.and(path("/v1/agent/agent-runtime-1/task/register"))
.respond_with(move |_request: &wiremock::Request| {
if response_count.fetch_add(1, Ordering::SeqCst) == 0 {
ResponseTemplate::new(500)
} else {
ResponseTemplate::new(200).set_body_json(json!({
"task_id": "task-run-1",
}))
}
}
}
})
.expect(2)
.mount(&server)
.await;
let record = agent_identity_record_with_generated_key();
AgentIdentityAuth::load(record.clone(), &server.uri())
.await
.expect_err("first registration should fail");
let auth = AgentIdentityAuth::load(record, &server.uri()).await?;
assert_eq!(request_count.load(Ordering::SeqCst), 2);
assert_eq!(auth.run_task_id(), "task-run-1");
Ok(())
}
}
codex-rs/login/src/auth/auth_tests.rs
+29 -16
View File
@@ -348,6 +348,7 @@ async fn pro_account_with_no_api_key_uses_chatgpt_auth() {
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::Direct,
/*agent_identity_authapi_base_url*/ None,
)
.await
.unwrap()
@@ -408,6 +409,7 @@ async fn loads_api_key_from_auth_json() {
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::Direct,
/*agent_identity_authapi_base_url*/ None,
)
.await
.unwrap()
@@ -500,6 +502,7 @@ async fn refresh_failure_is_scoped_to_the_matching_auth_snapshot() {
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::Direct,
/*agent_identity_authapi_base_url*/ None,
)
.await
.expect("load auth")
@@ -519,6 +522,7 @@ async fn refresh_failure_is_scoped_to_the_matching_auth_snapshot() {
AuthCredentialsStoreMode::File,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::Direct,
/*agent_identity_authapi_base_url*/ None,
)
.await
.expect("updated auth should parse");
@@ -824,6 +828,7 @@ async fn build_config(
forced_login_method,
forced_chatgpt_workspace_id,
chatgpt_base_url: None,
agent_identity_authapi_base_url: None,
}
}
@@ -886,7 +891,7 @@ async fn load_auth_reads_access_token_from_env() {
.mount(&server)
.await;
Mock::given(method("POST"))
.and(path("/backend-api/v1/agent/agent-runtime-id/task/register"))
.and(path("/v1/agent/agent-runtime-id/task/register"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"task_id": "task-123",
})))
@@ -895,9 +900,8 @@ async fn load_auth_reads_access_token_from_env() {
.await;
let _access_token_guard = EnvVarGuard::set(CODEX_ACCESS_TOKEN_ENV_VAR, &agent_identity);
let chatgpt_base_url = format!("{}/backend-api", server.uri());
let _authapi_guard =
EnvVarGuard::set("CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL", &chatgpt_base_url);
let authapi_base_url = server.uri();
let chatgpt_base_url = format!("{authapi_base_url}/backend-api");
let auth = super::load_auth(
codex_home.path(),
/*enable_codex_api_key_env*/ false,
@@ -905,6 +909,7 @@ async fn load_auth_reads_access_token_from_env() {
/*forced_chatgpt_workspace_id*/ None,
Some(&chatgpt_base_url),
AuthKeyringBackendKind::Direct,
Some(&authapi_base_url),
)
.await
.expect("env auth should load")
@@ -914,7 +919,7 @@ async fn load_auth_reads_access_token_from_env() {
panic!("env auth should load as agent identity");
};
assert_eq!(agent_identity.record(), &expected_record);
assert_eq!(agent_identity.process_task_id(), "task-123");
assert_eq!(agent_identity.run_task_id(), "task-123");
assert!(
!get_auth_file(codex_home.path()).exists(),
"env auth should not write auth.json"
@@ -951,6 +956,7 @@ async fn load_auth_reads_personal_access_token_from_env() {
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::default(),
/*agent_identity_authapi_base_url*/ None,
)
.await
.expect("env auth should load")
@@ -1119,6 +1125,7 @@ async fn load_auth_keeps_codex_api_key_env_precedence() {
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::Direct,
/*agent_identity_authapi_base_url*/ None,
)
.await
.expect("env auth should load")
@@ -1226,6 +1233,7 @@ async fn enforce_login_restrictions_logs_out_for_personal_access_token_workspace
forced_login_method: None,
forced_chatgpt_workspace_id: Some(vec![WORKSPACE_ID_ALLOWED.to_string()]),
chatgpt_base_url: None,
agent_identity_authapi_base_url: None,
};
let err = super::enforce_login_restrictions(&config)
@@ -1317,16 +1325,15 @@ async fn enforce_login_restrictions_logs_out_for_agent_identity_workspace_mismat
.mount(&server)
.await;
Mock::given(method("POST"))
.and(path("/backend-api/v1/agent/agent-runtime-id/task/register"))
.and(path("/v1/agent/agent-runtime-id/task/register"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"task_id": "task-123",
})))
.expect(1)
.mount(&server)
.await;
let chatgpt_base_url = format!("{}/backend-api", server.uri());
let _authapi_guard =
EnvVarGuard::set("CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL", &chatgpt_base_url);
let authapi_base_url = server.uri();
let chatgpt_base_url = format!("{authapi_base_url}/backend-api");
save_auth(
codex_home.path(),
&AuthDotJson {
@@ -1350,6 +1357,7 @@ async fn enforce_login_restrictions_logs_out_for_agent_identity_workspace_mismat
forced_login_method: None,
forced_chatgpt_workspace_id: Some(vec![WORKSPACE_ID_ALLOWED.to_string()]),
chatgpt_base_url: Some(chatgpt_base_url),
agent_identity_authapi_base_url: Some(authapi_base_url),
};
let err = super::enforce_login_restrictions(&config)
@@ -1563,19 +1571,19 @@ async fn assert_agent_identity_plan_alias(
.mount(&server)
.await;
Mock::given(method("POST"))
.and(path("/backend-api/v1/agent/agent-runtime-id/task/register"))
.and(path("/v1/agent/agent-runtime-id/task/register"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"task_id": "task-123",
})))
.expect(1)
.mount(&server)
.await;
let chatgpt_base_url = format!("{}/backend-api", server.uri());
let _authapi_guard =
EnvVarGuard::set("CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL", &chatgpt_base_url);
let auth = CodexAuth::from_agent_identity_jwt(&jwt, Some(&chatgpt_base_url))
.await
.expect("agent identity auth");
let authapi_base_url = server.uri();
let chatgpt_base_url = format!("{authapi_base_url}/backend-api");
let auth =
CodexAuth::from_agent_identity_jwt(&jwt, Some(&chatgpt_base_url), Some(&authapi_base_url))
.await
.expect("agent identity auth");
pretty_assertions::assert_eq!(auth.account_plan_type(), Some(expected_plan_type));
server.verify().await;
@@ -1603,6 +1611,7 @@ async fn plan_type_maps_known_plan() {
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::Direct,
/*agent_identity_authapi_base_url*/ None,
)
.await
.expect("load auth")
@@ -1633,6 +1642,7 @@ async fn plan_type_maps_self_serve_business_usage_based_plan() {
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::Direct,
/*agent_identity_authapi_base_url*/ None,
)
.await
.expect("load auth")
@@ -1666,6 +1676,7 @@ async fn plan_type_maps_enterprise_cbp_usage_based_plan() {
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::Direct,
/*agent_identity_authapi_base_url*/ None,
)
.await
.expect("load auth")
@@ -1699,6 +1710,7 @@ async fn plan_type_maps_unknown_to_unknown() {
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::Direct,
/*agent_identity_authapi_base_url*/ None,
)
.await
.expect("load auth")
@@ -1729,6 +1741,7 @@ async fn missing_plan_type_maps_to_unknown() {
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::Direct,
/*agent_identity_authapi_base_url*/ None,
)
.await
.expect("load auth")
codex-rs/login/src/auth/manager.rs
+48 -10
View File
@@ -18,6 +18,7 @@ use std::sync::atomic::Ordering;
use tokio::sync::Semaphore;
use tokio::sync::watch;
use codex_agent_identity::ChatGptEnvironment;
use codex_agent_identity::decode_agent_identity_jwt;
use codex_agent_identity::fetch_agent_identity_jwks;
use codex_app_server_protocol::AuthMode;
@@ -104,7 +105,6 @@ const REFRESH_TOKEN_INVALIDATED_MESSAGE: &str = "Your access token could not be
const REFRESH_TOKEN_UNKNOWN_MESSAGE: &str =
"Your access token could not be refreshed. Please log out and sign in again.";
const REFRESH_TOKEN_ACCOUNT_MISMATCH_MESSAGE: &str = "Your access token could not be refreshed because you have since logged out or signed in to another account. Please sign in again.";
const DEFAULT_CHATGPT_BACKEND_BASE_URL: &str = "https://chatgpt.com/backend-api";
const REFRESH_TOKEN_URL: &str = "https://auth.openai.com/oauth/token";
pub(super) const REVOKE_TOKEN_URL: &str = "https://auth.openai.com/oauth/revoke";
pub const REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR: &str = "CODEX_REFRESH_TOKEN_URL_OVERRIDE";
@@ -218,6 +218,7 @@ impl CodexAuth {
auth_credentials_store_mode: AuthCredentialsStoreMode,
chatgpt_base_url: Option<&str>,
keyring_backend_kind: AuthKeyringBackendKind,
agent_identity_authapi_base_url: Option<&str>,
) -> std::io::Result<Self> {
let auth_mode = auth_dot_json.resolved_mode();
let client = create_client();
@@ -233,7 +234,12 @@ impl CodexAuth {
"agent identity auth is missing an agent identity token.",
));
};
return Self::from_agent_identity_jwt(&agent_identity, chatgpt_base_url).await;
return Self::from_agent_identity_jwt(
&agent_identity,
chatgpt_base_url,
agent_identity_authapi_base_url,
)
.await;
}
if auth_mode == ApiAuthMode::PersonalAccessToken {
let Some(personal_access_token) = auth_dot_json.personal_access_token.as_deref() else {
@@ -292,6 +298,7 @@ impl CodexAuth {
/*forced_chatgpt_workspace_id*/ None,
chatgpt_base_url,
keyring_backend_kind,
/*agent_identity_authapi_base_url*/ None,
)
.await
}
@@ -299,13 +306,19 @@ impl CodexAuth {
pub async fn from_agent_identity_jwt(
jwt: &str,
chatgpt_base_url: Option<&str>,
agent_identity_authapi_base_url_override: Option<&str>,
) -> std::io::Result<Self> {
let base_url = chatgpt_base_url
.unwrap_or(DEFAULT_CHATGPT_BACKEND_BASE_URL)
.unwrap_or(ChatGptEnvironment::default().chatgpt_base_url())
.trim_end_matches('/')
.to_string();
let resolved_authapi_base_url = resolve_agent_identity_authapi_base_url(
Some(&base_url),
agent_identity_authapi_base_url_override,
)?;
let record = verified_agent_identity_record(jwt, &base_url).await?;
Ok(Self::AgentIdentity(AgentIdentityAuth::load(record).await?))
let auth = AgentIdentityAuth::load(record, &resolved_authapi_base_url).await?;
Ok(Self::AgentIdentity(auth))
}
pub async fn from_personal_access_token(access_token: &str) -> std::io::Result<Self> {
@@ -580,6 +593,22 @@ async fn verified_agent_identity_record(
Ok(claims.into())
}
fn resolve_agent_identity_authapi_base_url(
chatgpt_base_url: Option<&str>,
agent_identity_authapi_base_url_override: Option<&str>,
) -> std::io::Result<String> {
if let Some(base_url) = agent_identity_authapi_base_url_override {
return Ok(base_url.trim_end_matches('/').to_string());
}
let environment = match chatgpt_base_url {
Some(chatgpt_base_url) => ChatGptEnvironment::from_chatgpt_base_url(chatgpt_base_url)
.map_err(std::io::Error::other)?,
None => ChatGptEnvironment::default(),
};
Ok(environment.agent_identity_authapi_base_url().to_string())
}
/// Delete the auth.json file inside `codex_home` if it exists. Returns `Ok(true)`
/// if a file was removed, `Ok(false)` if no auth file was present.
pub fn logout(
@@ -672,7 +701,7 @@ pub async fn login_with_access_token(
}
CodexAccessToken::AgentIdentityJwt(jwt) => {
let base_url = chatgpt_base_url
.unwrap_or(DEFAULT_CHATGPT_BACKEND_BASE_URL)
.unwrap_or(ChatGptEnvironment::default().chatgpt_base_url())
.trim_end_matches('/')
.to_string();
verified_agent_identity_record(jwt, &base_url).await?;
@@ -763,6 +792,7 @@ pub struct AuthConfig {
pub keyring_backend_kind: AuthKeyringBackendKind,
pub forced_login_method: Option<ForcedLoginMethod>,
pub chatgpt_base_url: Option<String>,
pub agent_identity_authapi_base_url: Option<String>,
pub forced_chatgpt_workspace_id: Option<Vec<String>>,
}
@@ -774,6 +804,7 @@ pub async fn enforce_login_restrictions(config: &AuthConfig) -> std::io::Result<
/*forced_chatgpt_workspace_id*/ None,
config.chatgpt_base_url.as_deref(),
config.keyring_backend_kind,
config.agent_identity_authapi_base_url.as_deref(),
)
.await?
else {
@@ -916,6 +947,7 @@ async fn load_auth(
forced_chatgpt_workspace_id: Option<&[String]>,
chatgpt_base_url: Option<&str>,
keyring_backend_kind: AuthKeyringBackendKind,
agent_identity_authapi_base_url: Option<&str>,
) -> std::io::Result<Option<CodexAuth>> {
// API key via env var takes precedence over any other auth method.
if enable_codex_api_key_env && let Some(api_key) = read_codex_api_key_from_env() {
@@ -936,6 +968,7 @@ async fn load_auth(
AuthCredentialsStoreMode::Ephemeral,
chatgpt_base_url,
keyring_backend_kind,
agent_identity_authapi_base_url,
)
.await?;
if let CodexAuth::PersonalAccessToken(auth) = &auth {
@@ -951,11 +984,13 @@ async fn load_auth(
ensure_personal_access_token_workspace_allowed(forced_chatgpt_workspace_id, &auth)?;
Ok(Some(CodexAuth::PersonalAccessToken(auth)))
}
CodexAccessToken::AgentIdentityJwt(jwt) => {
CodexAuth::from_agent_identity_jwt(jwt, chatgpt_base_url)
.await
.map(Some)
}
CodexAccessToken::AgentIdentityJwt(jwt) => CodexAuth::from_agent_identity_jwt(
jwt,
chatgpt_base_url,
agent_identity_authapi_base_url,
)
.await
.map(Some),
};
}
@@ -981,6 +1016,7 @@ async fn load_auth(
auth_credentials_store_mode,
chatgpt_base_url,
keyring_backend_kind,
agent_identity_authapi_base_url,
)
.await?;
if let CodexAuth::PersonalAccessToken(auth) = &auth {
@@ -1573,6 +1609,7 @@ impl AuthManager {
forced_chatgpt_workspace_id.as_deref(),
chatgpt_base_url.as_deref(),
keyring_backend_kind,
/*agent_identity_authapi_base_url*/ None,
)
.await
.ok()
@@ -1799,6 +1836,7 @@ impl AuthManager {
forced_chatgpt_workspace_id.as_deref(),
self.chatgpt_base_url.as_deref(),
self.keyring_backend_kind,
/*agent_identity_authapi_base_url*/ None,
)
.await
.ok()
codex-rs/model-provider/src/auth.rs
+1 -5
View File
@@ -1,7 +1,6 @@
use std::sync::Arc;
use codex_agent_identity::AgentIdentityKey;
use codex_agent_identity::AgentTaskAuthorizationTarget;
use codex_agent_identity::authorization_header_for_agent_task;
use codex_api::AuthProvider;
use codex_api::SharedAuthProvider;
@@ -30,10 +29,7 @@ impl AuthProvider for AgentIdentityAuthProvider {
agent_runtime_id: &record.agent_runtime_id,
private_key_pkcs8_base64: &record.agent_private_key,
},
AgentTaskAuthorizationTarget {
agent_runtime_id: &record.agent_runtime_id,
task_id: self.auth.process_task_id(),
},
self.auth.run_task_id(),
)
.map_err(std::io::Error::other);
codex-rs/tui/src/lib.rs
+1
View File
@@ -1161,6 +1161,7 @@ pub async fn run_main(
forced_login_method: config.forced_login_method,
forced_chatgpt_workspace_id: config.forced_chatgpt_workspace_id.clone(),
chatgpt_base_url: Some(config.chatgpt_base_url.clone()),
agent_identity_authapi_base_url: None,
})
.await
{