1use std::sync::Arc;
  2
  3use anyhow::{Result, anyhow, bail};
  4use assistant_tool::{ActionLog, Tool, ToolResult, 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, format: LanguageModelToolSchemaFormat) -> Result<serde_json::Value> {
 57        let mut schema = self.tool.input_schema.clone();
 58        assistant_tool::adapt_schema_to_format(&mut schema, format)?;
 59        Ok(match schema {
 60            serde_json::Value::Null => {
 61                serde_json::json!({ "type": "object", "properties": [] })
 62            }
 63            serde_json::Value::Object(map) if map.is_empty() => {
 64                serde_json::json!({ "type": "object", "properties": [] })
 65            }
 66            _ => schema,
 67        })
 68    }
 69
 70    fn ui_text(&self, _input: &serde_json::Value) -> String {
 71        format!("Run MCP tool `{}`", self.tool.name)
 72    }
 73
 74    fn run(
 75        self: Arc<Self>,
 76        input: serde_json::Value,
 77        _messages: &[LanguageModelRequestMessage],
 78        _project: Entity<Project>,
 79        _action_log: Entity<ActionLog>,
 80        cx: &mut App,
 81    ) -> ToolResult {
 82        if let Some(server) = self.server_manager.read(cx).get_server(&self.server_id) {
 83            let tool_name = self.tool.name.clone();
 84            let server_clone = server.clone();
 85            let input_clone = input.clone();
 86
 87            cx.spawn(async move |_cx| {
 88                let Some(protocol) = server_clone.client() else {
 89                    bail!("Context server not initialized");
 90                };
 91
 92                let arguments = if let serde_json::Value::Object(map) = input_clone {
 93                    Some(map.into_iter().collect())
 94                } else {
 95                    None
 96                };
 97
 98                log::trace!(
 99                    "Running tool: {} with arguments: {:?}",
100                    tool_name,
101                    arguments
102                );
103                let response = protocol.run_tool(tool_name, arguments).await?;
104
105                let mut result = String::new();
106                for content in response.content {
107                    match content {
108                        types::ToolResponseContent::Text { text } => {
109                            result.push_str(&text);
110                        }
111                        types::ToolResponseContent::Image { .. } => {
112                            log::warn!("Ignoring image content from tool response");
113                        }
114                        types::ToolResponseContent::Resource { .. } => {
115                            log::warn!("Ignoring resource content from tool response");
116                        }
117                    }
118                }
119                Ok(result)
120            })
121            .into()
122        } else {
123            Task::ready(Err(anyhow!("Context server not found"))).into()
124        }
125    }
126}