assistant2: Add support for using tools provided by context servers (#21418)

This PR adds support to Assistant 2 for using tools provided by context
servers.

As part of this I introduced a new `ThreadStore`.

Release Notes:

- N/A

---------

Co-authored-by: Cole <cole@zed.dev>

crates/assistant2/src/thread_store.rs

@@ -0,0 +1,114 @@
+use std::sync::Arc;
+
+use anyhow::Result;
+use assistant_tool::{ToolId, ToolWorkingSet};
+use collections::HashMap;
+use context_server::manager::ContextServerManager;
+use context_server::{ContextServerFactoryRegistry, ContextServerTool};
+use gpui::{prelude::*, AppContext, Model, ModelContext, Task};
+use project::Project;
+use util::ResultExt as _;
+
+pub struct ThreadStore {
+    #[allow(unused)]
+    project: Model<Project>,
+    tools: Arc<ToolWorkingSet>,
+    context_server_manager: Model<ContextServerManager>,
+    context_server_tool_ids: HashMap<Arc<str>, Vec<ToolId>>,
+}
+
+impl ThreadStore {
+    pub fn new(
+        project: Model<Project>,
+        tools: Arc<ToolWorkingSet>,
+        cx: &mut AppContext,
+    ) -> Task<Result<Model<Self>>> {
+        cx.spawn(|mut cx| async move {
+            let this = cx.new_model(|cx: &mut ModelContext<Self>| {
+                let context_server_factory_registry =
+                    ContextServerFactoryRegistry::default_global(cx);
+                let context_server_manager = cx.new_model(|cx| {
+                    ContextServerManager::new(context_server_factory_registry, project.clone(), cx)
+                });
+
+                let this = Self {
+                    project,
+                    tools,
+                    context_server_manager,
+                    context_server_tool_ids: HashMap::default(),
+                };
+                this.register_context_server_handlers(cx);
+
+                this
+            })?;
+
+            Ok(this)
+        })
+    }
+
+    fn register_context_server_handlers(&self, cx: &mut ModelContext<Self>) {
+        cx.subscribe(
+            &self.context_server_manager.clone(),
+            Self::handle_context_server_event,
+        )
+        .detach();
+    }
+
+    fn handle_context_server_event(
+        &mut self,
+        context_server_manager: Model<ContextServerManager>,
+        event: &context_server::manager::Event,
+        cx: &mut ModelContext<Self>,
+    ) {
+        let tool_working_set = self.tools.clone();
+        match event {
+            context_server::manager::Event::ServerStarted { server_id } => {
+                if let Some(server) = context_server_manager.read(cx).get_server(server_id) {
+                    let context_server_manager = context_server_manager.clone();
+                    cx.spawn({
+                        let server = server.clone();
+                        let server_id = server_id.clone();
+                        |this, mut cx| async move {
+                            let Some(protocol) = server.client() else {
+                                return;
+                            };
+
+                            if protocol.capable(context_server::protocol::ServerCapability::Tools) {
+                                if let Some(tools) = protocol.list_tools().await.log_err() {
+                                    let tool_ids = tools
+                                        .tools
+                                        .into_iter()
+                                        .map(|tool| {
+                                            log::info!(
+                                                "registering context server tool: {:?}",
+                                                tool.name
+                                            );
+                                            tool_working_set.insert(Arc::new(
+                                                ContextServerTool::new(
+                                                    context_server_manager.clone(),
+                                                    server.id(),
+                                                    tool,
+                                                ),
+                                            ))
+                                        })
+                                        .collect::<Vec<_>>();
+
+                                    this.update(&mut cx, |this, _cx| {
+                                        this.context_server_tool_ids.insert(server_id, tool_ids);
+                                    })
+                                    .log_err();
+                                }
+                            }
+                        }
+                    })
+                    .detach();
+                }
+            }
+            context_server::manager::Event::ServerStopped { server_id } => {
+                if let Some(tool_ids) = self.context_server_tool_ids.remove(server_id) {
+                    tool_working_set.remove(&tool_ids);
+                }
+            }
+        }
+    }
+}