From bdaad6820cd884ea11787477f4c495e4de0a8be5 Mon Sep 17 00:00:00 2001
From: sayan-oai <sayan@openai.com>
Date: Wed, 8 Jul 2026 13:35:31 -0700
Subject: [PATCH] Reuse MCP tool snapshot within a sampling request (#31292)

Follow-up to #30226.

## Why

#30226 makes Apps World State inspect the MCP tool list, while
tool-router construction reads the same list again later in the sampling
request. `list_all_tools()` walks the MCP clients and may reconnect or
wait for tools, so doing that work twice adds latency and lets context
and tool construction observe different MCP states for one request.

## What

- Add a lazy MCP tool snapshot to `StepContext`.
- Reuse that snapshot for Apps World State and tool-router construction.
- Let each new `StepContext` refresh naturally for the next sampling
request, without manager-level caching or invalidation.

## Testing

- `just test -p codex-core apps_instructions`
- `just test -p codex-core
apps_guidance_appears_after_background_recovery_within_a_turn`
---
 codex-rs/core/src/session/step_context.rs | 11 +++++++++++
 codex-rs/core/src/session/turn.rs         |  8 ++++----
 codex-rs/core/src/session/world_state.rs  |  4 ++--
 3 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/codex-rs/core/src/session/step_context.rs b/codex-rs/core/src/session/step_context.rs
index 3210eb0286..1d7316d914 100644
--- a/codex-rs/core/src/session/step_context.rs
+++ b/codex-rs/core/src/session/step_context.rs
@@ -5,6 +5,8 @@ use crate::environment_selection::TurnEnvironmentSnapshot;
 use crate::session::McpRuntimeSnapshot;
 use crate::session::turn_context::TurnContext;
 use codex_exec_server::ResolvedSelectedCapabilityRoot;
+use codex_mcp::ToolInfo;
+use tokio::sync::OnceCell;
 
 /// Request-scoped state that may change between model sampling requests.
 #[derive(Debug)]
@@ -15,6 +17,8 @@ pub(crate) struct StepContext {
     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 fixed MCP tool list used for this exact sampling request.
+    mcp_tool_snapshot: OnceCell<Vec<ToolInfo>>,
     /// The canonical AGENTS.md value observed with this environment snapshot.
     pub(crate) loaded_agents_md: Option<Arc<LoadedAgentsMd>>,
 }
@@ -32,7 +36,14 @@ impl StepContext {
             environments,
             selected_capability_roots,
             mcp,
+            mcp_tool_snapshot: OnceCell::new(),
             loaded_agents_md,
         }
     }
+
+    pub(crate) async fn mcp_tools(&self) -> &[ToolInfo] {
+        self.mcp_tool_snapshot
+            .get_or_init(|| self.mcp.manager().list_all_tools())
+            .await
+    }
 }
diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs
index e231f6af86..6ee3ac2863 100644
--- a/codex-rs/core/src/session/turn.rs
+++ b/codex-rs/core/src/session/turn.rs
@@ -1222,8 +1222,8 @@ pub(crate) async fn built_tools(
     let turn_context = step_context.turn.as_ref();
     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()
+    let all_mcp_tools = step_context
+        .mcp_tools()
         .or_cancel(cancellation_token)
         .await?;
     let loaded_plugins = sess
@@ -1236,7 +1236,7 @@ pub(crate) async fn built_tools(
 
     let apps_enabled = turn_context.apps_enabled();
     let accessible_connectors =
-        apps_enabled.then(|| connectors::accessible_connectors_from_mcp_tools(&all_mcp_tools));
+        apps_enabled.then(|| connectors::accessible_connectors_from_mcp_tools(all_mcp_tools));
     let accessible_connectors_with_enabled_state =
         accessible_connectors.as_ref().map(|connectors| {
             connectors::with_app_enabled_state(connectors.clone(), &turn_context.config)
@@ -1329,7 +1329,7 @@ pub(crate) async fn built_tools(
             .await
         };
     let mcp_tool_exposure = build_mcp_tool_exposure(
-        &all_mcp_tools,
+        all_mcp_tools,
         connectors.as_deref(),
         &turn_context.config,
         search_tool_enabled(turn_context),
diff --git a/codex-rs/core/src/session/world_state.rs b/codex-rs/core/src/session/world_state.rs
index a0d93cac83..c7743c3fe1 100644
--- a/codex-rs/core/src/session/world_state.rs
+++ b/codex-rs/core/src/session/world_state.rs
@@ -41,9 +41,9 @@ impl Session {
         }
         let apps_available =
             if turn_context.config.include_apps_instructions && turn_context.apps_enabled() {
-                let tools = step_context.mcp.manager().list_all_tools().await;
+                let tools = step_context.mcp_tools().await;
                 connectors::with_app_enabled_state(
-                    connectors::accessible_connectors_from_mcp_tools(&tools),
+                    connectors::accessible_connectors_from_mcp_tools(tools),
                     &turn_context.config,
                 )
                 .into_iter()