diff --git a/Cargo.lock b/Cargo.lock index 0a2dd21184..b2b7650650 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -15653,7 +15653,7 @@ dependencies = [ [[package]] name = "warp_multi_agent_api" version = "0.0.0" -source = "git+https://github.com/warpdotdev/warp-proto-apis.git?rev=ac1af7303d2931b0fb485be650a1fbc8b80d5667#ac1af7303d2931b0fb485be650a1fbc8b80d5667" +source = "git+https://github.com/warpdotdev/warp-proto-apis.git?rev=45f9b1342b256b20f716aede64dc1b46a639e5e6#45f9b1342b256b20f716aede64dc1b46a639e5e6" dependencies = [ "prost", "prost-reflect", diff --git a/Cargo.toml b/Cargo.toml index 6b6d1111dc..ba3d280df4 100644 --- a/Cargo.toml +++ b/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 = [ diff --git a/app/src/ai/agent/api/convert_conversation_tests.rs b/app/src/ai/agent/api/convert_conversation_tests.rs index 29de963330..ef81ea67ab 100644 --- a/app/src/ai/agent/api/convert_conversation_tests.rs +++ b/app/src/ai/agent/api/convert_conversation_tests.rs @@ -360,6 +360,7 @@ fn test_into_exchanges_basic() { // Create minimal test data let messages = vec![ api::Message { + fetched_memories: vec![], id: "user_msg".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -375,6 +376,7 @@ fn test_into_exchanges_basic() { timestamp: None, }, api::Message { + fetched_memories: vec![], id: "agent_msg".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -388,6 +390,7 @@ fn test_into_exchanges_basic() { timestamp: None, }, api::Message { + fetched_memories: vec![], id: "user_msg2".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -403,6 +406,7 @@ fn test_into_exchanges_basic() { timestamp: None, }, api::Message { + fetched_memories: vec![], id: "agent_msg2".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -416,6 +420,7 @@ fn test_into_exchanges_basic() { timestamp: None, }, api::Message { + fetched_memories: vec![], id: "user_msg3".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -431,6 +436,7 @@ fn test_into_exchanges_basic() { timestamp: None, }, api::Message { + fetched_memories: vec![], id: "agent_msg3".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -475,6 +481,7 @@ fn test_invoke_skill_arguments_round_trip() { let query = "arg1 arg2".to_string(); let messages = vec![ api::Message { + fetched_memories: vec![], id: "invoke_skill_msg".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -495,6 +502,7 @@ fn test_invoke_skill_arguments_round_trip() { timestamp: None, }, api::Message { + fetched_memories: vec![], id: "agent_msg".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -542,6 +550,7 @@ fn test_invoke_skill_arguments_round_trip() { #[test] fn test_invoke_skill_missing_user_query_maps_to_none() { let messages = vec![api::Message { + fetched_memories: vec![], id: "invoke_skill_msg".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -588,6 +597,7 @@ fn test_into_exchanges_with_tool_calls_and_cancellation() { let messages = vec![ // User query api::Message { + fetched_memories: vec![], id: "user_query".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -604,6 +614,7 @@ fn test_into_exchanges_with_tool_calls_and_cancellation() { }, // Agent response api::Message { + fetched_memories: vec![], id: "agent_response".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -618,6 +629,7 @@ fn test_into_exchanges_with_tool_calls_and_cancellation() { }, // Tool call 1 api::Message { + fetched_memories: vec![], id: "tool_call_1".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -641,6 +653,7 @@ fn test_into_exchanges_with_tool_calls_and_cancellation() { }, // Tool call 2 api::Message { + fetched_memories: vec![], id: "tool_call_2".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -664,6 +677,7 @@ fn test_into_exchanges_with_tool_calls_and_cancellation() { }, // Tool call 3 api::Message { + fetched_memories: vec![], id: "tool_call_3".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -687,6 +701,7 @@ fn test_into_exchanges_with_tool_calls_and_cancellation() { }, // Tool call result - cancelled (call_2) api::Message { + fetched_memories: vec![], id: "result_cancelled".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -703,6 +718,7 @@ fn test_into_exchanges_with_tool_calls_and_cancellation() { }, // Tool call result - success (call_1) api::Message { + fetched_memories: vec![], id: "result_success_1".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -735,6 +751,7 @@ fn test_into_exchanges_with_tool_calls_and_cancellation() { }, // Tool call result - success (call_3) api::Message { + fetched_memories: vec![], id: "result_success_3".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -767,6 +784,7 @@ fn test_into_exchanges_with_tool_calls_and_cancellation() { }, // Final agent response api::Message { + fetched_memories: vec![], id: "final_response".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -781,6 +799,7 @@ fn test_into_exchanges_with_tool_calls_and_cancellation() { }, // Follow-up user query api::Message { + fetched_memories: vec![], id: "followup_query".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -797,6 +816,7 @@ fn test_into_exchanges_with_tool_calls_and_cancellation() { }, // Final agent response api::Message { + fetched_memories: vec![], id: "final_response2".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -898,6 +918,7 @@ fn test_into_exchanges_with_code_diffs() { let messages = vec![ // User query asking for code changes api::Message { + fetched_memories: vec![], id: "user_query".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -914,6 +935,7 @@ fn test_into_exchanges_with_code_diffs() { }, // Agent response api::Message { + fetched_memories: vec![], id: "agent_response".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -928,6 +950,7 @@ fn test_into_exchanges_with_code_diffs() { }, // File diff tool call api::Message { + fetched_memories: vec![], id: "diff_call".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -949,6 +972,7 @@ fn test_into_exchanges_with_code_diffs() { }, // User cancels the diff api::Message { + fetched_memories: vec![], id: "diff_cancelled".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -965,6 +989,7 @@ fn test_into_exchanges_with_code_diffs() { }, // User provides feedback api::Message { + fetched_memories: vec![], id: "user_feedback".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -981,6 +1006,7 @@ fn test_into_exchanges_with_code_diffs() { }, // Agent response api::Message { + fetched_memories: vec![], id: "agent_response_2".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -995,6 +1021,7 @@ fn test_into_exchanges_with_code_diffs() { }, // Second file diff tool call api::Message { + fetched_memories: vec![], id: "diff_call_2".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -1016,6 +1043,7 @@ fn test_into_exchanges_with_code_diffs() { }, // User accepts the diff api::Message { + fetched_memories: vec![], id: "diff_accepted".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -1043,6 +1071,7 @@ fn test_into_exchanges_with_code_diffs() { }, // Final agent response api::Message { + fetched_memories: vec![], id: "final_response".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -1057,6 +1086,7 @@ fn test_into_exchanges_with_code_diffs() { }, // Follow-up user query api::Message { + fetched_memories: vec![], id: "followup".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -1073,6 +1103,7 @@ fn test_into_exchanges_with_code_diffs() { }, // Final agent response api::Message { + fetched_memories: vec![], id: "final_response_2".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -1170,6 +1201,7 @@ fn test_into_exchanges_with_code_diffs() { fn test_user_query_mode_conversion() { // Test conversion with Plan mode let messages = vec![api::Message { + fetched_memories: vec![], id: "user_msg".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -1216,6 +1248,7 @@ fn test_user_query_mode_conversion() { // Test conversion with Normal mode (no type set) let messages_normal = vec![api::Message { + fetched_memories: vec![], id: "user_msg".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -1260,6 +1293,7 @@ fn test_user_query_mode_conversion() { // Test conversion with no mode field (should default to Normal) let messages_default = vec![api::Message { + fetched_memories: vec![], id: "user_msg".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -1313,6 +1347,7 @@ fn test_exchanges_grouped_by_request_id() { let messages = vec![ // Message 0: Server message (should be ignored or handled gracefully) api::Message { + fetched_memories: vec![], id: "2512077c-0ede-46b0-8f69-230c8792df07".to_string(), task_id: "d02463e1-2429-48de-ac8f-552df4acc4d0".to_string(), request_id: "78e236b8-84a2-45df-876e-ebfb86ceafc4".to_string(), @@ -1330,6 +1365,7 @@ fn test_exchanges_grouped_by_request_id() { }, // Message 1: User query with request_id 78e236b8 api::Message { + fetched_memories: vec![], id: "4d6c450d-3d54-446f-974c-5c414e6083e9".to_string(), task_id: "d02463e1-2429-48de-ac8f-552df4acc4d0".to_string(), request_id: "78e236b8-84a2-45df-876e-ebfb86ceafc4".to_string(), @@ -1346,6 +1382,7 @@ fn test_exchanges_grouped_by_request_id() { }, // Message 2: Agent output with same request_id api::Message { + fetched_memories: vec![], id: "10210d1a-5298-45ef-90ba-df6367805080".to_string(), task_id: "d02463e1-2429-48de-ac8f-552df4acc4d0".to_string(), request_id: "78e236b8-84a2-45df-876e-ebfb86ceafc4".to_string(), @@ -1360,6 +1397,7 @@ fn test_exchanges_grouped_by_request_id() { }, // Message 3: Tool call with same request_id api::Message { + fetched_memories: vec![], id: "936c7c86-eb4a-4edf-97c0-22f5c61b35a6".to_string(), task_id: "d02463e1-2429-48de-ac8f-552df4acc4d0".to_string(), request_id: "78e236b8-84a2-45df-876e-ebfb86ceafc4".to_string(), @@ -1383,6 +1421,7 @@ fn test_exchanges_grouped_by_request_id() { }, // Message 4: Tool call result with NEW request_id 59a3947f (starts new exchange) api::Message { + fetched_memories: vec![], id: "cbebf5fb-4dd8-4aef-be45-bb916eff552c".to_string(), task_id: "d02463e1-2429-48de-ac8f-552df4acc4d0".to_string(), request_id: "59a3947f-fc7e-413a-96b5-baecd7e406dc".to_string(), @@ -1417,6 +1456,7 @@ fn test_exchanges_grouped_by_request_id() { }, // Message 5: Agent output with same request_id api::Message { + fetched_memories: vec![], id: "7a89857d-fa33-4d45-88e3-5fa9cbce3f20".to_string(), task_id: "d02463e1-2429-48de-ac8f-552df4acc4d0".to_string(), request_id: "59a3947f-fc7e-413a-96b5-baecd7e406dc".to_string(), @@ -1431,6 +1471,7 @@ fn test_exchanges_grouped_by_request_id() { }, // Message 6: Write to long running command with NEW request_id 9f85acb2 (starts new exchange) api::Message { + fetched_memories: vec![], id: "dac6d336-9fcb-4e34-bc2b-b06e70f52ec5".to_string(), task_id: "d02463e1-2429-48de-ac8f-552df4acc4d0".to_string(), request_id: "9f85acb2-0b1f-41b1-a0de-3623e131758a".to_string(), @@ -1461,6 +1502,7 @@ fn test_exchanges_grouped_by_request_id() { }, // Message 7: Final tool call result with same request_id api::Message { + fetched_memories: vec![], id: "ad319d66-fac0-4169-8bf1-e6004aca1619".to_string(), task_id: "d02463e1-2429-48de-ac8f-552df4acc4d0".to_string(), request_id: "9f85acb2-0b1f-41b1-a0de-3623e131758a".to_string(), @@ -1493,6 +1535,7 @@ fn test_exchanges_grouped_by_request_id() { }, // Message 8: Final agent output with same request_id api::Message { + fetched_memories: vec![], id: "f15f8a59-2e9c-416e-b216-83b3bd52d6be".to_string(), task_id: "d02463e1-2429-48de-ac8f-552df4acc4d0".to_string(), request_id: "9f85acb2-0b1f-41b1-a0de-3623e131758a".to_string(), @@ -1581,6 +1624,7 @@ fn test_multiple_create_documents_get_default_version() { let messages = vec![ // User query api::Message { + fetched_memories: vec![], id: "user_msg".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -1597,6 +1641,7 @@ fn test_multiple_create_documents_get_default_version() { }, // Agent output api::Message { + fetched_memories: vec![], id: "agent_text".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -1611,6 +1656,7 @@ fn test_multiple_create_documents_get_default_version() { }, // First CreateDocuments tool call api::Message { + fetched_memories: vec![], id: "tool_call_create_a".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -1633,6 +1679,7 @@ fn test_multiple_create_documents_get_default_version() { }, // First CreateDocuments result api::Message { + fetched_memories: vec![], id: "result_create_a".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -1661,6 +1708,7 @@ fn test_multiple_create_documents_get_default_version() { }, // Agent output before second plan api::Message { + fetched_memories: vec![], id: "agent_text_2".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -1675,6 +1723,7 @@ fn test_multiple_create_documents_get_default_version() { }, // Second CreateDocuments tool call api::Message { + fetched_memories: vec![], id: "tool_call_create_b".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -1697,6 +1746,7 @@ fn test_multiple_create_documents_get_default_version() { }, // Second CreateDocuments result api::Message { + fetched_memories: vec![], id: "result_create_b".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -1791,6 +1841,7 @@ fn test_create_then_edit_then_create_version_tracking() { let messages = vec![ // User query api::Message { + fetched_memories: vec![], id: "user_msg".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -1807,6 +1858,7 @@ fn test_create_then_edit_then_create_version_tracking() { }, // Agent output api::Message { + fetched_memories: vec![], id: "agent_text".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -1821,6 +1873,7 @@ fn test_create_then_edit_then_create_version_tracking() { }, // Create doc A tool call api::Message { + fetched_memories: vec![], id: "tool_call_create_a".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -1843,6 +1896,7 @@ fn test_create_then_edit_then_create_version_tracking() { }, // Create doc A result api::Message { + fetched_memories: vec![], id: "result_create_a".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -1871,6 +1925,7 @@ fn test_create_then_edit_then_create_version_tracking() { }, // Agent output before edit api::Message { + fetched_memories: vec![], id: "agent_text_2".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -1885,6 +1940,7 @@ fn test_create_then_edit_then_create_version_tracking() { }, // Edit doc A tool call api::Message { + fetched_memories: vec![], id: "tool_call_edit_a".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -1906,6 +1962,7 @@ fn test_create_then_edit_then_create_version_tracking() { }, // Edit doc A result api::Message { + fetched_memories: vec![], id: "result_edit_a".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -1934,6 +1991,7 @@ fn test_create_then_edit_then_create_version_tracking() { }, // Agent output before second create api::Message { + fetched_memories: vec![], id: "agent_text_3".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -1948,6 +2006,7 @@ fn test_create_then_edit_then_create_version_tracking() { }, // Create doc B tool call api::Message { + fetched_memories: vec![], id: "tool_call_create_b".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -1970,6 +2029,7 @@ fn test_create_then_edit_then_create_version_tracking() { }, // Create doc B result api::Message { + fetched_memories: vec![], id: "result_create_b".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -2077,6 +2137,7 @@ fn test_handoff_rehydration_system_query_is_hidden() { let messages = vec![ // HandoffRehydration system query – should be hidden api::Message { + fetched_memories: vec![], id: "msg_handoff".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), @@ -2096,6 +2157,7 @@ fn test_handoff_rehydration_system_query_is_hidden() { }, // Agent output that follows the hidden system query api::Message { + fetched_memories: vec![], id: "msg_output".to_string(), task_id: "task1".to_string(), server_message_data: "".to_string(), diff --git a/app/src/ai/agent/api/convert_from_tests.rs b/app/src/ai/agent/api/convert_from_tests.rs index 6dff162662..8f5a32f306 100644 --- a/app/src/ai/agent/api/convert_from_tests.rs +++ b/app/src/ai/agent/api/convert_from_tests.rs @@ -21,6 +21,7 @@ fn start_agent_tool_call_message( lifecycle_subscription_event_types: Option>, ) -> 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>, ) -> 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(), diff --git a/app/src/ai/agent/conversation.rs b/app/src/ai/agent/conversation.rs index 3a6f4d610c..e13b226735 100644 --- a/app/src/ai/agent/conversation.rs +++ b/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 { + let mut memories: Vec = 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 diff --git a/app/src/ai/agent/conversation_tests.rs b/app/src/ai/agent/conversation_tests.rs index 434ce70e4c..3da7e7a4d1 100644 --- a/app/src/ai/agent/conversation_tests.rs +++ b/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::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 { + 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>, +) -> 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 = 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), + ] + ); +} diff --git a/app/src/ai/agent/conversation_yaml_tests.rs b/app/src/ai/agent/conversation_yaml_tests.rs index 88269b0fd7..022da5c5f1 100644 --- a/app/src/ai/agent/conversation_yaml_tests.rs +++ b/app/src/ai/agent/conversation_yaml_tests.rs @@ -21,6 +21,7 @@ fn list_dir_sorted(dir: &Path) -> Vec { 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(), diff --git a/app/src/ai/agent/linearization_tests.rs b/app/src/ai/agent/linearization_tests.rs index d28b4fff49..f988350644 100644 --- a/app/src/ai/agent/linearization_tests.rs +++ b/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(), diff --git a/app/src/ai/agent/task_tests.rs b/app/src/ai/agent/task_tests.rs index 40f7c644be..d1af6c1e16 100644 --- a/app/src/ai/agent/task_tests.rs +++ b/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(), diff --git a/app/src/ai/agent/telemetry.rs b/app/src/ai/agent/telemetry.rs index fd082a5162..372476260c 100644 --- a/app/src/ai/agent/telemetry.rs +++ b/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(), + }), } } } diff --git a/app/src/ai/blocklist/agent_view/orchestration_pill_bar_tests.rs b/app/src/ai/blocklist/agent_view/orchestration_pill_bar_tests.rs index d28e6d66b8..82765271c0 100644 --- a/app/src/ai/blocklist/agent_view/orchestration_pill_bar_tests.rs +++ b/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(), diff --git a/app/src/ai/blocklist/block.rs b/app/src/ai/blocklist/block.rs index 3338ac1554..39941d13db 100644 --- a/app/src/ai/blocklist/block.rs +++ b/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 { diff --git a/app/src/ai/blocklist/block/view_impl.rs b/app/src/ai/blocklist/block/view_impl.rs index 06cb2de3d7..a19bc6b1b5 100644 --- a/app/src/ai/blocklist/block/view_impl.rs +++ b/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. diff --git a/app/src/ai/blocklist/block/view_impl/orchestration_tests.rs b/app/src/ai/blocklist/block/view_impl/orchestration_tests.rs index f4b0c02550..125a5af582 100644 --- a/app/src/ai/blocklist/block/view_impl/orchestration_tests.rs +++ b/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(), diff --git a/app/src/ai/blocklist/block/view_impl/output.rs b/app/src/ai/blocklist/block/view_impl/output.rs index 0dd79a0e23..c3294d4f92 100644 --- a/app/src/ai/blocklist/block/view_impl/output.rs +++ b/app/src/ai/blocklist/block/view_impl/output.rs @@ -1098,9 +1098,30 @@ pub(super) fn render(props: Props, app: &AppContext) -> Box { } if should_render_references_section { - if let Some(references) = - render_references_footer(&output.citations, props, app) - { + let exchange_id = props.model.exchange_id(app); + let memory_citations: Vec = props + .model + .conversation(app) + .filter(|conv| { + // Only show memory citations on the first exchange. + conv.first_exchange().map(|e| Some(e.id)) == Some(exchange_id) + }) + .into_iter() + .flat_map(|conv| conv.fetched_memories()) + .filter(|m| !m.memory_store_id.is_empty() && !m.memory_id.is_empty()) + .map(|m| AIAgentCitation::AgentMemory { + memory_store_id: m.memory_store_id.clone(), + memory_id: m.memory_id.clone(), + content: m.content.clone(), + }) + .collect(); + let all_citations: Vec = output + .citations + .iter() + .cloned() + .chain(memory_citations) + .collect(); + if let Some(references) = render_references_footer(&all_citations, props, app) { output_items.add_child(references); } } diff --git a/app/src/ai/blocklist/history_model_tests.rs b/app/src/ai/blocklist/history_model_tests.rs index d7ef8c36e7..82ab304b23 100644 --- a/app/src/ai/blocklist/history_model_tests.rs +++ b/app/src/ai/blocklist/history_model_tests.rs @@ -72,6 +72,7 @@ fn create_user_query_message( query: &str, ) -> warp_multi_agent_api::Message { warp_multi_agent_api::Message { + fetched_memories: vec![], id: id.to_string(), task_id: task_id.to_string(), server_message_data: String::new(), diff --git a/app/src/ai/blocklist/orchestration_event_streamer_tests.rs b/app/src/ai/blocklist/orchestration_event_streamer_tests.rs index 80e16dd78e..9ffec121d4 100644 --- a/app/src/ai/blocklist/orchestration_event_streamer_tests.rs +++ b/app/src/ai/blocklist/orchestration_event_streamer_tests.rs @@ -128,6 +128,7 @@ fn ai_conversation_new_restored_preserves_last_event_sequence() { let task = api::Task { id: "root".to_string(), messages: vec![api::Message { + fetched_memories: vec![], id: "m1".to_string(), task_id: "root".to_string(), server_message_data: String::new(), diff --git a/app/src/ai/conversation_details_panel_tests.rs b/app/src/ai/conversation_details_panel_tests.rs index a192125752..0b33dd2f17 100644 --- a/app/src/ai/conversation_details_panel_tests.rs +++ b/app/src/ai/conversation_details_panel_tests.rs @@ -107,6 +107,7 @@ fn test_from_conversation_prefers_server_creator_profile() { fn create_message_with_directory(id: &str, task_id: &str, directory: &str) -> api::Message { api::Message { + fetched_memories: vec![], id: id.to_string(), task_id: task_id.to_string(), server_message_data: String::new(), @@ -132,6 +133,7 @@ fn create_message_with_directory(id: &str, task_id: &str, directory: &str) -> ap fn create_agent_output_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: String::new(), diff --git a/app/src/server/telemetry/events.rs b/app/src/server/telemetry/events.rs index 03bc5aa3a4..318e31b549 100644 --- a/app/src/server/telemetry/events.rs +++ b/app/src/server/telemetry/events.rs @@ -990,6 +990,12 @@ pub enum AgentModeCitation { #[serde(skip_serializing)] url: String, }, + /// A fetched memory surfaced as a citation so we can track whether memory-backed + /// responses are shown to users and whether users open those memory citations. + AgentMemory { + memory_store_id: String, + memory_id: String, + }, } #[derive(Clone, Copy, Debug, Serialize)] diff --git a/app/src/terminal/view.rs b/app/src/terminal/view.rs index 4d95363881..dcd92d4151 100644 --- a/app/src/terminal/view.rs +++ b/app/src/terminal/view.rs @@ -20107,6 +20107,19 @@ impl TerminalView { AIAgentCitation::WebPage { url } => { ctx.open_url(url); } + AIAgentCitation::AgentMemory { + memory_store_id, + memory_id, + .. + } => { + let oz_root_url = ChannelState::oz_root_url(); + let url = format!( + "{oz_root_url}/memory/{}/memories/{}", + urlencoding::encode(memory_store_id), + urlencoding::encode(memory_id) + ); + ctx.open_url(&url); + } }, AIBlockEvent::OpenAIFactCollection { sync_id } => { ctx.emit(Event::OpenAIFactCollection { sync_id: *sync_id }); diff --git a/app/src/terminal/view/shared_session/view_impl_tests.rs b/app/src/terminal/view/shared_session/view_impl_tests.rs index f9e207e492..09c9b85747 100644 --- a/app/src/terminal/view/shared_session/view_impl_tests.rs +++ b/app/src/terminal/view/shared_session/view_impl_tests.rs @@ -1912,6 +1912,7 @@ fn test_shared_followup_on_existing_conversation_converts_user_query_input() { api_client_action::AddMessagesToTask { task_id: root_task_id.to_string(), messages: vec![api::Message { + fetched_memories: vec![], id: "user-message".to_string(), task_id: root_task_id.to_string(), server_message_data: String::new(), diff --git a/app/src/test_util/ai_agent_tasks.rs b/app/src/test_util/ai_agent_tasks.rs index 925504ba9b..5e4d9266d3 100644 --- a/app/src/test_util/ai_agent_tasks.rs +++ b/app/src/test_util/ai_agent_tasks.rs @@ -2,6 +2,7 @@ use warp_multi_agent_api as api; pub 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: String::new(), @@ -23,6 +24,7 @@ pub fn create_subagent_tool_call_message( metadata: Option, ) -> api::Message { api::Message { + fetched_memories: vec![], id: id.to_string(), task_id: task_id.to_string(), server_message_data: String::new(), diff --git a/crates/ai/src/agent/citation.rs b/crates/ai/src/agent/citation.rs index 449e300fac..32fd5a75ca 100644 --- a/crates/ai/src/agent/citation.rs +++ b/crates/ai/src/agent/citation.rs @@ -5,9 +5,22 @@ use warp_multi_agent_api as api; /// A citation listed in an AI response. #[derive(Debug, Clone, Hash, Eq, PartialEq)] pub enum AIAgentCitation { - WarpDriveObject { uid: String }, - WarpDocumentation { path: String }, - WebPage { url: String }, + WarpDriveObject { + uid: String, + }, + WarpDocumentation { + path: String, + }, + WebPage { + url: String, + }, + /// A memory from an attached memory store. `content` is the raw memory + /// text shown as a preview in the chip. + AgentMemory { + memory_store_id: String, + memory_id: String, + content: String, + }, } impl Display for AIAgentCitation { @@ -22,6 +35,13 @@ impl Display for AIAgentCitation { AIAgentCitation::WebPage { url } => { write!(f, "Web Page: {url}") } + AIAgentCitation::AgentMemory { + memory_store_id, + memory_id, + .. + } => { + write!(f, "Agent Memory: {memory_store_id}/{memory_id}") + } } } } diff --git a/crates/integration/src/test/agent_mode.rs b/crates/integration/src/test/agent_mode.rs index cd5fecaebe..9e11a4efbe 100644 --- a/crates/integration/src/test/agent_mode.rs +++ b/crates/integration/src/test/agent_mode.rs @@ -136,6 +136,7 @@ fn restored_user_query_message(task_id: &str, request_id: &str, directory: &str) })), request_id: request_id.to_string(), timestamp: None, + fetched_memories: vec![], } } @@ -161,6 +162,7 @@ fn restored_agent_output_message(task_id: &str, request_id: &str) -> api::Messag )), request_id: request_id.to_string(), timestamp: None, + fetched_memories: vec![], } } diff --git a/crates/persistence/src/model_tests.rs b/crates/persistence/src/model_tests.rs index 947f42c876..df0dd52df2 100644 --- a/crates/persistence/src/model_tests.rs +++ b/crates/persistence/src/model_tests.rs @@ -11,6 +11,7 @@ fn parentless_task(id: &str, message_count: usize) -> api::Task { dependencies: None, messages: (0..message_count) .map(|i| api::Message { + fetched_memories: vec![], id: format!("{id}-msg-{i}"), task_id: id.to_string(), server_message_data: String::new(),