zeitlich 0.2.40 → 0.2.41

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "zeitlich",
-  "version": "0.2.40",
+  "version": "0.2.41",
   "description": "[EXPERIMENTAL] An opinionated AI agent implementation for Temporal",
   "main": "./dist/index.js",
   "module": "./dist/index.js",
@@ -227,8 +227,10 @@
     "@eslint/js": "^10.0.1",
     "@google/genai": "^1.44.0",
     "@langchain/core": "^1.1.30",
-    "@temporalio/envconfig": "^1.15.0",
-    "@temporalio/worker": "^1.15.0",
+    "@temporalio/common": "1.16.0",
+    "@temporalio/envconfig": "1.16.0",
+    "@temporalio/worker": "1.16.0",
+    "@temporalio/workflow": "1.16.0",
     "@types/node": "^25.3.3",
     "eslint": "^10.0.2",
     "husky": "^9.1.7",
@@ -247,6 +249,10 @@
     "@e2b/code-interpreter": "^2.3.3",
     "@google/genai": "^1.43.0",
     "@langchain/core": ">=1.0.0",
+    "@temporalio/common": ">=1.16.0 <2.0.0",
+    "@temporalio/envconfig": ">=1.16.0 <2.0.0",
+    "@temporalio/worker": ">=1.16.0 <2.0.0",
+    "@temporalio/workflow": ">=1.16.0 <2.0.0",
     "ioredis": ">=5.0.0",
     "just-bash": ">=2.0.0"
   },
@@ -266,6 +272,12 @@
     "@langchain/core": {
       "optional": true
     },
+    "@temporalio/envconfig": {
+      "optional": true
+    },
+    "@temporalio/worker": {
+      "optional": true
+    },
     "just-bash": {
       "optional": true
     }
@@ -276,9 +288,6 @@
   },
   "homepage": "https://github.com/bead-ai/zeitlich#readme",
   "dependencies": {
-    "@temporalio/common": "^1.15.0",
-    "@temporalio/plugin": "^1.15.0",
-    "@temporalio/workflow": "^1.15.0",
     "zod": "^4.3.6"
   }
 }

package/src/adapters/sandbox/e2b/README.md

@@ -0,0 +1,81 @@
+# E2B Sandbox Adapter
+
+Adapter that exposes [E2B](https://e2b.dev/) cloud sandboxes through the
+standard `SandboxProvider` interface used by the rest of Zeitlich.
+
+## Configuration
+
+`E2bSandboxProvider` accepts an `E2bSandboxConfig` (provider-level defaults).
+Per-create overrides may be passed to `create()` via `E2bSandboxCreateOptions`.
+
+```ts
+import { E2bSandboxProvider } from "zeitlich/adapters/sandbox/e2b";
+
+const provider = new E2bSandboxProvider({
+  template: "my-template",
+  timeoutMs: 15 * 60 * 1000, // kill-on-abandon safety net
+  keepAliveMs: 15 * 60 * 1000, // refreshed on every tool call
+});
+```
+
+## Keep-alive pattern
+
+E2B's `timeoutMs` on `Sandbox.create` is a **sandbox lifetime**, not an idle
+timeout: when it elapses, E2B kills the sandbox regardless of activity. Long
+agent loops (LLM thinking + many tool calls) can outlive that window, and the
+next tool call hits `Sandbox.connect(sandboxId)` and surfaces a
+`SandboxNotFoundError` mid-run.
+
+`keepAliveMs` solves this without giving up the kill-on-abandon safety net.
+When set, every call to `provider.get(sandboxId)` passes
+`{ timeoutMs: keepAliveMs }` to `Sandbox.connect()`. Per the E2B SDK's
+`SandboxConnectOpts.timeoutMs` JSDoc:
+
+> For running sandboxes, the timeout will update only if the new timeout is
+> longer than the existing one.
+
+So `connect()` with a `timeoutMs` is **monotonic**: it never shrinks the
+lifetime of a running sandbox. Pick `keepAliveMs` as the **full per-call
+refresh window** you want — passing a value smaller than the time remaining
+is a no-op rather than a shrink, but you should still pick the value with
+"every tool call should give me at least this much headroom" in mind, not
+"floor to add".
+
+`provider.get()` is invoked exactly once per tool call by `withSandbox`, so:
+
+- An active session's tool calls each refresh the lifetime to at least
+  `keepAliveMs`. The sandbox cannot be killed mid-run as long as tools are
+  still firing — conceptually this is the sandbox equivalent of a Temporal
+  activity heartbeat.
+- An abandoned sandbox still dies `keepAliveMs` after the last tool call. The
+  existing kill-on-timeout safety net is preserved.
+- Consumers can drop `timeoutMs` back down to short, safe values (e.g.
+  15 minutes) without tuning against worst-case run length.
+
+### Recommended usage
+
+- Set `timeoutMs` to a value that bounds how long a sandbox can sit unused
+  after the consumer abnormally terminates (worker crash, workflow terminated,
+  `workflowRunTimeout`). This is your **abandon safety net**.
+- Set `keepAliveMs` to your **per-call refresh window** — typically the same
+  value as `timeoutMs`, or shorter if you want sandboxes to be reaped sooner
+  after the last tool call.
+
+### Provider-level only
+
+`keepAliveMs` is a provider-construction-time config. There is intentionally
+no per-create override: every sandbox managed by the provider refreshes by
+the same amount on each `get()`. If a real use case for per-sandbox refresh
+windows ever shows up we can add it without breaking changes.
+
+### When connect-with-options is not enough
+
+If you ever need to **shrink** a sandbox's remaining lifetime (e.g. force an
+early reap), `connect()` won't do it because of the monotonic-extend rule
+above. Use `Sandbox.setTimeout(timeoutMs)` or the static
+`SandboxApi.setTimeout(sandboxId, timeoutMs)` instead — those can extend or
+reduce.
+
+If E2B ever changes the semantics of `Sandbox.connect(sandboxId, { timeoutMs })`
+so it stops extending a running sandbox's lifetime at all, `setTimeout` is
+also a drop-in replacement for the call site in `provider.get()`.

