zeitlich 0.2.36 → 0.2.38

package/src/lib/session/session.integration.test.ts

@@ -33,6 +33,16 @@ vi.mock("@temporalio/workflow", () => {
     }
   }
 
+  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 {
     proxyActivities: <T>() => ({}) as T,
     condition: async (fn: () => boolean) => fn(),
@@ -46,6 +56,8 @@ vi.mock("@temporalio/workflow", () => {
     uuid4: () =>
       `00000000-0000-0000-0000-${String(++idCounter).padStart(12, "0")}`,
     ApplicationFailure: MockApplicationFailure,
+    CancellationScope: MockCancellationScope,
+    isCancellation: (_err: unknown) => false,
     log: {
       trace: () => {},
       debug: () => {},
@@ -105,6 +117,9 @@ function createMockThreadOps() {
     forkThread: async (source, target) => {
       log.push({ op: "forkThread", args: [source, target] });
     },
+    truncateThread: async (threadId, length) => {
+      log.push({ op: "truncateThread", args: [threadId, length] });
+    },
   });
 
   return { ops, log };
@@ -123,12 +138,18 @@ function createScriptedRunAgent(
   return async () => {
     const turn = turns[call++];
     if (!turn) {
-      return { message: "done", rawToolCalls: [], usage: undefined };
+      return {
+        message: "done",
+        rawToolCalls: [],
+        usage: undefined,
+        threadLengthAtCall: 0,
+      };
     }
     return {
       message: turn.message,
       rawToolCalls: turn.toolCalls,
       usage: turn.usage,
+      threadLengthAtCall: 0,
     };
   };
 }
@@ -517,6 +538,8 @@ describe("createSession integration", () => {
         createdAt: new Date().toISOString(),
       }),
       forkSandbox: async () => "forked-sandbox-id",
+      restoreSandbox: async () => "restored-sandbox-id",
+      deleteSandboxSnapshot: async () => {},
       pauseSandbox: async () => {},
       resumeSandbox: async () => {},
     };
@@ -559,6 +582,8 @@ describe("createSession integration", () => {
         createdAt: new Date().toISOString(),
       }),
       forkSandbox: async () => "forked-sandbox-id",
+      restoreSandbox: async () => "restored-sandbox-id",
+      deleteSandboxSnapshot: async () => {},
       pauseSandbox: async () => {},
       resumeSandbox: async () => {},
     };
@@ -610,6 +635,8 @@ describe("createSession integration", () => {
         createdAt: new Date().toISOString(),
       }),
       forkSandbox: async () => "forked-sb",
+      restoreSandbox: async () => "restored-sb",
+      deleteSandboxSnapshot: async () => {},
     };
 
     const session = await createSession({
@@ -822,6 +849,8 @@ describe("createSession integration", () => {
         createdAt: new Date().toISOString(),
       }),
       forkSandbox: async () => "forked-sandbox-id",
+      restoreSandbox: async () => "restored-sandbox-id",
+      deleteSandboxSnapshot: async () => {},
       pauseSandbox: async () => {},
       resumeSandbox: async () => {},
     };
@@ -873,6 +902,8 @@ describe("createSession integration", () => {
         createdAt: new Date().toISOString(),
       }),
       forkSandbox: async () => "forked-sandbox-id",
+      restoreSandbox: async () => "restored-sandbox-id",
+      deleteSandboxSnapshot: async () => {},
       pauseSandbox: async () => {},
       resumeSandbox: async () => {},
     };
@@ -949,4 +980,133 @@ describe("createSession integration", () => {
     expect(result.usage.totalInputTokens).toBe(180);
     expect(result.usage.totalOutputTokens).toBe(90);
   });
