fix(memory): abort batch upload response reads (#95111)

vincentkoc · web-flow · commit c4d1f37d3326 · 2026-06-20T06:22:23.000+08:00
* fix(memory): abort batch upload response reads

* test(memory): stabilize batch upload abort proof
diff --git a/packages/memory-host-sdk/src/host/batch-upload.test.ts b/packages/memory-host-sdk/src/host/batch-upload.test.ts
@@ -31,6 +31,23 @@ function streamingTextResponse(params: {
   return new Response(stream, { status: params.status, headers: params.headers });
 }
 
+function stallingResponse(params: { status: number; onCancel: () => void }): Response {
+  const reader = {
+    read: () => new Promise<ReadableStreamReadResult<Uint8Array>>(() => {}),
+    cancel: async () => {
+      params.onCancel();
+    },
+    releaseLock: () => undefined,
+  } as ReadableStreamDefaultReader<Uint8Array>;
+
+  return {
+    status: params.status,
+    ok: params.status >= 200 && params.status < 300,
+    headers: new Headers(),
+    body: { getReader: () => reader },
+  } as Response;
+}
+
 describe("uploadBatchJsonlFile", () => {
   beforeEach(() => {
     vi.clearAllMocks();
@@ -108,4 +125,70 @@ describe("uploadBatchJsonlFile", () => {
     ).rejects.toThrow("file upload failed: response body too large: 64 bytes (limit: 8 bytes)");
     expect(canceled).toBe(true);
   });
+
+  it("passes caller abort signals through non-ok file-upload response snippets", async () => {
+    let canceled = false;
+    remoteHttpMock.mockImplementationOnce(async (params) => {
+      return await params.onResponse(
+        stallingResponse({
+          status: 500,
+          onCancel: () => {
+            canceled = true;
+          },
+        }),
+      );
+    });
+    const controller = new AbortController();
+    const upload = uploadBatchJsonlFile({
+      client: {
+        baseUrl: "https://memory.example/v1",
+        headers: { Authorization: "Bearer test" },
+      },
+      requests: [{ input: "one" }],
+      errorPrefix: "file upload failed",
+      signal: controller.signal,
+    });
+
+    await new Promise((resolve) => {
+      setTimeout(resolve, 0);
+    });
+    controller.abort(new Error("upload aborted"));
+
+    await expect(upload).rejects.toThrow("upload aborted");
+    expect(canceled).toBe(true);
+    expect(remoteHttpMock.mock.calls[0]?.[0].signal).toBe(controller.signal);
+  });
+
+  it("passes caller abort signals through successful file-upload JSON reads", async () => {
+    let canceled = false;
+    remoteHttpMock.mockImplementationOnce(async (params) => {
+      return await params.onResponse(
+        stallingResponse({
+          status: 200,
+          onCancel: () => {
+            canceled = true;
+          },
+        }),
+      );
+    });
+    const controller = new AbortController();
+    const upload = uploadBatchJsonlFile({
+      client: {
+        baseUrl: "https://memory.example/v1",
+        headers: { Authorization: "Bearer test" },
+      },
+      requests: [{ input: "one" }],
+      errorPrefix: "file upload failed",
+      signal: controller.signal,
+    });
+
+    await new Promise((resolve) => {
+      setTimeout(resolve, 0);
+    });
+    controller.abort(new Error("upload json aborted"));
+
+    await expect(upload).rejects.toThrow("upload json aborted");
+    expect(canceled).toBe(true);
+    expect(remoteHttpMock.mock.calls[0]?.[0].signal).toBe(controller.signal);
+  });
 });
diff --git a/packages/memory-host-sdk/src/host/batch-upload.ts b/packages/memory-host-sdk/src/host/batch-upload.ts
@@ -16,6 +16,7 @@ export async function uploadBatchJsonlFile(params: {
   requests: unknown[];
   errorPrefix: string;
   maxResponseBytes?: number;
+  signal?: AbortSignal;
 }): Promise<string> {
   const baseUrl = normalizeBatchBaseUrl(params.client);
   const jsonl = params.requests.map((request) => JSON.stringify(request)).join("\n");
@@ -31,19 +32,21 @@ export async function uploadBatchJsonlFile(params: {
     url: `${baseUrl}/files`,
     ssrfPolicy: params.client.ssrfPolicy,
     fetchImpl: params.client.fetchImpl,
+    signal: params.signal,
     init: {
       method: "POST",
       headers: buildBatchHeaders(params.client, { json: false }),
       body: form,
     },
     onResponse: async (fileRes) => {
       if (!fileRes.ok) {
-        const text = await readResponseTextSnippet(fileRes);
+        const text = await readResponseTextSnippet(fileRes, { signal: params.signal });
         throw new Error(`${params.errorPrefix}: ${fileRes.status} ${text}`);
       }
       return (await readResponseJsonWithLimit(fileRes, {
         errorPrefix: params.errorPrefix,
         maxBytes: params.maxResponseBytes,
+        signal: params.signal,
       })) as { id?: string };
     },
   });
