@desplega.ai/agent-swarm 2.5.3 → 2.5.4

package/src/tests/context-bound-note.test.ts

@@ -0,0 +1,269 @@
+import { afterAll, beforeAll, describe, expect, spyOn, test } from "bun:test";
+import * as fsPromises from "node:fs/promises";
+import { mkdtemp, rm } from "node:fs/promises";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import {
+  claimContextBoundNote,
+  contextBoundNoteThreshold,
+  contextStatePath,
+  createContextBoundNote,
+  recordContextState,
+} from "../hooks/context-bound-note";
+import { renderContextBoundNote } from "../prompts/context-bound-note";
+import { normalizeClaudeMessage } from "../providers/claude-session-events";
+
+let dir: string;
+let n = 0;
+
+beforeAll(async () => {
+  dir = await mkdtemp(join(tmpdir(), "context-bound-note-"));
+});
+
+afterAll(async () => {
+  await rm(dir, { recursive: true, force: true });
+});
+
+/** A fresh task file path with `contextTokens` recorded next to it. */
+function taskFileWith(contextTokens?: number): string {
+  const taskFile = join(dir, `task-${++n}.json`);
+  if (contextTokens !== undefined) recordContextState(taskFile, contextTokens);
+  return taskFile;
+}
+
+describe("contextBoundNoteThreshold", () => {
+  test("off when unset, blank, invalid or 0", () => {
+    expect(contextBoundNoteThreshold({})).toBe(0);
+    expect(contextBoundNoteThreshold({ CONTEXT_BOUND_NOTE_TOKENS: "" })).toBe(0);
+    expect(contextBoundNoteThreshold({ CONTEXT_BOUND_NOTE_TOKENS: " " })).toBe(0);
+    expect(contextBoundNoteThreshold({ CONTEXT_BOUND_NOTE_TOKENS: "-5" })).toBe(0);
+    expect(contextBoundNoteThreshold({ CONTEXT_BOUND_NOTE_TOKENS: "lots" })).toBe(0);
+    expect(contextBoundNoteThreshold({ CONTEXT_BOUND_NOTE_TOKENS: "1e5" })).toBe(0);
+    expect(contextBoundNoteThreshold({ CONTEXT_BOUND_NOTE_TOKENS: "99999999999999999999" })).toBe(
+      0,
+    );
+    expect(contextBoundNoteThreshold({ CONTEXT_BOUND_NOTE_TOKENS: "0" })).toBe(0);
+  });
+
+  test("a positive integer turns the note on at that threshold", () => {
+    expect(contextBoundNoteThreshold({ CONTEXT_BOUND_NOTE_TOKENS: "150000" })).toBe(150_000);
+    expect(contextBoundNoteThreshold({ CONTEXT_BOUND_NOTE_TOKENS: " 200000 " })).toBe(200_000);
+  });
+});
+
+describe("claimContextBoundNote", () => {
+  test("fires once at the threshold, not again on later calls", async () => {
+    const taskFile = taskFileWith(200_000);
+    expect(await claimContextBoundNote(taskFile, 200_000)).toEqual({ contextTokens: 200_000 });
+    expect(await claimContextBoundNote(taskFile, 200_000)).toBeNull();
+    recordContextState(taskFile, 400_000);
+    expect(await claimContextBoundNote(taskFile, 200_000)).toBeNull();
+  });
+
+  test("concurrent hooks fire the note once", async () => {
+    const taskFile = taskFileWith(300_000);
+    const results = await Promise.all(
+      Array.from({ length: 8 }, () => claimContextBoundNote(taskFile, 200_000)),
+    );
+    expect(results.filter((r) => r !== null)).toEqual([{ contextTokens: 300_000 }]);
+  });
+
+  test("no-op below the threshold, then fires once it is crossed", async () => {
+    const taskFile = taskFileWith(199_999);
+    expect(await claimContextBoundNote(taskFile, 200_000)).toBeNull();
+    recordContextState(taskFile, 210_000);
+    expect(await claimContextBoundNote(taskFile, 200_000)).toEqual({ contextTokens: 210_000 });
+  });
+
+  test("returns the model the session recorded", async () => {
+    const taskFile = join(dir, `task-${++n}.json`);
+    recordContextState(taskFile, 250_000, "claude-opus-5-5");
+    expect(await claimContextBoundNote(taskFile, 200_000)).toEqual({
+      contextTokens: 250_000,
+      model: "claude-opus-5-5",
+    });
+  });
+
+  test("disabled at 0", async () => {
+    expect(await claimContextBoundNote(taskFileWith(900_000), 0)).toBeNull();
+  });
+
+  test("fails open on missing task file, missing state or corrupt state", async () => {
+    expect(await claimContextBoundNote(undefined, 200_000)).toBeNull();
+    expect(await claimContextBoundNote(taskFileWith(), 200_000)).toBeNull();
+
+    const corrupt = taskFileWith();
+    await Bun.write(contextStatePath(corrupt), "{not json");
+    expect(await claimContextBoundNote(corrupt, 200_000)).toBeNull();
+
+    const wrongShape = taskFileWith();
+    await Bun.write(contextStatePath(wrongShape), JSON.stringify({ contextTokens: "300000" }));
+    expect(await claimContextBoundNote(wrongShape, 200_000)).toBeNull();
+  });
+
+  test("fails open when the marker cannot be written", async () => {
+    const taskFile = taskFileWith(300_000);
+    const writeFile = spyOn(fsPromises, "writeFile").mockRejectedValue(
+      Object.assign(new Error("EIO: i/o error"), { code: "EIO" }),
+    );
+    try {
+      expect(await claimContextBoundNote(taskFile, 200_000)).toBeNull();
+      expect(writeFile).toHaveBeenCalledTimes(1);
+    } finally {
+      writeFile.mockRestore();
+    }
+  });
+});
+
+describe("createContextBoundNote", () => {
+  function observer(
+    threshold: number,
+    inject: (note: string) => unknown = () => {},
+    canDeliver: () => boolean = () => true,
+    model?: () => string | undefined,
+  ) {
+    const notes: string[] = [];
+    const observe = createContextBoundNote({
+      threshold,
+      canDeliver,
+      model,
+      inject: (note) => {
+        notes.push(note);
+        return inject(note);
+      },
+    });
+    return { observe, notes };
+  }
+
+  test("fires once at the threshold, not again on later requests", async () => {
+    const { observe, notes } = observer(200_000);
+    await observe(200_000);
+    await observe(400_000);
+    expect(notes).toEqual([await renderContextBoundNote(200_000)]);
+  });
+
+  test("no-op below the threshold, then fires once it is crossed", async () => {
+    const { observe, notes } = observer(200_000);
+    await observe(199_999);
+    expect(notes).toEqual([]);
+    await observe(210_000);
+    expect(notes).toEqual([await renderContextBoundNote(210_000)]);
+  });
+
+  test("disabled at 0", async () => {
+    const { observe, notes } = observer(0);
+    await observe(900_000);
+    expect(notes).toEqual([]);
+  });
+
+  test("fails open on missing or corrupt sizes", async () => {
+    const { observe, notes } = observer(200_000);
+    await observe(undefined);
+    await observe(Number.NaN);
+    await observe(Number.POSITIVE_INFINITY);
+    await observe("300000" as unknown as number);
+    expect(notes).toEqual([]);
+  });
+
+  test("never rejects when injection fails, and does not retry", async () => {
+    const { observe, notes } = observer(200_000, () => Promise.reject(new Error("closed")));
+    await expect(observe(300_000)).resolves.toBeUndefined();
+    await observe(300_000);
+    expect(notes).toHaveLength(1);
+  });
+
+  test("holds a note rendered after its window closed, then delivers it on the next call", async () => {
+    let open = false;
+    const { observe, notes } = observer(200_000, undefined, () => open);
+    await observe(300_000);
+    expect(notes).toEqual([]);
+    open = true;
+    await observe(500_000);
+    // The held note, not a re-render at the later size.
+    expect(notes).toEqual([await renderContextBoundNote(300_000)]);
+    await observe(600_000);
+    expect(notes).toHaveLength(1);
+  });
+
+  test("never delivers when the window stays closed", async () => {
+    const { observe, notes } = observer(200_000, undefined, () => false);
+    await observe(300_000);
+    await observe(400_000);
+    expect(notes).toEqual([]);
+  });
+
+  test("renders with the session's model at render time", async () => {
+    let model = "made-up-model";
+    const { observe, notes } = observer(200_000, undefined, undefined, () => model);
+    model = "claude-opus-5-5";
+    await observe(210_000);
+    expect(notes).toEqual([
+      "Context is 210,000 tokens (21% of the 1,000,000-token window). Keep this task bounded: finish the current step, store-progress, and complete or defer instead of starting new exploration.",
+    ]);
+  });
+
+  test("checks the window after rendering, not before", async () => {
+    let open = true;
+    const { observe, notes } = observer(200_000, undefined, () => open);
+    const pending = observe(300_000);
+    open = false;
+    await pending;
+    expect(notes).toEqual([]);
+  });
+});
+
+describe("renderContextBoundNote", () => {
+  const rest =
+    "Keep this task bounded: finish the current step, store-progress, and complete or defer instead of starting new exploration.";
+
+  test("renders the code default with a formatted token count and no model", async () => {
+    expect(await renderContextBoundNote(213_456)).toBe(`Context is 213,456 tokens. ${rest}`);
+  });
+
+  test("adds the window share for a known model", async () => {
+    expect(await renderContextBoundNote(210_000, "claude-opus-5-5")).toBe(
+      `Context is 210,000 tokens (21% of the 1,000,000-token window). ${rest}`,
+    );
+    expect(await renderContextBoundNote(200_000, "gpt-5")).toBe(
+      `Context is 200,000 tokens (50% of the 400,000-token window). ${rest}`,
+    );
+  });
+
+  test("omits the share for an unknown model instead of using the 200k fallback", async () => {
+    expect(await renderContextBoundNote(210_000, "made-up-model")).toBe(
+      `Context is 210,000 tokens. ${rest}`,
+    );
+  });
+});
+
+describe("normalizeClaudeMessage requestContextTokens", () => {
+  const state = {
+    taskId: "task-fixture",
+    agentId: "agent-fixture",
+    model: "claude-opus-5-5",
+    contextWindowSize: 1_000_000,
+    transport: "cli" as const,
+  };
+  const usage = {
+    input_tokens: 10,
+    cache_read_input_tokens: 180_000,
+    cache_creation_input_tokens: 25_000,
+    output_tokens: 4_000,
+  };
+
+  test("is input + cache read + cache write of a main-thread request", () => {
+    const normalized = normalizeClaudeMessage(
+      { type: "assistant", message: { content: [], usage } },
+      state,
+    );
+    expect(normalized.requestContextTokens).toBe(205_010);
+  });
+
+  test("is unset for subagent messages", () => {
+    const normalized = normalizeClaudeMessage(
+      { type: "assistant", parent_tool_use_id: "toolu_1", message: { content: [], usage } },
+      state,
+    );
+    expect(normalized.requestContextTokens).toBeUndefined();
+  });
+});

