import { assert, assertEquals, assertRejects, assertStringIncludes } from "@std/assert";
import { BACKGROUND_CONTEXT as context } from "@earendil-works/chord/context";
import {
  type JsonObject,
  type Message,
  type ModelThinkingLevel,
  Type,
} from "@earendil-works/pi-ai";
import { getCurrentSystemPrompt, getCurrentTools } from "@earendil-works/pi-ai/utils/transcript";
import { createModels } from "@earendil-works/pi-ai/models";
import {
  fauxAssistantMessage,
  fauxProvider,
  fauxToolCall,
} from "@earendil-works/pi-ai/providers/faux";
import { AgentDoc, LiveDoc, MemoryStorage } from "@earendil-works/pi-durable";
import { openNodeSqliteStorage } from "@earendil-works/pi-durable/storage/sqlite/node";
import { defineExtension } from "@berry/api";
import { localEnvironment } from "@berry/api/environment.ts";
import { Conversations, openBerry } from "./berry.ts";
import {
  blocking,
  dir,
  lastToolResult,
  script,
  sectionsSeen,
  setup,
  toolsOffered,
  turn,
} from "./fixture.ts";
import { tempDir, until } from "@berry/util";

async function within<T>(promise: Promise<T>): Promise<T> {
  let timer: ReturnType<typeof setTimeout> | undefined;
  try {
    return await Promise.race([
      promise,
      new Promise<never>((_, reject) =>
        timer = setTimeout(() => reject(new Error("Timed out")), 5000)
      ),
    ]);
  } finally {
    clearTimeout(timer);
  }
}

const last = (messages: readonly Message[]) => messages.findLast((m) => m.role !== "system")!;
const textOf = (message: Message) =>
  typeof message.content === "string"
    ? message.content
    : message.content.flatMap((part) => part.type === "text" ? [part.text] : []).join("\n");

Deno.test("selecting subagent controls direct tools, the codemode catalog, and script invocation, and old conversations keep it", async () => {
  await using harness = await setup();
  const { berry, model } = harness;
  const tid = await berry.create({ extensions: {}, model });
  const [initial] = await turn(harness, tid, []);
  assert(!getCurrentTools(initial).some((tool) => tool.name === "subagent"));
  assert(!sectionsSeen(initial).some((section) => section.codemode?.includes("subagent:")));

  await berry.configure(tid, { extensions: ["subagent"] });
  const [enabled] = await turn(harness, tid, []);
  assert(getCurrentTools(enabled).some((tool) => tool.name === "subagent"));
  assert(sectionsSeen(enabled).some((section) => section.codemode?.includes("subagent:")));

  await berry.configure(tid, { extensions: [] });
  const disabled = await turn(harness, tid, [
    script(`return await tools.subagent({task:"must not run", model:${JSON.stringify(model)}});`),
  ]);
  assert(!getCurrentTools(disabled[0]).some((tool) => tool.name === "subagent"));
  const catalog =
    sectionsSeen(disabled[0]).findLast((section) => section.codemode !== undefined)?.codemode ?? "";
  assert(!catalog.includes("subagent:"));
  const result = lastToolResult(disabled[1]);
  assert(result.isError);
  assertStringIncludes(result.text, "subagent");
  assertEquals(Object.keys((await berry.harness.snapshot(Conversations, context))!), [tid]);

  // as conversations from before subagent was a choice stored their extensions, which keep it
  const selected = async (tid: string) =>
    (await (await berry.conversation(tid))!.agent(context)).extensions.some((ext) =>
      ext.name === "subagent"
    );
  const legacy = await berry.harness.createConversation({
    ownership: { kind: "ownerless" },
    agent: { extensions: { add: [] } },
  }, context);
  const old = await berry.tidOf(legacy.id);
  assertEquals(await selected(old), true);
  await berry.configure(old, { extensions: [] });
  assertEquals(await selected(old), false);
});

Deno.test("delegation needs a worker model enabled for it, and continues only the conversation's own subagents", async () => {
  await using harness = await setup(undefined, {
    models: [
      { id: "parent" },
      { id: "chosen-worker" },
      { id: "other-worker" },
      { id: "disabled-worker" },
    ],
  });
  const { berry, faux, model } = harness;
  const worker = { provider: "faux", modelId: "chosen-worker" };
  const other = { provider: "faux", modelId: "other-worker" };
  const ran: string[] = [];
  const delegate = async (tid: string, args: JsonObject) => {
    faux.setResponses(Array.from({ length: 4 }, () => (request, _, __, chosen) => {
      const message = last(request.messages);
      if (message.role === "toolResult") return fauxAssistantMessage("done");
      if (textOf(message) === "go") {
        return fauxAssistantMessage(fauxToolCall("subagent", args), { stopReason: "toolUse" });
      }
      ran.push(chosen.id);
      return fauxAssistantMessage("worker answer");
    }));
    const parent = (await berry.conversation(tid))!;
    await within((await parent.submit({ type: "input", content: "go" }, context)).wait(context));
    await within(parent.waitForIdle(context));
    const messages = (await parent.context(context)).messages;
    return { result: lastToolResult(messages), prompt: getCurrentSystemPrompt(messages) };
  };
  const children = async (tid: string) => {
    const index = (await berry.harness.snapshot(Conversations, context))!;
    return Object.keys(index).filter((key) => index[key].owner === tid);
  };
  await berry.enableModels([model, worker, other]);
  const owner = await berry.create({
    extensions: { subagent: { models: ["faux/chosen-worker"] } },
    model,
  });
  const refusals: [JsonObject, string][] = [
    [{ task: "must not run" }, "model"],
    [
      { task: "must not run", model: { provider: "faux", modelId: "disabled-worker" } },
      "not enabled",
    ],
    [{ task: "must not run", model: other }, "not enabled for delegation"],
  ];
  for (const [args, error] of refusals) {
    const { result, prompt } = await delegate(owner, args);
    assert(result.isError);
    assertStringIncludes(result.text, error);
    assertStringIncludes(prompt, "chosen-worker");
    assert(!prompt.includes("other-worker") && !prompt.includes("disabled-worker"));
  }
  assertEquals(await children(owner), []);
  const graph = await berry.harness.taskGraph(context);
  assertEquals(Object.keys(graph.value.tasks), []);
  graph.dispose();
  assert(!(await delegate(owner, { task: "start", model: worker })).result.isError);
  const [child] = await children(owner);

  const stranger = await berry.create({ extensions: { subagent: {} }, model });
  const foreign = (await delegate(stranger, { task: "steal", conversation: child })).result;
  assert(foreign.isError);
  assertStringIncludes(foreign.text, "not one of this conversation's subagents");

  await berry.configureExtension(owner, "subagent", { models: ["faux/other-worker"] });
  const stale = (await delegate(owner, { task: "again", conversation: child })).result;
  assert(stale.isError);
  assertStringIncludes(stale.text, "faux/chosen-worker is not enabled for delegation");
  const moved = await delegate(owner, { task: "again", conversation: child, model: other });
  assert(!moved.result.isError);
  assertEquals((await (await berry.conversation(child))!.agent(context)).model, other);

  // choosing none leaves every model enabled in settings
  assertEquals(await berry.configureExtension(owner, "subagent", { models: [] }), {
    models: [],
    limit: 4,
  });
  const fallback = await delegate(owner, { task: "more", model: worker });
  assert(!fallback.result.isError);
  assertStringIncludes(fallback.prompt, "other-worker");
  assertEquals(ran, ["chosen-worker", "other-worker", "chosen-worker"]);
  await assertRejects(
    () => berry.configureExtension(owner, "subagent", { models: ["faux/disabled-worker"] }),
    Error,
    "faux/disabled-worker is not enabled in settings",
  );
});

