zeitlich 0.2.47 → 0.2.49

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "zeitlich",
-  "version": "0.2.47",
+  "version": "0.2.49",
   "description": "[EXPERIMENTAL] An opinionated AI agent implementation for Temporal",
   "main": "./dist/index.js",
   "module": "./dist/index.js",
@@ -208,11 +208,11 @@
     "@e2b/code-interpreter": "^2.3.3",
     "@eslint/js": "^10.0.1",
     "@google/genai": "^1.44.0",
-    "@langchain/core": "^1.1.30",
-    "@temporalio/common": "^1.17.0",
-    "@temporalio/envconfig": "^1.17.0",
-    "@temporalio/worker": "^1.17.0",
-    "@temporalio/workflow": "^1.17.0",
+    "@langchain/core": "^1.1.48",
+    "@temporalio/common": "^1.17.2",
+    "@temporalio/envconfig": "^1.17.2",
+    "@temporalio/worker": "^1.17.2",
+    "@temporalio/workflow": "^1.17.2",
     "@types/node": "^25.3.3",
     "eslint": "^10.0.2",
     "husky": "^9.1.7",

package/src/lib/lifecycle.ts

@@ -11,11 +11,23 @@
  *   continue there. When the adapter has `onForkPrepareThread` and/or
  *   `onForkTransform` hooks configured, they are applied once to the forked
  *   thread before the session starts.
+ *
+ *   The optional `truncateAfterFork.fromMessageId` directs the session to
+ *   call `truncateThread` on the freshly forked thread immediately after
+ *   the fork, dropping that message and everything after. Used by
+ *   subagents that fork their parent's thread mid-tool-call to strip the
+ *   orphan assistant `tool_use` block (the one whose `tool_result` will
+ *   never arrive in the child's thread) so the first model call doesn't
+ *   reject on an unmatched tool-use/tool-result pair.
  */
 export type ThreadInit =
   | { mode: "new"; threadId?: string }
   | { mode: "continue"; threadId: string }
-  | { mode: "fork"; threadId: string };
+  | {
+      mode: "fork";
+      threadId: string;
+      truncateAfterFork?: { fromMessageId: string };
+    };
 
 // ============================================================================
 // Sandbox lifecycle

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

@@ -660,6 +660,50 @@ describe("createSession edge cases", () => {
     const forkOp = forkOps[0];
     if (!forkOp) throw new Error("expected fork op");
     expect(forkOp.args[0]).toBe("original-thread");
+
+    // No truncateAfterFork directive ⇒ no truncate call.
+    const truncateOps = log.filter((l) => l.op === "truncateThread");
+    expect(truncateOps).toHaveLength(0);
+  });
+
+  it("fork thread mode truncates the forked thread at fromMessageId when truncateAfterFork is set", async () => {
+    const { ops, log } = createMockThreadOps();
+
+    const session = await createSession({
+      agentName: "TestAgent",
+      thread: {
+        mode: "fork",
+        threadId: "parent-thread",
+        truncateAfterFork: { fromMessageId: "parent-asst-msg-1" },
+      },
+      runAgent: createScriptedRunAgent([
+        { message: "forked & continued", toolCalls: [] },
+      ]),
+      threadOps: ops,
+      buildContextMessage: () => "continue",
+    });
+
+    const stateManager = createAgentStateManager({
+      initialState: { systemPrompt: "test" },
+    });
+
+    const result = await session.runSession({ stateManager });
+
+    expect(result.exitReason).toBe("completed");
+    const forkedThreadId = result.threadId;
+    expect(forkedThreadId).not.toBe("parent-thread");
+
+    // Order matters: fork must happen before truncate, otherwise the
+    // truncate would no-op against an empty thread.
+    const forkIdx = log.findIndex((l) => l.op === "forkThread");
+    const truncIdx = log.findIndex((l) => l.op === "truncateThread");
+    expect(forkIdx).toBeGreaterThanOrEqual(0);
+    expect(truncIdx).toBeGreaterThan(forkIdx);
+
+    const truncOp = log[truncIdx];
+    if (!truncOp) throw new Error("expected truncate op");
+    expect(truncOp.args[0]).toBe(forkedThreadId);
+    expect(truncOp.args[1]).toBe("parent-asst-msg-1");
   });
 
   // --- maxTurns of 1 ---

package/src/lib/session/session.ts

@@ -196,6 +196,7 @@ export async function createSession<
     appendSystemMessage,
     appendAgentMessage,
     forkThread,
+    truncateThread,
     loadThreadState,
     saveThreadState,
     hydrateThread,
@@ -399,6 +400,17 @@ export async function createSession<
         // is already hot or when no cold tier is wired.
         await hydrateThread(sourceThreadId, threadKey);
         await forkThread(sourceThreadId, threadId, threadKey);
+        // If the caller asked to drop the tail of the forked thread
+        // (e.g. subagent forking its parent mid-tool-call needs to
+        // strip the orphan assistant `tool_use`), do it now — before
+        // any rehydration / state load so the truncated thread is
+        // what subsequent reads see.
+        const truncate = (
+          threadInit as { mode: "fork"; truncateAfterFork?: { fromMessageId: string } }
+        ).truncateAfterFork;
+        if (truncate?.fromMessageId) {
+          await truncateThread(threadId, truncate.fromMessageId, threadKey);
+        }
         const forkedSlice = await loadThreadState(threadId, threadKey);
         if (forkedSlice) rehydrateFromSlice(forkedSlice);
       } else if (threadMode === "continue") {
@@ -570,6 +582,20 @@ export async function createSession<
             {
               turn: currentTurn,
               ...(sandboxId !== undefined && { sandboxId }),
+              ...(assistantId !== undefined && {
+                assistantMessageId: assistantId,
+              }),
+              // Hand handlers a way to persist the parent's slice
+              // mid-loop (subagents that fork or continue the parent's
+              // thread need this — otherwise the child loads a stale
+              // snapshot from the prior session, since `saveThreadState`
+              // would otherwise only run in the `finally` below).
+              persistThreadState: () =>
+                saveThreadState(
+                  threadId,
+                  stateManager.getPersistedSlice(),
+                  threadKey
+                ),
             }
           );
 