+
+  // --- Snapshot-driven shutdown ---
+
+  it("captures base + exit snapshot and destroys sandbox on sandboxShutdown=snapshot", async () => {
+    const { ops } = createMockThreadOps();
+    const sandboxLog: string[] = [];
+    let snapCounter = 0;
+
+    const sandboxOps: SandboxOps = {
+      createSandbox: async () => {
+        sandboxLog.push("create");
+        return { sandboxId: "sb-snap" };
+      },
+      destroySandbox: async (id: string) => {
+        sandboxLog.push(`destroy:${id}`);
+      },
+      pauseSandbox: async () => {
+        sandboxLog.push("pause");
+      },
+      resumeSandbox: async () => {},
+      snapshotSandbox: async (id: string) => {
+        snapCounter += 1;
+        sandboxLog.push(`snapshot:${id}`);
+        return {
+          sandboxId: id,
+          providerId: "test",
+          data: { tag: `snap-${snapCounter}` },
+          createdAt: new Date().toISOString(),
+        };
+      },
+      restoreSandbox: async () => "restored-sb",
+      deleteSandboxSnapshot: async () => {},
+      forkSandbox: async () => "forked-sb",
+    };
+
+    const session = await createSession({
+      agentName: "TestAgent",
+      thread: { mode: "new", threadId: "thread-snap" },
+      runAgent: createScriptedRunAgent([{ message: "done", toolCalls: [] }]),
+      threadOps: ops,
+      buildContextMessage: () => "go",
+      sandboxOps,
+      sandboxShutdown: "snapshot",
+    });
+
+    const stateManager = createAgentStateManager({
+      initialState: { systemPrompt: "test" },
+    });
+
+    const result = await session.runSession({ stateManager });
+
+    expect(result.exitReason).toBe("completed");
+    expect(result.sandboxId).toBe("sb-snap");
+    expect(result.baseSnapshot?.data).toEqual({ tag: "snap-1" });
+    expect(result.snapshot?.data).toEqual({ tag: "snap-2" });
+    expect(sandboxLog).toEqual([
+      "create",
+      "snapshot:sb-snap",
+      "snapshot:sb-snap",
+      "destroy:sb-snap",
+    ]);
+    expect(sandboxLog).not.toContain("pause");
+  });
+
+  it("restores a sandbox when sandbox.mode=from-snapshot and skips base snapshot", async () => {
+    const { ops } = createMockThreadOps();
+    const sandboxLog: string[] = [];
+    const priorSnapshot = {
+      sandboxId: "sb-prior",
+      providerId: "test",
+      data: { tag: "prior" },
+      createdAt: new Date().toISOString(),
+    };
+
+    const sandboxOps: SandboxOps = {
+      createSandbox: async () => {
+        sandboxLog.push("create");
+        return { sandboxId: "sb-should-not-be-created" };
+      },
+      destroySandbox: async (id: string) => {
+        sandboxLog.push(`destroy:${id}`);
+      },
+      pauseSandbox: async () => {},
+      resumeSandbox: async () => {},
+      snapshotSandbox: async (id: string) => {
+        sandboxLog.push(`snapshot:${id}`);
+        return {
+          sandboxId: id,
+          providerId: "test",
+          data: { tag: "exit" },
+          createdAt: new Date().toISOString(),
+        };
+      },
+      restoreSandbox: async (snap) => {
+        sandboxLog.push(`restore:${(snap.data as { tag: string }).tag}`);
+        return "sb-restored";
+      },
+      deleteSandboxSnapshot: async () => {},
+      forkSandbox: async () => "forked-sb",
+    };
+
+    const session = await createSession({
+      agentName: "TestAgent",
+      thread: { mode: "new", threadId: "thread-restore" },
+      runAgent: createScriptedRunAgent([{ message: "done", toolCalls: [] }]),
+      threadOps: ops,
+      buildContextMessage: () => "go",
+      sandboxOps,
+      sandbox: { mode: "from-snapshot", snapshot: priorSnapshot },
+      sandboxShutdown: "snapshot",
+    });
+
+    const stateManager = createAgentStateManager({
+      initialState: { systemPrompt: "test" },
+    });
+
+    const result = await session.runSession({ stateManager });
+
+    expect(result.sandboxId).toBe("sb-restored");
+    // No base snapshot because the sandbox was restored, not freshly created.
+    expect(result.baseSnapshot).toBeUndefined();
+    expect(result.snapshot?.data).toEqual({ tag: "exit" });
+    expect(sandboxLog).toEqual([
+      "restore:prior",
+      "snapshot:sb-restored",
+      "destroy:sb-restored",
+    ]);
+    expect(sandboxLog).not.toContain("create");
+  });
 });

