lean_ctx/gateway_server/mcp/
proxy.rs

1//! `/mcp/{server}` — the governed MCP reverse proxy (GL#100).
2//!
3//! MCP Streamable HTTP has one endpoint per server: the client POSTs JSON-RPC
4//! frames (response arrives as `application/json` or as an SSE stream), GETs
5//! an optional server-push listen stream, and DELETEs its session. The
6//! gateway fronts each registered upstream under `/mcp/{id}`:
7//!
8//! - **Auth**: the proxy's Bearer guard runs first — org token or per-person
9//!   gateway key, exactly like the LLM channel. `/mcp/*` is deliberately not
10//!   a provider route, so the loopback provider-key fallback never applies.
11//! - **Credential isolation**: the caller's `Authorization` (their gateway
12//!   key) is always stripped; when the registry entry names an `auth_env`,
13//!   the gateway injects `Authorization: Bearer <env value>` upstream. Tool
14//!   credentials live in the gateway environment, never on laptops.
15//! - **Observe, don't touch**: request and response bytes pass through
16//!   verbatim (SSE responses are teed, never buffered-and-replayed). Only
17//!   POST exchanges are metered — `tools/call` is the billable unit; the GET
18//!   listen stream carries server-initiated traffic, not tool calls.
19//! - **Fail-open**: analysis/metering failures log and pass traffic through.
20//!
21//! Registry changes (config.toml edits) take effect on gateway restart, the
22//! same lifecycle as `gateway-keys.toml` (documented; live reload is an M4
23//! concern once enforcement makes it safety-relevant).
24
25use std::time::Instant;
26
27use axum::body::{Body, Bytes};
28use axum::extract::{Path, State};
29use axum::http::{HeaderMap, HeaderValue, Request, StatusCode};
30use axum::response::{IntoResponse, Response};
31use futures::StreamExt;
32
33use crate::core::config::ResolvedMcpServer;
34use crate::proxy::ProxyState;
35use crate::proxy::gateway_identity::GatewayTags;
36
37use super::frames::{self, ParsedRequest, RequestKind, ResponseInfo};
38use super::metering::{self, MeteredExchange};
39use super::store::McpEvent;
40
41/// Identity fallbacks, mirroring the LLM channel's honest defaults
42/// (`store::ANONYMOUS_PERSON`/`DEFAULT_PROJECT`): rows stay attributable in
43/// solo/loopback mode; strict gateways make keys mandatory anyway.
44const ANONYMOUS_PERSON: &str = "anonymous";
45const DEFAULT_PROJECT: &str = "default";
46
47/// Request-body ceiling for MCP POST frames. Tool *arguments* are small
48/// compared to LLM prompts; 8 MiB is generous without inviting abuse.
49const MAX_REQUEST_BODY: usize = 8 * 1024 * 1024;
50
51/// How many response bytes the analyzer will hold to find the JSON-RPC
52/// response frame. Beyond this the exchange is still passed through and
53/// metered, with tokens approximated from the byte count (documented in
54/// [`approx_tokens_from_bytes`]).
55const MAX_ANALYSIS_BYTES: usize = 8 * 1024 * 1024;
56
57/// The single entry point for every `/mcp/{server}` request. Mounted on the
58/// main proxy router (feature-gated in `proxy::start_proxy`), so it shares
59/// `ProxyState` — the upstream client and the registry snapshot.
60pub async fn handler(
61    State(state): State<ProxyState>,
62    Path(server_id): Path<String>,
63    req: Request<Body>,
64) -> Response {
65    let Some(server) = state
66        .mcp_servers
67        .iter()
68        .find(|s| s.id == server_id)
69        .cloned()
70    else {
71        return json_rpc_error(
72            StatusCode::NOT_FOUND,
73            &format!(
74                "unknown MCP server '{server_id}' — register it under [[gateway_server.mcp_servers]]"
75            ),
76        );
77    };
78
79    let tags = req
80        .extensions()
81        .get::<GatewayTags>()
82        .cloned()
83        .unwrap_or_default();
84
85    let method = req.method().clone();
86    let (parts, body) = req.into_parts();
87
88    // Upstream request: fixed registry URL (no path/query joining — the MCP
89    // endpoint is a single URL; not forwarding caller paths is the SSRF-
90    // narrowest possible surface), curated headers, gateway-held credential.
91    let mut upstream_headers = forwarded_request_headers(&parts.headers);
92    if let Err(resp) = inject_upstream_credential(&server, &mut upstream_headers) {
93        return *resp;
94    }
95
96    match method {
97        axum::http::Method::POST => {
98            let Ok(body_bytes) = axum::body::to_bytes(body, MAX_REQUEST_BODY).await else {
99                return json_rpc_error(
100                    StatusCode::PAYLOAD_TOO_LARGE,
101                    &format!("MCP request body exceeds {MAX_REQUEST_BODY} bytes"),
102                );
103            };
104            let parsed = frames::parse_request(&body_bytes);
105            let started = Instant::now();
106            let upstream = state
107                .client
108                .post(&server.url)
109                .headers(upstream_headers)
110                .body(body_bytes.to_vec())
111                .send()
112                .await;
113            relay_post_response(upstream, &server, parsed, tags, started).await
114        }
115        // GET opens the server-push listen stream; DELETE ends the session.
116        // Pure passthrough: no JSON-RPC exchange to meter here (tools/call
117        // always travels over POST).
118        axum::http::Method::GET | axum::http::Method::DELETE => {
119            let builder = if method == axum::http::Method::GET {
120                state.client.get(&server.url)
121            } else {
122                state.client.delete(&server.url)
123            };
124            match builder.headers(upstream_headers).send().await {
125                Ok(upstream) => passthrough_response(upstream),
126                Err(e) => upstream_unreachable(&server.id, &e),
127            }
128        }
129        _ => json_rpc_error(
130            StatusCode::METHOD_NOT_ALLOWED,
131            "MCP Streamable HTTP uses POST, GET and DELETE",
132        ),
133    }
134}
135
136/// Relays a POST response, teeing bytes into the frame analyzer so the
137/// exchange lands in `mcp_events` without delaying the client.
138async fn relay_post_response(
139    upstream: Result<reqwest::Response, reqwest::Error>,
140    server: &ResolvedMcpServer,
141    parsed: Option<ParsedRequest>,
142    tags: GatewayTags,
143    started: Instant,
144) -> Response {
145    let upstream = match upstream {
146        Ok(u) => u,
147        Err(e) => {
148            record_exchange(
149                server,
150                parsed.as_ref(),
151                &tags,
152                "upstream_error",
153                started.elapsed().as_millis(),
154                None,
155                0,
156            );
157            return upstream_unreachable(&server.id, &e);
158        }
159    };
160
161    let status = upstream.status();
162    let content_type = upstream
163        .headers()
164        .get(axum::http::header::CONTENT_TYPE)
165        .and_then(|v| v.to_str().ok())
166        .unwrap_or("")
167        .to_ascii_lowercase();
168    let headers = forwarded_response_headers(upstream.headers());
169
170    if metering::installed() && content_type.starts_with("text/event-stream") {
171        // SSE: tee the stream — verbatim passthrough, analyzer on the side,
172        // event recorded when the upstream closes the stream.
173        let analyzer = SseAnalyzer::new(server.clone(), parsed, tags, started, u16_status(status));
174        let teed = tee_sse(upstream.bytes_stream(), analyzer);
175        return build_response(status, headers, Body::from_stream(teed));
176    }
177
178    // JSON (or metering off, or an error body): buffer, analyze, relay.
179    // MCP JSON responses are single frames — buffering them is bounded by
180    // the upstream's own response discipline; the SSE path above covers the
181    // long-running case.
182    match upstream.bytes().await {
183        Ok(bytes) => {
184            if metering::installed() {
185                let info = parsed
186                    .as_ref()
187                    .and_then(|p| frames::analyze_response_json(&bytes, p.id.as_ref()));
188                let status_label = exchange_status(u16_status(status), info.as_ref());
189                record_exchange(
190                    server,
191                    parsed.as_ref(),
192                    &tags,
193                    status_label,
194                    started.elapsed().as_millis(),
195                    info,
196                    bytes.len() as u64,
197                );
198            }
199            build_response(status, headers, Body::from(bytes))
200        }
201        Err(e) => {
202            record_exchange(
203                server,
204                parsed.as_ref(),
205                &tags,
206                "upstream_error",
207                started.elapsed().as_millis(),
208                None,
209                0,
210            );
211            upstream_unreachable(&server.id, &e)
212        }
213    }
214}
215
216/// Streams an upstream response through untouched (GET listen / DELETE).
217fn passthrough_response(upstream: reqwest::Response) -> Response {
218    let status = upstream.status();
219    let headers = forwarded_response_headers(upstream.headers());
220    build_response(status, headers, Body::from_stream(upstream.bytes_stream()))
221}
222
223/// Observes teed SSE bytes and books the exchange when the stream ends.
224struct SseAnalyzer {
225    server: ResolvedMcpServer,
226    parsed: Option<ParsedRequest>,
227    tags: GatewayTags,
228    started: Instant,
229    http_status: u16,
230    /// Raw bytes held for frame analysis; dropped once the response frame is
231    /// found or the cap is passed (then only `total_bytes` keeps counting).
232    buffer: Option<Vec<u8>>,
233    total_bytes: u64,
234    found: Option<ResponseInfo>,
235}
236
237impl SseAnalyzer {
238    fn new(
239        server: ResolvedMcpServer,
240        parsed: Option<ParsedRequest>,
241        tags: GatewayTags,
242        started: Instant,
243        http_status: u16,
244    ) -> Self {
245        Self {
246            server,
247            parsed,
248            tags,
249            started,
250            http_status,
251            buffer: Some(Vec::new()),
252            total_bytes: 0,
253            found: None,
254        }
255    }
256
257    fn feed(&mut self, chunk: &[u8]) {
258        self.total_bytes += chunk.len() as u64;
259        if self.found.is_some() {
260            return;
261        }
262        let Some(buf) = self.buffer.as_mut() else {
263            return;
264        };
265        buf.extend_from_slice(chunk);
266        // Only re-scan when a complete SSE event boundary is in the buffer.
267        if chunk.windows(2).any(|w| w == b"\n\n") || buf.windows(2).any(|w| w == b"\n\n") {
268            let text = String::from_utf8_lossy(buf);
269            if let Some(info) = self
270                .parsed
271                .as_ref()
272                .and_then(|p| frames::analyze_response_sse(&text, p.id.as_ref()))
273            {
274                self.found = Some(info);
275                self.buffer = None;
276                return;
277            }
278        }
279        if buf.len() > MAX_ANALYSIS_BYTES {
280            self.buffer = None;
281        }
282    }
283
284    fn finish(self) {
285        let status_label = exchange_status(self.http_status, self.found.as_ref());
286        record_exchange(
287            &self.server,
288            self.parsed.as_ref(),
289            &self.tags,
290            status_label,
291            self.started.elapsed().as_millis(),
292            self.found,
293            self.total_bytes,
294        );
295    }
296}
297
298/// Byte-for-byte tee (same construction as `proxy::usage::tee_stream`): every
299/// chunk is forwarded unchanged; the analyzer observes on the side and books
300/// the exchange when the upstream ends the stream.
301fn tee_sse<S, E>(
302    inner: S,
303    analyzer: SseAnalyzer,
304) -> impl futures::Stream<Item = Result<Bytes, E>> + Send
305where
306    S: futures::Stream<Item = Result<Bytes, E>> + Send + Unpin + 'static,
307    E: Send + 'static,
308{
309    futures::stream::unfold(
310        (inner, Some(analyzer)),
311        |(mut inner, mut analyzer)| async move {
312            match inner.next().await {
313                Some(Ok(chunk)) => {
314                    if let Some(a) = analyzer.as_mut() {
315                        a.feed(&chunk);
316                    }
317                    Some((Ok(chunk), (inner, analyzer)))
318                }
319                Some(err) => Some((err, (inner, analyzer))),
320                None => {
321                    if let Some(a) = analyzer.take() {
322                        a.finish();
323                    }
324                    None
325                }
326            }
327        },
328    )
329}
330
331/// Books one exchange into the metering sink (fail-open, never blocks).
332fn record_exchange(
333    server: &ResolvedMcpServer,
334    parsed: Option<&ParsedRequest>,
335    tags: &GatewayTags,
336    status: &str,
337    duration_ms: u128,
338    info: Option<ResponseInfo>,
339    raw_bytes: u64,
340) {
341    if !metering::installed() {
342        return;
343    }
344    let (method, tool) = match parsed.map(|p| &p.kind) {
345        Some(RequestKind::ToolsCall { tool }) => ("tools/call".to_string(), Some(tool.clone())),
346        Some(kind) => (kind.method_label().to_string(), None),
347        // Notification / batch / non-JSON body: still a real exchange.
348        None => ("passthrough".to_string(), None),
349    };
350    let (result_bytes, result_tokens, inventory) = match info {
351        Some(i) => (i.result_bytes, i.result_tokens, i.tools),
352        // No parsed frame (oversized stream / foreign shape): honest byte
353        // count with the documented byte→token approximation.
354        None => (raw_bytes, approx_tokens_from_bytes(raw_bytes), None),
355    };
356    metering::record(MeteredExchange {
357        event: McpEvent {
358            person: tags
359                .person
360                .clone()
361                .unwrap_or_else(|| ANONYMOUS_PERSON.to_string()),
362            team: tags.team.clone(),
363            project: tags
364                .project
365                .clone()
366                .unwrap_or_else(|| DEFAULT_PROJECT.to_string()),
367            server_id: server.id.clone(),
368            method,
369            tool,
370            status: status.to_string(),
371            duration_ms: i64::try_from(duration_ms).unwrap_or(i64::MAX),
372            result_bytes: i64::try_from(result_bytes).unwrap_or(i64::MAX),
373            result_tokens: i64::try_from(result_tokens).unwrap_or(i64::MAX),
374            // Priced by the writer (one pricing table load per process).
375            context_cost_usd: 0.0,
376            reference_model: None,
377        },
378        inventory,
379    });
380}
381
382/// `ok` | `error` | `upstream_error` — the three-valued status column.
383fn exchange_status(http_status: u16, info: Option<&ResponseInfo>) -> &'static str {
384    if info.is_some_and(|i| i.is_error) {
385        "error"
386    } else if http_status >= 400 {
387        "upstream_error"
388    } else {
389        "ok"
390    }
391}
392
393/// Tokens from bytes when no frame could be parsed: the standard ≈4 bytes per
394/// token heuristic for o200k-family tokenizers, floor 1 for non-empty bodies.
395/// An approximation is honest here — the alternative (0) would silently erase
396/// real context volume from the cost story.
397fn approx_tokens_from_bytes(bytes: u64) -> u64 {
398    if bytes == 0 { 0 } else { (bytes / 4).max(1) }
399}
400
401/// Request headers the upstream needs — and nothing else. Notably absent:
402/// the caller's `Authorization`/`x-api-key` (their *gateway* credential must
403/// never reach a tool server) and `x-leanctx-project` (internal tag).
404fn forwarded_request_headers(incoming: &HeaderMap) -> HeaderMap {
405    const FORWARDED: &[&str] = &[
406        "content-type",
407        "accept",
408        "accept-encoding",
409        "user-agent",
410        "mcp-session-id",
411        "mcp-protocol-version",
412        "last-event-id",
413    ];
414    let mut out = HeaderMap::new();
415    for name in FORWARDED {
416        if let Some(v) = incoming.get(*name)
417            && let Ok(name) = axum::http::header::HeaderName::from_bytes(name.as_bytes())
418        {
419            out.insert(name, v.clone());
420        }
421    }
422    out
423}
424
425/// Response headers relayed to the caller: hop-by-hop headers and
426/// `content-length` stay behind (axum reframes the body), everything else —
427/// notably `mcp-session-id` and `content-type` — passes through.
428fn forwarded_response_headers(upstream: &HeaderMap) -> Vec<(String, HeaderValue)> {
429    const SKIP: &[&str] = &[
430        "connection",
431        "keep-alive",
432        "proxy-authenticate",
433        "proxy-authorization",
434        "te",
435        "trailer",
436        "transfer-encoding",
437        "upgrade",
438        "content-length",
439    ];
440    upstream
441        .iter()
442        .filter(|(name, _)| !SKIP.contains(&name.as_str()))
443        .map(|(name, value)| (name.as_str().to_string(), value.clone()))
444        .collect()
445}
446
447/// Injects the gateway-held upstream credential (`auth_env`). A configured-
448/// but-missing env var is a deployment error and surfaces as a loud 502 —
449/// the same contract as the LLM registry's `api_key_env`. (The `Err` response
450/// is boxed: it only exists on the misconfiguration path.)
451fn inject_upstream_credential(
452    server: &ResolvedMcpServer,
453    headers: &mut HeaderMap,
454) -> Result<(), Box<Response>> {
455    let Some(env_name) = server.auth_env.as_deref() else {
456        return Ok(());
457    };
458    let key = std::env::var(env_name)
459        .ok()
460        .filter(|k| !k.trim().is_empty());
461    let Some(key) = key else {
462        tracing::error!(
463            "mcp proxy: server '{}' configures auth_env='{env_name}' but the variable is \
464             unset/empty — cannot authenticate upstream (502)",
465            server.id
466        );
467        return Err(Box::new(json_rpc_error(
468            StatusCode::BAD_GATEWAY,
469            &format!(
470                "gateway misconfiguration: auth_env '{env_name}' for MCP server '{}' is unset",
471                server.id
472            ),
473        )));
474    };
475    if let Ok(v) = HeaderValue::from_str(&format!("Bearer {key}")) {
476        headers.insert(axum::http::header::AUTHORIZATION, v);
477        Ok(())
478    } else {
479        tracing::error!("mcp proxy: credential from {env_name} contains invalid header bytes");
480        Err(Box::new(json_rpc_error(
481            StatusCode::BAD_GATEWAY,
482            &format!("gateway misconfiguration: credential in '{env_name}' is not header-safe"),
483        )))
484    }
485}
486
487fn build_response(
488    status: reqwest::StatusCode,
489    headers: Vec<(String, HeaderValue)>,
490    body: Body,
491) -> Response {
492    let mut resp = Response::new(body);
493    *resp.status_mut() = StatusCode::from_u16(status.as_u16()).unwrap_or(StatusCode::BAD_GATEWAY);
494    for (name, value) in headers {
495        if let Ok(name) = axum::http::header::HeaderName::from_bytes(name.as_bytes()) {
496            resp.headers_mut().append(name, value);
497        }
498    }
499    resp
500}
501
502fn u16_status(status: reqwest::StatusCode) -> u16 {
503    status.as_u16()
504}
505
506fn upstream_unreachable(server_id: &str, err: &reqwest::Error) -> Response {
507    tracing::warn!("mcp proxy: upstream '{server_id}' unreachable: {err}");
508    json_rpc_error(
509        StatusCode::BAD_GATEWAY,
510        &format!("MCP server '{server_id}' is unreachable through the gateway"),
511    )
512}
513
514/// Error bodies stay in JSON-RPC shape so MCP clients surface them cleanly
515/// instead of choking on a bare-text proxy error.
516fn json_rpc_error(status: StatusCode, message: &str) -> Response {
517    (
518        status,
519        axum::Json(serde_json::json!({
520            "jsonrpc": "2.0",
521            "id": null,
522            "error": { "code": -32000, "message": message }
523        })),
524    )
525        .into_response()
526}
527
528#[cfg(test)]
529mod tests {
530    use super::*;
531
532    #[test]
533    fn request_header_curation_strips_caller_credentials() {
534        let mut incoming = HeaderMap::new();
535        incoming.insert("authorization", "Bearer gk-personal-key".parse().unwrap());
536        incoming.insert("x-api-key", "sk-something".parse().unwrap());
537        incoming.insert("x-leanctx-project", "secret-project".parse().unwrap());
538        incoming.insert("content-type", "application/json".parse().unwrap());
539        incoming.insert(
540            "accept",
541            "application/json, text/event-stream".parse().unwrap(),
542        );
543        incoming.insert("mcp-session-id", "abc123".parse().unwrap());
544        incoming.insert("mcp-protocol-version", "2025-06-18".parse().unwrap());
545        incoming.insert("last-event-id", "7".parse().unwrap());
546        incoming.insert("cookie", "session=steal-me".parse().unwrap());
547
548        let out = forwarded_request_headers(&incoming);
549        assert!(
550            out.get("authorization").is_none(),
551            "gateway key must not leak"
552        );
553        assert!(out.get("x-api-key").is_none());
554        assert!(out.get("x-leanctx-project").is_none());
555        assert!(out.get("cookie").is_none());
556        assert_eq!(out.get("mcp-session-id").unwrap(), "abc123");
557        assert_eq!(out.get("mcp-protocol-version").unwrap(), "2025-06-18");
558        assert_eq!(out.get("last-event-id").unwrap(), "7");
559        assert_eq!(out.get("content-type").unwrap(), "application/json");
560    }
561
562    #[test]
563    fn credential_injection_is_loud_on_missing_env_and_replaces_caller_auth() {
564        let _lock = crate::core::data_dir::test_env_lock();
565        let server = ResolvedMcpServer {
566            id: "github".into(),
567            url: "https://api.githubcopilot.com/mcp".into(),
568            auth_env: Some("LC_TEST_MCP_PAT".into()),
569        };
570
571        crate::test_env::remove_var("LC_TEST_MCP_PAT");
572        let mut headers = HeaderMap::new();
573        assert!(
574            inject_upstream_credential(&server, &mut headers).is_err(),
575            "missing env must 502, never forward the caller's key"
576        );
577
578        crate::test_env::set_var("LC_TEST_MCP_PAT", "ghp-upstream");
579        let mut headers = HeaderMap::new();
580        inject_upstream_credential(&server, &mut headers).expect("env present");
581        assert_eq!(headers.get("authorization").unwrap(), "Bearer ghp-upstream");
582        crate::test_env::remove_var("LC_TEST_MCP_PAT");
583
584        // No auth_env → no header injected (public upstream).
585        let open = ResolvedMcpServer {
586            id: "open".into(),
587            url: "https://mcp.example.com/mcp".into(),
588            auth_env: None,
589        };
590        let mut headers = HeaderMap::new();
591        inject_upstream_credential(&open, &mut headers).unwrap();
592        assert!(headers.get("authorization").is_none());
593    }
594
595    #[test]
596    fn response_header_relay_drops_hop_by_hop_and_length() {
597        let mut upstream = HeaderMap::new();
598        upstream.insert("content-type", "text/event-stream".parse().unwrap());
599        upstream.insert("mcp-session-id", "s-1".parse().unwrap());
600        upstream.insert("transfer-encoding", "chunked".parse().unwrap());
601        upstream.insert("content-length", "12".parse().unwrap());
602        upstream.insert("connection", "keep-alive".parse().unwrap());
603
604        let out = forwarded_response_headers(&upstream);
605        let names: Vec<&str> = out.iter().map(|(n, _)| n.as_str()).collect();
606        assert!(names.contains(&"content-type"));
607        assert!(names.contains(&"mcp-session-id"));
608        assert!(!names.contains(&"transfer-encoding"));
609        assert!(!names.contains(&"content-length"));
610        assert!(!names.contains(&"connection"));
611    }
612
613    #[test]
614    fn status_labels_and_byte_approximation_are_stable() {
615        assert_eq!(exchange_status(200, None), "ok");
616        assert_eq!(exchange_status(500, None), "upstream_error");
617        let err_info = ResponseInfo {
618            is_error: true,
619            result_bytes: 10,
620            result_tokens: 3,
621            tools: None,
622        };
623        assert_eq!(exchange_status(200, Some(&err_info)), "error");
624
625        assert_eq!(approx_tokens_from_bytes(0), 0);
626        assert_eq!(approx_tokens_from_bytes(2), 1, "non-empty floors at 1");
627        assert_eq!(approx_tokens_from_bytes(4000), 1000);
628    }
629
630    #[tokio::test]
631    async fn sse_tee_passes_bytes_through_verbatim() {
632        let server = ResolvedMcpServer {
633            id: "s".into(),
634            url: "https://mcp.example.com/mcp".into(),
635            auth_env: None,
636        };
637        let chunks: Vec<Result<Bytes, std::convert::Infallible>> = vec![
638            Ok(Bytes::from_static(b"data: {\"jsonrpc\":\"2.0\",\"id\":1,")),
639            Ok(Bytes::from_static(b"\"result\":{\"content\":[]}}\n\n")),
640        ];
641        let analyzer = SseAnalyzer::new(
642            server,
643            frames::parse_request(
644                br#"{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"t"}}"#,
645            ),
646            GatewayTags::default(),
647            Instant::now(),
648            200,
649        );
650        let teed = tee_sse(futures::stream::iter(chunks), analyzer);
651        let collected: Vec<_> = teed.collect().await;
652        assert_eq!(collected.len(), 2);
653        let all: Vec<u8> = collected
654            .into_iter()
655            .flat_map(|c| c.unwrap().to_vec())
656            .collect();
657        assert_eq!(
658            all, b"data: {\"jsonrpc\":\"2.0\",\"id\":1,\"result\":{\"content\":[]}}\n\n",
659            "tee must never mutate the byte stream"
660        );
661    }
662
663    #[test]
664    fn sse_analyzer_finds_the_frame_and_stops_buffering() {
665        let server = ResolvedMcpServer {
666            id: "s".into(),
667            url: "https://mcp.example.com/mcp".into(),
668            auth_env: None,
669        };
670        let mut a = SseAnalyzer::new(
671            server,
672            frames::parse_request(
673                br#"{"jsonrpc":"2.0","id":42,"method":"tools/call","params":{"name":"get_issue"}}"#,
674            ),
675            GatewayTags::default(),
676            Instant::now(),
677            200,
678        );
679        a.feed(b"data: {\"jsonrpc\":\"2.0\",\"id\":42,\"result\":{\"content\":[{\"type\":\"text\",\"text\":\"hi\"}]}}\n\n");
680        assert!(a.found.is_some(), "response frame must be detected");
681        assert!(a.buffer.is_none(), "buffer drops once the frame is found");
682        let before = a.total_bytes;
683        a.feed(b"data: {\"jsonrpc\":\"2.0\",\"method\":\"notifications/x\"}\n\n");
684        assert!(a.total_bytes > before, "bytes keep counting after the find");
685    }
686}
lean_ctx/gateway_server/mcp/proxy.rs

lean_ctx/gateway_server/mcp/
proxy.rs
