zeitlich 0.2.46 → 0.2.48

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>;
 }

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,8 @@ export function createToolRouter<T extends ToolMap>(
     toolCall: ParsedToolCallUnion<T>,
     turn: number,
     sandboxId?: string,
-    onRewindRequested?: (signal: RewindSignal) => void
+    onRewindRequested?: (signal: RewindSignal) => void,
+    assistantMessageId?: string
   ): Promise<ProcessedToolCall> {
     const startTime = Date.now();
     const tool = toolMap.get(toolCall.name);
@@ -263,6 +264,7 @@ export function createToolRouter<T extends ToolMap>(
           toolCallId: toolCall.id,
           toolName: toolCall.name,
           ...(sandboxId !== undefined && { sandboxId }),
+          ...(assistantMessageId !== undefined && { assistantMessageId }),
         };
         const response = await tool.handler(
           effectiveArgs as Parameters<typeof tool.handler>[0],
@@ -419,6 +421,7 @@ export function createToolRouter<T extends ToolMap>(
 
       const turn = context?.turn ?? 0;
       const sandboxId = context?.sandboxId;
+      const assistantMessageId = context?.assistantMessageId;
 
       let rewindSignal: RewindSignal | undefined;
 
@@ -435,7 +438,13 @@ 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
+              )
             )
           )
         );
