zeitlich 0.2.49 → 0.2.51

package/src/lib/thread/test-utils.ts

@@ -7,19 +7,35 @@
  * picks it up directly.
  */
 
-import type Redis from "ioredis";
+import type { RedisClientType } from "redis";
 import type { ColdThreadStore, ThreadSnapshot } from "./cold-store";
 
 type Value = string | string[];
 
+/** node-redis `SetOptions` subset the stub understands. */
+interface FakeSetOptions {
+  EX?: number;
+  NX?: boolean;
+  expiration?: { type: "EX" | "PX" | "EXAT" | "PXAT"; value: number } | "KEEPTTL";
+  condition?: "NX" | "XX";
+}
+
+/** node-redis accepts a single key or an array (`RedisVariadicArgument`). */
+type Keys = string | string[];
+const toKeys = (keys: Keys): string[] => (Array.isArray(keys) ? keys : [keys]);
+
 /**
- * Minimal in-memory Redis stub covering the commands the thread
+ * Minimal in-memory node-redis stub covering the commands the thread
  * manager + snapshot helpers use: get/set/del/exists/expire,
- * lrange/rpush/llen/ltrim, and the `eval`-based idempotent-append Lua
- * script. Behaviour matches Redis closely enough for unit tests; TTLs
- * are stored but never expire automatically.
+ * lRange/rPush/lLen/lTrim, and the `eval`-based idempotent-append Lua
+ * script. Mirrors the node-redis (`redis`) v4+ API surface — camelCase
+ * commands, an options object for `set`, variadic-or-array keys for
+ * `del`/`exists`, and a `multi().execAsPipeline()` pipeline that rejects
+ * with a `MultiErrorReply`-shaped error when a queued command fails.
+ * Behaviour matches Redis closely enough for unit tests; TTLs are stored
+ * but never expire automatically.
  */