Deno.test("delegation can override thinking without changing the parent", async (t) => {
  await using harness = await setup(undefined, {
    models: [
      { id: "parent", reasoning: true },
      { id: "worker", reasoning: true },
      { id: "plain-worker", reasoning: false },
    ],
  });
  const { berry, faux, model } = harness;
  const worker = { provider: "faux", modelId: "worker" };
  const plain = { provider: "faux", modelId: "plain-worker" };
  await berry.enableModels([model, worker, plain]);
  const cases: {
    name: string;
    requested?: ModelThinkingLevel;
    expected: ModelThinkingLevel;
    plain?: boolean;
  }[] = [
    { name: "omission inherits", expected: "high" },
    { name: "a request overrides", requested: "low", expected: "low" },
    {
      name: "a non-reasoning worker clamps to off",
      requested: "high",
      expected: "off",
      plain: true,
    },
  ];
  for (const scenario of cases) {
    await t.step(scenario.name, async () => {
      const selected = scenario.plain ? plain : worker;
      const tid = await berry.create({
        extensions: { subagent: {} },
        model,
        thinkingLevel: "high",
      });
      const parent = (await berry.conversation(tid))!;
      const args = {
        task: scenario.name,
        model: selected,
        ...(scenario.requested === undefined ? {} : { thinkingLevel: scenario.requested }),
      };
      const requests: { modelId: string; thinking: ModelThinkingLevel }[] = [];
      faux.setResponses([
        fauxAssistantMessage(fauxToolCall("subagent", args), { stopReason: "toolUse" }),
        (_, options, __, model) => {
          requests.push({ modelId: model.id, thinking: options?.reasoning ?? "off" });
          return fauxAssistantMessage("worker answer");
        },
        fauxAssistantMessage("done"),
      ]);
      await (await parent.submit({ type: "input", content: "go" }, context)).wait(context);
      await parent.waitForIdle(context);
      assertEquals(requests, [{ modelId: selected.modelId, thinking: scenario.expected }]);
      assert(!lastToolResult((await parent.context(context)).messages).isError);
      const index = (await berry.harness.snapshot(Conversations, context))!;
      const children = Object.keys(index).filter((child) => index[child].owner === tid);
      assertEquals(children.length, 1);
      const child = (await berry.conversation(children[0]))!;
      assertEquals((await child.agent(context)).thinkingLevel, scenario.expected);
      const unchanged = await parent.agent(context);
      assertEquals(unchanged.model, model);
      assertEquals(unchanged.thinkingLevel, "high");
    });
  }
});

Deno.test("a continued subagent keeps its conversation, context, model, and thinking, and queues follow-ups", async () => {
  await using harness = await setup(undefined, {
    models: [{ id: "parent" }, { id: "worker", reasoning: true }],
  });
  const { berry, faux, model } = harness;
  const worker = { provider: "faux", modelId: "worker" };
  const children: { model: string; thinking: unknown; users: string[] }[] = [];
  faux.setResponses(Array.from({ length: 8 }, () => (request, options, _, chosen) => {
    const users = request.messages.flatMap((m) => m.role === "user" ? [textOf(m)] : []);
    if (users[0] !== "go") {
      children.push({ model: chosen.id, thinking: options?.reasoning, users });
      return fauxAssistantMessage(`${users.at(-1)} after ${users.length - 1}`);
    }
    if (last(request.messages).role === "toolResult") return fauxAssistantMessage("done");
    return script(`
      const first = await tools.subagent({ task: "one", model: ${JSON.stringify(worker)},
        thinkingLevel: "low" });
      const more = await Promise.all(["two", "three"].map((task) =>
        tools.subagent({ task, conversation: first.conversation })));
      text(JSON.stringify([first, ...more]));
    `);
  }));
  await berry.enableModels([model, worker]);
  const tid = await berry.create({ extensions: { subagent: {} }, model });
  const parent = (await berry.conversation(tid))!;
  await within((await parent.submit({ type: "input", content: "go" }, context)).wait(context));
  await within(parent.waitForIdle(context));
  const result = lastToolResult((await parent.context(context)).messages);
  assert(!result.isError);
  const answers = JSON.parse(result.text.slice(result.text.indexOf("["))) as {
    conversation: string;
    answer: string;
  }[];
  const index = (await berry.harness.snapshot(Conversations, context))!;
  const [child] = Object.keys(index).filter((key) => index[key].owner === tid);
  assertEquals(answers.map((a) => a.conversation), [child, child, child]);
  const [task, after] = [0, 1].map((i) => answers.map((a) => a.answer.split(" after ")[i]));
  assertEquals(task, ["one", "two", "three"]);
  assertEquals([after[0], ...after.slice(1).sort()], ["0", "1", "2"]);
  assertEquals(children.map((c) => c.users.length), [1, 2, 3]);
  assertEquals(children[2].users[0], "one");
  for (const request of children) {
    assertEquals([request.model, request.thinking], ["worker", "low"]);
  }
});

