mirror of
https://github.com/openai/codex.git
synced 2026-11-14 06:48:26 +00:00
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:
1 parent
d60560f14e
commit
53691bad98
11 files changed
+520
-15
No files matched your search
@@ -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(¤t.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;
|
||||
}
|
||||
|
||||
Reference in new issue
Block a user