src/agents/tools/agents-wait-tool.test.ts

import { expectDefined } from "@openclaw/normalization-core";
import ts from "typescript";
import { beforeEach, describe, expect, it, onTestFinished, vi } from "vitest";
import { applyCodeModeCatalog } from "../code-mode.js";
import {
  createCodeModeHarness,
  resetCodeModeTestState,
  runUntilCompleted,
} from "../code-mode.test-support.js";
import type { SubagentRunRecord } from "../subagents/registry/subagent-registry.types.js";

const records = new Map<string, SubagentRunRecord>();
const registryEvents = vi.hoisted(() => ({ listeners: new Set<() => void>() }));

vi.mock("../subagents/registry/subagent-registry.js", () => ({
  getSubagentRunsByRunIds: (runIds: readonly string[]) => ({
    entries: new Map(
      runIds.flatMap((runId) => {
        const entry =
          records.get(runId) ??
          [...records.values()].find((candidate) => candidate.swarmRunId === runId);
        return entry ? [[runId, entry] as const] : [];
      }),
    ),
  }),
}));

vi.mock("../subagents/registry/subagent-registry-state.js", () => ({
  onSubagentRegistryPersisted: (listener: () => void) => {
    registryEvents.listeners.add(listener);
    return () => registryEvents.listeners.delete(listener);
  },
}));

import { isToolResultError } from "../tool-result-error.js";
import { createAgentsWaitTool, waitForCollectorCompletion } from "./agents-wait-tool.js";

function createMainSessionWaitTool() {
  return createAgentsWaitTool({
    agentSessionKey: "agent:main:main",
    agentId: "main",
    config: { tools: { swarm: true } },
  });
}

function collectorRun(
  runId: string,
  requesterSessionKey: string,
  completion?: SubagentRunRecord["collectorCompletion"],
): SubagentRunRecord {
  return {
    runId,
    childSessionKey: `agent:worker:subagent:${runId}`,
    controllerSessionKey: requesterSessionKey,
    requesterSessionKey,
    requesterDisplayKey: requesterSessionKey,
    task: runId,
    cleanup: "keep",
    createdAt: Date.now(),
    execution: { status: completion ? "terminal" : "running" },
    collect: true,
    swarmRequesterSessionKey: requesterSessionKey,
    groupId: "group",
    completion: { required: false, resultText: completion ? `result-${runId}` : undefined },
    collectorCompletion: completion,
  };
}