@@ -457,7 +466,13 @@ 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
+        );
         if (outcome.kind === "rewind") {
           rewindSignal = outcome.signal;
           break;
@@ -492,6 +507,9 @@ export function createToolRouter<T extends ToolMap>(
           ...(context?.sandboxId !== undefined && {
             sandboxId: context.sandboxId,
           }),
+          ...(context?.assistantMessageId !== undefined && {
+            assistantMessageId: context.assistantMessageId,
+          }),
         };
         const response = await handler(
           toolCall.args as ToolArgs<T, TName>,

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

@@ -178,6 +178,18 @@ 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;
 }
 
 /**
@@ -294,6 +306,14 @@ 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;
 }
 
 /**

package/src/tools/edit/handler.test.ts

@@ -0,0 +1,177 @@
+import { beforeEach, describe, expect, it } from "vitest";
+import { InMemorySandboxProvider } from "../../adapters/sandbox/inmemory/index";
+import type { Sandbox, SandboxCreateOptions } from "../../lib/sandbox";
+import { SandboxManager } from "../../lib/sandbox/manager";
+import type { RouterContext } from "../../lib/tool-router/types";
+import { withSandbox } from "../../lib/tool-router/with-sandbox";
+import { applyEditPlan, editHandler, multiEditHandler } from "./handler";
+
+describe("edit handlers", () => {
+  let manager: SandboxManager<SandboxCreateOptions, Sandbox, "inMemory">;
+  let sandboxId: string;
+
+  const ctx = (id: string): RouterContext => ({
+    sandboxId: id,
+    threadId: "test-thread",
+    toolCallId: "test-call",
+    toolName: "FileEdit",
+  });
+
+  beforeEach(async () => {
+    manager = new SandboxManager(new InMemorySandboxProvider());
+    const result = await manager.create({
+      initialFiles: {
+        "/src/app.ts": [
+          "export function greet(name: string) {",
+          '  return "hello " + name;',
+          "}",
+          "",
+          "export const status = 'draft';",
+          "export const repeated = 'draft';",
+          "",
+        ].join("\n"),
+      },
+    });
+    expect(result).not.toBeNull();
+    sandboxId = (result as NonNullable<typeof result>).sandboxId;
+  });
+
+  it("applies one unique exact replacement", async () => {
+    const handler = withSandbox(manager, editHandler);
+
+    const response = await handler(
+      {
+        file_path: "/src/app.ts",
+        old_string: '  return "hello " + name;',
+        new_string: "  return `hello ${name}`;",
+      },
+      ctx(sandboxId)
+    );
+
+    const sandbox = await manager.getSandbox(sandboxId);
+    await expect(sandbox.fs.readFile("/src/app.ts")).resolves.toContain(
+      "return `hello ${name}`;"
+    );
+    expect(response.data?.success).toBe(true);
+    expect(response.data?.replacements).toBe(1);
+    expect(response.data?.hunks?.[0]).toMatchObject({
+      oldStartLine: 2,
+      newStartLine: 2,
+      oldLines: ['  return "hello " + name;'],
+      newLines: ["  return `hello ${name}`;"],
+    });
+  });
+
+  it("refuses ambiguous single edits without replace_all", async () => {
+    const handler = withSandbox(manager, editHandler);
+
+    const response = await handler(
+      {
+        file_path: "/src/app.ts",
+        old_string: "draft",
+        new_string: "ready",
+      },
+      ctx(sandboxId)
+    );
+
+    const sandbox = await manager.getSandbox(sandboxId);
+    await expect(sandbox.fs.readFile("/src/app.ts")).resolves.toContain(
+      "status = 'draft'"
+    );
+    expect(response.data?.success).toBe(false);
+    expect(response.toolResponse).toContain("appears 2 times");
+  });
+
+  it("supports replace_all for one edit", async () => {
+    const handler = withSandbox(manager, editHandler);
+
+    const response = await handler(
+      {
+        file_path: "/src/app.ts",
+        old_string: "draft",
+        new_string: "ready",
+        replace_all: true,
+      },
+      ctx(sandboxId)
+    );
+
+    const sandbox = await manager.getSandbox(sandboxId);
+    const content = await sandbox.fs.readFile("/src/app.ts");
+    expect(content).toContain("status = 'ready'");
+    expect(content).toContain("repeated = 'ready'");
+    expect(response.data?.success).toBe(true);
+    expect(response.data?.replacements).toBe(2);
+  });
+
+  it("applies multiple edits sequentially and atomically", async () => {
+    const handler = withSandbox(manager, multiEditHandler);
+
+    const response = await handler(
+      {
+        file_path: "/src/app.ts",
+        edits: [
+          {
+            old_string: '  return "hello " + name;',
+            new_string: "  return `hello ${name}`;",
+          },
+          { old_string: "draft", new_string: "ready", replace_all: true },
+        ],
+      },
+      { ...ctx(sandboxId), toolName: "FileMultiEdit" }
+    );
+
+    const sandbox = await manager.getSandbox(sandboxId);
+    const content = await sandbox.fs.readFile("/src/app.ts");
+    expect(content).toContain("return `hello ${name}`;");
+    expect(content).toContain("status = 'ready'");
+    expect(content).toContain("repeated = 'ready'");
+    expect(response.data?.success).toBe(true);
+    expect(response.data?.replacements).toBe(3);
+    expect(response.data?.hunks).toHaveLength(3);
+  });
+
+  it("leaves the file unchanged when a later multi-edit fails", async () => {
+    const handler = withSandbox(manager, multiEditHandler);
+    const sandbox = await manager.getSandbox(sandboxId);
+    const before = await sandbox.fs.readFile("/src/app.ts");
+
+    const response = await handler(
+      {
+        file_path: "/src/app.ts",
+        edits: [
+          {
+            old_string: '  return "hello " + name;',
+            new_string: "  return `hello ${name}`;",
+          },
+          { old_string: "missing text", new_string: "replacement" },
+        ],
+      },
+      { ...ctx(sandboxId), toolName: "FileMultiEdit" }
+    );
+
+    await expect(sandbox.fs.readFile("/src/app.ts")).resolves.toBe(before);
+    expect(response.data?.success).toBe(false);
+    expect(response.toolResponse).toContain("edit 1");
+  });
+});
+
+describe("applyEditPlan", () => {
+  it("rejects empty old_string before mutating content", () => {
+    const result = applyEditPlan("abc", [{ old_string: "", new_string: "x" }]);
+
+    expect(result).toMatchObject({ ok: false, editIndex: 0 });
+  });
+
+  it("treats replacement text literally", () => {
+    const result = applyEditPlan("a.$^ b.$^", [
+      { old_string: ".$^", new_string: "literal" },
+    ]);
+
+    expect(result).toMatchObject({ ok: false });
+
+    const unique = applyEditPlan("a.$^", [
+      { old_string: ".$^", new_string: "literal" },
+    ]);
+    expect(unique).toMatchObject({ ok: true, content: "aliteral" });
+  });
+});