From d1ec95ec0ef273a6ec9b9880f6c64569b7f356ad Mon Sep 17 00:00:00 2001 From: Friel Date: Sun, 3 May 2026 22:33:18 +0000 Subject: [PATCH] Fix fork cache build callsites --- codex-rs/codex-mcp/src/connection_manager.rs | 45 ++++++++++++++++ codex-rs/core/src/agent/control.rs | 28 +++++----- codex-rs/core/src/agent/control_tests.rs | 56 ++++++++++---------- codex-rs/core/src/client.rs | 4 ++ codex-rs/exec/tests/suite/fork.rs | 32 ++--------- codex-rs/memories/write/src/runtime.rs | 1 + 6 files changed, 94 insertions(+), 72 deletions(-) diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 483a82796a..27655ab328 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -7,6 +7,7 @@ //! `codex-core`. use std::collections::HashMap; +use std::future::Future; use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; @@ -315,6 +316,34 @@ impl McpConnectionManager { failures } + pub fn required_startup_failures_future( + &self, + required_servers: Vec, + ) -> impl Future> + Send + 'static { + let clients = self.clients.clone(); + async move { + let mut failures = Vec::new(); + for server_name in required_servers { + let Some(async_managed_client) = clients.get(&server_name).cloned() else { + failures.push(McpStartupFailure { + server: server_name.clone(), + error: format!("required MCP server `{server_name}` was not initialized"), + }); + continue; + }; + + match async_managed_client.client().await { + Ok(_) => {} + Err(error) => failures.push(McpStartupFailure { + server: server_name.clone(), + error: startup_outcome_error_message(error), + }), + } + } + failures + } + } + /// Returns a single map that contains all tools. Each key is the /// fully-qualified name for the tool. #[instrument(level = "trace", skip_all)] @@ -329,6 +358,22 @@ impl McpConnectionManager { qualify_tools(tools) } + pub fn list_all_tools_future( + &self, + ) -> impl Future> + Send + 'static { + let clients = self.clients.values().cloned().collect::>(); + async move { + let mut tools = Vec::new(); + for managed_client in clients { + let Some(server_tools) = managed_client.listed_tools().await else { + continue; + }; + tools.extend(server_tools); + } + qualify_tools(tools) + } + } + /// Force-refresh codex apps tools by bypassing the in-process cache. /// /// On success, the refreshed tools replace the cache contents and the diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index e6299f225b..6f56328d03 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -5,10 +5,7 @@ use crate::agent::role::DEFAULT_ROLE_NAME; use crate::agent::role::resolve_role_config; use crate::agent::status::is_final; use crate::codex_thread::ThreadConfigSnapshot; -use crate::find_archived_thread_path_by_id_str; -use crate::find_thread_path_by_id_str; use crate::inherited_thread_state::InheritedThreadState; -use crate::rollout::RolloutRecorder; use crate::session::emit_subagent_session_started; use crate::session_prefix::format_subagent_context_line; use crate::session_prefix::format_subagent_notification_message; @@ -403,7 +400,7 @@ impl AgentControl { )) })?; -let response_continuation = inherited_thread_state.response_continuation(); + let response_continuation = inherited_thread_state.response_continuation(); let use_response_continuation_baseline = response_continuation.is_some() && matches!(fork_mode, SpawnAgentForkMode::FullHistory); let mut forked_rollout_items = if let (Some(response_continuation), true) = @@ -443,8 +440,7 @@ let response_continuation = inherited_thread_state.response_continuation(); Vec::new() }; forked_rollout_items.retain(|item| { - if let RolloutItem::ResponseItem(ResponseItem::Message { role, content, .. }) = - item + if let RolloutItem::ResponseItem(ResponseItem::Message { role, content, .. }) = item && role == "developer" && let [ContentItem::InputText { text }] = content.as_slice() && multi_agent_v2_usage_hint_texts_to_filter @@ -1291,15 +1287,17 @@ async fn parent_mcp_tool_snapshot_for_source( }; let parent_thread = state.get_thread(*parent_thread_id).await.ok()?; - let tools = parent_thread - .codex - .session - .services - .mcp_connection_manager - .read() - .await - .list_all_tools() - .await; + let list_all_tools = { + let mcp_connection_manager = parent_thread + .codex + .session + .services + .mcp_connection_manager + .read() + .await; + mcp_connection_manager.list_all_tools_future() + }; + let tools = list_all_tools.await; Some(McpToolSnapshot { tools }) } diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index 922c5b7723..b047202cd8 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -779,15 +779,17 @@ async fn spawn_agent_can_fork_parent_thread_history_with_sanitized_items() { .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 list_all_tools = { + let mcp_connection_manager = parent_thread + .codex + .session + .services + .mcp_connection_manager + .read() + .await; + mcp_connection_manager.list_all_tools_future() + }; + let parent_mcp_tools = list_all_tools.await; let mut snapshot_tool_names = mcp_tool_snapshot.tools.keys().cloned().collect::>(); snapshot_tool_names.sort(); let mut parent_tool_names = parent_mcp_tools.keys().cloned().collect::>(); @@ -951,26 +953,22 @@ while True: let parent = manager.start_thread(config.clone()).await?; let parent_thread_id = parent.thread_id; let parent_prompt_cache_key = parent.thread.codex.session.prompt_cache_key(); - let parent_mcp_tools = parent - .thread - .codex - .session - .services - .mcp_connection_manager - .read() - .await - .list_all_tools() - .await; - let startup_failures = parent - .thread - .codex - .session - .services - .mcp_connection_manager - .read() - .await - .required_startup_failures(&["rmcp".to_string()]) - .await; + let (list_all_tools, required_startup_failures) = { + let mcp_connection_manager = parent + .thread + .codex + .session + .services + .mcp_connection_manager + .read() + .await; + ( + mcp_connection_manager.list_all_tools_future(), + mcp_connection_manager.required_startup_failures_future(vec!["rmcp".to_string()]), + ) + }; + let parent_mcp_tools = list_all_tools.await; + let startup_failures = required_startup_failures.await; assert!( parent_mcp_tools.contains_key("mcp__rmcp__echo"), "parent MCP manager should expose live MCP tools before forking: tools={parent_mcp_tools:#?}; failures={startup_failures:#?}" diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index d362b3afe6..6c91c7d3db 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -1767,6 +1767,8 @@ fn map_response_stream( api_stream: codex_api::ResponseStream, session_telemetry: SessionTelemetry, inference_trace_attempt: InferenceTraceAttempt, + client_state: Option>, + request: Option, ) -> (ResponseStream, oneshot::Receiver) { let codex_api::ResponseStream { rx_event, @@ -1781,6 +1783,8 @@ fn map_response_stream( api_stream, session_telemetry, inference_trace_attempt, + client_state, + request, ) } diff --git a/codex-rs/exec/tests/suite/fork.rs b/codex-rs/exec/tests/suite/fork.rs index e706b2b558..01c8d1d9ae 100644 --- a/codex-rs/exec/tests/suite/fork.rs +++ b/codex-rs/exec/tests/suite/fork.rs @@ -81,29 +81,6 @@ fn extract_forked_from_id(path: &std::path::Path) -> Option { .map(ToString::to_string) } -fn rollout_contains_fork_reference(path: &std::path::Path) -> bool { - extract_fork_reference(path).is_some() -} - -fn extract_fork_reference(path: &std::path::Path) -> Option<(String, usize)> { - let Ok(content) = std::fs::read_to_string(path) else { - return None; - }; - content.lines().skip(1).find_map(|line| { - let item = serde_json::from_str::(line).ok()?; - if item.get("type").and_then(Value::as_str) != Some("fork_reference") { - return None; - } - let payload = item.get("payload")?; - let rollout_path = payload.get("rollout_path")?.as_str()?.to_string(); - let nth_user_message = payload - .get("nth_user_message")? - .as_u64() - .and_then(|value| usize::try_from(value).ok())?; - Some((rollout_path, nth_user_message)) - }) -} - fn exec_fixture() -> anyhow::Result { Ok(find_resource!("tests/fixtures/cli_responses_fixture.sse")?) } @@ -155,11 +132,10 @@ fn exec_fork_by_id_creates_new_session_with_copied_history() -> anyhow::Result<( extract_forked_from_id(&forked_path).as_deref(), Some(session_id.as_str()) ); - let fork_reference = - extract_fork_reference(&forked_path).context("forked rollout should record a reference")?; - assert_eq!(fork_reference.0, original_path.to_string_lossy().as_ref()); - assert_eq!(fork_reference.1, usize::MAX); - assert!(rollout_contains_fork_reference(&forked_path)); + assert!( + forked_content.contains(&marker), + "forked session should copy ancestor rollout history" + ); assert!(forked_content.contains(&marker2)); let original_content = std::fs::read_to_string(&original_path)?; diff --git a/codex-rs/memories/write/src/runtime.rs b/codex-rs/memories/write/src/runtime.rs index 737fb67870..0c6ecec087 100644 --- a/codex-rs/memories/write/src/runtime.rs +++ b/codex-rs/memories/write/src/runtime.rs @@ -175,6 +175,7 @@ impl MemoryStartupContext { Some(Arc::clone(&self.auth_manager)), self.thread_id, installation_id, + /*prompt_cache_key_override*/ None, config.model_provider.clone(), session_source, config.model_verbosity,