zeitlich 0.2.36 → 0.2.38

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

@@ -15,12 +15,19 @@ import type {
   ToolArgs,
   ToolResult,
   ProcessToolCallsContext,
+  ProcessToolCallsResult,
+  RewindSignal,
   ToolWithHandler,
 } from "./types";
 
 import type { JsonValue } from "../state/types";
 import type { z } from "zod";
-import { uuid4, log } from "@temporalio/workflow";
+import {
+  uuid4,
+  log,
+  CancellationScope,
+  isCancellation,
+} from "@temporalio/workflow";
 
 /**
  * Creates a tool router for declarative tool call processing.
@@ -162,7 +169,8 @@ export function createToolRouter<T extends ToolMap>(
 
     return {
       content: JSON.stringify({
-        error: "The tool encountered an error. Please try again or use a different approach.",
+        error:
+          "The tool encountered an error. Please try again or use a different approach.",
       }),
       result: { error: errorStr, suppressed: true },
     };
@@ -198,11 +206,22 @@ export function createToolRouter<T extends ToolMap>(
     }
   }
 
+  /**
+   * Internal per-tool-call outcome. `rewind` signals the caller that the
+   * handler requested a session-level rewind; when present, the result is
+   * not appended to the thread and siblings should be cancelled.
+   */
+  type ProcessedToolCall =
+    | { kind: "result"; value: ToolCallResultUnion<TResults> }
+    | { kind: "rewind"; signal: RewindSignal }
+    | { kind: "skipped" };
+
   async function processToolCall(
     toolCall: ParsedToolCallUnion<T>,
     turn: number,
     sandboxId?: string,
-  ): Promise<ToolCallResultUnion<TResults> | null> {
+    onRewindRequested?: (signal: RewindSignal) => void
+  ): Promise<ProcessedToolCall> {
     const startTime = Date.now();
     const tool = toolMap.get(toolCall.name);
 
@@ -219,7 +238,7 @@ export function createToolRouter<T extends ToolMap>(
           reason: "Skipped by PreToolUse hook",
         }),
       });
-      return null;
+      return { kind: "skipped" };
     }
     const effectiveArgs = preResult.args;
 
