tenant.openagents/omega

1//! This module implements parts of the Model Context Protocol.
2//!
3//! It handles the lifecycle messages, and provides a general interface to
4//! interacting with an MCP server. It uses the generic JSON-RPC client to
5//! read/write messages and the types from types.rs for serialization/deserialization
6//! of messages.
7
8use std::time::Duration;
9
10use anyhow::Result;
11use futures::{channel::oneshot, future::BoxFuture};
12use gpui::AsyncApp;
13use serde_json::Value;
14
15use crate::client::{Client, NotificationSubscription};
16use crate::oauth::WwwAuthenticate;
17use crate::types::{self, Notification, Request};
18
19pub struct ModelContextProtocol {
20    inner: Client,
21}
22
23impl ModelContextProtocol {
24    pub(crate) fn new(inner: Client) -> Self {
25        Self { inner }
26    }
27
28    fn supported_protocols() -> Vec<types::ProtocolVersion> {
29        vec![
30            types::ProtocolVersion(types::LATEST_PROTOCOL_VERSION.to_string()),
31            types::ProtocolVersion(types::VERSION_2025_06_18.to_string()),
32            types::ProtocolVersion(types::VERSION_2025_03_26.to_string()),
33            types::ProtocolVersion(types::VERSION_2024_11_05.to_string()),
34        ]
35    }
36
37    pub async fn initialize(
38        self,
39        client_info: types::Implementation,
40    ) -> Result<InitializedContextServerProtocol> {
41        let params = types::InitializeParams {
42            protocol_version: types::ProtocolVersion(types::LATEST_PROTOCOL_VERSION.to_string()),
43            capabilities: types::ClientCapabilities {
44                experimental: None,
45                sampling: None,
46                roots: None,
47            },
48            meta: None,
49            client_info,
50        };
51
52        let response: types::InitializeResponse = self
53            .inner
54            .request(types::requests::Initialize::METHOD, params)
55            .await?;
56
57        anyhow::ensure!(
58            Self::supported_protocols().contains(&response.protocol_version),
59            "Unsupported protocol version: {:?}",
60            response.protocol_version
61        );
62
63        log::trace!("mcp server info {:?}", response.server_info);
64
65        // Per MCP 2025-06-18, HTTP transport must attach the negotiated version
66        // as `MCP-Protocol-Version` on every post-initialize request.
67        self.inner
68            .set_protocol_version(&response.protocol_version.0);
69
70        let initialized_protocol = InitializedContextServerProtocol {
71            inner: self.inner,
72            initialize: response,
73        };
74
75        initialized_protocol.notify::<types::notifications::Initialized>(())?;
76
77        Ok(initialized_protocol)
78    }
79}
80
81pub struct InitializedContextServerProtocol {
82    inner: Client,
83    pub initialize: types::InitializeResponse,
84}
85
86#[derive(Debug, PartialEq, Clone, Copy)]
87pub enum ServerCapability {
88    Experimental,
89    Logging,
90    Prompts,
91    Resources,
92    Tools,
93}
94
95impl InitializedContextServerProtocol {
96    /// Check if the server supports a specific capability
97    pub fn capable(&self, capability: ServerCapability) -> bool {
98        match capability {
99            ServerCapability::Experimental => self.initialize.capabilities.experimental.is_some(),
100            ServerCapability::Logging => self.initialize.capabilities.logging.is_some(),
101            ServerCapability::Prompts => self.initialize.capabilities.prompts.is_some(),
102            ServerCapability::Resources => self.initialize.capabilities.resources.is_some(),
103            ServerCapability::Tools => self.initialize.capabilities.tools.is_some(),
104        }
105    }
106
107    pub async fn request<T: Request>(&self, params: T::Params) -> Result<T::Response> {
108        self.inner.request(T::METHOD, params).await
109    }
110
111    pub async fn request_with<T: Request>(
112        &self,
113        params: T::Params,
114        cancel_rx: Option<oneshot::Receiver<()>>,
115        timeout: Option<Duration>,
116    ) -> Result<T::Response> {
117        self.inner
118            .request_with(T::METHOD, params, cancel_rx, timeout)
119            .await
120    }
121
122    pub fn notify<T: Notification>(&self, params: T::Params) -> Result<()> {
123        self.inner.notify(T::METHOD, params)
124    }
125
126    /// A future that resolves once the underlying transport's output loop has
127    /// terminated — after a send failure, or when the client is dropped —
128    /// yielding the authentication challenge recorded by the transport if it
129    /// shut down on a `401 Unauthorized` response.
130    ///
131    /// Servers may accept `initialize` unauthenticated and only challenge a
132    /// later request or notification. Awaiting this is what lets the owner of
133    /// the connection notice such a challenge even when no request was in
134    /// flight to carry a typed error back. Returns `None` if the shutdown
135    /// signal was already claimed: there is a single signal per client.
136    pub fn wait_for_shutdown(&self) -> Option<BoxFuture<'static, Option<WwwAuthenticate>>> {
137        self.inner.wait_for_shutdown()
138    }
139
140    pub fn on_notification(
141        &self,
142        method: &'static str,
143        f: Box<dyn 'static + Send + FnMut(Value, AsyncApp)>,
144    ) -> NotificationSubscription {
145        self.inner.on_notification(method, f)
146    }
147}
148