clone/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-04-23 22:24:57 +00:00
Code Issues Packages Projects Releases Wiki Activity

Snapshot MCP tools for forked agents

Browse Source
This commit is contained in:
Friel
2026-04-09 19:08:07 +00:00
parent cc34536d13
commit bd8bb2e9c4
6 changed files with 77 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
codex-rs/core/src/agent/control.rs
View File

@@ -13,6 +13,7 @@ use crate::rollout::RolloutRecorder;
use crate::session_prefix::format_subagent_context_line;
use crate::session_prefix::format_subagent_notification_message;
use crate::shell_snapshot::ShellSnapshot;
use crate::state::McpToolSnapshot;
use crate::thread_manager::ThreadManagerState;
use crate::thread_rollout_truncation::truncate_rollout_to_last_n_fork_turns;
use codex_features::Feature;
@@ -1077,17 +1078,22 @@ impl AgentControl {
return InheritedThreadState::default();
};
let Some(prompt_cache_key) = state
.get_thread(*parent_thread_id)
.await
.ok()
.map(|parent_thread| parent_thread.codex.session.prompt_cache_key())
else {
let Some(parent_thread) = state.get_thread(*parent_thread_id).await.ok() else {
return InheritedThreadState::default();
};
let mcp_tools = parent_thread
.codex
.session
.services
.mcp_connection_manager
.read()
.await
.list_all_tools()
.await;
InheritedThreadState {
prompt_cache_key: Some(prompt_cache_key),
prompt_cache_key: Some(parent_thread.codex.session.prompt_cache_key()),
mcp_tool_snapshot: Some(McpToolSnapshot { tools: mcp_tools }),
}
}

27
codex-rs/core/src/agent/control_tests.rs
View File

