zeitlich 0.2.37 → 0.2.39

package/src/lib/thread/types.ts

@@ -1,5 +1,5 @@
 import type Redis from "ioredis";
-import type { JsonValue } from "../state/types";
+import type { JsonValue, PersistedThreadState } from "../state/types";
 export interface ThreadManagerConfig<T> {
   redis: Redis;
   threadId: string;
@@ -34,8 +34,50 @@ export interface BaseThreadManager<T> {
    * forks — each call creates an independent copy.
    */
   fork(newThreadId: string): Promise<BaseThreadManager<T>>;
+  /**
+   * Atomically replace the entire contents of the thread with `messages`.
+   * The existing list is cleared, the new messages are appended in order,
+   * and dedup markers from prior appends are cleared so future idempotent
+   * appends with ids that were removed aren't silently skipped.
+   *
+   * Requires the thread manager to be configured with `idOf`.
+   */
+  replaceAll(messages: T[]): Promise<void>;
   /** Delete the thread */
   delete(): Promise<void>;
+  /** Get the number of stored messages currently in the thread */
+  length(): Promise<number>;
+  /**
+   * Truncate the thread starting at the message with id `messageId`.
+   * That message and every message after it are removed. If `messageId`
+   * is not present in the thread this is a no-op — useful as the
+   * "truncate on entry" step of the `runAgent` activity, which becomes a
+   * no-op on the first attempt and a cleanup on Temporal workflow reset
+   * or in-workflow rewind retries.
+   *
+   * Dedup markers for removed single-message appends are also cleared so
+   * that appending the same id again (e.g. the same assistant message id
+   * on a rewind retry) is not silently skipped.
+   *
+   * Requires the thread manager to be configured with `idOf`.
+   */
+  truncateFromId(messageId: string): Promise<void>;
+  /**
+   * Load the persisted state slice associated with this thread, or
+   * `null` if none has been saved yet. Safe to call on any thread —
+   * treats a missing slice as a non-error.
+   */
+  loadState(): Promise<PersistedThreadState | null>;
+  /**
+   * Overwrite the persisted state slice for this thread. The thread
+   * itself must already exist (same TTL as the message list).
+   *
+   * Note: {@link BaseThreadManager.fork} already copies the slice to
+   * the new thread, so there's no separate `forkState` method.
+   */
+  saveState(state: PersistedThreadState): Promise<void>;
+  /** Delete just the persisted state slice, leaving messages intact. */
+  deleteState(): Promise<void>;
 }
 
 /**
@@ -72,6 +114,27 @@ export interface ThreadManagerHooks<TStored, TPrepared = TStored> {
     index: number,
     messages: readonly TPrepared[]
   ) => TPrepared;
+  /**
+   * One-shot list-level pre-pass applied once when a thread is forked with
+   * `transform: true`. Runs before {@link onForkTransform}. May filter,
+   * compact, prepend, or otherwise rewrite the whole forked thread — so the
+   * returned length need not match the input length. Async, so implementations
+   * may call an LLM or other I/O.
+   */
+  onForkPrepareThread?: (
+    messages: readonly TStored[]
+  ) => TStored[] | Promise<TStored[]>;
+  /**
+   * Per-message final pass applied once when a thread is forked with
+   * `transform: true`. Runs after {@link onForkPrepareThread}. Pure 1:1 map —
+   * must return a value for every input message; length cannot change. Same
+   * shape as {@link onPreparedMessage}.
+   */
+  onForkTransform?: (
+    message: TStored,
+    index: number,
+    messages: readonly TStored[]
+  ) => TStored;
 }
 
 export interface ProviderThreadManager<

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

@@ -22,6 +22,8 @@ export type {
   InferToolResults,
   ToolCallResultUnion,
   ProcessToolCallsContext,
+  ProcessToolCallsResult,
+  RewindSignal,
   PreToolUseHookResult,
   PostToolUseFailureHookResult,
   ToolHooks,

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

@@ -2,6 +2,7 @@ import type { TokenUsage, ToolResultConfig } from "../types";
 import type { JsonValue } from "../state/types";
 import type { z } from "zod";
 import type { ActivityFunctionWithOptions } from "@temporalio/workflow";
+import type { SandboxSnapshot } from "../sandbox/types";
 
 // ============================================================================
 // Tool Definition Types
@@ -139,12 +140,29 @@ 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 and the LLM call is retried. The session reuses
+   * the same `assistantMessageId` for the retry; the next `runAgent`
+   * activity truncates the thread from that id on entry, wiping the
+   * triggering assistant message and any tool results already appended
+   * before re-invoking the LLM.
+   *
+   * The `toolResponse` for a rewinding tool call is ignored (never
+   * appended) since the thread is rewound back to the pre-assistant
+   * state on the next invocation.
+   */
+  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) */
   threadId?: string;
   /** Sandbox ID created or used by the handler (e.g. child agent sandbox) */
   sandboxId?: string;
