zeitlich 0.2.46 → 0.2.48

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,9 @@ export async function createSession<
             {
               turn: currentTurn,
               ...(sandboxId !== undefined && { sandboxId }),
+              ...(assistantId !== undefined && {
+                assistantMessageId: assistantId,
+              }),
             }
           );
 

package/src/lib/subagent/define.ts

@@ -44,6 +44,7 @@ export function defineSubagent<
     enabled?: boolean | (() => boolean);
     taskQueue?: string;
     thread?: "new" | "fork" | "continue";
+    newThreadSource?: "new" | "from-parent";
     sandbox?: SubagentSandboxConfig;
   }
 ): SubagentConfig<TResult> {

package/src/lib/subagent/handler.ts

@@ -237,16 +237,51 @@ export function createSubagentHandler<
 
     const threadMode = config.thread ?? "new";
     const allowsContinuation = threadMode !== "new";
-    const continuationThreadId =
-      args.threadId && allowsContinuation ? args.threadId : undefined;
+    // The parent agent's tool call wins. When `threadId` is omitted,
+    // `newThreadSource` decides what to fall back to: `"new"` (default)
+    // starts fresh, `"from-parent"` continues/forks the parent's own
+    // thread via `context.threadId`. Both paths still require
+    // `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 ?? (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 ---

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

@@ -635,6 +635,184 @@ describe("createSubagentHandler", () => {
     expect(workflowInput.thread).toBeUndefined();
   });
 
+  // --- newThreadSource: "from-parent" ---
+
+  it("uses the parent's threadId when fork + newThreadSource 'from-parent' and args.threadId is absent", async () => {
+    const { executeChild } = await import("@temporalio/workflow");
+    const execMock = executeChild as ReturnType<typeof vi.fn>;
+
+    const subagent: SubagentConfig = {
+      agentName: "parent-fork",
+      description: "Forks parent thread by default",
+      workflow: mockWorkflow(),
+      thread: "fork",
+      newThreadSource: "from-parent",
+    };
+
+    const { handler } = createSubagentHandler([subagent]);
+
+    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" }
+    );
+
+    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",
+    });
+  });
+
+  it("uses the parent's threadId when continue + newThreadSource 'from-parent' and args.threadId is absent", async () => {
+    const { executeChild } = await import("@temporalio/workflow");
+    const execMock = executeChild as ReturnType<typeof vi.fn>;
+
+    const subagent: SubagentConfig = {
+      agentName: "parent-continue",
+      description: "Continues parent thread by default",
+      workflow: mockWorkflow(),
+      thread: "continue",
+      newThreadSource: "from-parent",
+    };
+
+    const { handler } = createSubagentHandler([subagent]);
+
+    await handler(
+      { subagent: "parent-continue", description: "test", prompt: "test" },
+      { threadId: "parent-t", toolCallId: "tc", toolName: "Subagent" }
+    );
+
+    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: "continue",
+      threadId: "parent-t",
+    });
+  });
+
+  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>;
+
+    const subagent: SubagentConfig = {
+      agentName: "parent-fork-explicit",
+      description: "Forks parent thread by default",
+      workflow: mockWorkflow(),
+      thread: "fork",
+      newThreadSource: "from-parent",
+    };
+
+    const { handler } = createSubagentHandler([subagent]);
+
+    await handler(
+      {
+        subagent: "parent-fork-explicit",
+        description: "test",
+        prompt: "test",
+        threadId: "explicit-prev",
+      },
+      {
+        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: "explicit-prev",
+    });
+  });
+
+  it("ignores newThreadSource 'from-parent' when thread is new", async () => {
+    const { executeChild } = await import("@temporalio/workflow");
+    const execMock = executeChild as ReturnType<typeof vi.fn>;
+
+    const subagent: SubagentConfig = {
+      agentName: "new-with-source",
+      description: "Should still start fresh",
+      workflow: mockWorkflow(),
+      newThreadSource: "from-parent",
+    };
+
+    const { handler } = createSubagentHandler([subagent]);
+
+    await handler(
+      { subagent: "new-with-source", description: "test", prompt: "test" },
+      { threadId: "parent-t", toolCallId: "tc", toolName: "Subagent" }
+    );
+
+    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).toBeUndefined();
+  });
+
+  it("preserves prior behavior when fork is set with default newThreadSource and args.threadId is absent", async () => {
+    const { executeChild } = await import("@temporalio/workflow");
+    const execMock = executeChild as ReturnType<typeof vi.fn>;
+
+    const subagent: SubagentConfig = {
+      agentName: "fork-default-source",
+      description: "Fork with default new-thread source",
+      workflow: mockWorkflow(),
+      thread: "fork",
+    };
+
+    const { handler } = createSubagentHandler([subagent]);
+
+    await handler(
+      { subagent: "fork-default-source", description: "test", prompt: "test" },
+      { threadId: "parent-t", toolCallId: "tc", toolName: "Subagent" }
+    );
+
+    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).toBeUndefined();
+  });
+
   // --- Sandbox continuation ---
 
   it("does not pass sandbox when thread is fork (own sandbox)", async () => {

package/src/lib/subagent/types.ts

@@ -632,6 +632,22 @@ export interface SubagentConfig<TResult extends z.ZodType = z.ZodType> {
    *   directly to the existing thread in-place.
    */
   thread?: "new" | "fork" | "continue";
+  /**
+   * Where the subagent's thread comes from when the parent's tool call
+   * omits `threadId`. Only meaningful in combination with
+   * `thread: "fork"` or `"continue"`.
+   *
+   * - `"new"` (default) — start a fresh thread (the prior behavior).
+   * - `"from-parent"` — use the parent's own `threadId` (from
+   *   `RouterContext`). With `thread: "fork"` the parent's conversation
+   *   is copied into a new thread; with `thread: "continue"` the
+   *   subagent appends to the parent's thread in-place.
+   *
+   * Has no effect when `thread` is `"new"` (or omitted). A `threadId`
+   * supplied by the parent agent always wins — `newThreadSource` only
+   * applies when none is provided.
+   */
+  newThreadSource?: "new" | "from-parent";
   /**
    * Sandbox strategy for this subagent.
    *

package/src/lib/thread/cold-store.test.ts

@@ -16,9 +16,9 @@ interface CommandInput {
 }
 
 /**
- * Minimal in-memory S3-shaped client. The `send()` method dispatches
- * by the command's constructor name, mirroring the contract that
- * `@aws-sdk/client-s3` exposes.
+ * Minimal S3-shaped fake. `send()` dispatches by command constructor
+ * name; `config` provides the fields `@aws-sdk/lib-storage`'s
+ * single-part upload path inspects.
  */
 function createFakeS3(): {
   s3: S3LikeClient;
@@ -31,7 +31,12 @@ function createFakeS3(): {
   const compositeKey = (bucket: string, key: string): string =>
     `${bucket}/${key}`;
 
-  const s3: S3LikeClient = {
+  const s3 = {
+    config: {
+      requestHandler: undefined,
+      forcePathStyle: false,
+      endpoint: async (): Promise<URL> => new URL("https://fake.s3.local"),
+    },
     async send<TInput, TOutput>(
       command: { input: TInput } & object
     ): Promise<TOutput> {
@@ -68,7 +73,7 @@ function createFakeS3(): {
       }
       throw new Error(`unknown command: ${name}`);
     },
-  };
+  } as unknown as S3LikeClient;
 
   return { s3, store, calls };
 }
@@ -190,4 +195,27 @@ describe("createS3ColdStore", () => {
     });
     expect(await cold.read("messages", "t-1")).toEqual(sampleSnapshot);
   });
+
+  it("round-trips a large payload through the async gzip path", async () => {
+    // ~1 MB payload — regression guard that large payloads still
+    // encode/decode correctly through the promisified gzip path.
+    const big: ThreadSnapshot = {
+      v: 1,
+      messages: Array.from({ length: 500 }, (_, i) =>
+        JSON.stringify({
+          id: `m${i}`,
+          text: "x".repeat(2048),
+        })
+      ),
+      state: null,
+      dedupIds: Array.from({ length: 500 }, (_, i) => `m${i}`),
+    };
+
+    const cold = createS3ColdStore({
+      s3: fake.s3,
+      bucket: "test-bucket",
+    });
+    await cold.write("messages", "big", big);
+    expect(await cold.read("messages", "big")).toEqual(big);
+  });
 });

