assistant2: Restore tool uses when loading saved threads (#25942)

This PR makes it so tool uses are restored when loading saved threads in
Assistant 2.

Release Notes:

- N/A

crates/assistant2/src/tool_use.rs

@@ -7,10 +7,11 @@ use futures::FutureExt as _;
 use gpui::{SharedString, Task};
 use language_model::{
     LanguageModelRequestMessage, LanguageModelToolResult, LanguageModelToolUse,
-    LanguageModelToolUseId, MessageContent,
+    LanguageModelToolUseId, MessageContent, Role,
 };
 
 use crate::thread::MessageId;
+use crate::thread_store::SavedMessage;
 
 #[derive(Debug)]
 pub struct ToolUse {
@@ -28,7 +29,6 @@ pub enum ToolUseStatus {
     Error(SharedString),
 }
 
-#[derive(Default)]
 pub struct ToolUseState {
     tool_uses_by_assistant_message: HashMap<MessageId, Vec<LanguageModelToolUse>>,
     tool_uses_by_user_message: HashMap<MessageId, Vec<LanguageModelToolUseId>>,
@@ -37,6 +37,65 @@ pub struct ToolUseState {
 }
 
 impl ToolUseState {
+    pub fn new() -> Self {
+        Self {
+            tool_uses_by_assistant_message: HashMap::default(),
+            tool_uses_by_user_message: HashMap::default(),
+            tool_results: HashMap::default(),
+            pending_tool_uses_by_id: HashMap::default(),
+        }
+    }
+
+    pub fn from_saved_messages(messages: &[SavedMessage]) -> Self {
+        let mut this = Self::new();
+
+        for message in messages {
+            match message.role {
+                Role::Assistant => {
+                    if !message.tool_uses.is_empty() {
+                        this.tool_uses_by_assistant_message.insert(
+                            message.id,
+                            message
+                                .tool_uses
+                                .iter()
+                                .map(|tool_use| LanguageModelToolUse {
+                                    id: tool_use.id.clone(),
+                                    name: tool_use.name.clone().into(),
+                                    input: tool_use.input.clone(),
+                                })
+                                .collect(),
+                        );
+                    }
+                }
+                Role::User => {
+                    if !message.tool_results.is_empty() {
+                        let tool_uses_by_user_message = this
+                            .tool_uses_by_user_message
+                            .entry(message.id)
+                            .or_default();
+
+                        for tool_result in &message.tool_results {
+                            let tool_use_id = tool_result.tool_use_id.clone();
+
+                            tool_uses_by_user_message.push(tool_use_id.clone());
+                            this.tool_results.insert(
+                                tool_use_id.clone(),
+                                LanguageModelToolResult {
+                                    tool_use_id,
+                                    is_error: tool_result.is_error,
+                                    content: tool_result.content.clone(),
+                                },
+                            );
+                        }
+                    }
+                }
+                Role::System => {}
+            }
+        }
+
+        this
+    }
+
     pub fn pending_tool_uses(&self) -> Vec<&PendingToolUse> {
         self.pending_tool_uses_by_id.values().collect()
     }
@@ -84,6 +143,17 @@ impl ToolUseState {
         tool_uses
     }
 
+    pub fn tool_results_for_message(&self, message_id: MessageId) -> Vec<&LanguageModelToolResult> {
+        let empty = Vec::new();
+
+        self.tool_uses_by_user_message
+            .get(&message_id)
+            .unwrap_or(&empty)
+            .iter()
+            .filter_map(|tool_use_id| self.tool_results.get(&tool_use_id))
+            .collect()
+    }
+
     pub fn message_has_tool_results(&self, message_id: MessageId) -> bool {
         self.tool_uses_by_user_message
             .get(&message_id)