858 lines · 32.4 KB · Rust
RawDownload
use anyhow::{Result, anyhow};use async_trait::async_trait;use collections::HashMap;use futures::{Stream, StreamExt};use gpui::BackgroundExecutor;use http_client::{AsyncBody, HttpClient, Request, Response, http::Method};use parking_lot::Mutex as SyncMutex;use std::{pin::Pin, sync::Arc}; use crate::oauth::{self, OAuthTokenProvider, WwwAuthenticate};use crate::transport::Transport;use crate::types; /// Typed errors returned by the HTTP transport that callers can downcast from/// `anyhow::Error` to handle specific failure modes.#[derive(Debug)]pub enum TransportError {    /// The server returned 401 and token refresh either wasn't possible or    /// failed. The caller should initiate the OAuth authorization flow.    AuthRequired { www_authenticate: WwwAuthenticate },} impl std::fmt::Display for TransportError {    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {        match self {            TransportError::AuthRequired { .. } => {                write!(f, "OAuth authorization required")            }        }    }} impl std::error::Error for TransportError {} // Constants from MCP specconst HEADER_SESSION_ID: &str = "Mcp-Session-Id";const HEADER_PROTOCOL_VERSION: &str = "MCP-Protocol-Version";const EVENT_STREAM_MIME_TYPE: &str = "text/event-stream";const JSON_MIME_TYPE: &str = "application/json"; /// HTTP Transport with session management and SSE supportpub struct HttpTransport {    http_client: Arc<dyn HttpClient>,    endpoint: String,    session_id: Arc<SyncMutex<Option<String>>>,    /// Negotiated MCP protocol version, populated by `set_protocol_version`    /// after the initialize handshake. From 2025-06-18 onward the server    /// requires clients to echo this in the `MCP-Protocol-Version` header on    /// every subsequent request.    protocol_version: Arc<SyncMutex<Option<String>>>,    executor: BackgroundExecutor,    response_tx: async_channel::Sender<String>,    response_rx: async_channel::Receiver<String>,    error_tx: async_channel::Sender<String>,    error_rx: async_channel::Receiver<String>,    /// Static headers to include in every request (e.g. from server config).    headers: HashMap<String, String>,    /// When set, the transport attaches `Authorization: Bearer` headers and    /// handles 401 responses with token refresh + retry.    token_provider: Option<Arc<dyn OAuthTokenProvider>>,    /// The challenge from the last 401 this transport gave up on; cleared at    /// the start of each send so it always describes the most recent attempt.    /// See [`Transport::auth_challenge`].    auth_challenge: SyncMutex<Option<WwwAuthenticate>>,} impl HttpTransport {    pub fn new(        http_client: Arc<dyn HttpClient>,        endpoint: String,        headers: HashMap<String, String>,        executor: BackgroundExecutor,    ) -> Self {        Self::new_with_token_provider(http_client, endpoint, headers, executor, None)    }     pub fn new_with_token_provider(        http_client: Arc<dyn HttpClient>,        endpoint: String,        headers: HashMap<String, String>,        executor: BackgroundExecutor,        token_provider: Option<Arc<dyn OAuthTokenProvider>>,    ) -> Self {        let (response_tx, response_rx) = async_channel::unbounded();        let (error_tx, error_rx) = async_channel::unbounded();         Self {            http_client,            executor,            endpoint,            session_id: Arc::new(SyncMutex::new(None)),            protocol_version: Arc::new(SyncMutex::new(None)),            response_tx,            response_rx,            error_tx,            error_rx,            headers,            token_provider,            auth_challenge: SyncMutex::new(None),        }    }     /// Build a POST request for the given message body, attaching all standard    /// headers (content-type, accept, session ID, static headers, and bearer    /// token if available).    fn build_request(&self, message: &[u8]) -> Result<http_client::Request<AsyncBody>> {        let mut request_builder = Request::builder()            .method(Method::POST)            .uri(&self.endpoint)            .header("Content-Type", JSON_MIME_TYPE)            .header(                "Accept",                format!("{}, {}", JSON_MIME_TYPE, EVENT_STREAM_MIME_TYPE),            );         for (key, value) in &self.headers {            request_builder = request_builder.header(key.as_str(), value.as_str());        }         // Attach bearer token when a token provider is present.        if let Some(token) = self.token_provider.as_ref().and_then(|p| p.access_token()) {            request_builder = request_builder.header("Authorization", format!("Bearer {}", token));        }         // Add session ID if we have one (except for initialize).        if let Some(ref session_id) = *self.session_id.lock() {            request_builder = request_builder.header(HEADER_SESSION_ID, session_id.as_str());        }         // Echo the negotiated protocol version once initialization has        // completed. Required by servers speaking MCP 2025-06-18 or later.        if let Some(ref version) = *self.protocol_version.lock()            && types::requires_protocol_version_header(version)        {            request_builder = request_builder.header(HEADER_PROTOCOL_VERSION, version.as_str());        }         Ok(request_builder.body(AsyncBody::from(message.to_vec()))?)    }     /// Record the challenge so it remains observable after the failed send    /// tears down the client (see [`Transport::auth_challenge`]), and build    /// the typed error for the send itself.    fn auth_required(&self, www_authenticate: WwwAuthenticate) -> anyhow::Error {        *self.auth_challenge.lock() = Some(www_authenticate.clone());        TransportError::AuthRequired { www_authenticate }.into()    }     /// Send a message and handle the response based on content type.    async fn send_message(&self, message: String) -> Result<()> {        // The same server instance can be restarted over this transport; a        // challenge recorded by a previous client generation must not be        // observed by the current one.        *self.auth_challenge.lock() = None;         let is_notification =            !message.contains("\"id\":") || message.contains("notifications/initialized");         // If we currently have no access token, try refreshing before sending        // the request so restored but expired sessions do not need an initial        // 401 round-trip before they can recover.        if let Some(ref provider) = self.token_provider {            if provider.access_token().is_none() {                provider.try_refresh().await.unwrap_or(false);            }        }         let request = self.build_request(message.as_bytes())?;        let mut response = self.http_client.send(request).await?;         // On 401, try refreshing the token and retry once.        if response.status().as_u16() == 401 {            let www_auth_header = response                .headers()                .get("www-authenticate")                .and_then(|v| v.to_str().ok())                .unwrap_or("Bearer");             let www_authenticate =                oauth::parse_www_authenticate(www_auth_header).unwrap_or(WwwAuthenticate {                    resource_metadata: None,                    scope: None,                    error: None,                    error_description: None,                });             if let Some(ref provider) = self.token_provider {                if provider.try_refresh().await.unwrap_or(false) {                    // Retry with the refreshed token.                    let retry_request = self.build_request(message.as_bytes())?;                    response = self.http_client.send(retry_request).await?;                     // If still 401 after refresh, give up.                    if response.status().as_u16() == 401 {                        return Err(self.auth_required(www_authenticate));                    }                } else {                    return Err(self.auth_required(www_authenticate));                }            } else {                return Err(self.auth_required(www_authenticate));            }        }         // Handle different response types based on status and content-type.        match response.status() {            status if status.is_success() => {                // Check content type                let content_type = response                    .headers()                    .get("content-type")                    .and_then(|v| v.to_str().ok());                 // Extract session ID from response headers if present                if let Some(session_id) = response                    .headers()                    .get(HEADER_SESSION_ID)                    .and_then(|v| v.to_str().ok())                {                    *self.session_id.lock() = Some(session_id.to_string());                    log::debug!("Session ID set: {}", session_id);                }                 match content_type {                    Some(ct) if ct.starts_with(JSON_MIME_TYPE) => {                        // JSON response - read and forward immediately                        let mut body = String::new();                        futures::AsyncReadExt::read_to_string(response.body_mut(), &mut body)                            .await?;                         // Only send non-empty responses                        if !body.is_empty() {                            self.response_tx                                .send(body)                                .await                                .map_err(|_| anyhow!("Failed to send JSON response"))?;                        }                    }                    Some(ct) if ct.starts_with(EVENT_STREAM_MIME_TYPE) => {                        // SSE stream - set up streaming                        self.setup_sse_stream(response).await?;                    }                    _ => {                        // For notifications, 202 Accepted with no content type is ok                        if is_notification && status.as_u16() == 202 {                            log::debug!("Notification accepted");                        } else {                            return Err(anyhow!("Unexpected content type: {:?}", content_type));                        }                    }                }            }            status if status.as_u16() == 202 => {                // Accepted - notification acknowledged, no response needed                log::debug!("Notification accepted");            }            _ => {                let mut error_body = String::new();                futures::AsyncReadExt::read_to_string(response.body_mut(), &mut error_body).await?;                 self.error_tx                    .send(format!("HTTP {}: {}", response.status(), error_body))                    .await                    .map_err(|_| anyhow!("Failed to send error"))?;            }        }         Ok(())    }     /// Set up SSE streaming from the response    async fn setup_sse_stream(&self, mut response: Response<AsyncBody>) -> Result<()> {        let response_tx = self.response_tx.clone();        let error_tx = self.error_tx.clone();         // Spawn a task to handle the SSE stream        self.executor            .spawn(async move {                let reader = futures::io::BufReader::new(response.body_mut());                let mut lines = futures::AsyncBufReadExt::lines(reader);                 let mut data_buffer = Vec::new();                let mut in_message = false;                 while let Some(line_result) = lines.next().await {                    match line_result {                        Ok(line) => {                            if line.is_empty() {                                // Empty line signals end of event                                if !data_buffer.is_empty() {                                    let message = data_buffer.join("\n");                                     // Filter out ping messages and empty data                                    if !message.trim().is_empty() && message != "ping" {                                        if let Err(e) = response_tx.send(message).await {                                            log::error!("Failed to send SSE message: {}", e);                                            break;                                        }                                    }                                    data_buffer.clear();                                }                                in_message = false;                            } else if let Some(data) = line.strip_prefix("data:") {                                // Handle data lines                                let data = data.trim();                                if !data.is_empty() {                                    // Check if this is a ping message                                    if data == "ping" {                                        log::trace!("Received SSE ping");                                        continue;                                    }                                    data_buffer.push(data.to_string());                                    in_message = true;                                }                            } else if line.starts_with("event:")                                || line.starts_with("id:")                                || line.starts_with("retry:")                            {                                // Ignore other SSE fields                                continue;                            } else if in_message {                                // Continuation of data                                data_buffer.push(line);                            }                        }                        Err(e) => {                            let _ = error_tx.send(format!("SSE stream error: {}", e)).await;                            break;                        }                    }                }            })            .detach();         Ok(())    }} #[async_trait]impl Transport for HttpTransport {    async fn send(&self, message: String) -> Result<()> {        self.send_message(message).await    }     fn receive(&self) -> Pin<Box<dyn Stream<Item = String> + Send>> {        Box::pin(self.response_rx.clone())    }     fn receive_err(&self) -> Pin<Box<dyn Stream<Item = String> + Send>> {        Box::pin(self.error_rx.clone())    }     fn set_protocol_version(&self, version: &str) {        *self.protocol_version.lock() = Some(version.to_string());    }     fn auth_challenge(&self) -> Option<WwwAuthenticate> {        self.auth_challenge.lock().clone()    }} impl Drop for HttpTransport {    fn drop(&mut self) {        // Try to cleanup session on drop        let http_client = self.http_client.clone();        let endpoint = self.endpoint.clone();        let session_id = self.session_id.lock().clone();        let protocol_version = self.protocol_version.lock().clone();        let headers = self.headers.clone();        let access_token = self.token_provider.as_ref().and_then(|p| p.access_token());         if let Some(session_id) = session_id {            self.executor                .spawn(async move {                    let mut request_builder = Request::builder()                        .method(Method::DELETE)                        .uri(&endpoint)                        .header(HEADER_SESSION_ID, &session_id);                     // Add static authentication headers.                    for (key, value) in headers {                        request_builder = request_builder.header(key.as_str(), value.as_str());                    }                     // Attach bearer token if available.                    if let Some(token) = access_token {                        request_builder =                            request_builder.header("Authorization", format!("Bearer {}", token));                    }                     // Stamp the negotiated MCP protocol version on the DELETE                    // too, matching what `build_request` does for POSTs.                    if let Some(ref version) = protocol_version                        && types::requires_protocol_version_header(version)                    {                        request_builder =                            request_builder.header(HEADER_PROTOCOL_VERSION, version.as_str());                    }                     let request = request_builder.body(AsyncBody::empty());                     if let Ok(request) = request {                        let _ = http_client.send(request).await;                    }                })                .detach();        }    }} #[cfg(test)]mod tests {    use super::*;    use async_trait::async_trait;    use futures::FutureExt as _;    use gpui::TestAppContext;    use parking_lot::Mutex as SyncMutex;    use std::{        sync::atomic::{AtomicBool, AtomicUsize, Ordering},        time::Duration,    };     /// A mock token provider that returns a configurable token and tracks    /// refresh attempts.    struct FakeTokenProvider {        token: SyncMutex<Option<String>>,        refreshed_token: SyncMutex<Option<String>>,        refresh_succeeds: AtomicBool,        refresh_count: AtomicUsize,    }     impl FakeTokenProvider {        fn new(token: Option<&str>, refresh_succeeds: bool) -> Arc<Self> {            Self::with_refreshed_token(token, None, refresh_succeeds)        }         fn with_refreshed_token(            token: Option<&str>,            refreshed_token: Option<&str>,            refresh_succeeds: bool,        ) -> Arc<Self> {            Arc::new(Self {                token: SyncMutex::new(token.map(String::from)),                refreshed_token: SyncMutex::new(refreshed_token.map(String::from)),                refresh_succeeds: AtomicBool::new(refresh_succeeds),                refresh_count: AtomicUsize::new(0),            })        }         fn set_token(&self, token: &str) {            *self.token.lock() = Some(token.to_string());        }         fn refresh_count(&self) -> usize {            self.refresh_count.load(Ordering::SeqCst)        }    }     #[async_trait]    impl OAuthTokenProvider for FakeTokenProvider {        fn access_token(&self) -> Option<String> {            self.token.lock().clone()        }         async fn try_refresh(&self) -> Result<bool> {            self.refresh_count.fetch_add(1, Ordering::SeqCst);             let refresh_succeeds = self.refresh_succeeds.load(Ordering::SeqCst);            if refresh_succeeds {                if let Some(token) = self.refreshed_token.lock().clone() {                    *self.token.lock() = Some(token);                }            }             Ok(refresh_succeeds)        }    }     fn make_fake_http_client(        handler: impl Fn(            http_client::Request<AsyncBody>,        ) -> std::pin::Pin<            Box<dyn std::future::Future<Output = anyhow::Result<Response<AsyncBody>>> + Send>,        > + Send        + Sync        + 'static,    ) -> Arc<dyn HttpClient> {        http_client::FakeHttpClient::create(handler) as Arc<dyn HttpClient>    }     fn json_response(status: u16, body: &str) -> anyhow::Result<Response<AsyncBody>> {        Ok(Response::builder()            .status(status)            .header("Content-Type", "application/json")            .body(AsyncBody::from(body.as_bytes().to_vec()))            .unwrap())    }     #[gpui::test]    async fn test_sse_data_field(cx: &mut TestAppContext) {

Lines 1–500 of 858.