From 05f0b4f590286536796200253252e3084b4719ff Mon Sep 17 00:00:00 2001 From: Celia Chen Date: Wed, 5 Nov 2025 13:52:50 -0800 Subject: [PATCH] [App-server] Implement v2 for `account/login/start` and `account/login/completed` (#6183) This PR implements `account/login/start` and `account/login/completed`. Instead of having separate endpoints for login with chatgpt and api, we have a single enum handling different login methods. For sync auth methods like sign in with api key, we still send a `completed` notification back to be compatible with the async login flow. --- .../src/protocol/common.rs | 19 +- .../app-server-protocol/src/protocol/v1.rs | 3 +- .../app-server-protocol/src/protocol/v2.rs | 46 ++- .../app-server/src/codex_message_processor.rs | 352 ++++++++++++------ codex-rs/app-server/src/outgoing_message.rs | 30 +- .../app-server/tests/common/mcp_process.rs | 20 + codex-rs/app-server/tests/suite/v2/account.rs | 187 +++++++++- 7 files changed, 521 insertions(+), 136 deletions(-) diff --git a/codex-rs/app-server-protocol/src/protocol/common.rs b/codex-rs/app-server-protocol/src/protocol/common.rs index 275d294a..17299fc3 100644 --- a/codex-rs/app-server-protocol/src/protocol/common.rs +++ b/codex-rs/app-server-protocol/src/protocol/common.rs @@ -141,8 +141,8 @@ client_request_definitions! { response: v2::ModelListResponse, }, - #[serde(rename = "account/login")] - #[ts(rename = "account/login")] + #[serde(rename = "account/login/start")] + #[ts(rename = "account/login/start")] LoginAccount { params: v2::LoginAccountParams, response: v2::LoginAccountResponse, @@ -514,8 +514,15 @@ server_notification_definitions! { AccountUpdated => "account/updated" (v2::AccountUpdatedNotification), AccountRateLimitsUpdated => "account/rateLimits/updated" (v2::AccountRateLimitsUpdatedNotification), + #[serde(rename = "account/login/completed")] + #[ts(rename = "account/login/completed")] + #[strum(serialize = "account/login/completed")] + AccountLoginCompleted(v2::AccountLoginCompletedNotification), + /// DEPRECATED NOTIFICATIONS below AuthStatusChange(v1::AuthStatusChangeNotification), + + /// Deprecated: use `account/login/completed` instead. LoginChatGptComplete(v1::LoginChatGptCompleteNotification), SessionConfigured(v1::SessionConfiguredNotification), } @@ -680,7 +687,7 @@ mod tests { }; assert_eq!( json!({ - "method": "account/login", + "method": "account/login/start", "id": 2, "params": { "type": "apiKey", @@ -696,11 +703,11 @@ mod tests { fn serialize_account_login_chatgpt() -> Result<()> { let request = ClientRequest::LoginAccount { request_id: RequestId::Integer(3), - params: v2::LoginAccountParams::ChatGpt, + params: v2::LoginAccountParams::Chatgpt, }; assert_eq!( json!({ - "method": "account/login", + "method": "account/login/start", "id": 3, "params": { "type": "chatgpt" @@ -756,7 +763,7 @@ mod tests { serde_json::to_value(&api_key)?, ); - let chatgpt = v2::Account::ChatGpt { + let chatgpt = v2::Account::Chatgpt { email: Some("user@example.com".to_string()), plan_type: PlanType::Plus, }; diff --git a/codex-rs/app-server-protocol/src/protocol/v1.rs b/codex-rs/app-server-protocol/src/protocol/v1.rs index 252fb041..bbd825ba 100644 --- a/codex-rs/app-server-protocol/src/protocol/v1.rs +++ b/codex-rs/app-server-protocol/src/protocol/v1.rs @@ -374,10 +374,11 @@ pub enum InputItem { LocalImage { path: PathBuf }, } -// Deprecated notifications (v1) +// Deprecated in favor of AccountLoginCompletedNotification. #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] #[serde(rename_all = "camelCase")] +/// Deprecated: use `account/login/completed` instead. pub struct LoginChatGptCompleteNotification { #[schemars(with = "String")] pub login_id: Uuid, diff --git a/codex-rs/app-server-protocol/src/protocol/v2.rs b/codex-rs/app-server-protocol/src/protocol/v2.rs index f96bbf43..a59f00d4 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2.rs @@ -13,7 +13,6 @@ use serde::Deserialize; use serde::Serialize; use serde_json::Value as JsonValue; use ts_rs::TS; -use uuid::Uuid; // Macro to declare a camelCased API v2 enum mirroring a core enum which // tends to use kebab-case. @@ -126,7 +125,7 @@ pub enum Account { #[serde(rename = "chatgpt", rename_all = "camelCase")] #[ts(rename = "chatgpt", rename_all = "camelCase")] - ChatGpt { + Chatgpt { email: Option, plan_type: PlanType, }, @@ -137,8 +136,8 @@ pub enum Account { #[ts(tag = "type")] #[ts(export_to = "v2/")] pub enum LoginAccountParams { - #[serde(rename = "apiKey")] - #[ts(rename = "apiKey")] + #[serde(rename = "apiKey", rename_all = "camelCase")] + #[ts(rename = "apiKey", rename_all = "camelCase")] ApiKey { #[serde(rename = "apiKey")] #[ts(rename = "apiKey")] @@ -146,20 +145,26 @@ pub enum LoginAccountParams { }, #[serde(rename = "chatgpt")] #[ts(rename = "chatgpt")] - ChatGpt, + Chatgpt, } #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] -#[serde(rename_all = "camelCase")] +#[serde(tag = "type", rename_all = "camelCase")] +#[ts(tag = "type")] #[ts(export_to = "v2/")] -pub struct LoginAccountResponse { - /// Only set if the login method is ChatGPT. - #[schemars(with = "String")] - pub login_id: Option, - - /// URL the client should open in a browser to initiate the OAuth flow. - /// Only set if the login method is ChatGPT. - pub auth_url: Option, +pub enum LoginAccountResponse { + #[serde(rename = "apiKey", rename_all = "camelCase")] + #[ts(rename = "apiKey", rename_all = "camelCase")] + ApiKey {}, + #[serde(rename = "chatgpt", rename_all = "camelCase")] + #[ts(rename = "chatgpt", rename_all = "camelCase")] + Chatgpt { + // Use plain String for identifiers to avoid TS/JSON Schema quirks around uuid-specific types. + // Convert to/from UUIDs at the application layer as needed. + login_id: String, + /// URL the client should open in a browser to initiate the OAuth flow. + auth_url: String, + }, } #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] @@ -335,7 +340,7 @@ pub struct Thread { #[serde(rename_all = "camelCase")] #[ts(export_to = "v2/")] pub struct AccountUpdatedNotification { - pub auth_method: Option, + pub auth_mode: Option, } #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] @@ -611,3 +616,14 @@ impl From for RateLimitWindow { } } } + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "v2/")] +pub struct AccountLoginCompletedNotification { + // Use plain String for identifiers to avoid TS/JSON Schema quirks around uuid-specific types. + // Convert to/from UUIDs at the application layer as needed. + pub login_id: Option, + pub success: bool, + pub error: Option, +} diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index 1d7c0a48..4c14ce3a 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -4,6 +4,7 @@ use crate::fuzzy_file_search::run_fuzzy_file_search; use crate::models::supported_models; use crate::outgoing_message::OutgoingMessageSender; use crate::outgoing_message::OutgoingNotification; +use codex_app_server_protocol::AccountLoginCompletedNotification; use codex_app_server_protocol::AccountRateLimitsUpdatedNotification; use codex_app_server_protocol::AccountUpdatedNotification; use codex_app_server_protocol::AddConversationListenerParams; @@ -36,6 +37,7 @@ use codex_app_server_protocol::InterruptConversationResponse; use codex_app_server_protocol::JSONRPCErrorError; use codex_app_server_protocol::ListConversationsParams; use codex_app_server_protocol::ListConversationsResponse; +use codex_app_server_protocol::LoginAccountParams; use codex_app_server_protocol::LoginApiKeyParams; use codex_app_server_protocol::LoginApiKeyResponse; use codex_app_server_protocol::LoginChatGptCompleteNotification; @@ -222,12 +224,8 @@ impl CodexMessageProcessor { ClientRequest::ModelList { request_id, params } => { self.list_models(request_id, params).await; } - ClientRequest::LoginAccount { - request_id, - params: _, - } => { - self.send_unimplemented_error(request_id, "account/login") - .await; + ClientRequest::LoginAccount { request_id, params } => { + self.login_v2(request_id, params).await; } ClientRequest::LogoutAccount { request_id, @@ -267,13 +265,13 @@ impl CodexMessageProcessor { self.git_diff_to_origin(request_id, params.cwd).await; } ClientRequest::LoginApiKey { request_id, params } => { - self.login_api_key(request_id, params).await; + self.login_api_key_v1(request_id, params).await; } ClientRequest::LoginChatGpt { request_id, params: _, } => { - self.login_chatgpt(request_id).await; + self.login_chatgpt_v1(request_id).await; } ClientRequest::CancelLoginChatGpt { request_id, params } => { self.cancel_login_chatgpt(request_id, params.login_id).await; @@ -335,20 +333,34 @@ impl CodexMessageProcessor { self.outgoing.send_error(request_id, error).await; } - async fn login_api_key(&mut self, request_id: RequestId, params: LoginApiKeyParams) { + async fn login_v2(&mut self, request_id: RequestId, params: LoginAccountParams) { + match params { + LoginAccountParams::ApiKey { api_key } => { + self.login_api_key_v2(request_id, LoginApiKeyParams { api_key }) + .await; + } + LoginAccountParams::Chatgpt => { + self.login_chatgpt_v2(request_id).await; + } + } + } + + async fn login_api_key_common( + &mut self, + params: &LoginApiKeyParams, + ) -> std::result::Result<(), JSONRPCErrorError> { if matches!( self.config.forced_login_method, Some(ForcedLoginMethod::Chatgpt) ) { - let error = JSONRPCErrorError { + return Err(JSONRPCErrorError { code: INVALID_REQUEST_ERROR_CODE, message: "API key login is disabled. Use ChatGPT login instead.".to_string(), data: None, - }; - self.outgoing.send_error(request_id, error).await; - return; + }); } + // Cancel any active login attempt. { let mut guard = self.active_login.lock().await; if let Some(active) = guard.take() { @@ -363,6 +375,19 @@ impl CodexMessageProcessor { ) { Ok(()) => { self.auth_manager.reload(); + Ok(()) + } + Err(err) => Err(JSONRPCErrorError { + code: INTERNAL_ERROR_CODE, + message: format!("failed to save api key: {err}"), + data: None, + }), + } + } + + async fn login_api_key_v1(&mut self, request_id: RequestId, params: LoginApiKeyParams) { + match self.login_api_key_common(¶ms).await { + Ok(()) => { self.outgoing .send_response(request_id, LoginApiKeyResponse {}) .await; @@ -374,31 +399,57 @@ impl CodexMessageProcessor { .send_server_notification(ServerNotification::AuthStatusChange(payload)) .await; } - Err(err) => { - let error = JSONRPCErrorError { - code: INTERNAL_ERROR_CODE, - message: format!("failed to save api key: {err}"), - data: None, - }; + Err(error) => { self.outgoing.send_error(request_id, error).await; } } } - async fn login_chatgpt(&mut self, request_id: RequestId) { + async fn login_api_key_v2(&mut self, request_id: RequestId, params: LoginApiKeyParams) { + match self.login_api_key_common(¶ms).await { + Ok(()) => { + let response = codex_app_server_protocol::LoginAccountResponse::ApiKey {}; + self.outgoing.send_response(request_id, response).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.auth().map(|auth| auth.mode), + }; + self.outgoing + .send_server_notification(ServerNotification::AccountUpdated(payload_v2)) + .await; + } + Err(error) => { + self.outgoing.send_error(request_id, error).await; + } + } + } + + // Build options for a ChatGPT login attempt; performs validation. + async fn login_chatgpt_common( + &self, + ) -> std::result::Result { let config = self.config.as_ref(); if matches!(config.forced_login_method, Some(ForcedLoginMethod::Api)) { - let error = JSONRPCErrorError { + return Err(JSONRPCErrorError { code: INVALID_REQUEST_ERROR_CODE, message: "ChatGPT login is disabled. Use API key login instead.".to_string(), data: None, - }; - self.outgoing.send_error(request_id, error).await; - return; + }); } - let opts = LoginServerOptions { + Ok(LoginServerOptions { open_browser: false, ..LoginServerOptions::new( config.codex_home.clone(), @@ -406,99 +457,190 @@ impl CodexMessageProcessor { config.forced_chatgpt_workspace_id.clone(), config.cli_auth_credentials_store_mode, ) - }; + }) + } - enum LoginChatGptReply { - Response(LoginChatGptResponse), - Error(JSONRPCErrorError), - } + // Deprecated in favor of login_chatgpt_v2. + async fn login_chatgpt_v1(&mut self, request_id: RequestId) { + match self.login_chatgpt_common().await { + Ok(opts) => match run_login_server(opts) { + Ok(server) => { + let login_id = Uuid::new_v4(); + let shutdown_handle = server.cancel_handle(); - let reply = match run_login_server(opts) { - Ok(server) => { - let login_id = Uuid::new_v4(); - let shutdown_handle = server.cancel_handle(); - - // Replace active login if present. - { - let mut guard = self.active_login.lock().await; - if let Some(existing) = guard.take() { - existing.drop(); - } - *guard = Some(ActiveLogin { - shutdown_handle: shutdown_handle.clone(), - login_id, - }); - } - - let response = LoginChatGptResponse { - login_id, - auth_url: server.auth_url.clone(), - }; - - // Spawn background task to monitor completion. - let outgoing_clone = self.outgoing.clone(); - let active_login = self.active_login.clone(); - let auth_manager = self.auth_manager.clone(); - tokio::spawn(async move { - let (success, error_msg) = match tokio::time::timeout( - LOGIN_CHATGPT_TIMEOUT, - server.block_until_done(), - ) - .await + // Replace active login if present. { - Ok(Ok(())) => (true, None), - Ok(Err(err)) => (false, Some(format!("Login server error: {err}"))), - Err(_elapsed) => { - // Timeout: cancel server and report - shutdown_handle.shutdown(); - (false, Some("Login timed out".to_string())) + let mut guard = self.active_login.lock().await; + if let Some(existing) = guard.take() { + existing.drop(); } - }; - let payload = LoginChatGptCompleteNotification { - login_id, - success, - error: error_msg, - }; - outgoing_clone - .send_server_notification(ServerNotification::LoginChatGptComplete(payload)) - .await; + *guard = Some(ActiveLogin { + shutdown_handle: shutdown_handle.clone(), + login_id, + }); + } - // Send an auth status change notification. - if success { - // Update in-memory auth cache now that login completed. - auth_manager.reload(); + // Spawn background task to monitor completion. + let outgoing_clone = self.outgoing.clone(); + let active_login = self.active_login.clone(); + let auth_manager = self.auth_manager.clone(); + let auth_url = server.auth_url.clone(); + tokio::spawn(async move { + let (success, error_msg) = match tokio::time::timeout( + LOGIN_CHATGPT_TIMEOUT, + server.block_until_done(), + ) + .await + { + Ok(Ok(())) => (true, None), + Ok(Err(err)) => (false, Some(format!("Login server error: {err}"))), + Err(_elapsed) => { + shutdown_handle.shutdown(); + (false, Some("Login timed out".to_string())) + } + }; - // Notify clients with the actual current auth mode. - let current_auth_method = auth_manager.auth().map(|a| a.mode); - let payload = AuthStatusChangeNotification { - auth_method: current_auth_method, + let payload = LoginChatGptCompleteNotification { + login_id, + success, + error: error_msg.clone(), }; outgoing_clone - .send_server_notification(ServerNotification::AuthStatusChange(payload)) + .send_server_notification(ServerNotification::LoginChatGptComplete( + payload, + )) .await; + + if success { + auth_manager.reload(); + + // Notify clients with the actual current auth mode. + let current_auth_method = auth_manager.auth().map(|a| a.mode); + let payload = AuthStatusChangeNotification { + auth_method: current_auth_method, + }; + outgoing_clone + .send_server_notification(ServerNotification::AuthStatusChange( + payload, + )) + .await; + } + + // Clear the active login if it matches this attempt. It may have been replaced or cancelled. + let mut guard = active_login.lock().await; + if guard.as_ref().map(|l| l.login_id) == Some(login_id) { + *guard = None; + } + }); + + let response = LoginChatGptResponse { login_id, auth_url }; + self.outgoing.send_response(request_id, response).await; + } + Err(err) => { + let error = JSONRPCErrorError { + code: INTERNAL_ERROR_CODE, + message: format!("failed to start login server: {err}"), + data: None, + }; + self.outgoing.send_error(request_id, error).await; + } + }, + Err(err) => { + self.outgoing.send_error(request_id, err).await; + } + } + } + + async fn login_chatgpt_v2(&mut self, request_id: RequestId) { + match self.login_chatgpt_common().await { + Ok(opts) => match run_login_server(opts) { + Ok(server) => { + let login_id = Uuid::new_v4(); + let shutdown_handle = server.cancel_handle(); + + // Replace active login if present. + { + let mut guard = self.active_login.lock().await; + if let Some(existing) = guard.take() { + existing.drop(); + } + *guard = Some(ActiveLogin { + shutdown_handle: shutdown_handle.clone(), + login_id, + }); } - // Clear the active login if it matches this attempt. It may have been replaced or cancelled. - let mut guard = active_login.lock().await; - if guard.as_ref().map(|l| l.login_id) == Some(login_id) { - *guard = None; - } - }); + // Spawn background task to monitor completion. + let outgoing_clone = self.outgoing.clone(); + let active_login = self.active_login.clone(); + let auth_manager = self.auth_manager.clone(); + let auth_url = server.auth_url.clone(); + tokio::spawn(async move { + let (success, error_msg) = match tokio::time::timeout( + LOGIN_CHATGPT_TIMEOUT, + server.block_until_done(), + ) + .await + { + Ok(Ok(())) => (true, None), + Ok(Err(err)) => (false, Some(format!("Login server error: {err}"))), + Err(_elapsed) => { + shutdown_handle.shutdown(); + (false, Some("Login timed out".to_string())) + } + }; - LoginChatGptReply::Response(response) - } - Err(err) => LoginChatGptReply::Error(JSONRPCErrorError { - code: INTERNAL_ERROR_CODE, - message: format!("failed to start login server: {err}"), - data: None, - }), - }; + let payload_v2 = AccountLoginCompletedNotification { + login_id: Some(login_id.to_string()), + success, + error: error_msg, + }; + outgoing_clone + .send_server_notification(ServerNotification::AccountLoginCompleted( + payload_v2, + )) + .await; - match reply { - LoginChatGptReply::Response(resp) => { - self.outgoing.send_response(request_id, resp).await + if success { + auth_manager.reload(); + + // Notify clients with the actual current auth mode. + let current_auth_method = auth_manager.auth().map(|a| a.mode); + let payload_v2 = AccountUpdatedNotification { + auth_mode: current_auth_method, + }; + outgoing_clone + .send_server_notification(ServerNotification::AccountUpdated( + payload_v2, + )) + .await; + } + + // Clear the active login if it matches this attempt. It may have been replaced or cancelled. + let mut guard = active_login.lock().await; + if guard.as_ref().map(|l| l.login_id) == Some(login_id) { + *guard = None; + } + }); + + let response = codex_app_server_protocol::LoginAccountResponse::Chatgpt { + login_id: login_id.to_string(), + auth_url, + }; + self.outgoing.send_response(request_id, response).await; + } + Err(err) => { + let error = JSONRPCErrorError { + code: INTERNAL_ERROR_CODE, + message: format!("failed to start login server: {err}"), + data: None, + }; + self.outgoing.send_error(request_id, error).await; + } + }, + Err(err) => { + self.outgoing.send_error(request_id, err).await; } - LoginChatGptReply::Error(err) => self.outgoing.send_error(request_id, err).await, } } @@ -581,7 +723,7 @@ impl CodexMessageProcessor { .await; let payload_v2 = AccountUpdatedNotification { - auth_method: current_auth_method, + auth_mode: current_auth_method, }; self.outgoing .send_server_notification(ServerNotification::AccountUpdated(payload_v2)) diff --git a/codex-rs/app-server/src/outgoing_message.rs b/codex-rs/app-server/src/outgoing_message.rs index d4dca94e..f0ee6cf9 100644 --- a/codex-rs/app-server/src/outgoing_message.rs +++ b/codex-rs/app-server/src/outgoing_message.rs @@ -141,6 +141,7 @@ pub(crate) struct OutgoingError { #[cfg(test)] mod tests { + use codex_app_server_protocol::AccountLoginCompletedNotification; use codex_app_server_protocol::AccountRateLimitsUpdatedNotification; use codex_app_server_protocol::AccountUpdatedNotification; use codex_app_server_protocol::AuthMode; @@ -178,6 +179,31 @@ mod tests { ); } + #[test] + fn verify_account_login_completed_notification_serialization() { + let notification = + ServerNotification::AccountLoginCompleted(AccountLoginCompletedNotification { + login_id: Some(Uuid::nil().to_string()), + success: true, + error: None, + }); + + let jsonrpc_notification = OutgoingMessage::AppServerNotification(notification); + assert_eq!( + json!({ + "method": "account/login/completed", + "params": { + "loginId": Uuid::nil().to_string(), + "success": true, + "error": null, + }, + }), + serde_json::to_value(jsonrpc_notification) + .expect("ensure the notification serializes correctly"), + "ensure the notification serializes correctly" + ); + } + #[test] fn verify_account_rate_limits_notification_serialization() { let notification = @@ -216,7 +242,7 @@ mod tests { #[test] fn verify_account_updated_notification_serialization() { let notification = ServerNotification::AccountUpdated(AccountUpdatedNotification { - auth_method: Some(AuthMode::ApiKey), + auth_mode: Some(AuthMode::ApiKey), }); let jsonrpc_notification = OutgoingMessage::AppServerNotification(notification); @@ -224,7 +250,7 @@ mod tests { json!({ "method": "account/updated", "params": { - "authMethod": "apikey" + "authMode": "apikey" }, }), serde_json::to_value(jsonrpc_notification) diff --git a/codex-rs/app-server/tests/common/mcp_process.rs b/codex-rs/app-server/tests/common/mcp_process.rs index 254f121a..298d9fd4 100644 --- a/codex-rs/app-server/tests/common/mcp_process.rs +++ b/codex-rs/app-server/tests/common/mcp_process.rs @@ -366,6 +366,26 @@ impl McpProcess { self.send_request("account/logout", None).await } + /// Send an `account/login/start` JSON-RPC request for API key login. + pub async fn send_login_account_api_key_request( + &mut self, + api_key: &str, + ) -> anyhow::Result { + let params = serde_json::json!({ + "type": "apiKey", + "apiKey": api_key, + }); + self.send_request("account/login/start", Some(params)).await + } + + /// Send an `account/login/start` JSON-RPC request for ChatGPT login. + pub async fn send_login_account_chatgpt_request(&mut self) -> anyhow::Result { + let params = serde_json::json!({ + "type": "chatgpt" + }); + self.send_request("account/login/start", Some(params)).await + } + /// Send a `fuzzyFileSearch` JSON-RPC request. pub async fn send_fuzzy_file_search_request( &mut self, diff --git a/codex-rs/app-server/tests/suite/v2/account.rs b/codex-rs/app-server/tests/suite/v2/account.rs index 5cc3719f..dfdc932b 100644 --- a/codex-rs/app-server/tests/suite/v2/account.rs +++ b/codex-rs/app-server/tests/suite/v2/account.rs @@ -2,15 +2,19 @@ use anyhow::Result; use anyhow::bail; use app_test_support::McpProcess; use app_test_support::to_response; +use codex_app_server_protocol::AuthMode; use codex_app_server_protocol::GetAuthStatusParams; use codex_app_server_protocol::GetAuthStatusResponse; +use codex_app_server_protocol::JSONRPCError; 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_core::auth::AuthCredentialsStoreMode; use codex_login::login_with_api_key; use pretty_assertions::assert_eq; +use serial_test::serial; use std::path::Path; use tempfile::TempDir; use tokio::time::timeout; @@ -18,14 +22,29 @@ use tokio::time::timeout; const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); // Helper to create a minimal config.toml for the app server -fn create_config_toml(codex_home: &Path) -> std::io::Result<()> { +fn create_config_toml( + codex_home: &Path, + forced_method: Option<&str>, + forced_workspace_id: Option<&str>, +) -> std::io::Result<()> { let config_toml = codex_home.join("config.toml"); - std::fs::write( - config_toml, + let forced_line = if let Some(method) = forced_method { + format!("forced_login_method = \"{method}\"\n") + } else { + String::new() + }; + let forced_workspace_line = if let Some(ws) = forced_workspace_id { + format!("forced_chatgpt_workspace_id = \"{ws}\"\n") + } else { + String::new() + }; + let contents = format!( r#" model = "mock-model" approval_policy = "never" sandbox_mode = "danger-full-access" +{forced_line} +{forced_workspace_line} model_provider = "mock_provider" @@ -35,14 +54,15 @@ base_url = "http://127.0.0.1:0/v1" wire_api = "chat" request_max_retries = 0 stream_max_retries = 0 -"#, - ) +"# + ); + std::fs::write(config_toml, contents) } #[tokio::test] async fn logout_account_removes_auth_and_notifies() -> Result<()> { let codex_home = TempDir::new()?; - create_config_toml(codex_home.path())?; + create_config_toml(codex_home.path(), None, None)?; login_with_api_key( codex_home.path(), @@ -72,7 +92,7 @@ async fn logout_account_removes_auth_and_notifies() -> Result<()> { bail!("unexpected notification: {parsed:?}"); }; assert!( - payload.auth_method.is_none(), + payload.auth_mode.is_none(), "auth_method should be None after logout" ); @@ -97,3 +117,156 @@ async fn logout_account_removes_auth_and_notifies() -> Result<()> { assert_eq!(status.auth_token, None); Ok(()) } + +#[tokio::test] +async fn login_account_api_key_succeeds_and_notifies() -> Result<()> { + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), None, None)?; + + let mut mcp = McpProcess::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let req_id = mcp + .send_login_account_api_key_request("sk-test-key") + .await?; + let resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(req_id)), + ) + .await??; + let login: LoginAccountResponse = to_response(resp)?; + assert_eq!(login, LoginAccountResponse::ApiKey {}); + + let note = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("account/login/completed"), + ) + .await??; + let parsed: ServerNotification = note.try_into()?; + let ServerNotification::AccountLoginCompleted(payload) = parsed else { + bail!("unexpected notification: {parsed:?}"); + }; + pretty_assertions::assert_eq!(payload.login_id, None); + pretty_assertions::assert_eq!(payload.success, true); + pretty_assertions::assert_eq!(payload.error, None); + + 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:?}"); + }; + pretty_assertions::assert_eq!(payload.auth_mode, Some(AuthMode::ApiKey)); + + assert!(codex_home.path().join("auth.json").exists()); + Ok(()) +} + +#[tokio::test] +async fn login_account_api_key_rejected_when_forced_chatgpt() -> Result<()> { + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), Some("chatgpt"), None)?; + + let mut mcp = McpProcess::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let request_id = mcp + .send_login_account_api_key_request("sk-test-key") + .await?; + let err: JSONRPCError = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_error_message(RequestId::Integer(request_id)), + ) + .await??; + + assert_eq!( + err.error.message, + "API key login is disabled. Use ChatGPT login instead." + ); + Ok(()) +} + +#[tokio::test] +async fn login_account_chatgpt_rejected_when_forced_api() -> Result<()> { + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), Some("api"), None)?; + + let mut mcp = McpProcess::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let request_id = mcp.send_login_account_chatgpt_request().await?; + let err: JSONRPCError = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_error_message(RequestId::Integer(request_id)), + ) + .await??; + + assert_eq!( + err.error.message, + "ChatGPT login is disabled. Use API key login instead." + ); + Ok(()) +} + +#[tokio::test] +// Serialize tests that launch the login server since it binds to a fixed port. +#[serial(login_port)] +async fn login_account_chatgpt_start() -> Result<()> { + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), None, None)?; + + let mut mcp = McpProcess::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + 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: _, + auth_url, + } = login + else { + bail!("unexpected login response: {login:?}"); + }; + assert!( + auth_url.contains("redirect_uri=http%3A%2F%2Flocalhost"), + "auth_url should contain a redirect_uri to localhost" + ); + Ok(()) +} + +#[tokio::test] +// Serialize tests that launch the login server since it binds to a fixed port. +#[serial(login_port)] +async fn login_account_chatgpt_includes_forced_workspace_query_param() -> Result<()> { + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), None, Some("ws-forced"))?; + + let mut mcp = McpProcess::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + 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 { auth_url, .. } = login else { + bail!("unexpected login response: {login:?}"); + }; + assert!( + auth_url.contains("allowed_workspace_id=ws-forced"), + "auth URL should include forced workspace" + ); + Ok(()) +}