package/src/lib/session/session.ts

@@ -7,7 +7,11 @@ import {
 } from "@temporalio/workflow";
 import type { SessionExitReason } from "../types";
 import type { SessionConfig, ZeitlichSession } from "./types";
-import type { SandboxOps } from "../sandbox/types";
+import type {
+  SandboxCreateOptions,
+  SandboxOps,
+  SandboxSnapshot,
+} from "../sandbox/types";
 import type {
   AgentState,
   AgentStateManager,
@@ -142,16 +146,19 @@ export async function createSession<
     appendSystemMessage,
     appendAgentMessage,
     forkThread,
+    truncateThread,
   } = threadOps;
 
   const plugins: ToolMap[string][] = [];
   let destroySubagentSandboxes: (() => Promise<void>) | undefined;
+  let cleanupSubagentSnapshots: (() => Promise<void>) | undefined;
 
   if (subagents) {
     const result = buildSubagentRegistration(subagents);
     if (result) {
       plugins.push(result.registration);
       destroySubagentSandboxes = result.destroySubagentSandboxes;
+      cleanupSubagentSnapshots = result.cleanupSubagentSnapshots;
     }
   }
   if (skills) {
@@ -210,10 +217,13 @@ export async function createSession<
         }
       );
 
-      // --- Sandbox lifecycle: create, continue, fork, or inherit ----------
+      // --- Sandbox lifecycle: create, continue, fork, from-snapshot, or inherit ---
       const sandboxMode = sandboxInit?.mode;
       let sandboxId: string | undefined;
       let sandboxOwned = false;
+      let baseSnapshot: SandboxSnapshot | undefined;
+      let exitSnapshot: SandboxSnapshot | undefined;
+      let freshlyCreated = false;
 
       if (sandboxMode === "inherit") {
         const inheritInit = sandboxInit as {
@@ -248,8 +258,31 @@ export async function createSession<
             nonRetryable: true,
           });
         }
+        const forkInit = sandboxInit as {
+          mode: "fork";
+          sandboxId: string;
+          options?: SandboxCreateOptions;
+        };
         sandboxId = await sandboxOps.forkSandbox(
-          (sandboxInit as { mode: "fork"; sandboxId: string }).sandboxId
+          forkInit.sandboxId,
+          forkInit.options
+        );
+        sandboxOwned = true;
+      } else if (sandboxMode === "from-snapshot") {
+        if (!sandboxOps) {
+          throw ApplicationFailure.create({
+            message: "No sandboxOps provided — cannot restore sandbox",
+            nonRetryable: true,
+          });
+        }
+        const restoreInit = sandboxInit as {
+          mode: "from-snapshot";
+          snapshot: SandboxSnapshot;
+          options?: SandboxCreateOptions;
+        };
+        sandboxId = await sandboxOps.restoreSandbox(
+          restoreInit.snapshot,
+          restoreInit.options
         );
         sandboxOwned = true;
       } else if (sandboxOps) {
@@ -263,10 +296,24 @@ export async function createSession<
         if (result) {
           sandboxId = result.sandboxId;
           sandboxOwned = true;
+          freshlyCreated = true;
         }
       }
 
-      if (sandboxId && onSandboxReady) {
+      // Capture a base snapshot immediately after seeding so it can be reused
+      // as a template for future runs that want to skip the (potentially
+      // expensive) seed step.
+      if (
+        sandboxId &&
+        sandboxOwned &&
+        freshlyCreated &&
+        sandboxShutdown === "snapshot" &&
+        sandboxOps
+      ) {
+        baseSnapshot = await sandboxOps.snapshotSandbox(sandboxId);
+      }
+
+      if (sandboxId && sandboxOwned && onSandboxReady) {
         onSandboxReady(sandboxId);
       }
 
@@ -347,6 +394,7 @@ export async function createSession<
       );
 
       let exitReason: SessionExitReason = "completed";
