zeitlich 0.2.46 → 0.2.47

package/src/lib/subagent/handler.ts

@@ -237,8 +237,17 @@ 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 continuationThreadId = !allowsContinuation
+      ? undefined
+      : (args.threadId ??
+        (newThreadSource === "from-parent" ? context.threadId : undefined));
 
     // --- Build thread init ---
     let thread: ThreadInit | undefined;

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

@@ -635,6 +635,145 @@ 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" }
+    );
+
+    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", 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" }
+    );
+
+    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(

package/src/lib/thread/proxy.ts

@@ -8,55 +8,105 @@ import {
   proxyActivities,
   workflowInfo,
   type ActivityInterfaceFor,
+  type ActivityOptions,
 } from "@temporalio/workflow";
 import type { ThreadOps } from "../session/types";
 
+type OpName = keyof ThreadOps;
+
+/** Tight `startToCloseTimeout` so a sick Redis surfaces quickly via retry. */
+const DEFAULT_OPTIONS: ActivityOptions = {
+  startToCloseTimeout: "10s",
+  retry: {
+    maximumAttempts: 6,
+    initialInterval: "5s",
+    maximumInterval: "15m",
+    backoffCoefficient: 4,
+  },
+};
+
 /**
- * Creates a workflow-safe Temporal activity proxy for {@link ThreadOps}.
+ * `heartbeatTimeout` assumes the built-in S3 cold store's progress
+ * events (multipart `Upload` + chunked stream read). Stalls trip via
+ * heartbeat rather than `startToCloseTimeout`. Custom backends without
+ * progress events should override via `perOp`. Harmless on Redis-only
+ * deployments — the activities no-op.
+ */
+const BUILTIN_PER_OP: Partial<Record<OpName, ActivityOptions>> = {
+  hydrateThread: { startToCloseTimeout: "60s", heartbeatTimeout: "15s" },
+  flushThread: { startToCloseTimeout: "60s", heartbeatTimeout: "15s" },
+};
+
+/**
+ * `perOp[op]` layers shallow-rightmost over `defaults` and the
+ * built-in cold-tier overlay (`hydrateThread` / `flushThread`).
+ * A bare {@link ActivityOptions} is also accepted (treated as `{ defaults }`).
  *
- * The proxy resolves activity names by combining the adapter prefix with
- * the workflow scope, so each adapter + workflow combination gets its own
- * namespace.
+ * @example
+ * ```typescript
+ * proxyAnthropicThreadOps(undefined, {
+ *   defaults: { startToCloseTimeout: "5s" },
+ *   perOp: {
+ *     flushThread: { startToCloseTimeout: "180s" }, // heartbeatTimeout still inherited
+ *   },
+ * });
+ * ```
+ */
+export interface ThreadOpsProxyOptions {
+  defaults?: ActivityOptions;
+  perOp?: Partial<Record<OpName, ActivityOptions>>;
+}
+
+function isProxyOptionsShape(o: object): o is ThreadOpsProxyOptions {
+  return "defaults" in o || "perOp" in o;
+}
+
+/**
+ * Creates a workflow-safe Temporal activity proxy for {@link ThreadOps}.
  *
  * @param adapterPrefix - Adapter identifier (e.g. "anthropic", "googleGenAI", "langChain")
- * @param scope - Optional workflow scope override. Defaults to `workflowInfo().workflowType`.
- * @param options - Optional Temporal `proxyActivities` options.
+ * @param scope - Workflow scope. Defaults to `workflowInfo().workflowType`.
+ * @param options - {@link ThreadOpsProxyOptions} or a bare {@link ActivityOptions}.
  */
 export function createThreadOpsProxy(
   adapterPrefix: string,
   scope?: string,
-  options?: Parameters<typeof proxyActivities>[0]
+  options?: ActivityOptions | ThreadOpsProxyOptions
 ): ActivityInterfaceFor<ThreadOps> {
   const resolvedScope = scope ?? workflowInfo().workflowType;
 
+  const opts: ThreadOpsProxyOptions =
+    options && isProxyOptionsShape(options) ? options : { defaults: options };
+
+  const base = opts.defaults ?? DEFAULT_OPTIONS;
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  const acts = proxyActivities<Record<string, (...args: any[]) => any>>(
-    options ?? {
-      startToCloseTimeout: "10s",
-      retry: {
-        maximumAttempts: 6,
-        initialInterval: "5s",
-        maximumInterval: "15m",
-        backoffCoefficient: 4,
-      },
-    }
-  );
+  const baseActs = proxyActivities<Record<string, (...args: any[]) => any>>(base);
 
   const prefix = `${adapterPrefix}${resolvedScope.charAt(0).toUpperCase()}${resolvedScope.slice(1)}`;
   const p = (key: string): string =>
     `${prefix}${key.charAt(0).toUpperCase()}${key.slice(1)}`;
 
+  const pick = (op: OpName): unknown => {
+    const overlay = { ...BUILTIN_PER_OP[op], ...opts.perOp?.[op] };
+    if (Object.keys(overlay).length === 0) return baseActs[p(op)];
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    return proxyActivities<Record<string, (...args: any[]) => any>>({
+      ...base,
+      ...overlay,
+    })[p(op)];
+  };
+
   return {
-    initializeThread: acts[p("initializeThread")],
-    appendHumanMessage: acts[p("appendHumanMessage")],
-    appendToolResult: acts[p("appendToolResult")],
-    appendAgentMessage: acts[p("appendAgentMessage")],
-    appendSystemMessage: acts[p("appendSystemMessage")],
-    forkThread: acts[p("forkThread")],
-    truncateThread: acts[p("truncateThread")],
-    loadThreadState: acts[p("loadThreadState")],
-    saveThreadState: acts[p("saveThreadState")],
-    hydrateThread: acts[p("hydrateThread")],
-    flushThread: acts[p("flushThread")],
+    initializeThread: pick("initializeThread"),
+    appendHumanMessage: pick("appendHumanMessage"),
+    appendToolResult: pick("appendToolResult"),
+    appendAgentMessage: pick("appendAgentMessage"),
+    appendSystemMessage: pick("appendSystemMessage"),
+    forkThread: pick("forkThread"),
+    truncateThread: pick("truncateThread"),
+    loadThreadState: pick("loadThreadState"),
+    saveThreadState: pick("saveThreadState"),
+    hydrateThread: pick("hydrateThread"),
+    flushThread: pick("flushThread"),
   } as ActivityInterfaceFor<ThreadOps>;
 }