fix(agents): bound Google prompt cache response reads

Alix-007 · Alix-007 · commit eedb29d85733 · 2026-06-21T06:28:53.000+08:00
The Google embedded-agent prompt-cache helpers parsed cachedContents
metadata with an unbounded `await response.json()` in both
createGooglePromptCache and updateGooglePromptCacheTtl. A buggy or
hostile Generative Language endpoint returning a 200 with a large or
never-ending body (especially with no Content-Length) would be fully
buffered into memory before parsing, with the existing
cancelUnreadResponseBody guard firing too late (json() already drained
the body).

Route both reads through the shared streaming byte-cap reader
(readResponseWithLimit) under a 1 MiB cap, cancelling the stream on
overflow instead of buffering it, then JSON.parse the bounded buffer.
This is the symmetric Google-endpoint counterpart to the Anthropic
error-stream and gateway pricing-catalog bounds.

Adds regressions that stream an oversized no-Content-Length body through
the real create and TTL-refresh paths and assert the body is cancelled.
diff --git a/src/agents/embedded-agent-runner/google-prompt-cache.test.ts b/src/agents/embedded-agent-runner/google-prompt-cache.test.ts
@@ -68,6 +68,30 @@ function createCacheFetchMock(params: { name: string; expireTime: string }) {
   );
 }
 
+// Builds a 200-OK response whose body streams more than 1 MiB with no
+// Content-Length, mirroring a buggy/hostile Google cachedContents endpoint.
+// The shared byte-cap reader must cancel the body before fully buffering it.
+function createOversizedJsonResponse(): { response: Response; cancel: ReturnType<typeof vi.fn> } {
+  const cancel = vi.fn(async () => undefined);
+  let pullCount = 0;
+  const response = new Response(
+    new ReadableStream<Uint8Array>({
+      pull(controller) {
+        pullCount += 1;
+        // First chunk already exceeds the 1 MiB cap so the reader truncates
+        // and cancels instead of waiting for the (never-ending) rest.
+        controller.enqueue(new Uint8Array(pullCount === 1 ? 1024 * 1024 + 1 : 1));
+      },
+      cancel,
+    }),
+    {
+      status: 200,
+      headers: { "content-type": "application/json" },
+    },
+  );
+  return { response, cancel };
+}
+
 function createCapturingStreamFn(result = "stream") {
   // The wrapper mutates payloads through onPayload before calling the real
   // stream; capture that final payload instead of mocking Google responses.
@@ -552,4 +576,91 @@ describe("google prompt cache", () => {
     expect(wrapped).toBeUndefined();
     expect(fetchMock).not.toHaveBeenCalled();
   });
+
+  it("bounds an oversized cache-creation response body instead of buffering it", async () => {
+    const now = 4_000_000;
+    const { response, cancel } = createOversizedJsonResponse();
+    const fetchMock = vi.fn(async () => response);
+    const sessionManager = makeSessionManager();
+    const innerStreamFn = vi.fn(() => "stream" as never);
+
+    const wrapped = await preparePromptCacheStream({
+      fetchMock,
+      now,
+      sessionManager,
+      streamFn: innerStreamFn,
+    });
+
+    await expect(
+      Promise.resolve(
+        wrapped?.(
+          makeGoogleModel(),
+          { systemPrompt: "Follow policy.", messages: [] } as never,
+          {} as never,
+        ),
+      ),
+    ).rejects.toThrow(/Google prompt cache response too large: \d+ bytes/);
+
+    expect(fetchMock).toHaveBeenCalledTimes(1);
+    expect(callArg(fetchMock, 0, 0)).toBe(
+      "https://generativelanguage.googleapis.com/v1beta/cachedContents",
+    );
+    expect(cancel).toHaveBeenCalledOnce();
+  });
+
+  it("bounds an oversized cache-refresh response body instead of buffering it", async () => {
+    const now = 4_500_000;
+    const expireSoon = new Date(now + 60_000).toISOString();
+    const systemPromptDigest = crypto.createHash("sha256").update("Follow policy.").digest("hex");
+    const sessionManager = makeSessionManager([
+      {
+        id: "entry-1",
+        parentId: null,
+        timestamp: new Date(now - 5_000).toISOString(),
+        type: "custom",
+        customType: "openclaw.google-prompt-cache",
+        data: {
+          status: "ready",
+          timestamp: now - 5_000,
+          provider: "google",
+          modelId: "gemini-3.1-pro-preview",
+          modelApi: "google-generative-ai",
+          baseUrl: "https://generativelanguage.googleapis.com/v1beta",
+          systemPromptDigest,
+          cacheRetention: "long",
+          cachedContent: "cachedContents/system-cache-overflow",
+          expireTime: expireSoon,
+        },
+      },
+    ]);
+    const { response, cancel } = createOversizedJsonResponse();
+    const fetchMock = vi.fn(async () => response);
+    const { streamFn: innerStreamFn, getCapturedPayload } = createCapturingStreamFn();
+
+    const wrapped = await preparePromptCacheStream({
+      fetchMock,
+      now,
+      sessionManager,
+      streamFn: innerStreamFn,
+    });
+
+    // The TTL-refresh read swallows errors (.catch(() => null)) and falls back
+    // to the still-valid cached content, so the oversized body must be cancelled
+    // by the byte cap rather than fully buffered.
+    await Promise.resolve(
+      wrapped?.(
+        makeGoogleModel(),
+        { systemPrompt: "Follow policy.", messages: [] } as never,
+        {} as never,
+      ),
+    );
+
+    expect(fetchMock).toHaveBeenCalledTimes(1);
+    expect(fetchUrl(fetchMock)).toBe(
+      "https://generativelanguage.googleapis.com/v1beta/cachedContents/system-cache-overflow?updateMask=ttl",
+    );
+    expect(fetchInit(fetchMock).method).toBe("PATCH");
+    expect(cancel).toHaveBeenCalledOnce();
+    expect(getCapturedPayload()?.cachedContent).toBe("cachedContents/system-cache-overflow");
+  });
 });
