zeitlich 0.2.35 → 0.2.37

package/src/lib/subagent/handler.ts

@@ -1,9 +1,10 @@
 import {
-  startChild,
   workflowInfo,
   setHandler,
   condition,
   log,
+  ApplicationFailure,
+  executeChild,
 } from "@temporalio/workflow";
 import { getShortId } from "../thread/id";
 import type { ToolHandlerResponse, RouterContext } from "../tool-router";
@@ -11,7 +12,6 @@ import type { JsonValue } from "../state/types";
 import type {
   InferSubagentResult,
   SubagentConfig,
-  SubagentHandlerResponse,
   SubagentSandboxConfig,
   SubagentWorkflowInput,
 } from "./types";
@@ -22,17 +22,14 @@ import type {
   SandboxInit,
   SubagentSandboxShutdown,
 } from "../lifecycle";
-import {
-  childResultSignal,
-  childSandboxReadySignal,
-  destroySandboxSignal,
-} from "./signals";
+import type { SandboxOps, SandboxSnapshot } from "../sandbox/types";
+import { childSandboxReadySignal } from "./signals";
 
 /** Normalized sandbox config after resolving the union. */
 interface ResolvedSandboxConfig {
   source: "none" | "inherit" | "own";
   init: "per-call" | "once";
-  continuation: "continue" | "fork";
+  continuation: "continue" | "fork" | "snapshot";
   shutdown?: SubagentSandboxShutdown;
 }
 
