diff --git a/src/jsonrpc.test.ts b/src/jsonrpc.test.ts index fb08dd4..0313a28 100644 --- a/src/jsonrpc.test.ts +++ b/src/jsonrpc.test.ts @@ -1011,7 +1011,7 @@ describe("JSON-RPC request cancellation", () => { }); describe("JSON-RPC malformed peer messages", () => { - it("rejects the pending request when a response's error member is not an object", async () => { + it("rejects a matching malformed response without replying to it", async () => { const [clientStream, serverStream] = memoryStreamPair(); const client = Connection.builder().connect(clientStream); const serverReader = serverStream.readable.getReader(); @@ -1029,20 +1029,97 @@ describe("JSON-RPC malformed peer messages", () => { await expect(response).rejects.toMatchObject({ code: -32600 }); + const notification = client.sendNotification("example/after-response", {}); + await expect(serverReader.read()).resolves.toMatchObject({ + value: { + jsonrpc: "2.0", + method: "example/after-response", + }, + }); + await notification; + + serverReader.releaseLock(); + serverWriter.releaseLock(); client.close(); await client.closed; }); - it("ignores non-object messages without tearing down the connection", async () => { + it("does not reply to unknown malformed response-shaped messages", async () => { const consoleError = vi .spyOn(console, "error") .mockImplementation(() => {}); const [clientStream, serverStream] = memoryStreamPair(); + const processed = Promise.withResolvers(); + const client = Connection.builder() + .onReceiveNotification( + "example/barrier", + (params) => params, + () => { + processed.resolve(); + }, + ) + .connect(clientStream); + const serverReader = serverStream.readable.getReader(); + const serverWriter = serverStream.writable.getWriter(); + + await serverWriter.write({ + jsonrpc: "2.0", + id: 999, + error: null, + } as unknown as AnyMessage); + await serverWriter.write({ + jsonrpc: "2.0", + result: true, + } as unknown as AnyMessage); + await serverWriter.write({ + jsonrpc: "2.0", + error: null, + } as unknown as AnyMessage); + await serverWriter.write({ + jsonrpc: "2.0", + method: "example/barrier", + }); + await processed.promise; + + const notification = client.sendNotification("example/after-responses", {}); + await expect(serverReader.read()).resolves.toMatchObject({ + value: { + jsonrpc: "2.0", + method: "example/after-responses", + }, + }); + await notification; + + serverReader.releaseLock(); + serverWriter.releaseLock(); + consoleError.mockRestore(); + client.close(); + await client.closed; + }); + + it("returns Invalid Request for malformed call-shaped values and stays open", async () => { + const [clientStream, serverStream] = memoryStreamPair(); const client = Connection.builder().connect(clientStream); const serverReader = serverStream.readable.getReader(); const serverWriter = serverStream.writable.getWriter(); - await serverWriter.write(42 as unknown as AnyMessage); + for (const malformed of [ + null, + 42, + {}, + { jsonrpc: "1.0", id: 1, method: "example/test" }, + { jsonrpc: "2.0", id: 1, method: 42 }, + { jsonrpc: "2.0", id: {}, method: "example/test" }, + ]) { + await serverWriter.write(malformed as unknown as AnyWireMessage); + await expect(serverReader.read()).resolves.toMatchObject({ + value: { + jsonrpc: "2.0", + id: null, + error: { code: -32600 }, + }, + }); + } const response = client.sendRequest("example/test", {}); const { value: request } = await serverReader.read(); @@ -1053,11 +1130,9 @@ describe("JSON-RPC malformed peer messages", () => { } as AnyMessage); await expect(response).resolves.toEqual({ ok: true }); - expect(consoleError).toHaveBeenCalledWith("Invalid message", { - message: 42, - }); - consoleError.mockRestore(); + serverReader.releaseLock(); + serverWriter.releaseLock(); client.close(); await client.closed; }); diff --git a/src/jsonrpc.ts b/src/jsonrpc.ts index 2c76d07..5f709cc 100644 --- a/src/jsonrpc.ts +++ b/src/jsonrpc.ts @@ -1217,9 +1217,6 @@ export class Connection { if (done) { break; } - if (!message) { - continue; - } this.receiveWireMessage(message); } @@ -1251,8 +1248,14 @@ export class Connection { return; } - if (!isRecord(message)) { - console.error("Invalid message", { message }); + if ( + !isRequestMessage(message) && + !isNotificationMessage(message) && + !isResponseShapedMessage(message) + ) { + void this.sendWireMessage( + protocolErrorResponse(RequestError.invalidRequest(message)), + ).catch(() => {}); return; } @@ -1261,11 +1264,9 @@ export class Connection { private receiveBatch(batch: unknown[]): void { if (batch.length === 0) { - void this.sendWireMessage({ - jsonrpc: "2.0", - id: null, - error: RequestError.invalidRequest(batch).toErrorResponse(), - }).catch(() => {}); + void this.sendWireMessage( + protocolErrorResponse(RequestError.invalidRequest(batch)), + ).catch(() => {}); return; } @@ -1314,11 +1315,9 @@ export class Connection { } if (!isRequestMessage(message) && !isNotificationMessage(message)) { - void collectResponse({ - jsonrpc: "2.0", - id: null, - error: RequestError.invalidRequest(message).toErrorResponse(), - }).catch(() => {}); + void collectResponse( + protocolErrorResponse(RequestError.invalidRequest(message)), + ).catch(() => {}); continue; } @@ -1800,3 +1799,16 @@ export class RequestError extends Error { }; } } + +/** + * Creates a JSON-RPC error response for a request whose ID cannot be known. + * + * @internal + */ +export function protocolErrorResponse(error: RequestError): AnyResponse { + return { + jsonrpc: "2.0", + id: null, + error: error.toErrorResponse(), + }; +} diff --git a/src/server-websocket-upgrade.test.ts b/src/server-websocket-upgrade.test.ts index ac096e3..40f5d07 100644 --- a/src/server-websocket-upgrade.test.ts +++ b/src/server-websocket-upgrade.test.ts @@ -83,6 +83,58 @@ function createProtocolAgent( } describe("AcpServer prepared WebSocket upgrades", () => { + it("returns JSON-RPC errors for malformed frames and remains usable", async () => { + const registry = new ConnectionRegistry(); + const agent = createProtocolAgent(1); + const socket = new FakeServerSocket(); + + try { + handleWebSocketConnection(socket, { registry, agent }); + + socket.receive("not json"); + await expect(readSentMessage(socket)).resolves.toMatchObject({ + jsonrpc: "2.0", + id: null, + error: { code: -32700 }, + }); + + socket.receive("42"); + await expect(readSentMessage(socket)).resolves.toMatchObject({ + jsonrpc: "2.0", + id: null, + error: { code: -32600 }, + }); + expect(socket.closeCount).toBe(0); + + socket.receive(JSON.stringify(initializeRequest)); + await expect(readSentMessage(socket)).resolves.toMatchObject({ + jsonrpc: "2.0", + id: initializeRequest.id, + result: { protocolVersion: 1 }, + }); + + socket.receive("{ malformed"); + await expect(readSentMessage(socket)).resolves.toMatchObject({ + jsonrpc: "2.0", + id: null, + error: { code: -32700 }, + }); + + socket.receive(JSON.stringify(sessionNewRequest)); + await expect(readSentMessage(socket)).resolves.toMatchObject({ + jsonrpc: "2.0", + id: sessionNewRequest.id, + result: { + sessionId: expect.stringMatching(/^[0-9a-f-]{36}$/), + }, + }); + expect(socket.closeCount).toBe(0); + } finally { + socket.close(); + await registry.closeAll(); + } + }); + it("uses the default factory when no per-upgrade override is provided", async () => { const createdBy: string[] = []; const server = new AcpServer({ diff --git a/src/stream.test.ts b/src/stream.test.ts index a1d8ea1..e3510ef 100644 --- a/src/stream.test.ts +++ b/src/stream.test.ts @@ -175,48 +175,32 @@ describe("ndJsonStream", () => { expect(messages).toEqual([msg]); }); - it("skips malformed lines and continues parsing", async () => { - const error = vi - .spyOn(console, "error") - .mockImplementation(() => undefined); - const msg1 = { jsonrpc: "2.0" as const, id: 1, method: "before" }; - const msg2 = { jsonrpc: "2.0" as const, id: 2, method: "after" }; + it("responds to malformed and primitive lines and continues parsing", async () => { + const outputChunks: Uint8Array[] = []; + const output = new WritableStream({ + write(chunk) { + outputChunks.push(chunk); + }, + }); + const msg = { jsonrpc: "2.0" as const, id: 2, method: "after" }; const input = streamFromChunks([ - JSON.stringify(msg1) + - "\n" + - "not valid json\n" + - JSON.stringify(msg2) + - "\n", + "not valid json\n42\n\n" + JSON.stringify(msg) + "\n", ]); - const { readable } = ndJsonStream(nullWritable, input); + const { readable } = ndJsonStream(output, input); const messages = await collectStream(readable); + const responses = outputChunks + .map((chunk) => new TextDecoder().decode(chunk)) + .join("") + .trim() + .split("\n") + .map((line) => JSON.parse(line)); - expect(messages).toEqual([msg1, msg2]); - expect(error).toHaveBeenCalledOnce(); - - error.mockRestore(); - }); - - it("skips non-object JSON lines that would break the connection layer", async () => { - const warn = vi.spyOn(console, "warn").mockImplementation(() => undefined); - const msg1 = { jsonrpc: "2.0" as const, id: 1, method: "before" }; - const msg2 = { jsonrpc: "2.0" as const, id: 2, method: "after" }; - const input = streamFromChunks([ - JSON.stringify(msg1) + - "\n" + - '42\n"str"\nnull\n' + - JSON.stringify(msg2) + - "\n", + expect(messages).toEqual([msg]); + expect(responses).toMatchObject([ + { jsonrpc: "2.0", id: null, error: { code: -32700 } }, + { jsonrpc: "2.0", id: null, error: { code: -32600 } }, ]); - - const { readable } = ndJsonStream(nullWritable, input); - const messages = await collectStream(readable); - - expect(messages).toEqual([msg1, msg2]); - expect(warn).toHaveBeenCalledTimes(3); - - warn.mockRestore(); }); it("passes through object messages without validating their shape", async () => { diff --git a/src/stream.ts b/src/stream.ts index f6663cd..009c39c 100644 --- a/src/stream.ts +++ b/src/stream.ts @@ -1,5 +1,5 @@ import type { AnyMessage, AnyWireMessage } from "./jsonrpc.js"; -import { isRecord } from "./jsonrpc.js"; +import { RequestError, isRecord, protocolErrorResponse } from "./jsonrpc.js"; import { LineBuffer } from "./line-buffer.js"; /** @@ -45,29 +45,46 @@ export function ndJsonStream( const textDecoder = new TextDecoder(); let cancelled = false; let inputReader: ReadableStreamDefaultReader | undefined; + let outputWrite: Promise = Promise.resolve(); + + const writeJson = (message: unknown): Promise => { + const content = JSON.stringify(message) + "\n"; + const write = outputWrite.then(async () => { + const writer = output.getWriter(); + try { + await writer.write(textEncoder.encode(content)); + } finally { + writer.releaseLock(); + } + }); + outputWrite = write.catch(() => {}); + return write; + }; const readable = new ReadableStream({ async start(controller) { const lines = new LineBuffer(); - const enqueueLine = (lineBytes: Uint8Array) => { + const enqueueLine = async (lineBytes: Uint8Array) => { const trimmedLine = textDecoder.decode(lineBytes).trim(); - if (trimmedLine) { - try { - const message: unknown = JSON.parse(trimmedLine); - // Skip primitive lines with a useful warning; individual objects - // and batch arrays are left for the connection layer to validate. - if (isRecord(message) || Array.isArray(message)) { - controller.enqueue(message as Message); - } else { - console.warn( - "Skipping JSON line that is not an object:", - trimmedLine, - ); - } - } catch (err) { - console.error("Failed to parse JSON message:", trimmedLine, err); - } + if (!trimmedLine) { + return; + } + + let message: unknown; + try { + message = JSON.parse(trimmedLine); + } catch { + await writeJson(protocolErrorResponse(RequestError.parseError())); + return; + } + + if (isRecord(message) || Array.isArray(message)) { + controller.enqueue(message as Message); + } else { + await writeJson( + protocolErrorResponse(RequestError.invalidRequest(message)), + ); } }; @@ -86,7 +103,7 @@ export function ndJsonStream( continue; } for (const line of lines.push(value)) { - enqueueLine(line); + await enqueueLine(line); if (cancelled) { return; } @@ -97,7 +114,7 @@ export function ndJsonStream( } const lastLine = lines.flush(); if (lastLine) { - enqueueLine(lastLine); + await enqueueLine(lastLine); } } catch (err) { if (cancelled) { @@ -123,14 +140,8 @@ export function ndJsonStream( }); const writable = new WritableStream({ - async write(message) { - const content = JSON.stringify(message) + "\n"; - const writer = output.getWriter(); - try { - await writer.write(textEncoder.encode(content)); - } finally { - writer.releaseLock(); - } + write(message) { + return writeJson(message); }, }); diff --git a/src/ws-server.ts b/src/ws-server.ts index 4f3397a..ea468e2 100644 --- a/src/ws-server.ts +++ b/src/ws-server.ts @@ -1,8 +1,10 @@ import { + RequestError, isNotificationMessage, isRecord, isRequestMessage, isResponseShapedMessage, + protocolErrorResponse, } from "./jsonrpc.js"; import { isInitializeRequest, @@ -128,18 +130,13 @@ class WebSocketServerSession implements WebSocketServerSessionHandle { let value: unknown; try { value = JSON.parse(text); - } catch (error) { - console.warn("Ignoring malformed ACP WebSocket JSON message:", error); - await this.shutdownIfUninitialized(1007, "Malformed JSON"); - + } catch { + this.send(protocolErrorResponse(RequestError.parseError())); return; } - // Skip non-object messages with a useful warning; anything object-shaped - // is left for the connection layer to validate. if (!Array.isArray(value) && !isRecord(value)) { - console.warn("Ignoring non-object ACP WebSocket message:", value); - await this.shutdownIfUninitialized(1002, "Invalid JSON-RPC message"); + this.send(protocolErrorResponse(RequestError.invalidRequest(value))); return; } @@ -159,6 +156,15 @@ class WebSocketServerSession implements WebSocketServerSessionHandle { const message = value as AnyWireMessage; if (!this.connection) { + if (isResponseShapedMessage(message)) { + return; + } + + if (!isRequestMessage(message) && !isNotificationMessage(message)) { + this.send(protocolErrorResponse(RequestError.invalidRequest(message))); + return; + } + await this.handleInitialize(message as AnyMessage, false); return; } diff --git a/src/ws-stream.test.ts b/src/ws-stream.test.ts index c83411d..9364796 100644 --- a/src/ws-stream.test.ts +++ b/src/ws-stream.test.ts @@ -424,8 +424,7 @@ describe("createWebSocketStream", () => { } }); - it("ignores binary, malformed JSON, and primitive messages, passing objects and batches through", async () => { - const warn = vi.spyOn(console, "warn").mockImplementation(() => undefined); + it("responds to malformed JSON and primitives while passing valid frames through", async () => { const instances: FakeWebSocket[] = []; const stream = createWebSocketStream("ws://agent.example/acp", { WebSocket: createFakeWebSocketConstructor(instances), @@ -453,10 +452,14 @@ describe("createWebSocketStream", () => { expect(await readMessage(reader)).toEqual({ hello: "world" }); expect(await readWireMessage(reader)).toEqual(batch); expect(await readMessage(reader)).toEqual(initializeResponse); - expect(warn).toHaveBeenCalledTimes(2); + await vi.waitFor(() => expect(socket.sent).toHaveLength(2)); + expect(socket.sent.map((message) => JSON.parse(message))).toMatchObject([ + { jsonrpc: "2.0", id: null, error: { code: -32700 } }, + { jsonrpc: "2.0", id: null, error: { code: -32600 } }, + ]); + expect(socket.readyState).toBe(1); } finally { reader.releaseLock(); - warn.mockRestore(); await closeStream(stream); } }); diff --git a/src/ws-stream.ts b/src/ws-stream.ts index 27e805f..25ab3d3 100644 --- a/src/ws-stream.ts +++ b/src/ws-stream.ts @@ -1,5 +1,5 @@ import { MemoryAcpCookieStore } from "./cookie-store.js"; -import { isRecord } from "./jsonrpc.js"; +import { RequestError, isRecord, protocolErrorResponse } from "./jsonrpc.js"; import { onWebSocket, webSocketMessageToString } from "./ws-utils.js"; import type { AcpCookieStore } from "./cookie-store.js"; import type { WebSocketLike } from "./ws-utils.js"; @@ -82,6 +82,7 @@ class WebSocketStreamTransport { private resolveOpen: (() => void) | undefined; private rejectOpen: ((error: unknown) => void) | undefined; private readonly detachListeners: Array<() => void> = []; + private sendQueue: Promise = Promise.resolve(); constructor(serverUrl: string, options: WebSocketStreamOptions) { const WebSocketCtor = resolveWebSocket(options.WebSocket); @@ -150,9 +151,7 @@ class WebSocketStreamTransport { }, }), writable: new WritableStream({ - write: async (message) => { - await this.sendMessage(message); - }, + write: (message) => this.queueMessage(message), close: () => { this.close(); }, @@ -163,7 +162,13 @@ class WebSocketStreamTransport { }; } - private async sendMessage(message: Message): Promise { + private queueMessage(message: AnyWireMessage): Promise { + const send = this.sendQueue.then(() => this.sendMessage(message)); + this.sendQueue = send.catch(() => {}); + return send; + } + + private async sendMessage(message: AnyWireMessage): Promise { if (this.isClosed) { throw new Error("ACP WebSocket stream is closed"); } @@ -201,21 +206,25 @@ class WebSocketStreamTransport { let value: unknown; try { value = JSON.parse(text); - } catch (error) { - console.warn("Ignoring malformed ACP WebSocket JSON message:", error); + } catch { + this.sendProtocolError(RequestError.parseError()); return; } - // Skip primitive messages with a useful warning; individual objects and - // batch arrays are left for the connection layer to validate. if (!isRecord(value) && !Array.isArray(value)) { - console.warn("Ignoring primitive ACP WebSocket message:", value); + this.sendProtocolError(RequestError.invalidRequest(value)); return; } this.readableController?.enqueue(value as Message); } + private sendProtocolError(error: RequestError): void { + void this.queueMessage(protocolErrorResponse(error)).catch((sendError) => { + this.errorReadable(sendError); + }); + } + private close(): void { this.closeSocket(); this.closeReadable();