fix(memory): bound remote JSON responses

packages/memory-host-sdk/src/host/batch-upload.test.ts

@@ -15,6 +15,7 @@ function textResponse(body: string, status: number): Response {
 function streamingTextResponse(params: {
   body: string;
   status: number;
+  headers?: HeadersInit;
   onCancel: () => void;
 }): Response {
   const encoded = new TextEncoder().encode(params.body);
@@ -26,7 +27,7 @@ function streamingTextResponse(params: {
       params.onCancel();
     },
   });
-  return new Response(stream, { status: params.status });
+  return new Response(stream, { status: params.status, headers: params.headers });
 }
 
 describe("uploadBatchJsonlFile", () => {
@@ -77,4 +78,33 @@ describe("uploadBatchJsonlFile", () => {
     ).rejects.toThrow(`file upload failed: 413 ${"x".repeat(1_000)}... [truncated]`);
     expect(canceled).toBe(true);
   });
+
+  it("rejects oversized successful file-upload JSON before parsing", async () => {
+    let canceled = false;
+    remoteHttpMock.mockImplementationOnce(async (params) => {
+      return await params.onResponse(
+        streamingTextResponse({
+          body: '{"id":"file_123"}',
+          status: 200,
+          headers: { "content-length": "64" },
+          onCancel: () => {
+            canceled = true;
+          },
+        }),
+      );
+    });
+
+    await expect(
+      uploadBatchJsonlFile({
+        client: {
+          baseUrl: "https://memory.example/v1",
+          headers: { Authorization: "Bearer test" },
+        },
+        requests: [{ input: "one" }],
+        errorPrefix: "file upload failed",
+        maxResponseBytes: 8,
+      }),
+    ).rejects.toThrow("file upload failed: response body too large: 64 bytes (limit: 8 bytes)");
+    expect(canceled).toBe(true);
+  });
 });

packages/memory-host-sdk/src/host/batch-upload.ts

@@ -5,12 +5,13 @@ import {
 } from "./batch-utils.js";
 import { hashText } from "./hash.js";
 import { withRemoteHttpResponse } from "./remote-http.js";
-import { readResponseTextSnippet } from "./response-snippet.js";
+import { readResponseJsonWithLimit, readResponseTextSnippet } from "./response-snippet.js";
 
 export async function uploadBatchJsonlFile(params: {
   client: BatchHttpClientConfig;
   requests: unknown[];
   errorPrefix: string;
+  maxResponseBytes?: number;
 }): Promise<string> {
   const baseUrl = normalizeBatchBaseUrl(params.client);
   const jsonl = params.requests.map((request) => JSON.stringify(request)).join("\n");
@@ -35,11 +36,10 @@ export async function uploadBatchJsonlFile(params: {
         const text = await readResponseTextSnippet(fileRes);
         throw new Error(`${params.errorPrefix}: ${fileRes.status} ${text}`);
       }
-      try {
-        return (await fileRes.json()) as { id?: string };
-      } catch (cause) {
-        throw new Error(`${params.errorPrefix}: malformed JSON response`, { cause });
-      }
+      return (await readResponseJsonWithLimit(fileRes, {
+        errorPrefix: params.errorPrefix,
+        maxBytes: params.maxResponseBytes,
+      })) as { id?: string };
     },
   });
   if (!filePayload.id) {

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

@@ -19,6 +19,7 @@ function textResponse(body: string, status: number): Response {
 function streamingTextResponse(params: {
   body: string;
   status: number;
+  headers?: HeadersInit;
   onCancel: () => void;
 }): Response {
   const encoded = new TextEncoder().encode(params.body);
@@ -30,7 +31,7 @@ function streamingTextResponse(params: {
       params.onCancel();
     },
   });
-  return new Response(stream, { status: params.status });
+  return new Response(stream, { status: params.status, headers: params.headers });
 }
 
 describe("postJson", () => {
@@ -136,4 +137,59 @@ describe("postJson", () => {
       }),
     ).rejects.toThrow("post failed: malformed JSON response");
   });
+
+  it("rejects successful JSON responses with oversized content-length", async () => {
+    let canceled = false;
+    remoteHttpMock.mockImplementationOnce(async (params) => {
+      return await params.onResponse(
+        streamingTextResponse({
+          body: "{}",
+          status: 200,
+          headers: { "content-length": "32" },
+          onCancel: () => {
+            canceled = true;
+          },
+        }),
+      );
+    });
+
+    await expect(
+      postJson({
+        url: "https://memory.example/v1/post",
+        headers: {},
+        body: {},
+        errorPrefix: "post failed",
+        maxResponseBytes: 8,
+        parse: () => ({}),
+      }),
+    ).rejects.toThrow("post failed: response body too large: 32 bytes (limit: 8 bytes)");
+    expect(canceled).toBe(true);
+  });
+
+  it("cancels successful JSON responses that exceed the streaming byte cap", async () => {
+    let canceled = false;
+    remoteHttpMock.mockImplementationOnce(async (params) => {
+      return await params.onResponse(
+        streamingTextResponse({
+          body: `{"data":"${"x".repeat(32)}"}`,
+          status: 200,
+          onCancel: () => {
+            canceled = true;
+          },
+        }),
+      );
+    });
+
+    await expect(
+      postJson({
+        url: "https://memory.example/v1/post",
+        headers: {},
+        body: {},
+        errorPrefix: "post failed",
+        maxResponseBytes: 16,
+        parse: () => ({}),
+      }),
+    ).rejects.toThrow("post failed: response body too large");
+    expect(canceled).toBe(true);
+  });
 });

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

@@ -1,5 +1,5 @@
 import { withRemoteHttpResponse } from "./remote-http.js";