-export function createFakeRedis(): Redis & {
+export function createFakeRedis(): RedisClientType & {
   _store: Map<string, Value>;
   _ttls: Map<string, number>;
 } {
@@ -48,56 +64,60 @@ export function createFakeRedis(): Redis & {
     async set(
       key: string,
       value: string,
-      ..._rest: (string | number)[]
-    ): Promise<"OK"> {
-      // NX guard: when the args contain "NX" and the key already exists,
+      options?: FakeSetOptions
+    ): Promise<"OK" | null> {
+      // NX guard: when the condition is NX and the key already exists,
       // Redis returns null. We follow the same contract for tests that
-      // need it; existing call sites use this for compare-and-set.
-      const rest = _rest.map((x) => (typeof x === "string" ? x.toUpperCase() : x));
-      if (rest.includes("NX") && store.has(key)) {
-        return null as unknown as "OK";
+      // need it.
+      const nx = options?.NX === true || options?.condition === "NX";
+      if (nx && store.has(key)) {
+        return null;
       }
       store.set(key, String(value));
-      const exIdx = rest.indexOf("EX");
-      if (exIdx >= 0 && typeof _rest[exIdx + 1] === "number") {
-        ttls.set(key, _rest[exIdx + 1] as number);
+      const ttl =
+        options?.EX ??
+        (options?.expiration && options.expiration !== "KEEPTTL"
+          ? options.expiration.value
+          : undefined);
+      if (typeof ttl === "number") {
+        ttls.set(key, ttl);
       }
       return "OK";
     },
-    async del(...keys: string[]): Promise<number> {
+    async del(keys: Keys): Promise<number> {
       let removed = 0;
-      for (const k of keys) {
+      for (const k of toKeys(keys)) {
         if (store.delete(k)) removed++;
         ttls.delete(k);
       }
       return removed;
     },
-    async exists(...keys: string[]): Promise<number> {
-      return keys.reduce((acc, k) => acc + (store.has(k) ? 1 : 0), 0);
+    async exists(keys: Keys): Promise<number> {
+      return toKeys(keys).reduce((acc, k) => acc + (store.has(k) ? 1 : 0), 0);
     },
     async expire(key: string, ttl: number): Promise<number> {
       if (!store.has(key)) return 0;
       ttls.set(key, ttl);
       return 1;
     },
-    async lrange(key: string, start: number, end: number): Promise<string[]> {
+    async lRange(key: string, start: number, end: number): Promise<string[]> {
       if (!store.has(key)) return [];
       if (!isList(key)) return [];
       const list = store.get(key) as string[];
       const last = end === -1 ? list.length - 1 : end;
       return list.slice(start, last + 1);
     },
-    async rpush(key: string, ...values: string[]): Promise<number> {
+    async rPush(key: string, element: Keys): Promise<number> {
       const list = ensureList(key);
-      list.push(...values);
+      list.push(...toKeys(element));
       return list.length;
     },
-    async llen(key: string): Promise<number> {
+    async lLen(key: string): Promise<number> {
       if (!store.has(key)) return 0;
       const list = store.get(key) as string[];
       return list.length;
     },
-    async ltrim(key: string, start: number, end: number): Promise<"OK"> {
+    async lTrim(key: string, start: number, end: number): Promise<"OK"> {
       if (!store.has(key)) return "OK";
       const list = store.get(key) as string[];
       const last = end === -1 ? list.length - 1 : end;
@@ -106,12 +126,11 @@ export function createFakeRedis(): Redis & {
     },
     async eval(
       _script: string,
-      numKeys: number,
-      ...args: (string | number)[]
+      options: { keys?: string[]; arguments?: string[] }
     ): Promise<number> {
       // Mirrors APPEND_IDEMPOTENT_SCRIPT in src/lib/thread/manager.ts.
-      const keys = args.slice(0, numKeys) as string[];
-      const argv = args.slice(numKeys) as string[];
+      const keys = options.keys ?? [];
+      const argv = options.arguments ?? [];
       const dedupKey = keys[0];
       const listKey = keys[1];
       const ttl = Number(argv[0]);
@@ -127,54 +146,64 @@ export function createFakeRedis(): Redis & {
       ttls.set(dedupKey, ttl);
       return 1;
     },
-    // Chainable pipeline stub. Defers each command to the underlying
-    // sync fake methods on `.exec()`, so TTL tracking and store
-    // semantics stay identical to the non-pipelined path. `fake` is
-    // typed as `Redis` after the cast below, so we narrow it back to
-    // the concrete impl shape here to avoid Redis's callback overloads.
-    pipeline(): FakePipeline {
+    // Chainable `multi()` stub. Defers each command to the underlying
+    // sync fake methods on `.execAsPipeline()`, so TTL tracking and store
+    // semantics stay identical to the non-pipelined path. Mirrors
+    // node-redis: per-command failures reject the pipeline with a
+    // `MultiErrorReply`-shaped error (`{ replies, errorIndexes }`).
+    multi(): FakeMulti {
       const impl = fake as unknown as {
-        set: (key: string, value: string, ...rest: (string | number)[]) => Promise<"OK">;
-        del: (...keys: string[]) => Promise<number>;
-        rpush: (key: string, ...values: string[]) => Promise<number>;
+        set: (
+          key: string,
+          value: string,
+          options?: FakeSetOptions
+        ) => Promise<"OK" | null>;
+        del: (keys: Keys) => Promise<number>;
+        rPush: (key: string, element: Keys) => Promise<number>;
         expire: (key: string, ttl: number) => Promise<number>;
       };
       const ops: Array<() => Promise<unknown>> = [];
-      const chain: FakePipeline = {
-        set: (...args) => {
-          const [key, value, ...rest] = args as [string, string, ...(string | number)[]];
-          ops.push(() => impl.set(key, value, ...rest));
+      const chain: FakeMulti = {
+        set: (key, value, options) => {
+          ops.push(() => impl.set(key, value, options));
           return chain;
         },
-        del: (...keys) => {
-          ops.push(() => impl.del(...keys));
+        del: (keys) => {
+          ops.push(() => impl.del(keys));
           return chain;
         },
-        rpush: (key, ...values) => {
-          ops.push(() => impl.rpush(key, ...values));
+        rPush: (key, element) => {
+          ops.push(() => impl.rPush(key, element));
           return chain;
         },
         expire: (key, ttl) => {
           ops.push(() => impl.expire(key, ttl));
           return chain;
         },
-        exec: async () => {
-          const results: Array<[Error | null, unknown]> = [];
+        execAsPipeline: async () => {
+          const replies: unknown[] = [];
+          const errorIndexes: number[] = [];
+          let i = 0;
           for (const op of ops) {
             try {
-              results.push([null, await op()]);
+              replies.push(await op());
             } catch (e) {
-              results.push([e as Error, null]);
+              replies.push(e);
+              errorIndexes.push(i);
             }
+            i++;
+          }
+          if (errorIndexes.length > 0) {
+            throw makeMultiError(replies, errorIndexes);
           }
-          return results;
+          return replies;
         },
       };
       return chain;
     },
     _store: store,
     _ttls: ttls,
-  } as unknown as Redis & {
+  } as unknown as RedisClientType & {
     _store: Map<string, Value>;
     _ttls: Map<string, number>;
   };
@@ -182,13 +211,31 @@ export function createFakeRedis(): Redis & {
   return fake;
 }
 
-/** Minimal chainable surface used by the fake-redis pipeline stub. */
-interface FakePipeline {
-  set: (...args: (string | number)[]) => FakePipeline;
-  del: (...keys: string[]) => FakePipeline;
-  rpush: (key: string, ...values: string[]) => FakePipeline;
-  expire: (key: string, ttl: number) => FakePipeline;
-  exec: () => Promise<Array<[Error | null, unknown]>>;
+/** Minimal chainable surface used by the fake-redis `multi()` stub. */
+interface FakeMulti {
+  set: (key: string, value: string, options?: FakeSetOptions) => FakeMulti;
+  del: (keys: Keys) => FakeMulti;
+  rPush: (key: string, element: Keys) => FakeMulti;
+  expire: (key: string, ttl: number) => FakeMulti;
+  execAsPipeline: () => Promise<unknown[]>;
+}
+
+/**
+ * Build a node-redis `MultiErrorReply`-shaped error: an `Error` carrying
+ * `replies` (per-command results, with failures as `Error`s) and
+ * `errorIndexes`. `applySnapshot` unwraps this to surface the first real
+ * error.
+ */
+export function makeMultiError(
+  replies: unknown[],
+  errorIndexes: number[]
+): Error & { replies: unknown[]; errorIndexes: number[] } {
+  return Object.assign(
+    new Error(
+      `${errorIndexes.length} commands failed, see .replies and .errorIndexes for more information`
+    ),
+    { replies, errorIndexes }
+  );
 }
 
 /**

package/src/lib/thread/tiered.test.ts

@@ -1,5 +1,5 @@
 import { describe, expect, it, beforeEach } from "vitest";
-import type Redis from "ioredis";
+import type { RedisClientType as Redis } from "redis";
 import { createTieredThreadManager } from "./tiered";
 import { createThreadManager } from "./manager";
 import {

package/src/lib/thread/types.ts

@@ -1,7 +1,7 @@
-import type Redis from "ioredis";
+import type { RedisClientType } from "redis";
 import type { JsonValue, PersistedThreadState } from "../state/types";
 export interface ThreadManagerConfig<T> {
-  redis: Redis;
+  redis: RedisClientType;
   threadId: string;
   /** Thread key, defaults to 'messages' */
   key?: string;

package/src/lib/tool-router/router.integration.test.ts

@@ -271,6 +271,50 @@ describe("createToolRouter integration", () => {
     expect(order[1]).toBe("start-echo-b");
   });
 
+  it("appends parallel results in original call order", async () => {
+    const slowEcho = defineTool({
+      name: "Echo" as const,
+      description: "slow echo with variable latency",
+      schema: z.object({ text: z.string(), delay: z.number() }),
+      handler: async (args: { text: string; delay: number }) => {
+        await new Promise((r) => setTimeout(r, args.delay));
+        return { toolResponse: args.text, data: { echoed: args.text } };
+      },
+    });
+
+    const router = createToolRouter({
+      tools: { Echo: slowEcho, Add: mathTool } as const,
+      threadId: "t-1",
+      appendToolResult: appendSpy.fn,
+      parallel: true,
+    });
+
+    const calls = [
+      router.parseToolCall({
+        id: "tc-1",
+        name: "Echo",
+        args: { text: "first", delay: 30 },
+      }),
+      router.parseToolCall({
+        id: "tc-2",
+        name: "Echo",
+        args: { text: "second", delay: 0 },
+      }),
+      router.parseToolCall({
+        id: "tc-3",
+        name: "Echo",
+        args: { text: "third", delay: 15 },
+      }),
+    ];
+
+    await router.processToolCalls(calls);
+
+    expect(appendSpy.calls).toHaveLength(3);
+    expect(at(appendSpy.calls, 0).toolCallId).toBe("tc-1");
+    expect(at(appendSpy.calls, 1).toolCallId).toBe("tc-2");
+    expect(at(appendSpy.calls, 2).toolCallId).toBe("tc-3");
+  });
+
   it("processes multiple tool calls sequentially", async () => {
     const order: string[] = [];
     const slowEcho = defineTool({

package/src/lib/tool-router/router.ts

@@ -211,10 +211,20 @@ export function createToolRouter<T extends ToolMap>(
    * handler requested a session-level rewind; when present, the result is
    * not appended to the thread and siblings should be cancelled.
    */
+  interface PendingAppend {
+    toolCallId: string;
+    toolName: string;
+    content: JsonValue;
+  }
+
   type ProcessedToolCall =
-    | { kind: "result"; value: ToolCallResultUnion<TResults> }
+    | {
+        kind: "result";
+        value: ToolCallResultUnion<TResults>;
+        pendingAppend?: PendingAppend;
+      }
     | { kind: "rewind"; signal: RewindSignal }
-    | { kind: "skipped" };
+    | { kind: "skipped"; pendingAppend?: PendingAppend };
 
   async function processToolCall(
     toolCall: ParsedToolCallUnion<T>,
@@ -222,7 +232,8 @@ export function createToolRouter<T extends ToolMap>(
     sandboxId?: string,
     onRewindRequested?: (signal: RewindSignal) => void,
     assistantMessageId?: string,
-    persistThreadState?: () => Promise<void>
+    persistThreadState?: () => Promise<void>,
+    deferAppend?: boolean
   ): Promise<ProcessedToolCall> {
     const startTime = Date.now();
     const tool = toolMap.get(toolCall.name);
@@ -230,15 +241,26 @@ export function createToolRouter<T extends ToolMap>(
     // --- Pre-hooks: may skip or modify args ---
     const preResult = await runPreHooks(toolCall, tool, turn);
     if (preResult.skip) {
+      const skipContent = JSON.stringify({
+        skipped: true,
+        reason: "Skipped by PreToolUse hook",
+      });
+      if (deferAppend) {
+        return {
+          kind: "skipped",
+          pendingAppend: {
+            toolCallId: toolCall.id,
+            toolName: toolCall.name,
+            content: skipContent,
+          },
+        };
+      }
       await appendToolResult(uuid4(), {
         threadId: options.threadId,
         threadKey: options.threadKey,
         toolCallId: toolCall.id,
         toolName: toolCall.name,
-        content: JSON.stringify({
-          skipped: true,
-          reason: "Skipped by PreToolUse hook",
-        }),
+        content: skipContent,
       });
       return { kind: "skipped" };
     }
@@ -314,19 +336,22 @@ export function createToolRouter<T extends ToolMap>(
     }
 
     // --- Append result to thread (unless handler already did) ---
-    if (!resultAppended) {
-      const config = {
-        threadId: options.threadId,
-        threadKey: options.threadKey,
-        toolCallId: toolCall.id,
-        toolName: toolCall.name,
-        content,
-      };
+    const needsAppend = !resultAppended;
+    if (needsAppend && !deferAppend) {
       await appendToolResult.executeWithOptions(
         {
           summary: `Append ${toolCall.name} result`,
         },
-        [uuid4(), config]
+        [
+          uuid4(),
+          {
+            threadId: options.threadId,
+            threadKey: options.threadKey,
+            toolCallId: toolCall.id,
+            toolName: toolCall.name,
+            content,
+          },
+        ]
       );
     }
 
@@ -356,7 +381,18 @@ export function createToolRouter<T extends ToolMap>(
       durationMs
     );
 
-    return { kind: "result", value: toolResult };
+    return {
+      kind: "result",
+      value: toolResult,
+      ...(needsAppend &&
+        deferAppend && {
+          pendingAppend: {
+            toolCallId: toolCall.id,
+            toolName: toolCall.name,
+            content,
+          },
+        }),
+    };
   }
 
   return {
@@ -409,7 +445,7 @@ export function createToolRouter<T extends ToolMap>(
     ): Promise<ProcessToolCallsResult<TResults>> {
       const attachRewind = (
         arr: ToolCallResultUnion<TResults>[],
-        rewind: RewindSignal | undefined,
+        rewind: RewindSignal | undefined
       ): ProcessToolCallsResult<TResults> => {
         if (rewind) {
           (arr as ProcessToolCallsResult<TResults>).rewind = rewind;
@@ -447,19 +483,55 @@ export function createToolRouter<T extends ToolMap>(
                 sandboxId,
                 onRewindRequested,
                 assistantMessageId,
-                persistThreadState
+                persistThreadState,
+                true
               )
             )
           )
         );
 
+        // Fail fast on non-cancellation rejections before appending
+        // anything, so the thread stays clean for retry/truncation.
+        for (const outcome of outcomes) {
+          if (
+            outcome.status === "rejected" &&
+            !isCancellation(outcome.reason)
+          ) {
+            throw outcome.reason;
+          }
+        }
+
+        // Append deferred results in original call order so positional
+        // correlation between function calls and responses is preserved.
+        if (!rewindSignal) {
+          for (const outcome of outcomes) {
+            if (
+              outcome.status === "fulfilled" &&
+              outcome.value.kind !== "rewind" &&
+              outcome.value.pendingAppend
+            ) {
+              const pa = outcome.value.pendingAppend;
+              await appendToolResult.executeWithOptions(
+                { summary: `Append ${pa.toolName} result` },
+                [
+                  uuid4(),
+                  {
+                    threadId: options.threadId,
+                    threadKey: options.threadKey,
+                    toolCallId: pa.toolCallId,
+                    toolName: pa.toolName,
+                    content: pa.content,
+                  },
+                ]
+              );
+            }
+          }
+        }
+
         const results: ToolCallResultUnion<TResults>[] = [];
         for (const outcome of outcomes) {
           if (outcome.status === "rejected") {
-            if (isCancellation(outcome.reason)) {
-              continue;
-            }
-            throw outcome.reason;
+            continue;
           }
           if (outcome.value.kind === "result") {
             results.push(outcome.value.value);
@@ -502,8 +574,12 @@ export function createToolRouter<T extends ToolMap>(
       }
 
       const processOne = async (
-        toolCall: ParsedToolCallUnion<T>
-      ): Promise<ToolCallResult<TName, TResult>> => {
+        toolCall: ParsedToolCallUnion<T>,
+        deferAppend?: boolean
+      ): Promise<{
+        result: ToolCallResult<TName, TResult>;
+        pendingAppend?: PendingAppend;
+      }> => {
         const routerContext: RouterContext = {
           threadId: options.threadId,
           ...(options.threadKey && { threadKey: options.threadKey }),
@@ -524,7 +600,8 @@ export function createToolRouter<T extends ToolMap>(
           routerContext as Parameters<typeof handler>[1]
         );
 
-        if (!response.resultAppended) {
+        const needsAppend = !response.resultAppended;
+        if (needsAppend && !deferAppend) {
           await appendToolResult.executeWithOptions(
             {
               summary: `Append ${toolCall.name} result`,
@@ -543,20 +620,51 @@ export function createToolRouter<T extends ToolMap>(
         }
 
         return {
-          toolCallId: toolCall.id,
-          name: toolCall.name as TName,
-          data: response.data,
-          ...(response.metadata && { metadata: response.metadata }),
+          result: {
+            toolCallId: toolCall.id,
+            name: toolCall.name as TName,
+            data: response.data,
+            ...(response.metadata && { metadata: response.metadata }),
+          },
+          ...(needsAppend &&
+            deferAppend && {
+              pendingAppend: {
+                toolCallId: toolCall.id,
+                toolName: toolCall.name,
+                content: response.toolResponse as JsonValue,
+              },
+            }),
         };
       };
 
       if (options.parallel) {
-        return Promise.all(matchingCalls.map(processOne));
+        const outcomes = await Promise.all(
+          matchingCalls.map((tc) => processOne(tc, true))
+        );
+        for (const { pendingAppend } of outcomes) {
+          if (pendingAppend) {
+            await appendToolResult.executeWithOptions(
+              { summary: `Append ${pendingAppend.toolName} result` },
+              [
+                uuid4(),
+                {
+                  threadId: options.threadId,
+                  threadKey: options.threadKey,
+                  toolCallId: pendingAppend.toolCallId,
+                  toolName: pendingAppend.toolName,
+                  content: pendingAppend.content,
+                },
+              ]
+            );
+          }
+        }
+        return outcomes.map((o) => o.result);
       }
 
       const results: ToolCallResult<TName, TResult>[] = [];
       for (const toolCall of matchingCalls) {
-        results.push(await processOne(toolCall));
+        const { result } = await processOne(toolCall);
+        results.push(result);
       }
       return results;
     },

package/src/lib/workflow.ts

@@ -1,4 +1,6 @@
 import type { ThreadInit, SandboxInit, SandboxShutdown } from "./lifecycle";
+import type { SandboxSnapshot } from "./sandbox/types";
+import type { TokenUsage } from "./types";
 
 /**
  * Session config fields derived from a main workflow input, ready to spread
@@ -13,6 +15,25 @@ export interface WorkflowSessionInput {
   sandbox?: SandboxInit;
   /** Sandbox shutdown policy (default: "destroy") */
   sandboxShutdown?: SandboxShutdown;
+  /**
+   * Called by the session right before `runSession` returns. Installed by
+   * `defineWorkflow` to capture sandbox / thread / usage outputs and forward
+   * them to the workflow's `onSessionExit` config hook. Spread into
+   * `createSession` via `...sessionInput`.
+   */
+  onSessionExit?: (result: {
+    sandboxId?: string;
+    snapshot?: SandboxSnapshot;
+    threadId: string;
+    usage: {
+      totalInputTokens: number;
+      totalOutputTokens: number;
+      totalCachedWriteTokens: number;
+      totalCachedReadTokens: number;
+      totalReasonTokens: number;
+      turns: number;
+    };
+  }) => void;
 }
 
 /** Raw workflow input fields that map into `WorkflowSessionInput`. */
@@ -34,6 +55,18 @@ export interface WorkflowConfig {
    * - `"keep"` — leave the sandbox running (no-op on exit).
    */
   sandboxShutdown?: SandboxShutdown;
+  /**
+   * Called right before the underlying session exits, with the sandbox /
+   * thread outputs and normalized token usage. Mirrors the capture logic in
+   * `defineSubagentWorkflow`; useful for emitting metrics or persisting
+   * sandbox / thread ids without threading them through the handler result.
+   */
+  onSessionExit?: (result: {
+    sandboxId?: string;
+    snapshot?: SandboxSnapshot;
+    threadId: string;
+    usage: TokenUsage;
+  }) => void;
 }
 
 /**
@@ -59,6 +92,22 @@ export function defineWorkflow<TInput, TResult>(
       sandboxShutdown: config.sandboxShutdown ?? "destroy",
       ...(workflowInput.thread && { thread: workflowInput.thread }),
       ...(workflowInput.sandbox && { sandbox: workflowInput.sandbox }),
+      ...(config.onSessionExit && {
+        onSessionExit: ({ sandboxId, snapshot, threadId, usage }): void => {
+          config.onSessionExit?.({
+            ...(sandboxId !== undefined && { sandboxId }),
+            ...(snapshot !== undefined && { snapshot }),
+            threadId,
+            usage: {
+              inputTokens: usage.totalInputTokens,
+              outputTokens: usage.totalOutputTokens,
+              cachedWriteTokens: usage.totalCachedWriteTokens,
+              cachedReadTokens: usage.totalCachedReadTokens,
+              reasonTokens: usage.totalReasonTokens,
+            },
+          });
+        },
+      }),
     };
     return fn(input, sessionInput);
   };