Deno.test("ending a codemode script cancels its unawaited subagents, new or continued, while the parent continues", async () => {
  await using harness = await setup();
  const { berry, faux, model } = harness;
  const blocker = blocking();
  await berry.install(blocker.extension);
  const tid = await berry.create({ extensions: { subagent: {}, blocking: {} }, model });
  const parent = (await berry.conversation(tid))!;
  const unawaited = async (args: JsonObject) => {
    faux.setResponses([
      script(
        `tools.subagent(${JSON.stringify(args)}); await tools.ready({}); return "moving on";`,
      ),
      fauxAssistantMessage(fauxToolCall("block", {}), { stopReason: "toolUse" }),
      fauxAssistantMessage("done"),
    ]);
    const cancelled = blocker.cancelled.length;
    const submission = await parent.submit({ type: "input", content: "go" }, context);
    await until(() => blocker.cancelled.length > cancelled);
    assertEquals((await within(submission.wait(context))).status, "done");
    await within(berry.harness.waitForIdle(context));
    const result = lastToolResult((await parent.context(context)).messages);
    assert(!result.isError);
    assertStringIncludes(result.text, "moving on");
    const graph = await berry.harness.taskGraph(context);
    assertEquals(Object.keys(graph.value.tasks), []);
    graph.dispose();
  };

  await unawaited({ task: "block", model });
  const index = (await berry.harness.snapshot(Conversations, context))!;
  const [child] = Object.keys(index).filter((key) => index[key].owner === tid);
  await unawaited({ task: "block", conversation: child });
  const childId = (await berry.conversation(child))!.id;
  assertEquals((await berry.harness.snapshot(LiveDoc, childId, context))?.run, undefined);
});

Deno.test("delegation fans out on a selected cheaper model with inherited configuration and no recursive delegation", async () => {
  const faux = fauxProvider({
    models: [
      { id: "parent-thinker", reasoning: true },
      { id: "cheap-worker", reasoning: false },
    ],
  });
  const models = createModels();
  models.setProvider(faux.provider);
  const berry = await openBerry(new MemoryStorage(), models, undefined, {
    open: (name, cwd) =>
      Promise.resolve(Object.assign(localEnvironment(cwd ?? Deno.cwd()), { name })),
    close: () => Promise.resolve(),
  });
  const model = { provider: "faux", modelId: "parent-thinker" };
  const worker = { provider: "faux", modelId: "cheap-worker" };
  await berry.enableModels([model, worker]);
  await berry.install(dir());
  const entered = Promise.withResolvers<void>();
  const release = Promise.withResolvers<void>();
  const held: { name: string; tid: string; cwd: string; machine: string }[] = [];
  const requests: (readonly Message[])[] = [];
  await berry.install(defineExtension({
    name: "workers",
    initial: () => null,
    setup(ext) {
      ext.tool({
        name: "hold",
        description: "hold until released",
        parameters: Type.Object({ name: Type.String() }),
        expose: "both",
        async execute({ name }, ctx) {
          held.push({ name, tid: ctx.tid, cwd: ctx.env.cwd, machine: ctx.env.name });
          if (held.length === 2) entered.resolve();
          await release.promise;
          return `answer ${name}`;
        },
      });
    },
  }));
  const tid = await berry.create({
    extensions: {
      subagent: {},
      dir: { cwd: "/chosen" },
      workers: {},
      environment: { environment: "remote", cwd: Deno.cwd() },
    },
    model,
    thinkingLevel: "high",
  });
  const parent = (await berry.conversation(tid))!;
  await (await berry.emit(tid, "dir", { type: "cd", cwd: "/updated" })).wait(context);
  const router = (request: { messages: readonly Message[] }) => {
    requests.push(request.messages);
    const message = last(request.messages);
    const text = textOf(message);
    if (message.role === "user" && text === "go") {
      return fauxAssistantMessage([
        fauxToolCall("subagent", { task: "alpha", model: worker }, { id: "alpha" }),
        fauxToolCall("subagent", { task: "beta", model: worker }, { id: "beta" }),
      ], { stopReason: "toolUse" });
    }
    if (message.role === "user") {
      assert(["alpha", "beta"].includes(text));
      assert(!request.messages.some((m) => m.role === "user" && textOf(m) === "go"));
      assert(!toolsOffered(request.messages).includes("subagent"));
      assertStringIncludes(
        sectionsSeen(request.messages).find((s) => s.dir)?.dir ?? "",
        "/updated",
      );
      return fauxAssistantMessage(fauxToolCall("hold", { name: text }), {
        stopReason: "toolUse",
      });
    }
    if (message.role === "toolResult" && message.toolName === "hold") {
      return fauxAssistantMessage(text);
    }
    return fauxAssistantMessage("done");
  };
  faux.setResponses(Array.from({ length: 15 }, () => (request, _options, _state, selected) => {
    const isParent = request.messages.some((m) => m.role === "user" && textOf(m) === "go");
    assertEquals(
      { provider: selected.provider, modelId: selected.id },
      isParent ? model : worker,
    );
    return router(request);
  }));
  try {
    const submission = await parent.submit({ type: "input", content: "go" }, context);
    await within(entered.promise);
    assert((await berry.harness.snapshot(LiveDoc, parent.id, context))?.run);
    assertEquals(held.map((item) => item.name).sort(), ["alpha", "beta"]);
    assertEquals(new Set(held.map((item) => item.tid)).size, 2);
    const index = (await berry.harness.snapshot(Conversations, context))!;
    const parentAgent = await parent.agent(context);
    assertEquals(parentAgent.model, model);
    assertEquals(parentAgent.thinkingLevel, "high");
    const regular = await berry.create({ extensions: { dir: {}, workers: {} }, model: worker });
    const regularAgent = await berry.harness.snapshot(
      AgentDoc,
      (await berry.conversation(regular))!.id,
      context,
    );
    for (const child of held) {
      assertEquals(index[child.tid].owner, tid);
      assertEquals(child.cwd, Deno.cwd());
      assertEquals(child.machine, "remote");
      assertEquals(await berry.state(child.tid, "dir"), { cwd: "/updated" });
      const conversation = (await berry.conversation(child.tid))!;
      const stored = await berry.harness.snapshot(AgentDoc, conversation.id, context);
      assertEquals(stored?.extensions, regularAgent?.extensions);
      const agent = await conversation.agent(context);
      assertEquals(agent.model, worker);
      assertEquals(agent.thinkingLevel, "off");
      assertEquals(
        agent.extensions.map((ext) => ext.name),
        parentAgent.extensions.map((ext) => ext.name).filter((name) => name !== "subagent"),
      );
    }
    release.resolve();
    assertEquals((await within(submission.wait(context))).status, "done");
    await within(parent.waitForIdle(context));
    assertEquals((await parent.agent(context)).model, model);
    assertEquals((await parent.agent(context)).thinkingLevel, "high");
    const resumed = requests.at(-1)!;
    const results = resumed.filter((m) => m.role === "toolResult");
    const text = results.map(textOf).join("\n");
    assertStringIncludes(text, "answer alpha");
    assertStringIncludes(text, "answer beta");
    for (const result of results) {
      const details = result.role === "toolResult" ? result.details as { tid?: string } : undefined;
      assert(held.some((child) => child.tid === details?.tid));
    }
    const graph = await berry.harness.taskGraph(context);
    assertEquals(Object.keys(graph.value.tasks), []);
    graph.dispose();
  } finally {
    release.resolve();
    await berry.close();
  }
});