-import { readResponseTextSnippet } from "./response-snippet.js";
+import { readResponseJsonWithLimit, readResponseTextSnippet } from "./response-snippet.js";
 import type { SsrFPolicy } from "./ssrf-policy.js";
 
 export async function postJson<T>(params: {
@@ -11,6 +11,7 @@ export async function postJson<T>(params: {
   body: unknown;
   errorPrefix: string;
   attachStatus?: boolean;
+  maxResponseBytes?: number;
   parse: (payload: unknown) => T | Promise<T>;
 }): Promise<T> {
   return await withRemoteHttpResponse({
@@ -34,15 +35,11 @@ export async function postJson<T>(params: {
         }
         throw err;
       }
-      return await params.parse(await readJsonResponse(res, params.errorPrefix));
+      const payload = await readResponseJsonWithLimit(res, {
+        errorPrefix: params.errorPrefix,
+        maxBytes: params.maxResponseBytes,
+      });
+      return await params.parse(payload);
     },
   });
 }
-
-async function readJsonResponse(res: Response, errorPrefix: string): Promise<unknown> {
-  try {
-    return await res.json();
-  } catch (cause) {
-    throw new Error(`${errorPrefix}: malformed JSON response`, { cause });
-  }
-}

packages/memory-host-sdk/src/host/response-snippet.ts

@@ -1,5 +1,6 @@
 const DEFAULT_ERROR_BODY_MAX_BYTES = 8 * 1024;
 const DEFAULT_ERROR_BODY_MAX_CHARS = 1_000;
+const DEFAULT_JSON_BODY_MAX_BYTES = 64 * 1024 * 1024;
 const TRUNCATED_SUFFIX = "... [truncated]";
 
 type ResponseTextSnippetOptions = {
@@ -7,6 +8,11 @@ type ResponseTextSnippetOptions = {
   maxChars?: number;
 };
 
+type ResponseJsonOptions = {
+  maxBytes?: number;
+  errorPrefix: string;
+};
+
 type ResponsePrefix = {
   bytes: Uint8Array[];
   length: number;
@@ -35,6 +41,26 @@ export async function readResponseTextSnippet(
   return collapsed;
 }
 
+export async function readResponseJsonWithLimit(
+  res: Response,
+  options: ResponseJsonOptions,
+): Promise<unknown> {
+  const maxBytes = options.maxBytes ?? DEFAULT_JSON_BODY_MAX_BYTES;
+  const contentLength = parseContentLength(res.headers.get("content-length"), options.errorPrefix);
+  if (typeof contentLength === "number" && contentLength > maxBytes) {
+    await cancelResponseBody(res);
+    throw responseTooLarge(options.errorPrefix, contentLength, maxBytes);
+  }
+
+  const text = await readResponseTextWithLimit(res, maxBytes, options.errorPrefix);
+
+  try {
+    return JSON.parse(text);
+  } catch (cause) {
+    throw new Error(`${options.errorPrefix}: malformed JSON response`, { cause });
+  }
+}
+
 async function readResponsePrefix(res: Response, maxBytes: number): Promise<ResponsePrefix> {
   const body = res.body;
   if (!body || typeof body.getReader !== "function") {
@@ -79,6 +105,79 @@ async function readResponsePrefix(res: Response, maxBytes: number): Promise<Resp
   return { bytes: chunks, length, truncated };
 }
 
+async function readResponseTextWithLimit(
+  res: Response,
+  maxBytes: number,
+  errorPrefix: string,
+): Promise<string> {
+  const body = res.body;
+  if (!body || typeof body.getReader !== "function") {
+    return "";
+  }
+
+  const reader = body.getReader();
+  const chunks: Uint8Array[] = [];
+  let length = 0;
+
+  try {
+    while (true) {
+      const { done, value } = await reader.read();
+      if (done) {
+        break;
+      }
+      if (!value?.length) {
+        continue;
+      }
+
+      const nextLength = length + value.length;
+      if (nextLength > maxBytes) {
+        await reader.cancel().catch(() => undefined);
+        throw responseTooLarge(errorPrefix, nextLength, maxBytes);
+      }
+
+      chunks.push(value);
+      length = nextLength;
+    }
+  } finally {
+    try {
+      reader.releaseLock();
+    } catch {}
+  }
+
+  return new TextDecoder().decode(joinChunks(chunks, length));
+}
+
+async function cancelResponseBody(res: Response): Promise<void> {
+  const body = res.body;
+  if (!body || typeof body.cancel !== "function") {
+    return;
+  }
+  await body.cancel().catch(() => undefined);
+}
+
+function parseContentLength(raw: string | null, errorPrefix: string): number | undefined {
+  const trimmed = raw?.trim();
+  if (!trimmed) {
+    return undefined;
+  }
+  if (!/^(0|[1-9]\d*)$/.test(trimmed)) {
+    throw new Error(`${errorPrefix}: invalid content-length header: ${raw}`);
+  }
+  const value = Number(trimmed);
+  if (!Number.isSafeInteger(value)) {
+    throw new Error(`${errorPrefix}: invalid content-length header: ${raw}`);
+  }
+  return value;
+}
+
+function responseTooLarge(errorPrefix: string, size: number, maxBytes: number): Error {
+  return new Error(responseTooLargeMessage(errorPrefix, size, maxBytes));
+}
+
+function responseTooLargeMessage(errorPrefix: string, size: number, maxBytes: number): string {
+  return `${errorPrefix}: response body too large: ${size} bytes (limit: ${maxBytes} bytes)`;
+}
+
 function joinChunks(chunks: Uint8Array[], length: number): Uint8Array {
   if (chunks.length === 1 && chunks[0]?.length === length) {
     return chunks[0];