From 529bcb2fdf11025b17709dc88b05d055626646a6 Mon Sep 17 00:00:00 2001 From: Eric Traut Date: Mon, 14 Sep 2026 23:39:53 +0000 Subject: [PATCH] Preserve streamed answers and plans when turns terminate (#45549) ## Why Interrupted or failed turns can discard buffered answer or plan source, leaving incomplete math content missing from the transcript or unable to reflow correctly. ## What changed Flush and consolidate both answer and plan streams before `finalize_turn` clears streaming state. Reuse the same helper for normal completion, preserving received source for terminal resizing and raw output mode changes. ## Testing Add regression coverage for interrupted answers and plans, failed plans, and policy errors with unfinished math blocks, with and without a trailing newline. Verify source preservation and stable content across width and raw output mode changes, with narrow transcript snapshots. Update the usage activity test to verify that history insertion waits for stream consolidation. GitOrigin-RevId: 1ae0c192e2fee00dc0b5ee2ac07f3d5f2313ad1b --- codex-rs/tui/src/app/tests.rs | 2 + .../src/app/tests/math_interruption_tests.rs | 112 ++++++++++++++++++ ...nterruption_tests__answer_math_narrow.snap | 11 ++ ...ruption_tests__plan_error_math_narrow.snap | 15 +++ ..._interruption_tests__plan_math_narrow.snap | 15 +++ codex-rs/tui/src/chatwidget/streaming.rs | 19 +++ .../src/chatwidget/tests/slash_commands.rs | 2 + codex-rs/tui/src/chatwidget/turn_runtime.rs | 18 +-- 8 files changed, 178 insertions(+), 16 deletions(-) create mode 100644 codex-rs/tui/src/app/tests/math_interruption_tests.rs create mode 100644 codex-rs/tui/src/app/tests/snapshots/codex_tui__app__tests__math_interruption_tests__answer_math_narrow.snap create mode 100644 codex-rs/tui/src/app/tests/snapshots/codex_tui__app__tests__math_interruption_tests__plan_error_math_narrow.snap create mode 100644 codex-rs/tui/src/app/tests/snapshots/codex_tui__app__tests__math_interruption_tests__plan_math_narrow.snap diff --git a/codex-rs/tui/src/app/tests.rs b/codex-rs/tui/src/app/tests.rs index 7e78acd8c7..afa25de17c 100644 --- a/codex-rs/tui/src/app/tests.rs +++ b/codex-rs/tui/src/app/tests.rs @@ -2,6 +2,8 @@ #[path = "tests/daybreak_tests.rs"] mod daybreak_tests; +#[path = "tests/math_interruption_tests.rs"] +mod math_interruption_tests; #[path = "tests/advanced_reasoning_tests.rs"] mod advanced_reasoning_tests; diff --git a/codex-rs/tui/src/app/tests/math_interruption_tests.rs b/codex-rs/tui/src/app/tests/math_interruption_tests.rs new file mode 100644 index 0000000000..1e33c0809f --- /dev/null +++ b/codex-rs/tui/src/app/tests/math_interruption_tests.rs @@ -0,0 +1,112 @@ +use super::*; +use pretty_assertions::assert_eq; + +#[tokio::test] +async fn unicode_math_termination_preserves_source_through_mode_and_resize() -> Result<()> { + for (kind, status) in [ + ("answer", TurnStatus::Interrupted), + ("plan", TurnStatus::Interrupted), + ("plan_error", TurnStatus::Failed), + ("plan_policy", TurnStatus::Failed), + ] { + for ending in ["\n", ""] { + let (mut app, mut events, _op_rx) = make_test_app_with_channels().await; + let mut app_server = start_config_write_test_app_server(&app).await?; + let mut tui = crate::tui::test_support::make_test_tui()?; + let thread_id = ThreadId::new(); + let plan = kind != "answer"; + if plan { + app.chat_widget + .set_feature_enabled(Feature::CollaborationModes, /*enabled*/ true); + app.chat_widget + .set_collaboration_mask(CollaborationModeMask { + name: "Plan".into(), + mode: Some(ModeKind::Plan), + model: None, + reasoning_effort: None, + developer_instructions: None, + }); + } + app.chat_widget.handle_server_notification( + turn_started_notification(thread_id, "math-turn"), + /*replay_kind*/ None, + ); + let source = format!("Intro.\n\n$$\n\\frac{{a+b+c+d+e+f}}{{g+h}}{ending}"); + for delta in ["Intro.\n\n", &source[8..]] { + let notification = if plan { + ServerNotification::PlanDelta( + codex_app_server_protocol::PlanDeltaNotification { + thread_id: thread_id.to_string(), + turn_id: "math-turn".into(), + item_id: "math-message".into(), + delta: delta.into(), + }, + ) + } else { + agent_message_delta_notification(thread_id, "math-turn", "math-message", delta) + }; + app.chat_widget + .handle_server_notification(notification, /*replay_kind*/ None); + for _ in 0..10 { + app.chat_widget.on_commit_tick(); + } + } + let mut completed = turn_completed_notification(thread_id, "math-turn", status.clone()); + if matches!(kind, "plan_error" | "plan_policy") + && let ServerNotification::TurnCompleted(notification) = &mut completed + { + notification.turn.error = Some(AppServerTurnError { + message: if kind == "plan_policy" { + r#"{"error":{"code":"bio_policy"}}"#.into() + } else { + "Server error".into() + }, + codex_error_info: None, + additional_details: None, + misalignment: None, + }); + } + app.chat_widget + .handle_server_notification(completed, /*replay_kind*/ None); + let mut saved = Vec::new(); + while let Ok(event) = events.try_recv() { + if let AppEvent::ConsolidateAgentMessage { source, .. } + | AppEvent::ConsolidateProposedPlan(source) = &event + { + saved.push(source.trim_end_matches('\n').to_owned()); + } + if matches!( + event, + AppEvent::InsertHistoryCell(_) + | AppEvent::ConsolidateAgentMessage { .. } + | AppEvent::ConsolidateProposedPlan(_) + ) { + app.handle_event(&mut tui, &mut app_server, event).await?; + } + } + assert_eq!(saved, vec![source.trim_end_matches('\n').to_owned()]); + let mut expected = [None, None]; + for raw in [false, true, false] { + app.chat_widget.set_raw_output_mode(raw); + for width in [80, 20, 80] { + let rendered = app + .render_transcript_lines_for_reflow(width) + .lines + .iter() + .map(rendered_line_text) + .collect::>() + .join("\n"); + let content = rendered.split_whitespace().collect::(); + assert_eq!( + expected[usize::from(raw)].get_or_insert_with(|| content.clone()), + &content + ); + if !raw && width == 20 && !ending.is_empty() && kind != "plan_policy" { + insta::assert_snapshot!(format!("{kind}_math_narrow"), rendered); + } + } + } + } + } + Ok(()) +} diff --git a/codex-rs/tui/src/app/tests/snapshots/codex_tui__app__tests__math_interruption_tests__answer_math_narrow.snap b/codex-rs/tui/src/app/tests/snapshots/codex_tui__app__tests__math_interruption_tests__answer_math_narrow.snap new file mode 100644 index 0000000000..7cf1280f79 --- /dev/null +++ b/codex-rs/tui/src/app/tests/snapshots/codex_tui__app__tests__math_interruption_tests__answer_math_narrow.snap @@ -0,0 +1,11 @@ +--- +source: tui/src/app/tests/math_interruption_tests.rs +expression: rendered +--- +• Intro. + + $$ + \frac{a+b+c+d+e+f} + {g+h} + +■ Conversation interrupted - tell the model what to do differently. Something went wrong? Hit `/feedback` to report the issue. diff --git a/codex-rs/tui/src/app/tests/snapshots/codex_tui__app__tests__math_interruption_tests__plan_error_math_narrow.snap b/codex-rs/tui/src/app/tests/snapshots/codex_tui__app__tests__math_interruption_tests__plan_error_math_narrow.snap new file mode 100644 index 0000000000..844b078cf6 --- /dev/null +++ b/codex-rs/tui/src/app/tests/snapshots/codex_tui__app__tests__math_interruption_tests__plan_error_math_narrow.snap @@ -0,0 +1,15 @@ +--- +source: tui/src/app/tests/math_interruption_tests.rs +expression: rendered +--- +• Proposed Plan + + + Intro. + + $$ + \frac{a+b+c+d+e+ + f}{g+h} + + +■ Server error diff --git a/codex-rs/tui/src/app/tests/snapshots/codex_tui__app__tests__math_interruption_tests__plan_math_narrow.snap b/codex-rs/tui/src/app/tests/snapshots/codex_tui__app__tests__math_interruption_tests__plan_math_narrow.snap new file mode 100644 index 0000000000..7ca58045f1 --- /dev/null +++ b/codex-rs/tui/src/app/tests/snapshots/codex_tui__app__tests__math_interruption_tests__plan_math_narrow.snap @@ -0,0 +1,15 @@ +--- +source: tui/src/app/tests/math_interruption_tests.rs +expression: rendered +--- +• Proposed Plan + + + Intro. + + $$ + \frac{a+b+c+d+e+ + f}{g+h} + + +■ Conversation interrupted - tell the model what to do differently. Something went wrong? Hit `/feedback` to report the issue. diff --git a/codex-rs/tui/src/chatwidget/streaming.rs b/codex-rs/tui/src/chatwidget/streaming.rs index 30204cfa8c..5c1cc08f56 100644 --- a/codex-rs/tui/src/chatwidget/streaming.rs +++ b/codex-rs/tui/src/chatwidget/streaming.rs @@ -60,6 +60,25 @@ impl ChatWidget { } } + /// Preserve received answer and plan source before ordinary turn termination. + pub(super) fn flush_answer_and_plan_streams(&mut self) { + self.flush_answer_stream_with_separator(); + if let Some(mut controller) = self.plan_stream_controller.take() { + let had_live_tail = controller.has_live_tail(); + self.clear_active_stream_tail(); + let (cell, source) = controller.finalize(); + if !had_live_tail && let Some(cell) = cell { + self.add_boxed_history(cell); + } + if let Some(source) = source { + self.note_stream_consolidation_queued(); + self.app_event_tx + .send(AppEvent::ConsolidateProposedPlan(source)); + } + self.request_pending_usage_output_insertion_after_stream_shutdown(); + } + } + pub(super) fn flush_answer_stream_with_separator(&mut self) { self.flush_answer_stream(/*completed_message*/ None); } diff --git a/codex-rs/tui/src/chatwidget/tests/slash_commands.rs b/codex-rs/tui/src/chatwidget/tests/slash_commands.rs index dbda4eb067..557002b211 100644 --- a/codex-rs/tui/src/chatwidget/tests/slash_commands.rs +++ b/codex-rs/tui/src/chatwidget/tests/slash_commands.rs @@ -1574,6 +1574,8 @@ async fn completed_token_activity_refresh_waits_for_active_stream() { ); chat.finalize_turn(); + assert!(chat.usage_history_insertion_blocked()); + chat.note_stream_consolidation_completed(); assert!(!chat.usage_history_insertion_blocked()); assert!( std::iter::from_fn(|| rx.try_recv().ok()) diff --git a/codex-rs/tui/src/chatwidget/turn_runtime.rs b/codex-rs/tui/src/chatwidget/turn_runtime.rs index 388af9dffc..6d7c32cecc 100644 --- a/codex-rs/tui/src/chatwidget/turn_runtime.rs +++ b/codex-rs/tui/src/chatwidget/turn_runtime.rs @@ -141,21 +141,7 @@ impl ChatWidget { .unwrap_or_default(); self.transcript.saw_copy_source_this_turn = false; // If a stream is currently active, finalize it. - self.flush_answer_stream_with_separator(); - if let Some(mut controller) = self.plan_stream_controller.take() { - let had_live_tail = controller.has_live_tail(); - self.clear_active_stream_tail(); - let (cell, source) = controller.finalize(); - if !had_live_tail && let Some(cell) = cell { - self.add_boxed_history(cell); - } - if let Some(source) = source { - self.note_stream_consolidation_queued(); - self.app_event_tx - .send(AppEvent::ConsolidateProposedPlan(source)); - } - self.request_pending_usage_output_insertion_after_stream_shutdown(); - } + self.flush_answer_and_plan_streams(); self.flush_unified_exec_wait_streak(); if !from_replay { self.collect_runtime_metrics_delta(); @@ -322,6 +308,7 @@ impl ChatWidget { /// This does not clear MCP startup tracking, because MCP startup can overlap with turn cleanup /// and should continue to drive the bottom-pane running indicator while it is in progress. pub(super) fn finalize_turn(&mut self) { + self.flush_answer_and_plan_streams(); if self.status_state.reasoning_resume_turn_id.is_some() { self.on_agent_reasoning_final(); } @@ -377,7 +364,6 @@ impl ChatWidget { fn on_error(&mut self, message: String) { self.input_queue.submit_pending_steers_after_interrupt = false; - self.flush_answer_stream_with_separator(); self.finalize_turn(); self.add_to_history(history_cell::new_error_event(message)); self.set_ambient_pet_notification(