zeitlich 0.2.38 → 0.2.39

package/src/lib/session/session.ts

@@ -16,6 +16,7 @@ import type {
   AgentState,
   AgentStateManager,
   JsonSerializable,
+  PersistedThreadState,
 } from "../state/types";
 import { createToolRouter } from "../tool-router/router";
 import type { ParsedToolCallUnion, ToolMap } from "../tool-router/types";
@@ -111,6 +112,7 @@ export async function createSession<
   sandbox: sandboxInit,
   sandboxShutdown = "destroy",
   onSandboxReady,
+  onSessionExit,
   virtualFs: virtualFsConfig,
   virtualFsOps,
 }: SessionConfig<T, M, TContent>): Promise<ZeitlichSession<M, boolean>> {
@@ -146,7 +148,8 @@ export async function createSession<
     appendSystemMessage,
     appendAgentMessage,
     forkThread,
-    truncateThread,
+    loadThreadState,
+    saveThreadState,
   } = threadOps;
 
   const plugins: ToolMap[string][] = [];
@@ -314,7 +317,10 @@ export async function createSession<
       }
 
       if (sandboxId && sandboxOwned && onSandboxReady) {
-        onSandboxReady(sandboxId);
+        onSandboxReady({
+          sandboxId,
+          ...(baseSnapshot && { baseSnapshot }),
+        });
       }
 
       // --- Virtual filesystem init (independent of sandbox) ----------------
@@ -366,10 +372,21 @@ export async function createSession<
       const systemPrompt = stateManager.getSystemPrompt();
 
       // --- Thread lifecycle: new, continue, or fork ----------------------
