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: opt ChatGPT auth into agent identity (#19049)

Browse Source 
## Stack

This is PR 2 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

This PR adds the disabled-by-default path for normal ChatGPT-login Codex
sessions to obtain Agent Identity runtime auth through the Codex
backend. Existing Agent Identity JWT startup mode remains a separate
path and does not require the feature flag.

What changed:

- adds the experimental `use_agent_identity` feature flag and config
schema entry
- adds an explicit `AgentIdentityAuthPolicy` so call sites choose
`JwtOnly` or `ChatGptAuth` instead of passing a bare boolean
- stores standalone Agent Identity JWT credentials separately from
backend-registered Agent Identity records
- persists the registered Agent Identity record, private key, and single
run task id in `auth.json` so process restarts reuse the same identity
- derives the agent/task registration base URL from ChatGPT/Codex auth
config while keeping JWT JWKS lookup separate
- provisions and caches ChatGPT-derived Agent Identity runtime auth when
`use_agent_identity` is enabled
- reuses the shared run-task registration helper from PR1 rather than
adding a second task-registration path

This PR intentionally does not switch model inference over to
`AgentAssertion` auth. The provider-auth integration lands in the next
PR.

## Testing

- `just test -p codex-login`
This commit is contained in:
Adrian
2026-06-18 14:05:27 -07:00
committed by GitHub
Unverified
parent 765309d5a6
commit ec848dde0e
24 changed files with 1359 additions and 172 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/Cargo.lock
Generated
+1
View File
@@ -2418,6 +2418,7 @@ version = "0.0.0"
dependencies = [
"base64 0.22.1",
"chrono",
"codex-agent-identity",
"codex-backend-client",
"codex-config",
"codex-core",
codex-rs/app-server-transport/src/transport/remote_control/tests.rs
+3
View File
@@ -1038,6 +1038,7 @@ async fn remote_control_start_allows_missing_auth_when_enabled() {
codex_home.path().to_path_buf(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::default(),
)
@@ -1866,6 +1867,7 @@ async fn remote_control_waits_for_account_id_before_enrolling() {
codex_home.path().to_path_buf(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::default(),
)
@@ -1961,6 +1963,7 @@ async fn persisted_enable_does_not_follow_auth_to_an_account_without_a_preferenc
codex_home.path().to_path_buf(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::default(),
)
codex-rs/app-server-transport/src/transport/remote_control/tests/clients_tests.rs
+2
View File
@@ -182,6 +182,7 @@ async fn list_remote_control_clients_recovers_auth_after_unauthorized() {
codex_home.path().to_path_buf(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::default(),
)
@@ -266,6 +267,7 @@ async fn list_remote_control_clients_retries_unauthorized_only_once() {
codex_home.path().to_path_buf(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::default(),
)
codex-rs/app-server-transport/src/transport/remote_control/tests/pairing_tests.rs
+2
View File
@@ -536,6 +536,7 @@ async fn remote_control_handle_recovers_auth_before_refreshing_pairing() {
codex_home.path().to_path_buf(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::default(),
)
@@ -804,6 +805,7 @@ async fn remote_control_handle_discards_pairing_response_after_auth_change() {
codex_home.path().to_path_buf(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::default(),
)
codex-rs/app-server-transport/src/transport/remote_control/websocket.rs
+3
View File
@@ -2211,6 +2211,7 @@ mod tests {
codex_home.path().to_path_buf(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::default(),
)
@@ -2308,6 +2309,7 @@ mod tests {
codex_home.path().to_path_buf(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::default(),
)
@@ -2430,6 +2432,7 @@ mod tests {
codex_home.path().to_path_buf(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::default(),
)
codex-rs/cli/src/doctor.rs
+2 -2
View File
@@ -1399,8 +1399,8 @@ fn stored_auth_issues(
codex_app_server_protocol::AuthMode::AgentIdentity => {
if auth
.agent_identity
.as_deref()
.is_none_or(|token| token.trim().is_empty())
.as_ref()
.is_none_or(|agent_identity| !agent_identity.has_auth_material())
{
issues.push("agent identity auth is missing an agent identity token");
}
codex-rs/cli/src/main.rs
+3 -6
View File
@@ -1738,12 +1738,9 @@ 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),
/*agent_identity_authapi_base_url_override*/ None,
)
.await?;
let auth =
CodexAuth::from_agent_identity_jwt(&agent_identity_jwt, Some(&config.chatgpt_base_url))
.await?;
return Ok(codex_model_provider::auth_provider_from_auth(&auth));
}
codex-rs/cloud-config/Cargo.toml
+1
View File
@@ -25,6 +25,7 @@ tokio = { workspace = true, features = ["fs", "rt", "sync", "time"] }
tracing = { workspace = true }
[dev-dependencies]
codex-agent-identity = { workspace = true }
pretty_assertions = { workspace = true }
tempfile = { workspace = true }
tokio = { workspace = true, features = ["macros", "rt", "test-util", "time"] }
codex-rs/cloud-config/src/bundle_loader.rs
+1
View File
@@ -63,6 +63,7 @@ pub async fn cloud_config_bundle_loader_for_storage(
codex_home.clone(),
enable_codex_api_key_env,
credentials_store_mode,
/*forced_chatgpt_workspace_id*/ None,
Some(chatgpt_base_url.clone()),
keyring_backend_kind,
)
codex-rs/cloud-config/src/service_tests.rs
+52
View File
@@ -17,6 +17,8 @@ use codex_config::CloudRequirementsFragment;
use codex_config::CloudRequirementsTomlBundle;
use codex_config::types::AuthCredentialsStoreMode;
use codex_login::AuthKeyringBackendKind;
use codex_login::auth::AgentIdentityAuth;
use codex_login::auth::AgentIdentityAuthRecord;
use pretty_assertions::assert_eq;
use serde_json::json;
use std::collections::VecDeque;
@@ -48,6 +50,7 @@ async fn auth_manager_with_api_key() -> Arc<AuthManager> {
tmp.path().to_path_buf(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::default(),
)
@@ -77,6 +80,7 @@ async fn auth_manager_with_plan_and_identity(
tmp.path().to_path_buf(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::default(),
)
@@ -88,6 +92,28 @@ async fn auth_manager_with_plan(plan_type: &str) -> Arc<AuthManager> {
auth_manager_with_plan_and_identity(plan_type, Some("user-12345"), Some("account-12345")).await
}
async fn auth_manager_with_agent_identity_business_plan() -> Arc<AuthManager> {
let key_material =
codex_agent_identity::generate_agent_key_material().expect("generate agent key material");
AuthManager::from_auth_for_testing(CodexAuth::AgentIdentity(
AgentIdentityAuth::from_record(
AgentIdentityAuthRecord {
agent_runtime_id: "agent-runtime-123".to_string(),
agent_private_key: key_material.private_key_pkcs8_base64,
account_id: "account-12345".to_string(),
chatgpt_user_id: "user-12345".to_string(),
email: "user@example.com".to_string(),
plan_type: PlanType::Business,
chatgpt_account_is_fedramp: false,
task_id: Some("task-123".to_string()),
},
"https://auth.openai.com/api/accounts",
)
.await
.expect("agent identity record should be complete"),
))
}
fn chatgpt_auth_json(
plan_type: &str,
chatgpt_user_id: Option<&str>,
@@ -408,6 +434,28 @@ async fn get_bundle_allows_eligible_workspace_plans_and_writes_cache() {
}
}
#[tokio::test]
async fn get_bundle_allows_agent_identity_business_plan() {
let bundle = test_bundle();
let fetcher = Arc::new(StaticBundleClient::new(bundle.clone()));
let codex_home = tempdir().expect("tempdir");
let service = CloudConfigBundleService::new(
auth_manager_with_agent_identity_business_plan().await,
fetcher.clone(),
codex_home.path().to_path_buf(),
CLOUD_CONFIG_BUNDLE_TIMEOUT,
);
assert_eq!(service.load_startup_bundle().await, Ok(Some(bundle)));
assert_eq!(fetcher.request_count.load(Ordering::SeqCst), 1);
assert!(
codex_home
.path()
.join(CLOUD_CONFIG_BUNDLE_CACHE_FILENAME)
.exists()
);
}
#[tokio::test]
async fn get_bundle_skips_team_like_usage_based_plan() {
let fetcher = Arc::new(StaticBundleClient::new(test_bundle()));
@@ -635,6 +683,7 @@ async fn get_bundle_recovers_after_unauthorized_reload() {
auth_home.path().to_path_buf(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::default(),
)
@@ -690,6 +739,7 @@ async fn get_bundle_recovers_after_unauthorized_reload_updates_cache_identity()
auth_home.path().to_path_buf(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::default(),
)
@@ -753,6 +803,7 @@ async fn get_bundle_surfaces_auth_recovery_message() {
auth_home.path().to_path_buf(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::default(),
)
@@ -818,6 +869,7 @@ async fn get_bundle_unauthorized_without_recovery_uses_generic_message() {
auth_home.path().to_path_buf(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::default(),
)
codex-rs/cloud-tasks/src/util.rs
+1
View File
@@ -49,6 +49,7 @@ pub async fn load_auth_manager(chatgpt_base_url: Option<String>) -> Option<AuthM
config.codex_home.to_path_buf(),
/*enable_codex_api_key_env*/ false,
config.cli_auth_credentials_store_mode,
config.forced_chatgpt_workspace_id.clone(),
chatgpt_base_url.or(Some(config.chatgpt_base_url.clone())),
config.auth_keyring_backend_kind(),
)
codex-rs/core/config.schema.json
+6
View File
@@ -692,6 +692,9 @@
"unified_exec_zsh_fork": {
"type": "boolean"
},
"use_agent_identity": {
"type": "boolean"
},
"use_legacy_landlock": {
"type": "boolean"
},
@@ -4932,6 +4935,9 @@
"unified_exec_zsh_fork": {
"type": "boolean"
},
"use_agent_identity": {
"type": "boolean"
},
"use_legacy_landlock": {
"type": "boolean"
},
codex-rs/exec/src/lib.rs
-1
View File
@@ -475,7 +475,6 @@ 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/features/src/lib.rs
+8
View File
@@ -228,6 +228,8 @@ pub enum Feature {
PreventIdleSleep,
/// Enable remote compaction v2 over the normal Responses API.
RemoteCompactionV2,
/// Use Agent Identity for ChatGPT-authenticated sessions.
UseAgentIdentity,
/// Enable workspace dependency support.
WorkspaceDependencies,
@@ -1314,6 +1316,12 @@ pub const FEATURES: &[FeatureSpec] = &[
stage: Stage::Stable,
default_enabled: true,
},
FeatureSpec {
id: Feature::UseAgentIdentity,
key: "use_agent_identity",
stage: Stage::UnderDevelopment,
default_enabled: false,
},
FeatureSpec {
id: Feature::WorkspaceDependencies,
key: "workspace_dependencies",
codex-rs/login/src/auth/agent_identity.rs
+379 -43
View File
@@ -1,83 +1,318 @@
use std::future::Future;
use std::sync::Arc;
use codex_agent_identity::AgentIdentityKey;
use codex_agent_identity::ChatGptEnvironment;
use codex_agent_identity::build_abom;
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::is_retryable_registration_error;
use codex_agent_identity::public_key_ssh_from_private_key_pkcs8_base64;
use codex_agent_identity::register_agent_identity;
use codex_agent_identity::register_agent_task;
use codex_protocol::account::PlanType as AccountPlanType;
use codex_protocol::protocol::SessionSource;
use thiserror::Error;
use crate::default_client::build_reqwest_client;
use super::storage::AgentIdentityAuthRecord;
#[derive(Clone, Debug)]
pub struct AgentIdentityAuth {
inner: Arc<AgentIdentityAuthInner>,
pub(super) const MAX_AGENT_IDENTITY_BOOTSTRAP_ATTEMPTS: usize = 3;
pub(super) fn agent_identity_authapi_base_url(
chatgpt_base_url: Option<&str>,
) -> std::io::Result<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())
}
#[derive(Debug)]
struct AgentIdentityAuthInner {
record: AgentIdentityAuthRecord,
run_task_id: String,
pub(super) fn require_agent_identity_authapi_base_url(
agent_identity_authapi_base_url: Option<&str>,
) -> std::io::Result<&str> {
agent_identity_authapi_base_url.ok_or_else(|| {
std::io::Error::other(
"Agent Identity only supports production and staging ChatGPT environments",
)
})
}
#[derive(Debug, Error)]
pub enum AgentIdentityAuthError {
#[error(
"agent identity bootstrap unavailable after {attempts} attempts during {operation}: {message}"
)]
BootstrapUnavailable {
operation: &'static str,
attempts: usize,
message: String,
},
}
impl AgentIdentityAuthError {
pub fn is_bootstrap_unavailable(error: &std::io::Error) -> bool {
matches!(
error
.get_ref()
.and_then(|source| source.downcast_ref::<Self>()),
Some(Self::BootstrapUnavailable { .. })
)
}
}
#[derive(Debug, Error)]
#[error("retryable agent identity registration failure: {message}")]
pub(super) struct RetryableAgentIdentityRegistrationError {
message: String,
}
impl RetryableAgentIdentityRegistrationError {
pub(super) fn new(message: String) -> Self {
Self { message }
}
}
#[derive(Clone, Debug)]
pub struct AgentIdentityAuth {
record: Arc<AgentIdentityAuthRecord>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(super) struct ManagedChatGptAgentIdentityBinding {
pub(super) account_id: String,
pub(super) chatgpt_user_id: String,
pub(super) email: String,
pub(super) plan_type: AccountPlanType,
pub(super) chatgpt_account_is_fedramp: bool,
pub(super) access_token: String,
}
impl AgentIdentityAuth {
pub async fn load(
record: AgentIdentityAuthRecord,
pub async fn from_record(
mut 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_for_record(&record),
)
.await
.map_err(std::io::Error::other)?;
public_key_ssh_from_private_key_pkcs8_base64(&record.agent_private_key)
.map_err(std::io::Error::other)?;
if record_needs_task_registration(&record) {
record.task_id = Some(
register_task_for_record_with_retries(&record, agent_identity_authapi_base_url)
.await?,
);
}
Ok(Self {
inner: Arc::new(AgentIdentityAuthInner {
record,
run_task_id,
}),
record: Arc::new(record),
})
}
pub async fn from_jwt(
jwt: &str,
chatgpt_base_url: &str,
agent_identity_authapi_base_url: &str,
) -> std::io::Result<Self> {
let record = verified_record_from_jwt(jwt, chatgpt_base_url).await?;
Self::from_record(record, agent_identity_authapi_base_url).await
}
#[cfg(test)]
fn from_initialized_record(record: AgentIdentityAuthRecord, run_task_id: String) -> Self {
fn from_initialized_record(mut record: AgentIdentityAuthRecord, run_task_id: String) -> Self {
record.task_id = Some(run_task_id);
Self {
inner: Arc::new(AgentIdentityAuthInner {
record,
run_task_id,
}),
record: Arc::new(record),
}
}
pub fn record(&self) -> &AgentIdentityAuthRecord {
&self.inner.record
self.record.as_ref()
}
pub fn run_task_id(&self) -> &str {
&self.inner.run_task_id
match self.record.task_id.as_deref() {
Some(task_id) => task_id,
None => unreachable!("AgentIdentityAuth should only be constructed with a task_id"),
}
}
pub fn account_id(&self) -> &str {
&self.inner.record.account_id
&self.record.account_id
}
pub fn chatgpt_user_id(&self) -> &str {
&self.inner.record.chatgpt_user_id
&self.record.chatgpt_user_id
}
pub fn email(&self) -> &str {
&self.inner.record.email
&self.record.email
}
pub fn plan_type(&self) -> AccountPlanType {
self.inner.record.plan_type
self.record.plan_type
}
pub fn is_fedramp_account(&self) -> bool {
self.inner.record.chatgpt_account_is_fedramp
self.record.chatgpt_account_is_fedramp
}
}
pub(super) async fn register_managed_chatgpt_agent_identity(
binding: ManagedChatGptAgentIdentityBinding,
agent_identity_authapi_base_url: &str,
session_source: SessionSource,
) -> std::io::Result<AgentIdentityAuth> {
let key_material = generate_agent_key_material().map_err(std::io::Error::other)?;
let client = build_reqwest_client();
let runtime_id = retry_registration(|| async {
register_agent_identity(
&client,
agent_identity_authapi_base_url,
&binding.access_token,
binding.chatgpt_account_is_fedramp,
&key_material,
build_abom(session_source.clone()),
vec!["responsesapi".to_string()],
)
.await
.map_err(|err| {
if is_retryable_registration_error(&err) {
std::io::Error::other(RetryableAgentIdentityRegistrationError::new(
err.to_string(),
))
} else {
std::io::Error::other(err)
}
})
})
.await
.map_err(|err| classify_bootstrap_error("agent identity registration", err))?;
let record = AgentIdentityAuthRecord {
agent_runtime_id: runtime_id,
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,
task_id: None,
};
AgentIdentityAuth::from_record(record, agent_identity_authapi_base_url)
.await
.map_err(|err| classify_bootstrap_error("agent task registration", err))
}
pub(super) async fn verified_record_from_jwt(
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())
}
pub(super) fn record_needs_task_registration(record: &AgentIdentityAuthRecord) -> bool {
record
.task_id
.as_deref()
.is_none_or(|task_id| task_id.trim().is_empty())
}
pub(super) fn record_matches_managed_chatgpt_binding(
record: &AgentIdentityAuthRecord,
binding: &ManagedChatGptAgentIdentityBinding,
) -> bool {
record.account_id == binding.account_id
&& record.chatgpt_user_id == binding.chatgpt_user_id
&& public_key_ssh_from_private_key_pkcs8_base64(&record.agent_private_key).is_ok()
}
pub(super) fn classify_bootstrap_error(
operation: &'static str,
err: std::io::Error,
) -> std::io::Error {
if is_retryable_io_registration_error(&err) {
std::io::Error::other(AgentIdentityAuthError::BootstrapUnavailable {
operation,
attempts: MAX_AGENT_IDENTITY_BOOTSTRAP_ATTEMPTS,
message: err.to_string(),
})
} else {
err
}
}
pub(super) fn is_retryable_io_registration_error(err: &std::io::Error) -> bool {
err.get_ref().is_some_and(
<dyn std::error::Error + std::marker::Send + std::marker::Sync + 'static>::is::<
RetryableAgentIdentityRegistrationError,
>,
)
}
pub(super) async fn retry_registration<T, F, Fut>(mut operation: F) -> std::io::Result<T>
where
F: FnMut() -> Fut,
Fut: Future<Output = std::io::Result<T>>,
{
let mut attempt = 1;
loop {
match operation().await {
Ok(value) => return Ok(value),
Err(err)
if attempt < MAX_AGENT_IDENTITY_BOOTSTRAP_ATTEMPTS
&& is_retryable_io_registration_error(&err) =>
{
tracing::warn!(
attempt,
max_attempts = MAX_AGENT_IDENTITY_BOOTSTRAP_ATTEMPTS,
error = %err,
"agent identity registration attempt failed; retrying"
);
attempt += 1;
}
Err(err) => return Err(err),
}
}
}
async fn register_task_for_record_with_retries(
record: &AgentIdentityAuthRecord,
agent_identity_authapi_base_url: &str,
) -> std::io::Result<String> {
retry_registration(|| async {
register_task_for_record(record, agent_identity_authapi_base_url).await
})
.await
}
async fn register_task_for_record(
record: &AgentIdentityAuthRecord,
agent_identity_authapi_base_url: &str,
) -> std::io::Result<String> {
register_agent_task(
&build_reqwest_client(),
agent_identity_authapi_base_url,
key_for_record(record),
)
.await
.map_err(|err| {
if is_retryable_registration_error(&err) {
std::io::Error::other(RetryableAgentIdentityRegistrationError::new(
err.to_string(),
))
} else {
std::io::Error::other(err)
}
})
}
fn key_for_record(record: &AgentIdentityAuthRecord) -> AgentIdentityKey<'_> {
AgentIdentityKey {
agent_runtime_id: &record.agent_runtime_id,
@@ -111,6 +346,7 @@ mod tests {
email: "agent@example.com".to_string(),
plan_type: AccountPlanType::Plus,
chatgpt_account_is_fedramp: false,
task_id: None,
}
}
@@ -120,7 +356,7 @@ mod tests {
}
#[tokio::test]
async fn load_registers_run_task() -> anyhow::Result<()> {
async fn from_record_registers_task() -> anyhow::Result<()> {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/v1/agent/agent-runtime-1/task/register"))
@@ -131,9 +367,11 @@ mod tests {
.mount(&server)
.await;
let auth =
AgentIdentityAuth::load(agent_identity_record_with_generated_key(), &server.uri())
.await?;
let auth = AgentIdentityAuth::from_record(
agent_identity_record_with_generated_key(),
&server.uri(),
)
.await?;
assert_eq!(auth.run_task_id(), "task-run-1");
let requests = server
@@ -152,6 +390,38 @@ mod tests {
Ok(())
}
#[tokio::test]
async fn from_jwt_registers_task() -> anyhow::Result<()> {
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("/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 record = agent_identity_record_with_generated_key();
let jwt = signed_agent_identity_jwt(&record)?;
let auth = AgentIdentityAuth::from_jwt(
&jwt,
&format!("{}/backend-api", server.uri()),
&server.uri(),
)
.await?;
assert_eq!(auth.record().agent_runtime_id, "agent-runtime-1");
assert_eq!(auth.run_task_id(), "task-run-1");
Ok(())
}
#[test]
fn run_task_is_shared_across_clones() {
let auth = AgentIdentityAuth::from_initialized_record(
@@ -160,12 +430,12 @@ mod tests {
);
let cloned = auth.clone();
assert!(Arc::ptr_eq(&auth.inner, &cloned.inner));
assert!(Arc::ptr_eq(&auth.record, &cloned.record));
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<()> {
async fn from_record_retries_transient_registration() -> anyhow::Result<()> {
let server = MockServer::start().await;
let request_count = Arc::new(AtomicUsize::new(0));
let response_count = Arc::clone(&request_count);
@@ -183,14 +453,80 @@ mod tests {
.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?;
let auth = AgentIdentityAuth::from_record(
agent_identity_record_with_generated_key(),
&server.uri(),
)
.await?;
assert_eq!(request_count.load(Ordering::SeqCst), 2);
assert_eq!(auth.run_task_id(), "task-run-1");
Ok(())
}
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-----"#;
}
codex-rs/login/src/auth/auth_tests.rs
+499 -24
View File
@@ -6,6 +6,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;
@@ -14,11 +15,14 @@ use pretty_assertions::assert_eq;
use serde::Serialize;
use serde_json::json;
use std::sync::Arc;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use tempfile::TempDir;
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;
@@ -95,7 +99,7 @@ fn login_with_api_key_overwrites_existing_auth_json() {
}
#[tokio::test]
async fn login_with_access_token_writes_only_token() {
async fn login_with_access_token_writes_agent_identity_jwt() {
let dir = tempdir().unwrap();
let auth_path = dir.path().join("auth.json");
let record = agent_identity_record(WORKSPACE_ID_ALLOWED);
@@ -108,7 +112,8 @@ async fn login_with_access_token_writes_only_token() {
.expect(1)
.mount(&server)
.await;
let chatgpt_base_url = format!("{}/backend-api", server.uri());
let authapi_base_url = server.uri();
let chatgpt_base_url = format!("{authapi_base_url}/backend-api");
super::login_with_access_token(
dir.path(),
@@ -127,14 +132,75 @@ 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(),
Some(agent_identity.as_str())
auth.agent_identity,
Some(AgentIdentityStorage::Jwt(agent_identity))
);
assert!(auth.tokens.is_none(), "tokens should be cleared");
assert!(auth.openai_api_key.is_none(), "API key should be cleared");
server.verify().await;
}
#[tokio::test]
#[serial(codex_auth_env)]
async fn stored_agent_identity_jwt_keeps_auth_json_unchanged() -> anyhow::Result<()> {
let _access_token_guard = remove_access_token_env_var();
let codex_home = tempdir()?;
let record = agent_identity_record(WORKSPACE_ID_ALLOWED);
let agent_identity =
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"))
.respond_with(ResponseTemplate::new(200).set_body_json(test_jwks_body()))
.expect(1)
.mount(&server)
.await;
mock_agent_task_registration(&server, "", &record.agent_runtime_id, "task-id").await;
let authapi_base_url = server.uri();
let chatgpt_base_url = format!("{authapi_base_url}/backend-api");
save_auth(
codex_home.path(),
&AuthDotJson {
auth_mode: Some(ApiAuthMode::AgentIdentity),
openai_api_key: None,
tokens: None,
last_refresh: None,
agent_identity: Some(AgentIdentityStorage::Jwt(agent_identity.clone())),
personal_access_token: None,
bedrock_api_key: None,
},
AuthCredentialsStoreMode::File,
AuthKeyringBackendKind::Direct,
)?;
let auth = super::load_auth(
codex_home.path(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*forced_chatgpt_workspace_id*/ None,
Some(&chatgpt_base_url),
AuthKeyringBackendKind::Direct,
Some(&authapi_base_url),
)
.await?
.expect("auth should load");
let CodexAuth::AgentIdentity(agent_identity_auth) = auth else {
panic!("stored JWT should load as agent identity auth");
};
assert_eq!(agent_identity_auth.run_task_id(), "task-id");
let storage = FileAuthStorage::new(codex_home.path().to_path_buf());
let auth = storage
.try_read_auth_json(&get_auth_file(codex_home.path()))
.expect("auth.json should parse");
assert_eq!(
auth.agent_identity,
Some(AgentIdentityStorage::Jwt(agent_identity))
);
server.verify().await;
Ok(())
}
#[tokio::test]
#[serial(codex_auth_env)]
async fn login_with_access_token_writes_only_personal_access_token() {
@@ -278,6 +344,387 @@ async fn login_with_access_token_rejects_invalid_jwt() {
);
}
#[tokio::test]
#[serial(codex_auth_env)]
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,
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::Direct,
/*agent_identity_authapi_base_url*/ None,
)
.await?
.expect("auth should load");
assert!(
auth.agent_identity_auth(
AgentIdentityAuthPolicy::JwtOnly,
/*agent_identity_authapi_base_url*/ None,
/*forced_chatgpt_workspace_id*/ None,
SessionSource::Cli,
)
.await?
.is_none()
);
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/v1/agent/register"))
.and(header("authorization", "Bearer test-access-token"))
.and(body_partial_json(json!({
"abom": {
"agent_harness_id": "codex-cli",
},
"capabilities": ["responsesapi"],
"ttl": null,
})))
.respond_with(ResponseTemplate::new(/*s*/ 200).set_body_json(json!({
"agent_runtime_id": "agent-runtime-123",
})))
.expect(/*r*/ 1)
.mount(&server)
.await;
mock_agent_task_registration(&server, "", "agent-runtime-123", "task-123").await;
let agent_auth = auth
.agent_identity_auth(
AgentIdentityAuthPolicy::ChatGptAuth,
Some(&server.uri()),
/*forced_chatgpt_workspace_id*/ None,
SessionSource::Cli,
)
.await?
.expect("agent identity should register");
let reused = auth
.agent_identity_auth(
AgentIdentityAuthPolicy::ChatGptAuth,
Some(&server.uri()),
/*forced_chatgpt_workspace_id*/ None,
SessionSource::Cli,
)
.await?
.expect("agent identity should be reused");
assert_eq!(
agent_auth.record().agent_runtime_id,
reused.record().agent_runtime_id
);
assert_eq!(agent_auth.run_task_id(), "task-123");
assert_eq!(reused.run_task_id(), "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!(agent_auth.record().task_id.as_deref(), Some("task-123"));
assert_eq!(reused.record().task_id.as_deref(), Some("task-123"));
let persisted = auth
.stored_managed_chatgpt_agent_identity_record("account-123")
.expect("identity should persist");
assert_eq!(persisted.agent_runtime_id, "agent-runtime-123");
assert_eq!(persisted.task_id.as_deref(), Some("task-123"));
let reloaded = super::load_auth(
codex_home.path(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::Direct,
/*agent_identity_authapi_base_url*/ None,
)
.await?
.expect("auth should reload");
let reloaded_agent_auth = reloaded
.agent_identity_auth(
AgentIdentityAuthPolicy::ChatGptAuth,
Some(&server.uri()),
/*forced_chatgpt_workspace_id*/ None,
SessionSource::Cli,
)
.await?
.expect("agent identity should reload from storage");
assert_eq!(
reloaded_agent_auth.record().agent_runtime_id,
"agent-runtime-123"
);
assert_eq!(reloaded_agent_auth.run_task_id(), "task-123");
Ok(())
}
#[tokio::test]
#[serial(codex_auth_env)]
async fn chatgpt_auth_retries_transient_agent_identity_registration() -> 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,
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::Direct,
/*agent_identity_authapi_base_url*/ None,
)
.await?
.expect("auth should load");
let server = MockServer::start().await;
let registration_count = Arc::new(AtomicUsize::new(0));
let response_count = Arc::clone(&registration_count);
Mock::given(method("POST"))
.and(path("/v1/agent/register"))
.respond_with(move |_request: &wiremock::Request| {
if response_count.fetch_add(1, Ordering::SeqCst) < 2 {
ResponseTemplate::new(/*status*/ 503)
} else {
ResponseTemplate::new(/*status*/ 200).set_body_json(json!({
"agent_runtime_id": "agent-runtime-123",
}))
}
})
.expect(/*requests*/ 3)
.mount(&server)
.await;
mock_agent_task_registration(&server, "", "agent-runtime-123", "task-123").await;
let agent_auth = auth
.agent_identity_auth(
AgentIdentityAuthPolicy::ChatGptAuth,
Some(&server.uri()),
/*forced_chatgpt_workspace_id*/ None,
SessionSource::Cli,
)
.await?
.expect("agent identity should register after retries");
assert_eq!(registration_count.load(Ordering::SeqCst), 3);
assert_eq!(agent_auth.record().agent_runtime_id, "agent-runtime-123");
assert_eq!(agent_auth.record().task_id.as_deref(), Some("task-123"));
assert_eq!(
auth.stored_managed_chatgpt_agent_identity_record("account-123")
.and_then(|record| record.task_id),
Some("task-123".to_string())
);
Ok(())
}
#[tokio::test]
#[serial(codex_auth_env)]
async fn chatgpt_auth_registration_retry_exhaustion_is_fallback_eligible() -> 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,
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::Direct,
/*agent_identity_authapi_base_url*/ None,
)
.await?
.expect("auth should load");
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/v1/agent/register"))
.respond_with(ResponseTemplate::new(/*status*/ 503))
.expect(/*requests*/ 3)
.mount(&server)
.await;
let err = auth
.agent_identity_auth(
AgentIdentityAuthPolicy::ChatGptAuth,
Some(&server.uri()),
/*forced_chatgpt_workspace_id*/ None,
SessionSource::Cli,
)
.await
.expect_err("retry exhaustion should return an error");
assert!(AgentIdentityAuthError::is_bootstrap_unavailable(&err));
assert!(
auth.stored_managed_chatgpt_agent_identity_record("account-123")
.is_none()
);
Ok(())
}
#[tokio::test]
#[serial(codex_auth_env)]
async fn chatgpt_auth_task_registration_retry_exhaustion_is_fallback_eligible() -> 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 mut record = agent_identity_record("account-123");
record.chatgpt_user_id = "user-12345".to_string();
record.email = "user@example.com".to_string();
let storage = FileAuthStorage::new(codex_home.path().to_path_buf());
let auth_path = get_auth_file(codex_home.path());
let mut auth_json = storage.try_read_auth_json(&auth_path)?;
auth_json.agent_identity = Some(AgentIdentityStorage::Record(record.clone()));
storage.save(&auth_json)?;
let auth = super::load_auth(
codex_home.path(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::Direct,
/*agent_identity_authapi_base_url*/ None,
)
.await?
.expect("auth should load");
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path(format!(
"/v1/agent/{}/task/register",
record.agent_runtime_id
)))
.respond_with(ResponseTemplate::new(/*status*/ 503))
.expect(/*requests*/ 3)
.mount(&server)
.await;
let err = auth
.agent_identity_auth(
AgentIdentityAuthPolicy::ChatGptAuth,
Some(&server.uri()),
/*forced_chatgpt_workspace_id*/ None,
SessionSource::Cli,
)
.await
.expect_err("task retry exhaustion should return an error");
assert!(AgentIdentityAuthError::is_bootstrap_unavailable(&err));
record.task_id = None;
assert_eq!(
auth.stored_managed_chatgpt_agent_identity_record("account-123"),
Some(record)
);
Ok(())
}
#[tokio::test]
#[serial(codex_auth_env)]
async fn chatgpt_auth_non_retryable_registration_error_is_hard_failure() -> 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,
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::Direct,
/*agent_identity_authapi_base_url*/ None,
)
.await?
.expect("auth should load");
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/v1/agent/register"))
.respond_with(ResponseTemplate::new(/*status*/ 403))
.expect(/*requests*/ 1)
.mount(&server)
.await;
let err = auth
.agent_identity_auth(
AgentIdentityAuthPolicy::ChatGptAuth,
Some(&server.uri()),
/*forced_chatgpt_workspace_id*/ None,
SessionSource::Cli,
)
.await
.expect_err("hard registration failure should return an error");
assert!(!AgentIdentityAuthError::is_bootstrap_unavailable(&err));
assert!(
auth.stored_managed_chatgpt_agent_identity_record("account-123")
.is_none()
);
Ok(())
}
#[tokio::test]
async fn agent_identity_jwt_task_registration_retry_exhaustion_is_strict() -> anyhow::Result<()> {
let record = agent_identity_record(WORKSPACE_ID_ALLOWED);
let agent_identity =
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"))
.respond_with(ResponseTemplate::new(200).set_body_json(test_jwks_body()))
.expect(1)
.mount(&server)
.await;
Mock::given(method("POST"))
.and(path(format!(
"/v1/agent/{}/task/register",
record.agent_runtime_id
)))
.respond_with(ResponseTemplate::new(/*status*/ 503))
.expect(/*requests*/ 3)
.mount(&server)
.await;
let authapi_base_url = server.uri();
let chatgpt_base_url = format!("{authapi_base_url}/backend-api");
let err = CodexAuth::from_agent_identity_jwt_with_authapi_base_url(
&agent_identity,
Some(&chatgpt_base_url),
&authapi_base_url,
)
.await
.expect_err("agent identity jwt task retry exhaustion should fail");
assert!(!AgentIdentityAuthError::is_bootstrap_unavailable(&err));
Ok(())
}
#[tokio::test]
async fn login_with_access_token_rejects_unsigned_jwt() {
let dir = tempdir().unwrap();
@@ -290,7 +737,8 @@ async fn login_with_access_token_rejects_unsigned_jwt() {
.expect(1)
.mount(&server)
.await;
let chatgpt_base_url = format!("{}/backend-api", server.uri());
let authapi_base_url = server.uri();
let chatgpt_base_url = format!("{authapi_base_url}/backend-api");
super::login_with_access_token(
dir.path(),
@@ -457,6 +905,7 @@ async fn unauthorized_recovery_reports_mode_and_step_names() {
dir.path().to_path_buf(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::default(),
)
@@ -828,7 +1277,6 @@ async fn build_config(
forced_login_method,
forced_chatgpt_workspace_id,
chatgpt_base_url: None,
agent_identity_authapi_base_url: None,
}
}
@@ -879,7 +1327,7 @@ fn remove_access_token_env_var() -> EnvVarGuard {
#[serial(codex_auth_env)]
async fn load_auth_reads_access_token_from_env() {
let codex_home = tempdir().unwrap();
let expected_record = agent_identity_record(WORKSPACE_ID_ALLOWED);
let mut expected_record = agent_identity_record(WORKSPACE_ID_ALLOWED);
let agent_identity =
signed_agent_identity_jwt(&expected_record, json!(expected_record.plan_type))
.expect("signed agent identity");
@@ -898,6 +1346,7 @@ async fn load_auth_reads_access_token_from_env() {
.expect(1)
.mount(&server)
.await;
expected_record.task_id = Some("task-123".to_string());
let _access_token_guard = EnvVarGuard::set(CODEX_ACCESS_TOKEN_ENV_VAR, &agent_identity);
let authapi_base_url = server.uri();
@@ -1003,11 +1452,10 @@ async fn auth_manager_rejects_env_personal_access_token_workspace_mismatch() {
let _access_token_guard =
EnvVarGuard::set(CODEX_ACCESS_TOKEN_ENV_VAR, "at-env-workspace-mismatch");
let manager = AuthManager::new_with_workspace_restriction(
let manager = AuthManager::new(
codex_home.path().to_path_buf(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*forced_chatgpt_workspace_id*/
Some(vec![WORKSPACE_ID_ALLOWED.to_string()]),
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::default(),
@@ -1054,11 +1502,10 @@ async fn auth_manager_rejects_stored_personal_access_token_workspace_mismatch()
.await
.expect("personal access token login should succeed");
let manager = AuthManager::new_with_workspace_restriction(
let manager = AuthManager::new(
codex_home.path().to_path_buf(),
/*enable_codex_api_key_env*/ false,
auth_credentials_store_mode,
/*forced_chatgpt_workspace_id*/
Some(vec![WORKSPACE_ID_ALLOWED.to_string()]),
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::default(),
@@ -1092,6 +1539,7 @@ async fn personal_access_token_does_not_offer_unauthorized_recovery() {
codex_home.path().to_path_buf(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::default(),
)
@@ -1233,7 +1681,6 @@ 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)
@@ -1341,7 +1788,7 @@ async fn enforce_login_restrictions_logs_out_for_agent_identity_workspace_mismat
openai_api_key: None,
tokens: None,
last_refresh: None,
agent_identity: Some(agent_identity),
agent_identity: Some(AgentIdentityStorage::Jwt(agent_identity)),
personal_access_token: None,
bedrock_api_key: None,
},
@@ -1357,15 +1804,21 @@ 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)
.await
.expect_err("expected workspace mismatch to error");
assert!(err.to_string().contains(&format!(
"current credentials belong to {WORKSPACE_ID_DISALLOWED}"
)));
let err = super::enforce_login_restrictions_with_agent_identity_authapi_base_url(
&config,
Some(&authapi_base_url),
)
.await
.expect_err("expected workspace mismatch to error");
let message = err.to_string();
assert!(
message.contains(&format!(
"current credentials belong to {WORKSPACE_ID_DISALLOWED}"
)),
"{message}"
);
assert!(
!codex_home.path().join("auth.json").exists(),
"auth.json should be removed on mismatch"
@@ -1437,9 +1890,28 @@ fn agent_identity_record(account_id: &str) -> AgentIdentityAuthRecord {
email: "user@example.com".to_string(),
plan_type: AccountPlanType::Pro,
chatgpt_account_is_fedramp: false,
task_id: None,
}
}
async fn mock_agent_task_registration(
server: &MockServer,
path_prefix: &str,
agent_runtime_id: &str,
task_id: &str,
) {
Mock::given(method("POST"))
.and(path(format!(
"{path_prefix}/v1/agent/{agent_runtime_id}/task/register"
)))
.respond_with(ResponseTemplate::new(/*s*/ 200).set_body_json(json!({
"task_id": task_id,
})))
.expect(/*r*/ 1)
.mount(server)
.await;
}
fn fake_agent_identity_jwt(record: &AgentIdentityAuthRecord) -> std::io::Result<String> {
fake_agent_identity_jwt_with_plan_type(record, serde_json::to_value(record.plan_type)?)
}
@@ -1580,10 +2052,13 @@ async fn assert_agent_identity_plan_alias(
.await;
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");
let auth = CodexAuth::from_agent_identity_jwt_with_authapi_base_url(
&jwt,
Some(&chatgpt_base_url),
&authapi_base_url,
)
.await
.expect("agent identity auth");
pretty_assertions::assert_eq!(auth.account_plan_type(), Some(expected_plan_type));
server.verify().await;
codex-rs/login/src/auth/bedrock_api_key_tests.rs
+3
View File
@@ -60,6 +60,7 @@ async fn login_with_bedrock_api_key_replaces_openai_auth() -> anyhow::Result<()>
codex_home.path().to_path_buf(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::default(),
)
@@ -107,6 +108,7 @@ async fn logout_removes_bedrock_auth() -> anyhow::Result<()> {
codex_home.path().to_path_buf(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::default(),
)
@@ -130,6 +132,7 @@ async fn bedrock_only_auth_storage_creates_primary_auth() -> anyhow::Result<()>
codex_home.path().to_path_buf(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::default(),
)
codex-rs/login/src/auth/manager.rs
+326 -91
View File
@@ -20,8 +20,6 @@ use tokio::sync::watch;
use tracing::instrument;
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;
use codex_app_server_protocol::AuthMode as ApiAuthMode;
use codex_protocol::config_types::ForcedLoginMethod;
@@ -29,18 +27,27 @@ use codex_protocol::config_types::ModelProviderAuthInfo;
use super::access_token::CodexAccessToken;
use super::access_token::classify_codex_access_token;
use super::agent_identity::ManagedChatGptAgentIdentityBinding;
use super::agent_identity::agent_identity_authapi_base_url;
use super::agent_identity::classify_bootstrap_error;
use super::agent_identity::record_matches_managed_chatgpt_binding;
use super::agent_identity::record_needs_task_registration;
use super::agent_identity::register_managed_chatgpt_agent_identity;
use super::agent_identity::require_agent_identity_authapi_base_url;
use super::agent_identity::verified_record_from_jwt;
use super::external_bearer::BearerTokenRefresher;
use super::revoke::revoke_auth_tokens;
pub use crate::auth::agent_identity::AgentIdentityAuth;
pub use crate::auth::agent_identity::AgentIdentityAuthError;
pub use crate::auth::bedrock_api_key::BedrockApiKeyAuth;
pub use crate::auth::personal_access_token::PersonalAccessTokenAuth;
pub use crate::auth::storage::AgentIdentityAuthRecord;
pub use crate::auth::storage::AgentIdentityStorage;
pub use crate::auth::storage::AuthDotJson;
pub use crate::auth::storage::AuthKeyringBackendKind;
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;
@@ -51,6 +58,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;
@@ -65,6 +73,15 @@ pub enum CodexAuth {
BedrockApiKey(BedrockApiKeyAuth),
}
/// Policy for resolving Agent Identity auth from a broader Codex auth snapshot.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AgentIdentityAuthPolicy {
/// Use Agent Identity auth only when the current auth is already Agent Identity.
JwtOnly,
/// Allow managed ChatGPT auth to register or reuse Agent Identity auth.
ChatGptAuth,
}
impl PartialEq for CodexAuth {
fn eq(&self, other: &Self) -> bool {
match (self, other) {
@@ -222,7 +239,6 @@ impl CodexAuth {
agent_identity_authapi_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."));
@@ -230,17 +246,34 @@ 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.clone() else {
return Err(std::io::Error::other(
"agent identity auth is missing an agent identity token.",
"agent identity auth is missing agent identity auth material.",
));
};
return Self::from_agent_identity_jwt(
&agent_identity,
chatgpt_base_url,
agent_identity_authapi_base_url,
)
.await;
let base_url = chatgpt_base_url
.unwrap_or(ChatGptEnvironment::default().chatgpt_base_url())
.trim_end_matches('/')
.to_string();
let agent_identity_authapi_base_url =
require_agent_identity_authapi_base_url(agent_identity_authapi_base_url)?;
match agent_identity {
AgentIdentityStorage::Jwt(jwt) => {
let auth = AgentIdentityAuth::from_jwt(
&jwt,
&base_url,
agent_identity_authapi_base_url,
)
.await?;
return Ok(Self::AgentIdentity(auth));
}
AgentIdentityStorage::Record(record) => {
let auth =
AgentIdentityAuth::from_record(record, agent_identity_authapi_base_url)
.await?;
return Ok(Self::AgentIdentity(auth));
}
}
}
if auth_mode == ApiAuthMode::PersonalAccessToken {
let Some(personal_access_token) = auth_dot_json.personal_access_token.as_deref() else {
@@ -262,7 +295,7 @@ impl CodexAuth {
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,
client: create_client(),
};
match auth_mode {
@@ -292,6 +325,8 @@ impl CodexAuth {
chatgpt_base_url: Option<&str>,
keyring_backend_kind: AuthKeyringBackendKind,
) -> std::io::Result<Option<Self>> {
let agent_identity_authapi_base_url =
agent_identity_authapi_base_url(chatgpt_base_url).ok();
load_auth(
codex_home,
/*enable_codex_api_key_env*/ false,
@@ -299,7 +334,7 @@ impl CodexAuth {
/*forced_chatgpt_workspace_id*/ None,
chatgpt_base_url,
keyring_backend_kind,
/*agent_identity_authapi_base_url*/ None,
agent_identity_authapi_base_url.as_deref(),
)
.await
}
@@ -307,19 +342,28 @@ 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 agent_identity_authapi_base_url = agent_identity_authapi_base_url(chatgpt_base_url)?;
Self::from_agent_identity_jwt_with_authapi_base_url(
jwt,
chatgpt_base_url,
&agent_identity_authapi_base_url,
)
.await
}
async fn from_agent_identity_jwt_with_authapi_base_url(
jwt: &str,
chatgpt_base_url: Option<&str>,
agent_identity_authapi_base_url: &str,
) -> std::io::Result<Self> {
let base_url = chatgpt_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?;
let auth = AgentIdentityAuth::load(record, &resolved_authapi_base_url).await?;
Ok(Self::AgentIdentity(auth))
Ok(Self::AgentIdentity(
AgentIdentityAuth::from_jwt(jwt, &base_url, agent_identity_authapi_base_url).await?,
))
}
pub async fn from_personal_access_token(access_token: &str) -> std::io::Result<Self> {
@@ -499,6 +543,89 @@ impl CodexAuth {
self.get_current_auth_json().and_then(|t| t.tokens)
}
fn stored_managed_chatgpt_agent_identity_record(
&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)
}
fn persist_managed_chatgpt_agent_identity_record(
&self,
record: AgentIdentityAuthRecord,
) -> std::io::Result<()> {
if let Self::Chatgpt(chatgpt_auth) = self {
chatgpt_auth.persist_agent_identity_record(record)?;
}
Ok(())
}
async fn agent_identity_auth(
&self,
policy: AgentIdentityAuthPolicy,
agent_identity_authapi_base_url: Option<&str>,
forced_chatgpt_workspace_id: Option<Vec<String>>,
session_source: SessionSource,
) -> std::io::Result<Option<AgentIdentityAuth>> {
match self {
Self::AgentIdentity(auth) => Ok(Some(auth.clone())),
Self::ApiKey(_)
| Self::ChatgptAuthTokens(_)
| Self::PersonalAccessToken(_)
| Self::BedrockApiKey(_) => Ok(None),
Self::Chatgpt(_) => {
if policy == AgentIdentityAuthPolicy::JwtOnly {
return Ok(None);
}
self.ensure_managed_chatgpt_agent_identity(
require_agent_identity_authapi_base_url(agent_identity_authapi_base_url)?,
forced_chatgpt_workspace_id,
session_source,
)
.await
.map(Some)
}
}
}
async fn ensure_managed_chatgpt_agent_identity(
&self,
agent_identity_authapi_base_url: &str,
forced_chatgpt_workspace_id: Option<Vec<String>>,
session_source: SessionSource,
) -> std::io::Result<AgentIdentityAuth> {
let binding =
ManagedChatGptAgentIdentityBinding::from_auth(self, forced_chatgpt_workspace_id)
.ok_or_else(|| std::io::Error::other("ChatGPT auth is unavailable"))?;
// JWT auth is loaded as CodexAuth::AgentIdentity; this path only reuses
// records created by the managed ChatGPT Agent Identity bootstrap.
if let Some(record) = self.stored_managed_chatgpt_agent_identity_record(&binding.account_id)
&& record_matches_managed_chatgpt_binding(&record, &binding)
{
let should_persist = record_needs_task_registration(&record);
let auth = AgentIdentityAuth::from_record(record, agent_identity_authapi_base_url)
.await
.map_err(|err| classify_bootstrap_error("agent task registration", err))?;
if should_persist {
self.persist_managed_chatgpt_agent_identity_record(auth.record().clone())?;
}
return Ok(auth);
}
let auth = register_managed_chatgpt_agent_identity(
binding,
agent_identity_authapi_base_url,
session_source,
)
.await?;
self.persist_managed_chatgpt_agent_identity_record(auth.record().clone())?;
Ok(auth)
}
/// Consider this private to integration tests.
pub fn create_dummy_chatgpt_auth_for_testing() -> Self {
let auth_dot_json = AuthDotJson {
@@ -516,10 +643,9 @@ impl CodexAuth {
bedrock_api_key: None,
};
let client = create_client();
let state = ChatgptAuthState {
auth_dot_json: Arc::new(Mutex::new(Some(auth_dot_json))),
client,
client: create_client(),
};
let dummy_auth_id = NEXT_DUMMY_AUTH_ID.fetch_add(1, Ordering::Relaxed);
let storage = create_auth_storage(
@@ -537,6 +663,43 @@ impl CodexAuth {
}
}
impl ManagedChatGptAgentIdentityBinding {
fn from_auth(auth: &CodexAuth, forced_workspace_id: Option<Vec<String>>) -> Option<Self> {
if !auth.is_chatgpt_auth() {
return None;
}
let token_data = auth.get_token_data().ok()?;
let forced_workspace_id =
forced_workspace_id
.as_deref()
.and_then(|workspace_ids| match workspace_ids {
[workspace_id] if !workspace_id.is_empty() => Some(workspace_id.clone()),
_ => None,
});
let account_id = forced_workspace_id
.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,
})
}
}
impl ChatgptAuth {
fn current_auth_json(&self) -> Option<AuthDotJson> {
#[expect(clippy::unwrap_used)]
@@ -554,6 +717,31 @@ impl ChatgptAuth {
fn client(&self) -> &CodexHttpClient {
&self.state.client
}
fn persist_agent_identity_record(
&self,
record: AgentIdentityAuthRecord,
) -> std::io::Result<()> {
persist_agent_identity_record(&self.state.auth_dot_json, &self.storage, record)
}
}
fn persist_agent_identity_record(
auth_dot_json: &Arc<Mutex<Option<AuthDotJson>>>,
storage: &Arc<dyn AuthStorageBackend>,
record: AgentIdentityAuthRecord,
) -> std::io::Result<()> {
let mut guard = auth_dot_json
.lock()
.map_err(|_| std::io::Error::other("failed to lock auth state"))?;
let mut auth = storage
.load()?
.or_else(|| guard.clone())
.ok_or_else(|| std::io::Error::other("auth data is not available"))?;
auth.agent_identity = Some(AgentIdentityStorage::Record(record));
storage.save(&auth)?;
*guard = Some(auth);
Ok(())
}
pub const OPENAI_API_KEY_ENV_VAR: &str = "OPENAI_API_KEY";
@@ -582,34 +770,6 @@ fn read_non_empty_env_var(key: &str) -> 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())
}
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(
@@ -705,13 +865,13 @@ pub async fn login_with_access_token(
.unwrap_or(ChatGptEnvironment::default().chatgpt_base_url())
.trim_end_matches('/')
.to_string();
verified_agent_identity_record(jwt, &base_url).await?;
verified_record_from_jwt(jwt, &base_url).await?;
AuthDotJson {
auth_mode: Some(ApiAuthMode::AgentIdentity),
openai_api_key: None,
tokens: None,
last_refresh: None,
agent_identity: Some(jwt.to_string()),
agent_identity: Some(AgentIdentityStorage::Jwt(jwt.to_string())),
personal_access_token: None,
bedrock_api_key: None,
}
@@ -793,11 +953,23 @@ 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>>,
}
pub async fn enforce_login_restrictions(config: &AuthConfig) -> std::io::Result<()> {
let agent_identity_authapi_base_url =
agent_identity_authapi_base_url(config.chatgpt_base_url.as_deref()).ok();
enforce_login_restrictions_with_agent_identity_authapi_base_url(
config,
agent_identity_authapi_base_url.as_deref(),
)
.await
}
async fn enforce_login_restrictions_with_agent_identity_authapi_base_url(
config: &AuthConfig,
agent_identity_authapi_base_url: Option<&str>,
) -> std::io::Result<()> {
let Some(auth) = load_auth(
&config.codex_home,
/*enable_codex_api_key_env*/ true,
@@ -805,7 +977,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(),
agent_identity_authapi_base_url,
)
.await?
else {
@@ -985,11 +1157,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,
agent_identity_authapi_base_url,
)
CodexAccessToken::AgentIdentityJwt(jwt) => {
CodexAuth::from_agent_identity_jwt_with_authapi_base_url(
jwt,
chatgpt_base_url,
require_agent_identity_authapi_base_url(agent_identity_authapi_base_url)?,
)
}
.await
.map(Some),
};
@@ -1524,7 +1698,9 @@ pub struct AuthManager {
keyring_backend_kind: AuthKeyringBackendKind,
forced_chatgpt_workspace_id: RwLock<Option<Vec<String>>>,
chatgpt_base_url: Option<String>,
agent_identity_authapi_base_url: Option<String>,
refresh_lock: Semaphore,
agent_identity_lock: Semaphore,
external_auth: RwLock<Option<Arc<dyn ExternalAuth>>>,
}
@@ -1572,30 +1748,16 @@ impl Debug for AuthManager {
}
}
fn default_agent_identity_authapi_base_url() -> Option<String> {
agent_identity_authapi_base_url(/*chatgpt_base_url*/ None).ok()
}
impl AuthManager {
/// Create a new manager loading the initial auth using the provided
/// preferred auth method. Errors loading auth are swallowed; `auth()` will
/// simply return `None` in that case so callers can treat it as an
/// unauthenticated state.
pub async fn new(
codex_home: PathBuf,
enable_codex_api_key_env: bool,
auth_credentials_store_mode: AuthCredentialsStoreMode,
chatgpt_base_url: Option<String>,
keyring_backend_kind: AuthKeyringBackendKind,
) -> Self {
Self::new_with_workspace_restriction(
codex_home,
enable_codex_api_key_env,
auth_credentials_store_mode,
/*forced_chatgpt_workspace_id*/ None,
chatgpt_base_url,
keyring_backend_kind,
)
.await
}
async fn new_with_workspace_restriction(
codex_home: PathBuf,
enable_codex_api_key_env: bool,
auth_credentials_store_mode: AuthCredentialsStoreMode,
@@ -1603,6 +1765,8 @@ impl AuthManager {
chatgpt_base_url: Option<String>,
keyring_backend_kind: AuthKeyringBackendKind,
) -> Self {
let agent_identity_authapi_base_url =
agent_identity_authapi_base_url(chatgpt_base_url.as_deref()).ok();
let managed_auth = load_auth(
&codex_home,
enable_codex_api_key_env,
@@ -1610,7 +1774,7 @@ impl AuthManager {
forced_chatgpt_workspace_id.as_deref(),
chatgpt_base_url.as_deref(),
keyring_backend_kind,
/*agent_identity_authapi_base_url*/ None,
agent_identity_authapi_base_url.as_deref(),
)
.await
.ok()
@@ -1628,7 +1792,9 @@ impl AuthManager {
keyring_backend_kind,
forced_chatgpt_workspace_id: RwLock::new(forced_chatgpt_workspace_id),
chatgpt_base_url,
agent_identity_authapi_base_url,
refresh_lock: Semaphore::new(/*permits*/ 1),
agent_identity_lock: Semaphore::new(/*permits*/ 1),
external_auth: RwLock::new(None),
}
}
@@ -1650,7 +1816,9 @@ impl AuthManager {
keyring_backend_kind: AuthKeyringBackendKind::default(),
forced_chatgpt_workspace_id: RwLock::new(None),
chatgpt_base_url: None,
agent_identity_authapi_base_url: default_agent_identity_authapi_base_url(),
refresh_lock: Semaphore::new(/*permits*/ 1),
agent_identity_lock: Semaphore::new(/*permits*/ 1),
external_auth: RwLock::new(None),
})
}
@@ -1671,7 +1839,40 @@ impl AuthManager {
keyring_backend_kind: AuthKeyringBackendKind::default(),
forced_chatgpt_workspace_id: RwLock::new(None),
chatgpt_base_url: None,
agent_identity_authapi_base_url: default_agent_identity_authapi_base_url(),
refresh_lock: Semaphore::new(/*permits*/ 1),
agent_identity_lock: Semaphore::new(/*permits*/ 1),
external_auth: RwLock::new(None),
})
}
/// Create an AuthManager with a specific CodexAuth and Agent Identity AuthAPI base URL, for testing only.
#[doc(hidden)]
pub fn from_auth_for_testing_with_agent_identity_authapi_base_url(
auth: CodexAuth,
agent_identity_authapi_base_url: String,
) -> Arc<Self> {
let cached = CachedAuth {
auth: Some(auth),
permanent_refresh_failure: None,
};
let (auth_change_tx, _auth_change_rx) = watch::channel(0);
Arc::new(Self {
codex_home: PathBuf::from("non-existent"),
inner: RwLock::new(cached),
auth_change_tx,
enable_codex_api_key_env: false,
auth_credentials_store_mode: AuthCredentialsStoreMode::File,
keyring_backend_kind: AuthKeyringBackendKind::default(),
forced_chatgpt_workspace_id: RwLock::new(None),
chatgpt_base_url: None,
agent_identity_authapi_base_url: Some(
agent_identity_authapi_base_url
.trim_end_matches('/')
.to_string(),
),
refresh_lock: Semaphore::new(/*permits*/ 1),
agent_identity_lock: Semaphore::new(/*permits*/ 1),
external_auth: RwLock::new(None),
})
}
@@ -1690,7 +1891,9 @@ impl AuthManager {
keyring_backend_kind: AuthKeyringBackendKind::default(),
forced_chatgpt_workspace_id: RwLock::new(None),
chatgpt_base_url: None,
agent_identity_authapi_base_url: default_agent_identity_authapi_base_url(),
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>
)),
@@ -1736,6 +1939,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::ChatGptAuth && matches!(auth, CodexAuth::Chatgpt(_)) {
let _bootstrap_permit = self
.agent_identity_lock
.acquire()
.await
.map_err(std::io::Error::other)?;
return auth
.agent_identity_auth(
policy,
self.agent_identity_authapi_base_url.as_deref(),
self.forced_chatgpt_workspace_id(),
session_source,
)
.await;
}
auth.agent_identity_auth(
policy,
self.agent_identity_authapi_base_url.as_deref(),
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 {
@@ -1838,7 +2073,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,
self.agent_identity_authapi_base_url.as_deref(),
)
.await
.ok()
@@ -1911,6 +2146,7 @@ impl AuthManager {
codex_home: PathBuf,
enable_codex_api_key_env: bool,
auth_credentials_store_mode: AuthCredentialsStoreMode,
forced_chatgpt_workspace_id: Option<Vec<String>>,
chatgpt_base_url: Option<String>,
keyring_backend_kind: AuthKeyringBackendKind,
) -> Arc<Self> {
@@ -1919,6 +2155,7 @@ impl AuthManager {
codex_home,
enable_codex_api_key_env,
auth_credentials_store_mode,
forced_chatgpt_workspace_id,
chatgpt_base_url,
keyring_backend_kind,
)
@@ -1931,17 +2168,15 @@ impl AuthManager {
config: &impl AuthManagerConfig,
enable_codex_api_key_env: bool,
) -> Arc<Self> {
Arc::new(
Self::new_with_workspace_restriction(
config.codex_home(),
enable_codex_api_key_env,
config.cli_auth_credentials_store_mode(),
config.forced_chatgpt_workspace_id(),
Some(config.chatgpt_base_url()),
config.auth_keyring_backend_kind(),
)
.await,
Self::shared(
config.codex_home(),
enable_codex_api_key_env,
config.cli_auth_credentials_store_mode(),
config.forced_chatgpt_workspace_id(),
Some(config.chatgpt_base_url()),
config.auth_keyring_backend_kind(),
)
.await
}
pub fn unauthorized_recovery(self: &Arc<Self>) -> UnauthorizedRecovery {
codex-rs/login/src/auth/storage.rs
+30 -1
View File
@@ -51,7 +51,7 @@ 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>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub personal_access_token: Option<String>,
@@ -60,6 +60,32 @@ pub struct AuthDotJson {
pub bedrock_api_key: Option<BedrockApiKeyAuth>,
}
#[derive(Deserialize, Serialize, Clone, Debug, PartialEq, Eq)]
#[serde(untagged)]
pub enum AgentIdentityStorage {
Jwt(String),
Record(AgentIdentityAuthRecord),
}
impl AgentIdentityStorage {
pub fn has_auth_material(&self) -> bool {
match self {
Self::Jwt(jwt) => !jwt.trim().is_empty(),
Self::Record(record) => {
!record.agent_runtime_id.trim().is_empty()
&& !record.agent_private_key.trim().is_empty()
}
}
}
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)]
pub struct AgentIdentityAuthRecord {
pub agent_runtime_id: String,
@@ -69,6 +95,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 task_id: Option<String>,
}
impl AgentIdentityAuthRecord {
@@ -90,6 +118,7 @@ impl From<AgentIdentityJwtClaims> for AgentIdentityAuthRecord {
email: claims.email,
plan_type: claims.plan_type.into(),
chatgpt_account_is_fedramp: claims.chatgpt_account_is_fedramp,
task_id: None,
}
}
}
codex-rs/login/src/auth/storage_tests.rs
+34 -3
View File
@@ -81,7 +81,38 @@ 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)),
personal_access_token: None,
bedrock_api_key: None,
};
storage.save(&auth_dot_json)?;
let loaded = storage.load()?;
assert_eq!(Some(auth_dot_json), loaded);
Ok(())
}
#[tokio::test]
async fn file_storage_round_trips_registered_agent_identity_auth() -> anyhow::Result<()> {
let codex_home = tempdir()?;
let storage = FileAuthStorage::new(codex_home.path().to_path_buf());
let record = AgentIdentityAuthRecord {
agent_runtime_id: "agent-runtime-id".to_string(),
agent_private_key: "private-key".to_string(),
account_id: "account-id".to_string(),
chatgpt_user_id: "user-id".to_string(),
email: "user@example.com".to_string(),
plan_type: AccountPlanType::Pro,
chatgpt_account_is_fedramp: false,
task_id: Some("task-id".to_string()),
};
let auth_dot_json = AuthDotJson {
auth_mode: Some(AuthMode::Chatgpt),
openai_api_key: None,
tokens: None,
last_refresh: None,
agent_identity: Some(AgentIdentityStorage::Record(record)),
personal_access_token: None,
bedrock_api_key: None,
};
@@ -139,8 +170,8 @@ 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(),
Some(agent_identity_jwt.as_str())
loaded.expect("auth should load").agent_identity,
Some(AgentIdentityStorage::Jwt(agent_identity_jwt))
);
Ok(())
}
codex-rs/login/src/lib.rs
+1
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::AuthKeyringBackendKind;
codex-rs/login/tests/suite/auth_refresh.rs
+1
View File
@@ -1217,6 +1217,7 @@ impl RefreshTokenTestContext {
codex_home.path().to_path_buf(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::default(),
)
codex-rs/login/tests/suite/logout.rs
+1
View File
@@ -201,6 +201,7 @@ async fn auth_manager_logout_with_revoke_uses_cached_auth() -> Result<()> {
codex_home.path().to_path_buf(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::default(),
)
codex-rs/tui/src/lib.rs
-1
View File
@@ -1163,7 +1163,6 @@ 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
{