clone/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-05-16 09:12:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: add chatgpt agent identity opt-in

Browse Source
This commit is contained in:
adrian
2026-04-22 15:55:49 -07:00
parent 7f4a791c6e
commit 08f17a2f51
9 changed files with 621 additions and 195 deletions
Download Patch File Download Diff File Expand all files Collapse all files

89
codex-rs/agent-identity/src/lib.rs
View File

@@ -234,14 +234,14 @@ struct RegisterTaskResponse {
}
#[derive(Debug, Serialize)]
struct RegisterAgentRequest {
abom: AgentBillOfMaterials,
struct CreateCodexAgentIdentityRequest {
agent_public_key: String,
capabilities: Vec<String>,
ttl: Option<u64>,
name: String,
}
#[derive(Debug, Deserialize)]
struct RegisterAgentResponse {
struct CreateCodexAgentIdentityResponse {
agent_runtime_id: String,
}
@@ -398,13 +398,13 @@ pub async fn register_agent_identity(
account_id: &str,
is_fedramp_account: bool,
key_material: &GeneratedAgentKeyMaterial,
abom: AgentBillOfMaterials,
name: &str,
) -> Result<AgentRuntimeId> {
let url = agent_registration_url(chatgpt_base_url);
let request = RegisterAgentRequest {
abom,
let request = CreateCodexAgentIdentityRequest {
agent_public_key: key_material.public_key_ssh.clone(),
capabilities: Vec::new(),
ttl: None,
name: name.to_string(),
};
let mut request_builder = client
@@ -423,7 +423,7 @@ pub async fn register_agent_identity(
.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>()
.json::<CreateCodexAgentIdentityResponse>()
.await
.with_context(|| format!("failed to parse agent identity response from {url}"))?;
@@ -496,21 +496,45 @@ pub fn curve25519_secret_key_from_private_key_pkcs8_base64(
}
pub fn agent_registration_url(chatgpt_base_url: &str) -> String {
let trimmed = chatgpt_base_url.trim_end_matches('/');
format!("{trimmed}/v1/agent/register")
format!("{}/agent-identities", codex_api_base_url(chatgpt_base_url))
}
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")
let api_path = format!("/v1/agent/{agent_runtime_id}/task/register");
if matches!(
trimmed,
"https://chatgpt.com/backend-api" | "https://chat.openai.com/backend-api"
) {
return format!("https://auth.openai.com/api/accounts{api_path}");
}
if trimmed == "https://chatgpt-staging.com/backend-api" {
return format!("https://auth.api.openai.org/api/accounts{api_path}");
}
if matches!(
trimmed,
"https://auth.openai.com" | "https://auth.api.openai.org"
) {
return format!("{trimmed}/api/accounts{api_path}");
}
format!("{trimmed}{api_path}")
}
pub fn agent_identity_jwks_url(chatgpt_base_url: &str) -> String {
format!(
"{}/agent-identities/jwks",
codex_api_base_url(chatgpt_base_url)
)
}
fn codex_api_base_url(chatgpt_base_url: &str) -> String {
let trimmed = chatgpt_base_url.trim_end_matches('/');
if trimmed.contains("/backend-api") {
format!("{trimmed}/wham/agent-identities/jwks")
if trimmed.ends_with("/api/codex") || trimmed.ends_with("/backend-api/codex") {
trimmed.to_string()
} else if trimmed.ends_with("/backend-api") {
format!("{trimmed}/codex")
} else {
format!("{trimmed}/agent-identities/jwks")
format!("{trimmed}/api/codex")
}
}
@@ -1009,15 +1033,46 @@ J1bwkqKZTB5dHolX9A58e/xXnfZ5P8f3Z83+Izap3FwqQulk7b1WO1MQcHuVg2NN
);
}
#[test]
fn agent_registration_url_uses_codex_backend_api_base_url() {
assert_eq!(
agent_registration_url("https://chatgpt.com/backend-api"),
"https://chatgpt.com/backend-api/codex/agent-identities"
);
assert_eq!(
agent_registration_url("http://localhost:8080/api/codex/"),
"http://localhost:8080/api/codex/agent-identities"
);
assert_eq!(
agent_registration_url("http://localhost:8080"),
"http://localhost:8080/api/codex/agent-identities"
);
}
#[test]
fn agent_task_registration_url_uses_public_authapi_for_chatgpt_base_url() {
assert_eq!(
agent_task_registration_url("https://chatgpt.com/backend-api", "agent-runtime-id"),
"https://auth.openai.com/api/accounts/v1/agent/agent-runtime-id/task/register"
);
assert_eq!(
agent_task_registration_url(
"https://chatgpt-staging.com/backend-api",
"agent-runtime-id"
),
"https://auth.api.openai.org/api/accounts/v1/agent/agent-runtime-id/task/register"
);
}
#[test]
fn agent_identity_jwks_url_uses_backend_api_base_url() {
assert_eq!(
agent_identity_jwks_url("https://chatgpt.com/backend-api"),
"https://chatgpt.com/backend-api/wham/agent-identities/jwks"
"https://chatgpt.com/backend-api/codex/agent-identities/jwks"
);
assert_eq!(
agent_identity_jwks_url("https://chatgpt.com/backend-api/"),
"https://chatgpt.com/backend-api/wham/agent-identities/jwks"
"https://chatgpt.com/backend-api/codex/agent-identities/jwks"
);
}

56
codex-rs/cloud-requirements/src/lib.rs
View File

@@ -838,34 +838,13 @@ mod tests {
use serde_json::json;
use std::collections::BTreeMap;
use std::collections::VecDeque;
use std::ffi::OsString;
use std::future::pending;
use std::io::Read;
use std::io::Write;
use std::net::TcpListener;
use std::path::Path;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::thread;
use tempfile::TempDir;
use tempfile::tempdir;
struct EnvVarGuard {
key: &'static str,
original: Option<OsString>,
}
impl Drop for EnvVarGuard {
fn drop(&mut self) {
unsafe {
match &self.original {
Some(value) => std::env::set_var(self.key, value),
None => std::env::remove_var(self.key),
}
}
}
}
fn write_auth_json(codex_home: &Path, value: serde_json::Value) -> std::io::Result<()> {
std::fs::write(codex_home.join("auth.json"), serde_json::to_string(&value)?)?;
Ok(())
@@ -1221,25 +1200,6 @@ mod tests {
#[tokio::test]
async fn cloud_requirements_eligible_auth_allows_agent_identity_business_plan() {
let listener = TcpListener::bind("127.0.0.1:0").expect("bind task registration server");
let addr = listener
.local_addr()
.expect("task registration server addr");
let server = thread::spawn(move || {
let (mut stream, _) = listener.accept().expect("accept task registration request");
let mut request = [0; 4096];
let _ = stream
.read(&mut request)
.expect("read task registration request");
let body = r#"{"task_id":"task-123"}"#;
write!(
stream,
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
body.len(),
body
)
.expect("write task registration response");
});
let record = AgentIdentityAuthRecord {
agent_runtime_id: "agent-runtime-123".to_string(),
agent_private_key: "MC4CAQAwBQYDK2VwBCIEIDQg14jybCLydjHQwXeBzsDM7oB6BSAenodx6oCovQ/D"
@@ -1249,21 +1209,9 @@ mod tests {
email: "user@example.com".to_string(),
plan_type: PlanType::Business,
chatgpt_account_is_fedramp: false,
registered_at: None,
};
let authapi_base_url = format!("http://{addr}/backend-api");
let original_authapi_base_url = std::env::var_os("CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL");
unsafe {
std::env::set_var("CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL", &authapi_base_url);
}
let _authapi_guard = EnvVarGuard {
key: "CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL",
original: original_authapi_base_url,
};
let auth = AgentIdentityAuth::load(record)
.await
.map(CodexAuth::AgentIdentity)
.expect("agent identity auth");
server.join().expect("task registration server joined");
let auth = CodexAuth::AgentIdentity(AgentIdentityAuth::new(record));
assert!(cloud_requirements_eligible_auth(&auth));
}

137
codex-rs/login/src/auth/agent_identity.rs
View File

@@ -1,43 +1,62 @@
use std::sync::Arc;
use codex_agent_identity::AgentIdentityKey;
use codex_agent_identity::normalize_chatgpt_base_url;
use codex_agent_identity::register_agent_task;
use codex_protocol::account::PlanType as AccountPlanType;
use std::env;
use tokio::sync::OnceCell;
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";
const DEFAULT_CHATGPT_BACKEND_BASE_URL: &str = "https://chatgpt.com/backend-api";
#[derive(Clone, Debug)]
#[derive(Debug)]
pub struct AgentIdentityAuth {
record: AgentIdentityAuthRecord,
process_task_id: String,
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),
}
}
}
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(
&build_reqwest_client(),
&agent_identity_authapi_base_url,
key(&record),
)
.await
.map_err(std::io::Error::other)?;
Ok(Self {
pub fn new(record: AgentIdentityAuthRecord) -> Self {
Self {
record,
process_task_id,
})
process_task_id: Arc::new(OnceCell::new()),
}
}
pub fn record(&self) -> &AgentIdentityAuthRecord {
&self.record
}
pub fn process_task_id(&self) -> &str {
&self.process_task_id
pub fn process_task_id(&self) -> Option<&str> {
self.process_task_id.get().map(String::as_str)
}
pub async fn ensure_runtime(&self, chatgpt_base_url: Option<String>) -> std::io::Result<()> {
self.process_task_id
.get_or_try_init(|| async {
let base_url = normalize_chatgpt_base_url(
chatgpt_base_url
.as_deref()
.unwrap_or(DEFAULT_CHATGPT_BACKEND_BASE_URL),
);
register_agent_task(&build_reqwest_client(), &base_url, self.key())
.await
.map_err(std::io::Error::other)
})
.await
.map(|_| ())
}
pub fn account_id(&self) -> &str {
@@ -59,82 +78,10 @@ impl AgentIdentityAuth {
pub fn is_fedramp_account(&self) -> bool {
self.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<'_> {
AgentIdentityKey {
agent_runtime_id: &record.agent_runtime_id,
private_key_pkcs8_base64: &record.agent_private_key,
}
}
#[cfg(test)]
mod tests {
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 }
}
}
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 key(&self) -> AgentIdentityKey<'_> {
AgentIdentityKey {
agent_runtime_id: &self.record.agent_runtime_id,
private_key_pkcs8_base64: &self.record.agent_private_key,
}
}
}

118
codex-rs/login/src/auth/auth_tests.rs
View File

@@ -1,4 +1,5 @@
use super::*;
use crate::auth::storage::AgentIdentityStorage;
use crate::auth::storage::FileAuthStorage;
use crate::auth::storage::get_auth_file;
use crate::token_data::IdTokenInfo;
@@ -6,6 +7,7 @@ use codex_app_server_protocol::AuthMode;
use codex_protocol::account::PlanType as AccountPlanType;
use codex_protocol::auth::KnownPlan as InternalKnownPlan;
use codex_protocol::auth::PlanType as InternalPlanType;
use codex_protocol::protocol::SessionSource;
use base64::Engine;
use codex_protocol::config_types::ForcedLoginMethod;
@@ -19,6 +21,8 @@ use tempfile::tempdir;
use wiremock::Mock;
use wiremock::MockServer;
use wiremock::ResponseTemplate;
use wiremock::matchers::body_partial_json;
use wiremock::matchers::header;
use wiremock::matchers::method;
use wiremock::matchers::path;
@@ -92,7 +96,7 @@ async fn login_with_access_token_writes_only_token() {
signed_agent_identity_jwt(&record, json!(record.plan_type)).expect("signed agent identity");
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/backend-api/wham/agent-identities/jwks"))
.and(path("/backend-api/codex/agent-identities/jwks"))
.respond_with(ResponseTemplate::new(200).set_body_json(test_jwks_body()))
.expect(1)
.mount(&server)
@@ -114,7 +118,9 @@ async fn login_with_access_token_writes_only_token() {
.expect("auth.json should parse");
assert_eq!(auth.auth_mode, Some(AuthMode::AgentIdentity));
assert_eq!(
auth.agent_identity.as_deref(),
auth.agent_identity
.as_ref()
.and_then(AgentIdentityStorage::as_jwt),
Some(agent_identity.as_str())
);
assert!(auth.tokens.is_none(), "tokens should be cleared");
@@ -142,6 +148,101 @@ async fn login_with_access_token_rejects_invalid_jwt() {
);
}
#[tokio::test]
async fn chatgpt_auth_registers_agent_identity_when_enabled() -> anyhow::Result<()> {
let codex_home = tempdir()?;
write_auth_file(
AuthFileParams {
openai_api_key: None,
chatgpt_plan_type: Some("pro".to_string()),
chatgpt_account_id: Some("account-123".to_string()),
},
codex_home.path(),
)?;
let auth = super::load_auth(
codex_home.path(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*chatgpt_base_url*/ None,
)
.await?
.expect("auth should load");
assert!(
auth.agent_identity_auth(
AgentIdentityAuthPolicy::JwtOnly,
/*chatgpt_base_url*/ None,
/*forced_chatgpt_workspace_id*/ None,
SessionSource::Cli,
)
.await?
.is_none()
);
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/api/codex/agent-identities"))
.and(header("authorization", "Bearer test-access-token"))
.and(header("chatgpt-account-id", "account-123"))
.and(body_partial_json(json!({
"name": "Codex CLI",
"ttl": null,
})))
.respond_with(ResponseTemplate::new(/*s*/ 200).set_body_json(json!({
"agent_runtime_id": "agent-runtime-123",
})))
.expect(/*r*/ 1)
.mount(&server)
.await;
let agent_auth = auth
.agent_identity_auth(
AgentIdentityAuthPolicy::JwtOrChatgpt,
Some(server.uri()),
/*forced_chatgpt_workspace_id*/ None,
SessionSource::Cli,
)
.await?
.expect("agent identity should register");
let reused = auth
.agent_identity_auth(
AgentIdentityAuthPolicy::JwtOrChatgpt,
Some(server.uri()),
/*forced_chatgpt_workspace_id*/ None,
SessionSource::Cli,
)
.await?
.expect("agent identity should be reused");
Mock::given(method("POST"))
.and(path("/v1/agent/agent-runtime-123/task/register"))
.respond_with(ResponseTemplate::new(/*s*/ 200).set_body_json(json!({
"task_id": "task-123",
})))
.expect(/*r*/ 1)
.mount(&server)
.await;
agent_auth.ensure_runtime(Some(server.uri())).await?;
reused.ensure_runtime(Some(server.uri())).await?;
assert_eq!(
agent_auth.record().agent_runtime_id,
reused.record().agent_runtime_id
);
assert_eq!(agent_auth.process_task_id(), Some("task-123"));
assert_eq!(reused.process_task_id(), Some("task-123"));
assert_eq!(agent_auth.record().agent_runtime_id, "agent-runtime-123");
assert_eq!(agent_auth.record().account_id, "account-123");
assert_eq!(agent_auth.record().chatgpt_user_id, "user-12345");
assert_eq!(
auth.get_agent_identity("account-123")
.expect("identity should persist")
.agent_runtime_id,
"agent-runtime-123"
);
Ok(())
}
#[tokio::test]
async fn login_with_access_token_rejects_unsigned_jwt() {
let dir = tempdir().unwrap();
@@ -149,7 +250,7 @@ async fn login_with_access_token_rejects_unsigned_jwt() {
let agent_identity = fake_agent_identity_jwt(&record).expect("fake agent identity");
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/backend-api/wham/agent-identities/jwks"))
.and(path("/backend-api/codex/agent-identities/jwks"))
.respond_with(ResponseTemplate::new(200).set_body_json(test_jwks_body()))
.expect(1)
.mount(&server)
@@ -718,7 +819,7 @@ async fn load_auth_reads_access_token_from_env() {
.expect("signed agent identity");
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/backend-api/wham/agent-identities/jwks"))
.and(path("/backend-api/codex/agent-identities/jwks"))
.respond_with(ResponseTemplate::new(200).set_body_json(test_jwks_body()))
.expect(1)
.mount(&server)
@@ -734,8 +835,6 @@ async fn load_auth_reads_access_token_from_env() {
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 auth = super::load_auth(
codex_home.path(),
/*enable_codex_api_key_env*/ false,
@@ -750,7 +849,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.process_task_id(), Some("task-123"));
assert!(
!get_auth_file(codex_home.path()).exists(),
"env auth should not write auth.json"
@@ -927,6 +1026,7 @@ fn agent_identity_record(account_id: &str) -> AgentIdentityAuthRecord {
email: "user@example.com".to_string(),
plan_type: AccountPlanType::Pro,
chatgpt_account_is_fedramp: false,
registered_at: None,
}
}
@@ -1045,7 +1145,7 @@ async fn assert_agent_identity_plan_alias(
let jwt = signed_agent_identity_jwt(&record, plan_type).expect("agent identity jwt");
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/backend-api/wham/agent-identities/jwks"))
.and(path("/backend-api/codex/agent-identities/jwks"))
.respond_with(ResponseTemplate::new(200).set_body_json(test_jwks_body()))
.expect(1)
.mount(&server)
@@ -1059,8 +1159,6 @@ async fn assert_agent_identity_plan_alias(
.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");

374
codex-rs/login/src/auth/manager.rs
View File

@@ -18,6 +18,10 @@ use tokio::sync::Semaphore;
use codex_agent_identity::decode_agent_identity_jwt;
use codex_agent_identity::fetch_agent_identity_jwks;
use codex_agent_identity::generate_agent_key_material;
use codex_agent_identity::normalize_chatgpt_base_url;
use codex_agent_identity::public_key_ssh_from_private_key_pkcs8_base64;
use codex_agent_identity::register_agent_identity;
use codex_app_server_protocol::AuthMode;
use codex_app_server_protocol::AuthMode as ApiAuthMode;
use codex_protocol::config_types::ForcedLoginMethod;
@@ -27,6 +31,7 @@ use super::external_bearer::BearerTokenRefresher;
use super::revoke::revoke_auth_tokens;
pub use crate::auth::agent_identity::AgentIdentityAuth;
pub use crate::auth::storage::AgentIdentityAuthRecord;
use crate::auth::storage::AgentIdentityStorage;
pub use crate::auth::storage::AuthDotJson;
use crate::auth::storage::AuthStorageBackend;
use crate::auth::storage::create_auth_storage;
@@ -42,6 +47,7 @@ use codex_protocol::account::PlanType as AccountPlanType;
use codex_protocol::auth::PlanType as InternalPlanType;
use codex_protocol::auth::RefreshTokenFailedError;
use codex_protocol::auth::RefreshTokenFailedReason;
use codex_protocol::protocol::SessionSource;
use serde_json::Value;
use thiserror::Error;
@@ -54,6 +60,15 @@ pub enum CodexAuth {
AgentIdentity(AgentIdentityAuth),
}
/// Policy for resolving Agent Identity auth from a broader Codex auth snapshot.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AgentIdentityAuthPolicy {
/// Only use an existing Agent Identity JWT/runtime auth.
JwtOnly,
/// Use an Agent Identity JWT/runtime auth, or register one from ChatGPT auth.
JwtOrChatgpt,
}
impl PartialEq for CodexAuth {
fn eq(&self, other: &Self) -> bool {
self.api_auth_mode() == other.api_auth_mode()
@@ -79,10 +94,38 @@ pub struct ChatgptAuthTokens {
#[derive(Debug, Clone)]
struct ChatgptAuthState {
auth_dot_json: Arc<Mutex<Option<AuthDotJson>>>,
agent_identity_auth: Arc<Mutex<Option<AgentIdentityAuth>>>,
client: CodexHttpClient,
}
impl ChatgptAuthState {
fn new(auth_dot_json: AuthDotJson) -> Self {
let agent_identity_auth = auth_dot_json
.agent_identity
.as_ref()
.and_then(AgentIdentityStorage::as_record)
.cloned()
.map(AgentIdentityAuth::new);
Self {
auth_dot_json: Arc::new(Mutex::new(Some(auth_dot_json))),
agent_identity_auth: Arc::new(Mutex::new(agent_identity_auth)),
client: create_client(),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct ChatgptAgentIdentityBinding {
account_id: String,
chatgpt_user_id: String,
email: String,
plan_type: AccountPlanType,
chatgpt_account_is_fedramp: bool,
access_token: String,
}
const TOKEN_REFRESH_INTERVAL: i64 = 8;
const DEFAULT_CHATGPT_BACKEND_BASE_URL: &str = "https://chatgpt.com/backend-api";
const REFRESH_TOKEN_EXPIRED_MESSAGE: &str = "Your access token could not be refreshed because your refresh token has expired. Please log out and sign in again.";
const REFRESH_TOKEN_REUSED_MESSAGE: &str = "Your access token could not be refreshed because your refresh token was already used. Please log out and sign in again.";
@@ -90,7 +133,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";
@@ -203,7 +245,6 @@ impl CodexAuth {
chatgpt_base_url: Option<&str>,
) -> std::io::Result<Self> {
let auth_mode = auth_dot_json.resolved_mode();
let client = create_client();
if auth_mode == ApiAuthMode::ApiKey {
let Some(api_key) = auth_dot_json.openai_api_key.as_deref() else {
return Err(std::io::Error::other("API key auth is missing a key."));
@@ -211,19 +252,20 @@ impl CodexAuth {
return Ok(Self::from_api_key(api_key));
}
if auth_mode == ApiAuthMode::AgentIdentity {
let Some(agent_identity) = auth_dot_json.agent_identity else {
let Some(agent_identity) = auth_dot_json
.agent_identity
.as_ref()
.and_then(AgentIdentityStorage::as_jwt)
else {
return Err(std::io::Error::other(
"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).await;
}
let storage_mode = auth_dot_json.storage_mode(auth_credentials_store_mode);
let state = ChatgptAuthState {
auth_dot_json: Arc::new(Mutex::new(Some(auth_dot_json))),
client,
};
let state = ChatgptAuthState::new(auth_dot_json);
match auth_mode {
ApiAuthMode::Chatgpt => {
@@ -261,7 +303,9 @@ impl CodexAuth {
.trim_end_matches('/')
.to_string();
let record = verified_agent_identity_record(jwt, &base_url).await?;
Ok(Self::AgentIdentity(AgentIdentityAuth::load(record).await?))
let auth = AgentIdentityAuth::new(record);
auth.ensure_runtime(Some(base_url)).await?;
Ok(Self::AgentIdentity(auth))
}
pub fn auth_mode(&self) -> AuthMode {
@@ -408,6 +452,196 @@ impl CodexAuth {
self.get_current_auth_json().and_then(|t| t.tokens)
}
pub fn get_agent_identity(&self, account_id: &str) -> Option<AgentIdentityAuthRecord> {
self.get_current_auth_json()
.and_then(|auth| auth.agent_identity)
.and_then(|identity| identity.as_record().cloned())
.filter(|identity| identity.account_id == account_id)
}
pub fn set_agent_identity(&self, record: AgentIdentityAuthRecord) -> std::io::Result<()> {
let agent_identity_auth = self.agent_identity_auth_for_record(record.clone())?;
match self {
Self::Chatgpt(auth) => auth
.update_auth_json(|auth_dot_json| {
auth_dot_json.agent_identity = Some(AgentIdentityStorage::Record(record));
true
})
.map(|_| ()),
Self::ChatgptAuthTokens(auth) => auth.update_auth_json_in_memory(|auth_dot_json| {
auth_dot_json.agent_identity = Some(AgentIdentityStorage::Record(record));
}),
Self::ApiKey(_) | Self::AgentIdentity(_) => Ok(()),
}?;
self.set_cached_agent_identity_auth(Some(agent_identity_auth))
}
pub fn remove_agent_identity(&self) -> std::io::Result<bool> {
let removed = match self {
Self::Chatgpt(auth) => {
auth.update_auth_json(|auth_dot_json| auth_dot_json.agent_identity.take().is_some())
}
Self::ChatgptAuthTokens(auth) => {
let mut removed = false;
auth.update_auth_json_in_memory(|auth_dot_json| {
removed = auth_dot_json.agent_identity.take().is_some();
})?;
Ok(removed)
}
Self::ApiKey(_) | Self::AgentIdentity(_) => Ok(false),
}?;
if removed {
self.set_cached_agent_identity_auth(/*auth*/ None)?;
}
Ok(removed)
}
fn cached_agent_identity_auth(
&self,
binding: &ChatgptAgentIdentityBinding,
) -> Option<AgentIdentityAuth> {
let auth = self.cached_agent_identity_auth_value()?;
if agent_identity_record_matches_binding(auth.record(), binding)
&& public_key_ssh_from_private_key_pkcs8_base64(&auth.record().agent_private_key)
.is_ok()
{
Some(auth)
} else {
None
}
}
fn agent_identity_auth_for_record(
&self,
record: AgentIdentityAuthRecord,
) -> std::io::Result<AgentIdentityAuth> {
if let Some(auth) = self.cached_agent_identity_auth_value()
&& auth.record() == &record
{
return Ok(auth);
}
Ok(AgentIdentityAuth::new(record))
}
fn cached_agent_identity_auth_value(&self) -> Option<AgentIdentityAuth> {
let state = self.chatgpt_state()?;
let auth = state.agent_identity_auth.lock().ok()?;
auth.clone()
}
fn set_cached_agent_identity_auth(
&self,
auth: Option<AgentIdentityAuth>,
) -> std::io::Result<()> {
let Some(state) = self.chatgpt_state() else {
return Ok(());
};
let mut cached = state
.agent_identity_auth
.lock()
.map_err(|_| std::io::Error::other("failed to lock agent identity cache"))?;
*cached = auth;
Ok(())
}
fn chatgpt_state(&self) -> Option<&ChatgptAuthState> {
match self {
Self::Chatgpt(auth) => Some(&auth.state),
Self::ChatgptAuthTokens(auth) => Some(&auth.state),
Self::ApiKey(_) | Self::AgentIdentity(_) => None,
}
}
pub async fn agent_identity_auth(
&self,
policy: AgentIdentityAuthPolicy,
chatgpt_base_url: Option<String>,
forced_chatgpt_workspace_id: Option<String>,
session_source: SessionSource,
) -> std::io::Result<Option<AgentIdentityAuth>> {
match self {
Self::AgentIdentity(auth) => Ok(Some(auth.clone())),
Self::ApiKey(_) => Ok(None),
Self::Chatgpt(_) | Self::ChatgptAuthTokens(_) => {
if policy == AgentIdentityAuthPolicy::JwtOnly {
return Ok(None);
}
self.ensure_chatgpt_agent_identity(
chatgpt_base_url,
forced_chatgpt_workspace_id,
session_source,
)
.await
.map(Some)
}
}
}
async fn ensure_chatgpt_agent_identity(
&self,
chatgpt_base_url: Option<String>,
forced_chatgpt_workspace_id: Option<String>,
session_source: SessionSource,
) -> std::io::Result<AgentIdentityAuth> {
let binding = ChatgptAgentIdentityBinding::from_auth(self, forced_chatgpt_workspace_id)
.ok_or_else(|| std::io::Error::other("ChatGPT auth is unavailable"))?;
if let Some(auth) = self.cached_agent_identity_auth(&binding) {
return Ok(auth);
}
if let Some(record) = self.get_agent_identity(&binding.account_id)
&& agent_identity_record_matches_binding(&record, &binding)
&& public_key_ssh_from_private_key_pkcs8_base64(&record.agent_private_key).is_ok()
{
let auth = self.agent_identity_auth_for_record(record)?;
self.set_cached_agent_identity_auth(Some(auth.clone()))?;
return Ok(auth);
}
let key_material = generate_agent_key_material().map_err(std::io::Error::other)?;
let base_url = normalize_chatgpt_base_url(
chatgpt_base_url
.as_deref()
.unwrap_or(DEFAULT_CHATGPT_BACKEND_BASE_URL),
);
let agent_name = match session_source {
SessionSource::VSCode => "Codex App",
SessionSource::Cli
| SessionSource::Exec
| SessionSource::Mcp
| SessionSource::Custom(_)
| SessionSource::Internal(_)
| SessionSource::SubAgent(_)
| SessionSource::Unknown => "Codex CLI",
};
let runtime_id = register_agent_identity(
&build_reqwest_client(),
&base_url,
&binding.access_token,
&binding.account_id,
binding.chatgpt_account_is_fedramp,
&key_material,
agent_name,
)
.await
.map_err(std::io::Error::other)?;
let record = AgentIdentityAuthRecord {
agent_runtime_id: runtime_id.into_string(),
agent_private_key: key_material.private_key_pkcs8_base64,
account_id: binding.account_id,
chatgpt_user_id: binding.chatgpt_user_id,
email: binding.email,
plan_type: binding.plan_type,
chatgpt_account_is_fedramp: binding.chatgpt_account_is_fedramp,
registered_at: Some(
Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Secs, /*use_z*/ true),
),
};
self.set_agent_identity(record.clone())?;
self.agent_identity_auth_for_record(record)
}
/// Consider this private to integration tests.
pub fn create_dummy_chatgpt_auth_for_testing() -> Self {
let auth_dot_json = AuthDotJson {
@@ -423,11 +657,7 @@ impl CodexAuth {
agent_identity: None,
};
let client = create_client();
let state = ChatgptAuthState {
auth_dot_json: Arc::new(Mutex::new(Some(auth_dot_json))),
client,
};
let state = ChatgptAuthState::new(auth_dot_json);
let dummy_auth_id = NEXT_DUMMY_AUTH_ID.fetch_add(1, Ordering::Relaxed);
let storage = create_auth_storage(
PathBuf::from(format!("dummy-chatgpt-auth-{dummy_auth_id}")),
@@ -443,6 +673,44 @@ impl CodexAuth {
}
}
impl ChatgptAgentIdentityBinding {
fn from_auth(auth: &CodexAuth, forced_workspace_id: Option<String>) -> Option<Self> {
if !auth.is_chatgpt_auth() {
return None;
}
let token_data = auth.get_token_data().ok()?;
let account_id = forced_workspace_id
.filter(|value| !value.is_empty())
.or(token_data
.account_id
.clone()
.filter(|value| !value.is_empty()))
.or(token_data.id_token.chatgpt_account_id.clone())?;
let chatgpt_user_id = token_data
.id_token
.chatgpt_user_id
.clone()
.filter(|value| !value.is_empty())?;
Some(Self {
account_id,
chatgpt_user_id,
email: token_data.id_token.email.clone().unwrap_or_default(),
plan_type: auth.account_plan_type().unwrap_or(AccountPlanType::Unknown),
chatgpt_account_is_fedramp: auth.is_fedramp_account(),
access_token: token_data.access_token,
})
}
}
fn agent_identity_record_matches_binding(
record: &AgentIdentityAuthRecord,
binding: &ChatgptAgentIdentityBinding,
) -> bool {
record.account_id == binding.account_id && record.chatgpt_user_id == binding.chatgpt_user_id
}
impl ChatgptAuth {
fn current_auth_json(&self) -> Option<AuthDotJson> {
#[expect(clippy::unwrap_used)]
@@ -460,6 +728,45 @@ impl ChatgptAuth {
fn client(&self) -> &CodexHttpClient {
&self.state.client
}
fn update_auth_json(
&self,
update: impl FnOnce(&mut AuthDotJson) -> bool,
) -> std::io::Result<bool> {
let mut guard = self
.state
.auth_dot_json
.lock()
.map_err(|_| std::io::Error::other("failed to lock auth state"))?;
let mut auth = guard
.clone()
.ok_or_else(|| std::io::Error::other("auth data is not available"))?;
let changed = update(&mut auth);
if changed {
self.storage.save(&auth)?;
*guard = Some(auth);
}
Ok(changed)
}
}
impl ChatgptAuthTokens {
fn update_auth_json_in_memory(
&self,
update: impl FnOnce(&mut AuthDotJson),
) -> std::io::Result<()> {
let mut guard = self
.state
.auth_dot_json
.lock()
.map_err(|_| std::io::Error::other("failed to lock auth state"))?;
let mut auth = guard
.clone()
.ok_or_else(|| std::io::Error::other("auth data is not available"))?;
update(&mut auth);
*guard = Some(auth);
Ok(())
}
}
pub const OPENAI_API_KEY_ENV_VAR: &str = "OPENAI_API_KEY";
@@ -558,7 +865,7 @@ pub async fn login_with_access_token(
openai_api_key: None,
tokens: None,
last_refresh: None,
agent_identity: Some(access_token.to_string()),
agent_identity: Some(AgentIdentityStorage::Jwt(access_token.to_string())),
};
save_auth(codex_home, &auth_dot_json, auth_credentials_store_mode)
}
@@ -1250,6 +1557,7 @@ pub struct AuthManager {
forced_chatgpt_workspace_id: RwLock<Option<String>>,
chatgpt_base_url: Option<String>,
refresh_lock: Semaphore,
agent_identity_lock: Semaphore,
external_auth: RwLock<Option<Arc<dyn ExternalAuth>>>,
}
@@ -1324,6 +1632,7 @@ impl AuthManager {
forced_chatgpt_workspace_id: RwLock::new(None),
chatgpt_base_url,
refresh_lock: Semaphore::new(/*permits*/ 1),
agent_identity_lock: Semaphore::new(/*permits*/ 1),
external_auth: RwLock::new(None),
}
}
@@ -1343,6 +1652,7 @@ impl AuthManager {
forced_chatgpt_workspace_id: RwLock::new(None),
chatgpt_base_url: None,
refresh_lock: Semaphore::new(/*permits*/ 1),
agent_identity_lock: Semaphore::new(/*permits*/ 1),
external_auth: RwLock::new(None),
})
}
@@ -1361,6 +1671,7 @@ impl AuthManager {
forced_chatgpt_workspace_id: RwLock::new(None),
chatgpt_base_url: None,
refresh_lock: Semaphore::new(/*permits*/ 1),
agent_identity_lock: Semaphore::new(/*permits*/ 1),
external_auth: RwLock::new(None),
})
}
@@ -1377,6 +1688,7 @@ impl AuthManager {
forced_chatgpt_workspace_id: RwLock::new(None),
chatgpt_base_url: None,
refresh_lock: Semaphore::new(/*permits*/ 1),
agent_identity_lock: Semaphore::new(/*permits*/ 1),
external_auth: RwLock::new(Some(
Arc::new(BearerTokenRefresher::new(config)) as Arc<dyn ExternalAuth>
)),
@@ -1416,6 +1728,38 @@ impl AuthManager {
self.auth_cached()
}
pub async fn agent_identity_auth(
&self,
policy: AgentIdentityAuthPolicy,
session_source: SessionSource,
) -> std::io::Result<Option<AgentIdentityAuth>> {
let Some(auth) = self.auth().await else {
return Ok(None);
};
if policy == AgentIdentityAuthPolicy::JwtOrChatgpt && auth.is_chatgpt_auth() {
let _permit = self
.agent_identity_lock
.acquire()
.await
.map_err(std::io::Error::other)?;
return auth
.agent_identity_auth(
policy,
self.chatgpt_base_url.clone(),
self.forced_chatgpt_workspace_id(),
session_source,
)
.await;
}
auth.agent_identity_auth(
policy,
self.chatgpt_base_url.clone(),
self.forced_chatgpt_workspace_id(),
session_source,
)
.await
}
/// Force a reload of the auth information from auth.json. Returns
/// whether the auth value changed.
pub async fn reload(&self) -> bool {

28
codex-rs/login/src/auth/storage.rs
View File

@@ -44,7 +44,30 @@ pub struct AuthDotJson {
pub last_refresh: Option<DateTime<Utc>>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub agent_identity: Option<String>,
pub agent_identity: Option<AgentIdentityStorage>,
}
#[derive(Deserialize, Serialize, Clone, Debug, PartialEq, Eq)]
#[serde(untagged)]
pub enum AgentIdentityStorage {
Jwt(String),
Record(AgentIdentityAuthRecord),
}
impl AgentIdentityStorage {
pub(crate) fn as_jwt(&self) -> Option<&str> {
match self {
Self::Jwt(jwt) => Some(jwt),
Self::Record(_) => None,
}
}
pub(crate) fn as_record(&self) -> Option<&AgentIdentityAuthRecord> {
match self {
Self::Jwt(_) => None,
Self::Record(record) => Some(record),
}
}
}
#[derive(Deserialize, Serialize, Clone, Debug, PartialEq, Eq)]
@@ -56,6 +79,8 @@ pub struct AgentIdentityAuthRecord {
pub email: String,
pub plan_type: AccountPlanType,
pub chatgpt_account_is_fedramp: bool,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub registered_at: Option<String>,
}
impl AgentIdentityAuthRecord {
@@ -77,6 +102,7 @@ impl From<AgentIdentityJwtClaims> for AgentIdentityAuthRecord {
email: claims.email,
plan_type: claims.plan_type.into(),
chatgpt_account_is_fedramp: claims.chatgpt_account_is_fedramp,
registered_at: None,
}
}
}

8
codex-rs/login/src/auth/storage_tests.rs
View File

@@ -72,7 +72,7 @@ async fn file_storage_round_trips_agent_identity_auth() -> anyhow::Result<()> {
openai_api_key: None,
tokens: None,
last_refresh: None,
agent_identity: Some(agent_identity),
agent_identity: Some(AgentIdentityStorage::Jwt(agent_identity)),
};
storage.save(&auth_dot_json)?;
@@ -107,7 +107,11 @@ async fn file_storage_loads_agent_identity_as_jwt() -> anyhow::Result<()> {
let loaded = storage.load()?;
assert_eq!(
loaded.expect("auth should load").agent_identity.as_deref(),
loaded
.expect("auth should load")
.agent_identity
.as_ref()
.and_then(AgentIdentityStorage::as_jwt),
Some(agent_identity_jwt.as_str())
);
Ok(())

1
codex-rs/login/src/lib.rs
View File

@@ -17,6 +17,7 @@ pub use server::ServerOptions;
pub use server::ShutdownHandle;
pub use server::run_login_server;
pub use auth::AgentIdentityAuthPolicy;
pub use auth::AuthConfig;
pub use auth::AuthDotJson;
pub use auth::AuthManager;

5
codex-rs/model-provider/src/auth.rs
View File

@@ -21,6 +21,9 @@ struct AgentIdentityAuthProvider {
impl AuthProvider for AgentIdentityAuthProvider {
fn add_auth_headers(&self, headers: &mut HeaderMap) {
let record = self.auth.record();
let Some(task_id) = self.auth.process_task_id() else {
return;
};
let header_value = authorization_header_for_agent_task(
AgentIdentityKey {
agent_runtime_id: &record.agent_runtime_id,
@@ -28,7 +31,7 @@ impl AuthProvider for AgentIdentityAuthProvider {
},
AgentTaskAuthorizationTarget {
agent_runtime_id: &record.agent_runtime_id,
task_id: self.auth.process_task_id(),
task_id,
},
)
.map_err(std::io::Error::other);