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:
jif
2026-06-26 00:53:07 +01:00
committed by GitHub
Unverified
parent 8ce931ab76
commit ee9e0f6387
27 changed files with 596 additions and 325 deletions
@@ -121,17 +121,11 @@ impl McpRequestProcessor {
} = params;
let auth = self.auth_manager.auth().await;
let (config, mcp_config) = match thread_id.as_deref() {
let (mcp_config, runtime_context) = match thread_id.as_deref() {
Some(thread_id) => {
let (_, thread) = self.load_thread(thread_id).await?;
let thread_config = thread.config().await;
let config = self
.config_manager
.load_latest_config_for_thread(thread_config.as_ref())
.await
.map_err(|err| internal_error(format!("failed to reload config: {err}")))?;
let mcp_config = thread.runtime_mcp_config(&config).await;
(config, mcp_config)
let runtime = thread.current_mcp_runtime();
(runtime.config().clone(), runtime.runtime_context().clone())
}
None => {
let config = self.load_latest_config(/*fallback_cwd*/ None).await?;
@@ -140,7 +134,11 @@ impl McpRequestProcessor {
.mcp_manager()
.runtime_config(&config)
.await;
(config, mcp_config)
let runtime_context = McpRuntimeContext::new(
self.thread_manager.environment_manager(),
config.cwd.to_path_buf(),
);
(mcp_config, runtime_context)
}
};
let effective_servers = codex_mcp::effective_mcp_servers(&mcp_config, auth.as_ref());
@@ -167,10 +165,6 @@ impl McpRequestProcessor {
}
};
let runtime_context = McpRuntimeContext::new(
self.thread_manager.environment_manager(),
config.cwd.to_path_buf(),
);
let http_client = runtime_context
.resolve_http_client(&name, server)
.map_err(|err| {
@@ -189,16 +183,16 @@ impl McpRequestProcessor {
let handle = perform_oauth_login_return_url_with_http_client(
&name,
&url,
config.mcp_oauth_credentials_store_mode,
config.auth_keyring_backend_kind(),
mcp_config.mcp_oauth_credentials_store_mode,
mcp_config.auth_keyring_backend_kind,
http_headers,
env_http_headers,
&resolved_scopes.scopes,
server.oauth_client_id(),
server.oauth_resource.as_deref(),
timeout_secs,
config.mcp_oauth_callback_port,
config.mcp_oauth_callback_url.as_deref(),
mcp_config.mcp_oauth_callback_port,
mcp_config.mcp_oauth_callback_url.as_deref(),
http_client,
)
.await
@@ -250,18 +244,16 @@ impl McpRequestProcessor {
None => (self.load_latest_config(/*fallback_cwd*/ None).await?, None),
};
let mcp_manager = self.thread_manager.mcp_manager();
let codex_apps_tools_cache = mcp_manager.codex_apps_tools_cache();
let auth = self.auth_manager.auth().await;
let mcp_config = match thread {
Some(thread) => thread.runtime_mcp_config(&config).await,
None => mcp_manager.runtime_config(&config).await,
};
let codex_apps_tools_cache = mcp_manager.codex_apps_tools_cache();
let auth = self.auth_manager.auth().await;
let environment_manager = self.thread_manager.environment_manager();
// This status path has no turn-selected environment. Use config cwd
// as the local stdio fallback; named environment stdio MCPs must
// declare their own absolute cwd.
let runtime_context =
McpRuntimeContext::new(Arc::clone(&environment_manager), config.cwd.to_path_buf());
let runtime_context = McpRuntimeContext::new(
self.thread_manager.environment_manager(),
config.cwd.to_path_buf(),
);
tokio::spawn(async move {
Self::list_mcp_server_status_task(
+24
View File
@@ -371,6 +371,30 @@ impl ResolvedMcpCatalog {
.collect()
}
/// Replaces the resolved server set while preserving known server sources.
///
/// Names not present in the existing catalog are treated as config-owned.
pub fn with_materialized_servers(&self, servers: HashMap<String, McpServerConfig>) -> Self {
let mut builder = Self::builder();
for (name, config) in servers {
let source = self
.server(&name)
.map(|server| server.source.clone())
.unwrap_or(McpServerSource::Config);
let precedence = match &source {
McpServerSource::Plugin(_) => RegistrationPrecedence::Plugin(Reverse(0)),
McpServerSource::SelectedPlugin(_) => {
RegistrationPrecedence::SelectedPlugin(Reverse(0))
}
McpServerSource::Config => RegistrationPrecedence::Config,
McpServerSource::Compatibility { .. } => RegistrationPrecedence::Compatibility,
McpServerSource::Extension { .. } => RegistrationPrecedence::Extension(0),
};
builder.register(McpServerRegistration::new(name, source, config, precedence));
}
builder.build()
}
/// Returns package attribution for each winning plugin-owned server.
pub fn plugin_attributions_by_server_name(&self) -> HashMap<String, McpPluginAttribution> {
self.servers
+11
View File
@@ -304,6 +304,17 @@ fn selected_plugins_override_discovered_plugins_but_not_config() {
}]
);
let refreshed = server("https://refreshed.example/mcp");
let catalog =
catalog.with_materialized_servers(HashMap::from([("docs".to_string(), refreshed.clone())]));
assert_eq!(
catalog.server("docs"),
Some(&super::ResolvedMcpServer {
source: selected_plugin_source("selected-alpha"),
config: refreshed,
})
);
let mut builder = catalog.to_builder();
let configured = server("https://config.example/mcp");
builder.register(McpServerRegistration::from_config(
+44 -18
View File
@@ -40,6 +40,7 @@ use crate::guardian::spawn_approval_request_review;
use crate::mcp_tool_call::MCP_TOOL_APPROVAL_ACCEPT;
use crate::mcp_tool_call::MCP_TOOL_APPROVAL_ACCEPT_FOR_SESSION;
use crate::mcp_tool_call::MCP_TOOL_APPROVAL_DECLINE_SYNTHETIC;
use crate::mcp_tool_call::McpToolApprovalMetadata;
use crate::mcp_tool_call::build_guardian_mcp_tool_review_request;
use crate::mcp_tool_call::is_mcp_tool_approval_question_id;
use crate::mcp_tool_call::lookup_mcp_tool_metadata;
@@ -60,6 +61,12 @@ use codex_protocol::protocol::MultiAgentVersion;
#[cfg(test)]
use crate::session::completed_session_loop_termination;
#[derive(Clone)]
struct PendingMcpInvocation {
invocation: McpInvocation,
metadata: Option<McpToolApprovalMetadata>,
}
/// Start an interactive sub-Codex thread and return IO channels.
///
/// The returned `events_rx` yields non-approval events emitted by the sub-agent.
@@ -149,10 +156,10 @@ pub(crate) async fn run_codex_thread_interactive(
let parent_session_clone = Arc::clone(&parent_session);
let parent_ctx_clone = Arc::clone(&parent_ctx);
let codex_for_events = Arc::clone(&codex);
// Cache delegated MCP invocations so guardian can recover the full tool call
// context when the later legacy RequestUserInput approval event only carries
// a call_id plus approval question metadata.
let pending_mcp_invocations = Arc::new(Mutex::new(HashMap::<String, McpInvocation>::new()));
// Cache the child call's MCP metadata at begin time. The later legacy
// RequestUserInput approval event only carries a call_id and question metadata.
let pending_mcp_invocations =
Arc::new(Mutex::new(HashMap::<String, PendingMcpInvocation>::new()));
tokio::spawn(async move {
forward_events(
codex_for_events,
@@ -270,7 +277,7 @@ async fn forward_events(
tx_sub: Sender<Event>,
parent_session: Arc<Session>,
parent_ctx: Arc<TurnContext>,
pending_mcp_invocations: Arc<Mutex<HashMap<String, McpInvocation>>>,
pending_mcp_invocations: Arc<Mutex<HashMap<String, PendingMcpInvocation>>>,
cancel_token: CancellationToken,
) {
let cancelled = cancel_token.cancelled();
@@ -357,10 +364,34 @@ async fn forward_events(
id,
msg: EventMsg::McpToolCallBegin(event),
} => {
// Runtime refreshes are published before a request step is captured, so
// the child runtime at call begin is the one executing this invocation.
// Cache its metadata now; the later approval event has only a call ID.
let metadata = if let Some(turn_context) =
codex.session.turn_context_for_sub_id(&id).await
{
let mcp = codex.session.services.latest_mcp_runtime();
lookup_mcp_tool_metadata(
codex.session.as_ref(),
turn_context.as_ref(),
mcp.manager(),
&event.invocation.server,
&event.invocation.tool,
)
.await
} else {
None
};
pending_mcp_invocations
.lock()
.await
.insert(event.call_id.clone(), event.invocation.clone());
.insert(
event.call_id.clone(),
PendingMcpInvocation {
invocation: event.invocation.clone(),
metadata,
},
);
if !forward_event_or_shutdown(
&codex,
&tx_sub,
@@ -648,7 +679,7 @@ async fn handle_request_user_input(
id: String,
parent_session: &Arc<Session>,
parent_ctx: &Arc<TurnContext>,
pending_mcp_invocations: &Arc<Mutex<HashMap<String, McpInvocation>>>,
pending_mcp_invocations: &Arc<Mutex<HashMap<String, PendingMcpInvocation>>>,
event: RequestUserInputEvent,
cancel_token: &CancellationToken,
) {
@@ -686,12 +717,12 @@ async fn handle_request_user_input(
/// programmatically after running the guardian review.
///
/// The RequestUserInput event only carries `call_id` plus approval question
/// metadata, so this helper joins it back to the cached `McpToolCallBegin`
/// invocation in order to rebuild the full guardian review request.
/// metadata, so this helper joins it back to the child runtime metadata cached at
/// `McpToolCallBegin` in order to rebuild the full guardian review request.
async fn maybe_auto_review_mcp_request_user_input(
parent_session: &Arc<Session>,
parent_ctx: &Arc<TurnContext>,
pending_mcp_invocations: &Arc<Mutex<HashMap<String, McpInvocation>>>,
pending_mcp_invocations: &Arc<Mutex<HashMap<String, PendingMcpInvocation>>>,
event: &RequestUserInputEvent,
cancel_token: &CancellationToken,
) -> Option<RequestUserInputResponse> {
@@ -702,18 +733,13 @@ async fn maybe_auto_review_mcp_request_user_input(
.questions
.iter()
.find(|question| is_mcp_tool_approval_question_id(&question.id))?;
let invocation = pending_mcp_invocations
let pending = pending_mcp_invocations
.lock()
.await
.get(&event.call_id)
.cloned()?;
let metadata = lookup_mcp_tool_metadata(
parent_session.as_ref(),
parent_ctx.as_ref(),
&invocation.server,
&invocation.tool,
)
.await;
let invocation = pending.invocation;
let metadata = pending.metadata;
let approvals_reviewer =
mcp_approvals_reviewer(parent_ctx, &invocation.server, metadata.as_ref());
if !routes_approval_to_guardian_with_reviewer(parent_ctx, approvals_reviewer) {
+14 -8
View File
@@ -411,10 +411,13 @@ async fn delegated_mcp_guardian_abort_returns_synthetic_decline_answer() {
let pending_mcp_invocations = Arc::new(Mutex::new(HashMap::from([(
"call-1".to_string(),
McpInvocation {
server: "custom_server".to_string(),
tool: "dangerous_tool".to_string(),
arguments: None,
PendingMcpInvocation {
invocation: McpInvocation {
server: "custom_server".to_string(),
tool: "dangerous_tool".to_string(),
arguments: None,
},
metadata: None,
},
)])));
let cancel_token = CancellationToken::new();
@@ -460,10 +463,13 @@ async fn delegated_mcp_user_reviewer_returns_none_without_metadata() {
crate::session::tests::make_session_and_context_with_rx().await;
let pending_mcp_invocations = Arc::new(Mutex::new(HashMap::from([(
"call-1".to_string(),
McpInvocation {
server: CODEX_APPS_MCP_SERVER_NAME.to_string(),
tool: "dangerous_tool".to_string(),
arguments: None,
PendingMcpInvocation {
invocation: McpInvocation {
server: CODEX_APPS_MCP_SERVER_NAME.to_string(),
tool: "dangerous_tool".to_string(),
arguments: None,
},
metadata: None,
},
)])));
let cancel_token = CancellationToken::new();
+5
View File
@@ -594,6 +594,11 @@ impl CodexThread {
self.codex.session.runtime_mcp_config(config).await
}
/// Returns the exact MCP config, environment bindings, and manager most recently published.
pub fn current_mcp_runtime(&self) -> Arc<crate::session::McpRuntimeSnapshot> {
self.codex.session.services.latest_mcp_runtime()
}
pub fn multi_agent_version(&self) -> Option<MultiAgentVersion> {
self.codex.session.multi_agent_version()
}
+2 -9
View File
@@ -199,15 +199,8 @@ pub(crate) async fn maybe_install_mcp_dependencies(
warn!("failed to refresh MCP dependencies for mentioned skills: {err}");
return;
}
let refresh_servers = sess.runtime_mcp_servers(&refresh_config).await;
sess.refresh_mcp_servers_now(
turn_context,
refresh_servers,
config.mcp_oauth_credentials_store_mode,
config.auth_keyring_backend_kind(),
elicitation_reviewer,
)
.await;
sess.refresh_mcp_servers_now(turn_context, &refresh_config, elicitation_reviewer)
.await;
}
async fn should_install_mcp_dependencies(
+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 {
+33 -21
View File
@@ -1,6 +1,7 @@
use super::*;
use crate::config::ConfigBuilder;
use crate::config::ManagedFeatures;
use crate::session::step_context::StepContext;
use crate::session::tests::make_session_and_context;
use crate::session::tests::make_session_and_context_with_rx;
use crate::session::turn_context::TurnEnvironment;
@@ -157,10 +158,12 @@ async fn execute_mcp_tool_call_records_replayable_correlation() -> anyhow::Resul
})
});
assert!(dispatch_trace.is_enabled());
let turn_context = Arc::new(turn_context);
let step_context = StepContext::for_test(Arc::clone(&turn_context));
let result = execute_mcp_tool_call(
&session,
&turn_context,
step_context.as_ref(),
"mcp-call",
&McpInvocation {
server: "docs".to_string(),
@@ -1147,7 +1150,8 @@ async fn mcp_sandbox_cwd_uses_matching_server_environment_uri() -> anyhow::Resul
/*shell*/ None,
));
let sandbox_cwd = sandbox_cwd_for_mcp_server(&turn_context, "remote");
let step_context = StepContext::for_test(Arc::new(turn_context));
let sandbox_cwd = sandbox_cwd_for_mcp_server(&step_context, "remote");
assert_eq!(sandbox_cwd, Some(secondary_cwd));
Ok(())
@@ -1157,7 +1161,8 @@ async fn mcp_sandbox_cwd_uses_matching_server_environment_uri() -> anyhow::Resul
async fn mcp_sandbox_cwd_is_none_for_unselected_server_environment() -> anyhow::Result<()> {
let (_, turn_context) = make_session_and_context().await;
let sandbox_cwd = sandbox_cwd_for_mcp_server(&turn_context, "remote");
let step_context = StepContext::for_test(Arc::new(turn_context));
let sandbox_cwd = sandbox_cwd_for_mcp_server(&step_context, "remote");
assert_eq!(sandbox_cwd, None);
Ok(())
@@ -1370,7 +1375,10 @@ fn codex_apps_auth_failure_metadata() -> McpToolApprovalMetadata {
)
}
async fn install_host_owned_codex_apps_manager(session: &Session, turn_context: &TurnContext) {
async fn host_owned_codex_apps_manager(
session: &Session,
turn_context: &TurnContext,
) -> Arc<codex_mcp::McpConnectionManager> {
let auth = session.services.auth_manager.auth().await;
let startup_cancellation_token = CancellationToken::new();
startup_cancellation_token.cancel();
@@ -1410,22 +1418,20 @@ async fn install_host_owned_codex_apps_manager(session: &Session, turn_context:
/*elicitation_reviewer*/ None,
)
.await;
session
.services
.mcp_connection_manager
.store(Arc::new(manager));
Arc::new(manager)
}
#[tokio::test]
async fn codex_apps_auth_elicitation_feature_disabled_returns_original_result() {
let (session, turn_context, rx_event) = make_session_and_context_with_rx().await;
install_host_owned_codex_apps_manager(&session, &turn_context).await;
let manager = host_owned_codex_apps_manager(&session, &turn_context).await;
let result = codex_apps_auth_failure_result();
let metadata = codex_apps_auth_failure_metadata();
let returned = maybe_request_codex_apps_auth_elicitation(
&session,
&turn_context,
manager.as_ref(),
"call_123",
CODEX_APPS_MCP_SERVER_NAME,
Some(&metadata),
@@ -1446,10 +1452,12 @@ async fn codex_apps_auth_elicitation_non_host_owned_server_returns_original_resu
Arc::make_mut(&mut turn_context.config).features = ManagedFeatures::from(features);
let result = codex_apps_auth_failure_result();
let metadata = codex_apps_auth_failure_metadata();
let manager = session.services.latest_mcp_runtime().manager_arc();
let returned = maybe_request_codex_apps_auth_elicitation(
&session,
turn_context,
manager.as_ref(),
"call_123",
CODEX_APPS_MCP_SERVER_NAME,
Some(&metadata),
@@ -1464,7 +1472,7 @@ async fn codex_apps_auth_elicitation_non_host_owned_server_returns_original_resu
#[tokio::test]
async fn codex_apps_auth_elicitation_disallowed_by_policy_returns_original_result() {
let (session, mut turn_context, rx_event) = make_session_and_context_with_rx().await;
install_host_owned_codex_apps_manager(&session, &turn_context).await;
let manager = host_owned_codex_apps_manager(&session, &turn_context).await;
let mut features = Features::with_defaults();
features.enable(Feature::AuthElicitation);
let turn_context = Arc::get_mut(&mut turn_context).expect("single turn context ref");
@@ -1479,6 +1487,7 @@ async fn codex_apps_auth_elicitation_disallowed_by_policy_returns_original_resul
let returned = maybe_request_codex_apps_auth_elicitation(
&session,
turn_context,
manager.as_ref(),
"call_123",
CODEX_APPS_MCP_SERVER_NAME,
Some(&metadata),
@@ -1493,7 +1502,7 @@ async fn codex_apps_auth_elicitation_disallowed_by_policy_returns_original_resul
#[tokio::test]
async fn codex_apps_auth_elicitation_granular_mcp_disabled_returns_original_result() {
let (session, mut turn_context, rx_event) = make_session_and_context_with_rx().await;
install_host_owned_codex_apps_manager(&session, &turn_context).await;
let manager = host_owned_codex_apps_manager(&session, &turn_context).await;
let mut features = Features::with_defaults();
features.enable(Feature::AuthElicitation);
let turn_context = Arc::get_mut(&mut turn_context).expect("single turn context ref");
@@ -1514,6 +1523,7 @@ async fn codex_apps_auth_elicitation_granular_mcp_disabled_returns_original_resu
let returned = maybe_request_codex_apps_auth_elicitation(
&session,
turn_context,
manager.as_ref(),
"call_123",
CODEX_APPS_MCP_SERVER_NAME,
Some(&metadata),
@@ -1528,7 +1538,7 @@ async fn codex_apps_auth_elicitation_granular_mcp_disabled_returns_original_resu
#[tokio::test]
async fn codex_apps_auth_elicitation_feature_enabled_requests_elicitation() {
let (session, mut turn_context, rx_event) = make_session_and_context_with_rx().await;
install_host_owned_codex_apps_manager(&session, &turn_context).await;
let manager = host_owned_codex_apps_manager(&session, &turn_context).await;
*session.active_turn.lock().await = Some(ActiveTurn::default());
let mut features = Features::with_defaults();
features.enable(Feature::AuthElicitation);
@@ -1542,10 +1552,12 @@ async fn codex_apps_auth_elicitation_feature_enabled_requests_elicitation() {
let request_task = tokio::spawn({
let session = Arc::clone(&session);
let turn_context = Arc::clone(&turn_context);
let manager = Arc::clone(&manager);
async move {
maybe_request_codex_apps_auth_elicitation(
&session,
&turn_context,
manager.as_ref(),
"call_123",
CODEX_APPS_MCP_SERVER_NAME,
Some(&metadata),
@@ -2466,7 +2478,7 @@ async fn approve_mode_skips_when_annotations_do_not_require_approval() {
let decision = maybe_request_mcp_tool_approval(
&session,
&turn_context,
&StepContext::for_test(Arc::clone(&turn_context)),
"call-1",
&invocation,
&HookToolName::new("mcp__test__tool"),
@@ -2542,7 +2554,7 @@ async fn guardian_mode_skips_auto_when_annotations_do_not_require_approval() {
let decision = maybe_request_mcp_tool_approval(
&session,
&turn_context,
&StepContext::for_test(Arc::clone(&turn_context)),
"call-guardian",
&invocation,
&HookToolName::new("mcp__test__tool"),
@@ -2601,7 +2613,7 @@ async fn permission_request_hook_allows_mcp_tool_call() {
let decision = maybe_request_mcp_tool_approval(
&session,
&turn_context,
&StepContext::for_test(Arc::clone(&turn_context)),
"call-mcp-hook",
&invocation,
&HookToolName::new("mcp__memory__create_entities"),
@@ -2663,7 +2675,7 @@ async fn permission_request_hook_uses_hook_tool_name_without_metadata() {
let decision = maybe_request_mcp_tool_approval(
&session,
&turn_context,
&StepContext::for_test(Arc::clone(&turn_context)),
"call-mcp-hook-no-metadata",
&invocation,
&HookToolName::new("mcp__memory__create_entities"),
@@ -2745,7 +2757,7 @@ async fn permission_request_hook_runs_after_remembered_mcp_approval() {
let turn_context = Arc::new(turn_context);
let decision = maybe_request_mcp_tool_approval(
&session,
&turn_context,
&StepContext::for_test(Arc::clone(&turn_context)),
"call-mcp-remembered",
&invocation,
&HookToolName::new("mcp__memory__create_entities"),
@@ -2828,7 +2840,7 @@ async fn guardian_mode_mcp_denial_returns_rationale_message() {
let decision = maybe_request_mcp_tool_approval(
&session,
&turn_context,
&StepContext::for_test(Arc::clone(&turn_context)),
"call-guardian-deny",
&invocation,
&HookToolName::new("mcp__test__tool"),
@@ -2888,7 +2900,7 @@ async fn prompt_mode_waits_for_approval_when_annotations_do_not_require_approval
tokio::spawn(async move {
maybe_request_mcp_tool_approval(
&session,
&turn_context,
&StepContext::for_test(Arc::clone(&turn_context)),
"call-prompt",
&invocation,
&HookToolName::new("mcp__test__tool"),
@@ -2946,7 +2958,7 @@ async fn full_access_mode_skips_mcp_tool_approval_for_all_approval_modes() {
] {
let decision = maybe_request_mcp_tool_approval(
&session,
&turn_context,
&StepContext::for_test(Arc::clone(&turn_context)),
"call-2",
&invocation,
&HookToolName::new("mcp__test__tool"),
@@ -3035,7 +3047,7 @@ async fn approve_mode_skips_guardian_in_every_permission_mode() {
let turn_context = Arc::new(turn_context);
let decision = maybe_request_mcp_tool_approval(
&session,
&turn_context,
&StepContext::for_test(Arc::clone(&turn_context)),
"call-3",
&invocation,
&HookToolName::new("mcp__test__tool"),
+2 -2
View File
@@ -598,8 +598,8 @@ async fn shutdown_session_runtime(sess: &Arc<Session>) {
warn!("failed to shutdown code mode session: {err}");
}
sess.services
.mcp_connection_manager
.load_full()
.latest_mcp_runtime()
.manager_arc()
.shutdown()
.await;
sess.guardian_review_session.shutdown().await;
+51 -73
View File
@@ -105,8 +105,8 @@ impl Session {
) -> McpServerElicitationOutcome {
if self
.services
.mcp_connection_manager
.load_full()
.latest_mcp_runtime()
.manager()
.elicitations_auto_deny()
{
return McpServerElicitationOutcome {
@@ -201,44 +201,20 @@ impl Session {
}
self.services
.mcp_connection_manager
.load_full()
.latest_mcp_runtime()
.manager_arc()
.resolve_elicitation(server_name, id, response)
.await
}
pub async fn list_resources(
&self,
server: &str,
params: Option<PaginatedRequestParams>,
) -> anyhow::Result<ListResourcesResult> {
self.services
.mcp_connection_manager
.load_full()
.list_resources(server, params)
.await
}
pub async fn list_resource_templates(
&self,
server: &str,
params: Option<PaginatedRequestParams>,
) -> anyhow::Result<ListResourceTemplatesResult> {
self.services
.mcp_connection_manager
.load_full()
.list_resource_templates(server, params)
.await
}
pub async fn read_resource(
&self,
server: &str,
params: ReadResourceRequestParams,
) -> anyhow::Result<ReadResourceResult> {
self.services
.mcp_connection_manager
.load_full()
.latest_mcp_runtime()
.manager_arc()
.read_resource(server, params)
.await
}
@@ -251,8 +227,8 @@ impl Session {
meta: Option<serde_json::Value>,
) -> anyhow::Result<CallToolResult> {
self.services
.mcp_connection_manager
.load_full()
.latest_mcp_runtime()
.manager_arc()
.call_tool(server, tool, arguments, meta)
.await
}
@@ -260,17 +236,17 @@ impl Session {
async fn refresh_mcp_servers_inner(
&self,
turn_context: &TurnContext,
mcp_servers: HashMap<String, McpServerConfig>,
store_mode: OAuthCredentialsStoreMode,
keyring_backend_kind: AuthKeyringBackendKind,
mut mcp_config: McpConfig,
configured_mcp_servers: HashMap<String, McpServerConfig>,
elicitation_reviewer: Option<ElicitationReviewerHandle>,
) {
mcp_config.mcp_server_catalog = mcp_config
.mcp_server_catalog
.with_materialized_servers(configured_mcp_servers);
let mcp_config = Arc::new(mcp_config);
let auth = self.services.auth_manager.auth().await;
let config = self.get_config().await;
let mcp_config = self.runtime_mcp_config(config.as_ref()).await;
let tool_plugin_provenance = codex_mcp::tool_plugin_provenance(&mcp_config);
let mcp_servers =
effective_mcp_servers_from_configured(mcp_servers, &mcp_config, auth.as_ref());
let mcp_servers = effective_mcp_servers(&mcp_config, auth.as_ref());
let environment_manager = self.services.turn_environments.environment_manager();
// TODO(anp): Migrate MCP runtime cwd plumbing to PathUri so foreign environment cwd
// values can be used without falling back to the legacy host cwd.
@@ -286,35 +262,36 @@ impl Session {
let mcp_runtime_context = McpRuntimeContext::new(environment_manager, cwd);
let auth_statuses = compute_auth_statuses(
mcp_servers.iter(),
store_mode,
keyring_backend_kind,
mcp_config.mcp_oauth_credentials_store_mode,
mcp_config.auth_keyring_backend_kind,
auth.as_ref(),
&mcp_runtime_context,
)
.await;
let mcp_startup_cancellation_token = {
let mut guard = self.services.mcp_startup_cancellation_token.lock().await;
guard.cancel();
// The previous runtime owns the old token and may still be serving an in-flight step.
// Its manager cancels that token when the last runtime handle is dropped.
let cancellation_token = CancellationToken::new();
*guard = cancellation_token.clone();
cancellation_token
};
let refreshed_manager = McpConnectionManager::new(
&mcp_servers,
store_mode,
keyring_backend_kind,
mcp_config.mcp_oauth_credentials_store_mode,
mcp_config.auth_keyring_backend_kind,
auth_statuses,
&turn_context.approval_policy,
turn_context.sub_id.clone(),
self.get_tx_event(),
mcp_startup_cancellation_token,
turn_context.permission_profile(),
mcp_runtime_context,
config.codex_home.to_path_buf(),
mcp_runtime_context.clone(),
mcp_config.codex_home.clone(),
self.services.mcp_manager.codex_apps_tools_cache(),
codex_apps_tools_cache_key(auth.as_ref()),
mcp_config.prefix_mcp_tool_names,
mcp_config.client_elicitation_capability,
mcp_config.client_elicitation_capability.clone(),
self.services
.supports_openai_form_elicitation
.load(std::sync::atomic::Ordering::Relaxed),
@@ -324,14 +301,12 @@ impl Session {
)
.await;
{
let current_manager = self.services.mcp_connection_manager.load_full();
refreshed_manager.set_elicitations_auto_deny(current_manager.elicitations_auto_deny());
let current_manager = self.services.latest_mcp_runtime();
refreshed_manager
.set_elicitations_auto_deny(current_manager.manager().elicitations_auto_deny());
}
let superseded_manager = self
.services
.mcp_connection_manager
.swap(Arc::new(refreshed_manager));
superseded_manager.shutdown().await;
self.services
.publish_mcp_runtime(mcp_config, mcp_runtime_context, refreshed_manager);
}
pub(crate) async fn refresh_mcp_servers_if_requested(
@@ -376,14 +351,23 @@ impl Session {
}
};
self.refresh_mcp_servers_inner(
turn_context,
mcp_servers,
store_mode,
keyring_backend_kind,
elicitation_reviewer,
)
.await;
let mut refresh_config = self.get_config().await.as_ref().clone();
refresh_config.mcp_oauth_credentials_store_mode = store_mode;
let secret_auth_storage_enabled = match keyring_backend_kind {
AuthKeyringBackendKind::Direct => false,
AuthKeyringBackendKind::Secrets => true,
};
if let Err(err) = refresh_config
.features
.set_enabled(Feature::SecretAuthStorage, secret_auth_storage_enabled)
{
warn!("failed to apply MCP auth keyring backend refresh config: {err}");
return;
}
let mcp_config = self.runtime_mcp_config(&refresh_config).await;
self.refresh_mcp_servers_inner(turn_context, mcp_config, mcp_servers, elicitation_reviewer)
.await;
}
pub(crate) async fn set_openai_form_elicitation_support(
@@ -417,19 +401,13 @@ impl Session {
pub(crate) async fn refresh_mcp_servers_now(
&self,
turn_context: &TurnContext,
mcp_servers: HashMap<String, McpServerConfig>,
store_mode: OAuthCredentialsStoreMode,
keyring_backend_kind: AuthKeyringBackendKind,
refresh_config: &Config,
elicitation_reviewer: Option<ElicitationReviewerHandle>,
) {
self.refresh_mcp_servers_inner(
turn_context,
mcp_servers,
store_mode,
keyring_backend_kind,
elicitation_reviewer,
)
.await;
let mcp_config = self.runtime_mcp_config(refresh_config).await;
let mcp_servers = codex_mcp::configured_mcp_servers(&mcp_config);
self.refresh_mcp_servers_inner(turn_context, mcp_config, mcp_servers, elicitation_reviewer)
.await;
}
#[cfg(test)]
+94
View File
@@ -0,0 +1,94 @@
use std::fmt;
use std::sync::Arc;
use codex_mcp::McpConfig;
use codex_mcp::McpConnectionManager;
use codex_mcp::McpRuntimeContext;
/// MCP config, exact environment bindings, and manager used by one model request.
pub struct McpRuntimeSnapshot {
config: Arc<McpConfig>,
manager: Arc<McpConnectionManager>,
runtime_context: McpRuntimeContext,
}
impl McpRuntimeSnapshot {
pub(crate) fn new(
config: Arc<McpConfig>,
manager: Arc<McpConnectionManager>,
runtime_context: McpRuntimeContext,
) -> Self {
Self {
config,
manager,
runtime_context,
}
}
pub fn config(&self) -> &McpConfig {
self.config.as_ref()
}
pub fn manager(&self) -> &McpConnectionManager {
self.manager.as_ref()
}
pub(crate) fn manager_arc(&self) -> Arc<McpConnectionManager> {
Arc::clone(&self.manager)
}
pub fn runtime_context(&self) -> &McpRuntimeContext {
&self.runtime_context
}
#[cfg(test)]
pub(crate) fn new_uninitialized_for_test(config: &crate::config::Config) -> Arc<Self> {
use codex_exec_server::EnvironmentManager;
use codex_features::Feature;
use codex_mcp::ResolvedMcpCatalog;
use rmcp::model::ElicitationCapability;
let mcp_config = McpConfig {
chatgpt_base_url: config.chatgpt_base_url.clone(),
apps_mcp_product_sku: config.apps_mcp_product_sku.clone(),
codex_home: config.codex_home.to_path_buf(),
mcp_oauth_credentials_store_mode: config.mcp_oauth_credentials_store_mode,
auth_keyring_backend_kind: config.auth_keyring_backend_kind(),
mcp_oauth_callback_port: config.mcp_oauth_callback_port,
mcp_oauth_callback_url: config.mcp_oauth_callback_url.clone(),
skill_mcp_dependency_install_enabled: config
.features
.enabled(Feature::SkillMcpDependencyInstall),
approval_policy: config.permissions.approval_policy.clone(),
codex_linux_sandbox_exe: config.codex_linux_sandbox_exe.clone(),
use_legacy_landlock: config.features.use_legacy_landlock(),
apps_enabled: config.features.enabled(Feature::Apps),
prefix_mcp_tool_names: config.prefix_mcp_tool_names(),
client_elicitation_capability: ElicitationCapability::default(),
mcp_server_catalog: ResolvedMcpCatalog::default(),
connector_snapshot: codex_connectors::ConnectorSnapshot::default(),
};
let manager = McpConnectionManager::new_uninitialized_with_permission_profile(
&config.permissions.approval_policy,
config.permissions.permission_profile(),
config.prefix_mcp_tool_names(),
);
let runtime_context = McpRuntimeContext::new(
Arc::new(EnvironmentManager::default_for_tests()),
config.cwd.to_path_buf(),
);
Arc::new(Self::new(
Arc::new(mcp_config),
Arc::new(manager),
runtime_context,
))
}
}
impl fmt::Debug for McpRuntimeSnapshot {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("McpRuntimeSnapshot")
.finish_non_exhaustive()
}
}
+29 -13
View File
@@ -158,15 +158,9 @@ use codex_utils_path_uri::PathUri;
use futures::future::BoxFuture;
use futures::future::Shared;
use futures::prelude::*;
use rmcp::model::ElicitationCapability;
use rmcp::model::FormElicitationCapability;
use rmcp::model::ListResourceTemplatesResult;
use rmcp::model::ListResourcesResult;
use rmcp::model::PaginatedRequestParams;
use rmcp::model::ReadResourceRequestParams;
use rmcp::model::ReadResourceResult;
use rmcp::model::RequestId;
use rmcp::model::UrlElicitationCapability;
use serde_json::Value;
use tokio::sync::Mutex;
use tokio::sync::RwLock;
@@ -214,6 +208,7 @@ mod handlers;
mod inject;
mod input_queue;
mod mcp;
mod mcp_runtime;
pub(crate) mod multi_agents;
mod review;
mod rollout_budget;
@@ -235,6 +230,7 @@ use self::handlers::submission_loop;
pub(crate) use self::input_queue::InputQueueActivity;
pub(crate) use self::input_queue::TurnInput;
pub(crate) use self::input_queue::TurnInputQueue;
pub use self::mcp_runtime::McpRuntimeSnapshot;
use self::review::spawn_review_thread;
use self::session::AppServerClientMetadata;
use self::session::Session;
@@ -340,7 +336,7 @@ use codex_core_plugins::RecommendedPluginCandidatesInput;
use codex_git_utils::get_git_repo_root;
use codex_mcp::McpConfig;
use codex_mcp::compute_auth_statuses;
use codex_mcp::effective_mcp_servers_from_configured;
use codex_mcp::effective_mcp_servers;
use codex_otel::SessionTelemetry;
use codex_otel::THREAD_STARTED_METRIC;
use codex_otel::TelemetryAuthMode;
@@ -840,8 +836,11 @@ impl Codex {
..Default::default()
})
.await?;
let mcp_connection_manager = self.session.services.mcp_connection_manager.load_full();
mcp_connection_manager.set_elicitations_auto_deny(mcp_elicitations_auto_deny);
self.session
.services
.latest_mcp_runtime()
.manager()
.set_elicitations_auto_deny(mcp_elicitations_auto_deny);
Ok(())
}
@@ -2855,10 +2854,12 @@ impl Session {
&environments.captured_environments(),
)
.await;
let mcp = self.services.latest_mcp_runtime();
Arc::new(StepContext::new(
turn_context,
environments,
selected_capability_roots,
mcp,
loaded_agents_md,
))
}
@@ -3125,6 +3126,17 @@ impl Session {
&self,
turn_context: &TurnContext,
world_state: &WorldState,
) -> Vec<ResponseItem> {
let mcp = self.services.latest_mcp_runtime();
self.build_initial_context_with_world_state_and_mcp(turn_context, world_state, &mcp)
.await
}
async fn build_initial_context_with_world_state_and_mcp(
&self,
turn_context: &TurnContext,
world_state: &WorldState,
mcp: &McpRuntimeSnapshot,
) -> Vec<ResponseItem> {
let mut developer_sections = Vec::<String>::with_capacity(8);
let mut contextual_user_sections = Vec::<String>::with_capacity(2);
@@ -3218,10 +3230,9 @@ impl Session {
}
}
if turn_context.config.include_apps_instructions && turn_context.apps_enabled() {
let mcp_connection_manager = self.services.mcp_connection_manager.load_full();
let accessible_and_enabled_connectors =
connectors::list_accessible_and_enabled_connectors_from_manager(
&mcp_connection_manager,
mcp.manager(),
&turn_context.config,
)
.await;
@@ -3328,7 +3339,8 @@ impl Session {
if turn_context.config.features.enabled(Feature::TokenBudget)
&& turn_context.model_context_window().is_some()
{
let mcp_result = self
let mcp_result = mcp
.manager()
.call_tool(
"notes",
"thread_hint",
@@ -3543,7 +3555,11 @@ impl Session {
// Full initial context resets the baseline; later turns persist only its changes.
let (mut context_items, world_state_item) = if should_inject_full_context {
let context_items = self
.build_initial_context_with_world_state(turn_context, world_state.as_ref())
.build_initial_context_with_world_state_and_mcp(
turn_context,
world_state.as_ref(),
step_context.mcp.as_ref(),
)
.await;
let snapshot = world_state.snapshot();
self.state
+25 -31
View File
@@ -661,10 +661,15 @@ impl Session {
let config_for_mcp = Arc::clone(&config);
let mcp_manager_for_mcp = Arc::clone(&mcp_manager);
let mcp_thread_init_for_startup = &mcp_thread_init;
let mcp_runtime_context_for_auth = McpRuntimeContext::new(
Arc::clone(&environment_manager),
session_configuration.cwd().to_path_buf(),
);
let mcp_runtime_cwd = session_configuration
.environment_selections()
.first()
.and_then(|environment| environment.cwd.to_abs_path().ok())
.map(|cwd| cwd.to_path_buf())
.unwrap_or_else(|| session_configuration.cwd().to_path_buf());
let mcp_runtime_context =
McpRuntimeContext::new(Arc::clone(&environment_manager), mcp_runtime_cwd);
let mcp_runtime_context_for_auth = mcp_runtime_context.clone();
let auth_and_mcp_fut = async move {
let auth = auth_manager_clone.auth().await;
let mcp_config = mcp_manager_for_mcp
@@ -680,7 +685,13 @@ impl Session {
&mcp_runtime_context_for_auth,
)
.await;
(auth, mcp_servers, auth_statuses, tool_plugin_provenance)
(
auth,
mcp_config,
mcp_servers,
auth_statuses,
tool_plugin_provenance,
)
}
.instrument(info_span!(
"session_init.auth_mcp",
@@ -691,7 +702,7 @@ impl Session {
let (
thread_persistence_result,
state_db_ctx,
(auth, mcp_servers, auth_statuses, tool_plugin_provenance),
(auth, mcp_config, mcp_servers, auth_statuses, tool_plugin_provenance),
) = tokio::join!(thread_persistence_fut, state_db_fut, auth_and_mcp_fut);
let mut live_thread_init =
@@ -1027,6 +1038,7 @@ impl Session {
// changing this to use Option or OnceCell, though the current
// setup is straightforward enough and performs well.
mcp_connection_manager,
mcp_runtime: arc_swap::ArcSwapOption::empty(),
mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()),
unified_exec_manager: UnifiedExecProcessManager::new(
config.background_terminal_max_timeout,
@@ -1159,14 +1171,6 @@ impl Session {
sess.send_event_raw(event).await;
}
let client_elicitation_capability = if config.features.enabled(Feature::AuthElicitation) {
ElicitationCapability {
form: Some(FormElicitationCapability::default()),
url: Some(UrlElicitationCapability::default()),
}
} else {
ElicitationCapability::default()
};
let mcp_startup_cancellation_token = {
let mut cancel_guard = sess.services.mcp_startup_cancellation_token.lock().await;
cancel_guard.cancel();
@@ -1174,20 +1178,6 @@ impl Session {
*cancel_guard = cancel_token.clone();
cancel_token
};
let mcp_runtime_context = {
let turn_environments = sess.services.turn_environments.snapshot().await;
// TODO(anp): Migrate MCP runtime cwd plumbing to PathUri so foreign environment
// cwd values can be used without falling back to the session host cwd.
let cwd = turn_environments
.primary()
.and_then(|turn_environment| turn_environment.cwd().to_abs_path().ok())
.map(|cwd| cwd.to_path_buf())
.unwrap_or_else(|| session_configuration.cwd().to_path_buf());
McpRuntimeContext::new(
sess.services.turn_environments.environment_manager(),
cwd,
)
};
let mcp_connection_manager = McpConnectionManager::new(
&mcp_servers,
config.mcp_oauth_credentials_store_mode,
@@ -1198,12 +1188,12 @@ impl Session {
tx_event.clone(),
mcp_startup_cancellation_token,
session_configuration.permission_profile(),
mcp_runtime_context,
mcp_runtime_context.clone(),
config.codex_home.to_path_buf(),
sess.services.mcp_manager.codex_apps_tools_cache(),
codex_apps_tools_cache_key(auth),
config.prefix_mcp_tool_names(),
client_elicitation_capability,
mcp_config.client_elicitation_capability.clone(),
sess.services
.supports_openai_form_elicitation
.load(std::sync::atomic::Ordering::Relaxed),
@@ -1217,7 +1207,11 @@ impl Session {
))
.await;
sess.services
.install_mcp_connection_manager(mcp_connection_manager)
.install_mcp_connection_manager(
Arc::new(mcp_config),
mcp_runtime_context,
mcp_connection_manager,
)
.await?;
sess.schedule_startup_prewarm(session_configuration.base_instructions.clone())
.await;
@@ -2,6 +2,7 @@ use std::sync::Arc;
use crate::agents_md::LoadedAgentsMd;
use crate::environment_selection::TurnEnvironmentSnapshot;
use crate::session::McpRuntimeSnapshot;
use crate::session::turn_context::TurnContext;
use codex_exec_server::ResolvedSelectedCapabilityRoot;
@@ -12,6 +13,8 @@ pub(crate) struct StepContext {
pub(crate) environments: TurnEnvironmentSnapshot,
/// Capability roots bound to ready environments in this exact step.
pub(crate) selected_capability_roots: Vec<ResolvedSelectedCapabilityRoot>,
/// The exact MCP config and manager used to advertise and execute tools for this step.
pub(crate) mcp: Arc<McpRuntimeSnapshot>,
/// The canonical AGENTS.md value observed with this environment snapshot.
pub(crate) loaded_agents_md: Option<Arc<LoadedAgentsMd>>,
}
@@ -21,12 +24,14 @@ impl StepContext {
turn: Arc<TurnContext>,
environments: TurnEnvironmentSnapshot,
selected_capability_roots: Vec<ResolvedSelectedCapabilityRoot>,
mcp: Arc<McpRuntimeSnapshot>,
loaded_agents_md: Option<Arc<LoadedAgentsMd>>,
) -> Self {
Self {
turn,
environments,
selected_capability_roots,
mcp,
loaded_agents_md,
}
}
+99 -22
View File
@@ -18,6 +18,7 @@ use crate::test_support::models_manager_with_provider;
use crate::tools::format_exec_output_str;
use codex_config::ConfigLayerStack;
use codex_config::ConfigLayerStackOrdering;
use codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID;
use codex_config::LoaderOverrides;
use codex_config::NetworkConstraints;
use codex_config::NetworkDomainPermissionToml;
@@ -25,6 +26,8 @@ use codex_config::NetworkDomainPermissionsToml;
use codex_config::RequirementSource;
use codex_config::Sourced;
use codex_config::loader::project_trust_key;
use codex_config::types::McpServerConfig;
use codex_config::types::McpServerTransportConfig;
use codex_config::types::ToolSuggestDisabledTool;
use codex_core_skills::HostSkillsSnapshot;
use core_test_support::test_codex::local_selections;
@@ -188,9 +191,10 @@ impl StepContext {
pub(crate) fn for_test(turn: Arc<TurnContext>) -> Arc<Self> {
let environments = turn.environments.clone();
Arc::new(Self::new(
turn,
Arc::clone(&turn),
environments,
Vec::new(),
crate::session::McpRuntimeSnapshot::new_uninitialized_for_test(&turn.config),
/*loaded_agents_md*/ None,
))
}
@@ -379,8 +383,8 @@ async fn request_mcp_server_elicitation_auto_accepts_when_auto_deny_is_enabled()
let (session, turn_context, rx) = make_session_and_context_with_rx().await;
session
.services
.mcp_connection_manager
.load_full()
.latest_mcp_runtime()
.manager()
.set_elicitations_auto_deny(/*auto_deny*/ true);
let response = session
@@ -5363,14 +5367,11 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
/*bundled_skills_enabled*/ true,
));
let network_approval = Arc::new(NetworkApprovalService::default());
let mcp_runtime =
crate::session::McpRuntimeSnapshot::new_uninitialized_for_test(config.as_ref());
let services = SessionServices {
mcp_connection_manager: Arc::new(arc_swap::ArcSwap::from_pointee(
McpConnectionManager::new_uninitialized_with_permission_profile(
&config.permissions.approval_policy,
config.permissions.permission_profile(),
config.prefix_mcp_tool_names(),
),
)),
mcp_connection_manager: Arc::new(arc_swap::ArcSwap::from(mcp_runtime.manager_arc())),
mcp_runtime: arc_swap::ArcSwapOption::from(Some(mcp_runtime)),
mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()),
unified_exec_manager: UnifiedExecProcessManager::new(
config.background_terminal_max_timeout,
@@ -7440,14 +7441,11 @@ where
/*bundled_skills_enabled*/ true,
));
let network_approval = Arc::new(NetworkApprovalService::default());
let mcp_runtime =
crate::session::McpRuntimeSnapshot::new_uninitialized_for_test(config.as_ref());
let services = SessionServices {
mcp_connection_manager: Arc::new(arc_swap::ArcSwap::from_pointee(
McpConnectionManager::new_uninitialized_with_permission_profile(
&config.permissions.approval_policy,
config.permissions.permission_profile(),
config.prefix_mcp_tool_names(),
),
)),
mcp_connection_manager: Arc::new(arc_swap::ArcSwap::from(mcp_runtime.manager_arc())),
mcp_runtime: arc_swap::ArcSwapOption::from(Some(mcp_runtime)),
mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()),
unified_exec_manager: UnifiedExecProcessManager::new(
config.background_terminal_max_timeout,
@@ -7603,8 +7601,14 @@ pub(crate) async fn make_session_and_context_with_rx() -> (
}
#[tokio::test]
async fn refresh_mcp_servers_is_deferred_until_next_turn() {
async fn refresh_mcp_servers_keeps_the_previous_runtime_alive() {
let (session, turn_context) = make_session_and_context().await;
let turn_context = Arc::new(turn_context);
let old_runtime = session.services.latest_mcp_runtime();
let step_context = session
.capture_step_context(Arc::clone(&turn_context))
.await;
assert!(Arc::ptr_eq(&step_context.mcp, &old_runtime));
let old_token = session.mcp_startup_cancellation_token().await;
assert!(!old_token.is_cancelled());
@@ -7612,8 +7616,16 @@ async fn refresh_mcp_servers_is_deferred_until_next_turn() {
serde_json::to_value(OAuthCredentialsStoreMode::Auto).expect("serialize store mode");
let auth_keyring_backend_kind =
serde_json::to_value(AuthKeyringBackendKind::Secrets).expect("serialize keyring backend");
let refreshed_mcp_servers = serde_json::from_value::<HashMap<String, McpServerConfig>>(json!({
"refreshed": {
"url": "https://refreshed.example/mcp",
"enabled": false
}
}))
.expect("parse refreshed MCP servers");
let refresh_config = McpServerRefreshConfig {
mcp_servers: json!({}),
mcp_servers: serde_json::to_value(&refreshed_mcp_servers)
.expect("serialize refreshed MCP servers"),
mcp_oauth_credentials_store_mode,
auth_keyring_backend_kind,
};
@@ -7635,7 +7647,7 @@ async fn refresh_mcp_servers_is_deferred_until_next_turn() {
.refresh_mcp_servers_if_requested(&turn_context, /*elicitation_reviewer*/ None)
.await;
assert!(old_token.is_cancelled());
assert!(!old_token.is_cancelled());
assert!(
session
.pending_mcp_server_refresh_config
@@ -7645,6 +7657,71 @@ async fn refresh_mcp_servers_is_deferred_until_next_turn() {
);
let new_token = session.mcp_startup_cancellation_token().await;
assert!(!new_token.is_cancelled());
let new_runtime = session.services.latest_mcp_runtime();
assert!(!Arc::ptr_eq(&old_runtime, &new_runtime));
assert!(Arc::ptr_eq(&step_context.mcp, &old_runtime));
assert_eq!(
codex_mcp::configured_mcp_servers(new_runtime.config()),
refreshed_mcp_servers
);
}
#[tokio::test]
async fn built_tools_uses_the_step_mcp_runtime() -> anyhow::Result<()> {
let (session, turn_context) = make_session_and_context().await;
let session = Arc::new(session);
let turn_context = Arc::new(turn_context);
let step_context = session.capture_step_context(turn_context).await;
let mut refresh_config = step_context.turn.config.as_ref().clone();
refresh_config.mcp_servers.set(HashMap::from([(
"newer".to_string(),
McpServerConfig {
auth: Default::default(),
transport: McpServerTransportConfig::Stdio {
command: "missing-test-mcp-server".to_string(),
args: Vec::new(),
env: None,
env_vars: Vec::new(),
cwd: None,
},
environment_id: DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(),
enabled: true,
required: false,
supports_parallel_tool_calls: false,
disabled_reason: None,
startup_timeout_sec: None,
tool_timeout_sec: None,
default_tools_approval_mode: None,
enabled_tools: None,
disabled_tools: None,
scopes: None,
oauth: None,
oauth_resource: None,
tools: HashMap::new(),
},
)]))?;
session
.refresh_mcp_servers_now(
step_context.turn.as_ref(),
&refresh_config,
/*elicitation_reviewer*/ None,
)
.await;
let router = crate::session::turn::built_tools(
session.as_ref(),
step_context.as_ref(),
&CancellationToken::new(),
)
.await?;
assert!(
!router
.registered_tool_names_for_test()
.iter()
.any(|name| name.to_string() == "list_mcp_resources")
);
Ok(())
}
#[tokio::test]
@@ -10161,8 +10238,8 @@ async fn fatal_tool_error_stops_turn_and_reports_error() {
let tools = {
session
.services
.mcp_connection_manager
.load_full()
.latest_mcp_runtime()
.manager()
.list_all_tools()
.await
};
+1 -1
View File
@@ -1177,7 +1177,7 @@ pub(crate) async fn built_tools(
cancellation_token: &CancellationToken,
) -> CodexResult<Arc<ToolRouter>> {
let turn_context = step_context.turn.as_ref();
let mcp_connection_manager = sess.services.mcp_connection_manager.load_full();
let mcp_connection_manager = step_context.mcp.manager();
let has_mcp_servers = mcp_connection_manager.has_servers();
let all_mcp_tools = mcp_connection_manager
.list_all_tools()
+2 -1
View File
@@ -686,7 +686,8 @@ impl Session {
.unwrap_or_else(|| session_configuration.cwd().clone());
let per_turn_config = Self::build_per_turn_config(&session_configuration, cwd.clone());
{
let mcp_connection_manager = self.services.mcp_connection_manager.load_full();
let mcp_runtime = self.services.latest_mcp_runtime();
let mcp_connection_manager = mcp_runtime.manager();
mcp_connection_manager.set_approval_policy(&session_configuration.approval_policy);
mcp_connection_manager
.set_permission_profile(session_configuration.permission_profile());
+32 -6
View File
@@ -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
}
}
+3 -2
View File
@@ -124,11 +124,12 @@ impl McpHandler {
) -> Result<Box<dyn crate::tools::context::ToolOutput>, FunctionCallError> {
let ToolInvocation {
session,
turn,
step_context,
call_id,
payload,
..
} = invocation;
let turn = Arc::clone(&step_context.turn);
let payload = match payload {
ToolPayload::Function { arguments } => arguments,
@@ -143,7 +144,7 @@ impl McpHandler {
// TODO(sayan): Use StepContext for MCP file arguments when MCP follows dynamic environments.
let result = handle_mcp_tool_call(
Arc::clone(&session),
&turn,
&step_context,
call_id.clone(),
self.tool_info.server_name.clone(),
self.tool_info.tool.name.to_string(),
@@ -53,11 +53,13 @@ impl ListMcpResourceTemplatesHandler {
) -> Result<Box<dyn crate::tools::context::ToolOutput>, FunctionCallError> {
let ToolInvocation {
session,
turn,
step_context,
call_id,
payload,
..
} = invocation;
let turn = std::sync::Arc::clone(&step_context.turn);
let manager = step_context.mcp.manager();
let arguments = match payload {
ToolPayload::Function { arguments } => arguments,
@@ -89,7 +91,7 @@ impl ListMcpResourceTemplatesHandler {
let params = cursor
.clone()
.map(|value| PaginatedRequestParams::default().with_cursor(Some(value)));
let result = session
let result = manager
.list_resource_templates(&server_name, params)
.await
.map_err(|err| {
@@ -108,10 +110,7 @@ impl ListMcpResourceTemplatesHandler {
));
}
let templates = session
.services
.mcp_connection_manager
.load_full()
let templates = manager
.list_all_resource_templates(|server_name| {
model_can_access_mcp_server(turn.as_ref(), server_name)
})
@@ -53,11 +53,13 @@ impl ListMcpResourcesHandler {
) -> Result<Box<dyn crate::tools::context::ToolOutput>, FunctionCallError> {
let ToolInvocation {
session,
turn,
step_context,
call_id,
payload,
..
} = invocation;
let turn = std::sync::Arc::clone(&step_context.turn);
let manager = step_context.mcp.manager();
let arguments = match payload {
ToolPayload::Function { arguments } => arguments,
@@ -89,7 +91,7 @@ impl ListMcpResourcesHandler {
let params = cursor
.clone()
.map(|value| PaginatedRequestParams::default().with_cursor(Some(value)));
let result = session
let result = manager
.list_resources(&server_name, params)
.await
.map_err(|err| {
@@ -106,10 +108,7 @@ impl ListMcpResourcesHandler {
));
}
let resources = session
.services
.mcp_connection_manager
.load_full()
let resources = manager
.list_all_resources(|server_name| {
model_can_access_mcp_server(turn.as_ref(), server_name)
})
@@ -52,11 +52,13 @@ impl ReadMcpResourceHandler {
) -> Result<Box<dyn crate::tools::context::ToolOutput>, FunctionCallError> {
let ToolInvocation {
session,
turn,
step_context,
call_id,
payload,
..
} = invocation;
let turn = std::sync::Arc::clone(&step_context.turn);
let manager = step_context.mcp.manager();
let arguments = match payload {
ToolPayload::Function { arguments } => arguments,
@@ -84,7 +86,7 @@ impl ReadMcpResourceHandler {
let payload_result: Result<ReadResourcePayload, FunctionCallError> = async {
ensure_model_can_access_mcp_server(turn.as_ref(), &server)?;
let result = session
let result = manager
.read_resource(&server, ReadResourceRequestParams::new(uri.clone()))
.await
.map_err(|err| {
@@ -1,4 +1,5 @@
use std::collections::HashSet;
use std::sync::Arc;
use codex_analytics::PluginInstallRequestSource;
use codex_analytics::PluginInstallRequested;
@@ -94,10 +95,12 @@ impl RequestPluginInstallHandler {
let ToolInvocation {
payload,
session,
turn,
step_context,
call_id,
..
} = invocation;
let turn = Arc::clone(&step_context.turn);
let manager = step_context.mcp.manager();
let arguments = match payload {
ToolPayload::Function { arguments } => arguments,
@@ -227,7 +230,8 @@ impl RequestPluginInstallHandler {
let auth = session.services.auth_manager.auth().await;
let completed = if user_confirmed {
verify_request_plugin_install_completed(&session, &turn, &tool, auth.as_ref()).await
verify_request_plugin_install_completed(&session, &turn, manager, &tool, auth.as_ref())
.await
} else {
false
};
@@ -345,13 +349,14 @@ fn disabled_install_request(tool: &DiscoverableTool) -> ToolSuggestDisabledTool
async fn verify_request_plugin_install_completed(
session: &crate::session::session::Session,
turn: &crate::session::turn_context::TurnContext,
manager: &codex_mcp::McpConnectionManager,
tool: &DiscoverableTool,
auth: Option<&codex_login::CodexAuth>,
) -> bool {
match tool {
DiscoverableTool::Connector(connector) => refresh_missing_requested_connectors(
session,
turn,
manager,
auth,
std::slice::from_ref(&connector.id),
connector.id.as_str(),
@@ -370,8 +375,8 @@ async fn verify_request_plugin_install_completed(
plugin.id.as_str(),
),
refresh_missing_requested_connectors(
session,
turn,
manager,
auth,
&plugin.app_connector_ids,
plugin.id.as_str(),
@@ -393,8 +398,8 @@ async fn verify_request_plugin_install_completed(
session.services.plugins_manager.as_ref(),
);
let _ = refresh_missing_requested_connectors(
session,
turn,
manager,
auth,
&plugin.app_connector_ids,
plugin.id.as_str(),
@@ -435,8 +440,8 @@ fn is_remote_plugin_install_suggestion(plugin_id: &str) -> bool {
}
async fn refresh_missing_requested_connectors(
session: &crate::session::session::Session,
turn: &crate::session::turn_context::TurnContext,
manager: &codex_mcp::McpConnectionManager,
auth: Option<&codex_login::CodexAuth>,
expected_connector_ids: &[String],
tool_id: &str,
@@ -445,7 +450,6 @@ async fn refresh_missing_requested_connectors(
return Some(Vec::new());
}
let manager = session.services.mcp_connection_manager.load_full();
let mcp_tools = manager.list_all_tools().await;
let accessible_connectors = connectors::with_app_enabled_state(
connectors::accessible_connectors_from_mcp_tools(&mcp_tools),
+2 -2
View File
@@ -110,8 +110,8 @@ async fn parallel_support_does_not_match_namespaced_local_tool_names() -> anyhow
let step_context = StepContext::for_test(Arc::clone(&turn));
let mcp_tools = session
.services
.mcp_connection_manager
.load_full()
.latest_mcp_runtime()
.manager()
.list_all_tools()
.await;
let router = ToolRouter::from_context(
+3 -1
View File
@@ -686,10 +686,12 @@ async fn environment_tools_follow_the_step_context() {
let environments = turn.environments.clone();
turn.environments.turn_environments.clear();
let turn = Arc::new(turn);
let step_context = Arc::new(StepContext::new(
Arc::new(turn),
Arc::clone(&turn),
environments,
Vec::new(),
crate::session::McpRuntimeSnapshot::new_uninitialized_for_test(&turn.config),
/*loaded_agents_md*/ None,
));
@@ -16,7 +16,7 @@ use core_test_support::test_codex::test_codex;
use core_test_support::wait_for_mcp_server;
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn refresh_shuts_down_superseded_mcp_stdio_server() -> anyhow::Result<()> {
async fn refresh_keeps_superseded_mcp_server_alive_for_in_flight_calls() -> anyhow::Result<()> {
skip_if_no_network!(Ok(()));
let server = responses::start_mock_server().await;
@@ -84,7 +84,7 @@ async fn refresh_shuts_down_superseded_mcp_stdio_server() -> anyhow::Result<()>
"sync",
Some(serde_json::json!({
"barrier": barrier,
"sleep_after_ms": 30_000
"sleep_after_ms": 300_000
})),
/*meta*/ None,
)
@@ -119,9 +119,16 @@ async fn refresh_shuts_down_superseded_mcp_stdio_server() -> anyhow::Result<()>
let replacement_pid = wait_for_pid_file(&pid_file).await?;
assert_ne!(replacement_pid, superseded_pid);
assert!(process_is_alive(&superseded_pid)?);
long_call.abort();
assert!(
long_call
.await
.expect_err("call should be aborted")
.is_cancelled()
);
wait_for_process_exit(&superseded_pid).await?;
assert!(process_is_alive(&replacement_pid)?);
assert!(long_call.await?.is_err());
fixture.codex.shutdown_and_wait().await?;
wait_for_process_exit(&replacement_pid).await