Skip to content

Commit 3ea2931

Browse files
hopelyndhopelynd
committed
fix(mcp): harden SSE parsing, abort, and fallback matching
1 parent 4dcec7d commit 3ea2931

7 files changed

Lines changed: 370 additions & 64 deletions

File tree

‎packages/commands/src/commands/mcp/activate-hint.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ import { mcpMarketplaceDetailPage } from "bailian-cli-runtime";
55
export function isMcpNotActivated(error: unknown): boolean {
66
if (!(error instanceof BailianError)) return false;
77
const message = error.message;
8-
if (!/MCP request failed:\s*404\b/i.test(message)) return false;
8+
if (!/^MCP request failed:\s*404\b/i.test(message)) return false;
99
return /未开通|MCP不存在|MCP_IS_INVALID/i.test(message);
1010
}
1111

‎packages/commands/tests/mcp-activate-hint.test.ts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,12 @@ describe("mcp-activate-hint", () => {
2626
false,
2727
);
2828
expect(isMcpNotActivated(new Error("MCP不存在或未开通"))).toBe(false);
29+
// Nested wrapper phrase must not match (anchored at start).
30+
expect(
31+
isMcpNotActivated(
32+
new BailianError("MCP error (-32000): MCP request failed: 404 Not Found - 未开通"),
33+
),
34+
).toBe(false);
2935
});
3036