@@ -234,6 +253,7 @@ export function createToolRouter<T extends ToolMap>(
     let content!: JsonValue;
     let resultAppended = false;
     let metadata: Record<string, unknown> | undefined;
+    let rewindRequested = false;
 
     try {
       if (tool) {
@@ -252,11 +272,15 @@ export function createToolRouter<T extends ToolMap>(
         content = response.toolResponse as JsonValue;
         resultAppended = response.resultAppended === true;
         metadata = response.metadata;
+        rewindRequested = response.rewind === true;
       } else {
         result = { error: `Unknown tool: ${toolCall.name}` };
         content = JSON.stringify(result, null, 2);
       }
     } catch (error) {
+      if (isCancellation(error)) {
+        throw error;
+      }
       log.warn("tool call failed", {
         toolName: toolCall.name,
         toolCallId: toolCall.id,
@@ -275,6 +299,16 @@ export function createToolRouter<T extends ToolMap>(
       content = recovery.content;
     }
 
+    if (rewindRequested) {
+      const signal: RewindSignal = {
+        toolCallId: toolCall.id,
+        toolName: toolCall.name,
+      };
+      log.info("tool requested rewind", { ...signal });
+      onRewindRequested?.(signal);
+      return { kind: "rewind", signal };
+    }
+
     // --- Append result to thread (unless handler already did) ---
     if (!resultAppended) {
       const config = {
@@ -318,7 +352,7 @@ export function createToolRouter<T extends ToolMap>(
       durationMs
     );
 
-    return toolResult;
+    return { kind: "result", value: toolResult };
   }
 
   return {
@@ -368,31 +402,71 @@ export function createToolRouter<T extends ToolMap>(
     async processToolCalls(
       toolCalls: ParsedToolCallUnion<T>[],
       context?: ProcessToolCallsContext
-    ): Promise<ToolCallResultUnion<TResults>[]> {
+    ): Promise<ProcessToolCallsResult<TResults>> {
+      const attachRewind = (
+        arr: ToolCallResultUnion<TResults>[],
+        rewind: RewindSignal | undefined,
+      ): ProcessToolCallsResult<TResults> => {
+        if (rewind) {
+          (arr as ProcessToolCallsResult<TResults>).rewind = rewind;
+        }
+        return arr as ProcessToolCallsResult<TResults>;
+      };
+
       if (toolCalls.length === 0) {
-        return [];
+        return attachRewind([], undefined);
       }
 
       const turn = context?.turn ?? 0;
       const sandboxId = context?.sandboxId;
 
+      let rewindSignal: RewindSignal | undefined;
+
       if (options.parallel) {
-        const results = await Promise.all(
-          toolCalls.map((tc) => processToolCall(tc, turn, sandboxId))
+        const scope = new CancellationScope({ cancellable: true });
+        const onRewindRequested = (signal: RewindSignal): void => {
+          if (!rewindSignal) {
+            rewindSignal = signal;
+            // Cancel all other in-flight tool calls in this batch.
+            scope.cancel();
+          }
+        };
+
+        const outcomes = await scope.run(async () =>
+          Promise.allSettled(
+            toolCalls.map((tc) =>
+              processToolCall(tc, turn, sandboxId, onRewindRequested)
+            )
+          )
         );
-        return results.filter(
-          (r): r is NonNullable<typeof r> => r !== null
-        ) as ToolCallResultUnion<TResults>[];
+
+        const results: ToolCallResultUnion<TResults>[] = [];
+        for (const outcome of outcomes) {
+          if (outcome.status === "rejected") {
+            if (isCancellation(outcome.reason)) {
+              continue;
+            }
+            throw outcome.reason;
+          }
+          if (outcome.value.kind === "result") {
+            results.push(outcome.value.value);
+          }
+        }
+        return attachRewind(results, rewindSignal);
       }
 
       const results: ToolCallResultUnion<TResults>[] = [];
       for (const toolCall of toolCalls) {
-        const result = await processToolCall(toolCall, turn, sandboxId);
-        if (result !== null) {
-          results.push(result);
+        const outcome = await processToolCall(toolCall, turn, sandboxId);
+        if (outcome.kind === "rewind") {
+          rewindSignal = outcome.signal;
+          break;
+        }
+        if (outcome.kind === "result") {
+          results.push(outcome.value);
         }
       }
-      return results;
+      return attachRewind(results, rewindSignal);
     },
 
     async processToolCallsByName<TName extends ToolNames<T>, TResult>(

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

@@ -124,7 +124,10 @@ export type AppendToolResultFn = ActivityFunctionWithOptions<
  * Tools that don't return additional data should use `data: null` (TResult defaults to null).
  * Tools that may fail to produce data should type TResult as `SomeType | null`.
  */
-export interface ToolHandlerResponse<TResult = null, TToolResponse = JsonValue> {
+export interface ToolHandlerResponse<
+  TResult = null,
+  TToolResponse = JsonValue,
+> {
   /** Content sent back to the LLM as the tool call response */
   toolResponse: TToolResponse;
   /** Data returned to the workflow and hooks for further processing */
@@ -136,6 +139,17 @@ export interface ToolHandlerResponse<TResult = null, TToolResponse = JsonValue>
    * payloads through Temporal's activity payload limit.
    */
   resultAppended?: boolean;
+  /**
+   * When true, the session will rewind: any in-flight parallel tool calls
+   * are cancelled, previously appended tool results and the triggering
+   * assistant message are removed from the thread, and the LLM call that
+   * produced the tool calls is retried.
+   *
+   * The `toolResponse` for a rewinding tool call is ignored (never
+   * appended) since the session truncates the thread back to the
+   * pre-invocation state.
+   */
+  rewind?: boolean;
   /** Token usage from the tool execution (e.g. child agent invocations) */
   usage?: TokenUsage;
   /** Thread ID used by the handler (surfaced to the LLM for subagent thread continuation) */
@@ -172,7 +186,9 @@ export type ToolHandler<
 > = (
   args: TArgs,
   context: TContext
-) => ToolHandlerResponse<TResult, TToolResponse> | Promise<ToolHandlerResponse<TResult, TToolResponse>>;
+) =>
+  | ToolHandlerResponse<TResult, TToolResponse>
+  | Promise<ToolHandlerResponse<TResult, TToolResponse>>;
 
 /**
  * Activity-compatible tool handler that always returns a Promise.
@@ -195,7 +211,10 @@ export type ActivityToolHandler<
   TResult,
   TContext extends RouterContext = RouterContext,
   TToolResponse = JsonValue,
-> = (args: TArgs, context: TContext) => Promise<ToolHandlerResponse<TResult, TToolResponse>>;
+> = (
+  args: TArgs,
+  context: TContext
+) => Promise<ToolHandlerResponse<TResult, TToolResponse>>;
 
 /**
  * Extract the args type for a specific tool name from a tool map.
@@ -270,6 +289,28 @@ export interface ProcessToolCallsContext {
   sandboxId?: string;
 }
 
+/**
+ * Signal that a tool handler requested a rewind. Attached to the
+ * {@link ProcessToolCallsResult} so the session can roll the thread
+ * back to the pre-invocation snapshot and retry the LLM call.
+ */
+export interface RewindSignal {
+  toolCallId: string;
+  toolName: string;
+}
+
+/**
+ * Result returned by {@link ToolRouter.processToolCalls}.
+ *
+ * The object is a standard array of tool call results for successful
+ * tool calls (cancelled or rewinding siblings are omitted), extended
+ * with a `rewind` property when any tool in the batch requested a
+ * rewind. Using an array-with-property lets existing code that treats
+ * the return value as `ToolCallResultUnion[]` continue to work.
+ */
+export type ProcessToolCallsResult<TResults extends Record<string, unknown>> =
+  ToolCallResultUnion<TResults>[] & { rewind?: RewindSignal };
+
 // ============================================================================
 // Hook Types
 // ============================================================================
@@ -461,7 +502,7 @@ export interface ToolRouter<T extends ToolMap> {
   processToolCalls(
     toolCalls: ParsedToolCallUnion<T>[],
     context?: ProcessToolCallsContext
-  ): Promise<ToolCallResultUnion<InferToolResults<T>>[]>;
+  ): Promise<ProcessToolCallsResult<InferToolResults<T>>>;
 
   /**
    * Process tool calls matching a specific name with a custom handler.

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

@@ -1,6 +1,10 @@
 import type { Sandbox } from "../sandbox/types";
 import type { JsonValue } from "../state/types";
-import type { ActivityToolHandler, RouterContext, ToolHandlerResponse } from "./types";
+import type {
+  ActivityToolHandler,
+  RouterContext,
+  ToolHandlerResponse,
+} from "./types";
 
 /**
  * Extended router context with a resolved {@link Sandbox} instance.
@@ -51,13 +55,23 @@ export interface SandboxContext extends RouterContext {
  * });
  * ```
  */
-export function withSandbox<TArgs, TResult, TSandbox extends Sandbox = Sandbox, TToolResponse = JsonValue>(
+export function withSandbox<
+  TArgs,
+  TResult,
+  TSandbox extends Sandbox = Sandbox,
+  TToolResponse = JsonValue,
+>(
   manager: { getSandbox(id: string): Promise<TSandbox> },
   handler: (
     args: TArgs,
-    context: RouterContext & { sandbox: TSandbox; sandboxId: string },
-  ) => Promise<ToolHandlerResponse<TResult, TToolResponse>>,
-): ActivityToolHandler<TArgs, TResult | null, RouterContext, TToolResponse | string> {
+    context: RouterContext & { sandbox: TSandbox; sandboxId: string }
+  ) => Promise<ToolHandlerResponse<TResult, TToolResponse>>
+): ActivityToolHandler<
+  TArgs,
+  TResult | null,
+  RouterContext,
+  TToolResponse | string
+> {
   return async (args, context) => {
     if (!context.sandboxId) {
       return {

package/src/lib/virtual-fs/filesystem.ts

@@ -67,7 +67,7 @@ export class VirtualFileSystem<
     private resolver: FileResolver<TCtx, TMeta>,
     private ctx: TCtx,
     workspaceBase = "/",
-    inlineFiles?: Record<string, string>,
+    inlineFiles?: Record<string, string>
   ) {
     this.workspaceBase = normalisePath(workspaceBase);
     this.entries = new Map(

package/src/lib/virtual-fs/index.ts

@@ -1,7 +1,11 @@
 export { VirtualFileSystem } from "./filesystem";
 export { withVirtualFs } from "./with-virtual-fs";
 export { createVirtualFsActivities } from "./manager";
-export { hasFileWithMimeType, filesWithMimeType, hasDirectory } from "./queries";
+export {
+  hasFileWithMimeType,
+  filesWithMimeType,
+  hasDirectory,
+} from "./queries";
 export { applyVirtualTreeMutations } from "./mutations";
 export { formatVirtualFileTree } from "./tree";
 export type { FileTreeAccessor } from "./queries";

package/src/lib/virtual-fs/mutations.ts

@@ -13,7 +13,7 @@ export function applyVirtualTreeMutations<TMeta = FileEntryMetadata>(
     get(key: "fileTree"): VirtualFileTree<TMeta>;
     set(key: "fileTree", value: VirtualFileTree<TMeta>): void;
   },
-  mutations: TreeMutation<TMeta>[],
+  mutations: TreeMutation<TMeta>[]
 ): VirtualFileTree<TMeta> {
   let tree = [...stateManager.get("fileTree")];
 
@@ -26,9 +26,7 @@ export function applyVirtualTreeMutations<TMeta = FileEntryMetadata>(
         tree = tree.filter((e) => e.path !== m.path);
         break;
       case "update":
-        tree = tree.map((e) =>
-          e.path === m.path ? { ...e, ...m.entry } : e
-        );
+        tree = tree.map((e) => (e.path === m.path ? { ...e, ...m.entry } : e));
         break;
     }
   }

package/src/lib/virtual-fs/queries.ts

@@ -25,10 +25,12 @@ export interface FileTreeAccessor<TMeta> {
  */
 export function hasFileWithMimeType<TMeta>(
   stateManager: FileTreeAccessor<TMeta>,
-  pattern: string | string[],
+  pattern: string | string[]
 ): boolean {
   const tree = stateManager.get("fileTree");
-  const matchers = (Array.isArray(pattern) ? pattern : [pattern]).map(buildMatcher);
+  const matchers = (Array.isArray(pattern) ? pattern : [pattern]).map(
+    buildMatcher
+  );
   return tree.some((entry) => {
     const meta = entry.metadata as Record<string, unknown> | undefined;
     const mime = meta?.mimeType;
@@ -41,7 +43,7 @@ export function hasFileWithMimeType<TMeta>(
  */
 export function filesWithMimeType<TMeta>(
   stateManager: FileTreeAccessor<TMeta>,
-  pattern: string,
+  pattern: string
 ): FileEntry<TMeta>[] {
   const tree = stateManager.get("fileTree");
   const match = buildMatcher(pattern);
@@ -66,7 +68,7 @@ export function filesWithMimeType<TMeta>(
  */
 export function hasDirectory<TMeta>(
   stateManager: FileTreeAccessor<TMeta>,
-  pattern: string,
+  pattern: string
 ): boolean {
   const tree = stateManager.get("fileTree");
   const match = buildGlobMatcher(pattern);
@@ -91,7 +93,9 @@ function buildMatcher(pattern: string): (value: string) => boolean {
 function buildGlobMatcher(pattern: string): (value: string) => boolean {
   if (!pattern.includes("*")) return (v) => v === pattern;
   const re = new RegExp(
-    "^" + pattern.replace(/[.+^${}()|[\]\\]/g, "\\$&").replace(/\*/g, ".*") + "$",
+    "^" +
+      pattern.replace(/[.+^${}()|[\]\\]/g, "\\$&").replace(/\*/g, ".*") +
+      "$"
   );
   return (v) => re.test(v);
 }

package/src/lib/virtual-fs/types.ts

@@ -104,7 +104,10 @@ export type PrefixedVirtualFsOps<
   [K in keyof VirtualFsOps<
     TCtx,
     TMeta
-  > as `virtualFs${Capitalize<TPrefix>}${Capitalize<K & string>}`]: VirtualFsOps<TCtx, TMeta>[K];
+  > as `virtualFs${Capitalize<TPrefix>}${Capitalize<K & string>}`]: VirtualFsOps<
+    TCtx,
+    TMeta
+  >[K];
 };
 
 // ============================================================================

package/src/lib/virtual-fs/virtual-fs.test.ts

@@ -53,7 +53,7 @@ function createMockResolver(): {
         id,
         typeof content === "string"
           ? content
-          : new TextDecoder().decode(content),
+          : new TextDecoder().decode(content)
       );
     },
 
@@ -63,7 +63,7 @@ function createMockResolver(): {
         id,
         typeof content === "string"
           ? content
-          : new TextDecoder().decode(content),
+          : new TextDecoder().decode(content)
       );
       const size =
         typeof content === "string"
@@ -286,9 +286,7 @@ describe("VirtualFileSystem", () => {
   });
 
   it("rm force does not throw for missing path", async () => {
-    await expect(
-      fs.rm("/nonexistent", { force: true }),
-    ).resolves.not.toThrow();
+    await expect(fs.rm("/nonexistent", { force: true })).resolves.not.toThrow();
   });
 
   // --- cp / mv ---
@@ -320,7 +318,7 @@ describe("VirtualFileSystem", () => {
 
   it("readlink throws SandboxNotSupportedError", async () => {
     await expect(fs.readlink("/src/index.ts")).rejects.toThrow(
-      SandboxNotSupportedError,
+      SandboxNotSupportedError
     );
   });
 });
@@ -349,7 +347,7 @@ describe("VirtualFileSystem — inlineFiles", () => {
       resolver,
       ctx,
       "/",
-      inlineFiles,
+      inlineFiles
     );
   }
 
@@ -418,7 +416,9 @@ describe("applyVirtualTreeMutations", () => {
     let fileTree = tree;
     return {
       get: (_key: "fileTree"): FileEntry[] => fileTree,
-      set: (_key: "fileTree", value: FileEntry[]): void => { fileTree = value; },
+      set: (_key: "fileTree", value: FileEntry[]): void => {
+        fileTree = value;
+      },
       current: (): FileEntry[] => fileTree,
     };
   }
@@ -490,9 +490,7 @@ describe("applyVirtualTreeMutations", () => {
   it("does not mutate the original array passed to the state manager", () => {
     const original = [...sampleTree];
     const sm = mockStateManager(sampleTree);
-    applyVirtualTreeMutations(sm, [
-      { type: "remove", path: "/src/index.ts" },
-    ]);
+    applyVirtualTreeMutations(sm, [{ type: "remove", path: "/src/index.ts" }]);
     expect(sampleTree).toEqual(original);
   });
 });

package/src/lib/workflow.test.ts

@@ -50,10 +50,13 @@ describe("defineWorkflow", () => {
       return { ok: true };
     });
 
-    await workflow({}, {
-      thread: { mode: "fork", threadId: "prev-1" },
-      sandbox: { mode: "continue", sandboxId: "sb-1" },
-    });
+    await workflow(
+      {},
+      {
+        thread: { mode: "fork", threadId: "prev-1" },
+        sandbox: { mode: "continue", sandboxId: "sb-1" },
+      }
+    );
 
     expect(capturedSession).toEqual({
       agentName: "test-workflow",

package/src/lib/workflow.ts

@@ -1,8 +1,4 @@
-import type {
-  ThreadInit,
-  SandboxInit,
-  SandboxShutdown,
-} from "./lifecycle";
+import type { ThreadInit, SandboxInit, SandboxShutdown } from "./lifecycle";
 
 /**
  * Session config fields derived from a main workflow input, ready to spread

package/src/tools/ask-user-question/tool.ts

@@ -42,6 +42,4 @@ Usage notes:
   strict: true,
 } satisfies ToolDefinition;
 
-export type AskUserQuestionArgs = z.infer<
-  typeof askUserQuestionTool.schema
->;
+export type AskUserQuestionArgs = z.infer<typeof askUserQuestionTool.schema>;

package/src/tools/glob/handler.ts

@@ -19,10 +19,7 @@ function matchGlob(pattern: string, path: string): boolean {
   return new RegExp(`^${regex}$`).test(path);
 }
 
-async function walk(
-  fs: Sandbox["fs"],
-  dir: string,
-): Promise<string[]> {
+async function walk(fs: Sandbox["fs"], dir: string): Promise<string[]> {
   const results: string[] = [];
   const entries = await fs.readdirWithFileTypes(dir);
   for (const entry of entries) {

package/src/tools/task-get/handler.ts

@@ -1,7 +1,4 @@
-import type {
-  AgentStateManager,
-  JsonSerializable,
-} from "../../lib/state";
+import type { AgentStateManager, JsonSerializable } from "../../lib/state";
 import type { ToolHandler } from "../../lib/tool-router";
 import type { WorkflowTask } from "../../lib/types";
 import type { TaskGetArgs } from "./tool";
@@ -20,7 +17,9 @@ export function createTaskGetHandler<TCustom extends JsonSerializable<TCustom>>(
 
     if (!task) {
       return {
-        toolResponse: JSON.stringify({ error: `Task not found: ${args.taskId}` }),
+        toolResponse: JSON.stringify({
+          error: `Task not found: ${args.taskId}`,
+        }),
         data: null,
       };
     }

package/src/tools/task-list/handler.ts

@@ -1,7 +1,4 @@
-import type {
-  AgentStateManager,
-  JsonSerializable,
-} from "../../lib/state";
+import type { AgentStateManager, JsonSerializable } from "../../lib/state";
 import type { ToolHandler } from "../../lib/tool-router";
 import type { WorkflowTask } from "../../lib/types";
 import type { TaskListArgs } from "./tool";

package/src/tools/task-update/handler.ts

@@ -1,7 +1,4 @@
-import type {
-  AgentStateManager,
-  JsonSerializable,
-} from "../../lib/state";
+import type { AgentStateManager, JsonSerializable } from "../../lib/state";
 import type { ToolHandler } from "../../lib/tool-router";
 import type { WorkflowTask } from "../../lib/types";
 import type { TaskUpdateArgs } from "./tool";
@@ -22,7 +19,9 @@ export function createTaskUpdateHandler<
 
     if (!task) {
       return {
-        toolResponse: JSON.stringify({ error: `Task not found: ${args.taskId}` }),
+        toolResponse: JSON.stringify({
+          error: `Task not found: ${args.taskId}`,
+        }),
         data: null,
       };
     }

package/src/workflow.ts

@@ -15,12 +15,28 @@
 
 // Session
 export { createSession } from "./lib/session";
-export type { ZeitlichSession, SessionResult, ThreadOps, PrefixedThreadOps, ScopedPrefix, SessionConfig } from "./lib/session";
+export type {
+  ZeitlichSession,
+  SessionResult,
+  ThreadOps,
+  PrefixedThreadOps,
+  ScopedPrefix,
+  SessionConfig,
+} from "./lib/session";
 export { defineWorkflow } from "./lib/workflow";
-export type { WorkflowConfig, WorkflowInput, WorkflowSessionInput } from "./lib/workflow";
+export type {
+  WorkflowConfig,
+  WorkflowInput,
+  WorkflowSessionInput,
+} from "./lib/workflow";
 
 // Lifecycle types
-export type { ThreadInit, SandboxInit, SandboxShutdown, SubagentSandboxShutdown } from "./lib/lifecycle";
+export type {
+  ThreadInit,
+  SandboxInit,
+  SandboxShutdown,
+  SubagentSandboxShutdown,
+} from "./lib/lifecycle";
 
 // Thread utilities
 export { getShortId } from "./lib/thread/id";
@@ -78,6 +94,8 @@ export type {
   // Other
   AppendToolResultFn,
   ProcessToolCallsContext,
+  ProcessToolCallsResult,
+  RewindSignal,
 } from "./lib/tool-router";
 
 // Session & message lifecycle hooks
@@ -94,10 +112,7 @@ export type {
 } from "./lib/hooks";
 
 // Observability
-export {
-  createObservabilityHooks,
-  composeHooks,
-} from "./lib/observability";
+export { createObservabilityHooks, composeHooks } from "./lib/observability";
 export type {
   ObservabilityHooks,
   ZeitlichObservabilitySinks,