From be7d86ed8005d4497d8702d52b41ad9ec86a3314 Mon Sep 17 00:00:00 2001 From: Vincent Koc Date: Fri, 19 Jun 2026 08:22:54 +0200 Subject: [PATCH] fix(memory): abort remote response body reads --- .../src/host/post-json.test.ts | 46 +++++++++++++++ .../memory-host-sdk/src/host/post-json.ts | 3 +- .../src/host/response-snippet.test.ts | 57 +++++++++++++++++- .../src/host/response-snippet.ts | 59 +++++++++++++++++-- 4 files changed, 158 insertions(+), 7 deletions(-) diff --git a/packages/memory-host-sdk/src/host/post-json.test.ts b/packages/memory-host-sdk/src/host/post-json.test.ts index 2fb031d6d06..1af1a98599d 100644 --- a/packages/memory-host-sdk/src/host/post-json.test.ts +++ b/packages/memory-host-sdk/src/host/post-json.test.ts @@ -35,6 +35,23 @@ function streamingTextResponse(params: { return new Response(stream, { status: params.status, headers: params.headers }); } +function stallingSuccessResponse(onCancel: () => void): Response { + const reader = { + read: () => new Promise>(() => {}), + cancel: async () => { + onCancel(); + }, + releaseLock: () => undefined, + } as ReadableStreamDefaultReader; + + return { + body: { getReader: () => reader }, + headers: new Headers(), + ok: true, + status: 200, + } as Response; +} + describe("postJson", () => { beforeEach(() => { vi.clearAllMocks(); @@ -73,6 +90,35 @@ describe("postJson", () => { }); }); + it("applies abort signals while reading successful response bodies", async () => { + let canceled = false; + const controller = new AbortController(); + remoteHttpMock.mockImplementationOnce(async (params) => { + return await params.onResponse( + stallingSuccessResponse(() => { + canceled = true; + }), + ); + }); + + const read = postJson({ + url: "https://memory.example/v1/post", + headers: {}, + body: {}, + signal: controller.signal, + errorPrefix: "post failed", + parse: () => ({}), + }); + + await new Promise((resolve) => { + setTimeout(resolve, 0); + }); + controller.abort(new Error("body aborted")); + + await expect(read).rejects.toThrow("body aborted"); + expect(canceled).toBe(true); + }); + it("attaches status to thrown error when requested", async () => { remoteHttpMock.mockImplementationOnce(async (params) => { return await params.onResponse(textResponse("bad gateway", 502)); diff --git a/packages/memory-host-sdk/src/host/post-json.ts b/packages/memory-host-sdk/src/host/post-json.ts index 1eb032aace1..2f7dcc250ae 100644 --- a/packages/memory-host-sdk/src/host/post-json.ts +++ b/packages/memory-host-sdk/src/host/post-json.ts @@ -30,7 +30,7 @@ export async function postJson(params: { }, onResponse: async (res) => { if (!res.ok) { - const text = await readResponseTextSnippet(res); + const text = await readResponseTextSnippet(res, { signal: params.signal }); const err = new Error(`${params.errorPrefix}: ${res.status} ${text}`) as Error & { status?: number; }; @@ -42,6 +42,7 @@ export async function postJson(params: { const payload = await readResponseJsonWithLimit(res, { errorPrefix: params.errorPrefix, maxBytes: params.maxResponseBytes, + signal: params.signal, }); return await params.parse(payload); }, diff --git a/packages/memory-host-sdk/src/host/response-snippet.test.ts b/packages/memory-host-sdk/src/host/response-snippet.test.ts index 7776b29bee6..bbfd1583bce 100644 --- a/packages/memory-host-sdk/src/host/response-snippet.test.ts +++ b/packages/memory-host-sdk/src/host/response-snippet.test.ts @@ -1,8 +1,23 @@ // Memory Host SDK tests cover response snippet behavior. import { describe, expect, it } from "vitest"; -import { readResponseTextSnippet } from "./response-snippet.js"; +import { readResponseJsonWithLimit, readResponseTextSnippet } from "./response-snippet.js"; describe("readResponseTextSnippet", () => { + function stallingResponse(onCancel: () => void): Response { + const reader = { + read: () => new Promise>(() => {}), + cancel: async () => { + onCancel(); + }, + releaseLock: () => undefined, + } as ReadableStreamDefaultReader; + + return { + body: { getReader: () => reader }, + headers: new Headers(), + } as Response; + } + it("does not wait for another chunk after reading the byte cap exactly", async () => { let canceled = false; const stream = new ReadableStream({ @@ -19,4 +34,44 @@ describe("readResponseTextSnippet", () => { ).resolves.toBe("abcd... [truncated]"); expect(canceled).toBe(true); }); + + it("cancels snippet body reads when the caller signal aborts", async () => { + let canceled = false; + const response = stallingResponse(() => { + canceled = true; + }); + const controller = new AbortController(); + const read = readResponseTextSnippet(response, { + maxBytes: 1024, + signal: controller.signal, + }); + + await new Promise((resolve) => { + setTimeout(resolve, 0); + }); + controller.abort(new Error("snippet aborted")); + + await expect(read).rejects.toThrow("snippet aborted"); + expect(canceled).toBe(true); + }); + + it("cancels JSON body reads when the caller signal aborts", async () => { + let canceled = false; + const response = stallingResponse(() => { + canceled = true; + }); + const controller = new AbortController(); + const read = readResponseJsonWithLimit(response, { + errorPrefix: "remote memory", + signal: controller.signal, + }); + + await new Promise((resolve) => { + setTimeout(resolve, 0); + }); + controller.abort(new Error("json aborted")); + + await expect(read).rejects.toThrow("json aborted"); + expect(canceled).toBe(true); + }); }); diff --git a/packages/memory-host-sdk/src/host/response-snippet.ts b/packages/memory-host-sdk/src/host/response-snippet.ts index b71eb3b1a45..4206e42f66b 100644 --- a/packages/memory-host-sdk/src/host/response-snippet.ts +++ b/packages/memory-host-sdk/src/host/response-snippet.ts @@ -8,11 +8,13 @@ const TRUNCATED_SUFFIX = "... [truncated]"; type ResponseTextSnippetOptions = { maxBytes?: number; maxChars?: number; + signal?: AbortSignal; }; type ResponseJsonOptions = { maxBytes?: number; errorPrefix: string; + signal?: AbortSignal; }; type ResponsePrefix = { @@ -28,7 +30,7 @@ export async function readResponseTextSnippet( ): Promise { const maxBytes = options.maxBytes ?? DEFAULT_ERROR_BODY_MAX_BYTES; const maxChars = options.maxChars ?? DEFAULT_ERROR_BODY_MAX_CHARS; - const prefix = await readResponsePrefix(res, maxBytes); + const prefix = await readResponsePrefix(res, maxBytes, options.signal); if (prefix.length === 0) { return ""; } @@ -56,7 +58,7 @@ export async function readResponseJsonWithLimit( throw responseTooLarge(options.errorPrefix, contentLength, maxBytes); } - const text = await readResponseTextWithLimit(res, maxBytes, options.errorPrefix); + const text = await readResponseTextWithLimit(res, maxBytes, options.errorPrefix, options.signal); try { return JSON.parse(text); @@ -65,7 +67,45 @@ export async function readResponseJsonWithLimit( } } -async function readResponsePrefix(res: Response, maxBytes: number): Promise { +function toAbortError(signal: AbortSignal, fallbackMessage: string): Error { + return signal.reason instanceof Error ? signal.reason : new Error(fallbackMessage); +} + +async function readChunkWithAbort( + reader: ReadableStreamDefaultReader, + signal: AbortSignal | undefined, + fallbackMessage: string, +): Promise> { + if (!signal) { + return await reader.read(); + } + if (signal.aborted) { + await reader.cancel().catch(() => undefined); + throw toAbortError(signal, fallbackMessage); + } + + let removeAbortListener: (() => void) | undefined; + const abortPromise = new Promise>((_resolve, reject) => { + const onAbort = () => { + void reader.cancel().catch(() => undefined); + reject(toAbortError(signal, fallbackMessage)); + }; + signal.addEventListener("abort", onAbort, { once: true }); + removeAbortListener = () => signal.removeEventListener("abort", onAbort); + }); + + try { + return await Promise.race([reader.read(), abortPromise]); + } finally { + removeAbortListener?.(); + } +} + +async function readResponsePrefix( + res: Response, + maxBytes: number, + signal?: AbortSignal, +): Promise { const body = res.body; if (!body || typeof body.getReader !== "function") { return { bytes: [], length: 0, truncated: false }; @@ -78,7 +118,11 @@ async function readResponsePrefix(res: Response, maxBytes: number): Promise { const body = res.body; if (!body || typeof body.getReader !== "function") { @@ -127,7 +172,11 @@ async function readResponseTextWithLimit( try { while (true) { - const { done, value } = await reader.read(); + const { done, value } = await readChunkWithAbort( + reader, + signal, + `${errorPrefix}: response body read aborted`, + ); if (done) { break; }