114 lines · 4.4 KB · Rust
RawDownload
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);                }            }        }    }}