Fix fork cache build callsites

This commit is contained in:
Friel
2026-05-04 00:18:21 +00:00
parent c2db9386aa
commit d1ec95ec0e
6 changed files with 94 additions and 72 deletions
@@ -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<String>,
) -> impl Future<Output = Vec<McpStartupFailure>> + 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<Output = HashMap<String, ToolInfo>> + Send + 'static {
let clients = self.clients.values().cloned().collect::<Vec<_>>();
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
+13 -15
View File
@@ -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 })
}
+27 -29
View File
@@ -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::<Vec<_>>();
snapshot_tool_names.sort();
let mut parent_tool_names = parent_mcp_tools.keys().cloned().collect::<Vec<_>>();
@@ -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:#?}"
+4
View File
@@ -1767,6 +1767,8 @@ fn map_response_stream(
api_stream: codex_api::ResponseStream,
session_telemetry: SessionTelemetry,
inference_trace_attempt: InferenceTraceAttempt,
client_state: Option<Arc<ModelClientState>>,
request: Option<ResponsesApiRequest>,
) -> (ResponseStream, oneshot::Receiver<LastResponse>) {
let codex_api::ResponseStream {
rx_event,
@@ -1781,6 +1783,8 @@ fn map_response_stream(
api_stream,
session_telemetry,
inference_trace_attempt,
client_state,
request,
)
}
+4 -28
View File
@@ -81,29 +81,6 @@ fn extract_forked_from_id(path: &std::path::Path) -> Option<String> {
.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::<Value>(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<std::path::PathBuf> {
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)?;
+1
View File
@@ -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,