From 896604a12d0c7edea37b59ed340dd4e282172fa5 Mon Sep 17 00:00:00 2001 From: jif Date: Mon, 19 Oct 2026 22:08:57 +0000 Subject: [PATCH] Record OAuth provenance and login identity for account tokens (#55397) ## Why Token refresh needs the issuer and client ID used at login. Unverified token claims cannot establish that authority, and OAuth error bodies may echo credentials. ## What changed - Save `oauth_provenance` and a new `login_owner_id` on browser and device-code login. Preserve the login identifier through token rotation. - Refresh using the recorded issuer and client ID, retaining the existing endpoint override and fallback behavior for legacy records. Clear provenance when an override changes the endpoint or when a response without token rotation encounters a concurrently replaced refresh token. - Keep both metadata fields optional for older records without inferring provenance from token claims. - Remove backend token error details from structured logs and refresh error messages. - Box large futures in token refresh, MCP initialization, and embedded app-server test startup to reduce surrounding async frame sizes. ## Testing Add coverage for legacy record compatibility and metadata round trips, browser and device-code login persistence, refresh using recorded OAuth settings, metadata preservation through rotation, and provenance clearing for a superseded refresh token. GitOrigin-RevId: b6b5d2fcbe2a284a171a242c1a17128af55f1024 --- .../src/transport/remote_control/tests.rs | 2 + .../src/transport/remote_control/websocket.rs | 2 + .../app-server/tests/common/auth_fixtures.rs | 2 + codex-rs/core/src/session/session.rs | 5 +- codex-rs/exec/tests/suite/apply_patch.rs | 2 + codex-rs/login/src/auth/auth_tests.rs | 50 +++++++++-- codex-rs/login/src/auth/manager.rs | 84 +++++++++++++------ .../login/src/auth/storage_policy_tests.rs | 2 + codex-rs/login/src/auth/storage_tests.rs | 2 + codex-rs/login/src/device_code_auth.rs | 11 +-- codex-rs/login/src/oauth/mod.rs | 1 - codex-rs/login/src/server.rs | 41 ++++----- codex-rs/login/src/token_data.rs | 15 ++++ codex-rs/login/src/token_data_tests.rs | 23 +++++ codex-rs/login/tests/suite/auth_refresh.rs | 17 +++- .../login/tests/suite/device_code_login.rs | 15 +++- .../login/tests/suite/login_server_e2e.rs | 8 +- codex-rs/login/tests/suite/logout.rs | 2 + codex-rs/models-manager/src/manager_tests.rs | 2 + codex-rs/tui/src/lib.rs | 5 +- codex-rs/tui/src/local_chatgpt_auth.rs | 2 + 21 files changed, 219 insertions(+), 74 deletions(-) diff --git a/codex-rs/app-server-transport/src/transport/remote_control/tests.rs b/codex-rs/app-server-transport/src/transport/remote_control/tests.rs index 490ac1e37b..8cd8793bdf 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/tests.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/tests.rs @@ -119,6 +119,8 @@ fn remote_control_auth_dot_json(account_id: Option<&str>) -> AuthDotJson { id_token: parse_chatgpt_jwt_claims(&fake_jwt).expect("fake jwt should parse"), access_token: "Access Token".to_string(), refresh_token: "refresh-token".to_string(), + oauth_provenance: None, + login_owner_id: None, account_id: account_id.map(str::to_string), }), last_refresh: Some(chrono::Utc::now()), diff --git a/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs b/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs index 7bede8dc65..d74d3dfd8a 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs @@ -2072,6 +2072,8 @@ mod tests { id_token: parse_chatgpt_jwt_claims(&fake_jwt).expect("fake jwt should parse"), access_token: access_token.to_string(), refresh_token: "refresh-token".to_string(), + oauth_provenance: None, + login_owner_id: None, account_id: Some("account_id".to_string()), }), last_refresh: Some(Utc::now()), diff --git a/codex-rs/app-server/tests/common/auth_fixtures.rs b/codex-rs/app-server/tests/common/auth_fixtures.rs index 2fa384210e..ac519a00ce 100644 --- a/codex-rs/app-server/tests/common/auth_fixtures.rs +++ b/codex-rs/app-server/tests/common/auth_fixtures.rs @@ -178,6 +178,8 @@ pub fn write_chatgpt_auth( id_token, access_token: fixture.access_token, refresh_token: fixture.refresh_token, + oauth_provenance: None, + login_owner_id: None, account_id: fixture.account_id, }; diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index a2a3415eed..05bc138da7 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1947,12 +1947,13 @@ impl Session { } else { mcp_projection }; - sess.install_initial_mcp_runtime( + // Keep MCP startup's large future out of the surrounding session initialization frame. + Box::pin(sess.install_initial_mcp_runtime( &session_configuration, latest_auth, mcp_projection, &resolved_environments, - ) + )) .await?; sess.start_mcp_prewarm_worker(mcp_prewarm_rx, mcp_auth_changes); sess.follow_inherited_environment_configurations( diff --git a/codex-rs/exec/tests/suite/apply_patch.rs b/codex-rs/exec/tests/suite/apply_patch.rs index 112e29122b..71ab2483de 100644 --- a/codex-rs/exec/tests/suite/apply_patch.rs +++ b/codex-rs/exec/tests/suite/apply_patch.rs @@ -406,6 +406,8 @@ async fn shutdown_flushes_completed_turn_and_file_diff() -> anyhow::Result<()> { }, access_token: "test-access-token".to_string(), refresh_token: "test-refresh-token".to_string(), + oauth_provenance: None, + login_owner_id: None, account_id: None, }), last_refresh: None, diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index ab15a33871..beec332bfd 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -8,6 +8,7 @@ use codex_protocol::auth::KnownPlan as InternalKnownPlan; use codex_protocol::auth::PlanType as InternalPlanType; use codex_protocol::protocol::SessionSource; +use anyhow::Context; use base64::Engine; use codex_protocol::config_types::ForcedLoginMethod; use codex_protocol::config_types::ModelProviderAuthInfo; @@ -59,8 +60,8 @@ fn header_auth_exposes_a_valid_chatgpt_account_id() { } #[tokio::test] -async fn refresh_without_id_token() { - let codex_home = tempdir().unwrap(); +async fn refresh_without_id_token() -> anyhow::Result<()> { + let codex_home = tempdir()?; let fake_jwt = write_auth_file( AuthFileParams { openai_api_key: None, @@ -68,26 +69,59 @@ async fn refresh_without_id_token() { chatgpt_account_id: None, }, codex_home.path(), - ) - .expect("failed to write auth file"); + )?; let storage = create_auth_storage( codex_home.path().to_path_buf(), AuthCredentialsStoreMode::File, AuthKeyringBackendKind::default(), ); + let provenance = AccountOAuthProvenance { + issuer: "https://idp.example".to_string(), + client_id: "codex-client".to_string(), + }; let updated = super::persist_tokens( &storage, + "test-refresh-token", + Some(provenance.clone()), /*id_token*/ None, Some("new-access-token".to_string()), Some("new-refresh-token".to_string()), - ) - .expect("update_tokens should succeed"); + )?; - let tokens = updated.tokens.expect("tokens should exist"); + let tokens = updated.tokens.context("persisted auth is missing tokens")?; assert_eq!(tokens.id_token.raw_jwt, fake_jwt); assert_eq!(tokens.access_token, "new-access-token"); assert_eq!(tokens.refresh_token, "new-refresh-token"); + assert_eq!(tokens.oauth_provenance, Some(provenance.clone())); + + let updated = super::persist_tokens( + &storage, + "new-refresh-token", + Some(provenance.clone()), + /*id_token*/ None, + Some("later-access-token".to_string()), + /*refresh_token*/ None, + )?; + let tokens = updated.tokens.context("persisted auth is missing tokens")?; + assert_eq!(tokens.refresh_token, "new-refresh-token"); + assert_eq!(tokens.oauth_provenance, Some(provenance.clone())); + + // A response without a rotated token cannot establish authority for a token that + // another writer installed after this refresh began. + let updated = super::persist_tokens( + &storage, + "superseded-refresh-token", + Some(provenance), + /*id_token*/ None, + Some("latest-access-token".to_string()), + /*refresh_token*/ None, + )?; + assert_eq!(storage.load()?, Some(updated.clone())); + let tokens = updated.tokens.context("persisted auth is missing tokens")?; + assert_eq!(tokens.refresh_token, "new-refresh-token"); + assert_eq!(tokens.oauth_provenance, None); + Ok(()) } #[test] @@ -1064,6 +1098,8 @@ async fn pro_account_with_no_api_key_uses_chatgpt_auth() { }, access_token: "test-access-token".to_string(), refresh_token: "test-refresh-token".to_string(), + oauth_provenance: None, + login_owner_id: None, account_id: None, }), last_refresh: Some(last_refresh), diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 88c80f5af2..7bfae6c8bb 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -66,8 +66,8 @@ use crate::oauth::RefreshReason; use crate::oauth::RefreshTokenGrant; use crate::oauth::TokenEncoding; use crate::oauth::TokenEndpoint; -use crate::oauth::TokenErrorDetail; use crate::outbound_proxy::AuthRouteConfig; +use crate::token_data::AccountOAuthProvenance; use crate::token_data::TokenData; use crate::token_data::parse_chatgpt_account_user_id; use crate::token_data::parse_chatgpt_jwt_claims; @@ -408,7 +408,12 @@ impl CodexAuth { } let storage_mode = auth_dot_json.storage_mode(auth_credentials_store_mode); - let client = create_default_auth_client(&refresh_token_endpoint(), auth_route_config)?; + let provenance = auth_dot_json + .tokens + .as_ref() + .and_then(|tokens| tokens.oauth_provenance.as_ref()); + let client = + create_default_auth_client(&refresh_token_endpoint(provenance), auth_route_config)?; let state = ChatgptAuthState { auth_dot_json: Arc::new(Mutex::new(Some(auth_dot_json))), client, @@ -823,6 +828,8 @@ impl CodexAuth { id_token: Default::default(), access_token: "Access Token".to_string(), refresh_token: "test".to_string(), + oauth_provenance: None, + login_owner_id: None, account_id: Some("account_id".to_string()), }), last_refresh: Some(Utc::now()), @@ -1600,6 +1607,8 @@ async fn load_auth( // Persist refreshed tokens into auth storage and update last_refresh. fn persist_tokens( storage: &Arc, + previous_refresh_token: &str, + provenance: Option, id_token: Option, access_token: Option, refresh_token: Option, @@ -1609,6 +1618,10 @@ fn persist_tokens( .ok_or(std::io::Error::other("Token data is not available."))?; let tokens = auth_dot_json.tokens.get_or_insert_with(TokenData::default); + // If a concurrent writer replaced the stored refresh token and the response did not + // rotate it, its authority is unknown. Do not attach the earlier grant's provenance. + tokens.oauth_provenance = provenance + .filter(|_| refresh_token.is_some() || tokens.refresh_token == previous_refresh_token); if let Some(id_token) = id_token { tokens.id_token = parse_chatgpt_jwt_claims(&id_token).map_err(std::io::Error::other)?; } @@ -1627,16 +1640,16 @@ fn persist_tokens( // The caller is responsible for persisting any returned tokens. async fn request_chatgpt_token_refresh( refresh_token: String, + endpoint: &str, + client_id: &str, client: &HttpClient, refresh_reason: RefreshReason, ) -> Result { - let client_id = oauth_client_id(); - let endpoint = refresh_token_endpoint(); let oauth = OAuthClient::new( client, TokenEndpoint { - url: &endpoint, - client_id: &client_id, + url: endpoint, + client_id, encoding: TokenEncoding::Json, timeout: None, error_body_limit: ErrorBodyLimit::Unlimited, @@ -1656,12 +1669,12 @@ async fn request_chatgpt_token_refresh( return Err(denied.into()); } let status = rejection.status; - let detail = &rejection.detail; - tracing::error!(%status, ?detail, "Failed to refresh token"); - let code = detail.error_code(); + // The IdP controls the response body and may echo credential material. + tracing::error!(%status, "Failed to refresh token"); + let code = rejection.detail.error_code(); let is_invalid_grant_bad_request = status == StatusCode::BAD_REQUEST && code.is_some_and(|code| code.eq_ignore_ascii_case("invalid_grant")); - let failed = classify_refresh_token_failure(code, detail, is_invalid_grant_bad_request); + let failed = classify_refresh_token_failure(code, is_invalid_grant_bad_request); if status == StatusCode::UNAUTHORIZED || failed.reason != RefreshTokenFailedReason::Other || is_invalid_grant_bad_request @@ -1669,7 +1682,7 @@ async fn request_chatgpt_token_refresh( Err(RefreshTokenError::Permanent(failed)) } else { Err(RefreshTokenError::Transient(std::io::Error::other( - format!("Failed to refresh token: {status}: {detail}"), + format!("Failed to refresh token: {status}"), ))) } } @@ -1682,7 +1695,6 @@ async fn request_chatgpt_token_refresh( fn classify_refresh_token_failure( code: Option<&str>, - detail: &TokenErrorDetail, is_invalid_grant_bad_request: bool, ) -> RefreshTokenFailedError { let normalized_code = code.map(str::to_ascii_lowercase); @@ -1694,10 +1706,7 @@ fn classify_refresh_token_failure( }; if reason == RefreshTokenFailedReason::Other && !is_invalid_grant_bad_request { - tracing::warn!( - backend_detail = ?detail, - "Encountered unknown response while refreshing token" - ); + tracing::warn!("Encountered unknown response while refreshing token"); } let message = match reason { @@ -1727,9 +1736,12 @@ pub fn oauth_client_id() -> String { .unwrap_or_else(|| CLIENT_ID.to_string()) } -fn refresh_token_endpoint() -> String { - std::env::var(REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR) - .unwrap_or_else(|_| REFRESH_TOKEN_URL.to_string()) +fn refresh_token_endpoint(provenance: Option<&AccountOAuthProvenance>) -> String { + std::env::var(REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR).unwrap_or_else(|_| { + provenance + .map(|provenance| format!("{}/oauth/token", provenance.issuer.trim_end_matches('/'))) + .unwrap_or_else(|| REFRESH_TOKEN_URL.to_string()) + }) } impl AuthDotJson { @@ -1749,6 +1761,8 @@ impl AuthDotJson { id_token: token_info, access_token: access_token.to_string(), refresh_token: String::new(), + oauth_provenance: None, + login_owner_id: None, account_id: Some(chatgpt_account_id.to_string()), }; @@ -2938,11 +2952,11 @@ impl AuthManager { "Token data is not available.", )) })?; - self.refresh_and_persist_chatgpt_token( + Box::pin(self.refresh_and_persist_chatgpt_token( chatgpt_auth, - token_data.refresh_token, + token_data, reason, - ) + )) .await } Some( @@ -3114,14 +3128,34 @@ impl AuthManager { async fn refresh_and_persist_chatgpt_token( &self, auth: &ChatgptAuth, - refresh_token: String, + tokens: TokenData, reason: RefreshReason, ) -> Result<(), RefreshTokenError> { - let refresh_response = - request_chatgpt_token_refresh(refresh_token, auth.client(), reason).await?; + let endpoint = refresh_token_endpoint(tokens.oauth_provenance.as_ref()); + let client_id = tokens + .oauth_provenance + .as_ref() + .map(|provenance| provenance.client_id.clone()) + .unwrap_or_else(oauth_client_id); + let refresh_response = request_chatgpt_token_refresh( + tokens.refresh_token.clone(), + &endpoint, + &client_id, + auth.client(), + reason, + ) + .await?; + + // An explicit override may send refresh to another authority. Retain provenance + // only when the endpoint used by this request matches the recorded issuer. + let provenance = tokens.oauth_provenance.filter(|provenance| { + endpoint == format!("{}/oauth/token", provenance.issuer.trim_end_matches('/')) + }); persist_tokens( auth.storage(), + &tokens.refresh_token, + provenance, refresh_response.id_token, refresh_response.access_token, refresh_response.refresh_token, diff --git a/codex-rs/login/src/auth/storage_policy_tests.rs b/codex-rs/login/src/auth/storage_policy_tests.rs index 0d389b756c..c4479d914f 100644 --- a/codex-rs/login/src/auth/storage_policy_tests.rs +++ b/codex-rs/login/src/auth/storage_policy_tests.rs @@ -427,6 +427,8 @@ fn secrets_keyring_auth_storage_save_persists_and_removes_fallback_file() -> any id_token: Default::default(), access_token: "access".to_string(), refresh_token: "refresh".to_string(), + oauth_provenance: None, + login_owner_id: None, account_id: Some("account".to_string()), }), last_refresh: Some(Utc::now()), diff --git a/codex-rs/login/src/auth/storage_tests.rs b/codex-rs/login/src/auth/storage_tests.rs index 01c2e6f575..dc35215d8a 100644 --- a/codex-rs/login/src/auth/storage_tests.rs +++ b/codex-rs/login/src/auth/storage_tests.rs @@ -433,6 +433,8 @@ fn auth_with_prefix(prefix: &str) -> AuthDotJson { id_token: id_token_with_prefix(prefix), access_token: format!("{prefix}-access"), refresh_token: format!("{prefix}-refresh"), + oauth_provenance: None, + login_owner_id: None, account_id: Some(format!("{prefix}-account-id")), }), last_refresh: None, diff --git a/codex-rs/login/src/device_code_auth.rs b/codex-rs/login/src/device_code_auth.rs index f9f8f1993b..14a8e57c2c 100644 --- a/codex-rs/login/src/device_code_auth.rs +++ b/codex-rs/login/src/device_code_auth.rs @@ -223,16 +223,7 @@ pub async fn complete_device_code_login( return Err(io::Error::new(io::ErrorKind::PermissionDenied, message)); } - crate::server::persist_tokens_async( - &opts.codex_home, - /*api_key*/ None, - tokens.id_token, - tokens.access_token, - tokens.refresh_token, - opts.cli_auth_credentials_store_mode, - opts.auth_keyring_backend_kind, - ) - .await + crate::server::persist_tokens_async(&opts, /*api_key*/ None, &tokens).await } pub async fn run_device_code_login(opts: ServerOptions) -> std::io::Result<()> { diff --git a/codex-rs/login/src/oauth/mod.rs b/codex-rs/login/src/oauth/mod.rs index 8ac9a559d0..dc292d267b 100644 --- a/codex-rs/login/src/oauth/mod.rs +++ b/codex-rs/login/src/oauth/mod.rs @@ -23,7 +23,6 @@ pub(crate) use client::TokenEndpoint; pub(crate) use diagnostics::sanitize_url_for_logging; pub(crate) use error::ErrorBodyLimit; pub(crate) use error::OAuthError; -pub(crate) use error::TokenErrorDetail; pub(crate) use error::TokenRejection; pub(crate) use pkce::PkceCodes; pub(crate) use pkce::generate_pkce; diff --git a/codex-rs/login/src/server.rs b/codex-rs/login/src/server.rs index d65c95287c..d4f65c1fd7 100644 --- a/codex-rs/login/src/server.rs +++ b/codex-rs/login/src/server.rs @@ -17,7 +17,6 @@ use std::io::Write; use std::io::{self}; use std::net::SocketAddr; use std::net::TcpStream; -use std::path::Path; use std::path::PathBuf; use std::sync::Arc; use std::sync::LazyLock; @@ -52,6 +51,7 @@ use crate::success_page::LoginSuccessPage; use crate::success_page::LoginSuccessRedirect; use crate::success_page::compose_success_url; use crate::success_page::jwt_auth_claims; +use crate::token_data::AccountOAuthProvenance; use crate::token_data::TokenData; use crate::token_data::parse_chatgpt_jwt_claims; use chrono::Utc; @@ -444,17 +444,7 @@ async fn process_request( obtain_api_key(&client, &opts.issuer, &opts.client_id, &tokens.id_token) .await .ok(); - if let Err(err) = persist_tokens_async( - &opts.codex_home, - api_key.clone(), - tokens.id_token.clone(), - tokens.access_token.clone(), - tokens.refresh_token.clone(), - opts.cli_auth_credentials_store_mode, - opts.auth_keyring_backend_kind, - ) - .await - { + if let Err(err) = persist_tokens_async(opts, api_key.clone(), &tokens).await { eprintln!("Persist error: {err}"); return login_error_response( "Sign-in completed but credentials could not be saved locally.", @@ -705,9 +695,8 @@ pub(crate) struct ExchangedTokens { /// Exchanges an authorization code for tokens and returns the client for further token exchanges. /// /// The returned error remains suitable for user-facing CLI/browser surfaces, so backend-provided -/// non-JSON error text is preserved there. Structured logging stays narrower: it logs reviewed -/// fields from parsed token responses and redacted transport errors, but does not log the final -/// callback-layer `%err` string. +/// non-JSON error text is preserved there. Structured logs only report status and redacted +/// transport errors, and do not log the final callback-layer `%err` string. pub(crate) async fn exchange_code_for_tokens( issuer: &str, client_id: &str, @@ -775,7 +764,6 @@ pub(crate) async fn exchange_code_for_tokens( } warn!( status = %rejection.status, - detail = ?rejection.detail, "oauth token exchange returned non-success status" ); Err(io::Error::other(rejection.to_string())) @@ -831,16 +819,21 @@ async fn send_code_exchange_request( /// Persists exchanged credentials using the configured local auth store. pub(crate) async fn persist_tokens_async( - codex_home: &Path, + opts: &ServerOptions, api_key: Option, - id_token: String, - access_token: String, - refresh_token: String, - auth_credentials_store_mode: AuthCredentialsStoreMode, - keyring_backend_kind: AuthKeyringBackendKind, + exchanged: &ExchangedTokens, ) -> io::Result<()> { // Reuse existing synchronous logic but run it off the async runtime. - let codex_home = codex_home.to_path_buf(); + let codex_home = opts.codex_home.clone(); + let auth_credentials_store_mode = opts.cli_auth_credentials_store_mode; + let keyring_backend_kind = opts.auth_keyring_backend_kind; + let id_token = exchanged.id_token.clone(); + let access_token = exchanged.access_token.clone(); + let refresh_token = exchanged.refresh_token.clone(); + let provenance = AccountOAuthProvenance { + issuer: opts.issuer.clone(), + client_id: opts.client_id.clone(), + }; // Blocking tasks don't inherit task-local context. // Preserve the initiating client for credential-storage metrics. let originator = AuthStorageOriginator::current(); @@ -851,6 +844,8 @@ pub(crate) async fn persist_tokens_async( id_token: parse_chatgpt_jwt_claims(&id_token).map_err(io::Error::other)?, access_token, refresh_token, + oauth_provenance: Some(provenance), + login_owner_id: Some(generate_state()), account_id: None, }; if let Some(acc) = jwt_auth_claims(&id_token) diff --git a/codex-rs/login/src/token_data.rs b/codex-rs/login/src/token_data.rs index af80bf3733..3bfdb9f167 100644 --- a/codex-rs/login/src/token_data.rs +++ b/codex-rs/login/src/token_data.rs @@ -21,9 +21,24 @@ pub struct TokenData { pub refresh_token: String, + /// OAuth authority recorded when this refresh token was issued. Older records have no + /// provenance; never infer authority from unverified token claims. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub oauth_provenance: Option, + + /// Identifies one completed login; token rotation preserves it. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub login_owner_id: Option, + pub account_id: Option, } +#[derive(Deserialize, Serialize, Clone, Debug, PartialEq, Eq)] +pub struct AccountOAuthProvenance { + pub issuer: String, + pub client_id: String, +} + /// Flat subset of useful claims in id_token from auth.json. #[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)] pub struct IdTokenInfo { diff --git a/codex-rs/login/src/token_data_tests.rs b/codex-rs/login/src/token_data_tests.rs index dc88aa648d..cbd25766e6 100644 --- a/codex-rs/login/src/token_data_tests.rs +++ b/codex-rs/login/src/token_data_tests.rs @@ -26,6 +26,29 @@ fn fake_jwt(payload: serde_json::Value) -> String { format!("{header_b64}.{payload_b64}.{signature_b64}") } +#[test] +fn legacy_account_token_record_has_no_inferred_oauth_provenance() -> serde_json::Result<()> { + let legacy = serde_json::json!({ + "id_token": fake_jwt(serde_json::json!({"iss": "https://untrusted.example"})), + "access_token": "account-access", + "refresh_token": "account-refresh", + "account_id": "workspace" + }); + let mut tokens: TokenData = serde_json::from_value(legacy.clone())?; + assert_eq!(tokens.oauth_provenance, None); + assert_eq!(serde_json::to_value(&tokens)?, legacy); + + let provenance = AccountOAuthProvenance { + issuer: "https://idp.example".to_string(), + client_id: "codex-client".to_string(), + }; + tokens.oauth_provenance = Some(provenance); + tokens.login_owner_id = Some("completed-login".into()); + let round_trip: TokenData = serde_json::from_value(serde_json::to_value(&tokens)?)?; + assert_eq!(round_trip, tokens); + Ok(()) +} + #[test] fn id_token_info_parses_email_and_plan() { let fake_jwt = fake_jwt(serde_json::json!({ diff --git a/codex-rs/login/tests/suite/auth_refresh.rs b/codex-rs/login/tests/suite/auth_refresh.rs index 60b6d2a3fe..11fb2d29a0 100644 --- a/codex-rs/login/tests/suite/auth_refresh.rs +++ b/codex-rs/login/tests/suite/auth_refresh.rs @@ -15,6 +15,7 @@ use codex_login::REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR; use codex_login::RefreshTokenError; use codex_login::load_auth_dot_json; use codex_login::save_auth; +use codex_login::token_data::AccountOAuthProvenance; use codex_login::token_data::IdTokenInfo; use codex_login::token_data::TokenData; use codex_protocol::auth::AuthMode; @@ -200,8 +201,18 @@ async fn refresh_token_succeeds_updates_storage() -> Result<()> { .await; let ctx = RefreshTokenTestContext::new(&server).await?; + // Exercise the recorded issuer instead of the test endpoint override. + // SAFETY: auth_env serializes these tests; ctx restores the environment on drop. + unsafe { + std::env::remove_var(REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR); + } let initial_last_refresh = Utc::now() - Duration::days(1); - let initial_tokens = build_tokens(INITIAL_ACCESS_TOKEN, INITIAL_REFRESH_TOKEN); + let mut initial_tokens = build_tokens(INITIAL_ACCESS_TOKEN, INITIAL_REFRESH_TOKEN); + initial_tokens.login_owner_id = Some("completed-login".into()); + initial_tokens.oauth_provenance = Some(AccountOAuthProvenance { + issuer: server.uri(), + client_id: "recorded-client".into(), + }); let initial_auth = AuthDotJson { auth_mode: Some(AuthMode::Chatgpt), openai_api_key: None, @@ -223,7 +234,7 @@ async fn refresh_token_succeeds_updates_storage() -> Result<()> { assert_eq!( serde_json::from_slice::(&requests[0].body)?, json!({ - "client_id": "staging-client", + "client_id": "recorded-client", "grant_type": "refresh_token", "refresh_token": INITIAL_REFRESH_TOKEN, "refresh_reason": "forced", @@ -1689,6 +1700,8 @@ fn build_tokens(access_token: &str, refresh_token: &str) -> TokenData { id_token, access_token: access_token.to_string(), refresh_token: refresh_token.to_string(), + oauth_provenance: None, + login_owner_id: None, account_id: Some("account-id".to_string()), } } diff --git a/codex-rs/login/tests/suite/device_code_login.rs b/codex-rs/login/tests/suite/device_code_login.rs index 50c1f12367..d35ca8cafb 100644 --- a/codex-rs/login/tests/suite/device_code_login.rs +++ b/codex-rs/login/tests/suite/device_code_login.rs @@ -8,6 +8,7 @@ use codex_login::AuthKeyringBackendKind; use codex_login::ServerOptions; use codex_login::auth::load_auth_dot_json; use codex_login::run_device_code_login; +use codex_login::token_data::AccountOAuthProvenance; use serde_json::json; use std::sync::Arc; use std::sync::atomic::AtomicUsize; @@ -144,7 +145,7 @@ async fn device_code_login_integration_succeeds() -> anyhow::Result<()> { mock_oauth_token_single(&mock_server, jwt.clone()).await; let issuer = mock_server.uri(); - let opts = server_opts(&codex_home, issuer, AuthCredentialsStoreMode::File); + let opts = server_opts(&codex_home, issuer.clone(), AuthCredentialsStoreMode::File); run_device_code_login(opts) .await @@ -163,6 +164,18 @@ async fn device_code_login_integration_succeeds() -> anyhow::Result<()> { assert_eq!(tokens.refresh_token, "refresh-token-123"); assert_eq!(tokens.id_token.raw_jwt, jwt); assert_eq!(tokens.account_id.as_deref(), Some(WORKSPACE_ID_ALLOWED)); + let login_owner_id = tokens + .login_owner_id + .as_ref() + .context("completed login identity")?; + assert!(!login_owner_id.is_empty()); + assert_eq!( + tokens.oauth_provenance, + Some(AccountOAuthProvenance { + issuer, + client_id: "client-id".to_string(), + }) + ); Ok(()) } diff --git a/codex-rs/login/tests/suite/login_server_e2e.rs b/codex-rs/login/tests/suite/login_server_e2e.rs index e1bc0fa65e..03d48769de 100644 --- a/codex-rs/login/tests/suite/login_server_e2e.rs +++ b/codex-rs/login/tests/suite/login_server_e2e.rs @@ -143,7 +143,7 @@ async fn end_to_end_login_flow_persists_auth_json() -> Result<()> { cli_auth_credentials_store_mode: AuthCredentialsStoreMode::File, auth_route_config: routes, client_id: codex_login::CLIENT_ID.to_string(), - issuer, + issuer: issuer.clone(), port: 0, open_browser: false, force_state: Some(state), @@ -230,6 +230,12 @@ async fn end_to_end_login_flow_persists_auth_json() -> Result<()> { assert_eq!(json["tokens"]["access_token"], "access-123"); assert_eq!(json["tokens"]["refresh_token"], "refresh-123"); assert_eq!(json["tokens"]["account_id"], chatgpt_account_id); + let login_owner_id = json["tokens"]["login_owner_id"].as_str().unwrap(); + assert!(!login_owner_id.is_empty()); + assert_eq!( + json["tokens"]["oauth_provenance"], + serde_json::json!({ "issuer": issuer, "client_id": codex_login::CLIENT_ID }) + ); // Stop mock issuer drop(issuer_handle); diff --git a/codex-rs/login/tests/suite/logout.rs b/codex-rs/login/tests/suite/logout.rs index 3ecdf0150f..ee33fa7537 100644 --- a/codex-rs/login/tests/suite/logout.rs +++ b/codex-rs/login/tests/suite/logout.rs @@ -257,6 +257,8 @@ fn chatgpt_auth_with_refresh_token(refresh_token: &str) -> AuthDotJson { }, access_token: ACCESS_TOKEN.to_string(), refresh_token: refresh_token.to_string(), + oauth_provenance: None, + login_owner_id: None, account_id: Some("account-id".to_string()), }), last_refresh: None, diff --git a/codex-rs/models-manager/src/manager_tests.rs b/codex-rs/models-manager/src/manager_tests.rs index 3bb8d71821..1bfb7e8771 100644 --- a/codex-rs/models-manager/src/manager_tests.rs +++ b/codex-rs/models-manager/src/manager_tests.rs @@ -611,6 +611,8 @@ c2ln", .expect("fake id token should parse"), access_token: "Access Token".to_string(), refresh_token: "test".to_string(), + oauth_provenance: None, + login_owner_id: None, account_id: Some("account_id".to_string()), }), last_refresh: Some(Utc::now()), diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index 8dd93b2cda..7780b168be 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -683,7 +683,8 @@ pub(crate) async fn start_embedded_app_server_for_picker( ) -> color_eyre::Result { let mut target = AppServerTarget::Embedded; let mut state_db = init_state_db_for_app_server_target(config, &target).await?; - let app_server = start_app_server( + // Keep the shared test fixture's frame independent of embedded startup size. + let app_server = Box::pin(start_app_server( &mut target, Arg0DispatchPaths::default(), config.clone(), @@ -696,7 +697,7 @@ pub(crate) async fn start_embedded_app_server_for_picker( &mut state_db, Arc::new(EnvironmentManager::default_for_tests()), Default::default(), - ) + )) .await?; Ok(AppServerSession::new( crate::backend_id::BackendId::new(), diff --git a/codex-rs/tui/src/local_chatgpt_auth.rs b/codex-rs/tui/src/local_chatgpt_auth.rs index f09ac2263f..6015f76e7d 100644 --- a/codex-rs/tui/src/local_chatgpt_auth.rs +++ b/codex-rs/tui/src/local_chatgpt_auth.rs @@ -110,6 +110,8 @@ mod tests { .expect("id token should parse"), access_token, refresh_token: "refresh-token".to_string(), + oauth_provenance: None, + login_owner_id: None, account_id: Some("workspace-1".to_string()), }), last_refresh: Some(Utc::now()),