Pin MCP runtime to model steps

This commit is contained in:
jif-oai
2026-06-25 17:42:14 +01:00
parent addf6b52e5
commit 10f9bfc4cb
24 changed files with 480 additions and 300 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
@@ -249,22 +243,25 @@ impl McpRequestProcessor {
}
None => (self.load_latest_config(/*fallback_cwd*/ None).await?, None),
};
let mcp_config = match thread {
Some(thread) => thread.runtime_mcp_config(&config).await,
let auth = self.auth_manager.auth().await;
let (mcp_config, runtime_context) = match thread {
Some(thread) => {
let runtime = thread.current_mcp_runtime();
(runtime.config().clone(), runtime.runtime_context().clone())
}
None => {
self.thread_manager
let mcp_config = self
.thread_manager
.mcp_manager()
.runtime_config(&config)
.await
.await;
let runtime_context = McpRuntimeContext::new(
self.thread_manager.environment_manager(),
config.cwd.to_path_buf(),
);
(mcp_config, runtime_context)
}
};
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());
tokio::spawn(async move {
Self::list_mcp_server_status_task(
+43 -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.
@@ -148,10 +155,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,
@@ -269,7 +276,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();
@@ -356,10 +363,33 @@ 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(
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,
@@ -647,7 +677,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,
) {
@@ -685,12 +715,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> {
@@ -701,18 +731,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
@@ -200,15 +200,8 @@ pub(crate) async fn maybe_install_mcp_dependencies(
warn!("failed to refresh MCP dependencies for mentioned skills: {err}");
return false;
}
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;
true
}
+55 -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() {
@@ -142,7 +146,7 @@ pub(crate) async fn handle_mcp_tool_call(
};
let metadata =
lookup_mcp_tool_metadata(sess.as_ref(), turn_context.as_ref(), &server, &tool_name).await;
lookup_mcp_tool_metadata(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 +173,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 +229,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 +244,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 +301,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 +343,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 +380,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 +420,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 +547,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 +570,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 +589,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 +601,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 +665,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 +692,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 +706,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 +741,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 +753,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 +762,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 +934,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 +981,7 @@ enum McpToolApprovalDecision {
Cancel,
}
#[derive(Clone)]
pub(crate) struct McpToolApprovalMetadata {
annotations: Option<ToolAnnotations>,
connector_id: Option<String>,
@@ -1176,13 +1175,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 +1202,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)
@@ -1450,12 +1446,11 @@ 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);
@@ -1562,16 +1557,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();
@@ -1409,22 +1417,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),
@@ -1445,10 +1451,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),
@@ -1463,7 +1471,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");
@@ -1478,6 +1486,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),
@@ -1492,7 +1501,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");
@@ -1513,6 +1522,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),
@@ -1527,7 +1537,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);
@@ -1541,10 +1551,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),
@@ -2465,7 +2477,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"),
@@ -2541,7 +2553,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"),
@@ -2600,7 +2612,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"),
@@ -2662,7 +2674,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"),
@@ -2744,7 +2756,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"),
@@ -2827,7 +2839,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"),
@@ -2887,7 +2899,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"),
@@ -2945,7 +2957,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"),
@@ -3034,7 +3046,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;
+52 -53
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,8 +201,8 @@ impl Session {
}
self.services
.mcp_connection_manager
.load_full()
.latest_mcp_runtime()
.manager_arc()
.resolve_elicitation(server_name, id, response)
.await
}
@@ -213,8 +213,8 @@ impl Session {
params: Option<PaginatedRequestParams>,
) -> anyhow::Result<ListResourcesResult> {
self.services
.mcp_connection_manager
.load_full()
.latest_mcp_runtime()
.manager_arc()
.list_resources(server, params)
.await
}
@@ -225,8 +225,8 @@ impl Session {
params: Option<PaginatedRequestParams>,
) -> anyhow::Result<ListResourceTemplatesResult> {
self.services
.mcp_connection_manager
.load_full()
.latest_mcp_runtime()
.manager_arc()
.list_resource_templates(server, params)
.await
}
@@ -237,8 +237,8 @@ impl Session {
params: ReadResourceRequestParams,
) -> anyhow::Result<ReadResourceResult> {
self.services
.mcp_connection_manager
.load_full()
.latest_mcp_runtime()
.manager_arc()
.read_resource(server, params)
.await
}
@@ -251,8 +251,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 +260,13 @@ impl Session {
async fn refresh_mcp_servers_inner(
&self,
turn_context: &TurnContext,
mcp_servers: HashMap<String, McpServerConfig>,
store_mode: OAuthCredentialsStoreMode,
keyring_backend_kind: AuthKeyringBackendKind,
refresh_config: &Config,
elicitation_reviewer: Option<ElicitationReviewerHandle>,
) {
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 mcp_config = Arc::new(self.runtime_mcp_config(refresh_config).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,34 +282,35 @@ 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(),
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),
@@ -323,14 +320,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(
@@ -375,14 +370,26 @@ 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();
if let Err(err) = refresh_config.mcp_servers.set(mcp_servers) {
warn!("failed to apply MCP server refresh config: {err}");
return;
}
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;
}
self.refresh_mcp_servers_inner(turn_context, &refresh_config, elicitation_reviewer)
.await;
}
pub(crate) async fn set_openai_form_elicitation_support(
@@ -416,19 +423,11 @@ 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;
self.refresh_mcp_servers_inner(turn_context, refresh_config, 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 -10
View File
@@ -154,15 +154,12 @@ 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;
@@ -210,6 +207,7 @@ mod handlers;
mod inject;
mod input_queue;
mod mcp;
mod mcp_runtime;
pub(crate) mod multi_agents;
mod review;
mod rollout_budget;
@@ -231,6 +229,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;
@@ -336,7 +335,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;
@@ -819,8 +818,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(())
}
@@ -2833,6 +2835,7 @@ impl Session {
&environments.captured_environment_availability(),
)
.await;
let mcp = self.services.latest_mcp_runtime();
let extra_skill_sources = self
.services
.thread_extension_data
@@ -2853,6 +2856,7 @@ impl Session {
environments,
selected_capability_roots,
skills,
mcp,
loaded_agents_md,
))
}
@@ -3118,6 +3122,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);
@@ -3211,10 +3226,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;
@@ -3298,7 +3312,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",
@@ -3513,7 +3528,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
+18 -28
View File
@@ -654,10 +654,11 @@ 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(
let mcp_runtime_context = McpRuntimeContext::new(
Arc::clone(&environment_manager),
session_configuration.cwd().to_path_buf(),
);
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
@@ -673,7 +674,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",
@@ -684,7 +691,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 =
@@ -1020,6 +1027,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,
@@ -1153,14 +1161,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();
@@ -1168,20 +1168,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,
@@ -1192,11 +1178,11 @@ 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(),
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),
@@ -1210,7 +1196,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_core_skills::SkillsSnapshot;
use codex_exec_server::ResolvedSelectedCapabilityRoot;
@@ -15,6 +16,8 @@ pub(crate) struct StepContext {
pub(crate) selected_capability_roots: Vec<ResolvedSelectedCapabilityRoot>,
/// Complete skill catalog and exact read routes captured for this step.
pub(crate) skills: Arc<SkillsSnapshot>,
/// 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>>,
}
@@ -25,6 +28,7 @@ impl StepContext {
environments: TurnEnvironmentSnapshot,
selected_capability_roots: Vec<ResolvedSelectedCapabilityRoot>,
skills: Arc<SkillsSnapshot>,
mcp: Arc<McpRuntimeSnapshot>,
loaded_agents_md: Option<Arc<LoadedAgentsMd>>,
) -> Self {
Self {
@@ -32,6 +36,7 @@ impl StepContext {
environments,
selected_capability_roots,
skills,
mcp,
loaded_agents_md,
}
}
+25 -21
View File
@@ -193,10 +193,11 @@ impl StepContext {
turn.model_context_window(),
));
Arc::new(Self::new(
turn,
Arc::clone(&turn),
environments,
Vec::new(),
skills,
crate::session::McpRuntimeSnapshot::new_uninitialized_for_test(&turn.config),
/*loaded_agents_md*/ None,
))
}
@@ -385,8 +386,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
@@ -5369,14 +5370,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,
@@ -7447,14 +7445,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,
@@ -7611,8 +7606,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());
@@ -7643,7 +7644,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
@@ -7653,6 +7654,9 @@ 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));
}
#[tokio::test]
@@ -10113,8 +10117,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
};
+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;
@@ -31,7 +32,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;
@@ -45,8 +48,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))]
@@ -103,12 +108,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
@@ -690,11 +690,13 @@ async fn environment_tools_follow_the_step_context() {
turn.model_context_window(),
));
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(),
skills,
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