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