feat: verify agent identity JWTs with JWKS (#19764)

This commit is contained in:
efrazer-oai
2026-04-28 09:56:20 -07:00
committed by GitHub
Unverified
parent 6138063656
commit f6797c3ac6
13 changed files with 517 additions and 126 deletions
+1
View File
@@ -2850,6 +2850,7 @@ dependencies = [
"codex-terminal-detection",
"codex-utils-template",
"core_test_support",
"jsonwebtoken",
"keyring",
"once_cell",
"os_info",
+186 -72
View File
@@ -19,6 +19,9 @@ use ed25519_dalek::pkcs8::EncodePrivateKey;
use jsonwebtoken::Algorithm;
use jsonwebtoken::DecodingKey;
use jsonwebtoken::Validation;
use jsonwebtoken::decode;
use jsonwebtoken::decode_header;
use jsonwebtoken::jwk::JwkSet;
use rand::TryRngCore;
use rand::rngs::OsRng;
use serde::Deserialize;
@@ -28,6 +31,9 @@ use sha2::Digest as _;
use sha2::Sha512;
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";
/// Stored key material for a registered agent identity.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
@@ -58,6 +64,10 @@ pub struct GeneratedAgentKeyMaterial {
/// Claims carried by an Agent Identity JWT.
#[derive(Clone, Debug, Deserialize, PartialEq, Eq)]
pub struct AgentIdentityJwtClaims {
pub iss: String,
pub aud: String,
pub iat: usize,
pub exp: usize,
pub agent_runtime_id: String,
pub agent_private_key: String,
pub account_id: String,
@@ -115,27 +125,49 @@ pub fn authorization_header_for_agent_task(
Ok(format!("AgentAssertion {serialized_assertion}"))
}
pub async fn fetch_agent_identity_jwks(
client: &reqwest::Client,
chatgpt_base_url: &str,
) -> Result<JwkSet> {
let response = client
.get(agent_identity_jwks_url(chatgpt_base_url))
.timeout(AGENT_IDENTITY_JWKS_TIMEOUT)
.send()
.await
.context("failed to request agent identity JWKS")?
.error_for_status()
.context("agent identity JWKS endpoint returned an error")?;
response
.json()
.await
.context("failed to decode agent identity JWKS")
}
pub fn decode_agent_identity_jwt(
jwt: &str,
public_key_base64: Option<&str>,
jwks: Option<&JwkSet>,
) -> Result<AgentIdentityJwtClaims> {
let Some(public_key_base64) = public_key_base64 else {
let Some(jwks) = jwks else {
return decode_agent_identity_jwt_payload(jwt);
};
let mut validation = Validation::new(Algorithm::EdDSA);
validation.required_spec_claims.clear();
validation.validate_exp = false;
validation.validate_aud = false;
let public_key = BASE64_STANDARD
.decode(public_key_base64)
.context("agent identity JWT public key is not valid base64")?;
let decoding_key = DecodingKey::from_ed_der(&public_key);
jsonwebtoken::decode::<AgentIdentityJwtClaims>(jwt, &decoding_key, &validation)
let header = decode_header(jwt).context("failed to decode agent identity JWT header")?;
let kid = header
.kid
.context("agent identity JWT header does not include a kid")?;
let jwk = jwks
.find(&kid)
.with_context(|| format!("agent identity JWT kid {kid} is not trusted"))?;
let decoding_key = DecodingKey::from_jwk(jwk).context("failed to build JWT decoding key")?;
let mut validation = Validation::new(Algorithm::RS256);
validation.set_audience(&[AGENT_IDENTITY_JWT_AUDIENCE]);
validation.set_issuer(&[AGENT_IDENTITY_JWT_ISSUER]);
validation.required_spec_claims.insert("iss".to_string());
validation.required_spec_claims.insert("aud".to_string());
decode::<AgentIdentityJwtClaims>(jwt, &decoding_key, &validation)
.map(|data| data.claims)
.context("failed to decode agent identity JWT")
.context("failed to verify agent identity JWT")
}
fn decode_agent_identity_jwt_payload<T: DeserializeOwned>(jwt: &str) -> Result<T> {
@@ -279,6 +311,15 @@ pub fn agent_identity_biscuit_url(chatgpt_base_url: &str) -> String {
format!("{trimmed}/authenticate_app_v2")
}
pub fn agent_identity_jwks_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")
} else {
format!("{trimmed}/agent-identities/jwks")
}
}
pub fn agent_identity_request_id() -> Result<String> {
let mut request_id_bytes = [0u8; 16];
OsRng
@@ -290,29 +331,6 @@ pub fn agent_identity_request_id() -> Result<String> {
))
}
pub fn normalize_chatgpt_base_url(chatgpt_base_url: &str) -> String {
let mut base_url = chatgpt_base_url.trim_end_matches('/').to_string();
for suffix in [
"/wham/remote/control/server/enroll",
"/wham/remote/control/server",
] {
if let Some(stripped) = base_url.strip_suffix(suffix) {
base_url = stripped.to_string();
break;
}
}
if let Some(stripped) = base_url.strip_suffix("/codex") {
base_url = stripped.to_string();
}
if (base_url.starts_with("https://chatgpt.com")
|| base_url.starts_with("https://chat.openai.com"))
&& !base_url.contains("/backend-api")
{
base_url = format!("{base_url}/backend-api");
}
base_url
}
pub fn build_abom(session_source: SessionSource) -> AgentBillOfMaterials {
AgentBillOfMaterials {
agent_version: env!("CARGO_PKG_VERSION").to_string(),
@@ -472,6 +490,10 @@ mod tests {
#[test]
fn decode_agent_identity_jwt_reads_claims() {
let jwt = jwt_with_payload(serde_json::json!({
"iss": AGENT_IDENTITY_JWT_ISSUER,
"aud": AGENT_IDENTITY_JWT_AUDIENCE,
"iat": 1_700_000_000usize,
"exp": 4_000_000_000usize,
"agent_runtime_id": "agent-runtime-id",
"agent_private_key": "private-key",
"account_id": "account-id",
@@ -481,12 +503,15 @@ mod tests {
"chatgpt_account_is_fedramp": false,
}));
let claims =
decode_agent_identity_jwt(&jwt, /*public_key_base64*/ None).expect("JWT should decode");
let claims = decode_agent_identity_jwt(&jwt, /*jwks*/ None).expect("JWT should decode");
assert_eq!(
claims,
AgentIdentityJwtClaims {
iss: AGENT_IDENTITY_JWT_ISSUER.to_string(),
aud: AGENT_IDENTITY_JWT_AUDIENCE.to_string(),
iat: 1_700_000_000,
exp: 4_000_000_000,
agent_runtime_id: "agent-runtime-id".to_string(),
agent_private_key: "private-key".to_string(),
account_id: "account-id".to_string(),
@@ -499,15 +524,13 @@ mod tests {
}
#[test]
fn decode_agent_identity_jwt_verifies_when_public_key_is_present() {
let mut secret_key_bytes = [0u8; 32];
secret_key_bytes[0] = 1;
let signing_key = SigningKey::from_bytes(&secret_key_bytes);
let private_key_pkcs8 = signing_key
.to_pkcs8_der()
.expect("private key should encode");
let public_key_base64 = BASE64_STANDARD.encode(signing_key.verifying_key().as_bytes());
fn decode_agent_identity_jwt_verifies_when_jwks_is_present() {
let jwks = test_jwks("test-key");
let claims = AgentIdentityJwtClaims {
iss: AGENT_IDENTITY_JWT_ISSUER.to_string(),
aud: AGENT_IDENTITY_JWT_AUDIENCE.to_string(),
iat: 1_700_000_000,
exp: 4_000_000_000,
agent_runtime_id: "agent-runtime-id".to_string(),
agent_private_key: "private-key".to_string(),
account_id: "account-id".to_string(),
@@ -517,8 +540,12 @@ mod tests {
chatgpt_account_is_fedramp: false,
};
let jwt = jsonwebtoken::encode(
&Header::new(Algorithm::EdDSA),
&test_jwt_header("test-key"),
&serde_json::json!({
"iss": claims.iss,
"aud": claims.aud,
"iat": claims.iat,
"exp": claims.exp,
"agent_runtime_id": claims.agent_runtime_id,
"agent_private_key": claims.agent_private_key,
"account_id": claims.account_id,
@@ -527,11 +554,15 @@ mod tests {
"plan_type": "pro",
"chatgpt_account_is_fedramp": claims.chatgpt_account_is_fedramp,
}),
&EncodingKey::from_ed_der(private_key_pkcs8.as_bytes()),
&test_rsa_encoding_key(),
)
.expect("JWT should encode");
let expected_claims = AgentIdentityJwtClaims {
iss: AGENT_IDENTITY_JWT_ISSUER.to_string(),
aud: AGENT_IDENTITY_JWT_AUDIENCE.to_string(),
iat: 1_700_000_000,
exp: 4_000_000_000,
agent_runtime_id: "agent-runtime-id".to_string(),
agent_private_key: "private-key".to_string(),
account_id: "account-id".to_string(),
@@ -541,31 +572,22 @@ mod tests {
chatgpt_account_is_fedramp: false,
};
assert_eq!(
decode_agent_identity_jwt(&jwt, Some(&public_key_base64)).expect("JWT should verify"),
decode_agent_identity_jwt(&jwt, Some(&jwks)).expect("JWT should verify"),
expected_claims
);
}
#[test]
fn decode_agent_identity_jwt_rejects_wrong_public_key() {
let mut signing_secret_key_bytes = [0u8; 32];
signing_secret_key_bytes[0] = 1;
let signing_key = SigningKey::from_bytes(&signing_secret_key_bytes);
let private_key_pkcs8 = signing_key
.to_pkcs8_der()
.expect("private key should encode");
let mut other_secret_key_bytes = [0u8; 32];
other_secret_key_bytes[0] = 2;
let other_public_key_base64 = BASE64_STANDARD.encode(
SigningKey::from_bytes(&other_secret_key_bytes)
.verifying_key()
.as_bytes(),
);
fn decode_agent_identity_jwt_rejects_untrusted_kid() {
let jwks = test_jwks("other-key");
let jwt = jsonwebtoken::encode(
&Header::new(Algorithm::EdDSA),
&test_jwt_header("test-key"),
&serde_json::json!({
"iss": AGENT_IDENTITY_JWT_ISSUER,
"aud": AGENT_IDENTITY_JWT_AUDIENCE,
"iat": 1_700_000_000,
"exp": 4_000_000_000usize,
"agent_runtime_id": "agent-runtime-id",
"agent_private_key": "private-key",
"account_id": "account-id",
@@ -574,19 +596,111 @@ mod tests {
"plan_type": "pro",
"chatgpt_account_is_fedramp": false,
}),
&EncodingKey::from_ed_der(private_key_pkcs8.as_bytes()),
&test_rsa_encoding_key(),
)
.expect("JWT should encode");
decode_agent_identity_jwt(&jwt, Some(&other_public_key_base64))
.expect_err("JWT should not verify");
decode_agent_identity_jwt(&jwt, Some(&jwks)).expect_err("JWT should not verify");
}
#[test]
fn normalize_chatgpt_base_url_strips_codex_before_backend_api() {
fn decode_agent_identity_jwt_requires_issuer_and_audience() {
let jwks = test_jwks("test-key");
let jwt = jsonwebtoken::encode(
&test_jwt_header("test-key"),
&serde_json::json!({
"iat": 1_700_000_000,
"exp": 4_000_000_000usize,
"agent_runtime_id": "agent-runtime-id",
"agent_private_key": "private-key",
"account_id": "account-id",
"chatgpt_user_id": "user-id",
"email": "user@example.com",
"plan_type": "pro",
"chatgpt_account_is_fedramp": false,
}),
&test_rsa_encoding_key(),
)
.expect("JWT should encode");
decode_agent_identity_jwt(&jwt, Some(&jwks)).expect_err("JWT should not verify");
}
fn test_jwt_header(kid: &str) -> Header {
let mut header = Header::new(Algorithm::RS256);
header.kid = Some(kid.to_string());
header
}
fn test_rsa_encoding_key() -> EncodingKey {
EncodingKey::from_rsa_pem(
br#"-----BEGIN PRIVATE KEY-----
MIIEvgIBADANBgkqhkiG9w0BAQEFAASCBKgwggSkAgEAAoIBAQDWpAXYypOsYAwO
bvBduMk/mxaoYDze0AZSzaSzLuIlcsl2EKDgC3AabhIWXh/qTGEJLOU3VB1e5mO9
FPbBlmIZSL3FQTbyt/hYutPFKfCou5PLmScw/TzILS3/RhT8UY9kxxZvXiEbTki9
mvxRuZFpVqDFJHwfitIjKZGhXDCYVKurPTrxetYZJg0h8sQBLKjkZ0BqqaTUkAsg
0eBgZAlXEzG3By8PGhUqYLt6W1Q3KYw0FmGy/gTyzH1g0ukGgSJvOd8SkNT8MbOs
zl5kKxDNqpuEE6UZ3jbuJ+5382d31w+rOAJRzbf7QVdI9+luCSwJcDACYPQ4WNBa
uCpV0ovpAgMBAAECggEAVu84LwZdqYN9XpswX8VoPYrjMm9IODapWQBRpQFoNyK2
1ksF3bjEPvA2Azk8U/l7k+vLKw22l6lY3EyRZPcz5GnB8xLm3ogE3mtNOp4yCyVu
RxhQ91aaN7mU17/a4BdorLi2LYVCg3zBmYociD1Q2AluNGsCmwPu+K7tfR2J0Sg8
NjqiTbDG1XDpR/icwgC9t6vh8lZpCHDhF4tbQfLLVLeA/OdcuzXDyMCXbmdVIdBQ
rm4aIFmr2e1/2ctTbCg85S6AGFTH+pSLjrwTzyvf+F6NW5uNjLQAQLFj+EznBDxj
Xdx90cySrjsKK6PVWQF4RiTvkSW8eWL7R6B2FZbGwQKBgQDuVQRj72hWloR7mbEL
aUEEv3pIXTMXWEsoMBNczos/1L1RnAN1AI44TurznasPZAWvQj+kVbLDR+TAeZrL
iA8HIWswQUI18hFmgKzSkwIXGtubcKVrgsKeS4lMDKCM/Ef6WAYdeq6ronoY5lCN
YrJFmGp81W5zcV7lyiycgbSiGwKBgQDmjWYf6pZjrK7Z+OJ3X1AZfi2vss15SCvL
3fPgzIDbViztpGyQhc3DQZIsBNIu0xZp/veGce9TEeTds2ro9NfdJFeou8+fC7Pq
sOsM3amGFFi+ZW/9BWyjZEM88bgWWAjqLHbpfHDxjAf5CSxddqxgHlbP0Ytyb1Vg
gmPDn9YKSwKBgQDbTi3hC35WFuDHn0/zcSHcDZmnFuOZeqyFyV83yfMGhGrEuqvP
sPgtRikajJ3IZsB4WZyYSidZXEFY/0z6NjOl2xF38MTNQPbT/FmK1q1Yt2UWrlv5
BvSwlk87RG9D7C0LZo4R+D7cPoDdgqjiwMvMEIkEX5zn641oI1ZTmWKuuwKBgQCD
KF+3unnRvHRAVoFnTZbA2fJdqMeRvogD04GhGlYX8V9f1hFY6nXTJaNlXVzA/J8c
r8ra9kgjJuPfZ+ljG58OFFW2DRohLcQtuHYPfK6rMzoFHqnl9EcIcMp7ijuionR3
29HOJFgQYgxLFXfit9d6WugiE+BTupiEbckZif13HwKBgE/lAlkVHP6YahOO2Ljc
J1bwkqKZTB5dHolX9A58e/xXnfZ5P8f3Z83+Izap3FwqQulk7b1WO1MQcHuVg2NN
5da0D4h2rYOXnbYIg0BVu4spQbaM6ewsp66b8+MzLOBvj8SzWdt1Oyw0q/MRyQAR
8U4M2TSWCKUY/A6sT4W8+mT9
-----END PRIVATE KEY-----"#,
)
.expect("test RSA key should parse")
}
fn test_jwks(kid: &str) -> jsonwebtoken::jwk::JwkSet {
serde_json::from_value(serde_json::json!({
"keys": [{
"kty": "RSA",
"kid": kid,
"use": "sig",
"alg": "RS256",
"n": "1qQF2MqTrGAMDm7wXbjJP5sWqGA83tAGUs2ksy7iJXLJdhCg4AtwGm4SFl4f6kxhCSzlN1QdXuZjvRT2wZZiGUi9xUE28rf4WLrTxSnwqLuTy5knMP08yC0t_0YU_FGPZMcWb14hG05IvZr8UbmRaVagxSR8H4rSIymRoVwwmFSrqz068XrWGSYNIfLEASyo5GdAaqmk1JALINHgYGQJVxMxtwcvDxoVKmC7eltUNymMNBZhsv4E8sx9YNLpBoEibznfEpDU_DGzrM5eZCsQzaqbhBOlGd427ifud_Nnd9cPqzgCUc23-0FXSPfpbgksCXAwAmD0OFjQWrgqVdKL6Q",
"e": "AQAB",
}]
}))
.expect("test JWKS should parse")
}
#[test]
fn agent_identity_jwks_url_uses_backend_api_base_url() {
assert_eq!(
normalize_chatgpt_base_url("https://chatgpt.com/codex"),
"https://chatgpt.com/backend-api"
agent_identity_jwks_url("https://chatgpt.com/backend-api"),
"https://chatgpt.com/backend-api/wham/agent-identities/jwks"
);
assert_eq!(
agent_identity_jwks_url("https://chatgpt.com/backend-api/"),
"https://chatgpt.com/backend-api/wham/agent-identities/jwks"
);
}
#[test]
fn agent_identity_jwks_url_uses_codex_api_base_url() {
assert_eq!(
agent_identity_jwks_url("http://localhost:8080/api/codex"),
"http://localhost:8080/api/codex/agent-identities/jwks"
);
assert_eq!(
agent_identity_jwks_url("http://localhost:8080/api/codex/"),
"http://localhost:8080/api/codex/agent-identities/jwks"
);
}
+10 -3
View File
@@ -207,7 +207,10 @@ pub async fn run_login_with_agent_identity(
&config.codex_home,
&agent_identity,
config.cli_auth_credentials_store_mode,
) {
Some(&config.chatgpt_base_url),
)
.await
{
Ok(_) => {
eprintln!("{LOGIN_SUCCESS_MESSAGE}");
std::process::exit(0);
@@ -362,8 +365,12 @@ pub async fn run_login_with_device_code_fallback_to_browser(
pub async fn run_login_status(cli_config_overrides: CliConfigOverrides) -> ! {
let config = load_config_or_exit(cli_config_overrides).await;
match CodexAuth::from_auth_storage(&config.codex_home, config.cli_auth_credentials_store_mode)
.await
match CodexAuth::from_auth_storage(
&config.codex_home,
config.cli_auth_credentials_store_mode,
Some(&config.chatgpt_base_url),
)
.await
{
Ok(Some(auth)) => match auth.auth_mode() {
AuthMode::ApiKey => match auth.get_token() {
+4 -12
View File
@@ -6,8 +6,6 @@ use pretty_assertions::assert_eq;
use serde_json::Value;
use tempfile::TempDir;
const FAKE_AGENT_IDENTITY_JWT: &str = "eyJhbGciOiJFZERTQSIsInR5cCI6IkpXVCJ9.eyJhZ2VudF9ydW50aW1lX2lkIjoiYWdlbnQtcnVudGltZS1pZCIsImFnZW50X3ByaXZhdGVfa2V5IjoicHJpdmF0ZS1rZXkiLCJhY2NvdW50X2lkIjoiYWNjb3VudC0xMjMiLCJjaGF0Z3B0X3VzZXJfaWQiOiJ1c2VyLWlkIiwiZW1haWwiOiJ1c2VyQGV4YW1wbGUuY29tIiwicGxhbl90eXBlIjoicHJvIiwiY2hhdGdwdF9hY2NvdW50X2lzX2ZlZHJhbXAiOmZhbHNlfQ.c2ln";
fn codex_command(codex_home: &Path) -> Result<assert_cmd::Command> {
let mut cmd = assert_cmd::Command::new(codex_utils_cargo_bin::cargo_bin("codex")?);
cmd.env("CODEX_HOME", codex_home);
@@ -53,22 +51,16 @@ fn login_with_api_key_reads_stdin_and_writes_auth_json() -> Result<()> {
}
#[test]
fn login_with_agent_identity_reads_stdin_and_writes_auth_json() -> Result<()> {
fn login_with_agent_identity_rejects_invalid_jwt() -> Result<()> {
let codex_home = TempDir::new()?;
write_file_auth_config(codex_home.path())?;
let mut cmd = codex_command(codex_home.path())?;
cmd.args(["login", "--with-agent-identity"])
.write_stdin(format!("{FAKE_AGENT_IDENTITY_JWT}\n"))
.write_stdin("not-a-jwt\n")
.assert()
.success()
.stderr(contains("Successfully logged in"));
let auth = read_auth_json(codex_home.path())?;
assert_eq!(auth["auth_mode"], "agentIdentity");
assert_eq!(auth["agent_identity"], FAKE_AGENT_IDENTITY_JWT);
assert!(auth["OPENAI_API_KEY"].is_null());
assert!(auth.get("tokens").is_none());
.failure()
.stderr(contains("Error logging in with Agent Identity"));
Ok(())
}
+11 -7
View File
@@ -1090,13 +1090,17 @@ async fn prefers_apikey_when_config_prefers_apikey_even_with_chatgpt_tokens() {
let mut config = load_default_config_for_test(&codex_home).await;
config.model_provider = model_provider;
let auth_manager =
match CodexAuth::from_auth_storage(codex_home.path(), AuthCredentialsStoreMode::File).await
{
Ok(Some(auth)) => codex_core::test_support::auth_manager_from_auth(auth),
Ok(None) => panic!("No CodexAuth found in codex_home"),
Err(e) => panic!("Failed to load CodexAuth: {e}"),
};
let auth_manager = match CodexAuth::from_auth_storage(
codex_home.path(),
AuthCredentialsStoreMode::File,
/*chatgpt_base_url*/ None,
)
.await
{
Ok(Some(auth)) => codex_core::test_support::auth_manager_from_auth(auth),
Ok(None) => panic!("No CodexAuth found in codex_home"),
Err(e) => panic!("Failed to load CodexAuth: {e}"),
};
let thread_manager = ThreadManager::new(
&config,
auth_manager,
+1
View File
@@ -439,6 +439,7 @@ pub async fn run_main(cli: Cli, arg0_paths: Arg0DispatchPaths) -> anyhow::Result
auth_credentials_store_mode: config.cli_auth_credentials_store_mode,
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()),
})
.await
{
+1
View File
@@ -45,6 +45,7 @@ webbrowser = { workspace = true }
[dev-dependencies]
anyhow = { workspace = true }
core_test_support = { workspace = true }
jsonwebtoken = { workspace = true }
keyring = { workspace = true }
pretty_assertions = { workspace = true }
regex-lite = { workspace = true }
+2 -2
View File
@@ -82,7 +82,7 @@ mod tests {
use serial_test::serial;
#[test]
#[serial(agent_identity_authapi_base_url_env)]
#[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,
@@ -95,7 +95,7 @@ mod tests {
}
#[test]
#[serial(agent_identity_authapi_base_url_env)]
#[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!(
+231 -13
View File
@@ -16,6 +16,11 @@ use serde_json::json;
use std::sync::Arc;
use tempfile::TempDir;
use tempfile::tempdir;
use wiremock::Mock;
use wiremock::MockServer;
use wiremock::ResponseTemplate;
use wiremock::matchers::method;
use wiremock::matchers::path;
#[tokio::test]
async fn refresh_without_id_token() {
@@ -78,15 +83,29 @@ fn login_with_api_key_overwrites_existing_auth_json() {
assert!(auth.tokens.is_none(), "tokens should be cleared");
}
#[test]
fn login_with_agent_identity_writes_only_token() {
#[tokio::test]
async fn login_with_agent_identity_writes_only_token() {
let dir = tempdir().unwrap();
let auth_path = dir.path().join("auth.json");
let record = agent_identity_record("account-123");
let agent_identity = fake_agent_identity_jwt(&record).expect("fake agent identity");
let agent_identity = signed_agent_identity_jwt(&record).expect("signed agent identity");
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/backend-api/wham/agent-identities/jwks"))
.respond_with(ResponseTemplate::new(200).set_body_json(test_jwks_body()))
.expect(1)
.mount(&server)
.await;
let chatgpt_base_url = format!("{}/backend-api", server.uri());
super::login_with_agent_identity(dir.path(), &agent_identity, AuthCredentialsStoreMode::File)
.expect("login_with_agent_identity should succeed");
super::login_with_agent_identity(
dir.path(),
&agent_identity,
AuthCredentialsStoreMode::File,
Some(&chatgpt_base_url),
)
.await
.expect("login_with_agent_identity should succeed");
let storage = FileAuthStorage::new(dir.path().to_path_buf());
let auth = storage
@@ -99,15 +118,21 @@ fn login_with_agent_identity_writes_only_token() {
);
assert!(auth.tokens.is_none(), "tokens should be cleared");
assert!(auth.openai_api_key.is_none(), "API key should be cleared");
server.verify().await;
}
#[test]
fn login_with_agent_identity_rejects_invalid_jwt() {
#[tokio::test]
async fn login_with_agent_identity_rejects_invalid_jwt() {
let dir = tempdir().unwrap();
let err =
super::login_with_agent_identity(dir.path(), "not-a-jwt", AuthCredentialsStoreMode::File)
.expect_err("invalid Agent Identity token should fail");
let err = super::login_with_agent_identity(
dir.path(),
"not-a-jwt",
AuthCredentialsStoreMode::File,
/*chatgpt_base_url*/ None,
)
.await
.expect_err("invalid Agent Identity token should fail");
assert_eq!(err.kind(), std::io::ErrorKind::Other);
assert!(
@@ -117,11 +142,47 @@ fn login_with_agent_identity_rejects_invalid_jwt() {
}
#[tokio::test]
async fn login_with_agent_identity_rejects_unsigned_jwt() {
let dir = tempdir().unwrap();
let record = agent_identity_record("account-123");
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"))
.respond_with(ResponseTemplate::new(200).set_body_json(test_jwks_body()))
.expect(1)
.mount(&server)
.await;
let chatgpt_base_url = format!("{}/backend-api", server.uri());
super::login_with_agent_identity(
dir.path(),
&agent_identity,
AuthCredentialsStoreMode::File,
Some(&chatgpt_base_url),
)
.await
.expect_err("unsigned Agent Identity token should fail");
assert!(
!get_auth_file(dir.path()).exists(),
"unsigned Agent Identity token should not write auth.json"
);
server.verify().await;
}
#[tokio::test]
#[serial(codex_auth_env)]
async fn missing_auth_json_returns_none() {
let dir = tempdir().unwrap();
let auth = CodexAuth::from_auth_storage(dir.path(), AuthCredentialsStoreMode::File)
.await
.expect("call should succeed");
let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
let auth = CodexAuth::from_auth_storage(
dir.path(),
AuthCredentialsStoreMode::File,
/*chatgpt_base_url*/ None,
)
.await
.expect("call should succeed");
assert_eq!(auth, None);
}
@@ -129,6 +190,7 @@ async fn missing_auth_json_returns_none() {
#[serial(codex_auth_env)]
async fn pro_account_with_no_api_key_uses_chatgpt_auth() {
let codex_home = tempdir().unwrap();
let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
let fake_jwt = write_auth_file(
AuthFileParams {
openai_api_key: None,
@@ -143,6 +205,7 @@ async fn pro_account_with_no_api_key_uses_chatgpt_auth() {
codex_home.path(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*chatgpt_base_url*/ None,
)
.await
.unwrap()
@@ -186,6 +249,7 @@ async fn pro_account_with_no_api_key_uses_chatgpt_auth() {
#[serial(codex_auth_env)]
async fn loads_api_key_from_auth_json() {
let dir = tempdir().unwrap();
let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
let auth_file = dir.path().join("auth.json");
std::fs::write(
auth_file,
@@ -197,6 +261,7 @@ async fn loads_api_key_from_auth_json() {
dir.path(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*chatgpt_base_url*/ None,
)
.await
.unwrap()
@@ -255,8 +320,10 @@ async fn unauthorized_recovery_reports_mode_and_step_names() {
}
#[tokio::test]
#[serial(codex_auth_env)]
async fn refresh_failure_is_scoped_to_the_matching_auth_snapshot() {
let codex_home = tempdir().unwrap();
let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
write_auth_file(
AuthFileParams {
openai_api_key: None,
@@ -271,6 +338,7 @@ async fn refresh_failure_is_scoped_to_the_matching_auth_snapshot() {
codex_home.path(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*chatgpt_base_url*/ None,
)
.await
.expect("load auth")
@@ -288,6 +356,7 @@ async fn refresh_failure_is_scoped_to_the_matching_auth_snapshot() {
codex_home.path(),
updated_auth_dot_json,
AuthCredentialsStoreMode::File,
/*chatgpt_base_url*/ None,
)
.await
.expect("updated auth should parse");
@@ -591,6 +660,7 @@ async fn build_config(
auth_credentials_store_mode: AuthCredentialsStoreMode::File,
forced_login_method,
forced_chatgpt_workspace_id,
chatgpt_base_url: None,
}
}
@@ -611,6 +681,14 @@ impl EnvVarGuard {
}
Self { key, original }
}
fn remove(key: &'static str) -> Self {
let original = env::var_os(key);
unsafe {
env::remove_var(key);
}
Self { key, original }
}
}
#[cfg(test)]
@@ -625,6 +703,55 @@ 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 =
signed_agent_identity_jwt(&expected_record).expect("signed agent identity");
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/backend-api/wham/agent-identities/jwks"))
.respond_with(ResponseTemplate::new(200).set_body_json(test_jwks_body()))
.expect(1)
.mount(&server)
.await;
Mock::given(method("POST"))
.and(path("/backend-api/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 _agent_guard = EnvVarGuard::set(CODEX_AGENT_IDENTITY_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,
AuthCredentialsStoreMode::File,
Some(&chatgpt_base_url),
)
.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_eq!(agent_identity.process_task_id(), "task-123");
assert!(
!get_auth_file(codex_home.path()).exists(),
"env auth should not write auth.json"
);
server.verify().await;
}
#[tokio::test]
#[serial(codex_auth_env)]
async fn load_auth_keeps_codex_api_key_env_precedence() {
@@ -638,6 +765,7 @@ async fn load_auth_keeps_codex_api_key_env_precedence() {
codex_home.path(),
/*enable_codex_api_key_env*/ true,
AuthCredentialsStoreMode::File,
/*chatgpt_base_url*/ None,
)
.await
.expect("env auth should load")
@@ -650,6 +778,7 @@ async fn load_auth_keeps_codex_api_key_env_precedence() {
#[serial(codex_auth_env)]
async fn enforce_login_restrictions_logs_out_for_method_mismatch() {
let codex_home = tempdir().unwrap();
let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
login_with_api_key(codex_home.path(), "sk-test", AuthCredentialsStoreMode::File)
.expect("seed api key");
@@ -674,6 +803,7 @@ async fn enforce_login_restrictions_logs_out_for_method_mismatch() {
#[serial(codex_auth_env)]
async fn enforce_login_restrictions_logs_out_for_workspace_mismatch() {
let codex_home = tempdir().unwrap();
let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
let _jwt = write_auth_file(
AuthFileParams {
openai_api_key: None,
@@ -705,6 +835,7 @@ async fn enforce_login_restrictions_logs_out_for_workspace_mismatch() {
#[serial(codex_auth_env)]
async fn enforce_login_restrictions_allows_matching_workspace() {
let codex_home = tempdir().unwrap();
let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
let _jwt = write_auth_file(
AuthFileParams {
openai_api_key: None,
@@ -736,6 +867,7 @@ async fn enforce_login_restrictions_allows_matching_workspace() {
async fn enforce_login_restrictions_allows_api_key_if_login_method_not_set_but_forced_chatgpt_workspace_id_is_set()
{
let codex_home = tempdir().unwrap();
let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
login_with_api_key(codex_home.path(), "sk-test", AuthCredentialsStoreMode::File)
.expect("seed api key");
@@ -759,6 +891,7 @@ async fn enforce_login_restrictions_allows_api_key_if_login_method_not_set_but_f
#[serial(codex_auth_env)]
async fn enforce_login_restrictions_blocks_env_api_key_when_chatgpt_required() {
let _guard = EnvVarGuard::set(CODEX_API_KEY_ENV_VAR, "sk-env");
let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
let codex_home = tempdir().unwrap();
let config = build_config(
@@ -795,6 +928,10 @@ fn fake_agent_identity_jwt(record: &AgentIdentityAuthRecord) -> std::io::Result<
let encode = |bytes: &[u8]| base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(bytes);
let header_b64 = encode(br#"{"alg":"EdDSA","typ":"JWT"}"#);
let payload = json!({
"iss": "https://chatgpt.com/codex-backend/agent-identity",
"aud": "codex-app-server",
"iat": 1_700_000_000usize,
"exp": 4_000_000_000usize,
"agent_runtime_id": record.agent_runtime_id,
"agent_private_key": record.agent_private_key,
"account_id": record.account_id,
@@ -808,9 +945,77 @@ fn fake_agent_identity_jwt(record: &AgentIdentityAuthRecord) -> std::io::Result<
Ok(format!("{header_b64}.{payload_b64}.{signature_b64}"))
}
fn signed_agent_identity_jwt(
record: &AgentIdentityAuthRecord,
) -> jsonwebtoken::errors::Result<String> {
let mut header = jsonwebtoken::Header::new(jsonwebtoken::Algorithm::RS256);
header.kid = Some("test-key".to_string());
jsonwebtoken::encode(
&header,
&json!({
"iss": "https://chatgpt.com/codex-backend/agent-identity",
"aud": "codex-app-server",
"iat": 1_700_000_000usize,
"exp": 4_000_000_000usize,
"agent_runtime_id": record.agent_runtime_id,
"agent_private_key": record.agent_private_key,
"account_id": record.account_id,
"chatgpt_user_id": record.chatgpt_user_id,
"email": record.email,
"plan_type": record.plan_type,
"chatgpt_account_is_fedramp": record.chatgpt_account_is_fedramp,
}),
&jsonwebtoken::EncodingKey::from_rsa_pem(TEST_AGENT_IDENTITY_RSA_PRIVATE_KEY_PEM)?,
)
}
fn test_jwks_body() -> serde_json::Value {
json!({
"keys": [{
"kty": "RSA",
"kid": "test-key",
"use": "sig",
"alg": "RS256",
"n": "1qQF2MqTrGAMDm7wXbjJP5sWqGA83tAGUs2ksy7iJXLJdhCg4AtwGm4SFl4f6kxhCSzlN1QdXuZjvRT2wZZiGUi9xUE28rf4WLrTxSnwqLuTy5knMP08yC0t_0YU_FGPZMcWb14hG05IvZr8UbmRaVagxSR8H4rSIymRoVwwmFSrqz068XrWGSYNIfLEASyo5GdAaqmk1JALINHgYGQJVxMxtwcvDxoVKmC7eltUNymMNBZhsv4E8sx9YNLpBoEibznfEpDU_DGzrM5eZCsQzaqbhBOlGd427ifud_Nnd9cPqzgCUc23-0FXSPfpbgksCXAwAmD0OFjQWrgqVdKL6Q",
"e": "AQAB",
}]
})
}
const TEST_AGENT_IDENTITY_RSA_PRIVATE_KEY_PEM: &[u8] = br#"-----BEGIN PRIVATE KEY-----
MIIEvgIBADANBgkqhkiG9w0BAQEFAASCBKgwggSkAgEAAoIBAQDWpAXYypOsYAwO
bvBduMk/mxaoYDze0AZSzaSzLuIlcsl2EKDgC3AabhIWXh/qTGEJLOU3VB1e5mO9
FPbBlmIZSL3FQTbyt/hYutPFKfCou5PLmScw/TzILS3/RhT8UY9kxxZvXiEbTki9
mvxRuZFpVqDFJHwfitIjKZGhXDCYVKurPTrxetYZJg0h8sQBLKjkZ0BqqaTUkAsg
0eBgZAlXEzG3By8PGhUqYLt6W1Q3KYw0FmGy/gTyzH1g0ukGgSJvOd8SkNT8MbOs
zl5kKxDNqpuEE6UZ3jbuJ+5382d31w+rOAJRzbf7QVdI9+luCSwJcDACYPQ4WNBa
uCpV0ovpAgMBAAECggEAVu84LwZdqYN9XpswX8VoPYrjMm9IODapWQBRpQFoNyK2
1ksF3bjEPvA2Azk8U/l7k+vLKw22l6lY3EyRZPcz5GnB8xLm3ogE3mtNOp4yCyVu
RxhQ91aaN7mU17/a4BdorLi2LYVCg3zBmYociD1Q2AluNGsCmwPu+K7tfR2J0Sg8
NjqiTbDG1XDpR/icwgC9t6vh8lZpCHDhF4tbQfLLVLeA/OdcuzXDyMCXbmdVIdBQ
rm4aIFmr2e1/2ctTbCg85S6AGFTH+pSLjrwTzyvf+F6NW5uNjLQAQLFj+EznBDxj
Xdx90cySrjsKK6PVWQF4RiTvkSW8eWL7R6B2FZbGwQKBgQDuVQRj72hWloR7mbEL
aUEEv3pIXTMXWEsoMBNczos/1L1RnAN1AI44TurznasPZAWvQj+kVbLDR+TAeZrL
iA8HIWswQUI18hFmgKzSkwIXGtubcKVrgsKeS4lMDKCM/Ef6WAYdeq6ronoY5lCN
YrJFmGp81W5zcV7lyiycgbSiGwKBgQDmjWYf6pZjrK7Z+OJ3X1AZfi2vss15SCvL
3fPgzIDbViztpGyQhc3DQZIsBNIu0xZp/veGce9TEeTds2ro9NfdJFeou8+fC7Pq
sOsM3amGFFi+ZW/9BWyjZEM88bgWWAjqLHbpfHDxjAf5CSxddqxgHlbP0Ytyb1Vg
gmPDn9YKSwKBgQDbTi3hC35WFuDHn0/zcSHcDZmnFuOZeqyFyV83yfMGhGrEuqvP
sPgtRikajJ3IZsB4WZyYSidZXEFY/0z6NjOl2xF38MTNQPbT/FmK1q1Yt2UWrlv5
BvSwlk87RG9D7C0LZo4R+D7cPoDdgqjiwMvMEIkEX5zn641oI1ZTmWKuuwKBgQCD
KF+3unnRvHRAVoFnTZbA2fJdqMeRvogD04GhGlYX8V9f1hFY6nXTJaNlXVzA/J8c
r8ra9kgjJuPfZ+ljG58OFFW2DRohLcQtuHYPfK6rMzoFHqnl9EcIcMp7ijuionR3
29HOJFgQYgxLFXfit9d6WugiE+BTupiEbckZif13HwKBgE/lAlkVHP6YahOO2Ljc
J1bwkqKZTB5dHolX9A58e/xXnfZ5P8f3Z83+Izap3FwqQulk7b1WO1MQcHuVg2NN
5da0D4h2rYOXnbYIg0BVu4spQbaM6ewsp66b8+MzLOBvj8SzWdt1Oyw0q/MRyQAR
8U4M2TSWCKUY/A6sT4W8+mT9
-----END PRIVATE KEY-----"#;
#[tokio::test]
#[serial(codex_auth_env)]
async fn plan_type_maps_known_plan() {
let codex_home = tempdir().unwrap();
let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
let _jwt = write_auth_file(
AuthFileParams {
openai_api_key: None,
@@ -825,6 +1030,7 @@ async fn plan_type_maps_known_plan() {
codex_home.path(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*chatgpt_base_url*/ None,
)
.await
.expect("load auth")
@@ -834,8 +1040,10 @@ async fn plan_type_maps_known_plan() {
}
#[tokio::test]
#[serial(codex_auth_env)]
async fn plan_type_maps_self_serve_business_usage_based_plan() {
let codex_home = tempdir().unwrap();
let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
let _jwt = write_auth_file(
AuthFileParams {
openai_api_key: None,
@@ -850,6 +1058,7 @@ async fn plan_type_maps_self_serve_business_usage_based_plan() {
codex_home.path(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*chatgpt_base_url*/ None,
)
.await
.expect("load auth")
@@ -862,8 +1071,10 @@ async fn plan_type_maps_self_serve_business_usage_based_plan() {
}
#[tokio::test]
#[serial(codex_auth_env)]
async fn plan_type_maps_enterprise_cbp_usage_based_plan() {
let codex_home = tempdir().unwrap();
let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
let _jwt = write_auth_file(
AuthFileParams {
openai_api_key: None,
@@ -878,6 +1089,7 @@ async fn plan_type_maps_enterprise_cbp_usage_based_plan() {
codex_home.path(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*chatgpt_base_url*/ None,
)
.await
.expect("load auth")
@@ -890,8 +1102,10 @@ async fn plan_type_maps_enterprise_cbp_usage_based_plan() {
}
#[tokio::test]
#[serial(codex_auth_env)]
async fn plan_type_maps_unknown_to_unknown() {
let codex_home = tempdir().unwrap();
let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
let _jwt = write_auth_file(
AuthFileParams {
openai_api_key: None,
@@ -906,6 +1120,7 @@ async fn plan_type_maps_unknown_to_unknown() {
codex_home.path(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*chatgpt_base_url*/ None,
)
.await
.expect("load auth")
@@ -915,8 +1130,10 @@ async fn plan_type_maps_unknown_to_unknown() {
}
#[tokio::test]
#[serial(codex_auth_env)]
async fn missing_plan_type_maps_to_unknown() {
let codex_home = tempdir().unwrap();
let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
let _jwt = write_auth_file(
AuthFileParams {
openai_api_key: None,
@@ -931,6 +1148,7 @@ async fn missing_plan_type_maps_to_unknown() {
codex_home.path(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*chatgpt_base_url*/ None,
)
.await
.expect("load auth")
+50 -9
View File
@@ -16,6 +16,8 @@ use std::sync::atomic::AtomicU64;
use std::sync::atomic::Ordering;
use tokio::sync::Semaphore;
use codex_agent_identity::decode_agent_identity_jwt;
use codex_agent_identity::fetch_agent_identity_jwks;
use codex_app_server_protocol::AuthMode;
use codex_app_server_protocol::AuthMode as ApiAuthMode;
use codex_protocol::config_types::ForcedLoginMethod;
@@ -29,6 +31,7 @@ pub use crate::auth::storage::AuthDotJson;
use crate::auth::storage::AuthStorageBackend;
use crate::auth::storage::create_auth_storage;
use crate::auth::util::try_parse_error_message;
use crate::default_client::build_reqwest_client;
use crate::default_client::create_client;
use crate::token_data::TokenData;
use crate::token_data::parse_chatgpt_jwt_claims;
@@ -88,6 +91,7 @@ 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";
@@ -197,6 +201,7 @@ impl CodexAuth {
codex_home: &Path,
auth_dot_json: AuthDotJson,
auth_credentials_store_mode: AuthCredentialsStoreMode,
chatgpt_base_url: Option<&str>,
) -> std::io::Result<Self> {
let auth_mode = auth_dot_json.resolved_mode();
let client = create_client();
@@ -212,7 +217,7 @@ impl CodexAuth {
"agent identity auth is missing an agent identity token.",
));
};
return Self::from_agent_identity_jwt(&agent_identity).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);
@@ -237,17 +242,26 @@ impl CodexAuth {
pub async fn from_auth_storage(
codex_home: &Path,
auth_credentials_store_mode: AuthCredentialsStoreMode,
chatgpt_base_url: Option<&str>,
) -> std::io::Result<Option<Self>> {
load_auth(
codex_home,
/*enable_codex_api_key_env*/ false,
auth_credentials_store_mode,
chatgpt_base_url,
)
.await
}
pub async fn from_agent_identity_jwt(jwt: &str) -> std::io::Result<Self> {
let record = AgentIdentityAuthRecord::from_agent_identity_jwt(jwt)?;
pub async fn from_agent_identity_jwt(
jwt: &str,
chatgpt_base_url: Option<&str>,
) -> std::io::Result<Self> {
let base_url = chatgpt_base_url
.unwrap_or(DEFAULT_CHATGPT_BACKEND_BASE_URL)
.trim_end_matches('/')
.to_string();
let record = verified_agent_identity_record(jwt, &base_url).await?;
Ok(Self::AgentIdentity(AgentIdentityAuth::load(record).await?))
}
@@ -493,6 +507,18 @@ pub fn read_codex_agent_identity_from_env() -> Option<String> {
.filter(|value| !value.is_empty())
}
async fn verified_agent_identity_record(
jwt: &str,
chatgpt_base_url: &str,
) -> std::io::Result<AgentIdentityAuthRecord> {
AgentIdentityAuthRecord::from_agent_identity_jwt(jwt)?;
let jwks = fetch_agent_identity_jwks(&build_reqwest_client(), chatgpt_base_url)
.await
.map_err(std::io::Error::other)?;
let claims = decode_agent_identity_jwt(jwt, Some(&jwks)).map_err(std::io::Error::other)?;
Ok(claims.into())
}
/// 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(
@@ -535,12 +561,17 @@ pub fn login_with_api_key(
}
/// Writes an `auth.json` that contains only the Agent Identity token.
pub fn login_with_agent_identity(
pub async fn login_with_agent_identity(
codex_home: &Path,
agent_identity: &str,
auth_credentials_store_mode: AuthCredentialsStoreMode,
chatgpt_base_url: Option<&str>,
) -> std::io::Result<()> {
AgentIdentityAuthRecord::from_agent_identity_jwt(agent_identity)?;
let base_url = chatgpt_base_url
.unwrap_or(DEFAULT_CHATGPT_BACKEND_BASE_URL)
.trim_end_matches('/')
.to_string();
verified_agent_identity_record(agent_identity, &base_url).await?;
let auth_dot_json = AuthDotJson {
auth_mode: Some(ApiAuthMode::AgentIdentity),
openai_api_key: None,
@@ -599,6 +630,7 @@ pub struct AuthConfig {
pub auth_credentials_store_mode: AuthCredentialsStoreMode,
pub forced_login_method: Option<ForcedLoginMethod>,
pub forced_chatgpt_workspace_id: Option<String>,
pub chatgpt_base_url: Option<String>,
}
pub async fn enforce_login_restrictions(config: &AuthConfig) -> std::io::Result<()> {
@@ -606,6 +638,7 @@ pub async fn enforce_login_restrictions(config: &AuthConfig) -> std::io::Result<
&config.codex_home,
/*enable_codex_api_key_env*/ true,
config.auth_credentials_store_mode,
config.chatgpt_base_url.as_deref(),
)
.await?
else {
@@ -711,6 +744,7 @@ async fn load_auth(
codex_home: &Path,
enable_codex_api_key_env: bool,
auth_credentials_store_mode: AuthCredentialsStoreMode,
chatgpt_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() {
@@ -728,6 +762,7 @@ async fn load_auth(
codex_home,
auth_dot_json,
AuthCredentialsStoreMode::Ephemeral,
chatgpt_base_url,
)
.await?;
return Ok(Some(auth));
@@ -739,7 +774,7 @@ async fn load_auth(
}
if let Some(agent_identity) = read_codex_agent_identity_from_env() {
return CodexAuth::from_agent_identity_jwt(&agent_identity)
return CodexAuth::from_agent_identity_jwt(&agent_identity, chatgpt_base_url)
.await
.map(Some);
}
@@ -751,9 +786,13 @@ async fn load_auth(
None => return Ok(None),
};
let auth =
CodexAuth::from_auth_dot_json(codex_home, auth_dot_json, auth_credentials_store_mode)
.await?;
let auth = CodexAuth::from_auth_dot_json(
codex_home,
auth_dot_json,
auth_credentials_store_mode,
chatgpt_base_url,
)
.await?;
Ok(Some(auth))
}
@@ -1288,6 +1327,7 @@ impl AuthManager {
&codex_home,
enable_codex_api_key_env,
auth_credentials_store_mode,
chatgpt_base_url.as_deref(),
)
.await
.ok()
@@ -1491,6 +1531,7 @@ impl AuthManager {
&self.codex_home,
self.enable_codex_api_key_env,
self.auth_credentials_store_mode,
self.chatgpt_base_url.as_deref(),
)
.await
.ok()
+11 -4
View File
@@ -19,6 +19,7 @@ use std::sync::Mutex;
use tracing::warn;
use crate::token_data::TokenData;
use codex_agent_identity::AgentIdentityJwtClaims;
use codex_agent_identity::decode_agent_identity_jwt;
use codex_app_server_protocol::AuthMode;
use codex_config::types::AuthCredentialsStoreMode;
@@ -59,10 +60,16 @@ pub struct AgentIdentityAuthRecord {
impl AgentIdentityAuthRecord {
pub(crate) fn from_agent_identity_jwt(jwt: &str) -> std::io::Result<Self> {
let claims = decode_agent_identity_jwt(jwt, /*public_key_base64*/ None)
.map_err(std::io::Error::other)?;
let claims =
decode_agent_identity_jwt(jwt, /*jwks*/ None).map_err(std::io::Error::other)?;
Ok(Self {
Ok(claims.into())
}
}
impl From<AgentIdentityJwtClaims> for AgentIdentityAuthRecord {
fn from(claims: AgentIdentityJwtClaims) -> Self {
Self {
agent_runtime_id: claims.agent_runtime_id,
agent_private_key: claims.agent_private_key,
account_id: claims.account_id,
@@ -70,7 +77,7 @@ impl AgentIdentityAuthRecord {
email: claims.email,
plan_type: claims.plan_type,
chatgpt_account_is_fedramp: claims.chatgpt_account_is_fedramp,
})
}
}
}
+8 -4
View File
@@ -228,10 +228,14 @@ c2ln",
)
.expect("auth.json should be written");
CodexAuth::from_auth_storage(codex_home, AuthCredentialsStoreMode::File)
.await
.expect("auth should load")
.expect("auth should be present")
CodexAuth::from_auth_storage(
codex_home,
AuthCredentialsStoreMode::File,
/*chatgpt_base_url*/ None,
)
.await
.expect("auth should load")
.expect("auth should be present")
}
#[tokio::test]
+1
View File
@@ -900,6 +900,7 @@ pub async fn run_main(
auth_credentials_store_mode: config.cli_auth_credentials_store_mode,
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()),
})
.await
{