+      let finalMessage: M | null = null;
 
       try {
         while (
@@ -361,13 +409,24 @@ export async function createSession<
 
           stateManager.setTools(toolRouter.getToolDefinitions());
 
-          const { message, rawToolCalls, usage } = await runAgent({
+          const {
+            message,
+            rawToolCalls,
+            usage,
+            threadLengthAtCall,
+          } = await runAgent({
             threadId,
             threadKey,
             agentName,
             metadata,
           });
 
+          // 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);
 
           if (usage) {
@@ -385,21 +444,8 @@ export async function createSession<
           if (!toolRouter.hasTools() || rawToolCalls.length === 0) {
             stateManager.complete();
             exitReason = "completed";
-            log.info("session ended", {
-              agentName,
-              threadId,
-              exitReason,
-              turns: currentTurn,
-              durationMs: Date.now() - sessionStartMs,
-              usage: stateManager.getTotalUsage(),
-            });
-            return {
-              threadId,
-              finalMessage: message,
-              exitReason,
-              usage: stateManager.getTotalUsage(),
-              sandboxId,
-            } as Awaited<ReturnType<ZeitlichSession<M, boolean>["runSession"]>>;
+            finalMessage = message;
+            break;
           }
 
           const parsedToolCalls: ParsedToolCallUnion<T>[] = [];
@@ -433,6 +479,33 @@ export async function createSession<
             }
           }
 
+          const rewind = toolCallResults.rewind;
+          if (rewind) {
+            log.info("rewinding turn", {
+              agentName,
+              threadId,
+              turn: currentTurn,
+              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
+            // rewind still consumes one of the `maxTurns` budget so a
+            // misbehaving tool cannot spin the session forever.
+            await truncateThread(threadId, preAssistantLength, threadKey);
+            continue;
+          }
+
           if (stateManager.getStatus() === "WAITING_FOR_INPUT") {
             const conditionMet = await condition(
               () => stateManager.getStatus() === "RUNNING",
@@ -480,12 +553,20 @@ export async function createSession<
             case "keep":
             case "keep-until-parent-close":
               break;
+            case "snapshot":
+              exitSnapshot = await sandboxOps.snapshotSandbox(sandboxId);
+              await sandboxOps.destroySandbox(sandboxId);
+              break;
           }
         }
 
         if (destroySubagentSandboxes) {
           await destroySubagentSandboxes();
         }
+
+        if (cleanupSubagentSnapshots) {
+          await cleanupSubagentSnapshots();
+        }
       }
 
       log.info("session ended", {
@@ -495,14 +576,18 @@ export async function createSession<
         turns: stateManager.getTurns(),
         durationMs: Date.now() - sessionStartMs,
         usage: stateManager.getTotalUsage(),
+        ...(baseSnapshot && { hasBaseSnapshot: true }),
+        ...(exitSnapshot && { hasExitSnapshot: true }),
       });
 
       return {
         threadId,
-        finalMessage: null,
+        finalMessage,
         exitReason,
         usage: stateManager.getTotalUsage(),
         sandboxId,
+        ...(baseSnapshot && { baseSnapshot }),
+        ...(exitSnapshot && { snapshot: exitSnapshot }),
       } as Awaited<ReturnType<ZeitlichSession<M, boolean>["runSession"]>>;
     },
   };

package/src/lib/session/types.ts

@@ -1,8 +1,5 @@
 import type { Duration } from "@temporalio/common";
-import type {
-  SessionExitReason,
-  ToolResultConfig,
-} from "../types";
+import type { SessionExitReason, ToolResultConfig } from "../types";
 import type {
   ToolMap,
   ToolCallResultUnion,
@@ -11,7 +8,7 @@ import type {
 import type { Hooks } from "../hooks/types";
 import type { SubagentConfig } from "../subagent/types";
 import type { Skill } from "../skills/types";
-import type { SandboxOps } from "../sandbox/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";
@@ -62,6 +59,18 @@ export interface ThreadOps<TContent = 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.
+   */
+  truncateThread(
+    threadId: string,
+    length: number,
+    threadKey?: string
+  ): Promise<void>;
 }
 
 /**
@@ -219,6 +228,17 @@ export type SessionResult<
   finalMessage: M | null;
   exitReason: SessionExitReason;
   usage: ReturnType<AgentStateManager<TState>["getTotalUsage"]>;
+  /**
+   * 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;
 } & (HasSandbox extends true
   ? { sandboxId: string }
   : { sandboxId?: undefined });

package/src/lib/skills/fs-provider.ts

@@ -24,7 +24,7 @@ import { parseSkillFile } from "./parse";
 export class FileSystemSkillProvider implements SkillProvider {
   constructor(
     private readonly fs: SandboxFileSystem,
-    private readonly baseDir: string,
+    private readonly baseDir: string
   ) {}
 
   async listSkills(): Promise<SkillMetadata[]> {
@@ -45,20 +45,21 @@ export class FileSystemSkillProvider implements SkillProvider {
   }
 
   async getSkill(name: string): Promise<Skill> {
-    const raw = await this.fs.readFile(
-      join(this.baseDir, name, "SKILL.md"),
-    );
+    const raw = await this.fs.readFile(join(this.baseDir, name, "SKILL.md"));
     const { frontmatter, body } = parseSkillFile(raw);
 
     if (frontmatter.name !== name) {
       throw new Error(
-        `Skill directory "${name}" contains SKILL.md with mismatched name "${frontmatter.name}"`,
+        `Skill directory "${name}" contains SKILL.md with mismatched name "${frontmatter.name}"`
       );
     }
 
     const location = join(this.baseDir, name);
     const resourcePaths = await this.discoverResources(name);
-    const resourceContents = await this.readResourceContents(location, resourcePaths);
+    const resourceContents = await this.readResourceContents(
+      location,
+      resourcePaths
+    );
     return {
       ...frontmatter,
       instructions: body,
@@ -80,7 +81,10 @@ export class FileSystemSkillProvider implements SkillProvider {
       const { frontmatter, body } = parseSkillFile(raw);
       const location = join(this.baseDir, dir);
       const resourcePaths = await this.discoverResources(dir);
-      const resourceContents = await this.readResourceContents(location, resourcePaths);
+      const resourceContents = await this.readResourceContents(
+        location,
+        resourcePaths
+      );
       skills.push({
         ...frontmatter,
         instructions: body,
@@ -119,7 +123,7 @@ export class FileSystemSkillProvider implements SkillProvider {
 
   private async readResourceContents(
     location: string,
-    resources: string[],
+    resources: string[]
   ): Promise<Record<string, string> | undefined> {
     if (resources.length === 0) return undefined;
     const contents: Record<string, string> = {};

package/src/lib/skills/handler.ts

@@ -18,7 +18,7 @@ function formatSkillResponse(skill: Skill): string {
   if (skill.location) {
     parts.push(`\nSkill directory: ${skill.location}`);
     parts.push(
-      "Relative paths in this skill resolve against the skill directory above.",
+      "Relative paths in this skill resolve against the skill directory above."
     );
   }
 

package/src/lib/skills/parse.ts

@@ -12,7 +12,9 @@ export function parseSkillFile(raw: string): {
   body: string;
 } {
   const trimmed = raw.replace(/^\uFEFF/, ""); // strip BOM
-  const match = trimmed.match(/^---[ \t]*\r?\n([\s\S]*?)\r?\n---[ \t]*\r?\n?([\s\S]*)$/);
+  const match = trimmed.match(
+    /^---[ \t]*\r?\n([\s\S]*?)\r?\n---[ \t]*\r?\n?([\s\S]*)$/
+  );
 
   if (!match) {
     throw new Error(

package/src/lib/skills/register.ts

@@ -11,9 +11,7 @@ function validateSkillNames(skills: SkillMetadata[]): void {
   const names = skills.map((s) => s.name);
   const dupes = names.filter((n, i) => names.indexOf(n) !== i);
   if (dupes.length > 0) {
-    throw new Error(
-      `Duplicate skill names: ${[...new Set(dupes)].join(", ")}`
-    );
+    throw new Error(`Duplicate skill names: ${[...new Set(dupes)].join(", ")}`);
   }
 }