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

+67 -5
View File
@@ -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<PublishedMcpRuntime>,
hosted_event_server_removals: watch::Sender<()>,
event_stream_cancellation: Mutex<EventStreamCancellation>,
reconnect_pending: AtomicBool,
elicitation_router: ElicitationRequestRouter,
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 {
connections: Arc<McpConnectionSet>,
config: Option<Arc<McpConfig>>,
@@ -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<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();
if server == CODEX_APPS_MCP_SERVER_NAME
&& !current
@@ -583,10 +610,45 @@ impl McpRuntime {
}
Ok((
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) {
self.latest_connections().shutdown().await;
}