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
This commit is contained in:
Evan Fannin authored and copyberry committed 2026-08-31 19:18:18 +00:00
1 parent d60560f14e
commit 53691bad98
11 files changed
+520 -15

No files matched your search

+5
View File
@@ -3646,12 +3646,15 @@ dependencies = [
name = "codex-mcp-extension" name = "codex-mcp-extension"
version = "0.0.0" version = "0.0.0"
dependencies = [ dependencies = [
"anyhow",
"axum",
"codex-config", "codex-config",
"codex-connectors", "codex-connectors",
"codex-connectors-extension", "codex-connectors-extension",
"codex-core", "codex-core",
"codex-core-plugins", "codex-core-plugins",
"codex-exec-server", "codex-exec-server",
"codex-exec-server-test-support",
"codex-extension-api", "codex-extension-api",
"codex-features", "codex-features",
"codex-login", "codex-login",
@@ -3660,11 +3663,13 @@ dependencies = [
"codex-protocol", "codex-protocol",
"codex-utils-absolute-path", "codex-utils-absolute-path",
"codex-utils-path-uri", "codex-utils-path-uri",
"futures",
"pretty_assertions", "pretty_assertions",
"serde_json", "serde_json",
"tempfile", "tempfile",
"thiserror 2.0.18", "thiserror 2.0.18",
"tokio", "tokio",
"tokio-util",
"tracing", "tracing",
] ]
+24 -2
View File
@@ -33,6 +33,7 @@ use crate::binding::call_tool_result_from_rmcp;
use crate::catalog::McpServerSource; use crate::catalog::McpServerSource;
use crate::elicitation::ElicitationRequestManager; use crate::elicitation::ElicitationRequestManager;
use crate::elicitation::ElicitationRequestRouter; use crate::elicitation::ElicitationRequestRouter;
use crate::event_stream::EventStreamConnectionSettings;
use crate::mcp::CODEX_APPS_MCP_SERVER_NAME; use crate::mcp::CODEX_APPS_MCP_SERVER_NAME;
use crate::mcp::ToolPluginProvenance; use crate::mcp::ToolPluginProvenance;
use crate::pagination::MAX_CODEX_APPS_TOOL_CATALOG_ITEMS; 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. /// A published view over a set of running MCP server connections.
pub(crate) struct McpConnectionSet { pub(crate) struct McpConnectionSet {
servers: HashMap<String, McpServerView>, servers: HashMap<String, McpServerView>,
pub(crate) event_stream_connection: Option<Arc<EventStreamConnectionSettings>>,
disabled_servers: Vec<String>, disabled_servers: Vec<String>,
protocol_mode: crate::McpProtocolMode, protocol_mode: crate::McpProtocolMode,
required_servers: Vec<String>, required_servers: Vec<String>,
@@ -237,6 +239,7 @@ impl McpConnectionSet {
let tool_plugin_provenance = crate::mcp::tool_plugin_provenance(&config); let tool_plugin_provenance = crate::mcp::tool_plugin_provenance(&config);
let auth = auth.as_ref(); let auth = auth.as_ref();
let mut servers = HashMap::new(); let mut servers = HashMap::new();
let mut event_stream_connection = None;
let disabled_servers = mcp_servers let disabled_servers = mcp_servers
.iter() .iter()
.filter(|(_, server)| !server.enabled()) .filter(|(_, server)| !server.enabled())
@@ -276,9 +279,12 @@ impl McpConnectionSet {
let static_chatgpt_auth_provider = auth let static_chatgpt_auth_provider = auth
.filter(|auth| auth.uses_codex_backend()) .filter(|auth| auth.uses_codex_backend())
.map(codex_model_provider::auth_provider_from_auth); .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| { 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 for (server_name, server) in mcp_servers
@@ -351,6 +357,20 @@ impl McpConnectionSet {
} else { } else {
chatgpt_auth_provider_for_server(&server, chatgpt_auth_provider) 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( let connection_identity = McpServerConnectionIdentity::new(
&server_name, &server_name,
&server, &server,
@@ -682,6 +702,7 @@ impl McpConnectionSet {
} }
let manager = Self { let manager = Self {
servers, servers,
event_stream_connection,
disabled_servers, disabled_servers,
protocol_mode, protocol_mode,
required_servers, required_servers,
@@ -743,6 +764,7 @@ impl McpConnectionSet {
pub fn empty(prefix_mcp_tool_names: bool) -> Self { pub fn empty(prefix_mcp_tool_names: bool) -> Self {
Self { Self {
servers: HashMap::new(), servers: HashMap::new(),
event_stream_connection: None,
disabled_servers: Vec::new(), disabled_servers: Vec::new(),
protocol_mode: crate::McpProtocolMode::Legacy, protocol_mode: crate::McpProtocolMode::Legacy,
required_servers: Vec::new(), required_servers: Vec::new(),
@@ -102,6 +102,7 @@ impl McpConnectionSet {
) -> Self { ) -> Self {
Self { Self {
servers: HashMap::new(), servers: HashMap::new(),
event_stream_connection: None,
disabled_servers: Vec::new(), disabled_servers: Vec::new(),
protocol_mode: crate::McpProtocolMode::Legacy, protocol_mode: crate::McpProtocolMode::Legacy,
required_servers: Vec::new(), required_servers: Vec::new(),
+156
View File
@@ -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<Option<Arc<Environment>>, String>,
pub auth_provider: Option<SharedAuthProvider>,
pub auth_manager: Option<Arc<AuthManager>>,
pub auth: Option<CodexAuth>,
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<EventStreamConnectionSettings>,
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<String, Value>>,
) -> Result<McpEventStream> {
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,
}
}
}
+2
View File
@@ -7,6 +7,7 @@ pub use elicitation::ElicitationLifecycle;
pub use elicitation::ElicitationReviewRequest; pub use elicitation::ElicitationReviewRequest;
pub use elicitation::ElicitationReviewer; pub use elicitation::ElicitationReviewer;
pub use elicitation::ElicitationReviewerHandle; pub use elicitation::ElicitationReviewerHandle;
pub use event_stream::McpEventStreamOpener;
pub use resource_client::McpEventCatalogSnapshot; pub use resource_client::McpEventCatalogSnapshot;
pub use resource_client::McpEventDefinition; pub use resource_client::McpEventDefinition;
pub use resource_client::McpEventNotification; pub use resource_client::McpEventNotification;
@@ -102,6 +103,7 @@ mod client_capabilities;
pub(crate) mod codex_apps; pub(crate) mod codex_apps;
pub(crate) mod connection_manager; pub(crate) mod connection_manager;
pub(crate) mod elicitation; pub(crate) mod elicitation;
mod event_stream;
mod executor_environment_http_client; mod executor_environment_http_client;
pub(crate) mod mcp; pub(crate) mod mcp;
mod openai_docs_source_attribution; mod openai_docs_source_attribution;
+17 -6
View File
@@ -21,6 +21,7 @@ use serde_json::json;
use tokio::runtime::Handle; use tokio::runtime::Handle;
use tokio::sync::watch; use tokio::sync::watch;
use crate::McpEventStreamOpener;
use crate::McpRuntime; use crate::McpRuntime;
use crate::connection_manager::McpConnectionSet; use crate::connection_manager::McpConnectionSet;
use crate::mcp::CODEX_APPS_MCP_SERVER_NAME; use crate::mcp::CODEX_APPS_MCP_SERVER_NAME;
@@ -70,13 +71,13 @@ pub struct McpEventStream {
request: Option<CancellableEventStreamRequest>, request: Option<CancellableEventStreamRequest>,
runtime_handle: Handle, runtime_handle: Handle,
client: Option<Arc<RmcpClient>>, client: Option<Arc<RmcpClient>>,
hosted_event_server_removals: watch::Receiver<()>, cancel_event_streams_on_server_removal: watch::Receiver<()>,
} }
impl McpEventStream { impl McpEventStream {
pub(crate) async fn open( pub(crate) async fn open(
client: Arc<RmcpClient>, client: Arc<RmcpClient>,
hosted_event_server_removals: watch::Receiver<()>, cancel_event_streams_on_server_removal: watch::Receiver<()>,
event_name: &str, event_name: &str,
arguments: &Value, arguments: &Value,
request_meta: Option<&Map<String, Value>>, request_meta: Option<&Map<String, Value>>,
@@ -93,7 +94,7 @@ impl McpEventStream {
request: Some(request), request: Some(request),
runtime_handle: Handle::current(), runtime_handle: Handle::current(),
client: Some(client), client: Some(client),
hosted_event_server_removals, cancel_event_streams_on_server_removal,
}) })
} }
@@ -106,7 +107,7 @@ impl McpEventStream {
tokio::select! { tokio::select! {
biased; biased;
Ok(()) = self.hosted_event_server_removals.changed() => { Ok(()) = self.cancel_event_streams_on_server_removal.changed() => {
self.cancel(); self.cancel();
Err(anyhow!("hosted MCP event server was removed")) Err(anyhow!("hosted MCP event server was removed"))
} }
@@ -286,7 +287,7 @@ impl McpResourceClient {
arguments: &Value, arguments: &Value,
request_meta: Option<&Map<String, Value>>, request_meta: Option<&Map<String, Value>>,
) -> Result<McpEventStream> { ) -> Result<McpEventStream> {
let (connections, hosted_event_server_removals) = self let (connections, cancel_event_streams_on_server_removal) = self
.runtime .runtime
.latest_connections_for_event_server(CODEX_APPS_MCP_SERVER_NAME)?; .latest_connections_for_event_server(CODEX_APPS_MCP_SERVER_NAME)?;
let (managed, _) = connections let (managed, _) = connections
@@ -294,13 +295,23 @@ impl McpResourceClient {
.await?; .await?;
McpEventStream::open( McpEventStream::open(
managed.client, managed.client,
hosted_event_server_removals, cancel_event_streams_on_server_removal,
event_name, event_name,
arguments, arguments,
request_meta, request_meta,
) )
.await .await
} }
/// Creates an event stream opener using the task's event server settings.
pub fn event_stream_opener(&self) -> Result<McpEventStreamOpener> {
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<Resource> { fn resource_from_rmcp(resource: rmcp::model::Resource) -> Result<Resource> {
+2 -2
View File
@@ -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_elicitation_capability: ElicitationCapability,
client_mcp_extensions: ClientMcpExtensions, client_mcp_extensions: ClientMcpExtensions,
) -> InitializeRequestParams { ) -> InitializeRequestParams {
@@ -1067,7 +1067,7 @@ struct StartServerTaskParams {
#[allow(clippy::too_many_arguments)] #[allow(clippy::too_many_arguments)]
#[instrument(level = "trace", skip_all, fields(server_name = %server_name))] #[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_name: &str,
server: EffectiveMcpServer, server: EffectiveMcpServer,
store_mode: OAuthCredentialsStoreMode, store_mode: OAuthCredentialsStoreMode,
+67 -5
View File
@@ -49,6 +49,7 @@ use crate::connection_manager::McpConnectionSet;
use crate::elicitation::ElicitationLifecycle; use crate::elicitation::ElicitationLifecycle;
use crate::elicitation::ElicitationRequestRouter; use crate::elicitation::ElicitationRequestRouter;
use crate::elicitation::ElicitationReviewerHandle; use crate::elicitation::ElicitationReviewerHandle;
use crate::event_stream::McpEventStreamOpener;
use crate::mcp::CODEX_APPS_MCP_SERVER_NAME; use crate::mcp::CODEX_APPS_MCP_SERVER_NAME;
use crate::resource_origin::ResourceOrigins; use crate::resource_origin::ResourceOrigins;
use crate::server::EffectiveMcpServer; use crate::server::EffectiveMcpServer;
@@ -91,12 +92,18 @@ pub struct McpRuntimeInput {
/// their exact connections and configuration for as long as they are needed. /// their exact connections and configuration for as long as they are needed.
pub struct McpRuntime { pub struct McpRuntime {
current: ArcSwap<PublishedMcpRuntime>, current: ArcSwap<PublishedMcpRuntime>,
hosted_event_server_removals: watch::Sender<()>, event_stream_cancellation: Mutex<EventStreamCancellation>,
reconnect_pending: AtomicBool, reconnect_pending: AtomicBool,
elicitation_router: ElicitationRequestRouter, elicitation_router: ElicitationRequestRouter,
resource_origins: Mutex<ResourceOrigins>, resource_origins: Mutex<ResourceOrigins>,
} }
struct EventStreamCancellation {
event_server_available: bool,
cancel_event_streams_on_server_removal: watch::Sender<()>,
retained_subscription_cancellation: Option<watch::Sender<()>>,
}
struct PublishedMcpRuntime { struct PublishedMcpRuntime {
connections: Arc<McpConnectionSet>, connections: Arc<McpConnectionSet>,
config: Option<Arc<McpConfig>>, config: Option<Arc<McpConfig>>,
@@ -178,7 +185,11 @@ impl McpRuntime {
selected_environments: HashMap::new(), selected_environments: HashMap::new(),
cached_binding: Mutex::new(None), 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), reconnect_pending: AtomicBool::new(false),
elicitation_router: ElicitationRequestRouter::default(), elicitation_router: ElicitationRequestRouter::default(),
resource_origins: Mutex::default(), resource_origins: Mutex::default(),
@@ -291,6 +302,10 @@ impl McpRuntime {
.source() .source()
.is_host_owned_apps(CODEX_APPS_MCP_SERVER_NAME, registration.config()) .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 { self.current.store(Arc::new(PublishedMcpRuntime {
connections, connections,
config: Some(config), config: Some(config),
@@ -302,8 +317,14 @@ impl McpRuntime {
cached_binding: Mutex::new(None), cached_binding: Mutex::new(None),
})); }));
let _ = publish.send(true); let _ = publish.send(true);
cancellation.event_server_available = hosted_event_server_retained;
if !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, &self,
server: &str, server: &str,
) -> anyhow::Result<(Arc<McpConnectionSet>, watch::Receiver<()>)> { ) -> anyhow::Result<(Arc<McpConnectionSet>, 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(); let current = self.current.load();
if server == CODEX_APPS_MCP_SERVER_NAME if server == CODEX_APPS_MCP_SERVER_NAME
&& !current && !current
@@ -583,10 +610,45 @@ impl McpRuntime {
} }
Ok(( Ok((
Arc::clone(&current.connections), Arc::clone(&current.connections),
hosted_event_server_removals, cancel_event_streams_on_server_removal,
)) ))
} }
pub(crate) fn event_stream_opener(&self) -> anyhow::Result<McpEventStreamOpener> {
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) { pub async fn shutdown(&self) {
self.latest_connections().shutdown().await; self.latest_connections().shutdown().await;
} }
+5
View File
@@ -32,7 +32,12 @@ tracing = { workspace = true }
tokio = { workspace = true, features = ["sync"] } tokio = { workspace = true, features = ["sync"] }
[dev-dependencies] [dev-dependencies]
anyhow = { workspace = true }
axum = { workspace = true }
codex-exec-server-test-support = { workspace = true }
codex-login = { workspace = true } codex-login = { workspace = true }
futures = { workspace = true }
pretty_assertions = { workspace = true } pretty_assertions = { workspace = true }
tempfile = { workspace = true } tempfile = { workspace = true }
tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
tokio-util = { workspace = true, features = ["rt"] }
+238
View File
@@ -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<Value>,
}
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<std::io::Result<()>>,
pub runtime: Arc<McpRuntime>,
pub auth: Arc<AuthManager>,
pub opened: mpsc::UnboundedReceiver<TestStream>,
}
impl Fixture {
pub async fn new() -> Result<Self> {
let (opened_tx, opened) = mpsc::unbounded_channel();
let router = Router::new().fallback(post(
move |headers: HeaderMap, Json(request): Json<Value>| {
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(())
}
+3
View File
@@ -7,6 +7,9 @@ use codex_extension_api::McpServerContributor;
use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME;
use codex_mcp::hosted_plugin_runtime_mcp_server_config; use codex_mcp::hosted_plugin_runtime_mcp_server_config;
#[cfg(test)]
#[path = "event_stream_tests.rs"]
mod event_stream_tests;
mod executor_plugin; mod executor_plugin;
struct HostedPluginRuntimeExtension; struct HostedPluginRuntimeExtension;