package/src/lib/subagent/handler.ts

@@ -244,18 +244,44 @@ export function createSubagentHandler<
     // `thread: "fork" | "continue"` — `thread: "new"` always starts
     // fresh regardless of the source.
     const newThreadSource = config.newThreadSource ?? "new";
+    const usingParentFallback =
+      allowsContinuation &&
+      !args.threadId &&
+      newThreadSource === "from-parent";
     const continuationThreadId = !allowsContinuation
       ? undefined
-      : (args.threadId ??
-        (newThreadSource === "from-parent" ? context.threadId : undefined));
+      : (args.threadId ?? (usingParentFallback ? context.threadId : undefined));
 
     // --- Build thread init ---
     let thread: ThreadInit | undefined;
     if (continuationThreadId) {
-      thread = {
-        mode: threadMode as "fork" | "continue",
-        threadId: continuationThreadId,
-      };
+      if (threadMode === "fork") {
+        // When falling back to the parent's thread for a fork, the
+        // parent is mid-tool-call: its assistant message containing the
+        // `Subagent` tool_use has already been persisted, but its
+        // matching tool_result hasn't. Forking verbatim would leave an
+        // orphan tool_use at the tail of the child thread, which most
+        // model APIs reject on the very next call. Have the child's
+        // session truncate that assistant message (and anything after)
+        // immediately after the fork so the first model call sees a
+        // well-formed history.
+        thread = {
+          mode: "fork",
+          threadId: continuationThreadId,
+          ...(usingParentFallback && context.assistantMessageId
+            ? {
+                truncateAfterFork: {
+                  fromMessageId: context.assistantMessageId,
+                },
+              }
+            : {}),
+        };
+      } else {
+        thread = {
+          mode: "continue",
+          threadId: continuationThreadId,
+        };
+      }
     }
 
     // --- Build sandbox init ---
@@ -413,6 +439,29 @@ export function createSubagentHandler<
       snapshotBaseCreatorAgent.set(childWorkflowId, config.agentName);
     }
 
