Merge branch 'main' into zuxin/citation_cli_trip

codex-rs/tui/src/chatwidget/tests.rs

@@ -7139,6 +7139,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::<String>();
+    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::<String>();
+    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.

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<usize> {
+    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<Line<'static>> = 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.