diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 4c321b3f6..35d312910 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1080,6 +1080,7 @@ version = "0.0.0" dependencies = [ "anyhow", "app_test_support", + "async-trait", "axum", "base64", "chrono", diff --git a/codex-rs/app-server-protocol/src/protocol/common.rs b/codex-rs/app-server-protocol/src/protocol/common.rs index cad6976b4..2aaf41944 100644 --- a/codex-rs/app-server-protocol/src/protocol/common.rs +++ b/codex-rs/app-server-protocol/src/protocol/common.rs @@ -28,6 +28,11 @@ impl GitSha { pub enum AuthMode { ApiKey, ChatGPT, + /// [UNSTABLE] FOR OPENAI INTERNAL USE ONLY - DO NOT USE. + #[serde(rename = "chatgptAuthTokens")] + #[ts(rename = "chatgptAuthTokens")] + #[strum(serialize = "chatgptAuthTokens")] + ChatgptAuthTokens, } /// Generates an `enum ClientRequest` where each variant is a request that the @@ -534,6 +539,11 @@ server_request_definitions! { response: v2::DynamicToolCallResponse, }, + ChatgptAuthTokensRefresh => "account/chatgptAuthTokens/refresh" { + params: v2::ChatgptAuthTokensRefreshParams, + response: v2::ChatgptAuthTokensRefreshResponse, + }, + /// DEPRECATED APIs below /// Request to approve a patch. /// This request is used for Turns started via the legacy APIs (i.e. SendUserTurn, SendUserMessage). @@ -753,6 +763,29 @@ mod tests { Ok(()) } + #[test] + fn serialize_chatgpt_auth_tokens_refresh_request() -> Result<()> { + let request = ServerRequest::ChatgptAuthTokensRefresh { + request_id: RequestId::Integer(8), + params: v2::ChatgptAuthTokensRefreshParams { + reason: v2::ChatgptAuthTokensRefreshReason::Unauthorized, + previous_account_id: Some("org-123".to_string()), + }, + }; + assert_eq!( + json!({ + "method": "account/chatgptAuthTokens/refresh", + "id": 8, + "params": { + "reason": "unauthorized", + "previousAccountId": "org-123" + } + }), + serde_json::to_value(&request)?, + ); + Ok(()) + } + #[test] fn serialize_get_account_rate_limits() -> Result<()> { let request = ClientRequest::GetAccountRateLimits { @@ -842,10 +875,34 @@ mod tests { Ok(()) } + #[test] + fn serialize_account_login_chatgpt_auth_tokens() -> Result<()> { + let request = ClientRequest::LoginAccount { + request_id: RequestId::Integer(5), + params: v2::LoginAccountParams::ChatgptAuthTokens { + access_token: "access-token".to_string(), + id_token: "id-token".to_string(), + }, + }; + assert_eq!( + json!({ + "method": "account/login/start", + "id": 5, + "params": { + "type": "chatgptAuthTokens", + "accessToken": "access-token", + "idToken": "id-token" + } + }), + serde_json::to_value(&request)?, + ); + Ok(()) + } + #[test] fn serialize_get_account() -> Result<()> { let request = ClientRequest::GetAccount { - request_id: RequestId::Integer(5), + request_id: RequestId::Integer(6), params: v2::GetAccountParams { refresh_token: false, }, @@ -853,7 +910,7 @@ mod tests { assert_eq!( json!({ "method": "account/read", - "id": 5, + "id": 6, "params": { "refreshToken": false } diff --git a/codex-rs/app-server-protocol/src/protocol/v2.rs b/codex-rs/app-server-protocol/src/protocol/v2.rs index 78cce916e..ac3df37bd 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2.rs @@ -835,6 +835,24 @@ pub enum LoginAccountParams { #[serde(rename = "chatgpt")] #[ts(rename = "chatgpt")] Chatgpt, + /// [UNSTABLE] FOR OPENAI INTERNAL USE ONLY - DO NOT USE. + /// The access token must contain the same scopes that Codex-managed ChatGPT auth tokens have. + #[serde(rename = "chatgptAuthTokens")] + #[ts(rename = "chatgptAuthTokens")] + ChatgptAuthTokens { + /// ID token (JWT) supplied by the client. + /// + /// This token is used for identity and account metadata (email, plan type, + /// workspace id). + #[serde(rename = "idToken")] + #[ts(rename = "idToken")] + id_token: String, + /// Access token (JWT) supplied by the client. + /// This token is used for backend API requests. + #[serde(rename = "accessToken")] + #[ts(rename = "accessToken")] + access_token: String, + }, } #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] @@ -854,6 +872,9 @@ pub enum LoginAccountResponse { /// URL the client should open in a browser to initiate the OAuth flow. auth_url: String, }, + #[serde(rename = "chatgptAuthTokens", rename_all = "camelCase")] + #[ts(rename = "chatgptAuthTokens", rename_all = "camelCase")] + ChatgptAuthTokens {}, } #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] @@ -884,6 +905,37 @@ pub struct CancelLoginAccountResponse { #[ts(export_to = "v2/")] pub struct LogoutAccountResponse {} +#[derive(Serialize, Deserialize, Debug, Clone, Copy, PartialEq, Eq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "v2/")] +pub enum ChatgptAuthTokensRefreshReason { + /// Codex attempted a backend request and received `401 Unauthorized`. + Unauthorized, +} + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "v2/")] +pub struct ChatgptAuthTokensRefreshParams { + pub reason: ChatgptAuthTokensRefreshReason, + /// Workspace/account identifier that Codex was previously using. + /// + /// Clients that manage multiple accounts/workspaces can use this as a hint + /// to refresh the token for the correct workspace. + /// + /// This may be `null` when the prior ID token did not include a workspace + /// identifier (`chatgpt_account_id`) or when the token could not be parsed. + pub previous_account_id: Option, +} + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "v2/")] +pub struct ChatgptAuthTokensRefreshResponse { + pub id_token: String, + pub access_token: String, +} + #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] #[serde(rename_all = "camelCase")] #[ts(export_to = "v2/")] @@ -895,6 +947,11 @@ pub struct GetAccountRateLimitsResponse { #[serde(rename_all = "camelCase")] #[ts(export_to = "v2/")] pub struct GetAccountParams { + /// When `true`, requests a proactive token refresh before returning. + /// + /// In managed auth mode this triggers the normal refresh-token flow. In + /// external auth mode this flag is ignored. Clients should refresh tokens + /// themselves and call `account/login/start` with `chatgptAuthTokens`. #[serde(default)] pub refresh_token: bool, } diff --git a/codex-rs/app-server/Cargo.toml b/codex-rs/app-server/Cargo.toml index af2f30c49..820edf59c 100644 --- a/codex-rs/app-server/Cargo.toml +++ b/codex-rs/app-server/Cargo.toml @@ -17,6 +17,7 @@ workspace = true [dependencies] anyhow = { workspace = true } +async-trait = { workspace = true } codex-arg0 = { workspace = true } codex-common = { workspace = true, features = ["cli"] } codex-core = { workspace = true } diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index caba8580a..e9ebcd966 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -659,10 +659,17 @@ $demo-app Pull the latest updates from the team. The JSON-RPC auth/account surface exposes request/response methods plus server-initiated notifications (no `id`). Use these to determine auth state, start or cancel logins, logout, and inspect ChatGPT rate limits. +### Authentication modes + +Codex supports these authentication modes. The current mode is surfaced in `account/updated` (`authMode`) and can be inferred from `account/read`. + +- **API key (`apiKey`)**: Caller supplies an OpenAI API key via `account/login/start` with `type: "apiKey"`. The API key is saved and used for API requests. +- **ChatGPT managed (`chatgpt`)** (recommended): Codex owns the ChatGPT OAuth flow and refresh tokens. Start via `account/login/start` with `type: "chatgpt"`; Codex persists tokens to disk and refreshes them automatically. + ### API Overview - `account/read` — fetch current account info; optionally refresh tokens. -- `account/login/start` — begin login (`apiKey` or `chatgpt`). +- `account/login/start` — begin login (`apiKey`, `chatgpt`). - `account/login/completed` (notify) — emitted when a login attempt finishes (success or error). - `account/login/cancel` — cancel a pending ChatGPT login by `loginId`. - `account/logout` — sign out; triggers `account/updated`. diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index e1c982a99..660b4b87d 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -57,6 +57,7 @@ use codex_app_server_protocol::ListConversationsResponse; use codex_app_server_protocol::ListMcpServerStatusParams; use codex_app_server_protocol::ListMcpServerStatusResponse; use codex_app_server_protocol::LoginAccountParams; +use codex_app_server_protocol::LoginAccountResponse; use codex_app_server_protocol::LoginApiKeyParams; use codex_app_server_protocol::LoginApiKeyResponse; use codex_app_server_protocol::LoginChatGptCompleteNotification; @@ -141,6 +142,7 @@ use codex_core::ThreadManager; use codex_core::ThreadSortKey as CoreThreadSortKey; use codex_core::auth::CLIENT_ID; use codex_core::auth::login_with_api_key; +use codex_core::auth::login_with_chatgpt_auth_tokens; use codex_core::config::Config; use codex_core::config::ConfigOverrides; use codex_core::config::ConfigService; @@ -169,6 +171,7 @@ use codex_core::read_session_meta_line; use codex_core::rollout_date_parts; use codex_core::sandboxing::SandboxPermissions; use codex_core::state_db::get_state_db; +use codex_core::token_data::parse_id_token; use codex_core::windows_sandbox::WindowsSandboxLevelExt; use codex_feedback::CodexFeedback; use codex_login::ServerOptions as LoginServerOptions; @@ -607,6 +610,22 @@ impl CodexMessageProcessor { LoginAccountParams::Chatgpt => { self.login_chatgpt_v2(request_id).await; } + LoginAccountParams::ChatgptAuthTokens { + id_token, + access_token, + } => { + self.login_chatgpt_auth_tokens(request_id, id_token, access_token) + .await; + } + } + } + + fn external_auth_active_error(&self) -> JSONRPCErrorError { + JSONRPCErrorError { + code: INVALID_REQUEST_ERROR_CODE, + message: "External auth is active. Use account/login/start (chatgptAuthTokens) to update it or account/logout to clear it." + .to_string(), + data: None, } } @@ -614,6 +633,10 @@ impl CodexMessageProcessor { &mut self, params: &LoginApiKeyParams, ) -> std::result::Result<(), JSONRPCErrorError> { + if self.auth_manager.is_external_auth_active() { + return Err(self.external_auth_active_error()); + } + if matches!( self.config.forced_login_method, Some(ForcedLoginMethod::Chatgpt) @@ -706,6 +729,10 @@ impl CodexMessageProcessor { ) -> std::result::Result { let config = self.config.as_ref(); + if self.auth_manager.is_external_auth_active() { + return Err(self.external_auth_active_error()); + } + if matches!(config.forced_login_method, Some(ForcedLoginMethod::Api)) { return Err(JSONRPCErrorError { code: INVALID_REQUEST_ERROR_CODE, @@ -964,6 +991,98 @@ impl CodexMessageProcessor { } } + async fn login_chatgpt_auth_tokens( + &mut self, + request_id: RequestId, + id_token: String, + access_token: String, + ) { + if matches!( + self.config.forced_login_method, + Some(ForcedLoginMethod::Api) + ) { + let error = JSONRPCErrorError { + code: INVALID_REQUEST_ERROR_CODE, + message: "External ChatGPT auth is disabled. Use API key login instead." + .to_string(), + data: None, + }; + self.outgoing.send_error(request_id, error).await; + return; + } + + // Cancel any active login attempt to avoid persisting managed auth state. + { + let mut guard = self.active_login.lock().await; + if let Some(active) = guard.take() { + drop(active); + } + } + + let id_token_info = match parse_id_token(&id_token) { + Ok(info) => info, + Err(err) => { + let error = JSONRPCErrorError { + code: INVALID_REQUEST_ERROR_CODE, + message: format!("invalid id token: {err}"), + data: None, + }; + self.outgoing.send_error(request_id, error).await; + return; + } + }; + + if let Some(expected_workspace) = self.config.forced_chatgpt_workspace_id.as_deref() + && id_token_info.chatgpt_account_id.as_deref() != Some(expected_workspace) + { + let account_id = id_token_info.chatgpt_account_id; + let error = JSONRPCErrorError { + code: INVALID_REQUEST_ERROR_CODE, + message: format!( + "External auth must use workspace {expected_workspace}, but received {account_id:?}." + ), + data: None, + }; + self.outgoing.send_error(request_id, error).await; + return; + } + + if let Err(err) = + login_with_chatgpt_auth_tokens(&self.config.codex_home, &id_token, &access_token) + { + let error = JSONRPCErrorError { + code: INTERNAL_ERROR_CODE, + message: format!("failed to set external auth: {err}"), + data: None, + }; + self.outgoing.send_error(request_id, error).await; + return; + } + self.auth_manager.reload(); + + self.outgoing + .send_response(request_id, LoginAccountResponse::ChatgptAuthTokens {}) + .await; + + let payload_login_completed = AccountLoginCompletedNotification { + login_id: None, + success: true, + error: None, + }; + self.outgoing + .send_server_notification(ServerNotification::AccountLoginCompleted( + payload_login_completed, + )) + .await; + + let payload_v2 = AccountUpdatedNotification { + auth_mode: self.auth_manager.get_auth_mode(), + }; + self.outgoing + .send_server_notification(ServerNotification::AccountUpdated(payload_v2)) + .await; + } + async fn logout_common(&mut self) -> std::result::Result, JSONRPCErrorError> { // Cancel any active login attempt. { @@ -1026,6 +1145,9 @@ impl CodexMessageProcessor { } async fn refresh_token_if_requested(&self, do_refresh: bool) { + if self.auth_manager.is_external_auth_active() { + return; + } if do_refresh && let Err(err) = self.auth_manager.refresh_token().await { tracing::warn!("failed to refresh token while getting account: {err}"); } @@ -1100,7 +1222,7 @@ impl CodexMessageProcessor { let account = match self.auth_manager.auth_cached() { Some(auth) => Some(match auth.mode { AuthMode::ApiKey => Account::ApiKey {}, - AuthMode::ChatGPT => { + AuthMode::ChatGPT | AuthMode::ChatgptAuthTokens => { let email = auth.get_account_email(); let plan_type = auth.account_plan_type(); @@ -1159,7 +1281,7 @@ impl CodexMessageProcessor { }); }; - if auth.mode != AuthMode::ChatGPT { + if !matches!(auth.mode, AuthMode::ChatGPT | AuthMode::ChatgptAuthTokens) { return Err(JSONRPCErrorError { code: INVALID_REQUEST_ERROR_CODE, message: "chatgpt authentication required to read rate limits".to_string(), diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index ca588e59d..5b3d39704 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -312,7 +312,7 @@ pub async fn run_main( JSONRPCMessage::Request(r) => processor.process_request(r).await, JSONRPCMessage::Response(r) => processor.process_response(r).await, JSONRPCMessage::Notification(n) => processor.process_notification(n).await, - JSONRPCMessage::Error(e) => processor.process_error(e), + JSONRPCMessage::Error(e) => processor.process_error(e).await, } } created = thread_created_rx.recv(), if listen_for_threads => { diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 95d62d3f9..ced93c3bd 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -5,6 +5,10 @@ use crate::codex_message_processor::CodexMessageProcessor; use crate::config_api::ConfigApi; use crate::error_code::INVALID_REQUEST_ERROR_CODE; use crate::outgoing_message::OutgoingMessageSender; +use async_trait::async_trait; +use codex_app_server_protocol::ChatgptAuthTokensRefreshParams; +use codex_app_server_protocol::ChatgptAuthTokensRefreshReason; +use codex_app_server_protocol::ChatgptAuthTokensRefreshResponse; use codex_app_server_protocol::ClientInfo; use codex_app_server_protocol::ClientRequest; use codex_app_server_protocol::ConfigBatchWriteParams; @@ -19,8 +23,13 @@ use codex_app_server_protocol::JSONRPCRequest; use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ServerNotification; +use codex_app_server_protocol::ServerRequestPayload; use codex_core::AuthManager; use codex_core::ThreadManager; +use codex_core::auth::ExternalAuthRefreshContext; +use codex_core::auth::ExternalAuthRefreshReason; +use codex_core::auth::ExternalAuthRefresher; +use codex_core::auth::ExternalAuthTokens; use codex_core::config::Config; use codex_core::config_loader::LoaderOverrides; use codex_core::default_client::SetOriginatorError; @@ -31,8 +40,64 @@ use codex_feedback::CodexFeedback; use codex_protocol::ThreadId; use codex_protocol::protocol::SessionSource; use tokio::sync::broadcast; +use tokio::time::Duration; +use tokio::time::timeout; use toml::Value as TomlValue; +const EXTERNAL_AUTH_REFRESH_TIMEOUT: Duration = Duration::from_secs(10); + +#[derive(Clone)] +struct ExternalAuthRefreshBridge { + outgoing: Arc, +} + +impl ExternalAuthRefreshBridge { + fn map_reason(reason: ExternalAuthRefreshReason) -> ChatgptAuthTokensRefreshReason { + match reason { + ExternalAuthRefreshReason::Unauthorized => ChatgptAuthTokensRefreshReason::Unauthorized, + } + } +} + +#[async_trait] +impl ExternalAuthRefresher for ExternalAuthRefreshBridge { + async fn refresh( + &self, + context: ExternalAuthRefreshContext, + ) -> std::io::Result { + let params = ChatgptAuthTokensRefreshParams { + reason: Self::map_reason(context.reason), + previous_account_id: context.previous_account_id, + }; + + let (request_id, rx) = self + .outgoing + .send_request_with_id(ServerRequestPayload::ChatgptAuthTokensRefresh(params)) + .await; + + let result = match timeout(EXTERNAL_AUTH_REFRESH_TIMEOUT, rx).await { + Ok(result) => result.map_err(|err| { + std::io::Error::other(format!("auth refresh request canceled: {err}")) + })?, + Err(_) => { + let _canceled = self.outgoing.cancel_request(&request_id).await; + return Err(std::io::Error::other(format!( + "auth refresh request timed out after {}s", + EXTERNAL_AUTH_REFRESH_TIMEOUT.as_secs() + ))); + } + }; + + let response: ChatgptAuthTokensRefreshResponse = + serde_json::from_value(result).map_err(std::io::Error::other)?; + + Ok(ExternalAuthTokens { + access_token: response.access_token, + id_token: response.id_token, + }) + } +} + pub(crate) struct MessageProcessor { outgoing: Arc, codex_message_processor: CodexMessageProcessor, @@ -59,6 +124,10 @@ impl MessageProcessor { false, config.cli_auth_credentials_store_mode, ); + auth_manager.set_forced_chatgpt_workspace_id(config.forced_chatgpt_workspace_id.clone()); + auth_manager.set_external_auth_refresher(Arc::new(ExternalAuthRefreshBridge { + outgoing: outgoing.clone(), + })); let thread_manager = Arc::new(ThreadManager::new( config.codex_home.clone(), auth_manager.clone(), @@ -236,8 +305,9 @@ impl MessageProcessor { } /// Handle an error object received from the peer. - pub(crate) fn process_error(&mut self, err: JSONRPCError) { + pub(crate) async fn process_error(&mut self, err: JSONRPCError) { tracing::error!("<- error: {:?}", err); + self.outgoing.notify_client_error(err.id, err.error).await; } async fn handle_config_read(&self, request_id: RequestId, params: ConfigReadParams) { diff --git a/codex-rs/app-server/src/outgoing_message.rs b/codex-rs/app-server/src/outgoing_message.rs index 7b4a599b0..be89775d8 100644 --- a/codex-rs/app-server/src/outgoing_message.rs +++ b/codex-rs/app-server/src/outgoing_message.rs @@ -39,6 +39,14 @@ impl OutgoingMessageSender { &self, request: ServerRequestPayload, ) -> oneshot::Receiver { + let (_id, rx) = self.send_request_with_id(request).await; + rx + } + + pub(crate) async fn send_request_with_id( + &self, + request: ServerRequestPayload, + ) -> (RequestId, oneshot::Receiver) { let id = RequestId::Integer(self.next_request_id.fetch_add(1, Ordering::Relaxed)); let outgoing_message_id = id.clone(); let (tx_approve, rx_approve) = oneshot::channel(); @@ -54,7 +62,7 @@ impl OutgoingMessageSender { let mut request_id_to_callback = self.request_id_to_callback.lock().await; request_id_to_callback.remove(&outgoing_message_id); } - rx_approve + (outgoing_message_id, rx_approve) } pub(crate) async fn notify_client_response(&self, id: RequestId, result: Result) { @@ -75,6 +83,30 @@ impl OutgoingMessageSender { } } + pub(crate) async fn notify_client_error(&self, id: RequestId, error: JSONRPCErrorError) { + let entry = { + let mut request_id_to_callback = self.request_id_to_callback.lock().await; + request_id_to_callback.remove_entry(&id) + }; + + match entry { + Some((id, _sender)) => { + warn!("client responded with error for {id:?}: {error:?}"); + } + None => { + warn!("could not find callback for {id:?}"); + } + } + } + + pub(crate) async fn cancel_request(&self, id: &RequestId) -> bool { + let entry = { + let mut request_id_to_callback = self.request_id_to_callback.lock().await; + request_id_to_callback.remove_entry(id) + }; + entry.is_some() + } + pub(crate) async fn send_response(&self, id: RequestId, response: T) { match serde_json::to_value(response) { Ok(result) => { diff --git a/codex-rs/app-server/tests/common/auth_fixtures.rs b/codex-rs/app-server/tests/common/auth_fixtures.rs index 9f1b62744..b78d5b105 100644 --- a/codex-rs/app-server/tests/common/auth_fixtures.rs +++ b/codex-rs/app-server/tests/common/auth_fixtures.rs @@ -6,6 +6,7 @@ use base64::Engine; use base64::engine::general_purpose::URL_SAFE_NO_PAD; use chrono::DateTime; use chrono::Utc; +use codex_app_server_protocol::AuthMode; use codex_core::auth::AuthCredentialsStoreMode; use codex_core::auth::AuthDotJson; use codex_core::auth::save_auth; @@ -158,6 +159,7 @@ pub fn write_chatgpt_auth( let last_refresh = fixture.last_refresh.unwrap_or_else(|| Some(Utc::now())); let auth = AuthDotJson { + auth_mode: Some(AuthMode::ChatGPT), openai_api_key: None, tokens: Some(tokens), last_refresh, diff --git a/codex-rs/app-server/tests/common/mcp_process.rs b/codex-rs/app-server/tests/common/mcp_process.rs index fa160b4c4..4eeab8f95 100644 --- a/codex-rs/app-server/tests/common/mcp_process.rs +++ b/codex-rs/app-server/tests/common/mcp_process.rs @@ -29,11 +29,13 @@ use codex_app_server_protocol::GetAuthStatusParams; use codex_app_server_protocol::InitializeParams; use codex_app_server_protocol::InterruptConversationParams; use codex_app_server_protocol::JSONRPCError; +use codex_app_server_protocol::JSONRPCErrorError; use codex_app_server_protocol::JSONRPCMessage; use codex_app_server_protocol::JSONRPCNotification; use codex_app_server_protocol::JSONRPCRequest; use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::ListConversationsParams; +use codex_app_server_protocol::LoginAccountParams; use codex_app_server_protocol::LoginApiKeyParams; use codex_app_server_protocol::ModelListParams; use codex_app_server_protocol::NewConversationParams; @@ -298,6 +300,20 @@ impl McpProcess { self.send_request("account/read", params).await } + /// Send an `account/login/start` JSON-RPC request with ChatGPT auth tokens. + pub async fn send_chatgpt_auth_tokens_login_request( + &mut self, + id_token: String, + access_token: String, + ) -> anyhow::Result { + let params = LoginAccountParams::ChatgptAuthTokens { + id_token, + access_token, + }; + let params = Some(serde_json::to_value(params)?); + self.send_request("account/login/start", params).await + } + /// Send a `feedback/upload` JSON-RPC request. pub async fn send_feedback_upload_request( &mut self, @@ -608,6 +624,15 @@ impl McpProcess { .await } + pub async fn send_error( + &mut self, + id: RequestId, + error: JSONRPCErrorError, + ) -> anyhow::Result<()> { + self.send_jsonrpc_message(JSONRPCMessage::Error(JSONRPCError { id, error })) + .await + } + pub async fn send_notification( &mut self, notification: ClientNotification, diff --git a/codex-rs/app-server/tests/suite/v2/account.rs b/codex-rs/app-server/tests/suite/v2/account.rs index cbbdad84c..d3145345e 100644 --- a/codex-rs/app-server/tests/suite/v2/account.rs +++ b/codex-rs/app-server/tests/suite/v2/account.rs @@ -4,28 +4,43 @@ use app_test_support::McpProcess; use app_test_support::to_response; use app_test_support::ChatGptAuthFixture; +use app_test_support::ChatGptIdTokenClaims; +use app_test_support::encode_id_token; use app_test_support::write_chatgpt_auth; +use app_test_support::write_models_cache; use codex_app_server_protocol::Account; use codex_app_server_protocol::AuthMode; use codex_app_server_protocol::CancelLoginAccountParams; use codex_app_server_protocol::CancelLoginAccountResponse; +use codex_app_server_protocol::CancelLoginAccountStatus; +use codex_app_server_protocol::ChatgptAuthTokensRefreshReason; +use codex_app_server_protocol::ChatgptAuthTokensRefreshResponse; use codex_app_server_protocol::GetAccountParams; use codex_app_server_protocol::GetAccountResponse; use codex_app_server_protocol::JSONRPCError; +use codex_app_server_protocol::JSONRPCErrorError; +use codex_app_server_protocol::JSONRPCNotification; use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::LoginAccountResponse; use codex_app_server_protocol::LogoutAccountResponse; use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ServerNotification; +use codex_app_server_protocol::ServerRequest; +use codex_app_server_protocol::TurnCompletedNotification; +use codex_app_server_protocol::TurnStatus; use codex_core::auth::AuthCredentialsStoreMode; use codex_login::login_with_api_key; use codex_protocol::account::PlanType as AccountPlanType; +use core_test_support::responses; use pretty_assertions::assert_eq; +use serde_json::json; use serial_test::serial; use std::path::Path; use std::time::Duration; use tempfile::TempDir; use tokio::time::timeout; +use wiremock::MockServer; +use wiremock::ResponseTemplate; const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); @@ -35,10 +50,14 @@ struct CreateConfigTomlParams { forced_method: Option, forced_workspace_id: Option, requires_openai_auth: Option, + base_url: Option, } fn create_config_toml(codex_home: &Path, params: CreateConfigTomlParams) -> std::io::Result<()> { let config_toml = codex_home.join("config.toml"); + let base_url = params + .base_url + .unwrap_or_else(|| "http://127.0.0.1:0/v1".to_string()); let forced_line = if let Some(method) = params.forced_method { format!("forced_login_method = \"{method}\"\n") } else { @@ -66,7 +85,7 @@ model_provider = "mock_provider" [model_providers.mock_provider] name = "Mock provider for test" -base_url = "http://127.0.0.1:0/v1" +base_url = "{base_url}" wire_api = "responses" request_max_retries = 0 stream_max_retries = 0 @@ -133,6 +152,627 @@ async fn logout_account_removes_auth_and_notifies() -> Result<()> { Ok(()) } +#[tokio::test] +async fn set_auth_token_updates_account_and_notifies() -> Result<()> { + let codex_home = TempDir::new()?; + let mock_server = MockServer::start().await; + create_config_toml( + codex_home.path(), + CreateConfigTomlParams { + requires_openai_auth: Some(true), + base_url: Some(format!("{}/v1", mock_server.uri())), + ..Default::default() + }, + )?; + write_models_cache(codex_home.path())?; + + let id_token = encode_id_token( + &ChatGptIdTokenClaims::new() + .email("embedded@example.com") + .plan_type("pro") + .chatgpt_account_id("org-embedded"), + )?; + let access_token = "access-embedded".to_string(); + + let mut mcp = McpProcess::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let set_id = mcp + .send_chatgpt_auth_tokens_login_request(id_token.clone(), access_token) + .await?; + let set_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(set_id)), + ) + .await??; + let response: LoginAccountResponse = to_response(set_resp)?; + assert_eq!(response, LoginAccountResponse::ChatgptAuthTokens {}); + + let note = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("account/updated"), + ) + .await??; + let parsed: ServerNotification = note.try_into()?; + let ServerNotification::AccountUpdated(payload) = parsed else { + bail!("unexpected notification: {parsed:?}"); + }; + assert_eq!(payload.auth_mode, Some(AuthMode::ChatgptAuthTokens)); + + let get_id = mcp + .send_get_account_request(GetAccountParams { + refresh_token: false, + }) + .await?; + let get_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(get_id)), + ) + .await??; + let account: GetAccountResponse = to_response(get_resp)?; + assert_eq!( + account, + GetAccountResponse { + account: Some(Account::Chatgpt { + email: "embedded@example.com".to_string(), + plan_type: AccountPlanType::Pro, + }), + requires_openai_auth: true, + } + ); + + Ok(()) +} + +#[tokio::test] +async fn account_read_refresh_token_is_noop_in_external_mode() -> Result<()> { + let codex_home = TempDir::new()?; + create_config_toml( + codex_home.path(), + CreateConfigTomlParams { + requires_openai_auth: Some(true), + ..Default::default() + }, + )?; + write_models_cache(codex_home.path())?; + + let id_token = encode_id_token( + &ChatGptIdTokenClaims::new() + .email("embedded@example.com") + .plan_type("pro") + .chatgpt_account_id("org-embedded"), + )?; + + let mut mcp = McpProcess::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let set_id = mcp + .send_chatgpt_auth_tokens_login_request(id_token, "access-embedded".to_string()) + .await?; + let set_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(set_id)), + ) + .await??; + let response: LoginAccountResponse = to_response(set_resp)?; + assert_eq!(response, LoginAccountResponse::ChatgptAuthTokens {}); + let _updated = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("account/updated"), + ) + .await??; + + let get_id = mcp + .send_get_account_request(GetAccountParams { + refresh_token: true, + }) + .await?; + let get_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(get_id)), + ) + .await??; + let account: GetAccountResponse = to_response(get_resp)?; + assert_eq!( + account, + GetAccountResponse { + account: Some(Account::Chatgpt { + email: "embedded@example.com".to_string(), + plan_type: AccountPlanType::Pro, + }), + requires_openai_auth: true, + } + ); + + let refresh_request = timeout( + Duration::from_millis(250), + mcp.read_stream_until_request_message(), + ) + .await; + assert!( + refresh_request.is_err(), + "external mode should not emit account/chatgptAuthTokens/refresh for refreshToken=true" + ); + + Ok(()) +} + +async fn respond_to_refresh_request( + mcp: &mut McpProcess, + access_token: &str, + id_token: &str, +) -> Result<()> { + let refresh_req: ServerRequest = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_request_message(), + ) + .await??; + let ServerRequest::ChatgptAuthTokensRefresh { request_id, params } = refresh_req else { + bail!("expected account/chatgptAuthTokens/refresh request, got {refresh_req:?}"); + }; + assert_eq!(params.reason, ChatgptAuthTokensRefreshReason::Unauthorized); + let response = ChatgptAuthTokensRefreshResponse { + access_token: access_token.to_string(), + id_token: id_token.to_string(), + }; + mcp.send_response(request_id, serde_json::to_value(response)?) + .await?; + Ok(()) +} + +#[tokio::test] +// 401 response triggers account/chatgptAuthTokens/refresh and retries with new tokens. +async fn external_auth_refreshes_on_unauthorized() -> Result<()> { + let codex_home = TempDir::new()?; + let mock_server = MockServer::start().await; + create_config_toml( + codex_home.path(), + CreateConfigTomlParams { + requires_openai_auth: Some(true), + base_url: Some(format!("{}/v1", mock_server.uri())), + ..Default::default() + }, + )?; + write_models_cache(codex_home.path())?; + + let success_sse = responses::sse(vec![ + responses::ev_response_created("resp-turn"), + responses::ev_assistant_message("msg-turn", "turn ok"), + responses::ev_completed("resp-turn"), + ]); + let unauthorized = ResponseTemplate::new(401).set_body_json(json!({ + "error": { "message": "unauthorized" } + })); + let responses_mock = responses::mount_response_sequence( + &mock_server, + vec![unauthorized, responses::sse_response(success_sse)], + ) + .await; + + let initial_id_token = encode_id_token( + &ChatGptIdTokenClaims::new() + .email("initial@example.com") + .plan_type("pro") + .chatgpt_account_id("org-initial"), + )?; + let refreshed_id_token = encode_id_token( + &ChatGptIdTokenClaims::new() + .email("refreshed@example.com") + .plan_type("pro") + .chatgpt_account_id("org-refreshed"), + )?; + let initial_access_token = "access-initial".to_string(); + let refreshed_access_token = "access-refreshed".to_string(); + + let mut mcp = McpProcess::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let set_id = mcp + .send_chatgpt_auth_tokens_login_request( + initial_id_token.clone(), + initial_access_token.clone(), + ) + .await?; + let set_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(set_id)), + ) + .await??; + let response: LoginAccountResponse = to_response(set_resp)?; + assert_eq!(response, LoginAccountResponse::ChatgptAuthTokens {}); + let _updated = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("account/updated"), + ) + .await??; + + let thread_req = mcp + .send_thread_start_request(codex_app_server_protocol::ThreadStartParams { + model: Some("mock-model".to_string()), + ..Default::default() + }) + .await?; + let thread_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), + ) + .await??; + let thread = to_response::(thread_resp)?; + + let turn_req = mcp + .send_turn_start_request(codex_app_server_protocol::TurnStartParams { + thread_id: thread.thread.id, + input: vec![codex_app_server_protocol::UserInput::Text { + text: "Hello".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }) + .await?; + respond_to_refresh_request(&mut mcp, &refreshed_access_token, &refreshed_id_token).await?; + let _turn_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), + ) + .await??; + let _turn_completed = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("turn/completed"), + ) + .await??; + + let requests = responses_mock.requests(); + assert_eq!(requests.len(), 2); + assert_eq!( + requests[0].header("authorization"), + Some(format!("Bearer {initial_access_token}")) + ); + assert_eq!( + requests[1].header("authorization"), + Some(format!("Bearer {refreshed_access_token}")) + ); + + Ok(()) +} + +#[tokio::test] +// Client returns JSON-RPC error to refresh; turn fails. +async fn external_auth_refresh_error_fails_turn() -> Result<()> { + let codex_home = TempDir::new()?; + let mock_server = MockServer::start().await; + create_config_toml( + codex_home.path(), + CreateConfigTomlParams { + requires_openai_auth: Some(true), + base_url: Some(format!("{}/v1", mock_server.uri())), + ..Default::default() + }, + )?; + write_models_cache(codex_home.path())?; + + let unauthorized = ResponseTemplate::new(401).set_body_json(json!({ + "error": { "message": "unauthorized" } + })); + let _responses_mock = + responses::mount_response_sequence(&mock_server, vec![unauthorized]).await; + + let initial_id_token = encode_id_token( + &ChatGptIdTokenClaims::new() + .email("initial@example.com") + .plan_type("pro") + .chatgpt_account_id("org-initial"), + )?; + + let mut mcp = McpProcess::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let set_id = mcp + .send_chatgpt_auth_tokens_login_request(initial_id_token, "access-initial".to_string()) + .await?; + let set_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(set_id)), + ) + .await??; + let response: LoginAccountResponse = to_response(set_resp)?; + assert_eq!(response, LoginAccountResponse::ChatgptAuthTokens {}); + let _updated = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("account/updated"), + ) + .await??; + + let thread_req = mcp + .send_thread_start_request(codex_app_server_protocol::ThreadStartParams { + model: Some("mock-model".to_string()), + ..Default::default() + }) + .await?; + let thread_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), + ) + .await??; + let thread = to_response::(thread_resp)?; + + let turn_req = mcp + .send_turn_start_request(codex_app_server_protocol::TurnStartParams { + thread_id: thread.thread.id.clone(), + input: vec![codex_app_server_protocol::UserInput::Text { + text: "Hello".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }) + .await?; + + let refresh_req: ServerRequest = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_request_message(), + ) + .await??; + let ServerRequest::ChatgptAuthTokensRefresh { request_id, .. } = refresh_req else { + bail!("expected account/chatgptAuthTokens/refresh request, got {refresh_req:?}"); + }; + + mcp.send_error( + request_id, + JSONRPCErrorError { + code: -32_000, + message: "refresh failed".to_string(), + data: None, + }, + ) + .await?; + + let _turn_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), + ) + .await??; + let completed_notif: JSONRPCNotification = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("turn/completed"), + ) + .await??; + let completed: TurnCompletedNotification = serde_json::from_value( + completed_notif + .params + .expect("turn/completed params must be present"), + )?; + assert_eq!(completed.turn.status, TurnStatus::Failed); + assert!(completed.turn.error.is_some()); + + Ok(()) +} + +#[tokio::test] +// Refresh returns tokens for the wrong workspace; turn fails. +async fn external_auth_refresh_mismatched_workspace_fails_turn() -> Result<()> { + let codex_home = TempDir::new()?; + let mock_server = MockServer::start().await; + create_config_toml( + codex_home.path(), + CreateConfigTomlParams { + forced_workspace_id: Some("org-expected".to_string()), + requires_openai_auth: Some(true), + base_url: Some(format!("{}/v1", mock_server.uri())), + ..Default::default() + }, + )?; + write_models_cache(codex_home.path())?; + + let unauthorized = ResponseTemplate::new(401).set_body_json(json!({ + "error": { "message": "unauthorized" } + })); + let _responses_mock = + responses::mount_response_sequence(&mock_server, vec![unauthorized]).await; + + let initial_id_token = encode_id_token( + &ChatGptIdTokenClaims::new() + .email("initial@example.com") + .plan_type("pro") + .chatgpt_account_id("org-expected"), + )?; + let refreshed_id_token = encode_id_token( + &ChatGptIdTokenClaims::new() + .email("refreshed@example.com") + .plan_type("pro") + .chatgpt_account_id("org-other"), + )?; + + let mut mcp = McpProcess::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let set_id = mcp + .send_chatgpt_auth_tokens_login_request(initial_id_token, "access-initial".to_string()) + .await?; + let set_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(set_id)), + ) + .await??; + let response: LoginAccountResponse = to_response(set_resp)?; + assert_eq!(response, LoginAccountResponse::ChatgptAuthTokens {}); + let _updated = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("account/updated"), + ) + .await??; + + let thread_req = mcp + .send_thread_start_request(codex_app_server_protocol::ThreadStartParams { + model: Some("mock-model".to_string()), + ..Default::default() + }) + .await?; + let thread_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), + ) + .await??; + let thread = to_response::(thread_resp)?; + + let turn_req = mcp + .send_turn_start_request(codex_app_server_protocol::TurnStartParams { + thread_id: thread.thread.id.clone(), + input: vec![codex_app_server_protocol::UserInput::Text { + text: "Hello".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }) + .await?; + + let refresh_req: ServerRequest = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_request_message(), + ) + .await??; + let ServerRequest::ChatgptAuthTokensRefresh { request_id, .. } = refresh_req else { + bail!("expected account/chatgptAuthTokens/refresh request, got {refresh_req:?}"); + }; + + mcp.send_response( + request_id, + serde_json::to_value(ChatgptAuthTokensRefreshResponse { + access_token: "access-refreshed".to_string(), + id_token: refreshed_id_token, + })?, + ) + .await?; + + let _turn_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), + ) + .await??; + let completed_notif: JSONRPCNotification = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("turn/completed"), + ) + .await??; + let completed: TurnCompletedNotification = serde_json::from_value( + completed_notif + .params + .expect("turn/completed params must be present"), + )?; + assert_eq!(completed.turn.status, TurnStatus::Failed); + assert!(completed.turn.error.is_some()); + + Ok(()) +} + +#[tokio::test] +// Refresh returns a malformed id_token; turn fails. +async fn external_auth_refresh_invalid_id_token_fails_turn() -> Result<()> { + let codex_home = TempDir::new()?; + let mock_server = MockServer::start().await; + create_config_toml( + codex_home.path(), + CreateConfigTomlParams { + requires_openai_auth: Some(true), + base_url: Some(format!("{}/v1", mock_server.uri())), + ..Default::default() + }, + )?; + write_models_cache(codex_home.path())?; + + let unauthorized = ResponseTemplate::new(401).set_body_json(json!({ + "error": { "message": "unauthorized" } + })); + let _responses_mock = + responses::mount_response_sequence(&mock_server, vec![unauthorized]).await; + + let initial_id_token = encode_id_token( + &ChatGptIdTokenClaims::new() + .email("initial@example.com") + .plan_type("pro") + .chatgpt_account_id("org-initial"), + )?; + + let mut mcp = McpProcess::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let set_id = mcp + .send_chatgpt_auth_tokens_login_request(initial_id_token, "access-initial".to_string()) + .await?; + let set_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(set_id)), + ) + .await??; + let response: LoginAccountResponse = to_response(set_resp)?; + assert_eq!(response, LoginAccountResponse::ChatgptAuthTokens {}); + let _updated = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("account/updated"), + ) + .await??; + + let thread_req = mcp + .send_thread_start_request(codex_app_server_protocol::ThreadStartParams { + model: Some("mock-model".to_string()), + ..Default::default() + }) + .await?; + let thread_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), + ) + .await??; + let thread = to_response::(thread_resp)?; + + let turn_req = mcp + .send_turn_start_request(codex_app_server_protocol::TurnStartParams { + thread_id: thread.thread.id.clone(), + input: vec![codex_app_server_protocol::UserInput::Text { + text: "Hello".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }) + .await?; + + let refresh_req: ServerRequest = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_request_message(), + ) + .await??; + let ServerRequest::ChatgptAuthTokensRefresh { request_id, .. } = refresh_req else { + bail!("expected account/chatgptAuthTokens/refresh request, got {refresh_req:?}"); + }; + + mcp.send_response( + request_id, + serde_json::to_value(ChatgptAuthTokensRefreshResponse { + access_token: "access-refreshed".to_string(), + id_token: "not-a-jwt".to_string(), + })?, + ) + .await?; + + let _turn_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), + ) + .await??; + let completed_notif: JSONRPCNotification = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("turn/completed"), + ) + .await??; + let completed: TurnCompletedNotification = serde_json::from_value( + completed_notif + .params + .expect("turn/completed params must be present"), + )?; + assert_eq!(completed.turn.status, TurnStatus::Failed); + assert!(completed.turn.error.is_some()); + + Ok(()) +} + #[tokio::test] async fn login_account_api_key_succeeds_and_notifies() -> Result<()> { let codex_home = TempDir::new()?; @@ -304,6 +944,71 @@ async fn login_account_chatgpt_start_can_be_cancelled() -> Result<()> { Ok(()) } +#[tokio::test] +// Serialize tests that launch the login server since it binds to a fixed port. +#[serial(login_port)] +async fn set_auth_token_cancels_active_chatgpt_login() -> Result<()> { + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), CreateConfigTomlParams::default())?; + + let mut mcp = McpProcess::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + // Initiate the ChatGPT login flow + let request_id = mcp.send_login_account_chatgpt_request().await?; + let resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + + let login: LoginAccountResponse = to_response(resp)?; + let LoginAccountResponse::Chatgpt { login_id, .. } = login else { + bail!("unexpected login response: {login:?}"); + }; + + let id_token = encode_id_token( + &ChatGptIdTokenClaims::new() + .email("embedded@example.com") + .plan_type("pro") + .chatgpt_account_id("org-embedded"), + )?; + // Set an external auth token instead of completing the ChatGPT login flow. + // This should cancel the active login attempt. + let set_id = mcp + .send_chatgpt_auth_tokens_login_request(id_token, "access-embedded".to_string()) + .await?; + let set_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(set_id)), + ) + .await??; + let response: LoginAccountResponse = to_response(set_resp)?; + assert_eq!(response, LoginAccountResponse::ChatgptAuthTokens {}); + let _updated = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("account/updated"), + ) + .await??; + + // Verify that the active login attempt was cancelled. + // We check this by trying to cancel it and expecting a not found error. + let cancel_id = mcp + .send_cancel_login_account_request(CancelLoginAccountParams { + login_id: login_id.clone(), + }) + .await?; + let cancel_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(cancel_id)), + ) + .await??; + let cancel: CancelLoginAccountResponse = to_response(cancel_resp)?; + assert_eq!(cancel.status, CancelLoginAccountStatus::NotFound); + + Ok(()) +} + #[tokio::test] // Serialize tests that launch the login server since it binds to a fixed port. #[serial(login_port)] diff --git a/codex-rs/cli/src/login.rs b/codex-rs/cli/src/login.rs index 45ee45969..0f4e4eb39 100644 --- a/codex-rs/cli/src/login.rs +++ b/codex-rs/cli/src/login.rs @@ -240,6 +240,10 @@ pub async fn run_login_status(cli_config_overrides: CliConfigOverrides) -> ! { eprintln!("Logged in using ChatGPT"); std::process::exit(0); } + AuthMode::ChatgptAuthTokens => { + eprintln!("Logged in using ChatGPT (external tokens)"); + std::process::exit(0); + } }, Ok(None) => { eprintln!("Not logged in"); diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index fe27f7e24..cdd2717fd 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -111,6 +111,13 @@ "auto" ], "type": "string" + }, + { + "description": "Store credentials in memory only for the current process.", + "enum": [ + "ephemeral" + ], + "type": "string" } ] }, diff --git a/codex-rs/core/src/auth.rs b/codex-rs/core/src/auth.rs index b74630927..d6cfec550 100644 --- a/codex-rs/core/src/auth.rs +++ b/codex-rs/core/src/auth.rs @@ -1,5 +1,6 @@ mod storage; +use async_trait::async_trait; use chrono::Utc; use reqwest::StatusCode; use serde::Deserialize; @@ -23,6 +24,7 @@ use crate::auth::storage::create_auth_storage; use crate::config::Config; use crate::error::RefreshTokenFailedError; use crate::error::RefreshTokenFailedReason; +use crate::token_data::IdTokenInfo; use crate::token_data::KnownPlan as InternalKnownPlan; use crate::token_data::PlanType as InternalPlanType; use crate::token_data::TokenData; @@ -68,6 +70,31 @@ pub enum RefreshTokenError { Transient(#[from] std::io::Error), } +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ExternalAuthTokens { + pub access_token: String, + pub id_token: String, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ExternalAuthRefreshReason { + Unauthorized, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ExternalAuthRefreshContext { + pub reason: ExternalAuthRefreshReason, + pub previous_account_id: Option, +} + +#[async_trait] +pub trait ExternalAuthRefresher: Send + Sync { + async fn refresh( + &self, + context: ExternalAuthRefreshContext, + ) -> std::io::Result; +} + impl RefreshTokenError { pub fn failed_reason(&self) -> Option { match self { @@ -110,9 +137,9 @@ impl CodexAuth { pub fn get_token(&self) -> Result { match self.mode { AuthMode::ApiKey => Ok(self.api_key.clone().unwrap_or_default()), - AuthMode::ChatGPT => { - let id_token = self.get_token_data()?.access_token; - Ok(id_token) + AuthMode::ChatGPT | AuthMode::ChatgptAuthTokens => { + let access_token = self.get_token_data()?.access_token; + Ok(access_token) } } } @@ -160,6 +187,7 @@ impl CodexAuth { /// Consider this private to integration tests. pub fn create_dummy_chatgpt_auth_for_testing() -> Self { let auth_dot_json = AuthDotJson { + auth_mode: Some(AuthMode::ChatGPT), openai_api_key: None, tokens: Some(TokenData { id_token: Default::default(), @@ -229,6 +257,7 @@ pub fn login_with_api_key( auth_credentials_store_mode: AuthCredentialsStoreMode, ) -> std::io::Result<()> { let auth_dot_json = AuthDotJson { + auth_mode: Some(AuthMode::ApiKey), openai_api_key: Some(api_key.to_string()), tokens: None, last_refresh: None, @@ -236,6 +265,20 @@ pub fn login_with_api_key( save_auth(codex_home, &auth_dot_json, auth_credentials_store_mode) } +/// Writes an in-memory auth payload for externally managed ChatGPT tokens. +pub fn login_with_chatgpt_auth_tokens( + codex_home: &Path, + id_token: &str, + access_token: &str, +) -> std::io::Result<()> { + let auth_dot_json = AuthDotJson::from_external_token_strings(id_token, access_token)?; + save_auth( + codex_home, + &auth_dot_json, + AuthCredentialsStoreMode::Ephemeral, + ) +} + /// Persist the provided auth payload using the specified backend. pub fn save_auth( codex_home: &Path, @@ -272,8 +315,8 @@ pub fn enforce_login_restrictions(config: &Config) -> std::io::Result<()> { if let Some(required_method) = config.forced_login_method { let method_violation = match (required_method, auth.mode) { (ForcedLoginMethod::Api, AuthMode::ApiKey) => None, - (ForcedLoginMethod::Chatgpt, AuthMode::ChatGPT) => None, - (ForcedLoginMethod::Api, AuthMode::ChatGPT) => Some( + (ForcedLoginMethod::Chatgpt, AuthMode::ChatGPT | AuthMode::ChatgptAuthTokens) => None, + (ForcedLoginMethod::Api, AuthMode::ChatGPT | AuthMode::ChatgptAuthTokens) => Some( "API key login is required, but ChatGPT is currently being used. Logging out." .to_string(), ), @@ -293,7 +336,7 @@ pub fn enforce_login_restrictions(config: &Config) -> std::io::Result<()> { } if let Some(expected_account_id) = config.forced_chatgpt_workspace_id.as_deref() { - if auth.mode != AuthMode::ChatGPT { + if !matches!(auth.mode, AuthMode::ChatGPT | AuthMode::ChatgptAuthTokens) { return Ok(()); } @@ -337,12 +380,26 @@ fn logout_with_message( message: String, auth_credentials_store_mode: AuthCredentialsStoreMode, ) -> std::io::Result<()> { - match logout(codex_home, auth_credentials_store_mode) { - Ok(_) => Err(std::io::Error::other(message)), - Err(err) => Err(std::io::Error::other(format!( - "{message}. Failed to remove auth.json: {err}" - ))), + // External auth tokens live in the ephemeral store, but persistent auth may still exist + // from earlier logins. Clear both so a forced logout truly removes all active auth. + let removal_result = logout_all_stores(codex_home, auth_credentials_store_mode); + let error_message = match removal_result { + Ok(_) => message, + Err(err) => format!("{message}. Failed to remove auth.json: {err}"), + }; + Err(std::io::Error::other(error_message)) +} + +fn logout_all_stores( + codex_home: &Path, + auth_credentials_store_mode: AuthCredentialsStoreMode, +) -> std::io::Result { + if auth_credentials_store_mode == AuthCredentialsStoreMode::Ephemeral { + return logout(codex_home, AuthCredentialsStoreMode::Ephemeral); } + let removed_ephemeral = logout(codex_home, AuthCredentialsStoreMode::Ephemeral)?; + let removed_managed = logout(codex_home, auth_credentials_store_mode)?; + Ok(removed_ephemeral || removed_managed) } fn load_auth( @@ -350,6 +407,12 @@ fn load_auth( enable_codex_api_key_env: bool, auth_credentials_store_mode: AuthCredentialsStoreMode, ) -> std::io::Result> { + let build_auth = |auth_dot_json: AuthDotJson, storage_mode| { + let client = crate::default_client::create_client(); + CodexAuth::from_auth_dot_json(codex_home, auth_dot_json, storage_mode, client) + }; + + // API key via env var takes precedence over any other auth method. if enable_codex_api_key_env && let Some(api_key) = read_codex_api_key_from_env() { let client = crate::default_client::create_client(); return Ok(Some(CodexAuth::from_api_key_with_client( @@ -358,39 +421,34 @@ fn load_auth( ))); } - let storage = create_auth_storage(codex_home.to_path_buf(), auth_credentials_store_mode); + // External ChatGPT auth tokens live in the in-memory (ephemeral) store. Always check this + // first so external auth takes precedence over any persisted credentials. + let ephemeral_storage = create_auth_storage( + codex_home.to_path_buf(), + AuthCredentialsStoreMode::Ephemeral, + ); + if let Some(auth_dot_json) = ephemeral_storage.load()? { + let auth = build_auth(auth_dot_json, AuthCredentialsStoreMode::Ephemeral)?; + return Ok(Some(auth)); + } - let client = crate::default_client::create_client(); + // If the caller explicitly requested ephemeral auth, there is no persisted fallback. + if auth_credentials_store_mode == AuthCredentialsStoreMode::Ephemeral { + return Ok(None); + } + + // Fall back to the configured persistent store (file/keyring/auto) for managed auth. + let storage = create_auth_storage(codex_home.to_path_buf(), auth_credentials_store_mode); let auth_dot_json = match storage.load()? { Some(auth) => auth, None => return Ok(None), }; - let AuthDotJson { - openai_api_key: auth_json_api_key, - tokens, - last_refresh, - } = auth_dot_json; - - // Prefer AuthMode.ApiKey if it's set in the auth.json. - if let Some(api_key) = &auth_json_api_key { - return Ok(Some(CodexAuth::from_api_key_with_client(api_key, client))); - } - - Ok(Some(CodexAuth { - api_key: None, - mode: AuthMode::ChatGPT, - storage: storage.clone(), - auth_dot_json: Arc::new(Mutex::new(Some(AuthDotJson { - openai_api_key: None, - tokens, - last_refresh, - }))), - client, - })) + let auth = build_auth(auth_dot_json, auth_credentials_store_mode)?; + Ok(Some(auth)) } -async fn update_tokens( +fn update_tokens( storage: &Arc, id_token: Option, access_token: Option, @@ -537,17 +595,108 @@ fn refresh_token_endpoint() -> String { .unwrap_or_else(|_| REFRESH_TOKEN_URL.to_string()) } +impl AuthDotJson { + fn from_external_tokens(external: &ExternalAuthTokens, id_token: IdTokenInfo) -> Self { + let account_id = id_token.chatgpt_account_id.clone(); + let tokens = TokenData { + id_token, + access_token: external.access_token.clone(), + refresh_token: String::new(), + account_id, + }; + + Self { + auth_mode: Some(AuthMode::ChatgptAuthTokens), + openai_api_key: None, + tokens: Some(tokens), + last_refresh: Some(Utc::now()), + } + } + + fn from_external_token_strings(id_token: &str, access_token: &str) -> std::io::Result { + let id_token_info = parse_id_token(id_token).map_err(std::io::Error::other)?; + let external = ExternalAuthTokens { + access_token: access_token.to_string(), + id_token: id_token.to_string(), + }; + Ok(Self::from_external_tokens(&external, id_token_info)) + } + + fn resolved_mode(&self) -> AuthMode { + if let Some(mode) = self.auth_mode { + return mode; + } + if self.openai_api_key.is_some() { + return AuthMode::ApiKey; + } + AuthMode::ChatGPT + } + + fn storage_mode( + &self, + auth_credentials_store_mode: AuthCredentialsStoreMode, + ) -> AuthCredentialsStoreMode { + if self.resolved_mode() == AuthMode::ChatgptAuthTokens { + AuthCredentialsStoreMode::Ephemeral + } else { + auth_credentials_store_mode + } + } +} + +impl CodexAuth { + fn from_auth_dot_json( + codex_home: &Path, + auth_dot_json: AuthDotJson, + auth_credentials_store_mode: AuthCredentialsStoreMode, + client: CodexHttpClient, + ) -> std::io::Result { + let auth_mode = auth_dot_json.resolved_mode(); + if auth_mode == AuthMode::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.")); + }; + return Ok(CodexAuth::from_api_key_with_client(api_key, client)); + } + + let storage_mode = auth_dot_json.storage_mode(auth_credentials_store_mode); + let storage = create_auth_storage(codex_home.to_path_buf(), storage_mode); + Ok(Self { + api_key: None, + mode: auth_mode, + storage, + auth_dot_json: Arc::new(Mutex::new(Some(auth_dot_json))), + client, + }) + } +} + use std::sync::RwLock; /// Internal cached auth state. -#[derive(Clone, Debug)] +#[derive(Clone)] struct CachedAuth { auth: Option, + /// Callback used to refresh external auth by asking the parent app for new tokens. + external_refresher: Option>, +} + +impl Debug for CachedAuth { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CachedAuth") + .field("auth_mode", &self.auth.as_ref().map(|auth| auth.mode)) + .field( + "external_refresher", + &self.external_refresher.as_ref().map(|_| "present"), + ) + .finish() + } } enum UnauthorizedRecoveryStep { Reload, RefreshToken, + ExternalRefresh, Done, } @@ -556,38 +705,63 @@ enum ReloadOutcome { Skipped, } +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum UnauthorizedRecoveryMode { + Managed, + External, +} + // UnauthorizedRecovery is a state machine that handles an attempt to refresh the authentication when requests // to API fail with 401 status code. // The client calls next() every time it encounters a 401 error, one time per retry. // For API key based authentication, we don't do anything and let the error bubble to the user. +// // For ChatGPT based authentication, we: // 1. Attempt to reload the auth data from disk. We only reload if the account id matches the one the current process is running as. // 2. Attempt to refresh the token using OAuth token refresh flow. // If after both steps the server still responds with 401 we let the error bubble to the user. +// +// For external ChatGPT auth tokens (chatgptAuthTokens), UnauthorizedRecovery does not touch disk or refresh +// tokens locally. Instead it calls the ExternalAuthRefresher (account/chatgptAuthTokens/refresh) to ask the +// parent app for new tokens, stores them in the ephemeral auth store, and retries once. pub struct UnauthorizedRecovery { manager: Arc, step: UnauthorizedRecoveryStep, expected_account_id: Option, + mode: UnauthorizedRecoveryMode, } impl UnauthorizedRecovery { fn new(manager: Arc) -> Self { - let expected_account_id = manager - .auth_cached() - .as_ref() - .and_then(CodexAuth::get_account_id); + let cached_auth = manager.auth_cached(); + let expected_account_id = cached_auth.as_ref().and_then(CodexAuth::get_account_id); + let mode = match cached_auth { + Some(auth) if auth.mode == AuthMode::ChatgptAuthTokens => { + UnauthorizedRecoveryMode::External + } + _ => UnauthorizedRecoveryMode::Managed, + }; + let step = match mode { + UnauthorizedRecoveryMode::Managed => UnauthorizedRecoveryStep::Reload, + UnauthorizedRecoveryMode::External => UnauthorizedRecoveryStep::ExternalRefresh, + }; Self { manager, - step: UnauthorizedRecoveryStep::Reload, + step, expected_account_id, + mode, } } pub fn has_next(&self) -> bool { - if !self - .manager - .auth_cached() - .is_some_and(|auth| auth.mode == AuthMode::ChatGPT) + if !self.manager.auth_cached().is_some_and(|auth| { + matches!(auth.mode, AuthMode::ChatGPT | AuthMode::ChatgptAuthTokens) + }) { + return false; + } + + if self.mode == UnauthorizedRecoveryMode::External + && !self.manager.has_external_auth_refresher() { return false; } @@ -622,6 +796,12 @@ impl UnauthorizedRecovery { self.manager.refresh_token().await?; self.step = UnauthorizedRecoveryStep::Done; } + UnauthorizedRecoveryStep::ExternalRefresh => { + self.manager + .refresh_external_auth(ExternalAuthRefreshReason::Unauthorized) + .await?; + self.step = UnauthorizedRecoveryStep::Done; + } UnauthorizedRecoveryStep::Done => {} } Ok(()) @@ -642,6 +822,7 @@ pub struct AuthManager { inner: RwLock, enable_codex_api_key_env: bool, auth_credentials_store_mode: AuthCredentialsStoreMode, + forced_chatgpt_workspace_id: RwLock>, } impl AuthManager { @@ -654,7 +835,7 @@ impl AuthManager { enable_codex_api_key_env: bool, auth_credentials_store_mode: AuthCredentialsStoreMode, ) -> Self { - let auth = load_auth( + let managed_auth = load_auth( &codex_home, enable_codex_api_key_env, auth_credentials_store_mode, @@ -663,34 +844,46 @@ impl AuthManager { .flatten(); Self { codex_home, - inner: RwLock::new(CachedAuth { auth }), + inner: RwLock::new(CachedAuth { + auth: managed_auth, + external_refresher: None, + }), enable_codex_api_key_env, auth_credentials_store_mode, + forced_chatgpt_workspace_id: RwLock::new(None), } } #[cfg(any(test, feature = "test-support"))] /// Create an AuthManager with a specific CodexAuth, for testing only. pub fn from_auth_for_testing(auth: CodexAuth) -> Arc { - let cached = CachedAuth { auth: Some(auth) }; + let cached = CachedAuth { + auth: Some(auth), + external_refresher: None, + }; Arc::new(Self { codex_home: PathBuf::from("non-existent"), inner: RwLock::new(cached), enable_codex_api_key_env: false, auth_credentials_store_mode: AuthCredentialsStoreMode::File, + forced_chatgpt_workspace_id: RwLock::new(None), }) } #[cfg(any(test, feature = "test-support"))] /// Create an AuthManager with a specific CodexAuth and codex home, for testing only. pub fn from_auth_for_testing_with_home(auth: CodexAuth, codex_home: PathBuf) -> Arc { - let cached = CachedAuth { auth: Some(auth) }; + let cached = CachedAuth { + auth: Some(auth), + external_refresher: None, + }; Arc::new(Self { codex_home, inner: RwLock::new(cached), enable_codex_api_key_env: false, auth_credentials_store_mode: AuthCredentialsStoreMode::File, + forced_chatgpt_workspace_id: RwLock::new(None), }) } @@ -715,7 +908,7 @@ impl AuthManager { pub fn reload(&self) -> bool { tracing::info!("Reloading auth"); let new_auth = self.load_auth_from_storage(); - self.set_auth(new_auth) + self.set_cached_auth(new_auth) } fn reload_if_account_id_matches(&self, expected_account_id: Option<&str>) -> ReloadOutcome { @@ -739,11 +932,11 @@ impl AuthManager { } tracing::info!("Reloading auth for account {expected_account_id}"); - self.set_auth(new_auth); + self.set_cached_auth(new_auth); ReloadOutcome::Reloaded } - fn auths_equal(a: &Option, b: &Option) -> bool { + fn auths_equal(a: Option<&CodexAuth>, b: Option<&CodexAuth>) -> bool { match (a, b) { (None, None) => true, (Some(a), Some(b)) => a == b, @@ -761,9 +954,10 @@ impl AuthManager { .flatten() } - fn set_auth(&self, new_auth: Option) -> bool { + fn set_cached_auth(&self, new_auth: Option) -> bool { if let Ok(mut guard) = self.inner.write() { - let changed = !AuthManager::auths_equal(&guard.auth, &new_auth); + let previous = guard.auth.as_ref(); + let changed = !AuthManager::auths_equal(previous, new_auth.as_ref()); tracing::info!("Reloaded auth, changed: {changed}"); guard.auth = new_auth; changed @@ -772,6 +966,38 @@ impl AuthManager { } } + pub fn set_external_auth_refresher(&self, refresher: Arc) { + if let Ok(mut guard) = self.inner.write() { + guard.external_refresher = Some(refresher); + } + } + + pub fn set_forced_chatgpt_workspace_id(&self, workspace_id: Option) { + if let Ok(mut guard) = self.forced_chatgpt_workspace_id.write() { + *guard = workspace_id; + } + } + + pub fn forced_chatgpt_workspace_id(&self) -> Option { + self.forced_chatgpt_workspace_id + .read() + .ok() + .and_then(|guard| guard.clone()) + } + + pub fn has_external_auth_refresher(&self) -> bool { + self.inner + .read() + .ok() + .map(|guard| guard.external_refresher.is_some()) + .unwrap_or(false) + } + + pub fn is_external_auth_active(&self) -> bool { + self.auth_cached() + .is_some_and(|auth| auth.mode == AuthMode::ChatgptAuthTokens) + } + /// Convenience constructor returning an `Arc` wrapper. pub fn shared( codex_home: PathBuf, @@ -799,6 +1025,11 @@ impl AuthManager { Some(auth) => auth, None => return Ok(()), }; + if auth.mode == AuthMode::ChatgptAuthTokens { + return self + .refresh_external_auth(ExternalAuthRefreshReason::Unauthorized) + .await; + } let token_data = auth.get_current_token_data().ok_or_else(|| { RefreshTokenError::Transient(std::io::Error::other("Token data is not available.")) })?; @@ -813,7 +1044,7 @@ impl AuthManager { /// reloads the in‑memory auth cache so callers immediately observe the /// unauthenticated state. pub fn logout(&self) -> std::io::Result { - let removed = super::auth::logout(&self.codex_home, self.auth_credentials_store_mode)?; + let removed = logout_all_stores(&self.codex_home, self.auth_credentials_store_mode)?; // Always reload to clear any cached auth (even if file absent). self.reload(); Ok(removed) @@ -848,6 +1079,59 @@ impl AuthManager { Ok(true) } + async fn refresh_external_auth( + &self, + reason: ExternalAuthRefreshReason, + ) -> Result<(), RefreshTokenError> { + let forced_chatgpt_workspace_id = self.forced_chatgpt_workspace_id(); + let refresher = match self.inner.read() { + Ok(guard) => guard.external_refresher.clone(), + Err(_) => { + return Err(RefreshTokenError::Transient(std::io::Error::other( + "failed to read external auth state", + ))); + } + }; + + let Some(refresher) = refresher else { + return Err(RefreshTokenError::Transient(std::io::Error::other( + "external auth refresher is not configured", + ))); + }; + + let previous_account_id = self + .auth_cached() + .as_ref() + .and_then(CodexAuth::get_account_id); + let context = ExternalAuthRefreshContext { + reason, + previous_account_id, + }; + + let refreshed = refresher.refresh(context).await?; + let id_token = parse_id_token(&refreshed.id_token) + .map_err(|err| RefreshTokenError::Transient(std::io::Error::other(err)))?; + if let Some(expected_workspace_id) = forced_chatgpt_workspace_id.as_deref() { + let actual_workspace_id = id_token.chatgpt_account_id.as_deref(); + if actual_workspace_id != Some(expected_workspace_id) { + return Err(RefreshTokenError::Transient(std::io::Error::other( + format!( + "external auth refresh returned workspace {actual_workspace_id:?}, expected {expected_workspace_id:?}", + ), + ))); + } + } + let auth_dot_json = AuthDotJson::from_external_tokens(&refreshed, id_token); + save_auth( + &self.codex_home, + &auth_dot_json, + AuthCredentialsStoreMode::Ephemeral, + ) + .map_err(RefreshTokenError::Transient)?; + self.reload(); + Ok(()) + } + async fn refresh_tokens( &self, auth: &CodexAuth, @@ -861,7 +1145,6 @@ impl AuthManager { refresh_response.access_token, refresh_response.refresh_token, ) - .await .map_err(RefreshTokenError::from)?; Ok(()) @@ -910,7 +1193,6 @@ mod tests { Some("new-access-token".to_string()), Some("new-refresh-token".to_string()), ) - .await .expect("update_tokens should succeed"); let tokens = updated.tokens.expect("tokens should exist"); @@ -991,6 +1273,7 @@ mod tests { assert_eq!( &AuthDotJson { + auth_mode: None, openai_api_key: None, tokens: Some(TokenData { id_token: IdTokenInfo { @@ -1034,6 +1317,7 @@ mod tests { fn logout_removes_auth_file() -> Result<(), std::io::Error> { let dir = tempdir()?; let auth_dot_json = AuthDotJson { + auth_mode: Some(AuthMode::ApiKey), openai_api_key: Some("sk-test-key".to_string()), tokens: None, last_refresh: None, diff --git a/codex-rs/core/src/auth/storage.rs b/codex-rs/core/src/auth/storage.rs index 48b67aca0..b4f4bb8e7 100644 --- a/codex-rs/core/src/auth/storage.rs +++ b/codex-rs/core/src/auth/storage.rs @@ -5,6 +5,7 @@ use serde::Deserialize; use serde::Serialize; use sha2::Digest; use sha2::Sha256; +use std::collections::HashMap; use std::fmt::Debug; use std::fs::File; use std::fs::OpenOptions; @@ -15,11 +16,14 @@ use std::os::unix::fs::OpenOptionsExt; use std::path::Path; use std::path::PathBuf; use std::sync::Arc; +use std::sync::Mutex; use tracing::warn; use crate::token_data::TokenData; +use codex_app_server_protocol::AuthMode; use codex_keyring_store::DefaultKeyringStore; use codex_keyring_store::KeyringStore; +use once_cell::sync::Lazy; /// Determine where Codex should store CLI auth credentials. #[derive(Debug, Default, Copy, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] @@ -32,11 +36,16 @@ pub enum AuthCredentialsStoreMode { Keyring, /// Use keyring when available; otherwise, fall back to a file in CODEX_HOME. Auto, + /// Store credentials in memory only for the current process. + Ephemeral, } /// Expected structure for $CODEX_HOME/auth.json. #[derive(Deserialize, Serialize, Clone, Debug, PartialEq)] pub struct AuthDotJson { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub auth_mode: Option, + #[serde(rename = "OPENAI_API_KEY")] pub openai_api_key: Option, @@ -76,8 +85,8 @@ impl FileAuthStorage { Self { codex_home } } - /// Attempt to read and refresh the `auth.json` file in the given `CODEX_HOME` directory. - /// Returns the full AuthDotJson structure after refreshing if necessary. + /// Attempt to read and parse the `auth.json` file in the given `CODEX_HOME` directory. + /// Returns the full AuthDotJson structure. pub(super) fn try_read_auth_json(&self, auth_file: &Path) -> std::io::Result { let mut file = File::open(auth_file)?; let mut contents = String::new(); @@ -256,6 +265,49 @@ impl AuthStorageBackend for AutoAuthStorage { } } +// A global in-memory store for mapping codex_home -> AuthDotJson. +static EPHEMERAL_AUTH_STORE: Lazy>> = + Lazy::new(|| Mutex::new(HashMap::new())); + +#[derive(Clone, Debug)] +struct EphemeralAuthStorage { + codex_home: PathBuf, +} + +impl EphemeralAuthStorage { + fn new(codex_home: PathBuf) -> Self { + Self { codex_home } + } + + fn with_store(&self, action: F) -> std::io::Result + where + F: FnOnce(&mut HashMap, String) -> std::io::Result, + { + let key = compute_store_key(&self.codex_home)?; + let mut store = EPHEMERAL_AUTH_STORE + .lock() + .map_err(|_| std::io::Error::other("failed to lock ephemeral auth storage"))?; + action(&mut store, key) + } +} + +impl AuthStorageBackend for EphemeralAuthStorage { + fn load(&self) -> std::io::Result> { + self.with_store(|store, key| Ok(store.get(&key).cloned())) + } + + fn save(&self, auth: &AuthDotJson) -> std::io::Result<()> { + self.with_store(|store, key| { + store.insert(key, auth.clone()); + Ok(()) + }) + } + + fn delete(&self) -> std::io::Result { + self.with_store(|store, key| Ok(store.remove(&key).is_some())) + } +} + pub(super) fn create_auth_storage( codex_home: PathBuf, mode: AuthCredentialsStoreMode, @@ -275,6 +327,7 @@ fn create_auth_storage_with_keyring_store( Arc::new(KeyringAuthStorage::new(codex_home, keyring_store)) } AuthCredentialsStoreMode::Auto => Arc::new(AutoAuthStorage::new(codex_home, keyring_store)), + AuthCredentialsStoreMode::Ephemeral => Arc::new(EphemeralAuthStorage::new(codex_home)), } } @@ -296,6 +349,7 @@ mod tests { let codex_home = tempdir()?; let storage = FileAuthStorage::new(codex_home.path().to_path_buf()); let auth_dot_json = AuthDotJson { + auth_mode: Some(AuthMode::ApiKey), openai_api_key: Some("test-key".to_string()), tokens: None, last_refresh: Some(Utc::now()), @@ -315,6 +369,7 @@ mod tests { let codex_home = tempdir()?; let storage = FileAuthStorage::new(codex_home.path().to_path_buf()); let auth_dot_json = AuthDotJson { + auth_mode: Some(AuthMode::ApiKey), openai_api_key: Some("test-key".to_string()), tokens: None, last_refresh: Some(Utc::now()), @@ -336,6 +391,7 @@ mod tests { fn file_storage_delete_removes_auth_file() -> anyhow::Result<()> { let dir = tempdir()?; let auth_dot_json = AuthDotJson { + auth_mode: Some(AuthMode::ApiKey), openai_api_key: Some("sk-test-key".to_string()), tokens: None, last_refresh: None, @@ -350,6 +406,32 @@ mod tests { Ok(()) } + #[test] + fn ephemeral_storage_save_load_delete_is_in_memory_only() -> anyhow::Result<()> { + let dir = tempdir()?; + let storage = create_auth_storage( + dir.path().to_path_buf(), + AuthCredentialsStoreMode::Ephemeral, + ); + let auth_dot_json = AuthDotJson { + auth_mode: Some(AuthMode::ApiKey), + openai_api_key: Some("sk-ephemeral".to_string()), + tokens: None, + last_refresh: Some(Utc::now()), + }; + + storage.save(&auth_dot_json)?; + let loaded = storage.load()?; + assert_eq!(Some(auth_dot_json), loaded); + + let removed = storage.delete()?; + assert!(removed); + let loaded = storage.load()?; + assert_eq!(None, loaded); + assert!(!get_auth_file(dir.path()).exists()); + Ok(()) + } + fn seed_keyring_and_fallback_auth_file_for_delete( mock_keyring: &MockKeyringStore, codex_home: &Path, @@ -425,6 +507,7 @@ mod tests { fn auth_with_prefix(prefix: &str) -> AuthDotJson { AuthDotJson { + auth_mode: Some(AuthMode::ApiKey), openai_api_key: Some(format!("{prefix}-api-key")), tokens: Some(TokenData { id_token: id_token_with_prefix(prefix), @@ -445,6 +528,7 @@ mod tests { Arc::new(mock_keyring.clone()), ); let expected = AuthDotJson { + auth_mode: Some(AuthMode::ApiKey), openai_api_key: Some("sk-test".to_string()), tokens: None, last_refresh: None, @@ -481,6 +565,7 @@ mod tests { let auth_file = get_auth_file(codex_home.path()); std::fs::write(&auth_file, "stale")?; let auth = AuthDotJson { + auth_mode: Some(AuthMode::ChatGPT), openai_api_key: None, tokens: Some(TokenData { id_token: Default::default(), diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 4cb37a46d..0f05f58c4 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -469,7 +469,9 @@ impl ModelClientSession { .config .features .enabled(Feature::EnableRequestCompression) - && auth.is_some_and(|auth| auth.mode == AuthMode::ChatGPT) + && auth.is_some_and(|auth| { + matches!(auth.mode, AuthMode::ChatGPT | AuthMode::ChatgptAuthTokens) + }) && self.state.provider.is_openai() { Compression::Zstd diff --git a/codex-rs/core/src/model_provider_info.rs b/codex-rs/core/src/model_provider_info.rs index 4f69f02d9..0f3551480 100644 --- a/codex-rs/core/src/model_provider_info.rs +++ b/codex-rs/core/src/model_provider_info.rs @@ -137,7 +137,10 @@ impl ModelProviderInfo { &self, auth_mode: Option, ) -> crate::error::Result { - let default_base_url = if matches!(auth_mode, Some(AuthMode::ChatGPT)) { + let default_base_url = if matches!( + auth_mode, + Some(AuthMode::ChatGPT | AuthMode::ChatgptAuthTokens) + ) { "https://chatgpt.com/backend-api/codex" } else { "https://api.openai.com/v1" diff --git a/codex-rs/core/src/models_manager/manager.rs b/codex-rs/core/src/models_manager/manager.rs index 77fb44cd5..c993032ea 100644 --- a/codex-rs/core/src/models_manager/manager.rs +++ b/codex-rs/core/src/models_manager/manager.rs @@ -204,7 +204,8 @@ impl ModelsManager { let _timer = codex_otel::start_global_timer("codex.remote_models.fetch_update.duration_ms", &[]); let auth = self.auth_manager.auth().await; - let api_provider = self.provider.to_api_provider(Some(AuthMode::ChatGPT))?; + let auth_mode = self.auth_manager.get_auth_mode(); + let api_provider = self.provider.to_api_provider(auth_mode)?; let api_auth = auth_provider_from_auth(auth.clone(), &self.provider)?; let transport = ReqwestTransport::new(build_reqwest_client()); let client = ModelsClient::new(transport, api_provider, api_auth); @@ -271,7 +272,10 @@ impl ModelsManager { let remote_presets: Vec = remote_models.into_iter().map(Into::into).collect(); let existing_presets = self.local_models.clone(); let mut merged_presets = ModelPreset::merge(remote_presets, existing_presets); - let chatgpt_mode = self.auth_manager.get_auth_mode() == Some(AuthMode::ChatGPT); + let chatgpt_mode = matches!( + self.auth_manager.get_auth_mode(), + Some(AuthMode::ChatGPT | AuthMode::ChatgptAuthTokens) + ); merged_presets = ModelPreset::filter_by_auth(merged_presets, chatgpt_mode); for preset in &mut merged_presets { diff --git a/codex-rs/core/src/token_data.rs b/codex-rs/core/src/token_data.rs index 744cc73d9..2010af16f 100644 --- a/codex-rs/core/src/token_data.rs +++ b/codex-rs/core/src/token_data.rs @@ -76,10 +76,18 @@ pub(crate) enum KnownPlan { struct IdClaims { #[serde(default)] email: Option, + #[serde(rename = "https://api.openai.com/profile", default)] + profile: Option, #[serde(rename = "https://api.openai.com/auth", default)] auth: Option, } +#[derive(Deserialize)] +struct ProfileClaims { + #[serde(default)] + email: Option, +} + #[derive(Deserialize)] struct AuthClaims { #[serde(default)] @@ -112,17 +120,20 @@ pub fn parse_id_token(id_token: &str) -> Result { let payload_bytes = base64::engine::general_purpose::URL_SAFE_NO_PAD.decode(payload_b64)?; let claims: IdClaims = serde_json::from_slice(&payload_bytes)?; + let email = claims + .email + .or_else(|| claims.profile.and_then(|profile| profile.email)); match claims.auth { Some(auth) => Ok(IdTokenInfo { - email: claims.email, + email, raw_jwt: id_token.to_string(), chatgpt_plan_type: auth.chatgpt_plan_type, chatgpt_user_id: auth.chatgpt_user_id.or(auth.user_id), chatgpt_account_id: auth.chatgpt_account_id, }), None => Ok(IdTokenInfo { - email: claims.email, + email, raw_jwt: id_token.to_string(), chatgpt_plan_type: None, chatgpt_user_id: None, diff --git a/codex-rs/core/tests/suite/auth_refresh.rs b/codex-rs/core/tests/suite/auth_refresh.rs index d0b8d2738..4ef3b82ee 100644 --- a/codex-rs/core/tests/suite/auth_refresh.rs +++ b/codex-rs/core/tests/suite/auth_refresh.rs @@ -3,6 +3,7 @@ use anyhow::Result; use base64::Engine; use chrono::Duration; use chrono::Utc; +use codex_app_server_protocol::AuthMode; use codex_core::AuthManager; use codex_core::auth::AuthCredentialsStoreMode; use codex_core::auth::AuthDotJson; @@ -50,6 +51,7 @@ async fn refresh_token_succeeds_updates_storage() -> Result<()> { let initial_last_refresh = Utc::now() - Duration::days(1); let initial_tokens = build_tokens(INITIAL_ACCESS_TOKEN, INITIAL_REFRESH_TOKEN); let initial_auth = AuthDotJson { + auth_mode: Some(AuthMode::ChatGPT), openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(initial_last_refresh), @@ -111,6 +113,7 @@ async fn returns_fresh_tokens_as_is() -> Result<()> { let initial_last_refresh = Utc::now() - Duration::days(1); let initial_tokens = build_tokens(INITIAL_ACCESS_TOKEN, INITIAL_REFRESH_TOKEN); let initial_auth = AuthDotJson { + auth_mode: Some(AuthMode::ChatGPT), openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(initial_last_refresh), @@ -156,6 +159,7 @@ async fn refreshes_token_when_last_refresh_is_stale() -> Result<()> { let stale_refresh = Utc::now() - Duration::days(9); let initial_tokens = build_tokens(INITIAL_ACCESS_TOKEN, INITIAL_REFRESH_TOKEN); let initial_auth = AuthDotJson { + auth_mode: Some(AuthMode::ChatGPT), openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(stale_refresh), @@ -214,6 +218,7 @@ async fn refresh_token_returns_permanent_error_for_expired_refresh_token() -> Re let initial_last_refresh = Utc::now() - Duration::days(1); let initial_tokens = build_tokens(INITIAL_ACCESS_TOKEN, INITIAL_REFRESH_TOKEN); let initial_auth = AuthDotJson { + auth_mode: Some(AuthMode::ChatGPT), openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(initial_last_refresh), @@ -263,6 +268,7 @@ async fn refresh_token_returns_transient_error_on_server_failure() -> Result<()> let initial_last_refresh = Utc::now() - Duration::days(1); let initial_tokens = build_tokens(INITIAL_ACCESS_TOKEN, INITIAL_REFRESH_TOKEN); let initial_auth = AuthDotJson { + auth_mode: Some(AuthMode::ChatGPT), openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(initial_last_refresh), @@ -314,6 +320,7 @@ async fn unauthorized_recovery_reloads_then_refreshes_tokens() -> Result<()> { let initial_last_refresh = Utc::now() - Duration::days(1); let initial_tokens = build_tokens(INITIAL_ACCESS_TOKEN, INITIAL_REFRESH_TOKEN); let initial_auth = AuthDotJson { + auth_mode: Some(AuthMode::ChatGPT), openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(initial_last_refresh), @@ -322,6 +329,7 @@ async fn unauthorized_recovery_reloads_then_refreshes_tokens() -> Result<()> { let disk_tokens = build_tokens("disk-access-token", "disk-refresh-token"); let disk_auth = AuthDotJson { + auth_mode: Some(AuthMode::ChatGPT), openai_api_key: None, tokens: Some(disk_tokens.clone()), last_refresh: Some(initial_last_refresh), @@ -404,6 +412,7 @@ async fn unauthorized_recovery_skips_reload_on_account_mismatch() -> Result<()> let initial_last_refresh = Utc::now() - Duration::days(1); let initial_tokens = build_tokens(INITIAL_ACCESS_TOKEN, INITIAL_REFRESH_TOKEN); let initial_auth = AuthDotJson { + auth_mode: Some(AuthMode::ChatGPT), openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(initial_last_refresh), @@ -418,6 +427,7 @@ async fn unauthorized_recovery_skips_reload_on_account_mismatch() -> Result<()> ..disk_tokens.clone() }; let disk_auth = AuthDotJson { + auth_mode: Some(AuthMode::ChatGPT), openai_api_key: None, tokens: Some(disk_tokens), last_refresh: Some(initial_last_refresh), @@ -481,6 +491,7 @@ async fn unauthorized_recovery_requires_chatgpt_auth() -> Result<()> { let server = MockServer::start().await; let ctx = RefreshTokenTestContext::new(&server)?; let auth = AuthDotJson { + auth_mode: Some(AuthMode::ApiKey), openai_api_key: Some("sk-test".to_string()), tokens: None, last_refresh: None, diff --git a/codex-rs/login/src/server.rs b/codex-rs/login/src/server.rs index 999c19072..5272c84ea 100644 --- a/codex-rs/login/src/server.rs +++ b/codex-rs/login/src/server.rs @@ -14,6 +14,7 @@ use crate::pkce::PkceCodes; use crate::pkce::generate_pkce; use base64::Engine; use chrono::Utc; +use codex_app_server_protocol::AuthMode; use codex_core::auth::AuthCredentialsStoreMode; use codex_core::auth::AuthDotJson; use codex_core::auth::save_auth; @@ -559,6 +560,7 @@ pub(crate) async fn persist_tokens_async( tokens.account_id = Some(acc.to_string()); } let auth = AuthDotJson { + auth_mode: Some(AuthMode::ChatGPT), openai_api_key: api_key, tokens: Some(tokens), last_refresh: Some(Utc::now()), diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index c13213051..61006d80c 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -3457,7 +3457,10 @@ impl ChatWidget { fn prefetch_rate_limits(&mut self) { self.stop_rate_limit_poller(); - if self.auth_manager.auth_cached().map(|auth| auth.mode) != Some(AuthMode::ChatGPT) { + if !matches!( + self.auth_manager.auth_cached().map(|auth| auth.mode), + Some(AuthMode::ChatGPT | AuthMode::ChatgptAuthTokens) + ) { return; } @@ -3470,7 +3473,7 @@ impl ChatWidget { loop { if let Some(auth) = auth_manager.auth().await - && auth.mode == AuthMode::ChatGPT + && matches!(auth.mode, AuthMode::ChatGPT | AuthMode::ChatgptAuthTokens) && let Some(snapshot) = fetch_rate_limits(base_url.clone(), auth).await { app_event_tx.send(AppEvent::RateLimitSnapshotFetched(snapshot)); diff --git a/codex-rs/tui/src/onboarding/auth.rs b/codex-rs/tui/src/onboarding/auth.rs index 551fa9180..daf64f36d 100644 --- a/codex-rs/tui/src/onboarding/auth.rs +++ b/codex-rs/tui/src/onboarding/auth.rs @@ -661,7 +661,10 @@ impl AuthModeWidget { } fn handle_existing_chatgpt_login(&mut self) -> bool { - if matches!(self.login_status, LoginStatus::AuthMode(AuthMode::ChatGPT)) { + if matches!( + self.login_status, + LoginStatus::AuthMode(AuthMode::ChatGPT | AuthMode::ChatgptAuthTokens) + ) { *self.sign_in_state.write().unwrap() = SignInState::ChatGptSuccess; self.request_frame.schedule_frame(); true diff --git a/codex-rs/tui/src/status/helpers.rs b/codex-rs/tui/src/status/helpers.rs index 78a956b27..91f958a9c 100644 --- a/codex-rs/tui/src/status/helpers.rs +++ b/codex-rs/tui/src/status/helpers.rs @@ -91,7 +91,7 @@ pub(crate) fn compose_account_display( let auth = auth_manager.auth_cached()?; match auth.mode { - AuthMode::ChatGPT => { + AuthMode::ChatGPT | AuthMode::ChatgptAuthTokens => { let email = auth.get_account_email(); let plan = plan .map(|plan_type| title_case(format!("{plan_type:?}").as_str()))