github/codex
4
0 Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-13 14:46:14 +00:00
Code Issues Packages Projects Releases Wiki Activity

Pin MCP runtimes to model steps (#30101)

Browse Source 
## 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:
jif authored and GitHub committed 2026-06-26 00:53:07 +01:00
1 parent 8ce931ab76
commit ee9e0f6387
27 files changed
+596 -325
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

codex-rs/core/src/mcp_tool_call.rs
+62 -65
View File
@@ -18,6 +18,7 @@ use crate::mcp_openai_file::rewrite_mcp_tool_arguments_for_openai_files;
use crate::mcp_tool_approval_templates::RenderedMcpToolApprovalParam;
use crate::mcp_tool_approval_templates::render_mcp_tool_approval_template;
use crate::session::session::Session;
use crate::session::step_context::StepContext;
use crate::session::turn_context::TurnContext;
use crate::tools::hook_names::HookToolName;
use crate::tools::sandboxing::PermissionRequestPayload;
@@ -35,6 +36,7 @@ use codex_features::Feature;
use codex_hooks::PermissionRequestDecision;
use codex_mcp::CODEX_APPS_MCP_SERVER_NAME;
use codex_mcp::MCP_TOOL_CODEX_APPS_META_KEY;
use codex_mcp::McpConnectionManager;
use codex_mcp::McpPermissionPromptAutoApproveContext;
use codex_mcp::SandboxState;
use codex_mcp::auth_elicitation_completed_result;
@@ -111,13 +113,15 @@ const MCP_TOOL_CALL_EVENT_RESULT_MAX_BYTES: usize = DEFAULT_OUTPUT_BYTES_CAP;
/// item lifecycle events to the `Session`.
pub(crate) async fn handle_mcp_tool_call(
sess: Arc<Session>,
turn_context: &Arc<TurnContext>,
step_context: &Arc<StepContext>,
call_id: String,
server: String,
tool_name: String,
hook_tool_name: HookToolName,
arguments: String,
) -> HandledMcpToolCall {
let turn_context = &step_context.turn;
let manager = step_context.mcp.manager();
// Parse the `arguments` as JSON. An empty string is OK, but invalid JSON
// is not.
let arguments_value = if arguments.trim().is_empty() {
@@ -141,8 +145,14 @@ pub(crate) async fn handle_mcp_tool_call(
arguments: arguments_value.clone(),
};
let metadata =
lookup_mcp_tool_metadata(sess.as_ref(), turn_context.as_ref(), &server, &tool_name).await;
let metadata = lookup_mcp_tool_metadata(
sess.as_ref(),
turn_context.as_ref(),
manager,
&server,
&tool_name,
)
.await;
let item_metadata = McpToolCallItemMetadata::from_tool_metadata(&server, metadata.as_ref());
let app_tool_policy = if server == CODEX_APPS_MCP_SERVER_NAME {
let annotations = metadata
@@ -169,7 +179,6 @@ pub(crate) async fn handle_mcp_tool_call(
} else if let Some(approval_mode) = {
// Selected-plugin registrations are absent from config.toml and the legacy plugin manager,
// so their resolved catalog entry is the authoritative source for tool approval policy.
let manager = sess.services.mcp_connection_manager.load();
manager
.is_selected_plugin_mcp_server(&server)
.then(|| manager.tool_approval_mode(&server, &tool_name))
@@ -226,7 +235,7 @@ pub(crate) async fn handle_mcp_tool_call(
if let Some(decision) = maybe_request_mcp_tool_approval(
&sess,
turn_context,
step_context,
&call_id,
&invocation,
&hook_tool_name,
@@ -241,7 +250,7 @@ pub(crate) async fn handle_mcp_tool_call(
| McpToolApprovalDecision::AcceptAndRemember => {
return handle_approved_mcp_tool_call(
sess.as_ref(),
turn_context.as_ref(),
step_context.as_ref(),
&call_id,
invocation,
metadata.as_ref(),
@@ -298,7 +307,7 @@ pub(crate) async fn handle_mcp_tool_call(
handle_approved_mcp_tool_call(
sess.as_ref(),
turn_context.as_ref(),
step_context.as_ref(),
&call_id,
invocation,
metadata.as_ref(),
@@ -340,24 +349,21 @@ impl McpToolCallItemMetadata {
async fn handle_approved_mcp_tool_call(
sess: &Session,
turn_context: &TurnContext,
step_context: &StepContext,
call_id: &str,
invocation: McpInvocation,
metadata: Option<&McpToolApprovalMetadata>,
item_metadata: McpToolCallItemMetadata,
) -> HandledMcpToolCall {
let turn_context = step_context.turn.as_ref();
let manager = step_context.mcp.manager();
let server = invocation.server.clone();
maybe_mark_thread_memory_mode_polluted(sess, turn_context, &server).await;
maybe_mark_thread_memory_mode_polluted(sess, turn_context, manager, &server).await;
let tool_name = invocation.tool.clone();
let arguments_value = invocation.arguments.clone();
let connector_id = metadata.and_then(|metadata| metadata.connector_id.as_deref());
let connector_name = metadata.and_then(|metadata| metadata.connector_name.as_deref());
let server_origin = {
let mcp_connection_manager = sess.services.mcp_connection_manager.load_full();
mcp_connection_manager
.server_origin(&server)
.map(str::to_string)
};
let server_origin = manager.server_origin(&server).map(str::to_string);
let start = Instant::now();
let rewrite = rewrite_mcp_tool_arguments_for_openai_files(
@@ -380,7 +386,7 @@ async fn handle_approved_mcp_tool_call(
build_mcp_tool_call_request_meta(turn_context, &server, call_id, metadata);
execute_mcp_tool_call(
sess,
turn_context,
step_context,
call_id,
&invocation,
rewritten_arguments,
@@ -420,7 +426,7 @@ async fn handle_approved_mcp_tool_call(
truncate_mcp_tool_result_for_event(&result),
)
.await;
maybe_track_codex_app_used(sess, turn_context, &server, &tool_name).await;
maybe_track_codex_app_used(sess, turn_context, manager, &server, &tool_name).await;
let outcome = mcp_call_metric_outcome(&result);
emit_mcp_call_metrics(
@@ -547,17 +553,19 @@ fn truncate_str_to_char_boundary(value: &str, max_chars: usize) -> &str {
async fn execute_mcp_tool_call(
sess: &Session,
turn_context: &TurnContext,
step_context: &StepContext,
call_id: &str,
invocation: &McpInvocation,
rewritten_arguments: Option<JsonValue>,
metadata: Option<&McpToolApprovalMetadata>,
request_meta: Option<JsonValue>,
) -> Result<CallToolResult, String> {
let turn_context = step_context.turn.as_ref();
let manager = step_context.mcp.manager();
let request_meta = with_mcp_tool_call_thread_id_meta(request_meta, &sess.thread_id.to_string());
let request_meta = augment_mcp_tool_request_meta_with_sandbox_state(
sess,
turn_context,
step_context,
manager,
&invocation.server,
request_meta,
)
@@ -568,7 +576,7 @@ async fn execute_mcp_tool_call(
.rollout_thread_trace
.start_mcp_call_trace(call_id);
let request_meta = mcp_call_trace.add_request_meta(request_meta);
let result = sess
let result = manager
.call_tool(
&invocation.server,
&invocation.tool,
@@ -587,6 +595,7 @@ async fn execute_mcp_tool_call(
Ok(maybe_request_codex_apps_auth_elicitation(
sess,
turn_context,
manager,
call_id,
&invocation.server,
metadata,
@@ -598,17 +607,13 @@ async fn execute_mcp_tool_call(
async fn maybe_request_codex_apps_auth_elicitation(
sess: &Session,
turn_context: &TurnContext,
manager: &McpConnectionManager,
call_id: &str,
server: &str,
metadata: Option<&McpToolApprovalMetadata>,
result: CallToolResult,
) -> CallToolResult {
if !sess
.services
.mcp_connection_manager
.load_full()
.is_host_owned_codex_apps_server(server)
{
if !manager.is_host_owned_codex_apps_server(server) {
return result;
}
@@ -666,15 +671,16 @@ async fn maybe_request_codex_apps_auth_elicitation(
return result;
}
refresh_codex_apps_after_connector_auth(sess, turn_context).await;
refresh_codex_apps_after_connector_auth(sess, turn_context, manager).await;
auth_elicitation_completed_result(&plan.auth_failure, result.meta)
}
async fn refresh_codex_apps_after_connector_auth(sess: &Session, turn_context: &TurnContext) {
let mcp_tools_result = {
let manager = sess.services.mcp_connection_manager.load_full();
manager.hard_refresh_codex_apps_tools_cache().await
};
async fn refresh_codex_apps_after_connector_auth(
sess: &Session,
turn_context: &TurnContext,
manager: &McpConnectionManager,
) {
let mcp_tools_result = manager.hard_refresh_codex_apps_tools_cache().await;
match mcp_tools_result {
Ok(mcp_tools) => {
@@ -692,13 +698,13 @@ async fn refresh_codex_apps_after_connector_auth(sess: &Session, turn_context: &
}
async fn augment_mcp_tool_request_meta_with_sandbox_state(
sess: &Session,
turn_context: &TurnContext,
step_context: &StepContext,
manager: &McpConnectionManager,
server: &str,
mut meta: Option<serde_json::Value>,
) -> anyhow::Result<Option<serde_json::Value>> {
let mcp_connection_manager = sess.services.mcp_connection_manager.load_full();
let supports_sandbox_state_meta = mcp_connection_manager
let turn_context = step_context.turn.as_ref();
let supports_sandbox_state_meta = manager
.server_supports_sandbox_state_meta_capability(server)
.await
.unwrap_or(false);
@@ -706,18 +712,18 @@ async fn augment_mcp_tool_request_meta_with_sandbox_state(
return Ok(meta);
}
let server_environment_id = mcp_connection_manager
let server_environment_id = manager
.server_environment_id(server)
.unwrap_or(codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID);
let Some(sandbox_cwd) = sandbox_cwd_for_mcp_server(turn_context, server_environment_id) else {
let Some(sandbox_cwd) = sandbox_cwd_for_mcp_server(step_context, server_environment_id) else {
return Ok(meta);
};
let permission_profile = turn_context.permission_profile();
let sandbox_state = serde_json::to_value(SandboxState {
permission_profile,
codex_linux_sandbox_exe: turn_context.config.codex_linux_sandbox_exe.clone(),
codex_linux_sandbox_exe: step_context.mcp.config().codex_linux_sandbox_exe.clone(),
sandbox_cwd,
use_legacy_landlock: turn_context.config.features.use_legacy_landlock(),
use_legacy_landlock: step_context.mcp.config().use_legacy_landlock,
})?;
match meta.as_mut() {
@@ -741,8 +747,8 @@ async fn augment_mcp_tool_request_meta_with_sandbox_state(
Ok(meta)
}
fn sandbox_cwd_for_mcp_server(turn_context: &TurnContext, environment_id: &str) -> Option<PathUri> {
if let Some(environment) = turn_context
fn sandbox_cwd_for_mcp_server(step_context: &StepContext, environment_id: &str) -> Option<PathUri> {
if let Some(environment) = step_context
.environments
.turn_environments
.iter()
@@ -753,7 +759,7 @@ fn sandbox_cwd_for_mcp_server(turn_context: &TurnContext, environment_id: &str)
if environment_id == codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID {
#[allow(deprecated)]
return Some(PathUri::from_abs_path(&turn_context.cwd));
return Some(PathUri::from_abs_path(&step_context.turn.cwd));
}
None
@@ -762,16 +768,13 @@ fn sandbox_cwd_for_mcp_server(turn_context: &TurnContext, environment_id: &str)
async fn maybe_mark_thread_memory_mode_polluted(
sess: &Session,
turn_context: &TurnContext,
manager: &McpConnectionManager,
server: &str,
) {
if !turn_context.config.memories.disable_on_external_context {
return;
}
let pollutes_memory = sess
.services
.mcp_connection_manager
.load_full()
.server_pollutes_memory(server);
let pollutes_memory = manager.server_pollutes_memory(server);
if !pollutes_memory {
return;
}
@@ -937,13 +940,14 @@ struct McpAppUsageMetadata {
async fn maybe_track_codex_app_used(
sess: &Session,
turn_context: &TurnContext,
manager: &McpConnectionManager,
server: &str,
tool_name: &str,
) {
if server != CODEX_APPS_MCP_SERVER_NAME {
return;
}
let metadata = lookup_mcp_app_usage_metadata(sess, server, tool_name).await;
let metadata = lookup_mcp_app_usage_metadata(manager, server, tool_name).await;
let (connector_id, app_name) = metadata
.map(|metadata| (metadata.connector_id, metadata.app_name))
.unwrap_or((None, None));
@@ -983,6 +987,7 @@ enum McpToolApprovalDecision {
Cancel,
}
#[derive(Clone)]
pub(crate) struct McpToolApprovalMetadata {
annotations: Option<ToolAnnotations>,
connector_id: Option<String>,
@@ -1176,13 +1181,15 @@ fn mcp_tool_approval_prompt_options(
async fn maybe_request_mcp_tool_approval(
sess: &Arc<Session>,
turn_context: &Arc<TurnContext>,
step_context: &Arc<StepContext>,
call_id: &str,
invocation: &McpInvocation,
hook_tool_name: &HookToolName,
metadata: Option<&McpToolApprovalMetadata>,
approval_mode: AppToolApproval,
) -> Option<McpToolApprovalDecision> {
let turn_context = &step_context.turn;
let manager = step_context.mcp.manager();
let approvals_reviewer = mcp_approvals_reviewer(turn_context, &invocation.server, metadata);
if mcp_permission_prompt_is_auto_approved(
turn_context.approval_policy.value(),
@@ -1201,12 +1208,7 @@ async fn maybe_request_mcp_tool_approval(
}
let session_approval_key = session_mcp_tool_approval_key(invocation, metadata, approval_mode);
let persistent_approval_key = if sess
.services
.mcp_connection_manager
.load()
.is_selected_plugin_mcp_server(&invocation.server)
{
let persistent_approval_key = if manager.is_selected_plugin_mcp_server(&invocation.server) {
None
} else {
persistent_mcp_tool_approval_key(invocation, metadata, approval_mode)
@@ -1452,10 +1454,10 @@ async fn mcp_tool_approval_decision_from_guardian(
pub(crate) async fn lookup_mcp_tool_metadata(
sess: &Session,
turn_context: &TurnContext,
manager: &McpConnectionManager,
server: &str,
tool_name: &str,
) -> Option<McpToolApprovalMetadata> {
let manager = sess.services.mcp_connection_manager.load_full();
let plugin_id = manager
.plugin_id_for_mcp_server_name(server)
.map(str::to_string);
@@ -1568,16 +1570,11 @@ fn get_mcp_app_resource_uri(
}
async fn lookup_mcp_app_usage_metadata(
sess: &Session,
manager: &McpConnectionManager,
server: &str,
tool_name: &str,
) -> Option<McpAppUsageMetadata> {
let tools = sess
.services
.mcp_connection_manager
.load_full()
.list_all_tools()
.await;
let tools = manager.list_all_tools().await;
tools.into_iter().find_map(|tool_info| {
if tool_info.server_name == server && tool_info.tool.name == tool_name {