+    // The parent's `PersistedThreadState` slice (`tasks` + custom
+    // state) only lands in storage in the session's `finally`. When
+    // the child reads from the parent's thread (`from-parent`
+    // fallback or an explicit args.threadId pointing at the parent),
+    // that `loadThreadState` would otherwise see the prior session's
+    // snapshot. Flush the live slice now via the session-supplied
+    // callback so the child sees the parent's current state.
+    if (
+      continuationThreadId &&
+      continuationThreadId === context.threadId &&
+      context.persistThreadState
+    ) {
+      try {
+        await context.persistThreadState();
+      } catch (err) {
+        log.warn("failed to persist parent thread state for subagent", {
+          subagent: config.agentName,
+          childWorkflowId,
+          error: err instanceof Error ? err.message : String(err),
+        });
+      }
+    }
+
     log.info("subagent spawned", {
       subagent: config.agentName,
       childWorkflowId,

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

@@ -653,6 +653,40 @@ describe("createSubagentHandler", () => {
 
     await handler(
       { subagent: "parent-fork", description: "test", prompt: "test" },
+      {
+        threadId: "parent-t",
+        toolCallId: "tc",
+        toolName: "Subagent",
+        assistantMessageId: "parent-asst-msg-1",
+      }
+    );
+
+    const lastCall = execMock.mock.calls[execMock.mock.calls.length - 1];
+    if (!lastCall) throw new Error("expected executeChild call");
+    const workflowInput = lastCall[1].args[1] as SubagentWorkflowInput;
+    expect(workflowInput.thread).toEqual({
+      mode: "fork",
+      threadId: "parent-t",
+      truncateAfterFork: { fromMessageId: "parent-asst-msg-1" },
+    });
+  });
+
+  it("omits truncateAfterFork on parent-fallback fork when assistantMessageId is absent from context", async () => {
+    const { executeChild } = await import("@temporalio/workflow");
+    const execMock = executeChild as ReturnType<typeof vi.fn>;
+
+    const subagent: SubagentConfig = {
+      agentName: "parent-fork-no-asst",
+      description: "Forks parent thread by default",
+      workflow: mockWorkflow(),
+      thread: "fork",
+      newThreadSource: "from-parent",
+    };
+
+    const { handler } = createSubagentHandler([subagent]);
+
+    await handler(
+      { subagent: "parent-fork-no-asst", description: "test", prompt: "test" },
       { threadId: "parent-t", toolCallId: "tc", toolName: "Subagent" }
     );
 
@@ -693,7 +727,7 @@ describe("createSubagentHandler", () => {
     });
   });
 
-  it("prefers args.threadId over the parent source when both are available", async () => {
+  it("prefers args.threadId over the parent source when both are available and skips truncateAfterFork", async () => {
     const { executeChild } = await import("@temporalio/workflow");
     const execMock = executeChild as ReturnType<typeof vi.fn>;
 
@@ -714,7 +748,12 @@ describe("createSubagentHandler", () => {
         prompt: "test",
         threadId: "explicit-prev",
       },
-      { threadId: "parent-t", toolCallId: "tc", toolName: "Subagent" }
+      {
+        threadId: "parent-t",
+        toolCallId: "tc",
+        toolName: "Subagent",
+        assistantMessageId: "parent-asst-msg-1",
+      }
     );
 
     const lastCall = execMock.mock.calls[execMock.mock.calls.length - 1];
@@ -774,6 +813,204 @@ describe("createSubagentHandler", () => {
     expect(workflowInput.thread).toBeUndefined();
   });
 
+  // --- persistThreadState: parent slice flush before child loads ---
+
+  it("calls persistThreadState before executeChild when forking parent's thread", async () => {
+    const { executeChild } = await import("@temporalio/workflow");
+    const execMock = executeChild as ReturnType<typeof vi.fn>;
+    const persistThreadState = vi.fn(async () => undefined);
+
+    const subagent: SubagentConfig = {
+      agentName: "parent-fork-persist",
+      description: "Forks parent thread",
+      workflow: mockWorkflow(),
+      thread: "fork",
+      newThreadSource: "from-parent",
+    };
+
+    const { handler } = createSubagentHandler([subagent]);
+
+    await handler(
+      {
+        subagent: "parent-fork-persist",
+        description: "test",
+        prompt: "test",
+      },
+      {
+        threadId: "parent-t",
+        toolCallId: "tc",
+        toolName: "Subagent",
+        persistThreadState,
+      }
+    );
+
+    expect(persistThreadState).toHaveBeenCalledTimes(1);
+    const lastExecOrder =
+      execMock.mock.invocationCallOrder[
+        execMock.mock.invocationCallOrder.length - 1
+      ] ?? Infinity;
+    expect(persistThreadState.mock.invocationCallOrder[0]).toBeLessThan(
+      lastExecOrder
+    );
+  });
+
+  it("calls persistThreadState before executeChild when continuing parent's thread", async () => {
+    const { executeChild } = await import("@temporalio/workflow");
+    const execMock = executeChild as ReturnType<typeof vi.fn>;
+    const persistThreadState = vi.fn(async () => undefined);
+
+    const subagent: SubagentConfig = {
+      agentName: "parent-continue-persist",
+      description: "Continues parent thread",
+      workflow: mockWorkflow(),
+      thread: "continue",
+      newThreadSource: "from-parent",
+    };
+
+    const { handler } = createSubagentHandler([subagent]);
+
+    await handler(
+      {
+        subagent: "parent-continue-persist",
+        description: "test",
+        prompt: "test",
+      },
+      {
+        threadId: "parent-t",
+        toolCallId: "tc",
+        toolName: "Subagent",
+        persistThreadState,
+      }
+    );
+
+    expect(persistThreadState).toHaveBeenCalledTimes(1);
+    const lastExecOrder =
+      execMock.mock.invocationCallOrder[
+        execMock.mock.invocationCallOrder.length - 1
+      ] ?? Infinity;
+    expect(persistThreadState.mock.invocationCallOrder[0]).toBeLessThan(
+      lastExecOrder
+    );
+  });
+
+  it("calls persistThreadState when args.threadId points at the parent's thread", async () => {
+    const persistThreadState = vi.fn(async () => undefined);
+
+    const subagent: SubagentConfig = {
+      agentName: "explicit-parent",
+      description: "Explicit parent threadId",
+      workflow: mockWorkflow(),
+      thread: "continue",
+    };
+
+    const { handler } = createSubagentHandler([subagent]);
+
+    await handler(
+      {
+        subagent: "explicit-parent",
+        description: "test",
+        prompt: "test",
+        threadId: "parent-t",
+      },
+      {
+        threadId: "parent-t",
+        toolCallId: "tc",
+        toolName: "Subagent",
+        persistThreadState,
+      }
+    );
+
+    expect(persistThreadState).toHaveBeenCalledTimes(1);
+  });
+
+  it("does not call persistThreadState when child thread is independent of the parent", async () => {
+    const persistThreadState = vi.fn(async () => undefined);
+
+    const subagent: SubagentConfig = {
+      agentName: "independent-thread",
+      description: "Continues a sibling thread",
+      workflow: mockWorkflow(),
+      thread: "continue",
+    };
+
+    const { handler } = createSubagentHandler([subagent]);
+
+    await handler(
+      {
+        subagent: "independent-thread",
+        description: "test",
+        prompt: "test",
+        threadId: "sibling-thread",
+      },
+      {
+        threadId: "parent-t",
+        toolCallId: "tc",
+        toolName: "Subagent",
+        persistThreadState,
+      }
+    );
+
+    expect(persistThreadState).not.toHaveBeenCalled();
+  });
+
+  it("does not call persistThreadState when the child starts a fresh thread", async () => {
+    const persistThreadState = vi.fn(async () => undefined);
+
+    const subagent: SubagentConfig = {
+      agentName: "fresh-thread",
+      description: "Always starts fresh",
+      workflow: mockWorkflow(),
+    };
+
+    const { handler } = createSubagentHandler([subagent]);
+
+    await handler(
+      { subagent: "fresh-thread", description: "test", prompt: "test" },
+      {
+        threadId: "parent-t",
+        toolCallId: "tc",
+        toolName: "Subagent",
+        persistThreadState,
+      }
+    );
+
+    expect(persistThreadState).not.toHaveBeenCalled();
+  });
+
+  it("still spawns the child when persistThreadState throws", async () => {
+    const { executeChild } = await import("@temporalio/workflow");
+    const execMock = executeChild as ReturnType<typeof vi.fn>;
+    const persistThreadState = vi.fn(async () => {
+      throw new Error("redis down");
+    });
+    const callsBefore = execMock.mock.calls.length;
+
+    const subagent: SubagentConfig = {
+      agentName: "persist-fails",
+      description: "Persist failure should not block child",
+      workflow: mockWorkflow(),
+      thread: "fork",
+      newThreadSource: "from-parent",
+    };
+
+    const { handler } = createSubagentHandler([subagent]);
+
+    await expect(
+      handler(
+        { subagent: "persist-fails", description: "test", prompt: "test" },
+        {
+          threadId: "parent-t",
+          toolCallId: "tc",
+          toolName: "Subagent",
+          persistThreadState,
+        }
+      )
+    ).resolves.toBeDefined();
+
+    expect(persistThreadState).toHaveBeenCalledTimes(1);
+    expect(execMock.mock.calls.length).toBe(callsBefore + 1);
+  });
+
   // --- Sandbox continuation ---
 
   it("does not pass sandbox when thread is fork (own sandbox)", async () => {

package/src/lib/tool-router/router-edge-cases.integration.test.ts

@@ -134,6 +134,42 @@ describe("createToolRouter edge cases", () => {
     expect(appendSpy.calls).toHaveLength(0);
   });
 
+  // --- assistantMessageId propagation into RouterContext ---
+
+  it("forwards assistantMessageId from ProcessToolCallsContext into RouterContext", async () => {
+    let capturedAssistantMessageId: string | undefined;
+    const captureTool = defineTool({
+      name: "Capture" as const,
+      description: "captures router context",
+      schema: z.object({}),
+      handler: async (
+        _args: Record<string, never>,
+        ctx: { assistantMessageId?: string }
+      ): Promise<ToolHandlerResponse<null>> => {
+        capturedAssistantMessageId = ctx.assistantMessageId;
+        return { toolResponse: "ok", data: null };
+      },
+    });
+
+    const router = createToolRouter({
+      tools: { Capture: captureTool } as const,
+      threadId: "t-1",
+      appendToolResult: appendSpy.fn,
+    });
+
+    const parsed = router.parseToolCall({
+      id: "tc-1",
+      name: "Capture",
+      args: {},
+    });
+    await router.processToolCalls([parsed], {
+      turn: 1,
+      assistantMessageId: "asst-msg-42",
+    });
+
+    expect(capturedAssistantMessageId).toBe("asst-msg-42");
+  });
+
   // --- Both global and per-tool pre-hooks run in order ---
 
   it("global pre-hook runs before per-tool pre-hook", async () => {

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

@@ -220,7 +220,9 @@ export function createToolRouter<T extends ToolMap>(
     toolCall: ParsedToolCallUnion<T>,
     turn: number,
     sandboxId?: string,
-    onRewindRequested?: (signal: RewindSignal) => void
+    onRewindRequested?: (signal: RewindSignal) => void,
+    assistantMessageId?: string,
+    persistThreadState?: () => Promise<void>
   ): Promise<ProcessedToolCall> {
     const startTime = Date.now();
     const tool = toolMap.get(toolCall.name);
@@ -263,6 +265,8 @@ export function createToolRouter<T extends ToolMap>(
           toolCallId: toolCall.id,
           toolName: toolCall.name,
           ...(sandboxId !== undefined && { sandboxId }),
+          ...(assistantMessageId !== undefined && { assistantMessageId }),
+          ...(persistThreadState !== undefined && { persistThreadState }),
         };
         const response = await tool.handler(
           effectiveArgs as Parameters<typeof tool.handler>[0],
@@ -419,6 +423,8 @@ export function createToolRouter<T extends ToolMap>(
 
       const turn = context?.turn ?? 0;
       const sandboxId = context?.sandboxId;
+      const assistantMessageId = context?.assistantMessageId;
+      const persistThreadState = context?.persistThreadState;
 
       let rewindSignal: RewindSignal | undefined;
 
@@ -435,7 +441,14 @@ export function createToolRouter<T extends ToolMap>(
         const outcomes = await scope.run(async () =>
           Promise.allSettled(
             toolCalls.map((tc) =>
-              processToolCall(tc, turn, sandboxId, onRewindRequested)
+              processToolCall(
+                tc,
+                turn,
+                sandboxId,
+                onRewindRequested,
+                assistantMessageId,
+                persistThreadState
+              )
             )
           )
         );
@@ -457,7 +470,14 @@ export function createToolRouter<T extends ToolMap>(
 
       const results: ToolCallResultUnion<TResults>[] = [];
       for (const toolCall of toolCalls) {
-        const outcome = await processToolCall(toolCall, turn, sandboxId);
+        const outcome = await processToolCall(
+          toolCall,
+          turn,
+          sandboxId,
+          undefined,
+          assistantMessageId,
+          persistThreadState
+        );
         if (outcome.kind === "rewind") {
           rewindSignal = outcome.signal;
           break;
@@ -492,6 +512,12 @@ export function createToolRouter<T extends ToolMap>(
           ...(context?.sandboxId !== undefined && {
             sandboxId: context.sandboxId,
           }),
+          ...(context?.assistantMessageId !== undefined && {
+            assistantMessageId: context.assistantMessageId,
+          }),
+          ...(context?.persistThreadState !== undefined && {
+            persistThreadState: context.persistThreadState,
+          }),
         };
         const response = await handler(
           toolCall.args as ToolArgs<T, TName>,

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

@@ -178,6 +178,32 @@ export interface RouterContext {
   toolCallId: string;
   toolName: string;
   sandboxId?: string;
+  /**
+   * Id of the assistant message that issued this tool call (the message
+   * the session passed as `assistantMessageId` into `runAgent`). Present
+   * for any tool call processed through `processToolCalls` from a
+   * session; may be absent when the router is driven manually (e.g.
+   * tests, custom orchestrators).
+   *
+   * Subagent handlers that fork the parent's thread mid-call use this
+   * to truncate the orphan trailing assistant message from the forked
+   * thread so the child's first model call sees a well-formed history.
+   */
+  assistantMessageId?: string;
+  /**
+   * Persist the parent session's current `PersistedThreadState` slice
+   * (tasks + custom state) to the durable thread store. Wired up by
+   * the session — absent for manually-driven routers (tests, custom
+   * orchestrators).
+   *
+   * Subagent handlers invoke this before spawning a child that will
+   * read the parent's thread (`newThreadSource: "from-parent"` or an
+   * explicit parent threadId): the parent's slice otherwise only
+   * lands in storage at session-exit time, so the child would load a
+   * stale (or empty) snapshot. Best-effort — failures are logged by
+   * the session but never thrown.
+   */
+  persistThreadState?: () => Promise<void>;
 }
 
 /**
@@ -294,6 +320,23 @@ export interface ProcessToolCallsContext {
   turn?: number;
   /** Active sandbox ID (when a sandbox is configured for this session) */
   sandboxId?: string;
+  /**
+   * Id of the assistant message that produced these tool calls. The
+   * router forwards it into every handler's {@link RouterContext} so
+   * handlers can reference the message they were issued from (e.g.
+   * subagent forks that need to truncate the orphan assistant message
+   * out of a parent-forked thread).
+   */
+  assistantMessageId?: string;
+  /**
+   * Optional callback that flushes the session's in-memory
+   * `PersistedThreadState` slice to the durable thread store. The
+   * router forwards it into every handler's {@link RouterContext}
+   * verbatim. The session uses this to let mid-loop tool handlers
+   * (notably subagents that fork or continue the parent's thread)
+   * persist the parent's slice before the child reads it.
+   */
+  persistThreadState?: () => Promise<void>;
 }
 
 /**