@@ -61,9 +58,9 @@ function resolveSandboxConfig(
 /**
  * Creates a Subagent tool handler that spawns child workflows for configured subagents.
  *
- * Child workflows signal their result back via `childResultSignal` instead of
- * returning it as the workflow return value. The handler awaits the signal
- * before continuing.
+ * Sandbox and snapshot cleanup happens inside the parent via each subagent's
+ * `sandbox.proxy` — the proxy factory is invoked once per subagent with
+ * `scope = agentName` so it resolves to the same activities the child uses.
  *
  * @param subagents - Array of subagent configurations
  * @returns A tool handler function that can be used with the tool router
@@ -78,13 +75,27 @@ export function createSubagentHandler<
     context: RouterContext
   ) => Promise<ToolHandlerResponse<InferSubagentResult<T[number]> | null>>;
   destroySubagentSandboxes: () => Promise<void>;
+  cleanupSubagentSnapshots: () => Promise<void>;
 } {
   const { taskQueue: parentTaskQueue } = workflowInfo();
 
-  const childResults = new Map<string, SubagentHandlerResponse>();
+  /** Sandbox ops proxy per subagent, built eagerly from `sandbox.proxy` factories. */
+  const agentSandboxOps = new Map<string, SandboxOps>();
+  for (const cfg of subagents) {
+    if (cfg.sandbox && cfg.sandbox !== "none") {
+      agentSandboxOps.set(cfg.agentName, cfg.sandbox.proxy(cfg.agentName));
+    }
+  }
+
+  /**
+   * Sandboxes that outlived their child session and must be destroyed by the
+   * parent at shutdown (shutdown = `pause-until-parent-close` /
+   * `keep-until-parent-close`). Keyed by `persistent:<agent>` for lazy
+   * shared sandboxes and by childWorkflowId otherwise.
+   */
   const pendingDestroys = new Map<
     string,
-    Awaited<ReturnType<typeof startChild>>
+    { agentName: string; sandboxId: string }
   >();
   /** Maps childThreadId → sandboxId for sandbox continuation across invocations (init: per-call) */
   const threadSandboxes = new Map<string, string>();
@@ -94,21 +105,26 @@ export function createSubagentHandler<
   const persistentSandboxCreating = new Set<string>();
   /** Reverse lookup: childWorkflowId → agentName for in-flight lazy creators */
   const lazyCreatorAgent = new Map<string, string>();
-
-  setHandler(childResultSignal, ({ childWorkflowId, result }) => {
-    childResults.set(childWorkflowId, result);
-  });
-
-  setHandler(
-    childSandboxReadySignal,
-    ({ childWorkflowId, sandboxId }) => {
-      const agentName = lazyCreatorAgent.get(childWorkflowId);
-      if (agentName && !persistentSandboxes.has(agentName)) {
-        persistentSandboxes.set(agentName, sandboxId);
-        lazyCreatorAgent.delete(childWorkflowId);
-      }
+  /** Maps childThreadId → latest snapshot for sandbox continuation via snapshots */
+  const threadSnapshots = new Map<
+    string,
+    {
+      agentName: string;
+      snapshot: SandboxSnapshot;
     }
-  );
+  >();
+  /** Maps agentName → reusable base snapshot captured on first-ever call (init: once + continuation: "snapshot") */
+  const persistentBaseSnapshot = new Map<string, SandboxSnapshot>();
+  /** Tracks agents whose first snapshot-backed sandbox creation is in-flight */
+  const persistentBaseSnapshotCreating = new Set<string>();
+
+  setHandler(childSandboxReadySignal, ({ childWorkflowId, sandboxId }) => {
+    const agentName = lazyCreatorAgent.get(childWorkflowId);
+    if (agentName && !persistentSandboxes.has(agentName)) {
+      persistentSandboxes.set(agentName, sandboxId);
+      lazyCreatorAgent.delete(childWorkflowId);
+    }
+  });
 
   const handler = async (
     args: SubagentArgs,
@@ -127,6 +143,16 @@ export function createSubagentHandler<
     const { sandboxId: parentSandboxId } = context;
     const sandboxCfg = resolveSandboxConfig(config.sandbox);
 
+    if (
+      sandboxCfg.source !== "none" &&
+      !agentSandboxOps.has(config.agentName)
+    ) {
+      throw ApplicationFailure.create({
+        message: `Subagent "${config.agentName}" uses a sandbox but no \`sandbox.proxy\` is configured on its SubagentConfig`,
+        nonRetryable: true,
+      });
+    }
+
     if (sandboxCfg.source === "inherit" && !parentSandboxId) {
       throw new Error(
         `Subagent "${config.agentName}" is configured with sandbox: "inherit" but the parent has no sandbox`
@@ -151,13 +177,50 @@ export function createSubagentHandler<
     let sandbox: SandboxInit | undefined;
     let sandboxShutdownOverride: SubagentSandboxShutdown | undefined;
     let isLazyCreator = false;
+    let isSnapshotBaseCreator = false;
 
     if (sandboxCfg.source === "inherit" && parentSandboxId) {
       if (sandboxCfg.continuation === "fork") {
         sandbox = { mode: "fork", sandboxId: parentSandboxId };
+      } else if (sandboxCfg.continuation === "snapshot") {
+        throw new Error(
+          `Subagent "${config.agentName}" has sandbox source "inherit" with continuation "snapshot" — snapshot continuation is only supported for source "own"`
+        );
       } else {
         sandbox = { mode: "inherit", sandboxId: parentSandboxId };
       }
+    } else if (
+      sandboxCfg.source === "own" &&
+      sandboxCfg.continuation === "snapshot"
+    ) {
+      // Snapshot-driven continuation: each call boots a fresh sandbox from a
+      // stored snapshot (per-thread, or a per-agent base for new threads with
+      // init: "once"). The session destroys its sandbox inline on exit;
+      // stored snapshot IDs are cleaned up by the parent at shutdown.
+      const isLazy = sandboxCfg.init === "once";
+
+      let baseSnap: SandboxSnapshot | undefined;
+      if (continuationThreadId) {
+        baseSnap = threadSnapshots.get(continuationThreadId)?.snapshot;
+      }
+
+      if (!baseSnap && isLazy) {
+        baseSnap = persistentBaseSnapshot.get(config.agentName);
+        if (!baseSnap) {
+          if (persistentBaseSnapshotCreating.has(config.agentName)) {
+            await condition(() => persistentBaseSnapshot.has(config.agentName));
+            baseSnap = persistentBaseSnapshot.get(config.agentName);
+          } else {
+            persistentBaseSnapshotCreating.add(config.agentName);
+            isSnapshotBaseCreator = true;
+          }
+        }
+      }
+
+      if (baseSnap) {
+        sandbox = { mode: "from-snapshot", snapshot: baseSnap };
+      }
+      sandboxShutdownOverride = "snapshot";
     } else if (sandboxCfg.source === "own") {
       const isLazy = sandboxCfg.init === "once";
 
@@ -246,45 +309,11 @@ export function createSubagentHandler<
       sandboxSource: sandboxCfg.source,
     });
 
-    const childHandle = await startChild(config.workflow, childOpts);
+    const childResult = await executeChild(config.workflow, childOpts);
 
-    // Track child handles that need signaling at parent shutdown.
     const effectiveShutdown =
       sandboxShutdownOverride ?? sandboxCfg.shutdown ?? "destroy";
 
-    if (
-      effectiveShutdown === "pause-until-parent-close" ||
-      effectiveShutdown === "keep-until-parent-close"
-    ) {
-      const key = isLazyCreator
-        ? `persistent:${config.agentName}`
-        : childWorkflowId;
-      pendingDestroys.set(key, childHandle);
-    }
-
-    // Wait for signal from child; race with child completion to propagate failures
-    await Promise.race([
-      condition(() => childResults.has(childWorkflowId)),
-      childHandle.result(),
-    ]);
-    if (!childResults.has(childWorkflowId)) {
-      await condition(() => childResults.has(childWorkflowId));
-    }
-
-    const childResult = childResults.get(childWorkflowId);
-    childResults.delete(childWorkflowId);
-
-    if (!childResult) {
-      log.warn("subagent returned no result", {
-        subagent: config.agentName,
-        childWorkflowId,
-      });
-      return {
-        toolResponse: "Subagent workflow did not signal a result",
-        data: null,
-      };
-    }
-
     log.info("subagent completed", {
       subagent: config.agentName,
       childWorkflowId,
@@ -297,27 +326,71 @@ export function createSubagentHandler<
       usage,
       threadId: childThreadId,
       sandboxId: childSandboxId,
+      snapshot: childSnapshot,
+      baseSnapshot: childBaseSnapshot,
       metadata,
     } = childResult;
 
-    // Store sandbox ID for future continuation/fork
     if (childSandboxId) {
       if (
         sandboxCfg.source === "own" &&
         sandboxCfg.init === "once" &&
+        sandboxCfg.continuation !== "snapshot" &&
         !persistentSandboxes.has(config.agentName)
       ) {
         // Fallback: signal may have already set this via childSandboxReadySignal
         persistentSandboxes.set(config.agentName, childSandboxId);
-      } else if (allowsContinuation && childThreadId) {
+      } else if (
+        allowsContinuation &&
+        childThreadId &&
+        sandboxCfg.source === "own" &&
+        sandboxCfg.continuation !== "snapshot"
+      ) {
         threadSandboxes.set(childThreadId, childSandboxId);
       }
     }
 
+    // Track sandboxes that must be destroyed by the parent at shutdown.
+    if (
+      childSandboxId &&
+      (effectiveShutdown === "pause-until-parent-close" ||
+        effectiveShutdown === "keep-until-parent-close")
+    ) {
+      const key = isLazyCreator
+        ? `persistent:${config.agentName}`
+        : childWorkflowId;
+      pendingDestroys.set(key, {
+        agentName: config.agentName,
+        sandboxId: childSandboxId,
+      });
+    }
+
+    // Store snapshots for future snapshot-driven continuation and final sweep.
+    // Tag each with `agentName` so `cleanupSubagentSnapshots` knows which
+    // sandbox ops to call for deletion.
+    if (sandboxCfg.source === "own" && sandboxCfg.continuation === "snapshot") {
+      if (childSnapshot && childThreadId) {
+        threadSnapshots.set(childThreadId, {
+          agentName: config.agentName,
+          snapshot: childSnapshot,
+        });
+      }
+      if (
+        isSnapshotBaseCreator &&
+        childBaseSnapshot &&
+        !persistentBaseSnapshot.has(config.agentName)
+      ) {
+        persistentBaseSnapshot.set(config.agentName, childBaseSnapshot);
+      }
+    }
+
     if (isLazyCreator) {
       persistentSandboxCreating.delete(config.agentName);
       lazyCreatorAgent.delete(childWorkflowId);
     }
+    if (isSnapshotBaseCreator) {
+      persistentBaseSnapshotCreating.delete(config.agentName);
+    }
 
     if (!toolResponse) {
       return {
@@ -355,7 +428,9 @@ export function createSubagentHandler<
 
     return {
       toolResponse: finalToolResponse,
-      data: validated ? validated.data : data,
+      data: validated
+        ? validated.data
+        : (data as InferSubagentResult<T[number]> | null),
       ...(usage && { usage }),
       ...(childSandboxId && { sandboxId: childSandboxId }),
       ...(metadata && { metadata }),
@@ -363,15 +438,55 @@ export function createSubagentHandler<
   };
 
   const destroySubagentSandboxes = async (): Promise<void> => {
-    const handles = [...pendingDestroys.values()];
+    const entries = [...pendingDestroys.values()];
     pendingDestroys.clear();
     await Promise.all(
-      handles.map(async (handle) => {
+      entries.map(async ({ agentName, sandboxId }) => {
+        const ops = agentSandboxOps.get(agentName);
+        if (!ops) {
+          log.warn(
+            "Skipping sandbox destroy — no sandbox.proxy registered for agent",
+            { agentName, sandboxId }
+          );
+          return;
+        }
+        try {
+          await ops.destroySandbox(sandboxId);
+        } catch (err) {
+          log.warn("Failed to destroy subagent sandbox", {
+            agentName,
+            sandboxId,
+            error: err,
+          });
+        }
+      })
+    );
+  };
+
+  const cleanupSubagentSnapshots = async (): Promise<void> => {
+    const tagged = [];
+    for (const entry of threadSnapshots.values()) tagged.push(entry);
+    for (const [agentName, snapshot] of persistentBaseSnapshot.entries()) {
+      tagged.push({ agentName, snapshot });
+    }
+    threadSnapshots.clear();
+    persistentBaseSnapshot.clear();
+
+    await Promise.all(
+      tagged.map(async ({ agentName, snapshot }) => {
+        const ops = agentSandboxOps.get(agentName);
+        if (!ops) {
+          log.warn(
+            "Skipping snapshot delete — no sandbox.proxy registered for agent",
+            { agentName }
+          );
+          return;
+        }
         try {
-          await handle.signal(destroySandboxSignal);
-          await handle.result();
+          await ops.deleteSandboxSnapshot(snapshot);
         } catch (err) {
-          log.warn("Failed to signal destroySandbox to child workflow", {
+          log.warn("Failed to delete subagent snapshot", {
+            agentName,
             error: err,
           });
         }
@@ -379,5 +494,5 @@ export function createSubagentHandler<
     );
   };
 
-  return { handler, destroySubagentSandboxes };
+  return { handler, destroySubagentSandboxes, cleanupSubagentSnapshots };
 }

package/src/lib/subagent/index.ts

@@ -17,8 +17,4 @@ export { createSubagentHandler } from "./handler";
 export { defineSubagent } from "./define";
 export { defineSubagentWorkflow } from "./workflow";
 export { buildSubagentRegistration } from "./register";
-export {
-  childResultSignal,
-  childSandboxReadySignal,
-  destroySandboxSignal,
-} from "./signals";
+export { childSandboxReadySignal } from "./signals";

package/src/lib/subagent/register.ts

@@ -26,6 +26,7 @@ import { createSubagentHandler } from "./handler";
 export function buildSubagentRegistration(subagents: SubagentConfig[]): {
   registration: ToolMap[string];
   destroySubagentSandboxes: () => Promise<void>;
+  cleanupSubagentSnapshots: () => Promise<void>;
 } | null {
   if (subagents.length === 0) return null;
 
@@ -42,7 +43,7 @@ export function buildSubagentRegistration(subagents: SubagentConfig[]): {
   const resolveSubagentName = (args: unknown): string =>
     (args as SubagentArgs).subagent;
 
-  const { handler, destroySubagentSandboxes } =
+  const { handler, destroySubagentSandboxes, cleanupSubagentSnapshots } =
     createSubagentHandler(subagents);
 
   const registration: ToolMap[string] = {
@@ -72,5 +73,5 @@ export function buildSubagentRegistration(subagents: SubagentConfig[]): {
     }),
   };
 
-  return { registration, destroySubagentSandboxes };
+  return { registration, destroySubagentSandboxes, cleanupSubagentSnapshots };
 }

package/src/lib/subagent/signals.ts

@@ -1,15 +1,6 @@
 import { defineSignal } from "@temporalio/workflow";
-import type {
-  ChildResultSignalPayload,
-  ChildSandboxReadySignalPayload,
-} from "./types";
-
-export const childResultSignal =
-  defineSignal<[ChildResultSignalPayload]>("childResult");
+import type { ChildSandboxReadySignalPayload } from "./types";
 
 /** Sent by a child workflow as soon as its sandbox is created, before the agent loop starts. */
 export const childSandboxReadySignal =
   defineSignal<[ChildSandboxReadySignalPayload]>("childSandboxReady");
-
-/** Sent by the parent to tell a subagent it may destroy its sandbox. */
-export const destroySandboxSignal = defineSignal("destroySandbox");