From 0d46c252b3f29f10bacf0ef58a17a1aa5d17ead3 Mon Sep 17 00:00:00 2001 From: ningyi-oai Date: Wed, 9 Sep 2026 02:11:11 +0000 Subject: [PATCH] Encapsulate executed tool call metadata recording (#44002) ## What changed Introduce `ExecutedToolCalls` to own optional shared recorder state and centralize feature checks, call recording, and accepted result source lookup. Update session, tool runtime, MCP, and Code Mode callers to use this interface. Move prompt attachment, retry handling, and metadata budgeting into the `request_metadata` module, preserving existing recording and request behavior. ## Testing Add coverage for all session and turn feature combinations, lazy result source lookup, and unchanged prompts when recording is disabled. Retain the existing metadata budget and cell completeness tests in the new module. GitOrigin-RevId: bc19f34317232b6599c7ec71c15db4a80123ff6f --- codex-rs/core/src/session/session.rs | 5 +- codex-rs/core/src/session/tests.rs | 10 +- codex-rs/core/src/session/turn.rs | 9 +- codex-rs/core/src/state/mod.rs | 2 +- codex-rs/core/src/state/service.rs | 4 +- codex-rs/core/src/tools/code_mode/delegate.rs | 16 +- .../src/tools/code_mode/delegate_tests.rs | 2 +- .../src/tools/code_mode/execute_handler.rs | 7 +- codex-rs/core/src/tools/code_mode/mod.rs | 4 +- .../core/src/tools/code_mode/wait_handler.rs | 9 +- .../core/src/tools/executed_tool_calls.rs | 302 +++++------------- .../executed_tool_calls/request_metadata.rs | 220 +++++++++++++ .../request_metadata_tests.rs} | 127 +++++++- codex-rs/core/src/tools/handlers/mcp.rs | 11 +- codex-rs/core/src/tools/mod.rs | 2 +- codex-rs/core/src/tools/parallel.rs | 19 +- 16 files changed, 450 insertions(+), 299 deletions(-) create mode 100644 codex-rs/core/src/tools/executed_tool_calls/request_metadata.rs rename codex-rs/core/src/tools/{executed_tool_calls_tests.rs => executed_tool_calls/request_metadata_tests.rs} (82%) diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index ae95d17a37..261e180167 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1470,10 +1470,7 @@ impl Session { }).await; } - let executed_tool_calls = config - .features - .enabled(Feature::ExecutedToolCallMetadata) - .then(|| Arc::new(crate::state::ExecutedToolCallRecorder::default())); + let executed_tool_calls = crate::state::ExecutedToolCalls::new(&config.features); let services = SessionServices { // Start with an empty connection set. The initialized set is // published after SessionConfigured so MCP events follow it. diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index e262983e0e..b5969c64ce 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -6544,10 +6544,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { )); let network_approval = Arc::new(NetworkApprovalService::default()); let mcp_runtime = Arc::new(codex_mcp::McpRuntime::empty(config.prefix_mcp_tool_names())); - let executed_tool_calls = config - .features - .enabled(Feature::ExecutedToolCallMetadata) - .then(|| Arc::new(crate::state::ExecutedToolCallRecorder::default())); + let executed_tool_calls = crate::state::ExecutedToolCalls::new(&config.features); let (hooks, async_hook_results) = Hooks::new( HooksConfig { legacy_notify_argv: config.notify.clone(), @@ -8855,10 +8852,7 @@ where )); let network_approval = Arc::new(NetworkApprovalService::default()); let mcp_runtime = Arc::new(codex_mcp::McpRuntime::empty(config.prefix_mcp_tool_names())); - let executed_tool_calls = config - .features - .enabled(Feature::ExecutedToolCallMetadata) - .then(|| Arc::new(crate::state::ExecutedToolCallRecorder::default())); + let executed_tool_calls = crate::state::ExecutedToolCalls::new(&config.features); let (hooks, async_hook_results) = Hooks::new( HooksConfig { legacy_notify_argv: config.notify.clone(), diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index f9bac7c153..eab4fbab2e 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -1460,12 +1460,9 @@ async fn run_sampling_request( .for_prompt(&step_context.settings.model_info.input_modalities) }; let mut prompt_input = prompt_input; - if let Some(executed_tool_calls) = sess.services.executed_tool_calls.as_ref() - && executed_tool_calls - .attach_pending_to_prompt(&mut prompt_input, &mut executed_tool_calls_by_output) - { - codex_protocol::models::bound_executed_tool_calls_for_prompt(&mut prompt_input); - } + sess.services + .executed_tool_calls + .attach_to_prompt(&mut prompt_input, &mut executed_tool_calls_by_output); let prompt = build_prompt( prompt_input, step_context.as_ref(), diff --git a/codex-rs/core/src/state/mod.rs b/codex-rs/core/src/state/mod.rs index 71d9616e52..b5100bbcef 100644 --- a/codex-rs/core/src/state/mod.rs +++ b/codex-rs/core/src/state/mod.rs @@ -4,7 +4,7 @@ mod service; mod session; mod turn; -pub(crate) use crate::tools::ExecutedToolCallRecorder; +pub(crate) use crate::tools::ExecutedToolCalls; pub(crate) use additional_context::AdditionalContextStore; pub(crate) use auto_compact_window::AutoCompactWindowIds; pub(crate) use auto_compact_window::AutoCompactWindowSnapshot; diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index e5f96eaf18..6524bbfbb3 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -14,7 +14,7 @@ use crate::exec_policy::ExecPolicyManager; use crate::guardian::GuardianRejectionCircuitBreaker; use crate::mcp::McpManager; use crate::mcp_tool_exposure::McpHandlerCache; -use crate::tools::ExecutedToolCallRecorder; +use crate::tools::ExecutedToolCalls; use crate::tools::code_mode::CodeModeService; use crate::tools::handlers::ToolSearchHandlerCache; use crate::tools::network_approval::NetworkApprovalService; @@ -95,7 +95,7 @@ pub(crate) struct SessionServices { pub(crate) time_provider: Arc, /// Session-scoped model client shared across turns. pub(crate) model_client: ModelClient, - pub(crate) executed_tool_calls: Option>, + pub(crate) executed_tool_calls: ExecutedToolCalls, pub(crate) code_mode_service: CodeModeService, pub(crate) tool_search_handler_cache: ToolSearchHandlerCache, pub(crate) turn_environments: Arc, diff --git a/codex-rs/core/src/tools/code_mode/delegate.rs b/codex-rs/core/src/tools/code_mode/delegate.rs index bc9cb2f913..331622fd1c 100644 --- a/codex-rs/core/src/tools/code_mode/delegate.rs +++ b/codex-rs/core/src/tools/code_mode/delegate.rs @@ -20,7 +20,7 @@ use super::ExecContext; use super::PUBLIC_TOOL_NAME; use super::submit_nested_tool; use crate::session::step_context::StepContext; -use crate::tools::ExecutedToolCallRecorder; +use crate::tools::ExecutedToolCalls; use crate::tools::context::SharedTurnDiffTracker; use crate::tools::parallel::ToolCallRuntime; @@ -28,7 +28,7 @@ pub(super) struct CodeModeDispatchBroker { dispatch_tx: async_channel::Sender, dispatch_rx: async_channel::Receiver, dispatch_gates: Arc>>, - executed_tool_calls: Option>, + executed_tool_calls: ExecutedToolCalls, } struct CellDispatchGate { @@ -38,7 +38,7 @@ struct CellDispatchGate { } impl CodeModeDispatchBroker { - pub(super) fn new(executed_tool_calls: Option>) -> Self { + pub(super) fn new(executed_tool_calls: ExecutedToolCalls) -> Self { let (dispatch_tx, dispatch_rx) = async_channel::unbounded(); Self { dispatch_tx, @@ -84,9 +84,7 @@ impl CodeModeDispatchBroker { .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); dispatch_gates.remove(cell_id); - if let Some(recorder) = &self.executed_tool_calls { - recorder.finish_cell_recording(cell_id); - } + self.executed_tool_calls.finish_cell_recording(cell_id); } pub(super) fn active_cell_ids(&self) -> Vec { @@ -104,11 +102,7 @@ impl CodeModeDispatchBroker { step_context: Arc, tracker: SharedTurnDiffTracker, ) -> CodeModeDispatchWorker { - let track_completeness = exec - .turn - .config - .features - .enabled(codex_features::Feature::ExecutedToolCallMetadata); + let track_completeness = ExecutedToolCalls::is_enabled(&exec.turn.config.features); let tool_runtime = ToolCallRuntime::new(Arc::clone(&exec.session), step_context, tracker); let host = Arc::new(CoreTurnHost { exec, tool_runtime }); let dispatch_rx = self.dispatch_rx.clone(); diff --git a/codex-rs/core/src/tools/code_mode/delegate_tests.rs b/codex-rs/core/src/tools/code_mode/delegate_tests.rs index 07c9a0bb8b..de84f82761 100644 --- a/codex-rs/core/src/tools/code_mode/delegate_tests.rs +++ b/codex-rs/core/src/tools/code_mode/delegate_tests.rs @@ -116,7 +116,7 @@ async fn detached_code_mode_callback_keeps_thread_id_on_dispatch_span() -> anyho )); let step = StepContext::for_test(Arc::clone(&turn)).with_tool_router_for_test(router); let broker = Arc::new(CodeModeDispatchBroker::new( - /*executed_tool_calls*/ None, + /*executed_tool_calls*/ Default::default(), )); let cell_id = CellId::new("audit-cell".to_string()); broker.mark_cell_ready_for_dispatch(&cell_id, /*originating_item_id*/ None); diff --git a/codex-rs/core/src/tools/code_mode/execute_handler.rs b/codex-rs/core/src/tools/code_mode/execute_handler.rs index 4012d972f6..e0cc305048 100644 --- a/codex-rs/core/src/tools/code_mode/execute_handler.rs +++ b/codex-rs/core/src/tools/code_mode/execute_handler.rs @@ -87,9 +87,10 @@ impl CodeModeExecuteHandler { call_id: call_id.clone(), cell_id: cell_id.to_string(), }); - if let Some(executed_tool_calls) = exec.session.services.executed_tool_calls.as_ref() { - executed_tool_calls.start_cell(&cell_id, &call_id); - } + exec.session + .services + .executed_tool_calls + .start_cell(&cell_id, &call_id); let runtime_cell_id = cell_id.to_string(); let code_cell_trace = exec .session diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index 142312c99e..7bb9459097 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/codex-rs/core/src/tools/code_mode/mod.rs @@ -30,7 +30,7 @@ use crate::original_image_detail::sanitize_original_image_detail as sanitize_ima use crate::session::session::Session; use crate::session::step_context::StepContext; use crate::session::turn_context::TurnContext; -use crate::tools::ExecutedToolCallRecorder; +use crate::tools::ExecutedToolCalls; use crate::tools::context::FunctionToolOutput; use crate::tools::context::SharedTurnDiffTracker; use crate::tools::context::ToolPayload; @@ -80,7 +80,7 @@ impl CodeModeService { pub(crate) fn new( session_provider: Arc, config: &CodeModeConfig, - executed_tool_calls: Option>, + executed_tool_calls: ExecutedToolCalls, ) -> Self { let dispatch_broker = Arc::new(CodeModeDispatchBroker::new(executed_tool_calls)); let availability = session_provider.availability(); diff --git a/codex-rs/core/src/tools/code_mode/wait_handler.rs b/codex-rs/core/src/tools/code_mode/wait_handler.rs index e2c31e4bbe..c6880799c9 100644 --- a/codex-rs/core/src/tools/code_mode/wait_handler.rs +++ b/codex-rs/core/src/tools/code_mode/wait_handler.rs @@ -123,11 +123,10 @@ impl CodeModeWaitHandler { | codex_code_mode::RuntimeResponse::Result { cell_id, .. } => cell_id, }; telemetry.cell_id = Some(runtime_cell_id.to_string()); - if let Some(executed_tool_calls) = - exec.session.services.executed_tool_calls.as_ref() - { - executed_tool_calls.register_cell(runtime_cell_id, &call_id); - } + exec.session + .services + .executed_tool_calls + .register_cell(runtime_cell_id, &call_id); if !matches!(response, codex_code_mode::RuntimeResponse::Yielded { .. }) { exec.session .services diff --git a/codex-rs/core/src/tools/executed_tool_calls.rs b/codex-rs/core/src/tools/executed_tool_calls.rs index ec4843b5aa..c30184d3ff 100644 --- a/codex-rs/core/src/tools/executed_tool_calls.rs +++ b/codex-rs/core/src/tools/executed_tool_calls.rs @@ -1,17 +1,29 @@ +//! Records attempted calls at existing execution boundaries. Request metadata policy +//! lives in the private request metadata module; this recorder must never dispatch or await tools. + +mod request_metadata; + use std::collections::HashMap; use std::collections::HashSet; +use std::sync::Arc; +use std::sync::Mutex; use codex_code_mode::CellId; +use codex_features::Feature; +use codex_features::Features; use codex_protocol::models::ExecutedToolCall; use codex_protocol::models::ExecutedToolCallArguments; use codex_protocol::models::ResponseItem; +use codex_protocol::models::bound_executed_tool_calls_for_prompt; use codex_protocol::models::bound_executed_tool_calls_for_prompt_prioritizing_recent; use codex_protocol::models::executed_tool_call_metadata_bytes; use codex_protocol::openai_models::ToolMode; use indexmap::IndexMap; use serde_json::Value as JsonValue; +use crate::session::step_context::StepContext; use crate::tools::context::ToolCallSource; +use crate::tools::context::ToolOutput; use crate::tools::context::ToolPayload; use crate::tools::router::ToolCall; use crate::utils::json::serialized_json_bytes; @@ -20,11 +32,14 @@ const MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES: usize = 8 * 1024; const MAX_EXECUTED_TOOL_CALL_FULL_ARGUMENT_BYTES_PER_OUTPUT: usize = 32 * 1024; const MAX_PENDING_EXECUTED_TOOL_CALLS: usize = 256; -/// Best-effort, session-scoped attempted-tool metadata; cancellation, compaction, -/// and yielded cells without another wait can leave pending calls unreported. -#[derive(Default)] -pub(crate) struct ExecutedToolCallRecorder { - state: std::sync::Mutex, +type ExecutedToolCallCache = + HashMap<(std::mem::Discriminant, String), Vec>; + +/// Best-effort recording shared by a session and its Code Mode broker. Disabled +/// sessions allocate no state; missing calls and truncated evidence remain incomplete. +#[derive(Clone, Default)] +pub(crate) struct ExecutedToolCalls { + state: Option>>, } #[derive(Default)] @@ -101,13 +116,47 @@ impl ExecutedToolCallRecorderState { } } -impl ExecutedToolCallRecorder { +impl ExecutedToolCalls { + pub(crate) fn new(features: &Features) -> Self { + Self { + state: Self::is_enabled(features) + .then(|| Arc::new(Mutex::new(ExecutedToolCallRecorderState::default()))), + } + } + + /// The turn's feature policy is independent of whether this session has a recorder. + pub(crate) fn is_enabled(features: &Features) -> bool { + features.enabled(Feature::ExecutedToolCallMetadata) + } + pub(crate) fn record_tool_call( &self, call: &ToolCall, source: &ToolCallSource, - tool_mode: ToolMode, + step_context: &StepContext, ) { + if Self::is_enabled(&step_context.turn.config.features) && self.state.is_some() { + self.record_call(call, source, step_context.tool_router.tool_mode()); + } + } + + pub(crate) fn record_accepted_result( + &self, + source: &ToolCallSource, + call_id: &str, + result: &dyn ToolOutput, + ) { + if self.state.is_some() + && let Some(sources) = result.tool_result_sources() + { + self.record_tool_result_sources(source, call_id, sources); + } + } + + fn record_call(&self, call: &ToolCall, source: &ToolCallSource, tool_mode: ToolMode) { + let Some(state) = &self.state else { + return; + }; if matches!(source, ToolCallSource::Direct) && matches!(tool_mode, ToolMode::CodeMode | ToolMode::CodeModeOnly) && call.tool_name.is_default_namespace() @@ -148,8 +197,7 @@ impl ExecutedToolCallRecorder { }; match source { ToolCallSource::Direct | ToolCallSource::DirectPlaintextMessage => { - let mut state = self - .state + let mut state = state .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); if state.direct_calls.len() < MAX_PENDING_EXECUTED_TOOL_CALLS { @@ -188,8 +236,10 @@ impl ExecutedToolCallRecorder { call: ExecutedToolCall, original_bytes: usize, ) { - let mut state = self - .state + let Some(state) = &self.state else { + return; + }; + let mut state = state .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); if state.pending_nested_calls > MAX_PENDING_EXECUTED_TOOL_CALLS @@ -232,14 +282,16 @@ impl ExecutedToolCallRecorder { state.pending_nested_calls += 1; } - pub(crate) fn record_tool_result_sources( + fn record_tool_result_sources( &self, source: &ToolCallSource, call_id: &str, result_sources: codex_protocol::models::ToolResultSources, ) -> bool { - let mut state = self - .state + let Some(state) = &self.state else { + return false; + }; + let mut state = state .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); let call = match source { @@ -271,217 +323,21 @@ impl ExecutedToolCallRecorder { retained.calls[index].set_tool_result_sources(result_sources) } - pub(crate) fn attach_pending_to_prompt( - &self, - items: &mut [ResponseItem], - retry_cache: &mut HashMap< - (std::mem::Discriminant, String), - Vec, - >, - ) -> bool { - let mut state = self - .state - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - if state.direct_calls.is_empty() - && state.output_cells.is_empty() - && state.retained_calls.is_empty() - && retry_cache.is_empty() - { - return false; - } - - // Updated records supersede older retry snapshots. - retry_cache.retain(|key, _| { - !state - .retained_calls - .get(key) - .is_some_and(|retained| retained.sources_updated) - }); - let mut pending_retry_outputs = retry_cache.keys().cloned().collect::>(); - let mut pending_retained_outputs = - state.retained_calls.keys().cloned().collect::>(); - let mut attached = false; - let mut retained_bytes = 0_usize; - for item in items.iter_mut().rev() { - if state.direct_calls.is_empty() - && state.output_cells.is_empty() - && pending_retry_outputs.is_empty() - && pending_retained_outputs.is_empty() - { - break; - } - let call_id = match &*item { - ResponseItem::FunctionCallOutput { - call_id: Some(call_id), - .. - } - | ResponseItem::CustomToolCallOutput { call_id, .. } - | ResponseItem::ToolSearchOutput { - call_id: Some(call_id), - .. - } => call_id, - _ => continue, - }; - let key = (std::mem::discriminant(&*item), call_id.clone()); - let retained = state.retained_calls.get(&key); - let mut complete = retained.is_some_and(|retained| retained.complete); - let mut cell_id = retained.and_then(|retained| retained.cell_id.clone()); - let mut runtime_cell_id = - retained.and_then(|retained| retained.runtime_cell_id.clone()); - let calls = if let Some(cached) = retry_cache.get(&key) { - if !pending_retry_outputs.remove(&key) { - continue; - } - pending_retained_outputs.remove(&key); - cached.clone() - } else if let Some(retained) = retained { - if !pending_retained_outputs.remove(&key) { - continue; - } - retained.calls.clone() - } else { - let mut calls = state - .direct_calls - .remove(call_id) - .into_iter() - .collect::>(); - let mut call_index_by_id = HashMap::new(); - if let Some(output_cell_id) = state.output_cells.remove(call_id) - && let Some(cell) = state.cells.get_mut(&output_cell_id) - { - cell_id = cell.originating_call_id.clone(); - runtime_cell_id = Some(output_cell_id.clone()); - let pending_calls = cell.pending_calls.len(); - for (call_id, call) in cell.pending_calls.drain(..) { - if matches!(call.arguments(), ExecutedToolCallArguments::Raw(_)) { - call_index_by_id.insert(call_id, calls.len()); - } - calls.push(call); - } - cell.pending_full_argument_bytes = 0; - complete = cell.completion == CellCompletion::Complete; - if matches!( - cell.completion, - CellCompletion::Complete | CellCompletion::Incomplete - ) { - state.cells.remove(&output_cell_id); - } - state.pending_nested_calls = - state.pending_nested_calls.saturating_sub(pending_calls); - state - .output_cells - .retain(|_, registered_cell_id| registered_cell_id != &output_cell_id); - } - if calls.is_empty() && !complete { - continue; - } - retry_cache.insert(key.clone(), calls.clone()); - state.retained_calls.insert( - key, - RetainedToolCalls { - calls: calls.clone(), - complete, - cell_id: cell_id.clone(), - runtime_cell_id, - call_index_by_id, - sources_updated: false, - }, - ); - calls - }; - item.append_executed_tool_calls(calls); - if let Some(cell_id) = cell_id { - item.set_tool_call_cell_id(&cell_id); - } - if complete { - item.mark_tool_calls_complete(); - } - retained_bytes = retained_bytes.saturating_add(executed_tool_call_metadata_bytes(item)); - attached = true; - } - if !pending_retained_outputs.is_empty() { - state - .retained_calls - .retain(|key, _| !pending_retained_outputs.contains(key)); - } - if retained_bytes > MAX_EXECUTED_TOOL_CALL_FULL_ARGUMENT_BYTES_PER_OUTPUT { - bound_executed_tool_calls_for_prompt_prioritizing_recent(items); - let retained_before_bounding = std::mem::take(&mut state.retained_calls); - let mut bounded_outputs = HashSet::new(); - for item in items { - let call_id = match &*item { - ResponseItem::FunctionCallOutput { - call_id: Some(call_id), - .. - } - | ResponseItem::CustomToolCallOutput { call_id, .. } - | ResponseItem::ToolSearchOutput { - call_id: Some(call_id), - .. - } => call_id, - _ => continue, - }; - let key = (std::mem::discriminant(&*item), call_id.clone()); - let metadata = item.executed_tool_call_metadata(); - let unique_output = bounded_outputs.insert(key.clone()); - if !unique_output && let Some(retained) = state.retained_calls.get_mut(&key) { - retained.call_index_by_id.clear(); - } - let previous = retained_before_bounding.get(&key); - if let Some(retained) = previous - && let Some(runtime_cell_id) = &retained.runtime_cell_id - && metadata - .is_none_or(|metadata| !metadata.has_same_tool_calls(&retained.calls)) - && let Some(cell) = state.cells.get_mut(runtime_cell_id) - { - cell.completion = CellCompletion::Incomplete; - } - if let Some(metadata) = metadata - && (metadata - .executed_tool_calls - .as_ref() - .is_some_and(|calls| !calls.is_empty()) - || metadata.tool_calls_complete.is_some()) - { - let runtime_cell_id = - previous.and_then(|retained| retained.runtime_cell_id.clone()); - // Budget rewriting or duplicate outputs invalidate the original index slots. - let call_index_by_id = previous - .filter(|retained| { - unique_output - && metadata.executed_tool_calls.as_ref() == Some(&retained.calls) - }) - .map(|retained| retained.call_index_by_id.clone()) - .unwrap_or_default(); - let retained = state.retained_calls.entry(key).or_default(); - retained.runtime_cell_id = runtime_cell_id; - retained.calls = metadata.executed_tool_calls.clone().unwrap_or_default(); - if let Some(cell_id) = metadata.cell_id.as_ref() { - retained.cell_id = Some(cell_id.clone()); - } - retained.complete |= metadata.tool_calls_complete == Some(true); - retained.call_index_by_id = call_index_by_id; - retained.sources_updated = - previous.is_some_and(|retained| retained.sources_updated); - } - } - } - - attached - } - pub(crate) fn register_cell(&self, cell_id: &CellId, output_call_id: &str) { - let mut state = self - .state + let Some(state) = &self.state else { + return; + }; + let mut state = state .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); state.register_cell(cell_id, output_call_id); } pub(crate) fn start_cell(&self, cell_id: &CellId, output_call_id: &str) { - let mut state = self - .state + let Some(state) = &self.state else { + return; + }; + let mut state = state .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); state.register_cell(cell_id, output_call_id); @@ -491,8 +347,10 @@ impl ExecutedToolCallRecorder { } pub(crate) fn finish_cell_recording(&self, cell_id: &CellId) { - let mut state = self - .state + let Some(state) = &self.state else { + return; + }; + let mut state = state .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); if let Some(cell) = state.cells.get_mut(cell_id) { @@ -504,7 +362,3 @@ impl ExecutedToolCallRecorder { } } } - -#[cfg(test)] -#[path = "executed_tool_calls_tests.rs"] -mod tests; diff --git a/codex-rs/core/src/tools/executed_tool_calls/request_metadata.rs b/codex-rs/core/src/tools/executed_tool_calls/request_metadata.rs new file mode 100644 index 0000000000..0e75d15e0c --- /dev/null +++ b/codex-rs/core/src/tools/executed_tool_calls/request_metadata.rs @@ -0,0 +1,220 @@ +//! Attaches host observations to request metadata and applies its byte budget. +//! Final request budgeting runs after the recorder lock is released; it never changes tool execution. + +use super::*; + +impl ExecutedToolCalls { + /// Attaches host-recorded calls within the existing request metadata budget. + pub(crate) fn attach_to_prompt( + &self, + items: &mut [ResponseItem], + retry_cache: &mut ExecutedToolCallCache, + ) { + if self.attach_pending_to_prompt(items, retry_cache) { + bound_executed_tool_calls_for_prompt(items); + } + } + + fn attach_pending_to_prompt( + &self, + items: &mut [ResponseItem], + retry_cache: &mut ExecutedToolCallCache, + ) -> bool { + let Some(state) = &self.state else { + return false; + }; + let mut state = state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state.direct_calls.is_empty() + && state.output_cells.is_empty() + && state.retained_calls.is_empty() + && retry_cache.is_empty() + { + return false; + } + + // Updated records supersede older retry snapshots. + retry_cache.retain(|key, _| { + !state + .retained_calls + .get(key) + .is_some_and(|retained| retained.sources_updated) + }); + let mut pending_retry_outputs = retry_cache.keys().cloned().collect::>(); + let mut pending_retained_outputs = + state.retained_calls.keys().cloned().collect::>(); + let mut attached = false; + let mut retained_bytes = 0_usize; + for item in items.iter_mut().rev() { + if state.direct_calls.is_empty() + && state.output_cells.is_empty() + && pending_retry_outputs.is_empty() + && pending_retained_outputs.is_empty() + { + break; + } + let call_id = match &*item { + ResponseItem::FunctionCallOutput { + call_id: Some(call_id), + .. + } + | ResponseItem::CustomToolCallOutput { call_id, .. } + | ResponseItem::ToolSearchOutput { + call_id: Some(call_id), + .. + } => call_id, + _ => continue, + }; + let key = (std::mem::discriminant(&*item), call_id.clone()); + let retained = state.retained_calls.get(&key); + let mut complete = retained.is_some_and(|retained| retained.complete); + let mut cell_id = retained.and_then(|retained| retained.cell_id.clone()); + let mut runtime_cell_id = + retained.and_then(|retained| retained.runtime_cell_id.clone()); + let calls = if let Some(cached) = retry_cache.get(&key) { + if !pending_retry_outputs.remove(&key) { + continue; + } + pending_retained_outputs.remove(&key); + cached.clone() + } else if let Some(retained) = retained { + if !pending_retained_outputs.remove(&key) { + continue; + } + retained.calls.clone() + } else { + let mut calls = state + .direct_calls + .remove(call_id) + .into_iter() + .collect::>(); + let mut call_index_by_id = HashMap::new(); + if let Some(output_cell_id) = state.output_cells.remove(call_id) + && let Some(cell) = state.cells.get_mut(&output_cell_id) + { + cell_id = cell.originating_call_id.clone(); + runtime_cell_id = Some(output_cell_id.clone()); + let pending_calls = cell.pending_calls.len(); + for (call_id, call) in cell.pending_calls.drain(..) { + if matches!(call.arguments(), ExecutedToolCallArguments::Raw(_)) { + call_index_by_id.insert(call_id, calls.len()); + } + calls.push(call); + } + cell.pending_full_argument_bytes = 0; + complete = cell.completion == CellCompletion::Complete; + if matches!( + cell.completion, + CellCompletion::Complete | CellCompletion::Incomplete + ) { + state.cells.remove(&output_cell_id); + } + state.pending_nested_calls = + state.pending_nested_calls.saturating_sub(pending_calls); + state + .output_cells + .retain(|_, registered_cell_id| registered_cell_id != &output_cell_id); + } + if calls.is_empty() && !complete { + continue; + } + retry_cache.insert(key.clone(), calls.clone()); + state.retained_calls.insert( + key, + RetainedToolCalls { + calls: calls.clone(), + complete, + cell_id: cell_id.clone(), + runtime_cell_id, + call_index_by_id, + sources_updated: false, + }, + ); + calls + }; + item.append_executed_tool_calls(calls); + if let Some(cell_id) = cell_id { + item.set_tool_call_cell_id(&cell_id); + } + if complete { + item.mark_tool_calls_complete(); + } + retained_bytes = retained_bytes.saturating_add(executed_tool_call_metadata_bytes(item)); + attached = true; + } + if !pending_retained_outputs.is_empty() { + state + .retained_calls + .retain(|key, _| !pending_retained_outputs.contains(key)); + } + if retained_bytes > MAX_EXECUTED_TOOL_CALL_FULL_ARGUMENT_BYTES_PER_OUTPUT { + bound_executed_tool_calls_for_prompt_prioritizing_recent(items); + let retained_before_bounding = std::mem::take(&mut state.retained_calls); + let mut bounded_outputs = HashSet::new(); + for item in items { + let call_id = match &*item { + ResponseItem::FunctionCallOutput { + call_id: Some(call_id), + .. + } + | ResponseItem::CustomToolCallOutput { call_id, .. } + | ResponseItem::ToolSearchOutput { + call_id: Some(call_id), + .. + } => call_id, + _ => continue, + }; + let key = (std::mem::discriminant(&*item), call_id.clone()); + let metadata = item.executed_tool_call_metadata(); + let unique_output = bounded_outputs.insert(key.clone()); + if !unique_output && let Some(retained) = state.retained_calls.get_mut(&key) { + retained.call_index_by_id.clear(); + } + let previous = retained_before_bounding.get(&key); + if let Some(retained) = previous + && let Some(runtime_cell_id) = &retained.runtime_cell_id + && metadata + .is_none_or(|metadata| !metadata.has_same_tool_calls(&retained.calls)) + && let Some(cell) = state.cells.get_mut(runtime_cell_id) + { + cell.completion = CellCompletion::Incomplete; + } + if let Some(metadata) = metadata + && (metadata + .executed_tool_calls + .as_ref() + .is_some_and(|calls| !calls.is_empty()) + || metadata.tool_calls_complete.is_some()) + { + let runtime_cell_id = + previous.and_then(|retained| retained.runtime_cell_id.clone()); + // Budget rewriting or duplicate outputs invalidate the original index slots. + let call_index_by_id = previous + .filter(|retained| { + unique_output + && metadata.executed_tool_calls.as_ref() == Some(&retained.calls) + }) + .map(|retained| retained.call_index_by_id.clone()) + .unwrap_or_default(); + let retained = state.retained_calls.entry(key).or_default(); + retained.runtime_cell_id = runtime_cell_id; + retained.calls = metadata.executed_tool_calls.clone().unwrap_or_default(); + if let Some(cell_id) = metadata.cell_id.as_ref() { + retained.cell_id = Some(cell_id.clone()); + } + retained.complete |= metadata.tool_calls_complete == Some(true); + retained.call_index_by_id = call_index_by_id; + retained.sources_updated = + previous.is_some_and(|retained| retained.sources_updated); + } + } + } + + attached + } +} + +#[cfg(test)] +#[path = "request_metadata_tests.rs"] +mod tests; diff --git a/codex-rs/core/src/tools/executed_tool_calls_tests.rs b/codex-rs/core/src/tools/executed_tool_calls/request_metadata_tests.rs similarity index 82% rename from codex-rs/core/src/tools/executed_tool_calls_tests.rs rename to codex-rs/core/src/tools/executed_tool_calls/request_metadata_tests.rs index 441ac44dfb..e06e675fa2 100644 --- a/codex-rs/core/src/tools/executed_tool_calls_tests.rs +++ b/codex-rs/core/src/tools/executed_tool_calls/request_metadata_tests.rs @@ -6,6 +6,12 @@ use serde_json::json; use super::*; +fn new_recorder() -> ExecutedToolCalls { + let mut features = Features::default(); + features.enable(Feature::ExecutedToolCallMetadata); + ExecutedToolCalls::new(&features) +} + fn output(call_id: &str) -> ResponseItem { ResponseItem::from(ResponseInputItem::FunctionCallOutput { call_id: call_id.to_string(), @@ -13,12 +19,103 @@ fn output(call_id: &str) -> ResponseItem { }) } +#[tokio::test] +async fn recorder_preserves_session_and_turn_gates_and_lazy_result_sources() { + struct SourceLookup(std::cell::Cell); + + impl ToolOutput for SourceLookup { + fn log_output(&self) -> String { + panic!("recording must not read the diagnostic output") + } + + fn success_for_logging(&self) -> bool { + panic!("recording must not inspect execution success") + } + + fn to_response_item(&self, _call_id: &str, _payload: &ToolPayload) -> ResponseInputItem { + panic!("recording must not rebuild the tool result") + } + + fn tool_result_sources(&self) -> Option { + self.0.set(self.0.get() + 1); + Some(ToolResultSources::new(Vec::new())) + } + } + + let (_session, turn) = crate::session::tests::make_session_and_context().await; + let mut turn = Arc::new(turn); + assert!(ExecutedToolCalls::default().state.is_none()); + for session_enabled in [false, true] { + let mut features = Features::default(); + features.set_enabled(Feature::ExecutedToolCallMetadata, session_enabled); + for turn_enabled in [false, true] { + let calls = ExecutedToolCalls::new(&features); + assert_eq!(calls.state.is_some(), session_enabled); + Arc::make_mut( + &mut Arc::get_mut(&mut turn) + .expect("the previous step must have released its turn") + .config, + ) + .features + .set_enabled(Feature::ExecutedToolCallMetadata, turn_enabled) + .expect("test feature must be configurable"); + let step = StepContext::for_test(Arc::clone(&turn)); + assert_eq!( + ExecutedToolCalls::is_enabled(&step.turn.config.features), + turn_enabled + ); + + let call = ToolCall { + tool_name: codex_tools::ToolName::plain("test_tool"), + call_id: "call".to_string(), + payload: ToolPayload::Function { + arguments: r#"{"value":7}"#.to_string(), + }, + encrypted_function_args: None, + }; + calls.record_tool_call(&call, &ToolCallSource::Direct, &step); + let result = SourceLookup(std::cell::Cell::new(0)); + calls.record_accepted_result(&ToolCallSource::Direct, &call.call_id, &result); + // Accepted sources follow session availability even when this turn stops recording calls. + assert_eq!(result.0.get(), usize::from(session_enabled)); + + let input = serde_json::from_value(json!({ + "type": "function_call", + "call_id": "call", + "name": "test_tool", + "arguments": "{\"value\":7}", + })) + .expect("tool input must deserialize"); + let mut prompt = vec![input, output("call")]; + if !session_enabled { + prompt[1].mark_tool_calls_complete(); + } + let original = prompt.clone(); + calls.attach_to_prompt(&mut prompt, &mut HashMap::new()); + if !session_enabled { + assert_eq!(prompt, original); + } + let mut expected_call = + ExecutedToolCall::new("test_tool".to_string(), json!({"value": 7})); + assert!(expected_call.set_tool_result_sources(ToolResultSources::new(Vec::new()))); + let expected_calls = vec![expected_call]; + assert_eq!( + prompt[1] + .executed_tool_call_metadata() + .and_then(|metadata| metadata.executed_tool_calls.as_ref()), + (session_enabled && turn_enabled).then_some(&expected_calls), + "session_enabled={session_enabled}, turn_enabled={turn_enabled}", + ); + } + } +} + #[test] fn executed_tool_call_recorder_bounds_pending_calls_and_preserves_overflow() { - let recorder = ExecutedToolCallRecorder::default(); + let recorder = new_recorder(); for index in 0..MAX_PENDING_EXECUTED_TOOL_CALLS + 2 { - recorder.record_tool_call( + recorder.record_call( &ToolCall { tool_name: codex_tools::ToolName::plain("direct_tool"), call_id: format!("direct-{index}"), @@ -58,6 +155,8 @@ fn executed_tool_call_recorder_bounds_pending_calls_and_preserves_overflow() { { let state = recorder .state + .as_ref() + .unwrap() .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); assert_eq!( @@ -131,6 +230,8 @@ fn executed_tool_call_recorder_bounds_pending_calls_and_preserves_overflow() { { let state = recorder .state + .as_ref() + .unwrap() .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); assert_eq!(state.pending_nested_calls, 0); @@ -159,6 +260,8 @@ fn executed_tool_call_recorder_bounds_pending_calls_and_preserves_overflow() { assert!(!recorder.attach_pending_to_prompt(&mut [], &mut compacted_retry_cache)); let state = recorder .state + .as_ref() + .unwrap() .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); assert!(state.retained_calls.is_empty()); @@ -166,7 +269,7 @@ fn executed_tool_call_recorder_bounds_pending_calls_and_preserves_overflow() { #[test] fn executed_tool_call_recorder_bounds_retained_history_and_reports_omissions() { - let recorder = ExecutedToolCallRecorder::default(); + let recorder = new_recorder(); let mut history = Vec::new(); let arguments = serde_json::to_string(&json!({ "payload": "x".repeat(1024) })) .expect("tool arguments must serialize"); @@ -174,7 +277,7 @@ fn executed_tool_call_recorder_bounds_retained_history_and_reports_omissions() { for index in 0..512 { let call_id = format!("retained-{index}"); - recorder.record_tool_call( + recorder.record_call( &ToolCall { tool_name: codex_tools::ToolName::plain(format!("retained_tool_{index}")), call_id: call_id.clone(), @@ -215,6 +318,8 @@ fn executed_tool_call_recorder_bounds_retained_history_and_reports_omissions() { let state = recorder .state + .as_ref() + .unwrap() .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); let retained_bytes = state @@ -245,7 +350,7 @@ fn executed_tool_call_recorder_bounds_retained_history_and_reports_omissions() { #[test] fn tool_call_completeness_requires_finished_lossless_recording() { for scenario in ["empty", "unobserved", "late_call"] { - let recorder = ExecutedToolCallRecorder::default(); + let recorder = new_recorder(); let cell_id = CellId::new(scenario.to_string()); if scenario == "unobserved" { recorder.register_cell(&cell_id, "output"); @@ -289,7 +394,7 @@ fn tool_call_completeness_survives_waits_without_changing_deltas() { (false, ToolResultSources::parse_failed()), (true, ToolResultSources::parse_failed()), ] { - let recorder = ExecutedToolCallRecorder::default(); + let recorder = new_recorder(); let cell_id = CellId::new("multi-wait".to_string()); recorder.start_cell(&cell_id, "exec"); let mut history = Vec::new(); @@ -368,7 +473,7 @@ fn tool_call_completeness_survives_waits_without_changing_deltas() { assert_eq!(prompt, expected); } } - let state = recorder.state.lock().unwrap(); + let state = recorder.state.as_ref().unwrap().lock().unwrap(); assert!(state.cells.is_empty()); assert_eq!(state.pending_nested_calls, 0); } @@ -379,7 +484,7 @@ fn cell_correlation_uses_originating_exec_across_runtime_restarts() { let runtime_cell_id = CellId::new("1".to_string()); for originating_call_id in ["first-exec", "resumed-exec"] { - let recorder = ExecutedToolCallRecorder::default(); + let recorder = new_recorder(); recorder.start_cell(&runtime_cell_id, originating_call_id); recorder.record_nested_tool_call( runtime_cell_id.clone(), @@ -411,7 +516,7 @@ fn cell_correlation_uses_originating_exec_across_runtime_restarts() { #[test] fn request_truncation_prevents_completion_after_compaction() { - let recorder = ExecutedToolCallRecorder::default(); + let recorder = new_recorder(); let cell_id = CellId::new("compacted-cell".to_string()); recorder.start_cell(&cell_id, "exec"); @@ -465,7 +570,7 @@ fn request_truncation_prevents_completion_after_compaction() { #[test] fn source_shedding_preserves_completion_after_compaction() { - let recorder = ExecutedToolCallRecorder::default(); + let recorder = new_recorder(); let cell_id = CellId::new("compacted-cell".to_string()); recorder.start_cell(&cell_id, "exec"); @@ -547,7 +652,7 @@ fn source_shedding_preserves_completion_after_compaction() { #[test] fn finished_cells_without_more_waits_do_not_block_new_calls() { - let recorder = ExecutedToolCallRecorder::default(); + let recorder = new_recorder(); let call = ExecutedToolCall::new("nested_tool".to_string(), json!({})); for index in 0..MAX_PENDING_EXECUTED_TOOL_CALLS { let cell = CellId::new(format!("cell-{index}")); diff --git a/codex-rs/core/src/tools/handlers/mcp.rs b/codex-rs/core/src/tools/handlers/mcp.rs index 48e828fcfa..7ddbb2f921 100644 --- a/codex-rs/core/src/tools/handlers/mcp.rs +++ b/codex-rs/core/src/tools/handlers/mcp.rs @@ -284,12 +284,11 @@ impl CoreToolRuntime for McpHandler { } fn on_tool_result_accepted(&self, invocation: &ToolInvocation, result: &dyn ToolOutput) { - // Direct calls also record sources, before the Code Mode-only evidence path below. - if let Some(recorder) = invocation.session.services.executed_tool_calls.as_ref() - && let Some(sources) = result.tool_result_sources() - { - recorder.record_tool_result_sources(&invocation.source, &invocation.call_id, sources); - } + invocation + .session + .services + .executed_tool_calls + .record_accepted_result(&invocation.source, &invocation.call_id, result); let ToolCallSource::CodeMode { cell_id, .. } = &invocation.source else { return; }; diff --git a/codex-rs/core/src/tools/mod.rs b/codex-rs/core/src/tools/mod.rs index ff1b6d063d..7a65bc95a3 100644 --- a/codex-rs/core/src/tools/mod.rs +++ b/codex-rs/core/src/tools/mod.rs @@ -31,7 +31,7 @@ use codex_tools::ToolName; use codex_utils_output_truncation::TruncationPolicy; use codex_utils_output_truncation::formatted_truncate_text; use codex_utils_output_truncation::truncate_text; -pub(crate) use executed_tool_calls::ExecutedToolCallRecorder; +pub(crate) use executed_tool_calls::ExecutedToolCalls; pub use router::ToolRouter; /// Legacy boundaries such as hook payloads, telemetry tags, and Responses tool diff --git a/codex-rs/core/src/tools/parallel.rs b/codex-rs/core/src/tools/parallel.rs index 1f5576cb71..4b5f962bf6 100644 --- a/codex-rs/core/src/tools/parallel.rs +++ b/codex-rs/core/src/tools/parallel.rs @@ -98,20 +98,11 @@ impl ToolCallRuntime { source: ToolCallSource, cancellation_token: CancellationToken, ) -> impl std::future::Future> { - if self - .step_context - .turn - .config - .features - .enabled(codex_features::Feature::ExecutedToolCallMetadata) - && let Some(executed_tool_calls) = self.session.services.executed_tool_calls.as_ref() - { - executed_tool_calls.record_tool_call( - &call, - &source, - self.step_context.tool_router.tool_mode(), - ); - } + self.session.services.executed_tool_calls.record_tool_call( + &call, + &source, + &self.step_context, + ); let router = &self.step_context.tool_router; let supports_parallel = router.tool_supports_parallel(&call); let tool_runtime = router.tool_runtime(&call.tool_name);