+  /** Snapshot captured on exit when `sandboxShutdown === "snapshot"`. */
+  snapshot?: SandboxSnapshot;
+  /** Snapshot captured immediately after sandbox seeding (before the agent loop starts) when `sandbox.mode === "new"` and `sandboxShutdown === "snapshot"`. Intended as a reusable "base" for new threads that want to skip re-seeding. */
+  baseSnapshot?: SandboxSnapshot;
   /** Unvalidated metadata passthrough from handler to hooks (e.g. infrastructure state) */
   metadata?: Record<string, unknown>;
 }
@@ -278,6 +296,29 @@ export interface ProcessToolCallsContext {
   sandboxId?: string;
 }
 
+/**
+ * Signal that a tool handler requested a rewind. Attached to the
+ * {@link ProcessToolCallsResult} so the session can reuse the same
+ * `assistantMessageId` for the retry; the next `runAgent` activity
+ * then truncates the thread from that id on entry.
+ */
+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 +510,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/types.ts

@@ -92,6 +92,18 @@ export interface RunAgentConfig extends AgentConfig {
   threadKey?: string;
   /** Metadata for the session */
   metadata?: Record<string, unknown>;
+  /**
+   * The id under which the assistant message produced by this call will
+   * be appended. The activity truncates the thread from this id on
+   * entry (no-op on the first attempt) so that:
+   *
+   * - Rewind retries can reuse the same id and the previous (bad)
+   *   assistant + its tool results are wiped before the retry LLM call.
+   * - Resetting the Temporal workflow to this activity restores the
+   *   pre-call thread state: replay re-truncates, re-invokes, and
+   *   appends under the same id.
+   */
+  assistantMessageId: string;
 }
 
 /**

package/src/workflow.ts

@@ -40,6 +40,11 @@ export type {
 
 // Thread utilities
 export { getShortId } from "./lib/thread/id";
+export {
+  THREAD_TTL_SECONDS,
+  getThreadListKey,
+  getThreadMetaKey,
+} from "./lib/thread/keys";
 
 // State management
 export { createAgentStateManager } from "./lib/state";
@@ -49,6 +54,7 @@ export type {
   JsonSerializable,
   JsonValue,
   JsonPrimitive,
+  PersistedThreadState,
 } from "./lib/state";
 
 // Tool router (includes registry functionality)
@@ -57,7 +63,11 @@ export {
   hasNoOtherToolCalls,
   defineTool,
 } from "./lib/tool-router";
-export { defineSubagent, defineSubagentWorkflow } from "./lib/subagent";
+export {
+  defineSubagent,
+  defineSubagentWorkflow,
+  DEFAULT_SUBAGENT_WORKFLOW_RUN_TIMEOUT,
+} from "./lib/subagent";
 export type {
   // Tool definition types
   ToolDefinition,
@@ -94,6 +104,8 @@ export type {
   // Other
   AppendToolResultFn,
   ProcessToolCallsContext,
+  ProcessToolCallsResult,
+  RewindSignal,
 } from "./lib/tool-router";
 
 // Session & message lifecycle hooks
@@ -151,6 +163,7 @@ export { proxyRunAgent } from "./lib/model/proxy";
 // Subagent types
 export type {
   SubagentConfig,
+  SubagentChildWorkflowOptions,
   SubagentDefinition,
   SubagentFnResult,
   SubagentHooks,

package/tsup.config.ts

@@ -4,6 +4,7 @@ export default defineConfig({
   entry: {
     index: "src/index.ts",
     workflow: "src/workflow.ts",
+    "adapters/thread/index": "src/adapters/thread/index.ts",
     "adapters/thread/langchain/index": "src/adapters/thread/langchain/index.ts",
     "adapters/thread/langchain/workflow":
       "src/adapters/thread/langchain/proxy.ts",