zeitlich 0.2.38 → 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";
@@ -73,7 +74,6 @@ type TurnScript = {
   message: unknown;
   toolCalls: RawToolCall[];
   usage?: TokenUsage;
-  threadLengthAtCall?: number;
 };
 
 /**
@@ -96,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] });
@@ -114,33 +115,43 @@ 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, length) => {
-      log.push({ op: "truncateThread", args: [threadId, length] });
+    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,
-        threadLengthAtCall: 0,
       };
     }
     return {
       message: turn.message,
       rawToolCalls: turn.toolCalls,
       usage: turn.usage,
-      threadLengthAtCall: turn.threadLengthAtCall ?? 0,
     };
   };
 }
@@ -797,8 +808,15 @@ describe("createSession edge cases", () => {
       forkThread: async (source, target) => {
         log.push({ op: "forkThread", args: [source, target] });
       },
-      truncateThread: async (threadId, length) => {
-        log.push({ op: "truncateThread", args: [threadId, length] });
+      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] });
       },
     });
 
@@ -1686,6 +1704,14 @@ describe("createSession edge cases", () => {
   });
 
   // --- 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();
@@ -1708,20 +1734,24 @@ describe("createSession edge cases", () => {
       },
     });
 
+    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: [] },
-      ]),
+      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",
@@ -1737,8 +1767,16 @@ describe("createSession edge cases", () => {
     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(1);
+    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;
@@ -1749,9 +1787,15 @@ describe("createSession edge cases", () => {
 
     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("truncates the thread back to the pre-assistant state so sibling tool results are dropped on rewind", async () => {
+  it("reuses the assistantMessageId on rewind even with sibling tool calls", async () => {
     const { ops, log } = createMockThreadOps();
 
     let rewindFired = false;
@@ -1776,22 +1820,23 @@ describe("createSession edge cases", () => {
       },
     });
 
+    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: {} },
-          ],
-          // Invoker reports 2 stored messages (system + human) at the
-          // moment the LLM was called.
-          threadLengthAtCall: 2,
-        },
-        { message: "done", toolCalls: [], threadLengthAtCall: 2 },
-      ]),
+      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",
@@ -1805,13 +1850,15 @@ describe("createSession edge cases", () => {
 
     expect(result.exitReason).toBe("completed");
 
-    // Exactly one truncate fired — back to the pre-assistant-message
-    // length that runAgent reported.
+    // 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(1);
-    const truncateOp = truncateOps[0];
-    if (!truncateOp) throw new Error("expected truncate op");
-    expect(truncateOp.args[1]).toBe(2);
+    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) => {
@@ -1821,10 +1868,11 @@ describe("createSession edge cases", () => {
     });
     expect(rewindResultAppends).toHaveLength(0);
 
-    // Two assistant messages expected: one from the rewound turn, one from
-    // the successful retry.
+    // 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 () => {
@@ -1844,21 +1892,25 @@ describe("createSession edge cases", () => {
       },
     });
 
+    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: [] },
-      ]),
+      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",
@@ -1878,7 +1930,12 @@ describe("createSession edge cases", () => {
     expect(attempts).toBe(2);
 
     const truncateOps = log.filter((l) => l.op === "truncateThread");
-    expect(truncateOps).toHaveLength(1);
+    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 () => {
@@ -1895,16 +1952,20 @@ describe("createSession edge cases", () => {
       },
     });
 
+    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: {} }] },
-      ]),
+      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",
@@ -1920,7 +1981,14 @@ describe("createSession edge cases", () => {
     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(3);
+    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";
@@ -97,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) => {
@@ -116,13 +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, length) => {
-      log.push({ op: "truncateThread", args: [threadId, length] });
+    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 = {
@@ -142,14 +154,12 @@ function createScriptedRunAgent(
         message: "done",
         rawToolCalls: [],
         usage: undefined,
-        threadLengthAtCall: 0,
       };
     }
     return {
       message: turn.message,
       rawToolCalls: turn.toolCalls,
       usage: turn.usage,
-      threadLengthAtCall: 0,
     };
   };
 }
@@ -1109,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);
+  });
 });