zeitlich 0.2.37 → 0.2.39

package/src/lib/session/session-edge-cases.integration.test.ts

@@ -2,6 +2,7 @@ import { describe, expect, it, vi, beforeEach } from "vitest";
 import { z } from "zod";
 import type { ToolResultConfig, TokenUsage } from "../types";
 import type { ThreadOps } from "./types";
+import type { PersistedThreadState } from "../state/types";
 import type { RunAgentActivity } from "../model/types";
 import type { RawToolCall } from "../tool-router/types";
 import type { SandboxOps } from "../sandbox/types";
@@ -29,6 +30,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(),
@@ -42,6 +53,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: () => {},
@@ -83,6 +96,7 @@ function toActivityInterface<TContent = string>(
 
 function createMockThreadOps() {
   const log: { op: string; args: unknown[] }[] = [];
+  const stateStore = new Map<string, PersistedThreadState>();
   const ops = toActivityInterface({
     initializeThread: async (threadId) => {
       log.push({ op: "initializeThread", args: [threadId] });
@@ -101,19 +115,38 @@ function createMockThreadOps() {
     },
     forkThread: async (source, target) => {
       log.push({ op: "forkThread", args: [source, target] });
+      const src = stateStore.get(source);
+      if (src) stateStore.set(target, src);
+    },
+    truncateThread: async (threadId, messageId) => {
+      log.push({ op: "truncateThread", args: [threadId, messageId] });
+    },
+    loadThreadState: async (threadId) => {
+      log.push({ op: "loadThreadState", args: [threadId] });
+      return stateStore.get(threadId) ?? null;
+    },
+    saveThreadState: async (threadId, state) => {
+      log.push({ op: "saveThreadState", args: [threadId, state] });
+      stateStore.set(threadId, state);
     },
   });
-  return { ops, log };
+  return { ops, log, stateStore };
 }
 
 function createScriptedRunAgent(
-  turns: TurnScript[]
+  turns: TurnScript[],
+  assistantIdLog?: string[]
 ): RunAgentActivity<unknown> {
   let call = 0;
-  return async () => {
+  return async (config) => {
+    assistantIdLog?.push(config.assistantMessageId);
     const turn = turns[call++];
     if (!turn) {
-      return { message: "done", rawToolCalls: [], usage: undefined };
+      return {
+        message: "done",
+        rawToolCalls: [],
+        usage: undefined,
+      };
     }
     return {
       message: turn.message,
@@ -775,6 +808,16 @@ describe("createSession edge cases", () => {
       forkThread: async (source, target) => {
         log.push({ op: "forkThread", args: [source, target] });
       },
+      truncateThread: async (threadId, messageId) => {
+        log.push({ op: "truncateThread", args: [threadId, messageId] });
+      },
+      loadThreadState: async (threadId) => {
+        log.push({ op: "loadThreadState", args: [threadId] });
+        return null;
+      },
+      saveThreadState: async (threadId, state) => {
+        log.push({ op: "saveThreadState", args: [threadId, state] });
+      },
     });
 
     const session = await createSession<
@@ -1659,4 +1702,293 @@ describe("createSession edge cases", () => {
     expect(sandboxLog).not.toContain("pause:kept-sb");
     expect(sandboxLog).not.toContain("destroy:kept-sb");
   });
+
+  // --- Rewind flow: tool requests rewind and turn is retried -------------
+  //
+  // The session no longer issues an explicit truncateThread on rewind.
+  // Instead it reuses the pre-generated assistantMessageId for the retry,
+  // and the runAgent activity itself truncates the thread from that id
+  // on entry. These tests assert the observable behaviour: the rewinding
+  // tool's result is not appended, turns are consumed as expected, and
+  // the retry invocation receives the same assistantMessageId so the
+  // invoker can wipe the prior attempt.
+
+  it("rewinds the turn when a tool handler returns rewind:true", async () => {
+    const { ops, log } = createMockThreadOps();
+
+    let rewindAttempts = 0;
+    const rewindTool = defineTool({
+      name: "Rewind" as const,
+      description: "rewinds once then succeeds",
+      schema: z.object({}),
+      handler: async () => {
+        rewindAttempts += 1;
+        if (rewindAttempts === 1) {
+          return {
+            toolResponse: "ignored",
+            data: null,
+            rewind: true,
+          };
+        }
+        return { toolResponse: "ok", data: null };
+      },
+    });
+
+    const assistantIds: string[] = [];
+    const session = await createSession({
+      agentName: "TestAgent",
+      thread: { mode: "new", threadId: "thread-1" },
+      runAgent: createScriptedRunAgent(
+        [
+          {
+            message: "attempt-1",
+            toolCalls: [{ id: "tc-1", name: "Rewind", args: {} }],
+          },
+          {
+            message: "attempt-2",
+            toolCalls: [{ id: "tc-2", name: "Rewind", args: {} }],
+          },
+          { message: "done", toolCalls: [] },
+        ],
+        assistantIds
+      ),
+      threadOps: ops,
+      tools: { Rewind: rewindTool },
+      buildContextMessage: () => "go",
+    });
+
+    const stateManager = createAgentStateManager({
+      initialState: { systemPrompt: "test" },
+    });
+
+    const result = await session.runSession({ stateManager });
+
+    expect(result.exitReason).toBe("completed");
+    expect(result.finalMessage).toBe("done");
+    expect(rewindAttempts).toBe(2);
+
+    // Session does not call truncateThread directly on rewind — the
+    // invoker truncates on entry via the reused assistantMessageId.
+    const truncateOps = log.filter((l) => l.op === "truncateThread");
+    expect(truncateOps).toHaveLength(0);
+
+    // The first and second calls reuse the same assistantMessageId
+    // (rewind retry), then the third uses a fresh id.
+    expect(assistantIds).toHaveLength(3);
+    expect(assistantIds[0]).toBe(assistantIds[1]);
+    expect(assistantIds[1]).not.toBe(assistantIds[2]);
+
+    const noRewindToolResult = log.filter((l) => {
+      if (l.op !== "appendToolResult") return false;
+      const config = l.args[1] as ToolResultConfig;
+      return config.toolCallId === "tc-1";
+    });
+    expect(noRewindToolResult).toHaveLength(0);
+
+    const agentAppends = log.filter((l) => l.op === "appendAgentMessage");
+    expect(agentAppends).toHaveLength(3);
+    // The first two assistant appends reuse the same id — the second
+    // will be a no-op in the real adapter because truncateFromId clears
+    // the dedup marker for the old one before the retry invocation.
+    const asstIds = agentAppends.map((l) => l.args[1]);
+    expect(asstIds[0]).toBe(asstIds[1]);
+    expect(asstIds[1]).not.toBe(asstIds[2]);
+  });
+
+  it("reuses the assistantMessageId on rewind even with sibling tool calls", async () => {
+    const { ops, log } = createMockThreadOps();
+
+    let rewindFired = false;
+
+    const siblingTool = defineTool({
+      name: "Sibling" as const,
+      description: "sibling",
+      schema: z.object({}),
+      handler: async () => ({ toolResponse: "sibling-ok", data: null }),
+    });
+
+    const rewindTool = defineTool({
+      name: "Rewind" as const,
+      description: "rewinds",
+      schema: z.object({}),
+      handler: async () => {
+        if (!rewindFired) {
+          rewindFired = true;
+          return { toolResponse: "ignored", data: null, rewind: true };
+        }
+        return { toolResponse: "ok", data: null };
+      },
+    });
+
+    const assistantIds: string[] = [];
+    const session = await createSession({
+      agentName: "TestAgent",
+      thread: { mode: "new", threadId: "thread-1" },
+      runAgent: createScriptedRunAgent(
+        [
+          {
+            message: "parallel",
+            toolCalls: [
+              { id: "tc-sibling", name: "Sibling", args: {} },
+              { id: "tc-rewind", name: "Rewind", args: {} },
+            ],
+          },
+          { message: "done", toolCalls: [] },
+        ],
+        assistantIds
+      ),
+      threadOps: ops,
+      tools: { Rewind: rewindTool, Sibling: siblingTool },
+      buildContextMessage: () => "go",
+    });
+
+    const stateManager = createAgentStateManager({
+      initialState: { systemPrompt: "test" },
+    });
+
+    const result = await session.runSession({ stateManager });
+
+    expect(result.exitReason).toBe("completed");
+
+    // No explicit truncate from the session — the invoker will do it
+    // on entry using the reused assistantMessageId.
+    const truncateOps = log.filter((l) => l.op === "truncateThread");
+    expect(truncateOps).toHaveLength(0);
+
+    // The rewound turn and its retry share one assistantMessageId; the
+    // final `done` turn gets a fresh one.
+    expect(assistantIds).toHaveLength(2);
+    expect(assistantIds[0]).toBe(assistantIds[1]);
+
+    // Rewinding tool never appends its own result.
+    const rewindResultAppends = log.filter((l) => {
+      if (l.op !== "appendToolResult") return false;
+      const config = l.args[1] as ToolResultConfig;
+      return config.toolCallId === "tc-rewind";
+    });
+    expect(rewindResultAppends).toHaveLength(0);
+
+    // Two assistant messages expected: one from the rewound turn, one
+    // from the successful retry — sharing the same id.
+    const agentAppends = log.filter((l) => l.op === "appendAgentMessage");
+    expect(agentAppends).toHaveLength(2);
+    expect(agentAppends[0]?.args[1]).toBe(agentAppends[1]?.args[1]);
+  });
+
+  it("does not rewind when the rewinding tool is no longer present after retry", async () => {
+    const { ops, log } = createMockThreadOps();
+
+    let attempts = 0;
+    const rewindOnce = defineTool({
+      name: "RewindOnce" as const,
+      description: "rewinds once",
+      schema: z.object({}),
+      handler: async () => {
+        attempts += 1;
+        if (attempts === 1) {
+          return { toolResponse: "ignored", data: null, rewind: true };
+        }
+        return { toolResponse: "ok", data: null };
+      },
+    });
+
+    const assistantIds: string[] = [];
+    const session = await createSession({
+      agentName: "TestAgent",
+      thread: { mode: "new", threadId: "thread-1" },
+      maxTurns: 5,
+      runAgent: createScriptedRunAgent(
+        [
+          {
+            message: "call-1",
+            toolCalls: [{ id: "tc-1", name: "RewindOnce", args: {} }],
+          },
+          {
+            message: "call-2",
+            toolCalls: [{ id: "tc-2", name: "RewindOnce", args: {} }],
+          },
+          { message: "done", toolCalls: [] },
+        ],
+        assistantIds
+      ),
+      threadOps: ops,
+      tools: { RewindOnce: rewindOnce },
+      buildContextMessage: () => "go",
+    });
+
+    const stateManager = createAgentStateManager({
+      initialState: { systemPrompt: "test" },
+    });
+
+    const result = await session.runSession({ stateManager });
+
+    expect(result.exitReason).toBe("completed");
+    expect(result.finalMessage).toBe("done");
+    // Each rewind still consumes a turn from the `maxTurns` budget:
+    // turn 1 (rewound) + turn 2 (successful tool call) + turn 3 (done) = 3.
+    expect(result.usage.turns).toBe(3);
+    expect(attempts).toBe(2);
+
+    const truncateOps = log.filter((l) => l.op === "truncateThread");
+    expect(truncateOps).toHaveLength(0);
+
+    // Turn 1 rewound → call 1 & 2 share an id, call 3 fresh.
+    expect(assistantIds).toHaveLength(3);
+    expect(assistantIds[0]).toBe(assistantIds[1]);
+    expect(assistantIds[1]).not.toBe(assistantIds[2]);
+  });
+
+  it("bails out with max_turns when a tool keeps requesting rewind", async () => {
+    const { ops, log } = createMockThreadOps();
+
+    let attempts = 0;
+    const alwaysRewind = defineTool({
+      name: "AlwaysRewind" as const,
+      description: "always rewinds",
+      schema: z.object({}),
+      handler: async () => {
+        attempts += 1;
+        return { toolResponse: "ignored", data: null, rewind: true };
+      },
+    });
+
+    const assistantIds: string[] = [];
+    const session = await createSession({
+      agentName: "TestAgent",
+      thread: { mode: "new", threadId: "thread-1" },
+      maxTurns: 3,
+      runAgent: createScriptedRunAgent(
+        [
+          { message: "t1", toolCalls: [{ id: "tc-1", name: "AlwaysRewind", args: {} }] },
+          { message: "t2", toolCalls: [{ id: "tc-2", name: "AlwaysRewind", args: {} }] },
+          { message: "t3", toolCalls: [{ id: "tc-3", name: "AlwaysRewind", args: {} }] },
+          { message: "t4", toolCalls: [{ id: "tc-4", name: "AlwaysRewind", args: {} }] },
+        ],
+        assistantIds
+      ),
+      threadOps: ops,
+      tools: { AlwaysRewind: alwaysRewind },
+      buildContextMessage: () => "go",
+    });
+
+    const stateManager = createAgentStateManager({
+      initialState: { systemPrompt: "test" },
+    });
+
+    const result = await session.runSession({ stateManager });
+
+    expect(result.exitReason).toBe("max_turns");
+    expect(result.usage.turns).toBe(3);
+    expect(attempts).toBe(3);
+
+    // Session does not issue explicit truncates; invoker-side
+    // truncation isn't visible here because runAgent is mocked.
+    const truncateOps = log.filter((l) => l.op === "truncateThread");
+    expect(truncateOps).toHaveLength(0);
+
+    // Every attempt reuses the same assistantMessageId — the LLM call
+    // truncates-from-id on each replay.
+    expect(assistantIds).toHaveLength(3);
+    expect(new Set(assistantIds).size).toBe(1);
+  });
 });

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

@@ -2,6 +2,7 @@ import { describe, expect, it, vi, beforeEach } from "vitest";
 import { z } from "zod";
 import type { ToolResultConfig, TokenUsage } from "../types";
 import type { ThreadOps } from "./types";
+import type { PersistedThreadState } from "../state/types";
 import type { RunAgentActivity } from "../model/types";
 import type { RawToolCall } from "../tool-router/types";
 import type { SandboxOps } from "../sandbox/types";
@@ -33,6 +34,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 +57,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: () => {},
@@ -85,6 +98,7 @@ function toActivityInterface(raw: ThreadOps): ActivityInterfaceFor<ThreadOps> {
 
 function createMockThreadOps() {
   const log: { op: string; args: unknown[] }[] = [];
+  const stateStore = new Map<string, PersistedThreadState>();
 
   const ops = toActivityInterface({
     initializeThread: async (threadId) => {
@@ -104,10 +118,23 @@ function createMockThreadOps() {
     },
     forkThread: async (source, target) => {
       log.push({ op: "forkThread", args: [source, target] });
+      const src = stateStore.get(source);
+      if (src) stateStore.set(target, src);
+    },
+    truncateThread: async (threadId, messageId) => {
+      log.push({ op: "truncateThread", args: [threadId, messageId] });
+    },
+    loadThreadState: async (threadId) => {
+      log.push({ op: "loadThreadState", args: [threadId] });
+      return stateStore.get(threadId) ?? null;
+    },
+    saveThreadState: async (threadId, state) => {
+      log.push({ op: "saveThreadState", args: [threadId, state] });
+      stateStore.set(threadId, state);
     },
   });
 
-  return { ops, log };
+  return { ops, log, stateStore };
 }
 
 type TurnScript = {
@@ -123,7 +150,11 @@ function createScriptedRunAgent(
   return async () => {
     const turn = turns[call++];
     if (!turn) {
-      return { message: "done", rawToolCalls: [], usage: undefined };
+      return {
+        message: "done",
+        rawToolCalls: [],
+        usage: undefined,
+      };
     }
     return {
       message: turn.message,
@@ -1088,4 +1119,163 @@ describe("createSession integration", () => {
     ]);
     expect(sandboxLog).not.toContain("create");
   });
+
+  // --- Persistent thread state ---
+
+  it("saves tasks + custom state to the thread store on session exit", async () => {
+    const { ops, log, stateStore } = createMockThreadOps();
+
+    const writeTasks = defineTool({
+      name: "WriteTasks" as const,
+      description: "create tasks via state manager",
+      schema: z.object({}),
+      handler: async (
+        _args: Record<string, never>,
+        _ctx: RouterContext
+      ): Promise<ToolHandlerResponse<null>> => ({
+        toolResponse: "ok",
+        data: null,
+      }),
+    });
+
+    const session = await createSession({
+      agentName: "TestAgent",
+      thread: { mode: "new", threadId: "thread-save" },
+      runAgent: createScriptedRunAgent([
+        {
+          message: "doing work",
+          toolCalls: [{ id: "tc-1", name: "WriteTasks", args: {} }],
+        },
+        { message: "done", toolCalls: [] },
+      ]),
+      threadOps: ops,
+      tools: { WriteTasks: writeTasks },
+      buildContextMessage: () => "go",
+    });
+
+    const stateManager = createAgentStateManager<{ note: string }>({
+      initialState: { systemPrompt: "test", note: "hello" },
+    });
+
+    stateManager.setTask({
+      id: "task-A",
+      subject: "A",
+      description: "A",
+      activeForm: "doing A",
+      status: "in_progress",
+      metadata: { priority: "high" },
+      blockedBy: [],
+      blocks: [],
+    });
+
+    const result = await session.runSession({ stateManager });
+    expect(result.exitReason).toBe("completed");
+
+    const saves = log.filter((l) => l.op === "saveThreadState");
+    expect(saves).toHaveLength(1);
+    const saved = stateStore.get("thread-save");
+    expect(saved).toBeDefined();
+    expect(saved?.tasks).toHaveLength(1);
+    if (saved) {
+      expect(at(saved.tasks, 0)[0]).toBe("task-A");
+    }
+    expect(saved?.custom).toEqual({ note: "hello" });
+  });
+
+  it("rehydrates tasks + custom state on continue before the agent loop runs", async () => {
+    const { ops, stateStore } = createMockThreadOps();
+
+    stateStore.set("thread-cont", {
+      tasks: [
+        [
+          "task-restored",
+          {
+            id: "task-restored",
+            subject: "restored",
+            description: "restored",
+            activeForm: "restoring",
+            status: "pending",
+            metadata: {},
+            blockedBy: [],
+            blocks: [],
+          },
+        ],
+      ],
+      custom: { label: "from-prior-run" },
+    });
+
+    type State = { label: string };
+    let observedTasksBeforeFirstTurn: string[] = [];
+    let observedLabelBeforeFirstTurn: string | undefined;
+
+    const session = await createSession({
+      agentName: "TestAgent",
+      thread: { mode: "continue", threadId: "thread-cont" },
+      runAgent: async () => {
+        observedTasksBeforeFirstTurn = stateManager.getTasks().map((t) => t.id);
+        observedLabelBeforeFirstTurn = stateManager.get("label");
+        return { message: "done", rawToolCalls: [], usage: undefined };
+      },
+      threadOps: ops,
+      buildContextMessage: () => "continue please",
+    });
+
+    const stateManager = createAgentStateManager<State>({
+      initialState: { systemPrompt: "test", label: "initial" },
+    });
+
+    await session.runSession({ stateManager });
+
+    expect(observedTasksBeforeFirstTurn).toEqual(["task-restored"]);
+    expect(observedLabelBeforeFirstTurn).toBe("from-prior-run");
+  });
+
+  it("fork copies the source thread's state slice into the new thread", async () => {
+    const { ops, log, stateStore } = createMockThreadOps();
+
+    stateStore.set("source-thread", {
+      tasks: [
+        [
+          "task-src",
+          {
+            id: "task-src",
+            subject: "src",
+            description: "src",
+            activeForm: "src",
+            status: "completed",
+            metadata: {},
+            blockedBy: [],
+            blocks: [],
+          },
+        ],
+      ],
+      custom: { counter: 3 },
+    });
+
+    const session = await createSession({
+      agentName: "TestAgent",
+      thread: { mode: "fork", threadId: "source-thread" },
+      runAgent: createScriptedRunAgent([{ message: "done", toolCalls: [] }]),
+      threadOps: ops,
+      buildContextMessage: () => "continue",
+    });
+
+    type State = { counter: number };
+    const stateManager = createAgentStateManager<State>({
+      initialState: { systemPrompt: "test", counter: 0 },
+    });
+
+    const result = await session.runSession({ stateManager });
+    expect(result.exitReason).toBe("completed");
+
+    const forkOps = log.filter((l) => l.op === "forkThread");
+    expect(forkOps).toHaveLength(1);
+    expect(at(forkOps, 0).args[0]).toBe("source-thread");
+
+    expect(stateManager.getTask("task-src")).toBeDefined();
+    expect(stateManager.get("counter")).toBe(3);
+
+    const newThreadSlice = stateStore.get(result.threadId);
+    expect(newThreadSlice?.tasks).toHaveLength(1);
+  });
 });