1use std::sync::Arc;
2
3use anyhow::{Result, anyhow, bail};
4use assistant_tool::{ActionLog, Tool, ToolResult, ToolSource};
5use gpui::{AnyWindowHandle, 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 _window: Option<AnyWindowHandle>,
81 cx: &mut App,
82 ) -> ToolResult {
83 if let Some(server) = self.server_manager.read(cx).get_server(&self.server_id) {
84 let tool_name = self.tool.name.clone();
85 let server_clone = server.clone();
86 let input_clone = input.clone();
87
88 cx.spawn(async move |_cx| {
89 let Some(protocol) = server_clone.client() else {
90 bail!("Context server not initialized");
91 };
92
93 let arguments = if let serde_json::Value::Object(map) = input_clone {
94 Some(map.into_iter().collect())
95 } else {
96 None
97 };
98
99 log::trace!(
100 "Running tool: {} with arguments: {:?}",
101 tool_name,
102 arguments
103 );
104 let response = protocol.run_tool(tool_name, arguments).await?;
105
106 let mut result = String::new();
107 for content in response.content {
108 match content {
109 types::ToolResponseContent::Text { text } => {
110 result.push_str(&text);
111 }
112 types::ToolResponseContent::Image { .. } => {
113 log::warn!("Ignoring image content from tool response");
114 }
115 types::ToolResponseContent::Resource { .. } => {
116 log::warn!("Ignoring resource content from tool response");
117 }
118 }
119 }
120 Ok(result)
121 })
122 .into()
123 } else {
124 Task::ready(Err(anyhow!("Context server not found"))).into()
125 }
126 }
127}