diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 8ec0a86608..e8b1096596 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2144,6 +2144,7 @@ dependencies = [ "futures", "hmac 0.12.1", "http 1.4.0", + "keyring", "opentelemetry", "opentelemetry_sdk", "pretty_assertions", @@ -4395,6 +4396,7 @@ dependencies = [ "thiserror 2.0.18", "tiny_http", "tokio", + "tokio-util", "tracing", "tracing-subscriber", "tracing-test", diff --git a/codex-rs/app-server-protocol/schema/json/ServerNotification.json b/codex-rs/app-server-protocol/schema/json/ServerNotification.json index 1e20ed648d..4ae27cf806 100644 --- a/codex-rs/app-server-protocol/schema/json/ServerNotification.json +++ b/codex-rs/app-server-protocol/schema/json/ServerNotification.json @@ -2791,6 +2791,13 @@ "null" ] }, + "loginId": { + "description": "Identifies the explicit login attempt. Older servers omit this field.", + "type": [ + "string", + "null" + ] + }, "name": { "type": "string" }, diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json index 5a70f96f1a..a4dc6548af 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json @@ -14838,6 +14838,13 @@ "null" ] }, + "loginId": { + "description": "Identifies the explicit login attempt. Older servers omit this field.", + "type": [ + "string", + "null" + ] + }, "name": { "type": "string" }, @@ -14909,6 +14916,14 @@ "properties": { "authorizationUrl": { "type": "string" + }, + "loginId": { + "default": null, + "description": "Identifies this login attempt across the response and completion notification. Older servers omit this field; current servers always return it.", + "type": [ + "string", + "null" + ] } }, "required": [ diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json index 76c7ae81fc..75ecacf9c2 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json @@ -10870,6 +10870,13 @@ "null" ] }, + "loginId": { + "description": "Identifies the explicit login attempt. Older servers omit this field.", + "type": [ + "string", + "null" + ] + }, "name": { "type": "string" }, @@ -10941,6 +10948,14 @@ "properties": { "authorizationUrl": { "type": "string" + }, + "loginId": { + "default": null, + "description": "Identifies this login attempt across the response and completion notification. Older servers omit this field; current servers always return it.", + "type": [ + "string", + "null" + ] } }, "required": [ diff --git a/codex-rs/app-server-protocol/schema/json/v2/McpServerOauthLoginCompletedNotification.json b/codex-rs/app-server-protocol/schema/json/v2/McpServerOauthLoginCompletedNotification.json index 6204ff672b..6cc9fd96f2 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/McpServerOauthLoginCompletedNotification.json +++ b/codex-rs/app-server-protocol/schema/json/v2/McpServerOauthLoginCompletedNotification.json @@ -7,6 +7,13 @@ "null" ] }, + "loginId": { + "description": "Identifies the explicit login attempt. Older servers omit this field.", + "type": [ + "string", + "null" + ] + }, "name": { "type": "string" }, diff --git a/codex-rs/app-server-protocol/schema/json/v2/McpServerOauthLoginResponse.json b/codex-rs/app-server-protocol/schema/json/v2/McpServerOauthLoginResponse.json index efeb612dd3..7672a53739 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/McpServerOauthLoginResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/McpServerOauthLoginResponse.json @@ -3,6 +3,14 @@ "properties": { "authorizationUrl": { "type": "string" + }, + "loginId": { + "default": null, + "description": "Identifies this login attempt across the response and completion notification. Older servers omit this field; current servers always return it.", + "type": [ + "string", + "null" + ] } }, "required": [ diff --git a/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-experimental.json.zst b/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-experimental.json.zst index 0a35c5522f..5838b537c5 100644 Binary files a/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-experimental.json.zst and b/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-experimental.json.zst differ diff --git a/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst b/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst index f18fae87d1..e8bfb39944 100644 Binary files a/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst and b/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst differ diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/McpServerOauthLoginCompletedNotification.ts b/codex-rs/app-server-protocol/schema/typescript/v2/McpServerOauthLoginCompletedNotification.ts index cfa6603094..1a45a947dd 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/McpServerOauthLoginCompletedNotification.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/McpServerOauthLoginCompletedNotification.ts @@ -2,4 +2,8 @@ // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. -export type McpServerOauthLoginCompletedNotification = { name: string, threadId: string | null, success: boolean, error?: string, }; +export type McpServerOauthLoginCompletedNotification = { name: string, threadId: string | null, +/** + * Identifies the explicit login attempt. Older servers omit this field. + */ +loginId?: string | null, success: boolean, error?: string, }; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/McpServerOauthLoginResponse.ts b/codex-rs/app-server-protocol/schema/typescript/v2/McpServerOauthLoginResponse.ts index 5933574765..a9f67df430 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/McpServerOauthLoginResponse.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/McpServerOauthLoginResponse.ts @@ -2,4 +2,9 @@ // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. -export type McpServerOauthLoginResponse = { authorizationUrl: string, }; +export type McpServerOauthLoginResponse = { authorizationUrl: string, +/** + * Identifies this login attempt across the response and completion notification. + * Older servers omit this field; current servers always return it. + */ +loginId?: string, }; diff --git a/codex-rs/app-server-protocol/src/export.rs b/codex-rs/app-server-protocol/src/export.rs index db5ff718b6..0d8ce4aaa4 100644 --- a/codex-rs/app-server-protocol/src/export.rs +++ b/codex-rs/app-server-protocol/src/export.rs @@ -2193,6 +2193,12 @@ mod tests { .ok_or_else(|| anyhow::anyhow!("missing account usage response fixture"))?, )?; assert!(account_usage_response_ts.contains("threadUsage?: ThreadUsage | null")); + let mcp_login_completion_ts = std::str::from_utf8( + fixture_tree + .get(Path::new("v2/McpServerOauthLoginCompletedNotification.ts")) + .ok_or_else(|| anyhow::anyhow!("missing MCP login completion fixture"))?, + )?; + assert!(mcp_login_completion_ts.contains("loginId?: string | null")); let server_request_ts = std::str::from_utf8( fixture_tree .get(Path::new("ServerRequest.ts")) @@ -2386,14 +2392,16 @@ mod tests { // If the last non-whitespace before ':' is '?', then this is an // optional field with a nullable type (i.e., "?: T | null"). - // These are only allowed in *Params types, except the additive stable usage - // response field, which older servers omit and newer servers return as null. - let legacy_account_usage_response = path + // These are only allowed in *Params types, except additive stable fields + // that older servers omit and newer servers may return as null. + let legacy_optional_nullable_field = (path == Path::new("v2/GetAccountTokenUsageResponse.ts") - && field_prefix.trim() == "threadUsage?"; + && field_prefix.trim() == "threadUsage?") + || (path == Path::new("v2/McpServerOauthLoginCompletedNotification.ts") + && field_prefix.trim() == "loginId?"); if field_prefix.chars().rev().find(|c| !c.is_whitespace()) == Some('?') && !allow_optional_nullable - && !legacy_account_usage_response + && !legacy_optional_nullable_field { let line_number = contents[..abs_idx].chars().filter(|c| *c == '\n').count() + 1; diff --git a/codex-rs/app-server-protocol/src/protocol/v2/mcp.rs b/codex-rs/app-server-protocol/src/protocol/v2/mcp.rs index be16f57fd4..e1ed3b624b 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/mcp.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/mcp.rs @@ -321,6 +321,11 @@ pub enum McpServerOauthClientRegistration { #[ts(export_to = "v2/")] pub struct McpServerOauthLoginResponse { pub authorization_url: String, + /// Identifies this login attempt across the response and completion notification. + /// Older servers omit this field; current servers always return it. + #[serde(default)] + #[ts(optional)] + pub login_id: Option, } #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] #[serde(rename_all = "camelCase")] @@ -338,6 +343,9 @@ pub struct McpToolCallProgressNotification { pub struct McpServerOauthLoginCompletedNotification { pub name: String, pub thread_id: Option, + /// Identifies the explicit login attempt. Older servers omit this field. + #[ts(optional, as = "Option>")] + pub login_id: Option, pub success: bool, #[serde(default, skip_serializing_if = "Option::is_none")] #[ts(optional)] diff --git a/codex-rs/app-server-protocol/src/protocol/v2/tests.rs b/codex-rs/app-server-protocol/src/protocol/v2/tests.rs index 5bc54053e4..ff34043464 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/tests.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/tests.rs @@ -99,6 +99,22 @@ fn managed_hooks_requirements_default_interrupt_to_empty() { assert_eq!(parsed.interrupt, Vec::new()); } +#[test] +fn mcp_oauth_login_response_accepts_older_servers_without_login_id() { + let response = serde_json::from_value::(json!({ + "authorizationUrl": "https://example.com/authorize", + })) + .expect("older login response should deserialize"); + + assert_eq!( + response, + McpServerOauthLoginResponse { + authorization_url: "https://example.com/authorize".to_string(), + login_id: None, + } + ); +} + #[test] fn external_agent_config_detect_response_defaults_connectors_for_older_servers() { let response = serde_json::from_value::(json!({ diff --git a/codex-rs/app-server/Cargo.toml b/codex-rs/app-server/Cargo.toml index 8834c54b7e..12418814aa 100644 --- a/codex-rs/app-server/Cargo.toml +++ b/codex-rs/app-server/Cargo.toml @@ -140,6 +140,7 @@ core_test_support = { workspace = true } flate2 = { workspace = true } hmac = { workspace = true } http = { workspace = true } +keyring = { workspace = true } opentelemetry = { workspace = true } opentelemetry_sdk = { workspace = true } pretty_assertions = { workspace = true } diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index 5331d42b0a..268eee9022 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -352,6 +352,15 @@ layers with enterprise MCP disabled. Other planned thread refreshes are processe before the rejection is reported, so an error does not imply that no changes were applied. Correct the policy before retrying the reload. +# Enterprise sign-in + +Call `mcpServer/oauth/login` with a directly configured server's `name` and its +connected `threadId`. Open the returned `authorizationUrl` and match +`mcpServer/oauthLogin/completed` by `loginId`. Starting another enterprise sign-in +cancels the previous attempt and waits for its callback listener to close. Use +`account/login/cancel` to cancel explicitly. Start a fresh session after success to +use the saved grant. + # Selected workspace routing The experimental `account/read.workspaceRouting` response field returns the selected ChatGPT workspace's `chatgptAccountId`, resolved HTTPS `backendOrigin`, and backend-provided `accountRoutingOverride`. The routing value is `us`, `us_cr`, or the explicit `NO_CONSTRAINT` value. API-only and signed-out accounts return `null` and do not need `accounts/check`. diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 9e9a90437b..13a28a1053 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -473,6 +473,7 @@ impl MessageProcessor { thread_state_manager.clone(), outgoing.clone(), config_manager.clone(), + account_processor.clone(), ); let plugin_processor = PluginRequestProcessor::new( auth_manager.clone(), @@ -1680,7 +1681,8 @@ impl MessageProcessor { self.turn_processor.review_start(&request_id, params).await } ClientRequest::McpServerOauthLogin { params, .. } => { - self.mcp_processor.mcp_server_oauth_login(params).await + // Keep MCP discovery's large future from inflating every request. + Box::pin(self.mcp_processor.mcp_server_oauth_login(params)).await } ClientRequest::McpServerRefresh { params, .. } => { self.mcp_processor.mcp_server_refresh(params).await diff --git a/codex-rs/app-server/src/request_processors/account_processor.rs b/codex-rs/app-server/src/request_processors/account_processor.rs index 95ed6d571d..e1a0cd33d2 100644 --- a/codex-rs/app-server/src/request_processors/account_processor.rs +++ b/codex-rs/app-server/src/request_processors/account_processor.rs @@ -11,13 +11,19 @@ use codex_app_server_protocol::DesktopOnboardingEntrypoint; use codex_app_server_protocol::GetAccountRateLimitsParams; use codex_login::LoginOnboardingEntrypoint; use codex_login::login_with_bedrock_access_keys; +use codex_mcp::ema_auth_scope; use codex_model_provider::is_supported_amazon_bedrock_region; +use codex_rmcp_client::EnterpriseOAuthCredentialGuard; mod bedrock_setup; +mod enterprise_login; mod gateway_oauth; mod rate_limit_resets; mod workspace_routing; +pub(super) use enterprise_login::EnterpriseLoginCompletion; +pub(super) use enterprise_login::EnterpriseLoginTarget; + // Duration before a browser ChatGPT login attempt is abandoned. const LOGIN_CHATGPT_TIMEOUT: Duration = Duration::from_secs(10 * 60); const ACCOUNT_TOKEN_USAGE_FETCH_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 10); @@ -100,6 +106,7 @@ pub(crate) struct AccountRequestProcessor { gateway_login: Arc>>, gateway_client: Arc>>>, _gateway_notifications: Arc>, + pub(super) enterprise_login: Arc, } impl AccountRequestProcessor { @@ -115,6 +122,11 @@ impl AccountRequestProcessor { config_manager.clone(), Arc::clone(&outgoing), ); + let enterprise_login = Arc::new(enterprise_login::EnterpriseLoginState::new( + Arc::clone(&auth_manager), + Arc::clone(&thread_manager), + config_manager.clone(), + )); let processor = Arc::new(Self { _gateway_notifications: Arc::new(gateway_notifications), auth_manager, @@ -128,6 +140,7 @@ impl AccountRequestProcessor { workspace_routing: Arc::new(Mutex::new(None)), workspace_routing_fetches: Arc::new(Mutex::new(HashMap::new())), workspace_routing_shutdown: CancellationToken::new(), + enterprise_login, }); let resolver: Arc = processor.clone(); processor @@ -210,6 +223,7 @@ impl AccountRequestProcessor { pub(crate) async fn cancel_active_login(&self) { self.cancel_gateway_login(); + self.enterprise_login.cancel(/*login_id*/ None).await; let mut guard = self.active_login.lock().await; if let Some(active_login) = guard.take() { drop(active_login); @@ -318,6 +332,7 @@ impl AccountRequestProcessor { if self.auth_manager.is_workload_identity_selected() { return Err(self.configured_auth_owned_by_host_error()); } + self.enterprise_login.cancel(/*login_id*/ None).await; match params { LoginAccountParams::ApiKey { api_key } => { self.login_api_key_v2(request_id, LoginApiKeyParams { api_key }) @@ -801,6 +816,11 @@ impl AccountRequestProcessor { let login_id = params.login_id; let uuid = Uuid::parse_str(&login_id) .map_err(|_| invalid_request(format!("invalid login id: {login_id}")))?; + if self.enterprise_login.cancel(Some(&login_id)).await { + return Ok(CancelLoginAccountResponse { + status: CancelLoginAccountStatus::Canceled, + }); + } let status = match self.cancel_login_chatgpt_common(uuid).await { Ok(()) => CancelLoginAccountStatus::Canceled, Err(CancelLoginError::NotFound) => CancelLoginAccountStatus::NotFound, @@ -884,6 +904,7 @@ impl AccountRequestProcessor { } async fn send_login_success_notifications(&self, login_id: Option) { + self.thread_manager.invalidate_mcp_runtimes().await; self.send_account_login_notifications(AccountLoginCompletedNotification { login_id: login_id.map(|id| id.to_string()), success: true, @@ -945,6 +966,7 @@ impl AccountRequestProcessor { self.auth_manager.reload().await; let auth_changes = self.auth_manager.auth_change_state_receiver(); let owner_generation = auth_changes.borrow().owner_generation; + self.thread_manager.invalidate_mcp_runtimes().await; self.config_manager.replace_cloud_config_bundle_loader( self.auth_manager.clone(), self.config.chatgpt_base_url.clone(), @@ -965,16 +987,65 @@ impl AccountRequestProcessor { if self.auth_manager.is_workload_identity_selected() { return Err(self.configured_auth_owned_by_host_error()); } - let config = self.load_latest_config().await; + // Another process may have changed the persisted workspace. Reload both + // account authority and its policy before selecting a grant to remove. + self.auth_manager.reload().await; + self.config_manager.replace_cloud_config_bundle_loader( + Arc::clone(&self.auth_manager), + self.config.chatgpt_base_url.clone(), + self.config.http_client_factory(), + ); + let config = self + .config_manager + .load_latest_config(/*fallback_cwd*/ None) + .await; + let scope = ema_auth_scope(self.auth_manager.auth_cached().as_ref()); + let enterprise_policy_failed = scope.is_some() && config.is_err(); + // Startup policy may belong to another workspace. Never use that fallback + // to select a credential; a policy failure must not block primary logout. + let enterprise_grant = config.as_ref().ok().and_then(|config| { + scope + .as_ref() + .zip(config.mcp_enterprise_managed_auth.as_ref()) + .map(|(scope, profile)| { + ( + profile.idp.credential_name(scope), + profile.idp.issuer.clone(), + config.auth_keyring_backend_kind(), + ) + }) + }); + let config = config.unwrap_or_else(|_| self.config.as_ref().clone()); self.cancel_active_login().await; + // Retain the credential lock through primary logout. Otherwise a second + // process can commit after deletion but before the account is removed. + let enterprise_guard = + if let Some((credential_name, issuer, keyring_backend)) = enterprise_grant { + EnterpriseOAuthCredentialGuard::acquire(&credential_name, &issuer, keyring_backend) + .await + .map(Some) + } else { + Ok(None) + }; + let cleanup_failed = match &enterprise_guard { + Ok(Some(guard)) => guard.delete_tokens().is_err(), + Ok(None) => false, + Err(_) => true, + }; + if enterprise_policy_failed || cleanup_failed { + tracing::warn!("Failed to remove enterprise authorization; continuing account logout"); + } + match self.auth_manager.logout_with_revoke().await { Ok(_) => {} Err(err) => { return Err(internal_error(format!("logout failed: {err}"))); } } + drop(enterprise_guard); + self.thread_manager.invalidate_mcp_runtimes().await; self.config_manager.clear_cloud_config_bundle_loader(); diff --git a/codex-rs/app-server/src/request_processors/account_processor/enterprise_login.rs b/codex-rs/app-server/src/request_processors/account_processor/enterprise_login.rs new file mode 100644 index 0000000000..cc66120377 --- /dev/null +++ b/codex-rs/app-server/src/request_processors/account_processor/enterprise_login.rs @@ -0,0 +1,274 @@ +//! Owns one explicit enterprise sign-in attempt for a configured MCP server. +//! Start, replacement, cancellation, and local commit admission all pass through this state. +//! The credential store separately fences logins invalidated by another process. + +use super::*; +use codex_config::McpEmaRegistration; +use codex_config::types::AuthKeyringBackendKind; +use codex_mcp::ema_auth_scope; +use codex_rmcp_client::EnterpriseOAuthLoginRequest; +use codex_rmcp_client::perform_enterprise_oauth_login_return_url; +use futures::FutureExt; +use tokio::sync::watch; + +pub(crate) struct EnterpriseLoginTarget { + pub(crate) thread_id: String, + pub(crate) server_name: String, + pub(crate) registration: McpEmaRegistration, +} + +#[derive(Clone, Copy, Eq, PartialEq)] +pub(crate) enum EnterpriseLoginCompletion { + Pending, + Succeeded, + Failed, +} + +impl EnterpriseLoginCompletion { + pub(crate) fn is_complete(self) -> bool { + matches!(self, Self::Succeeded | Self::Failed) + } +} + +pub(crate) struct EnterpriseLogin { + pub(crate) login_id: String, + pub(crate) authorization_url: String, + pub(crate) completion: watch::Receiver, +} + +pub(crate) struct EnterpriseLoginState { + auth_manager: Arc, + thread_manager: Arc, + config_manager: ConfigManager, + // Cancellation and commit admission use this same lock. A cancelled attempt cannot commit. + active: Mutex>, + // Start and cancel hold this through callback shutdown; the worker never acquires it. + start_gate: Mutex<()>, + // Account changes cancel queued and running setup before waiting for start_gate. + setup_cancel: Mutex, +} + +struct EnterpriseLoginAttempt { + login_id: String, + cancel: CancellationToken, + completion: watch::Receiver, +} + +impl Drop for EnterpriseLoginAttempt { + fn drop(&mut self) { + self.cancel.cancel(); + } +} + +impl EnterpriseLoginState { + pub(super) fn new( + auth_manager: Arc, + thread_manager: Arc, + config_manager: ConfigManager, + ) -> Self { + Self { + auth_manager, + thread_manager, + config_manager, + active: Mutex::new(None), + start_gate: Mutex::new(()), + setup_cancel: Mutex::new(CancellationToken::new()), + } + } + + #[expect( + clippy::await_holding_invalid_type, + reason = "serialize start/cancel through callback shutdown; the worker uses a separate state lock" + )] + pub(super) async fn cancel(&self, login_id: Option<&str>) -> bool { + if login_id.is_none() { + let mut setup_cancel = self.setup_cancel.lock().await; + setup_cancel.cancel(); + *setup_cancel = CancellationToken::new(); + } + let (login_id, mut completion) = { + let active = self.active.lock().await; + if login_id.is_some_and(|id| active.as_ref().is_none_or(|a| a.login_id != id)) { + return false; + } + let Some(attempt) = active.as_ref() else { + return false; + }; + attempt.cancel.cancel(); + (attempt.login_id.clone(), attempt.completion.clone()) + }; + // A partial callback request can stall the blocking HTTP worker. Cancellation + // already fences commit; primary logout must not wait indefinitely for I/O. + if tokio::time::timeout(std::time::Duration::from_secs(10), async { + let _start = self.start_gate.lock().await; + completion.wait_for(|state| state.is_complete()).await + }) + .await + .is_ok() + { + self.clear(&login_id).await; + } + // Retain a timed-out attempt so a replacement still waits for port cleanup. + true + } + + async fn clear(&self, login_id: &str) { + let mut active = self.active.lock().await; + if active + .as_ref() + .is_some_and(|attempt| attempt.login_id == login_id) + { + *active = None; + } + } + + async fn configuration_is_current( + &self, + target: &EnterpriseLoginTarget, + keyring_backend: AuthKeyringBackendKind, + ) -> bool { + let Ok(thread_id) = ThreadId::from_string(&target.thread_id) else { + return false; + }; + let Ok(thread) = self.thread_manager.get_thread(thread_id).await else { + return false; + }; + let current = thread.config().await; + let Ok(config) = self + .config_manager + .load_latest_config_with_session_layers(¤t.config_layer_stack, ¤t.cwd) + .await + else { + return false; + }; + let config = thread.runtime_mcp_config_and_context(&config).await.0; + let auth = self.auth_manager.auth_cached(); + let servers = codex_mcp::effective_mcp_servers(&config, auth.as_ref()); + config.xaa_enabled + && config.auth_keyring_backend_kind == keyring_backend + && servers.get(&target.server_name).is_some_and(|server| { + server.config().enabled + && server.config().ema_registration() == Some(&target.registration) + }) + } + + #[expect( + clippy::await_holding_invalid_type, + reason = "serialize start/cancel through callback shutdown; the worker uses a separate state lock" + )] + pub(crate) async fn start( + self: &Arc, + request: EnterpriseOAuthLoginRequest<'_>, + scope: codex_config::McpEmaAuthScope, + target: EnterpriseLoginTarget, + ) -> Result { + // MCP requests are not in the account request queue. Serialize start and + // cancellation here; the separate commit gate stays responsive to the worker. + let setup_cancel = self.setup_cancel.lock().await.child_token(); + let _start = tokio::select! { + biased; + _ = setup_cancel.cancelled() => return Err(internal_error("enterprise sign-in was cancelled")), + guard = self.start_gate.lock() => guard, + }; + let previous = { + let active = self.active.lock().await; + // Keep the cancelled attempt installed until its callback worker closes. + // A cancelled replacement request leaves the next start responsible for waiting. + active.as_ref().map(|attempt| { + attempt.cancel.cancel(); + (attempt.login_id.clone(), attempt.completion.clone()) + }) + }; + if let Some((login_id, mut completion)) = previous { + // An abandoned setup request drops its flow and closes this channel. + let _ = tokio::select! { + biased; + _ = setup_cancel.cancelled() => return Err(internal_error("enterprise sign-in was cancelled")), + result = tokio::time::timeout( + std::time::Duration::from_secs(10), + completion.wait_for(|state| state.is_complete()), + ) => result.map_err(|_| internal_error("previous enterprise sign-in did not close"))?, + }; + self.clear(&login_id).await; + } + let keyring_backend = request.keyring_backend_kind; + // Keep commit cancellation under `active`; cancelling the setup epoch must + // not race an already-admitted synchronous credential write. + let cancel = CancellationToken::new(); + let login_id = Uuid::new_v4().to_string(); + let (completed, completion) = watch::channel(EnterpriseLoginCompletion::Pending); + *self.active.lock().await = Some(EnterpriseLoginAttempt { + login_id: login_id.clone(), + cancel: cancel.clone(), + completion: completion.clone(), + }); + let handle = tokio::select! { + biased; + _ = setup_cancel.cancelled() => Err(internal_error("enterprise sign-in was cancelled")), + _ = cancel.cancelled() => Err(internal_error("enterprise sign-in was cancelled")), + result = perform_enterprise_oauth_login_return_url(request) => + result.map_err(|_| internal_error("failed to start enterprise sign-in")), + }; + let handle = match handle { + Ok(handle) => handle, + Err(error) => { + self.clear(&login_id).await; + completed.send_replace(EnterpriseLoginCompletion::Failed); + return Err(error); + } + }; + let callback_closed = handle.callback_closed().shared(); + let authorization_url = handle.authorization_url(); + let lifecycle = Arc::clone(self); + let completion_id = login_id.clone(); + tokio::spawn(async move { + let result = tokio::select! { + biased; + _ = cancel.cancelled() => Err(anyhow::anyhow!("Enterprise sign-in was cancelled.")), + result = async { + let credentials = handle.wait().await?; + callback_closed.clone().await; + let mut active = credentials.commit_if(|| async { + lifecycle.auth_manager.reload().await; + let eligible = lifecycle.configuration_is_current(&target, keyring_backend).await; + // Reread persisted account authority after acquiring the credential + // lock, including changes made by another app-server process. + lifecycle.auth_manager.reload().await; + if !eligible { + return None; + } + // Keep cancellation responsive during the async checks above. + // Once admitted, retain this guard through the synchronous write. + let active = lifecycle.active.lock().await; + (active.as_ref().is_some_and(|attempt| { + attempt.login_id == completion_id && !attempt.cancel.is_cancelled() + }) + && ema_auth_scope(lifecycle.auth_manager.auth_cached().as_ref()).as_ref() == Some(&scope)) + .then_some(active) + }).await?; + // A successful commit is no longer cancelable. Clear it while + // still holding the commit gate. + *active = None; + Ok(()) + } => result, + }; + // The canceled wait future has dropped its listener guard. Also wait + // for the blocking callback worker before allowing a fixed-port retry. + callback_closed.await; + if result.is_err() { + lifecycle.clear(&completion_id).await; + } + let success = result.is_ok(); + completed.send_replace(if success { + EnterpriseLoginCompletion::Succeeded + } else { + EnterpriseLoginCompletion::Failed + }); + }); + Ok(EnterpriseLogin { + login_id, + authorization_url, + completion, + }) + } +} diff --git a/codex-rs/app-server/src/request_processors/mcp_processor.rs b/codex-rs/app-server/src/request_processors/mcp_processor.rs index ed9a640c2c..ac2a55ba79 100644 --- a/codex-rs/app-server/src/request_processors/mcp_processor.rs +++ b/codex-rs/app-server/src/request_processors/mcp_processor.rs @@ -1,8 +1,13 @@ +use super::account_processor::EnterpriseLoginCompletion; +use super::account_processor::EnterpriseLoginTarget; use super::thread_input::ensure_direct_input_allowed; use super::*; +use codex_config::types::McpServerAuth; use codex_mcp::McpServerSource; use codex_mcp::ReadResourceRequestParams; +use codex_mcp::ema_auth_scope; use codex_mcp::resolve_oauth_callback; +use codex_rmcp_client::EnterpriseOAuthLoginRequest; use std::future::Future; use crate::thread_state::ThreadStateManager; @@ -15,6 +20,7 @@ pub(crate) struct McpRequestProcessor { thread_manager: Arc, pub(super) outgoing: Arc, config_manager: ConfigManager, + account_processor: Arc, pub(super) thread_state_manager: ThreadStateManager, } @@ -25,12 +31,14 @@ impl McpRequestProcessor { thread_state_manager: ThreadStateManager, outgoing: Arc, config_manager: ConfigManager, + account_processor: Arc, ) -> Self { Self { auth_manager, thread_manager, outgoing, config_manager, + account_processor, thread_state_manager, } } @@ -134,12 +142,6 @@ impl McpRequestProcessor { scopes, timeout_secs, } = params; - let client_registration = match client_registration.unwrap_or_default() { - McpServerOauthClientRegistration::Auto => McpOAuthClientRegistration::Auto, - McpServerOauthClientRegistration::Cimd => McpOAuthClientRegistration::Cimd, - McpServerOauthClientRegistration::Dcr => McpOAuthClientRegistration::Dcr, - }; - let auth = self.auth_manager.auth().await; let (mcp_config, runtime_context) = match thread_id.as_deref() { Some(thread_id) => { @@ -174,84 +176,175 @@ impl McpRequestProcessor { StreamableHttpRedirectMode::Legacy }; let server = server.config(); - if matches!(server.auth, codex_config::McpServerAuth::EmaAuth) { - return Err(invalid_request( - "EMA MCP connections are not enabled in this version", - )); - } - - let (url, http_headers, env_http_headers) = match &server.transport { - McpServerTransportConfig::StreamableHttp { - url, - http_headers, - env_http_headers, - .. - } => (url.clone(), http_headers.clone(), env_http_headers.clone()), - _ => { + let enterprise = matches!(server.auth, McpServerAuth::EmaAuth); + let (authorization_url, login_id, completion): ( + String, + String, + futures::future::BoxFuture<'static, anyhow::Result<()>>, + ) = if enterprise { + let thread_id = thread_id.clone().ok_or_else(|| { + invalid_request("enterprise MCP login requires a connected thread") + })?; + if !mcp_config.xaa_enabled { return Err(invalid_request( - "OAuth login is only supported for streamable HTTP servers.", + "enterprise MCP authentication requires `[features].use_xaa = true`", )); } - }; - - let http_client = runtime_context - .resolve_http_client(&name, server) - .map_err(|err| { - internal_error(format!("failed to resolve MCP server runtime: {err}")) + if !server.enabled { + return Err(invalid_request(format!("MCP server '{name}' is disabled."))); + } + server + .validate_ema_auth_transport() + .map_err(|reason| invalid_request(format!("MCP server '{name}' {reason}")))?; + if client_registration.is_some() { + return Err(invalid_request( + "enterprise MCP login uses the configured enterprise client and does not support clientRegistration overrides", + )); + } + // Resource scopes and per-server callbacks do not alter the shared IdP registration. + let scope = ema_auth_scope(auth.as_ref()).ok_or_else(|| { + invalid_request( + "enterprise MCP authentication requires a signed-in Codex user and workspace", + ) })?; + let idp = server.oauth_idp().ok_or_else(|| { + invalid_request(format!( + "MCP server '{name}' has no configured enterprise IdP" + )) + })?; + let registration = server.ema_registration().cloned().ok_or_else(|| { + invalid_request(format!( + "MCP server '{name}' has no validated enterprise registration" + )) + })?; + let credential_name = idp.credential_name(&scope); + let http_client = + runtime_context + .resolve_http_client(&name, server) + .map_err(|err| { + internal_error(format!("failed to resolve MCP server runtime: {err}")) + })?; + let login = self + .account_processor + .enterprise_login + .start( + EnterpriseOAuthLoginRequest { + credential_name: &credential_name, + issuer: &idp.issuer, + client_id: &idp.client_id, + keyring_backend_kind: mcp_config.auth_keyring_backend_kind, + callback_port: mcp_config.mcp_oauth_callback_port, + callback_url: mcp_config.mcp_oauth_callback_url.as_deref(), + timeout_secs, + http_client, + redirect_mode, + }, + scope, + EnterpriseLoginTarget { + thread_id, + server_name: name.clone(), + registration, + }, + ) + .await?; + let mut completion = login.completion; + ( + login.authorization_url, + login.login_id, + Box::pin(async move { + match completion.wait_for(|state| state.is_complete()).await { + Ok(result) if *result == EnterpriseLoginCompletion::Succeeded => Ok(()), + _ => Err(anyhow::anyhow!("Enterprise sign-in failed.")), + } + }), + ) + } else { + let (url, http_headers, env_http_headers) = match &server.transport { + McpServerTransportConfig::StreamableHttp { + url, + http_headers, + env_http_headers, + .. + } => (url.clone(), http_headers.clone(), env_http_headers.clone()), + _ => { + return Err(invalid_request( + "OAuth login is only supported for streamable HTTP servers.", + )); + } + }; - let discovered_scopes = if scopes.is_none() && server.scopes.is_none() { - discover_supported_scopes( - &server.transport, - Arc::clone(&http_client), - codex_rmcp_client::OAuthDiscoveryTimeout::Requested, + let http_client = + runtime_context + .resolve_http_client(&name, server) + .map_err(|err| { + internal_error(format!("failed to resolve MCP server runtime: {err}")) + })?; + + let discovered_scopes = if scopes.is_none() && server.scopes.is_none() { + discover_supported_scopes( + &server.transport, + Arc::clone(&http_client), + codex_rmcp_client::OAuthDiscoveryTimeout::Requested, + redirect_mode, + ) + .await + } else { + None + }; + let resolved_scopes = + resolve_oauth_scopes(scopes, server.scopes.clone(), discovered_scopes); + let oauth_credential_name = server.oauth_credential_name(&name); + let client_registration = match client_registration.unwrap_or_default() { + McpServerOauthClientRegistration::Auto => McpOAuthClientRegistration::Auto, + McpServerOauthClientRegistration::Cimd => McpOAuthClientRegistration::Cimd, + McpServerOauthClientRegistration::Dcr => McpOAuthClientRegistration::Dcr, + }; + let callback_url = + resolve_oauth_callback(server, &url, mcp_config.mcp_oauth_callback_url.as_deref()) + .map_err(|err| { + internal_error(format!("failed to resolve MCP OAuth callback: {err}")) + })?; + + let handle = perform_oauth_login_return_url( + oauth_credential_name.as_ref(), + &url, + mcp_config.mcp_oauth_credentials_store_mode, + mcp_config.auth_keyring_backend_kind, + http_headers, + env_http_headers, + &resolved_scopes.scopes, + server.oauth.as_ref(), + client_registration, + server.oauth_resource.as_deref(), + timeout_secs, + server.oauth_callback_port(mcp_config.mcp_oauth_callback_port), + callback_url.as_deref(), + mcp_config.mcp_oauth_callback_url.as_deref(), + http_client, redirect_mode, ) .await - } else { - None + .map_err(|err| { + internal_error(format!("failed to login to MCP server '{name}': {err}")) + })?; + ( + handle.authorization_url().to_string(), + Uuid::now_v7().to_string(), + Box::pin(handle.wait()), + ) }; - let resolved_scopes = - resolve_oauth_scopes(scopes, server.scopes.clone(), discovered_scopes); - let oauth_credential_name = server.oauth_credential_name(&name); - let callback_url = - resolve_oauth_callback(server, &url, mcp_config.mcp_oauth_callback_url.as_deref()) - .map_err(|err| { - internal_error(format!("failed to resolve MCP OAuth callback: {err}")) - })?; - - let handle = perform_oauth_login_return_url( - oauth_credential_name.as_ref(), - &url, - mcp_config.mcp_oauth_credentials_store_mode, - mcp_config.auth_keyring_backend_kind, - http_headers, - env_http_headers, - &resolved_scopes.scopes, - server.oauth.as_ref(), - client_registration, - server.oauth_resource.as_deref(), - timeout_secs, - server.oauth_callback_port(mcp_config.mcp_oauth_callback_port), - callback_url.as_deref(), - mcp_config.mcp_oauth_callback_url.as_deref(), - http_client, - redirect_mode, - ) - .await - .map_err(|err| internal_error(format!("failed to login to MCP server '{name}': {err}")))?; - let authorization_url = handle.authorization_url().to_string(); let notification_name = name.clone(); let notification_thread_id = thread_id; + let notification_login_id = login_id.clone(); let outgoing = Arc::clone(&self.outgoing); let thread_manager = Arc::clone(&self.thread_manager); tokio::spawn(async move { - let (success, error) = match handle.wait().await { + let (success, error) = match completion.await { Ok(()) => (true, None), Err(err) => (false, Some(err.to_string())), }; - if success { + if success && !enterprise { thread_manager.invalidate_mcp_runtimes().await; } @@ -259,6 +352,7 @@ impl McpRequestProcessor { McpServerOauthLoginCompletedNotification { name: notification_name, thread_id: notification_thread_id, + login_id: Some(notification_login_id), success, error, }, @@ -266,7 +360,10 @@ impl McpRequestProcessor { outgoing.send_server_notification(notification).await; }); - Ok(McpServerOauthLoginResponse { authorization_url }) + Ok(McpServerOauthLoginResponse { + authorization_url, + login_id: Some(login_id), + }) } async fn list_mcp_server_status( diff --git a/codex-rs/app-server/src/request_processors/plugins.rs b/codex-rs/app-server/src/request_processors/plugins.rs index 713794c253..63702b6e50 100644 --- a/codex-rs/app-server/src/request_processors/plugins.rs +++ b/codex-rs/app-server/src/request_processors/plugins.rs @@ -1954,6 +1954,7 @@ impl PluginRequestProcessor { McpServerOauthLoginCompletedNotification { name: notification_name, thread_id: None, + login_id: None, success, error, }, diff --git a/codex-rs/app-server/tests/suite/v2/account.rs b/codex-rs/app-server/tests/suite/v2/account.rs index b2d9e51425..239ed7097f 100644 --- a/codex-rs/app-server/tests/suite/v2/account.rs +++ b/codex-rs/app-server/tests/suite/v2/account.rs @@ -92,6 +92,9 @@ fn expected_workspace_routing( }) } +#[path = "account_enterprise_tests.rs"] +mod enterprise_tests; + // Helper to create a minimal config.toml for the app server #[derive(Default)] struct CreateConfigTomlParams { @@ -281,23 +284,62 @@ async fn mock_oauth_token(server: &MockServer, id_token: &str) { .await; } +#[test_case(false; "api_key")] +#[test_case(true; "enterprise_cleanup_failure")] #[tokio::test] -async fn logout_account_removes_auth_and_notifies() -> Result<()> { +async fn logout_account_removes_auth_and_notifies(enterprise_cleanup_failure: bool) -> Result<()> { + let oauth_server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/oauth/revoke")) + .respond_with(ResponseTemplate::new(200)) + .expect(u64::from(enterprise_cleanup_failure)) + .mount(&oauth_server) + .await; let codex_home = TempDir::new()?; create_config_toml(codex_home.path(), CreateConfigTomlParams::default())?; - login_with_api_key( - codex_home.path(), - "sk-test-key", - AuthCredentialsStoreMode::File, - AuthKeyringBackendKind::default(), - )?; + if enterprise_cleanup_failure { + write_chatgpt_auth( + codex_home.path(), + ChatGptAuthFixture::new("access-token") + .account_id(WORKSPACE_ID_INITIAL) + .chatgpt_user_id("enterprise-user"), + AuthCredentialsStoreMode::File, + )?; + let mut config = read_config_toml(codex_home.path())?; + config.as_table_mut().expect("config table").insert( + "mcp_enterprise_managed_auth".into(), + toml::toml! { + [idp] + issuer = "https://idp.example" + client_id = "enterprise-client" + } + .into(), + ); + std::fs::write( + codex_home.path().join("config.toml"), + toml::to_string(&config)?, + )?; + // Fail credential cleanup before touching the platform keyring. + std::fs::write(codex_home.path().join("mcp-oauth-locks"), "not a directory")?; + } else { + login_with_api_key( + codex_home.path(), + "sk-test-key", + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + } assert!(codex_home.path().join("auth.json").exists()); + let refresh_url = format!("{}/oauth/token", oauth_server.uri()); let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .with_env_overrides(&[("OPENAI_API_KEY", None)]) + .with_env_overrides(&[ + ("OPENAI_API_KEY", None), + (REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR, Some(&refresh_url)), + ]) .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; diff --git a/codex-rs/app-server/tests/suite/v2/account_enterprise_tests.rs b/codex-rs/app-server/tests/suite/v2/account_enterprise_tests.rs new file mode 100644 index 0000000000..e55ac863c6 --- /dev/null +++ b/codex-rs/app-server/tests/suite/v2/account_enterprise_tests.rs @@ -0,0 +1,814 @@ +//! Account API regressions for enterprise policy, persisted authority, and login attempts. + +use super::*; +use std::any::Any; +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::Mutex; + +use anyhow::Context; +use app_test_support::MockResponsesConfig; +use base64::Engine; +use base64::engine::general_purpose::URL_SAFE_NO_PAD; +use codex_app_server::in_process; +use codex_app_server::in_process::InProcessClientHandle; +use codex_app_server::in_process::InProcessServerEvent; +use codex_app_server::in_process::InProcessStartArgs; +use codex_app_server_protocol::InitializeParams; +use codex_arg0::Arg0DispatchPaths; +use codex_config::CloudConfigBundleLoader; +use codex_config::LoaderOverrides; +use codex_config::McpEmaAuthScope; +use codex_config::McpServerIdpOAuthConfig; +use codex_config::NoopThreadConfigLoader; +use codex_config::test_support::CloudConfigBundleFixture; +use codex_config::types::OAuthCredentialsStoreMode; +use codex_core::config::ConfigBuilder; +use codex_exec_server::EnvironmentManager; +use codex_features::Feature; +use codex_feedback::CodexFeedback; +use codex_login::CODEX_ACCESS_TOKEN_ENV_VAR; +use codex_protocol::protocol::SessionSource; +use codex_rmcp_client::stored_oauth_credentials; +use core_test_support::skip_if_no_network; +use keyring::credential::Credential; +use keyring::credential::CredentialApi; +use keyring::credential::CredentialBuilderApi; +use keyring::credential::CredentialPersistence; +use keyring::mock::MockCredential; +use pretty_assertions::assert_eq; +use serde_json::Value; +use test_case::test_case; +use tokio::io::AsyncReadExt; +use tokio::io::AsyncWriteExt; +use tokio::process::Command; +use wiremock::matchers::body_partial_json; +use wiremock::matchers::header; + +const WORKSPACE_ID_REFRESHED: &str = "123e4567-e89b-42d3-a456-426614174012"; + +#[test_case(false; "current_policy")] +#[test_case(true; "policy_load_failure")] +#[tokio::test] +async fn logout_reloads_persisted_workspace_and_its_enterprise_policy( + policy_load_failure: bool, +) -> Result<()> { + let server = MockServer::start().await; + for (workspace, issuer) in [ + (WORKSPACE_ID_INITIAL, "https://initial-idp.example"), + (WORKSPACE_ID_REFRESHED, "https://next-idp.example"), + ] { + let contents = if policy_load_failure && workspace == WORKSPACE_ID_REFRESHED { + "invalid = [".to_string() + } else { + format!( + r#"[mcp_enterprise_managed_auth.idp] +issuer = "{issuer}" +client_id = "enterprise-client""# + ) + }; + Mock::given(method("GET")) + .and(path("/backend-api/wham/config/bundle")) + .and(header("chatgpt-account-id", workspace)) + .respond_with( + ResponseTemplate::new(200).set_body_json( + CloudConfigBundleFixture::enterprise_config(contents).into_bundle(), + ), + ) + .expect(1) + .mount(&server) + .await; + } + Mock::given(method("POST")) + .and(path("/oauth/revoke")) + .and(body_partial_json(json!({ + "token": "next-refresh-token", + "token_type_hint": "refresh_token", + }))) + .respond_with(ResponseTemplate::new(200)) + .expect(1) + .mount(&server) + .await; + let codex_home = TempDir::new()?; + create_config_toml( + codex_home.path(), + CreateConfigTomlParams { + chatgpt_base_url: Some(format!("{}/backend-api", server.uri())), + ..Default::default() + }, + )?; + write_chatgpt_auth( + codex_home.path(), + ChatGptAuthFixture::new("initial-token") + .refresh_token("initial-refresh-token") + .account_id(WORKSPACE_ID_INITIAL) + .chatgpt_user_id("enterprise-user") + .plan_type("business"), + AuthCredentialsStoreMode::File, + )?; + // Keep the real platform keyring out of this API test. Enterprise cleanup + // failure must still permit revocation and removal of the selected account. + if !policy_load_failure { + std::fs::write(codex_home.path().join("mcp-oauth-locks"), "not a directory")?; + } + let refresh_url = format!("{}/oauth/token", server.uri()); + let mut app = TestAppServer::builder() + .with_codex_home(codex_home.path()) + .without_auto_env() + .with_env_overrides(&[ + ("OPENAI_API_KEY", None), + (REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR, Some(&refresh_url)), + ]) + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) + .await?; + + // Simulate a second process switching workspaces while this server caches A. + write_chatgpt_auth( + codex_home.path(), + ChatGptAuthFixture::new("next-token") + .refresh_token("next-refresh-token") + .account_id(WORKSPACE_ID_REFRESHED) + .chatgpt_user_id("enterprise-user") + .plan_type("business"), + AuthCredentialsStoreMode::File, + )?; + let id = app.send_logout_account_request().await?; + let response: LogoutAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, app.read_response(id)).await??; + assert_eq!(response, LogoutAccountResponse {}); + let notification = timeout( + DEFAULT_READ_TIMEOUT, + app.read_stream_until_notification_message("account/updated"), + ) + .await??; + let notification: ServerNotification = notification.try_into()?; + let ServerNotification::AccountUpdated(account) = notification else { + bail!("unexpected notification: {notification:?}"); + }; + assert_eq!( + account, + AccountUpdatedNotification { + auth_mode: None, + plan_type: None, + } + ); + assert!(!codex_home.path().join("auth.json").exists()); + if policy_load_failure { + assert!( + !codex_home.path().join("mcp-oauth-locks").exists(), + "invalid current policy must not select any enterprise credential for deletion" + ); + } + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn enterprise_login_attempt_lifecycle() -> Result<()> { + skip_if_no_network!(Ok(())); + const CHILD: &str = "CODEX_ENTERPRISE_ATTEMPT_TEST_CHILD"; + if std::env::var_os(CHILD).is_none() { + // The app server and callback worker share only this child's mock keyring. + let home = TempDir::new()?; + let oauth_server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/oauth/revoke")) + .respond_with(ResponseTemplate::new(200)) + .expect(1) + .mount(&oauth_server) + .await; + let output = Command::new(std::env::current_exe()?) + .args([ + "--exact", + "suite::v2::account::enterprise_tests::enterprise_login_attempt_lifecycle", + "--nocapture", + ]) + .env(CHILD, "1") + .env("CODEX_HOME", home.path()) + .env( + REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR, + format!("{}/oauth/token", oauth_server.uri()), + ) + .env_remove("OPENAI_API_KEY") + .env_remove("CODEX_API_KEY") + .env_remove(CODEX_ACCESS_TOKEN_ENV_VAR) + .current_dir(home.path()) + .output() + .await?; + let stdout = String::from_utf8_lossy(&output.stdout); + assert!( + output.status.success() && stdout.contains("1 passed; 0 failed"), + "{stdout}\n{}", + String::from_utf8_lossy(&output.stderr) + ); + return Ok(()); + } + + keyring::set_default_credential_builder(Box::::default()); + let server = MockServer::builder().start().await; + let origin = server.uri(); + let issuer = format!("{origin}/idp"); + Mock::given(method("GET")) + .and(path("/backend-api/wham/config/bundle")) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(CloudConfigBundleFixture::default().into_bundle()), + ) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/.well-known/oauth-authorization-server/idp")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "issuer": issuer, + "authorization_endpoint": format!("{issuer}/authorize"), + "token_endpoint": format!("{issuer}/token"), + "token_endpoint_auth_methods_supported": ["none"], + }))) + .mount(&server) + .await; + let assertion = format!( + "{}.{}.signature", + URL_SAFE_NO_PAD.encode(r#"{"alg":"ES256"}"#), + URL_SAFE_NO_PAD.encode( + json!({"iss":issuer,"sub":"enterprise-user","aud":"idp-client","exp":4102444800_u64}) + .to_string() + ) + ); + Mock::given(method("POST")) + .and(path("/idp/token")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token":"enterprise-access", "refresh_token":"enterprise-refresh", + "id_token":assertion, "token_type":"Bearer", + }))) + .mount(&server) + .await; + + let home = TempDir::new()?; + let listener = std::net::TcpListener::bind("127.0.0.1:0")?; + let callback_port = listener.local_addr()?.port(); + MockResponsesConfig::new(&origin) + .enable_feature(Feature::UseXaa) + .disable_feature(Feature::SecretAuthStorage) + .disable_feature(Feature::Apps) + .with_root_config(&format!( + r#"cli_auth_credentials_store = "file" +chatgpt_base_url = "{origin}/backend-api" +mcp_oauth_callback_port = {callback_port} +analytics = {{ enabled = false }}"# + )) + .with_provider_config("requires_openai_auth = true") + .with_extra_config(&format!( + r#"[mcp_enterprise_managed_auth.idp] +issuer = "{issuer}" +client_id = "idp-client" + +[mcp_servers.enterprise] +url = "{origin}/mcp" +auth = "ema_auth" +scopes = ["files.read"] +oauth = {{ client_id = "mcp-client", authorization_server_issuer = "{origin}/as" }} + +[mcp_servers.disabled-enterprise] +url = "{origin}/mcp" +enabled = false +auth = "ema_auth" +oauth = {{ client_id = "mcp-client", authorization_server_issuer = "{origin}/as" }}"# + )) + .write(home.path())?; + write_chatgpt_auth( + home.path(), + ChatGptAuthFixture::new("account-access") + .account_id(WORKSPACE_ID_INITIAL) + .chatgpt_user_id("enterprise-user") + .plan_type("business"), + AuthCredentialsStoreMode::File, + )?; + let loader_overrides = LoaderOverrides::without_managed_config_for_tests(); + let config = Arc::new( + ConfigBuilder::default() + .codex_home(home.path().to_path_buf()) + .fallback_cwd(Some(home.path().to_path_buf())) + .loader_overrides(loader_overrides.clone()) + .build() + .await?, + ); + let mut client = in_process::start(InProcessStartArgs { + arg0_paths: Arg0DispatchPaths::default(), + config, + cli_overrides: Vec::new(), + loader_overrides, + strict_config: false, + cloud_config_bundle: CloudConfigBundleLoader::default(), + thread_config_loader: Arc::new(NoopThreadConfigLoader), + feedback: CodexFeedback::new(), + log_db: None, + state_db: None, + environment_manager: Arc::new(EnvironmentManager::default_for_tests()), + config_warnings: Vec::new(), + embedded_network_policy: Default::default(), + session_source: SessionSource::Cli, + enable_codex_api_key_env: false, + initialize: InitializeParams { + client_info: ClientInfo { + name: "codex-app-server-tests".into(), + title: None, + version: "0.1.0".into(), + }, + capabilities: Some(InitializeCapabilities { + experimental_api: true, + ..Default::default() + }), + }, + channel_capacity: in_process::DEFAULT_IN_PROCESS_CHANNEL_CAPACITY, + }) + .await?; + drop(listener); + let thread = enterprise_rpc(&client, "thread/start", json!({})).await?; + let thread_id = thread["thread"]["id"].as_str().context("thread ID")?; + + assert_eq!( + enterprise_rpc( + &client, + "mcpServer/oauth/login", + json!({"name":"disabled-enterprise","threadId":thread_id}), + ) + .await + .expect_err("disabled enterprise MCP login must reject immediately") + .to_string(), + "mcpServer/oauth/login: MCP server 'disabled-enterprise' is disabled." + ); + let first = enterprise_rpc( + &client, + "mcpServer/oauth/login", + json!({"name":"enterprise","threadId":thread_id}), + ) + .await?; + enterprise_callback(&first, &issuer, callback_port)?; + let replacement = enterprise_rpc( + &client, + "mcpServer/oauth/login", + json!({"name":"enterprise","threadId":thread_id}), + ) + .await?; + assert_ne!(replacement["loginId"], first["loginId"]); + enterprise_callback(&replacement, &issuer, callback_port)?; + assert_eq!( + enterprise_completion(&mut client).await?, + json!({"name":"enterprise","threadId":thread_id,"loginId":first["loginId"],"success":false,"error":"Enterprise sign-in failed."}) + ); + assert_eq!( + enterprise_rpc( + &client, + "account/login/cancel", + json!({"loginId":replacement["loginId"]}), + ) + .await?, + json!({"status":"canceled"}) + ); + assert_eq!( + enterprise_completion(&mut client).await?, + json!({"name":"enterprise","threadId":thread_id,"loginId":replacement["loginId"],"success":false,"error":"Enterprise sign-in failed."}) + ); + let idp = McpServerIdpOAuthConfig { + issuer: issuer.clone(), + client_id: "idp-client".into(), + }; + let initial_scope = McpEmaAuthScope::new("enterprise-user".into(), WORKSPACE_ID_INITIAL.into()) + .expect("initial scope"); + assert_eq!( + stored_oauth_credentials( + &idp.credential_name(&initial_scope), + &issuer, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )?, + None + ); + + let previous = enterprise_rpc( + &client, + "mcpServer/oauth/login", + json!({"name":"enterprise","threadId":thread_id}), + ) + .await?; + let previous_callback = enterprise_callback(&previous, &issuer, callback_port)?; + write_chatgpt_auth( + home.path(), + ChatGptAuthFixture::new("account-access") + .account_id(WORKSPACE_ID_REFRESHED) + .chatgpt_user_id("enterprise-user") + .plan_type("business"), + AuthCredentialsStoreMode::File, + )?; + let browser = HttpClientBuilder::new() + .without_request_logging() + .without_redirects() + .build_direct()?; + // Each attempt owns a one-shot listener on this port. Do not reuse an HTTP + // connection from the previous callback after that listener has retired. + let response = timeout( + DEFAULT_READ_TIMEOUT, + browser + .get(previous_callback) + .header(http::header::CONNECTION, "close") + .send(), + ) + .await??; + assert_eq!(response.status(), http::StatusCode::OK); + assert_eq!( + enterprise_completion(&mut client).await?, + json!({"name":"enterprise","threadId":thread_id,"loginId":previous["loginId"],"success":false,"error":"Enterprise sign-in failed."}) + ); + + let retry = enterprise_rpc( + &client, + "mcpServer/oauth/login", + json!({"name":"enterprise","threadId":thread_id}), + ) + .await?; + let retry_callback = enterprise_callback(&retry, &issuer, callback_port)?; + let response = timeout( + DEFAULT_READ_TIMEOUT, + browser + .get(retry_callback) + .header(http::header::CONNECTION, "close") + .send(), + ) + .await??; + assert_eq!(response.status(), http::StatusCode::OK); + assert_eq!( + enterprise_completion(&mut client).await?, + json!({"name":"enterprise","threadId":thread_id,"loginId":retry["loginId"],"success":true}) + ); + let replacement_scope = + McpEmaAuthScope::new("enterprise-user".into(), WORKSPACE_ID_REFRESHED.into()) + .expect("replacement scope"); + for (scope, present) in [(initial_scope, false), (replacement_scope, true)] { + assert_eq!( + stored_oauth_credentials( + &idp.credential_name(&scope), + &issuer, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )? + .is_some(), + present + ); + } + + let stalled = enterprise_rpc( + &client, + "mcpServer/oauth/login", + json!({"name":"enterprise","threadId":thread_id}), + ) + .await?; + let mut unfinished = tokio::net::TcpStream::connect(("127.0.0.1", callback_port)).await?; + unfinished + .write_all(b"POST /unfinished HTTP/1.1\r\nHost: 127.0.0.1\r\nContent-Length: 2048\r\n\r\n") + .await?; + // Reading the response proves the callback worker has accepted this request. + // It cannot finish draining the declared body while this socket stays open. + let mut response = [0; 256]; + assert!(timeout(DEFAULT_READ_TIMEOUT, unfinished.read(&mut response)).await?? > 0); + let logout = timeout( + std::time::Duration::from_secs(15), + enterprise_rpc(&client, "account/logout", json!({})), + ) + .await; + // Release the worker even when the regression fails, so child shutdown cannot hang. + drop(unfinished); + logout.context("an unfinished callback must not prevent primary logout")??; + assert!(!home.path().join("auth.json").exists()); + assert!( + stored_oauth_credentials( + &idp.credential_name( + &McpEmaAuthScope::new("enterprise-user".into(), WORKSPACE_ID_REFRESHED.into(),) + .expect("replacement scope") + ), + &issuer, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )? + .is_none() + ); + assert_eq!( + enterprise_completion(&mut client).await?, + json!({"name":"enterprise","threadId":thread_id,"loginId":stalled["loginId"],"success":false,"error":"Enterprise sign-in failed."}) + ); + client.shutdown().await?; + Ok(()) +} + +fn enterprise_callback(start: &Value, issuer: &str, callback_port: u16) -> Result { + let query = Url::parse( + start["authorizationUrl"] + .as_str() + .context("authorization URL")?, + )? + .query_pairs() + .into_owned() + .collect::>(); + let mut callback = Url::parse(query.get("redirect_uri").context("redirect_uri")?)?; + assert_eq!( + (callback.scheme(), callback.host_str(), callback.port()), + ("http", Some("127.0.0.1"), Some(callback_port)) + ); + callback + .query_pairs_mut() + .append_pair("code", "enterprise-code") + .append_pair("state", query.get("state").context("state")?) + .append_pair("iss", issuer); + Ok(callback) +} + +async fn enterprise_rpc( + client: &InProcessClientHandle, + method: &str, + params: Value, +) -> Result { + let request = serde_json::from_value( + json!({"id":uuid::Uuid::new_v4().to_string(),"method":method,"params":params}), + )?; + let response = timeout(DEFAULT_READ_TIMEOUT, client.request(request)) + .await + .with_context(|| format!("timed out waiting for {method}"))??; + response.map_err(|error| anyhow::anyhow!("{method}: {}", error.message)) +} + +async fn enterprise_completion(client: &mut InProcessClientHandle) -> Result { + timeout(DEFAULT_READ_TIMEOUT, async { + loop { + match client + .next_event() + .await + .context("app-server notification")? + { + InProcessServerEvent::ServerNotification(notification) => { + if let ServerNotification::McpServerOauthLoginCompleted(notification) = + *notification + { + return Ok(serde_json::to_value(notification)?); + } + } + event => bail!("unexpected app-server event: {event:?}"), + } + } + }) + .await + .context("timed out waiting for enterprise login completion")? +} + +#[derive(Default)] +struct TestKeyring(Mutex>>); + +struct TestCredential(Arc); + +impl CredentialApi for TestCredential { + fn set_secret(&self, secret: &[u8]) -> keyring::Result<()> { + self.0.set_secret(secret) + } + fn get_secret(&self) -> keyring::Result> { + self.0.get_secret() + } + fn delete_credential(&self) -> keyring::Result<()> { + self.0.delete_credential() + } + fn as_any(&self) -> &dyn Any { + self + } +} + +impl CredentialBuilderApi for TestKeyring { + fn build( + &self, + _target: Option<&str>, + service: &str, + user: &str, + ) -> keyring::Result> { + let credential = self + .0 + .lock() + .expect("mock keyring") + .entry((service.into(), user.into())) + .or_default() + .clone(); + Ok(Box::new(TestCredential(credential))) + } + fn as_any(&self) -> &dyn Any { + self + } + fn persistence(&self) -> CredentialPersistence { + CredentialPersistence::ProcessOnly + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn enterprise_login_setup_is_cancelled_by_account_changes() -> Result<()> { + skip_if_no_network!(Ok(())); + const CHILD: &str = "CODEX_ENTERPRISE_SETUP_TEST_CHILD"; + if std::env::var_os(CHILD).is_none() { + // The app server and callback worker share only this child's mock keyring. + let home = TempDir::new()?; + let oauth_server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/oauth/revoke")) + .respond_with(ResponseTemplate::new(200)) + .expect(0..=1) + .mount(&oauth_server) + .await; + let output = Command::new(std::env::current_exe()?) + .args([ + "--exact", + "suite::v2::account::enterprise_tests::enterprise_login_setup_is_cancelled_by_account_changes", + "--nocapture", + ]) + .env(CHILD, "1") + .env("CODEX_HOME", home.path()) + .env( + REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR, + format!("{}/oauth/token", oauth_server.uri()), + ) + .env_remove("OPENAI_API_KEY") + .env_remove("CODEX_API_KEY") + .env_remove(CODEX_ACCESS_TOKEN_ENV_VAR) + .current_dir(home.path()) + .output() + .await?; + let stdout = String::from_utf8_lossy(&output.stdout); + assert!( + output.status.success() && stdout.contains("1 passed; 0 failed"), + "{stdout}\n{}", + String::from_utf8_lossy(&output.stderr) + ); + return Ok(()); + } + + keyring::set_default_credential_builder(Box::::default()); + for (account_method, account_params) in [ + ("account/logout", json!({})), + ( + "account/login/start", + json!({"type":"apiKey","apiKey":"replacement-api-key"}), + ), + ] { + let server = MockServer::builder().start().await; + let origin = server.uri(); + let issuer = format!("{origin}/idp"); + Mock::given(method("GET")) + .and(path("/backend-api/wham/config/bundle")) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(CloudConfigBundleFixture::default().into_bundle()), + ) + .mount(&server) + .await; + let discovery_started = Arc::new(tokio::sync::Notify::new()); + let entered = Arc::clone(&discovery_started); + Mock::given(method("GET")) + .and(path("/.well-known/oauth-authorization-server/idp")) + .respond_with(move |_: &wiremock::Request| { + entered.notify_one(); + ResponseTemplate::new(200).set_delay(std::time::Duration::from_secs(60)) + }) + .mount(&server) + .await; + let assertion = format!( + "{}.{}.signature", + URL_SAFE_NO_PAD.encode(r#"{"alg":"ES256"}"#), + URL_SAFE_NO_PAD.encode( + json!({"iss":issuer,"sub":"enterprise-user","aud":"idp-client","exp":4102444800_u64}) + .to_string() + ) + ); + Mock::given(method("POST")) + .and(path("/idp/token")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token":"enterprise-access", "refresh_token":"enterprise-refresh", + "id_token":assertion, "token_type":"Bearer", + }))) + .mount(&server) + .await; + + let home = TempDir::new()?; + let listener = std::net::TcpListener::bind("127.0.0.1:0")?; + let callback_port = listener.local_addr()?.port(); + MockResponsesConfig::new(&origin) + .enable_feature(Feature::UseXaa) + .disable_feature(Feature::SecretAuthStorage) + .disable_feature(Feature::Apps) + .with_root_config(&format!( + r#"cli_auth_credentials_store = "file" +chatgpt_base_url = "{origin}/backend-api" +mcp_oauth_callback_port = {callback_port} +analytics = {{ enabled = false }}"# + )) + .with_provider_config("requires_openai_auth = true") + .with_extra_config(&format!( + r#"[mcp_enterprise_managed_auth.idp] +issuer = "{issuer}" +client_id = "idp-client" + +[mcp_servers.enterprise] +url = "{origin}/mcp" +auth = "ema_auth" +scopes = ["files.read"] +oauth = {{ client_id = "mcp-client", authorization_server_issuer = "{origin}/as" }} + +[mcp_servers.disabled-enterprise] +url = "{origin}/mcp" +enabled = false +auth = "ema_auth" +oauth = {{ client_id = "mcp-client", authorization_server_issuer = "{origin}/as" }}"# + )) + .write(home.path())?; + write_chatgpt_auth( + home.path(), + ChatGptAuthFixture::new("account-access") + .account_id(WORKSPACE_ID_INITIAL) + .chatgpt_user_id("enterprise-user") + .plan_type("business"), + AuthCredentialsStoreMode::File, + )?; + let loader_overrides = LoaderOverrides::without_managed_config_for_tests(); + let config = Arc::new( + ConfigBuilder::default() + .codex_home(home.path().to_path_buf()) + .fallback_cwd(Some(home.path().to_path_buf())) + .loader_overrides(loader_overrides.clone()) + .build() + .await?, + ); + let client = in_process::start(InProcessStartArgs { + arg0_paths: Arg0DispatchPaths::default(), + config, + cli_overrides: Vec::new(), + loader_overrides, + strict_config: false, + cloud_config_bundle: CloudConfigBundleLoader::default(), + thread_config_loader: Arc::new(NoopThreadConfigLoader), + feedback: CodexFeedback::new(), + log_db: None, + state_db: None, + environment_manager: Arc::new(EnvironmentManager::default_for_tests()), + config_warnings: Vec::new(), + embedded_network_policy: Default::default(), + session_source: SessionSource::Cli, + enable_codex_api_key_env: false, + initialize: InitializeParams { + client_info: ClientInfo { + name: "codex-app-server-tests".into(), + title: None, + version: "0.1.0".into(), + }, + capabilities: Some(InitializeCapabilities { + experimental_api: true, + ..Default::default() + }), + }, + channel_capacity: in_process::DEFAULT_IN_PROCESS_CHANNEL_CAPACITY, + }) + .await?; + drop(listener); + let thread = enterprise_rpc(&client, "thread/start", json!({})).await?; + let thread_id = thread["thread"]["id"].as_str().context("thread ID")?; + let (setup_result, account_result) = tokio::join!( + enterprise_rpc( + &client, + "mcpServer/oauth/login", + json!({"name":"enterprise","threadId":thread_id}), + ), + async { + timeout(DEFAULT_READ_TIMEOUT, discovery_started.notified()).await?; + timeout( + std::time::Duration::from_secs(15), + enterprise_rpc(&client, account_method, account_params), + ) + .await + .context("account changes must not wait for IdP discovery")? + }, + ); + account_result.context("stalled enterprise setup must not block account changes")?; + assert_eq!( + setup_result + .expect_err("account changes cancel pending enterprise setup") + .to_string(), + "mcpServer/oauth/login: enterprise sign-in was cancelled", + ); + let idp = McpServerIdpOAuthConfig { + issuer: issuer.clone(), + client_id: "idp-client".into(), + }; + let scope = + McpEmaAuthScope::new("enterprise-user".into(), WORKSPACE_ID_INITIAL.into()).unwrap(); + assert_eq!( + stored_oauth_credentials( + &idp.credential_name(&scope), + &issuer, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )?, + None + ); + client.shutdown().await?; + } + Ok(()) +} diff --git a/codex-rs/app-server/tests/suite/v2/executor_mcp.rs b/codex-rs/app-server/tests/suite/v2/executor_mcp.rs index 314aeb5d88..ce5e1e314b 100644 --- a/codex-rs/app-server/tests/suite/v2/executor_mcp.rs +++ b/codex-rs/app-server/tests/suite/v2/executor_mcp.rs @@ -1310,6 +1310,11 @@ startup_timeout_sec = 10 McpServerOauthLoginCompletedNotification { name: OAUTH_MCP_SERVER_NAME.to_string(), thread_id: Some(selected_thread.clone()), + login_id: Some( + response + .login_id + .expect("login response should contain an ID") + ), success: true, error: None, } diff --git a/codex-rs/app-server/tests/suite/v2/mcp_server_status.rs b/codex-rs/app-server/tests/suite/v2/mcp_server_status.rs index 9bb9243f2b..09daa053f3 100644 --- a/codex-rs/app-server/tests/suite/v2/mcp_server_status.rs +++ b/codex-rs/app-server/tests/suite/v2/mcp_server_status.rs @@ -372,6 +372,15 @@ async fn oauth_login_validates_callback_issuer_and_uses_http_headers_helper( mcp.read_notification("mcpServer/oauthLogin/completed"), ) .await??; + assert_eq!( + completed.login_id.as_deref(), + Some( + response + .login_id + .as_deref() + .expect("login response should contain an ID") + ) + ); assert_eq!( ( completed.name.as_str(), @@ -629,11 +638,21 @@ async fn oauth_login_automatically_selects_callback_specific_cimd_without_metada app_server.read_notification("mcpServer/oauthLogin/completed"), ) .await??; + assert_eq!( + completed.login_id.as_deref(), + Some( + response + .login_id + .as_deref() + .expect("login response should contain an ID") + ) + ); assert_eq!( completed, McpServerOauthLoginCompletedNotification { name: "cimd".to_string(), thread_id: None, + login_id: response.login_id, success: true, error: None, } diff --git a/codex-rs/app-server/tests/suite/v2/plugin_install.rs b/codex-rs/app-server/tests/suite/v2/plugin_install.rs index 3d62550c2f..3450633f47 100644 --- a/codex-rs/app-server/tests/suite/v2/plugin_install.rs +++ b/codex-rs/app-server/tests/suite/v2/plugin_install.rs @@ -2037,11 +2037,21 @@ async fn plugin_oauth_login_preserves_registered_callbacks_or_uses_legacy_fallba mcp.read_notification("mcpServer/oauthLogin/completed"), ) .await??; + assert_eq!( + completed.login_id.as_deref(), + Some( + response + .login_id + .as_deref() + .expect("login response should contain an ID") + ) + ); assert_eq!( completed, McpServerOauthLoginCompletedNotification { name: "sample-mcp".to_string(), thread_id: None, + login_id: response.login_id, success: true, error: None, } diff --git a/codex-rs/cli/src/login.rs b/codex-rs/cli/src/login.rs index 22e607c6e6..8b9f081dad 100644 --- a/codex-rs/cli/src/login.rs +++ b/codex-rs/cli/src/login.rs @@ -22,8 +22,10 @@ use codex_login::login_with_api_key; use codex_login::logout_with_revoke; use codex_login::run_device_code_login; use codex_login::run_login_server; +use codex_mcp::ema_auth_scope; use codex_protocol::auth::AuthMode; use codex_protocol::config_types::ForcedLoginMethod; +use codex_rmcp_client::EnterpriseOAuthCredentialGuard; use codex_utils_cli::CliConfigOverrides; use std::fs::OpenOptions; use std::io::IsTerminal; @@ -511,6 +513,37 @@ pub async fn run_logout(cli_config_overrides: CliConfigOverrides) -> ! { let config = load_config_or_exit(cli_config_overrides).await; let auth_route_config = config.auth_route_config(); + // Cleanup is independent of the current feature gate. Retain the grant lock + // through account removal so another process cannot commit a late login. + let enterprise_guard = async { + let Some(profile) = &config.mcp_enterprise_managed_auth else { + return Ok::<_, anyhow::Error>(None); + }; + let auth = config + .auth_config() + .load_auth(/*enable_codex_api_key_env*/ false) + .await?; + let Some(scope) = ema_auth_scope(auth.as_ref()) else { + return Ok(None); + }; + EnterpriseOAuthCredentialGuard::acquire( + &profile.idp.credential_name(&scope), + &profile.idp.issuer, + config.auth_keyring_backend_kind(), + ) + .await + .map(Some) + } + .await; + let cleanup_failed = match &enterprise_guard { + Ok(Some(guard)) => guard.delete_tokens().is_err(), + Ok(None) => false, + Err(_) => true, + }; + if cleanup_failed { + eprintln!("Warning: failed to remove enterprise authorization; continuing account logout"); + } + let logged_out = match logout_with_revoke( &config.codex_home, config.cli_auth_credentials_store_mode, @@ -525,6 +558,7 @@ pub async fn run_logout(cli_config_overrides: CliConfigOverrides) -> ! { std::process::exit(1); } }; + drop(enterprise_guard); let cleared_bedrock_config = if let Some(paths) = ConfigEditsBuilder::bedrock_provider_config_paths_to_clear(&config) { diff --git a/codex-rs/cli/src/mcp_cmd.rs b/codex-rs/cli/src/mcp_cmd.rs index 8fda684d73..ac4b0154b3 100644 --- a/codex-rs/cli/src/mcp_cmd.rs +++ b/codex-rs/cli/src/mcp_cmd.rs @@ -8,6 +8,7 @@ use anyhow::bail; use clap::ArgGroup; use clap::builder::TypedValueParser; use codex_config::types::AppToolApproval; +use codex_config::types::McpServerAuth; use codex_config::types::McpServerConfig; use codex_config::types::McpServerOAuthConfig; use codex_config::types::McpServerTransportConfig; @@ -28,6 +29,7 @@ use codex_mcp::McpRuntimeContext; use codex_mcp::apply_http_headers_helper; use codex_mcp::compute_auth_statuses; use codex_mcp::discover_supported_scopes; +use codex_mcp::ema_auth_scope; use codex_mcp::oauth_login_support; use codex_mcp::resolve_oauth_callback; use codex_mcp::resolve_oauth_scopes; @@ -36,6 +38,7 @@ use codex_rmcp_client::McpOAuthCallbackMode; use codex_rmcp_client::McpOAuthClientRegistration; use codex_rmcp_client::OAuthDiscoveryTimeout; use codex_rmcp_client::StreamableHttpRedirectMode; +use codex_rmcp_client::delete_enterprise_oauth_tokens; use codex_rmcp_client::delete_oauth_tokens; use codex_rmcp_client::resolve_mcp_oauth_callback_url; use codex_utils_cli::CliConfigOverrides; @@ -639,6 +642,36 @@ async fn run_logout(config: &Config, logout_args: LogoutArgs) -> Result<()> { .get(&name) .ok_or_else(|| anyhow!("No MCP server named '{name}' found in configuration."))?; + if matches!(server.auth, McpServerAuth::EmaAuth) { + let auth_manager = + AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ true).await?; + let auth = auth_manager.auth().await; + let scope = ema_auth_scope(auth.as_ref()) + .context("Sign in to your Codex account before removing enterprise authorization")?; + // Cleanup must remain available even when this server is no longer eligible + // for EMA (for example, after disabling the feature or changing its transport). + let profile = config + .mcp_enterprise_managed_auth + .as_ref() + .context("EMA logout requires a trusted enterprise IdP profile")?; + let idp = &profile.idp; + let credential_name = idp.credential_name(&scope); + match delete_enterprise_oauth_tokens( + &credential_name, + &idp.issuer, + config.auth_keyring_backend_kind(), + ) + .await + { + Ok(true) => println!( + "Removed the shared enterprise authorization used by EMA MCP servers (selected via '{name}')." + ), + Ok(false) => println!("No shared enterprise authorization is stored."), + Err(_) => return Err(anyhow!("failed to delete enterprise authorization")), + } + return Ok(()); + } + let url = match &server.transport { McpServerTransportConfig::StreamableHttp { url, .. } => url.clone(), _ => bail!("OAuth logout is only supported for streamable_http transports."), diff --git a/codex-rs/cli/tests/login.rs b/codex-rs/cli/tests/login.rs index f631057e19..810ecadaed 100644 --- a/codex-rs/cli/tests/login.rs +++ b/codex-rs/cli/tests/login.rs @@ -183,6 +183,79 @@ fn logout_clears_only_the_selected_bedrock_provider() -> Result<()> { Ok(()) } +#[tokio::test] +async fn logout_survives_enterprise_cleanup_failure_with_xaa_disabled() -> Result<()> { + let server = MockServer::start().await; + Mock::given(path("/backend-api/wham/config/bundle")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "config_toml": {"enterprise_managed": []}, + "requirements_toml": {"enterprise_managed": []}, + }))) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/oauth/revoke")) + .respond_with(ResponseTemplate::new(200)) + .expect(1) + .mount(&server) + .await; + let home = TempDir::new()?; + std::fs::write( + home.path().join("config.toml"), + format!( + "cli_auth_credentials_store = \"file\"\nchatgpt_base_url = \"{}/backend-api\"\n[features]\nuse_xaa = false\n[mcp_enterprise_managed_auth.idp]\nissuer = \"https://idp.example\"\nclient_id = \"enterprise-client\"\n[mcp_servers.enterprise]\nurl = \"https://resource.example/mcp\"\nauth = \"ema_auth\"\nbearer_token_env_var = \"UNUSED_TOKEN\"\n", + server.uri() + ), + )?; + write_chatgpt_auth( + home.path(), + ChatGptAuthFixture::new("account-access") + .account_id("workspace") + .chatgpt_user_id("user"), + AuthCredentialsStoreMode::File, + )?; + // Fail before touching the real keyring; the subprocess owns this isolated home. + std::fs::write(home.path().join("mcp-oauth-locks"), "not a directory")?; + for enabled in [false, true] { + codex_command(home.path())? + .current_dir(home.path()) + .env("NO_PROXY", "127.0.0.1,localhost") + .env("no_proxy", "127.0.0.1,localhost") + .env_remove("CODEX_API_KEY") + .env_remove("OPENAI_API_KEY") + .env_remove(CODEX_ACCESS_TOKEN_ENV_VAR) + .args([ + "-c", + &format!("features.use_xaa={enabled}"), + "mcp", + "logout", + "enterprise", + ]) + .assert() + .failure() + .stderr(contains("failed to delete enterprise authorization")); + assert!(home.path().join("auth.json").exists()); + } + codex_command(home.path())? + .current_dir(home.path()) + .env("NO_PROXY", "127.0.0.1,localhost") + .env("no_proxy", "127.0.0.1,localhost") + .env( + REVOKE_TOKEN_URL_OVERRIDE_ENV_VAR, + format!("{}/oauth/revoke", server.uri()), + ) + .env_remove("CODEX_API_KEY") + .env_remove("OPENAI_API_KEY") + .env_remove(CODEX_ACCESS_TOKEN_ENV_VAR) + .args(["logout"]) + .assert() + .success() + .stderr(contains("continuing account logout")) + .stderr(contains("Successfully logged out")); + assert!(!home.path().join("auth.json").exists()); + Ok(()) +} + #[test] fn login_with_access_token_rejects_invalid_jwt() -> Result<()> { let codex_home = TempDir::new()?; diff --git a/codex-rs/codex-mcp/src/ema_auth.rs b/codex-rs/codex-mcp/src/ema_auth.rs new file mode 100644 index 0000000000..80b3a0e55d --- /dev/null +++ b/codex-rs/codex-mcp/src/ema_auth.rs @@ -0,0 +1,9 @@ +//! Derives the enterprise credential owner from the active Codex account. + +use codex_config::McpEmaAuthScope; +use codex_login::CodexAuth; + +pub fn ema_auth_scope(auth: Option<&CodexAuth>) -> Option { + let auth = auth?; + McpEmaAuthScope::new(auth.get_chatgpt_user_id()?, auth.get_account_id()?) +} diff --git a/codex-rs/codex-mcp/src/lib.rs b/codex-rs/codex-mcp/src/lib.rs index 6eca43bcb3..d1b3a4b041 100644 --- a/codex-rs/codex-mcp/src/lib.rs +++ b/codex-rs/codex-mcp/src/lib.rs @@ -8,6 +8,7 @@ pub use elicitation::ElicitationLifecycle; pub use elicitation::ElicitationReviewRequest; pub use elicitation::ElicitationReviewer; pub use elicitation::ElicitationReviewerHandle; +pub use ema_auth::ema_auth_scope; pub use event_stream::McpEventStreamOpener; pub use resource_client::CodexAppsResourceListParams; pub use resource_client::McpEventCatalogSnapshot; @@ -110,6 +111,7 @@ mod client_tool_catalog; pub(crate) mod codex_apps; pub(crate) mod connection_manager; pub(crate) mod elicitation; +mod ema_auth; mod event_stream; mod executor_environment_http_client; pub(crate) mod mcp; diff --git a/codex-rs/config/src/lib.rs b/codex-rs/config/src/lib.rs index 0893930ba0..d90d57ae87 100644 --- a/codex-rs/config/src/lib.rs +++ b/codex-rs/config/src/lib.rs @@ -171,6 +171,7 @@ pub use marketplace_edit::record_user_marketplace; pub use marketplace_edit::remove_user_marketplace; pub use marketplace_edit::remove_user_marketplace_config; pub use mcp_edit::load_global_mcp_servers; +pub use mcp_ema::McpEmaAuthScope; pub use mcp_ema::McpEmaRegistration; pub use mcp_ema::McpEnterpriseManagedAuthConfig; pub use mcp_ema::McpServerIdpOAuthConfig; diff --git a/codex-rs/config/src/mcp_ema.rs b/codex-rs/config/src/mcp_ema.rs index 327b471d35..da354b23ed 100644 --- a/codex-rs/config/src/mcp_ema.rs +++ b/codex-rs/config/src/mcp_ema.rs @@ -1,13 +1,17 @@ -//! Enterprise MCP configuration and registration provenance. +//! Enterprise MCP configuration and account-scoped credential names. //! //! Only host, user, or managed configuration selects the IdP. Plugin declarations //! and project settings may not redirect the enterprise credential source. use std::io; +use base64::Engine; +use base64::engine::general_purpose::URL_SAFE_NO_PAD; use schemars::JsonSchema; use serde::Deserialize; use serde::Serialize; +use sha2::Digest; +use sha2::Sha256; use codex_features::Feature; @@ -70,6 +74,40 @@ impl McpEmaRegistration { } } +/// Account identity that owns an enterprise OAuth credential. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct McpEmaAuthScope { + user_id: String, + workspace_id: String, +} + +impl McpEmaAuthScope { + pub fn new(user_id: String, workspace_id: String) -> Option { + (!user_id.trim().is_empty() && !workspace_id.trim().is_empty()).then_some(Self { + user_id, + workspace_id, + }) + } +} + +impl McpServerIdpOAuthConfig { + /// Shared by allowed MCP resources, but never by another account or IdP client. + pub fn credential_name(&self, scope: &McpEmaAuthScope) -> String { + let mut digest = Sha256::new(); + digest.update(b"codex-ema-idp-credential-v1\0"); + for part in [ + &scope.user_id, + &scope.workspace_id, + &self.issuer, + &self.client_id, + ] { + digest.update((part.len() as u64).to_be_bytes()); + digest.update(part.as_bytes()); + } + format!("ema-idp:{}", URL_SAFE_NO_PAD.encode(digest.finalize())) + } +} + impl McpEnterpriseManagedAuthConfig { /// Resolves the trusted profile and checks provenance at a config load boundary. pub fn resolve( diff --git a/codex-rs/config/src/mcp_ema_tests.rs b/codex-rs/config/src/mcp_ema_tests.rs index 836bf37197..e196db41b5 100644 --- a/codex-rs/config/src/mcp_ema_tests.rs +++ b/codex-rs/config/src/mcp_ema_tests.rs @@ -1,4 +1,4 @@ -//! Enterprise registration provenance checks. +//! Enterprise registration provenance and credential-boundary checks. use super::*; use crate::AbsolutePathBuf; @@ -215,6 +215,44 @@ fn xaa_opt_in_requires_a_non_project_source() { ); } +#[test] +fn ema_credential_names_bind_user_workspace_issuer_and_client() { + let mut names = std::collections::HashSet::new(); + for user in ["a", "b"] { + for workspace in ["a", "b"] { + for issuer in ["https://one.example", "https://two.example"] { + for client in ["a", "b"] { + let scope = McpEmaAuthScope::new(user.into(), workspace.into()).unwrap(); + let idp = McpServerIdpOAuthConfig { + issuer: issuer.into(), + client_id: client.into(), + }; + assert!(names.insert(idp.credential_name(&scope))); + } + } + } + } + assert_eq!(McpEmaAuthScope::new(" ".into(), "workspace".into()), None); + + let identifiers = [ + "sensitive-user@example.com", + "sensitive-workspace-id", + "https://sensitive-idp.example", + "sensitive-oauth-client", + ]; + let scope = McpEmaAuthScope::new(identifiers[0].into(), identifiers[1].into()).unwrap(); + let name = McpServerIdpOAuthConfig { + issuer: identifiers[2].into(), + client_id: identifiers[3].into(), + } + .credential_name(&scope); + assert!(name.starts_with("ema-idp:")); + for identifier in identifiers { + assert!(!name.contains(identifier)); + assert!(!name.contains(&URL_SAFE_NO_PAD.encode(identifier))); + } +} + #[test] fn ema_rejects_alternate_credentials_and_executor_custody() { for extra in [ diff --git a/codex-rs/config/src/types.rs b/codex-rs/config/src/types.rs index 23fe035f10..e16333c4a1 100644 --- a/codex-rs/config/src/types.rs +++ b/codex-rs/config/src/types.rs @@ -3,6 +3,7 @@ // Note this file should generally be restricted to simple struct/enum // definitions that do not contain business logic. +pub use crate::mcp_ema::McpEmaAuthScope; pub use crate::mcp_ema::McpEnterpriseManagedAuthConfig; pub use crate::mcp_ema::McpServerIdpOAuthConfig; pub use crate::mcp_types::AppToolApproval; diff --git a/codex-rs/rmcp-client/Cargo.toml b/codex-rs/rmcp-client/Cargo.toml index c661aa0ed8..b133ddaf46 100644 --- a/codex-rs/rmcp-client/Cargo.toml +++ b/codex-rs/rmcp-client/Cargo.toml @@ -61,6 +61,7 @@ tokio = { workspace = true, features = [ "io-std", "time", ] } +tokio-util = { workspace = true, features = ["rt"] } tracing = { workspace = true, features = ["log"] } url = { workspace = true } urlencoding = { workspace = true } diff --git a/codex-rs/rmcp-client/src/auth_status.rs b/codex-rs/rmcp-client/src/auth_status.rs index 479ae184f4..9caea78dcd 100644 --- a/codex-rs/rmcp-client/src/auth_status.rs +++ b/codex-rs/rmcp-client/src/auth_status.rs @@ -336,10 +336,14 @@ mod tests { handle: JoinHandle<()>, } - fn test_http_client() -> Arc { - Arc::new(RouteAwareHttpClient::new(HttpClientFactory::new( - OutboundProxyPolicy::ReqwestDefault, - ))) + async fn test_http_client() -> Arc { + let client: Arc = Arc::new(RouteAwareHttpClient::new( + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + )); + crate::oauth::test_support::warm_http_client(client.as_ref()) + .await + .expect("warm production HTTP client before discovery deadlines"); + client } impl Drop for TestServer { @@ -487,7 +491,7 @@ mod tests { /*env_http_headers*/ None, OAuthCredentialsStoreMode::Keyring, AuthKeyringBackendKind::default(), - test_http_client(), + Arc::new(RecordingHttpClient::default()), OAuthDiscoveryTimeout::Requested, StreamableHttpRedirectMode::Legacy, ) @@ -512,7 +516,7 @@ mod tests { )])), OAuthCredentialsStoreMode::Keyring, AuthKeyringBackendKind::default(), - test_http_client(), + Arc::new(RecordingHttpClient::default()), OAuthDiscoveryTimeout::Requested, StreamableHttpRedirectMode::Legacy, ) @@ -567,7 +571,7 @@ mod tests { &url, /*http_headers*/ None, /*env_http_headers*/ None, - test_http_client(), + test_http_client().await, OAuthDiscoveryTimeout::LOCAL, StreamableHttpRedirectMode::Legacy, ) @@ -610,7 +614,7 @@ mod tests { "sensitive-key".to_string(), )])), /*env_http_headers*/ None, - test_http_client(), + test_http_client().await, OAuthDiscoveryTimeout::LOCAL, StreamableHttpRedirectMode::Legacy, ) @@ -632,6 +636,7 @@ mod tests { #[tokio::test] async fn determine_auth_status_preserves_transient_http_errors() { + let _home = crate::oauth::test_support::TempCodexHome::new(); for status in [ StatusCode::REQUEST_TIMEOUT, StatusCode::TOO_EARLY, @@ -653,7 +658,7 @@ mod tests { /*env_http_headers*/ None, OAuthCredentialsStoreMode::File, AuthKeyringBackendKind::default(), - test_http_client(), + test_http_client().await, OAuthDiscoveryTimeout::LOCAL, StreamableHttpRedirectMode::Legacy, ) @@ -685,7 +690,7 @@ mod tests { &server.url, /*http_headers*/ None, /*env_http_headers*/ None, - test_http_client(), + test_http_client().await, OAuthDiscoveryTimeout::LOCAL, StreamableHttpRedirectMode::Legacy, ) @@ -716,7 +721,7 @@ mod tests { &server.url, /*http_headers*/ None, /*env_http_headers*/ None, - test_http_client(), + test_http_client().await, OAuthDiscoveryTimeout::LOCAL, StreamableHttpRedirectMode::Legacy, ) @@ -872,7 +877,7 @@ mod tests { &resource_server.url, /*http_headers*/ None, /*env_http_headers*/ None, - test_http_client(), + test_http_client().await, OAuthDiscoveryTimeout::LOCAL, StreamableHttpRedirectMode::Legacy, ) @@ -899,7 +904,7 @@ mod tests { &server.url, /*http_headers*/ None, /*env_http_headers*/ None, - test_http_client(), + test_http_client().await, OAuthDiscoveryTimeout::LOCAL, StreamableHttpRedirectMode::Legacy, ) @@ -922,7 +927,7 @@ mod tests { &server.url, /*http_headers*/ None, /*env_http_headers*/ None, - test_http_client(), + test_http_client().await, OAuthDiscoveryTimeout::LOCAL, StreamableHttpRedirectMode::Legacy, ) diff --git a/codex-rs/rmcp-client/src/enterprise_oauth_login.rs b/codex-rs/rmcp-client/src/enterprise_oauth_login.rs index 92bf031e76..a956420547 100644 --- a/codex-rs/rmcp-client/src/enterprise_oauth_login.rs +++ b/codex-rs/rmcp-client/src/enterprise_oauth_login.rs @@ -119,6 +119,12 @@ impl EnterpriseOAuthLoginHandle { self.flow.authorization_url() } + /// Capture before waiting; after the handle or wait future is dropped, await + /// this to ensure the callback worker released its listener before rebinding. + pub fn callback_closed(&self) -> impl Future + Send + 'static { + self.flow.callback_closed() + } + pub async fn wait(self) -> Result { let stored = self .flow @@ -187,22 +193,17 @@ impl EnterpriseOAuthCredentials { pub async fn perform_enterprise_oauth_login_return_url( request: EnterpriseOAuthLoginRequest<'_>, ) -> Result { - // Capture before discovery or browser setup, then release the lock while the user signs in. - let generation = { + let observed_generation = { let guard = EnterpriseOAuthCredentialGuard::acquire( request.credential_name, request.issuer, request.keyring_backend_kind, ) .await?; - match guard.generation_file.current() { - Ok(Some(generation)) => generation, - Ok(None) => guard - .generation_file - .replace() - .map_err(|_| anyhow!("failed to initialize enterprise login generation"))?, - Err(_) => bail!("failed to read enterprise login generation"), - } + guard + .generation_file + .current() + .map_err(|_| anyhow!("failed to read enterprise login generation"))? }; let flow = OauthLoginFlow::new( request.credential_name, @@ -232,6 +233,37 @@ pub async fn perform_enterprise_oauth_login_return_url( .with_subscriber(tracing::subscriber::NoSubscriber::default()) .await .map_err(|_| anyhow!("failed to start enterprise IdP authorization"))?; + // Keep the current attempt valid until replacement setup succeeds. A logout + // or another completed setup during this work must still invalidate this start. + let generation = async { + let guard = EnterpriseOAuthCredentialGuard::acquire( + request.credential_name, + request.issuer, + request.keyring_backend_kind, + ) + .await?; + let current_generation = guard + .generation_file + .current() + .map_err(|_| anyhow!("failed to read enterprise login generation"))?; + if current_generation != observed_generation { + bail!("enterprise login changed during replacement setup"); + } + guard + .generation_file + .replace() + .map_err(|_| anyhow!("failed to start enterprise login generation")) + } + .await; + let generation = match generation { + Ok(generation) => generation, + Err(error) => { + let callback_closed = flow.callback_closed(); + drop(flow); + callback_closed.await; + return Err(error); + } + }; Ok(EnterpriseOAuthLoginHandle { flow, keyring_backend: request.keyring_backend_kind, diff --git a/codex-rs/rmcp-client/src/enterprise_oauth_login_tests.rs b/codex-rs/rmcp-client/src/enterprise_oauth_login_tests.rs index b651d76e11..1487fd6446 100644 --- a/codex-rs/rmcp-client/src/enterprise_oauth_login_tests.rs +++ b/codex-rs/rmcp-client/src/enterprise_oauth_login_tests.rs @@ -19,7 +19,6 @@ use keyring::mock::MockCredential; use oauth2::TokenResponse; use pretty_assertions::assert_eq; use serde_json::json; -use sha2::Digest; use tokio::io::AsyncReadExt; use tokio::io::AsyncWriteExt; use tracing_test::traced_test; @@ -239,12 +238,6 @@ async fn enterprise_public_api_storage_and_privacy() -> Result<()> { let server = MockServer::start().await; let issuer = format!("{}/idp", server.uri()); metadata(&server, &issuer).await; - assert!( - login(&issuer, Some(&format!("{}/callback", server.uri()))) - .await - .is_err(), - "an occupied registered callback port must not silently move" - ); let assertion = format!( "{}.{}.signature", URL_SAFE_NO_PAD.encode(r#"{"alg":"ES256"}"#), @@ -284,31 +277,31 @@ async fn enterprise_public_api_storage_and_privacy() -> Result<()> { }) .collect::>(); assert_eq!(token_requests.len(), login_index + 1); - let token_form = url::form_urlencoded::parse(&token_requests[login_index].body) + let token_request = &token_requests[login_index]; + assert!(!token_request.headers.contains_key("authorization")); + let mut token_form = url::form_urlencoded::parse(&token_request.body) .into_owned() .collect::>(); - assert_eq!( - token_form.get("grant_type").map(String::as_str), - Some("authorization_code") - ); - assert!(!token_form.contains_key("resource")); - assert_eq!( - token_form.get("redirect_uri"), - Some( - authorization_query - .get("redirect_uri") - .expect("authorization redirect URI") - ) + let verifier = oauth2::PkceCodeVerifier::new( + token_form.remove("code_verifier").expect("PKCE verifier"), ); assert_eq!( - authorization_query - .get("code_challenge_method") - .map(String::as_str), - Some("S256") + oauth2::PkceCodeChallenge::from_code_verifier_sha256(&verifier).as_str(), + authorization_query["code_challenge"], + ); + assert_eq!(authorization_query["code_challenge_method"], "S256"); + assert_eq!( + token_form, + HashMap::from([ + ("grant_type".to_string(), "authorization_code".to_string()), + ("code".to_string(), SECRET.to_string()), + ("client_id".to_string(), "enterprise-client".to_string()), + ( + "redirect_uri".to_string(), + authorization_query["redirect_uri"].clone() + ), + ]) ); - let verifier = token_form.get("code_verifier").expect("PKCE verifier"); - let challenge = URL_SAFE_NO_PAD.encode(sha2::Sha256::digest(verifier.as_bytes())); - assert_eq!(authorization_query.get("code_challenge"), Some(&challenge)); assert!( keyring .values @@ -381,13 +374,23 @@ async fn enterprise_public_api_storage_and_privacy() -> Result<()> { .all(|value| value.get_secret().is_err()) ); - // Rejected old attempts neither write nor delete a newer grant. + // Failed replacement setup leaves the existing attempt valid. + let retained = complete_login(&issuer).await?; + assert!( + login(&issuer, Some(&format!("{}/callback", server.uri()))) + .await + .is_err(), + "an occupied registered callback port must not silently move" + ); + retained.commit_if(|| async { Some(()) }).await?; + + // A newer start fences a delayed callback even across Codex processes. let old = complete_login(&issuer).await?; complete_login(&issuer) .await? .commit_if(|| async { Some(()) }) .await?; - assert!(old.commit_if(|| async { None::<()> }).await.is_err()); + assert!(old.commit_if(|| async { Some(()) }).await.is_err()); assert!( keyring .values diff --git a/codex-rs/rmcp-client/src/enterprise_oauth_logout_tests.rs b/codex-rs/rmcp-client/src/enterprise_oauth_logout_tests.rs index 974ba72695..3ca54d2120 100644 --- a/codex-rs/rmcp-client/src/enterprise_oauth_logout_tests.rs +++ b/codex-rs/rmcp-client/src/enterprise_oauth_logout_tests.rs @@ -3,9 +3,16 @@ use std::fs; use std::path::PathBuf; +use codex_exec_server::ExecServerError; +use codex_exec_server::HttpRequestParams; +use codex_exec_server::HttpRequestResponse; +use codex_exec_server::HttpResponseBodyStream; +use futures::FutureExt; +use futures::future::BoxFuture; use pretty_assertions::assert_eq; use sha2::Digest; use sha2::Sha256; +use tokio::sync::Notify; use super::*; @@ -46,6 +53,48 @@ async fn logout_invalidates_pending_login_across_processes() -> Result<()> { .mount(&server) .await; + // A logout during provider setup must invalidate the in-flight start. + let paused_client = Arc::new(PausedSetupHttpClient { + inner: Arc::new(RouteAwareHttpClient::new(HttpClientFactory::new( + OutboundProxyPolicy::ReqwestDefault, + ))), + should_pause: AtomicBool::new(true), + entered: Notify::new(), + release: Notify::new(), + }); + let starting_client = Arc::clone(&paused_client); + let starting_issuer = issuer.clone(); + let starting = tokio::spawn(async move { + perform_enterprise_oauth_login_return_url(EnterpriseOAuthLoginRequest { + credential_name: CREDENTIAL_NAME, + issuer: &starting_issuer, + client_id: "enterprise-client", + keyring_backend_kind: AuthKeyringBackendKind::Direct, + callback_port: None, + callback_url: None, + timeout_secs: Some(5), + http_client: starting_client, + redirect_mode: StreamableHttpRedirectMode::Legacy, + }) + .await + }); + tokio::time::timeout( + std::time::Duration::from_secs(5), + paused_client.entered.notified(), + ) + .await?; + delete_enterprise_oauth_tokens(CREDENTIAL_NAME, &issuer, AuthKeyringBackendKind::Direct) + .await?; + paused_client.release.notify_one(); + let error = tokio::time::timeout(std::time::Duration::from_secs(5), starting) + .await?? + .err() + .expect("logout must invalidate a start during provider setup"); + assert_eq!( + error.to_string(), + "enterprise login changed during replacement setup" + ); + // The first logout has no grant to delete. The second deletes the fresh grant // from the previous iteration. Both must invalidate pending callbacks and staged grants. for _ in 0..2 { @@ -127,6 +176,46 @@ async fn logout_invalidates_pending_login_across_processes() -> Result<()> { Ok(()) } +struct PausedSetupHttpClient { + inner: Arc, + should_pause: AtomicBool, + entered: Notify, + release: Notify, +} + +impl PausedSetupHttpClient { + async fn pause_once(&self) { + if self.should_pause.swap(false, Ordering::SeqCst) { + self.entered.notify_one(); + self.release.notified().await; + } + } +} + +impl HttpClient for PausedSetupHttpClient { + fn http_request( + &self, + params: HttpRequestParams, + ) -> BoxFuture<'_, Result> { + async move { + self.pause_once().await; + self.inner.http_request(params).await + } + .boxed() + } + + fn http_request_stream( + &self, + params: HttpRequestParams, + ) -> BoxFuture<'_, Result<(HttpRequestResponse, HttpResponseBodyStream), ExecServerError>> { + async move { + self.pause_once().await; + self.inner.http_request_stream(params).await + } + .boxed() + } +} + fn stored(issuer: &str) -> Result> { crate::stored_oauth_credentials( CREDENTIAL_NAME, diff --git a/codex-rs/rmcp-client/src/oauth/enterprise_generation.rs b/codex-rs/rmcp-client/src/oauth/enterprise_generation.rs index cd227cfbeb..9845f79c98 100644 --- a/codex-rs/rmcp-client/src/oauth/enterprise_generation.rs +++ b/codex-rs/rmcp-client/src/oauth/enterprise_generation.rs @@ -1,4 +1,4 @@ -//! Persistent, non-secret logout generations for staged enterprise logins. +//! Persistent, non-secret generations for staged enterprise logins and logout. //! All reads and writes require the credential lock; missing state never admits an old attempt. use std::fs::File; diff --git a/codex-rs/rmcp-client/src/oauth_http_client_security_tests.rs b/codex-rs/rmcp-client/src/oauth_http_client_security_tests.rs index 6ce9f0a34e..ee3f86de02 100644 --- a/codex-rs/rmcp-client/src/oauth_http_client_security_tests.rs +++ b/codex-rs/rmcp-client/src/oauth_http_client_security_tests.rs @@ -386,6 +386,10 @@ async fn oauth_registration_redirects_never_forward_resource_only_headers() -> R #[tokio::test] async fn same_origin_redirects_preserve_timeout_and_response_body_limits() -> Result<()> { + let http_client: Arc = Arc::new(RouteAwareHttpClient::new( + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + )); + crate::oauth::test_support::warm_http_client(http_client.as_ref()).await?; for oversized_redirect_body in [false, true] { let server = MockServer::start().await; Mock::given(method("GET")) @@ -417,9 +421,7 @@ async fn same_origin_redirects_preserve_timeout_and_response_body_limits() -> Re .await; let adapter = OAuthHttpClientAdapter::new( - Arc::new(RouteAwareHttpClient::new(HttpClientFactory::new( - OutboundProxyPolicy::ReqwestDefault, - ))), + Arc::clone(&http_client), build_default_headers( Some(HashMap::from([( "X-Api-Key".to_string(), diff --git a/codex-rs/rmcp-client/src/perform_oauth_login.rs b/codex-rs/rmcp-client/src/perform_oauth_login.rs index ce7839cb8a..a271ed342e 100644 --- a/codex-rs/rmcp-client/src/perform_oauth_login.rs +++ b/codex-rs/rmcp-client/src/perform_oauth_login.rs @@ -1,4 +1,5 @@ use std::collections::HashMap; +use std::future::Future; use std::net::SocketAddr; use std::string::String; use std::sync::Arc; @@ -282,9 +283,13 @@ fn spawn_callback_server( server: Arc, tx: oneshot::Sender, expected_callback_path: String, + closed: tokio_util::sync::CancellationToken, ) { tokio::task::spawn_blocking(move || { - while let Ok(request) = server.recv() { + // Drop the listener before publishing closure, including if this task unwinds. + let _closed = closed.drop_guard(); + let listener = server; + while let Ok(request) = listener.recv() { let path = request.url().to_string(); match parse_oauth_callback(&path, &expected_callback_path) { CallbackOutcome::Success(OauthCallbackResult { @@ -442,6 +447,7 @@ pub(crate) struct OauthLoginFlow { authorization_server_issuer: Option, rx: oneshot::Receiver, guard: CallbackServerGuard, + callback_closed: tokio_util::sync::CancellationToken, server_name: String, server_url: String, store_mode: OAuthCredentialsStoreMode, @@ -639,7 +645,25 @@ impl OauthLoginFlow { // Port zero asks the OS for a free ephemeral port; the resolved // redirect receives that port after the listener has been bound. let bind_addr = SocketAddr::new(bind_ip, callback_port.unwrap_or(0)); - let server = Arc::new(Server::http(bind_addr).map_err(|err| anyhow!(err))?); + let bind_deadline = tokio::time::Instant::now() + Duration::from_secs(/*secs*/ 1); + let server = loop { + match Server::http(bind_addr) { + Ok(server) => break Arc::new(server), + Err(err) + if is_enterprise_idp + && callback_port.is_some() + && err + .downcast_ref::() + .is_some_and(|err| err.kind() == std::io::ErrorKind::AddrInUse) + && tokio::time::Instant::now() < bind_deadline => + { + // tiny_http drops its accept loop asynchronously after Server::drop. + // Only retry the fixed listener bind, never OAuth requests or callbacks. + tokio::time::sleep(Duration::from_millis(/*millis*/ 10)).await; + } + Err(err) => return Err(anyhow!(err)), + } + }; let guard = CallbackServerGuard { server: Arc::clone(&server), }; @@ -701,8 +725,6 @@ impl OauthLoginFlow { .await? }; let callback_path = callback_path_from_redirect_uri(&redirect_uri)?; - let (tx, rx) = oneshot::channel(); - spawn_callback_server(server, tx, callback_path); let auth_url = append_query_param( &oauth_state.get_authorization_url().await?, "resource", @@ -710,6 +732,11 @@ impl OauthLoginFlow { ); let timeout_secs = timeout_secs.unwrap_or(DEFAULT_OAUTH_TIMEOUT_SECS).max(1); let timeout = Duration::from_secs(timeout_secs as u64); + // No fallible work or suspension after spawning: every callback worker + // must be owned by a returned flow, including during constructor cancellation. + let (tx, rx) = oneshot::channel(); + let callback_closed = tokio_util::sync::CancellationToken::new(); + spawn_callback_server(server, tx, callback_path, callback_closed.clone()); Ok(Self { auth_url, @@ -718,6 +745,7 @@ impl OauthLoginFlow { authorization_server_issuer, rx, guard, + callback_closed, server_name: server_name.to_string(), server_url: server_url.to_string(), store_mode, @@ -731,6 +759,10 @@ impl OauthLoginFlow { self.auth_url.clone() } + pub(crate) fn callback_closed(&self) -> impl Future + Send + 'static { + self.callback_closed.clone().cancelled_owned() + } + async fn finish(self, emit_browser_url: bool) -> Result<()> { let store_mode = self.store_mode; let keyring_backend_kind = self.keyring_backend_kind; diff --git a/sdk/python/src/openai_codex/generated/v2_all.py b/sdk/python/src/openai_codex/generated/v2_all.py index 4d7d54aeb4..f9ee2be07f 100644 --- a/sdk/python/src/openai_codex/generated/v2_all.py +++ b/sdk/python/src/openai_codex/generated/v2_all.py @@ -2584,6 +2584,13 @@ class McpServerOauthLoginCompletedNotification(BaseModel): populate_by_name=True, ) error: str | None = None + login_id: Annotated[ + str | None, + Field( + alias="loginId", + description="Identifies the explicit login attempt. Older servers omit this field.", + ), + ] = None name: str success: bool thread_id: Annotated[str | None, Field(alias="threadId")] = None @@ -2611,6 +2618,13 @@ class McpServerOauthLoginResponse(BaseModel): populate_by_name=True, ) authorization_url: Annotated[str, Field(alias="authorizationUrl")] + login_id: Annotated[ + str | None, + Field( + alias="loginId", + description="Identifies this login attempt across the response and completion notification. Older servers omit this field; current servers always return it.", + ), + ] = None class McpServerRefreshResponse(BaseModel):