zeitlich 0.2.48 → 0.2.50

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,17 +211,29 @@ 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>,
     turn: number,
     sandboxId?: string,
     onRewindRequested?: (signal: RewindSignal) => void,
-    assistantMessageId?: string
+    assistantMessageId?: string,
+    persistThreadState?: () => Promise<void>,
+    deferAppend?: boolean
   ): Promise<ProcessedToolCall> {
     const startTime = Date.now();
     const tool = toolMap.get(toolCall.name);
@@ -229,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" };
     }
@@ -265,6 +288,7 @@ export function createToolRouter<T extends ToolMap>(
           toolName: toolCall.name,
           ...(sandboxId !== undefined && { sandboxId }),
           ...(assistantMessageId !== undefined && { assistantMessageId }),
+          ...(persistThreadState !== undefined && { persistThreadState }),
         };
         const response = await tool.handler(
           effectiveArgs as Parameters<typeof tool.handler>[0],
@@ -312,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,
+          },
+        ]
       );
     }
 
@@ -354,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 {
@@ -407,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;
@@ -422,6 +460,7 @@ export function createToolRouter<T extends ToolMap>(
       const turn = context?.turn ?? 0;
       const sandboxId = context?.sandboxId;
       const assistantMessageId = context?.assistantMessageId;
+      const persistThreadState = context?.persistThreadState;
 
       let rewindSignal: RewindSignal | undefined;
 
@@ -443,19 +482,56 @@ export function createToolRouter<T extends ToolMap>(
                 turn,
                 sandboxId,
                 onRewindRequested,
-                assistantMessageId
+                assistantMessageId,
+                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);
@@ -471,7 +547,8 @@ export function createToolRouter<T extends ToolMap>(
           turn,
           sandboxId,
           undefined,
-          assistantMessageId
+          assistantMessageId,
+          persistThreadState
         );
         if (outcome.kind === "rewind") {
           rewindSignal = outcome.signal;
@@ -497,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 }),
@@ -510,13 +591,17 @@ export function createToolRouter<T extends ToolMap>(
           ...(context?.assistantMessageId !== undefined && {
             assistantMessageId: context.assistantMessageId,
           }),
+          ...(context?.persistThreadState !== undefined && {
+            persistThreadState: context.persistThreadState,
+          }),
         };
         const response = await handler(
           toolCall.args as ToolArgs<T, TName>,
           routerContext as Parameters<typeof handler>[1]
         );
 
-        if (!response.resultAppended) {
+        const needsAppend = !response.resultAppended;
+        if (needsAppend && !deferAppend) {
           await appendToolResult.executeWithOptions(
             {
               summary: `Append ${toolCall.name} result`,
@@ -535,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/tool-router/types.ts

@@ -190,6 +190,20 @@ export interface RouterContext {
    * thread so the child's first model call sees a well-formed history.
    */
   assistantMessageId?: string;
+  /**
+   * Persist the parent session's current `PersistedThreadState` slice
+   * (tasks + custom state) to the durable thread store. Wired up by
+   * the session — absent for manually-driven routers (tests, custom
+   * orchestrators).
+   *
+   * Subagent handlers invoke this before spawning a child that will
+   * read the parent's thread (`newThreadSource: "from-parent"` or an
+   * explicit parent threadId): the parent's slice otherwise only
+   * lands in storage at session-exit time, so the child would load a
+   * stale (or empty) snapshot. Best-effort — failures are logged by
+   * the session but never thrown.
+   */
+  persistThreadState?: () => Promise<void>;
 }
 
 /**
@@ -314,6 +328,15 @@ export interface ProcessToolCallsContext {
    * out of a parent-forked thread).
    */
   assistantMessageId?: string;
+  /**
+   * Optional callback that flushes the session's in-memory
+   * `PersistedThreadState` slice to the durable thread store. The
+   * router forwards it into every handler's {@link RouterContext}
+   * verbatim. The session uses this to let mid-loop tool handlers
+   * (notably subagents that fork or continue the parent's thread)
+   * persist the parent's slice before the child reads it.
+   */
+  persistThreadState?: () => Promise<void>;
 }
 
 /**

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);
   };

package/src/{adapters/sandbox/inmemory/proxy.ts → test-utils/in-memory-sandbox-proxy.ts}

@@ -1,26 +1,15 @@
 /**
- * Workflow-safe proxy for in-memory sandbox operations.
+ * Test-only workflow-safe proxy for the in-memory sandbox fixture.
  *
- * Import this from `zeitlich/adapters/sandbox/inmemory/workflow`
- * in your Temporal workflow files.
- *
- * By default the scope is derived from `workflowInfo().workflowType`,
- * so activities are automatically namespaced per workflow.
- *
- * @example
- * ```typescript
- * import { proxyInMemorySandboxOps } from 'zeitlich/adapters/sandbox/inmemory/workflow';
- *
- * // Auto-scoped to the current workflow name
- * const sandbox = proxyInMemorySandboxOps();
- * ```
+ * Not part of the shipped public surface — used by the capability-type
+ * fixtures to exercise the full-capability proxy shape.
  */
 import { proxyActivities, workflowInfo } from "@temporalio/workflow";
 import type {
   SandboxCreateOptions,
   SandboxOps,
-} from "../../../lib/sandbox/types";
-import type { InMemoryCaps } from "./index";
+} from "../lib/sandbox/types";
+import type { InMemoryCaps } from "./in-memory-sandbox";
 
 const ADAPTER_PREFIX = "inMemory";
 

package/src/{adapters/sandbox/inmemory/index.ts → test-utils/in-memory-sandbox.ts}

@@ -1,3 +1,11 @@
+/**
+ * Test-only in-memory {@link Sandbox} provider backed by `just-bash`.
+ *
+ * This is **not** part of the shipped public surface — it lives under
+ * `src/test-utils` purely so the sandbox manager, tool handlers, and
+ * capability-type fixtures have a lightweight, fully-featured backend to
+ * exercise against. `just-bash` is a dev dependency for the same reason.
+ */
 import {
   Bash,
   InMemoryFs,
@@ -18,9 +26,9 @@ import type {
   ExecResult,
   DirentEntry,
   FileStat,
-} from "../../../lib/sandbox/types";
-import { SandboxNotFoundError } from "../../../lib/sandbox/types";
-import { getShortId } from "../../../lib/thread/id";
+} from "../lib/sandbox/types";
+import { SandboxNotFoundError } from "../lib/sandbox/types";
+import { getShortId } from "../lib/thread/id";
 
 // ============================================================================
 // Adapter: IFileSystem → SandboxFileSystem

package/src/tools/bash/bash.test.ts

@@ -2,7 +2,7 @@ import { describe, expect, it, beforeEach } from "vitest";
 import { bashHandler } from "./handler";
 import { withSandbox } from "../../lib/tool-router/with-sandbox";
 import { SandboxManager } from "../../lib/sandbox/manager";
-import { InMemorySandboxProvider } from "../../adapters/sandbox/inmemory/index";
+import { InMemorySandboxProvider } from "../../test-utils/in-memory-sandbox";
 import type { RouterContext } from "../../lib/tool-router/types";
 import type { Sandbox, SandboxCreateOptions } from "../../lib/sandbox";
 

package/src/tools/edit/handler.test.ts

@@ -1,5 +1,5 @@
 import { beforeEach, describe, expect, it } from "vitest";
-import { InMemorySandboxProvider } from "../../adapters/sandbox/inmemory/index";
+import { InMemorySandboxProvider } from "../../test-utils/in-memory-sandbox";
 import type { Sandbox, SandboxCreateOptions } from "../../lib/sandbox";
 import { SandboxManager } from "../../lib/sandbox/manager";
 import type { RouterContext } from "../../lib/tool-router/types";

package/tsup.config.ts

@@ -15,9 +15,6 @@ export default defineConfig({
     "adapters/thread/anthropic/index": "src/adapters/thread/anthropic/index.ts",
     "adapters/thread/anthropic/workflow":
       "src/adapters/thread/anthropic/proxy.ts",
-    "adapters/sandbox/inmemory/index": "src/adapters/sandbox/inmemory/index.ts",
-    "adapters/sandbox/inmemory/workflow":
-      "src/adapters/sandbox/inmemory/proxy.ts",
     "adapters/sandbox/daytona/index": "src/adapters/sandbox/daytona/index.ts",
     "adapters/sandbox/daytona/workflow":
       "src/adapters/sandbox/daytona/proxy.ts",
@@ -38,7 +35,8 @@ export default defineConfig({
     /^@anthropic-ai\//,
     /^@daytonaio\//,
     /^@e2b\//,
-    "ioredis",
+    "redis",
+    /^@redis\//,
     "@mongodb-js/zstd",
     "node-liblzma",
   ],