diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index e0250aee6d..c5082fda91 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1428,10 +1428,11 @@ impl Session { session_configuration.parent_thread_id, ), ), - executed_tool_calls, + executed_tool_calls: executed_tool_calls.clone(), code_mode_service: crate::tools::code_mode::CodeModeService::new( Arc::clone(&code_mode_session_provider), &config.code_mode, + executed_tool_calls, ), tool_search_handler_cache: Default::default(), turn_environments: Arc::clone(&turn_environments), diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 9aa49c1a19..e9b3227aee 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -6336,10 +6336,11 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { /*attestation_provider*/ None, config.http_client_factory(), ), - executed_tool_calls, + executed_tool_calls: executed_tool_calls.clone(), code_mode_service: crate::tools::code_mode::CodeModeService::new( Arc::new(codex_code_mode::DisabledCodeModeSessionProvider), &config.code_mode, + executed_tool_calls, ), tool_search_handler_cache: Default::default(), turn_environments: Arc::clone(&turn_environments), @@ -8614,10 +8615,11 @@ where /*attestation_provider*/ None, config.http_client_factory(), ), - executed_tool_calls, + executed_tool_calls: executed_tool_calls.clone(), code_mode_service: crate::tools::code_mode::CodeModeService::new( Arc::new(codex_code_mode::DisabledCodeModeSessionProvider), &config.code_mode, + executed_tool_calls, ), tool_search_handler_cache: Default::default(), turn_environments: Arc::clone(&turn_environments), diff --git a/codex-rs/core/src/tools/code_mode/delegate.rs b/codex-rs/core/src/tools/code_mode/delegate.rs index 56389b0f3b..0319d13d36 100644 --- a/codex-rs/core/src/tools/code_mode/delegate.rs +++ b/codex-rs/core/src/tools/code_mode/delegate.rs @@ -18,8 +18,9 @@ use tracing::Instrument; use super::ExecContext; use super::PUBLIC_TOOL_NAME; -use super::call_nested_tool; +use super::submit_nested_tool; use crate::session::step_context::StepContext; +use crate::tools::ExecutedToolCallRecorder; use crate::tools::context::SharedTurnDiffTracker; use crate::tools::parallel::ToolCallRuntime; @@ -27,6 +28,7 @@ pub(super) struct CodeModeDispatchBroker { dispatch_tx: async_channel::Sender, dispatch_rx: async_channel::Receiver, dispatch_gates: Arc>>, + executed_tool_calls: Option>, } struct CellDispatchGate { @@ -36,12 +38,13 @@ struct CellDispatchGate { } impl CodeModeDispatchBroker { - pub(super) fn new() -> Self { + pub(super) fn new(executed_tool_calls: Option>) -> Self { let (dispatch_tx, dispatch_rx) = async_channel::unbounded(); Self { dispatch_tx, dispatch_rx, dispatch_gates: Arc::new(Mutex::new(HashMap::new())), + executed_tool_calls, } } @@ -76,7 +79,14 @@ impl CodeModeDispatchBroker { } pub(super) fn close_cell(&self, cell_id: &CellId) { - remove_dispatch_gate(&self.dispatch_gates, cell_id); + let mut dispatch_gates = self + .dispatch_gates + .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); + } } pub(super) fn active_cell_ids(&self) -> Vec { @@ -94,6 +104,11 @@ impl CodeModeDispatchBroker { step_context: Arc, tracker: SharedTurnDiffTracker, ) -> CodeModeDispatchWorker { + let track_completeness = exec + .turn + .config + .features + .enabled(codex_features::Feature::ExecutedToolCallMetadata); 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(); @@ -148,10 +163,26 @@ impl CodeModeDispatchBroker { continue; } let host = Arc::clone(&host); + let dispatch_gates = Arc::clone(&dispatch_gates); tokio::spawn(async move { - let invocation = host - .invoke_tool(invocation, cancellation_token.clone()) - .instrument(span); + let invocation = { + let dispatch_gate = track_completeness.then(|| { + dispatch_gates + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + }); + if dispatch_gate.as_ref().is_some_and(|gates| { + cancellation_token.is_cancelled() + || !gates.contains_key(&cell_id) + }) { + return; + } + // Submission and cell closure share this gate. + span.in_scope(|| { + host.submit_tool(invocation, cancellation_token.clone()) + }) + .instrument(span) + }; tokio::pin!(invocation); let response = tokio::select! { biased; @@ -324,19 +355,19 @@ struct CoreTurnHost { } impl CoreTurnHost { - async fn invoke_tool( + fn submit_tool( &self, invocation: CodeModeNestedToolCall, cancellation_token: CancellationToken, - ) -> Result { - call_nested_tool( + ) -> impl std::future::Future> + Send + 'static { + let invocation = submit_nested_tool( self.exec.clone(), self.tool_runtime.clone(), invocation, cancellation_token, ) - .await - .map_err(|error| error.to_string()) + .map_err(|error| error.to_string()); + async move { invocation?.await.map_err(|error| error.to_string()) } } async fn notify(&self, call_id: String, cell_id: CellId, text: String) -> Result<(), String> { 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 c8a497fac1..313290da5a 100644 --- a/codex-rs/core/src/tools/code_mode/execute_handler.rs +++ b/codex-rs/core/src/tools/code_mode/execute_handler.rs @@ -88,7 +88,7 @@ impl CodeModeExecuteHandler { cell_id: cell_id.to_string(), }); if let Some(executed_tool_calls) = exec.session.services.executed_tool_calls.as_ref() { - executed_tool_calls.register_cell(&cell_id, &call_id); + executed_tool_calls.start_cell(&cell_id, &call_id); } let runtime_cell_id = cell_id.to_string(); let code_cell_trace = exec diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index 26559df62e..6cbc1a2d1c 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/codex-rs/core/src/tools/code_mode/mod.rs @@ -30,6 +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::context::FunctionToolOutput; use crate::tools::context::SharedTurnDiffTracker; use crate::tools::context::ToolPayload; @@ -80,8 +81,9 @@ impl CodeModeService { pub(crate) fn new( session_provider: Arc, config: &CodeModeConfig, + executed_tool_calls: Option>, ) -> Self { - let dispatch_broker = Arc::new(CodeModeDispatchBroker::new()); + let dispatch_broker = Arc::new(CodeModeDispatchBroker::new(executed_tool_calls)); let availability = session_provider.availability(); Self { session: OnceCell::new(), @@ -341,12 +343,16 @@ fn truncate_code_mode_result( truncate_function_output_items_with_policy(&items, policy, estimate_audio_token_count) } -async fn call_nested_tool( +// Submit synchronously so the recorder sees the call before the cell's dispatch gate closes. +fn submit_nested_tool( exec: ExecContext, tool_runtime: ToolCallRuntime, invocation: CodeModeNestedToolCall, cancellation_token: CancellationToken, -) -> Result { +) -> Result< + impl std::future::Future> + Send + 'static, + FunctionCallError, +> { let CodeModeNestedToolCall { cell_id, runtime_tool_call_id, @@ -380,17 +386,15 @@ async fn call_nested_tool( call_id: call.call_id.clone(), cell_id: cell_id.to_string(), }); - let result = tool_runtime - .handle_tool_call_with_source( - call, - ToolCallSource::CodeMode { - cell_id: cell_id.to_string(), - runtime_tool_call_id, - }, - cancellation_token, - ) - .await?; - Ok(result.code_mode_result()) + let result = tool_runtime.handle_tool_call_with_source( + call, + ToolCallSource::CodeMode { + cell_id: cell_id.to_string(), + runtime_tool_call_id, + }, + cancellation_token, + ); + Ok(async move { Ok(result.await?.code_mode_result()) }) } fn build_nested_tool_payload( diff --git a/codex-rs/core/src/tools/executed_tool_calls.rs b/codex-rs/core/src/tools/executed_tool_calls.rs index f16363fe83..cda42c5a13 100644 --- a/codex-rs/core/src/tools/executed_tool_calls.rs +++ b/codex-rs/core/src/tools/executed_tool_calls.rs @@ -3,6 +3,7 @@ use std::collections::HashSet; use codex_code_mode::CellId; 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_prioritizing_recent; use codex_protocol::models::executed_tool_call_metadata_bytes; @@ -30,14 +31,70 @@ struct ExecutedToolCallRecorderState { direct_calls: HashMap, cells: HashMap, output_cells: HashMap, - retained_calls: HashMap<(std::mem::Discriminant, String), Vec>, + retained_calls: HashMap<(std::mem::Discriminant, String), RetainedToolCalls>, pending_nested_calls: usize, } +/// Keep each output's calls and completion marker together through replay and pruning. +#[derive(Default)] +struct RetainedToolCalls { + calls: Vec, + complete: bool, + cell_id: Option, + runtime_cell_id: Option, +} + +#[derive(Default, PartialEq, Eq)] +enum CellCompletion { + #[default] + Unobserved, + Started, + Recording, + Incomplete, + Complete, +} + #[derive(Default)] struct RecordedCell { pending_calls: Vec, pending_full_argument_bytes: usize, + completion: CellCompletion, + originating_call_id: Option, +} + +impl ExecutedToolCallRecorderState { + fn register_cell(&mut self, cell_id: &CellId, output_call_id: &str) { + if self.cells.len() >= MAX_PENDING_EXECUTED_TOOL_CALLS && !self.cells.contains_key(cell_id) + { + let output_cells = self.output_cells.values().collect::>(); + let finished_cell = self.cells.iter().find_map(|(id, cell)| { + // A finished cell can still have missing or truncated tool call records. + (matches!( + cell.completion, + CellCompletion::Complete | CellCompletion::Incomplete + ) && cell.pending_calls.is_empty() + && !output_cells.contains(id)) + .then(|| id.clone()) + }); + if let Some(id) = finished_cell { + self.cells.remove(&id); + } + } + if (self.cells.len() >= MAX_PENDING_EXECUTED_TOOL_CALLS + && !self.cells.contains_key(cell_id)) + || (self.output_cells.len() >= MAX_PENDING_EXECUTED_TOOL_CALLS + && !self.output_cells.contains_key(output_call_id)) + { + return; + } + self.cells + .entry(cell_id.clone()) + .or_default() + .originating_call_id + .get_or_insert_with(|| output_call_id.to_string()); + self.output_cells + .insert(output_call_id.to_string(), cell_id.clone()); + } } impl ExecutedToolCallRecorder { @@ -133,6 +190,9 @@ impl ExecutedToolCallRecorder { || (state.cells.len() >= MAX_PENDING_EXECUTED_TOOL_CALLS && !state.cells.contains_key(&cell_id)) { + if let Some(cell) = state.cells.get_mut(&cell_id) { + cell.completion = CellCompletion::Incomplete; + } return; } let at_pending_call_limit = state.pending_nested_calls == MAX_PENDING_EXECUTED_TOOL_CALLS; @@ -151,6 +211,17 @@ impl ExecutedToolCallRecorder { } else { ExecutedToolCall::truncated(call.name, original_bytes, max_bytes) }; + cell.completion = if matches!( + cell.completion, + CellCompletion::Started | CellCompletion::Recording + ) && !matches!( + call.arguments(), + ExecutedToolCallArguments::Truncated { .. } + ) { + CellCompletion::Recording + } else { + CellCompletion::Incomplete + }; cell.pending_calls.push(call); state.pending_nested_calls += 1; } @@ -201,42 +272,71 @@ impl ExecutedToolCallRecorder { _ => 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) = state.retained_calls.get(&key) { + } else if let Some(retained) = retained { if !pending_retained_outputs.remove(&key) { continue; } - retained.clone() + retained.calls.clone() } else { let mut calls = state .direct_calls .remove(call_id) .into_iter() .collect::>(); - if let Some(cell_id) = state.output_cells.remove(call_id) - && let Some(mut cell) = state.cells.remove(&cell_id) + if let Some(output_cell_id) = state.output_cells.remove(call_id) + && let Some(cell) = state.cells.get_mut(&output_cell_id) { - state.pending_nested_calls = state - .pending_nested_calls - .saturating_sub(cell.pending_calls.len()); + cell_id = cell.originating_call_id.clone(); + runtime_cell_id = Some(output_cell_id.clone()); + let pending_calls = cell.pending_calls.len(); + calls.append(&mut cell.pending_calls); + 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(|_, output_cell_id| output_cell_id != &cell_id); - calls.append(&mut cell.pending_calls); + .retain(|_, registered_cell_id| registered_cell_id != &output_cell_id); } - if calls.is_empty() { + if calls.is_empty() && !complete { continue; } retry_cache.insert(key.clone(), calls.clone()); - state.retained_calls.insert(key, calls.clone()); + state.retained_calls.insert( + key, + RetainedToolCalls { + calls: calls.clone(), + complete, + cell_id: cell_id.clone(), + runtime_cell_id, + }, + ); 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; } @@ -247,7 +347,7 @@ impl ExecutedToolCallRecorder { } if retained_bytes > MAX_EXECUTED_TOOL_CALL_FULL_ARGUMENT_BYTES_PER_OUTPUT { bound_executed_tool_calls_for_prompt_prioritizing_recent(items); - state.retained_calls.clear(); + let retained_before_bounding = std::mem::take(&mut state.retained_calls); for item in items { let call_id = match &*item { ResponseItem::FunctionCallOutput { @@ -261,15 +361,33 @@ impl ExecutedToolCallRecorder { } => call_id, _ => continue, }; - if let Some(calls) = item - .executed_tool_call_metadata() - .and_then(|metadata| metadata.executed_tool_calls.as_ref()) - .filter(|calls| !calls.is_empty()) + let key = (std::mem::discriminant(&*item), call_id.clone()); + let metadata = item.executed_tool_call_metadata(); + if let Some(retained) = retained_before_bounding.get(&key) + && let Some(runtime_cell_id) = &retained.runtime_cell_id + && metadata.and_then(|metadata| metadata.executed_tool_calls.as_ref()) + != Some(&retained.calls) + && let Some(cell) = state.cells.get_mut(runtime_cell_id) { - state.retained_calls.insert( - (std::mem::discriminant(&*item), call_id.clone()), - calls.clone(), - ); + 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 = retained_before_bounding + .get(&key) + .and_then(|retained| retained.runtime_cell_id.clone()); + 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); } } } @@ -282,17 +400,32 @@ impl ExecutedToolCallRecorder { .state .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); - if (state.cells.len() >= MAX_PENDING_EXECUTED_TOOL_CALLS - && !state.cells.contains_key(cell_id)) - || (state.output_cells.len() >= MAX_PENDING_EXECUTED_TOOL_CALLS - && !state.output_cells.contains_key(output_call_id)) - { - return; + 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 + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + state.register_cell(cell_id, output_call_id); + if let Some(cell) = state.cells.get_mut(cell_id) { + cell.completion = CellCompletion::Started; + } + } + + pub(crate) fn finish_cell_recording(&self, cell_id: &CellId) { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(cell) = state.cells.get_mut(cell_id) { + if cell.completion == CellCompletion::Recording { + cell.completion = CellCompletion::Complete; + } else if cell.completion != CellCompletion::Complete && cell.pending_calls.is_empty() { + state.cells.remove(cell_id); + } } - state.cells.entry(cell_id.clone()).or_default(); - state - .output_cells - .insert(output_call_id.to_string(), cell_id.clone()); } } diff --git a/codex-rs/core/src/tools/executed_tool_calls_tests.rs b/codex-rs/core/src/tools/executed_tool_calls_tests.rs index 6a4a8d6c8a..3a2011aaad 100644 --- a/codex-rs/core/src/tools/executed_tool_calls_tests.rs +++ b/codex-rs/core/src/tools/executed_tool_calls_tests.rs @@ -1,9 +1,17 @@ use codex_protocol::models::FunctionCallOutputPayload; +use codex_protocol::models::ResponseInputItem; use pretty_assertions::assert_eq; use serde_json::json; use super::*; +fn output(call_id: &str) -> ResponseItem { + ResponseItem::from(ResponseInputItem::FunctionCallOutput { + call_id: call_id.to_string(), + output: FunctionCallOutputPayload::from_text(String::new()), + }) +} + #[test] fn executed_tool_call_recorder_bounds_pending_calls_and_preserves_overflow() { let recorder = ExecutedToolCallRecorder::default(); @@ -24,7 +32,7 @@ fn executed_tool_call_recorder_bounds_pending_calls_and_preserves_overflow() { } let cell_id = CellId::new("bounded-cell".to_string()); - recorder.register_cell(&cell_id, "bounded-output"); + recorder.start_cell(&cell_id, "bounded-output"); for _ in 0..MAX_PENDING_EXECUTED_TOOL_CALLS + 2 { recorder.record_nested_tool_call( cell_id.clone(), @@ -84,8 +92,15 @@ fn executed_tool_call_recorder_bounds_pending_calls_and_preserves_overflow() { internal_chat_message_metadata_passthrough: None, }]; let mut retry_cache = HashMap::new(); + recorder.finish_cell_recording(&cell_id); recorder.attach_pending_to_prompt(&mut items, &mut retry_cache); + assert_eq!( + items[0] + .executed_tool_call_metadata() + .and_then(|metadata| metadata.tool_calls_complete), + None, + ); let calls = items[0] .executed_tool_call_metadata() .and_then(|metadata| metadata.executed_tool_calls.as_ref()) @@ -198,7 +213,7 @@ fn executed_tool_call_recorder_bounds_retained_history_and_reports_omissions() { let retained_bytes = state .retained_calls .values() - .map(serialized_json_bytes) + .map(|retained| serialized_json_bytes(&retained.calls)) .sum::>() .expect("retained calls must serialize"); assert!(retained_bytes <= MAX_EXECUTED_TOOL_CALL_FULL_ARGUMENT_BYTES_PER_OUTPUT); @@ -219,3 +234,189 @@ fn executed_tool_call_recorder_bounds_retained_history_and_reports_omissions() { assert!(omitted_calls > 0); assert_eq!(metadata.len() as u64 + omitted_calls, 512); } + +#[test] +fn tool_call_completeness_requires_finished_lossless_recording() { + for scenario in ["empty", "unobserved", "late_call"] { + let recorder = ExecutedToolCallRecorder::default(); + let cell_id = CellId::new(scenario.to_string()); + if scenario == "unobserved" { + recorder.register_cell(&cell_id, "output"); + } else { + recorder.start_cell(&cell_id, "output"); + } + if scenario == "late_call" { + recorder.finish_cell_recording(&cell_id); + } + if scenario != "empty" { + recorder.record_nested_tool_call( + cell_id.clone(), + ExecutedToolCall::new("nested_tool".to_string(), json!({})), + /*original_bytes*/ 2, + ); + } + recorder.finish_cell_recording(&cell_id); + + let mut items = [output("output")]; + recorder.attach_pending_to_prompt(&mut items, &mut HashMap::new()); + assert_eq!( + items[0] + .executed_tool_call_metadata() + .and_then(|metadata| metadata.tool_calls_complete), + None, + "{scenario} must not claim complete recording", + ); + } +} + +#[test] +fn tool_call_completeness_survives_waits_without_changing_deltas() { + for truncated in [false, true] { + let recorder = ExecutedToolCallRecorder::default(); + let cell_id = CellId::new("multi-wait".to_string()); + recorder.start_cell(&cell_id, "exec"); + let mut history = Vec::new(); + let mut expected = Vec::new(); + for (index, call_id) in ["exec", "wait-1", "wait-2", "wait-3"] + .into_iter() + .enumerate() + { + recorder.register_cell(&cell_id, call_id); + let mut expected_output = output(call_id); + history.push(expected_output.clone()); + if index < 2 { + // Identical calls are distinct submissions; truncation stays sticky across waits. + let original_bytes = if truncated && index == 1 { 9_000 } else { 2 }; + let call = if original_bytes > MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES { + ExecutedToolCall::truncated( + "nested_tool".to_string(), + original_bytes, + MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES, + ) + } else { + ExecutedToolCall::new("nested_tool".to_string(), json!({})) + }; + recorder.record_nested_tool_call(cell_id.clone(), call.clone(), original_bytes); + expected_output.append_executed_tool_calls(vec![call]); + } else if index == 3 { + recorder.finish_cell_recording(&cell_id); + } + if index < 2 || index == 3 && !truncated { + expected_output.set_tool_call_cell_id("exec"); + } + if index == 3 && !truncated { + expected_output.mark_tool_calls_complete(); + } + expected.push(expected_output); + let mut retry_cache = HashMap::new(); + for _ in 0..2 { + let mut prompt = history.clone(); + assert!(recorder.attach_pending_to_prompt(&mut prompt, &mut retry_cache)); + assert_eq!(prompt, expected); + } + } + let state = recorder.state.lock().unwrap(); + assert!(state.cells.is_empty()); + assert_eq!(state.pending_nested_calls, 0); + } +} + +#[test] +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(); + recorder.start_cell(&runtime_cell_id, originating_call_id); + recorder.record_nested_tool_call( + runtime_cell_id.clone(), + ExecutedToolCall::new("nested_tool".to_string(), json!({})), + /*original_bytes*/ 2, + ); + + let mut initial = [output(originating_call_id)]; + recorder.attach_pending_to_prompt(&mut initial, &mut HashMap::new()); + assert_eq!( + initial[0] + .executed_tool_call_metadata() + .and_then(|metadata| metadata.cell_id.as_deref()), + Some(originating_call_id), + ); + + recorder.register_cell(&runtime_cell_id, "wait"); + recorder.finish_cell_recording(&runtime_cell_id); + let mut final_output = [output("wait")]; + recorder.attach_pending_to_prompt(&mut final_output, &mut HashMap::new()); + let metadata = final_output[0] + .executed_tool_call_metadata() + .expect("completed cell must have metadata"); + assert_eq!(metadata.cell_id.as_deref(), Some(originating_call_id)); + assert_eq!(metadata.tool_calls_complete, Some(true)); + } +} + +#[test] +fn request_truncation_prevents_completion_after_compaction() { + let recorder = ExecutedToolCallRecorder::default(); + let cell_id = CellId::new("compacted-cell".to_string()); + recorder.start_cell(&cell_id, "exec"); + + let arguments = json!({ + "payload": "x".repeat(MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES - r#"{"payload":""}"#.len()), + }); + for _ in 0..4 { + recorder.record_nested_tool_call( + cell_id.clone(), + ExecutedToolCall::new("nested_tool".to_string(), arguments.clone()), + MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES, + ); + } + + let mut initial = [output("exec")]; + assert!(recorder.attach_pending_to_prompt(&mut initial, &mut HashMap::new())); + assert!( + initial[0] + .executed_tool_call_metadata() + .and_then(|metadata| metadata.executed_tool_calls.as_ref()) + .is_some_and(|calls| calls.iter().any(|call| matches!( + call.arguments(), + ExecutedToolCallArguments::Truncated { .. } + ))) + ); + + recorder.attach_pending_to_prompt(&mut [], &mut HashMap::new()); + recorder.register_cell(&cell_id, "wait"); + recorder.finish_cell_recording(&cell_id); + let mut final_output = [output("wait")]; + recorder.attach_pending_to_prompt(&mut final_output, &mut HashMap::new()); + assert_eq!( + final_output[0] + .executed_tool_call_metadata() + .and_then(|metadata| metadata.tool_calls_complete), + None, + ); +} + +#[test] +fn finished_cells_without_more_waits_do_not_block_new_calls() { + let recorder = ExecutedToolCallRecorder::default(); + 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}")); + recorder.start_cell(&cell, cell.as_str()); + recorder.record_nested_tool_call(cell.clone(), call.clone(), /*original_bytes*/ 2); + recorder.attach_pending_to_prompt(&mut [output(cell.as_str())], &mut HashMap::new()); + recorder.finish_cell_recording(&cell); + } + let fresh = CellId::new("fresh".to_string()); + recorder.start_cell(&fresh, "fresh-output"); + recorder.record_nested_tool_call(fresh.clone(), call.clone(), /*original_bytes*/ 2); + recorder.finish_cell_recording(&fresh); + let mut expected = output("fresh-output"); + expected.append_executed_tool_calls(vec![call]); + expected.set_tool_call_cell_id("fresh-output"); + expected.mark_tool_calls_complete(); + let mut items = [output("fresh-output")]; + assert!(recorder.attach_pending_to_prompt(&mut items, &mut HashMap::new())); + assert_eq!(items, [expected]); +} diff --git a/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs b/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs index a06c87d3d3..44c03df82c 100644 --- a/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs +++ b/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs @@ -283,6 +283,7 @@ async fn missing_code_mode_wait_traces_only_the_wait_tool_call() -> anyhow::Resu session.services.code_mode_service = CodeModeService::new( Arc::new(MissingCellCodeModeSessionProvider), &turn.config.code_mode, + session.services.executed_tool_calls.clone(), ); attach_test_trace(&mut session, &turn, temp.path())?; diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index 9c9c195a33..90eb2c6d2e 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -1203,6 +1203,118 @@ async fn mcp_code_mode_exclusion_does_not_change_direct_mode_tool_exposure() -> Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[test_case(false, false, false; "disabled")] +#[test_case(true, false, false; "completed")] +#[test_case(true, true, false; "yielded")] +#[test_case(true, false, true; "oversized")] +async fn code_mode_tool_call_completeness_is_private_and_opt_in( + metadata_enabled: bool, + yielded: bool, + oversized: bool, +) -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let code = if oversized { + r#" +const args = { barrier: { id: "", participants: 1 } }; +args.barrier.id = "x".repeat(8192 - JSON.stringify(args).length); +for (let index = 0; index < 4; index++) await tools.test_sync_tool(args); +text("done"); +yield_control(); +await new Promise(() => {}); +"# + } else if yielded { + r#"await tools.test_sync_tool({}); text("done"); yield_control(); await new Promise(() => {});"# + } else { + r#"await tools.test_sync_tool({}); text("done");"# + }; + // run_code_mode_turn_with_config enables this feature; production keeps it off by default. + let (test, follow_up) = + run_code_mode_turn_with_config(&server, "Record a nested tool call", code, move |config| { + if !metadata_enabled { + config + .features + .disable(Feature::ExecutedToolCallMetadata) + .expect("tool call metadata should be disabled"); + } + if oversized { + let _ = config.features.disable(Feature::RemoteCompactionV2); + } + }) + .await?; + + let request = follow_up.single_request(); + let first_output = request.custom_tool_call_output("call-1"); + let metadata = &first_output["internal_chat_message_metadata_passthrough"]; + let calls = serde_json::json!([{ "name": "test_sync_tool", "arguments": {} }]); + let complete = Value::Bool(true); + if oversized { + assert!( + metadata["executed_tool_calls"] + .as_array() + .is_some_and(|calls| calls.iter().any(|call| call["arguments"] + .get("_codex_executed_tool_call_truncated") + .is_some())) + ); + } else { + assert_eq!( + metadata.get("executed_tool_calls"), + metadata_enabled.then_some(&calls), + ); + } + assert_eq!( + metadata.get("tool_calls_complete"), + (metadata_enabled && !yielded && !oversized).then_some(&complete), + ); + assert_eq!( + metadata.get("cell_id").and_then(Value::as_str), + metadata_enabled.then_some("call-1"), + ); + if yielded || oversized { + let first_items = custom_tool_output_items(&request, "call-1"); + let cell_id = extract_running_cell_id(text_item(&first_items, /*index*/ 0)); + + if oversized { + responses::mount_compact_user_history_with_summary_once(&server, "compacted history") + .await; + test.codex.submit(Op::Compact).await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + let wait = responses::mount_function_call_agent_response( + &server, + "call-2", + &serde_json::to_string(&serde_json::json!({ + "cell_id": cell_id, + "terminate": true, + }))?, + "wait", + ) + .await; + test.submit_turn("Finish the compacted cell").await?; + let final_request = wait.completion.single_request(); + assert!( + final_request + .input() + .iter() + .all(|item| item["call_id"] != "call-1") + ); + assert!( + final_request.function_call_output("call-2") + ["internal_chat_message_metadata_passthrough"] + .get("tool_calls_complete") + .is_none() + ); + } + } + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn code_mode_only_guides_all_tools_search_and_calls_deferred_app_tools() -> Result<()> { skip_if_no_network!(Ok(())); @@ -1608,6 +1720,7 @@ async fn code_mode_nested_tool_calls_can_run_in_parallel() -> Result<()> { .with_model("test-gpt-5.1-codex") .with_config(move |config| { let _ = config.features.enable(Feature::CodeMode); + let _ = config.features.enable(Feature::ExecutedToolCallMetadata); }); let test = builder.build(&server).await?; @@ -2074,13 +2187,16 @@ try { .await?; let request = second_mock.single_request(); + let metadata = + &request.custom_tool_call_output("call-1")["internal_chat_message_metadata_passthrough"]; assert_eq!( - request.custom_tool_call_output("call-1")["internal_chat_message_metadata_passthrough"]["executed_tool_calls"], + metadata["executed_tool_calls"], serde_json::json!([ {"name": "exec_command", "arguments": {}}, ]), "failed nested tool attempts remain private request metadata", ); + assert_eq!(metadata["tool_calls_complete"], Value::Bool(true)); let (output, success) = custom_tool_output_body_and_success(&request, "call-1"); assert_ne!( success, @@ -2478,6 +2594,7 @@ async fn code_mode_can_run_multiple_yielded_sessions() -> Result<()> { let server = responses::start_mock_server().await; let mut builder = test_codex().with_config(move |config| { let _ = config.features.enable(Feature::CodeMode); + let _ = config.features.enable(Feature::ExecutedToolCallMetadata); }); let test = builder.build(&server).await?; let session_a_gate = test.workspace_path("code-mode-session-a.ready"); @@ -2487,6 +2604,7 @@ async fn code_mode_can_run_multiple_yielded_sessions() -> Result<()> { let session_a_code = format!( r#" +await tools.exec_command({{ cmd: "printf cell-a" }}); text("session a start"); yield_control(); {session_a_wait} @@ -2495,6 +2613,7 @@ text("session a done"); ); let session_b_code = format!( r#" +await tools.exec_command({{ cmd: "printf cell-b" }}); text("session b start"); yield_control(); {session_b_wait} @@ -2565,7 +2684,7 @@ text("session b done"); "wait", &serde_json::to_string(&serde_json::json!({ "cell_id": session_a_id.clone(), - "yield_time_ms": 1_000, + "yield_time_ms": 5_000, }))?, ), ev_completed("resp-5"), @@ -2605,7 +2724,7 @@ text("session b done"); "wait", &serde_json::to_string(&serde_json::json!({ "cell_id": session_b_id.clone(), - "yield_time_ms": 1_000, + "yield_time_ms": 5_000, }))?, ), ev_completed("resp-7"), @@ -2634,6 +2753,44 @@ text("session b done"); text_item(&fourth_items, /*index*/ 0), ); assert_eq!(text_item(&fourth_items, /*index*/ 1), "session b done"); + for (output, originating_call_id, expected_command) in [ + ( + first_request.custom_tool_call_output("call-1"), + "call-1", + Some("printf cell-a"), + ), + ( + second_request.custom_tool_call_output("call-2"), + "call-2", + Some("printf cell-b"), + ), + (third_request.function_call_output("call-3"), "call-1", None), + ( + fourth_request.function_call_output("call-4"), + "call-2", + None, + ), + ] { + let metadata = &output["internal_chat_message_metadata_passthrough"]; + assert_eq!(metadata["cell_id"], serde_json::json!(originating_call_id)); + assert_eq!( + metadata.get("tool_calls_complete").and_then(Value::as_bool), + expected_command.is_none().then_some(true), + ); + assert_eq!( + metadata + .get("executed_tool_calls") + .and_then(Value::as_array) + .and_then(|calls| { + calls.iter().find_map(|call| { + call["arguments"]["cmd"] + .as_str() + .filter(|command| command.starts_with("printf cell-")) + }) + }), + expected_command, + ); + } Ok(()) } @@ -3333,6 +3490,7 @@ async fn code_mode_interrupt_terminates_active_cells_and_nested_tools() -> Resul .with_config(|config| { let _ = config.features.enable(Feature::CodeMode); let _ = config.features.enable(Feature::CodeModeInterrupt); + let _ = config.features.enable(Feature::ExecutedToolCallMetadata); }); let test = builder.build_with_auto_env(&server).await?; @@ -3453,6 +3611,12 @@ async fn code_mode_interrupt_terminates_active_cells_and_nested_tools() -> Resul test.submit_turn("verify interrupted cells and reuse their session") .await?; let requests = recovery.requests(); + let metadata = &requests[0].custom_tool_call_output("call-interrupted")["internal_chat_message_metadata_passthrough"]; + assert_eq!( + metadata["executed_tool_calls"], + serde_json::json!([{"name": "test_sync_tool", "arguments": {"sleep_after_ms": 60_000}}]), + ); + assert_eq!(metadata["tool_calls_complete"], Value::Bool(true)); let background_output = function_tool_output_items(&requests[1], "call-wait-background"); assert!(text_item(&background_output, /*index*/ 1).contains("not found")); let active_output = function_tool_output_items(&requests[2], "call-wait-active"); diff --git a/codex-rs/protocol/src/models.rs b/codex-rs/protocol/src/models.rs index c7fb57ee61..9005f811a0 100644 --- a/codex-rs/protocol/src/models.rs +++ b/codex-rs/protocol/src/models.rs @@ -937,11 +937,23 @@ pub struct InternalChatMessageMetadataPassthrough { #[schemars(skip)] #[ts(skip)] pub content_item_kinds: Option>, + // Ignore input values so requests cannot fake tool call records. + /// Host-owned Code Mode cell shared by its `exec` and subsequent `wait` outputs. + #[serde(default, skip_deserializing, skip_serializing_if = "Option::is_none")] + #[schemars(skip)] + #[ts(skip)] + pub cell_id: Option, /// Warehouse-only Responses metadata, not part of the public app-server protocol. #[serde(default, skip_deserializing, skip_serializing_if = "Option::is_none")] #[schemars(skip)] #[ts(skip)] pub executed_tool_calls: Option>, + /// Whether the host finished recording this cell's calls without losing calls or arguments. + /// This describes the call inventory across the cell's outputs, not tool success. + #[serde(default, skip_deserializing, skip_serializing_if = "Option::is_none")] + #[schemars(skip)] + #[ts(skip)] + pub tool_calls_complete: Option, } impl InternalChatMessageMetadataPassthrough { diff --git a/codex-rs/protocol/src/models/executed_tool_calls.rs b/codex-rs/protocol/src/models/executed_tool_calls.rs index 668b06e11d..8f504ff05c 100644 --- a/codex-rs/protocol/src/models/executed_tool_calls.rs +++ b/codex-rs/protocol/src/models/executed_tool_calls.rs @@ -16,12 +16,28 @@ const INTERNAL_CHAT_MESSAGE_METADATA_PASSTHROUGH_FIELD_BYTES: usize = fn executed_tool_call_metadata_field_bytes( metadata: &InternalChatMessageMetadataPassthrough, ) -> usize { - EXECUTED_TOOL_CALL_METADATA_FIELD_BYTES - + if metadata.turn_id.is_some() { - 1 - } else { - INTERNAL_CHAT_MESSAGE_METADATA_PASSTHROUGH_FIELD_BYTES + 3 - } + let fields = InternalChatMessageMetadataPassthrough { + cell_id: metadata.cell_id.clone(), + tool_calls_complete: metadata.tool_calls_complete, + ..Default::default() + }; + let mut bytes = + serde_json::to_vec(&fields).map_or(usize::MAX, |fields| fields.len().saturating_sub(2)); + if metadata.executed_tool_calls.is_some() { + bytes = bytes + .saturating_add(usize::from(bytes > 0)) + .saturating_add(EXECUTED_TOOL_CALL_METADATA_FIELD_BYTES); + } + if bytes == 0 { + 0 + } else if metadata.turn_id.is_some() + || metadata.create_time.is_some() + || metadata.content_item_kinds.is_some() + { + bytes + 1 + } else { + bytes + INTERNAL_CHAT_MESSAGE_METADATA_PASSTHROUGH_FIELD_BYTES + 3 + } } /// Returns the exact serialized wire size of an item's attempted-tool metadata. @@ -29,20 +45,15 @@ pub fn executed_tool_call_metadata_bytes(item: &ResponseItem) -> usize { let Some(metadata) = item.executed_tool_call_metadata() else { return 0; }; - let Some(calls) = metadata + metadata .executed_tool_calls .as_ref() - .filter(|calls| !calls.is_empty()) - else { - return 0; - }; - serde_json::to_vec(calls) - .map(|bytes| { - bytes - .len() - .saturating_add(executed_tool_call_metadata_field_bytes(metadata)) + .map_or(0, |calls| { + serde_json::to_vec(calls) + .map(|calls| calls.len()) + .unwrap_or(usize::MAX) }) - .unwrap_or(usize::MAX) + .saturating_add(executed_tool_call_metadata_field_bytes(metadata)) } /// Bounds attempted-tool metadata fairly across the complete serialized request. @@ -64,6 +75,7 @@ fn bound_executed_tool_calls_for_prompt_with_priority( let mut remaining_items = 0_usize; let mut original_calls = 0_usize; let mut original_metadata_bytes = 0_usize; + let mut truncated = false; for item in items.iter_mut() { let Some(calls) = item @@ -72,6 +84,8 @@ fn bound_executed_tool_calls_for_prompt_with_priority( .and_then(|metadata| metadata.executed_tool_calls.as_mut()) .filter(|calls| !calls.is_empty()) else { + original_metadata_bytes = + original_metadata_bytes.saturating_add(executed_tool_call_metadata_bytes(item)); continue; }; for call in calls { @@ -86,6 +100,7 @@ fn bound_executed_tool_calls_for_prompt_with_priority( /*omitted_calls*/ None, ); } + truncated |= call.truncation().is_some(); original_calls = original_calls.saturating_add(1).saturating_add( call.truncation() .and_then(|truncation| truncation.omitted_calls) @@ -97,15 +112,51 @@ fn bound_executed_tool_calls_for_prompt_with_priority( original_metadata_bytes.saturating_add(executed_tool_call_metadata_bytes(item)); } + // A terminal marker can be on an empty wait output, separate from the lost calls. + if truncated || original_metadata_bytes > MAX_EXECUTED_TOOL_CALL_METADATA_BYTES { + for item in items.iter_mut() { + let item_bytes = executed_tool_call_metadata_bytes(item); + if let Some(metadata) = item + .internal_chat_message_metadata_passthrough_mut() + .and_then(Option::as_mut) + && metadata.tool_calls_complete == Some(true) + { + metadata.tool_calls_complete = None; + original_metadata_bytes = original_metadata_bytes.saturating_sub( + item_bytes.saturating_sub(executed_tool_call_metadata_bytes(item)), + ); + } + } + } + if original_metadata_bytes <= MAX_EXECUTED_TOOL_CALL_METADATA_BYTES { + return; + } + + // Drop marker-only waits before spending the budget on recorded calls. + for item in items.iter_mut() { + if item.executed_tool_call_metadata().is_some_and(|metadata| { + metadata + .executed_tool_calls + .as_ref() + .is_none_or(Vec::is_empty) + }) { + original_metadata_bytes = + original_metadata_bytes.saturating_sub(executed_tool_call_metadata_bytes(item)); + item.clear_executed_tool_calls(); + } + } if original_metadata_bytes <= MAX_EXECUTED_TOOL_CALL_METADATA_BYTES { return; } let overflow_fallback = items.iter().enumerate().find_map(|(index, item)| { - item.executed_tool_call_metadata() - .and_then(|metadata| metadata.executed_tool_calls.as_ref()) - .and_then(|calls| calls.first()) - .map(|call| (index, call.clone())) + item.executed_tool_call_metadata().and_then(|metadata| { + metadata + .executed_tool_calls + .as_ref() + .and_then(|calls| calls.first()) + .map(|call| (index, call.clone(), metadata.cell_id.clone())) + }) }); let omitted_call_reservation = serde_json::to_vec(&serde_json::json!({ @@ -162,7 +213,7 @@ fn bound_executed_tool_calls_for_prompt_with_priority( } if represented_calls == 0 { - if let Some((index, mut call)) = overflow_fallback { + if let Some((index, mut call, cell_id)) = overflow_fallback { let original_bytes = call .truncation() .map(|truncation| truncation.original_bytes) @@ -183,6 +234,9 @@ fn bound_executed_tool_calls_for_prompt_with_priority( (name_boundary < original_name_bytes).then_some(original_name_bytes), ); items[index].append_executed_tool_calls(vec![call]); + if let Some(cell_id) = cell_id { + items[index].set_tool_call_cell_id(&cell_id); + } } return; } @@ -326,21 +380,39 @@ impl ExecutedToolCall { } impl ResponseItem { + fn ensure_tool_call_metadata(&mut self) -> Option<&mut InternalChatMessageMetadataPassthrough> { + self.internal_chat_message_metadata_passthrough_mut() + .map(Option::get_or_insert_default) + } + + /// Associates host-recorded calls and completeness with their Code Mode cell. + pub fn set_tool_call_cell_id(&mut self, cell_id: &str) { + if let Some(metadata) = self.ensure_tool_call_metadata() { + metadata.cell_id = Some(cell_id.to_string()); + } + } + /// Attaches model-attempted tool invocations without replacing existing item metadata. pub fn append_executed_tool_calls(&mut self, calls: Vec) { if calls.is_empty() { return; } - let Some(metadata) = self.internal_chat_message_metadata_passthrough_mut() else { + let Some(metadata) = self.ensure_tool_call_metadata() else { return; }; metadata - .get_or_insert_with(InternalChatMessageMetadataPassthrough::default) .executed_tool_calls .get_or_insert_with(Vec::new) .extend(calls); } + /// Marks a host-owned cell's recorded tool calls as complete. + pub fn mark_tool_calls_complete(&mut self) { + if let Some(metadata) = self.ensure_tool_call_metadata() { + metadata.tool_calls_complete = Some(true); + } + } + /// Returns warehouse-only attempted-tool metadata for any supported item variant. pub fn executed_tool_call_metadata(&self) -> Option<&InternalChatMessageMetadataPassthrough> { self.internal_chat_message_metadata_passthrough() @@ -404,10 +476,7 @@ impl ResponseItem { }); if calls.is_empty() { - passthrough.executed_tool_calls = None; - if *passthrough == InternalChatMessageMetadataPassthrough::default() { - *metadata = None; - } + self.clear_executed_tool_calls(); } } @@ -419,7 +488,9 @@ impl ResponseItem { let Some(passthrough) = metadata.as_mut() else { return; }; + passthrough.cell_id = None; passthrough.executed_tool_calls = None; + passthrough.tool_calls_complete = None; if *passthrough == InternalChatMessageMetadataPassthrough::default() { *metadata = None; } diff --git a/codex-rs/protocol/src/models/executed_tool_calls_tests.rs b/codex-rs/protocol/src/models/executed_tool_calls_tests.rs index e541f56b95..b49443fd11 100644 --- a/codex-rs/protocol/src/models/executed_tool_calls_tests.rs +++ b/codex-rs/protocol/src/models/executed_tool_calls_tests.rs @@ -3,6 +3,7 @@ use pretty_assertions::assert_eq; use super::super::FunctionCallOutputBody; use super::super::FunctionCallOutputPayload; +use super::super::ResponseInputItem; use super::*; fn passthrough_metadata(turn_id: &str) -> InternalChatMessageMetadataPassthrough { @@ -43,6 +44,8 @@ fn executed_tool_call_prompt_budget_includes_metadata_fields() -> Result<()> { String::new(), serde_json::Value::Null, )]); + item.set_tool_call_cell_id("cell-\"\\"); + item.mark_tool_calls_complete(); item }) .collect::>(); @@ -50,6 +53,12 @@ fn executed_tool_call_prompt_budget_includes_metadata_fields() -> Result<()> { assert!(metadata_bytes(&items)? > MAX_EXECUTED_TOOL_CALL_METADATA_BYTES); bound_executed_tool_calls_for_prompt(&mut items); assert!(metadata_bytes(&items)? <= MAX_EXECUTED_TOOL_CALL_METADATA_BYTES); + for metadata in items + .iter() + .filter_map(ResponseItem::executed_tool_call_metadata) + { + assert_eq!(metadata.tool_calls_complete, None); + } let calls = items .iter() @@ -94,6 +103,7 @@ fn executed_tool_call_prompt_budget_includes_metadata_fields() -> Result<()> { oversized_name.clone(), serde_json::Value::Null, )]); + item.set_tool_call_cell_id("cell-\"\\"); item }) .collect::>(); @@ -108,6 +118,13 @@ fn executed_tool_call_prompt_budget_includes_metadata_fields() -> Result<()> { .collect::>(); assert_eq!(calls.len(), 1); let call = calls[0]; + assert_eq!( + oversized_items + .iter() + .filter_map(ResponseItem::executed_tool_call_metadata) + .find_map(|metadata| metadata.cell_id.as_deref()), + Some("cell-\"\\"), + ); assert_eq!(call.name.len(), MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES / 2); assert!(oversized_name.starts_with(&call.name)); let truncation = call.truncation().expect("trusted omission marker"); @@ -122,6 +139,29 @@ fn executed_tool_call_prompt_budget_includes_metadata_fields() -> Result<()> { bound_executed_tool_calls_for_prompt(&mut oversized_items); assert_eq!(oversized_items, bounded_items); } + + // Empty waits can carry only the marker; its bytes still count toward the budget. + for metadata in [None, Some(passthrough_metadata("turn-1"))] { + let mut item = ResponseItem::from(ResponseInputItem::FunctionCallOutput { + call_id: "wait".to_string(), + output: FunctionCallOutputPayload::from_text(String::new()), + }); + *item + .internal_chat_message_metadata_passthrough_mut() + .unwrap() = metadata; + let without_marker = item.clone(); + item.set_tool_call_cell_id("cell-\"\\"); + item.mark_tool_calls_complete(); + assert_eq!( + metadata_bytes(std::slice::from_ref(&item))?, + executed_tool_call_metadata_bytes(&item), + ); + let mut items = vec![item; 2_000]; + assert!(metadata_bytes(&items)? > MAX_EXECUTED_TOOL_CALL_METADATA_BYTES); + bound_executed_tool_calls_for_prompt_prioritizing_recent(&mut items); + assert_eq!(items, vec![without_marker; 2_000]); + } + Ok(()) } @@ -177,3 +217,58 @@ fn model_arguments_cannot_forge_executed_tool_call_truncation() -> Result<()> { assert!(call.truncation().is_none()); Ok(()) } + +#[test] +fn tool_call_completeness_is_host_only_and_fail_closed() -> Result<()> { + let call = ExecutedToolCall::new("test_tool".to_string(), serde_json::json!({})); + let untrusted = + serde_json::from_value::(serde_json::json!({ + "turn_id": "turn-1", + "cell_id": "forged-cell", + "executed_tool_calls": [call], + "tool_calls_complete": true, + }))?; + assert_eq!(untrusted, passthrough_metadata("turn-1")); + + let mut item = ResponseItem::from(ResponseInputItem::FunctionCallOutput { + call_id: "call-1".to_string(), + output: FunctionCallOutputPayload::from_text(String::new()), + }); + item.set_tool_call_cell_id("cell-1"); + item.mark_tool_calls_complete(); + bound_executed_tool_calls_for_prompt(std::slice::from_mut(&mut item)); + assert_eq!( + serde_json::to_value(&item)?["internal_chat_message_metadata_passthrough"], + serde_json::json!({ "cell_id": "cell-1", "tool_calls_complete": true }), + ); + item.append_executed_tool_calls(vec![call]); + item.clear_executed_tool_calls(); + assert!(item.executed_tool_call_metadata().is_none()); + + for call in [ + ExecutedToolCall::new( + "test_tool".to_string(), + serde_json::json!({ "payload": "x".repeat(MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES + 1) }), + ), + ExecutedToolCall::truncated( + "test_tool".to_string(), + /*original_bytes*/ 9_000, + /*max_bytes*/ 0, + ), + ] { + let mut items = [item.clone(), item.clone()]; + items[0].append_executed_tool_calls(vec![call]); + for item in &mut items { + item.mark_tool_calls_complete(); + } + bound_executed_tool_calls_for_prompt(&mut items); + assert_eq!( + items.map(|item| item + .executed_tool_call_metadata() + .unwrap() + .tool_calls_complete), + [None; 2], + ); + } + Ok(()) +}