diff --git a/src/agents/embedded-agent-runner/google-prompt-cache.ts b/src/agents/embedded-agent-runner/google-prompt-cache.ts
@@ -2,6 +2,7 @@
  * Prepares Google prompt-cache payloads for embedded-agent stream calls.
  */
 import crypto from "node:crypto";
+import { readResponseWithLimit } from "@openclaw/media-core/read-response-with-limit";
 import {
   asDateTimestampMs,
   isFutureDateTimestampMs,
@@ -23,6 +24,9 @@ import { isGooglePromptCacheEligible, resolveCacheRetention } from "./prompt-cac
 import { EmbeddedAttemptSessionTakeoverError } from "./run/attempt.session-lock.js";
 
 const GOOGLE_PROMPT_CACHE_CUSTOM_TYPE = "openclaw.google-prompt-cache";
+// CachedContent metadata responses are tiny (name + expireTime); cap the read so
+// a buggy/hostile Google endpoint cannot stream an unbounded body into memory.
+const GOOGLE_PROMPT_CACHE_RESPONSE_MAX_BYTES = 1024 * 1024;
 const GOOGLE_PROMPT_CACHE_RETRY_BACKOFF_MS = 10 * 60_000;
 const GOOGLE_PROMPT_CACHE_SHORT_REFRESH_WINDOW_MS = 30_000;
 const GOOGLE_PROMPT_CACHE_LONG_REFRESH_WINDOW_MS = 5 * 60_000;
@@ -278,6 +282,19 @@ async function cancelUnreadResponseBody(response: Response | undefined): Promise
   }
 }
 
+/**
+ * Reads a Google cachedContents JSON body under a byte cap and parses it.
+ * Streams through the shared limiter so an oversized response is cancelled
+ * mid-flight instead of being fully buffered by `response.json()`.
+ */
+async function readGooglePromptCacheJson<T>(response: Response): Promise<T> {
+  const buffer = await readResponseWithLimit(response, GOOGLE_PROMPT_CACHE_RESPONSE_MAX_BYTES, {
+    onOverflow: ({ size, maxBytes }) =>
+      new Error(`Google prompt cache response too large: ${size} bytes (limit: ${maxBytes} bytes)`),
+  });
+  return JSON.parse(buffer.toString("utf8")) as T;
+}
+
 async function updateGooglePromptCacheTtl(params: {
   apiKey: string;
   baseUrl: string;
@@ -300,7 +317,7 @@ async function updateGooglePromptCacheTtl(params: {
     if (!response.ok) {
       return null;
     }
-    const json = (await response.json()) as { expireTime?: string };
+    const json = await readGooglePromptCacheJson<{ expireTime?: string }>(response);
     return json;
   } finally {
     await cancelUnreadResponseBody(response);
@@ -338,7 +355,7 @@ async function createGooglePromptCache(params: {
     if (!response.ok) {
       return null;
     }
-    const json = (await response.json()) as { name?: string; expireTime?: string };
+    const json = await readGooglePromptCacheJson<{ name?: string; expireTime?: string }>(response);
     const cachedContent = normalizeOptionalString(json.name) ?? "";
     return cachedContent ? { cachedContent, expireTime: json.expireTime } : null;
   } finally {