package/src/adapters/sandbox/e2b/index.ts

@@ -1,4 +1,8 @@
-import { Sandbox as E2bSdkSandbox } from "@e2b/code-interpreter";
+import {
+  NotFoundError as E2bNotFoundError,
+  Sandbox as E2bSdkSandbox,
+  SandboxNotFoundError as E2bSandboxNotFoundError,
+} from "@e2b/code-interpreter";
 import type {
   Sandbox,
   SandboxCapabilities,
@@ -19,6 +23,20 @@ import type {
   E2bSandboxCreateOptions,
 } from "./types";
 
+/**
+ * True iff `err` is the E2B SDK's "this sandbox doesn't exist (anymore)"
+ * signal. We narrow to `SandboxNotFoundError` (the canonical class) and to
+ * its deprecated parent `NotFoundError` as a defensive fallback — older
+ * SDK paths still throw the parent for sandbox-not-found cases. Any other
+ * error (auth failure, network blip, 5xx, validation) is propagated
+ * unchanged so callers can react to it specifically.
+ */
+function isE2bSandboxNotFound(err: unknown): boolean {
+  return (
+    err instanceof E2bSandboxNotFoundError || err instanceof E2bNotFoundError
+  );
+}
+
 // ============================================================================
 // E2bSandbox
 // ============================================================================
@@ -76,6 +94,7 @@ export class E2bSandboxProvider implements SandboxProvider<
   private readonly defaultTemplate?: string;
   private readonly defaultWorkspaceBase: string;
   private readonly defaultTimeoutMs?: number;
+  private readonly defaultKeepAliveMs?: number;
   private readonly defaultAllowInternetAccess?: boolean;
   private readonly defaultNetwork?: E2bSandboxConfig["network"];
   private readonly defaultMetadata?: E2bSandboxConfig["metadata"];
@@ -85,6 +104,7 @@ export class E2bSandboxProvider implements SandboxProvider<
     this.defaultTemplate = config?.template;
     this.defaultWorkspaceBase = config?.workspaceBase ?? "/home/user";
     this.defaultTimeoutMs = config?.timeoutMs;
+    this.defaultKeepAliveMs = config?.keepAliveMs;
     this.defaultAllowInternetAccess = config?.allowInternetAccess;
     this.defaultNetwork = config?.network;
     this.defaultMetadata = config?.metadata;
@@ -120,15 +140,22 @@ export class E2bSandboxProvider implements SandboxProvider<
   }
 
   async get(sandboxId: string): Promise<E2bSandbox> {
+    const keepAliveMs = this.defaultKeepAliveMs;
     try {
-      const sdkSandbox = await E2bSdkSandbox.connect(sandboxId);
+      const sdkSandbox =
+        keepAliveMs !== undefined
+          ? await E2bSdkSandbox.connect(sandboxId, { timeoutMs: keepAliveMs })
+          : await E2bSdkSandbox.connect(sandboxId);
       return new E2bSandboxImpl(
         sandboxId,
         sdkSandbox,
         this.defaultWorkspaceBase
       );
-    } catch {
-      throw new SandboxNotFoundError(sandboxId);
+    } catch (err) {
+      if (isE2bSandboxNotFound(err)) {
+        throw new SandboxNotFoundError(sandboxId);
+      }
+      throw err;
     }
   }
 