@@ -667,6 +667,33 @@ async fn spawn_agent_can_fork_parent_thread_history_with_sanitized_items() {
child_thread.codex.session.prompt_cache_key(),
parent_thread.codex.session.prompt_cache_key(),
);
assert!(!Arc::ptr_eq(
&child_thread.codex.session.services.mcp_connection_manager,
&parent_thread.codex.session.services.mcp_connection_manager,
));
let mcp_tool_snapshot = child_thread
.codex
.session
.services
.mcp_tool_snapshot
.lock()
.await
.clone()
.expect("forked child should inherit an MCP tool snapshot");
let parent_mcp_tools = parent_thread
.codex
.session
.services
.mcp_connection_manager
.read()
.await
.list_all_tools()
.await;
let mut snapshot_tool_names = mcp_tool_snapshot.tools.keys().cloned().collect::<Vec<_>>();
snapshot_tool_names.sort();
let mut parent_tool_names = parent_mcp_tools.keys().cloned().collect::<Vec<_>>();
parent_tool_names.sort();
assert_eq!(snapshot_tool_names, parent_tool_names);
let history = child_thread.codex.session.clone_history().await;
let expected_history = [
ResponseItem::Message {

40
codex-rs/core/src/codex.rs
View File

@@ -297,6 +297,7 @@ use crate::skills_watcher::SkillsWatcher;
use crate::skills_watcher::SkillsWatcherEvent;
use crate::state::ActiveTurn;
use crate::state::MailboxDeliveryPhase;
use crate::state::McpToolSnapshot;
use crate::state::SessionServices;
use crate::state::SessionState;
use crate::tasks::GhostSnapshotTask;
@@ -437,9 +438,10 @@ pub(crate) struct CodexSpawnArgs {
pub(crate) parent_trace: Option<W3cTraceContext>,
}
#[derive(Clone, Copy, Debug, Default)]
#[derive(Clone, Default)]
pub(crate) struct InheritedThreadState {
pub(crate) prompt_cache_key: Option<ThreadId>,
pub(crate) mcp_tool_snapshot: Option<McpToolSnapshot>,
}
pub(crate) const INITIAL_SUBMIT_ID: &str = "";
@@ -1654,9 +1656,11 @@ impl Session {
),
),
};
let prompt_cache_key = inherited_thread_state
.prompt_cache_key
.unwrap_or(conversation_id);
let InheritedThreadState {
prompt_cache_key,
mcp_tool_snapshot,
} = inherited_thread_state;
let prompt_cache_key = prompt_cache_key.unwrap_or(conversation_id);
let window_generation = match &initial_history {
InitialHistory::Resumed(resumed_history) => u64::try_from(
resumed_history
@@ -2008,6 +2012,7 @@ impl Session {
&config.permissions.approval_policy,
&config.permissions.sandbox_policy,
))),
mcp_tool_snapshot: Mutex::new(mcp_tool_snapshot),
mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()),
unified_exec_manager: UnifiedExecProcessManager::new(
config.background_terminal_max_timeout,
@@ -4526,8 +4531,12 @@ impl Session {
*guard = cancel_token;
}
let mut manager = self.services.mcp_connection_manager.write().await;
*manager = refreshed_manager;
{
let mut manager = self.services.mcp_connection_manager.write().await;
*manager = refreshed_manager;
}
let mut snapshot = self.services.mcp_tool_snapshot.lock().await;
*snapshot = None;
}
async fn refresh_mcp_servers_if_requested(&self, turn_context: &TurnContext) {
@@ -6946,13 +6955,18 @@ pub(crate) async fn built_tools(
skills_outcome: Option<&SkillLoadOutcome>,
cancellation_token: &CancellationToken,
) -> CodexResult<Arc<ToolRouter>> {
let mcp_connection_manager = sess.services.mcp_connection_manager.read().await;
let has_mcp_servers = mcp_connection_manager.has_servers();
let all_mcp_tools = mcp_connection_manager
.list_all_tools()
.or_cancel(cancellation_token)
.await?;
drop(mcp_connection_manager);
let inherited_mcp_tools = sess.services.mcp_tool_snapshot.lock().await.clone();
let (has_mcp_servers, all_mcp_tools) = if let Some(snapshot) = inherited_mcp_tools {
(!snapshot.tools.is_empty(), snapshot.tools)
} else {
let mcp_connection_manager = sess.services.mcp_connection_manager.read().await;
let has_mcp_servers = mcp_connection_manager.has_servers();
let all_mcp_tools = mcp_connection_manager
.list_all_tools()
.or_cancel(cancellation_token)
.await?;
(has_mcp_servers, all_mcp_tools)
};
let loaded_plugins = sess
.services
.plugins_manager

2
codex-rs/core/src/codex_tests.rs
View File

@@ -2857,6 +2857,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
&config.permissions.approval_policy,
&config.permissions.sandbox_policy,
))),
mcp_tool_snapshot: Mutex::new(None),
mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()),
unified_exec_manager: UnifiedExecProcessManager::new(
config.background_terminal_max_timeout,
@@ -3703,6 +3704,7 @@ pub(crate) async fn make_session_and_context_with_dynamic_tools_and_rx(
&config.permissions.approval_policy,
&config.permissions.sandbox_policy,
))),
mcp_tool_snapshot: Mutex::new(None),
mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()),
unified_exec_manager: UnifiedExecProcessManager::new(
config.background_terminal_max_timeout,

1
codex-rs/core/src/state/mod.rs
View File

@@ -2,6 +2,7 @@ mod service;
mod session;
mod turn;
pub(crate) use service::McpToolSnapshot;
pub(crate) use service::SessionServices;
pub(crate) use session::SessionState;
pub(crate) use turn::ActiveTurn;

7
codex-rs/core/src/state/service.rs
View File

@@ -19,6 +19,7 @@ use codex_exec_server::Environment;
use codex_hooks::Hooks;
use codex_login::AuthManager;
use codex_mcp::McpConnectionManager;
use codex_mcp::ToolInfo as McpToolInfo;
use codex_models_manager::manager::ModelsManager;
use codex_otel::SessionTelemetry;
use codex_rollout::state_db::StateDbHandle;
@@ -28,8 +29,14 @@ use tokio::sync::RwLock;
use tokio::sync::watch;
use tokio_util::sync::CancellationToken;
#[derive(Clone, Default)]
pub(crate) struct McpToolSnapshot {
pub(crate) tools: HashMap<String, McpToolInfo>,
}
pub(crate) struct SessionServices {
pub(crate) mcp_connection_manager: Arc<RwLock<McpConnectionManager>>,
pub(crate) mcp_tool_snapshot: Mutex<Option<McpToolSnapshot>>,
pub(crate) mcp_startup_cancellation_token: Mutex<CancellationToken>,
pub(crate) unified_exec_manager: UnifiedExecProcessManager,
#[cfg_attr(not(unix), allow(dead_code))]