tenant.openagents/omega

1use crate::{AgentToolOutput, AnyAgentTool, ToolCallEventStream, ToolInput};
2use agent_client_protocol::schema::v1 as acp;
3use anyhow::Result;
4use collections::{BTreeMap, HashMap};
5use context_server::{ContextServerId, client::NotificationSubscription};
6use futures::FutureExt as _;
7use gpui::{App, AppContext, AsyncApp, Context, Entity, EventEmitter, SharedString, Task};
8use language_model::{LanguageModelImage, LanguageModelImageExt, LanguageModelToolResultContent};
9use project::context_server_store::{ContextServerStatus, ContextServerStore};
10use std::sync::Arc;
11use util::{ResultExt, markdown::MarkdownEscaped};
12
13/// Maximum number of characters to show from a tool argument in the
14/// collapsed tool-call header. Longer values are truncated with an ellipsis.
15const MAX_INLINE_ARG_LEN: usize = 120;
16
17/// Generates a tool ID for an MCP tool that can be used in settings.
18///
19/// The format is `mcp:<server_id>:<tool_name>` to avoid collisions with built-in tools.
20pub fn mcp_tool_id(server_id: &str, tool_name: &str) -> String {
21    format!("mcp:{}:{}", server_id, tool_name)
22}
23
24pub struct ContextServerPrompt {
25    pub server_id: ContextServerId,
26    pub prompt: context_server::types::Prompt,
27}
28
29pub enum ContextServerRegistryEvent {
30    ToolsChanged,
31    PromptsChanged,
32}
33
34impl EventEmitter<ContextServerRegistryEvent> for ContextServerRegistry {}
35
36pub struct ContextServerRegistry {
37    server_store: Entity<ContextServerStore>,
38    registered_servers: HashMap<ContextServerId, RegisteredContextServer>,
39    _subscription: gpui::Subscription,
40}
41
42struct RegisteredContextServer {
43    tools: BTreeMap<SharedString, Arc<dyn AnyAgentTool>>,
44    prompts: BTreeMap<SharedString, ContextServerPrompt>,
45    load_tools: Task<Result<()>>,
46    load_prompts: Task<Result<()>>,
47    _tools_updated_subscription: Option<NotificationSubscription>,
48}
49
50impl ContextServerRegistry {
51    pub fn new(server_store: Entity<ContextServerStore>, cx: &mut Context<Self>) -> Self {
52        let mut this = Self {
53            server_store: server_store.clone(),
54            registered_servers: HashMap::default(),
55            _subscription: cx.subscribe(&server_store, Self::handle_context_server_store_event),
56        };
57        for server in server_store.read(cx).running_servers() {
58            this.reload_tools_for_server(server.id(), cx);
59            this.reload_prompts_for_server(server.id(), cx);
60        }
61        this
62    }
63
64    pub fn tools_for_server(
65        &self,
66        server_id: &ContextServerId,
67    ) -> impl Iterator<Item = &Arc<dyn AnyAgentTool>> {
68        self.registered_servers
69            .get(server_id)
70            .map(|server| server.tools.values())
71            .into_iter()
72            .flatten()
73    }
74
75    pub fn servers(
76        &self,
77    ) -> impl Iterator<
78        Item = (
79            &ContextServerId,
80            &BTreeMap<SharedString, Arc<dyn AnyAgentTool>>,
81        ),
82    > {
83        self.registered_servers
84            .iter()
85            .map(|(id, server)| (id, &server.tools))
86    }
87
88    pub fn prompts(&self) -> impl Iterator<Item = &ContextServerPrompt> {
89        self.registered_servers
90            .values()
91            .flat_map(|server| server.prompts.values())
92    }
93
94    pub fn find_prompt(
95        &self,
96        server_id: Option<&ContextServerId>,
97        name: &str,
98    ) -> Option<&ContextServerPrompt> {
99        if let Some(server_id) = server_id {
100            self.registered_servers
101                .get(server_id)
102                .and_then(|server| server.prompts.get(name))
103        } else {
104            self.registered_servers
105                .values()
106                .find_map(|server| server.prompts.get(name))
107        }
108    }
109
110    pub fn server_store(&self) -> &Entity<ContextServerStore> {
111        &self.server_store
112    }
113
114    fn get_or_register_server(
115        &mut self,
116        server_id: &ContextServerId,
117        cx: &mut Context<Self>,
118    ) -> &mut RegisteredContextServer {
119        self.registered_servers
120            .entry(server_id.clone())
121            .or_insert_with(|| Self::init_registered_server(server_id, &self.server_store, cx))
122    }
123
124    fn init_registered_server(
125        server_id: &ContextServerId,
126        server_store: &Entity<ContextServerStore>,
127        cx: &mut Context<Self>,
128    ) -> RegisteredContextServer {
129        let tools_updated_subscription = server_store
130            .read(cx)
131            .get_running_server(server_id)
132            .and_then(|server| {
133                let client = server.client()?;
134
135                if !client.capable(context_server::protocol::ServerCapability::Tools) {
136                    return None;
137                }
138
139                let server_id = server.id();
140                let this = cx.entity().downgrade();
141
142                Some(client.on_notification(
143                    "notifications/tools/list_changed",
144                    Box::new(move |_params, cx: AsyncApp| {
145                        let server_id = server_id.clone();
146                        let this = this.clone();
147                        cx.spawn(async move |cx| {
148                            this.update(cx, |this, cx| {
149                                log::info!(
150                                    "Received tools/list_changed notification for server {}",
151                                    server_id
152                                );
153                                this.reload_tools_for_server(server_id, cx);
154                            })
155                        })
156                        .detach();
157                    }),
158                ))
159            });
160
161        RegisteredContextServer {
162            tools: BTreeMap::default(),
163            prompts: BTreeMap::default(),
164            load_tools: Task::ready(Ok(())),
165            load_prompts: Task::ready(Ok(())),
166            _tools_updated_subscription: tools_updated_subscription,
167        }
168    }
169
170    fn reload_tools_for_server(&mut self, server_id: ContextServerId, cx: &mut Context<Self>) {
171        let Some(server) = self.server_store.read(cx).get_running_server(&server_id) else {
172            return;
173        };
174        let Some(client) = server.client() else {
175            return;
176        };
177
178        if !client.capable(context_server::protocol::ServerCapability::Tools) {
179            return;
180        }
181
182        let registered_server = self.get_or_register_server(&server_id, cx);
183        registered_server.load_tools = cx.spawn(async move |this, cx| {
184            let response = client
185                .request::<context_server::types::requests::ListTools>(())
186                .await;
187
188            this.update(cx, |this, cx| {
189                let Some(registered_server) = this.registered_servers.get_mut(&server_id) else {
190                    return;
191                };
192
193                registered_server.tools.clear();
194                if let Some(response) = response.log_err() {
195                    for tool in response.tools {
196                        let tool = Arc::new(ContextServerTool::new(
197                            this.server_store.clone(),
198                            server.id(),
199                            tool,
200                        ));
201                        registered_server.tools.insert(tool.name(), tool);
202                    }
203                    cx.emit(ContextServerRegistryEvent::ToolsChanged);
204                    cx.notify();
205                }
206            })
207        });
208    }
209
210    fn reload_prompts_for_server(&mut self, server_id: ContextServerId, cx: &mut Context<Self>) {
211        let Some(server) = self.server_store.read(cx).get_running_server(&server_id) else {
212            return;
213        };
214        let Some(client) = server.client() else {
215            return;
216        };
217        if !client.capable(context_server::protocol::ServerCapability::Prompts) {
218            return;
219        }
220
221        let registered_server = self.get_or_register_server(&server_id, cx);
222
223        registered_server.load_prompts = cx.spawn(async move |this, cx| {
224            let response = client
225                .request::<context_server::types::requests::PromptsList>(())
226                .await;
227
228            this.update(cx, |this, cx| {
229                let Some(registered_server) = this.registered_servers.get_mut(&server_id) else {
230                    return;
231                };
232
233                registered_server.prompts.clear();
234                if let Some(response) = response.log_err() {
235                    for prompt in response.prompts {
236                        let name: SharedString = prompt.name.clone().into();
237                        registered_server.prompts.insert(
238                            name,
239                            ContextServerPrompt {
240                                server_id: server_id.clone(),
241                                prompt,
242                            },
243                        );
244                    }
245                    cx.emit(ContextServerRegistryEvent::PromptsChanged);
246                    cx.notify();
247                }
248            })
249        });
250    }
251
252    fn handle_context_server_store_event(
253        &mut self,
254        _: Entity<ContextServerStore>,
255        event: &project::context_server_store::ServerStatusChangedEvent,
256        cx: &mut Context<Self>,
257    ) {
258        let project::context_server_store::ServerStatusChangedEvent { server_id, status } = event;
259
260        match status {
261            ContextServerStatus::Starting | ContextServerStatus::Authenticating => {}
262            ContextServerStatus::Running => {
263                self.reload_tools_for_server(server_id.clone(), cx);
264                self.reload_prompts_for_server(server_id.clone(), cx);
265            }
266            ContextServerStatus::Stopped
267            | ContextServerStatus::Error(_)
268            | ContextServerStatus::AuthRequired
269            | ContextServerStatus::ClientSecretRequired { .. } => {
270                if let Some(registered_server) = self.registered_servers.remove(server_id) {
271                    if !registered_server.tools.is_empty() {
272                        cx.emit(ContextServerRegistryEvent::ToolsChanged);
273                    }
274                    if !registered_server.prompts.is_empty() {
275                        cx.emit(ContextServerRegistryEvent::PromptsChanged);
276                    }
277                }
278                cx.notify();
279            }
280        };
281    }
282}
283
284struct ContextServerTool {
285    store: Entity<ContextServerStore>,
286    server_id: ContextServerId,
287    tool: context_server::types::Tool,
288}
289
290impl ContextServerTool {
291    fn new(
292        store: Entity<ContextServerStore>,
293        server_id: ContextServerId,
294        tool: context_server::types::Tool,
295    ) -> Self {
296        Self {
297            store,
298            server_id,
299            tool,
300        }
301    }
302}
303
304impl AnyAgentTool for ContextServerTool {
305    fn name(&self) -> SharedString {
306        self.tool.name.clone().into()
307    }
308
309    fn description(&self) -> SharedString {
310        self.tool.description.clone().unwrap_or_default().into()
311    }
312
313    fn kind(&self) -> acp::ToolKind {
314        acp::ToolKind::Other
315    }
316
317    fn initial_title(&self, input: serde_json::Value, _cx: &mut App) -> SharedString {
318        format_mcp_initial_title(&self.tool.name, &input).into()
319    }
320
321    fn input_schema(
322        &self,
323        format: language_model::LanguageModelToolSchemaFormat,
324    ) -> Result<serde_json::Value> {
325        let mut schema = self.tool.input_schema.clone();
326        language_model::tool_schema::adapt_schema_to_format(&mut schema, format)?;
327        Ok(match schema {
328            serde_json::Value::Null => {
329                serde_json::json!({ "type": "object", "properties": [] })
330            }
331            serde_json::Value::Object(map) if map.is_empty() => {
332                serde_json::json!({ "type": "object", "properties": [] })
333            }
334            _ => schema,
335        })
336    }
337
338    fn run(
339        self: Arc<Self>,
340        input: ToolInput<serde_json::Value>,
341        event_stream: ToolCallEventStream,
342        cx: &mut App,
343    ) -> Task<Result<AgentToolOutput, AgentToolOutput>> {
344        let Some(server) = self.store.read(cx).get_running_server(&self.server_id) else {
345            return Task::ready(Err(anyhow::anyhow!("Context server not found").into()));
346        };
347        let tool_name = self.tool.name.clone();
348        let tool_id = mcp_tool_id(&self.server_id.0, &self.tool.name);
349        let display_name = self.tool.name.clone();
350        let initial_title = self.initial_title(serde_json::Value::Null, cx);
351        let authorize =
352            event_stream.authorize_third_party_tool(initial_title, tool_id, display_name, cx);
353
354        cx.spawn(async move |cx| {
355            let input = input
356                .recv()
357                .await
358                .map_err(|e| anyhow::anyhow!(e.to_string()))?;
359
360            authorize
361                .await
362                .map_err(|e| anyhow::anyhow!(e.to_string()))?;
363
364            let Some(protocol) = server.client() else {
365                return Err(anyhow::anyhow!("Context server not initialized").into());
366            };
367
368            let arguments = if let serde_json::Value::Object(map) = input {
369                Some(map.into_iter().collect())
370            } else {
371                None
372            };
373
374            log::trace!(
375                "Running tool: {} with arguments: {:?}",
376                tool_name,
377                arguments
378            );
379
380            let request = protocol.request::<context_server::types::requests::CallTool>(
381                context_server::types::CallToolParams {
382                    name: tool_name,
383                    arguments,
384                    meta: None,
385                },
386            );
387
388            let response = futures::select! {
389                response = request.fuse() => response?,
390                _ = event_stream.cancelled_by_user().fuse() => {
391                    return Err(anyhow::anyhow!("MCP tool cancelled by user").into());
392                }
393            };
394
395            if response.is_error == Some(true) {
396                let error_message: String =
397                    response.content.iter().filter_map(|c| c.text()).collect();
398                return Err(anyhow::anyhow!(error_message).into());
399            }
400
401            let mut llm_output = Vec::new();
402            let mut tool_call_content = Vec::new();
403            let mut concatenated_text = String::new();
404            for content in response.content {
405                match content {
406                    context_server::types::ToolResponseContent::Text { text } => {
407                        concatenated_text.push_str(&text);
408                        tool_call_content.push(acp::ToolCallContent::Content(acp::Content::new(
409                            acp::ContentBlock::Text(acp::TextContent::new(text.clone())),
410                        )));
411                        llm_output.push(LanguageModelToolResultContent::Text(text.into()));
412                    }
413                    context_server::types::ToolResponseContent::Image { data, mime_type } => {
414                        tool_call_content.push(acp::ToolCallContent::Content(acp::Content::new(
415                            acp::ContentBlock::Image(acp::ImageContent::new(
416                                data.clone(),
417                                mime_type.clone(),
418                            )),
419                        )));
420                        let language_model_image = cx
421                            .background_spawn({
422                                let mime_type = mime_type.clone();
423                                async move {
424                                    LanguageModelImage::from_base64_image(&data, &mime_type)
425                                }
426                            })
427                            .await;
428                        match language_model_image {
429                            Ok(Some(image)) => {
430                                llm_output.push(LanguageModelToolResultContent::Image(image));
431                            }
432                            Ok(None) => {
433                                log::warn!(
434                                    "Skipping MCP tool response image with MIME type `{}` because it cannot be converted for language model input",
435                                    mime_type
436                                );
437                            }
438                            Err(error) => {
439                                log::warn!(
440                                    "Failed to convert MCP tool response image with MIME type `{}` for language model input: {:#}",
441                                    mime_type,
442                                    error
443                                );
444                            }
445                        }
446                    }
447                    context_server::types::ToolResponseContent::Audio { .. } => {
448                        log::warn!("Ignoring audio content from tool response");
449                    }
450                    context_server::types::ToolResponseContent::Resource { .. } => {
451                        log::warn!("Ignoring resource content from tool response");
452                    }
453                    context_server::types::ToolResponseContent::ResourceLink { .. } => {
454                        log::warn!("Ignoring resource link content from tool response");
455                    }
456                }
457            }
458            if !tool_call_content.is_empty() {
459                event_stream
460                    .update_fields(acp::ToolCallUpdateFields::new().content(tool_call_content));
461            }
462            let raw_output = serde_json::Value::String(concatenated_text);
463            Ok(AgentToolOutput {
464                raw_output,
465                llm_output,
466            })
467        })
468    }
469
470    fn replay(
471        &self,
472        _input: serde_json::Value,
473        _output: serde_json::Value,
474        _event_stream: ToolCallEventStream,
475        _cx: &mut App,
476    ) -> Result<()> {
477        Ok(())
478    }
479}
480
481/// Builds the header label shown for an MCP tool call. When the input is an
482/// object with a single string-valued field, the value is inlined next to the
483/// tool name so the primary argument (e.g. a URL, path, or query) is visible
484/// without expanding the input block — matching the UX of built-in tools like
485/// `Fetch`. All other shapes fall back to the tool name alone.
486fn format_mcp_initial_title(tool_name: &str, input: &serde_json::Value) -> String {
487    if let Some(value) = single_string_arg(input) {
488        let preview = truncate_chars(value, MAX_INLINE_ARG_LEN);
489        format!("Run MCP tool `{}` {}", tool_name, MarkdownEscaped(&preview))
490    } else {
491        format!("Run MCP tool `{}`", tool_name)
492    }
493}
494
495fn single_string_arg(input: &serde_json::Value) -> Option<&str> {
496    let obj = input.as_object()?;
497    if obj.len() != 1 {
498        return None;
499    }
500    obj.values().next()?.as_str()
501}
502
503fn truncate_chars(s: &str, max: usize) -> String {
504    if s.chars().count() <= max {
505        s.to_string()
506    } else {
507        let mut out: String = s.chars().take(max).collect();
508        out.push('…');
509        out
510    }
511}
512
513pub fn get_prompt(
514    server_store: &Entity<ContextServerStore>,
515    server_id: &ContextServerId,
516    prompt_name: &str,
517    arguments: HashMap<String, String>,
518    cx: &mut AsyncApp,
519) -> Task<Result<context_server::types::PromptsGetResponse>> {
520    let server = cx.update(|cx| server_store.read(cx).get_running_server(server_id));
521    let Some(server) = server else {
522        return Task::ready(Err(anyhow::anyhow!("Context server not found")));
523    };
524
525    let Some(protocol) = server.client() else {
526        return Task::ready(Err(anyhow::anyhow!("Context server not initialized")));
527    };
528
529    let prompt_name = prompt_name.to_string();
530
531    cx.background_spawn(async move {
532        let response = protocol
533            .request::<context_server::types::requests::PromptsGet>(
534                context_server::types::PromptsGetParams {
535                    name: prompt_name,
536                    arguments: (!arguments.is_empty()).then(|| arguments),
537                    meta: None,
538                },
539            )
540            .await?;
541
542        Ok(response)
543    })
544}
545
546#[cfg(test)]
547mod tests {
548    use super::*;
549
550    #[test]
551    fn test_mcp_tool_id_format() {
552        assert_eq!(
553            mcp_tool_id("filesystem", "read_file"),
554            "mcp:filesystem:read_file"
555        );
556        assert_eq!(
557            mcp_tool_id("github", "create_issue"),
558            "mcp:github:create_issue"
559        );
560        assert_eq!(
561            mcp_tool_id("my-custom-server", "do_something"),
562            "mcp:my-custom-server:do_something"
563        );
564        // Underscores in names
565        assert_eq!(mcp_tool_id("my_server", "my_tool"), "mcp:my_server:my_tool");
566    }
567
568    // Note: Tests for MCP tool ID collision with built-in tools and permission
569    // decisions are in crates/agent/src/tool_permissions.rs to avoid duplication.
570
571    #[test]
572    fn test_format_mcp_initial_title_inlines_single_string_arg() {
573        let input = serde_json::json!({ "url": "https://example.com/page" });
574        assert_eq!(
575            format_mcp_initial_title("open_url_in_browser", &input),
576            "Run MCP tool `open_url_in_browser` https://example.com/page"
577        );
578    }
579
580    #[test]
581    fn test_format_mcp_initial_title_no_args() {
582        let input = serde_json::json!({});
583        assert_eq!(
584            format_mcp_initial_title("cleanup", &input),
585            "Run MCP tool `cleanup`"
586        );
587    }
588
589    #[test]
590    fn test_format_mcp_initial_title_null_input() {
591        assert_eq!(
592            format_mcp_initial_title("cleanup", &serde_json::Value::Null),
593            "Run MCP tool `cleanup`"
594        );
595    }
596
597    #[test]
598    fn test_format_mcp_initial_title_multiple_fields_falls_back() {
599        let input = serde_json::json!({ "x": "a", "y": "b" });
600        assert_eq!(
601            format_mcp_initial_title("do_thing", &input),
602            "Run MCP tool `do_thing`"
603        );
604    }
605
606    #[test]
607    fn test_format_mcp_initial_title_non_string_field_falls_back() {
608        let input = serde_json::json!({ "count": 42 });
609        assert_eq!(
610            format_mcp_initial_title("tick", &input),
611            "Run MCP tool `tick`"
612        );
613    }
614
615    #[test]
616    fn test_format_mcp_initial_title_truncates_long_values() {
617        let long = "x".repeat(MAX_INLINE_ARG_LEN + 50);
618        let input = serde_json::json!({ "q": long });
619        let title = format_mcp_initial_title("search", &input);
620        assert!(
621            title.ends_with('…'),
622            "expected truncation ellipsis, got: {title}"
623        );
624        // Prefix + backticked name + space + MAX chars + ellipsis — no full 170-char value.
625        assert!(title.chars().count() < MAX_INLINE_ARG_LEN + 50);
626    }
627
628    #[test]
629    fn test_format_mcp_initial_title_escapes_markdown_in_value() {
630        let input = serde_json::json!({ "q": "**bold** _italic_" });
631        let title = format_mcp_initial_title("search", &input);
632        // Asterisks and underscores must be escaped so the header renders literally.
633        assert!(title.contains("\\*"), "expected \\*, got: {title}");
634        assert!(title.contains("\\_"), "expected \\_, got: {title}");
635    }
636
637    #[test]
638    fn test_truncate_chars_boundary() {
639        assert_eq!(truncate_chars("abc", 3), "abc");
640        assert_eq!(truncate_chars("abcd", 3), "abc…");
641    }
642
643    #[test]
644    fn test_truncate_chars_handles_multibyte() {
645        // "café" is 4 chars but 5 bytes — byte-based truncation would panic.
646        assert_eq!(truncate_chars("café", 4), "café");
647        assert_eq!(truncate_chars("café", 3), "caf…");
648    }
649
650    #[test]
651    fn test_single_string_arg_ignores_empty_string() {
652        // An empty string is still a string — we inline it rather than fall back,
653        // which lets callers tell "the server sent an empty arg" apart from
654        // "no args at all".
655        let input = serde_json::json!({ "q": "" });
656        assert_eq!(single_string_arg(&input), Some(""));
657    }
658}
659