From 0ac62697a8d16cb087cbf0944cef5025c0ce0d4c Mon Sep 17 00:00:00 2001
From: Joey Krug <joeykrug@gmail.com>
Date: Sun, 15 Mar 2026 00:42:45 -0400
Subject: [PATCH] Tests: cover yield spawn announce gating

---
 ...subagent-announce.yield-spawn-race.test.ts | 124 ++++++++++++++++++
 1 file changed, 124 insertions(+)
 create mode 100644 src/agents/subagent-announce.yield-spawn-race.test.ts

diff --git a/src/agents/subagent-announce.yield-spawn-race.test.ts b/src/agents/subagent-announce.yield-spawn-race.test.ts
new file mode 100644
index 00000000000..1917265e733
--- /dev/null
+++ b/src/agents/subagent-announce.yield-spawn-race.test.ts
@@ -0,0 +1,124 @@
+import { beforeEach, describe, expect, it, vi } from "vitest";
+
+type GatewayCall = {
+  method?: string;
+  timeoutMs?: number;
+  expectFinal?: boolean;
+  params?: Record<string, unknown>;
+};
+
+const gatewayCalls: GatewayCall[] = [];
+let callGatewayImpl: (request: GatewayCall) => Promise<unknown> = async () => ({});
+let sessionStore: Record<string, Record<string, unknown>> = {};
+let configOverride: ReturnType<(typeof import("../config/config.js"))["loadConfig"]> = {
+  session: {
+    mainKey: "main",
+    scope: "per-sender",
+  },
+};
+
+vi.mock("../gateway/call.js", () => ({
+  callGateway: vi.fn(async (request: GatewayCall) => {
+    gatewayCalls.push(request);
+    return await callGatewayImpl(request);
+  }),
+}));
+
+vi.mock("../config/config.js", async (importOriginal) => {
+  const actual = await importOriginal<typeof import("../config/config.js")>();
+  return {
+    ...actual,
+    loadConfig: () => configOverride,
+  };
+});
+
+vi.mock("../config/sessions.js", () => ({
+  loadSessionStore: vi.fn(() => sessionStore),
+  resolveAgentIdFromSessionKey: () => "main",
+  resolveStorePath: () => "/tmp/sessions-main.json",
+  resolveMainSessionKey: () => "agent:main:main",
+}));
+
+vi.mock("./pi-embedded.js", () => ({
+  isEmbeddedPiRunActive: () => false,
+  queueEmbeddedPiMessage: () => false,
+  waitForEmbeddedPiRunEnd: async () => true,
+}));
+
+import { runSubagentAnnounceFlow } from "./subagent-announce.js";
+import {
+  addSubagentRunForTests,
+  countPendingDescendantRuns,
+  resetSubagentRegistryForTests,
+} from "./subagent-registry.js";
+
+function findFinalDirectAgentCall(): GatewayCall | undefined {
+  return gatewayCalls.find((call) => call.method === "agent" && call.expectFinal === true);
+}
+
+describe("subagent announce yield + spawn race", () => {
+  beforeEach(() => {
+    gatewayCalls.length = 0;
+    callGatewayImpl = async () => ({});
+    sessionStore = {};
+    configOverride = {
+      session: {
+        mainKey: "main",
+        scope: "per-sender",
+      },
+    };
+    resetSubagentRegistryForTests({ persist: false });
+  });
+
+  it("defers announce when a yield-aborted parent still has a concurrently spawned pending child", async () => {
+    const parentSessionKey = "agent:main:subagent:orchestrator-race";
+
+    addSubagentRunForTests({
+      runId: "run-orchestrator-race",
+      childSessionKey: parentSessionKey,
+      controllerSessionKey: "agent:main:main",
+      requesterSessionKey: "agent:main:main",
+      requesterDisplayKey: "main",
+      task: "orchestrate",
+      cleanup: "keep",
+      createdAt: 100,
+      startedAt: 100,
+      endedAt: 110,
+      expectsCompletionMessage: true,
+    });
+    addSubagentRunForTests({
+      runId: "run-worker-race",
+      childSessionKey: `${parentSessionKey}:subagent:worker`,
+      controllerSessionKey: parentSessionKey,
+      requesterSessionKey: parentSessionKey,
+      requesterDisplayKey: parentSessionKey,
+      task: "child task",
+      cleanup: "keep",
+      createdAt: 111,
+      startedAt: 111,
+      expectsCompletionMessage: true,
+    });
+
+    // Regression guard: when sessions_spawn commits before the announce check,
+    // the parent must still see the pending child and defer its completion.
+    expect(countPendingDescendantRuns(parentSessionKey)).toBe(1);
+
+    const didAnnounce = await runSubagentAnnounceFlow({
+      childSessionKey: parentSessionKey,
+      childRunId: "run-orchestrator-race",
+      requesterSessionKey: "agent:main:main",
+      requesterDisplayKey: "main",
+      task: "orchestrate",
+      timeoutMs: 1_000,
+      cleanup: "keep",
+      roundOneReply: "Yielded after concurrent sessions_spawn.",
+      waitForCompletion: false,
+      expectsCompletionMessage: true,
+      outcome: { status: "ok" },
+    });
+
+    expect(didAnnounce).toBe(false);
+    expect(countPendingDescendantRuns(parentSessionKey)).toBe(1);
+    expect(findFinalDirectAgentCall()).toBeUndefined();
+  });
+});