zeitlich 0.2.49 → 0.2.51

package/src/adapters/thread/google-genai/model-invoker.test.ts

@@ -0,0 +1,386 @@
+import { describe, expect, it, vi } from "vitest";
+import {
+  FunctionCallingConfigMode,
+  type Content,
+  type GenerateContentResponse,
+  type Part,
+} from "@google/genai";
+import { createGoogleGenAIModelInvoker } from "./model-invoker";
+import type { StoredContent } from "./thread-manager";
+import type { AgentResponse } from "../../../lib/model";
+import { THREAD_TTL_SECONDS } from "../../../lib/thread/keys";
+
+const textReply: Part[] = [{ text: "ok" }];
+
+function createMockRedis(
+  stored: StoredContent[],
+  extra?: Record<string, string>
+) {
+  return {
+    exists: vi.fn().mockResolvedValue(1),
+    lRange: vi.fn().mockResolvedValue(stored.map((m) => JSON.stringify(m))),
+    lTrim: vi.fn().mockResolvedValue("OK"),
+    get: vi
+      .fn()
+      .mockImplementation((key: string) =>
+        Promise.resolve(extra?.[key] ?? null)
+      ),
+    del: vi.fn().mockResolvedValue(1),
+    set: vi.fn().mockResolvedValue("OK"),
+    rPush: vi.fn().mockResolvedValue(1),
+    expire: vi.fn().mockResolvedValue(1),
+    eval: vi.fn().mockResolvedValue(1),
+  };
+}
+
+function createMockClient(parts: Part[] = textReply) {
+  const chunk: Partial<GenerateContentResponse> = {
+    candidates: [{ content: { role: "model", parts } }],
+    usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 5 },
+  };
+  return {
+    models: {
+      generateContentStream: vi.fn().mockResolvedValue({
+        async *[Symbol.asyncIterator]() {
+          yield chunk;
+        },
+      }),
+    },
+    caches: {
+      create: vi.fn().mockResolvedValue({ name: "cached-content-ref" }),
+    },
+  };
+}
+
+const defaultStored: StoredContent[] = [
+  {
+    id: "msg-1",
+    content: { role: "user", parts: [{ text: "classify these files" }] },
+  },
+];
+
+const invokerConfig = {
+  threadId: "thread-1",
+  assistantMessageId: "assistant-1",
+  state: { tools: [] } as never,
+  agentName: "TestAgent",
+};
+
+function invoke(parts: Part[]): Promise<AgentResponse<Content>> {
+  const redis = createMockRedis(defaultStored);
+  const client = createMockClient(parts);
+
+  const invoker = createGoogleGenAIModelInvoker({
+    redis: redis as never,
+    client: client as never,
+    model: "gemini-2.5-flash",
+  });
+
+  return invoker(invokerConfig);
+}
+
+describe("Google GenAI model invoker — function call IDs", () => {
+  it("assigns synthetic IDs when Gemini omits them", async () => {
+    const result = await invoke([
+      { functionCall: { name: "classifyFile", args: { index: 0 } } },
+      { functionCall: { name: "classifyFile", args: { index: 1 } } },
+    ]);
+
+    expect(result.rawToolCalls).toHaveLength(2);
+    for (const tc of result.rawToolCalls) {
+      expect(tc.id).toBeDefined();
+      expect(tc.id).not.toBe("");
+    }
+  });
+
+  it("preserves existing IDs from Gemini when present", async () => {
+    const result = await invoke([
+      {
+        functionCall: {
+          id: "gemini-abc123",
+          name: "lookupFile",
+          args: { path: "/a" },
+        },
+      },
+    ]);
+
+    expect(result.rawToolCalls[0]?.id).toBe("gemini-abc123");
+  });
+
+  it("generates unique IDs across multiple function calls", async () => {
+    const parts: Part[] = Array.from({ length: 5 }, (_, i) => ({
+      functionCall: { name: "inspect", args: { index: i } },
+    }));
+
+    const result = await invoke(parts);
+
+    const ids = result.rawToolCalls.map((tc) => tc.id);
+    expect(new Set(ids).size).toBe(5);
+  });
+
+  it("matches IDs between message parts and rawToolCalls", async () => {
+    const result = await invoke([
+      { functionCall: { name: "toolA", args: {} } },
+      { functionCall: { name: "toolB", args: {} } },
+    ]);
+
+    const partIds = result.message.parts
+      ?.filter((p) => p.functionCall)
+      .map((p) => p.functionCall?.id);
+    const rawIds = result.rawToolCalls.map((tc) => tc.id);
+
+    expect(partIds).toEqual(rawIds);
+  });
+
+  it("handles a mix of parts with and without existing IDs", async () => {
+    const result = await invoke([
+      { functionCall: { id: "existing-id", name: "toolA", args: {} } },
+      { functionCall: { name: "toolB", args: {} } },
+      { text: "some reasoning text" },
+    ]);
+
+    expect(result.rawToolCalls).toHaveLength(2);
+    expect(result.rawToolCalls[0]?.id).toBe("existing-id");
+    expect(result.rawToolCalls[1]?.id).toBeDefined();
+    expect(result.rawToolCalls[1]?.id).not.toBe("");
+    expect(result.rawToolCalls[1]?.id).not.toBe("existing-id");
+  });
+});
+
+describe("Google GenAI model invoker — context caching", () => {
+  const multiMessageThread: StoredContent[] = [
+    {
+      id: "msg-1",
+      content: {
+        role: "user",
+        parts: [{ inlineData: { data: "base64img", mimeType: "image/png" } }],
+      },
+    },
+    {
+      id: "msg-2",
+      content: { role: "model", parts: [{ text: "I see the image" }] },
+    },
+    {
+      id: "msg-3",
+      content: { role: "user", parts: [{ text: "classify it" }] },
+    },
+  ];
+
+  it("creates a cache and sends only live contents when contents exceed splitIndex", async () => {
+    const redis = createMockRedis(multiMessageThread);
+    const client = createMockClient();
+
+    const invoker = createGoogleGenAIModelInvoker({
+      redis: redis as never,
+      client: client as never,
+      model: "gemini-2.5-flash",
+      cache: { splitIndex: 1 },
+    });
+
+    await invoker(invokerConfig);
+
+    expect(client.caches.create).toHaveBeenCalledOnce();
+    const cacheCall = client.caches.create.mock.calls[0]?.[0];
+    expect(cacheCall.model).toBe("gemini-2.5-flash");
+    expect(cacheCall.config.contents).toHaveLength(1);
+    expect(cacheCall.config.ttl).toBe("300s");
+
+    const streamCall = client.models.generateContentStream.mock.calls[0]?.[0];
+    expect(streamCall.contents).toHaveLength(2);
+    expect(streamCall.config.cachedContent).toBe("cached-content-ref");
+    expect(streamCall.config.systemInstruction).toBeUndefined();
+    expect(streamCall.config.tools).toBeUndefined();
+  });
+
+  it("skips caching when contents.length <= splitIndex", async () => {
+    const redis = createMockRedis(defaultStored);
+    const client = createMockClient();
+
+    const invoker = createGoogleGenAIModelInvoker({
+      redis: redis as never,
+      client: client as never,
+      model: "gemini-2.5-flash",
+      cache: { splitIndex: 1 },
+    });
+
+    await invoker(invokerConfig);
+
+    expect(client.caches.create).not.toHaveBeenCalled();
+    const streamCall = client.models.generateContentStream.mock.calls[0]?.[0];
+    expect(streamCall.contents).toHaveLength(1);
+    expect(streamCall.config.cachedContent).toBeUndefined();
+  });
+
+  it("uses custom TTL", async () => {
+    const redis = createMockRedis(multiMessageThread);
+    const client = createMockClient();
+
+    const invoker = createGoogleGenAIModelInvoker({
+      redis: redis as never,
+      client: client as never,
+      model: "gemini-2.5-flash",
+      cache: { splitIndex: 1, ttlSeconds: 600 },
+    });
+
+    await invoker(invokerConfig);
+
+    const cacheCall = client.caches.create.mock.calls[0]?.[0];
+    expect(cacheCall.config.ttl).toBe("600s");
+  });
+
+  it("moves toolConfig into cache and clears it from live request", async () => {
+    const redis = createMockRedis(multiMessageThread);
+    const client = createMockClient();
+
+    const toolConfig = {
+      functionCallingConfig: { mode: FunctionCallingConfigMode.ANY },
+    };
+
+    const invoker = createGoogleGenAIModelInvoker({
+      redis: redis as never,
+      client: client as never,
+      model: "gemini-2.5-flash",
+      cache: { splitIndex: 1 },
+      config: { toolConfig },
+    });
+
+    await invoker(invokerConfig);
+
+    const cacheCall = client.caches.create.mock.calls[0]?.[0];
+    expect(cacheCall.config.toolConfig).toEqual(toolConfig);
+
+    const streamCall = client.models.generateContentStream.mock.calls[0]?.[0];
+    expect(streamCall.config.toolConfig).toBeUndefined();
+  });
+
+  it("skips caching when splitIndex is 0", async () => {
+    const redis = createMockRedis(multiMessageThread);
+    const client = createMockClient();
+
+    const invoker = createGoogleGenAIModelInvoker({
+      redis: redis as never,
+      client: client as never,
+      model: "gemini-2.5-flash",
+      cache: { splitIndex: 0 },
+    });
+
+    await invoker(invokerConfig);
+
+    expect(client.caches.create).not.toHaveBeenCalled();
+    const streamCall = client.models.generateContentStream.mock.calls[0]?.[0];
+    expect(streamCall.config.cachedContent).toBeUndefined();
+  });
+
+  it("reuses cached content name from Redis instead of creating a new cache", async () => {
+    const redis = createMockRedis(multiMessageThread, {
+      "messages:gemini-cache:gemini-2.5-flash:1:thread:thread-1":
+        "cachedContents/existing",
+    });
+    const client = createMockClient();
+
+    const invoker = createGoogleGenAIModelInvoker({
+      redis: redis as never,
+      client: client as never,
+      model: "gemini-2.5-flash",
+      cache: { splitIndex: 1 },
+    });
+
+    await invoker(invokerConfig);
+
+    expect(client.caches.create).not.toHaveBeenCalled();
+    const streamCall = client.models.generateContentStream.mock.calls[0]?.[0];
+    expect(streamCall.config.cachedContent).toBe("cachedContents/existing");
+    expect(streamCall.contents).toHaveLength(2);
+  });
+
+  it("stores cache name in Redis after creation", async () => {
+    const redis = createMockRedis(multiMessageThread);
+    const client = createMockClient();
+
+    const invoker = createGoogleGenAIModelInvoker({
+      redis: redis as never,
+      client: client as never,
+      model: "gemini-2.5-flash",
+      cache: { splitIndex: 1, ttlSeconds: 600 },
+    });
+
+    await invoker(invokerConfig);
+
+    expect(client.caches.create).toHaveBeenCalledOnce();
+    const setCall = redis.set.mock.calls.find(
+      (c: string[]) =>
+        c[0] === "messages:gemini-cache:gemini-2.5-flash:1:thread:thread-1"
+    );
+    expect(setCall).toBeDefined();
+    expect(setCall?.[1]).toBe("cached-content-ref");
+    expect(setCall?.[2]).toEqual({ EX: 595 });
+  });
+
+  it("reports cachedWriteTokens from cache creation", async () => {
+    const redis = createMockRedis(multiMessageThread);
+    const client = createMockClient();
+    client.caches.create.mockResolvedValue({
+      name: "cached-content-ref",
+      usageMetadata: { totalTokenCount: 4200 },
+    });
+
+    const invoker = createGoogleGenAIModelInvoker({
+      redis: redis as never,
+      client: client as never,
+      model: "gemini-2.5-flash",
+      cache: { splitIndex: 1 },
+    });
+
+    const result = await invoker(invokerConfig);
+
+    expect(result.usage?.cachedWriteTokens).toBe(4200);
+  });
+});
+
+describe("Google GenAI model invoker — thread TTL", () => {
+  // A thread whose tail is a prior attempt's assistant message stored
+  // under `assistant-1`, so the invoker's `truncateFromId(assistant-1)`
+  // trims it and re-stamps the surviving list key's TTL.
+  const retriedThread: StoredContent[] = [
+    { id: "msg-1", content: { role: "user", parts: [{ text: "hi" }] } },
+    {
+      id: "assistant-1",
+      content: { role: "model", parts: [{ text: "prior attempt" }] },
+    },
+  ];
+  const listKey = "messages:thread:thread-1";
+
+  it("re-stamps trimmed hot keys at the configured ttlSeconds", async () => {
+    const redis = createMockRedis(retriedThread);
+    const client = createMockClient();
+
+    const invoker = createGoogleGenAIModelInvoker({
+      redis: redis as never,
+      client: client as never,
+      model: "gemini-2.5-flash",
+      ttlSeconds: 3600,
+    });
+
+    await invoker(invokerConfig);
+
+    expect(redis.lTrim).toHaveBeenCalledWith(listKey, 0, 0);
+    expect(redis.expire).toHaveBeenCalledWith(listKey, 3600);
+    expect(redis.expire).not.toHaveBeenCalledWith(listKey, THREAD_TTL_SECONDS);
+  });
+
+  it("defaults to THREAD_TTL_SECONDS when ttlSeconds is omitted", async () => {
+    const redis = createMockRedis(retriedThread);
+    const client = createMockClient();
+
+    const invoker = createGoogleGenAIModelInvoker({
+      redis: redis as never,
+      client: client as never,
+      model: "gemini-2.5-flash",
+    });
+
+    await invoker(invokerConfig);
+
+    expect(redis.lTrim).toHaveBeenCalledWith(listKey, 0, 0);
+    expect(redis.expire).toHaveBeenCalledWith(listKey, THREAD_TTL_SECONDS);
+  });
+});