describe("agents_wait", () => {
  beforeEach(() => {
    records.clear();
    registryEvents.listeners.clear();
  });

  it("composes real collector outputs through discovery, describe, and generated declarations", async () => {
    onTestFinished(resetCodeModeTestState);
    const entry = collectorRun("ready", "agent:main:main", {
      status: "done",
      structured: { answer: 42 },
    });
    records.set(entry.runId, entry);
    const h = createCodeModeHarness();
    const tool = createMainSessionWaitTool();
    applyCodeModeCatalog({ ...h.ctx, tools: [...h.tools, tool] });
    const execTool = expectDefined(h.tools[0], "Code Mode exec");
    const waitTool = expectDefined(h.tools[1], "Code Mode wait");
    expect(execTool.description).toContain("completed: Array<");
    const listed = await runUntilCompleted({
      execTool,
      waitTool,
      code: 'return await API.list("tools");',
    });
    expect(listed).toMatchObject({
      status: "completed",
      value: { files: [{ path: "tools/agents_wait.d.ts" }] },
      telemetry: { describeCount: 0, callCount: 0 },
    });
    const result = await runUntilCompleted({
      execTool,
      waitTool,
      code: 'const [wait] = await catalog.search("agents_wait"); const description = await wait.describe(); const file = await API.read("tools/agents_wait.d.ts"); const result = await wait({ids:["ready"],timeoutSeconds:0}); return {description, file, ids:result.completed.map(item=>item.runId), structured:result.completed[0].structured};',
    });
    expect(result).toMatchObject({
      status: "completed",
      telemetry: { describeCount: 2, callCount: 1 },
      value: {
        ids: ["ready"],
        structured: { answer: 42 },
        description: { outputSchema: tool.outputSchema },
      },
    });
    const { file } = result.value as { file: { content: string } };
    // Consume the intact guest declaration, with opaque collector-specific output.
    expect(file.content).not.toContain("truncated: true");
    const fileName = "/collector-consumer.ts";
    const source = ts.createSourceFile(
      fileName,
      file.content +
        "\n" +
        [
          "async function consume() {",
          'const result = await agents_wait({ids:["ready"]});',
          "const ids: string[] = result.completed.map(item => item.runId);",
          "// @ts-expect-error No invented builds field.",
          "result.builds.map(item => item.id);",
          "// @ts-expect-error Collector structured output is unknown without its own schema.",
          "result.completed[0].structured.answer;",
          "return ids;",
          "}",
          "// @ts-expect-error Required ids stay required.",
          "agents_wait({});",
        ].join("\n"),
      ts.ScriptTarget.ESNext,
      true,
    );
    const options = { noEmit: true, strict: true, types: [], target: ts.ScriptTarget.ESNext };
    const host = ts.createCompilerHost(options);
    const original = host.getSourceFile.bind(host);
    host.getSourceFile = (name, ...args) => (name === fileName ? source : original(name, ...args));
    const program = ts.createProgram([fileName], options, host);
    expect(
      ts
        .getPreEmitDiagnostics(program)
        .map((diagnostic) => ts.flattenDiagnosticMessageText(diagnostic.messageText, "\n")),
    ).toEqual([]);
  });

  it("settles a parked collector bridge from a registry write event", async () => {
    const entry = collectorRun("event-driven", "agent:main:main");
    records.set(entry.runId, entry);
    const completion = waitForCollectorCompletion({
      runId: entry.runId,
      currentSessionKeys: new Set(["agent:main:main"]),
    });

    entry.completion = { required: false, resultText: "event result" };
    entry.collectorCompletion = { status: "done" };
    for (const listener of registryEvents.listeners) {
      listener();
    }

    await expect(completion).resolves.toMatchObject({
      runId: "event-driven",
      status: "done",
      result: "event result",
    });
    expect(registryEvents.listeners.size).toBe(0);
  });

  it("returns retained visible output when a successful collector ends with NO_REPLY", async () => {
    const entry = collectorRun("retained", "agent:main:main", { status: "done" });
    entry.execution = { status: "terminal", outcome: { status: "ok" } };
    entry.completion = {
      required: false,
      resultText: "NO_REPLY",
      fallbackResultText: "retained collector result",
    };
    records.set(entry.runId, entry);

    await expect(
      waitForCollectorCompletion({
        runId: entry.runId,
        currentSessionKeys: new Set(["agent:main:main"]),
      }),
    ).resolves.toMatchObject({ result: "retained collector result" });
  });

  it("rejects when abort wins the listener-registration race", async () => {
    const entry = collectorRun("abort-race", "agent:main:main");
    records.set(entry.runId, entry);
    const controller = new AbortController();
    const originalAddEventListener = controller.signal.addEventListener.bind(controller.signal);
    vi.spyOn(controller.signal, "addEventListener").mockImplementation((...args) => {
      controller.abort();
      originalAddEventListener(...args);
    });

    await expect(
      waitForCollectorCompletion({
        runId: entry.runId,
        currentSessionKeys: new Set(["agent:main:main"]),
        signal: controller.signal,
      }),
    ).rejects.toThrow("agents.run wait aborted");
    expect(registryEvents.listeners.size).toBe(0);
  });

  it("returns the first completed child and leaves siblings pending", async () => {
    records.set("one", collectorRun("one", "agent:main:main"));
    records.set("two", collectorRun("two", "agent:main:main"));
    const tool = createMainSessionWaitTool();
    setTimeout(() => {
      const entry = records.get("two");
      if (!entry) {
        return;
      }
      entry.completion = { required: false, resultText: "result-two" };
      entry.collectorCompletion = {
        status: "done",
        structured: { winner: 2 },
      };
      for (const listener of registryEvents.listeners) {
        listener();
      }
    }, 5);

    const result = await tool.execute("call", { ids: ["one", "two"], timeoutSeconds: 1 });
    expect(result.details).toEqual({
      completed: [
        {
          runId: "two",
          status: "done",
          result: "result-two",
          structured: { winner: 2 },
          sessionKey: "agent:worker:subagent:two",
        },
      ],
      pending: ["one"],
    });
  });

  it("parks without reading until a registry mutation wakes it", async () => {
    vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout", "performance"] });
    const entry = collectorRun("local-wake", "agent:main:main");
    records.set(entry.runId, entry);
    const controller = new AbortController();
    const tool = createMainSessionWaitTool();
    const reads = vi.spyOn(records, "get");
    let result: unknown;
    const waiting = tool
      .execute("call", { ids: [entry.runId], timeoutSeconds: 1 }, controller.signal)
      .then((value) => {
        result = value.details;
      });
    try {
      const initialReads = reads.mock.calls.length;
      await vi.advanceTimersByTimeAsync(750);
      expect(reads).toHaveBeenCalledTimes(initialReads);
      entry.collectorCompletion = { status: "done" };
      for (const listener of registryEvents.listeners) {
        listener();
      }
      await vi.advanceTimersByTimeAsync(0);
      expect(result).toMatchObject({ completed: [{ runId: entry.runId }], pending: [] });
      expect(registryEvents.listeners.size).toBe(0);
      expect(vi.getTimerCount()).toBe(0);
    } finally {
      reads.mockRestore();
      controller.abort();
      await waiting.catch(() => {});
      vi.useRealTimers();
    }
  });

  it("projects an authorized collector failure without failing a mixed batch", async () => {
    const failed = collectorRun("failed", "agent:main:main", {
      status: "failed",
      structured: { partial: true },
    });
    failed.execution = {
      status: "terminal",
      outcome: { status: "error", error: "provider failed after tool output" },
    };
    failed.completion = { required: false, resultText: null, capturedAt: 10 };
    records.set(failed.runId, failed);
    records.set("pending", collectorRun("pending", "agent:main:main"));
    const tool = createMainSessionWaitTool();

    const result = await tool.execute("call", {
      ids: [failed.runId, "pending"],
      timeoutSeconds: 0,
    });

    expect(result.details).toEqual({
      completed: [
        {
          runId: failed.runId,
          status: "failed",
          result: "",
          structured: { partial: true },
          error: "provider failed after tool output",
          sessionKey: failed.childSessionKey,
        },
      ],
      pending: ["pending"],
    });
    expect(isToolResultError(result)).toBe(false);
  });

  it.each([-60_000, 60_000])(
    "keeps its elapsed deadline after a %d ms clock step",
    async (step) => {
      vi.useFakeTimers({ toFake: ["Date", "performance", "setTimeout", "clearTimeout"] });
      vi.setSystemTime(100_000);
      const controller = new AbortController();
      records.set("clock-step", collectorRun("clock-step", "agent:main:main"));
      const tool = createMainSessionWaitTool();
      let result: unknown;
      const waiting = tool
        .execute("clock", { ids: ["clock-step"], timeoutSeconds: 0.1 }, controller.signal)
        .then((value) => {
          result = value.details;
        });
      try {
        await vi.advanceTimersByTimeAsync(25);
        vi.setSystemTime(Date.now() + step);
        await vi.advanceTimersByTimeAsync(74);
        expect(result).toBeUndefined();
        await vi.advanceTimersByTimeAsync(1);
        expect(result).toEqual({ completed: [], pending: ["clock-step"] });
      } finally {
        controller.abort();
        await waiting.catch(() => {});
        vi.useRealTimers();
      }
    },
  );

  it("orders completions by durable capture time with input-order ties", async () => {
    const later = collectorRun("later", "agent:main:main", { status: "done" });
    later.completion = { required: false, resultText: "later", capturedAt: 10 };
    const earlier = collectorRun("earlier", "agent:main:main", { status: "done" });
    earlier.completion = { required: false, resultText: "earlier", capturedAt: 5 };
    records.set(later.runId, later);
    records.set(earlier.runId, earlier);
    const tool = createMainSessionWaitTool();

    const result = await tool.execute("call", {
      ids: ["later", "earlier"],
      timeoutSeconds: 0,
    });

    expect(result.details).toMatchObject({
      completed: [{ runId: "earlier" }, { runId: "later" }],
      pending: [],
    });

    const tied = collectorRun("tied", "agent:main:main", { status: "done" });
    tied.completion = { required: false, resultText: "tied", capturedAt: 5 };
    records.set(tied.runId, tied);
    records.set("foreign", collectorRun("foreign", "agent:other:main", { status: "done" }));
    records.set("pending-one", collectorRun("pending-one", "agent:main:main"));
    records.set("pending-two", collectorRun("pending-two", "agent:main:main"));

    const mixed = await tool.execute("mixed", {
      ids: ["later", "missing", "tied", "pending-two", "foreign", "earlier", "pending-one"],
      timeoutSeconds: 0,
    });

    expect(mixed.details).toMatchObject({
      completed: [{ runId: "tied" }, { runId: "earlier" }, { runId: "later" }],
      pending: ["pending-two", "pending-one"],
      errors: [
        { runId: "missing", error: "not_found" },
        { runId: "foreign", error: "not_owner" },
      ],
    });
    expect(isToolResultError(mixed)).toBe(false);
  });

  it("is idempotent and returns per-id ownership and unknown errors", async () => {
    const done = collectorRun("done", "agent:worker:subagent:owner", { status: "done" });
    done.swarmWaitOwnerSessionKeys = ["agent:worker:subagent:owner", "agent:main:main"];
    records.set("done", done);
    records.set("owner", collectorRun("owner", "agent:main:main"));
    records.set("foreign", collectorRun("foreign", "agent:other:main", { status: "failed" }));
    const tool = createAgentsWaitTool({
      agentSessionKey: "agent:main:main",
      agentId: "main",
      config: { tools: { swarm: { enabled: true, waitTimeoutSecondsMax: 1 } } },
    });

    const first = await tool.execute("call", {
      ids: ["done", "foreign", "missing"],
      timeoutSeconds: 5,
    });
    const second = await tool.execute("call", {
      ids: ["done", "foreign", "missing"],
      timeoutSeconds: 5,
    });
    expect(second.details).toEqual(first.details);
    expect(first.details).toMatchObject({
      completed: [{ runId: "done", status: "done" }],
      pending: [],
      errors: [
        { runId: "foreign", error: "not_owner" },
        { runId: "missing", error: "not_found" },
      ],
    });
    expect(isToolResultError(first)).toBe(false);
  });

  it("authorizes a snapshotted ancestor after the ordinary spawner row is archived", async () => {
    const ownerSessionKey = "agent:worker:subagent:ordinary-owner";
    records.set("ordinary-owner", {
      runId: "ordinary-owner",
      childSessionKey: ownerSessionKey,
      controllerSessionKey: "agent:main:main",
      requesterSessionKey: "agent:main:main",
      requesterDisplayKey: "main",
      task: "spawn collector",
      cleanup: "delete",
      createdAt: Date.now(),
      execution: { status: "running" },
    });
    const completed = collectorRun("nested", ownerSessionKey, { status: "done" });
    completed.swarmWaitOwnerSessionKeys = [ownerSessionKey, "agent:main:main"];
    records.set(completed.runId, completed);
    records.delete("ordinary-owner");
    const tool = createMainSessionWaitTool();

    const result = await tool.execute("call", { ids: ["nested"], timeoutSeconds: 0 });

    expect(result.details).toMatchObject({
      completed: [{ runId: "nested", status: "done" }],
      pending: [],
    });
  });

  it("keeps the public collector id after gateway run replacement", async () => {
    const remapped = collectorRun("gateway-run", "agent:main:main", { status: "done" });
    remapped.swarmRunId = "collector-run";
    records.set(remapped.runId, remapped);
    const tool = createMainSessionWaitTool();

    const result = await tool.execute("call", { ids: ["collector-run"], timeoutSeconds: 0 });

    expect(result.details).toMatchObject({
      completed: [{ runId: "collector-run", status: "done" }],
      pending: [],
    });
  });

  it("re-resolves a collector replaced while waiting", async () => {
    const pending = collectorRun("old-gateway-run", "agent:main:main");
    pending.swarmRunId = "collector-run";
    records.set(pending.runId, pending);
    const tool = createMainSessionWaitTool();
    setTimeout(() => {
      records.delete(pending.runId);
      const completed = collectorRun("new-gateway-run", "agent:main:main", { status: "done" });
      completed.swarmRunId = "collector-run";
      records.set(completed.runId, completed);
      for (const listener of registryEvents.listeners) {
        listener();
      }
    }, 5);

    const result = await tool.execute("call", { ids: ["collector-run"], timeoutSeconds: 1 });
    expect(result.details).toMatchObject({
      completed: [{ runId: "collector-run", status: "done" }],
      pending: [],
    });
  });

  it.each(["tool", "bridge"] as const)(
    "rejects a foreign replacement while the %s is parked",
    async (boundary) => {
      vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout", "performance"] });
      const pending = collectorRun("old-gateway-run", "agent:main:main");
      pending.swarmRunId = "collector-run";
      records.set(pending.runId, pending);
      const controller = new AbortController();
      const waiting =
        boundary === "tool"
          ? createMainSessionWaitTool()
              .execute("wait", { ids: ["collector-run"], timeoutSeconds: 1 }, controller.signal)
              .then((result) => result.details)
          : waitForCollectorCompletion({
              runId: "collector-run",
              currentSessionKeys: new Set(["agent:main:main"]),
              currentAgentId: "main",
              signal: controller.signal,
            });
      const observed = waiting.then(
        (value) => ({ value }),
        (error: unknown) => ({ error }),
      );
      try {
        expect(registryEvents.listeners.size).toBe(1);
        records.delete(pending.runId);
        const foreign = collectorRun("new-gateway-run", "agent:other:main", { status: "done" });
        foreign.swarmRunId = "collector-run";
        foreign.completion = { required: false, resultText: "foreign result" };
        records.set(foreign.runId, foreign);
        for (const listener of registryEvents.listeners) {
          listener();
        }

        expect(await observed).toEqual(
          boundary === "tool"
            ? {
                value: {
                  completed: [],
                  pending: [],
                  errors: [{ runId: "collector-run", error: "not_owner" }],
                  success: false,
                },
              }
            : {
                error: expect.objectContaining({ message: "agents.run not_owner: collector-run" }),
              },
        );
        expect(registryEvents.listeners.size).toBe(0);
        expect(vi.getTimerCount()).toBe(0);
      } finally {
        controller.abort();
        await observed;
        vi.useRealTimers();
      }
    },
  );

  it("does not treat a routed completion owner as the spawning session", async () => {
    const routed = collectorRun("routed", "agent:main:main", { status: "done" });
    routed.controllerSessionKey = "agent:worker:route-a";
    routed.swarmRequesterSessionKey = "agent:worker:route-a";
    records.set(routed.runId, routed);
    const routedOwner = createAgentsWaitTool({
      agentSessionKey: "agent:worker:route-a",
      agentId: "worker",
      config: { tools: { swarm: true } },
    });
    const completionOwner = createMainSessionWaitTool();

    const allowed = await routedOwner.execute("owner", { ids: ["routed"], timeoutSeconds: 0 });
    const denied = await completionOwner.execute("proxy", {
      ids: ["routed"],
      timeoutSeconds: 0,
    });

    expect(allowed.details).toMatchObject({ completed: [{ runId: "routed" }] });
    expect(denied.details).toEqual({
      completed: [],
      pending: [],
      errors: [{ runId: "routed", error: "not_owner" }],
      success: false,
    });
    expect(isToolResultError(denied)).toBe(true);
  });

  it("rejects a foreign collector with the same bare requester key", async () => {
    const foreign = collectorRun("foreign-global", "global", { status: "done" });
    foreign.requesterAgentId = "ops";
    records.set(foreign.runId, foreign);
    const tool = createAgentsWaitTool({
      agentSessionKey: "global",
      agentId: "research",
      config: {
        agents: { ownership: "explicit", entries: { research: {}, ops: {} } },
        tools: { swarm: true },
      },
    });

    const result = await tool.execute("wait", {
      ids: [foreign.runId],
      timeoutSeconds: 0,
    });

    expect(result.details).toMatchObject({
      errors: [{ runId: foreign.runId, error: "not_owner" }],
      success: false,
    });
  });

  it("marks entirely missing collector batches as failures without losing per-id errors", async () => {
    const tool = createMainSessionWaitTool();

    const result = await tool.execute("call", { ids: ["missing"], timeoutSeconds: 0 });

    expect(result.details).toEqual({
      completed: [],
      pending: [],
      errors: [{ runId: "missing", error: "not_found" }],
      success: false,
    });
    expect(isToolResultError(result)).toBe(true);
  });

  it("rejects collector batches containing only blank run ids", async () => {
    const tool = createMainSessionWaitTool();

    await expect(tool.execute("call", { ids: [" ", "\t"], timeoutSeconds: 0 })).rejects.toThrow(
      "at least one non-empty run id",
    );
  });

  it.each(["before", "during", "registration"] as const)(
    "rejects when the wait is aborted %s collector waiting",
    async (abortTiming) => {
      records.set("pending", collectorRun("pending", "agent:main:main"));
      const tool = createMainSessionWaitTool();
      const controller = new AbortController();
      if (abortTiming === "before") {
        controller.abort();
      }
      if (abortTiming === "registration") {
        const addEventListener = controller.signal.addEventListener.bind(controller.signal);
        vi.spyOn(controller.signal, "addEventListener").mockImplementation((...args) => {
          controller.abort();
          addEventListener(...args);
        });
      }

      const result = tool.execute(
        "call",
        { ids: ["pending"], timeoutSeconds: 1 },
        controller.signal,
      );
      if (abortTiming === "during") {
        controller.abort();
      }

      await expect(result).rejects.toMatchObject({
        name: "AbortError",
        message: "agents_wait aborted.",
      });
    },
  );

  it("rejects oversized wait batches before waiting", async () => {
    const tool = createMainSessionWaitTool();

    await expect(
      tool.execute("call", {
        ids: Array.from({ length: 1_001 }, (_, index) => `run-${index}`),
      }),
    ).rejects.toThrow("at most 1000 ids");
  });
});