feat: add memory citation to agent message (#14821)

Client side to come

codex-rs/core/src/memories/citations.rs

@@ -1,36 +1,89 @@
 use codex_protocol::ThreadId;
+use codex_protocol::memory_citation::MemoryCitation;
+use codex_protocol::memory_citation::MemoryCitationEntry;
+use std::collections::HashSet;
+
+pub fn parse_memory_citation(citations: Vec<String>) -> Option<MemoryCitation> {
+    let mut entries = Vec::new();
+    let mut rollout_ids = Vec::new();
+    let mut seen_rollout_ids = HashSet::new();
 
-pub fn get_thread_id_from_citations(citations: Vec<String>) -> Vec<ThreadId> {
-    let mut result = Vec::new();
     for citation in citations {
-        let mut ids_block = None;
-        for (open, close) in [
-            ("<thread_ids>", "</thread_ids>"),
-            ("<rollout_ids>", "</rollout_ids>"),
-        ] {
-            if let Some((_, rest)) = citation.split_once(open)
-                && let Some((ids, _)) = rest.split_once(close)
-            {
-                ids_block = Some(ids);
-                break;
-            }
+        if let Some(entries_block) =
+            extract_block(&citation, "<citation_entries>", "</citation_entries>")
+        {
+            entries.extend(
+                entries_block
+                    .lines()
+                    .filter_map(parse_memory_citation_entry),
+            );
         }
 
-        if let Some(ids_block) = ids_block {
+        if let Some(ids_block) = extract_ids_block(&citation) {
             for id in ids_block
                 .lines()
                 .map(str::trim)
                 .filter(|line| !line.is_empty())
             {
-                if let Ok(thread_id) = ThreadId::try_from(id) {
-                    result.push(thread_id);
+                if seen_rollout_ids.insert(id.to_string()) {
+                    rollout_ids.push(id.to_string());
                 }
             }
         }
     }
+
+    if entries.is_empty() && rollout_ids.is_empty() {
+        None
+    } else {
+        Some(MemoryCitation {
+            entries,
+            rollout_ids,
+        })
+    }
+}
+
+pub fn get_thread_id_from_citations(citations: Vec<String>) -> Vec<ThreadId> {
+    let mut result = Vec::new();
+    if let Some(memory_citation) = parse_memory_citation(citations) {
+        for rollout_id in memory_citation.rollout_ids {
+            if let Ok(thread_id) = ThreadId::try_from(rollout_id.as_str()) {
+                result.push(thread_id);
+            }
+        }
+    }
     result
 }
 
+fn parse_memory_citation_entry(line: &str) -> Option<MemoryCitationEntry> {
+    let line = line.trim();
+    if line.is_empty() {
+        return None;
+    }
+
+    let (location, note) = line.rsplit_once("|note=[")?;
+    let note = note.strip_suffix(']')?.trim().to_string();
+    let (path, line_range) = location.rsplit_once(':')?;
+    let (line_start, line_end) = line_range.split_once('-')?;
+
+    Some(MemoryCitationEntry {
+        path: path.trim().to_string(),
+        line_start: line_start.trim().parse().ok()?,
+        line_end: line_end.trim().parse().ok()?,
+        note,
+    })
+}
+
+fn extract_block<'a>(text: &'a str, open: &str, close: &str) -> Option<&'a str> {
+    let (_, rest) = text.split_once(open)?;
+    let (body, _) = rest.split_once(close)?;
+    Some(body)
+}
+
+fn extract_ids_block(text: &str) -> Option<&str> {
+    extract_block(text, "<rollout_ids>", "</rollout_ids>")
+        .or_else(|| extract_block(text, "<thread_ids>", "</thread_ids>"))
+}
+
 #[cfg(test)]
 #[path = "citations_tests.rs"]
 mod tests;

codex-rs/core/src/memories/citations_tests.rs

@@ -1,4 +1,5 @@
 use super::get_thread_id_from_citations;
+use super::parse_memory_citation;
 use codex_protocol::ThreadId;
 use pretty_assertions::assert_eq;
 
@@ -24,3 +25,40 @@ fn get_thread_id_from_citations_supports_legacy_rollout_ids() {
 
     assert_eq!(get_thread_id_from_citations(citations), vec![thread_id]);
 }
+
+#[test]
+fn parse_memory_citation_extracts_entries_and_rollout_ids() {
+    let first = ThreadId::new();
+    let second = ThreadId::new();
+    let citations = vec![format!(
+        "<citation_entries>\nMEMORY.md:1-2|note=[summary]\nrollout_summaries/foo.md:10-12|note=[details]\n</citation_entries>\n<rollout_ids>\n{first}\n{second}\n{first}\n</rollout_ids>"
+    )];
+
+    let parsed = parse_memory_citation(citations).expect("memory citation should parse");
+
+    assert_eq!(
+        parsed
+            .entries
+            .iter()
+            .map(|entry| (
+                entry.path.clone(),
+                entry.line_start,
+                entry.line_end,
+                entry.note.clone(),
+            ))
+            .collect::<Vec<_>>(),
+        vec![
+            ("MEMORY.md".to_string(), 1, 2, "summary".to_string()),
+            (
+                "rollout_summaries/foo.md".to_string(),
+                10,
+                12,
+                "details".to_string()
+            ),
+        ]
+    );
+    assert_eq!(
+        parsed.rollout_ids,
+        vec![first.to_string(), second.to_string()]
+    );
+}