diff --git a/codex-rs/tui/src/app/replay_filter.rs b/codex-rs/tui/src/app/replay_filter.rs index df5a90b8f2..707192cb06 100644 --- a/codex-rs/tui/src/app/replay_filter.rs +++ b/codex-rs/tui/src/app/replay_filter.rs @@ -2,6 +2,8 @@ use codex_app_server_protocol::ServerNotification; use codex_app_server_protocol::ServerRequest; +use codex_app_server_protocol::ThreadItem; +use std::collections::HashSet; use super::ThreadBufferedEvent; use super::ThreadEventSnapshot; @@ -35,3 +37,41 @@ pub(super) fn event_is_notice(event: &ThreadBufferedEvent) -> bool { ) ) } + +/// A completed item's full text replaces its earlier streaming deltas during replay. +/// Keep deltas without a later completion so an in-progress or truncated stream still renders. +/// Other events are barriers: streaming text may need to flush before a tool or prompt. +/// This only changes the replay snapshot; the live notification store remains untouched. +pub(super) fn omit_completed_agent_deltas(events: &mut Vec) { + let mut completed = HashSet::new(); + events.reverse(); + events.retain(|event| { + if let ThreadBufferedEvent::Notification(notification) = event { + match notification.as_ref() { + ServerNotification::ItemCompleted(notification) => { + if let ThreadItem::AgentMessage { id, .. } = ¬ification.item { + completed.insert(( + notification.thread_id.clone(), + notification.turn_id.clone(), + id.clone(), + )); + } else { + completed.clear(); + } + } + ServerNotification::AgentMessageDelta(notification) => { + return !completed.contains(&( + notification.thread_id.clone(), + notification.turn_id.clone(), + notification.item_id.clone(), + )); + } + _ => completed.clear(), + } + } else { + completed.clear(); + } + true + }); + events.reverse(); +} diff --git a/codex-rs/tui/src/app/tests.rs b/codex-rs/tui/src/app/tests.rs index 971adc1f08..e57a84f6e0 100644 --- a/codex-rs/tui/src/app/tests.rs +++ b/codex-rs/tui/src/app/tests.rs @@ -12,6 +12,8 @@ mod backend_banner_recovery_tests; mod backend_banner_startup_tests; #[path = "tests/background_exit_tests.rs"] mod background_exit_tests; +#[path = "tests/buffered_replay.rs"] +mod buffered_replay; #[path = "tests/connector_policy.rs"] mod connector_policy; #[path = "tests/disconnect_tests.rs"] diff --git a/codex-rs/tui/src/app/tests/buffered_replay.rs b/codex-rs/tui/src/app/tests/buffered_replay.rs new file mode 100644 index 0000000000..99e11149d5 --- /dev/null +++ b/codex-rs/tui/src/app/tests/buffered_replay.rs @@ -0,0 +1,121 @@ +//! Real buffered notifications must use finalized rendering without losing live stream tails. + +use super::*; +use codex_app_server_protocol::ItemCompletedNotification; +use pretty_assertions::assert_eq; + +fn delta(thread: &str, turn: &str, item: &str) -> ServerNotification { + ServerNotification::AgentMessageDelta(AgentMessageDeltaNotification { + thread_id: thread.into(), + turn_id: turn.into(), + item_id: item.into(), + delta: "Partial **answer**\n".into(), + }) +} + +fn completed(thread: &str) -> ServerNotification { + ServerNotification::ItemCompleted(ItemCompletedNotification { + thread_id: thread.into(), + turn_id: "turn".into(), + completed_at_ms: 0, + item: ThreadItem::AgentMessage { + id: "answer".into(), + text: "Final **answer**".into(), + phase: None, + memory_citation: None, + delivery: None, + }, + }) +} + +#[test] +fn buffered_replay_keeps_unfinished_and_unrelated_deltas_in_order() { + let completion = completed("thread"); + let unfinished = delta("thread", "turn", "unfinished"); + let other_turn = delta("thread", "other-turn", "answer"); + let other_thread = delta("other-thread", "turn", "answer"); + let later_delta = delta("thread", "turn", "answer"); + let mut store = ThreadEventStore::new(/*capacity*/ 16); + for event in [ + delta("thread", "turn", "answer"), + unfinished.clone(), + other_turn.clone(), + other_thread.clone(), + completion.clone(), + later_delta.clone(), + ] { + store.push_notification(event); + } + let mut events = store.snapshot().events; + replay_filter::omit_completed_agent_deltas(&mut events); + let actual = events + .into_iter() + .map(|event| match event { + ThreadBufferedEvent::Notification(notification) => *notification, + other => panic!("unexpected event: {other:?}"), + }) + .collect::>(); + assert_eq!( + serde_json::to_value(actual).unwrap(), + serde_json::to_value(vec![ + unfinished, + other_turn, + other_thread, + completion, + later_delta + ]) + .unwrap() + ); + assert_eq!(store.snapshot().events.len(), 6); +} + +#[tokio::test] +async fn buffered_replay_renders_completed_text_without_streaming_again() { + let (mut app, mut events, _ops) = make_test_app_with_channels().await; + let mut store = ThreadEventStore::new(/*capacity*/ 16); + store.push_notification(delta("thread", "turn", "answer")); + store.push_notification(completed("thread")); + app.replay_thread_snapshot(store.snapshot(), /*resume_restored_queue*/ false); + let mut lines = Vec::new(); + while let Ok(event) = events.try_recv() { + match event { + AppEvent::InsertHistoryCell(cell) => lines.extend(cell.display_lines(/*width*/ 80)), + AppEvent::StartCommitAnimation | AppEvent::ConsolidateAgentMessage { .. } => { + panic!("completed replay must not reconstruct a stream") + } + _ => {} + } + } + let text = lines + .iter() + .map(ToString::to_string) + .collect::>() + .join("\n"); + insta::assert_snapshot!(text.trim(), @" + • Final answer + "); +} + +#[test] +fn buffered_replay_preserves_text_before_interactive_requests() { + let mut store = ThreadEventStore::new(/*capacity*/ 16); + store.push_notification(delta("thread", "turn", "answer")); + store.push_request( + codex_app_server_protocol::ServerRequest::ToolRequestUserInput { + request_id: codex_app_server_protocol::RequestId::Integer(1), + params: codex_app_server_protocol::ToolRequestUserInputParams { + thread_id: "thread".into(), + turn_id: "turn".into(), + item_id: "tool".into(), + questions: Vec::new(), + is_blocking: true, + auto_resolution_ms: None, + }, + }, + ); + store.push_notification(completed("thread")); + let mut events = store.snapshot().events; + let before = format!("{events:?}"); + replay_filter::omit_completed_agent_deltas(&mut events); + assert_eq!(format!("{events:?}"), before); +} diff --git a/codex-rs/tui/src/app/thread_routing.rs b/codex-rs/tui/src/app/thread_routing.rs index 5ba12650d3..98b443f264 100644 --- a/codex-rs/tui/src/app/thread_routing.rs +++ b/codex-rs/tui/src/app/thread_routing.rs @@ -1555,6 +1555,7 @@ impl App { mut snapshot: ThreadEventSnapshot, resume_restored_queue: bool, ) { + replay_filter::omit_completed_agent_deltas(&mut snapshot.events); let request_changes = snapshot .events .iter() diff --git a/codex-rs/tui/src/chatwidget/replay.rs b/codex-rs/tui/src/chatwidget/replay.rs index 1d085e6e73..1660b702e9 100644 --- a/codex-rs/tui/src/chatwidget/replay.rs +++ b/codex-rs/tui/src/chatwidget/replay.rs @@ -6,6 +6,24 @@ use super::*; impl ChatWidget { + /// Flush prior activity and preserve its separator before live or replayed assistant text. + pub(super) fn prepare_assistant_message(&mut self) { + // Before starting an agent stream, flush any active exec cell group. + self.flush_unified_exec_wait_streak(); + self.flush_active_cell(); + // If the previous turn inserted non-stream history (exec output, patch status, MCP + // calls), render a separator before starting the next streamed assistant message. + if self.transcript.needs_final_message_separator && self.transcript.had_work_activity { + self.add_to_history(history_cell::FinalMessageSeparator::new( + /*elapsed_seconds*/ None, /*runtime_metrics*/ None, + )); + self.transcript.needs_final_message_separator = false; + } else if self.transcript.needs_final_message_separator { + // Reset the flag even if we don't show separator (no work was done) + self.transcript.needs_final_message_separator = false; + } + } + /// Replay a subset of initial events into the UI to seed the transcript when /// resuming an existing session. This approximates the live event flow and /// is intentionally conservative: only safe-to-replay items are rendered to diff --git a/codex-rs/tui/src/chatwidget/streaming.rs b/codex-rs/tui/src/chatwidget/streaming.rs index b596ac3352..f5b73e7f35 100644 --- a/codex-rs/tui/src/chatwidget/streaming.rs +++ b/codex-rs/tui/src/chatwidget/streaming.rs @@ -330,7 +330,28 @@ impl ChatWidget { } } let parsed = parse_assistant_markdown(&message, self.config.cwd.as_path()); - self.finalize_completed_assistant_message(Some(parsed.visible_markdown.as_str())); + if from_replay && self.stream_controller.is_none() && !parsed.visible_markdown.is_empty() { + self.prepare_assistant_message(); + self.mark_safety_buffering_agent_message_started(); + self.bottom_pane.hide_status_indicator(); + let context = self.thread_id.and_then(|thread_id| { + crate::inline_visualization::InlineVisualizationContext::from_config( + &self.config, + thread_id, + ) + }); + self.add_to_history( + history_cell::AgentMarkdownCell::new_with_inline_visualizations( + parsed.visible_markdown.clone(), + self.config.cwd.as_path(), + context, + ), + ); + self.handle_stream_finished(); + self.request_redraw(); + } else { + self.finalize_completed_assistant_message(Some(parsed.visible_markdown.as_str())); + } if matches!(item.phase, Some(MessagePhase::FinalAnswer) | None) && !parsed.visible_markdown.is_empty() { @@ -453,20 +474,7 @@ impl ChatWidget { self.mark_safety_buffering_agent_message_started(); } if self.stream_controller.is_none() { - // Before starting an agent stream, flush any active exec cell group. - self.flush_unified_exec_wait_streak(); - self.flush_active_cell(); - // If the previous turn inserted non-stream history (exec output, patch status, MCP - // calls), render a separator before starting the next streamed assistant message. - if self.transcript.needs_final_message_separator && self.transcript.had_work_activity { - self.add_to_history(history_cell::FinalMessageSeparator::new( - /*elapsed_seconds*/ None, /*runtime_metrics*/ None, - )); - self.transcript.needs_final_message_separator = false; - } else if self.transcript.needs_final_message_separator { - // Reset the flag even if we don't show separator (no work was done) - self.transcript.needs_final_message_separator = false; - } + self.prepare_assistant_message(); let inline_visualization_context = self.thread_id.and_then(|thread_id| { crate::inline_visualization::InlineVisualizationContext::from_config( &self.config, diff --git a/codex-rs/tui/src/chatwidget/tests.rs b/codex-rs/tui/src/chatwidget/tests.rs index cd27b0fcf0..33efb56033 100644 --- a/codex-rs/tui/src/chatwidget/tests.rs +++ b/codex-rs/tui/src/chatwidget/tests.rs @@ -258,6 +258,8 @@ mod plugin_catalog; mod popups_and_settings; #[path = "tests/rate_limit_recovery_tests.rs"] mod rate_limit_recovery_tests; +#[path = "tests/replay_render_tests.rs"] +mod replay_render_tests; mod review_mode; mod side; mod slash_commands; diff --git a/codex-rs/tui/src/chatwidget/tests/replay_render_tests.rs b/codex-rs/tui/src/chatwidget/tests/replay_render_tests.rs new file mode 100644 index 0000000000..b155e5a585 --- /dev/null +++ b/codex-rs/tui/src/chatwidget/tests/replay_render_tests.rs @@ -0,0 +1,39 @@ +//! Completed replay messages preserve order, formatting, and composer state. + +use super::*; +use pretty_assertions::assert_eq; + +#[tokio::test] +async fn completed_replay_preserves_messages_and_draft_across_reconstruction() { + let mut outputs = Vec::new(); + for _ in 0..2 { + let (mut chat, mut rx, _ops) = make_chatwidget_manual(/*model_override*/ None).await; + chat.note_rendered_width(/*width*/ 80); + chat.handle_paste("Unsent draft".to_string()); + replay_user_message_text(&mut chat, "user", "Question", ReplayKind::ThreadSnapshot); + replay_agent_message( + &mut chat, + "first", + "First **answer**", + ReplayKind::ThreadSnapshot, + ); + replay_agent_message( + &mut chat, + "second", + "Second `answer`", + ReplayKind::ThreadSnapshot, + ); + let cells = drain_insert_history(&mut rx); + outputs.push(cells.into_iter().flatten().collect::>()); + assert_eq!(chat.composer_text_with_pending(), "Unsent draft"); + } + assert_eq!(outputs[0], outputs[1]); + insta::assert_snapshot!(outputs[0].iter().map(ToString::to_string).collect::>().join("\n").trim(), @" + › Question + + + • First answer + + • Second answer + "); +}