diff --git a/codex-rs/app-server-protocol/src/export.rs b/codex-rs/app-server-protocol/src/export.rs index c8facffd..3affa7af 100644 --- a/codex-rs/app-server-protocol/src/export.rs +++ b/codex-rs/app-server-protocol/src/export.rs @@ -44,6 +44,7 @@ macro_rules! for_each_schema_type { $macro!(crate::ArchiveConversationParams); $macro!(crate::ArchiveConversationResponse); $macro!(crate::AuthMode); + $macro!(crate::AccountUpdatedNotification); $macro!(crate::AuthStatusChangeNotification); $macro!(crate::CancelLoginChatGptParams); $macro!(crate::CancelLoginChatGptResponse); diff --git a/codex-rs/app-server-protocol/src/protocol/common.rs b/codex-rs/app-server-protocol/src/protocol/common.rs index dcab5eae..efc60830 100644 --- a/codex-rs/app-server-protocol/src/protocol/common.rs +++ b/codex-rs/app-server-protocol/src/protocol/common.rs @@ -372,6 +372,11 @@ pub struct FuzzyFileSearchResponse { #[strum(serialize_all = "camelCase")] pub enum ServerNotification { /// NEW NOTIFICATIONS + #[serde(rename = "account/updated")] + #[ts(rename = "account/updated")] + #[strum(serialize = "account/updated")] + AccountUpdated(v2::AccountUpdatedNotification), + #[serde(rename = "account/rateLimits/updated")] #[ts(rename = "account/rateLimits/updated")] #[strum(serialize = "account/rateLimits/updated")] @@ -391,6 +396,7 @@ pub enum ServerNotification { impl ServerNotification { pub fn to_params(self) -> Result { match self { + ServerNotification::AccountUpdated(params) => serde_json::to_value(params), ServerNotification::AccountRateLimitsUpdated(params) => serde_json::to_value(params), ServerNotification::AuthStatusChange(params) => serde_json::to_value(params), ServerNotification::LoginChatGptComplete(params) => serde_json::to_value(params), diff --git a/codex-rs/app-server-protocol/src/protocol/v1.rs b/codex-rs/app-server-protocol/src/protocol/v1.rs index 8fbcf6b7..252fb041 100644 --- a/codex-rs/app-server-protocol/src/protocol/v1.rs +++ b/codex-rs/app-server-protocol/src/protocol/v1.rs @@ -400,6 +400,7 @@ pub struct SessionConfiguredNotification { #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] #[serde(rename_all = "camelCase")] +/// Deprecated notification. Use AccountUpdatedNotification instead. pub struct AuthStatusChangeNotification { pub auth_method: Option, } diff --git a/codex-rs/app-server-protocol/src/protocol/v2.rs b/codex-rs/app-server-protocol/src/protocol/v2.rs index ccd89e6a..8c396284 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2.rs @@ -1,3 +1,4 @@ +use crate::protocol::common::AuthMode; use codex_protocol::ConversationId; use codex_protocol::account::PlanType; use codex_protocol::config_types::ReasoningEffort; @@ -120,3 +121,9 @@ pub struct UploadFeedbackParams { pub struct UploadFeedbackResponse { pub thread_id: String, } + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +pub struct AccountUpdatedNotification { + pub auth_method: Option, +} diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index 7cd60eae..8ea4cef8 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::AccountUpdatedNotification; use codex_app_server_protocol::AddConversationListenerParams; use codex_app_server_protocol::AddConversationSubscriptionResponse; use codex_app_server_protocol::ApplyPatchApprovalParams; @@ -200,8 +201,7 @@ impl CodexMessageProcessor { request_id, params: _, } => { - self.send_unimplemented_error(request_id, "account/logout") - .await; + self.logout_v2(request_id).await; } ClientRequest::GetAccount { request_id, @@ -250,7 +250,7 @@ impl CodexMessageProcessor { request_id, params: _, } => { - self.logout_chatgpt(request_id).await; + self.logout_v1(request_id).await; } ClientRequest::GetAuthStatus { request_id, params } => { self.get_auth_status(request_id, params).await; @@ -494,9 +494,9 @@ impl CodexMessageProcessor { } } - async fn logout_chatgpt(&mut self, request_id: RequestId) { + async fn logout_common(&mut self) -> std::result::Result, JSONRPCErrorError> { + // Cancel any active login attempt. { - // Cancel any active login attempt. let mut guard = self.active_login.lock().await; if let Some(active) = guard.take() { active.drop(); @@ -504,31 +504,61 @@ impl CodexMessageProcessor { } if let Err(err) = self.auth_manager.logout() { - let error = JSONRPCErrorError { + return Err(JSONRPCErrorError { code: INTERNAL_ERROR_CODE, message: format!("logout failed: {err}"), data: None, - }; - self.outgoing.send_error(request_id, error).await; - return; + }); } - self.outgoing - .send_response( - request_id, - codex_app_server_protocol::LogoutChatGptResponse {}, - ) - .await; + // Reflect the current auth method after logout (likely None). + Ok(self.auth_manager.auth().map(|auth| auth.mode)) + } - // Send auth status change notification reflecting the current auth mode - // after logout. - let current_auth_method = self.auth_manager.auth().map(|auth| auth.mode); - let payload = AuthStatusChangeNotification { - auth_method: current_auth_method, - }; - self.outgoing - .send_server_notification(ServerNotification::AuthStatusChange(payload)) - .await; + async fn logout_v1(&mut self, request_id: RequestId) { + match self.logout_common().await { + Ok(current_auth_method) => { + self.outgoing + .send_response( + request_id, + codex_app_server_protocol::LogoutChatGptResponse {}, + ) + .await; + + let payload = AuthStatusChangeNotification { + auth_method: current_auth_method, + }; + self.outgoing + .send_server_notification(ServerNotification::AuthStatusChange(payload)) + .await; + } + Err(error) => { + self.outgoing.send_error(request_id, error).await; + } + } + } + + async fn logout_v2(&mut self, request_id: RequestId) { + match self.logout_common().await { + Ok(current_auth_method) => { + self.outgoing + .send_response( + request_id, + codex_app_server_protocol::LogoutAccountResponse {}, + ) + .await; + + let payload_v2 = AccountUpdatedNotification { + auth_method: current_auth_method, + }; + self.outgoing + .send_server_notification(ServerNotification::AccountUpdated(payload_v2)) + .await; + } + Err(error) => { + self.outgoing.send_error(request_id, error).await; + } + } } async fn get_auth_status( diff --git a/codex-rs/app-server/src/outgoing_message.rs b/codex-rs/app-server/src/outgoing_message.rs index 3f4497fb..3d56c9bd 100644 --- a/codex-rs/app-server/src/outgoing_message.rs +++ b/codex-rs/app-server/src/outgoing_message.rs @@ -141,6 +141,8 @@ pub(crate) struct OutgoingError { #[cfg(test)] mod tests { + use codex_app_server_protocol::AccountUpdatedNotification; + use codex_app_server_protocol::AuthMode; use codex_app_server_protocol::LoginChatGptCompleteNotification; use codex_protocol::protocol::RateLimitSnapshot; use codex_protocol::protocol::RateLimitWindow; @@ -204,4 +206,24 @@ mod tests { "ensure the notification serializes correctly" ); } + + #[test] + fn verify_account_updated_notification_serialization() { + let notification = ServerNotification::AccountUpdated(AccountUpdatedNotification { + auth_method: Some(AuthMode::ApiKey), + }); + + let jsonrpc_notification = OutgoingMessage::AppServerNotification(notification); + assert_eq!( + json!({ + "method": "account/updated", + "params": { + "authMethod": "apikey" + }, + }), + serde_json::to_value(jsonrpc_notification) + .expect("ensure the notification serializes correctly"), + "ensure the notification serializes correctly" + ); + } } diff --git a/codex-rs/app-server/tests/common/mcp_process.rs b/codex-rs/app-server/tests/common/mcp_process.rs index 90c7645d..3c1fbf77 100644 --- a/codex-rs/app-server/tests/common/mcp_process.rs +++ b/codex-rs/app-server/tests/common/mcp_process.rs @@ -321,6 +321,11 @@ impl McpProcess { self.send_request("logoutChatGpt", None).await } + /// Send an `account/logout` JSON-RPC request. + pub async fn send_logout_account_request(&mut self) -> anyhow::Result { + self.send_request("account/logout", None).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/mod.rs b/codex-rs/app-server/tests/suite/mod.rs index c8763cd0..ce8b69a0 100644 --- a/codex-rs/app-server/tests/suite/mod.rs +++ b/codex-rs/app-server/tests/suite/mod.rs @@ -13,3 +13,4 @@ mod send_message; mod set_default_model; mod user_agent; mod user_info; +mod v2; diff --git a/codex-rs/app-server/tests/suite/v2/account.rs b/codex-rs/app-server/tests/suite/v2/account.rs new file mode 100644 index 00000000..5cc3719f --- /dev/null +++ b/codex-rs/app-server/tests/suite/v2/account.rs @@ -0,0 +1,99 @@ +use anyhow::Result; +use anyhow::bail; +use app_test_support::McpProcess; +use app_test_support::to_response; +use codex_app_server_protocol::GetAuthStatusParams; +use codex_app_server_protocol::GetAuthStatusResponse; +use codex_app_server_protocol::JSONRPCResponse; +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 std::path::Path; +use tempfile::TempDir; +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<()> { + let config_toml = codex_home.join("config.toml"); + std::fs::write( + config_toml, + r#" +model = "mock-model" +approval_policy = "never" +sandbox_mode = "danger-full-access" + +model_provider = "mock_provider" + +[model_providers.mock_provider] +name = "Mock provider for test" +base_url = "http://127.0.0.1:0/v1" +wire_api = "chat" +request_max_retries = 0 +stream_max_retries = 0 +"#, + ) +} + +#[tokio::test] +async fn logout_account_removes_auth_and_notifies() -> Result<()> { + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path())?; + + login_with_api_key( + codex_home.path(), + "sk-test-key", + AuthCredentialsStoreMode::File, + )?; + assert!(codex_home.path().join("auth.json").exists()); + + let mut mcp = McpProcess::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let id = mcp.send_logout_account_request().await?; + let resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(id)), + ) + .await??; + let _ok: LogoutAccountResponse = to_response(resp)?; + + 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!( + payload.auth_method.is_none(), + "auth_method should be None after logout" + ); + + assert!( + !codex_home.path().join("auth.json").exists(), + "auth.json should be deleted" + ); + + let status_id = mcp + .send_get_auth_status_request(GetAuthStatusParams { + include_token: Some(true), + refresh_token: Some(false), + }) + .await?; + let status_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(status_id)), + ) + .await??; + let status: GetAuthStatusResponse = to_response(status_resp)?; + assert_eq!(status.auth_method, None); + assert_eq!(status.auth_token, None); + Ok(()) +} diff --git a/codex-rs/app-server/tests/suite/v2/mod.rs b/codex-rs/app-server/tests/suite/v2/mod.rs new file mode 100644 index 00000000..bd7a1467 --- /dev/null +++ b/codex-rs/app-server/tests/suite/v2/mod.rs @@ -0,0 +1,2 @@ +// v2 test suite modules +mod account; diff --git a/codex-rs/core/tests/common/lib.rs b/codex-rs/core/tests/common/lib.rs index 0064d0e5..3f75ed18 100644 --- a/codex-rs/core/tests/common/lib.rs +++ b/codex-rs/core/tests/common/lib.rs @@ -244,7 +244,7 @@ pub mod fs_wait { if path.exists() { Ok(path) } else { - Err(anyhow!("timed out waiting for {:?}", path)) + Err(anyhow!("timed out waiting for {path:?}")) } } @@ -284,7 +284,7 @@ pub mod fs_wait { if let Some(found) = scan_for_match(&root, predicate) { Ok(found) } else { - Err(anyhow!("timed out waiting for matching file in {:?}", root)) + Err(anyhow!("timed out waiting for matching file in {root:?}")) } }