zeitlich 0.2.48 → 0.2.50

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,18 @@ export interface GoogleGenAIModelInvokerConfig {
   client: GoogleGenAI;
   model: string;
   hooks?: GoogleGenAIThreadManagerHooks;
+  /** 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 +46,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 +69,8 @@ export function createGoogleGenAIModelInvoker({
   client,
   model,
   hooks,
+  config: generationConfig,
+  cache: cacheConfig,
 }: GoogleGenAIModelInvokerConfig) {
   return async function invokeGoogleGenAIModel(
     config: ModelInvokerConfig
@@ -78,19 +89,77 @@ export function createGoogleGenAIModelInvoker({
     // 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 +176,63 @@ 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,
   config,
+  generationConfig,
+  cache,
 }: {
   redis: Redis;
   client: GoogleGenAI;
   model: string;
   hooks?: GoogleGenAIThreadManagerHooks;
   config: ModelInvokerConfig;
+  generationConfig?: GenerateContentConfig;
+  cache?: GoogleGenAIModelInvokerConfig["cache"];
 }): Promise<AgentResponse<Content>> {
   const invoker = createGoogleGenAIModelInvoker({
     redis,
     client,
     model,
     hooks,
+    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 {

package/src/adapters/thread/langchain/activities.ts

@@ -1,4 +1,4 @@
-import type Redis from "ioredis";
+import type { RedisClientType as Redis } from "redis";
 import type { ToolResultConfig } from "../../../lib/types";
 import type { PersistedThreadState } from "../../../lib/state/types";
 import type { MessageContent } from "@langchain/core/messages";

package/src/adapters/thread/langchain/fork-transform.test.ts

@@ -11,27 +11,27 @@ function createStatefulRedis() {
   const strings = new Map<string, string>();
 
   return {
-    exists: vi.fn(async (...keys: string[]) =>
-      keys.reduce(
+    exists: vi.fn(async (keys: string | string[]) =>
+      (Array.isArray(keys) ? keys : [keys]).reduce(
         (acc, k) => acc + (lists.has(k) || strings.has(k) ? 1 : 0),
         0
       )
     ),
-    lrange: vi.fn(async (key: string, start: number, stop: number) => {
+    lRange: vi.fn(async (key: string, start: number, stop: number) => {
       const list = lists.get(key) ?? [];
       const end = stop === -1 ? list.length : stop + 1;
       return list.slice(start, end);
     }),
-    rpush: vi.fn(async (key: string, ...values: string[]) => {
+    rPush: vi.fn(async (key: string, element: string | string[]) => {
       const list = lists.get(key) ?? [];
-      list.push(...values);
+      list.push(...(Array.isArray(element) ? element : [element]));
       lists.set(key, list);
       return list.length;
     }),
-    ltrim: vi.fn(async () => "OK"),
-    del: vi.fn(async (...keys: string[]) => {
+    lTrim: vi.fn(async () => "OK"),
+    del: vi.fn(async (keys: string | string[]) => {
       let removed = 0;
-      for (const k of keys) {
+      for (const k of Array.isArray(keys) ? keys : [keys]) {
         if (lists.delete(k)) removed++;
         if (strings.delete(k)) removed++;
       }
@@ -43,10 +43,16 @@ function createStatefulRedis() {
     }),
     get: vi.fn(async (key: string) => strings.get(key) ?? null),
     expire: vi.fn(async () => 1),
-    llen: vi.fn(async (key: string) => (lists.get(key) ?? []).length),
+    lLen: vi.fn(async (key: string) => (lists.get(key) ?? []).length),
     eval: vi.fn(
-      async (_script: string, _numKeys: number, ...args: string[]) => {
-        const [dedupKey, listKey, , ...serialised] = args;
+      async (
+        _script: string,
+        options: { keys?: string[]; arguments?: string[] }
+      ) => {
+        const keys = options.keys ?? [];
+        const argv = options.arguments ?? [];
+        const [dedupKey, listKey] = keys;
+        const serialised = argv.slice(1);
         if (!dedupKey || !listKey) return 0;
         if (strings.has(dedupKey)) return 0;
         const list = lists.get(listKey) ?? [];

package/src/adapters/thread/langchain/model-invoker.ts

@@ -1,4 +1,4 @@
-import type Redis from "ioredis";
+import type { RedisClientType as Redis } from "redis";
 import type { AgentResponse, ModelInvokerConfig } from "../../../lib/model";
 import type { StoredMessage } from "@langchain/core/messages";
 import { v4 as uuidv4 } from "uuid";

package/src/adapters/thread/langchain/thread-manager.test.ts

@@ -9,10 +9,10 @@ import { createLangChainThreadManager } from "./thread-manager";
 function createMockRedis(stored: StoredMessage[]) {
   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/langchain/thread-manager.ts

@@ -1,4 +1,4 @@
-import type Redis from "ioredis";
+import type { RedisClientType as Redis } from "redis";
 import type { JsonValue } from "../../../lib/state/types";
 import {
   AIMessage,

package/src/index.ts

@@ -17,8 +17,8 @@
  *   toTree,
  * } from 'zeitlich';
  *
- * // In-memory sandbox adapter
- * import { InMemorySandboxProvider } from 'zeitlich/adapters/sandbox/inmemory';
+ * // Sandbox adapter
+ * import { DaytonaSandboxProvider } from 'zeitlich/adapters/sandbox/daytona';
  *
  * // LangChain adapter
  * import { createLangChainAdapter } from 'zeitlich/adapters/thread/langchain';

package/src/lib/sandbox/capability-types.test.ts

@@ -27,7 +27,7 @@ import type {
   SandboxProvider,
   SandboxSnapshot,
 } from "./types";
-import { InMemorySandboxProvider } from "../../adapters/sandbox/inmemory/index";
+import { InMemorySandboxProvider } from "../../test-utils/in-memory-sandbox";
 import { DaytonaSandboxProvider } from "../../adapters/sandbox/daytona/index";
 import type { E2bSandboxProvider } from "../../adapters/sandbox/e2b/index";
 
@@ -303,7 +303,7 @@ class _ImplWithoutDeclProvider {
 import type { SubagentSandboxConfig } from "../subagent/types";
 import { proxyDaytonaSandboxOps } from "../../adapters/sandbox/daytona/proxy";
 import { proxyE2bSandboxOps } from "../../adapters/sandbox/e2b/proxy";
-import { proxyInMemorySandboxOps } from "../../adapters/sandbox/inmemory/proxy";
+import { proxyInMemorySandboxOps } from "../../test-utils/in-memory-sandbox-proxy";
 
 // Helper that pins the matrix cell type to `SubagentSandboxConfig` so
 // `@ts-expect-error` directives consistently land on the call line. The

package/src/lib/sandbox/manager.ts

@@ -102,12 +102,12 @@ export interface SandboxManagerHooks<
  *
  * @example
  * ```typescript
- * const manager = new SandboxManager(new InMemorySandboxProvider());
+ * const manager = new SandboxManager(new DaytonaSandboxProvider(config));
  * const activities = {
  *   ...manager.createActivities("CodingAgent"),
  *   bashHandler: withSandbox(manager, bashHandler),
  * };
- * // registers: inMemoryCodingAgentCreateSandbox, …
+ * // registers: daytonaCodingAgentCreateSandbox, …
  * ```
  *
  * @example
@@ -336,10 +336,6 @@ export class SandboxManager<
    *
    * @example
    * ```typescript
-   * const manager = new SandboxManager(new InMemorySandboxProvider());
-   * manager.createActivities("CodingAgent");
-   * // registers: inMemoryCodingAgentCreateSandbox, inMemoryCodingAgentDestroySandbox, …
-   *
    * const dmgr = new SandboxManager(new DaytonaSandboxProvider(config));
    * dmgr.createActivities("CodingAgent");
    * // registers: daytonaCodingAgentCreateSandbox, daytonaCodingAgentDestroySandbox

package/src/lib/sandbox/sandbox.test.ts

@@ -1,6 +1,6 @@
 import { describe, expect, it, beforeEach } from "vitest";
 import { SandboxManager } from "./manager";
-import { InMemorySandboxProvider } from "../../adapters/sandbox/inmemory/index";
+import { InMemorySandboxProvider } from "../../test-utils/in-memory-sandbox";
 import {
   SandboxNotFoundError,
   type Sandbox,

package/src/lib/sandbox/types.ts

@@ -342,8 +342,8 @@ export type SandboxOps<
  *
  * @example
  * ```typescript
- * type InMemOps = PrefixedSandboxOps<"inMemory">;
- * // → { inMemoryCreateSandbox, inMemoryDestroySandbox, inMemorySnapshotSandbox, … }
+ * type E2bOps = PrefixedSandboxOps<"e2b">;
+ * // → { e2bCreateSandbox, e2bDestroySandbox, e2bSnapshotSandbox, … }
  * ```
  */
 export type PrefixedSandboxOps<

package/src/lib/session/session.integration.test.ts

@@ -681,6 +681,98 @@ describe("createSession integration", () => {
     expect(capturedSandboxId).toBe("my-sandbox");
   });
 
+  // --- persistThreadState dedupe ---
+
+  it("memoizes persistThreadState across parallel tool calls in one turn", async () => {
+    const { ops, log } = createMockThreadOps();
+
+    const persistTool = defineTool({
+      name: "Persist" as const,
+      description: "calls persistThreadState",
+      schema: z.object({}),
+      handler: async (_args: Record<string, never>, ctx: RouterContext) => {
+        await ctx.persistThreadState?.();
+        return { toolResponse: "ok", data: null };
+      },
+    });
+
+    const session = await createSession({
+      agentName: "TestAgent",
+      thread: { mode: "new", threadId: "thread-1" },
+      runAgent: createScriptedRunAgent([
+        {
+          message: "fan out",
+          toolCalls: [
+            { id: "tc-1", name: "Persist", args: {} },
+            { id: "tc-2", name: "Persist", args: {} },
+            { id: "tc-3", name: "Persist", args: {} },
+          ],
+        },
+        { message: "done", toolCalls: [] },
+      ]),
+      threadOps: ops,
+      tools: { Persist: persistTool },
+      buildContextMessage: () => "go",
+    });
+
+    const stateManager = createAgentStateManager({
+      initialState: { systemPrompt: "test" },
+    });
+
+    await session.runSession({ stateManager });
+
+    const saves = log.filter((l) => l.op === "saveThreadState");
+    expect(saves).toHaveLength(2);
+  });
+
+  it("re-persists across separate turns even when handlers call persistThreadState", async () => {
+    const { ops, log } = createMockThreadOps();
+
+    const persistTool = defineTool({
+      name: "Persist" as const,
+      description: "calls persistThreadState",
+      schema: z.object({}),
+      handler: async (_args: Record<string, never>, ctx: RouterContext) => {
+        await ctx.persistThreadState?.();
+        return { toolResponse: "ok", data: null };
+      },
+    });
+
+    const session = await createSession({
+      agentName: "TestAgent",
+      thread: { mode: "new", threadId: "thread-1" },
+      runAgent: createScriptedRunAgent([
+        {
+          message: "turn 1",
+          toolCalls: [
+            { id: "tc-1", name: "Persist", args: {} },
+            { id: "tc-2", name: "Persist", args: {} },
+          ],
+        },
+        {
+          message: "turn 2",
+          toolCalls: [
+            { id: "tc-3", name: "Persist", args: {} },
+            { id: "tc-4", name: "Persist", args: {} },
+          ],
+        },
+        { message: "done", toolCalls: [] },
+      ]),
+      threadOps: ops,
+      tools: { Persist: persistTool },
+      buildContextMessage: () => "go",
+    });
+
+    const stateManager = createAgentStateManager({
+      initialState: { systemPrompt: "test" },
+    });
+
+    await session.runSession({ stateManager });
+
+    const saves = log.filter((l) => l.op === "saveThreadState");
+    expect(saves).toHaveLength(3);
+  });
+
   // --- Error propagation ---
 
   it("propagates runAgent errors and calls onSessionEnd with failed reason", async () => {

package/src/lib/session/session.ts

@@ -577,6 +577,28 @@ export async function createSession<
             }
           }
 
+          // Hand handlers a way to persist the parent's slice mid-loop
+          // (subagents that fork or continue the parent's thread need
+          // this — otherwise the child loads a stale snapshot from the
+          // prior session, since `saveThreadState` would otherwise only
+          // run in the `finally` below).
+          //
+          // Memoized per-batch so a single assistant message that emits
+          // N parallel subagent calls only writes the slice once.
+          // Persisting again later in the same turn is a no-op anyway
+          // (the slice doesn't mutate between handler dispatch and the
+          // batch's last `executeChild`), and Redis/cold-store writes
+          // aren't free.
+          let persistInflight: Promise<void> | undefined;
+          const persistThreadStateOnce = (): Promise<void> => {
+            persistInflight ??= saveThreadState(
+              threadId,
+              stateManager.getPersistedSlice(),
+              threadKey
+            );
+            return persistInflight;
+          };
+
           const toolCallResults = await toolRouter.processToolCalls(
             parsedToolCalls,
             {
@@ -585,6 +607,7 @@ export async function createSession<
               ...(assistantId !== undefined && {
                 assistantMessageId: assistantId,
               }),
+              persistThreadState: persistThreadStateOnce,
             }
           );
 

package/src/lib/subagent/handler.ts

@@ -439,6 +439,29 @@ export function createSubagentHandler<
       snapshotBaseCreatorAgent.set(childWorkflowId, config.agentName);
     }
 
+    // The parent's `PersistedThreadState` slice (`tasks` + custom
+    // state) only lands in storage in the session's `finally`. When
+    // the child reads from the parent's thread (`from-parent`
+    // fallback or an explicit args.threadId pointing at the parent),
+    // that `loadThreadState` would otherwise see the prior session's
+    // snapshot. Flush the live slice now via the session-supplied
+    // callback so the child sees the parent's current state.
+    if (
+      continuationThreadId &&
+      continuationThreadId === context.threadId &&
+      context.persistThreadState
+    ) {
+      try {
+        await context.persistThreadState();
+      } catch (err) {
+        log.warn("failed to persist parent thread state for subagent", {
+          subagent: config.agentName,
+          childWorkflowId,
+          error: err instanceof Error ? err.message : String(err),
+        });
+      }
+    }
+
     log.info("subagent spawned", {
       subagent: config.agentName,
       childWorkflowId,