zeitlich 0.2.37 → 0.2.38

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

@@ -23,8 +23,20 @@ vi.mock("@temporalio/workflow", () => {
     }
   }
   const noop = () => {};
+  class MockCancellationScope {
+    cancellable: boolean;
+    constructor(opts?: { cancellable?: boolean }) {
+      this.cancellable = opts?.cancellable ?? true;
+    }
+    async run<T>(fn: () => Promise<T>): Promise<T> {
+      return fn();
+    }
+    cancel(): void {}
+  }
   return {
     ApplicationFailure: MockApplicationFailure,
+    CancellationScope: MockCancellationScope,
+    isCancellation: (_err: unknown) => false,
     uuid4: () => "00000000-0000-0000-0000-000000000000",
     log: { trace: noop, debug: noop, info: noop, warn: noop, error: noop },
   };
@@ -648,6 +660,86 @@ describe("createToolRouter edge cases", () => {
       recovered: true,
     });
   });
+
+  // --- Rewind signal -------------------------------------------------------
+
+  it("attaches a rewind signal and skips result append when handler returns rewind:true", async () => {
+    const rewindTool = defineTool({
+      name: "Rewind" as const,
+      description: "rewinds",
+      schema: z.object({}),
+      handler: async () => ({
+        toolResponse: "ignored",
+        data: null,
+        rewind: true,
+      }),
+    });
+
+    const router = createToolRouter({
+      tools: { Rewind: rewindTool } as const,
+      threadId: "t-1",
+      appendToolResult: appendSpy.fn,
+    });
+
+    const parsed = router.parseToolCall({
+      id: "tc-1",
+      name: "Rewind",
+      args: {},
+    });
+
+    const results = await router.processToolCalls([parsed]);
+
+    expect(results).toHaveLength(0);
+    expect(results.rewind).toEqual({
+      toolCallId: "tc-1",
+      toolName: "Rewind",
+    });
+    expect(appendSpy.calls).toHaveLength(0);
+  });
+
+  it("short-circuits further sequential tool calls when one requests rewind", async () => {
+    let laterCalled = false;
+    const laterTool = defineTool({
+      name: "Later" as const,
+      description: "runs after rewind",
+      schema: z.object({}),
+      handler: async () => {
+        laterCalled = true;
+        return { toolResponse: "ok", data: null };
+      },
+    });
+    const rewindTool = defineTool({
+      name: "Rewind" as const,
+      description: "rewinds",
+      schema: z.object({}),
+      handler: async () => ({
+        toolResponse: "ignored",
+        data: null,
+        rewind: true,
+      }),
+    });
+
+    const router = createToolRouter({
+      tools: { Rewind: rewindTool, Later: laterTool } as const,
+      threadId: "t-1",
+      appendToolResult: appendSpy.fn,
+      parallel: false,
+    });
+
+    const calls = [
+      router.parseToolCall({ id: "tc-1", name: "Rewind", args: {} }),
+      router.parseToolCall({ id: "tc-2", name: "Later", args: {} }),
+    ];
+
+    const results = await router.processToolCalls(calls);
+
+    expect(results).toHaveLength(0);
+    expect(results.rewind).toEqual({
+      toolCallId: "tc-1",
+      toolName: "Rewind",
+    });
+    expect(laterCalled).toBe(false);
+  });
 });
 
 describe("hasNoOtherToolCalls", () => {

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

@@ -23,8 +23,20 @@ vi.mock("@temporalio/workflow", () => {
     }
   }
   const noop = () => {};
+  class MockCancellationScope {
+    cancellable: boolean;
+    constructor(opts?: { cancellable?: boolean }) {
+      this.cancellable = opts?.cancellable ?? true;
+    }
+    async run<T>(fn: () => Promise<T>): Promise<T> {
+      return fn();
+    }
+    cancel(): void {}
+  }
   return {
     ApplicationFailure: MockApplicationFailure,
+    CancellationScope: MockCancellationScope,
+    isCancellation: (_err: unknown) => false,
     uuid4: () => "00000000-0000-0000-0000-000000000000",
     log: { trace: noop, debug: noop, info: noop, warn: noop, error: noop },
   };

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.
@@ -199,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> {
+    sandboxId?: string,
+    onRewindRequested?: (signal: RewindSignal) => void
+  ): Promise<ProcessedToolCall> {
     const startTime = Date.now();
     const tool = toolMap.get(toolCall.name);
 
@@ -220,7 +238,7 @@ export function createToolRouter<T extends ToolMap>(
           reason: "Skipped by PreToolUse hook",
         }),
       });
-      return null;
+      return { kind: "skipped" };
     }
     const effectiveArgs = preResult.args;
 
@@ -235,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) {
@@ -253,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,
@@ -276,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 = {
@@ -319,7 +352,7 @@ export function createToolRouter<T extends ToolMap>(
       durationMs
     );
 
-    return toolResult;
+    return { kind: "result", value: toolResult };
   }
 
   return {
@@ -369,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

@@ -139,6 +139,17 @@ export interface ToolHandlerResponse<
    * 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) */
@@ -278,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
 // ============================================================================
@@ -469,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/workflow.ts

@@ -94,6 +94,8 @@ export type {
   // Other
   AppendToolResultFn,
   ProcessToolCallsContext,
+  ProcessToolCallsResult,
+  RewindSignal,
 } from "./lib/tool-router";
 
 // Session & message lifecycle hooks