From e6743eb783f9503137c63b86f36fa110677e62bf Mon Sep 17 00:00:00 2001 From: Vincent Koc Date: Fri, 19 Jun 2026 10:30:04 +0200 Subject: [PATCH] fix(agents): cancel prompt cache error bodies --- .../google-prompt-cache.test.ts | 85 +++++++++++++++++++ .../google-prompt-cache.ts | 69 +++++++++------ 2 files changed, 126 insertions(+), 28 deletions(-) diff --git a/src/agents/embedded-agent-runner/google-prompt-cache.test.ts b/src/agents/embedded-agent-runner/google-prompt-cache.test.ts index c33fe461bbea..bfd518d90aec 100644 --- a/src/agents/embedded-agent-runner/google-prompt-cache.test.ts +++ b/src/agents/embedded-agent-runner/google-prompt-cache.test.ts @@ -256,6 +256,40 @@ describe("google prompt cache", () => { ]); }); + it("cancels failed cache creation response bodies", async () => { + const now = 1_500_000; + const response = new Response("permission denied", { status: 403 }); + const cancel = vi.spyOn(response.body!, "cancel").mockResolvedValue(undefined); + const fetchMock = vi.fn(async () => response); + const entries: SessionCustomEntry[] = []; + const sessionManager = makeSessionManager(entries); + const innerStreamFn = vi.fn(() => "stream" as never); + + const wrapped = await preparePromptCacheStream({ + fetchMock, + now, + sessionManager, + streamFn: innerStreamFn, + }); + + await Promise.resolve( + wrapped?.( + makeGoogleModel(), + { systemPrompt: "Follow policy.", messages: [] } as never, + {} as never, + ), + ); + + expect(cancel).toHaveBeenCalledOnce(); + expect(innerStreamFn).toHaveBeenCalledTimes(1); + expect(streamContext(innerStreamFn).systemPrompt).toBe("Follow policy."); + expect(entries[0]?.data).toMatchObject({ + status: "failed", + provider: "google", + modelId: "gemini-3.1-pro-preview", + }); + }); + it("reuses a persisted cache entry without creating a second cache", async () => { const now = 2_000_000; const entries: SessionCustomEntry[] = []; @@ -394,6 +428,57 @@ describe("google prompt cache", () => { expect(getCapturedPayload()?.cachedContent).toBe("cachedContents/system-cache-3"); }); + it("cancels failed cache refresh response bodies", async () => { + const now = 3_500_000; + const expireSoon = new Date(now + 60_000).toISOString(); + const systemPromptDigest = crypto.createHash("sha256").update("Follow policy.").digest("hex"); + const entries: SessionCustomEntry[] = [ + { + 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-4", + expireTime: expireSoon, + }, + }, + ]; + const response = new Response("refresh denied", { status: 403 }); + const cancel = vi.spyOn(response.body!, "cancel").mockResolvedValue(undefined); + const fetchMock = vi.fn(async () => response); + const sessionManager = makeSessionManager(entries); + const { streamFn: innerStreamFn, getCapturedPayload } = createCapturingStreamFn(); + + const wrapped = await preparePromptCacheStream({ + fetchMock, + now, + sessionManager, + streamFn: innerStreamFn, + }); + + await Promise.resolve( + wrapped?.( + makeGoogleModel(), + { systemPrompt: "Follow policy.", messages: [] } as never, + {} as never, + ), + ); + + expect(cancel).toHaveBeenCalledOnce(); + expect(entries).toHaveLength(1); + expect(getCapturedPayload()?.cachedContent).toBe("cachedContents/system-cache-4"); + }); + it("does not bypass failed-cache backoff when the process clock is invalid", async () => { const systemPromptDigest = crypto.createHash("sha256").update("Follow policy.").digest("hex"); const sessionManager = makeSessionManager([ diff --git a/src/agents/embedded-agent-runner/google-prompt-cache.ts b/src/agents/embedded-agent-runner/google-prompt-cache.ts index d50c87f37925..7d4c47918f76 100644 --- a/src/agents/embedded-agent-runner/google-prompt-cache.ts +++ b/src/agents/embedded-agent-runner/google-prompt-cache.ts @@ -272,6 +272,12 @@ function buildManagedContextForCachedContent(context: GooglePromptCacheContext) }; } +async function cancelUnreadResponseBody(response: Response | undefined): Promise { + if (response && !response.bodyUsed) { + await response.body?.cancel().catch(() => undefined); + } +} + async function updateGooglePromptCacheTtl(params: { apiKey: string; baseUrl: string; @@ -281,22 +287,24 @@ async function updateGooglePromptCacheTtl(params: { headers?: Record; signal?: AbortSignal; }): Promise<{ expireTime?: string } | null> { - const response = await params.fetchImpl( - `${params.baseUrl}/${params.cachedContent}?updateMask=ttl`, - { + let response: Response | undefined; + try { + response = await params.fetchImpl(`${params.baseUrl}/${params.cachedContent}?updateMask=ttl`, { method: "PATCH", headers: mergeTransportHeaders(parseGeminiAuth(params.apiKey).headers, params.headers), body: JSON.stringify({ ttl: resolveGooglePromptCacheTtl(params.cacheRetention), }), signal: params.signal, - }, - ); - if (!response.ok) { - return null; + }); + if (!response.ok) { + return null; + } + const json = (await response.json()) as { expireTime?: string }; + return json; + } finally { + await cancelUnreadResponseBody(response); } - const json = (await response.json()) as { expireTime?: string }; - return json; } async function createGooglePromptCache(params: { @@ -311,26 +319,31 @@ async function createGooglePromptCache(params: { tools?: unknown; toolConfig?: unknown; }): Promise<{ cachedContent: string; expireTime?: string } | null> { - const response = await params.fetchImpl(`${params.baseUrl}/cachedContents`, { - method: "POST", - headers: mergeTransportHeaders(parseGeminiAuth(params.apiKey).headers, params.headers), - body: JSON.stringify({ - model: params.modelId.startsWith("models/") ? params.modelId : `models/${params.modelId}`, - ttl: resolveGooglePromptCacheTtl(params.cacheRetention), - systemInstruction: { - parts: [{ text: params.systemPrompt }], - }, - ...(params.tools ? { tools: params.tools } : {}), - ...(params.toolConfig ? { toolConfig: params.toolConfig } : {}), - }), - signal: params.signal, - }); - if (!response.ok) { - return null; + let response: Response | undefined; + try { + response = await params.fetchImpl(`${params.baseUrl}/cachedContents`, { + method: "POST", + headers: mergeTransportHeaders(parseGeminiAuth(params.apiKey).headers, params.headers), + body: JSON.stringify({ + model: params.modelId.startsWith("models/") ? params.modelId : `models/${params.modelId}`, + ttl: resolveGooglePromptCacheTtl(params.cacheRetention), + systemInstruction: { + parts: [{ text: params.systemPrompt }], + }, + ...(params.tools ? { tools: params.tools } : {}), + ...(params.toolConfig ? { toolConfig: params.toolConfig } : {}), + }), + signal: params.signal, + }); + if (!response.ok) { + return null; + } + const json = (await response.json()) as { name?: string; expireTime?: string }; + const cachedContent = normalizeOptionalString(json.name) ?? ""; + return cachedContent ? { cachedContent, expireTime: json.expireTime } : null; + } finally { + await cancelUnreadResponseBody(response); } - const json = (await response.json()) as { name?: string; expireTime?: string }; - const cachedContent = normalizeOptionalString(json.name) ?? ""; - return cachedContent ? { cachedContent, expireTime: json.expireTime } : null; } async function ensureGooglePromptCache(