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

Client side to come

codex-rs/core/src/codex.rs

@@ -6856,6 +6856,7 @@ async fn emit_agent_message_in_plan_mode(
                     id: agent_message_id.clone(),
                     content: Vec::new(),
                     phase: None,
+                    memory_citation: None,
                 })
             });
         sess.emit_turn_item_started(turn_context, &start_item).await;

codex-rs/core/src/event_mapping.rs

@@ -83,7 +83,12 @@ fn parse_agent_message(
         }
     }
     let id = id.cloned().unwrap_or_else(|| Uuid::new_v4().to_string());
-    AgentMessageItem { id, content, phase }
+    AgentMessageItem {
+        id,
+        content,
+        phase,
+        memory_citation: None,
+    }
 }
 
 pub fn parse_turn_item(item: &ResponseItem) -> Option<TurnItem> {

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()]
+    );
+}

codex-rs/core/src/rollout/recorder_tests.rs

@@ -85,6 +85,7 @@ async fn recorder_materializes_only_after_explicit_persist() -> std::io::Result<
             AgentMessageEvent {
                 message: "buffered-event".to_string(),
                 phase: None,
+                memory_citation: None,
             },
         ))])
         .await?;
@@ -201,6 +202,7 @@ async fn metadata_irrelevant_events_touch_state_db_updated_at() -> std::io::Resu
             AgentMessageEvent {
                 message: "assistant text".to_string(),
                 phase: None,
+                memory_citation: None,
             },
         ))])
         .await?;
@@ -251,6 +253,7 @@ async fn metadata_irrelevant_events_fall_back_to_upsert_when_thread_missing() ->
         AgentMessageEvent {
             message: "assistant text".to_string(),
             phase: None,
+            memory_citation: None,
         },
     ))];
 

codex-rs/core/src/stream_events_utils.rs

@@ -15,6 +15,7 @@ use crate::error::CodexErr;
 use crate::error::Result;
 use crate::function_tool::FunctionCallError;
 use crate::memories::citations::get_thread_id_from_citations;
+use crate::memories::citations::parse_memory_citation;
 use crate::parse_turn_item;
 use crate::state_db;
 use crate::tools::parallel::ToolCallRuntime;
@@ -38,6 +39,22 @@ fn strip_hidden_assistant_markup(text: &str, plan_mode: bool) -> String {
     }
 }
 
+fn strip_hidden_assistant_markup_and_parse_memory_citation(
+    text: &str,
+    plan_mode: bool,
+) -> (
+    String,
+    Option<codex_protocol::memory_citation::MemoryCitation>,
+) {
+    let (without_citations, citations) = strip_citations(text);
+    let visible_text = if plan_mode {
+        strip_proposed_plan_blocks(&without_citations)
+    } else {
+        without_citations
+    };
+    (visible_text, parse_memory_citation(citations))
+}
+
 pub(crate) fn raw_assistant_output_text_from_item(item: &ResponseItem) -> Option<String> {
     if let ResponseItem::Message { role, content, .. } = item
         && role == "assistant"
@@ -297,9 +314,11 @@ pub(crate) async fn handle_non_tool_response_item(
                         codex_protocol::items::AgentMessageContent::Text { text } => text.as_str(),
                     })
                     .collect::<String>();
-                let stripped = strip_hidden_assistant_markup(&combined, plan_mode);
+                let (stripped, memory_citation) =
+                    strip_hidden_assistant_markup_and_parse_memory_citation(&combined, plan_mode);
                 agent_message.content =
                     vec![codex_protocol::items::AgentMessageContent::Text { text: stripped }];
+                agent_message.memory_citation = memory_citation;
             }
             if let TurnItem::ImageGeneration(image_item) = &mut turn_item {
                 match save_image_generation_result(&image_item.id, &image_item.result).await {

codex-rs/core/src/stream_events_utils_tests.rs

@@ -23,7 +23,9 @@ fn assistant_output_text(text: &str) -> ResponseItem {
 #[tokio::test]
 async fn handle_non_tool_response_item_strips_citations_from_assistant_message() {
     let (session, turn_context) = make_session_and_context().await;
-    let item = assistant_output_text("hello<oai-mem-citation>doc1</oai-mem-citation> world");
+    let item = assistant_output_text(
+        "hello<oai-mem-citation><citation_entries>\nMEMORY.md:1-2|note=[x]\n</citation_entries>\n<rollout_ids>\n019cc2ea-1dff-7902-8d40-c8f6e5d83cc4\n</rollout_ids></oai-mem-citation> world",
+    );
 
     let turn_item = handle_non_tool_response_item(&session, &turn_context, &item, false)
         .await
@@ -40,6 +42,15 @@ async fn handle_non_tool_response_item_strips_citations_from_assistant_message()
         })
         .collect::<String>();
     assert_eq!(text, "hello world");
+    let memory_citation = agent_message
+        .memory_citation
+        .expect("memory citation should be parsed");
+    assert_eq!(memory_citation.entries.len(), 1);
+    assert_eq!(memory_citation.entries[0].path, "MEMORY.md");
+    assert_eq!(
+        memory_citation.rollout_ids,
+        vec!["019cc2ea-1dff-7902-8d40-c8f6e5d83cc4".to_string()]
+    );
 }
 
 #[test]

codex-rs/core/src/turn_timing_tests.rs

@@ -58,6 +58,7 @@ async fn turn_timing_state_records_ttfm_independently_of_ttft() {
                 id: "msg-1".to_string(),
                 content: Vec::new(),
                 phase: None,
+                memory_citation: None,
             }))
             .await
             .is_some()
@@ -68,6 +69,7 @@ async fn turn_timing_state_records_ttfm_independently_of_ttft() {
                 id: "msg-2".to_string(),
                 content: Vec::new(),
                 phase: None,
+                memory_citation: None,
             }))
             .await,
         None