@@ -137,7 +164,7 @@ export class E2bSandboxProvider implements SandboxProvider<
       const sdkSandbox = await E2bSdkSandbox.connect(sandboxId);
       await sdkSandbox.kill();
     } catch {
-      // Already gone or not found
+      // Already gone or not found — destroy is idempotent.
     }
   }
 

package/src/adapters/sandbox/e2b/keep-alive.test.ts

@@ -0,0 +1,115 @@
+import { describe, expect, it, vi, beforeEach } from "vitest";
+import {
+  Sandbox as E2bSdkSandbox,
+  SandboxNotFoundError as E2bSandboxNotFoundError,
+} from "@e2b/code-interpreter";
+import { E2bSandboxProvider } from "./index";
+import { SandboxNotFoundError } from "../../../lib/sandbox/types";
+
+vi.mock("@e2b/code-interpreter", () => {
+  class FakeSdkSandbox {
+    static create = vi.fn();
+    static connect = vi.fn();
+    static createSnapshot = vi.fn();
+    static deleteSnapshot = vi.fn();
+    sandboxId: string;
+    constructor(sandboxId: string) {
+      this.sandboxId = sandboxId;
+    }
+    commands = { run: vi.fn() };
+    files = {};
+    async kill() {}
+    async pause() {}
+  }
+  // Mirror the real SDK error class hierarchy: SandboxNotFoundError extends
+  // (deprecated) NotFoundError extends SandboxError extends Error.
+  class FakeSandboxError extends Error {}
+  class FakeNotFoundError extends FakeSandboxError {}
+  class FakeSandboxNotFoundError extends FakeNotFoundError {}
+  return {
+    Sandbox: FakeSdkSandbox,
+    SandboxError: FakeSandboxError,
+    NotFoundError: FakeNotFoundError,
+    SandboxNotFoundError: FakeSandboxNotFoundError,
+  };
+});
+
+const sdk = E2bSdkSandbox as unknown as {
+  create: ReturnType<typeof vi.fn>;
+  connect: ReturnType<typeof vi.fn>;
+  createSnapshot: ReturnType<typeof vi.fn>;
+};
+
+function makeFakeSdkSandbox(id = "sbx-1") {
+  return {
+    sandboxId: id,
+    commands: { run: vi.fn() },
+    files: {},
+    kill: vi.fn(),
+    pause: vi.fn(),
+  };
+}
+
+describe("E2bSandboxProvider keep-alive", () => {
+  beforeEach(() => {
+    sdk.create.mockReset();
+    sdk.connect.mockReset();
+    sdk.createSnapshot.mockReset();
+  });
+
+  it("forwards timeoutMs to connect() when keepAliveMs is configured at the provider level", async () => {
+    const fake = makeFakeSdkSandbox();
+    sdk.connect.mockResolvedValue(fake);
+
+    const provider = new E2bSandboxProvider({ keepAliveMs: 15 * 60 * 1000 });
+    const sandbox = await provider.get("sbx-1");
+
+    expect(sandbox.id).toBe("sbx-1");
+    expect(sdk.connect).toHaveBeenCalledTimes(1);
+    expect(sdk.connect).toHaveBeenCalledWith("sbx-1", {
+      timeoutMs: 15 * 60 * 1000,
+    });
+  });
+
+  it("omits timeoutMs from connect() when keepAliveMs is not configured", async () => {
+    const fake = makeFakeSdkSandbox();
+    sdk.connect.mockResolvedValue(fake);
+
+    const provider = new E2bSandboxProvider();
+    await provider.get("sbx-1");
+
+    expect(sdk.connect).toHaveBeenCalledTimes(1);
+    expect(sdk.connect).toHaveBeenCalledWith("sbx-1");
+  });
+
+  it("uses provider-level keepAliveMs for every sandbox managed by the provider", async () => {
+    const fake = makeFakeSdkSandbox("sbx-default");
+    sdk.connect.mockResolvedValue(fake);
+
+    const provider = new E2bSandboxProvider({ keepAliveMs: 60_000 });
+    await provider.get("sbx-default");
+
+    expect(sdk.connect).toHaveBeenCalledWith("sbx-default", {
+      timeoutMs: 60_000,
+    });
+  });
+
+  it("translates the SDK's SandboxNotFoundError into our SandboxNotFoundError", async () => {
+    sdk.connect.mockRejectedValue(
+      new E2bSandboxNotFoundError("sandbox missing-sbx not found")
+    );
+
+    const provider = new E2bSandboxProvider({ keepAliveMs: 60_000 });
+    await expect(provider.get("missing-sbx")).rejects.toBeInstanceOf(
+      SandboxNotFoundError
+    );
+  });
+
+  it("propagates non-not-found connect() errors unchanged (auth, network, 5xx)", async () => {
+    const transient = new Error("ECONNRESET: socket hang up");
+    sdk.connect.mockRejectedValue(transient);
+
+    const provider = new E2bSandboxProvider({ keepAliveMs: 60_000 });
+    await expect(provider.get("sbx-1")).rejects.toBe(transient);
+  });
+});