+      const rehydrateFromSlice = (slice: PersistedThreadState): void => {
+        stateManager.mergeUpdate({
+          tasks: new Map(slice.tasks),
+          ...slice.custom,
+        } as Partial<AgentState<TState>>);
+      };
+
       if (threadMode === "fork" && sourceThreadId) {
         await forkThread(sourceThreadId, threadId, threadKey);
+        const forkedSlice = await loadThreadState(threadId, threadKey);
+        if (forkedSlice) rehydrateFromSlice(forkedSlice);
       } else if (threadMode === "continue") {
         // "continue" — thread already exists, just append the new message
+        const continuedSlice = await loadThreadState(threadId, threadKey);
+        if (continuedSlice) rehydrateFromSlice(continuedSlice);
       } else {
         if (appendSystemPrompt) {
           if (
@@ -397,6 +414,13 @@ export async function createSession<
       let finalMessage: M | null = null;
 
       try {
+        // Per-turn assistant message id. Pre-generated in the workflow
+        // so the runAgent activity can truncate the thread from this id
+        // on entry (deterministic rewind + time-travel via Temporal
+        // workflow reset). On a rewind retry we keep the same id so the
+        // prior attempt's assistant + tool results are wiped by the next
+        // runAgent call.
+        let assistantId: string | undefined;
         while (
           stateManager.isRunning() &&
           !stateManager.isTerminal() &&
@@ -409,25 +433,17 @@ export async function createSession<
 
           stateManager.setTools(toolRouter.getToolDefinitions());
 
-          const {
-            message,
-            rawToolCalls,
-            usage,
-            threadLengthAtCall,
-          } = await runAgent({
+          assistantId ??= uuid4();
+
+          const { message, rawToolCalls, usage } = await runAgent({
             threadId,
             threadKey,
             agentName,
             metadata,
+            assistantMessageId: assistantId,
           });
 
-          // The invoker loaded the thread right before calling the LLM,
-          // so it already knows how many messages were stored at that
-          // point — we use that directly as the rewind snapshot instead
-          // of a separate activity round-trip.
-          const preAssistantLength = threadLengthAtCall;
-
-          await appendAgentMessage(threadId, uuid4(), message, threadKey);
+          await appendAgentMessage(threadId, assistantId, message, threadKey);
 
           if (usage) {
             stateManager.updateUsage(usage);
@@ -488,24 +504,19 @@ export async function createSession<
               toolCallId: rewind.toolCallId,
               toolName: rewind.toolName,
             });
-            if (preAssistantLength === undefined) {
-              throw ApplicationFailure.create({
-                message:
-                  "Rewind requested but runAgent did not report " +
-                  "`threadLengthAtCall`; the adapter must populate it to " +
-                  "support rewinds.",
-                nonRetryable: true,
-              });
-            }
-            // Drop the assistant message + any already-saved tool results
-            // so the LLM call can be retried from the pre-assistant state.
-            // The turn counter is intentionally NOT rolled back — each
+            // Keep the same assistantId for the retry. The next
+            // runAgent call will call truncateFromId(assistantId) on
+            // entry, wiping the bad assistant message + any already
+            // appended tool results before re-invoking the LLM. The
+            // turn counter is intentionally NOT rolled back — each
             // rewind still consumes one of the `maxTurns` budget so a
             // misbehaving tool cannot spin the session forever.
-            await truncateThread(threadId, preAssistantLength, threadKey);
             continue;
           }
 
+          // Turn committed: fresh id for the next turn.
+          assistantId = undefined;
+
           if (stateManager.getStatus() === "WAITING_FOR_INPUT") {
             const conditionMet = await condition(
               () => stateManager.getStatus() === "RUNNING",
@@ -539,6 +550,28 @@ export async function createSession<
         });
         throw ApplicationFailure.fromError(error);
       } finally {
+        // Persist the task map + custom state slice alongside the thread so
+        // a future `continue` / `fork` run can rehydrate it. Runs on every
+        // exit path (completed, failed, cancelled, max_turns,
+        // waiting_for_input timeout). Best-effort: failures here must not
+        // mask the original exit reason / error.
+        try {
+          await saveThreadState(
+            threadId,
+            stateManager.getPersistedSlice(),
+            threadKey
+          );
+        } catch (persistError) {
+          log.warn("failed to persist thread state", {
+            agentName,
+            threadId,
+            error:
+              persistError instanceof Error
+                ? persistError.message
+                : String(persistError),
+          });
+        }
+
         await callSessionEnd(exitReason, stateManager.getTurns());
 
         if (sandboxOwned && sandboxId && sandboxOps) {
@@ -580,6 +613,13 @@ export async function createSession<
         ...(exitSnapshot && { hasExitSnapshot: true }),
       });
 
+      if (onSessionExit) {
+        onSessionExit({
+          ...(sandboxId && { sandboxId }),
+          ...(exitSnapshot && { snapshot: exitSnapshot }),
+        });
+      }
+
       return {
         threadId,
         finalMessage,

package/src/lib/session/types.ts

@@ -11,7 +11,11 @@ import type { Skill } from "../skills/types";
 import type { SandboxOps, SandboxSnapshot } from "../sandbox/types";
 import type { VirtualFsOps } from "../virtual-fs/types";
 import type { RunAgentActivity } from "../model/types";
-import type { AgentStateManager, JsonSerializable } from "../state/types";
+import type {
+  AgentStateManager,
+  JsonSerializable,
+  PersistedThreadState,
+} from "../state/types";
 import type { ActivityInterfaceFor } from "@temporalio/workflow";
 import type {
   ThreadInit,
@@ -53,22 +57,52 @@ export interface ThreadOps<TContent = string> {
     content: unknown,
     threadKey?: string
   ): Promise<void>;
-  /** Copy all messages from sourceThreadId into a new thread at targetThreadId */
+  /**
+   * Copy all messages AND the persisted state slice (tasks + custom
+   * state) from `sourceThreadId` into a new thread at `targetThreadId`.
+   * Adapters that have `onForkPrepareThread` and/or `onForkTransform`
+   * hooks configured apply them once to the new thread's messages
+   * before returning.
+   */
   forkThread(
     sourceThreadId: string,
     targetThreadId: string,
     threadKey?: string
   ): Promise<void>;
   /**
-   * Truncate the thread back to `length` messages. Used by the session's
-   * rewind flow to roll the thread back before retrying a turn. The
-   * session obtains `length` from `AgentResponse.threadLengthAtCall`,
-   * which the model invoker computes for free from the messages it
-   * loaded before invoking the LLM.
+   * Truncate the thread starting at `messageId`: that message and every
+   * message after it are removed. If `messageId` is not present the call
+   * is a no-op.
+   *
+   * The `runAgent` activity invokes this on entry with the pre-generated
+   * `assistantMessageId`. On the happy path the id is not yet in the
+   * thread and the call is a no-op. On a rewind retry (same assistant
+   * id reused) or a Temporal workflow reset-to-this-activity the id is
+   * present, so the bad assistant + any tool results it produced are
+   * wiped and the call is then replayable.
    */
   truncateThread(
     threadId: string,
-    length: number,
+    messageId: string,
+    threadKey?: string
+  ): Promise<void>;
+  /**
+   * Load the persisted state slice (tasks + custom state) associated with
+   * the thread, or `null` if none has been saved yet. Called on session
+   * start for `continue`/`fork` threads to rehydrate {@link AgentStateManager}.
+   */
+  loadThreadState(
+    threadId: string,
+    threadKey?: string
+  ): Promise<PersistedThreadState | null>;
+  /**
+   * Overwrite the persisted state slice for the thread. Called once from
+   * the session's `finally` block on every exit path so that "finish,
+   * store, continue later" works regardless of exit reason.
+   */
+  saveThreadState(
+    threadId: string,
+    state: PersistedThreadState,
     threadKey?: string
   ): Promise<void>;
 }
@@ -195,8 +229,25 @@ export interface SessionConfig<
   /**
    * Called as soon as the sandbox is created (or resumed/forked), before the
    * agent loop starts. Useful for signalling sandbox readiness to a parent.
+   *
+   * `baseSnapshot` is only populated when the sandbox was freshly created
+   * this run and `sandboxShutdown === "snapshot"` — i.e. when the session
+   * captured a seed snapshot intended for reuse.
+   */
+  onSandboxReady?: (args: {
+    sandboxId: string;
+    baseSnapshot?: SandboxSnapshot;
+  }) => void;
+  /**
+   * Called right before `runSession` returns, with the session's sandbox
+   * outputs. Useful for callers (e.g. `defineSubagentWorkflow`) that want to
+   * forward these fields to their own return value without requiring user
+   * code to manually thread them through.
    */
-  onSandboxReady?: (sandboxId: string) => void;
+  onSessionExit?: (result: {
+    sandboxId?: string;
+    snapshot?: SandboxSnapshot;
+  }) => void;
 
   // ---------------------------------------------------------------------------
   // Virtual filesystem

package/src/lib/state/index.ts

@@ -6,4 +6,5 @@ export type {
   JsonSerializable,
   AgentState,
   AgentStateManager,
+  PersistedThreadState,
 } from "./types";

package/src/lib/state/manager.integration.test.ts

@@ -350,4 +350,113 @@ describe("createAgentStateManager integration", () => {
     sm.incrementTurns();
     expect(sm.getTotalUsage().turns).toBe(2);
   });
+
+  // --- Persisted slice round-trip ---
+
+  it("getPersistedSlice snapshots tasks + custom state, omits runtime bookkeeping", () => {
+    const sm = createAgentStateManager<{ label: string; count: number }>({
+      initialState: {
+        systemPrompt: "test",
+        label: "original",
+        count: 1,
+      },
+    });
+
+    const task: WorkflowTask = {
+      id: "task-1",
+      subject: "subject",
+      description: "description",
+      activeForm: "doing",
+      status: "pending",
+      metadata: {},
+      blockedBy: [],
+      blocks: [],
+    };
+    sm.setTask(task);
+    sm.incrementTurns();
+    sm.updateUsage({ inputTokens: 42 });
+
+    const slice = sm.getPersistedSlice();
+    expect(slice.tasks).toEqual([["task-1", task]]);
+    expect(slice.custom).toEqual({ label: "original", count: 1 });
+    expect("status" in slice.custom).toBe(false);
+    expect("version" in slice.custom).toBe(false);
+    expect("turns" in slice.custom).toBe(false);
+    expect("tools" in slice.custom).toBe(false);
+  });
+
+  it("mergeUpdate replaces the task map when given a tasks field", () => {
+    const sm = createAgentStateManager<{ label: string; extra?: string }>({
+      initialState: {
+        systemPrompt: "test",
+        label: "original",
+      },
+    });
+
+    const existing: WorkflowTask = {
+      id: "task-pre",
+      subject: "pre",
+      description: "pre",
+      activeForm: "pre",
+      status: "pending",
+      metadata: {},
+      blockedBy: [],
+      blocks: [],
+    };
+    sm.setTask(existing);
+    const versionBefore = sm.getVersion();
+
+    const incoming: WorkflowTask = {
+      id: "task-new",
+      subject: "new",
+      description: "new",
+      activeForm: "new",
+      status: "in_progress",
+      metadata: { foo: "bar" },
+      blockedBy: [],
+      blocks: [],
+    };
+
+    sm.mergeUpdate({
+      tasks: new Map([["task-new", incoming]]),
+      label: "restored",
+      extra: "extra-value",
+    });
+
+    expect(sm.getTasks()).toEqual([incoming]);
+    expect(sm.getTask("task-pre")).toBeUndefined();
+    expect(sm.get("label")).toBe("restored");
+    expect(sm.get("extra")).toBe("extra-value");
+    expect(sm.getVersion()).toBe(versionBefore + 1);
+  });
+
+  it("getPersistedSlice + mergeUpdate round-trips tasks + custom state", () => {
+    const original = createAgentStateManager<{ answer: number }>({
+      initialState: { systemPrompt: "test", answer: 42 },
+    });
+    const task: WorkflowTask = {
+      id: "t",
+      subject: "s",
+      description: "d",
+      activeForm: "doing",
+      status: "completed",
+      metadata: { a: "b" },
+      blockedBy: ["x"],
+      blocks: ["y"],
+    };
+    original.setTask(task);
+
+    const slice = original.getPersistedSlice();
+
+    const restored = createAgentStateManager<{ answer: number }>({
+      initialState: { systemPrompt: "test", answer: 0 },
+    });
+    restored.mergeUpdate({
+      tasks: new Map(slice.tasks),
+      ...slice.custom,
+    });
+
+    expect(restored.getTasks()).toEqual([task]);
+    expect(restored.get("answer")).toBe(42);
+  });
 });

package/src/lib/state/manager.ts

@@ -11,7 +11,13 @@ import {
   isTerminalStatus,
 } from "../types";
 import type { ToolDefinition } from "../tool-router/types";
-import type { AgentState, AgentStateManager, JsonSerializable } from "./types";
+import type {
+  AgentState,
+  AgentStateManager,
+  JsonSerializable,
+  JsonValue,
+  PersistedThreadState,
+} from "./types";
 import { z } from "zod";
 
 /**
@@ -63,11 +69,19 @@ export function createAgentStateManager<
   const tasks = new Map<string, WorkflowTask>(initialState?.tasks);
 
   const {
-    status: _,
-    version: __,
-    turns: ___,
-    tasks: ____,
-    tools: _____,
+    status: _status,
+    version: _version,
+    turns: _turns,
+    tasks: _tasks,
+    tools: _tools,
+    systemPrompt: _systemPrompt,
+    fileTree: _fileTree,
+    inlineFiles: _inlineFiles,
+    virtualFsCtx: _virtualFsCtx,
+    totalInputTokens: _totalInputTokens,
+    totalOutputTokens: _totalOutputTokens,
+    cachedWriteTokens: _cachedWriteTokens,
+    cachedReadTokens: _cachedReadTokens,
     ...custom
   } = initialState ?? {};
   const customState = custom as TCustom;
@@ -166,8 +180,17 @@ export function createAgentStateManager<
       version++;
     },
 
-    mergeUpdate(update: Partial<TCustom>): void {
-      Object.assign(customState as object, update);
+    mergeUpdate(update: Partial<AgentState<TCustom>>): void {
+      const { tasks: nextTasks, ...rest } = update as Partial<
+        AgentState<TCustom>
+      > & { tasks?: Map<string, WorkflowTask> };
+      if (nextTasks) {
+        tasks.clear();
+        for (const [id, task] of nextTasks) {
+          tasks.set(id, task);
+        }
+      }
+      Object.assign(customState as object, rest);
       version++;
     },
 
@@ -214,6 +237,13 @@ export function createAgentStateManager<
       return deleted;
     },
 
+    getPersistedSlice(): PersistedThreadState {
+      return {
+        tasks: Array.from(tasks.entries()),
+        custom: { ...(customState as unknown as Record<string, JsonValue>) },
+      };
+    },
+
     updateUsage(usage: {
       inputTokens?: number;
       outputTokens?: number;

package/src/lib/state/types.ts

@@ -48,6 +48,23 @@ export type JsonSerializable<T> = {
 export type AgentState<TCustom extends JsonSerializable<TCustom>> =
   BaseAgentState & TCustom;
 
+/**
+ * The slice of agent state that is persisted alongside the thread in the
+ * thread store (e.g. Redis) so that a workflow can terminate, store its
+ * state, and be continued or forked later with that state rehydrated.
+ *
+ * Only fields that make sense to carry across workflow runs belong here.
+ * Runtime bookkeeping like status, version, turns, tools, fileTree, token
+ * counters, and the system prompt is intentionally NOT persisted — each run
+ * rebuilds those from scratch.
+ */
+export interface PersistedThreadState {
+  /** Task map serialized as entries so it round-trips through JSON. */
+  tasks: [string, WorkflowTask][];
+  /** All custom state fields declared by the caller. */
+  custom: Record<string, JsonValue>;
+}
+
 /**
  * Agent state manager interface
  * Note: Temporal handlers must be set up in the workflow file due to
@@ -122,6 +139,14 @@ export interface AgentStateManager<TCustom extends JsonSerializable<TCustom>> {
   /** Set the tools (converts Zod schemas to JSON Schema for serialization) */
   setTools(newTools: ToolDefinition[]): void;
 
+  /**
+   * Snapshot the fields that should survive across workflow runs
+   * (tasks + all custom state). Safe to pass directly to
+   * {@link ThreadOps.saveThreadState}. Rehydrate on the next run with
+   * `mergeUpdate({ tasks: new Map(slice.tasks), ...slice.custom })`.
+   */
+  getPersistedSlice(): PersistedThreadState;
+
   /** Update the usage */
   updateUsage(usage: TokenUsage): void;