diff --git a/codex-rs/codex-mcp/src/auth_changes.rs b/codex-rs/codex-mcp/src/auth_changes.rs new file mode 100644 index 0000000000..78498bd071 --- /dev/null +++ b/codex-rs/codex-mcp/src/auth_changes.rs @@ -0,0 +1,62 @@ +//! Forwards auth invalidations without credentials. The managed client owns the watcher. + +use std::sync::Arc; +use std::time::Duration; + +use anyhow::Result; +use codex_login::AuthChangeState; +use codex_rmcp_client::RmcpClient; +use rmcp::model::ServerCapabilities; +use serde_json::json; +use tokio::sync::watch; +use tokio_util::task::AbortOnDropHandle; + +pub(crate) const CAPABILITY: &str = "codex/auth-change"; +const NOTIFICATION: &str = "notifications/codex/authChanged"; +const SEND_TIMEOUT: Duration = Duration::from_secs(5); + +pub(crate) async fn start( + client: Arc, + capabilities: &ServerCapabilities, + changes: Option>, +) -> Result>>> { + let Some(mut changes) = changes.filter(|_| { + capabilities + .experimental + .as_ref() + .is_some_and(|capabilities| capabilities.contains_key(CAPABILITY)) + }) else { + return Ok(None); + }; + + notify(&client, &mut changes).await?; + let task = tokio::spawn(async move { + while changes.changed().await.is_ok() { + if notify(&client, &mut changes).await.is_err() { + tracing::warn!("MCP auth invalidation delivery failed; closing connection"); + client.shutdown().await; + break; + } + } + }); + Ok(Some(Arc::new(AbortOnDropHandle::new(task)))) +} + +async fn notify(client: &RmcpClient, changes: &mut watch::Receiver) -> Result<()> { + let state = *changes.borrow_and_update(); + tokio::time::timeout( + SEND_TIMEOUT, + client.send_custom_notification( + NOTIFICATION, + Some(json!({ + "generation": state.generation, + "ownerGeneration": state.owner_generation, + })), + ), + ) + .await? +} + +#[cfg(test)] +#[path = "auth_changes_tests.rs"] +mod tests; diff --git a/codex-rs/codex-mcp/src/auth_changes_tests.rs b/codex-rs/codex-mcp/src/auth_changes_tests.rs new file mode 100644 index 0000000000..36de21ba0c --- /dev/null +++ b/codex-rs/codex-mcp/src/auth_changes_tests.rs @@ -0,0 +1,104 @@ +use super::*; +use codex_rmcp_client::InProcessTransportFactory; +use futures::FutureExt; +use futures::future::BoxFuture; +use pretty_assertions::assert_eq; +use rmcp::ServiceExt; +use rmcp::model::ClientCapabilities; +use rmcp::model::CustomNotification; +use rmcp::model::Implementation; +use rmcp::model::InitializeRequestParams; +use tokio::sync::mpsc; +use tokio::time::timeout; + +#[derive(Clone)] +struct NotificationServer(mpsc::Sender); + +impl rmcp::ServerHandler for NotificationServer { + async fn on_custom_notification( + &self, + notification: CustomNotification, + _context: rmcp::service::NotificationContext, + ) { + self.0.send(json!(notification)).await.unwrap(); + } +} + +impl InProcessTransportFactory for NotificationServer { + fn open(&self) -> BoxFuture<'static, std::io::Result> { + let server = self.clone(); + async move { + let (client, transport) = tokio::io::duplex(/*max_buf_size*/ 4096); + tokio::spawn(async move { + let service = server.serve(transport).await.unwrap(); + service.waiting().await.unwrap(); + }); + Ok(client) + } + .boxed() + } +} + +#[tokio::test] +async fn auth_notifications_require_opt_in_and_follow_client_lifetime() -> Result<()> { + let (notifications, mut received) = mpsc::channel(/*buffer*/ 8); + let client = Arc::new( + RmcpClient::new_in_process_client(Arc::new(NotificationServer(notifications))).await?, + ); + client + .initialize( + InitializeRequestParams::new( + ClientCapabilities::default(), + Implementation::new("test", "1"), + ), + Some(SEND_TIMEOUT), + Box::new(|_, _| async { anyhow::bail!("unexpected elicitation") }.boxed()), + ) + .await?; + let (changes, receiver) = watch::channel(AuthChangeState::default()); + let mut capabilities = ServerCapabilities::default(); + assert!( + start(Arc::clone(&client), &capabilities, Some(receiver.clone())) + .await? + .is_none() + ); + capabilities.experimental = Some([(CAPABILITY.to_string(), Default::default())].into()); + assert!( + start(Arc::clone(&client), &capabilities, /*changes*/ None) + .await? + .is_none() + ); + assert_eq!(received.try_recv(), Err(mpsc::error::TryRecvError::Empty)); + let watcher = start(Arc::clone(&client), &capabilities, Some(receiver)) + .await? + .unwrap(); + assert_eq!( + timeout(SEND_TIMEOUT, received.recv()).await?, + Some( + json!({"method": NOTIFICATION, "params": {"_meta": {}, "generation": 0, "ownerGeneration": 0}}) + ), + ); + changes.send_modify(|state| state.generation += 1); + assert_eq!( + timeout(SEND_TIMEOUT, received.recv()).await?, + Some( + json!({"method": NOTIFICATION, "params": {"_meta": {}, "generation": 1, "ownerGeneration": 0}}) + ), + ); + for _ in 0..2 { + changes.send_modify(|state| { + state.generation += 1; + state.owner_generation += 1; + }); + } + assert_eq!( + timeout(SEND_TIMEOUT, received.recv()).await?, + Some( + json!({"method": NOTIFICATION, "params": {"_meta": {}, "generation": 3, "ownerGeneration": 2}}) + ), + ); + drop(watcher); + timeout(SEND_TIMEOUT, changes.closed()).await?; + client.shutdown().await; + Ok(()) +} diff --git a/codex-rs/codex-mcp/src/binding_tests.rs b/codex-rs/codex-mcp/src/binding_tests.rs index 9ba610159f..be36e56a98 100644 --- a/codex-rs/codex-mcp/src/binding_tests.rs +++ b/codex-rs/codex-mcp/src/binding_tests.rs @@ -79,6 +79,7 @@ async fn test_step( ); let tool_catalog = Arc::new(ClientToolCatalog::new(vec![tool.clone()])); let managed_client = Arc::new(ManagedClient { + _auth_change_notifications: None, client: Arc::clone(&client), server_info: McpServerInfo { name: label.to_string(), diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index adbac51863..4ac5c1e689 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -548,6 +548,15 @@ impl McpConnectionSet { runtime_auth_provider, client_elicitation_capability.clone(), client_mcp_extensions.clone(), + auth_manager + .as_ref() + .filter(|_| { + matches!( + &configured_config.transport, + McpServerTransportConfig::Stdio { .. } + ) + }) + .map(|manager| manager.auth_change_state_receiver()), protocol_mode, catalog_item_limit, ); diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 45deec37ea..bae8d5207f 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -451,6 +451,7 @@ async fn legacy_tool_catalog_does_not_follow_pagination_cursor() -> anyhow::Resu async fn create_test_managed_client(tools: Vec) -> ManagedClient { ManagedClient { + _auth_change_notifications: None, client: Arc::new( RmcpClient::new_in_process_client(Arc::new(TestInProcessTransportFactory)) .await @@ -754,6 +755,7 @@ pub(crate) async fn create_test_manager_with_ready_apps_client( .await?; let managed_client = ManagedClient { + _auth_change_notifications: None, client, server_info: create_test_server_info("Codex Apps"), tool_catalog: Arc::new(ClientToolCatalog::new(vec![tool])), @@ -5078,6 +5080,7 @@ async fn reconciliation_reuses_connection_without_relisting_regular_tools() -> a ) .await?; let managed_client = ManagedClient { + _auth_change_notifications: None, client, server_info: create_test_server_info("Mutable tools"), tool_catalog: Arc::new(ClientToolCatalog::new(initial_tools)), diff --git a/codex-rs/codex-mcp/src/lib.rs b/codex-rs/codex-mcp/src/lib.rs index bbf52e638e..1d8f90105d 100644 --- a/codex-rs/codex-mcp/src/lib.rs +++ b/codex-rs/codex-mcp/src/lib.rs @@ -96,6 +96,7 @@ pub use mcp::McpPermissionPromptAutoApproveContext; pub use mcp::mcp_permission_prompt_is_auto_approved; pub use mcp::qualified_mcp_tool_name_prefix; +mod auth_changes; pub(crate) mod auth_elicitation; mod binding; pub(crate) mod binding_clients; diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index 0034d36179..90b30e5caf 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -53,6 +53,7 @@ use codex_config::types::OAuthCredentialsStoreMode; use codex_connectors::ConnectorRuntimeContext; use codex_connectors::ConnectorRuntimeFetchSource; use codex_exec_server::Environment; +use codex_login::AuthChangeState; use codex_protocol::mcp::ClientMcpExtensions; use codex_protocol::mcp::McpServerInfo; use codex_protocol::protocol::Event; @@ -79,8 +80,10 @@ use rmcp::model::InitializeRequestParams; use rmcp::model::ProtocolVersion; use rmcp::model::ServerPeerInfo; use rmcp::model::Tool as RmcpTool; +use tokio::sync::watch; use tokio::time::Instant as TokioInstant; use tokio_util::sync::CancellationToken; +use tokio_util::task::AbortOnDropHandle; use tracing::Instrument; use tracing::instrument; use tracing::warn; @@ -112,6 +115,7 @@ const UNTRUSTED_CONNECTOR_META_KEYS: &[&str] = &[ #[derive(Clone)] pub(crate) struct ManagedClient { + pub(crate) _auth_change_notifications: Option>>, pub(crate) client: Arc, pub(crate) server_info: McpServerInfo, pub(crate) tool_catalog: Arc, @@ -293,6 +297,7 @@ struct ManagedClientStartup { runtime_auth_provider: Option, client_elicitation_capability: ElicitationCapability, client_mcp_extensions: ClientMcpExtensions, + auth_changes: Option>, protocol_mode: McpProtocolMode, catalog_item_limit: usize, cancel_token: CancellationToken, @@ -316,6 +321,7 @@ impl ManagedClientStartup { runtime_auth_provider, client_elicitation_capability, client_mcp_extensions, + auth_changes, protocol_mode, catalog_item_limit, cancel_token, @@ -373,6 +379,7 @@ impl ManagedClientStartup { tool_catalog_fetch_ticket, client_elicitation_capability, client_mcp_extensions, + auth_changes, catalog_item_limit, }, ) @@ -441,6 +448,7 @@ impl AsyncManagedClient { runtime_auth_provider: Option, client_elicitation_capability: ElicitationCapability, client_mcp_extensions: ClientMcpExtensions, + auth_changes: Option>, protocol_mode: McpProtocolMode, catalog_item_limit: usize, ) -> Self { @@ -470,6 +478,7 @@ impl AsyncManagedClient { runtime_auth_provider, client_elicitation_capability, client_mcp_extensions, + auth_changes, protocol_mode, catalog_item_limit, cancel_token: cancel_token.clone(), @@ -895,12 +904,23 @@ async fn start_server_task( tool_catalog_fetch_ticket, client_elicitation_capability, client_mcp_extensions, + auth_changes, catalog_item_limit, } = params; let send_elicitation = elicitation_requests.make_sender(server_name.clone(), tx_event, &client_mcp_extensions); - let params = + let mut params = mcp_initialize_request_params(client_elicitation_capability, client_mcp_extensions); + if auth_changes.is_some() { + params + .capabilities + .experimental + .get_or_insert_default() + .insert( + crate::auth_changes::CAPABILITY.to_string(), + Default::default(), + ); + } let started_at = Instant::now(); let initialize_result = client @@ -914,6 +934,14 @@ async fn start_server_task( ); let initialize_result = initialize_result.map_err(StartupOutcomeError::from)?; + let auth_change_notifications = crate::auth_changes::start( + Arc::clone(&client), + &initialize_result.capabilities, + auth_changes, + ) + .await + .map_err(StartupOutcomeError::from)?; + let server_disables_tool_catalog_cache = initialize_result .capabilities .experimental @@ -974,6 +1002,7 @@ async fn start_server_task( ); } let managed = ManagedClient { + _auth_change_notifications: auth_change_notifications, client: Arc::clone(&client), server_info, tool_catalog: Arc::new(ClientToolCatalog::new(client_tools)), @@ -1072,6 +1101,7 @@ struct StartServerTaskParams { tool_catalog_fetch_ticket: Option, client_elicitation_capability: ElicitationCapability, client_mcp_extensions: ClientMcpExtensions, + auth_changes: Option>, catalog_item_limit: usize, } diff --git a/codex-rs/login/src/auth/change_state.rs b/codex-rs/login/src/auth/change_state.rs new file mode 100644 index 0000000000..27d13992bd --- /dev/null +++ b/codex-rs/login/src/auth/change_state.rs @@ -0,0 +1,30 @@ +//! Tracks credential changes separately from ownership changes, before notifications coalesce. + +use super::CodexAuth; + +/// Opaque revisions local to one auth manager. Consumers must reset on reconnect. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct AuthChangeState { + /// Advances whenever cached credentials change. + pub generation: u64, + /// Advances on login, logout, or changes of user, workspace, or auth mode. + /// Credential changes with incomplete owner identity also advance this revision. + pub owner_generation: u64, +} + +pub(super) fn same_owner(previous: Option<&CodexAuth>, current: Option<&CodexAuth>) -> bool { + let (Some(previous), Some(current)) = (previous, current) else { + return false; + }; + if previous.api_auth_mode() != current.api_auth_mode() { + return false; + } + let (Some(user), Some(workspace)) = (previous.get_chatgpt_user_id(), previous.get_account_id()) + else { + return false; + }; + !user.trim().is_empty() + && !workspace.trim().is_empty() + && current.get_chatgpt_user_id().as_ref() == Some(&user) + && current.get_account_id().as_ref() == Some(&workspace) +} diff --git a/codex-rs/login/src/auth/change_state_tests.rs b/codex-rs/login/src/auth/change_state_tests.rs new file mode 100644 index 0000000000..4022fbbfe8 --- /dev/null +++ b/codex-rs/login/src/auth/change_state_tests.rs @@ -0,0 +1,76 @@ +use super::*; +use base64::Engine; +use pretty_assertions::assert_eq; +use serde_json::json; + +fn chatgpt_auth(user: &str, workspace: &str, token: &str) -> CodexAuth { + let claims = json!({ + "jti": token, + "https://api.openai.com/auth": {"chatgpt_user_id": user}, + }); + let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(claims.to_string()); + CodexAuth::from_external_chatgpt_tokens( + &format!("header.{payload}.signature"), + workspace, + /*chatgpt_plan_type*/ None, + ) + .unwrap() +} + +#[test] +fn auth_owner_generation_distinguishes_refreshes_from_identity_changes() { + let manager = AuthManager::from_optional_auth_for_testing(/*auth*/ None); + let mut changes = manager.auth_change_state_receiver(); + let mut legacy_changes = manager.auth_change_receiver(); + let initial = chatgpt_auth("user-a", "workspace-a", "token-1"); + for (auth, generation, owner_generation) in [ + (Some(initial.clone()), 1, 1), + (Some(initial), 1, 1), + (Some(chatgpt_auth("user-a", "workspace-a", "token-2")), 2, 1), + (Some(chatgpt_auth("user-b", "workspace-a", "token-3")), 3, 2), + (Some(chatgpt_auth("user-b", "workspace-b", "token-4")), 4, 3), + (Some(chatgpt_auth("", "workspace-b", "token-5")), 5, 4), + (Some(chatgpt_auth("", "workspace-b", "token-6")), 6, 5), + (Some(CodexAuth::from_api_key("key-1")), 7, 6), + (Some(CodexAuth::from_api_key("key-2")), 8, 7), + ] { + manager.set_cached_auth(auth); + assert_eq!( + *changes.borrow_and_update(), + AuthChangeState { + generation, + owner_generation, + }, + ); + assert_eq!(*legacy_changes.borrow_and_update(), generation); + } +} + +#[tokio::test] +async fn auth_owner_generation_preserves_logout_and_switches_when_coalesced() { + let home = tempfile::tempdir().unwrap(); + let initial = chatgpt_auth("user-a", "workspace-a", "token-1"); + let manager = + AuthManager::from_auth_for_testing_with_home(initial.clone(), home.path().to_path_buf()); + let mut changes = manager.auth_change_state_receiver(); + + manager.logout().await.unwrap(); + manager.set_cached_auth(Some(initial.clone())); + assert_eq!( + *changes.borrow_and_update(), + AuthChangeState { + generation: 2, + owner_generation: 2, + }, + ); + + manager.set_cached_auth(Some(chatgpt_auth("user-b", "workspace-a", "token-2"))); + manager.set_cached_auth(Some(initial)); + assert_eq!( + *changes.borrow_and_update(), + AuthChangeState { + generation: 4, + owner_generation: 4, + }, + ); +} diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 2949b7039a..18b94ef930 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -36,6 +36,8 @@ use super::agent_identity::record_needs_task_registration; use super::agent_identity::register_managed_chatgpt_agent_identity; use super::agent_identity::require_agent_identity_authapi_base_url; use super::agent_identity::verified_record_from_jwt; +use super::change_state::AuthChangeState; +use super::change_state::same_owner; use super::external_bearer::BearerTokenRefresher; use super::revoke::revoke_auth_tokens; use super::workload_identity::WorkloadIdentityExternalAuth; @@ -2034,6 +2036,7 @@ pub struct AuthManager { codex_home: PathBuf, inner: RwLock, auth_change_tx: watch::Sender, + auth_change_state_tx: watch::Sender, enable_codex_api_key_env: bool, auth_credentials_store_mode: AuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, @@ -2170,6 +2173,7 @@ impl AuthManager { permanent_refresh_failure: None, }), auth_change_tx, + auth_change_state_tx: watch::channel(AuthChangeState::default()).0, enable_codex_api_key_env, auth_credentials_store_mode, keyring_backend_kind, @@ -2204,6 +2208,7 @@ impl AuthManager { codex_home: PathBuf::from("non-existent"), inner: RwLock::new(cached), auth_change_tx, + auth_change_state_tx: watch::channel(AuthChangeState::default()).0, enable_codex_api_key_env: false, auth_credentials_store_mode: AuthCredentialsStoreMode::File, keyring_backend_kind: AuthKeyringBackendKind::default(), @@ -2232,6 +2237,7 @@ impl AuthManager { codex_home, inner: RwLock::new(cached), auth_change_tx, + auth_change_state_tx: watch::channel(AuthChangeState::default()).0, enable_codex_api_key_env: false, auth_credentials_store_mode: AuthCredentialsStoreMode::File, keyring_backend_kind: AuthKeyringBackendKind::default(), @@ -2264,6 +2270,7 @@ impl AuthManager { codex_home: PathBuf::from("non-existent"), inner: RwLock::new(cached), auth_change_tx, + auth_change_state_tx: watch::channel(AuthChangeState::default()).0, enable_codex_api_key_env: false, auth_credentials_store_mode: AuthCredentialsStoreMode::File, keyring_backend_kind: AuthKeyringBackendKind::default(), @@ -2294,6 +2301,7 @@ impl AuthManager { permanent_refresh_failure: None, }), auth_change_tx, + auth_change_state_tx: watch::channel(AuthChangeState::default()).0, enable_codex_api_key_env: false, auth_credentials_store_mode: AuthCredentialsStoreMode::File, keyring_backend_kind: AuthKeyringBackendKind::default(), @@ -2328,6 +2336,11 @@ impl AuthManager { self.auth_change_tx.subscribe() } + /// Subscribes to credential and owner revisions published together, including when changes coalesce. + pub fn auth_change_state_receiver(&self) -> watch::Receiver { + self.auth_change_state_tx.subscribe() + } + pub fn refresh_failure_for_auth(&self, auth: &CodexAuth) -> Option { self.inner.read().ok().and_then(|cached| { cached @@ -2575,12 +2588,20 @@ impl AuthManager { let changed = !AuthManager::auths_equal(previous, new_auth.as_ref()); let auth_changed_for_refresh = !Self::auths_equal_for_refresh(previous, new_auth.as_ref()); + let owner_changed = + auth_changed_for_refresh && !same_owner(previous, new_auth.as_ref()); if auth_changed_for_refresh { guard.permanent_refresh_failure = None; } tracing::info!("Reloaded auth, changed: {changed}"); guard.auth = new_auth; if auth_changed_for_refresh { + self.auth_change_state_tx.send_modify(|state| { + state.generation += 1; + if owner_changed { + state.owner_generation += 1; + } + }); self.auth_change_tx.send_modify(|revision| *revision += 1); } changed @@ -3045,3 +3066,7 @@ fn auth_config_from(config: &impl AuthManagerConfig) -> AuthConfig { #[cfg(test)] #[path = "auth_tests.rs"] mod tests; + +#[cfg(test)] +#[path = "change_state_tests.rs"] +mod change_state_tests; diff --git a/codex-rs/login/src/auth/mod.rs b/codex-rs/login/src/auth/mod.rs index e4898f0976..226255f39e 100644 --- a/codex-rs/login/src/auth/mod.rs +++ b/codex-rs/login/src/auth/mod.rs @@ -3,6 +3,7 @@ mod agent_identity; mod auth_headers; mod bedrock_access_keys; mod bedrock_api_key; +mod change_state; pub mod default_client; pub mod error; mod personal_access_token; @@ -19,6 +20,7 @@ pub use bedrock_access_keys::BedrockAccessKeysAuth; pub use bedrock_access_keys::login_with_bedrock_access_keys; pub use bedrock_api_key::BedrockApiKeyAuth; pub use bedrock_api_key::login_with_bedrock_api_key; +pub use change_state::AuthChangeState; pub use error::RefreshTokenFailedError; pub use error::RefreshTokenFailedReason; pub use manager::*; diff --git a/codex-rs/login/src/lib.rs b/codex-rs/login/src/lib.rs index 33659722fd..e33e587121 100644 --- a/codex-rs/login/src/lib.rs +++ b/codex-rs/login/src/lib.rs @@ -27,6 +27,7 @@ pub use success_page::LoginSuccessPage; pub use success_page::LoginSuccessPageBrand; pub use auth::AgentIdentityAuthPolicy; +pub use auth::AuthChangeState; pub use auth::AuthConfig; pub use auth::AuthDotJson; pub use auth::AuthHeaders;