From 294fe805645820cb4cda1f4d137c2e94d40b2373 Mon Sep 17 00:00:00 2001 From: wassname <1103714+wassname@users.noreply.github.com> Date: Sun, 6 Sep 2026 22:36:37 +0800 Subject: [PATCH] Use pi-supervise acknowledgement for visible workers Co-Authored-By: PI[gpt-5.6-sol] <288921227+claudypoo@users.noreply.github.com> --- src/herdr.ts | 3 +- src/index.ts | 11 ++-- src/intercom.ts | 100 -------------------------------- src/supervise.ts | 38 +++++++----- src/supervisor-session.ts | 18 +++--- test/goals-flow.test.ts | 36 +++++++++--- test/herdr.test.ts | 5 +- test/supervise.test.ts | 34 +++++++++++ test/supervisor-session.test.ts | 25 +++++--- 9 files changed, 120 insertions(+), 150 deletions(-) delete mode 100644 src/intercom.ts create mode 100644 test/supervise.test.ts diff --git a/src/herdr.ts b/src/herdr.ts index 056880c..e014ead 100644 --- a/src/herdr.ts +++ b/src/herdr.ts @@ -56,13 +56,14 @@ export function supervisorCommand(input: LaunchSupervisorInput): string { const args = [ "pi", "--no-extensions", - "-e", input.extensionPath, "-e", "npm:pi-intercom", "-e", process.env.PI_GOALS_SUPERVISE_EXTENSION ?? "npm:@wassname2/pi-supervise@0.0.4", + "-e", input.extensionPath, "--fork", input.sourceSessionFile, "--name", `goals-supervisor-${input.workerSessionId.slice(0, 8)}`, ]; if (input.model) args.push("--model", input.model); + args.push("Initialize supervision startup."); return `env ${[...env, ...args].map(shellQuote).join(" ")}`; } diff --git a/src/index.ts b/src/index.ts index abdfbc6..1cce31c 100644 --- a/src/index.ts +++ b/src/index.ts @@ -22,8 +22,8 @@ import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-a import { Type } from "typebox"; import { approvalMatches, approvalPath, goalBlock, hashGoalBlock, readApproval, repositoryState } from "./approval.js"; import { closeSupervisorPane, openSupervisorPane } from "./herdr.js"; -import { registerGoalsIntercom } from "./intercom.js"; import { completeGoalDescription, completeGoalParamDescription, planDrafting, planningState, resync } from "./prompts.js"; +import { workerPiSupervise } from "./supervise.js"; import { isVisibleSupervisor, registerVisibleSupervisor } from "./supervisor-session.js"; const STATE = "pi-goals-state"; @@ -100,11 +100,10 @@ interface PlanState { export default function piGoalsExtension(pi: ExtensionAPI): void { if (isVisibleSupervisor()) { - registerVisibleSupervisor(pi, registerGoalsIntercom(pi)); + registerVisibleSupervisor(pi); return; } if (!isMainSession()) return; - const intercom = registerGoalsIntercom(pi); let state: PlanState = { phase: null, supervisorModel: null, @@ -148,7 +147,7 @@ export default function piGoalsExtension(pi: ExtensionAPI): void { repositoryRoot(ctx.cwd); const sourceSessionFile = ctx.sessionManager.getSessionFile(); if (!sourceSessionFile) throw new Error("The current session is not persisted, so it cannot be forked."); - const workerIntercomId = await intercom.workerIntercomId(); + const worker = await workerPiSupervise(pi); beginReview(ctx); let paneId: string | null = null; try { @@ -156,13 +155,13 @@ export default function piGoalsExtension(pi: ExtensionAPI): void { cwd: ctx.cwd, sourceSessionFile, workerSessionId: ctx.sessionManager.getSessionId(), - workerIntercomId, + workerIntercomId: worker.intercomId, planPath: planPath(ctx), approvalId: state.approvalId!, extensionPath: fileURLToPath(import.meta.url), model: state.supervisorModel, }); - await intercom.waitForSupervisorReady(state.approvalId!); + await worker.paired; } catch (error) { if (paneId) await closeSupervisorPane(paneId).catch(() => {}); throw error; diff --git a/src/intercom.ts b/src/intercom.ts deleted file mode 100644 index 4c314c5..0000000 --- a/src/intercom.ts +++ /dev/null @@ -1,100 +0,0 @@ -import type { ExtensionAPI } from "@earendil-works/pi-coding-agent"; - -const INTERCOM_REGISTER_EVENT = "intercom:extension-register"; -const NAMESPACE = "pi-goals/visible-supervisor/v1"; -const READY_TIMEOUT_MS = 15_000; - -type Channel = { - snapshot(): { connected: boolean }; - publish(payload: unknown, options?: { audience?: "owner" | "capable"; ownerOnly?: boolean }): void; - listSessions(): Promise>; -}; - -type IntercomEvent = { type: string; connected?: boolean; fromSessionId?: string; payload?: unknown }; - -type ReadyMessage = { type: "supervisor-ready"; to: string; approvalId: string }; - -function timeout(promise: Promise, message: string): Promise { - return new Promise((resolve, reject) => { - const timer = setTimeout(() => reject(new Error(message)), READY_TIMEOUT_MS); - promise.then( - (value) => { - clearTimeout(timer); - resolve(value); - }, - (error) => { - clearTimeout(timer); - reject(error); - }, - ); - }); -} - -function isReadyMessage(value: unknown): value is ReadyMessage { - if (!value || typeof value !== "object") return false; - const record = value as Record; - return record.type === "supervisor-ready" && typeof record.to === "string" && typeof record.approvalId === "string"; -} - -export interface GoalsIntercom { - workerIntercomId(): Promise; - waitForSupervisorReady(approvalId: string): Promise; - announceSupervisorReady(workerIntercomId: string, approvalId: string): Promise; -} - -export function registerGoalsIntercom(pi: ExtensionAPI): GoalsIntercom { - let channel: Channel | undefined; - let resolveIntercomConnected!: () => void; - const intercomConnected = new Promise((resolve) => { - resolveIntercomConnected = resolve; - }); - const ready = new Map void>(); - const announced = new Set(); - let ownId = ""; - - const currentId = async (): Promise => { - await intercomConnected; - if (ownId) return ownId; - const sessions = await channel!.listSessions(); - const session = sessions.find((item) => item.pid === process.pid); - if (!session) throw new Error("pi-goals could not find this Pi session in pi-intercom."); - ownId = session.id; - return ownId; - }; - - (pi as unknown as { events: { emit(name: string, value: unknown): void } }).events.emit(INTERCOM_REGISTER_EVENT, { - namespace: NAMESPACE, - ownerEligible: false, - onReady(value: Channel) { - channel = value; - if (value.snapshot().connected) resolveIntercomConnected(); - }, - onEvent(event: IntercomEvent) { - if (event.type === "connection" && event.connected) { - resolveIntercomConnected(); - return; - } - if (event.type !== "message" || !isReadyMessage(event.payload)) return; - if (event.payload.to !== ownId) return; - const resolve = ready.get(event.payload.approvalId); - if (!resolve) { - announced.add(event.payload.approvalId); - return; - } - ready.delete(event.payload.approvalId); - resolve(); - }, - }); - - return { - workerIntercomId: () => timeout(currentId(), "pi-goals needs pi-intercom before it can start a visible supervisor."), - waitForSupervisorReady(approvalId) { - if (announced.delete(approvalId)) return Promise.resolve(); - return timeout(new Promise((resolve) => ready.set(approvalId, resolve)), "The visible supervisor did not acknowledge pairing with this worker."); - }, - async announceSupervisorReady(workerIntercomId, approvalId) { - await timeout(intercomConnected, "pi-goals needs pi-intercom before it can confirm visible-supervisor pairing."); - channel!.publish({ type: "supervisor-ready", to: workerIntercomId, approvalId }, { audience: "capable" }); - }, - }; -} diff --git a/src/supervise.ts b/src/supervise.ts index bd73ebc..07a6ffb 100644 --- a/src/supervise.ts +++ b/src/supervise.ts @@ -1,21 +1,31 @@ import type { ExtensionAPI } from "@earendil-works/pi-coding-agent"; const PAIR_EVENT = "pi-supervise:pair:v1"; -const PAIR_TIMEOUT_MS = 15_000; +const WORKER_STATE_EVENT = "pi-supervise:worker-state:v1"; +const WORKER_PAIRED_EVENT = "pi-supervise:worker-paired:v1"; +const API_READY_EVENT = "pi-supervise:api-ready:v1"; +const TIMEOUT_MS = 15_000; -export function pairWithPiSupervise(pi: ExtensionAPI, workerIntercomId: string, goal: string): Promise { +type Events = { emit(name: string, value: unknown): boolean; on(name: string, handler: (value: any) => void): void }; + +function wait(start: (resolve: (value: T) => void, reject: (error: Error) => void) => void, message: string): Promise { return new Promise((resolve, reject) => { - const timer = setTimeout(() => reject(new Error("pi-supervise did not accept the visible-supervisor pairing request.")), PAIR_TIMEOUT_MS); - const settle = (callback: () => void) => { - clearTimeout(timer); - callback(); - }; - (pi as unknown as { events: { emit(name: string, value: unknown): void } }).events.emit(PAIR_EVENT, { - version: 1, - workerIntercomId, - goal, - resolve: () => settle(resolve), - reject: (error: Error) => settle(() => reject(error)), - }); + const timer = setTimeout(() => reject(new Error(message)), TIMEOUT_MS); + start((value) => { clearTimeout(timer); resolve(value); }, (error) => { clearTimeout(timer); reject(error); }); }); } + +export function pairWithPiSupervise(pi: ExtensionAPI, workerIntercomId: string, goal: string): Promise { + const events = (pi as unknown as { events: Events }).events; + return wait((resolve, reject) => events.emit(PAIR_EVENT, { version: 1, workerIntercomId, goal, resolve, reject }), "pi-supervise did not accept the visible-supervisor pairing request."); +} + +export function workerPiSupervise(pi: ExtensionAPI): Promise<{ intercomId: string; paired: Promise }> { + const events = (pi as unknown as { events: Events }).events; + return wait((resolve, _reject) => { + const paired = new Promise((pairedResolve) => events.on(WORKER_PAIRED_EVENT, () => pairedResolve())); + const request = () => events.emit(WORKER_STATE_EVENT, (state: { intercomId: string }) => resolve({ intercomId: state.intercomId, paired })); + events.on(API_READY_EVENT, request); + request(); + }, "pi-supervise did not publish this worker's intercom state."); +} diff --git a/src/supervisor-session.ts b/src/supervisor-session.ts index 738b817..d9944ca 100644 --- a/src/supervisor-session.ts +++ b/src/supervisor-session.ts @@ -3,7 +3,6 @@ import { resolve } from "node:path"; import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent"; import { Type } from "typebox"; import { approvalPath, goalBlock, hashGoalBlock, repositoryState, writeApproval } from "./approval.js"; -import type { GoalsIntercom } from "./intercom.js"; import { pairWithPiSupervise } from "./supervise.js"; const BOOTSTRAPPED = "pi-goals-visible-supervisor-v1"; @@ -76,17 +75,21 @@ export function isVisibleSupervisor(): boolean { return process.env.PI_GOALS_ROLE === "supervisor"; } -export function registerVisibleSupervisor(pi: ExtensionAPI, intercom: GoalsIntercom): void { +export function registerVisibleSupervisor(pi: ExtensionAPI): void { const settings = config(); let compacting = false; + let bootstrapping = false; - pi.on("before_agent_start", async (_event, ctx) => ({ - systemPrompt: `${ctx.getSystemPrompt()}\n\n${supervisorPrompt(settings)}`, - })); + pi.on("before_agent_start", async (_event, ctx) => { + await bootstrap(ctx); + return { systemPrompt: `${ctx.getSystemPrompt()}\n\n${supervisorPrompt(settings)}` }; + }); - pi.on("session_start", async (_event, ctx) => { + const bootstrap = async (ctx: ExtensionContext): Promise => { + if (bootstrapping) return; const entries = ctx.sessionManager.getEntries(); if (entries.some((entry: { type?: string; customType?: string }) => entry.type === "custom" && entry.customType === BOOTSTRAPPED)) return; + bootstrapping = true; compacting = true; try { await new Promise((resolve, reject) => { @@ -97,14 +100,13 @@ export function registerVisibleSupervisor(pi: ExtensionAPI, intercom: GoalsInter }); }); await pairWithPiSupervise(pi, settings.workerIntercomId, settings.planPath); - await intercom.announceSupervisorReady(settings.workerIntercomId, settings.approvalId); pi.appendEntry(BOOTSTRAPPED, { version: 1, workerSessionId: settings.workerSessionId, planPath: settings.planPath }); } catch (error) { ctx.ui.notify(`Supervisor startup failed: ${error instanceof Error ? error.message : String(error)}`, "error"); } finally { compacting = false; } - }); + }; pi.on("agent_settled", async (_event, ctx) => { if (compacting || (ctx.getContextUsage()?.tokens ?? 0) < COMPACT_AT_TOKENS) return; diff --git a/test/goals-flow.test.ts b/test/goals-flow.test.ts index 8ca1876..21df95b 100644 --- a/test/goals-flow.test.ts +++ b/test/goals-flow.test.ts @@ -1,4 +1,5 @@ import { execFileSync } from "node:child_process"; +import { EventEmitter } from "node:events"; import { mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; @@ -9,14 +10,6 @@ import { approvalPath, goalBlock, hashGoalBlock, repositoryState, writeApproval const openSupervisorPane = vi.fn(async () => "pane-2"); const closeSupervisorPane = vi.fn(async () => undefined); vi.mock("../src/herdr.js", () => ({ openSupervisorPane, closeSupervisorPane })); -vi.mock("../src/intercom.js", () => ({ - registerGoalsIntercom: () => ({ - workerIntercomId: async () => "worker-intercom", - waitForSupervisorReady: async () => {}, - announceSupervisorReady: async () => {}, - }), -})); - const { default: piGoalsExtension, isMainSession } = await import("../src/index.js"); function setup(selectChoices: Array, editorChoices: Array = []) { @@ -49,7 +42,14 @@ function setup(selectChoices: Array, editorChoices: Array editorChoices.shift(), }, }; + const events = new EventEmitter(); + events.on("pi-supervise:worker-state:v1", (reply) => reply({ intercomId: "worker-intercom" })); + openSupervisorPane.mockImplementation(async () => { + queueMicrotask(() => events.emit("pi-supervise:worker-paired:v1", { supervisorIntercomId: "supervisor-intercom" })); + return "pane-2"; + }); const pi = { + events, registerCommand: (name: string, command: any) => commands.set(name, command), on: (name: string, handler: any) => hooks.set(name, handler), appendEntry: (customType: string, data: unknown) => entries.push({ type: "custom", customType, data }), @@ -58,7 +58,7 @@ function setup(selectChoices: Array, editorChoices: Array messages.push({ content }), }; piGoalsExtension(pi as unknown as ExtensionAPI); - return { commands, ctx, cwd, entries, hooks, messages, notifications, tools }; + return { commands, ctx, cwd, entries, events, hooks, messages, notifications, tools }; } function writePlan(cwd: string, content: string): string { @@ -121,6 +121,24 @@ describe("/goals flow", () => { } }); + it("waits for the worker's real paired acknowledgement before beginning work", async () => { + const flow = setup(["Ready"]); + try { + openSupervisorPane.mockImplementationOnce(async () => "pane-2"); + await flow.commands.get("goals").handler("make the file", flow.ctx); + approvedPlan(flow.cwd); + const ready = flow.hooks.get("agent_settled")({}, flow.ctx); + await new Promise((resolve) => setImmediate(resolve)); + expect(flow.entries.at(-1)?.data).toMatchObject({ phase: "planning" }); + expect(flow.messages.some((message) => message.content === "The plan is approved. Begin implementation as the worker.")).toBe(false); + flow.events.emit("pi-supervise:worker-paired:v1", { supervisorIntercomId: "supervisor-intercom" }); + await ready; + expect(flow.entries.at(-1)?.data).toMatchObject({ phase: "working", supervisorPaneId: "pane-2" }); + } finally { + rmSync(flow.cwd, { recursive: true, force: true }); + } + }); + it("closes the supervisor on clear but keeps the plan file", async () => { const flow = setup(["Ready"]); try { diff --git a/test/herdr.test.ts b/test/herdr.test.ts index 2859083..0b3fe37 100644 --- a/test/herdr.test.ts +++ b/test/herdr.test.ts @@ -24,10 +24,9 @@ describe("supervisor pane command", () => { const command = supervisorCommand(input()); expect(command).toContain("'PI_GOALS_ROLE=supervisor'"); expect(command).toContain("'PI_GOALS_WORKER_INTERCOM_ID=intercom-12345678'"); - expect(command).toContain("'pi' '--no-extensions' '-e' '/repo/src/index.ts'"); - expect(command).toContain("'-e' 'npm:pi-intercom' '-e' 'npm:@wassname2/pi-supervise@0.0.4'"); + expect(command).toContain("'pi' '--no-extensions' '-e' 'npm:pi-intercom' '-e' 'npm:@wassname2/pi-supervise@0.0.4' '-e' '/repo/src/index.ts'"); expect(command).toContain("'--fork' '/sessions/worker.jsonl'"); - expect(command).toContain("'--model' 'provider/supervisor'"); + expect(command).toContain("'--model' 'provider/supervisor' 'Initialize supervision startup.'"); expect(command).not.toContain("pi-subagents"); }); diff --git a/test/supervise.test.ts b/test/supervise.test.ts new file mode 100644 index 0000000..01cfe4c --- /dev/null +++ b/test/supervise.test.ts @@ -0,0 +1,34 @@ +import { EventEmitter } from "node:events"; +import type { ExtensionAPI } from "@earendil-works/pi-coding-agent"; +import { describe, expect, it } from "vitest"; +import { workerPiSupervise } from "../src/supervise.js"; + +const API_READY = "pi-supervise:api-ready:v1"; +const WORKER_STATE = "pi-supervise:worker-state:v1"; +const WORKER_PAIRED = "pi-supervise:worker-paired:v1"; + +function pi(events: EventEmitter): ExtensionAPI { + return { events } as unknown as ExtensionAPI; +} + +describe("pi-supervise worker API", () => { + it("discovers pi-supervise when it loads after pi-goals", async () => { + const events = new EventEmitter(); + const worker = workerPiSupervise(pi(events)); + events.on(WORKER_STATE, (reply) => reply({ intercomId: "worker-id", paired: false })); + events.emit(API_READY); + expect((await worker).intercomId).toBe("worker-id"); + }); + + it("discovers an already-loaded pi-supervise and accepts duplicate paired events once", async () => { + const events = new EventEmitter(); + events.on(WORKER_STATE, (reply) => reply({ intercomId: "worker-id", paired: false })); + const worker = await workerPiSupervise(pi(events)); + let acknowledgements = 0; + void worker.paired.then(() => { acknowledgements += 1; }); + events.emit(WORKER_PAIRED, { supervisorIntercomId: "supervisor-id" }); + events.emit(WORKER_PAIRED, { supervisorIntercomId: "supervisor-id" }); + await worker.paired; + expect(acknowledgements).toBe(1); + }); +}); diff --git a/test/supervisor-session.test.ts b/test/supervisor-session.test.ts index ed650bf..9807b1b 100644 --- a/test/supervisor-session.test.ts +++ b/test/supervisor-session.test.ts @@ -17,7 +17,6 @@ function setup(cwd: string, planPath: string) { const tools = new Map(); const entries: any[] = []; const paired: Array<{ workerIntercomId: string; goal: string }> = []; - const announced: Array<{ workerIntercomId: string; approvalId: string }> = []; let branch: any[] = []; const ctx = { cwd, @@ -44,12 +43,8 @@ function setup(cwd: string, planPath: string) { registerTool: (tool: any) => tools.set(tool.name, tool), appendEntry: (customType: string, data: unknown) => entries.push({ type: "custom", customType, data }), }; - registerVisibleSupervisor(pi as unknown as ExtensionAPI, { - workerIntercomId: async () => "worker-intercom", - waitForSupervisorReady: async () => {}, - announceSupervisorReady: async (workerIntercomId, approvalId) => { announced.push({ workerIntercomId, approvalId }); }, - }); - return { announced, branch: (value: any[]) => { branch = value; }, ctx, entries, hooks, paired, tools }; + registerVisibleSupervisor(pi as unknown as ExtensionAPI); + return { branch: (value: any[]) => { branch = value; }, ctx, entries, hooks, paired, tools }; } afterEach(() => vi.unstubAllEnvs()); @@ -59,11 +54,23 @@ describe("visible supervisor session", () => { const cwd = mkdtempSync(join(tmpdir(), "pi-goals-supervisor-")); try { const runtime = setup(cwd, join(cwd, ".pi/plan/worker-v1.md")); - await runtime.hooks.get("session_start")({}, runtime.ctx); + await runtime.hooks.get("before_agent_start")({}, runtime.ctx); expect(runtime.ctx.compact).toHaveBeenCalledOnce(); expect(runtime.entries.at(-1)).toMatchObject({ customType: "pi-goals-visible-supervisor-v1" }); expect(runtime.paired).toEqual([{ workerIntercomId: "worker-intercom", goal: join(cwd, ".pi/plan/worker-v1.md") }]); - expect(runtime.announced).toEqual([{ workerIntercomId: "worker-intercom", approvalId: "approval-1" }]); + } finally { + rmSync(cwd, { recursive: true, force: true }); + } + }); + + it("does not pair twice when startup reaches a second worker turn", async () => { + const cwd = mkdtempSync(join(tmpdir(), "pi-goals-supervisor-")); + try { + const runtime = setup(cwd, join(cwd, ".pi/plan/worker-v1.md")); + await runtime.hooks.get("before_agent_start")({}, runtime.ctx); + await runtime.hooks.get("before_agent_start")({}, runtime.ctx); + expect(runtime.ctx.compact).toHaveBeenCalledOnce(); + expect(runtime.paired).toHaveLength(1); } finally { rmSync(cwd, { recursive: true, force: true }); }