package/src/adapters/sandbox/e2b/types.ts

@@ -19,8 +19,35 @@ export interface E2bSandboxConfig {
   template?: string;
   /** Default working directory inside the sandbox */
   workspaceBase?: string;
-  /** Sandbox idle timeout in milliseconds */
+  /**
+   * Sandbox lifetime in milliseconds. Despite the name, this is **not** an
+   * idle timeout: E2B kills the sandbox once this many milliseconds elapse
+   * from creation regardless of activity. Pair with {@link keepAliveMs} to
+   * refresh the lifetime on every `provider.get()` call so that this value
+   * acts as a kill-on-abandon safety net rather than a hard cap on run
+   * length.
+   */
   timeoutMs?: number;
+  /**
+   * If set, every call to `provider.get(sandboxId)` passes
+   * `{ timeoutMs: keepAliveMs }` to `Sandbox.connect()`, refreshing the
+   * sandbox lifetime on each tool invocation. The provider-level
+   * `timeoutMs` then acts as a kill-on-abandon safety net rather than a
+   * hard cap on run length.
+   *
+   * E2B's `Sandbox.connect()` is monotonic for running sandboxes: per the
+   * SDK's `SandboxConnectOpts.timeoutMs` doc, "the timeout will update
+   * only if the new timeout is longer than the existing one". Pick
+   * `keepAliveMs` as the full per-call refresh window you want; passing a
+   * value smaller than the time remaining is a no-op rather than a
+   * shrink. (If you ever need to shrink, use `Sandbox.setTimeout` /
+   * `SandboxApi.setTimeout`, which can extend or reduce.)
+   *
+   * Per-sandbox overrides are intentionally not exposed — this is a
+   * provider-level config only. Every sandbox managed by the provider
+   * refreshes by the same amount on each `get()`.
+   */
+  keepAliveMs?: number;
   /** Default outbound internet access policy */
   allowInternetAccess?: boolean;
   /** Default outbound network allow/deny rules */
@@ -34,6 +61,11 @@ export interface E2bSandboxConfig {
 export interface E2bSandboxCreateOptions extends SandboxCreateOptions {
   /** Sandbox template name or ID — overrides the provider default */
   template?: string;
-  /** Sandbox idle timeout in milliseconds — overrides the provider default */
+  /**
+   * Sandbox lifetime in milliseconds — overrides the provider default. See
+   * {@link E2bSandboxConfig.timeoutMs} for the full semantics; pair with
+   * the provider-level `keepAliveMs` to refresh on every `provider.get()`
+   * call.
+   */
   timeoutMs?: number;
 }

package/src/index.ts

@@ -45,7 +45,7 @@ export type { ModelInvoker, ModelInvokerConfig } from "./lib/model";
 
 // Activity-side handler wrappers
 export { withAutoAppend, withSandbox } from "./lib/tool-router";
-export type { SandboxContext } from "./lib/tool-router";
+export type { SandboxContext, WithSandboxOptions } from "./lib/tool-router";
 
 // Activity-side wrappers (requires Temporal client)
 export {

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

@@ -895,6 +895,64 @@ describe("createSession integration", () => {
     });
   });
 
+  it("embeds skill resourceContents on synthetic file tree entries via inlineContent", async () => {
+    const { ops } = createMockThreadOps();
+
+    const session = await createSession({
+      agentName: "TestAgent",
+      thread: { mode: "new", threadId: "thread-1" },
+      runAgent: createScriptedRunAgent([{ message: "done", toolCalls: [] }]),
+      threadOps: ops,
+      buildContextMessage: () => "go",
+      virtualFs: { ctx: { projectId: "p" } },
+      virtualFsOps: {
+        resolveFileTree: async () => ({ fileTree: [] }),
+      },
+      skills: [
+        {
+          name: "test-skill",
+          description: "Test",
+          instructions: "Do test",
+          location: "/skills/test-skill",
+          resourceContents: {
+            "references/alpha.md": "# Alpha doc",
+            "references/beta.md": "# Beta doc",
+          },
+        },
+      ],
+      hooks: {
+        onSessionStart: async () => {},
+      },
+    });
+
+    const stateManager = createAgentStateManager({
+      initialState: { systemPrompt: "test" },
+    });
+
+    await session.runSession({ stateManager });
+
+    const capturedFileTree = stateManager.getCurrentState().fileTree;
+    expect(Array.isArray(capturedFileTree)).toBe(true);
+    const entries = capturedFileTree as Array<{
+      path: string;
+      inlineContent?: string;
+    }>;
+
+    const alpha = entries.find(
+      (e) => e.path === "/skills/test-skill/references/alpha.md"
+    );
+    const beta = entries.find(
+      (e) => e.path === "/skills/test-skill/references/beta.md"
+    );
+    expect(alpha?.inlineContent).toBe("# Alpha doc");
+    expect(beta?.inlineContent).toBe("# Beta doc");
+
+    expect(stateManager.getCurrentState().inlineFiles).toEqual({
+      "/skills/test-skill/references/alpha.md": "# Alpha doc",
+      "/skills/test-skill/references/beta.md": "# Beta doc",
+    });
+  });
+
   it("does not pass initialFiles when skills have no resourceContents", async () => {
     const { ops } = createMockThreadOps();
     let capturedOptions: Record<string, unknown> | undefined;

package/src/lib/session/session.ts

@@ -1,5 +1,4 @@
 import {
-  condition,
   defineUpdate,
   setHandler,
   ApplicationFailure,
@@ -105,7 +104,6 @@ export async function createSession<
   processToolsInParallel = true,
   hooks = {},
   appendSystemPrompt = true,
-  waitForInputTimeout = "48h",
   threadKey,
   sandboxOps,
   thread: threadInit,
@@ -342,12 +340,23 @@ export async function createSession<
                 size: content.length,
                 mtime: new Date().toISOString(),
                 metadata: {},
+                // Carry the content directly on the entry so any handler that
+                // constructs a VirtualFileSystem from `fileTree` can read it
+                // without needing to also wire up `inlineFiles` from state.
+                inlineContent: content,
               })),
             ]
           : result.fileTree;
         stateManager.mergeUpdate({
           fileTree,
           virtualFsCtx: virtualFsConfig.ctx,
+          // `inlineFiles` is still the source of truth at read time:
+          // VirtualFileSystem checks the inlineFiles map first and only
+          // falls through to entry.inlineContent. Embedding the content on
+          // the entry is the migration target so that handlers building a
+          // VirtualFileSystem from `fileTree` alone (without forwarding
+          // `inlineFiles` from state) can read skill resources. Until a
+          // follow-up drops `inlineFiles`, both fields are populated.
           ...(skillFiles && { inlineFiles: skillFiles }),
         } as Partial<AgentState<TState>>);
       }