package/src/lib/thread/cold-store.ts

@@ -18,9 +18,17 @@
  * last-writer-wins; no compare-and-swap is required.
  */
 
-import { gzipSync, gunzipSync } from "node:zlib";
+import { gunzip as gunzipCb, gzip as gzipCb } from "node:zlib";
+import { promisify } from "node:util";
 import type { PersistedThreadState } from "../state/types";
-import { DeleteObjectCommand, GetObjectCommand, PutObjectCommand } from "@aws-sdk/client-s3";
+import { DeleteObjectCommand, GetObjectCommand, type S3Client } from "@aws-sdk/client-s3";
+import { Upload } from "@aws-sdk/lib-storage";
+import { getActivityContext } from "../activity";
+
+// Async zlib so gzip/gunzip don't block the worker's event loop
+// during compression of large snapshots.
+const gzipAsync = promisify(gzipCb);
+const gunzipAsync = promisify(gunzipCb);
 
 /**
  * Serialized form of a thread that can be written to and read from a
@@ -69,20 +77,16 @@ export interface ColdThreadStore {
 }
 
 /**
- * Compact, duck-typed shape of an S3 client. Zeitlich only needs the
- * `send(...)` method; declaring this locally avoids forcing
- * `@aws-sdk/client-s3` to be installed when the consumer is using a
- * different cold-store backend.
+ * Alias for `@aws-sdk/client-s3`'s `S3Client`. The built-in store
+ * calls `send(...)` and accesses `client.config` (read by
+ * `@aws-sdk/lib-storage`'s `Upload`) — a duck-type with just `send`
+ * is not sufficient.
  */
