diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index d96ad4eca1..d997aa1527 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -1932,6 +1932,11 @@ impl Session { } pub(crate) async fn refresh_runtime_config(&self, next_config: Config) { + self.refresh_runtime_config_inner(next_config, /*refresh_recording*/ true) + .await; + } + + async fn refresh_runtime_config_inner(&self, next_config: Config, refresh_recording: bool) { // Refresh only the user layer from the incoming snapshot. Preserve thread-local // layers such as request/session overrides that were present when this session // was created. @@ -1950,6 +1955,13 @@ impl Session { config.mcp_servers = next_config.mcp_servers.clone(); config.mcp_optional_startup_grace = next_config.mcp_optional_startup_grace; config.mcp_oauth_credentials_store_mode = next_config.mcp_oauth_credentials_store_mode; + // Recording can follow rollout changes without changing the session's + // execution features (including Code Mode's dispatch gate). + if refresh_recording { + self.services + .executed_tool_calls + .refresh(&next_config.features); + } if let Err(err) = config.features.set_enabled( Feature::Mcp20260728, next_config.features.enabled(Feature::Mcp20260728), @@ -2144,7 +2156,10 @@ impl Session { }; self.services.skills_service.clear_cache(); self.services.plugins_manager.clear_cache(); - self.refresh_runtime_config(next_config).await; + // This legacy snapshot still has the original execution features, not + // the host's latest rollout settings. Leave the live recorder alone. + self.refresh_runtime_config_inner(next_config, /*refresh_recording*/ false) + .await; } /// Record a terminal CodexErr before the app-server completion notification is reduced. diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index a4487d7516..fa3835d601 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1478,7 +1478,10 @@ impl Session { }).await; } - let executed_tool_calls = crate::state::ExecutedToolCalls::new(&config.features); + let executed_tool_calls = crate::state::ExecutedToolCalls::new( + &config.features, + &initial_history, + ); 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 cfaec60b17..be539ca4f4 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -6592,7 +6592,8 @@ 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 = crate::state::ExecutedToolCalls::new(&config.features); + let executed_tool_calls = + crate::state::ExecutedToolCalls::new(&config.features, &InitialHistory::Forked(Vec::new())); let (hooks, async_hook_results) = Hooks::new( HooksConfig { legacy_notify_argv: config.notify.clone(), @@ -8764,7 +8765,8 @@ 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 = crate::state::ExecutedToolCalls::new(&config.features); + let executed_tool_calls = + crate::state::ExecutedToolCalls::new(&config.features, &InitialHistory::Forked(Vec::new())); let (hooks, async_hook_results) = Hooks::new( HooksConfig { legacy_notify_argv: config.notify.clone(), diff --git a/codex-rs/core/src/tools/executed_tool_calls.rs b/codex-rs/core/src/tools/executed_tool_calls.rs index 959699e093..bbbefeb741 100644 --- a/codex-rs/core/src/tools/executed_tool_calls.rs +++ b/codex-rs/core/src/tools/executed_tool_calls.rs @@ -7,10 +7,12 @@ use std::collections::HashMap; use std::collections::HashSet; use std::sync::Arc; use std::sync::Mutex; +use std::sync::MutexGuard; use codex_code_mode::CellId; use codex_features::Feature; use codex_features::Features; +use codex_history::InitialHistory; use codex_protocol::models::ExecutedToolCall; use codex_protocol::models::ExecutedToolCallArguments; use codex_protocol::models::ResponseItem; @@ -28,6 +30,10 @@ use crate::tools::context::ToolPayload; use crate::tools::router::ToolCall; use crate::utils::json::serialized_json_bytes; +mod seen_ids; + +use seen_ids::SeenIds; + 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; @@ -36,10 +42,10 @@ 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. +/// sessions allocate no recorder state; missing evidence remains incomplete. #[derive(Clone, Default)] pub(crate) struct ExecutedToolCalls { - state: Option>>, + state: Arc>>, } #[derive(Default)] @@ -49,6 +55,9 @@ struct ExecutedToolCallRecorderState { output_cells: HashMap, retained_calls: HashMap<(std::mem::Discriminant, String), RetainedToolCalls>, pending_nested_calls: usize, + seen_ids: SeenIds, + can_prove_wait_completion: bool, + pending_wrapper_origins: HashSet, } /// Keep each output's calls and completion marker together through replay and pruning. @@ -82,6 +91,41 @@ struct RecordedCell { } impl ExecutedToolCallRecorderState { + fn invalidate_origin(&mut self, origin: &str) { + self.pending_wrapper_origins.remove(origin); + for cell in self.cells.values_mut() { + if cell.originating_call_id.as_deref() == Some(origin) { + cell.completion = CellCompletion::Incomplete; + } + } + for ((_, output_id), retained) in &mut self.retained_calls { + if output_id == origin || retained.cell_id.as_deref() == Some(origin) { + retained.complete = false; + } + } + } + + // A successful callback consumes the freshness observed at wrapper submission. + fn observe_cell_origin(&mut self, origin: &str) -> bool { + let fresh = + self.pending_wrapper_origins.remove(origin) || self.seen_ids.observe_call_id(origin); + if !fresh { + self.invalidate_origin(origin); + } + fresh + } + + fn invalidate_cell(&mut self, cell_id: &CellId) { + if let Some(cell) = self.cells.get_mut(cell_id) { + cell.completion = CellCompletion::Incomplete; + } + for retained in self.retained_calls.values_mut() { + if retained.runtime_cell_id.as_ref() == Some(cell_id) { + retained.complete = false; + } + } + } + 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) { @@ -117,10 +161,38 @@ impl ExecutedToolCallRecorderState { } impl ExecutedToolCalls { - pub(crate) fn new(features: &Features) -> Self { + pub(crate) fn new(features: &Features, history: &InitialHistory) -> Self { Self { - state: Self::is_enabled(features) - .then(|| Arc::new(Mutex::new(ExecutedToolCallRecorderState::default()))), + state: Arc::new(Mutex::new(Self::is_enabled(features).then(|| { + ExecutedToolCallRecorderState { + seen_ids: SeenIds::from_history(history), + // Runtime cell IDs can restart after resume/fork; inherited wait + // handles cannot be distinguished from newly allocated handles. + can_prove_wait_completion: matches!( + history, + InitialHistory::New | InitialHistory::Cleared + ), + ..Default::default() + } + }))), + } + } + + fn lock_state(&self) -> MutexGuard<'_, Option> { + self.state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + } + + /// Called under the session config lock; this never changes execution features. + pub(crate) fn refresh(&self, features: &Features) { + let enabled = Self::is_enabled(features); + let mut state = self.lock_state(); + if enabled != state.is_some() { + *state = enabled.then(|| ExecutedToolCallRecorderState { + seen_ids: SeenIds::unobserved_history(), + ..Default::default() + }); } } @@ -135,9 +207,7 @@ impl ExecutedToolCalls { source: &ToolCallSource, 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()); - } + self.record_call(call, source, step_context.tool_router.tool_mode()); } pub(crate) fn record_accepted_result( @@ -146,17 +216,17 @@ impl ExecutedToolCalls { call_id: &str, result: &dyn ToolOutput, ) { - if self.state.is_some() - && let Some(metadata) = result.tool_result_metadata() - { + // Release the lock before calling the output's trait method. + let recording = self.lock_state().is_some(); + if recording && let Some(metadata) = result.tool_result_metadata() { self.record_tool_result_metadata(source, call_id, metadata); } } fn record_call(&self, call: &ToolCall, source: &ToolCallSource, tool_mode: ToolMode) { - let Some(state) = &self.state else { + if self.lock_state().is_none() { return; - }; + } if matches!(source, ToolCallSource::Direct) && matches!(tool_mode, ToolMode::CodeMode | ToolMode::CodeModeOnly) && call.tool_name.is_default_namespace() @@ -171,6 +241,17 @@ impl ExecutedToolCalls { ) ) { + let mut state = self.lock_state(); + let Some(state) = state.as_mut() else { + return; + }; + let fresh = state.seen_ids.observe_call_id(&call.call_id); + if !fresh { + state.invalidate_origin(&call.call_id); + } + if fresh && state.pending_wrapper_origins.len() < MAX_PENDING_EXECUTED_TOOL_CALLS { + state.pending_wrapper_origins.insert(call.call_id.clone()); + } return; } @@ -197,9 +278,13 @@ impl ExecutedToolCalls { }; match source { ToolCallSource::Direct | ToolCallSource::DirectPlaintextMessage => { - let mut state = state - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); + let mut state = self.lock_state(); + let Some(state) = state.as_mut() else { + return; + }; + if !state.seen_ids.observe_call_id(&call.call_id) { + state.invalidate_origin(&call.call_id); + } if state.direct_calls.len() < MAX_PENDING_EXECUTED_TOOL_CALLS { state .direct_calls @@ -236,12 +321,13 @@ impl ExecutedToolCalls { call: ExecutedToolCall, original_bytes: usize, ) { - let Some(state) = &self.state else { + let mut state = self.lock_state(); + let Some(state) = state.as_mut() else { return; }; - let mut state = state - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); + if !state.cells.contains_key(&cell_id) { + state.invalidate_cell(&cell_id); + } if state.pending_nested_calls > MAX_PENDING_EXECUTED_TOOL_CALLS || (state.cells.len() >= MAX_PENDING_EXECUTED_TOOL_CALLS && !state.cells.contains_key(&cell_id)) @@ -253,6 +339,7 @@ impl ExecutedToolCalls { } let at_pending_call_limit = state.pending_nested_calls == MAX_PENDING_EXECUTED_TOOL_CALLS; let cell = state.cells.entry(cell_id).or_default(); + let duplicate_call_id = cell.pending_calls.contains_key(&call_id); let max_bytes = MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES.min( MAX_EXECUTED_TOOL_CALL_FULL_ARGUMENT_BYTES_PER_OUTPUT .saturating_sub(cell.pending_full_argument_bytes), @@ -270,10 +357,11 @@ impl ExecutedToolCalls { cell.completion = if matches!( cell.completion, CellCompletion::Started | CellCompletion::Recording - ) && !matches!( - call.arguments(), - ExecutedToolCallArguments::Truncated { .. } - ) { + ) && !duplicate_call_id + && !matches!( + call.arguments(), + ExecutedToolCallArguments::Truncated { .. } + ) { CellCompletion::Recording } else { CellCompletion::Incomplete @@ -288,14 +376,12 @@ impl ExecutedToolCalls { call_id: &str, metadata: &JsonValue, ) -> bool { - let Some(state) = &self.state else { - return false; - }; let metadata = codex_protocol::models::ToolResultMetadata::new(metadata); let has_metadata = metadata.is_some(); - let mut state = state - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); + let mut state = self.lock_state(); + let Some(state) = state.as_mut() else { + return false; + }; let call = match source { ToolCallSource::Direct | ToolCallSource::DirectPlaintextMessage => { state.direct_calls.get_mut(call_id) @@ -328,35 +414,43 @@ impl ExecutedToolCalls { } pub(crate) fn register_cell(&self, cell_id: &CellId, output_call_id: &str) { - let Some(state) = &self.state else { + let mut state = self.lock_state(); + let Some(state) = state.as_mut() else { return; }; - let mut state = state - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); + if !state.observe_cell_origin(output_call_id) { + state.invalidate_cell(cell_id); + } state.register_cell(cell_id, output_call_id); } pub(crate) fn start_cell(&self, cell_id: &CellId, output_call_id: &str) { - let Some(state) = &self.state else { + let mut state = self.lock_state(); + let Some(state) = state.as_mut() else { return; }; - let mut state = state - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); + let unique_cell = state.seen_ids.observe_runtime_cell_id(cell_id); + let unique_origin = state.observe_cell_origin(output_call_id); + if !unique_cell { + state.invalidate_cell(cell_id); + } + let history_ids_indexed = state.seen_ids.history_ids_indexed(); state.register_cell(cell_id, output_call_id); if let Some(cell) = state.cells.get_mut(cell_id) { - cell.completion = CellCompletion::Started; + // Failed indexing must not make a known historical ID look fresh. + cell.completion = if unique_cell && unique_origin && history_ids_indexed { + CellCompletion::Started + } else { + CellCompletion::Incomplete + }; } } pub(crate) fn finish_cell_recording(&self, cell_id: &CellId) { - let Some(state) = &self.state else { + let mut state = self.lock_state(); + let Some(state) = state.as_mut() 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) { if cell.completion == CellCompletion::Recording { cell.completion = CellCompletion::Complete; @@ -366,3 +460,22 @@ impl ExecutedToolCalls { } } } + +fn input_call_id(item: &ResponseItem) -> Option<&str> { + match item { + ResponseItem::FunctionCall { call_id, .. } + | ResponseItem::CustomToolCall { call_id, .. } => Some(call_id), + ResponseItem::ToolSearchCall { call_id, .. } + | ResponseItem::LocalShellCall { call_id, .. } => call_id.as_deref(), + _ => None, + } +} + +fn output_call_id(item: &ResponseItem) -> Option<&str> { + match item { + ResponseItem::FunctionCallOutput { call_id, .. } + | ResponseItem::ToolSearchOutput { call_id, .. } => call_id.as_deref(), + ResponseItem::CustomToolCallOutput { call_id, .. } => Some(call_id), + _ => None, + } +} 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 index 5db02d71e7..7ebabace18 100644 --- a/codex-rs/core/src/tools/executed_tool_calls/request_metadata.rs +++ b/codex-rs/core/src/tools/executed_tool_calls/request_metadata.rs @@ -4,7 +4,7 @@ use super::*; impl ExecutedToolCalls { - /// Attaches host-recorded calls within the existing request metadata budget. + /// Attaches trusted completeness; request budgeting only revokes damaged inventories. pub(crate) fn attach_to_prompt( &self, items: &mut [ResponseItem], @@ -20,12 +20,12 @@ impl ExecutedToolCalls { items: &mut [ResponseItem], retry_cache: &mut ExecutedToolCallCache, ) -> bool { - let Some(state) = &self.state else { + let mut state = self.lock_state(); + let Some(state) = state.as_mut() else { return false; }; - let mut state = state - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); + // Failed wrappers have no callback; only their bounded bitmap observation survives. + state.pending_wrapper_origins.clear(); if state.direct_calls.is_empty() && state.output_cells.is_empty() && state.retained_calls.is_empty() @@ -34,6 +34,20 @@ impl ExecutedToolCalls { return false; } + let mut output_counts = HashMap::new(); + let mut input_indices = HashMap::new(); + for (index, item) in items.iter().enumerate() { + if let Some(call_id) = output_call_id(item) { + *output_counts.entry(call_id.to_string()).or_insert(0_usize) += 1; + } + if let Some(call_id) = input_call_id(item) { + input_indices + .entry(call_id.to_string()) + .and_modify(|entry: &mut Option| *entry = None) + .or_insert(Some(index)); + } + } + // Updated records supersede older retry snapshots. retry_cache.retain(|key, _| { !state @@ -41,60 +55,69 @@ impl ExecutedToolCalls { .get(key) .is_some_and(|retained| retained.result_metadata_updated) }); - let mut pending_retry_outputs = retry_cache.keys().cloned().collect::>(); - let mut pending_retained_outputs = - state.retained_calls.keys().cloned().collect::>(); + let mut pending_outputs = retry_cache + .keys() + .chain(state.retained_calls.keys()) + .cloned() + .collect::>(); let mut attached = false; - let mut retained_bytes = 0_usize; - for item in items.iter_mut().rev() { + for index in (0..items.len()).rev() { if state.direct_calls.is_empty() && state.output_cells.is_empty() - && pending_retry_outputs.is_empty() - && pending_retained_outputs.is_empty() + && pending_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 item = &items[index]; + let Some(call_id) = output_call_id(item) else { + continue; }; - let key = (std::mem::discriminant(&*item), call_id.clone()); + let key = (std::mem::discriminant(item), call_id.to_string()); 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) { + let calls = if let Some(cached) = + retry_cache.get(&key).or_else(|| retained.map(|r| &r.calls)) + { + if !pending_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 runtime_cell_id = None; let mut calls = state .direct_calls .remove(call_id) .into_iter() .collect::>(); + let direct_call_collision = !calls.is_empty(); 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()); + // Validate the first delta too: a later wait cannot repair a + // missing or ambiguous exec/output association. + let matches_input = + input_indices + .get(call_id) + .copied() + .flatten() + .is_some_and(|input_index| { + input_index < index + && cell_id.as_deref().is_some_and(|origin| { + code_mode_input_matches_output( + &items[input_index], + item, + origin, + &output_cell_id, + ) + }) + }); + if output_counts.get(call_id) != Some(&1) || !matches_input { + cell.completion = CellCompletion::Incomplete; + } let pending_calls = cell.pending_calls.len(); for (call_id, call) in cell.pending_calls.drain(..) { if matches!(call.arguments(), ExecutedToolCallArguments::Raw(_)) { @@ -103,7 +126,10 @@ impl ExecutedToolCalls { calls.push(call); } cell.pending_full_argument_bytes = 0; - complete = cell.completion == CellCompletion::Complete; + complete = cell.completion == CellCompletion::Complete + && !direct_call_collision + && (state.can_prove_wait_completion + || matches!(item, ResponseItem::CustomToolCallOutput { .. })); if matches!( cell.completion, CellCompletion::Complete | CellCompletion::Incomplete @@ -133,49 +159,117 @@ impl ExecutedToolCalls { ); calls }; + let item = &mut items[index]; 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(); + } else { + item.clear_tool_calls_complete(); } - retained_bytes = retained_bytes.saturating_add(executed_tool_call_metadata_bytes(item)); attached = true; } - if !pending_retained_outputs.is_empty() { + if !pending_outputs.is_empty() { state .retained_calls - .retain(|key, _| !pending_retained_outputs.contains(key)); + .retain(|key, _| !pending_outputs.contains(key)); } - if retained_bytes > MAX_EXECUTED_TOOL_CALL_FULL_ARGUMENT_BYTES_PER_OUTPUT { + let mut invalid_cells = HashSet::new(); + for index in 0..items.len() { + let item = &items[index]; + let Some(call_id) = output_call_id(item) else { + continue; + }; + let key = (std::mem::discriminant(item), call_id.to_string()); + let Some(retained) = state.retained_calls.get_mut(&key) else { + continue; + }; + let matches_inventory = item.executed_tool_call_metadata().is_some_and(|metadata| { + metadata.cell_id == retained.cell_id + && (metadata.has_same_tool_calls(&retained.calls) + || (retained.calls.is_empty() && metadata.executed_tool_calls.is_none())) + }); + // A verified output may outlive its input after compaction. Any input + // still present must continue to identify the same exec or wait. + let matches_input = input_indices.get(call_id).is_none_or(|input_index| { + input_index.is_some_and(|input_index| { + input_index < index + && match (&retained.cell_id, &retained.runtime_cell_id) { + (Some(origin), Some(runtime_cell)) => code_mode_input_matches_output( + &items[input_index], + item, + origin, + runtime_cell, + ), + _ => true, + } + }) + }); + // Lost or ambiguous evidence cannot become complete again on a retry. + let intact = + output_counts.get(call_id) == Some(&1) && matches_inventory && matches_input; + retained.complete &= intact; + if !intact && let Some(cell) = retained.runtime_cell_id.clone() { + invalid_cells.insert(cell); + } + } + for cell_id in &invalid_cells { + if let Some(cell) = state.cells.get_mut(cell_id) { + cell.completion = CellCompletion::Incomplete; + } + } + if !invalid_cells.is_empty() { + for retained in state.retained_calls.values_mut() { + if retained + .runtime_cell_id + .as_ref() + .is_some_and(|cell_id| invalid_cells.contains(cell_id)) + { + retained.complete = false; + } + } + } + for item in items.iter_mut() { + let Some(call_id) = output_call_id(item) else { + continue; + }; + let key = (std::mem::discriminant(&*item), call_id.to_string()); + let Some(retained) = state.retained_calls.get(&key) else { + continue; + }; + if retained.complete { + item.mark_tool_calls_complete(); + } else { + item.clear_tool_calls_complete(); + } + } + + let metadata_bytes = items.iter().fold(0_usize, |bytes, item| { + bytes.saturating_add(executed_tool_call_metadata_bytes(item)) + }); + if metadata_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, + for item in items.iter_mut() { + let Some(call_id) = output_call_id(item) else { + continue; + }; + let key = (std::mem::discriminant(&*item), call_id.to_string()); + // History already carries unmanaged records; retaining them would append them again. + let Some(previous) = retained_before_bounding.get(&key) else { + 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 + if let Some(runtime_cell_id) = &previous.runtime_cell_id && metadata - .is_none_or(|metadata| !metadata.has_same_tool_calls(&retained.calls)) + .is_none_or(|metadata| !metadata.has_same_tool_calls(&previous.calls)) && let Some(cell) = state.cells.get_mut(runtime_cell_id) { cell.completion = CellCompletion::Incomplete; @@ -187,25 +281,22 @@ impl ExecutedToolCalls { .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()); // Metadata-only shedding preserves slots; changed calls or duplicate outputs do not. - let call_index_by_id = previous - .filter(|retained| { - unique_output && metadata.has_same_tool_calls(&retained.calls) - }) - .map(|retained| retained.call_index_by_id.clone()) - .unwrap_or_default(); + let call_index_by_id = + if unique_output && metadata.has_same_tool_calls(&previous.calls) { + previous.call_index_by_id.clone() + } else { + HashMap::new() + }; let retained = state.retained_calls.entry(key).or_default(); - retained.runtime_cell_id = runtime_cell_id; + retained.runtime_cell_id = previous.runtime_cell_id.clone(); 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.result_metadata_updated = - previous.is_some_and(|retained| retained.result_metadata_updated); + retained.result_metadata_updated = previous.result_metadata_updated; } } } @@ -214,6 +305,50 @@ impl ExecutedToolCalls { } } +fn code_mode_input_matches_output( + input: &ResponseItem, + output: &ResponseItem, + origin: &str, + runtime_cell: &CellId, +) -> bool { + match (input, output) { + ( + ResponseItem::CustomToolCall { + name, + namespace, + call_id, + .. + }, + ResponseItem::CustomToolCallOutput { .. }, + ) => { + call_id == origin + && crate::tools::code_mode::is_exec_tool_name(&codex_tools::ToolName::new( + namespace.clone(), + name, + )) + } + ( + ResponseItem::FunctionCall { + name, + namespace, + call_id, + arguments, + .. + }, + ResponseItem::FunctionCallOutput { .. }, + ) => { + call_id != origin + && name == crate::tools::code_mode::WAIT_TOOL_NAME + && codex_tools::ToolName::new(namespace.clone(), name).is_default_namespace() + && serde_json::from_str::(arguments).is_ok_and(|arguments| { + arguments.get("cell_id").and_then(JsonValue::as_str) + == Some(runtime_cell.as_str()) + }) + } + _ => false, + } +} + #[cfg(test)] #[path = "request_metadata_tests.rs"] mod tests; diff --git a/codex-rs/core/src/tools/executed_tool_calls/request_metadata_tests.rs b/codex-rs/core/src/tools/executed_tool_calls/request_metadata_tests.rs index e13a0679d6..2513022719 100644 --- a/codex-rs/core/src/tools/executed_tool_calls/request_metadata_tests.rs +++ b/codex-rs/core/src/tools/executed_tool_calls/request_metadata_tests.rs @@ -1,15 +1,19 @@ +use codex_history::RolloutItem; use codex_protocol::models::FunctionCallOutputPayload; use codex_protocol::models::ResponseInputItem; use codex_protocol::models::ToolResultMetadata; +use codex_protocol::protocol::SessionMeta; +use codex_protocol::protocol::SessionMetaLine; +use codex_protocol::protocol::ThreadHistoryMode; use pretty_assertions::assert_eq; use serde_json::json; use super::*; -fn new_recorder() -> ExecutedToolCalls { +fn new_recorder(history: InitialHistory) -> ExecutedToolCalls { let mut features = Features::default(); features.enable(Feature::ExecutedToolCallMetadata); - ExecutedToolCalls::new(&features) + ExecutedToolCalls::new(&features, &history) } fn output(call_id: &str) -> ResponseItem { @@ -19,11 +23,54 @@ fn output(call_id: &str) -> ResponseItem { }) } -#[tokio::test] -async fn recorder_preserves_session_and_turn_gates_and_lazy_result_metadata() { - struct MetadataLookup(std::cell::Cell, JsonValue); +fn exec_output(call_id: &str) -> ResponseItem { + ResponseItem::from(ResponseInputItem::CustomToolCallOutput { + call_id: call_id.to_string(), + name: None, + output: FunctionCallOutputPayload::from_text(String::new()), + }) +} - impl ToolOutput for MetadataLookup { +fn exec_input(call_id: &str) -> ResponseItem { + serde_json::from_value(json!({ + "type": "custom_tool_call", "call_id": call_id, "name": "exec", "input": "", + })) + .expect("exec input must deserialize") +} + +fn wait_input(call_id: &str, cell: &CellId) -> ResponseItem { + serde_json::from_value(json!({ + "type": "function_call", "call_id": call_id, "name": "wait", + "arguments": json!({"cell_id": cell.as_str()}).to_string(), + })) + .expect("wait input must deserialize") +} + +fn record_nested_call( + recorder: &ExecutedToolCalls, + cell: &CellId, + call_id: &str, +) -> ExecutedToolCall { + let call = ExecutedToolCall::new("nested_tool".to_string(), json!({})); + recorder.record_nested_tool_call( + cell.clone(), + call_id.to_string(), + call.clone(), + /*original_bytes*/ 2, + ); + call +} + +fn tool_calls_complete(item: &ResponseItem) -> Option { + item.executed_tool_call_metadata() + .and_then(|metadata| metadata.tool_calls_complete) +} + +#[tokio::test] +async fn recorder_refreshes_without_changing_execution_features_or_claiming_missing_history() { + struct MetadataLookup<'a>(std::cell::Cell, JsonValue, &'a ExecutedToolCalls); + + impl ToolOutput for MetadataLookup<'_> { fn log_output(&self) -> String { panic!("recording must not read the diagnostic output") } @@ -37,82 +84,107 @@ async fn recorder_preserves_session_and_turn_gates_and_lazy_result_metadata() { } fn tool_result_metadata(&self) -> Option<&JsonValue> { + assert!( + self.2.state.try_lock().is_ok(), + "result callbacks must run unlocked" + ); self.0.set(self.0.get() + 1); Some(&self.1) } } - 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, - ) + let (session, turn) = crate::session::tests::make_session_and_context().await; + let step = StepContext::for_test(Arc::new(turn)); + // The broker holds a clone before rollout enablement reaches the session. + let calls = session.services.executed_tool_calls.clone(); + assert!(calls.lock_state().is_none()); + let mut next_config = (*session.get_config().await).clone(); + let mut retry_cache = HashMap::new(); + for (index, enabled) in [false, true, false, true].into_iter().enumerate() { + next_config .features - .set_enabled(Feature::ExecutedToolCallMetadata, turn_enabled) + .set_enabled(Feature::ExecutedToolCallMetadata, 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 - ); + session.refresh_runtime_config(next_config.clone()).await; + assert_eq!(calls.lock_state().is_some(), enabled); + // A legacy file reload has no new rollout snapshot and must keep this setting. + session.reload_user_config_layer().await; + assert_eq!(calls.lock_state().is_some(), enabled); + // An already-running turn and future turns keep their execution feature policy. + assert!(!ExecutedToolCalls::is_enabled(&step.turn.config.features)); + assert!(!ExecutedToolCalls::is_enabled( + &session.get_config().await.features + )); - 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 = MetadataLookup(std::cell::Cell::new(0), json!({})); - calls.record_accepted_result(&ToolCallSource::Direct, &call.call_id, &result); - // Accepted metadata follows session availability even when this turn stops recording calls. - assert_eq!(result.0.get(), usize::from(session_enabled)); + let call = ToolCall { + tool_name: codex_tools::ToolName::plain("test_tool"), + call_id: format!("call-{index}"), + payload: ToolPayload::Function { + arguments: r#"{"value":7}"#.to_string(), + }, + encrypted_function_args: None, + }; + calls.record_tool_call(&call, &ToolCallSource::Direct, &step); + let result = MetadataLookup(std::cell::Cell::new(0), json!({}), &calls); + calls.record_accepted_result(&ToolCallSource::Direct, &call.call_id, &result); + assert_eq!(result.0.get(), usize::from(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})); - expected_call.set_tool_result_metadata(ToolResultMetadata::new(&json!({}))); - 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}", - ); + let input = serde_json::from_value(json!({ + "type": "function_call", + "call_id": call.call_id, + "name": "test_tool", + "arguments": "{\"value\":7}", + })) + .expect("tool input must deserialize"); + let mut prompt = vec![input, output(&call.call_id)]; + if !enabled { + prompt[1].mark_tool_calls_complete(); } + let original = prompt.clone(); + calls.attach_to_prompt(&mut prompt, &mut retry_cache); + if !enabled { + assert_eq!(prompt, original); + } else { + assert_eq!(tool_calls_complete(&prompt[1]), None); + } + let mut expected_call = ExecutedToolCall::new("test_tool".to_string(), json!({"value": 7})); + expected_call.set_tool_result_metadata(ToolResultMetadata::new(&json!({}))); + let expected_calls = vec![expected_call]; + assert_eq!( + prompt[1] + .executed_tool_call_metadata() + .and_then(|metadata| metadata.executed_tool_calls.as_ref()), + enabled.then_some(&expected_calls), + ); + + if enabled { + // Reapplying an enabled config must preserve already-recorded calls. + session.refresh_runtime_config(next_config.clone()).await; + let mut replay = original; + calls.attach_to_prompt(&mut replay, &mut HashMap::new()); + assert_eq!(replay, prompt); + } + + let cell = CellId::new(format!("cell-{index}")); + let origin = format!("exec-{index}"); + calls.start_cell(&cell, &origin); + let nested = record_nested_call(&calls, &cell, "nested"); + calls.finish_cell_recording(&cell); + let mut prompt = vec![exec_input(&origin), exec_output(&origin)]; + calls.attach_to_prompt(&mut prompt, &mut retry_cache); + assert_eq!(tool_calls_complete(&prompt[1]), None); + assert_eq!( + prompt[1] + .executed_tool_call_metadata() + .and_then(|metadata| metadata.executed_tool_calls.as_ref()), + enabled.then_some(&vec![nested]), + ); } } #[test] fn executed_tool_call_recorder_bounds_pending_calls_and_preserves_overflow() { - let recorder = new_recorder(); + let recorder = new_recorder(InitialHistory::Forked(Vec::new())); for index in 0..MAX_PENDING_EXECUTED_TOOL_CALLS + 2 { recorder.record_call( @@ -127,6 +199,18 @@ fn executed_tool_call_recorder_bounds_pending_calls_and_preserves_overflow() { &ToolCallSource::Direct, ToolMode::Direct, ); + recorder.record_call( + &ToolCall { + tool_name: codex_tools::ToolName::plain(crate::tools::code_mode::PUBLIC_TOOL_NAME), + call_id: format!("failed-wrapper-{index}"), + payload: ToolPayload::Custom { + input: "".to_string(), + }, + encrypted_function_args: None, + }, + &ToolCallSource::Direct, + ToolMode::CodeMode, + ); } let metadata = json!({ "arbitrary-key": ["R1", null] }); assert!(recorder.record_tool_result_metadata(&ToolCallSource::Direct, "direct-0", &metadata)); @@ -150,12 +234,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); + let state = recorder.lock_state(); + let state = state.as_ref().unwrap(); assert_eq!( state.direct_calls.len(), MAX_PENDING_EXECUTED_TOOL_CALLS + 1 @@ -186,29 +266,26 @@ fn executed_tool_call_recorder_bounds_pending_calls_and_preserves_overflow() { state.pending_nested_calls, MAX_PENDING_EXECUTED_TOOL_CALLS + 1 ); + assert_eq!( + state.pending_wrapper_origins.len(), + MAX_PENDING_EXECUTED_TOOL_CALLS + ); assert_eq!(state.cells.len(), MAX_PENDING_EXECUTED_TOOL_CALLS); assert_eq!(state.output_cells.len(), MAX_PENDING_EXECUTED_TOOL_CALLS); } - let mut items = [ResponseItem::FunctionCallOutput { - id: None, - call_id: Some("bounded-output".to_string()), - name: None, - namespace: None, - output: FunctionCallOutputPayload::from_text(String::new()), - internal_chat_message_metadata_passthrough: None, - }]; + let mut items = [exec_input("bounded-output"), exec_output("bounded-output")]; 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] + items[1] .executed_tool_call_metadata() .and_then(|metadata| metadata.tool_calls_complete), None, ); - let calls = items[0] + let calls = items[1] .executed_tool_call_metadata() .and_then(|metadata| metadata.executed_tool_calls.as_ref()) .expect("bounded nested calls must attach to their own output"); @@ -229,25 +306,15 @@ fn executed_tool_call_recorder_bounds_pending_calls_and_preserves_overflow() { let expected_calls = calls.clone(); { - let state = recorder - .state - .as_ref() - .unwrap() - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); + let state = recorder.lock_state(); + let state = state.as_ref().unwrap(); + assert!(state.pending_wrapper_origins.is_empty()); assert_eq!(state.pending_nested_calls, 0); assert!(!state.cells.contains_key(&cell_id)); assert_eq!(retry_cache.len(), 1); } - let mut replayed_items = [ResponseItem::FunctionCallOutput { - id: None, - call_id: Some("bounded-output".to_string()), - name: None, - namespace: None, - output: FunctionCallOutputPayload::from_text(String::new()), - internal_chat_message_metadata_passthrough: None, - }]; + let mut replayed_items = [exec_output("bounded-output")]; let mut replay_retry_cache = HashMap::new(); assert!(recorder.attach_pending_to_prompt(&mut replayed_items, &mut replay_retry_cache)); assert_eq!( @@ -259,18 +326,14 @@ fn executed_tool_call_recorder_bounds_pending_calls_and_preserves_overflow() { let mut compacted_retry_cache = HashMap::new(); 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); + let state = recorder.lock_state(); + let state = state.as_ref().unwrap(); assert!(state.retained_calls.is_empty()); } #[test] -fn executed_tool_call_recorder_bounds_retained_history_and_reports_omissions() { - let recorder = new_recorder(); +fn executed_tool_call_recorder_bounds_retained_history_and_keeps_latest_calls() { + let recorder = new_recorder(InitialHistory::Forked(Vec::new())); let mut history = Vec::new(); let arguments = serde_json::to_string(&json!({ "payload": "x".repeat(1024) })) .expect("tool arguments must serialize"); @@ -317,12 +380,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 state = recorder.lock_state(); + let state = state.as_ref().unwrap(); let retained_bytes = state .retained_calls .values() @@ -338,20 +397,17 @@ fn executed_tool_call_recorder_bounds_retained_history_and_reports_omissions() { .flatten() .map(|call| serde_json::to_value(call).expect("retained call must serialize")) .collect::>(); - let omitted_calls = metadata - .iter() - .filter_map(|call| { - call["arguments"]["_codex_executed_tool_call_truncated"]["omitted_calls"].as_u64() - }) - .sum::(); - assert!(omitted_calls > 0); - assert_eq!(metadata.len() as u64 + omitted_calls, 512); + assert!(metadata.len() < 512); + // Dropped outputs do not add their omission counts to unrelated surviving calls. + assert!(metadata.iter().all(|call| { + call["arguments"]["_codex_executed_tool_call_truncated"]["omitted_calls"].is_null() + })); } #[test] fn tool_call_completeness_requires_finished_lossless_recording() { for scenario in ["empty", "unobserved", "late_call"] { - let recorder = new_recorder(); + let recorder = new_recorder(InitialHistory::Forked(Vec::new())); let cell_id = CellId::new(scenario.to_string()); if scenario == "unobserved" { recorder.register_cell(&cell_id, "output"); @@ -371,10 +427,10 @@ fn tool_call_completeness_requires_finished_lossless_recording() { } recorder.finish_cell_recording(&cell_id); - let mut items = [output("output")]; + let mut items = [exec_input("output"), exec_output("output")]; recorder.attach_pending_to_prompt(&mut items, &mut HashMap::new()); assert_eq!( - items[0] + items[1] .executed_tool_call_metadata() .and_then(|metadata| metadata.tool_calls_complete), None, @@ -387,7 +443,7 @@ fn tool_call_completeness_requires_finished_lossless_recording() { fn tool_call_completeness_survives_waits_without_changing_deltas() { let metadata = json!({ "provider": { "ids": ["CLATE"] } }); for truncated in [false, true] { - let recorder = new_recorder(); + let recorder = new_recorder(InitialHistory::New); let cell_id = CellId::new("multi-wait".to_string()); recorder.start_cell(&cell_id, "exec"); let mut history = Vec::new(); @@ -397,9 +453,36 @@ fn tool_call_completeness_survives_waits_without_changing_deltas() { .into_iter() .enumerate() { - recorder.register_cell(&cell_id, call_id); - let mut expected_output = output(call_id); + if index > 0 { + recorder.record_call( + &ToolCall { + tool_name: codex_tools::ToolName::plain( + crate::tools::code_mode::WAIT_TOOL_NAME, + ), + call_id: call_id.to_string(), + payload: ToolPayload::Function { + arguments: json!({"cell_id": cell_id.as_str()}).to_string(), + }, + encrypted_function_args: None, + }, + &ToolCallSource::Direct, + ToolMode::CodeModeOnly, + ); + recorder.register_cell(&cell_id, call_id); + } + let input = if index == 0 { + exec_input(call_id) + } else { + wait_input(call_id, &cell_id) + }; + let mut expected_output = if index == 0 { + exec_output(call_id) + } else { + output(call_id) + }; + history.push(input.clone()); history.push(expected_output.clone()); + expected.push(input); if index < 2 { // Identical calls are distinct submissions; truncation stays sticky across waits. let original_bytes = if truncated && index == 1 { 9_000 } else { 2 }; @@ -434,10 +517,10 @@ fn tool_call_completeness_survives_waits_without_changing_deltas() { let mut delayed_call = ExecutedToolCall::new("nested_tool".to_string(), json!({})); delayed_call.set_tool_result_metadata(ToolResultMetadata::new(&metadata)); - let mut delayed_output = output("exec"); + let mut delayed_output = exec_output("exec"); delayed_output.append_executed_tool_calls(vec![delayed_call]); delayed_output.set_tool_call_cell_id("exec"); - expected[0] = delayed_output; + expected[1] = delayed_output; assert!(recorder.record_tool_result_metadata(&source, "nested-1", &json!({}))); call.set_tool_result_metadata(ToolResultMetadata::new(&json!({}))); } @@ -458,50 +541,261 @@ fn tool_call_completeness_survives_waits_without_changing_deltas() { assert_eq!(prompt, expected); } } - let state = recorder.state.as_ref().unwrap().lock().unwrap(); + let independent = CellId::new("independent-runtime".to_string()); + recorder.start_cell(&independent, "independent-exec"); + record_nested_call(&recorder, &independent, "independent-call"); + recorder.finish_cell_recording(&independent); + history.extend([ + exec_input("independent-exec"), + exec_output("independent-exec"), + ]); + // Damaged deltas revoke their terminal sibling, not an independent cell. + history[2] = wait_input("wait-1", &independent); + history[0] = wait_input("exec", &independent); + recorder.attach_pending_to_prompt(&mut history, &mut retry_cache); + assert_eq!(tool_calls_complete(&history[7]), None); + assert_eq!(tool_calls_complete(&history[9]), Some(true)); + + let state = recorder.lock_state(); + let state = state.as_ref().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()); +fn unrelated_raw_overflow_retains_a_complete_exec_within_the_request_budget() { + let request_budget = 32 * 1024; + let recorder = new_recorder(InitialHistory::New); + let cell = CellId::new("clean-overflow-cell".to_string()); + recorder.start_cell(&cell, "exec-clean"); + let clean_call = record_nested_call(&recorder, &cell, "nested-clean"); + recorder.finish_cell_recording(&cell); - for originating_call_id in ["first-exec", "resumed-exec"] { - let recorder = new_recorder(); - recorder.start_cell(&runtime_cell_id, originating_call_id); - recorder.record_nested_tool_call( - runtime_cell_id.clone(), - "nested-call".to_string(), - 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] + let mut unrelated = output("direct-oversized"); + let arguments = json!("x".repeat(7 * 1024)); + assert!(serialized_json_bytes(&arguments).unwrap() <= MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES); + unrelated.append_executed_tool_calls( + (0..6) + .map(|index| ExecutedToolCall::new(format!("direct-{index}"), arguments.clone())) + .collect(), + ); + let original = [ + unrelated, + exec_input("exec-clean"), + exec_output("exec-clean"), + ]; + let mut prompt = original.clone(); + assert!( + prompt + .iter() + .map(executed_tool_call_metadata_bytes) + .sum::() + > request_budget + ); + recorder.attach_to_prompt(&mut prompt, &mut HashMap::new()); + let metadata = prompt[2] + .executed_tool_call_metadata() + .expect("the later clean output must retain its calls"); + assert_eq!(metadata.executed_tool_calls, Some(vec![clean_call])); + assert_eq!(metadata.tool_calls_complete, Some(true)); + assert!( + prompt[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)); + .and_then(|metadata| metadata.executed_tool_calls.as_ref()) + .is_some_and(|calls| calls.iter().any(|call| matches!( + call.arguments(), + ExecutedToolCallArguments::Truncated { .. } + ))) + ); + assert!( + prompt + .iter() + .map(executed_tool_call_metadata_bytes) + .sum::() + <= request_budget + ); + + let mut retry = original; + recorder.attach_to_prompt(&mut retry, &mut HashMap::new()); + assert_eq!(retry, prompt); +} + +#[test] +fn paginated_history_allows_fresh_origins_but_rejects_known_ids() { + for origin in ["fresh-exec", "old-exec"] { + let history = InitialHistory::Forked(vec![ + RolloutItem::SessionMeta(SessionMetaLine { + meta: SessionMeta { + history_mode: ThreadHistoryMode::Paginated, + ..SessionMeta::default() + }, + git: None, + }), + RolloutItem::ResponseItem(exec_input("old-exec").into()), + RolloutItem::ResponseItem(exec_output("old-exec").into()), + ]); + let recorder = new_recorder(history); + let cell = CellId::new("1".to_string()); + recorder.start_cell(&cell, origin); + let call = record_nested_call(&recorder, &cell, "nested"); + recorder.finish_cell_recording(&cell); + + // The old call is only in inherited history, so duplicate prompt IDs cannot reject it. + let mut prompt = [exec_input(origin), exec_output(origin)]; + recorder.attach_to_prompt(&mut prompt, &mut HashMap::new()); + assert_eq!( + tool_calls_complete(&prompt[1]), + (origin == "fresh-exec").then_some(true), + ); + assert_eq!( + prompt[1] + .executed_tool_call_metadata() + .unwrap() + .executed_tool_calls, + Some(vec![call]), + ); + + // Compaction may drop the input after the output has already been authenticated. + let mut retry = [exec_output(origin)]; + recorder.attach_to_prompt(&mut retry, &mut HashMap::new()); + assert_eq!(retry[0], prompt[1]); } } +#[test] +fn forked_wait_keeps_calls_without_completion() { + let recorder = new_recorder(InitialHistory::Forked(vec![RolloutItem::ResponseItem( + exec_input("old-exec").into(), + )])); + let cell = CellId::new("1".to_string()); + recorder.start_cell(&cell, "new-exec"); + let call = record_nested_call(&recorder, &cell, "nested"); + recorder.register_cell(&cell, "wait"); + recorder.finish_cell_recording(&cell); + + let mut prompt = [wait_input("wait", &cell), output("wait")]; + recorder.attach_to_prompt(&mut prompt, &mut HashMap::new()); + let mut expected = output("wait"); + expected.append_executed_tool_calls(vec![call]); + expected.set_tool_call_cell_id("new-exec"); + assert_eq!(prompt[1], expected); + + let mut compacted = [output("wait")]; + recorder.attach_to_prompt(&mut compacted, &mut HashMap::new()); + assert_eq!(compacted[0], expected); +} + +#[test] +fn wrapper_mismatch_cannot_become_complete_on_a_later_wait() { + for scenario in [ + "missing_exec", + "wrong_output", + "wrong_wait_cell", + "changed_exec", + ] { + let recorder = new_recorder(InitialHistory::New); + let cell = CellId::new("runtime".to_string()); + recorder.start_cell(&cell, "exec"); + let call = record_nested_call(&recorder, &cell, "nested"); + let mut prompt = match scenario { + "missing_exec" => vec![exec_output("exec")], + "changed_exec" => vec![exec_input("exec"), exec_output("exec")], + "wrong_wait_cell" => { + recorder.register_cell(&cell, "first-wait"); + vec![ + wait_input("first-wait", &CellId::new("other-runtime".to_string())), + output("first-wait"), + ] + } + // An exec custom call cannot authenticate a function output with the same ID. + _ => vec![exec_input("exec"), output("exec")], + }; + recorder.attach_to_prompt(&mut prompt, &mut HashMap::new()); + assert_eq!( + prompt + .last() + .unwrap() + .executed_tool_call_metadata() + .unwrap() + .executed_tool_calls, + Some(vec![call]), + ); + if scenario == "changed_exec" { + // Invalidate already-retained evidence while its cell is still active. + prompt[0] = wait_input("exec", &cell); + recorder.attach_to_prompt(&mut prompt, &mut HashMap::new()); + } + + recorder.register_cell(&cell, "last-wait"); + recorder.finish_cell_recording(&cell); + let mut last = [wait_input("last-wait", &cell), output("last-wait")]; + recorder.attach_to_prompt(&mut last, &mut HashMap::new()); + assert_eq!(tool_calls_complete(&last[1]), None); + } +} + +#[test] +fn completeness_rejects_reused_origin() { + let recorder = new_recorder(InitialHistory::New); + let first = CellId::new("first-runtime".to_string()); + recorder.start_cell(&first, "reused-exec"); + record_nested_call(&recorder, &first, "nested-first"); + recorder.finish_cell_recording(&first); + let mut original = [exec_input("reused-exec"), exec_output("reused-exec")]; + recorder.attach_pending_to_prompt(&mut original, &mut HashMap::new()); + assert_eq!(tool_calls_complete(&original[1]), Some(true)); + recorder.attach_pending_to_prompt(&mut [], &mut HashMap::new()); + + let second = CellId::new("second-runtime".to_string()); + recorder.start_cell(&second, "reused-exec"); + record_nested_call(&recorder, &second, "nested-second"); + recorder.finish_cell_recording(&second); + let mut prompt = [exec_input("reused-exec"), exec_output("reused-exec")]; + recorder.attach_to_prompt(&mut prompt, &mut HashMap::new()); + assert_eq!(tool_calls_complete(&prompt[1]), None); +} + +#[test] +fn completeness_rejects_direct_id_collision_and_late_calls() { + let recorder = new_recorder(InitialHistory::New); + let cell = CellId::new("collision-runtime".to_string()); + recorder.start_cell(&cell, "exec-collision"); + recorder.record_call( + &ToolCall { + tool_name: codex_tools::ToolName::plain("direct_tool"), + call_id: "exec-collision".to_string(), + payload: ToolPayload::Function { + arguments: "{}".to_string(), + }, + encrypted_function_args: None, + }, + &ToolCallSource::Direct, + ToolMode::Direct, + ); + record_nested_call(&recorder, &cell, "nested"); + recorder.finish_cell_recording(&cell); + let mut collision = [exec_input("exec-collision"), exec_output("exec-collision")]; + recorder.attach_to_prompt(&mut collision, &mut HashMap::new()); + assert_eq!(tool_calls_complete(&collision[1]), None); + + let recorder = new_recorder(InitialHistory::New); + let cell = CellId::new("late-runtime".to_string()); + recorder.start_cell(&cell, "exec-late"); + record_nested_call(&recorder, &cell, "nested-first"); + recorder.finish_cell_recording(&cell); + let mut first = [exec_input("exec-late"), exec_output("exec-late")]; + recorder.attach_to_prompt(&mut first, &mut HashMap::new()); + assert_eq!(tool_calls_complete(&first[1]), Some(true)); + record_nested_call(&recorder, &cell, "nested-late"); + let mut late = [exec_input("exec-late"), exec_output("exec-late")]; + recorder.attach_to_prompt(&mut late, &mut HashMap::new()); + assert_eq!(tool_calls_complete(&late[1]), None); +} + #[test] fn request_truncation_prevents_completion_after_compaction() { - let recorder = new_recorder(); + let recorder = new_recorder(InitialHistory::New); let cell_id = CellId::new("compacted-cell".to_string()); recorder.start_cell(&cell_id, "exec"); @@ -517,10 +811,10 @@ fn request_truncation_prevents_completion_after_compaction() { ); } - let mut initial = [output("exec")]; + let mut initial = [exec_input("exec"), exec_output("exec")]; assert!(recorder.attach_pending_to_prompt(&mut initial, &mut HashMap::new())); assert!( - initial[0] + initial[1] .executed_tool_call_metadata() .and_then(|metadata| metadata.executed_tool_calls.as_ref()) .is_some_and(|calls| calls.iter().any(|call| matches!( @@ -543,10 +837,10 @@ fn request_truncation_prevents_completion_after_compaction() { 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")]; + let mut final_output = [wait_input("wait", &cell_id), output("wait")]; recorder.attach_pending_to_prompt(&mut final_output, &mut HashMap::new()); assert_eq!( - final_output[0] + final_output[1] .executed_tool_call_metadata() .and_then(|metadata| metadata.tool_calls_complete), None, @@ -555,7 +849,7 @@ fn request_truncation_prevents_completion_after_compaction() { #[test] fn result_metadata_updates_the_exact_retained_call_and_marks_oversized_retries() { - let recorder = new_recorder(); + let recorder = new_recorder(InitialHistory::Forked(Vec::new())); let cell_id = CellId::new("metadata-cell".to_string()); recorder.start_cell(&cell_id, "exec"); let call = ExecutedToolCall::new("nested_tool".to_string(), json!({})); @@ -569,7 +863,7 @@ fn result_metadata_updates_the_exact_retained_call_and_marks_oversized_retries() } recorder.finish_cell_recording(&cell_id); let mut retry_cache = HashMap::new(); - let mut initial = [output("exec")]; + let mut initial = [exec_input("exec"), exec_output("exec")]; assert!(recorder.attach_pending_to_prompt(&mut initial, &mut retry_cache)); let metadata = json!({ "all-keys": { "ids": ["R1"] }, "other": true }); @@ -583,11 +877,11 @@ fn result_metadata_updates_the_exact_retained_call_and_marks_oversized_retries() let mut expected_call = call.clone(); expected_call.set_tool_result_metadata(ToolResultMetadata::new(&metadata)); - let mut expected = output("exec"); + let mut expected = exec_output("exec"); expected.append_executed_tool_calls(vec![expected_call, call]); expected.set_tool_call_cell_id("exec"); expected.mark_tool_calls_complete(); - let mut retry = [output("exec")]; + let mut retry = [exec_output("exec")]; assert!(recorder.attach_pending_to_prompt(&mut retry, &mut retry_cache)); assert_eq!(retry, [expected]); @@ -596,17 +890,48 @@ fn result_metadata_updates_the_exact_retained_call_and_marks_oversized_retries() "first", &json!({ "oversized": "x".repeat(32 * 1024) }), )); - let mut expected = serde_json::to_value(&initial).unwrap(); + let mut expected = serde_json::to_value(&initial[1..]).unwrap(); expected[0]["internal_chat_message_metadata_passthrough"]["executed_tool_calls"][0]["tool_result_metadata"] = json!("omitted_due_to_size_limit"); - let mut retry = [output("exec")]; + let mut retry = [exec_output("exec")]; assert!(recorder.attach_pending_to_prompt(&mut retry, &mut retry_cache)); assert_eq!(serde_json::to_value(&retry).unwrap(), expected); } +#[test] +fn untracked_history_budget_loss_keeps_later_wait_incomplete() { + let recorder = new_recorder(InitialHistory::New); + let cell = CellId::new("running-cell".to_string()); + recorder.start_cell(&cell, "exec"); + recorder.record_nested_tool_call( + cell.clone(), + "nested".to_string(), + ExecutedToolCall::new("nested_tool".to_string(), json!("x".repeat(8190))), + /*original_bytes*/ 8192, + ); + let mut prompt = vec![exec_input("exec"), exec_output("exec")]; + for index in 0..64 { + let mut item = output(&format!("untracked-{index}")); + item.append_executed_tool_calls(vec![ExecutedToolCall::new( + "other_tool".to_string(), + json!("x".repeat(1024)), + )]); + prompt.push(item); + } + recorder.attach_to_prompt(&mut prompt, &mut HashMap::new()); + assert!(prompt[1].executed_tool_call_metadata().is_none()); + + recorder.attach_to_prompt(&mut [], &mut HashMap::new()); + recorder.register_cell(&cell, "wait"); + recorder.finish_cell_recording(&cell); + let mut wait = [wait_input("wait", &cell), output("wait")]; + recorder.attach_to_prompt(&mut wait, &mut HashMap::new()); + assert_eq!(tool_calls_complete(&wait[1]), None); +} + #[test] fn result_metadata_shedding_preserves_completion_after_compaction() { - let recorder = new_recorder(); + let recorder = new_recorder(InitialHistory::New); let cell_id = CellId::new("compacted-cell".to_string()); recorder.start_cell(&cell_id, "exec"); @@ -627,7 +952,10 @@ fn result_metadata_shedding_preserves_completion_after_compaction() { runtime_tool_call_id: "runtime-call".to_string(), }; let mut retry_cache = HashMap::new(); - recorder.attach_pending_to_prompt(&mut [output("exec")], &mut retry_cache); + recorder.attach_pending_to_prompt( + &mut [exec_input("exec"), exec_output("exec")], + &mut retry_cache, + ); // This snapshot fits alone, but pushes the four calls over the aggregate budget. assert!(recorder.record_tool_result_metadata( &source, @@ -640,7 +968,7 @@ fn result_metadata_shedding_preserves_completion_after_compaction() { "nested-1", &json!({ "status": "unavailable" }) )); - let mut initial = [output("exec")]; + let mut initial = [exec_output("exec")]; assert!(recorder.attach_pending_to_prompt(&mut initial, &mut retry_cache)); let calls = initial[0] .executed_tool_call_metadata() @@ -651,7 +979,7 @@ fn result_metadata_shedding_preserves_completion_after_compaction() { call.set_tool_result_metadata(ToolResultMetadata::new(&json!("omitted_due_to_size_limit"))); } assert_eq!(calls, &expected_calls); - let mut retry = [output("exec")]; + let mut retry = [exec_output("exec")]; assert!(recorder.attach_pending_to_prompt(&mut retry, &mut retry_cache)); assert_eq!(retry, initial); @@ -660,10 +988,10 @@ fn result_metadata_shedding_preserves_completion_after_compaction() { assert!(recorder.record_tool_result_metadata(&source, "nested-2", &metadata)); let mut expected_calls = calls.clone(); expected_calls[2].set_tool_result_metadata(ToolResultMetadata::new(&metadata)); - let mut expected = output("exec"); + let mut expected = exec_output("exec"); expected.append_executed_tool_calls(expected_calls); expected.set_tool_call_cell_id("exec"); - let mut retry = [output("exec")]; + let mut retry = [exec_output("exec")]; assert!(recorder.attach_pending_to_prompt(&mut retry, &mut retry_cache)); assert_eq!(retry, [expected]); @@ -672,10 +1000,10 @@ fn result_metadata_shedding_preserves_completion_after_compaction() { assert!(!recorder.record_tool_result_metadata(&source, "nested-2", &metadata)); recorder.register_cell(&cell_id, "wait"); recorder.finish_cell_recording(&cell_id); - let mut final_output = [output("wait")]; + let mut final_output = [wait_input("wait", &cell_id), output("wait")]; recorder.attach_pending_to_prompt(&mut final_output, &mut HashMap::new()); assert_eq!( - final_output[0] + final_output[1] .executed_tool_call_metadata() .and_then(|metadata| metadata.tool_calls_complete), Some(true), @@ -684,7 +1012,7 @@ fn result_metadata_shedding_preserves_completion_after_compaction() { #[test] fn finished_cells_without_more_waits_do_not_block_new_calls() { - let recorder = new_recorder(); + let recorder = new_recorder(InitialHistory::Forked(Vec::new())); 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}")); @@ -695,7 +1023,10 @@ fn finished_cells_without_more_waits_do_not_block_new_calls() { call.clone(), /*original_bytes*/ 2, ); - recorder.attach_pending_to_prompt(&mut [output(cell.as_str())], &mut HashMap::new()); + recorder.attach_pending_to_prompt( + &mut [exec_input(cell.as_str()), exec_output(cell.as_str())], + &mut HashMap::new(), + ); recorder.finish_cell_recording(&cell); } let fresh = CellId::new("fresh".to_string()); @@ -707,11 +1038,11 @@ fn finished_cells_without_more_waits_do_not_block_new_calls() { /*original_bytes*/ 2, ); recorder.finish_cell_recording(&fresh); - let mut expected = output("fresh-output"); + let mut expected = exec_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")]; + let mut items = [exec_input("fresh-output"), exec_output("fresh-output")]; assert!(recorder.attach_pending_to_prompt(&mut items, &mut HashMap::new())); - assert_eq!(items, [expected]); + assert_eq!(items, [exec_input("fresh-output"), expected]); } diff --git a/codex-rs/core/src/tools/executed_tool_calls/seen_ids.rs b/codex-rs/core/src/tools/executed_tool_calls/seen_ids.rs new file mode 100644 index 0000000000..217839c614 --- /dev/null +++ b/codex-rs/core/src/tools/executed_tool_calls/seen_ids.rs @@ -0,0 +1,133 @@ +//! Tracks seen call and runtime cell IDs with bounded memory. Bits are never cleared: +//! collisions can withhold a proof, but cannot make an observed ID appear fresh again. + +use std::collections::hash_map::DefaultHasher; +use std::hash::Hash; +use std::hash::Hasher; + +use codex_code_mode::CellId; +use codex_history::InitialHistory; +use codex_history::RolloutItem; +use codex_protocol::models::ResponseItem; + +const ID_FILTER_WORDS: usize = 8 * 1024; +const MAX_HISTORY_ENTRIES: usize = 100_000; +const MAX_HISTORY_ID_BYTES: usize = 4 * 1024 * 1024; +// IDs are untrusted strings. A longer ID cannot prove freshness without unbounded hashing. +const MAX_ID_BYTES: usize = 1024; + +pub(super) struct SeenIds { + bits: Box<[u64]>, + history_ids_indexed: bool, +} + +impl Default for SeenIds { + fn default() -> Self { + Self { + bits: vec![0; ID_FILTER_WORDS].into_boxed_slice(), + history_ids_indexed: true, + } + } +} + +impl SeenIds { + pub(super) fn unobserved_history() -> Self { + // Recording started after session initialization, so earlier IDs are unknown. + Self { + history_ids_indexed: false, + ..Self::default() + } + } + + pub(super) fn from_history(history: &InitialHistory) -> Self { + // Index the supplied history, including paginated context. Missing older pages + // do not invalidate a new cell whose calls and outputs are observed locally. + let mut tracker = Self::default(); + let mut entries_left = MAX_HISTORY_ENTRIES; + let mut id_bytes_left = MAX_HISTORY_ID_BYTES; + 'history: for item in history.get_rollout_items() { + if entries_left == 0 { + tracker.history_ids_indexed = false; + break; + } + entries_left -= 1; + match item { + RolloutItem::ResponseItem(item) + if !tracker.observe_history_item(item, &mut id_bytes_left) => + { + tracker.history_ids_indexed = false; + break; + } + RolloutItem::Compacted(item) => { + for item in item.replacement_history.as_deref().unwrap_or_default() { + if entries_left == 0 { + tracker.history_ids_indexed = false; + break 'history; + } + entries_left -= 1; + if !tracker.observe_history_item(item, &mut id_bytes_left) { + tracker.history_ids_indexed = false; + break 'history; + } + } + } + _ => {} + } + } + tracker + } + + /// Fresh relative to IDs observed in the supplied history and this recorder. + pub(super) fn observe_call_id(&mut self, id: &str) -> bool { + self.observe(/*namespace*/ 0, id) + } + + pub(super) fn observe_runtime_cell_id(&mut self, id: &CellId) -> bool { + // Runtime handles are tracked only within this recorder's lifetime. + self.observe(/*namespace*/ 1, id.as_str()) + } + + pub(super) fn history_ids_indexed(&self) -> bool { + self.history_ids_indexed + } + + fn observe_history_item(&mut self, item: &ResponseItem, bytes_left: &mut usize) -> bool { + let call_id = super::input_call_id(item).or_else(|| super::output_call_id(item)); + // Metadata cell_id identifies the originating exec call, not a runtime handle. + let exec_origin = item + .executed_tool_call_metadata() + .and_then(|metadata| metadata.cell_id.as_deref()); + for id in [call_id, exec_origin].into_iter().flatten() { + if id.len() > MAX_ID_BYTES || id.len() > *bytes_left { + return false; + } + *bytes_left -= id.len(); + self.observe_call_id(id); + } + true + } + + fn observe(&mut self, namespace: u8, id: &str) -> bool { + if id.len() > MAX_ID_BYTES { + return false; + } + let mut hasher = DefaultHasher::new(); + (namespace, id).hash(&mut hasher); + let hash = hasher.finish(); + let first = hash as usize; + let step = (hash >> 32) as usize | 1; + let mut fresh = false; + for probe in 0_usize..4 { + let bit = first.wrapping_add(probe.wrapping_mul(step)) % (ID_FILTER_WORDS * 64); + let word = &mut self.bits[bit / 64]; + let mask = 1_u64 << (bit % 64); + fresh |= *word & mask == 0; + *word |= mask; + } + fresh + } +} + +#[cfg(test)] +#[path = "seen_ids_tests.rs"] +mod tests; diff --git a/codex-rs/core/src/tools/executed_tool_calls/seen_ids_tests.rs b/codex-rs/core/src/tools/executed_tool_calls/seen_ids_tests.rs new file mode 100644 index 0000000000..47ecb7c0b1 --- /dev/null +++ b/codex-rs/core/src/tools/executed_tool_calls/seen_ids_tests.rs @@ -0,0 +1,14 @@ +use super::*; + +#[test] +fn id_namespaces_are_separate_and_saturated_filters_reject_new_ids() { + let mut tracker = SeenIds::default(); + assert!(tracker.observe_call_id("shared-id")); + assert!(tracker.observe_runtime_cell_id(&CellId::new("shared-id".to_string()))); + assert!(!tracker.observe_call_id("shared-id")); + assert!(!tracker.observe_runtime_cell_id(&CellId::new("shared-id".to_string()))); + + tracker.bits.fill(u64::MAX); + assert!(!tracker.observe_call_id("another-origin")); + assert!(!tracker.observe_runtime_cell_id(&CellId::new("another-runtime".to_string()))); +} diff --git a/codex-rs/core/src/tools/handlers/mcp.rs b/codex-rs/core/src/tools/handlers/mcp.rs index f7ebcdd98f..027f62f463 100644 --- a/codex-rs/core/src/tools/handlers/mcp.rs +++ b/codex-rs/core/src/tools/handlers/mcp.rs @@ -709,8 +709,10 @@ mod tests { .features .enable(Feature::ExecutedToolCallMetadata) .expect("test feature must be configurable"); - let recorder = - crate::tools::executed_tool_calls::ExecutedToolCalls::new(&turn.config.features); + let recorder = crate::tools::executed_tool_calls::ExecutedToolCalls::new( + &turn.config.features, + &codex_history::InitialHistory::New, + ); session.services.executed_tool_calls = recorder.clone(); let turn = Arc::new(turn); let handler = McpHandler::new(tool_info("filesystem", "filesystem", "read_file")) diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index 76bf9cd661..4243ae75da 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -49,6 +49,7 @@ use codex_protocol::openai_models::ToolMode; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::Op; +use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::ThreadSettingsOverrides; use codex_protocol::user_input::UserInput; use codex_tools::FreeformTool; @@ -1343,6 +1344,133 @@ await new Promise(() => {}); Ok(()) } +#[test_case("call-wait", None; "reused_wait_id")] +#[test_case("call-fresh-wait", Some(true); "fresh_wait_id")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn code_mode_wait_id_stays_known_after_compaction( + terminal_call_id: &str, + expected_complete: Option, +) -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let mut builder = test_codex() + .with_model("test-gpt-5.1-codex") + .with_config(|config| { + let _ = config.features.enable(Feature::CodeMode); + let _ = config.features.enable(Feature::CodeModeHost); + let _ = config.features.enable(Feature::ExecutedToolCallMetadata); + config.code_mode.disable_in_process_fallback = true; + }); + let test = builder.build_with_auto_env(&server).await?; + let started = responses::mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_custom_tool_call( + "call-exec", + "exec", + r#"await tools.test_sync_tool({}); text("started"); yield_control(); await new Promise(() => {});"#, + ), + ev_completed("resp-start"), + ]), + sse(vec![ + ev_assistant_message("msg-start", "running"), + ev_completed("resp-start-done"), + ]), + ], + ) + .await; + test.submit_turn("Start a cell").await?; + let requests = started.requests(); + assert_eq!(requests.len(), 2); + let first_request = &requests[1]; + let first_items = custom_tool_output_items(first_request, "call-exec"); + let cell_id = extract_running_cell_id(text_item(&first_items, /*index*/ 0)); + assert_eq!(text_item(&first_items, /*index*/ 1), "started"); + let first_output = first_request.custom_tool_call_output("call-exec"); + let metadata = &first_output["internal_chat_message_metadata_passthrough"]; + assert_eq!(metadata["cell_id"], "call-exec"); + assert_eq!( + metadata["executed_tool_calls"], + serde_json::json!([{"name": "test_sync_tool", "arguments": {}}]), + ); + assert!(metadata.get("tool_calls_complete").is_none()); + + let yielded = responses::mount_function_call_agent_response( + &server, + "call-wait", + &serde_json::json!({"cell_id": cell_id, "yield_time_ms": 1}).to_string(), + "wait", + ) + .await; + test.submit_turn("Wait once").await?; + yielded.function_call.single_request(); + let yielded_items = + function_tool_output_items(&yielded.completion.single_request(), "call-wait"); + assert_eq!( + extract_running_cell_id(text_item(&yielded_items, /*index*/ 0)), + cell_id, + ); + + let compact = responses::mount_sse_once( + &server, + sse(vec![ + serde_json::json!({ + "type": "response.output_item.done", + "item": {"type": "compaction", "encrypted_content": "compacted history"}, + }), + ev_completed("resp-compact"), + ]), + ) + .await; + test.codex.submit(Op::Compact).await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + assert_eq!( + compact + .single_request() + .inputs_of_type("compaction_trigger") + .len(), + 1, + ); + + let terminal = responses::mount_function_call_agent_response( + &server, + terminal_call_id, + &serde_json::json!({"cell_id": cell_id, "terminate": true}).to_string(), + "wait", + ) + .await; + test.submit_turn("Finish the compacted cell").await?; + assert!( + terminal + .function_call + .single_request() + .input() + .iter() + .all(|item| { item["call_id"] != "call-exec" && item["call_id"] != "call-wait" }) + ); + let final_request = terminal.completion.single_request(); + let final_items = function_tool_output_items(&final_request, terminal_call_id); + assert_eq!(final_items.len(), 1); + assert_regex_match( + r"^Script terminated\nWall time \d+\.\d seconds\nOutput:\n\z", + text_item(&final_items, /*index*/ 0), + ); + let output = final_request.function_call_output(terminal_call_id); + let metadata = &output["internal_chat_message_metadata_passthrough"]; + assert_eq!( + metadata.get("tool_calls_complete").and_then(Value::as_bool), + expected_complete, + ); + assert!(metadata.get("executed_tool_calls").is_none()); + test.codex.shutdown_and_wait().await?; + Ok(()) +} + const RESULT_METADATA_TOOL: &str = "message_search"; const RESULT_METADATA_PRIVATE_RESULT: &str = "connector result text is not result metadata"; @@ -1603,6 +1731,103 @@ async fn code_mode_result_metadata_follows_call_binding( Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn code_mode_result_metadata_follows_runtime_recording_enablement() -> Result<()> { + skip_if_no_network!(Ok(())); + let server = responses::start_mock_server().await; + let result_metadata = serde_json::json!({ "provider": { "id": "refresh" } }); + let apps_server = mount_result_metadata_app( + &server, + Some(result_metadata.clone()), + /*is_error*/ false, + ) + .await?; + let mut builder = + result_metadata_apps_builder(apps_server.chatgpt_base_url, "employee@openai.com") + .with_config(|config| { + config + .features + .disable(Feature::ExecutedToolCallMetadata) + .unwrap(); + }); + let test = builder.build_with_auto_env(&server).await?; + let arguments = serde_json::json!({ "search": "launch plan" }); + let code = format!( + "const tool = ALL_TOOLS.find(({{ name }}) => name.endsWith(\"{RESULT_METADATA_TOOL}\")); \ + const result = await tools[tool.name]({arguments}); \ + text(JSON.stringify({{ isError: Boolean(result.isError), hasMeta: Object.hasOwn(result, \"_meta\") }}));" + ); + for (call_id, enabled) in [("call-off", false), ("call-on", true)] { + if enabled { + let mut config = test.config.clone(); + config + .features + .enable(Feature::ExecutedToolCallMetadata) + .unwrap(); + test.codex.refresh_runtime_config(config).await; + // Runtime recording changes without updating the session's execution features. + assert!( + !test + .codex + .config() + .await + .features + .enabled(Feature::ExecutedToolCallMetadata) + ); + } + let initial = responses::mount_sse_once( + &server, + sse(vec![ + ev_custom_tool_call(call_id, "exec", &code), + ev_completed(call_id), + ]), + ) + .await; + let follow_up = responses::mount_sse_once( + &server, + sse(vec![ + ev_assistant_message(&format!("{call_id}-message"), "done"), + ev_completed(&format!("{call_id}-done")), + ]), + ) + .await; + test.submit_turn("Search a connected app").await?; + initial.single_request(); + let request = follow_up.single_request(); + let (body, success) = custom_tool_output_body_and_success(&request, call_id); + assert_ne!(success, Some(false), "Code Mode failed: {body}"); + assert_eq!( + serde_json::from_str::(&body)?, + serde_json::json!({ "isError": false, "hasMeta": false }), + ); + assert_eq!( + result_metadata_fixture_calls(&request.input()).count(), + usize::from(enabled) + ); + let output = request.custom_tool_call_output(call_id); + assert!( + output["internal_chat_message_metadata_passthrough"] + .get("tool_calls_complete") + .is_none() + ); + if enabled { + assert_result_metadata_call(&output, &arguments, /*expected_metadata*/ None); + } + } + assert_eq!(recorded_apps_tool_calls(&server).await.len(), 2); + let captured = codex_core::test_support::history_with_tool_call_metadata(&test.codex).await; + let captured = serde_json::to_value(captured)?; + let captured_output = captured + .as_array() + .unwrap() + .iter() + .find(|item| item["type"] == "custom_tool_call_output" && item["call_id"] == "call-on") + .expect("captured exec output after runtime enablement"); + let expected_metadata = None; + assert_result_metadata_call(captured_output, &arguments, expected_metadata); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn code_mode_result_metadata_keeps_prepared_call_binding_across_runtime_refresh() -> Result<()> { @@ -1759,6 +1984,230 @@ async fn code_mode_result_metadata_keeps_prepared_call_binding_across_runtime_re Ok(()) } +#[cfg_attr(windows, ignore = "no exec_command on Windows")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn code_mode_resumed_wait_does_not_certify_a_reused_runtime_cell() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let configure = |config: &mut Config| { + let _ = config.features.enable(Feature::CodeMode); + let _ = config.features.enable(Feature::CodeModeHost); + let _ = config.features.enable(Feature::ExecutedToolCallMetadata); + config.code_mode.disable_in_process_fallback = true; + }; + let mut builder = test_codex() + .with_model("test-gpt-5.1-codex") + .with_history_mode(ThreadHistoryMode::Legacy) + .with_config(configure); + let initial = builder.build(&server).await?; + let seed = responses::mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_custom_tool_call( + "call-a", + "exec", + r#"text("A"); yield_control(); await new Promise(() => {});"#, + ), + ev_completed("resp-a"), + ]), + sse(vec![ + ev_assistant_message("msg-a", "A is running"), + ev_completed("resp-a-done"), + ]), + ], + ) + .await; + initial.submit_turn("Start A").await?; + let old_output = seed.requests()[1].custom_tool_call_output("call-a"); + let old_items = custom_tool_output_items(&seed.requests()[1], "call-a"); + let old_cell_id = extract_running_cell_id(text_item(&old_items, /*index*/ 0)); + assert_eq!(old_cell_id, "1"); + assert_eq!(text_item(&old_items, /*index*/ 1), "A"); + + // Restart loads the old running-cell output, but starts a new host whose IDs reset. + let resumed = test_codex() + .with_model("test-gpt-5.1-codex") + .with_config(configure) + .restart(&server, &initial) + .await?; + assert_eq!( + resumed.session_configured.thread_id, + initial.session_configured.thread_id + ); + let gate = resumed.workspace_path("release-cell-b"); + let quoted_gate = shlex::try_join([gate.to_string_lossy().as_ref()])?; + let arguments = serde_json::json!({ + "cmd": format!( + "attempt=0; while [ ! -f {quoted_gate} ] && [ \"$attempt\" -lt 1000 ]; \ + do sleep 0.01; attempt=$((attempt + 1)); done; [ -f {quoted_gate} ] && printf B" + ), + "login": false, + "yield_time_ms": 30_000, + "max_output_tokens": 1000, + }); + let code = format!("yield_control(); text((await tools.exec_command({arguments})).output);"); + let continued = responses::mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_custom_tool_call("call-b", "exec", &code), + ev_completed("resp-b"), + ]), + sse(vec![ + ev_assistant_message("msg-b", "B is running"), + ev_completed("resp-b-done"), + ]), + sse(vec![ + responses::ev_function_call( + "call-stale-wait", + "wait", + &serde_json::json!({"cell_id": old_cell_id, "yield_time_ms": 10_000}) + .to_string(), + ), + ev_completed("resp-wait"), + ]), + sse(vec![ + ev_assistant_message("msg-wait", "done"), + ev_completed("resp-wait-done"), + ]), + sse(vec![ + ev_custom_tool_call( + "call-fresh", + "exec", + r#"await tools.test_sync_tool({}); text("fresh");"#, + ), + ev_completed("resp-fresh"), + ]), + sse(vec![ + ev_assistant_message("msg-fresh", "done"), + ev_completed("resp-fresh-done"), + ]), + ], + ) + .await; + resumed.submit_turn("Start B").await?; + let requests = continued.requests(); + assert_eq!( + requests[0].custom_tool_call_output("call-a")["output"], + old_output["output"] + ); + let new_items = custom_tool_output_items(&requests[1], "call-b"); + assert_eq!( + extract_running_cell_id(text_item(&new_items, /*index*/ 0)), + old_cell_id + ); + + fs::write(&gate, "ready")?; + resumed.submit_turn("Wait using A's old cell ID").await?; + let requests = continued.requests(); + assert_eq!(requests.len(), 4); + let items = function_tool_output_items(&requests[3], "call-stale-wait"); + assert_eq!(items.len(), 2); + assert!(text_item(&items, /*index*/ 0).starts_with("Script completed\n")); + assert_eq!(text_item(&items, /*index*/ 1), "B"); + let wait = requests[3].function_call_output("call-stale-wait"); + let metadata = &wait["internal_chat_message_metadata_passthrough"]; + assert!(metadata.get("tool_calls_complete").is_none()); + let mut recorded = Vec::new(); + for output in [requests[3].custom_tool_call_output("call-b"), wait] { + let metadata = &output["internal_chat_message_metadata_passthrough"]; + if let Some(calls) = metadata + .get("executed_tool_calls") + .and_then(Value::as_array) + && !calls.is_empty() + { + assert_eq!(metadata["cell_id"], "call-b"); + recorded.extend_from_slice(calls); + } + } + assert_eq!( + recorded, + vec![serde_json::json!({"name": "exec_command", "arguments": arguments})] + ); + + resumed.submit_turn("Run a fresh immediate exec").await?; + let requests = continued.requests(); + assert_eq!(requests.len(), 6); + assert_eq!( + custom_tool_output_last_non_empty_text(&requests[5], "call-fresh").as_deref(), + Some("fresh") + ); + let fresh = requests[5].custom_tool_call_output("call-fresh"); + let metadata = &fresh["internal_chat_message_metadata_passthrough"]; + assert_eq!( + metadata["executed_tool_calls"], + serde_json::json!([{"name": "test_sync_tool", "arguments": {}}]) + ); + assert_eq!(metadata["cell_id"], "call-fresh"); + assert_eq!(metadata["tool_calls_complete"], true); + resumed.codex.shutdown_and_wait().await?; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn code_mode_complete_call_survives_unrelated_truncation() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let (test, _) = run_code_mode_turn_with_config( + &server, + "Record an oversized call", + r#"await tools.test_sync_tool({ barrier: { id: "x".repeat(8192), participants: 1 } }); text("done");"#, + |config| { + let _ = config.features.disable(Feature::RemoteCompactionV2); + }, + ) + .await?; + + responses::mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-3"), + ev_custom_tool_call( + "call-2", + "exec", + r#"await tools.test_sync_tool({ wait_for_git_enrichment: false }); text("done");"#, + ), + ev_completed("resp-3"), + ]), + ) + .await; + let follow_up = responses::mount_sse_once( + &server, + sse(vec![ + ev_assistant_message("msg-2", "done"), + ev_completed("resp-4"), + ]), + ) + .await; + test.submit_turn("Record a complete call").await?; + test.codex.shutdown_and_wait().await?; + + let request = follow_up.single_request(); + let overflow = + &request.custom_tool_call_output("call-1")["internal_chat_message_metadata_passthrough"]; + assert!( + overflow["executed_tool_calls"] + .as_array() + .is_some_and(|calls| calls.iter().any(|call| call["arguments"] + .get("_codex_executed_tool_call_truncated") + .is_some())) + ); + assert!(overflow.get("tool_calls_complete").is_none()); + + let complete = + &request.custom_tool_call_output("call-2")["internal_chat_message_metadata_passthrough"]; + assert_eq!( + complete["executed_tool_calls"], + serde_json::json!([{ "name": "test_sync_tool", "arguments": { "wait_for_git_enrichment": false } }]), + ); + assert_eq!(complete["cell_id"], "call-2"); + assert_eq!(complete["tool_calls_complete"], true); + 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(())); diff --git a/codex-rs/protocol/src/models/executed_tool_calls.rs b/codex-rs/protocol/src/models/executed_tool_calls.rs index fa4d1aaf4e..ffca9bbfb5 100644 --- a/codex-rs/protocol/src/models/executed_tool_calls.rs +++ b/codex-rs/protocol/src/models/executed_tool_calls.rs @@ -1,3 +1,5 @@ +use std::collections::HashSet; + use schemars::JsonSchema; use serde::Deserialize; use serde::Serialize; @@ -88,23 +90,16 @@ fn bound_executed_tool_calls_for_prompt_with_priority( items: &mut [ResponseItem], prioritize_recent: bool, ) { - let mut remaining_items = 0_usize; - let mut original_calls = 0_usize; - let mut original_metadata_bytes = 0_usize; - let mut truncated = false; - + let mut damaged_cells = HashSet::new(); for item in items.iter_mut() { - let Some(calls) = item + let Some(metadata) = item .internal_chat_message_metadata_passthrough_mut() .and_then(Option::as_mut) - .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 { + let mut truncated = false; + for call in metadata.executed_tool_calls.iter_mut().flatten() { let argument_bytes = serde_json::to_vec(&call.arguments) .map(|bytes| bytes.len()) .unwrap_or(usize::MAX); @@ -117,229 +112,103 @@ fn bound_executed_tool_calls_for_prompt_with_priority( ); } truncated |= call.truncation().is_some(); - original_calls = original_calls.saturating_add(1).saturating_add( - call.truncation() - .and_then(|truncation| truncation.omitted_calls) - .unwrap_or_default(), - ); } - remaining_items += 1; - original_metadata_bytes = - original_metadata_bytes.saturating_add(executed_tool_call_metadata_bytes(item)); - } - - // Raw result metadata must not displace existing source evidence, calls or completion proof. - if original_metadata_bytes > MAX_EXECUTED_TOOL_CALL_METADATA_BYTES { - original_metadata_bytes = 0; - for item in items.iter_mut() { - if let Some(metadata) = item - .internal_chat_message_metadata_passthrough_mut() - .and_then(Option::as_mut) - { - for call in metadata.executed_tool_calls.iter_mut().flatten() { - if call.tool_result_metadata.is_some() { - call.tool_result_metadata = ToolResultMetadata::omitted_due_to_size_limit(); - } - } - } - original_metadata_bytes = - original_metadata_bytes.saturating_add(executed_tool_call_metadata_bytes(item)); + if truncated { + metadata.tool_calls_complete = None; + damaged_cells.extend(metadata.cell_id.clone()); } } + clear_damaged_cell_completeness(items, &damaged_cells); - // Omission markers are optional too; keep the original call budget if they cannot fit. - if original_metadata_bytes > MAX_EXECUTED_TOOL_CALL_METADATA_BYTES { - original_metadata_bytes = 0; - for item in items.iter_mut() { - item.clear_tool_result_metadata(); - original_metadata_bytes = - original_metadata_bytes.saturating_add(executed_tool_call_metadata_bytes(item)); - } - } - - // Source evidence is optional; dropping it must not discard calls or their completion proof. - if original_metadata_bytes > MAX_EXECUTED_TOOL_CALL_METADATA_BYTES { - original_metadata_bytes = 0; - for item in items.iter_mut() { - if let Some(metadata) = item - .internal_chat_message_metadata_passthrough_mut() - .and_then(Option::as_mut) - { - for call in metadata.executed_tool_calls.iter_mut().flatten() { - call.tool_result_sources = None; - } - } - original_metadata_bytes = - 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(), metadata.cell_id.clone())) + let metadata_bytes = |items: &[ResponseItem]| { + items.iter().fold(0_usize, |bytes, item| { + bytes.saturating_add(executed_tool_call_metadata_bytes(item)) }) - }); - - let omitted_call_reservation = serde_json::to_vec(&serde_json::json!({ - "_codex_executed_tool_call_truncated": ExecutedToolCallTruncation { - original_bytes: usize::MAX, - max_bytes: usize::MAX, - omitted_calls: Some(usize::MAX), - original_name_bytes: Some(usize::MAX), - }, - })) - .map(|bytes| bytes.len()) - .unwrap_or(usize::MAX); - let mut remaining_bytes = - MAX_EXECUTED_TOOL_CALL_METADATA_BYTES.saturating_sub(omitted_call_reservation); + }; + if metadata_bytes(items) <= MAX_EXECUTED_TOOL_CALL_METADATA_BYTES { + return; + } + // Raw result metadata must not displace existing source evidence, calls or completion proof. for item in items.iter_mut() { - let Some(metadata) = item.executed_tool_call_metadata() else { - continue; - }; - if metadata - .executed_tool_calls - .as_ref() - .is_none_or(Vec::is_empty) + if let Some(metadata) = item + .internal_chat_message_metadata_passthrough_mut() + .and_then(Option::as_mut) { + for call in metadata.executed_tool_calls.iter_mut().flatten() { + if call.tool_result_metadata.is_some() { + call.tool_result_metadata = ToolResultMetadata::omitted_due_to_size_limit(); + } + } + } + } + + if metadata_bytes(items) <= MAX_EXECUTED_TOOL_CALL_METADATA_BYTES { + return; + } + // Omission markers are optional too; keep the original call budget if they cannot fit. + for item in items.iter_mut() { + item.clear_tool_result_metadata(); + } + + if metadata_bytes(items) <= MAX_EXECUTED_TOOL_CALL_METADATA_BYTES { + return; + } + // Source evidence is optional; dropping it must not discard calls or their completion proof. + for item in items.iter_mut() { + if let Some(metadata) = item + .internal_chat_message_metadata_passthrough_mut() + .and_then(Option::as_mut) + { + for call in metadata.executed_tool_calls.iter_mut().flatten() { + call.tool_result_sources = None; + } + } + } + if metadata_bytes(items) <= MAX_EXECUTED_TOOL_CALL_METADATA_BYTES { + return; + } + + let mut remaining_items = items + .iter() + .filter(|item| executed_tool_call_metadata_bytes(item) > 0) + .count(); + let mut remaining_bytes = MAX_EXECUTED_TOOL_CALL_METADATA_BYTES; + for item in items.iter_mut() { + let item_bytes = executed_tool_call_metadata_bytes(item); + if item_bytes == 0 { continue; } - - let metadata_field_bytes = executed_tool_call_metadata_field_bytes(metadata); let item_budget = if prioritize_recent { remaining_bytes } else { remaining_bytes / remaining_items }; - item.bound_executed_tool_calls_with_budget( - item_budget.saturating_sub(metadata_field_bytes), - ); + if item_bytes > item_budget { + // Remember the cell before a too-small share removes its metadata entirely. + damaged_cells.extend( + item.executed_tool_call_metadata() + .and_then(|metadata| metadata.cell_id.clone()), + ); + item.clear_tool_calls_complete(); + item.bound_executed_tool_calls_with_budget(item_budget); + } remaining_bytes = remaining_bytes.saturating_sub(executed_tool_call_metadata_bytes(item)); remaining_items -= 1; } - - let represented_calls = items - .iter() - .filter_map(ResponseItem::executed_tool_call_metadata) - .filter_map(|metadata| metadata.executed_tool_calls.as_ref()) - .flatten() - .map(|call| { - 1 + call - .truncation() - .and_then(|truncation| truncation.omitted_calls) - .unwrap_or_default() - }) - .sum::(); - if represented_calls == original_calls { - return; - } - - if represented_calls == 0 { - if let Some((index, mut call, cell_id)) = overflow_fallback { - let original_bytes = call - .truncation() - .map(|truncation| truncation.original_bytes) - .unwrap_or_else(|| { - serde_json::to_vec(&call.arguments) - .map(|bytes| bytes.len()) - .unwrap_or(usize::MAX) - }); - let original_name_bytes = call.name.len(); - let name_boundary = call.name.floor_char_boundary( - original_name_bytes.min(MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES / 2), - ); - call.name.truncate(name_boundary); - call.set_truncation_with_name( - original_bytes, - /*max_bytes*/ 0, - Some(original_calls.saturating_sub(1)), - (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; - } - - let omission_call = if prioritize_recent { - items.iter_mut().rev().find_map(first_executed_tool_call) - } else { - items.iter_mut().find_map(first_executed_tool_call) - }; - if let Some(call) = omission_call { - let original_bytes = call - .truncation() - .map(|truncation| truncation.original_bytes) - .unwrap_or_else(|| { - serde_json::to_vec(&call.arguments) - .map(|bytes| bytes.len()) - .unwrap_or(usize::MAX) - }); - let max_bytes = call - .truncation() - .map(|truncation| truncation.max_bytes) - .unwrap_or_default(); - let previous_omissions = call - .truncation() - .and_then(|truncation| truncation.omitted_calls) - .unwrap_or_default(); - call.set_truncation( - original_bytes, - max_bytes, - Some( - previous_omissions.saturating_add(original_calls.saturating_sub(represented_calls)), - ), - ); - } + clear_damaged_cell_completeness(items, &damaged_cells); } -fn first_executed_tool_call(item: &mut ResponseItem) -> Option<&mut ExecutedToolCall> { - item.internal_chat_message_metadata_passthrough_mut() - .and_then(Option::as_mut) - .and_then(|metadata| metadata.executed_tool_calls.as_mut()) - .and_then(|calls| calls.first_mut()) +fn clear_damaged_cell_completeness(items: &mut [ResponseItem], damaged_cells: &HashSet) { + for item in items { + if item.executed_tool_call_metadata().is_some_and(|metadata| { + metadata + .cell_id + .as_ref() + .is_some_and(|cell_id| damaged_cells.contains(cell_id)) + }) { + item.clear_tool_calls_complete(); + } + } } /// Raw model arguments or trusted truncation metadata for an attempted tool call. @@ -601,6 +470,16 @@ impl ResponseItem { } } + /// Discards a completion claim when the host cannot retain its full evidence. + pub fn clear_tool_calls_complete(&mut self) { + if let Some(metadata) = self + .internal_chat_message_metadata_passthrough_mut() + .and_then(Option::as_mut) + { + metadata.tool_calls_complete = None; + } + } + /// 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() @@ -618,64 +497,68 @@ impl ResponseItem { } } - /// Bounds one request item's attempted calls to its share of the prompt budget. + /// Replaces an over-budget output's calls with its own omission marker. fn bound_executed_tool_calls_with_budget(&mut self, max_metadata_bytes: usize) { - let Some(metadata) = self.internal_chat_message_metadata_passthrough_mut() else { + let Some(metadata) = self.executed_tool_call_metadata() else { return; }; - let Some(passthrough) = metadata.as_mut() else { + let max_call_bytes = + max_metadata_bytes.saturating_sub(executed_tool_call_metadata_field_bytes(metadata)); + let Some(calls) = self + .internal_chat_message_metadata_passthrough_mut() + .and_then(Option::as_mut) + .and_then(|metadata| metadata.executed_tool_calls.as_mut()) + .filter(|calls| !calls.is_empty()) + else { + self.clear_executed_tool_calls(); return; }; - let Some(calls) = passthrough.executed_tool_calls.as_mut() else { - return; - }; - - let mut serialized_bytes = 2; - let mut retained_calls = 0; - calls.retain_mut(|call| { - let separator_bytes = usize::from(retained_calls > 0); - let remaining_bytes = max_metadata_bytes - .saturating_sub(serialized_bytes) - .saturating_sub(separator_bytes); - let argument_bytes = serde_json::to_vec(&call.arguments) - .map(|bytes| bytes.len()) - .unwrap_or(usize::MAX); - let call_bytes = serde_json::to_vec(&*call) - .map(|bytes| bytes.len()) - .unwrap_or(usize::MAX); - - if call_bytes > remaining_bytes - || argument_bytes > MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES - { - let original_bytes = call - .truncation() - .map(|truncation| truncation.original_bytes) - .unwrap_or(argument_bytes); - let omitted_calls = call - .truncation() - .and_then(|truncation| truncation.omitted_calls); - call.set_truncation( - original_bytes, - remaining_bytes.min(MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES), - omitted_calls, - ); - } - - let call_bytes = serde_json::to_vec(&*call) - .map(|bytes| bytes.len()) - .unwrap_or(usize::MAX); - if call_bytes > remaining_bytes { - return false; - } - - serialized_bytes = serialized_bytes - .saturating_add(separator_bytes) - .saturating_add(call_bytes); - retained_calls += 1; - true + let represented_calls = calls.iter().fold(0_usize, |count, call| { + count.saturating_add(1).saturating_add( + call.truncation() + .and_then(|truncation| truncation.omitted_calls) + .unwrap_or_default(), + ) }); - - if calls.is_empty() { + calls.truncate(1); + let call = &mut calls[0]; + let original_bytes = call + .truncation() + .map(|truncation| truncation.original_bytes) + .unwrap_or_else(|| { + serde_json::to_vec(&call.arguments) + .map(|bytes| bytes.len()) + .unwrap_or(usize::MAX) + }); + let original_name_bytes = call + .truncation() + .and_then(|truncation| truncation.original_name_bytes); + let omitted_calls = (represented_calls > 1).then_some(represented_calls - 1); + call.set_truncation_with_name( + original_bytes, + max_call_bytes.min(MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES), + omitted_calls, + original_name_bytes, + ); + let serialized_bytes = |calls: &[ExecutedToolCall]| { + serde_json::to_vec(calls) + .map(|bytes| bytes.len()) + .unwrap_or(usize::MAX) + }; + if serialized_bytes(calls) > max_call_bytes { + let call = &mut calls[0]; + call.set_truncation_with_name( + original_bytes, + max_call_bytes.min(MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES), + omitted_calls, + Some(original_name_bytes.unwrap_or(call.name.len())), + ); + // Removing UTF-8 name bytes saves at least that many serialized JSON bytes. + let excess_bytes = serialized_bytes(calls).saturating_sub(max_call_bytes); + let name = &mut calls[0].name; + name.truncate(name.floor_char_boundary(name.len().saturating_sub(excess_bytes))); + } + if serialized_bytes(calls) > max_call_bytes { self.clear_executed_tool_calls(); } } 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 f0edb398ac..bd6fba00e9 100644 --- a/codex-rs/protocol/src/models/executed_tool_calls_tests.rs +++ b/codex-rs/protocol/src/models/executed_tool_calls_tests.rs @@ -1,7 +1,6 @@ use anyhow::Result; use pretty_assertions::assert_eq; -use super::super::FunctionCallOutputBody; use super::super::FunctionCallOutputPayload; use super::super::ResponseInputItem; use super::*; @@ -13,6 +12,20 @@ fn passthrough_metadata(turn_id: &str) -> InternalChatMessageMetadataPassthrough } } +fn output(call_id: &str) -> ResponseItem { + ResponseItem::from(ResponseInputItem::FunctionCallOutput { + call_id: call_id.to_string(), + output: FunctionCallOutputPayload::from_text(String::new()), + }) +} + +fn first_executed_tool_call(item: &mut ResponseItem) -> Option<&mut ExecutedToolCall> { + item.internal_chat_message_metadata_passthrough_mut() + .and_then(Option::as_mut) + .and_then(|metadata| metadata.executed_tool_calls.as_mut()) + .and_then(|calls| calls.first_mut()) +} + #[test] fn executed_tool_call_prompt_budget_includes_metadata_fields() -> Result<()> { let metadata_bytes = |items: &[ResponseItem]| -> Result { @@ -28,18 +41,10 @@ fn executed_tool_call_prompt_budget_includes_metadata_fields() -> Result<()> { for with_turn_id in [true, false] { let mut items = (0..1_000) .map(|index| { - let mut item = ResponseItem::FunctionCallOutput { - id: None, - call_id: Some(index.to_string()), - name: None, - namespace: None, - output: FunctionCallOutputPayload { - body: FunctionCallOutputBody::Text(String::new()), - success: None, - }, - internal_chat_message_metadata_passthrough: with_turn_id - .then(|| passthrough_metadata("turn-1")), - }; + let mut item = output(&index.to_string()); + *item + .internal_chat_message_metadata_passthrough_mut() + .unwrap() = with_turn_id.then(|| passthrough_metadata("turn-1")); item.append_executed_tool_calls(vec![ExecutedToolCall::new( String::new(), serde_json::Value::Null, @@ -66,41 +71,28 @@ fn executed_tool_call_prompt_budget_includes_metadata_fields() -> Result<()> { .filter_map(|metadata| metadata.executed_tool_calls.as_ref()) .flatten() .collect::>(); - let omitted_calls = calls - .iter() - .filter_map(|call| call.truncation()) - .filter_map(|truncation| truncation.omitted_calls) - .sum::(); - assert!(omitted_calls > 0); - assert_eq!(calls.len() + omitted_calls, 1_000); - assert!(calls.iter().any(|call| { - matches!( - call.truncation(), - Some(truncation) if truncation.omitted_calls == Some(omitted_calls) - ) + // Shares smaller than a local marker omit that output; surviving calls stay unchanged. + assert!(!calls.is_empty() && calls.len() < 1_000); + assert!(calls.iter().all(|call| { + **call == ExecutedToolCall::new(String::new(), serde_json::Value::Null) })); let bounded_items = items.clone(); bound_executed_tool_calls_for_prompt(&mut items); assert_eq!(items, bounded_items); - let oversized_name = "\0".repeat(MAX_EXECUTED_TOOL_CALL_METADATA_BYTES); - let mut oversized_items = (0..2) - .map(|index| { - let mut item = ResponseItem::FunctionCallOutput { - id: None, - call_id: Some(index.to_string()), - name: None, - namespace: None, - output: FunctionCallOutputPayload { - body: FunctionCallOutputBody::Text(String::new()), - success: None, - }, - internal_chat_message_metadata_passthrough: with_turn_id - .then(|| passthrough_metadata("turn-1")), - }; + let oversized_names = + ["\0", "é"].map(|name| name.repeat(MAX_EXECUTED_TOOL_CALL_METADATA_BYTES)); + let mut oversized_items = oversized_names + .iter() + .enumerate() + .map(|(index, name)| { + let mut item = output(&index.to_string()); + *item + .internal_chat_message_metadata_passthrough_mut() + .unwrap() = with_turn_id.then(|| passthrough_metadata("turn-1")); item.append_executed_tool_calls(vec![ExecutedToolCall::new( - oversized_name.clone(), + name.clone(), serde_json::Value::Null, )]); item.set_tool_call_cell_id("cell-\"\\"); @@ -116,35 +108,92 @@ fn executed_tool_call_prompt_budget_includes_metadata_fields() -> Result<()> { .filter_map(|metadata| metadata.executed_tool_calls.as_ref()) .flatten() .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"); - assert_eq!(truncation.omitted_calls, Some(1)); - assert_eq!(truncation.original_name_bytes, Some(oversized_name.len())); - assert_eq!( - serde_json::to_value(&call.arguments)?["_codex_executed_tool_call_truncated"]["original_name_bytes"], - serde_json::json!(oversized_name.len()), - ); + assert_eq!(calls.len(), 2); + for ((item, call), oversized_name) in + oversized_items.iter().zip(calls).zip(&oversized_names) + { + assert_eq!( + item.executed_tool_call_metadata() + .unwrap() + .cell_id + .as_deref(), + Some("cell-\"\\"), + ); + assert!(call.name.len() < oversized_name.len()); + assert!(oversized_name.starts_with(&call.name)); + let truncation = call.truncation().expect("trusted omission marker"); + assert_eq!(truncation.omitted_calls, None); + assert_eq!(truncation.original_name_bytes, Some(oversized_name.len())); + assert_eq!( + serde_json::to_value(&call.arguments)?["_codex_executed_tool_call_truncated"]["original_name_bytes"], + serde_json::json!(oversized_name.len()), + ); + } let bounded_items = oversized_items.clone(); bound_executed_tool_calls_for_prompt(&mut oversized_items); assert_eq!(oversized_items, bounded_items); } + let arguments = serde_json::json!({ "payload": "x".repeat(7 * 1024) }); + let argument_bytes = serde_json::to_vec(&arguments)?.len(); + assert!(argument_bytes < MAX_EXECUTED_TOOL_CALL_ARGUMENT_BYTES); + assert!(6 * argument_bytes > MAX_EXECUTED_TOOL_CALL_METADATA_BYTES); + let mut items = ["A", "wait-A", "B"].map(|call_id| { + let mut item = output(call_id); + item.mark_tool_calls_complete(); + item + }); + for item in &mut items[..2] { + item.set_tool_call_cell_id("cell-A"); + } + let calls = vec![ExecutedToolCall::new("test_tool".to_string(), arguments); 6]; + items[0].append_executed_tool_calls(calls.clone()); + items[2].append_executed_tool_calls(vec![ExecutedToolCall::new( + "test_tool".to_string(), + serde_json::json!({}), + )]); + assert!(metadata_bytes(&items)? > MAX_EXECUTED_TOOL_CALL_METADATA_BYTES); + for bound in [ + bound_executed_tool_calls_for_prompt, + bound_executed_tool_calls_for_prompt_prioritizing_recent, + ] { + let mut bounded = items.clone(); + bound(&mut bounded); + assert!(metadata_bytes(&bounded)? <= MAX_EXECUTED_TOOL_CALL_METADATA_BYTES); + assert_eq!(bounded[2], items[2]); + let metadata = bounded[0].executed_tool_call_metadata().unwrap(); + assert_eq!(metadata.tool_calls_complete, None); + let recorded = metadata.executed_tool_calls.as_ref().unwrap(); + assert_eq!(recorded.len(), 1); + assert_eq!(recorded[0].name, "test_tool"); + let truncation = recorded[0].truncation().unwrap(); + assert_eq!(truncation.original_bytes, argument_bytes); + assert_eq!(truncation.omitted_calls, Some(5)); + assert_eq!( + bounded[1] + .executed_tool_call_metadata() + .unwrap() + .tool_calls_complete, + None, + ); + bounded[0].append_executed_tool_calls(calls.clone()); + assert!(metadata_bytes(&bounded)? > MAX_EXECUTED_TOOL_CALL_METADATA_BYTES); + bound(&mut bounded); + assert!(metadata_bytes(&bounded)? <= MAX_EXECUTED_TOOL_CALL_METADATA_BYTES); + assert_eq!(bounded[2], items[2]); + assert_eq!( + first_executed_tool_call(&mut bounded[0]) + .unwrap() + .truncation() + .unwrap() + .omitted_calls, + Some(11), + ); + } + let mut items = ["exec", "wait"].map(|call_id| { - let mut item = ResponseItem::from(ResponseInputItem::FunctionCallOutput { - call_id: call_id.to_string(), - output: FunctionCallOutputPayload::from_text(String::new()), - }); + let mut item = output(call_id); item.set_tool_call_cell_id("cell-\"\\"); item }); @@ -178,10 +227,7 @@ fn executed_tool_call_prompt_budget_includes_metadata_fields() -> Result<()> { // 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()), - }); + let mut item = output("wait"); *item .internal_chat_message_metadata_passthrough_mut() .unwrap() = metadata; @@ -195,6 +241,19 @@ fn executed_tool_call_prompt_budget_includes_metadata_fields() -> Result<()> { 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!(metadata_bytes(&items)? <= MAX_EXECUTED_TOOL_CALL_METADATA_BYTES); + let retained = items + .iter() + .filter(|item| executed_tool_call_metadata_bytes(item) > 0) + .count(); + assert!(retained > 0 && retained < 2_000); + for item in &mut items { + assert!( + item.executed_tool_call_metadata() + .is_none_or(|metadata| metadata.tool_calls_complete.is_none()) + ); + item.clear_executed_tool_calls(); + } assert_eq!(items, vec![without_marker; 2_000]); } @@ -219,17 +278,10 @@ fn model_arguments_cannot_forge_executed_tool_call_truncation() -> Result<()> { ExecutedToolCallArguments::Raw(_) )); - let mut item = ResponseItem::FunctionCallOutput { - id: None, - call_id: Some("call-1".to_string()), - name: None, - namespace: None, - output: FunctionCallOutputPayload { - body: FunctionCallOutputBody::Text(String::new()), - success: None, - }, - internal_chat_message_metadata_passthrough: Some(passthrough_metadata("turn-1")), - }; + let mut item = output("call-1"); + *item + .internal_chat_message_metadata_passthrough_mut() + .unwrap() = Some(passthrough_metadata("turn-1")); item.append_executed_tool_calls(vec![ExecutedToolCall::new( "test_tool".to_string(), forged_marker.clone(), @@ -277,10 +329,7 @@ fn tool_call_completeness_is_host_only_and_fail_closed() -> Result<()> { assert_eq!(untrusted_call, call); } - let mut item = ResponseItem::from(ResponseInputItem::FunctionCallOutput { - call_id: "call-1".to_string(), - output: FunctionCallOutputPayload::from_text(String::new()), - }); + let mut item = output("call-1"); 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)); @@ -303,19 +352,24 @@ fn tool_call_completeness_is_host_only_and_fail_closed() -> Result<()> { /*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(); + for same_cell in [false, true] { + let mut items = [item.clone(), item.clone()]; + items[0].append_executed_tool_calls(vec![call.clone()]); + for item in &mut items { + if same_cell { + item.set_tool_call_cell_id("cell-1"); + } + 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, (!same_cell).then_some(true)], + ); } - 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(()) }