From 4d573c03fc5235e066a3800a2a1ec29aa79423f9 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Fri, 26 Jun 2026 19:29:40 -0700 Subject: [PATCH] core: persist repaired call outputs --- codex-rs/core/src/compact.rs | 2 + codex-rs/core/src/compact_remote.rs | 2 + codex-rs/core/src/compact_remote_v2.rs | 2 + codex-rs/core/src/context_manager/mod.rs | 1 + .../core/src/context_manager/normalize.rs | 92 ++++++++++++++----- codex-rs/core/src/prompt_debug.rs | 5 +- codex-rs/core/src/session/mod.rs | 75 +++++++++++++-- codex-rs/core/src/session/tests.rs | 92 +++++++++++++++++++ codex-rs/core/src/session/turn.rs | 15 +-- codex-rs/core/src/tasks/mod.rs | 5 + codex-rs/core/tests/suite/client.rs | 92 +++++++++++++++++++ 11 files changed, 338 insertions(+), 45 deletions(-) diff --git a/codex-rs/core/src/compact.rs b/codex-rs/core/src/compact.rs index d2a9fd5222..49996e2ddd 100644 --- a/codex-rs/core/src/compact.rs +++ b/codex-rs/core/src/compact.rs @@ -229,6 +229,8 @@ async fn run_compact_task_inner_impl( .await; let initial_input_for_turn: ResponseInputItem = ResponseInputItem::from(input); + sess.record_missing_call_outputs(turn_context.as_ref()) + .await; let mut history = sess.clone_history().await; history.record_items( &[initial_input_for_turn.into()], diff --git a/codex-rs/core/src/compact_remote.rs b/codex-rs/core/src/compact_remote.rs index 6b251562f7..b2c80bac61 100644 --- a/codex-rs/core/src/compact_remote.rs +++ b/codex-rs/core/src/compact_remote.rs @@ -193,6 +193,8 @@ async fn run_remote_compact_task_inner_impl( let compaction_item = TurnItem::ContextCompaction(context_compaction_item); sess.emit_turn_item_started(turn_context, &compaction_item) .await; + sess.record_missing_call_outputs(turn_context.as_ref()) + .await; let mut history = sess.clone_history().await; let base_instructions = sess.get_base_instructions().await; let (rewritten_outputs, estimated_deleted_tokens) = diff --git a/codex-rs/core/src/compact_remote_v2.rs b/codex-rs/core/src/compact_remote_v2.rs index 6e110045bb..0a72da80ff 100644 --- a/codex-rs/core/src/compact_remote_v2.rs +++ b/codex-rs/core/src/compact_remote_v2.rs @@ -203,6 +203,8 @@ async fn run_remote_compact_task_inner_impl( sess.emit_turn_item_started(turn_context, &compaction_item) .await; + sess.record_missing_call_outputs(turn_context.as_ref()) + .await; let mut history = sess.clone_history().await; let base_instructions = sess.get_base_instructions().await; let (rewritten_outputs, estimated_deleted_tokens) = diff --git a/codex-rs/core/src/context_manager/mod.rs b/codex-rs/core/src/context_manager/mod.rs index f811c8b373..e5d2a7906e 100644 --- a/codex-rs/core/src/context_manager/mod.rs +++ b/codex-rs/core/src/context_manager/mod.rs @@ -5,3 +5,4 @@ pub(crate) mod updates; pub(crate) use history::ContextManager; pub(crate) use history::is_user_turn_boundary; pub(crate) use history::truncate_function_output_payload; +pub(crate) use normalize::missing_call_outputs; diff --git a/codex-rs/core/src/context_manager/normalize.rs b/codex-rs/core/src/context_manager/normalize.rs index d11cffe4e9..fb3fa4f0cf 100644 --- a/codex-rs/core/src/context_manager/normalize.rs +++ b/codex-rs/core/src/context_manager/normalize.rs @@ -12,6 +12,55 @@ const IMAGE_CONTENT_OMITTED_PLACEHOLDER: &str = "image content omitted because you do not support image input"; pub(crate) fn ensure_call_outputs_present(items: &mut Vec) { + let missing_outputs_to_insert = missing_call_outputs_with_positions(items); + + for (idx, _) in &missing_outputs_to_insert { + match &items[*idx] { + ResponseItem::FunctionCall { call_id, .. } => { + info!("Function call output is missing for call id: {call_id}"); + } + ResponseItem::ToolSearchCall { + call_id: Some(call_id), + .. + } => { + info!("Tool search output is missing for call id: {call_id}"); + } + ResponseItem::CustomToolCall { call_id, .. } => { + error_or_panic(format!( + "Custom tool call output is missing for call id: {call_id}" + )); + } + ResponseItem::LocalShellCall { + call_id: Some(call_id), + .. + } => { + error_or_panic(format!( + "Local shell call output is missing for call id: {call_id}" + )); + } + _ => {} + } + } + + // Insert synthetic outputs in reverse index order to avoid re-indexing. + for (idx, output_item) in missing_outputs_to_insert.into_iter().rev() { + items.insert(idx + 1, output_item); + } +} + +/// Returns terminal outputs for calls that have no matching output in `items`. +/// +/// Callers that own canonical conversation history should record these items before exposing the +/// history to another model request. `ensure_call_outputs_present` remains as a defensive repair +/// for detached prompt snapshots. +pub(crate) fn missing_call_outputs(items: &[ResponseItem]) -> Vec { + missing_call_outputs_with_positions(items) + .into_iter() + .map(|(_, item)| item) + .collect() +} + +fn missing_call_outputs_with_positions(items: &[ResponseItem]) -> Vec<(usize, ResponseItem)> { let mut function_output_ids = HashSet::new(); let mut tool_search_output_ids = HashSet::new(); let mut custom_tool_output_ids = HashSet::new(); @@ -40,25 +89,27 @@ pub(crate) fn ensure_call_outputs_present(items: &mut Vec) { for (idx, item) in items.iter().enumerate() { match item { - ResponseItem::FunctionCall { call_id, .. } - if !function_output_ids.contains(call_id.as_str()) => - { - info!("Function call output is missing for call id: {call_id}"); + ResponseItem::FunctionCall { + call_id, + internal_chat_message_metadata_passthrough, + .. + } if !function_output_ids.contains(call_id.as_str()) => { missing_outputs_to_insert.push(( idx, ResponseItem::FunctionCallOutput { id: None, call_id: call_id.clone(), output: FunctionCallOutputPayload::from_text("aborted".to_string()), - internal_chat_message_metadata_passthrough: None, + internal_chat_message_metadata_passthrough: + internal_chat_message_metadata_passthrough.clone(), }, )); } ResponseItem::ToolSearchCall { call_id: Some(call_id), + internal_chat_message_metadata_passthrough, .. } if !tool_search_output_ids.contains(call_id.as_str()) => { - info!("Tool search output is missing for call id: {call_id}"); missing_outputs_to_insert.push(( idx, ResponseItem::ToolSearchOutput { @@ -67,16 +118,16 @@ pub(crate) fn ensure_call_outputs_present(items: &mut Vec) { status: "completed".to_string(), execution: "client".to_string(), tools: Vec::new(), - internal_chat_message_metadata_passthrough: None, + internal_chat_message_metadata_passthrough: + internal_chat_message_metadata_passthrough.clone(), }, )); } - ResponseItem::CustomToolCall { call_id, .. } - if !custom_tool_output_ids.contains(call_id.as_str()) => - { - error_or_panic(format!( - "Custom tool call output is missing for call id: {call_id}" - )); + ResponseItem::CustomToolCall { + call_id, + internal_chat_message_metadata_passthrough, + .. + } if !custom_tool_output_ids.contains(call_id.as_str()) => { missing_outputs_to_insert.push(( idx, ResponseItem::CustomToolCallOutput { @@ -84,25 +135,25 @@ pub(crate) fn ensure_call_outputs_present(items: &mut Vec) { call_id: call_id.clone(), name: None, output: FunctionCallOutputPayload::from_text("aborted".to_string()), - internal_chat_message_metadata_passthrough: None, + internal_chat_message_metadata_passthrough: + internal_chat_message_metadata_passthrough.clone(), }, )); } // LocalShellCall is represented in upstream streams by a FunctionCallOutput ResponseItem::LocalShellCall { call_id: Some(call_id), + internal_chat_message_metadata_passthrough, .. } if !function_output_ids.contains(call_id.as_str()) => { - error_or_panic(format!( - "Local shell call output is missing for call id: {call_id}" - )); missing_outputs_to_insert.push(( idx, ResponseItem::FunctionCallOutput { id: None, call_id: call_id.clone(), output: FunctionCallOutputPayload::from_text("aborted".to_string()), - internal_chat_message_metadata_passthrough: None, + internal_chat_message_metadata_passthrough: + internal_chat_message_metadata_passthrough.clone(), }, )); } @@ -115,10 +166,7 @@ pub(crate) fn ensure_call_outputs_present(items: &mut Vec) { custom_tool_output_ids, )); - // Insert synthetic outputs in reverse index order to avoid re-indexing. - for (idx, output_item) in missing_outputs_to_insert.into_iter().rev() { - items.insert(idx + 1, output_item); - } + missing_outputs_to_insert } pub(crate) fn remove_orphan_outputs(items: &mut Vec) { diff --git a/codex-rs/core/src/prompt_debug.rs b/codex-rs/core/src/prompt_debug.rs index cac9bb41ae..5685d38079 100644 --- a/codex-rs/core/src/prompt_debug.rs +++ b/codex-rs/core/src/prompt_debug.rs @@ -88,10 +88,7 @@ pub(crate) async fn build_prompt_input_from_session( .await; } - let prompt_input = sess - .clone_history() - .await - .for_prompt(&turn_context.model_info.input_modalities); + let prompt_input = sess.prompt_input_from_history(turn_context.as_ref()).await; let router = built_tools(sess, step_context.as_ref(), &CancellationToken::new()).await?; let base_instructions = sess.get_base_instructions().await; let prompt = build_prompt( diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index f45616eed8..d750c726b0 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -188,6 +188,7 @@ use crate::config::PermissionProfileState; use crate::config::StartedNetworkProxy; use crate::config::resolve_web_search_mode_for_turn; use crate::context_manager::ContextManager; +use crate::context_manager::missing_call_outputs; use crate::thread_rollout_truncation::initial_history_has_prior_user_turns; use codex_config::CONFIG_TOML_FILE; use codex_config::ConfigLayerSource; @@ -1361,6 +1362,13 @@ impl Session { } } } + + // Copy the source rollout before reconstruction appends any repairs for legacy + // unmatched calls, preserving their causal order in the forked rollout. + if !rollout_items.is_empty() { + self.persist_rollout_items(&rollout_items).await; + } + self.apply_rollout_reconstruction(&turn_context, &rollout_items) .await; @@ -1371,11 +1379,6 @@ impl Session { state.set_token_info(Some(info)); } - // If persisting, persist all rollout items as-is (the store filters). - if !rollout_items.is_empty() { - self.persist_rollout_items(&rollout_items).await; - } - // Forked threads should remain file-backed immediately after startup. self.ensure_rollout_materialized().await; @@ -1412,11 +1415,18 @@ impl Session { } = self .reconstruct_history_from_rollout(turn_context, rollout_items) .await; - // Keep the recorded rollout unchanged. Prepare its reconstructed history before - // installing it, so legacy images are processed once for this resume or fork and - // will be processed again if the rollout is reconstructed in a future session. - // This meets image resizing requirements without modifying persisted rollouts. + // Prepare reconstructed items before installing them. Legacy images are processed once + // for this reconstruction. Legacy unmatched calls are terminal, so close them in the + // canonical history and persist those repairs to keep their identities stable on future + // prompt builds and resumes. prepare_response_items(&mut history); + let mut repaired_outputs = missing_call_outputs(&history); + prepare_response_items(&mut repaired_outputs); + if turn_context.config.features.enabled(Feature::ItemIds) { + repaired_outputs = + Self::assign_missing_response_item_ids(Cow::Owned(repaired_outputs)).into_owned(); + } + history.extend(repaired_outputs.iter().cloned()); { let mut state = self.state.lock().await; state.replace_history(history, reference_context_item); @@ -1435,6 +1445,9 @@ impl Session { ); state.set_previous_turn_settings(previous_turn_settings.clone()); } + if !repaired_outputs.is_empty() { + self.persist_rollout_response_items(&repaired_outputs).await; + } let prefix_tokens = if matches!( turn_context.config.model_auto_compact_token_limit_scope, AutoCompactTokenLimitScope::BodyAfterPrefix @@ -3472,6 +3485,50 @@ impl Session { state.clone_history() } + /// Returns canonical history projected for a model request. + /// + /// Missing outputs should normally be closed when a task terminates or a rollout is + /// reconstructed. Repairing here is a final ownership boundary for callers that build a + /// standalone request from session history. + pub(crate) async fn prompt_input_from_history( + &self, + turn_context: &TurnContext, + ) -> Vec { + self.record_missing_call_outputs(turn_context).await; + self.clone_history() + .await + .for_prompt(&turn_context.model_info.input_modalities) + } + + pub(crate) async fn record_missing_call_outputs(&self, turn_context: &TurnContext) { + let repaired_outputs = { + let mut state = self.state.lock().await; + let repaired_outputs = missing_call_outputs(state.history.raw_items()); + if repaired_outputs.is_empty() { + return; + } + let repaired_outputs = self + .prepare_conversation_items_for_history(turn_context, &repaired_outputs) + .into_owned(); + state + .current_time_reminder + .note_recorded_items(&repaired_outputs); + state.record_items( + repaired_outputs.iter(), + turn_context.model_info.truncation_policy.into(), + ); + repaired_outputs + }; + + warn!( + repaired_output_count = repaired_outputs.len(), + "recording terminal outputs for unmatched calls" + ); + self.persist_rollout_response_items(&repaired_outputs).await; + self.send_raw_response_items(turn_context, &repaired_outputs) + .await; + } + pub(crate) async fn current_window_id(&self) -> String { let state = self.state.lock().await; let thread_id = self.thread_id; diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index d6336acd57..9c71e52365 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -9304,6 +9304,43 @@ struct NeverEndingTask { listen_to_cancellation_token: bool, } +struct UnmatchedFunctionCallTask; + +impl SessionTask for UnmatchedFunctionCallTask { + fn kind(&self) -> TaskKind { + TaskKind::Regular + } + + fn span_name(&self) -> &'static str { + "session_task.unmatched_function_call" + } + + async fn run( + self: Arc, + session: Arc, + ctx: Arc, + _input: Vec, + cancellation_token: CancellationToken, + ) -> SessionTaskResult { + session + .clone_session() + .record_conversation_items( + ctx.as_ref(), + &[ResponseItem::FunctionCall { + id: None, + name: "do_it".to_string(), + namespace: None, + arguments: "{}".to_string(), + call_id: "unmatched-call".to_string(), + internal_chat_message_metadata_passthrough: None, + }], + ) + .await; + cancellation_token.cancelled().await; + Ok(None) + } +} + impl SessionTask for NeverEndingTask { fn kind(&self) -> TaskKind { self.kind @@ -9611,6 +9648,61 @@ async fn abort_gracefully_emits_marker_before_turn_aborted() { assert!(rx.try_recv().is_err()); } +#[tokio::test] +async fn abort_records_terminal_output_for_unmatched_call() { + let (sess, tc, rx) = make_session_and_context_with_auth_and_config_and_rx( + CodexAuth::from_api_key("Test API Key"), + Vec::new(), + |config| { + let _ = config.features.enable(Feature::ItemIds); + }, + ) + .await; + sess.spawn_task(Arc::clone(&tc), Vec::new(), UnmatchedFunctionCallTask) + .await; + + timeout(StdDuration::from_secs(2), async { + loop { + let event = rx.recv().await.expect("event"); + if matches!( + event.msg, + EventMsg::RawResponseItem(ref raw) + if matches!( + raw.item, + ResponseItem::FunctionCall { ref call_id, .. } + if call_id == "unmatched-call" + ) + ) { + break; + } + } + }) + .await + .expect("timed out waiting for unmatched function call"); + sess.abort_all_tasks(TurnAbortReason::Interrupted).await; + + let history = sess.clone_history().await; + let outputs = history + .raw_items() + .iter() + .filter_map(|item| match item { + ResponseItem::FunctionCallOutput { + id, + call_id, + output, + .. + } if call_id == "unmatched-call" => Some((id.as_deref(), output)), + _ => None, + }) + .collect::>(); + assert_eq!(outputs.len(), 1); + assert!(outputs[0].0.is_some_and(|id| id.starts_with("fco_"))); + assert_eq!( + outputs[0].1, + &FunctionCallOutputPayload::from_text("aborted".to_string()) + ); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn task_finish_emits_turn_item_lifecycle_for_leftover_pending_user_input() { let (sess, tc, rx) = make_session_and_context_with_rx().await; diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 2e948b9c1e..fbb9d6aff6 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -268,13 +268,10 @@ pub(crate) async fn run_turn( } // Construct the input that we will send to the model. - let sampling_request_input: Vec = async { - sess.clone_history() - .await - .for_prompt(&turn_context.model_info.input_modalities) - } - .instrument(trace_span!("run_turn.prepare_sampling_request_input")) - .await; + let sampling_request_input: Vec = + async { sess.prompt_input_from_history(turn_context.as_ref()).await } + .instrument(trace_span!("run_turn.prepare_sampling_request_input")) + .await; let responses_metadata = turn_context.turn_metadata_state.to_responses_metadata( sess.installation_id.clone(), @@ -1104,9 +1101,7 @@ async fn run_sampling_request( let prompt_input = if let Some(input) = initial_input.take() { input } else { - sess.clone_history() - .await - .for_prompt(&turn_context.model_info.input_modalities) + sess.prompt_input_from_history(turn_context.as_ref()).await }; let prompt = build_prompt( prompt_input, diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index c6d9d2990c..415cc4d7ee 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -588,6 +588,8 @@ impl Session { let Some(turn_state) = turn_state else { return; }; + self.record_missing_call_outputs(turn_context.as_ref()) + .await; let pending_input = self .input_queue .take_pending_input_for_turn_state(turn_state.as_ref()) @@ -857,6 +859,9 @@ impl Session { .abort(session_ctx, Arc::clone(&task.turn_context)) .await; + self.record_missing_call_outputs(task.turn_context.as_ref()) + .await; + if reason == TurnAbortReason::Interrupted && let Some(marker) = interrupted_turn_history_marker( InterruptedTurnHistoryMarker::from_config_and_version( diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index 4df2bca1c1..4e09e8a141 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -363,6 +363,98 @@ async fn response_item_ids_persist_across_resume_and_preserve_server_ids() -> an Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn reconstructed_call_output_id_persists_across_prompt_builds_and_resumes() +-> anyhow::Result<()> { + let function_call_id = "missing-output-call"; + let thread_id = ThreadId::default(); + let rollout = vec![ + RolloutLine { + timestamp: "2024-01-01T00:00:00.000Z".to_string(), + item: RolloutItem::SessionMeta(SessionMetaLine { + meta: SessionMeta { + session_id: thread_id.into(), + id: thread_id, + parent_thread_id: None, + timestamp: "2024-01-01T00:00:00Z".to_string(), + cwd: ".".into(), + originator: "test_originator".to_string(), + cli_version: "test_version".to_string(), + model_provider: Some("test-provider".to_string()), + ..Default::default() + }, + git: None, + }), + }, + RolloutLine { + timestamp: "2024-01-01T00:00:01.000Z".to_string(), + item: RolloutItem::ResponseItem(ResponseItem::FunctionCall { + id: Some("fc_existing".to_string()), + name: "do_it".to_string(), + namespace: None, + arguments: "{}".to_string(), + call_id: function_call_id.to_string(), + internal_chat_message_metadata_passthrough: None, + }), + }, + ]; + let tmpdir = TempDir::new()?; + let session_path = tmpdir.path().join("normalized-call-output-item-id.jsonl"); + let mut file = std::fs::File::create(&session_path)?; + for line in rollout { + writeln!(file, "{}", serde_json::to_string(&line)?)?; + } + + let server = MockServer::start().await; + let response_mock = mount_sse_sequence( + &server, + vec![ + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]), + ], + ) + .await; + let codex_home = Arc::new(TempDir::new()?); + let mut builder = test_codex().with_config(|config| { + let _ = config.features.enable(Feature::ItemIds); + }); + let first = builder + .resume(&server, Arc::clone(&codex_home), session_path.clone()) + .await?; + + first.submit_turn("first resume").await?; + first.codex.submit(Op::Shutdown).await?; + wait_for_event(&first.codex, |event| { + matches!(event, EventMsg::ShutdownComplete) + }) + .await; + + builder = builder.with_config(|config| { + let _ = config.features.enable(Feature::ItemIds); + }); + let second = builder.resume(&server, codex_home, session_path).await?; + second.submit_turn("second resume").await?; + + let requests = response_mock.requests(); + assert_eq!(requests.len(), 2); + let first_output = requests[0].function_call_output(function_call_id); + let first_output_id = first_output + .get("id") + .and_then(serde_json::Value::as_str) + .expect("reconstructed output should have an item ID") + .to_string(); + assert!(first_output_id.starts_with("fco_")); + assert_eq!( + requests[1] + .function_call_output(function_call_id) + .get("id") + .and_then(serde_json::Value::as_str), + Some(first_output_id.as_str()) + ); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn response_item_ids_are_sent_for_all_remote_v2_compaction_requests() -> anyhow::Result<()> { let server = MockServer::start().await;