test(e2e): cover a denied worker write approved once and retried

Author: bwachman

e2e/fleet.ts

@@ -32,11 +32,44 @@ export const WaitAgentsResult = type({
     "error?": "string",
     "question?": "string",
     "question_id?": "string",
+    "grant_request_id?": "string",
     "report?": "string",
   }).array(),
   timed_out: "boolean",
 });
 
+/** Extract the public worker-grant correlation from serialized tool output. */
+export function waitAgentsGrantRequestId(body: string): string | undefined {
+  const payload = JSON.parse(body) as unknown;
+  const contents = serializedContents(payload);
+  for (const content of contents) {
+    try {
+      const result = WaitAgentsResult.assert(content);
+      const grantRequestId = result.results.find(
+        (lane) => lane.grant_request_id !== undefined,
+      )?.grant_request_id;
+      if (grantRequestId !== undefined) return grantRequestId;
+    } catch {
+      // Request messages that are not wait_agents results are irrelevant here.
+    }
+  }
+  return undefined;
+}
+
+function serializedContents(value: unknown): unknown[] {
+  if (typeof value === "string") {
+    try {
+      return [JSON.parse(value)];
+    } catch {
+      return [];
+    }
+  }
+  if (Array.isArray(value)) return value.flatMap(serializedContents);
+  if (value !== null && typeof value === "object")
+    return Object.values(value).flatMap(serializedContents);
+  return [];
+}
+
 /** Every wait_agents tool result in the event stream, in call order. */
 export function waitAgentsResults(
   events: ReactorEmittedEvent[],

e2e/harness.ts

@@ -14,6 +14,11 @@ import { dirname, join } from "node:path";
 
 import type { ReactorEmittedEvent } from "@intx/inference";
 import type { InboundMessage } from "@intx/types/runtime";
+import {
+  wire,
+  type HarnessRequest,
+  type RequestPredicate,
+} from "@intx/inference-testing";
 
 import { OPERATOR_ORIGINATED_FLAG } from "../src/agent/message-provenance.js";
 import {
@@ -154,6 +159,50 @@ export interface ScriptedReply {
   toolCalls?: { name: string; args: Record<string, unknown> }[];
 }
 
+export interface RuntimeReplyRequest {
+  body: string;
+  request: HarnessRequest;
+}
+
+/**
+ * Register one response whose payload is built when its matching request
+ * arrives. This keeps correlation values discovered during a scenario out of
+ * eagerly serialized `replyOnce` fixtures.
+ */
+export function replyOnceAtRequest(
+  session: E2ESession,
+  provider: Parameters<typeof wire.completeResponse>[0],
+  predicate: RequestPredicate,
+  reply: (request: RuntimeReplyRequest) => ScriptedReply | undefined,
+): void {
+  const stream = session.harness.scenario.createStream();
+  session.harness.scenario.whenRequestBodyMatches((body, request) => {
+    if (!predicate(request)) return false;
+    const response = reply({ body, request });
+    if (response === undefined) return false;
+    const chunks = wire.completeResponse(provider, {
+      text: response.text ?? "",
+      ...(response.toolCalls === undefined
+        ? {}
+        : {
+            toolCalls: response.toolCalls.map((call, index) => ({
+              callId: `runtime_call_${crypto.randomUUID()}_${index}`,
+              name: call.name,
+              argsJSON: JSON.stringify(call.args),
+            })),
+          }),
+    });
+    stream.enqueueAll(chunks, { startAt: session.harness.clock.now() + 1 });
+    // A request-time reply can land after runUntilDone's harness.run has
+    // gone quiescent (e.g. a worker request while the parent waits on the
+    // wall clock), so drive the clock for the chunks just scheduled.
+    void session.harness
+      .run({ wallClockBudgetMs: Infinity })
+      .catch(() => undefined);
+    return true;
+  }, stream);
+}
+
 /**
  * Queue one model response per inference request, in order — the common
  * "model calls tools, then answers" spine. Use

e2e/subagent-orchestration.test.ts

@@ -1,10 +1,14 @@
 import { describe, expect, test } from "bun:test";
+import { join } from "node:path";
 import { type } from "arktype";
 
+import { createPermissionGate } from "../src/permission/gate.js";
+
 import { createSubAgentSessionStore } from "../src/subagent/index.js";
 import { withMockedModuleDuring } from "../testkit/mock-module.js";
 import {
   fromHost,
+  waitAgentsGrantRequestId,
   waitAgentsResults,
   WORKER_HOST,
   WORKER_PROVIDER,
@@ -13,17 +17,43 @@ import {
   closeE2ESession,
   e2ePermissionGate,
   openE2ESession,
+  replyOnceAtRequest,
   runUntilDone,
 } from "./harness.js";
 import { DIRECTOR_RESOLVABLE_ESCALATION } from "../testkit/escalation.js";
 
 const RequestURL = type({ url: "string" });
+type AuthorizationCheckpoint =
+  | "denied-write"
+  | "asked-for-approval"
+  | "approval-issued"
+  | "authorized-retry";
+
+test("wait_agents grant correlation accepts pretty serialized tool output", () => {
+  const result = JSON.stringify(
+    {
+      results: [
+        {
+          agent_id: "lane-1",
+          grant_request_id: "grant-1",
+          status: "awaiting_director",
+        },
+      ],
+      timed_out: false,
+    },
+    null,
+    2,
+  );
+  const request = JSON.stringify({ messages: [{ content: result }] });
+  expect(waitAgentsGrantRequestId(request)).toBe("grant-1");
+});
 
 async function openFleetSession(
   sessions: ReturnType<typeof createSubAgentSessionStore>,
+  permissionGate = e2ePermissionGate(),
 ) {
   return openE2ESession({
-    permissionGate: e2ePermissionGate(),
+    permissionGate,
     subAgent: {
       provider: WORKER_PROVIDER,
       sessions,
@@ -140,6 +170,194 @@ describe("e2e — fleet orchestration", () => {
     60000,
   );
 
+  test.serial(
+    "a real denied worker write is approved once with live correlation IDs and retried once",
+    async () => {
+      const fleetSessions = createSubAgentSessionStore();
+      const session = await openFleetSession(
+        fleetSessions,
+        createPermissionGate({
+          approvals: [
+            { tool: "spawn_agent", pattern: "*" },
+            { tool: "wait_agents", pattern: "*" },
+            { tool: "approve_worker_call", pattern: "*" },
+          ],
+          interactive: true,
+          skipPermissions: false,
+          reactorGated: true,
+          requestApproval: async () => ({ allow: true }),
+        }),
+      );
+      try {
+        const parent = fromHost("api.anthropic.com");
+        const worker = fromHost(WORKER_HOST);
+        const output = "approved.txt";
+        const checkpoints: AuthorizationCheckpoint[] = [];
+        let approvedRequestId: string | undefined;
+        session.harness.scenario.replyOnce("anthropic", {
+          predicate: parent,
+          toolCalls: [
+            {
+              name: "spawn_agent",
+              args: {
+                description: "authorized write lane",
+                prompt:
+                  "Write approved.txt with approved, ask for authorization if denied, then report.",
+                intent: "implement",
+                success_criteria: ["approved.txt contains approved"],
+              },
+            },
+          ],
+        });
+        session.harness.scenario.replyOnce("openai", {
+          predicate: worker,
+          toolCalls: [
+            { name: "write_file", args: { path: output, content: "approved" } },
+          ],
+        });
+        replyOnceAtRequest(session, "openai", worker, ({ body }) => {
+          const grantRequestId = /grant request ([\w-]+)/.exec(body)?.[1];
+          if (grantRequestId === undefined) {
+            throw new Error(
+              `worker did not receive the denied-call grant marker: ${body}`,
+            );
+          }
+          checkpoints.push("denied-write");
+          return {
+            toolCalls: [
+              {
+                name: "ask_director",
+                args: {
+                  question: "Please approve the denied write.",
+                  grant_request_id: grantRequestId,
+                  escalation: DIRECTOR_RESOLVABLE_ESCALATION,
+                },
+              },
+            ],
+          };
+        });
+        replyOnceAtRequest(session, "openai", worker, ({ body }) =>
+          body.includes("Parent approved the exact denied call")
+            ? (checkpoints.push("authorized-retry"),
+              {
+                toolCalls: [
+                  {
+                    name: "write_file",
+                    args: { path: output, content: "approved" },
+                  },
+                ],
+              })
+            : undefined,
+        );
+        replyOnceAtRequest(session, "openai", worker, ({ body }) =>
+          body.includes("wrote 8 bytes to")
+            ? {
+                text: "## Summary\nWrite completed.\n## Findings\nAuthorized retry succeeded.\n## Blockers\nNone.\n## Paths\napproved.txt",
+              }
+            : undefined,
+        );
+        session.harness.scenario.replyOnce("anthropic", {
+          predicate: parent,
+          toolCalls: [{ name: "wait_agents", args: { timeout_ms: 30_000 } }],
+        });
+        session.harness.scenario.replyOnce("anthropic", {
+          predicate: parent,
+          toolCalls: [{ name: "wait_agents", args: { timeout_ms: 30_000 } }],
+        });
+        replyOnceAtRequest(session, "anthropic", parent, ({ body }) => {
+          const lane = fleetSessions.list()[0];
+          const ask =
+            lane === undefined ? undefined : fleetSessions.peekAsk(lane.id);
+          const requestId = ask?.deniedCall?.requestId;
+          const publicRequestId = waitAgentsGrantRequestId(body);
+          if (
+            lane === undefined ||
+            ask?.questionId === undefined ||
+            requestId === undefined ||
+            publicRequestId === undefined
+          ) {
+            return undefined;
+          }
+          if (publicRequestId !== requestId) {
+            throw new Error(
+              `wait_agents grant_request_id ${publicRequestId} did not match pending ask ${requestId}`,
+            );
+          }
+          checkpoints.push("asked-for-approval");
+          approvedRequestId = publicRequestId;
+          const args = {
+            target: lane.id,
+            question_id: ask.questionId,
+            request_id: publicRequestId,
+          };
+          return { toolCalls: [{ name: "approve_worker_call", args }] };
+        });
+        replyOnceAtRequest(session, "anthropic", parent, ({ body }) =>
+          /\\?"status\\?":\s*\\?"authorized/.test(body)
+            ? (checkpoints.push("approval-issued"),
+              {
+                toolCalls: [
+                  { name: "wait_agents", args: { timeout_ms: 30_000 } },
+                ],
+              })
+            : undefined,
+        );
+        replyOnceAtRequest(session, "anthropic", parent, ({ body }) =>
+          body.includes("Write completed")
+            ? { text: "The approved worker write completed." }
+            : undefined,
+        );
+
+        const result = await withMockedModuleDuring(
+          import.meta.resolve("../src/session/assemble-runtime.js"),
+          (real: typeof import("../src/session/assemble-runtime.js")) => ({
+            ...real,
+            assembleInferenceBase: async () => session.harness.deps,
+          }),
+          () => runUntilDone(session, "Run the authorized write"),
+        );
+
+        const waits = waitAgentsResults(result.events);
+        const parked = waits[0];
+        const settled = waits.at(-1);
+        const lane = parked?.results[0];
+        expect(checkpoints.slice(0, 2)).toEqual([
+          "denied-write",
+          "asked-for-approval",
+        ]);
+        expect(checkpoints.slice(2).sort()).toEqual([
+          "approval-issued",
+          "authorized-retry",
+        ]);
+        expect(lane?.status).toBe("awaiting_director");
+        expect(lane?.question_id).toBeString();
+        expect(lane?.grant_request_id).toBeString();
+        expect(lane?.grant_request_id).toBe(approvedRequestId);
+        expect(settled?.results[0]?.status).toBe("done");
+        expect(await Bun.file(join(session.cwd, output)).text()).toBe(
+          "approved",
+        );
+        const approvals = result.events.filter(
+          (event) =>
+            event.type === "tool.done" &&
+            typeof event.data.result.content === "string" &&
+            /"status":\s*"authorized"/.test(event.data.result.content),
+        );
+        expect(approvals).toHaveLength(1);
+        const workerRequests = session.harness.scenario
+          .matchedRequests()
+          .filter((request) =>
+            RequestURL.assert(request).url.includes(WORKER_HOST),
+          );
+        expect(workerRequests).toHaveLength(4);
+        expect(result.reply).toContain("completed");
+      } finally {
+        await closeE2ESession(session);
+      }
+    },
+    60000,
+  );
+
   test.serial(
     "a retryable first send triggers the outer retry and the worker recovers",
     async () => {