Surface fetched memories as citations in agent responses

Cargo.toml

@@ -338,7 +338,7 @@ version-compare = "0.1"
 vte = { git = "https://github.com/warpdotdev/vte.git", rev = "4b399c87b63ba88f45709edaa6383fc519f6c900", default-features = false }
 walkdir = "2"
 warp-workflows = { git = "https://github.com/warpdotdev/workflows", rev = "793a98ddda6ef19682aed66364faebd2829f0e01" }
-warp_multi_agent_api = { git = "https://github.com/warpdotdev/warp-proto-apis.git", rev = "ac1af7303d2931b0fb485be650a1fbc8b80d5667" }
+warp_multi_agent_api = { git = "https://github.com/warpdotdev/warp-proto-apis.git", rev = "45f9b1342b256b20f716aede64dc1b46a639e5e6" }
 wasm-bindgen = "0.2.89"
 wasm-bindgen-futures = "0.4.42"
 web-sys = { version = "0.3.69", features = [

app/src/ai/agent/api/convert_from_tests.rs

@@ -21,6 +21,7 @@ fn start_agent_tool_call_message(
     lifecycle_subscription_event_types: Option<Vec<i32>>,
 ) -> api::Message {
     api::Message {
+        fetched_memories: vec![],
         id: "message-id".to_string(),
         task_id: "task-id".to_string(),
         server_message_data: String::new(),
@@ -67,6 +68,7 @@ fn start_agent_v2_tool_call_message(
     lifecycle_subscription_event_types: Option<Vec<i32>>,
 ) -> api::Message {
     api::Message {
+        fetched_memories: vec![],
         id: "message-id".to_string(),
         task_id: "task-id".to_string(),
         server_message_data: String::new(),
@@ -91,6 +93,7 @@ fn start_agent_v2_tool_call_message(
 
 fn upload_artifact_tool_call_message(path: &str, description: &str) -> api::Message {
     api::Message {
+        fetched_memories: vec![],
         id: "message-id".to_string(),
         task_id: "task-id".to_string(),
         server_message_data: String::new(),
@@ -144,6 +147,7 @@ fn remote_start_agent_v2_execution_mode(
 
 fn file_artifact_created_message(filepath: &str, description: &str) -> api::Message {
     api::Message {
+        fetched_memories: vec![],
         id: "message-id".to_string(),
         task_id: "task-id".to_string(),
         server_message_data: String::new(),
@@ -616,6 +620,7 @@ fn transfer_control_tool_call_converts_to_action_message() {
     let task_id = TaskId::new("task".to_string());
     let reason = "Please finish the interactive flow".to_string();
     let message = api::Message {
+        fetched_memories: vec![],
         id: "message".to_string(),
         task_id: "task".to_string(),
         server_message_data: String::new(),

app/src/ai/agent/conversation.rs

@@ -1184,6 +1184,28 @@ impl AIConversation {
         self.task_store.all_linearized_messages()
     }
 
+    /// Returns the memories the server fetched for this conversation, in the order they first
+    /// appeared across messages (server-side rank order within each message). Re-fetched
+    /// memories are deduped by `(memory_store_id, memory_id)`: the first appearance keeps its
+    /// position while the entry's content/source are updated to the latest occurrence.
+    pub fn fetched_memories(&self) -> Vec<api::message::FetchedMemory> {
+        let mut memories: Vec<api::message::FetchedMemory> = Vec::new();
+        let mut index_by_id: HashMap<(String, String), usize> = HashMap::new();
+        for message in self.task_store.all_linearized_messages() {
+            for memory in &message.fetched_memories {
+                let key = (memory.memory_store_id.clone(), memory.memory_id.clone());
+                match index_by_id.get(&key) {
+                    Some(index) => memories[*index] = memory.clone(),
+                    None => {
+                        index_by_id.insert(key, memories.len());
+                        memories.push(memory.clone());
+                    }
+                }
+            }
+        }
+        memories
+    }
+
     /// Returns all the tasks in this conversation.
     ///
     /// Note that until we've fully migrated to the multi-agent endpoint, in reality, each

app/src/ai/agent/conversation_tests.rs

@@ -53,6 +53,7 @@ fn restored_conversation_with_root_description(description: &str) -> AIConversat
 
 fn user_query_message(id: &str, request_id: &str, query: &str) -> api::Message {
     api::Message {
+        fetched_memories: vec![],
         id: id.to_string(),
         task_id: "root-task".to_string(),
         server_message_data: String::new(),
@@ -71,6 +72,7 @@ fn user_query_message(id: &str, request_id: &str, query: &str) -> api::Message {
 
 fn agent_output_message(id: &str, request_id: &str) -> api::Message {
     api::Message {
+        fetched_memories: vec![],
         id: id.to_string(),
         task_id: "root-task".to_string(),
         server_message_data: String::new(),
@@ -741,3 +743,115 @@ fn restored_conversation_does_not_re_enter_waiting_for_events() {
 
     assert_eq!(conversation.status(), &ConversationStatus::Success);
 }
+
+fn fetched_memory(
+    memory_id: &str,
+    content: &str,
+    memory_store_id: &str,
+    source: Option<api::message::fetched_memory::Source>,
+) -> api::message::FetchedMemory {
+    api::message::FetchedMemory {
+        memory_id: memory_id.to_string(),
+        content: content.to_string(),
+        memory_store_id: memory_store_id.to_string(),
+        source,
+    }
+}
+
+fn conversation_source(conversation_id: &str) -> Option<api::message::fetched_memory::Source> {
+    Some(api::message::fetched_memory::Source::Conversation(
+        api::message::fetched_memory::Conversation {
+            conversation_id: conversation_id.to_string(),
+        },
+    ))
+}
+
+fn restored_conversation_with_memories_per_query(
+    memories_per_query: Vec<Vec<api::message::FetchedMemory>>,
+) -> AIConversation {
+    let messages = memories_per_query
+        .into_iter()
+        .enumerate()
+        .flat_map(|(index, memories)| {
+            let request_id = format!("request-{index}");
+            let query = api::Message {
+                fetched_memories: memories,
+                ..user_query_message(&format!("user-{index}"), &request_id, "query")
+            };
+            [
+                query,
+                agent_output_message(&format!("agent-{index}"), &request_id),
+            ]
+        })
+        .collect();
+
+    AIConversation::new_restored(
+        AIConversationId::new(),
+        vec![api::Task {
+            id: "root-task".to_string(),
+            messages,
+            ..Default::default()
+        }],
+        None,
+    )
+    .unwrap()
+}
+
+#[test]
+fn fetched_memories_is_empty_when_no_message_has_memories() {
+    let conversation = restored_conversation_with_memories_per_query(vec![vec![]]);
+
+    assert_eq!(conversation.fetched_memories(), vec![]);
+}
+
+#[test]
+fn fetched_memories_preserves_order_across_and_within_messages() {
+    let conversation = restored_conversation_with_memories_per_query(vec![
+        vec![
+            fetched_memory("m1", "first", "store-1", None),
+            fetched_memory("m2", "second", "store-1", None),
+        ],
+        vec![fetched_memory("m3", "third", "store-2", None)],
+    ]);
+
+    let ids: Vec<String> = conversation
+        .fetched_memories()
+        .into_iter()
+        .map(|memory| memory.memory_id)
+        .collect();
+    assert_eq!(ids, vec!["m1", "m2", "m3"]);
+}
+
+#[test]
+fn fetched_memories_dedupes_keeping_first_position_and_latest_data() {
+    let conversation = restored_conversation_with_memories_per_query(vec![
+        vec![
+            fetched_memory("m1", "old content", "store-1", None),
+            fetched_memory("m2", "other", "store-1", None),
+        ],
+        vec![
+            fetched_memory(
+                "m1",
+                "new content",
+                "store-1",
+                conversation_source("conversation-1"),
+            ),
+            fetched_memory("m1", "same memory id different store", "store-2", None),
+        ],
+    ]);
+
+    let memories = conversation.fetched_memories();
+    assert_eq!(
+        memories,
+        vec![
+            fetched_memory(
+                "m1",
+                "new content",
+                "store-1",
+                conversation_source("conversation-1"),
+            ),
+            fetched_memory("m2", "other", "store-1", None),
+            fetched_memory("m1", "same memory id different store", "store-2", None),
+        ]
+    );
+}

app/src/ai/agent/conversation_yaml_tests.rs

@@ -21,6 +21,7 @@ fn list_dir_sorted(dir: &Path) -> Vec<String> {
 
 fn make_user_query_message(id: &str, task_id: &str, query: &str) -> api::Message {
     api::Message {
+        fetched_memories: vec![],
         id: id.to_string(),
         task_id: task_id.to_string(),
         server_message_data: String::new(),
@@ -44,6 +45,7 @@ fn make_tool_call_message(
     tool: api::message::tool_call::Tool,
 ) -> api::Message {
     api::Message {
+        fetched_memories: vec![],
         id: id.to_string(),
         task_id: task_id.to_string(),
         server_message_data: String::new(),
@@ -64,6 +66,7 @@ fn make_tool_call_result_message(
     result: api::message::tool_call_result::Result,
 ) -> api::Message {
     api::Message {
+        fetched_memories: vec![],
         id: id.to_string(),
         task_id: task_id.to_string(),
         server_message_data: String::new(),

app/src/ai/agent/linearization_tests.rs

@@ -7,6 +7,7 @@ use super::*;
 // Helper function to create a basic message
 fn create_message(id: &str, task_id: &str) -> api::Message {
     api::Message {
+        fetched_memories: vec![],
         id: id.to_string(),
         task_id: task_id.to_string(),
         server_message_data: "server_data".to_string(),
@@ -23,6 +24,7 @@ fn create_message(id: &str, task_id: &str) -> api::Message {
 
 fn create_subagent_tool_call_message(id: &str, task_id: &str, subtask_id: &str) -> api::Message {
     api::Message {
+        fetched_memories: vec![],
         id: id.to_string(),
         task_id: task_id.to_string(),
         server_message_data: "server_data".to_string(),
@@ -45,6 +47,7 @@ fn create_subagent_tool_call_message(id: &str, task_id: &str, subtask_id: &str)
 // Helper function to create a tool call result message.
 fn create_tool_call_result_message(id: &str, task_id: &str, tool_call_id: &str) -> api::Message {
     api::Message {
+        fetched_memories: vec![],
         id: id.to_string(),
         task_id: task_id.to_string(),
         server_message_data: "server_data".to_string(),

app/src/ai/agent/task_tests.rs

@@ -48,6 +48,7 @@ fn create_start_agent_tool_call_message(
     prompt: &str,
 ) -> api::Message {
     api::Message {
+        fetched_memories: vec![],
         id: id.to_string(),
         task_id: task_id.to_string(),
         server_message_data: String::new(),

app/src/ai/agent/telemetry.rs

@@ -33,6 +33,14 @@ impl ForTelemetry for AIAgentCitation {
                 Some(CitationForTelemetry::WarpDocs { page: path.clone() })
             }
             Self::WebPage { url } => Some(CitationForTelemetry::WebPage { url: url.clone() }),
+            Self::AgentMemory {
+                memory_store_id,
+                memory_id,
+                ..
+            } => Some(CitationForTelemetry::AgentMemory {
+                memory_store_id: memory_store_id.clone(),
+                memory_id: memory_id.clone(),
+            }),
         }
     }
 }

app/src/ai/blocklist/agent_view/orchestration_pill_bar_tests.rs

@@ -62,6 +62,7 @@ fn pill_bar_data_layer_finds_restored_children_before_pane_creation() {
                 tasks: vec![warp_multi_agent_api::Task {
                     id: format!("task-{child_id}"),
                     messages: vec![warp_multi_agent_api::Message {
+                        fetched_memories: vec![],
                         id: "child-msg".to_string(),
                         task_id: format!("task-{child_id}"),
                         server_message_data: String::new(),
@@ -111,6 +112,7 @@ fn pill_bar_data_layer_finds_restored_children_before_pane_creation() {
                 tasks: vec![warp_multi_agent_api::Task {
                     id: format!("task-{parent_id}"),
                     messages: vec![warp_multi_agent_api::Message {
+                        fetched_memories: vec![],
                         id: "parent-msg".to_string(),
                         task_id: format!("task-{parent_id}"),
                         server_message_data: String::new(),

app/src/ai/blocklist/block.rs

@@ -2149,6 +2149,30 @@ impl AIBlock {
                 .entry(citation.clone())
                 .or_default();
         }
+        // Also register handles for memory citations derived from fetched_memories,
+        // which are synthesized at render time and never go through output.citations.
+        // Only register for the first exchange since that's the only one that shows them.
+        if let Some(conversation) = self.model.conversation(ctx) {
+            let is_first_exchange = conversation
+                .first_exchange()
+                .map(|e| Some(e.id) == self.model.exchange_id(ctx))
+                .unwrap_or(false);
+            if is_first_exchange {
+                for memory in conversation.fetched_memories() {
+                    if memory.memory_store_id.is_empty() || memory.memory_id.is_empty() {
+                        continue;
+                    }
+                    self.state_handles
+                        .footer_citation_chip_handles
+                        .entry(AIAgentCitation::AgentMemory {
+                            memory_store_id: memory.memory_store_id.clone(),
+                            memory_id: memory.memory_id.clone(),
+                            content: memory.content.clone(),
+                        })
+                        .or_default();
+                }
+            }
+        }
 
         // Register element state for reasoning messages and track summarization timing.
         for message in &output.messages {

app/src/ai/blocklist/block/view_impl.rs

@@ -678,6 +678,15 @@ pub fn render_citation(
             let name = url.clone();
             (Some(icon), name)
         }
+        AIAgentCitation::AgentMemory { content, .. } => {
+            let icon = Icon::Cognition.to_warpui_icon(theme.foreground()).finish();
+            let name = if content.is_empty() {
+                String::from("Memory")
+            } else {
+                content.clone()
+            };
+            (Some(icon), name)
+        }
     };
 
     // Shorten the name to 30 chars.

app/src/ai/blocklist/block/view_impl/orchestration_tests.rs

@@ -306,6 +306,7 @@ fn participant_for_restored_child_run_id_resolves_to_agent_name() {
             tasks: vec![warp_multi_agent_api::Task {
                 id: format!("task-{child_id}"),
                 messages: vec![warp_multi_agent_api::Message {
+                    fetched_memories: vec![],
                     id: "child-msg".to_string(),
                     task_id: format!("task-{child_id}"),
                     server_message_data: String::new(),
@@ -358,6 +359,7 @@ fn participant_for_restored_child_run_id_resolves_to_agent_name() {
             tasks: vec![warp_multi_agent_api::Task {
                 id: format!("task-{parent_id}"),
                 messages: vec![warp_multi_agent_api::Message {
+                    fetched_memories: vec![],
                     id: "parent-msg".to_string(),
                     task_id: format!("task-{parent_id}"),
                     server_message_data: String::new(),