Deno.test("delegation is reserved and an owned conversation cannot re-enable it to delegate", async () => {
  await using harness = await setup();
  const { berry, faux, model } = harness;
  await assertRejects(
    () => berry.install(defineExtension({ name: "subagent", initial: () => null, setup() {} })),
    Error,
    "harness's own",
  );
  const tid = await berry.create({ extensions: { subagent: {} }, model });
  const parent = (await berry.conversation(tid))!;
  const anchor = Promise.withResolvers<void>();
  let childTid = "";
  await berry.install(defineExtension({
    name: "capture",
    initial: () => null,
    setup(ext) {
      ext.tool({
        name: "capture",
        description: "capture child",
        parameters: Type.Object({}),
        expose: "model",
        execute(_, ctx) {
          childTid = ctx.tid;
          anchor.resolve();
          return "captured";
        },
      });
    },
  }));
  await berry.configure(tid, { extensions: ["subagent", "capture"] });
  faux.setResponses([
    fauxAssistantMessage(fauxToolCall("subagent", { task: "child", model }), {
      stopReason: "toolUse",
    }),
    fauxAssistantMessage(fauxToolCall("capture", {}), { stopReason: "toolUse" }),
    fauxAssistantMessage("child answer"),
    fauxAssistantMessage("done"),
  ]);
  await within((await parent.submit({ type: "input", content: "go" }, context)).wait(context));
  await within(anchor.promise);
  await berry.configure(childTid, { extensions: ["subagent"] });
  faux.setResponses([
    fauxAssistantMessage(fauxToolCall("subagent", { task: "grandchild", model }), {
      stopReason: "toolUse",
    }),
    fauxAssistantMessage("done"),
  ]);
  const child = (await berry.conversation(childTid))!;
  await within(
    (await child.submit({ type: "input", content: "try again" }, context)).wait(context),
  );
  const entries = (await child.entries({ order: "ascending" }, 256, undefined, context)).items;
  const result = entries.flatMap((entry) => entry.model ?? []).findLast((m) =>
    m.role === "toolResult"
  );
  assert(result?.role === "toolResult" && result.isError);
  assertStringIncludes(textOf(result), "Only top-level agents");
  assertEquals(Object.keys((await berry.harness.snapshot(Conversations, context))!).length, 2);
});

/** Answers "go" with two subagent calls, alpha and beta, and has each child block under its task's name. */
const fanOut = (
  model: { provider: string; modelId: string },
  request: { messages: readonly Message[] },
) => {
  const message = last(request.messages);
  if (message.role === "user" && textOf(message) === "go") {
    return fauxAssistantMessage([
      fauxToolCall("subagent", { task: "alpha", model }, { id: "alpha" }),
      fauxToolCall("subagent", { task: "beta", model }, { id: "beta" }),
    ], { stopReason: "toolUse" });
  }
  if (message.role === "user") {
    return fauxAssistantMessage(fauxToolCall("block", { name: textOf(message) }), {
      stopReason: "toolUse",
    });
  }
  if (message.role === "toolResult" && message.toolName === "block") {
    return fauxAssistantMessage(textOf(message));
  }
  return fauxAssistantMessage("done");
};

Deno.test("aborting the parent stops every foreground child", async () => {
  await using harness = await setup();
  const { berry, faux, model } = harness;
  const blocker = blocking();
  await berry.install(blocker.extension);
  const tid = await berry.create({ extensions: { subagent: {}, blocking: {} }, model });
  const parent = (await berry.conversation(tid))!;
  faux.setResponses(Array.from({ length: 10 }, () => (request) => fanOut(model, request)));
  const submission = await parent.submit({ type: "input", content: "go" }, context);
  await until(() => blocker.running === 2);
  await within(parent.abort(context));
  const settled = await within(submission.wait(context));
  assert(settled.status === "unanswered" && settled.reason === "aborted");
  assertEquals(blocker.cancelled.sort(), ["alpha", "beta"]);
  const graph = await berry.harness.taskGraph(context);
  assertEquals(Object.keys(graph.value.tasks), []);
  graph.dispose();
});

