返回 CodeWhale
wire.rs
根目录 / crates / tui / src / mcp / wire.rs
1 //! MCP wire-format helpers shared by the HTTP, SSE, streamable-HTTP, and
2 //! stdio transports: frame/response size ceilings, SSE event framing and
3 //! field parsing, and the error-text classifiers that decide whether a
4 //! failure is a stale session or a closed connection.
5 /// Hard ceiling on the SSE frame-assembly buffer. A server that never emits a
6 /// frame separator would otherwise grow it without bound (OOM DoS).
7 pub(super) const MAX_SSE_FRAME_BYTES: usize = 8 * 1024 * 1024;
8
9 /// Hard ceiling on a single MCP HTTP response body / stdio line. A misbehaving
10 /// or malicious server could otherwise stream an unbounded body (or a
11 /// newline-free multi-GB "line") and OOM the process at transport-read time,
12 /// before any transcript-level spillover applies.
13 pub(super) const MAX_MCP_RESPONSE_BYTES: usize = 16 * 1024 * 1024;
14
15 pub(super) fn is_mcp_stale_session_body(body: &str) -> bool {
16 let body = body.to_ascii_lowercase();
17 body.contains("session") && (body.contains("expired") || body.contains("invalid"))
18 }
19
20 /// A tool call worth replaying after drop→reconnect: either the server
21 /// rejected the session id, or the transport itself is gone (dead
22 /// pipe/socket) rather than merely idle.
23 pub(super) fn is_retriable_mcp_call_error(err: &anyhow::Error) -> bool {
24 if is_mcp_stale_session_error(err) {
25 return true;
26 }
27 let lower = format!("{err:#}").to_ascii_lowercase();
28 is_connection_closed_error_text(&lower)
29 }
30
31 pub(super) fn is_mcp_stale_session_error(err: &anyhow::Error) -> bool {
32 let err = format!("{err:#}");
33 let lower_err = err.to_ascii_lowercase();
34 err.contains("MCP Streamable HTTP session expired")
35 || err.contains("MCP session expired")
36 || err.contains("SSE transport closed")
37 // The exact bail text of a stdio transport whose child died (the
38 // EOF arm of `StdioTransport::recv`); without this arm a dead-child
39 // error missed the drop→reconnect→retry path that SSE closes get.
40 || err.contains("Stdio transport closed")
41 || (err.contains("MCP SSE POST send failed") && is_connection_closed_error_text(&lower_err))
42 || is_mcp_stale_session_body(&err)
43 }
44
45 pub(super) fn is_connection_closed_error_text(err: &str) -> bool {
46 err.contains("connection closed")
47 || err.contains("connection reset")
48 || err.contains("broken pipe")
49 || err.contains("unexpected eof")
50 || err.contains("forcibly closed")
51 }
52
53 pub(super) fn parse_sse_message_data(body: &str) -> Vec<Vec<u8>> {
54 let normalized = body.replace("\r\n", "\n");
55 let mut messages = Vec::new();
56
57 for block in normalized.split("\n\n") {
58 let mut event_type = "message";
59 let mut data = String::new();
60
61 for line in block.lines() {
62 if let Some(value) = sse_field_value(line, "event:") {
63 event_type = value;
64 } else if let Some(value) = sse_field_value(line, "data:") {
65 if !data.is_empty() {
66 data.push('\n');
67 }
68 data.push_str(value);
69 }
70 }
71
72 if event_type != "message" || data.trim().is_empty() {
73 continue;
74 }
75
76 messages.push(data.trim().as_bytes().to_vec());
77 }
78
79 messages
80 }
81
82 // Retained for tests; the SSE transport now uses the byte-oriented twin.
83 #[cfg(test)]
84 pub(super) fn find_sse_event_separator(buffer: &str) -> Option<(usize, usize)> {
85 match (buffer.find("\n\n"), buffer.find("\r\n\r\n")) {
86 (Some(lf), Some(crlf)) if crlf < lf => Some((crlf, 4)),
87 (Some(lf), _) => Some((lf, 2)),
88 (_, Some(crlf)) => Some((crlf, 4)),
89 _ => None,
90 }
91 }
92
93 /// Byte-oriented twin of `find_sse_event_separator`. Used by the SSE
94 /// transport so it can accumulate RAW bytes and decode only complete event
95 /// blocks — a multi-byte UTF-8 char split across two network reads is never
96 /// corrupted to U+FFFD (the `\n`/`\r` separators are ASCII and can never fall
97 /// inside a multi-byte sequence).
98 pub(super) fn find_sse_event_separator_bytes(buffer: &[u8]) -> Option<(usize, usize)> {
99 let lf = buffer.windows(2).position(|w| w == b"\n\n");
100 let crlf = buffer.windows(4).position(|w| w == b"\r\n\r\n");
101 match (lf, crlf) {
102 (Some(lf), Some(crlf)) if crlf < lf => Some((crlf, 4)),
103 (Some(lf), _) => Some((lf, 2)),
104 (_, Some(crlf)) => Some((crlf, 4)),
105 _ => None,
106 }
107 }
108
109 pub(super) fn sse_field_value<'a>(line: &'a str, field: &str) -> Option<&'a str> {
110 let value = line.strip_prefix(field)?;
111 Some(value.strip_prefix(' ').unwrap_or(value))
112 }
113
113 lines RUST