@@ -516,19 +525,6 @@ export async function createSession<
 
           // Turn committed: fresh id for the next turn.
           assistantId = undefined;
-
-          if (stateManager.getStatus() === "WAITING_FOR_INPUT") {
-            const conditionMet = await condition(
-              () => stateManager.getStatus() === "RUNNING",
-              waitForInputTimeout
-            );
-            if (!conditionMet) {
-              stateManager.cancel();
-              exitReason = "cancelled";
-              await condition(() => false, "2s");
-              break;
-            }
-          }
         }
 
         if (stateManager.getTurns() >= maxTurns && stateManager.isRunning()) {
@@ -618,6 +614,7 @@ export async function createSession<
           ...(sandboxId && { sandboxId }),
           ...(exitSnapshot && { snapshot: exitSnapshot }),
           threadId,
+          usage: stateManager.getTotalUsage(),
         });
       }
 

package/src/lib/session/types.ts

@@ -1,4 +1,3 @@
-import type { Duration } from "@temporalio/common";
 import type { SessionExitReason, ToolResultConfig } from "../types";
 import type {
   ToolMap,
@@ -177,8 +176,6 @@ export interface SessionConfig<
    * Returns SDK-native content for the initial human message.
    */
   buildContextMessage: () => TContent | Promise<TContent>;
-  /** How long to wait for input before cancelling the workflow */
-  waitForInputTimeout?: Duration;
 
   // ---------------------------------------------------------------------------
   // Thread lifecycle
@@ -248,6 +245,14 @@ export interface SessionConfig<
     threadId: string;
     sandboxId?: string;
     snapshot?: SandboxSnapshot;
+    usage: {
+      totalInputTokens: number;
+      totalOutputTokens: number;
+      totalCachedWriteTokens: number;
+      totalCachedReadTokens: number;
+      totalReasonTokens: number;
+      turns: number;
+    };
   }) => void;
 
   // ---------------------------------------------------------------------------

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