Deno.test("a failed delegation leaves its sibling working", async () => {
  await using harness = await setup();
  const { berry, faux, model } = harness;
  const blocker = blocking();
  await berry.install(blocker.extension);
  let resumed: readonly Message[] = [];
  const tid = await berry.create({ extensions: { subagent: {}, blocking: {} }, model });
  const parent = (await berry.conversation(tid))!;
  faux.setResponses(Array.from({ length: 15 }, () => (request) => {
    const message = last(request.messages);
    if (message.role === "toolResult" && message.toolName === "subagent") {
      resumed = request.messages;
    }
    return fanOut(model, request);
  }));
  try {
    const submission = await parent.submit({ type: "input", content: "go" }, context);
    await until(() => blocker.running === 2);
    await within((await berry.conversation(blocker.started.get("alpha")!))!.abort(context));
    assertEquals(blocker.cancelled, ["alpha"]);
    assert((await berry.harness.snapshot(LiveDoc, parent.id, context))?.run);
    blocker.release.resolve("beta finished");
    assertEquals((await within(submission.wait(context))).status, "done");
    await within(parent.waitForIdle(context));
    const results = resumed.filter((m) => m.role === "toolResult").map(textOf).join("\n");
    assertStringIncludes(results, "Subagent failed: aborted");
    const failure = resumed.find((m) => m.role === "toolResult" && m.isError)!;
    assertStringIncludes(textOf(failure), blocker.started.get("alpha")!);
    assertStringIncludes(results, "beta finished");
    assertEquals(blocker.cancelled, ["alpha"]);
  } finally {
    blocker.release.resolve("cleanup");
  }
});

Deno.test("an aborted codemode subagent reports a conversation the parent can continue", async () => {
  await using harness = await setup();
  const { berry, faux, model } = harness;
  const blocker = blocking();
  await berry.install(blocker.extension);
  const tid = await berry.create({ extensions: { subagent: {}, blocking: {} }, model });
  const parent = (await berry.conversation(tid))!;
  faux.setResponses([
    script(`return await tools.subagent({ task: "block", model: ${JSON.stringify(model)} });`),
    fauxAssistantMessage(fauxToolCall("block", {}), { stopReason: "toolUse" }),
    fauxAssistantMessage("paused"),
  ]);
  const submission = await parent.submit({ type: "input", content: "go" }, context);
  await until(() => blocker.running === 1);
  const child = blocker.started.get("")!;
  await within((await berry.conversation(child))!.abort(context));
  assertEquals((await within(submission.wait(context))).status, "done");
  const result = lastToolResult((await parent.context(context)).messages);
  assert(result.isError);
  assertStringIncludes(result.text, "aborted");
  assertStringIncludes(result.text, child);

  const reported = result.text.match(/subagent conversation ([^;]+);/)?.[1];
  assertEquals(reported, child);
  const seen = await turn(harness, tid, [
    script(
      `return await tools.subagent({ task: "continue", conversation: ${
        JSON.stringify(reported)
      } });`,
    ),
    fauxAssistantMessage("recovered answer"),
  ]);
  const continued = lastToolResult(seen.at(-1)!);
  assert(!continued.isError);
  assertStringIncludes(continued.text, "recovered answer");
  assertStringIncludes(continued.text, child);
});

// a codemode script cannot resume after a restart, so recovery differs by how the call was made
for (const mode of ["direct", "codemode"] as const) {
  Deno.test(`${mode} delegation recovers without duplicating children or submissions`, async () => {
    await using temporary = await tempDir();
    const directory = temporary.path;
    const faux = fauxProvider({
      models: [{ id: "parent", reasoning: true }, { id: "worker", reasoning: true }],
    });
    const models = createModels();
    models.setProvider(faux.provider);
    const model = { provider: "faux", modelId: "parent" };
    const worker = { provider: "faux", modelId: "worker" };
    const open = () =>
      openNodeSqliteStorage(`${directory}/session.sqlite`).then((storage) =>
        openBerry(storage, models)
      );
    let berry = await open();
    const blocker = blocking();
    try {
      await berry.install(blocker.extension);
      await berry.enableModels([model, worker]);
      const tid = await berry.create({
        extensions: { subagent: {}, blocking: {} },
        model,
        thinkingLevel: "high",
      });
      const parent = (await berry.conversation(tid))!;
      const args = { task: "child", model: worker, thinkingLevel: "low" };
      faux.setResponses([
        mode === "codemode"
          ? script(`return await tools.subagent(${JSON.stringify(args)});`)
          : fauxAssistantMessage(fauxToolCall("subagent", args), { stopReason: "toolUse" }),
        fauxAssistantMessage(fauxToolCall("block", {}), { stopReason: "toolUse" }),
      ]);
      const submission = await parent.submit({ type: "input", content: "go" }, context);
      await until(() => blocker.running > 0);
      const before = (await berry.harness.snapshot(Conversations, context))!;
      const [childTid] = Object.keys(before).filter((key) => key !== tid);
      assertEquals((await (await berry.conversation(childTid))!.agent(context)).model, worker);
      assertEquals(
        (await (await berry.conversation(childTid))!.agent(context)).thinkingLevel,
        "low",
      );
      const resumed: (readonly Message[])[] = [];
      await within(berry.close());
      faux.setResponses(Array.from({ length: 10 }, () => (request) => {
        resumed.push(request.messages);
        const message = last(request.messages);
        assert(message.role === "toolResult");
        return fauxAssistantMessage(
          message.toolName === "block" ? "recovered child answer" : "done",
        );
      }));
      berry = await open();
      await berry.install(blocker.extension);
      berry.harness.resume();
      const receipt = (await berry.harness.submission(submission.id, context))!;
      assertEquals((await within(receipt.wait(context))).status, "done");
      await within(berry.harness.waitForIdle(context));
      assert(
        (await (await berry.conversation(tid))!.agent(context)).extensions.some((ext) =>
          ext.name === "subagent"
        ),
      );
      const after = (await berry.harness.snapshot(Conversations, context))!;
      assertEquals(Object.keys(after).sort(), Object.keys(before).sort());
      const child = (await berry.conversation(childTid))!;
      assertEquals((await child.agent(context)).model, worker);
      assertEquals((await child.agent(context)).thinkingLevel, "low");
      assertEquals((await (await berry.conversation(tid))!.agent(context)).model, model);
      assertEquals((await (await berry.conversation(tid))!.agent(context)).thinkingLevel, "high");
      const entries = (await child.entries({ order: "ascending" }, 256, undefined, context)).items;
      assertEquals(
        entries.flatMap((entry) => entry.model ?? []).filter((m) => m.role === "user").length,
        1,
      );
      const result = lastToolResult(resumed.at(-1)!);
      if (mode === "direct") {
        assertEquals(result.text.split("\n")[0], "recovered child answer");
        assertStringIncludes(result.text, childTid);
        assert(!result.isError);
        assertEquals(resumed.length, 2);
      } else {
        assert(result.isError);
        assertStringIncludes(result.text.toLowerCase(), "interrupted");
        assertEquals(resumed.length, 1);
        assertEquals((await berry.harness.snapshot(LiveDoc, child.id, context))?.run, undefined);
      }
    } finally {
      await berry.close();
    }
  });
}

