diff --git a/src/tests/event-loop-yield.test.ts b/src/tests/event-loop-yield.test.ts index 21b12b45..92737d21 100644 --- a/src/tests/event-loop-yield.test.ts +++ b/src/tests/event-loop-yield.test.ts @@ -174,6 +174,43 @@ describe("consumer under a buffered backlog", () => { }); }); +describe("consumer teardown under a buffered backlog", () => { + it("settles the turn when the session closes mid-backlog", async () => { + const chunks = Array.from({ length: 30 }, (_, i) => `chunk-${i} `); + const updates: SessionNotification[] = []; + let closeScheduled = false; + let chunksAtClose: number | undefined; + let closed: Promise | undefined; + const agent = createAgent(bufferedAnswer(chunks), (notification) => { + updates.push(notification); + if (notification.update.sessionUpdate !== "agent_message_chunk") return; + if (!closeScheduled) { + closeScheduled = true; + setImmediate(() => { + chunksAtClose = chunkTexts(updates).length; + closed = agent.closeSession({ sessionId }); + }); + } + spin(2); + }); + + const response = await agent.prompt({ sessionId, prompt: [{ type: "text", text: "go" }] }); + + expect(chunksAtClose).toBeDefined(); + expect(chunksAtClose!).toBeLessThan(chunks.length); + expect(response.stopReason).toBe("cancelled"); + await closed; + expect(agent.sessions[sessionId]).toBeUndefined(); + // Whatever reached the client is a prefix of the answer, in order, and + // the rest of the backlog is not forwarded after the close. + const received = chunkTexts(updates); + expect(received.length).toBeLessThan(chunks.length); + expect(received).toEqual(chunks.slice(0, received.length)); + for (let i = 0; i < 5; i++) await new Promise((resolve) => setImmediate(resolve)); + expect(chunkTexts(updates)).toEqual(received); + }); +}); + describe("nodeToWebWritable", () => { it("writes the bytes of each chunk, including a view into a larger buffer", async () => { const received: Buffer[] = []; @@ -190,6 +227,22 @@ describe("nodeToWebWritable", () => { expect(Buffer.concat(received).toString()).toBe("hello\nworld\n"); }); + it("gives an object-mode stream a Buffer copy", async () => { + const received: unknown[] = []; + const sink = new Writable({ + objectMode: true, + write(chunk, _encoding, callback) { + received.push(chunk); + callback(); + }, + }); + const bytes = new TextEncoder().encode("hi"); + await nodeToWebWritable(sink).getWriter().write(bytes); + expect(Buffer.isBuffer(received[0])).toBe(true); + expect(received[0]).not.toBe(bytes); + expect((received[0] as Buffer).toString()).toBe("hi"); + }); + it("rejects the write that fails", async () => { const sink = new Writable({ write(_chunk, _encoding, callback) { diff --git a/src/utils.ts b/src/utils.ts index e6d5c47f..d93b4e21 100644 --- a/src/utils.ts +++ b/src/utils.ts @@ -45,15 +45,20 @@ export class Pushable implements AsyncIterable { } } -// Helper to convert Node.js streams to Web Streams +// Helper to convert Node.js streams to Web Streams. A byte stream gets each +// chunk without a copy, so the caller must not change a chunk until its write +// resolves. export function nodeToWebWritable(nodeStream: Writable): WritableStream { return new WritableStream({ write(chunk) { return new Promise((resolve, reject) => { // A Uint8Array is written as it is: Node wraps it in a Buffer view - // without a copy. The encoder hands each message a fresh array, so - // nothing changes the bytes while the write is pending. - nodeStream.write(chunk, (err) => { + // without a copy. The ACP encoder hands each message a fresh array, + // so nothing changes the bytes while the write is pending. An + // object-mode stream would get the array itself, so it keeps + // getting a Buffer copy, as before. + const data = nodeStream.writableObjectMode ? Buffer.from(chunk) : chunk; + nodeStream.write(data, (err) => { if (err) { reject(err); } else { @@ -100,8 +105,11 @@ export function unreachable(value: never, logger: Logger = console) { * * The budget is measured from the first call since the event loop last * reached its check phase, so a loop that already waits on real I/O never - * yields, and all loops that share one instance share the budget of the - * macrotask that runs them. + * yields, and all loops that share one instance share one budget per + * event-loop iteration. + * + * A test that fakes `setImmediate` and drives such a loop for longer than the + * budget must advance the fake timers, or the loop waits for them. */ export class EventLoopYielder { /** Whether a check-phase marker is pending. While it is, the event loop has