3137
test("hint 含对应 server 的 MCP 广场深链", () => {

‎packages/core/src/client/mcp-sse.ts‎

Lines changed: 63 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -114,8 +114,7 @@ export class McpSseClient {
114114
private async openSse(): Promise<void> {
115115
if (this.abortController) return;
116116

117-
// use shared abortController：header wait use timer abort；after getting header, clearTimeout,
118-
// the long-lived stream is only ended by close()/session abort (compatible with Node 18, no AbortSignal.any).
117+
// One abortController for header/error-body wait; clear timer before the long-lived stream.
119118
this.abortController = new AbortController();
120119
const timeoutMs = this.deps.settings.timeout * 1000;
121120
let headerTimedOut = false;
@@ -146,6 +145,8 @@ export class McpSseClient {
146145
});
147146
} catch (error) {
148147
clearTimeout(headerTimer);
148+
// Allow a later initialize() to openSse again on this instance.
149+
this.abortController = undefined;
149150
if (this.closed) {
150151
throw new BailianError("MCP SSE session closed.", ExitCode.GENERAL);
151152
}
@@ -155,27 +156,43 @@ export class McpSseClient {
155156
throw new BailianError(
156157
`MCP SSE request failed: ${error instanceof Error ? error.message : String(error)}`,
157158
ExitCode.NETWORK,
159+
undefined,
160+
{ cause: error },
158161
);
159162
}
160-
// 已收到响应头：取消 header 等待，后续仅由 abortController 结束流。
161-
clearTimeout(headerTimer);
162163

163164
if (this.deps.settings.verbose) {
164165
console.error(`< ${response.status} ${response.statusText}`);
165166
}
166167

167168
if (!response.ok) {
169+
// Keep headerTimer until error body is read (or times out).
168170
let errMsg = `MCP request failed: ${response.status} ${response.statusText}`;
169171
try {
170172
const errBody = await response.text();
171173
if (errBody) errMsg += ` - ${errBody.slice(0, 500)}`;
172-
} catch {
173-
/* ignore */
174+
} catch (error) {
175+
clearTimeout(headerTimer);
176+
this.abortController = undefined;
177+
if (this.closed) {
178+
throw new BailianError("MCP SSE session closed.", ExitCode.GENERAL);
179+
}
180+
if (headerTimedOut) {
181+
throw new BailianError(
182+
"MCP SSE timed out reading error response body.",
183+
ExitCode.TIMEOUT,
184+
);
185+
}
186+
throw new BailianError(errMsg, ExitCode.GENERAL, undefined, { cause: error });
174187
}
175-
// Throw only — do not rejectEndpoint; this path never awaits endpointReady.
188+
clearTimeout(headerTimer);
189+
this.abortController = undefined;
190+
// Do not rejectEndpoint — openSse never awaits endpointReady on this path.
176191
throw new BailianError(errMsg, ExitCode.GENERAL);
177192
}
178193

194+
clearTimeout(headerTimer);
195+
179196
void this.consumeSse(response).catch((error) => {
180197
if (this.closed) return;
181198
const reason =
@@ -247,8 +264,7 @@ export class McpSseClient {
247264
throw error;
248265
}
249266

250-
// Stream ended after endpoint: mark session dead and wake pending; do not throw,
251-
// so void consumeSse().catch does not surface an extra unhandled rejection.
267+
// After endpoint: mark dead and wake pending; don't throw (avoid unhandledRejection).
252268
this.markStreamEnded(new BailianError("MCP SSE stream ended unexpectedly.", ExitCode.GENERAL));
253269
}
254270

@@ -332,12 +348,24 @@ export class McpSseClient {
332348
}
333349

334350
const timeoutMs = this.deps.settings.timeout * 1000;
335-
const res = await fetch(this.messageUrl, {
336-
method: "POST",
337-
headers,
338-
body: JSON.stringify(body),
339-
signal: AbortSignal.timeout(timeoutMs),
340-
});
351+
// Combine per-RPC timeout with session abort so close() cancels in-flight POSTs.
352+
const requestSignal = createLinkedAbortSignal(timeoutMs, this.abortController?.signal);
353+
let res: Response;
354+
try {
355+
res = await fetch(this.messageUrl, {
356+
method: "POST",
357+
headers,
358+
body: JSON.stringify(body),
359+
signal: requestSignal.signal,
360+
});
361+
} catch (error) {
362+
if (this.closed) {
363+
throw new BailianError("MCP SSE session closed.", ExitCode.GENERAL);
364+
}
365+
throw error;
366+
} finally {
367+
requestSignal.cleanup();
368+
}
341369

342370
if (this.deps.settings.verbose) {
343371
console.error(`< ${res.status} ${res.statusText}`);
@@ -406,3 +434,23 @@ function cancellableTimeoutReject(
406434
},
407435
};
408436
}
437+
438+
/** Timeout + optional parent abort without AbortSignal.any (Node 18). */
439+
function createLinkedAbortSignal(
440+
timeoutMs: number,
441+
parentSignal?: AbortSignal,
442+
): { signal: AbortSignal; cleanup: () => void } {
443+
const controller = new AbortController();
444+
const timeout = setTimeout(() => controller.abort(), timeoutMs);
445+
const abortFromParent = () => controller.abort(parentSignal?.reason);
446+
const cleanup = () => {
447+
clearTimeout(timeout);
448+
parentSignal?.removeEventListener("abort", abortFromParent);
449+
};
450+
451+
if (parentSignal?.aborted) abortFromParent();
452+
else parentSignal?.addEventListener("abort", abortFromParent, { once: true });
453+
controller.signal.addEventListener("abort", cleanup, { once: true });
454+
455+
return { signal: controller.signal, cleanup };
456+
}

‎packages/core/src/client/mcp.ts‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -70,20 +70,19 @@ export function bailianMcpSsePath(serverCode: string): string {
7070

7171
/**
7272
* True when Streamable HTTP is unsupported and classic SSE fallback should be tried.
73-
* Match HTTP wrapper text `MCP request failed: 405` only — not JSON-RPC `MCP error (405)`.
74-
* Bailian HTTP 404 (not activated) is intentionally excluded.
73+
* Anchored to HTTP wrapper text only (not JSON-RPC / nested copies). Bailian 404 excluded.
7574
*/
7675
export function isStreamableHttpUnsupported(error: unknown): boolean {
7776
if (!(error instanceof BailianError)) return false;
78-
return /MCP request failed:\s*405\b/i.test(error.message);
77+
return /^MCP request failed:\s*405\b/i.test(error.message);
7978
}
8079

8180
/**
8281
* SSE fallback for `--url` (official backwards-compat: same URL, HTTP 405/404 then GET SSE).
8382
*/
8483
export function isUrlOverrideSseFallbackCandidate(error: unknown): boolean {
8584
if (!(error instanceof BailianError)) return false;
86-
return /MCP request failed:\s*(405|404)\b/i.test(error.message);
85+
return /^MCP request failed:\s*(405|404)\b/i.test(error.message);
8786
}
8887

8988
export type McpConnectedClient = {

‎packages/core/src/client/stream.ts‎

Lines changed: 95 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -7,76 +7,127 @@ export interface ServerSentEvent {
77
id?: string;
88
}
99

10+
/** Normalize CRLF/CR to LF; hold a trailing `\r` so a split CRLF is not double-broken. */
11+
function takeNormalizedSseLines(buffer: string): { lines: string[]; rest: string } {
12+
let text = buffer;
13+
let holdTrailingCr = false;
14+
if (text.endsWith("\r")) {
15+
holdTrailingCr = true;
16+
text = text.slice(0, -1);
17+
}
18+
19+
text = text.replace(/\r\n/g, "\n").replace(/\r/g, "\n");
20+
const parts = text.split("\n");
21+
const incomplete = parts.pop() ?? "";
22+
return {
23+
lines: parts,
24+
rest: holdTrailingCr ? `${incomplete}\r` : incomplete,
25+
};
26+
}
27+
28+
function applySseLine(
29+
line: string,
30+
event: Partial<ServerSentEvent>,
31+
maxBuffer: number,
32+
): { event: Partial<ServerSentEvent>; completed?: ServerSentEvent } {
33+
if (line === "") {
34+
if (event.data === undefined) {
35+
return { event: {} };
36+
}
37+
return {
38+
event: {},
39+
completed: { data: event.data, event: event.event, id: event.id },
40+
};
41+
}
42+
43+
if (line.startsWith(":")) {
44+
return { event };
45+
}
46+
47+
const colonIndex = line.indexOf(":");
48+
if (colonIndex === -1) {
49+
return { event };
50+
}
51+
52+
const field = line.slice(0, colonIndex);
53+
const fieldValue = line.slice(colonIndex + 1).trimStart();
54+
const nextEvent: Partial<ServerSentEvent> = { ...event };
55+
56+
switch (field) {
57+
case "data":
58+
nextEvent.data =
59+
nextEvent.data !== undefined ? `${nextEvent.data}\n${fieldValue}` : fieldValue;
60+
if (nextEvent.data.length > maxBuffer) {
61+
throw new BailianError("SSE event exceeded the maximum buffer size.", ExitCode.GENERAL);
62+
}
63+
break;
64+
case "event":
65+
nextEvent.event = fieldValue;
66+
break;
67+
case "id":
68+
nextEvent.id = fieldValue;
69+
break;
70+
}
71+
72+
return { event: nextEvent };
73+
}
74+
1075
export async function* parseSSE(response: Response): AsyncGenerator<ServerSentEvent> {
1176
const reader = response.body?.getReader();
1277
if (!reader) return;
1378

1479
const decoder = new TextDecoder();
1580
let buffer = "";
1681

17-
// Guard against a hostile or malfunctioning stream that never emits a newline
18-
// (or builds a single absurdly large event): bound the in-memory buffer so the
19-
// parser cannot be driven to exhaust process memory.
2082
const MAX_SSE_BUFFER = 16 * 1024 * 1024; // 16 MiB
2183

2284
try {
85+
// Keep partial event fields across chunks.
2386
let event: Partial<ServerSentEvent> = {};
2487

2588
while (true) {
2689
const { done, value } = await reader.read();
27-
if (done) break;
90+
if (done) {
91+
// EOF: treat any held `\r` as a line ending.
92+
if (buffer.length > 0) {
93+
const finalText = buffer.replace(/\r\n/g, "\n").replace(/\r/g, "\n");
94+
const parts = finalText.split("\n");
95+
buffer = parts.pop() ?? "";
96+
for (const line of parts) {
97+
const applied = applySseLine(line, event, MAX_SSE_BUFFER);
98+
event = applied.event;
99+
if (applied.completed) {
100+
yield applied.completed;
101+
}
102+
}
103+
}
104+
break;
105+
}
28106

29107
buffer += decoder.decode(value, { stream: true });
30108
if (buffer.length > MAX_SSE_BUFFER) {
31109
throw new BailianError("SSE stream exceeded the maximum buffer size.", ExitCode.GENERAL);
32110
}
33111

34-
const lines = buffer.split("\n");
35-
buffer = lines.pop() || "";
112+
const { lines, rest } = takeNormalizedSseLines(buffer);
113+
buffer = rest;
36114

37115
for (const line of lines) {
38-
if (line === "") {
39-
if (event.data !== undefined) {
40-
yield { data: event.data, event: event.event, id: event.id };
41-
}
42-
event = {};
43-
continue;
44-
}
45-
46-
if (line.startsWith(":")) continue; // comment
47-
48-
const colonIndex = line.indexOf(":");
49-
if (colonIndex === -1) continue;
50-
51-
const field = line.slice(0, colonIndex);
52-
const value = line.slice(colonIndex + 1).trimStart();
53-
54-
switch (field) {
55-
case "data":
56-
event.data = event.data !== undefined ? `${event.data}\n${value}` : value;
57-
if (event.data.length > MAX_SSE_BUFFER) {
58-
throw new BailianError(
59-
"SSE event exceeded the maximum buffer size.",
60-
ExitCode.GENERAL,
61-
);
62-
}
63-
break;
64-
case "event":
65-
event.event = value;
66-
break;
67-
case "id":
68-
event.id = value;
69-
break;
116+
const applied = applySseLine(line, event, MAX_SSE_BUFFER);
117+
event = applied.event;
118+
if (applied.completed) {
119+
yield applied.completed;
70120
}
71121
}
72122
}
73123

74-
// Flush remaining
75-
if (buffer.trim() && buffer.includes("data:")) {
76-
const colonIndex = buffer.indexOf(":");
77-
if (colonIndex !== -1) {
78-
yield { data: buffer.slice(colonIndex + 1).trimStart() };
79-
}
124+
// Legacy EOF flush: apply trailing field line and dispatch with event/id intact.
125+
if (buffer.length > 0) {
126+
const applied = applySseLine(buffer, event, MAX_SSE_BUFFER);
127+
event = applied.event;
128+
}
129+
if (event.data !== undefined) {
130+
yield { data: event.data, event: event.event, id: event.id };
80131
}
81132
} finally {
82133
reader.releaseLock();

0 commit comments

Comments
 (0)