Deno.test("subagent guidance is editable prose with enabled models always present", async () => {
  await using harness = await setup(undefined, {
    models: [{ id: "parent" }, { id: "worker" }],
  });
  const { berry, model } = harness;
  const worker = { provider: "faux", modelId: "worker" };
  await berry.enableModels([model, worker]);
  const def = berry.promptProseDefault("subagent");
  assert(def?.trim());

  const tid = await berry.create({ extensions: { subagent: {} }, model });
  const [first] = await turn(harness, tid, []);
  const seen = sectionsSeen(first)[0];
  assertStringIncludes(seen.subagent!, def!);
  assertStringIncludes(seen.subagent!, "Enabled models:");
  assertStringIncludes(seen.subagent!, "worker");

  await berry.setSystemPrompt({
    userName: "Sam",
    promptProse: { subagent: "Custom guidance for $USER." },
  });
  const [overrode] = await turn(harness, tid, []);
  const patch = sectionsSeen(overrode).at(-1)!;
  assertStringIncludes(patch.subagent!, "Custom guidance for Sam.");
  assertStringIncludes(patch.subagent!, "Enabled models:");
  assertStringIncludes(patch.subagent!, "worker");
  assert(!patch.subagent!.includes(def!));

  await berry.setSystemPrompt({ promptProse: { subagent: "" } });
  const [blanked] = await turn(harness, tid, []);
  const cleared = sectionsSeen(blanked).at(-1)!;
  assertStringIncludes(cleared.subagent!, "Enabled models:");
  assertStringIncludes(cleared.subagent!, "worker");
  assert(!cleared.subagent!.includes("Custom guidance"));
});

Deno.test("a background spawn returns at once and its report wakes an idle parent", async () => {
  await using harness = await setup();
  const { berry, faux, model } = harness;
  const blocker = blocking();
  await berry.install(blocker.extension);
  const tid = await berry.create({ extensions: { subagent: {}, blocking: {} }, model });
  const parent = (await berry.conversation(tid))!;
  let spawned = false;
  faux.setResponses(Array.from({ length: 20 }, () => (request) => {
    const message = last(request.messages);
    const text = textOf(message);
    if (message.role === "user" && text === "go" && !spawned) {
      spawned = true;
      return fauxAssistantMessage(
        fauxToolCall("codemode", {
          code: `return await tools.subagent({task:"child work", model:${
            JSON.stringify(model)
          }, background:true})`,
        }),
        { stopReason: "toolUse" },
      );
    }
    if (message.role === "user" && text === "child work") {
      return fauxAssistantMessage(fauxToolCall("block", {}), { stopReason: "toolUse" });
    }
    if (message.role === "toolResult" && message.toolName === "block") {
      return fauxAssistantMessage(text);
    }
    return fauxAssistantMessage("done");
  }));
  const submission = await parent.submit({ type: "input", content: "go" }, context);
  await until(() => blocker.running > 0);
  // the spawn returns without waiting for the child, so the parent idles while it stays blocked
  await within(parent.waitForIdle(context));
  assertEquals(blocker.running, 1);
  assertEquals((await within(submission.wait(context))).status, "done");
  const index = (await berry.harness.snapshot(Conversations, context))!;
  const [child] = Object.keys(index).filter((key) => index[key].owner === tid);
  assert(child);
  blocker.release.resolve("child answer");
  const reports = await within((async () => {
    for (;;) {
      const items = (await parent.entries({ order: "ascending" }, 256, undefined, context)).items;
      const found = items.filter((entry) => entry.kind === "berry.subagent-report");
      if (found.length > 0) return found;
      await new Promise((resolve) => setTimeout(resolve, 20));
    }
  })());
  assertEquals(reports.length, 1);
  const data = reports[0].data as { tid: string; task: string; status: string; answer?: string };
  assertEquals(data.tid, child);
  assertEquals(data.task, "child work");
  assertEquals(data.status, "completed");
  assertEquals(data.answer, "child answer");
  const messages = (await parent.context(context)).messages;
  assert(messages.some((message) => textOf(message).includes("child answer")));
});

Deno.test("the cap refuses a spawn and allows one after a stop", async () => {
  await using harness = await setup();
  const { berry, faux, model } = harness;
  const blocker = blocking();
  await berry.install(blocker.extension);
  const tid = await berry.create({ extensions: { subagent: {}, blocking: {} }, model });
  await berry.configureExtension(tid, "subagent", { models: [], limit: 1 });
  const parent = (await berry.conversation(tid))!;
  const scripts = [
    `return await tools.subagent({task:"one", model:${JSON.stringify(model)}, background:true})`,
    `return await tools.subagent({task:"two", model:${JSON.stringify(model)}, background:true})`,
    `return await tools.subagent_stop({})`,
    `return await tools.subagent({task:"two", model:${JSON.stringify(model)}, background:true})`,
  ];
  faux.setResponses(Array.from({ length: 30 }, () => (request) => {
    const message = last(request.messages);
    const text = textOf(message);
    if (message.role === "user" && text === "go") {
      const code = scripts.shift() ?? "return 0;";
      return fauxAssistantMessage(fauxToolCall("codemode", { code }), { stopReason: "toolUse" });
    }
    if (message.role === "user" && (text === "one" || text === "two")) {
      return fauxAssistantMessage(fauxToolCall("block", {}), { stopReason: "toolUse" });
    }
    return fauxAssistantMessage("done");
  }));
  const run = async () => {
    await within((await parent.submit({ type: "input", content: "go" }, context)).wait(context));
    await within(parent.waitForIdle(context));
    return lastToolResult((await parent.context(context)).messages);
  };
  const first = await run();
  assert(!first.isError);
  assertStringIncludes(first.text, "running");
  await until(() => blocker.running > 0);
  const refused = await run();
  assert(refused.isError);
  assertStringIncludes(refused.text, "Subagent limit reached");
  assertStringIncludes(refused.text, "1 running");
  assertStringIncludes(refused.text, "subagent_wait");
  assertStringIncludes(refused.text, "subagent_stop");
  const stopped = await run();
  assert(!stopped.isError);
  assertStringIncludes(stopped.text, "stopped");
  const retry = await run();
  assert(!retry.isError);
  assertStringIncludes(retry.text, "running");
});

