diff --git a/codex-rs/tui/src/chatwidget/tests.rs b/codex-rs/tui/src/chatwidget/tests.rs index bfae9581b2..dd4891d07d 100644 --- a/codex-rs/tui/src/chatwidget/tests.rs +++ b/codex-rs/tui/src/chatwidget/tests.rs @@ -7132,6 +7132,98 @@ async fn deltas_then_same_final_message_are_rendered_snapshot() { assert_snapshot!(combined); } +#[tokio::test] +async fn streamed_message_hides_trailing_memory_footer() { + let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(None).await; + + chat.handle_codex_event(Event { + id: "s1".into(), + msg: EventMsg::TurnStarted(TurnStartedEvent { + turn_id: "turn-1".to_string(), + model_context_window: None, + collaboration_mode_kind: ModeKind::Default, + }), + }); + chat.handle_codex_event(Event { + id: "s1".into(), + msg: EventMsg::AgentMessageDelta(AgentMessageDeltaEvent { + delta: "Final answer line\n\nMemory".into(), + }), + }); + chat.handle_codex_event(Event { + id: "s1".into(), + msg: EventMsg::AgentMessageDelta(AgentMessageDeltaEvent { + delta: " used: `notes.md:1-2`, `MEMORY.md:10-20`\n".into(), + }), + }); + chat.handle_codex_event(Event { + id: "s1".into(), + msg: EventMsg::TurnComplete(TurnCompleteEvent { + turn_id: "turn-1".to_string(), + last_agent_message: None, + }), + }); + + let cells = drain_insert_history(&mut rx); + let combined = cells + .iter() + .map(|lines| lines_to_single_string(lines)) + .collect::(); + assert!( + combined.contains("Final answer line"), + "expected assistant content to render" + ); + assert!( + !combined.contains("Memory used:"), + "memory citation footer should be stripped from TUI rendering: {combined}" + ); +} + +#[tokio::test] +async fn replayed_agent_message_hides_trailing_memory_footer() { + let (mut chat, mut rx, _ops) = make_chatwidget_manual(None).await; + + let conversation_id = ThreadId::new(); + let rollout_file = NamedTempFile::new().unwrap(); + let configured = codex_core::protocol::SessionConfiguredEvent { + session_id: conversation_id, + forked_from_id: None, + thread_name: None, + model: "test-model".to_string(), + model_provider_id: "test-provider".to_string(), + approval_policy: AskForApproval::Never, + sandbox_policy: SandboxPolicy::new_read_only_policy(), + cwd: PathBuf::from("/home/user/project"), + reasoning_effort: Some(ReasoningEffortConfig::default()), + history_log_id: 0, + history_entry_count: 0, + initial_messages: Some(vec![EventMsg::AgentMessage(AgentMessageEvent { + message: "Replay answer\nMemory used: notes.md:1-2".to_string(), + })]), + network_proxy: None, + rollout_path: Some(rollout_file.path().to_path_buf()), + }; + + chat.handle_codex_event(Event { + id: "initial".into(), + msg: EventMsg::SessionConfigured(configured), + }); + + let cells = drain_insert_history(&mut rx); + let combined = cells + .iter() + .map(|lines| lines_to_single_string(lines)) + .collect::(); + assert!( + combined.contains("Replay answer"), + "expected replayed answer content" + ); + assert!( + !combined.contains("Memory used:"), + "memory citation footer should be stripped from replay rendering: {combined}" + ); +} + // Combined visual snapshot using vt100 for history + direct buffer overlay for UI. // This renders the final visual as seen in a terminal: history above, then a blank line, // then the exec block, another blank line, the status line, a blank line, and the composer. diff --git a/codex-rs/tui/src/markdown_stream.rs b/codex-rs/tui/src/markdown_stream.rs index 6ac457eee2..47fa6a4c48 100644 --- a/codex-rs/tui/src/markdown_stream.rs +++ b/codex-rs/tui/src/markdown_stream.rs @@ -2,6 +2,61 @@ use ratatui::text::Line; use crate::markdown; +const MEMORY_CITATION_PREFIX: &str = "Memory used:"; + +fn strip_trailing_memory_citation_section(input: &str) -> String { + let Some(citation_start) = find_trailing_memory_citation_start(input) else { + return input.to_string(); + }; + trim_trailing_blank_lines(&input[..citation_start]).to_string() +} + +fn find_trailing_memory_citation_start(input: &str) -> Option { + let mut line_start = 0usize; + let mut candidate_start = None; + + loop { + let rest = &input[line_start..]; + let next_newline = rest.find('\n'); + let line_end = next_newline.map_or(input.len(), |offset| line_start + offset); + let line = &input[line_start..line_end]; + let line = line.strip_suffix('\r').unwrap_or(line); + if line.trim_start().starts_with(MEMORY_CITATION_PREFIX) { + candidate_start = Some(line_start); + } + + match next_newline { + Some(_) => line_start = line_end + 1, + None => break, + } + } + + let citation_start = candidate_start?; + let citation_line_end = input[citation_start..] + .find('\n') + .map_or(input.len(), |offset| citation_start + offset); + let after_citation_line = &input[citation_line_end..]; + after_citation_line + .chars() + .all(char::is_whitespace) + .then_some(citation_start) +} + +fn trim_trailing_blank_lines(input: &str) -> &str { + let mut end = input.len(); + while end > 0 { + let line_start = input[..end].rfind('\n').map_or(0, |idx| idx + 1); + let line = &input[line_start..end]; + let line = line.strip_suffix('\r').unwrap_or(line); + if line.trim().is_empty() { + end = line_start.saturating_sub(1); + } else { + break; + } + } + &input[..end] +} + /// Newline-gated accumulator that renders markdown and commits only fully /// completed logical lines. pub(crate) struct MarkdownStreamCollector { @@ -40,6 +95,7 @@ impl MarkdownStreamCollector { } else { return Vec::new(); }; + let source = strip_trailing_memory_citation_section(&source); let mut rendered: Vec> = Vec::new(); markdown::append_markdown(&source, self.width, &mut rendered); let mut complete_line_count = rendered.len(); @@ -72,6 +128,7 @@ impl MarkdownStreamCollector { if !source.ends_with('\n') { source.push('\n'); } + source = strip_trailing_memory_citation_section(&source); tracing::debug!( raw_len = raw_buffer.len(), source_len = source.len(), @@ -118,8 +175,33 @@ pub(crate) fn simulate_stream_markdown_for_tests( #[cfg(test)] mod tests { use super::*; + use pretty_assertions::assert_eq; use ratatui::style::Color; + #[test] + fn strip_trailing_memory_citation_canonical_format() { + let input = "Answer body\n\nMemory used: `notes.md:1-2`, `MEMORY.md:10-20`\n"; + assert_eq!(strip_trailing_memory_citation_section(input), "Answer body"); + } + + #[test] + fn strip_trailing_memory_citation_plain_text_with_blank_tail() { + let input = "Answer body\n\nMemory used: notes.md:1-2, MEMORY.md:10-20\n\n"; + assert_eq!(strip_trailing_memory_citation_section(input), "Answer body"); + } + + #[test] + fn preserve_memory_used_when_not_trailing_footer() { + let input = "Answer body\nMemory used: notes.md:1-2\nAdditional details\n"; + assert_eq!(strip_trailing_memory_citation_section(input), input); + } + + #[test] + fn no_op_when_memory_footer_absent() { + let input = "Answer body\nNo footer here\n"; + assert_eq!(strip_trailing_memory_citation_section(input), input); + } + #[tokio::test] async fn no_commit_until_newline() { let mut c = super::MarkdownStreamCollector::new(None); @@ -249,6 +331,16 @@ mod tests { } } + #[tokio::test] + async fn streamed_memory_footer_arriving_in_chunks_is_filtered() { + let out = super::simulate_stream_markdown_for_tests( + &["Final answer line\n", "Memory", " used: `notes.md:1-2`\n"], + true, + ); + let rendered = lines_to_plain_strings(&out); + assert_eq!(rendered, vec!["Final answer line"]); + } + #[tokio::test] async fn heading_starts_on_new_line_when_following_paragraph() { // Stream a paragraph line, then a heading on the next line.