fix(memory): abort remote response body reads

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<ReadableStreamReadResult<Uint8Array>>(() => {}),
+    cancel: async () => {
+      onCancel();
+    },
+    releaseLock: () => undefined,
+  } as ReadableStreamDefaultReader<Uint8Array>;
+
+  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));

packages/memory-host-sdk/src/host/post-json.ts

@@ -30,7 +30,7 @@ export async function postJson<T>(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<T>(params: {
       const payload = await readResponseJsonWithLimit(res, {
         errorPrefix: params.errorPrefix,
         maxBytes: params.maxResponseBytes,
+        signal: params.signal,
       });
       return await params.parse(payload);
     },

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<ReadableStreamReadResult<Uint8Array>>(() => {}),
+      cancel: async () => {
+        onCancel();
+      },
+      releaseLock: () => undefined,
+    } as ReadableStreamDefaultReader<Uint8Array>;
+
+    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<Uint8Array>({
@@ -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);
+  });
 });

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<string> {
   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<ResponsePrefix> {
+function toAbortError(signal: AbortSignal, fallbackMessage: string): Error {
+  return signal.reason instanceof Error ? signal.reason : new Error(fallbackMessage);
+}
+
+async function readChunkWithAbort(
+  reader: ReadableStreamDefaultReader<Uint8Array>,
+  signal: AbortSignal | undefined,
+  fallbackMessage: string,
+): Promise<ReadableStreamReadResult<Uint8Array>> {
+  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<ReadableStreamReadResult<Uint8Array>>((_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<ResponsePrefix> {
   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<Resp
 
   try {
     while (true) {
-      const { done, value } = await reader.read();
+      const { done, value } = await readChunkWithAbort(
+        reader,
+        signal,
+        "Response snippet body read aborted",
+      );
       if (done) {
         break;
       }
@@ -115,6 +159,7 @@ async function readResponseTextWithLimit(
   res: Response,
   maxBytes: number,
   errorPrefix: string,
+  signal?: AbortSignal,
 ): Promise<string> {
   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;
       }