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);