package/src/tests/extensions-pre-task.test.ts

@@ -22,9 +22,12 @@ import * as dbClient from "../be/db-client";
 import type { InstallExtensionArgs as ExtensionInstallBody } from "../be/extensions/db";
 import { installExtension, listExtensionRuns } from "../be/extensions/db";
 import { getCatalogEntry } from "../extensions/catalog";
+import type { TaskBlockedEvent } from "../extensions/contract";
 import { disableExtension, enableExtension, stopExtensionRuntime } from "../extensions/lifecycle";
+import { initExtensionPostBridge, teardownExtensionPostBridge } from "../extensions/post-bridge";
 import { handleTasks } from "../http/tasks";
 import { dispatchScheduleTarget } from "../scheduler/scheduler";
+import { bufferThreadMessage, instantFlush } from "../slack/thread-buffer";
 import { createTaskWithSiblingAwareness } from "../tasks/sibling-awareness";
 import { createWorkerTaskFollowUp } from "../tasks/worker-follow-up";
 import { registerSendTaskTool, sendTaskHandler } from "../tools/send-task";
@@ -145,6 +148,276 @@ describe("extension task boundaries", () => {
     await client.run("DELETE FROM agents");
   });
 
+  describe("post.task.blocked", () => {
+    const events: unknown[] = [];
+    const capture = (event: unknown) => {
+      events.push(event);
+    };
+
+    beforeAll(() => {
+      initExtensionPostBridge();
+      workflowEventBus.on("task.create_blocked", capture);
+    });
+
+    afterAll(() => {
+      workflowEventBus.off("task.create_blocked", capture);
+      teardownExtensionPostBridge();
+    });
+
+    beforeEach(() => {
+      events.length = 0;
+    });
+
+    async function waitForCounter(
+      extension: Awaited<ReturnType<typeof installAndEnable>>,
+      key: string,
+      count = 1,
+    ): Promise<void> {
+      const deadline = Date.now() + 3_000;
+      while (Date.now() < deadline) {
+        const entry = await getKv(`task:agent:${extension.agentId}`, key);
+        if (entry?.value === count) return;
+        await Bun.sleep(10);
+      }
+      throw new Error(`post.task.blocked was not recorded ${count} time(s)`);
+    }
+
+    async function recordedBlock(
+      extension: Awaited<ReturnType<typeof installAndEnable>>,
+      count = 1,
+    ): Promise<TaskBlockedEvent> {
+      await waitForCounter(extension, "ext:record-blocks:count", count);
+      return (await getKv(`task:agent:${extension.agentId}`, "ext:record-blocks:last"))
+        ?.value as TaskBlockedEvent;
+    }
+
+    test.each([
+      "rest",
+      "mcp",
+      "slack",
+      "schedule",
+    ] as const)("%s delivers one blocked event with its source and requester", async (origin) => {
+      const blocker = await installAndEnable("block-tasks-from-source", { source: origin });
+      const recorder = await installAndEnable("record-blocks");
+      let requester: TaskBlockedEvent["requester"] = {};
+      let scheduleId: string | undefined;
+
+      if (origin === "rest") {
+        const user = await createUser({ name: "Blocked REST requester" });
+        requester = { userId: user.id };
+        expect(
+          (await postTask({ task: "blocked REST task", requestedByUserId: user.id })).status,
+        ).toBe(422);
+      } else if (origin === "mcp") {
+        const caller = await createAgent({ name: "blocked-mcp-caller", status: "idle" });
+        requester = { agentId: caller.id };
+        const server = new McpServer({ name: "blocked-event-send", version: "1.0.0" });
+        registerSendTaskTool(server);
+        const result = await registeredSendTask(server).handler(sendTaskArgs("blocked MCP task"), {
+          sessionId: "blocked-event",
+          requestInfo: { headers: { "x-agent-id": caller.id } },
+        });
+        expect(result.isError).toBe(true);
+      } else if (origin === "slack") {
+        requester = { slackUserId: "U_BLOCKED_EVENT" };
+        await createAgent({ name: "blocked-slack-lead", isLead: true, status: "idle" });
+        bufferThreadMessage(
+          "C_BLOCKED_EVENT",
+          "1000.0001",
+          "blocked Slack task",
+          "U_BLOCKED_EVENT",
+          "1000.0002",
+        );
+        await instantFlush("C_BLOCKED_EVENT:1000.0001");
+      } else {
+        const schedule = await createScheduledTask({
+          name: "blocked-event-schedule",
+          taskTemplate: "blocked schedule task",
+          intervalMs: 60_000,
+        });
+        scheduleId = schedule.id;
+        await expect(dispatchScheduleTarget(schedule)).rejects.toThrow(
+          "Task source schedule is blocked",
+        );
+      }
+
+      const event = await recordedBlock(recorder);
+      expect(events).toHaveLength(1);
+      expect(event).toMatchObject({
+        extension: { id: blocker.id, name: "block-tasks-from-source" },
+        origin,
+        source: origin === "rest" ? "api" : origin,
+        requester,
+      });
+      expect(event.requester).toEqual(requester);
+      if (scheduleId) expect(event.context?.scheduleId).toBe(scheduleId);
+      expect(event).not.toHaveProperty("task");
+      expect(event).not.toHaveProperty("taskId");
+    });
+
+    test("scrubs and bounds the payload, preserving correlation on each retry", async () => {
+      const secret = `ghp_${"Z".repeat(36)}`;
+      await installSource(
+        "secret-blocker",
+        `
+import { block, type SwarmExtension } from "swarm-extension";
+const extension: SwarmExtension = (api) => {
+  api.on("pre.task.create", () => block(${JSON.stringify(`Ticket required: ${secret}`)}));
+};
+export default extension;
+`,
+      );
+      const recorder = await installAndEnable("record-blocks");
+      const context = {
+        contextKey: "github:acme/project:7",
+        scheduleId: crypto.randomUUID(),
+        workflowRunId: crypto.randomUUID(),
+        workflowRunStepId: crypto.randomUUID(),
+        parentTaskId: crypto.randomUUID(),
+        vcsRepo: "acme/project",
+        vcsNumber: 7,
+        slackChannelId: "C_CORRELATED_BLOCK",
+        slackThreadTs: "2000.0001",
+      };
+      const requester = {
+        userId: crypto.randomUUID(),
+        agentId: crypto.randomUUID(),
+        slackUserId: "U_CORRELATED_BLOCK",
+      };
+      for (const attempt of [1, 2]) {
+        await expect(
+          createTaskWithSiblingAwareness(
+            `${"x".repeat(480)} ${secret}`,
+            {
+              ...context,
+              source: "github",
+              requestedByUserId: requester.userId,
+              creatorAgentId: requester.agentId,
+              slackUserId: requester.slackUserId,
+            },
+            { origin: "webhook" },
+          ),
+        ).rejects.toThrow("Ticket required:");
+        const event = await recordedBlock(recorder, attempt);
+        expect(events).toHaveLength(attempt);
+        expect(event.reason).toContain("[REDACTED:");
+        expect(event.reason).not.toContain(secret);
+        expect(event.description).toHaveLength(500);
+        expect(event.description).not.toContain("ghp_");
+        expect(event.context).toEqual(context);
+        expect(event.requester).toEqual(requester);
+      }
+    });
+
+    test("a rejected model rewrite emits a block attributed to its editor", async () => {
+      const codex = await createAgent({
+        name: "blocked-rewrite-codex",
+        harnessProvider: "codex",
+        status: "idle",
+      });
+      const editor = await installSource(
+        "incompatible-model",
+        `
+import { modify, type SwarmExtension } from "swarm-extension";
+const extension: SwarmExtension = (api) => {
+  api.on("pre.task.create", () => modify({ model: "claude-opus-5-5" }));
+};
+export default extension;
+`,
+      );
+      const recorder = await installAndEnable("record-blocks");
+      const result = await postTask({
+        task: "incompatible rewrite",
+        agentId: codex.id,
+        routingReason: "human_pinned",
+      });
+      expect(result.status).toBe(422);
+      const event = await recordedBlock(recorder);
+      expect(events).toHaveLength(1);
+      expect(event.extension).toEqual({ id: editor.id, name: "incompatible-model" });
+      expect(event.reason).toContain("rewrite");
+      expect(event.reason).toContain("codex");
+    });
+
+    test("the creator is skipped while the blocking extension receives its own block", async () => {
+      const creator = await installAndEnable("record-blocks");
+      const blocker = await installSource(
+        "blocking-recorder",
+        `
+import { block, type SwarmExtension } from "swarm-extension";
+const extension: SwarmExtension = (api) => {
+  api.on("pre.task.create", () => block("blocked by recorder"));
+  api.on("post.task.blocked", async (_event, ctx) => { await ctx.state.incr("count"); });
+};
+export default extension;
+`,
+      );
+      const observerBundle = await loadBundleFixture("record-blocks");
+      const installedObserver = await installExtension({
+        ...observerBundle,
+        manifest: { ...observerBundle.manifest, name: "last-block-observer" },
+        priority: 1_000,
+      });
+      const observer = await enableExtension(installedObserver.extension.id);
+      const server = new McpServer({ name: "extension-blocked-event", version: "1.0.0" });
+      registerSendTaskTool(server);
+      const result = await registeredSendTask(server).handler(
+        sendTaskArgs("extension task"),
+        markExtensionRequestOrigin({
+          sessionId: "blocked-extension",
+          requestInfo: { headers: { "x-agent-id": creator.agentId! } },
+        }),
+      );
+      expect(result.isError).toBe(true);
+      // The final subscriber acknowledges only after every earlier recipient has finished.
+      await waitForCounter(observer, "ext:last-block-observer:count");
+      expect(
+        (await getKv(`task:agent:${blocker.agentId}`, "ext:blocking-recorder:count"))?.value,
+      ).toBe(1);
+      expect(events).toHaveLength(1);
+      expect(events[0]).toMatchObject({ origin: "extension:record-blocks" });
+      expect(await getKv(`task:agent:${creator.agentId}`, "ext:record-blocks:count")).toBeNull();
+      expect(await listExtensionRuns(creator.id)).toEqual([]);
+      const runs = await listExtensionRuns(blocker.id);
+      expect(runs).toContainEqual(
+        expect.objectContaining({
+          event: "post.task.blocked",
+          agentId: creator.agentId,
+          subject: "extension:record-blocks: blocking-recorder",
+        }),
+      );
+    });
+
+    test("continue and valid rewrites emit no blocked event", async () => {
+      const recorder = await installAndEnable("record-blocks");
+      expect((await postTask({ task: "continued task" })).status).toBe(201);
+      await installAndEnable("rewrite-task-priority");
+      const modified = await postTask({ task: "rewritten task" });
+      expect(modified.status).toBe(201);
+      expect(modified.body.priority).toBe(1);
+      expect(events).toEqual([]);
+      expect(await getKv(`task:agent:${recorder.agentId}`, "ext:record-blocks:count")).toBeNull();
+    });
+
+    test("a throwing bus observer does not replace the blocked response", async () => {
+      await installAndEnable("block-tasks-from-source", { source: "rest" });
+      let calls = 0;
+      const observer = () => {
+        calls++;
+        throw new Error("broken block observer");
+      };
+      const error = spyOn(console, "error").mockImplementation(() => {});
+      workflowEventBus.on("task.create_blocked", observer);
+      try {
+        expect((await postTask({ task: "still blocked" })).status).toBe(422);
+        expect(calls).toBe(1);
+      } finally {
+        workflowEventBus.off("task.create_blocked", observer);
+        error.mockRestore();
+      }
+    });
+  });
+
   test("REST task creation applies rewrites and reports blocks as 422", async () => {
     const rewrite = await installAndEnable("rewrite-task-priority");
     const created = await postTask({ task: "rewrite this task" });

package/src/tests/fetch-resolved-env.test.ts

@@ -14,6 +14,7 @@ import {
   fetchResolvedEnv,
   RELOADABLE_ENV_KEYS,
 } from "../commands/runner";
+import { contextBoundNoteThreshold } from "../hooks/context-bound-note";
 import { checkCodexCredentials } from "../providers/codex-adapter";
 import { MASKED_SECRET_VALUE } from "../utils/config-secrets";
 import { registerVolatileSecret } from "../utils/secret-scrubber";
@@ -795,6 +796,22 @@ describe("agent-scoped PI_CODEMODE_MODELS reaches the worker env", () => {
   });
 });
 
