1use std::sync::Arc;
 2
 3use anyhow::{anyhow, bail};
 4use assistant_tool::Tool;
 5use gpui::{App, Entity, Task};
 6
 7use crate::manager::ContextServerManager;
 8use crate::types;
 9
10pub struct ContextServerTool {
11    server_manager: Entity<ContextServerManager>,
12    server_id: Arc<str>,
13    tool: types::Tool,
14}
15
16impl ContextServerTool {
17    pub fn new(
18        server_manager: Entity<ContextServerManager>,
19        server_id: impl Into<Arc<str>>,
20        tool: types::Tool,
21    ) -> Self {
22        Self {
23            server_manager,
24            server_id: server_id.into(),
25            tool,
26        }
27    }
28}
29
30impl Tool for ContextServerTool {
31    fn name(&self) -> String {
32        self.tool.name.clone()
33    }
34
35    fn description(&self) -> String {
36        self.tool.description.clone().unwrap_or_default()
37    }
38
39    fn input_schema(&self) -> serde_json::Value {
40        match &self.tool.input_schema {
41            serde_json::Value::Null => {
42                serde_json::json!({ "type": "object", "properties": [] })
43            }
44            serde_json::Value::Object(map) if map.is_empty() => {
45                serde_json::json!({ "type": "object", "properties": [] })
46            }
47            _ => self.tool.input_schema.clone(),
48        }
49    }
50
51    fn run(
52        self: std::sync::Arc<Self>,
53        input: serde_json::Value,
54        _project: gpui::WeakEntity<project::Project>,
55        cx: &mut App,
56    ) -> gpui::Task<gpui::Result<String>> {
57        if let Some(server) = self.server_manager.read(cx).get_server(&self.server_id) {
58            cx.foreground_executor().spawn({
59                let tool_name = self.tool.name.clone();
60                async move {
61                    let Some(protocol) = server.client() else {
62                        bail!("Context server not initialized");
63                    };
64
65                    let arguments = if let serde_json::Value::Object(map) = input {
66                        Some(map.into_iter().collect())
67                    } else {
68                        None
69                    };
70
71                    log::trace!(
72                        "Running tool: {} with arguments: {:?}",
73                        tool_name,
74                        arguments
75                    );
76                    let response = protocol.run_tool(tool_name, arguments).await?;
77
78                    let mut result = String::new();
79                    for content in response.content {
80                        match content {
81                            types::ToolResponseContent::Text { text } => {
82                                result.push_str(&text);
83                            }
84                            types::ToolResponseContent::Image { .. } => {
85                                log::warn!("Ignoring image content from tool response");
86                            }
87                            types::ToolResponseContent::Resource { .. } => {
88                                log::warn!("Ignoring resource content from tool response");
89                            }
90                        }
91                    }
92                    Ok(result)
93                }
94            })
95        } else {
96            Task::ready(Err(anyhow!("Context server not found")))
97        }
98    }
99}