Deno.test("subagent_wait collects a settled child and times out without killing work", async () => {
  await using harness = await setup();
  const { berry, faux, model } = harness;
  const blocker = blocking();
  await berry.install(blocker.extension);
  const tid = await berry.create({ extensions: { subagent: {}, blocking: {} }, model });
  const parent = (await berry.conversation(tid))!;
  const scripts: string[] = [
    `const spawned = await tools.subagent({task:"waited", model:${
      JSON.stringify(model)
    }, background:true}); store("child", spawned.conversation); return spawned;`,
    `return await tools.subagent_wait({timeoutMs:50})`,
    `const child = load("child"); return await tools.subagent_wait({conversations:[child]})`,
  ];
  faux.setResponses(Array.from({ length: 40 }, () => (request) => {
    const message = last(request.messages);
    const text = textOf(message);
    if (message.role === "user" && text === "go") {
      const code = scripts.shift() ?? "return 0;";
      return fauxAssistantMessage(fauxToolCall("codemode", { code }), { stopReason: "toolUse" });
    }
    if (message.role === "user" && text === "waited") {
      return fauxAssistantMessage(fauxToolCall("block", {}), { stopReason: "toolUse" });
    }
    if (message.role === "toolResult" && message.toolName === "block") {
      return fauxAssistantMessage(text);
    }
    return fauxAssistantMessage("done");
  }));
  const run = async () => {
    await within((await parent.submit({ type: "input", content: "go" }, context)).wait(context));
    await within(parent.waitForIdle(context));
    return lastToolResult((await parent.context(context)).messages);
  };
  const spawned = await run();
  assert(!spawned.isError);
  await until(() => blocker.running > 0);
  const index = (await berry.harness.snapshot(Conversations, context))!;
  const [child] = Object.keys(index).filter((key) => index[key].owner === tid);
  assert(child);
  // the wait times out while the child stays blocked, so the work survives it
  const timedOut = await run();
  assert(!timedOut.isError);
  assertStringIncludes(timedOut.text, "running");
  assertStringIncludes(timedOut.text, child);
  assertEquals(blocker.running, 1);
  blocker.release.resolve("waited answer");
  const collected = await run();
  assert(!collected.isError);
  assertStringIncludes(collected.text, "done");
  assertStringIncludes(collected.text, "waited answer");
  assertStringIncludes(collected.text, child);
});

Deno.test("subagent_stop stops the child and leaves a note without starting a run", async () => {
  await using harness = await setup();
  const { berry, faux, model } = harness;
  const blocker = blocking();
  await berry.install(blocker.extension);
  const tid = await berry.create({ extensions: { subagent: {}, blocking: {} }, model });
  const parent = (await berry.conversation(tid))!;
  const scripts = [
    `return await tools.subagent({task:"stop me", model:${
      JSON.stringify(model)
    }, background:true})`,
    `return await tools.subagent_stop({})`,
  ];
  faux.setResponses(Array.from({ length: 30 }, () => (request) => {
    const message = last(request.messages);
    const text = textOf(message);
    if (message.role === "user" && text === "go") {
      const code = scripts.shift() ?? "return 0;";
      return fauxAssistantMessage(fauxToolCall("codemode", { code }), { stopReason: "toolUse" });
    }
    if (message.role === "user" && text === "stop me") {
      return fauxAssistantMessage(fauxToolCall("block", { name: "stop me" }), {
        stopReason: "toolUse",
      });
    }
    return fauxAssistantMessage("done");
  }));
  const run = async () => {
    await within((await parent.submit({ type: "input", content: "go" }, context)).wait(context));
    await within(parent.waitForIdle(context));
    return lastToolResult((await parent.context(context)).messages);
  };
  const spawned = await run();
  assert(!spawned.isError);
  await until(() => blocker.running > 0);
  const index = (await berry.harness.snapshot(Conversations, context))!;
  const [child] = Object.keys(index).filter((key) => index[key].owner === tid);
  assert(child);
  const stopped = await run();
  assert(!stopped.isError);
  assertStringIncludes(stopped.text, "stopped");
  assertStringIncludes(stopped.text, child);
  assertEquals(blocker.running, 0);
  assertEquals(blocker.cancelled, ["stop me"]);
  const entries = (await parent.entries({ order: "ascending" }, 256, undefined, context)).items;
  const notes = entries.filter((entry) => entry.kind === "berry.subagent-report");
  assertEquals(notes.length, 1);
  const data = notes[0].data as { tid: string; task: string; status: string };
  assertEquals(data.tid, child);
  assertEquals(data.task, "stop me");
  assertEquals(data.status, "stopped");
  // the note does not wake the parent: no new run starts on its own
  const count = entries.length;
  await new Promise((resolve) => setTimeout(resolve, 100));
  await within(parent.waitForIdle(context));
  assertEquals((await berry.harness.snapshot(LiveDoc, parent.id, context))?.run, undefined);
  const later = (await parent.entries({ order: "ascending" }, 256, undefined, context)).items;
  assertEquals(later.length, count);
});

