From 53691bad9878d6e26ba546bc5959d2d49eaecf86 Mon Sep 17 00:00:00 2001 From: Evan Fannin Date: Mon, 31 Aug 2026 18:44:27 +0000 Subject: [PATCH] Keep MCP event subscriptions alive after task unloading (#41899) ## Why Event subscriptions may need to outlive the MCP runtime associated with a task. ## What changed - Add `McpEventStreamOpener`, which retains the hosted Apps MCP connection settings and opens each event stream with an independent client. - Keep subscription cancellation tied to account changes and removal of the task's event server, including when cancellation ownership is forwarded across runtime replacement. - Expose the opener and cancellation forwarding through `McpResourceClient`. ## Testing Add an HTTP transport test that shuts down the runtime, refreshes credentials for the same account, and verifies that a new event stream connects with the current authorization and receives notifications. GitOrigin-RevId: e102aa382e63121661ff0567f7523515e459922d --- codex-rs/Cargo.lock | 5 + codex-rs/codex-mcp/src/connection_manager.rs | 26 +- .../codex-mcp/src/connection_manager_tests.rs | 1 + codex-rs/codex-mcp/src/event_stream.rs | 156 ++++++++++++ codex-rs/codex-mcp/src/lib.rs | 2 + codex-rs/codex-mcp/src/resource_client.rs | 23 +- codex-rs/codex-mcp/src/rmcp_client.rs | 4 +- codex-rs/codex-mcp/src/runtime.rs | 72 +++++- codex-rs/ext/mcp/Cargo.toml | 5 + codex-rs/ext/mcp/src/event_stream_tests.rs | 238 ++++++++++++++++++ codex-rs/ext/mcp/src/lib.rs | 3 + 11 files changed, 520 insertions(+), 15 deletions(-) create mode 100644 codex-rs/codex-mcp/src/event_stream.rs create mode 100644 codex-rs/ext/mcp/src/event_stream_tests.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 2c22dde85b..1d54ec0619 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3646,12 +3646,15 @@ dependencies = [ name = "codex-mcp-extension" version = "0.0.0" dependencies = [ + "anyhow", + "axum", "codex-config", "codex-connectors", "codex-connectors-extension", "codex-core", "codex-core-plugins", "codex-exec-server", + "codex-exec-server-test-support", "codex-extension-api", "codex-features", "codex-login", @@ -3660,11 +3663,13 @@ dependencies = [ "codex-protocol", "codex-utils-absolute-path", "codex-utils-path-uri", + "futures", "pretty_assertions", "serde_json", "tempfile", "thiserror 2.0.18", "tokio", + "tokio-util", "tracing", ] diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index cd9a4f378f..a755947b60 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -33,6 +33,7 @@ use crate::binding::call_tool_result_from_rmcp; use crate::catalog::McpServerSource; use crate::elicitation::ElicitationRequestManager; use crate::elicitation::ElicitationRequestRouter; +use crate::event_stream::EventStreamConnectionSettings; use crate::mcp::CODEX_APPS_MCP_SERVER_NAME; use crate::mcp::ToolPluginProvenance; use crate::pagination::MAX_CODEX_APPS_TOOL_CATALOG_ITEMS; @@ -184,6 +185,7 @@ impl McpServerView { /// A published view over a set of running MCP server connections. pub(crate) struct McpConnectionSet { servers: HashMap, + pub(crate) event_stream_connection: Option>, disabled_servers: Vec, protocol_mode: crate::McpProtocolMode, required_servers: Vec, @@ -237,6 +239,7 @@ impl McpConnectionSet { let tool_plugin_provenance = crate::mcp::tool_plugin_provenance(&config); let auth = auth.as_ref(); let mut servers = HashMap::new(); + let mut event_stream_connection = None; let disabled_servers = mcp_servers .iter() .filter(|(_, server)| !server.enabled()) @@ -276,9 +279,12 @@ impl McpConnectionSet { let static_chatgpt_auth_provider = auth .filter(|auth| auth.uses_codex_backend()) .map(codex_model_provider::auth_provider_from_auth); - let codex_apps_auth_provider = auth_manager.and_then(|auth_manager| { + let codex_apps_auth_provider = auth_manager.as_ref().and_then(|auth_manager| { auth.filter(|auth| auth.uses_codex_backend()).map(|auth| { - codex_model_provider::auth_provider_from_auth_manager(auth_manager, auth) + codex_model_provider::auth_provider_from_auth_manager( + Arc::clone(auth_manager), + auth, + ) }) }); for (server_name, server) in mcp_servers @@ -351,6 +357,20 @@ impl McpConnectionSet { } else { chatgpt_auth_provider_for_server(&server, chatgpt_auth_provider) }; + if is_host_owned_codex_apps { + event_stream_connection = Some(Arc::new(EventStreamConnectionSettings { + server: server.clone(), + store_mode, + keyring_backend_kind, + runtime_context: runtime_context.clone(), + resolved_environment: resolved_environment.clone(), + auth_provider: runtime_auth_provider.clone(), + auth_manager: auth_manager.clone(), + auth: auth.cloned(), + protocol_mode, + client_mcp_extensions: client_mcp_extensions.clone(), + })); + } let connection_identity = McpServerConnectionIdentity::new( &server_name, &server, @@ -682,6 +702,7 @@ impl McpConnectionSet { } let manager = Self { servers, + event_stream_connection, disabled_servers, protocol_mode, required_servers, @@ -743,6 +764,7 @@ impl McpConnectionSet { pub fn empty(prefix_mcp_tool_names: bool) -> Self { Self { servers: HashMap::new(), + event_stream_connection: None, disabled_servers: Vec::new(), protocol_mode: crate::McpProtocolMode::Legacy, required_servers: Vec::new(), diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 47aa905f6b..ecc006097e 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -102,6 +102,7 @@ impl McpConnectionSet { ) -> Self { Self { servers: HashMap::new(), + event_stream_connection: None, disabled_servers: Vec::new(), protocol_mode: crate::McpProtocolMode::Legacy, required_servers: Vec::new(), diff --git a/codex-rs/codex-mcp/src/event_stream.rs b/codex-rs/codex-mcp/src/event_stream.rs new file mode 100644 index 0000000000..7ad58dd626 --- /dev/null +++ b/codex-rs/codex-mcp/src/event_stream.rs @@ -0,0 +1,156 @@ +//! Opens MCP event streams with clients that remain connected after task unloading. + +use std::sync::Arc; + +use anyhow::Result; +use anyhow::bail; +use codex_api::SharedAuthProvider; +use codex_config::types::AuthKeyringBackendKind; +use codex_config::types::OAuthCredentialsStoreMode; +use codex_exec_server::Environment; +use codex_login::AuthManager; +use codex_login::CodexAuth; +use codex_protocol::mcp::ClientMcpExtensions; +use codex_rmcp_client::ElicitationResponse; +use rmcp::model::ElicitationAction; +use rmcp::model::ElicitationCapability; +use serde_json::Map; +use serde_json::Value; +use tokio::sync::watch; + +use crate::CODEX_APPS_MCP_SERVER_NAME; +use crate::EffectiveMcpServer; +use crate::McpEventStream; +use crate::McpProtocolMode; +use crate::McpRuntimeContext; +use crate::rmcp_client::DEFAULT_STARTUP_TIMEOUT; +use crate::rmcp_client::make_rmcp_client; +use crate::rmcp_client::mcp_initialize_request_params; + +pub(crate) struct EventStreamConnectionSettings { + pub server: EffectiveMcpServer, + pub store_mode: OAuthCredentialsStoreMode, + pub keyring_backend_kind: AuthKeyringBackendKind, + pub runtime_context: McpRuntimeContext, + pub resolved_environment: std::result::Result>, String>, + pub auth_provider: Option, + pub auth_manager: Option>, + pub auth: Option, + pub protocol_mode: McpProtocolMode, + pub client_mcp_extensions: ClientMcpExtensions, +} + +/// Opens each event stream with its own MCP client. +/// Owners watch `wait_for_access_change` to cancel streams when access changes. +#[derive(Clone)] +pub struct McpEventStreamOpener { + pub(crate) connection: Arc, + pub(crate) cancellation_receiver: watch::Receiver<()>, + pub(crate) cancel_event_streams_on_server_removal: watch::Sender<()>, +} + +impl McpEventStreamOpener { + /// Retains cancellation for this task's subscriptions across MCP runtime replacement. + pub fn event_stream_cancellation_sender(&self) -> watch::Sender<()> { + self.cancel_event_streams_on_server_removal.clone() + } + + /// Creates an MCP client and opens an event stream. + pub async fn open( + &self, + event_name: &str, + arguments: &Value, + request_meta: Option<&Map>, + ) -> Result { + tokio::select! { + biased; + () = self.wait_for_access_change() => bail!("event subscription access changed"), + result = async { + let connection = &self.connection; + if let Some(manager) = &connection.auth_manager { + let auth = manager.auth().await; + if !self.matches_auth(auth.as_ref()) { + bail!("event subscription account changed"); + } + } + + let startup_timeout = connection.server.config().startup_timeout_sec + .unwrap_or(DEFAULT_STARTUP_TIMEOUT); + let client = Arc::new(tokio::time::timeout(startup_timeout, make_rmcp_client( + CODEX_APPS_MCP_SERVER_NAME, + connection.server.clone(), + connection.store_mode, + connection.keyring_backend_kind, + connection.runtime_context.clone(), + connection.resolved_environment.clone(), + connection.auth_provider.clone(), + connection.protocol_mode, + )).await??); + client.initialize( + mcp_initialize_request_params( + ElicitationCapability::default(), + connection.client_mcp_extensions.clone(), + ), + Some(startup_timeout), + Box::new(|_, _| Box::pin(async { + Ok(ElicitationResponse { + action: ElicitationAction::Decline, + content: None, + meta: None, + }) + })), + ).await?; + McpEventStream::open( + client, + self.cancellation_receiver.clone(), + event_name, + arguments, + request_meta, + ).await + } => result, + } + } + + /// Waits for an account change or removal of the event server from the task. + pub async fn wait_for_access_change(&self) { + let mut cancellation_receiver = self.cancellation_receiver.clone(); + let auth_change = async { + let Some(manager) = &self.connection.auth_manager else { + return std::future::pending().await; + }; + let mut changes = manager.auth_change_receiver(); + loop { + if !self.matches_auth(manager.auth_cached().as_ref()) { + return; + } + if changes.changed().await.is_err() { + return; + } + } + }; + tokio::select! { + Ok(()) = cancellation_receiver.changed() => {}, + () = auth_change => {}, + } + } + + fn matches_auth(&self, current: Option<&CodexAuth>) -> bool { + match (self.connection.auth.as_ref(), current) { + (Some(CodexAuth::AgentIdentity(expected)), Some(CodexAuth::AgentIdentity(current))) => { + expected.record() == current.record() + } + (Some(CodexAuth::AgentIdentity(_)), _) | (_, Some(CodexAuth::AgentIdentity(_))) => { + false + } + (Some(expected), Some(current)) => { + expected.get_account_id() == current.get_account_id() + && expected.get_chatgpt_user_id() == current.get_chatgpt_user_id() + && expected.is_workspace_account() == current.is_workspace_account() + && expected.is_fedramp_account() == current.is_fedramp_account() + && (expected.get_account_id().is_some() || expected == current) + } + (None, None) => true, + _ => false, + } + } +} diff --git a/codex-rs/codex-mcp/src/lib.rs b/codex-rs/codex-mcp/src/lib.rs index f2bf80fbb4..32e80e099f 100644 --- a/codex-rs/codex-mcp/src/lib.rs +++ b/codex-rs/codex-mcp/src/lib.rs @@ -7,6 +7,7 @@ pub use elicitation::ElicitationLifecycle; pub use elicitation::ElicitationReviewRequest; pub use elicitation::ElicitationReviewer; pub use elicitation::ElicitationReviewerHandle; +pub use event_stream::McpEventStreamOpener; pub use resource_client::McpEventCatalogSnapshot; pub use resource_client::McpEventDefinition; pub use resource_client::McpEventNotification; @@ -102,6 +103,7 @@ mod client_capabilities; pub(crate) mod codex_apps; pub(crate) mod connection_manager; pub(crate) mod elicitation; +mod event_stream; mod executor_environment_http_client; pub(crate) mod mcp; mod openai_docs_source_attribution; diff --git a/codex-rs/codex-mcp/src/resource_client.rs b/codex-rs/codex-mcp/src/resource_client.rs index b9a9bd4ff2..43d82a8378 100644 --- a/codex-rs/codex-mcp/src/resource_client.rs +++ b/codex-rs/codex-mcp/src/resource_client.rs @@ -21,6 +21,7 @@ use serde_json::json; use tokio::runtime::Handle; use tokio::sync::watch; +use crate::McpEventStreamOpener; use crate::McpRuntime; use crate::connection_manager::McpConnectionSet; use crate::mcp::CODEX_APPS_MCP_SERVER_NAME; @@ -70,13 +71,13 @@ pub struct McpEventStream { request: Option, runtime_handle: Handle, client: Option>, - hosted_event_server_removals: watch::Receiver<()>, + cancel_event_streams_on_server_removal: watch::Receiver<()>, } impl McpEventStream { pub(crate) async fn open( client: Arc, - hosted_event_server_removals: watch::Receiver<()>, + cancel_event_streams_on_server_removal: watch::Receiver<()>, event_name: &str, arguments: &Value, request_meta: Option<&Map>, @@ -93,7 +94,7 @@ impl McpEventStream { request: Some(request), runtime_handle: Handle::current(), client: Some(client), - hosted_event_server_removals, + cancel_event_streams_on_server_removal, }) } @@ -106,7 +107,7 @@ impl McpEventStream { tokio::select! { biased; - Ok(()) = self.hosted_event_server_removals.changed() => { + Ok(()) = self.cancel_event_streams_on_server_removal.changed() => { self.cancel(); Err(anyhow!("hosted MCP event server was removed")) } @@ -286,7 +287,7 @@ impl McpResourceClient { arguments: &Value, request_meta: Option<&Map>, ) -> Result { - let (connections, hosted_event_server_removals) = self + let (connections, cancel_event_streams_on_server_removal) = self .runtime .latest_connections_for_event_server(CODEX_APPS_MCP_SERVER_NAME)?; let (managed, _) = connections @@ -294,13 +295,23 @@ impl McpResourceClient { .await?; McpEventStream::open( managed.client, - hosted_event_server_removals, + cancel_event_streams_on_server_removal, event_name, arguments, request_meta, ) .await } + + /// Creates an event stream opener using the task's event server settings. + pub fn event_stream_opener(&self) -> Result { + self.runtime.event_stream_opener() + } + + /// Forwards event server removal to the owner of the task's subscriptions. + pub fn forward_event_server_removals_to(&self, cancellation: watch::Sender<()>) { + self.runtime.forward_event_server_removals_to(cancellation); + } } fn resource_from_rmcp(resource: rmcp::model::Resource) -> Result { diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index 1bf0cfc5b7..7acd6b3b7a 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -1007,7 +1007,7 @@ fn record_protocol_discovery_metrics( ); } -fn mcp_initialize_request_params( +pub(crate) fn mcp_initialize_request_params( client_elicitation_capability: ElicitationCapability, client_mcp_extensions: ClientMcpExtensions, ) -> InitializeRequestParams { @@ -1067,7 +1067,7 @@ struct StartServerTaskParams { #[allow(clippy::too_many_arguments)] #[instrument(level = "trace", skip_all, fields(server_name = %server_name))] -async fn make_rmcp_client( +pub(crate) async fn make_rmcp_client( server_name: &str, server: EffectiveMcpServer, store_mode: OAuthCredentialsStoreMode, diff --git a/codex-rs/codex-mcp/src/runtime.rs b/codex-rs/codex-mcp/src/runtime.rs index bdd07128e1..dc30070e90 100644 --- a/codex-rs/codex-mcp/src/runtime.rs +++ b/codex-rs/codex-mcp/src/runtime.rs @@ -49,6 +49,7 @@ use crate::connection_manager::McpConnectionSet; use crate::elicitation::ElicitationLifecycle; use crate::elicitation::ElicitationRequestRouter; use crate::elicitation::ElicitationReviewerHandle; +use crate::event_stream::McpEventStreamOpener; use crate::mcp::CODEX_APPS_MCP_SERVER_NAME; use crate::resource_origin::ResourceOrigins; use crate::server::EffectiveMcpServer; @@ -91,12 +92,18 @@ pub struct McpRuntimeInput { /// their exact connections and configuration for as long as they are needed. pub struct McpRuntime { current: ArcSwap, - hosted_event_server_removals: watch::Sender<()>, + event_stream_cancellation: Mutex, reconnect_pending: AtomicBool, elicitation_router: ElicitationRequestRouter, resource_origins: Mutex, } +struct EventStreamCancellation { + event_server_available: bool, + cancel_event_streams_on_server_removal: watch::Sender<()>, + retained_subscription_cancellation: Option>, +} + struct PublishedMcpRuntime { connections: Arc, config: Option>, @@ -178,7 +185,11 @@ impl McpRuntime { selected_environments: HashMap::new(), cached_binding: Mutex::new(None), }), - hosted_event_server_removals: watch::channel(()).0, + event_stream_cancellation: Mutex::new(EventStreamCancellation { + event_server_available: false, + cancel_event_streams_on_server_removal: watch::channel(()).0, + retained_subscription_cancellation: None, + }), reconnect_pending: AtomicBool::new(false), elicitation_router: ElicitationRequestRouter::default(), resource_origins: Mutex::default(), @@ -291,6 +302,10 @@ impl McpRuntime { .source() .is_host_owned_apps(CODEX_APPS_MCP_SERVER_NAME, registration.config()) }); + let mut cancellation = self + .event_stream_cancellation + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); self.current.store(Arc::new(PublishedMcpRuntime { connections, config: Some(config), @@ -302,8 +317,14 @@ impl McpRuntime { cached_binding: Mutex::new(None), })); let _ = publish.send(true); + cancellation.event_server_available = hosted_event_server_retained; if !hosted_event_server_retained { - self.hosted_event_server_removals.send_replace(()); + cancellation + .cancel_event_streams_on_server_removal + .send_replace(()); + if let Some(retained) = &cancellation.retained_subscription_cancellation { + retained.send_replace(()); + } } } @@ -566,7 +587,13 @@ impl McpRuntime { &self, server: &str, ) -> anyhow::Result<(Arc, watch::Receiver<()>)> { - let hosted_event_server_removals = self.hosted_event_server_removals.subscribe(); + let cancellation = self + .event_stream_cancellation + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let cancel_event_streams_on_server_removal = cancellation + .cancel_event_streams_on_server_removal + .subscribe(); let current = self.current.load(); if server == CODEX_APPS_MCP_SERVER_NAME && !current @@ -583,10 +610,45 @@ impl McpRuntime { } Ok(( Arc::clone(¤t.connections), - hosted_event_server_removals, + cancel_event_streams_on_server_removal, )) } + pub(crate) fn event_stream_opener(&self) -> anyhow::Result { + let cancellation = self + .event_stream_cancellation + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let connection = self + .current + .load() + .connections + .event_stream_connection + .clone() + .ok_or_else(|| anyhow::anyhow!("Event subscriptions are unavailable for this task"))?; + let cancel_event_streams_on_server_removal = cancellation + .retained_subscription_cancellation + .as_ref() + .unwrap_or(&cancellation.cancel_event_streams_on_server_removal) + .clone(); + Ok(McpEventStreamOpener { + connection, + cancellation_receiver: cancel_event_streams_on_server_removal.subscribe(), + cancel_event_streams_on_server_removal, + }) + } + + pub(crate) fn forward_event_server_removals_to(&self, retained: watch::Sender<()>) { + let mut cancellation = self + .event_stream_cancellation + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if !cancellation.event_server_available { + retained.send_replace(()); + } + cancellation.retained_subscription_cancellation = Some(retained); + } + pub async fn shutdown(&self) { self.latest_connections().shutdown().await; } diff --git a/codex-rs/ext/mcp/Cargo.toml b/codex-rs/ext/mcp/Cargo.toml index e639081358..8a97a7554b 100644 --- a/codex-rs/ext/mcp/Cargo.toml +++ b/codex-rs/ext/mcp/Cargo.toml @@ -32,7 +32,12 @@ tracing = { workspace = true } tokio = { workspace = true, features = ["sync"] } [dev-dependencies] +anyhow = { workspace = true } +axum = { workspace = true } +codex-exec-server-test-support = { workspace = true } codex-login = { workspace = true } +futures = { workspace = true } pretty_assertions = { workspace = true } tempfile = { workspace = true } tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } +tokio-util = { workspace = true, features = ["rt"] } diff --git a/codex-rs/ext/mcp/src/event_stream_tests.rs b/codex-rs/ext/mcp/src/event_stream_tests.rs new file mode 100644 index 0000000000..784f316ad3 --- /dev/null +++ b/codex-rs/ext/mcp/src/event_stream_tests.rs @@ -0,0 +1,238 @@ +//! Exercises independent event connections through the real MCP HTTP transport. +#![allow(clippy::unwrap_used)] + +use std::collections::HashMap; +use std::convert::Infallible; +use std::sync::Arc; +use std::time::Duration; + +use anyhow::Result; +use axum::Json; +use axum::Router; +use axum::http::HeaderMap; +use axum::http::StatusCode; +use axum::response::IntoResponse; +use axum::response::sse::Event; +use axum::response::sse::Sse; +use axum::routing::post; +use codex_config::McpServerTransportConfig; +use codex_core::config::ConfigBuilder; +use codex_core::plugins_manager_for_config; +use codex_exec_server_test_support::environment_manager_without_environments; +use codex_login::AuthManager; +use codex_login::CodexAuth; +use codex_login::ExternalAuth; +use codex_login::ExternalAuthFuture; +use codex_login::ExternalAuthRefreshContext; +use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; +use codex_mcp::EffectiveMcpServer; +use codex_mcp::McpEventNotification; +use codex_mcp::McpResourceClient; +use codex_mcp::McpRuntime; +use codex_mcp::McpRuntimeContext; +use codex_mcp::McpRuntimeInput; +use codex_mcp::McpServerRegistration; +use codex_mcp::McpStartupPolicy; +use codex_mcp::ResolvedMcpCatalog; +use futures::StreamExt; +use pretty_assertions::assert_eq; +use serde_json::Value; +use serde_json::json; +use tokio::sync::mpsc; +use tokio::time::timeout; +use tokio_util::sync::CancellationToken; + +pub(super) const WAIT: Duration = Duration::from_secs(/*secs*/ 5); + +pub(super) struct TestStream { + pub request: Value, + pub headers: HeaderMap, + pub notifications: mpsc::UnboundedSender, +} + +impl TestStream { + pub fn notify(&self, method: &str) -> McpEventNotification { + let params = json!({ + "data": {"text": "hello"}, + "_meta": {"io.modelcontextprotocol/subscriptionId": self.request["id"], "connector": "test"}, + }); + self.notifications + .send(json!({"jsonrpc": "2.0", "method": method, "params": params})) + .unwrap(); + McpEventNotification { + method: method.into(), + params: Some(params), + } + } +} + +pub(super) struct Fixture { + _home: tempfile::TempDir, + _server: tokio_util::task::AbortOnDropHandle>, + pub runtime: Arc, + pub auth: Arc, + pub opened: mpsc::UnboundedReceiver, +} + +impl Fixture { + pub async fn new() -> Result { + let (opened_tx, opened) = mpsc::unbounded_channel(); + let router = Router::new().fallback(post( + move |headers: HeaderMap, Json(request): Json| { + let opened = opened_tx.clone(); + async move { + let result = match request["method"].as_str() { + Some("initialize") => json!({ + "protocolVersion": "2025-06-18", "capabilities": {}, + "serverInfo": {"name": "test", "version": "1"}, + }), + Some("tools/list") => json!({"tools": []}), + Some("notifications/initialized") => { + return StatusCode::ACCEPTED.into_response(); + } + Some("events/stream") => { + let (notifications, receive) = mpsc::unbounded_channel(); + opened + .send(TestStream { + request, + headers, + notifications, + }) + .unwrap(); + let events = futures::stream::unfold(receive, |mut receive| async { + receive.recv().await.map(|value| (value, receive)) + }) + .map(|value: Value| { + Ok::<_, Infallible>( + Event::default().event("message").data(value.to_string()), + ) + }); + return Sse::new(events).into_response(); + } + method => panic!("unexpected request: {method:?}"), + }; + Json(json!({"jsonrpc": "2.0", "id": request["id"], "result": result})) + .into_response() + } + }, + )); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?; + let url = format!("http://{}", listener.local_addr()?); + let server = tokio::spawn(async move { axum::serve(listener, router).await }); + let home = tempfile::tempdir()?; + let auth_value = CodexAuth::from_external_chatgpt_tokens( + "header.e30.first", + "account", + /*chatgpt_plan_type*/ None, + )?; + let auth = + AuthManager::from_auth_for_testing_with_home(auth_value.clone(), home.path().into()); + let config = ConfigBuilder::default() + .codex_home(home.path().into()) + .fallback_cwd(Some(home.path().into())) + .build() + .await?; + let plugins = plugins_manager_for_config(&config, Arc::clone(&auth)); + let mut config = config.to_mcp_config(&plugins).await; + let mut server_config = codex_mcp::hosted_plugin_runtime_mcp_server_config( + &url, /*apps_mcp_product_sku*/ None, /*originator*/ None, + ); + let McpServerTransportConfig::StreamableHttp { + bearer_token_env_var, + .. + } = &mut server_config.transport + else { + panic!("HTTP server"); + }; + *bearer_token_env_var = None; + let mut catalog = ResolvedMcpCatalog::builder(); + catalog.register(McpServerRegistration::from_hosted_apps( + "test", + /*contribution_order*/ 0, + server_config.clone(), + )); + config.mcp_server_catalog = catalog.build(); + let config = Arc::new(config); + let runtime_input = McpRuntimeInput { + startup_policy: McpStartupPolicy::Eager, + config: Arc::clone(&config), + plugins_available: false, + ready_selected_capability_roots: Vec::new(), + mcp_servers: HashMap::from([( + CODEX_APPS_MCP_SERVER_NAME.into(), + EffectiveMcpServer::configured(server_config.clone()), + )]), + submit_id: "test".into(), + tx_event: None, + startup_cancellation_token: CancellationToken::new(), + runtime_context: McpRuntimeContext::new( + Arc::new(environment_manager_without_environments()), + home.path().into(), + ), + codex_apps_tools_cache: Default::default(), + tool_catalog_cache: Default::default(), + codex_apps_tools_cache_key: codex_mcp::codex_apps_tools_cache_key(Some(&auth_value)), + client_mcp_extensions: Default::default(), + auth: Some(auth_value.clone()), + auth_manager: Some(Arc::clone(&auth)), + elicitation_reviewer: None, + elicitation_lifecycle: None, + }; + let runtime = Arc::new(McpRuntime::new(runtime_input).await); + assert!( + runtime + .latest_wait_for_server_ready(CODEX_APPS_MCP_SERVER_NAME, WAIT) + .await + ); + Ok(Self { + _home: home, + _server: tokio_util::task::AbortOnDropHandle::new(server), + runtime, + auth, + opened, + }) + } +} + +pub(super) struct StaticAuth(pub CodexAuth); + +impl ExternalAuth for StaticAuth { + fn resolve(&self) -> ExternalAuthFuture<'_, CodexAuth> { + Box::pin(async { Ok(self.0.clone()) }) + } + fn refresh(&self, _: ExternalAuthRefreshContext) -> ExternalAuthFuture<'_, CodexAuth> { + self.resolve() + } +} + +#[tokio::test] +async fn event_connections_open_after_runtime_shutdown_with_current_auth() -> Result<()> { + let mut fixture = Fixture::new().await?; + let opener = McpResourceClient::new(Arc::clone(&fixture.runtime)).event_stream_opener()?; + fixture.runtime.shutdown().await; + drop(fixture.runtime); + + fixture + .auth + .set_external_auth(Arc::new(StaticAuth( + CodexAuth::from_external_chatgpt_tokens( + "header.e30.refreshed", + "account", + /*chatgpt_plan_type*/ None, + )?, + ))) + .await?; + let mut stream = opener + .open("test.event", &json!({}), /*request_meta*/ None) + .await?; + let remote = timeout(WAIT, fixture.opened.recv()).await?.unwrap(); + assert_eq!( + remote.headers["authorization"], + "Bearer header.e30.refreshed" + ); + let expected = remote.notify("notifications/events/active"); + assert_eq!(timeout(WAIT, stream.recv()).await??, Some(expected)); + drop(stream); + timeout(WAIT, remote.notifications.closed()).await?; + Ok(()) +} diff --git a/codex-rs/ext/mcp/src/lib.rs b/codex-rs/ext/mcp/src/lib.rs index e417794671..5cbf1eb1c9 100644 --- a/codex-rs/ext/mcp/src/lib.rs +++ b/codex-rs/ext/mcp/src/lib.rs @@ -7,6 +7,9 @@ use codex_extension_api::McpServerContributor; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_mcp::hosted_plugin_runtime_mcp_server_config; +#[cfg(test)] +#[path = "event_stream_tests.rs"] +mod event_stream_tests; mod executor_plugin; struct HostedPluginRuntimeExtension;