@@ -2602,6 +2602,7 @@ describe("defineSubagentWorkflow", () => {
           sandboxId: "sb-1",
           snapshot,
           threadId: "t",
+          usage: { totalInputTokens: 0, totalOutputTokens: 0, totalCachedWriteTokens: 0, totalCachedReadTokens: 0, totalReasonTokens: 0, turns: 0 },
         });
         return { toolResponse: "ok", data: null, threadId: "t" };
       }
@@ -2625,6 +2626,7 @@ describe("defineSubagentWorkflow", () => {
             createdAt: new Date().toISOString(),
           },
           threadId: "t",
+          usage: { totalInputTokens: 0, totalOutputTokens: 0, totalCachedWriteTokens: 0, totalCachedReadTokens: 0, totalReasonTokens: 0, turns: 0 },
         });
         return {
           toolResponse: "ok",

package/src/lib/subagent/types.ts

@@ -278,5 +278,13 @@ export interface SubagentSessionInput {
     sandboxId?: string;
     snapshot?: SandboxSnapshot;
     threadId: string;
+    usage: {
+      totalInputTokens: number;
+      totalOutputTokens: number;
+      totalCachedWriteTokens: number;
+      totalCachedReadTokens: number;
+      totalReasonTokens: number;
+      turns: number;
+    };
   }) => void;
 }

package/src/lib/subagent/workflow.ts

@@ -13,6 +13,7 @@ import type {
 } from "./types";
 import type { SubagentSandboxShutdown } from "../lifecycle";
 import type { SandboxSnapshot } from "../sandbox/types";
+import type { TokenUsage } from "../types";
 import { childSandboxReadySignal } from "./signals";
 
 /**
@@ -132,6 +133,7 @@ export function defineSubagentWorkflow(
     let capturedSnapshot: SandboxSnapshot | undefined;
     let capturedBaseSnapshot: SandboxSnapshot | undefined;
     let capturedThreadId: string | undefined;
+    let capturedUsage: TokenUsage | undefined;
     const sessionInput: SubagentSessionInput = {
       agentName: config.name,
       sandboxShutdown: effectiveShutdown,
@@ -148,10 +150,17 @@ export function defineSubagentWorkflow(
           });
         }
       },
-      onSessionExit: ({ sandboxId, snapshot, threadId }) => {
+      onSessionExit: ({ sandboxId, snapshot, threadId, usage }) => {
         capturedSandboxId = sandboxId;
         capturedSnapshot = snapshot;
         capturedThreadId = threadId;
+        capturedUsage = {
+          inputTokens: usage.totalInputTokens,
+          outputTokens: usage.totalOutputTokens,
+          cachedWriteTokens: usage.totalCachedWriteTokens,
+          cachedReadTokens: usage.totalCachedReadTokens,
+          reasonTokens: usage.totalReasonTokens,
+        };
       },
     };
 
@@ -168,6 +177,7 @@ export function defineSubagentWorkflow(
       ...(capturedBaseSnapshot !== undefined && {
         baseSnapshot: capturedBaseSnapshot,
       }),
+      ...(capturedUsage !== undefined && { usage: capturedUsage }),
     };
   };