package/src/adapters/thread/google-genai/model-invoker.ts

@@ -1,8 +1,10 @@
-import type Redis from "ioredis";
+import type { RedisClientType as Redis } from "redis";
+import { randomBytes } from "node:crypto";
 import type {
   GoogleGenAI,
   Content,
   FunctionDeclaration,
+  GenerateContentConfig,
   Part,
   GenerateContentResponse,
 } from "@google/genai";
@@ -19,6 +21,24 @@ export interface GoogleGenAIModelInvokerConfig {
   client: GoogleGenAI;
   model: string;
   hooks?: GoogleGenAIThreadManagerHooks;
+  /**
+   * Redis TTL for the thread's keys; defaults to 90 days. Use a shorter
+   * value (hours) with a cold tier. Distinct from `cache.ttlSeconds`
+   * (server-side context caching).
+   */
+  ttlSeconds?: number;
+  /** Passed through to `generateContentStream().config`.
+   *  `systemInstruction`, `tools`, and `abortSignal` are managed by the
+   *  invoker and will override any values set here. */
+  config?: GenerateContentConfig;
+  /** Caches the first `splitIndex` messages server-side (with
+   *  `systemInstruction`, `tools`, and `toolConfig`). Skipped when
+   *  `contents.length <= splitIndex`. */
+  cache?: {
+    splitIndex: number;
+    /** Default: 300. */
+    ttlSeconds?: number;
+  };
 }
 
 function toFunctionDeclarations(
@@ -32,12 +52,7 @@ function toFunctionDeclarations(
 }
 
 /**
- * Creates a Google GenAI model invoker that satisfies the generic
- * `ModelInvoker<Content>` contract.
- *
- * Internally streams the response and emits Temporal heartbeats on each
- * chunk so that long-running LLM calls remain visible to the scheduler.
- * The caller is responsible for appending the response to the thread.
+ * The caller is responsible for appending the returned response to the thread.
  *
  * @example
  * ```typescript
@@ -60,6 +75,9 @@ export function createGoogleGenAIModelInvoker({
   client,
   model,
   hooks,
+  ttlSeconds,
+  config: generationConfig,
+  cache: cacheConfig,
 }: GoogleGenAIModelInvokerConfig) {
   return async function invokeGoogleGenAIModel(
     config: ModelInvokerConfig
@@ -72,25 +90,84 @@ export function createGoogleGenAIModelInvoker({
       threadId,
       key: threadKey,
       hooks,
+      ...(ttlSeconds !== undefined && { ttlSeconds }),
     });
     // Truncate the thread starting at the id the assistant message
     // will be stored under. No-op on the first attempt; on rewind
     // retry / Temporal reset it wipes the prior attempt's assistant
     // + tool results so the LLM sees the original pre-call state.
     await thread.truncateFromId(assistantMessageId);
-    const { contents, systemInstruction } =
-      await thread.prepareForInvocation();
+    const { contents, systemInstruction } = await thread.prepareForInvocation();
 
     const functionDeclarations = toFunctionDeclarations(state.tools);
     const tools =
       functionDeclarations.length > 0 ? [{ functionDeclarations }] : undefined;
 
+    const {
+      systemInstruction: _si,
+      tools: _t,
+      abortSignal: _as,
+      cachedContent: callerCachedContent,
+      toolConfig: callerToolConfig,
+      ...callerConfig
+    } = generationConfig ?? {};
+
+    let liveContents = contents;
+    let cachedContentName: string | undefined;
+    let cachedWriteTokens: number | undefined;
+
+    if (
+      cacheConfig &&
+      cacheConfig.splitIndex > 0 &&
+      contents.length > cacheConfig.splitIndex
+    ) {
+      liveContents = contents.slice(cacheConfig.splitIndex);
+      const ttl = cacheConfig.ttlSeconds ?? 300;
+      const cacheRedisKey = `${threadKey ?? "messages"}:gemini-cache:${model}:${cacheConfig.splitIndex}:thread:${threadId}`;
+
+      cachedContentName = (await redis.get(cacheRedisKey)) ?? undefined;
+
+      if (!cachedContentName) {
+        const cacheInstance = await client.caches.create({
+          model,
+          config: {
+            contents: contents.slice(0, cacheConfig.splitIndex),
+            ...(systemInstruction ? { systemInstruction } : {}),
+            ...(tools ? { tools } : {}),
+            ...(callerToolConfig ? { toolConfig: callerToolConfig } : {}),
+            ttl: `${ttl}s`,
+            abortSignal: signal,
+          },
+        });
+        if (!cacheInstance?.name) {
+          throw new Error("Gemini cache creation did not return a cache name");
+        }
+        cachedContentName = cacheInstance.name;
+        cachedWriteTokens =
+          cacheInstance.usageMetadata?.totalTokenCount ?? undefined;
+        const redisTtl = ttl - 5;
+        if (redisTtl > 0) {
+          await redis.set(cacheRedisKey, cachedContentName, { EX: redisTtl });
+        }
+      }
+    }
+
     const stream = await client.models.generateContentStream({
       model,
-      contents,
+      contents: liveContents,
       config: {
-        ...(systemInstruction ? { systemInstruction } : {}),
-        ...(tools ? { tools } : {}),
+        ...callerConfig,
+        ...(cachedContentName
+          ? { cachedContent: cachedContentName }
+          : {
+              ...(callerCachedContent
+                ? { cachedContent: callerCachedContent }
+                : {
+                    ...(systemInstruction ? { systemInstruction } : {}),
+                    ...(tools ? { tools } : {}),
+                  }),
+              ...(callerToolConfig ? { toolConfig: callerToolConfig } : {}),
+            }),
         abortSignal: signal,
       },
     });
@@ -107,48 +184,66 @@ export function createGoogleGenAIModelInvoker({
       throw new Error("Google GenAI stream ended without producing any chunks");
     }
 
+    for (const part of allParts) {
+      if (part.functionCall && !part.functionCall.id) {
+        part.functionCall.id = randomBytes(8).toString("hex");
+      }
+    }
+
     const modelContent: Content = { role: "model", parts: allParts };
-    const functionCalls = lastChunk.functionCalls ?? [];
 
     return {
       message: modelContent,
-      rawToolCalls: functionCalls.map((fc) => ({
-        id: fc.id,
-        name: fc.name ?? "",
-        args: fc.args ?? {},
-      })),
+      rawToolCalls: allParts
+        .filter(
+          (
+            p
+          ): p is Part & { functionCall: NonNullable<Part["functionCall"]> } =>
+            !!p.functionCall
+        )
+        .map((p) => ({
+          id: p.functionCall.id,
+          name: p.functionCall.name ?? "",
+          args: p.functionCall.args ?? {},
+        })),
       usage: {
         inputTokens: lastChunk.usageMetadata?.promptTokenCount,
         outputTokens: lastChunk.usageMetadata?.candidatesTokenCount,
+        cachedWriteTokens,
         cachedReadTokens: lastChunk.usageMetadata?.cachedContentTokenCount,
+        reasonTokens: lastChunk.usageMetadata?.thoughtsTokenCount,
       },
     };
   };
 }
 
-/**
- * Standalone function for one-shot Google GenAI model invocation.
- * Convenience wrapper around createGoogleGenAIModelInvoker for cases
- * where you don't need to reuse the invoker.
- */
 export async function invokeGoogleGenAIModel({
   redis,
   client,
   model,
   hooks,
+  ttlSeconds,
   config,
+  generationConfig,
+  cache,
 }: {
   redis: Redis;
   client: GoogleGenAI;
   model: string;
   hooks?: GoogleGenAIThreadManagerHooks;
+  ttlSeconds?: number;
   config: ModelInvokerConfig;
+  generationConfig?: GenerateContentConfig;
+  cache?: GoogleGenAIModelInvokerConfig["cache"];
 }): Promise<AgentResponse<Content>> {
   const invoker = createGoogleGenAIModelInvoker({
     redis,
     client,
     model,
     hooks,
+    ...(ttlSeconds !== undefined && { ttlSeconds }),
+    config: generationConfig,
+    cache,
   });
   return invoker(config);
 }

package/src/adapters/thread/google-genai/thread-manager.test.ts

@@ -6,10 +6,10 @@ import { createGoogleGenAIThreadManager } from "./thread-manager";
 function createMockRedis(stored: StoredContent[]) {
   return {
     exists: vi.fn().mockResolvedValue(1),
-    lrange: vi.fn().mockResolvedValue(stored.map((m) => JSON.stringify(m))),
+    lRange: vi.fn().mockResolvedValue(stored.map((m) => JSON.stringify(m))),
     del: vi.fn().mockResolvedValue(1),
     set: vi.fn().mockResolvedValue("OK"),
-    rpush: vi.fn().mockResolvedValue(1),
+    rPush: vi.fn().mockResolvedValue(1),
     expire: vi.fn().mockResolvedValue(1),
     eval: vi.fn().mockResolvedValue(1),
   };

package/src/adapters/thread/google-genai/thread-manager.ts

@@ -1,4 +1,4 @@
-import type Redis from "ioredis";
+import type { RedisClientType as Redis } from "redis";
 import type { Content, Part } from "@google/genai";
 import { createThreadManager } from "../../../lib/thread/manager";
 import type {
@@ -32,9 +32,8 @@ export interface GoogleGenAIThreadManagerConfig {
   key?: string;
   hooks?: GoogleGenAIThreadManagerHooks;
   /**
-   * Override the default thread TTL (90 days). When pairing the
-   * adapter with a durable cold tier, a shorter TTL (hours) is
-   * typically more appropriate.
+   * Redis TTL for the thread's keys; defaults to 90 days. Use a shorter
+   * value (hours) with a cold tier.
    */
   ttlSeconds?: number;
 }