rmcp_proxy/
sse_client.rs

1use http::HeaderName;
2/**
3 * Create a local server that proxies requests to a remote server over SSE.
4 */
5use rmcp::{
6    model::{ClientCapabilities, ClientInfo},
7    transport::{sse::SseTransport, stdio},
8    ServiceExt,
9};
10use std::{collections::HashMap, error::Error as StdError, str::FromStr};
11use tracing::info;
12
13use crate::proxy_handler::ProxyHandler;
14
15/// Configuration for the SSE client
16pub struct SseClientConfig {
17    pub url: String,
18    pub headers: HashMap<String, String>,
19}
20
21/// Run the SSE client
22///
23/// This function connects to a remote SSE server and exposes it as a stdio server.
24pub async fn run_sse_client(config: SseClientConfig) -> Result<(), Box<dyn StdError>> {
25    info!("Running SSE client with URL: {}", config.url);
26
27    // Create the header map
28    let mut headers = reqwest::header::HeaderMap::new();
29    for (key, value) in config.headers {
30        headers.insert(HeaderName::from_str(&key)?, value.parse()?);
31    }
32
33    // Create the reqwest client to be by the SSE client.
34    let client = reqwest::Client::builder()
35        .default_headers(headers)
36        .build()?;
37
38    // Create SSE transport
39    let transport = SseTransport::start_with_client(&config.url, client).await?;
40
41    // Create client info with full capabilities to ensure we can use all the server's features
42    let client_info = ClientInfo {
43        protocol_version: Default::default(),
44        capabilities: ClientCapabilities::builder()
45            .enable_experimental()
46            .enable_roots()
47            .enable_roots_list_changed()
48            .enable_sampling()
49            .build(),
50        ..Default::default()
51    };
52
53    // Create client service with transport
54    let client = client_info.serve(transport).await?;
55
56    // Get server info
57    let server_info = client.peer_info();
58    info!("Connected to server: {}", server_info.server_info.name);
59
60    // Create proxy handler
61    let proxy_handler = ProxyHandler::new(client);
62
63    // Create stdio transport
64    let stdio_transport = stdio();
65
66    // Create server with proxy handler and stdio transport
67    let server = proxy_handler.serve(stdio_transport).await?;
68
69    // Wait for completion
70    server.waiting().await?;
71
72    Ok(())
73}