diff --git a/packages/test/src/test/ai-provider/OpenAIGeminiCheckpointParams.test.ts b/packages/test/src/test/ai-provider/OpenAIGeminiCheckpointParams.test.ts index eaf784656..8f9d9f521 100644 --- a/packages/test/src/test/ai-provider/OpenAIGeminiCheckpointParams.test.ts +++ b/packages/test/src/test/ai-provider/OpenAIGeminiCheckpointParams.test.ts @@ -7,21 +7,41 @@ import type { AiProviderRunFn, AiSessionContext, + CacheCheckpointTaskInput, + TextGenerationTaskInput, ToolCallingTaskInput, ToolDefinition, } from "@workglow/ai"; import { GOOGLE_GEMINI, _testOnly } from "@workglow/google-gemini/ai"; +import * as GeminiRuntime from "@workglow/google-gemini/ai-runtime"; import { buildGeminiPrefixedContents, deleteGeminiCachedContent, geminiCachedToolsMatch, - getGeminiCachedContent, _testOnly as runtimeTestOnly, - setGeminiCachedContent, } from "@workglow/google-gemini/ai-runtime"; import { mergeOpenAICheckpointPrefix } from "@workglow/openai/ai-runtime"; import { afterEach, beforeEach, describe, expect, it } from "vitest"; +/** + * The ai-runtime bundle carries its own copy of the cache-store map, + * independent of the ai bundle's. The pre-existing store-lifecycle test uses + * the ai-runtime copy because `deleteGeminiCachedContent` is imported off + * ai-runtime and clears its own map. + */ +const runtimeTestOnlyStore = { + setGeminiCachedContent: GeminiRuntime.setGeminiCachedContent, + getGeminiCachedContent: GeminiRuntime.getGeminiCachedContent, +}; + +// The ai bundle carries its own copy of the cache-store map (independent of +// ai-runtime's). Route reads and writes through the ai bundle's `_testOnly` +// helpers so the state the run-fns (also imported off the ai bundle) look at +// is exactly what the test seeds and inspects. +const setGeminiCachedContent = _testOnly.setGeminiCachedContent; +const getGeminiCachedContent = _testOnly.getGeminiCachedContent; +const _cacheStoreTestOnly = _testOnly.cacheStoreTestOnly; + const geminiRequests: Array> = []; // Inject a fake Gemini client through the runtime's own test seam rather than @@ -50,11 +70,13 @@ beforeEach(() => { geminiRequests.length = 0; _testOnly.setGeminiClientForTests(fakeGeminiClient); runtimeTestOnly.setGeminiClientForTests(fakeGeminiClient); + _cacheStoreTestOnly.clearForTests(); }); afterEach(() => { _testOnly.setGeminiClientForTests(undefined); runtimeTestOnly.setGeminiClientForTests(undefined); + _cacheStoreTestOnly.clearForTests(); }); function getGeminiRequests(): Array> { @@ -471,8 +493,13 @@ describe("Gemini tool-calling cached-content parity", () => { describe("Gemini cached-content store", () => { it("stores, retrieves, and idempotently deletes entries", async () => { const id = "test-ckpt-store"; - expect(getGeminiCachedContent(id)).toBeUndefined(); - setGeminiCachedContent(id, { + // Use the ai-runtime bundle's own set/get here — the ai-runtime delete path + // clears the ai-runtime map, so seeding via the same bundle keeps this + // test focused on store lifecycle without cross-bundle plumbing. + const set = runtimeTestOnlyStore.setGeminiCachedContent; + const get = runtimeTestOnlyStore.getGeminiCachedContent; + expect(get(id)).toBeUndefined(); + set(id, { name: "cachedContents/abc", // Deletion lazily builds a client from this config; the entry is removed // from the map before the API call, and API failures are swallowed, so a @@ -480,10 +507,316 @@ describe("Gemini cached-content store", () => { model: { provider_config: { api_key: "test", model_name: "gemini-x" } } as never, systemPrompt: "sys", }); - expect(getGeminiCachedContent(id)?.name).toBe("cachedContents/abc"); + expect(get(id)?.name).toBe("cachedContents/abc"); await deleteGeminiCachedContent(id); - expect(getGeminiCachedContent(id)).toBeUndefined(); + expect(get(id)).toBeUndefined(); // second delete is a no-op await deleteGeminiCachedContent(id); }); }); + +/** + * Overrides only the client methods the test cares about, keeping the shared + * `fakeGeminiClient` shape and the request-log seam intact so every override + * still exercises the same run-fn wiring. + */ +function installGeminiClient(overrides: { + cachesCreate?: (...args: unknown[]) => Promise>; + cachesDelete?: (arg: { name: string }) => Promise; + generateContentStream?: ( + request: Record + ) => Promise>>; +}): { deletedNames: string[]; createCalls: number } { + const deletedNames: string[] = []; + let createCalls = 0; + const client = { + models: { + generateContentStream: + overrides.generateContentStream ?? + (async (request: Record) => { + geminiRequests.push(request); + return { async *[Symbol.asyncIterator]() {} }; + }), + }, + caches: { + create: async (...args: unknown[]) => { + createCalls += 1; + if (overrides.cachesCreate) { + return await overrides.cachesCreate(...args); + } + return {}; + }, + delete: async (arg: { name: string }) => { + deletedNames.push(arg.name); + if (overrides.cachesDelete) await overrides.cachesDelete(arg); + }, + }, + } as never; + _testOnly.setGeminiClientForTests(client); + runtimeTestOnly.setGeminiClientForTests(client); + return { + deletedNames, + get createCalls() { + return createCalls; + }, + } as { deletedNames: string[]; createCalls: number }; +} + +/** Registration lookup — the `serves` list of the run-fn we want to drive. */ +function findGeminiRunFn(capability: string): AiProviderRunFn { + const registration = _testOnly.GEMINI_RUN_FNS.find(({ serves }) => + (serves as readonly string[]).includes(capability) + ); + expect(registration).toBeDefined(); + return registration!.runFn as AiProviderRunFn; +} + +const testModel = { + provider: GOOGLE_GEMINI, + provider_config: { api_key: "test-key", model_name: "gemini-test" }, +} as never; + +describe("Gemini cache checkpoint warm-up error classification", () => { + const cacheCheckpointPrefix = { + systemPrompt: "sys", + messages: [{ role: "user" as const, content: [{ type: "text" as const, text: "hi" }] }], + }; + const cacheCheckpointInput: CacheCheckpointTaskInput = { + model: "gemini-test", + }; + + it("rethrows an abort and best-effort deletes any resource created before the abort", async () => { + const checkpointId = "abort-before-create"; + const controller = new AbortController(); + const helper = installGeminiClient({ + cachesCreate: async () => { + controller.abort(); + const err = new Error("The user aborted the request."); + err.name = "AbortError"; + throw err; + }, + }); + const runFn = findGeminiRunFn("cache.checkpoint"); + const events: unknown[] = []; + await expect( + runFn( + cacheCheckpointInput, + testModel, + controller.signal, + (event) => events.push(event), + undefined, + { sessionId: checkpointId, prefix: cacheCheckpointPrefix } + ) + ).rejects.toThrow(/abort/i); + expect(events).toEqual([]); + expect(getGeminiCachedContent(checkpointId)).toBeUndefined(); + // No resource name was ever returned, so nothing to server-delete. + expect(helper.deletedNames).toEqual([]); + }); + + it("preserves the degrade-to-inline path on a prefix-too-small 400", async () => { + const checkpointId = "degrade-preserved"; + installGeminiClient({ + cachesCreate: async () => { + throw Object.assign(new Error("cached content prefix too small"), { status: 400 }); + }, + }); + const runFn = findGeminiRunFn("cache.checkpoint"); + const events: Array> = []; + await runFn( + cacheCheckpointInput, + testModel, + new AbortController().signal, + (event) => events.push(event as Record), + undefined, + { sessionId: checkpointId, prefix: cacheCheckpointPrefix } + ); + // finish emitted, store empty, no rethrow + expect(events).toHaveLength(1); + expect(events[0].type).toBe("finish"); + expect(getGeminiCachedContent(checkpointId)).toBeUndefined(); + }); + + it("rethrows a quota 429 without attempting to delete (no resource name)", async () => { + const checkpointId = "throws-429"; + const helper = installGeminiClient({ + cachesCreate: async () => { + throw Object.assign(new Error("quota"), { status: 429 }); + }, + }); + const runFn = findGeminiRunFn("cache.checkpoint"); + await expect( + runFn(cacheCheckpointInput, testModel, new AbortController().signal, () => {}, undefined, { + sessionId: checkpointId, + prefix: cacheCheckpointPrefix, + }) + ).rejects.toMatchObject({ status: 429 }); + expect(helper.deletedNames).toEqual([]); + expect(getGeminiCachedContent(checkpointId)).toBeUndefined(); + }); + + it("cleans up the created resource when a post-create step throws", async () => { + const checkpointId = "partial-delete"; + const helper = installGeminiClient({ + cachesCreate: async () => ({ name: "cachedContents/abc" }), + }); + // Inject a bookkeeping failure at the store-insert step so the classifier + // sees a non-abort, non-degrade throw *after* a resource has been minted. + _cacheStoreTestOnly.setPreSetHook(() => { + throw new Error("bookkeeping failed"); + }); + try { + const runFn = findGeminiRunFn("cache.checkpoint"); + await expect( + runFn(cacheCheckpointInput, testModel, new AbortController().signal, () => {}, undefined, { + sessionId: checkpointId, + prefix: cacheCheckpointPrefix, + }) + ).rejects.toThrow(/bookkeeping failed/); + expect(helper.deletedNames).toEqual(["cachedContents/abc"]); + } finally { + _cacheStoreTestOnly.setPreSetHook(undefined); + } + }); +}); +describe("Gemini cachedContent NOT_FOUND fallback (text.generation)", () => { + it("evicts and retries inline once on a 404 NOT_FOUND", async () => { + const checkpointId = "not-found-text"; + setGeminiCachedContent(checkpointId, { + name: "cachedContents/text", + model: testModel, + systemPrompt: "sys", + }); + const requests: Record[] = []; + let call = 0; + installGeminiClient({ + generateContentStream: async (request) => { + requests.push(request); + call += 1; + if (call === 1) { + throw Object.assign(new Error("cachedContents/text NOT_FOUND"), { + status: 404, + code: "NOT_FOUND", + }); + } + return { async *[Symbol.asyncIterator]() {} }; + }, + }); + const runFn = findGeminiRunFn("text.generation"); + const input: TextGenerationTaskInput = { model: "gemini-test", prompt: "tail" }; + const session: AiSessionContext = { + sessionId: checkpointId, + prefix: { + systemPrompt: "sys", + messages: [{ role: "user", content: [{ type: "text", text: "prefix" }] }], + }, + }; + const events: Array> = []; + await runFn( + input, + testModel, + new AbortController().signal, + (event) => events.push(event as Record), + undefined, + session + ); + expect(requests).toHaveLength(2); + expect((requests[0].config as Record).cachedContent).toBe( + "cachedContents/text" + ); + expect((requests[1].config as Record).cachedContent).toBeUndefined(); + // Retry replays the prefix inline (prefix + tail messages). + const retryContents = requests[1].contents as Array<{ parts: Array<{ text?: string }> }>; + const retryTexts = retryContents.map((c) => c.parts[0]?.text); + expect(retryTexts).toContain("prefix"); + expect(retryTexts).toContain("tail"); + expect(getGeminiCachedContent(checkpointId)).toBeUndefined(); + expect(events.some((e) => e.type === "finish")).toBe(true); + }); +}); + +describe("Gemini cachedContent NOT_FOUND fallback (tool-use)", () => { + it("evicts and retries inline once on a NOT_FOUND", async () => { + const checkpointId = "not-found-tool"; + setGeminiCachedContent(checkpointId, { + name: "cachedContents/tool", + model: testModel, + systemPrompt: undefined, + }); + const requests: Record[] = []; + let call = 0; + installGeminiClient({ + generateContentStream: async (request) => { + requests.push(request); + call += 1; + if (call === 1) { + throw Object.assign(new Error("cache not found"), { + status: 404, + code: "NOT_FOUND", + }); + } + return { async *[Symbol.asyncIterator]() {} }; + }, + }); + const runFn = findGeminiRunFn("tool-use"); + const tools = cachedTools; + const input: ToolCallingTaskInput = { + model: "gemini-test", + prompt: "run tool", + tools, + toolChoice: "auto", + }; + const session: AiSessionContext = { + sessionId: checkpointId, + prefix: { + tools, + messages: [{ role: "user", content: [{ type: "text", text: "prefix" }] }], + }, + }; + await runFn(input, testModel, new AbortController().signal, () => {}, undefined, session); + expect(requests).toHaveLength(2); + expect((requests[0].config as Record).cachedContent).toBe( + "cachedContents/tool" + ); + const retryConfig = requests[1].config as { + cachedContent?: string; + tools?: Array>; + }; + expect(retryConfig.cachedContent).toBeUndefined(); + expect(retryConfig.tools).toBeDefined(); + expect(getGeminiCachedContent(checkpointId)).toBeUndefined(); + }); +}); + +describe("Gemini cachedContent proactive stale fallback", () => { + it("skips the cached-content handle on a stale entry and does not re-issue", async () => { + const checkpointId = "proactive-stale"; + // 3.6M ms > 3.5M default stale horizon + setGeminiCachedContent(checkpointId, { + name: "cachedContents/stale", + model: testModel, + systemPrompt: "sys", + createdAtMs: Date.now() - 3_600_000, + }); + const requests: Record[] = []; + installGeminiClient({ + generateContentStream: async (request) => { + requests.push(request); + return { async *[Symbol.asyncIterator]() {} }; + }, + }); + const runFn = findGeminiRunFn("text.generation"); + const input: TextGenerationTaskInput = { model: "gemini-test", prompt: "tail" }; + const session: AiSessionContext = { + sessionId: checkpointId, + prefix: { + systemPrompt: "sys", + messages: [{ role: "user", content: [{ type: "text", text: "prefix" }] }], + }, + }; + await runFn(input, testModel, new AbortController().signal, () => {}, undefined, session); + expect(requests).toHaveLength(1); + expect((requests[0].config as Record).cachedContent).toBeUndefined(); + expect(getGeminiCachedContent(checkpointId)).toBeUndefined(); + }); +}); diff --git a/providers/google-gemini/src/ai/common/Gemini_CacheCheckpoint.ts b/providers/google-gemini/src/ai/common/Gemini_CacheCheckpoint.ts index a24c7b250..499b60642 100644 --- a/providers/google-gemini/src/ai/common/Gemini_CacheCheckpoint.ts +++ b/providers/google-gemini/src/ai/common/Gemini_CacheCheckpoint.ts @@ -245,6 +245,43 @@ function promptTailMessages(prompt: unknown): ChatMessage[] { * applies. The consumed cache also expires by TTL, which the inline-replay * fallback covers as well. */ +/** + * Bucket a cache-creation (or subsequent) error into one of three fates: + * - `"abort"` — the caller cancelled (via the AbortSignal or a bubbled-up + * AbortError); the run-fn must rethrow so the abort surfaces to the run. + * - `"degrade"` — the model rejected the prefix as too small / unsupported for + * explicit caching (`400 INVALID_ARGUMENT` with a matching message); the + * warm-up degrades to no entry and the consumer replays inline. + * - `"throw"` — every other class (auth, quota, transport, server 5xx) is a + * real failure; the run-fn rethrows so the caller sees it and can retry. + */ +function classifyGeminiCacheError( + err: unknown, + signal: AbortSignal | undefined +): "abort" | "degrade" | "throw" { + const anyErr = err as { name?: unknown; message?: unknown; status?: unknown; code?: unknown }; + const message = String(anyErr?.message ?? err ?? ""); + const name = String(anyErr?.name ?? ""); + if ( + signal?.aborted || + (typeof DOMException !== "undefined" && + err instanceof DOMException && + err.name === "AbortError") || + name === "AbortError" || + /aborted|AbortError/i.test(message) + ) { + return "abort"; + } + const status = anyErr?.status; + const code = anyErr?.code; + const looksLikePrefixTooSmall = + /prefix.*too.*small|cached.*content.*not.*supported|minimum.*token/i.test(message); + if ((status === 400 || code === "INVALID_ARGUMENT") && looksLikePrefixTooSmall) { + return "degrade"; + } + return "throw"; +} + export const Gemini_CacheCheckpoint_Stream: AiProviderRunFn< CacheCheckpointTaskInput, CacheCheckpointTaskOutput, @@ -264,6 +301,10 @@ export const Gemini_CacheCheckpoint_Stream: AiProviderRunFn< ? buildGeminiContents(prefix.messages, "") : [{ role: "user", parts: [{ text: "." }] }]; + // Track the created resource name across the try / catch so a downstream + // failure (e.g. a bookkeeping throw from `setGeminiCachedContent`) can still + // cleanup the server-side entry it just minted. + let createdName: string | undefined; try { signal?.throwIfAborted?.(); const cached = await ai.caches.create({ @@ -277,6 +318,7 @@ export const Gemini_CacheCheckpoint_Stream: AiProviderRunFn< ttl: GEMINI_CACHE_TTL, }, } as Parameters[0]); + createdName = cached?.name ?? undefined; if (cached?.name) { setGeminiCachedContent(checkpointId, { name: cached.name, @@ -285,6 +327,15 @@ export const Gemini_CacheCheckpoint_Stream: AiProviderRunFn< }); } } catch (err) { + const fate = classifyGeminiCacheError(err, signal); + if (fate === "abort" || fate === "throw") { + if (createdName) { + await ai.caches + .delete({ name: createdName } as Parameters[0]) + .catch(() => {}); + } + throw err; + } getLogger().warn( `Gemini cache checkpoint warm-up degraded to inline replay: ${ err instanceof Error ? err.message : String(err) diff --git a/providers/google-gemini/src/ai/common/Gemini_CacheStore.ts b/providers/google-gemini/src/ai/common/Gemini_CacheStore.ts index 70faa4597..6a6611ea4 100644 --- a/providers/google-gemini/src/ai/common/Gemini_CacheStore.ts +++ b/providers/google-gemini/src/ai/common/Gemini_CacheStore.ts @@ -22,16 +22,76 @@ export interface GeminiCachedContentEntry { readonly model: GeminiModelConfig; /** The prefix system prompt baked into the cache (consumption must not resend it). */ readonly systemPrompt: string | undefined; + /** + * Wall-clock timestamp (ms since epoch) recorded when the entry was inserted, + * used by the proactive stale check so consumers can evict an entry before + * the server-side TTL burns them into a NOT_FOUND. + */ + readonly createdAtMs: number; } +/** + * Default staleness horizon (ms). Cache creation uses a 3600s TTL; treat + * entries older than ~58 minutes as stale so consumers proactively fall back to + * inline replay before the server-side entry actually expires. + */ +const GEMINI_CACHE_DEFAULT_MAX_AGE_MS = 3_500_000; + const geminiCachedContents = new Map(); +let _preSetHook: ((id: string) => void) | undefined; + export function getGeminiCachedContent(id: string): GeminiCachedContentEntry | undefined { return geminiCachedContents.get(id); } -export function setGeminiCachedContent(id: string, entry: GeminiCachedContentEntry): void { - geminiCachedContents.set(id, entry); +export function setGeminiCachedContent( + id: string, + entry: Omit & + Partial> +): void { + _preSetHook?.(id); + const createdAtMs = entry.createdAtMs ?? Date.now(); + geminiCachedContents.set(id, { ...entry, createdAtMs }); +} + +/** + * @internal Test-only seam that lets `@workglow/test` inject a hook fired + * immediately before every `setGeminiCachedContent` insertion — used to + * simulate a bookkeeping failure so the run-fn's partial-delete cleanup path + * can be exercised. Pass `undefined` to clear. Not part of the stable API. + */ +export const _cacheStoreTestOnly = { + setPreSetHook(hook: ((id: string) => void) | undefined): void { + _preSetHook = hook; + }, + clearForTests(): void { + geminiCachedContents.clear(); + _preSetHook = undefined; + }, +} as const; + +/** + * Returns `true` when the entry is older than `maxAgeMs`. Consumers call this + * before referencing a cache handle so a soon-to-expire entry falls back to + * inline replay instead of erroring the request. + */ +export function isGeminiCacheEntryStale( + entry: GeminiCachedContentEntry, + maxAgeMs: number = GEMINI_CACHE_DEFAULT_MAX_AGE_MS +): boolean { + return Date.now() - entry.createdAtMs > maxAgeMs; +} + +/** + * Removes only the runtime-local map entry for `id`, leaving the server-side + * CachedContent alone. Consumers use this on a proactive stale eviction — the + * server-side entry is about to TTL out on its own, so a delete round-trip is + * unnecessary. Use {@link deleteGeminiCachedContent} when the server-side + * resource must also be released. + */ +export function deleteGeminiCachedContentLocal(id: string): void { + geminiCachedContents.delete(id); } /** diff --git a/providers/google-gemini/src/ai/common/Gemini_CachedContentFallback.ts b/providers/google-gemini/src/ai/common/Gemini_CachedContentFallback.ts new file mode 100644 index 000000000..5fb57e2bf --- /dev/null +++ b/providers/google-gemini/src/ai/common/Gemini_CachedContentFallback.ts @@ -0,0 +1,66 @@ +/** + * @license + * Copyright 2026 Steven Roussey + * SPDX-License-Identifier: Apache-2.0 + */ + +import { getLogger } from "@workglow/util/worker"; +import { deleteGeminiCachedContent } from "./Gemini_CacheStore"; + +/** + * Reactive NOT_FOUND signature. A CachedContent that TTL-expires (or was + * disposed elsewhere) between the consumer's proactive check and the API call + * surfaces as a 404 / `NOT_FOUND` / "not found" from `generateContentStream`. + * When the caller was referencing that entry, the fallback is the same as the + * stale path: evict, rebuild the request without `cachedContent`, and retry + * once. Any other error propagates untouched. + */ +export function isGeminiCachedContentNotFoundError(err: unknown): boolean { + const anyErr = err as { status?: unknown; code?: unknown; message?: unknown }; + if (anyErr?.status === 404) return true; + if (anyErr?.code === "NOT_FOUND") return true; + const message = String(anyErr?.message ?? err ?? ""); + return /NOT_FOUND|not.*found/i.test(message); +} + +interface ExecuteWithFallbackParams { + /** Whether the pending request references a `cachedContent` handle. */ + readonly useCachedContent: boolean; + /** Checkpoint id whose cache entry the pending request references. */ + readonly checkpointId: string | undefined; + /** Build the current request — called once up front and again on retry. */ + readonly buildRequest: (useCachedContent: boolean) => Record; + /** Kick the request off (`ai.models.generateContentStream(request)`). */ + readonly runStream: (request: Record) => Promise; +} + +/** + * Runs `generateContentStream` with the reactive NOT_FOUND fallback: if the + * initial request references a `cachedContent` handle and the API returns a + * NOT_FOUND, evict the store entry (locally + best-effort server delete), + * rebuild the request inline, and retry **once**. All other errors — including + * a NOT_FOUND on a request that was never using cached content — propagate. + * + * The proactive stale check lives in the callers (they own the request-shape + * choice and want to log a different debug line); this helper covers only the + * reactive path. + */ +export async function generateGeminiStreamWithCacheFallback( + params: ExecuteWithFallbackParams +): Promise { + const { useCachedContent, checkpointId, buildRequest, runStream } = params; + const request = buildRequest(useCachedContent); + try { + return await runStream(request); + } catch (err) { + if (!useCachedContent || !checkpointId || !isGeminiCachedContentNotFoundError(err)) { + throw err; + } + getLogger().debug("Gemini cachedContent NOT_FOUND; replaying inline"); + // Best-effort: also releases the server-side handle if it happens to still + // exist; on a genuine NOT_FOUND this is a no-op the helper swallows. + await deleteGeminiCachedContent(checkpointId); + const retryRequest = buildRequest(false); + return await runStream(retryRequest); + } +} diff --git a/providers/google-gemini/src/ai/common/Gemini_TextGeneration.ts b/providers/google-gemini/src/ai/common/Gemini_TextGeneration.ts index 06d8dc8c9..2531afb25 100644 --- a/providers/google-gemini/src/ai/common/Gemini_TextGeneration.ts +++ b/providers/google-gemini/src/ai/common/Gemini_TextGeneration.ts @@ -11,7 +11,12 @@ import type { } from "@workglow/ai"; import { getLogger } from "@workglow/util/worker"; import { buildGeminiPrefixedContents } from "./Gemini_CacheCheckpoint"; -import { getGeminiCachedContent } from "./Gemini_CacheStore"; +import { generateGeminiStreamWithCacheFallback } from "./Gemini_CachedContentFallback"; +import { + deleteGeminiCachedContentLocal, + getGeminiCachedContent, + isGeminiCacheEntryStale, +} from "./Gemini_CacheStore"; import { createGeminiClient, getModelName, resolveThinkingConfig } from "./Gemini_Client"; import type { GeminiModelConfig } from "./Gemini_ModelSchema"; import { emitGeminiRefusal, geminiRefusalCategory } from "./Gemini_Refusal"; @@ -86,50 +91,91 @@ export const Gemini_TextGeneration_Stream: AiProviderRunFn< // TTL expiry — replay the prefix content inline; implicit caching still // applies there. const prefix = sessionContext?.prefix; - const cachedEntry = sessionContext?.sessionId - ? getGeminiCachedContent(sessionContext.sessionId) - : undefined; + const checkpointId = sessionContext?.sessionId; + const cachedEntry = checkpointId ? getGeminiCachedContent(checkpointId) : undefined; const ownSystemPrompt = hasMessages ? unified.systemPrompt || undefined : undefined; - const useCachedContent = + let useCachedContent = prefix !== undefined && cachedEntry !== undefined && (ownSystemPrompt === undefined || ownSystemPrompt === cachedEntry.systemPrompt); - let contents: any[]; - let systemInstruction: string | undefined; - if (prefix && !useCachedContent) { - contents = buildGeminiPrefixedContents( - prefix, - hasMessages ? (unified.messages as Parameters[0]) : undefined, - unified.prompt - ); - systemInstruction = ownSystemPrompt ?? prefix.systemPrompt; - } else { - contents = hasMessages + // Proactive stale check. Explicit CachedContent is TTL-bound (~1h), and a + // consumer reaching a nearly-expired entry would eat a reactive NOT_FOUND + // that costs a round-trip. Evict the runtime-local entry (leave the server + // side to its own TTL) and fall back to inline replay up front. + if (useCachedContent && cachedEntry && isGeminiCacheEntryStale(cachedEntry) && checkpointId) { + logger.debug("Gemini cache entry stale; falling back to inline replay"); + deleteGeminiCachedContentLocal(checkpointId); + useCachedContent = false; + } + + // Thinking is opt-in here (no default budget); when a budget is configured, + // the output cap is padded so reasoning can't starve the visible answer. + const { thinkingConfig, maxOutputTokens } = resolveThinkingConfig(model, input.maxTokens); + + /** Build the tail-only request that references the CachedContent handle. */ + const buildCachedRequest = (): Record => { + const contents = hasMessages ? buildGeminiContents( unified.messages as Parameters[0], unified.prompt ?? "" ) : [{ role: "user", parts: [{ text: input.prompt }] }]; - systemInstruction = useCachedContent ? undefined : ownSystemPrompt; - } + return { + model: getModelName(model), + contents, + config: { + abortSignal: signal ?? undefined, + systemInstruction: undefined, + cachedContent: cachedEntry!.name, + ...buildGenerationConfig(input), + maxOutputTokens, + thinkingConfig, + }, + }; + }; - // Thinking is opt-in here (no default budget); when a budget is configured, - // the output cap is padded so reasoning can't starve the visible answer. - const { thinkingConfig, maxOutputTokens } = resolveThinkingConfig(model, input.maxTokens); + /** Build the full inline-replay request (prefix messages + tail). */ + const buildInlineReplayRequest = (): Record => { + let contents: any[]; + let systemInstruction: string | undefined; + if (prefix) { + contents = buildGeminiPrefixedContents( + prefix, + hasMessages ? (unified.messages as Parameters[0]) : undefined, + unified.prompt + ); + systemInstruction = ownSystemPrompt ?? prefix.systemPrompt; + } else { + contents = hasMessages + ? buildGeminiContents( + unified.messages as Parameters[0], + unified.prompt ?? "" + ) + : [{ role: "user", parts: [{ text: input.prompt }] }]; + systemInstruction = ownSystemPrompt; + } + return { + model: getModelName(model), + contents, + config: { + abortSignal: signal ?? undefined, + systemInstruction, + ...buildGenerationConfig(input), + maxOutputTokens, + thinkingConfig, + }, + }; + }; - const result = await ai.models.generateContentStream({ - model: getModelName(model), - contents, - config: { - abortSignal: signal ?? undefined, - systemInstruction, - ...(useCachedContent ? { cachedContent: cachedEntry!.name } : {}), - ...buildGenerationConfig(input), - // Override maxOutputTokens from buildGenerationConfig with the thinking-aware value. - maxOutputTokens, - thinkingConfig, - }, + const result = await generateGeminiStreamWithCacheFallback({ + useCachedContent, + checkpointId, + buildRequest: (useCached) => (useCached ? buildCachedRequest() : buildInlineReplayRequest()), + runStream: (request) => + ai.models.generateContentStream( + request as unknown as Parameters[0] + ), }); let refusalCategory: string | undefined; diff --git a/providers/google-gemini/src/ai/common/Gemini_ToolCalling.ts b/providers/google-gemini/src/ai/common/Gemini_ToolCalling.ts index 001955640..d075837b5 100644 --- a/providers/google-gemini/src/ai/common/Gemini_ToolCalling.ts +++ b/providers/google-gemini/src/ai/common/Gemini_ToolCalling.ts @@ -12,12 +12,18 @@ import type { ToolCallingTaskOutput, } from "@workglow/ai"; import { filterValidToolCalls, sanitizeToolArgs } from "@workglow/ai/worker"; +import { getLogger } from "@workglow/util/worker"; import { buildGeminiFunctionDeclarations, buildGeminiPrefixedContents, geminiCachedToolsMatch, } from "./Gemini_CacheCheckpoint"; -import { getGeminiCachedContent } from "./Gemini_CacheStore"; +import { generateGeminiStreamWithCacheFallback } from "./Gemini_CachedContentFallback"; +import { + deleteGeminiCachedContentLocal, + getGeminiCachedContent, + isGeminiCacheEntryStale, +} from "./Gemini_CacheStore"; import { createGeminiClient, getModelName, resolveThinkingConfig } from "./Gemini_Client"; import type { GeminiModelConfig } from "./Gemini_ModelSchema"; import { emitGeminiRefusal, geminiRefusalCategory } from "./Gemini_Refusal"; @@ -137,11 +143,10 @@ export const Gemini_ToolCalling_Stream: AiProviderRunFn< // ("auto") tool choice. Anything else — including after the cache's TTL // expiry — replays the prefix content inline; implicit caching still applies. const prefix = sessionContext?.prefix; - const cachedEntry = sessionContext?.sessionId - ? getGeminiCachedContent(sessionContext.sessionId) - : undefined; + const checkpointId = sessionContext?.sessionId; + const cachedEntry = checkpointId ? getGeminiCachedContent(checkpointId) : undefined; const defaultToolChoice = input.toolChoice === undefined || input.toolChoice === "auto"; - const useCachedContent = + let useCachedContent = prefix !== undefined && cachedEntry !== undefined && defaultToolChoice && @@ -152,34 +157,66 @@ export const Gemini_ToolCalling_Stream: AiProviderRunFn< input.systemPrompt === "" || input.systemPrompt === cachedEntry.systemPrompt); - const contents = - prefix && !useCachedContent - ? buildGeminiPrefixedContents(prefix, input.messages, input.prompt) - : buildGeminiContents(input.messages, input.prompt); - const systemInstruction = useCachedContent - ? undefined - : input.systemPrompt || (prefix ? prefix.systemPrompt : undefined); + // Proactive stale check. Explicit CachedContent is TTL-bound (~1h), and a + // consumer reaching a nearly-expired entry would eat a reactive NOT_FOUND + // that costs a round-trip. Evict the runtime-local entry (leave the server + // side to its own TTL) and fall back to inline replay up front. + if (useCachedContent && cachedEntry && isGeminiCacheEntryStale(cachedEntry) && checkpointId) { + getLogger().debug("Gemini cache entry stale; falling back to inline replay"); + deleteGeminiCachedContentLocal(checkpointId); + useCachedContent = false; + } // Thinking is opt-in here (no default budget): the model uses its own default // reasoning unless `provider_config.thinking_budget` is set, in which case the // output cap is padded so reasoning can't starve the tool call / answer. const { thinkingConfig, maxOutputTokens } = resolveThinkingConfig(model, input.maxTokens); - const result = await ai.models.generateContentStream({ + /** Build the tail-only request that references the CachedContent handle. */ + const buildCachedRequest = (): Record => ({ model: getModelName(model), - contents, + contents: buildGeminiContents(input.messages, input.prompt), config: { abortSignal: signal ?? undefined, - systemInstruction, + systemInstruction: undefined, maxOutputTokens, temperature: input.temperature, - ...(useCachedContent - ? { cachedContent: cachedEntry!.name } - : { tools: [{ functionDeclarations }], toolConfig: toolConfig as any }), + cachedContent: cachedEntry!.name, thinkingConfig, }, }); + /** Build the full inline-replay request (prefix messages + tail + tools). */ + const buildInlineReplayRequest = (): Record => { + const contents = prefix + ? buildGeminiPrefixedContents(prefix, input.messages, input.prompt) + : buildGeminiContents(input.messages, input.prompt); + const systemInstruction = input.systemPrompt || (prefix ? prefix.systemPrompt : undefined); + return { + model: getModelName(model), + contents, + config: { + abortSignal: signal ?? undefined, + systemInstruction, + maxOutputTokens, + temperature: input.temperature, + tools: [{ functionDeclarations }], + toolConfig: toolConfig as any, + thinkingConfig, + }, + }; + }; + + const result = await generateGeminiStreamWithCacheFallback({ + useCachedContent, + checkpointId, + buildRequest: (useCached) => (useCached ? buildCachedRequest() : buildInlineReplayRequest()), + runStream: (request) => + ai.models.generateContentStream( + request as unknown as Parameters[0] + ), + }); + let callIndex = 0; let refusalCategory: string | undefined; diff --git a/providers/google-gemini/src/ai/index.ts b/providers/google-gemini/src/ai/index.ts index 3c2213acb..1c7f9c7f4 100644 --- a/providers/google-gemini/src/ai/index.ts +++ b/providers/google-gemini/src/ai/index.ts @@ -14,13 +14,21 @@ export * from "./registerGemini"; import { GEMINI_RUN_FN_SPECS } from "./common/Gemini_Capabilities"; import { _testOnly as clientTestOnly } from "./common/Gemini_Client"; +import { + _cacheStoreTestOnly, + getGeminiCachedContent, + setGeminiCachedContent, +} from "./common/Gemini_CacheStore"; import { GEMINI_RUN_FNS } from "./common/Gemini_JobRunFns"; import { emitGeminiRefusal, geminiRefusalCategory } from "./common/Gemini_Refusal"; import { buildGeminiContents } from "./common/Gemini_ToolCalling"; import { GoogleGeminiQueuedProvider } from "./GoogleGeminiQueuedProvider"; /** - * @internal Symbols exported only for use by `@workglow/test`. Not part of the stable public API. + * @internal Symbols exported only for use by `@workglow/test`. Not part of the + * stable public API. The cache-store helpers are re-exported off this barrel + * (in addition to `ai-runtime`) so tests that drive the ai-bundle run-fns can + * seed / read the same runtime-local map the run-fns look at. */ export const _testOnly = { GoogleGeminiQueuedProvider, @@ -30,4 +38,7 @@ export const _testOnly = { geminiRefusalCategory, emitGeminiRefusal, setGeminiClientForTests: clientTestOnly.setGeminiClientForTests, + setGeminiCachedContent, + getGeminiCachedContent, + cacheStoreTestOnly: _cacheStoreTestOnly, } as const;