zeitlich 0.2.48 → 0.2.49

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "zeitlich",
-  "version": "0.2.48",
+  "version": "0.2.49",
   "description": "[EXPERIMENTAL] An opinionated AI agent implementation for Temporal",
   "main": "./dist/index.js",
   "module": "./dist/index.js",

package/src/lib/session/session.ts

@@ -585,6 +585,17 @@ export async function createSession<
               ...(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

@@ -439,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

@@ -813,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.ts

@@ -221,7 +221,8 @@ export function createToolRouter<T extends ToolMap>(
     turn: number,
     sandboxId?: string,
     onRewindRequested?: (signal: RewindSignal) => void,
-    assistantMessageId?: string
+    assistantMessageId?: string,
+    persistThreadState?: () => Promise<void>
   ): Promise<ProcessedToolCall> {
     const startTime = Date.now();
     const tool = toolMap.get(toolCall.name);
@@ -265,6 +266,7 @@ export function createToolRouter<T extends ToolMap>(
           toolName: toolCall.name,
           ...(sandboxId !== undefined && { sandboxId }),
           ...(assistantMessageId !== undefined && { assistantMessageId }),
+          ...(persistThreadState !== undefined && { persistThreadState }),
         };
         const response = await tool.handler(
           effectiveArgs as Parameters<typeof tool.handler>[0],
@@ -422,6 +424,7 @@ 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;
 
@@ -443,7 +446,8 @@ export function createToolRouter<T extends ToolMap>(
                 turn,
                 sandboxId,
                 onRewindRequested,
-                assistantMessageId
+                assistantMessageId,
+                persistThreadState
               )
             )
           )
@@ -471,7 +475,8 @@ export function createToolRouter<T extends ToolMap>(
           turn,
           sandboxId,
           undefined,
-          assistantMessageId
+          assistantMessageId,
+          persistThreadState
         );
         if (outcome.kind === "rewind") {
           rewindSignal = outcome.signal;
@@ -510,6 +515,9 @@ export function createToolRouter<T extends ToolMap>(
           ...(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

@@ -190,6 +190,20 @@ export interface RouterContext {
    * 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>;
 }
 
 /**
@@ -314,6 +328,15 @@ export interface ProcessToolCallsContext {
    * 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>;
 }
 
 /**