Skip to content

Reuse MCP tool snapshot within a sampling request #31292

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sayan-oai merged 1 commit into from
Jul 8, 2026
Merged

Reuse MCP tool snapshot within a sampling request #31292

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions codex-rs/core/src/session/step_context.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand All @@ -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>>,
}
Expand All @@ -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] {

@pakrym-oai pakrym-oai Jul 7, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If this field is purely computable based on McpRuntimeSnapshot it should stay on McpRuntimeSnapshot and be invalidated when it is changed.

self.mcp_tool_snapshot
.get_or_init(|| self.mcp.manager().list_all_tools())
.await
}
}
8 changes: 4 additions & 4 deletions codex-rs/core/src/session/turn.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1182,8 +1182,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
Expand All @@ -1196,7 +1196,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)
Expand Down Expand Up @@ -1289,7 +1289,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),
Expand Down
4 changes: 2 additions & 2 deletions codex-rs/core/src/session/world_state.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
Loading