+describe("agent-scoped CONTEXT_BOUND_NOTE_TOKENS turns the note on for one agent", () => {
+  // Each session's env is this agent's /api/config/resolved result (global,
+  // then agent, then repo rows), and every adapter reads the threshold from it.
+  test("the agent with the row gets the threshold; an agent without it stays off", async () => {
+    mockResponsesByAgentId.set("bound-note-pilot", {
+      status: 200,
+      body: { configs: [{ key: "CONTEXT_BOUND_NOTE_TOKENS", value: "150000" }] },
+    });
+    const baseEnv: Record<string, string | undefined> = {};
+    const pilot = await fetchResolvedEnv(testUrl, "key", "bound-note-pilot", baseEnv);
+    const other = await fetchResolvedEnv(testUrl, "key", "bound-note-other", baseEnv);
+    expect(contextBoundNoteThreshold(pilot.env)).toBe(150_000);
+    expect(contextBoundNoteThreshold(other.env)).toBe(0);
+  });
+});
+
 describe("deleting a reloadable config row", () => {
   // One key per test: the runner remembers each overridden key's boot value
   // for the life of the process.

package/src/tests/fixtures/extensions/record-blocks/hooks.ts

@@ -0,0 +1,10 @@
+import type { SwarmExtension } from "swarm-extension";
+
+const extension: SwarmExtension = (api) => {
+  api.on("post.task.blocked", async (event, ctx) => {
+    await ctx.state.set("last", event);
+    await ctx.state.incr("count");
+  });
+};
+
+export default extension;

package/src/tests/fixtures/extensions/record-blocks/manifest.json

@@ -0,0 +1,7 @@
+{
+  "name": "record-blocks",
+  "description": "Records blocked task attempts",
+  "version": "1.0.0",
+  "runtime": "api",
+  "assets": { "hooks": "hooks.ts" }
+}

package/src/tests/fixtures/fake-amp.ts

@@ -104,12 +104,19 @@ async function main() {
   );
 
   const sessionId = "T-fake-session";
-  const usage = {
-    input_tokens: 4,
-    output_tokens: 5,
-    cache_creation_input_tokens: 1000,
-    cache_read_input_tokens: 200,
-  };
+  // `AMP_TEST_REQUEST_USAGE`: "missing" drops per-request usage, "corrupt" sends non-numbers.
+  const requestUsage = process.env.AMP_TEST_REQUEST_USAGE;
+  const usage =
+    requestUsage === "missing"
+      ? undefined
+      : requestUsage === "corrupt"
+        ? { input_tokens: "lots", cache_creation_input_tokens: null, cache_read_input_tokens: {} }
+        : {
+            input_tokens: 4,
+            output_tokens: 5,
+            cache_creation_input_tokens: 1000,
+            cache_read_input_tokens: 200,
+          };
   // Amp runs shell commands in a session of its own, outside the process group.
   // Return only once the child has left: a group kill that lands before its
   // setsid() would take it down and hide a leak.
@@ -217,8 +224,13 @@ async function main() {
       session_id: sessionId,
     });
     turns += 1;
-    const steerTool = mode === "steer" || mode.startsWith("steer-crash");
-    if (mode === "tool" || (steerTool && index === 0)) {
+    const steerTool = mode === "steer" || mode === "steer-twice" || mode.startsWith("steer-crash");
+    const toolCalls = mode === "steer-twice" ? 2 : 1;
+    for (
+      let call = 0;
+      call < toolCalls && (mode === "tool" || (steerTool && index === 0));
+      call++
+    ) {
       assistant(
         [{ type: "tool_use", id: "TU-1", name: "code_exec", input: { code: "1+1" } }],
         "tool_use",
@@ -261,6 +273,8 @@ async function main() {
     if (mode === "success-child") await startDetachedChild();
     // A turn that thinks for a while with no tool call.
     if (mode === "steer-no-tool") await Bun.sleep(600);
+    // `AMP_TEST_POST_RESULT_DELAY_MS`: the final reply takes this long after the last tool result.
+    await Bun.sleep(Number(process.env.AMP_TEST_POST_RESULT_DELAY_MS) || 0);
     lastText = mode.startsWith("steer") ? `DONE${index + 1}` : "Done ✓";
     assistant([{ type: "text", text: lastText }], "end_turn");
   };