-export interface S3LikeClient {
-  send<TInput, TOutput>(
-    command: { input: TInput } & object
-  ): Promise<TOutput>;
-}
+export type S3LikeClient = S3Client;
 
 /** Configuration for the built-in S3 cold store. */
 export interface S3ColdStoreConfig {
-  /** An `@aws-sdk/client-s3` `S3Client` (or duck-typed equivalent). */
+  /** An `@aws-sdk/client-s3` `S3Client`. */
   s3: S3LikeClient;
   /** S3 bucket that holds the archive. */
   bucket: string;
@@ -126,10 +130,26 @@ function buildKey(
 }
 
 async function streamToBuffer(
-  body: unknown
+  body: unknown,
+  onChunk?: () => void
 ): Promise<Buffer> {
   if (body == null) return Buffer.alloc(0);
   if (body instanceof Uint8Array) return Buffer.from(body);
+  // Prefer async iteration so `onChunk` fires per chunk. Node S3
+  // bodies (`SdkStream<Readable>`) iterate; bulk-read fallbacks
+  // below cover browser body shapes.
+  if (
+    typeof (body as { [Symbol.asyncIterator]?: unknown })[
+      Symbol.asyncIterator
+    ] === "function"
+  ) {
+    const chunks: Buffer[] = [];
+    for await (const chunk of body as AsyncIterable<Buffer | Uint8Array>) {
+      chunks.push(Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk));
+      onChunk?.();
+    }
+    return Buffer.concat(chunks);
+  }
   if (typeof (body as { transformToByteArray?: () => Promise<Uint8Array> })
     .transformToByteArray === "function") {
     const bytes = await (
@@ -144,12 +164,7 @@ async function streamToBuffer(
     ).arrayBuffer();
     return Buffer.from(ab);
   }
-  // Node.js Readable stream fallback
-  const chunks: Buffer[] = [];
-  for await (const chunk of body as AsyncIterable<Buffer | Uint8Array>) {
-    chunks.push(Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk));
-  }
-  return Buffer.concat(chunks);
+  return Buffer.alloc(0);
 }
 
 /**
@@ -191,9 +206,10 @@ export function createS3ColdStore(
         const resp = (await s3.send(
           new GetObjectCommand({ Bucket: bucket, Key })
         )) as { Body?: unknown };
-        const buf = await streamToBuffer(resp.Body);
+        const { heartbeat } = getActivityContext();
+        const buf = await streamToBuffer(resp.Body, heartbeat);
         const json = gzip
-          ? gunzipSync(buf).toString("utf8")
+          ? (await gunzipAsync(buf)).toString("utf8")
           : buf.toString("utf8");
         return JSON.parse(json) as ThreadSnapshot;
       } catch (err) {
@@ -209,16 +225,19 @@ export function createS3ColdStore(
     ): Promise<void> {
       const Key = buildKey(prefix, threadKey, threadId, gzip);
       const json = JSON.stringify(snapshot);
-      const body = gzip ? gzipSync(Buffer.from(json, "utf8")) : json;
-
-      await s3.send(
-        new PutObjectCommand({
-          Bucket: bucket,
-          Key,
-          Body: body,
-          ContentType: contentType,
-        })
-      );
+      const body = gzip ? await gzipAsync(Buffer.from(json, "utf8")) : json;
+
+      const upload = new Upload({
+        client: s3,
+        params: { Bucket: bucket, Key, Body: body, ContentType: contentType },
+      });
+
+      // Heartbeat per S3 part completion so a stalled upload trips
+      // `heartbeatTimeout` instead of `startToCloseTimeout`.
+      const { heartbeat } = getActivityContext();
+      if (heartbeat) upload.on("httpUploadProgress", heartbeat);
+
+      await upload.done();
     },
 
     async delete(