Deno.test("a background delegation survives a restart and reports exactly once", async () => {
  await using temporary = await tempDir();
  const directory = temporary.path;
  const faux = fauxProvider({ models: [{ id: "parent" }, { id: "worker" }] });
  const models = createModels();
  models.setProvider(faux.provider);
  const model = { provider: "faux", modelId: "parent" };
  const worker = { provider: "faux", modelId: "worker" };
  const open = () =>
    openNodeSqliteStorage(`${directory}/session.sqlite`).then((storage) =>
      openBerry(storage, models)
    );
  let berry = await open();
  const blocker = blocking();
  try {
    await berry.install(blocker.extension);
    await berry.enableModels([model, worker]);
    const tid = await berry.create({ extensions: { subagent: {}, blocking: {} }, model });
    const parent = (await berry.conversation(tid))!;
    const respond = () => (request: { messages: readonly Message[] }) => {
      const message = last(request.messages);
      const text = textOf(message);
      if (message.role === "user" && text === "go") {
        return fauxAssistantMessage(
          fauxToolCall("codemode", {
            code: `return await tools.subagent({task:"restart work", model:${
              JSON.stringify(worker)
            }, background:true})`,
          }),
          { stopReason: "toolUse" },
        );
      }
      if (message.role === "user" && text === "restart work") {
        return fauxAssistantMessage(fauxToolCall("block", {}), { stopReason: "toolUse" });
      }
      if (message.role === "toolResult" && message.toolName === "block") {
        return fauxAssistantMessage(text);
      }
      return fauxAssistantMessage("done");
    };
    faux.setResponses(Array.from({ length: 20 }, respond));
    await parent.submit({ type: "input", content: "go" }, context);
    await until(() => blocker.running > 0);
    const before = (await berry.harness.snapshot(Conversations, context))!;
    const [childTid] = Object.keys(before).filter((key) => key !== tid);
    assert(childTid);
    await within(berry.close());
    // after a restart the interrupted block settles and the child answers anew
    faux.setResponses(
      Array.from({ length: 20 }, () => (request: { messages: readonly Message[] }) => {
        const message = last(request.messages);
        if (message.role === "toolResult" && message.toolName === "block") {
          return fauxAssistantMessage("restarted answer");
        }
        return fauxAssistantMessage("done");
      }),
    );
    berry = await open();
    await berry.install(blocker.extension);
    berry.harness.resume();
    const parentAfter = (await berry.conversation(tid))!;
    await within(berry.harness.waitForIdle(context));
    const after = (await berry.harness.snapshot(Conversations, context))!;
    assertEquals(Object.keys(after).sort(), Object.keys(before).sort());
    const entries =
      (await parentAfter.entries({ order: "ascending" }, 256, undefined, context)).items;
    const reports = entries.filter((entry) => entry.kind === "berry.subagent-report");
    assertEquals(reports.length, 1);
    const data = reports[0].data as { tid: string; task: string; status: string; answer?: string };
    assertEquals(data.tid, childTid);
    assertEquals(data.task, "restart work");
    assertEquals(data.status, "completed");
    assertEquals(data.answer, "restarted answer");
  } finally {
    await berry.close();
  }
});

Deno.test("continuing a subagent in the background reports the new answer and stops", async () => {
  await using harness = await setup();
  const { berry, faux, model } = harness;
  const blocker = blocking();
  await berry.install(blocker.extension);
  const tid = await berry.create({ extensions: { subagent: {}, blocking: {} }, model });
  const parent = (await berry.conversation(tid))!;
  const scripts = [
    `const first = await tools.subagent({task:"one", model:${
      JSON.stringify(model)
    }, background:true}); store("child", first.conversation); return first;`,
    `const child = load("child"); return await tools.subagent({task:"two", conversation: child, background:true});`,
    `const child = load("child"); return await tools.subagent({task:"three", conversation: child, background:true});`,
    `const child = load("child"); return await tools.subagent_stop({conversations:[child]});`,
  ];
  faux.setResponses(Array.from({ length: 40 }, () => (request) => {
    const message = last(request.messages);
    const text = textOf(message);
    if (message.role === "user" && text === "go") {
      const code = scripts.shift() ?? "return 0;";
      return fauxAssistantMessage(fauxToolCall("codemode", { code }), { stopReason: "toolUse" });
    }
    if (message.role === "user" && (text === "one" || text === "two")) {
      return fauxAssistantMessage(text === "one" ? "ans1" : "ans2");
    }
    if (message.role === "user" && text === "three") {
      return fauxAssistantMessage(fauxToolCall("block", { name: "three" }), {
        stopReason: "toolUse",
      });
    }
    if (message.role === "toolResult" && message.toolName === "block") {
      return fauxAssistantMessage(text);
    }
    return fauxAssistantMessage("done");
  }));
  const run = async () => {
    await within((await parent.submit({ type: "input", content: "go" }, context)).wait(context));
    await within(parent.waitForIdle(context));
    return lastToolResult((await parent.context(context)).messages);
  };
  const reportWith = (answer: string) =>
    within((async () => {
      for (;;) {
        const items = (await parent.entries({ order: "ascending" }, 256, undefined, context)).items;
        const match = items.find((entry) =>
          entry.kind === "berry.subagent-report" &&
          (entry.data as { answer?: string }).answer === answer
        );
        if (match) return match;
        await new Promise((resolve) => setTimeout(resolve, 20));
      }
    })());
  const spawned = await run();
  assert(!spawned.isError);
  const first = await reportWith("ans1");
  assertEquals((first.data as { task: string }).task, "one");
  const childTid = (first.data as { tid: string }).tid;
  assertEquals((await berry.harness.snapshot(Conversations, context))![childTid].title, "one");
  const continued = await run();
  assert(!continued.isError);
  assertStringIncludes(continued.text, "running");
  const second = await reportWith("ans2");
  assertEquals((second.data as { task: string }).task, "two");
  assertEquals((await berry.harness.snapshot(Conversations, context))![childTid].title, "two");
  assertEquals((second.data as { tid: string }).tid, (first.data as { tid: string }).tid);
  const live = await run();
  assert(!live.isError);
  await until(() => blocker.running > 0);
  const index = (await berry.harness.snapshot(Conversations, context))!;
  const [child] = Object.keys(index).filter((key) => index[key].owner === tid);
  assert(child);
  const stopped = await run();
  assert(!stopped.isError);
  assertStringIncludes(stopped.text, "stopped");
  assertEquals(blocker.running, 0);
  assertEquals(blocker.cancelled, ["three"]);
  assertEquals((await berry.harness.snapshot(Conversations, context))![childTid].title, "three");
  const entries = (await parent.entries({ order: "ascending" }, 256, undefined, context)).items;
  const note = entries.find((entry) =>
    entry.kind === "berry.subagent-report" &&
    (entry.data as { status?: string }).status === "stopped"
  );
  assert(note);
  assertEquals((note!.data as { tid: string }).tid, child);
  assertEquals((note!.data as { task: string }).task, "three");
});