mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
Pin MCP runtimes to model steps (#30101)
## Why
An MCP refresh can replace the session's current manager while a model
step is still running. The step must execute calls through the same
manager whose tools it advertised.
## Boundary
```text
current session MCP runtime
|
| capture once for this model step
v
StepContext.mcp
- exact MCP config
- exact connection manager
- exact runtime environment context
```
```rust
pub struct McpRuntimeSnapshot {
config: Arc<McpConfig>,
manager: Arc<McpConnectionManager>,
runtime_context: McpRuntimeContext,
}
```
## Example
```text
step A captures runtime A and advertises A's tools
refresh publishes runtime B
step A tool call -> runtime A
next step -> runtime B
```
Capturing the snapshot is only an `Arc` clone. It does not restart MCPs
or make an RPC.
## What changes
- Captures one MCP runtime in `StepContext`.
- Uses it for tool planning, tool calls, resources, approvals, connector
attribution, and elicitation.
- Publishes replacement runtimes atomically.
- Lets an old runtime live only while an in-flight step or request still
holds its `Arc`.
Most of this diff is mechanical routing from the session-global manager
to `step_context.mcp`; it does not introduce selected-plugin discovery
yet.
## What does not change
- No plugin or extension migration.
- No new MCP cache policy.
- No environment file watching.
- No client sharing between separate managers.
## Stack
1. Extension-owned World State sections.
2. Project executor skills through World State.
3. **This PR:** pin one MCP runtime to each model step.
4. Project selected MCP/app/connector metadata by environment
availability.
5. One end-to-end integration scenario.
This commit is contained in:
@@ -15,6 +15,7 @@ use crate::exec_policy::ExecPolicyManager;
|
||||
use crate::guardian::GuardianRejection;
|
||||
use crate::guardian::GuardianRejectionCircuitBreaker;
|
||||
use crate::mcp::McpManager;
|
||||
use crate::session::McpRuntimeSnapshot;
|
||||
use crate::tools::code_mode::CodeModeService;
|
||||
use crate::tools::handlers::ToolSearchHandlerCache;
|
||||
use crate::tools::network_approval::NetworkApprovalService;
|
||||
@@ -30,7 +31,9 @@ use codex_extension_api::ExtensionDataInit;
|
||||
use codex_extension_api::ExtensionRegistry;
|
||||
use codex_hooks::Hooks;
|
||||
use codex_login::AuthManager;
|
||||
use codex_mcp::McpConfig;
|
||||
use codex_mcp::McpConnectionManager;
|
||||
use codex_mcp::McpRuntimeContext;
|
||||
use codex_models_manager::manager::SharedModelsManager;
|
||||
use codex_otel::SessionTelemetry;
|
||||
use codex_protocol::capabilities::SelectedCapabilityRoot;
|
||||
@@ -44,8 +47,10 @@ use tokio::sync::Mutex;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
pub(crate) struct SessionServices {
|
||||
/// The latest manager; callers retain an owned handle while performing MCP I/O.
|
||||
/// Mirror of the latest manager for extension resource clients that predate runtime snapshots.
|
||||
pub(crate) mcp_connection_manager: Arc<ArcSwap<McpConnectionManager>>,
|
||||
/// The latest atomically published MCP config and manager pair.
|
||||
pub(crate) mcp_runtime: ArcSwapOption<McpRuntimeSnapshot>,
|
||||
pub(crate) mcp_startup_cancellation_token: Mutex<CancellationToken>,
|
||||
pub(crate) unified_exec_manager: UnifiedExecProcessManager,
|
||||
#[cfg_attr(not(unix), allow(dead_code))]
|
||||
@@ -99,12 +104,33 @@ impl SessionServices {
|
||||
/// resolve through the session's manager while validation waits.
|
||||
pub(crate) async fn install_mcp_connection_manager(
|
||||
&self,
|
||||
config: Arc<McpConfig>,
|
||||
runtime_context: McpRuntimeContext,
|
||||
manager: McpConnectionManager,
|
||||
) -> Result<()> {
|
||||
self.mcp_connection_manager.store(Arc::new(manager));
|
||||
self.mcp_connection_manager
|
||||
.load_full()
|
||||
.validate_required_servers()
|
||||
.await
|
||||
let runtime = self.publish_mcp_runtime(config, runtime_context, manager);
|
||||
runtime.manager().validate_required_servers().await
|
||||
}
|
||||
|
||||
pub(crate) fn publish_mcp_runtime(
|
||||
&self,
|
||||
config: Arc<McpConfig>,
|
||||
runtime_context: McpRuntimeContext,
|
||||
manager: McpConnectionManager,
|
||||
) -> Arc<McpRuntimeSnapshot> {
|
||||
let manager = Arc::new(manager);
|
||||
// Publish the manager for legacy resource clients first. Once the paired snapshot is
|
||||
// visible, every model-scoped consumer observes this exact manager.
|
||||
self.mcp_connection_manager.store(Arc::clone(&manager));
|
||||
let runtime = Arc::new(McpRuntimeSnapshot::new(config, manager, runtime_context));
|
||||
self.mcp_runtime.store(Some(Arc::clone(&runtime)));
|
||||
runtime
|
||||
}
|
||||
|
||||
pub(crate) fn latest_mcp_runtime(&self) -> Arc<McpRuntimeSnapshot> {
|
||||
let Some(runtime) = self.mcp_runtime.load_full() else {
|
||||
unreachable!("MCP runtime must be installed before handling requests");
|
||||
};
|
||||
runtime
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user