From a5e86733866684b24211546ad275a45f98ff0165 Mon Sep 17 00:00:00 2001 From: elrond <171308700+elrond298@users.noreply.github.com> Date: Fri, 21 Aug 2026 11:12:57 +0800 Subject: [PATCH 1/5] fix(pi-plugin): support pi-web multi-session and RPC hosts Scope in-process child detection to lifecycle AsyncLocalStorage so independent sessions initialize normally, and release lifecycle subscriptions on shutdown. Run startup maintenance once per process and defer session-history reads until the backfill lease is acquired. Resolve the child Pi CLI independently from embedded host argv, present command output through RPC notifications and dialogs, and ensure Dreamer registration before manual runs. --- .../src/commands/ctx-commands.test.ts | 10 +- packages/pi-plugin/src/commands/ctx-dream.ts | 3 + packages/pi-plugin/src/commands/ctx-embed.ts | 1 + packages/pi-plugin/src/commands/ctx-recomp.ts | 1 + .../src/commands/ctx-session-upgrade.ts | 3 + packages/pi-plugin/src/commands/ctx-status.ts | 7 +- .../src/commands/pi-command-utils.test.ts | 74 ++++ .../src/commands/pi-command-utils.ts | 73 +++- packages/pi-plugin/src/dreamer/index.test.ts | 131 ++++++-- packages/pi-plugin/src/dreamer/index.ts | 28 +- .../pi-plugin/src/index-env-guard.test.ts | 7 +- .../src/index-in-process-latch.test.ts | 203 +++++++---- packages/pi-plugin/src/index.ts | 315 +++++++++++------- .../src/session-cleanup-wiring.test.ts | 2 +- .../pi-plugin/src/subagent-runner.test.ts | 29 +- packages/pi-plugin/src/subagent-runner.ts | 49 ++- .../session-project-backfill.test.ts | 8 +- 17 files changed, 686 insertions(+), 258 deletions(-) diff --git a/packages/pi-plugin/src/commands/ctx-commands.test.ts b/packages/pi-plugin/src/commands/ctx-commands.test.ts index 31381b87e..9d7fb5d64 100644 --- a/packages/pi-plugin/src/commands/ctx-commands.test.ts +++ b/packages/pi-plugin/src/commands/ctx-commands.test.ts @@ -259,17 +259,21 @@ describe("Pi Magic Context commands", () => { it("registers /ctx-dream and starts a run (Dreamer v2 manual path)", async () => { const db = createDb(); const { pi, handlers, sent } = createMockPi(); + const registrationCwds: string[] = []; registerCtxDreamCommand(pi as never, { db, projectDir: "/tmp/project", projectIdentity: "/tmp/project", + ensureRegistered: (ctx) => { + registrationCwds.push(ctx.cwd); + }, }); - // Not registered with the dreamer timer in this unit test, so runManual - // throws "not registered" → the handler reports the failure. We only - // assert the command is wired and emits a /ctx-dream status message. + // This unit test isolates the command from the real dreamer registry. The + // injected registration sync runs immediately before the manual dream call. await handlers.get("ctx-dream")?.("", createCtx()); + expect(registrationCwds).toEqual(["/tmp/project"]); expect(sent[0]?.customType).toBe("ctx-status"); expect(sent[0]?.data.text).toContain("/ctx-dream"); }); diff --git a/packages/pi-plugin/src/commands/ctx-dream.ts b/packages/pi-plugin/src/commands/ctx-dream.ts index 27f653eca..ac162b4b9 100644 --- a/packages/pi-plugin/src/commands/ctx-dream.ts +++ b/packages/pi-plugin/src/commands/ctx-dream.ts @@ -24,6 +24,7 @@ export function registerCtxDreamCommand( dreamerEnabled?: boolean; resolveDreamerEnabled?: (ctx: { cwd: string }) => boolean | undefined; onProjectSeen?: (projectIdentity: string) => void; + ensureRegistered?: (ctx: { cwd: string }) => void | Promise; }, ): void { pi.registerCommand("ctx-dream", { @@ -108,6 +109,7 @@ export function registerCtxDreamCommand( // Dreamer v2: run due/forced tasks now via the per-task scheduler. try { + await deps.ensureRegistered?.(ctx); const result = await runPiDreamForProject( project.projectIdentity, task, @@ -146,6 +148,7 @@ export function registerCtxDreamCommand( title: "/ctx-dream", text: ["## /ctx-dream", "", ...lines].join("\n"), level: result.ran.length > 0 ? "success" : "info", + rpcDisplay: "dialog", }, { projectDir: project.projectDir, diff --git a/packages/pi-plugin/src/commands/ctx-embed.ts b/packages/pi-plugin/src/commands/ctx-embed.ts index c7d5e7911..78bff2ad3 100644 --- a/packages/pi-plugin/src/commands/ctx-embed.ts +++ b/packages/pi-plugin/src/commands/ctx-embed.ts @@ -240,6 +240,7 @@ export function registerCtxEmbedCommand( title: "/ctx-embed", text: `## Embedding Status\n\n${statusText}`, level: "info", + rpcDisplay: "dialog", }); }, }); diff --git a/packages/pi-plugin/src/commands/ctx-recomp.ts b/packages/pi-plugin/src/commands/ctx-recomp.ts index 3f81b2bc1..d467df3b7 100644 --- a/packages/pi-plugin/src/commands/ctx-recomp.ts +++ b/packages/pi-plugin/src/commands/ctx-recomp.ts @@ -104,6 +104,7 @@ export function registerCtxRecompCommand( title: "/ctx-recomp", text: executeRecompUpgradeStub(currentDeps.db, sessionId), level: "info", + rpcDisplay: "dialog", }); return; } diff --git a/packages/pi-plugin/src/commands/ctx-session-upgrade.ts b/packages/pi-plugin/src/commands/ctx-session-upgrade.ts index 71936fa3c..7621a0f0a 100644 --- a/packages/pi-plugin/src/commands/ctx-session-upgrade.ts +++ b/packages/pi-plugin/src/commands/ctx-session-upgrade.ts @@ -198,6 +198,7 @@ export function registerCtxSessionUpgradeCommand( : "This session's compartments are already in the current format.", ].join("\n"), level: "info", + rpcDisplay: "dialog", }); return; } @@ -225,6 +226,7 @@ export function registerCtxSessionUpgradeCommand( title: "/ctx-session-upgrade", text: ["## Session Upgrade — Complete", "", summary].join("\n"), level: "info", + rpcDisplay: "dialog", }); }, }); @@ -375,6 +377,7 @@ export function registerCtxSessionUpgradeCommand( recompResult.message, ].join("\n"), level: "info", + rpcDisplay: "dialog", }); }, }); diff --git a/packages/pi-plugin/src/commands/ctx-status.ts b/packages/pi-plugin/src/commands/ctx-status.ts index b65be4a93..394f412cc 100644 --- a/packages/pi-plugin/src/commands/ctx-status.ts +++ b/packages/pi-plugin/src/commands/ctx-status.ts @@ -146,7 +146,12 @@ export function registerCtxStatusCommand( const details = buildStatusDetails(currentDeps, sessionId); sendCtxStatusMessage( pi, - { title: "/ctx-status", text: statusText, level: "info" }, + { + title: "/ctx-status", + text: statusText, + level: "info", + rpcDisplay: "dialog", + }, details, ); } catch (error) { diff --git a/packages/pi-plugin/src/commands/pi-command-utils.test.ts b/packages/pi-plugin/src/commands/pi-command-utils.test.ts index e627a14ef..085b13d3f 100644 --- a/packages/pi-plugin/src/commands/pi-command-utils.test.ts +++ b/packages/pi-plugin/src/commands/pi-command-utils.test.ts @@ -6,11 +6,15 @@ import { type PiMessageSender, registerCtxStatusEntryRenderer, sendCtxStatusMessage, + setCtxStatusPresenter, + shouldShowCtxStatusDialog, + showCtxStatusDialog, } from "./pi-command-utils"; describe("ctx-status entries", () => { it("appends model-invisible entry data instead of sending a message", () => { const appended: Array<{ customType: string; data: unknown }> = []; + const presented: CtxStatusEntryData[] = []; let sent = 0; const pi = { registerEntryRenderer() {}, @@ -23,6 +27,7 @@ describe("ctx-status entries", () => { } as unknown as PiMessageSender; registerCtxStatusEntryRenderer(pi); + setCtxStatusPresenter(pi, (content) => presented.push(content)); sendCtxStatusMessage( pi, { title: "Magic Embed", text: "Embedding history…", level: "info" }, @@ -41,6 +46,7 @@ describe("ctx-status entries", () => { }, ]); expect(sent).toBe(0); + expect(presented).toEqual([appended[0]?.data]); }); it("registers one ctx-status entry renderer and ignores malformed data", () => { @@ -91,4 +97,72 @@ describe("ctx-status entries", () => { ]); expect(sent).toBe(0); }); + + it("keeps progress notifications short and routes detailed results to a dialog", () => { + expect( + shouldShowCtxStatusDialog({ + title: "/ctx-dream", + text: "Starting…", + level: "info", + }), + ).toBe(false); + expect( + shouldShowCtxStatusDialog({ + title: "/ctx-flush", + text: "Complete", + level: "success", + }), + ).toBe(true); + expect( + shouldShowCtxStatusDialog({ + title: "/ctx-status", + text: "Detailed status", + level: "info", + rpcDisplay: "dialog", + }), + ).toBe(true); + }); + + it("renders RPC detail output through Pi custom UI", async () => { + let rendered: string[] = []; + let closed = false; + let options: unknown; + const ctx = { + ui: { + async custom(factory: unknown, customOptions: unknown) { + options = customOptions; + const create = factory as (...args: unknown[]) => { + render: (width: number) => string[]; + handleInput: (data: string) => void; + }; + const component = create( + {}, + { + fg: (_name: string, text: string) => text, + bold: (text: string) => text, + }, + {}, + () => { + closed = true; + }, + ); + rendered = component.render(92); + component.handleInput("\r"); + }, + }, + }; + + await showCtxStatusDialog(ctx as never, { + title: "/ctx-flush", + text: "## /ctx-flush\n\nDetailed result", + level: "success", + }); + + expect(rendered.join("\n")).toContain("Detailed result"); + expect(closed).toBe(true); + expect(options).toEqual({ + overlay: true, + overlayOptions: { anchor: "center", width: 92 }, + }); + }); }); diff --git a/packages/pi-plugin/src/commands/pi-command-utils.ts b/packages/pi-plugin/src/commands/pi-command-utils.ts index ee22b8471..a34c9f328 100644 --- a/packages/pi-plugin/src/commands/pi-command-utils.ts +++ b/packages/pi-plugin/src/commands/pi-command-utils.ts @@ -4,7 +4,7 @@ import type { ExtensionCommandContext, Theme, } from "@earendil-works/pi-coding-agent"; -import { Box, type Component, Text } from "@earendil-works/pi-tui"; +import { Box, type Component, matchesKey, Text } from "@earendil-works/pi-tui"; import { sessionLog } from "@magic-context/core/shared/logger"; export const CTX_STATUS_CUSTOM_TYPE = "ctx-status"; @@ -15,6 +15,7 @@ export interface CtxStatusEntryData { title: string; text: string; level?: CtxStatusLevel; + rpcDisplay?: "notification" | "dialog"; details?: unknown; } @@ -40,6 +41,28 @@ type PiEntryRendererRegistration = { export type PiMessageSender = Pick & PiEntryRendererRegistration; +type CtxStatusPresenter = (content: CtxStatusMessageContent) => void; + +const ctxStatusPresenters = new WeakMap(); + +export function setCtxStatusPresenter( + pi: PiMessageSender, + present: CtxStatusPresenter, +): void { + ctxStatusPresenters.set(pi, present); +} + +export function shouldShowCtxStatusDialog( + content: CtxStatusMessageContent, +): boolean { + return ( + content.rpcDisplay === "dialog" || + (content.rpcDisplay !== "notification" && + content.level !== undefined && + content.level !== "info") + ); +} + export function resolveSessionId( ctx: ExtensionCommandContext, ): string | undefined { @@ -55,6 +78,52 @@ export function resolveSessionId( } } +export async function showCtxStatusDialog( + ctx: Pick, + content: CtxStatusMessageContent, +): Promise { + await ctx.ui.custom( + (_tui, theme, _keybindings, done) => + new CtxStatusDialog(content, theme, done), + { overlay: true, overlayOptions: { anchor: "center", width: 92 } }, + ); +} + +class CtxStatusDialog implements Component { + constructor( + private readonly content: CtxStatusMessageContent, + private readonly theme: Theme, + private readonly done: (value: undefined) => void, + ) {} + + handleInput(data: string): void { + if ( + matchesKey(data, "escape") || + matchesKey(data, "ctrl+c") || + matchesKey(data, "return") + ) { + this.done(undefined); + } + } + + invalidate(): void {} + + render(_width: number): string[] { + return [ + this.theme.bold( + this.theme.fg( + statusTitleColor(this.content.level), + `[${this.content.title}]`, + ), + ), + "", + ...this.content.text.split("\n"), + "", + this.theme.fg("dim", "Press Enter or Escape to close"), + ]; + } +} + function statusTitleColor(level: CtxStatusLevel | undefined) { switch (level) { case "success": @@ -125,6 +194,8 @@ export function sendCtxStatusMessage( if (typeof pi.appendEntry === "function") { pi.appendEntry(CTX_STATUS_CUSTOM_TYPE, data); } + ctxStatusPresenters.get(pi)?.(data); + // Minimal non-interactive API shims may omit appendEntry; logging remains the // safe fallback and status text must never be routed through sendMessage. sessionLog("pi-status", `${content.title}: ${content.text}`); diff --git a/packages/pi-plugin/src/dreamer/index.test.ts b/packages/pi-plugin/src/dreamer/index.test.ts index 120f32745..2833b7fa6 100644 --- a/packages/pi-plugin/src/dreamer/index.test.ts +++ b/packages/pi-plugin/src/dreamer/index.test.ts @@ -56,6 +56,7 @@ function dreamerOptions(args: { database: Database; projectIdentity: string; projectDir?: string; + registrationOwner?: object; config?: DreamerConfig; language?: string; onAdjunctsRefreshNeeded?: (projectIdentity: string) => void; @@ -66,6 +67,7 @@ function dreamerOptions(args: { args.projectDir ?? `/tmp/${args.projectIdentity.replace(/[^a-z0-9-]/gi, "-")}`, projectIdentity: args.projectIdentity, + registrationOwner: args.registrationOwner ?? {}, config: args.config ?? enabledConfig(), embeddingConfig: { provider: "off" as const }, memoryEnabled: true, @@ -295,6 +297,45 @@ describe("Pi dreamer wiring", () => { expect(timerCleanup).not.toHaveBeenCalled(); }); + test("one session shutdown keeps a same-project sibling registered", async () => { + db = createDb(); + const firstCleanup = mock(() => {}); + const secondCleanup = mock(() => {}); + const cleanups = [firstCleanup, secondCleanup]; + __test.setStartDreamScheduleTimerFactory( + async () => cleanups.shift() ?? mock(() => {}), + ); + + const firstOpts = dreamerOptions({ + database: db, + projectDir: "/tmp/pi-shared-project", + projectIdentity: "git:pi-shared-project", + }); + const secondOpts = dreamerOptions({ + database: db, + projectDir: "/tmp/pi-shared-project", + projectIdentity: "git:pi-shared-project", + }); + registerPiDreamerProject(firstOpts); + await flushMicrotasks(); + registerPiDreamerProject(secondOpts); + + unregisterPiDreamerProject({ + projectIdentity: "git:pi-shared-project", + registrationOwner: firstOpts.registrationOwner, + }); + await flushMicrotasks(); + expect(__test.registeredProjectCount()).toBe(1); + expect(firstCleanup).toHaveBeenCalledTimes(1); + + unregisterPiDreamerProject({ + projectIdentity: "git:pi-shared-project", + registrationOwner: secondOpts.registrationOwner, + }); + expect(__test.registeredProjectCount()).toBe(0); + expect(secondCleanup).toHaveBeenCalledTimes(1); + }); + test("re-registering the same identity with a DIFFERENT dir rebuilds (worktree switch)", async () => { db = createDb(); const firstCleanup = mock(() => {}); @@ -307,22 +348,20 @@ describe("Pi dreamer wiring", () => { }); // Worktree A of the same repo → identity X. - registerPiDreamerProject( - dreamerOptions({ - database: db, - projectDir: "/tmp/worktree-A", - projectIdentity: "git:pi-worktree", - }), - ); + const firstOpts = dreamerOptions({ + database: db, + projectDir: "/tmp/worktree-A", + projectIdentity: "git:pi-worktree", + }); + registerPiDreamerProject(firstOpts); await flushMicrotasks(); // Worktree B of the SAME repo (same identity, different dir). - registerPiDreamerProject( - dreamerOptions({ - database: db, - projectDir: "/tmp/worktree-B", - projectIdentity: "git:pi-worktree", - }), - ); + const secondOpts = dreamerOptions({ + database: db, + projectDir: "/tmp/worktree-B", + projectIdentity: "git:pi-worktree", + }); + registerPiDreamerProject(secondOpts); await flushMicrotasks(); // Still one registration, but rebuilt: first timer torn down, second @@ -330,19 +369,36 @@ describe("Pi dreamer wiring", () => { expect(__test.registeredProjectCount()).toBe(1); expect(firstCleanup).toHaveBeenCalledTimes(1); expect(dirs).toEqual(["/tmp/worktree-A", "/tmp/worktree-B"]); + + // When the active worktree owner leaves, keep the sibling owner alive + // and restore its registration instead of deleting the project timer. + unregisterPiDreamerProject({ + projectIdentity: "git:pi-worktree", + registrationOwner: secondOpts.registrationOwner, + }); + await flushMicrotasks(); + expect(__test.registeredProjectCount()).toBe(1); + expect(secondCleanup).toHaveBeenCalledTimes(1); + expect(dirs).toEqual([ + "/tmp/worktree-A", + "/tmp/worktree-B", + "/tmp/worktree-A", + ]); }); test("unregister removes the project", () => { db = createDb(); - registerPiDreamerProject( - dreamerOptions({ - database: db, - projectDir: "/tmp/pi-project-unregister", - projectIdentity: "git:pi-unregister", - }), - ); + const opts = dreamerOptions({ + database: db, + projectDir: "/tmp/pi-project-unregister", + projectIdentity: "git:pi-unregister", + }); + registerPiDreamerProject(opts); - unregisterPiDreamerProject({ projectIdentity: "git:pi-unregister" }); + unregisterPiDreamerProject({ + projectIdentity: "git:pi-unregister", + registrationOwner: opts.registrationOwner, + }); expect(__test.registeredProjectCount()).toBe(0); }); @@ -498,10 +554,15 @@ describe("Pi dreamer wiring", () => { const timer = deferred<() => void>(); __test.setStartDreamScheduleTimerFactory(() => timer.promise); - registerPiDreamerProject( - dreamerOptions({ database: db, projectIdentity: "git:pi-g12-race" }), - ); - unregisterPiDreamerProject({ projectIdentity: "git:pi-g12-race" }); + const opts = dreamerOptions({ + database: db, + projectIdentity: "git:pi-g12-race", + }); + registerPiDreamerProject(opts); + unregisterPiDreamerProject({ + projectIdentity: "git:pi-g12-race", + registrationOwner: opts.registrationOwner, + }); expect(timerCleanup).not.toHaveBeenCalled(); timer.resolve(timerCleanup); @@ -516,14 +577,22 @@ describe("Pi dreamer wiring", () => { const timer = deferred<() => void>(); __test.setStartDreamScheduleTimerFactory(() => timer.promise); - registerPiDreamerProject( - dreamerOptions({ database: db, projectIdentity: "git:pi-g12-normal" }), - ); + const opts = dreamerOptions({ + database: db, + projectIdentity: "git:pi-g12-normal", + }); + registerPiDreamerProject(opts); timer.resolve(timerCleanup); await flushMicrotasks(); - unregisterPiDreamerProject({ projectIdentity: "git:pi-g12-normal" }); - unregisterPiDreamerProject({ projectIdentity: "git:pi-g12-normal" }); + unregisterPiDreamerProject({ + projectIdentity: "git:pi-g12-normal", + registrationOwner: opts.registrationOwner, + }); + unregisterPiDreamerProject({ + projectIdentity: "git:pi-g12-normal", + registrationOwner: opts.registrationOwner, + }); expect(timerCleanup).toHaveBeenCalledTimes(1); }); diff --git a/packages/pi-plugin/src/dreamer/index.ts b/packages/pi-plugin/src/dreamer/index.ts index f54d1e45d..ced669dc0 100644 --- a/packages/pi-plugin/src/dreamer/index.ts +++ b/packages/pi-plugin/src/dreamer/index.ts @@ -24,6 +24,8 @@ export interface PiDreamerOptions { db: ContextDatabase; projectDir: string; projectIdentity: string; + /** One stable token per full Pi extension instance. */ + registrationOwner: object; /** Resolved runnable DreamerConfig from loadPiConfig(). When disable=true, the caller does not register. */ config: DreamerConfig; /** @@ -83,6 +85,8 @@ type SessionDeleteArgs = SessionMessagesArgs; interface ProjectRegistration { cleanup: () => void; + activeOwner: object; + owners: Map; /** Run dream tasks for this project IMMEDIATELY (Dreamer v2 manual path). * `task` forces one task ignoring its gate; omitted runs all enabled. The * registered dreamer timer also runs due tasks on its own schedule. */ @@ -122,8 +126,12 @@ export function registerPiDreamerProject(opts: PiDreamerOptions): void { } const existing = registeredProjects.get(opts.projectIdentity); + const owners = existing?.owners ?? new Map(); + owners.set(opts.registrationOwner, opts); if (existing) { // Same identity, same directory → genuinely already registered, no-op. + // Keep this extension instance as an owner so another session cannot + // deregister the shared timer while it is still active. if (existing.projectDir === opts.projectDir) { return; } @@ -209,6 +217,8 @@ export function registerPiDreamerProject(opts: PiDreamerOptions): void { }); registeredProjects.set(opts.projectIdentity, { + activeOwner: opts.registrationOwner, + owners, cleanup: () => { cancelled = true; cleanup?.(); @@ -244,12 +254,26 @@ export async function runPiDreamForProject( return registration.runManual(task); } -/** Cleanup hook — call from session_shutdown to deregister this project. */ +/** Cleanup hook — call from session_shutdown to release this session's ownership. */ export function unregisterPiDreamerProject(opts: { projectIdentity: string; + registrationOwner: object; }): void { const registration = registeredProjects.get(opts.projectIdentity); - if (!registration) { + if (!registration?.owners.delete(opts.registrationOwner)) { + return; + } + + if (registration.owners.size > 0) { + if (registration.activeOwner !== opts.registrationOwner) return; + // The active worktree owner left while sibling sessions still use this + // project. Rebuild from the remaining owners so the shared timer follows + // a live session instead of retaining the departed session's directory. + const remaining = [...registration.owners.values()]; + registration.cleanup(); + registeredProjects.delete(opts.projectIdentity); + for (const remainingOptions of remaining) + registerPiDreamerProject(remainingOptions); return; } diff --git a/packages/pi-plugin/src/index-env-guard.test.ts b/packages/pi-plugin/src/index-env-guard.test.ts index fce4505de..8c0b6059c 100644 --- a/packages/pi-plugin/src/index-env-guard.test.ts +++ b/packages/pi-plugin/src/index-env-guard.test.ts @@ -32,6 +32,7 @@ function createCountingPi() { const commands: string[] = []; const entryRenderers: string[] = []; const pi = { + events: { on: mock(() => () => undefined) }, on: mock((event: string) => { events.push(event); }), @@ -56,9 +57,9 @@ function createCountingPi() { afterEach(() => { restoreEnv(); - // Clear the process-global init latch so one test's full init does not - // leak into the next (the latch lives on globalThis, not module state). - __test.clearPiMagicContextActive(); + // Clear the process-global marker context so one test's full init does not + // leak into the next (the holder lives on globalThis, not module state). + __test.clearPiInProcessSubagentInitContext(); }); describe("Pi full extension subagent env guard", () => { diff --git a/packages/pi-plugin/src/index-in-process-latch.test.ts b/packages/pi-plugin/src/index-in-process-latch.test.ts index 0cd8428e0..24f133ea1 100644 --- a/packages/pi-plugin/src/index-in-process-latch.test.ts +++ b/packages/pi-plugin/src/index-in-process-latch.test.ts @@ -22,15 +22,15 @@ function restoreEnv() { function isolateXdgEnv() { const root = mkdtempSync(join(tmpdir(), "magic-context-pi-latch-test-")); process.env.XDG_CONFIG_HOME = join(root, "config"); - process.env.XDG_DATA_HOME = join(root, "data"); + // Use the preload's migration-safe test database; isolate only configuration. + delete process.env.XDG_DATA_HOME; } /** - * Counting ExtensionAPI seam. Every registration method pushes the name onto - * a list, so a test can assert that a second init registered NOTHING (no - * duplicate tools, events, commands, timers, or watchers). The `on` mock is - * the key seam for the latch: a second init that no-ops must not register any - * event handlers, because those handlers would wire timers / background scans. + * Counting ExtensionAPI seam. Every ordinary registration method pushes the name onto + * a list, so a test can assert that a child init registered NOTHING (no + * duplicate tools, events, commands, timers, or watchers). The custom event + * bus drives the in-process child lifecycle signal. */ function createCountingPi() { const events: string[] = []; @@ -38,10 +38,28 @@ function createCountingPi() { const flags: string[] = []; const commands: string[] = []; const entryRenderers: string[] = []; + const eventBusHandlers = new Map void>>(); + const piEventHandlers = new Map< + string, + Set<(event: unknown, ctx: unknown) => unknown> + >(); const pi = { - on: mock((event: string) => { - events.push(event); - }), + events: { + on(channel: string, handler: (data: unknown) => void) { + const handlers = eventBusHandlers.get(channel) ?? new Set(); + handlers.add(handler); + eventBusHandlers.set(channel, handlers); + return () => handlers.delete(handler); + }, + }, + on: mock( + (event: string, handler: (event: unknown, ctx: unknown) => unknown) => { + events.push(event); + const handlers = piEventHandlers.get(event) ?? new Set(); + handlers.add(handler); + piEventHandlers.set(event, handlers); + }, + ), registerTool: mock((tool: { name?: string }) => { tools.push(tool.name ?? ""); }), @@ -58,73 +76,122 @@ function createCountingPi() { sendMessage: mock(() => undefined), sendUserMessage: mock(() => undefined), } as unknown as ExtensionAPI; - return { pi, events, tools, flags, commands, entryRenderers }; + return { + pi, + events, + tools, + flags, + commands, + entryRenderers, + eventBusHandlerCount(channel: string) { + return eventBusHandlers.get(channel)?.size ?? 0; + }, + emitEvent(channel: string, data: unknown = {}) { + for (const handler of eventBusHandlers.get(channel) ?? []) handler(data); + }, + async emitPiEvent(event: string, data: unknown = {}, ctx: unknown = {}) { + for (const handler of piEventHandlers.get(event) ?? []) { + await handler(data, ctx); + } + }, + }; } afterEach(() => { restoreEnv(); - // The latch lives on globalThis (process-global by design), so clear it - // between tests or one test's init would suppress the next. - __test.clearPiMagicContextActive(); + // The marker context lives on globalThis (process-global by design), so clear it + // between tests or one test's child state could suppress the next. + __test.clearPiInProcessSubagentInitContext(); + __test.clearPiStartupMaintenanceClaim(); }); -describe("Pi in-process re-init latch (#247)", () => { - it("second init in the same process is a no-op (no duplicate registrations)", async () => { +describe("Pi in-process child guard (#247)", () => { + it("claims process-wide startup maintenance from the full runtime", async () => { isolateXdgEnv(); delete process.env[MAGIC_CONTEXT_PI_SUBAGENT_ENV]; - __test.clearPiMagicContextActive(); const first = createCountingPi(); await magicContextPiExtension(first.pi); + expect(__test.claimPiStartupMaintenance()).toBe(false); + + const second = createCountingPi(); + await magicContextPiExtension(second.pi); + expect(__test.claimPiStartupMaintenance()).toBe(false); + }, 15_000); + it("registers independent sessions in the same process", async () => { + isolateXdgEnv(); + delete process.env[MAGIC_CONTEXT_PI_SUBAGENT_ENV]; + const first = createCountingPi(); + await magicContextPiExtension(first.pi); // Sanity: the first init registered the full runtime. expect(first.events.length).toBeGreaterThan(0); - expect(first.tools.length).toBeGreaterThan(0); - expect(first.commands.length).toBeGreaterThan(0); + expect(first.tools).toContain("ctx_search"); + expect(first.commands).toContain("ctx-status"); expect(first.entryRenderers).toEqual(["ctx-status"]); - // The latch is now set in this process. - expect(__test.isPiMagicContextActiveInProcess()).toBe(true); - - // Second init in the SAME process (the in-process child case). - // It must register nothing — same contract as a spawned subagent. const second = createCountingPi(); await magicContextPiExtension(second.pi); - - expect(second.events).toEqual([]); - expect(second.tools).toEqual([]); - expect(second.flags).toEqual([]); - expect(second.commands).toEqual([]); - expect(second.entryRenderers).toEqual([]); + expect(second.events.length).toBeGreaterThan(0); + expect(second.tools).toContain("ctx_search"); + expect(second.commands).toContain("ctx-status"); + expect(second.entryRenderers).toEqual(["ctx-status"]); }, 15_000); - it("clearing the latch (dispose) allows a full re-init", async () => { + it("unsubscribes child lifecycle listeners on session shutdown", async () => { isolateXdgEnv(); delete process.env[MAGIC_CONTEXT_PI_SUBAGENT_ENV]; - __test.clearPiMagicContextActive(); - const first = createCountingPi(); - await magicContextPiExtension(first.pi); - expect(first.tools.length).toBeGreaterThan(0); + const runtime = createCountingPi(); + await magicContextPiExtension(runtime.pi); + expect( + runtime.eventBusHandlerCount("subagents:child:session-created"), + ).toBe(1); + expect(runtime.eventBusHandlerCount("subagents:child:disposed")).toBe(1); + + await runtime.emitPiEvent( + "session_shutdown", + {}, + { + sessionManager: { getSessionId: () => undefined }, + ui: { setStatus: () => undefined }, + }, + ); + expect( + runtime.eventBusHandlerCount("subagents:child:session-created"), + ).toBe(0); + expect(runtime.eventBusHandlerCount("subagents:child:disposed")).toBe(0); + }, 15_000); - // Simulate the session_shutdown dispose path clearing the latch. - __test.clearPiMagicContextActive(); - expect(__test.isPiMagicContextActiveInProcess()).toBe(false); + it("skips only the marked in-process child", async () => { + isolateXdgEnv(); + delete process.env[MAGIC_CONTEXT_PI_SUBAGENT_ENV]; - // A subsequent init re-registers the full runtime. - const second = createCountingPi(); - await magicContextPiExtension(second.pi); + const parent = createCountingPi(); + await magicContextPiExtension(parent.pi); + parent.emitEvent("subagents:child:spawning"); + parent.emitEvent("subagents:child:session-created"); - expect(second.events.length).toBeGreaterThan(0); - expect(second.tools.length).toBeGreaterThan(0); - expect(second.commands.length).toBeGreaterThan(0); - expect(second.entryRenderers).toEqual(["ctx-status"]); + // Second init in the SAME process (the in-process child case). + // It must register nothing — same contract as a spawned subagent. + const child = createCountingPi(); + await magicContextPiExtension(child.pi); + expect(child.events).toEqual([]); + expect(child.tools).toEqual([]); + expect(child.commands).toEqual([]); + + // Simulate the child dispose path clearing its lifecycle marker. + parent.emitEvent("subagents:child:disposed"); + // A subsequent independent init re-registers the full runtime. + const sibling = createCountingPi(); + await magicContextPiExtension(sibling.pi); + expect(sibling.tools).toContain("ctx_search"); + expect(sibling.commands).toContain("ctx-status"); }, 15_000); - it("spawned-child env guard still no-ops even when the latch is clear", async () => { + it("keeps the spawned-child environment guard", async () => { isolateXdgEnv(); process.env[MAGIC_CONTEXT_PI_SUBAGENT_ENV] = "1"; - __test.clearPiMagicContextActive(); const registrations = createCountingPi(); await magicContextPiExtension(registrations.pi); @@ -134,35 +201,35 @@ describe("Pi in-process re-init latch (#247)", () => { expect(registrations.flags).toEqual([]); expect(registrations.commands).toEqual([]); expect(registrations.entryRenderers).toEqual([]); - // The env guard returns BEFORE setting the latch, so a later in-process - // init in the same process would still initialize fully. This pins the - // spawned-child contract: the env guard is a separate, earlier gate. - expect(__test.isPiMagicContextActiveInProcess()).toBe(false); + // The env guard returns BEFORE registering lifecycle markers, so a later + // independent init in the same process would still initialize fully. + delete process.env[MAGIC_CONTEXT_PI_SUBAGENT_ENV]; + const later = createCountingPi(); + await magicContextPiExtension(later.pi); + expect(later.tools).toContain("ctx_search"); }); - it("mutation direction: removing the latch makes the double-init test fail", async () => { - // This test documents the regression guard: if the latch check is - // removed from the entry, a second init would re-register everything. - // We simulate the "latch removed" state by clearing it between the two - // inits and asserting the second init then registers the full runtime - // — proving the latch is what suppresses it. + it("mutation direction: clearing the marker makes the child-init test fail", async () => { + // This test documents the regression guard: if the marker check is + // removed from the entry, a child init would register everything. + // We simulate the marker being absent before the child init and assert + // that it then registers the full runtime — proving the marker suppresses it. isolateXdgEnv(); delete process.env[MAGIC_CONTEXT_PI_SUBAGENT_ENV]; - __test.clearPiMagicContextActive(); - const first = createCountingPi(); - await magicContextPiExtension(first.pi); - expect(first.tools.length).toBeGreaterThan(0); + const parent = createCountingPi(); + await magicContextPiExtension(parent.pi); + parent.emitEvent("subagents:child:session-created"); - // Simulate the latch being absent: clear it before the second init. - __test.clearPiMagicContextActive(); + // Simulate the marker being absent: clear it before the child init. + __test.clearPiInProcessSubagentInitContext(); - const second = createCountingPi(); - await magicContextPiExtension(second.pi); + const child = createCountingPi(); + await magicContextPiExtension(child.pi); - // Without the latch suppressing it, the second init re-registers. - expect(second.events.length).toBeGreaterThan(0); - expect(second.tools.length).toBeGreaterThan(0); - expect(second.commands.length).toBeGreaterThan(0); + // Without the marker suppressing it, the child init registers. + expect(child.events.length).toBeGreaterThan(0); + expect(child.tools.length).toBeGreaterThan(0); + expect(child.commands.length).toBeGreaterThan(0); }, 15_000); }); diff --git a/packages/pi-plugin/src/index.ts b/packages/pi-plugin/src/index.ts index 5f61a181d..2e09c381c 100644 --- a/packages/pi-plugin/src/index.ts +++ b/packages/pi-plugin/src/index.ts @@ -20,6 +20,7 @@ * Falls back to schema defaults when neither file exists. */ +import { AsyncLocalStorage } from "node:async_hooks"; import { createRequire } from "node:module"; import { join, resolve } from "node:path"; import type { ExtensionAPI } from "@earendil-works/pi-coding-agent"; @@ -120,6 +121,9 @@ import { registerCtxWrapupCommand } from "./commands/ctx-wrapup"; import { registerCtxStatusEntryRenderer, sendCtxStatusMessage, + setCtxStatusPresenter, + shouldShowCtxStatusDialog, + showCtxStatusDialog, } from "./commands/pi-command-utils"; import { loadPiConfig } from "./config"; import { @@ -184,7 +188,7 @@ import { const PREFIX = "[magic-context][pi]"; // --------------------------------------------------------------------------- -// Process-global init latch (issue #247) +// In-process child guard (issue #247) // // `@gotgenes/pi-subagents` runs child agent sessions IN-PROCESS inside the // parent Pi process. Each child inherits the parent's user packages, so Pi @@ -197,41 +201,70 @@ const PREFIX = "[magic-context][pi]"; // parallel children fanned out concurrent `SessionManager.listAll` scans over // ~392 JSONL sessions and crashed the parent with heap OOM. // -// The latch below is a `Symbol.for` key on `globalThis` so it survives the +// The marker below is a `Symbol.for` key on `globalThis` so it survives the // duplicate module instances Pi's jiti loader creates per session // (`moduleCache: false` resets module-level state on every re-import, but a -// Symbol.for key is process-global). The first init in this process sets it; -// every later init in the same process (in-process child, or a second factory -// call from any source) sees it set and no-ops with the SAME contract as a -// spawned subagent child — no watchers, no timers, no background scans. The -// parent's already-registered extension instance keeps serving its session. +// Symbol.for key is process-global). The `session-created` lifecycle event sets +// it only in the child's async context; that child factory sees it and no-ops +// with the SAME contract as a spawned subagent child — no watchers, no timers, +// no background scans. The parent's already-registered extension instance keeps +// serving its session, while independent same-process sessions initialize normally. // -// Dispose / re-arm: Pi fires `session_shutdown` (reason "reload") before a -// `/reload` re-imports extensions, and (reason "shutdown") when the user -// leaves the session. Each AgentSession owns its own ExtensionRunner, so a -// child session's `session_shutdown` only fires handlers the CHILD registered -// (none, because the child no-op'd) — it cannot clear the parent's latch. -// We clear the latch in the parent's `session_shutdown` handler so a `/reload` -// legitimately re-initializes, while ephemeral in-process children never touch -// it. +// Dispose / re-arm: `subagents:child:disposed` clears the child's marker after +// its run. Since the marker is scoped with AsyncLocalStorage, a child cannot +// suppress unrelated sessions hosted by pi-web. // --------------------------------------------------------------------------- -const PI_ACTIVE_LATCH = Symbol.for("magic-context.pi.active"); +const PI_CHILD_INIT_CONTEXT = Symbol.for("magic-context.pi.child-init-context"); +const SUBAGENT_CHILD_SESSION_CREATED = "subagents:child:session-created"; +const SUBAGENT_CHILD_DISPOSED = "subagents:child:disposed"; +const PI_STARTUP_MAINTENANCE_SCHEDULED = Symbol.for( + "magic-context.pi.startup-maintenance-scheduled", +); + +function getPiChildInitContext(): AsyncLocalStorage { + const globals = globalThis as Record; + const existing = globals[PI_CHILD_INIT_CONTEXT]; + if (existing instanceof AsyncLocalStorage) return existing; + const context = new AsyncLocalStorage(); + globals[PI_CHILD_INIT_CONTEXT] = context; + return context; +} -function isPiMagicContextActiveInProcess(): boolean { - return (globalThis as Record)[PI_ACTIVE_LATCH] === true; +function isPiInProcessSubagentInit(): boolean { + return getPiChildInitContext().getStore() === true; } -function markPiMagicContextActive(): void { - (globalThis as Record)[PI_ACTIVE_LATCH] = true; +function clearPiInProcessSubagentInitContext(): void { + getPiChildInitContext().enterWith(false); } -function clearPiMagicContextActive(): void { - try { - delete (globalThis as Record)[PI_ACTIVE_LATCH]; - } catch { - // Some runtimes disallow delete on globalThis; fall back to overwrite. - (globalThis as Record)[PI_ACTIVE_LATCH] = undefined; - } +function registerPiSubagentInitContext(pi: ExtensionAPI): () => void { + const context = getPiChildInitContext(); + // session-created fires after child creation has its own async branch but before + // bindExtensions(); marking on spawning would leak into the parent's call chain. + const unsubscribeCreated = pi.events.on(SUBAGENT_CHILD_SESSION_CREATED, () => + context.enterWith(true), + ); + const unsubscribeDisposed = pi.events.on(SUBAGENT_CHILD_DISPOSED, () => + context.enterWith(false), + ); + return () => { + unsubscribeCreated(); + unsubscribeDisposed(); + }; +} + +function claimPiStartupMaintenance(): boolean { + const globals = globalThis as Record; + if (globals[PI_STARTUP_MAINTENANCE_SCHEDULED] === true) return false; + globals[PI_STARTUP_MAINTENANCE_SCHEDULED] = true; + return true; +} + +function clearPiStartupMaintenanceClaim(): void { + delete (globalThis as Record)[ + PI_STARTUP_MAINTENANCE_SCHEDULED + ]; } function resolveCurrentProject( @@ -465,9 +498,9 @@ export const __test = { resetLoggedPiConfigDirs(): void { loggedPiConfigDirs.clear(); }, - isPiMagicContextActiveInProcess, - markPiMagicContextActive, - clearPiMagicContextActive, + clearPiInProcessSubagentInitContext, + claimPiStartupMaintenance, + clearPiStartupMaintenanceClaim, }; function formatTokens(value: number): string { @@ -732,18 +765,17 @@ export default async function (pi: ExtensionAPI): Promise { // In-process child guard (issue #247): `@gotgenes/pi-subagents` runs child // agent sessions in the SAME process as the parent. They share the parent's // env (so the spawned-child env guard above never fires) and re-trigger this - // factory for every child session. The process-global latch marks that the - // full Magic Context runtime is already active in this process; a second - // init no-ops with the same contract as a spawned subagent (no watchers, no - // timers, no background scans). The parent's registered instance keeps - // serving. See the latch block above for the dispose / `/reload` re-arm path. - if (isPiMagicContextActiveInProcess()) { + // factory for every child session. The lifecycle marker scopes the no-op to + // that child: no database, watchers, timers, or background scans. Independent + // same-process sessions remain unmarked and initialize normally. + if (isPiInProcessSubagentInit()) { log( - `${PREFIX} in-process re-init detected (Magic Context already active in this process); skipping full extension registration`, + `${PREFIX} in-process subagent child detected; skipping full extension registration`, ); return; } - markPiMagicContextActive(); + const unregisterPiSubagentInitContext = registerPiSubagentInitContext(pi); + registerPiSubagentInitContextCleanup(pi, unregisterPiSubagentInitContext); beginBootQuietPeriod(); // Resolve the user-tier storage policy before opening the shared database. @@ -862,39 +894,69 @@ async function startPiMagicContextRuntime( // v22 deferred legacy-memory identity backfill. openDatabase() has already // run migrations; the runner is fire-and-forget and logs failures without - // blocking Pi startup. - scheduleAfterBootQuiet(() => { - runDeferredV22Backfill(db).catch((err) => { - warn(`[v22-backfill] background runner failed: ${err}`); + // blocking Pi startup. Multiple independent AgentSessions share one process, so + // only the first full runtime schedules process-wide startup maintenance. + if (claimPiStartupMaintenance()) { + scheduleAfterBootQuiet(() => { + runDeferredV22Backfill(db).catch((err) => { + warn(`[v22-backfill] background runner failed: ${err}`); + }); }); - }); - scheduleAfterBootQuiet(() => { - void (async () => { - try { - const api = await loadDefaultPiSessionApi(); - const sessions = (await api.listSessions()) as Array<{ - id?: unknown; - cwd?: unknown; - }>; - await runSessionProjectBackfill( - database, - sessions.map((session) => ({ - sessionId: typeof session?.id === "string" ? session.id : "", - directory: typeof session?.cwd === "string" ? session.cwd : "", - })), - ); - } catch (err) { - warn(`[session-projects] background runner failed: ${err}`); - } - })(); - }, 0); + scheduleAfterBootQuiet(() => { + void (async () => { + try { + let sessions: + | Array<{ sessionId: string; directory: string }> + | undefined; + await runSessionProjectBackfill( + database, + async (afterSessionId, limit) => { + if (!sessions) { + const api = await loadDefaultPiSessionApi(); + const sessionsById = new Map< + string, + { sessionId: string; directory: string } + >(); + for (const session of (await api.listSessions()) as Array<{ + id?: unknown; + cwd?: unknown; + }>) { + const sessionId = + typeof session?.id === "string" ? session.id : ""; + const directory = + typeof session?.cwd === "string" ? session.cwd : ""; + if ( + sessionId && + (!sessionsById.has(sessionId) || directory) + ) { + sessionsById.set(sessionId, { sessionId, directory }); + } + } + sessions = [...sessionsById.values()]; + } + const offset = + afterSessionId === null + ? 0 + : sessions.findIndex( + (session) => session.sessionId === afterSessionId, + ) + 1; + return sessions.slice(offset, offset + limit); + }, + ); + } catch (err) { + warn(`[session-projects] background runner failed: ${err}`); + } + })(); + }, 0); + } // Capture boot project for initial config load and logging only. Runtime // identity/path resolution uses ctx.cwd per hook/command so session cwd // switches follow the active project without reloading config. const projectDir = process.cwd(); const seenDreamerProjectIdentities = new Set(); + const dreamerRegistrationOwner = {}; // Step 5b: load the user's full magic-context.jsonc config. The loader // reads the shared CortexKit project/user paths, validates them through the // shared Zod schema, falls back to Pi-owned legacy files only while migration @@ -1129,6 +1191,37 @@ async function startPiMagicContextRuntime( const bootProjectDeps = buildProjectDeps(projectDir, projectIdentity, config); projectDepsByDir.set(projectDir, bootProjectDeps); + + function syncDreamerProjectRegistration( + current: ResolvedPiProjectDeps, + ): void { + seenDreamerProjectIdentities.add(current.projectIdentity); + if (!current.dreamerConfig) { + unregisterPiDreamerProject({ + projectIdentity: current.projectIdentity, + registrationOwner: dreamerRegistrationOwner, + }); + return; + } + registerPiDreamerProject({ + db, + projectDir: current.projectDir, + projectIdentity: current.projectIdentity, + registrationOwner: dreamerRegistrationOwner, + config: current.dreamerConfig, + // Council finding #7: thread real embedding + memory config so + // dreamer can do semantic dedup AND can write memory updates. + // Previously hardcoded to off/false, making most dreamer tasks + // useless on Pi. + embeddingConfig: current.config.embedding, + memoryEnabled: current.config.memory.enabled, + retinaHandoff: current.config.smart_notes.retina_handoff, + mural: current.config.mural, + language: current.config.language, + gitCommitIndexing: current.config.memory.git_commit_indexing, + onAdjunctsRefreshNeeded: signalPiSystemPromptRefreshForProject, + }); + } const todowriteEnabled = bootProjectDeps.config.todowrite.enabled !== false; const todowriteOverlayEnabled = todowriteEnabled && bootProjectDeps.config.todowrite.overlay !== false; @@ -1188,6 +1281,23 @@ async function startPiMagicContextRuntime( ); pi.on("session_start", async (event, ctx) => { + if (ctx.mode === "rpc") { + setCtxStatusPresenter(pi, (content) => { + const type = + content.level === "error" || content.level === "warning" + ? content.level + : "info"; + const useDialog = shouldShowCtxStatusDialog(content); + if (!useDialog) { + ctx.ui.notify(content.text, type); + return; + } + void showCtxStatusDialog(ctx, content).catch((err) => { + warn("ctx status dialog failed:", err); + ctx.ui.notify(content.text, type); + }); + }); + } await handlePiCloneSessionStart(event, ctx, { db, signalPendingMarker: signalPiDeferredCompactionMarkerDrain, @@ -1438,6 +1548,8 @@ async function startPiMagicContextRuntime( resolveDreamerEnabled: (ctx) => resolveCurrentProjectDeps(ctx).dreamerEnabled, onProjectSeen: (identity) => seenDreamerProjectIdentities.add(identity), + ensureRegistered: (ctx) => + syncDreamerProjectRegistration(resolveCurrentProjectDeps(ctx)), }); info("registered /ctx-dream"); @@ -1464,23 +1576,7 @@ async function startPiMagicContextRuntime( // PiSubagentRunner to spawn child sessions for each task. const dreamerConfig = bootProjectDeps.dreamerConfig; if (dreamerConfig) { - registerPiDreamerProject({ - db, - projectDir, - projectIdentity, - config: dreamerConfig, - // Council finding #7: thread real embedding + memory config so - // dreamer can do semantic dedup AND can write memory updates. - // Previously hardcoded to off/false, making most dreamer tasks - // useless on Pi. - embeddingConfig: bootProjectDeps.config.embedding, - memoryEnabled: bootProjectDeps.config.memory.enabled, - retinaHandoff: bootProjectDeps.config.smart_notes.retina_handoff, - mural: bootProjectDeps.config.mural, - language: bootProjectDeps.config.language, - gitCommitIndexing: bootProjectDeps.config.memory.git_commit_indexing, - onAdjunctsRefreshNeeded: signalPiSystemPromptRefreshForProject, - }); + syncDreamerProjectRegistration(bootProjectDeps); info(`registered dreamer (${summarizeDreamSchedule(dreamerConfig)})`); } else { info( @@ -1555,7 +1651,6 @@ async function startPiMagicContextRuntime( projectIdentity: effectiveProjectDeps.projectIdentity, }; const effectiveConfig = effectiveProjectDeps.config; - seenDreamerProjectIdentities.add(currentProject.projectIdentity); // Re-register the dreamer for the CURRENT project. The boot-time // registration above used process.cwd(), but Pi can switch projects @@ -1570,35 +1665,13 @@ async function startPiMagicContextRuntime( // pipeline. A switched-into project may carry its own config (different // model/schedule, or its own `dreamer.disable`), so boot config must not // leak into this registration. - const effectiveDreamerConfig = effectiveProjectDeps.dreamerConfig; - if (effectiveDreamerConfig) { - try { - registerPiDreamerProject({ - db, - projectDir: currentProject.projectDir, - projectIdentity: currentProject.projectIdentity, - config: effectiveDreamerConfig, - embeddingConfig: effectiveConfig.embedding, - memoryEnabled: effectiveConfig.memory.enabled, - retinaHandoff: effectiveConfig.smart_notes.retina_handoff, - language: effectiveConfig.language, - gitCommitIndexing: effectiveConfig.memory.git_commit_indexing, - onAdjunctsRefreshNeeded: signalPiSystemPromptRefreshForProject, - }); - } catch (err) { - warn("before_agent_start: registerPiDreamerProject threw:", err); - } - } else { - // The current checkout disables the dreamer. Any existing registration - // for this identity may have been created while another checkout's - // config was active, so tear it down explicitly here. - try { - unregisterPiDreamerProject({ - projectIdentity: currentProject.projectIdentity, - }); - } catch (err) { - warn("before_agent_start: unregisterPiDreamerProject threw:", err); - } + // The current checkout may also disable the dreamer. This instance may own a + // registration created while another checkout's config was active, so the + // shared helper releases that ownership explicitly. + try { + syncDreamerProjectRegistration(effectiveProjectDeps); + } catch (err) { + warn("before_agent_start: dreamer registration sync threw:", err); } // Pi exposes `sessionManager.getSessionId()` once a session is // active. We resolve it here defensively because before_agent_start @@ -2294,7 +2367,10 @@ async function startPiMagicContextRuntime( } try { for (const identity of seenDreamerProjectIdentities) { - unregisterPiDreamerProject({ projectIdentity: identity }); + unregisterPiDreamerProject({ + projectIdentity: identity, + registrationOwner: dreamerRegistrationOwner, + }); } } catch (err) { warn("shutdown: unregisterPiDreamerProject threw:", err); @@ -2325,16 +2401,6 @@ async function startPiMagicContextRuntime( } catch { // best-effort cleanup } - // Re-arm the process-global init latch (issue #247). Pi fires - // `session_shutdown` (reason "reload") before a `/reload` re-imports - // extensions, and (reason "shutdown") when the user leaves the - // session. Each AgentSession owns its own ExtensionRunner, so an - // in-process child's `session_shutdown` only fires handlers the - // CHILD registered — and a child that no-op'd via the latch - // registered none, so it cannot clear the parent's latch. Clearing - // here lets a `/reload` legitimately re-initialize the full runtime, - // while ephemeral in-process children never touch it. - clearPiMagicContextActive(); }); // Pi has no `session_deleted` event, but `session_before_switch` @@ -2375,6 +2441,13 @@ async function startPiMagicContextRuntime( }); } +function registerPiSubagentInitContextCleanup( + pi: ExtensionAPI, + unsubscribe: () => void, +): void { + pi.on("session_shutdown", unsubscribe); +} + /** * Format `execute_threshold_percentage` for the boot log. The config accepts * either a bare number or a per-model map (`{ default: 65, "provider/model": 50 }`); diff --git a/packages/pi-plugin/src/session-cleanup-wiring.test.ts b/packages/pi-plugin/src/session-cleanup-wiring.test.ts index 099699693..b07516a45 100644 --- a/packages/pi-plugin/src/session-cleanup-wiring.test.ts +++ b/packages/pi-plugin/src/session-cleanup-wiring.test.ts @@ -85,7 +85,7 @@ describe("session_before_switch handler wiring", () => { describe("session_shutdown handler also drains per-session maps", () => { const handler = INDEX_SRC.match( - /pi\.on\("session_shutdown"[\s\S]*?\n\s*\}\);/, + /pi\.on\("session_shutdown"[\s\S]*?\n\t\}\);(?=\n\n\t\/\/ Pi has no `session_deleted` event)/, ); test("session_shutdown handler exists", () => { diff --git a/packages/pi-plugin/src/subagent-runner.test.ts b/packages/pi-plugin/src/subagent-runner.test.ts index 624185703..f12df46f1 100644 --- a/packages/pi-plugin/src/subagent-runner.test.ts +++ b/packages/pi-plugin/src/subagent-runner.test.ts @@ -202,6 +202,7 @@ function nextTick() { return new Promise((resolve) => setTimeout(resolve, 0)); } +const originalTestDataDir = process.env.MAGIC_CONTEXT_TEST_DATA_DIR; const originalXdgDataHome = process.env.XDG_DATA_HOME; describe("subagent-runner pure helpers", () => { @@ -234,6 +235,15 @@ describe("subagent-runner pure helpers", () => { ).toEqual({ text: null, stopReason: null, errorMessage: null }); }); + it("distinguishes Pi from embedded Node hosts", () => { + expect(__test.isGenericRuntimeExecutable("/usr/bin/node24")).toBe(true); + expect(__test.isPiCliScript("/app/node_modules/.bin/next")).toBe(false); + expect( + __test.isPiCliScript( + "/app/node_modules/@earendil-works/pi-coding-agent/dist/cli.js", + ), + ).toBe(true); + }); it("builds argv with system prompt, primary model, and prompt last", () => { expect( buildArgsForTest({ @@ -1057,10 +1067,8 @@ describe("PiSubagentRunner spawn lifecycle", () => { }); }); - it("with no piBinary override, spawns the host runtime + cli.js (Windows-safe, #177)", async () => { - // Default resolution must NOT spawn a bare "pi" (which ENOENTs on Windows - // because npm installs a pi.cmd shim, not a literal pi). It re-invokes the - // exact host CLI: process.execPath + process.argv[1], with no shell. + it("with no piBinary override, does not re-run an embedded host", async () => { + // The Bun test file stands in for pi-web's Next.js argv[1]. const child = createMockChild(); const spawnImpl = mock(() => child as never); const { PiSubagentRunner } = await import("./subagent-runner"); @@ -1084,17 +1092,12 @@ describe("PiSubagentRunner spawn lifecycle", () => { const [command, spawnArgs, opts] = ( spawnImpl.mock.calls as unknown[][] )[0] as [string, string[], { shell?: boolean }]; - // In this test runner argv[1] is a real on-disk script (bun/node test - // file), so the host-CLI branch fires: command is the runtime, the first - // arg is the running script, and the child is spawned without a shell. - expect(command).toBe(process.execPath); - expect(spawnArgs[0]).toBe(process.argv[1]); + expect(spawnArgs[0]).not.toBe(process.argv[1]); expect(spawnArgs).toContain("--no-session"); + expect(command.length).toBeGreaterThan(0); // Never spawned through a shell (no cmd.exe in the path = no arg-escaping // or injection on the untrusted prompt/task text). expect(opts.shell).toBeFalsy(); - // Crucially, never a bare "pi". - expect(command).not.toBe("pi"); }); it("returns model_failed promptly for live terminal error stopReason", async () => { @@ -2479,6 +2482,7 @@ describe("Pi subagent schema-fence probe", () => { it("does not spawn a Pi child when the shared database is newer than this build", async () => { const dataHome = mkdtempSync(join(tmpdir(), "mc-pi-fence-probe-")); try { + process.env.MAGIC_CONTEXT_TEST_DATA_DIR = dataHome; process.env.XDG_DATA_HOME = dataHome; closeDatabase(); __resetSchemaFenceStateForTests(); @@ -2505,6 +2509,9 @@ describe("Pi subagent schema-fence probe", () => { } finally { closeDatabase(); __resetSchemaFenceStateForTests(); + if (originalTestDataDir === undefined) + delete process.env.MAGIC_CONTEXT_TEST_DATA_DIR; + else process.env.MAGIC_CONTEXT_TEST_DATA_DIR = originalTestDataDir; if (originalXdgDataHome === undefined) delete process.env.XDG_DATA_HOME; else process.env.XDG_DATA_HOME = originalXdgDataHome; rmSync(dataHome, { recursive: true, force: true }); diff --git a/packages/pi-plugin/src/subagent-runner.ts b/packages/pi-plugin/src/subagent-runner.ts index ae4f68540..753af5c27 100644 --- a/packages/pi-plugin/src/subagent-runner.ts +++ b/packages/pi-plugin/src/subagent-runner.ts @@ -88,17 +88,20 @@ interface PiInvocation { * a literal `pi`), and Node's `spawn("pi")` without a shell looks for a file * named exactly `pi`, so it ENOENTs; and Windows ignores the `#!/usr/bin/env * node` shebang entirely, so spawning `dist/cli.js` "directly" only works on - * POSIX. The reliable, cross-platform approach is to re-invoke the EXACT host - * CLI the user is already running: `process.execPath` (the node/bun binary) plus - * `process.argv[1]` (the absolute path to the running `cli.js`). That sidesteps - * shim resolution completely and pins the child to the same Pi version/runtime. + * POSIX. When the host itself is Pi, the reliable, cross-platform approach is + * to re-invoke the EXACT host CLI the user is already running: + * `process.execPath` (the node/bun binary) plus `process.argv[1]` (the absolute + * path to the running `cli.js`). That sidesteps shim resolution completely and + * pins the child to the same Pi version/runtime. Embedded hosts such as pi-web + * must not reuse their unrelated `argv[1]`. * - * Mirrors Pi's own `getPiInvocation` reference. MUST be evaluated in the host Pi - * process (extensions load in-process, so `argv[1]` is the host `cli.js`). + * Mirrors Pi's own `getPiInvocation` reference. MUST be evaluated in the host + * process: a Pi host has its `cli.js` in `argv[1]`; embedded hosts fall through + * to bundled-Pi or PATH resolution. * * Resolution order: - * 1. argv[1] is a real on-disk script (not a bun-compiled `/$bunfs/root/` - * virtual path) -> `execPath cli.js ...` (node + absolute cli.js). + * 1. argv[1] belongs to an on-disk Pi package (and is not a bun-compiled + * `/$bunfs/root/` virtual path) -> `execPath cli.js ...`. * 2. execPath is a packaged binary (basename not node/bun) -> `execPath ...` * (the compiled binary IS pi; no script arg). * 3. A bundled `@earendil-works/pi-coding-agent/dist/cli.js` resolves -> @@ -106,25 +109,41 @@ interface PiInvocation { * 4. Last resort: bare `pi` on PATH. * * Everything is spawned WITHOUT a shell. The primary path (execPath + argv[1]) - * covers every real runtime because the extension loads in-process, so argv[1] - * is the host cli.js; the bare-`pi` step is a near-unreachable backstop. We do + * covers every real Pi CLI runtime; embedded hosts fall through rather than + * accidentally re-running themselves. We do * NOT fall back to a shell for it (which on Windows would resolve the .cmd shim * but pass the prompt/task text through cmd.exe, exposing arg-escaping and * injection), and we don't pull in cross-spawn just for a dead path. */ +function isPiCliScript(scriptPath: string): boolean { + const normalized = scriptPath.replaceAll("\\", "/"); + return /\/@(?:earendil-works|oh-my-pi)\/pi-coding-agent\/dist\/cli\.js$/.test( + normalized, + ); +} + +function isGenericRuntimeExecutable(execPath: string): boolean { + return /^(node(?:js)?\d*|bun)(\.exe)?$/.test( + basename(execPath).toLowerCase(), + ); +} + function resolvePiInvocation(): PiInvocation { const execPath = process.execPath; const currentScript = process.argv[1]; const isBunVirtualScript = currentScript?.startsWith("/$bunfs/root/") ?? false; - if (currentScript && !isBunVirtualScript && existsSync(currentScript)) { + if ( + currentScript && + !isBunVirtualScript && + existsSync(currentScript) && + isPiCliScript(currentScript) + ) { return { command: execPath, prefixArgs: [currentScript] }; } - const execName = basename(execPath).toLowerCase(); - const isGenericRuntime = /^(node|bun)(\.exe)?$/.test(execName); - if (!isGenericRuntime) { + if (!isGenericRuntimeExecutable(execPath)) { // A packaged single-file binary: execPath itself is pi. return { command: execPath, prefixArgs: [] }; } @@ -1920,6 +1939,8 @@ function terminateChild(child: ReturnType) { export const __test = { buildArgs, extractFinalAssistant, + isGenericRuntimeExecutable, + isPiCliScript, parsePiEventLine, terminateChild, DREAMER_ACTION_AGENTS, diff --git a/packages/plugin/src/features/magic-context/session-project-backfill.test.ts b/packages/plugin/src/features/magic-context/session-project-backfill.test.ts index 141f82bc5..691c023fe 100644 --- a/packages/plugin/src/features/magic-context/session-project-backfill.test.ts +++ b/packages/plugin/src/features/magic-context/session-project-backfill.test.ts @@ -150,7 +150,7 @@ describe("runSessionProjectBackfill", () => { const db = createDb(); const directory = makeTempDir("session-project-backfill-live-"); let resolverCalls = 0; - + let sourceCalls = 0; const first = await runSessionProjectBackfill(db, [{ sessionId: "ses-first", directory }], { resolveIdentity: () => { resolverCalls += 1; @@ -160,7 +160,10 @@ describe("runSessionProjectBackfill", () => { }); const second = await runSessionProjectBackfill( db, - [{ sessionId: "ses-second", directory }], + async () => { + sourceCalls += 1; + return [{ sessionId: "ses-second", directory }]; + }, { resolveIdentity: () => { resolverCalls += 1; @@ -175,6 +178,7 @@ describe("runSessionProjectBackfill", () => { expect(second.status).toBe("already_completed"); expect(second.backfilledSessions).toBe(0); expect(resolverCalls).toBe(1); + expect(sourceCalls).toBe(0); expect(getStoredProjectPath(db, "ses-first")).toBe("git:first"); expect(getStoredProjectPath(db, "ses-second")).toBeNull(); expect(_getSessionProjectBackfillState(db)?.status).toBe("completed"); From fee1fdf0024d573d61eb0792eea1d9bfe39d3f87 Mon Sep 17 00:00:00 2001 From: elrond <171308700+elrond298@users.noreply.github.com> Date: Fri, 21 Aug 2026 12:06:21 +0800 Subject: [PATCH 2/5] fix(pi-plugin): stabilize Dreamer owner handoff --- packages/pi-plugin/src/dreamer/index.test.ts | 65 ++++++++++++++++++++ packages/pi-plugin/src/dreamer/index.ts | 19 ++++-- 2 files changed, 80 insertions(+), 4 deletions(-) diff --git a/packages/pi-plugin/src/dreamer/index.test.ts b/packages/pi-plugin/src/dreamer/index.test.ts index 2833b7fa6..d73f2d75b 100644 --- a/packages/pi-plugin/src/dreamer/index.test.ts +++ b/packages/pi-plugin/src/dreamer/index.test.ts @@ -386,6 +386,71 @@ describe("Pi dreamer wiring", () => { ]); }); + test("active-owner handoff starts one timer when remaining worktree dirs repeat", async () => { + db = createDb(); + const dirs: string[] = []; + __test.setStartDreamScheduleTimerFactory(async (registration) => { + dirs.push((registration as { directory: string }).directory); + return mock(() => {}); + }); + + const firstA = dreamerOptions({ + database: db, + projectDir: "/tmp/worktree-A", + projectIdentity: "git:pi-handoff", + }); + const ownerB = dreamerOptions({ + database: db, + projectDir: "/tmp/worktree-B", + projectIdentity: "git:pi-handoff", + }); + const secondA = dreamerOptions({ + database: db, + projectDir: "/tmp/worktree-A", + projectIdentity: "git:pi-handoff", + }); + const activeC = dreamerOptions({ + database: db, + projectDir: "/tmp/worktree-C", + projectIdentity: "git:pi-handoff", + }); + for (const owner of [firstA, ownerB, secondA, activeC]) { + registerPiDreamerProject(owner); + await flushMicrotasks(); + } + + unregisterPiDreamerProject({ + projectIdentity: "git:pi-handoff", + registrationOwner: activeC.registrationOwner, + }); + await flushMicrotasks(); + expect(dirs).toEqual([ + "/tmp/worktree-A", + "/tmp/worktree-B", + "/tmp/worktree-A", + "/tmp/worktree-C", + "/tmp/worktree-A", + ]); + + unregisterPiDreamerProject({ + projectIdentity: "git:pi-handoff", + registrationOwner: secondA.registrationOwner, + }); + await flushMicrotasks(); + unregisterPiDreamerProject({ + projectIdentity: "git:pi-handoff", + registrationOwner: ownerB.registrationOwner, + }); + await flushMicrotasks(); + expect(dirs.slice(-2)).toEqual(["/tmp/worktree-B", "/tmp/worktree-A"]); + + unregisterPiDreamerProject({ + projectIdentity: "git:pi-handoff", + registrationOwner: firstA.registrationOwner, + }); + expect(__test.registeredProjectCount()).toBe(0); + }); + test("unregister removes the project", () => { db = createDb(); const opts = dreamerOptions({ diff --git a/packages/pi-plugin/src/dreamer/index.ts b/packages/pi-plugin/src/dreamer/index.ts index ced669dc0..1a5160400 100644 --- a/packages/pi-plugin/src/dreamer/index.ts +++ b/packages/pi-plugin/src/dreamer/index.ts @@ -267,13 +267,24 @@ export function unregisterPiDreamerProject(opts: { if (registration.owners.size > 0) { if (registration.activeOwner !== opts.registrationOwner) return; // The active worktree owner left while sibling sessions still use this - // project. Rebuild from the remaining owners so the shared timer follows - // a live session instead of retaining the departed session's directory. + // project. Rebuild once from the most recently registered remaining owner + // so the shared timer follows a live session, then retain every sibling + // owner without repeatedly replacing the timer. const remaining = [...registration.owners.values()]; + const replacementOptions = remaining[remaining.length - 1]; + if (!replacementOptions) return; registration.cleanup(); registeredProjects.delete(opts.projectIdentity); - for (const remainingOptions of remaining) - registerPiDreamerProject(remainingOptions); + registerPiDreamerProject(replacementOptions); + const replacement = registeredProjects.get(opts.projectIdentity); + if (replacement) { + for (const remainingOptions of remaining) { + replacement.owners.set( + remainingOptions.registrationOwner, + remainingOptions, + ); + } + } return; } From 93e3ed5776acade55078d46ac2d8b4dd27d2b72c Mon Sep 17 00:00:00 2001 From: elrond <171308700+elrond298@users.noreply.github.com> Date: Fri, 21 Aug 2026 18:52:35 +0800 Subject: [PATCH 3/5] fix(pi-plugin): harden multi-session lifecycle in embedded hosts Handle multiple Pi sessions running in the same embedded host. - share Dreamer registration across plugin instances and hand its timer to the most recently registered remaining session - wait only for each session's historian, recomp, and Dreamer jobs during session_shutdown - keep agent_end synchronous so it does not delay turn completion - use the current command context to display status in Pi RPC mode - keep standalone Pi and embedded-host subagent launch regression tests - document the behavior in PARITY.md Tests: - NODE_ENV=test bun test src (813 passed) - NODE_ENV=test bun test session-project-backfill.test.ts (9 passed) - bun run build - bun run typecheck - bun run format:check --- packages/pi-plugin/PARITY.md | 41 +++-- .../pi-plugin/src/agent-end-handler.test.ts | 17 +- .../src/commands/ctx-commands.test.ts | 36 ++++- packages/pi-plugin/src/commands/ctx-dream.ts | 20 ++- packages/pi-plugin/src/commands/ctx-embed.ts | 17 +- packages/pi-plugin/src/commands/ctx-flush.ts | 10 +- packages/pi-plugin/src/commands/ctx-recomp.ts | 25 +-- .../src/commands/ctx-session-upgrade.ts | 27 ++-- packages/pi-plugin/src/commands/ctx-status.ts | 10 +- packages/pi-plugin/src/commands/ctx-wrapup.ts | 22 +-- .../src/commands/pi-command-utils.test.ts | 4 - .../src/commands/pi-command-utils.ts | 45 ++++-- .../pi-plugin/src/context-handler.test.ts | 38 +++++ packages/pi-plugin/src/context-handler.ts | 22 ++- packages/pi-plugin/src/dreamer/index.test.ts | 147 ++++++++++++++++++ packages/pi-plugin/src/dreamer/index.ts | 110 +++++++++---- .../src/index-in-process-latch.test.ts | 38 +++++ packages/pi-plugin/src/index.ts | 61 ++++---- .../pi-plugin/src/pi-recomp-runner.test.ts | 37 +++++ packages/pi-plugin/src/pi-recomp-runner.ts | 20 ++- .../pi-plugin/src/subagent-runner.test.ts | 52 +++++++ 21 files changed, 619 insertions(+), 180 deletions(-) create mode 100644 packages/pi-plugin/src/pi-recomp-runner.test.ts diff --git a/packages/pi-plugin/PARITY.md b/packages/pi-plugin/PARITY.md index 25a14ae10..40990bddb 100644 --- a/packages/pi-plugin/PARITY.md +++ b/packages/pi-plugin/PARITY.md @@ -23,14 +23,20 @@ plugin process and reach `experimental.chat.messages.transform`. OpenCode gates historian / m[0]m[1] injection / nudges / auto-search behind `fullFeatureMode` (i.e. `!isSubagent`), and detects subagents via OpenCode's `session.parent_id`. -**Pi:** Pi has **no native subagent concept**. The *only* subagents that exist -are the ones Magic Context itself spawns (historian, dreamer, sidekick), and each -runs as a **separate `pi --print` process** loading only the lean -`subagent-entry.js`, whose recursion guard **never wires `pi.on("context")`** -(see `subagent-entry.ts` header). A Pi subagent therefore *cannot* reach the -context-handler pipeline at all. - -**Consequence:** `is_subagent` is **never written `true`** for any Pi session. +**Pi:** Pi has **no native subagent concept**. The subagents Magic Context itself +spawns (historian, dreamer, sidekick) each run as a **separate `pi --print` process** +loading only the lean `subagent-entry.js`, whose recursion guard **never wires +`pi.on("context")`** (see `subagent-entry.ts` header). A Magic Context subagent +therefore *cannot* reach the context-handler pipeline at all. +`@gotgenes/pi-subagents`, however, can initialize a child session inside the same +process. The full extension uses Pi's public child-session lifecycle events plus +process-shared `AsyncLocalStorage` to suppress only the child while allowing +unrelated same-process sessions to initialize normally. + +**Consequence:** `is_subagent` is **never written `true`** for any Pi session +that reaches the context-handler pipeline. Separate child processes load the lean +entry, while in-process child initialization is suppressed before the normal +context pipeline is registered. There is nothing to gate, so Pi does NOT need OpenCode's `fullFeatureMode` reduced-mode enforcement in `context-handler.ts`. The vestigial `!isSubagent` checks that exist in the Pi context handler are harmless (always take the @@ -115,7 +121,7 @@ the source array for dirty indices only. --- -## 6. Transient UI: Pi uses `ctx.ui.notify` toasts, not persistent dialogs +## 6. Transient UI: Pi uses `ctx.ui.notify` toasts and RPC dialogs **OpenCode:** TUI dialogs (upgrade prompt, `/ctx-status`, `/ctx-recomp`, `/ctx-embed`, `/ctx-flush`) via RPC, with an ignored-message fallback for Desktop/Web. Notification drain is @@ -123,7 +129,12 @@ with an ignored-message fallback for Desktop/Web. Notification drain is another) because one process can serve multiple sessions and TUI port discovery is newest-pid-wins. -**Pi:** transient terminal notifications. The upgrade reminder passes +**Pi:** command status is appended as a model-invisible custom entry. Interactive +terminals render that entry through the registered entry renderer. In Pi RPC +mode, each command uses its live `ctx`: `ctx.ui.notify` presents short progress +as toasts and `ctx.ui.custom` presents detailed results as dialogs. A context +captured by `session_start` cannot be reused because pi-web can host multiple +sessions in one process. The upgrade reminder passes `deliveryPersists=false` on Pi, so a missed toast does not honor the old explicit- dismissal stamp. Both harnesses persist the 24-hour reminder cooldown and three- delivery cap, preventing repeated startup toasts while `/ctx-status` still reports @@ -155,6 +166,13 @@ shared resolver's log-only dubious-ownership warning while still using the same **stdin** (Pi concatenates stdin + positional) to avoid Linux `MAX_ARG_STRLEN` / E2BIG; the positional is omitted when piping. - `--no-session` keeps subagent JSONL out of the user's session picker. +- In pi-web, multiple sessions can share one process. Startup maintenance runs + once per process, while each session wires its own hooks. Dreamer registration + is process-shared and tracks sibling ownership, so one session's shutdown cannot + deregister another session's project timer. +- `session_shutdown` drains only that session's in-flight historian and recomp work + and only the shutting-down extension instance's Dreamer work. Child-session + lifecycle listeners are detached only for that extension instance. --- @@ -378,7 +396,8 @@ mechanism differs because the process models differ: inline `await` froze all input. Pi instead spawns the recomp via `spawnPiRecompRun` (mirroring `spawnPiHistorianRun`): the handler returns immediately after the ack message, the run is tracked in an in-flight map for - `session_shutdown` drain, and progress surfaces through `[ctx-status]` + `session_shutdown` drain (keyed by session id so one session does not drain + another), and progress surfaces through `[ctx-status]` messages + the `recomp` status-line flag. Because Pi's recomp runs in the background (not inside the user's turn), its diff --git a/packages/pi-plugin/src/agent-end-handler.test.ts b/packages/pi-plugin/src/agent-end-handler.test.ts index ef60b6083..16af9aefe 100644 --- a/packages/pi-plugin/src/agent-end-handler.test.ts +++ b/packages/pi-plugin/src/agent-end-handler.test.ts @@ -107,15 +107,22 @@ describe("session_shutdown handler (drain location)", () => { const body = extractSessionShutdownHandlerBody(INDEX_SRC); test("drains in-flight historians through withTimeout", () => { - expect(body).toContain("awaitInFlightHistorians"); - expect(body).toContain( - "withTimeout(awaitInFlightHistorians(), SHUTDOWN_DRAIN_MS)", + expect(body).toMatch( + /withTimeout\(\s*awaitInFlightHistorians\(sessionId\),\s*SHUTDOWN_DRAIN_MS,?\s*\)/, ); expect(body).not.toContain("Promise.race"); }); - test("drains in-flight dreamers (Promise.race with timeout)", () => { - expect(body).toContain("awaitInFlightDreamers"); + test("drains the shutting-down session's recomp through withTimeout", () => { + expect(body).toMatch( + /withTimeout\(\s*awaitInFlightRecomps\(sessionId\),\s*SHUTDOWN_DRAIN_MS,?\s*\)/, + ); + }); + + test("drains the current extension owner's dreamers through withTimeout", () => { + expect(body).toMatch( + /withTimeout\(\s*awaitInFlightDreamers\(dreamerRegistrationOwner\),\s*SHUTDOWN_DRAIN_MS,?\s*\)/, + ); }); test("drain timeout uses unref/clear helper", () => { diff --git a/packages/pi-plugin/src/commands/ctx-commands.test.ts b/packages/pi-plugin/src/commands/ctx-commands.test.ts index 9d7fb5d64..93330c8e4 100644 --- a/packages/pi-plugin/src/commands/ctx-commands.test.ts +++ b/packages/pi-plugin/src/commands/ctx-commands.test.ts @@ -29,8 +29,10 @@ interface AppendedEntry { interface MockCommandContext { cwd: string; hasUI?: boolean; + mode?: "rpc"; ui: { custom: (factory: unknown, options?: unknown) => Promise; + notify?: (text: string, type?: string) => void; setStatus?: (key: string, text: string) => void; }; model?: { @@ -152,6 +154,34 @@ describe("Pi Magic Context commands", () => { expect(sent[0]?.data.text).toContain("## Magic Status"); }); + it("presents /ctx-status through the live RPC command context", async () => { + const db = createDb(); + const { pi, handlers } = createMockPi(); + const shownA: unknown[] = []; + const shownB: unknown[] = []; + const rpcCtx = (sessionId: string, shown: unknown[]) => ({ + ...createCtx(sessionId), + mode: "rpc" as const, + ui: { + async custom(factory: unknown) { + shown.push(factory); + return undefined; + }, + notify() {}, + }, + }); + registerCtxStatusCommand(pi as never, { + db, + projectIdentity: "/tmp/project", + }); + + await handlers.get("ctx-status")?.("", rpcCtx("ses-a", shownA)); + await handlers.get("ctx-status")?.("", rpcCtx("ses-b", shownB)); + + expect(shownA).toHaveLength(1); + expect(shownB).toHaveLength(1); + }); + it("/ctx-status keeps the persisted usable limit when command context omits maxTokens", async () => { const db = createDb(); const sessionId = "ses-status-persisted-reserve"; @@ -269,8 +299,10 @@ describe("Pi Magic Context commands", () => { registrationCwds.push(ctx.cwd); }, }); - // This unit test isolates the command from the real dreamer registry. The - // injected registration sync runs immediately before the manual dream call. + // Not registered with the dreamer timer in this unit test, so runManual + // throws "not registered" → the handler reports the failure. We only + // assert the command is wired and emits a /ctx-dream status message. + // The injected registration sync runs immediately before runManual. await handlers.get("ctx-dream")?.("", createCtx()); expect(registrationCwds).toEqual(["/tmp/project"]); diff --git a/packages/pi-plugin/src/commands/ctx-dream.ts b/packages/pi-plugin/src/commands/ctx-dream.ts index ac162b4b9..67f886425 100644 --- a/packages/pi-plugin/src/commands/ctx-dream.ts +++ b/packages/pi-plugin/src/commands/ctx-dream.ts @@ -9,7 +9,7 @@ import { import type { ContextDatabase } from "@magic-context/core/features/magic-context/storage"; import { sessionLog } from "@magic-context/core/shared/logger"; import { runPiDreamForProject } from "../dreamer"; -import { sendCtxStatusMessage } from "./pi-command-utils"; +import { createCtxStatusSender } from "./pi-command-utils"; export function registerCtxDreamCommand( pi: ExtensionAPI, @@ -25,11 +25,13 @@ export function registerCtxDreamCommand( resolveDreamerEnabled?: (ctx: { cwd: string }) => boolean | undefined; onProjectSeen?: (projectIdentity: string) => void; ensureRegistered?: (ctx: { cwd: string }) => void | Promise; + registrationOwner?: object; }, ): void { pi.registerCommand("ctx-dream", { description: "Run Magic Context dreamer tasks for this project now", handler: async (args, ctx) => { + const sendStatus = createCtxStatusSender(pi, ctx); const project = deps.resolveProject?.(ctx) ?? { projectDir: deps.projectDir, projectIdentity: deps.projectIdentity, @@ -44,8 +46,7 @@ export function registerCtxDreamCommand( let task: DreamTaskName | undefined; if (requested) { if (!isCanonicalDreamTask(requested)) { - sendCtxStatusMessage( - pi, + sendStatus( { title: "/ctx-dream", text: `## /ctx-dream\n\nUnknown task "${requested}".`, @@ -61,8 +62,7 @@ export function registerCtxDreamCommand( task = requested; } if (dreamerEnabled === false) { - sendCtxStatusMessage( - pi, + sendStatus( { title: "/ctx-dream", text: "## /ctx-dream\n\nDreamer is disabled for this project (`dreamer.disable=true`).", @@ -84,8 +84,7 @@ export function registerCtxDreamCommand( // Tell the user we're starting a real run, including the read-only count // captured before the task acquires its lease. - sendCtxStatusMessage( - pi, + sendStatus( { title: "/ctx-dream", text: [ @@ -113,6 +112,7 @@ export function registerCtxDreamCommand( const result = await runPiDreamForProject( project.projectIdentity, task, + deps.registrationOwner, ); const lines: string[] = []; if (result.ran.length > 0) lines.push(`Ran: ${result.ran.join(", ")}`); @@ -142,8 +142,7 @@ export function registerCtxDreamCommand( } if (lines.length === 0) lines.push("No enabled dream tasks to run."); - sendCtxStatusMessage( - pi, + sendStatus( { title: "/ctx-dream", text: ["## /ctx-dream", "", ...lines].join("\n"), @@ -158,8 +157,7 @@ export function registerCtxDreamCommand( } catch (error) { const message = error instanceof Error ? error.message : String(error); sessionLog(project.projectIdentity, `/ctx-dream failed: ${message}`); - sendCtxStatusMessage( - pi, + sendStatus( { title: "/ctx-dream", text: [ diff --git a/packages/pi-plugin/src/commands/ctx-embed.ts b/packages/pi-plugin/src/commands/ctx-embed.ts index 78bff2ad3..ef24af5f3 100644 --- a/packages/pi-plugin/src/commands/ctx-embed.ts +++ b/packages/pi-plugin/src/commands/ctx-embed.ts @@ -12,7 +12,7 @@ import { } from "@magic-context/core/hooks/magic-context/embed-session-state"; import { formatEmbedStatusText } from "@magic-context/core/hooks/magic-context/format-embed-status"; import { ensureProjectRegisteredFromPiDirectory } from "../embedding-bootstrap"; -import { resolveSessionId, sendCtxStatusMessage } from "./pi-command-utils"; +import { createCtxStatusSender, resolveSessionId } from "./pi-command-utils"; const EMBED_PROGRESS_COMPARTMENT_STEP = 8; const EMBED_PROGRESS_MIN_INTERVAL_MS = 10_000; @@ -159,9 +159,10 @@ export function registerCtxEmbedCommand( description: "Embedding status, or start/pause history compartment embedding (start | pause)", handler: async (args, ctx) => { + const sendStatus = createCtxStatusSender(pi, ctx); const sessionId = resolveSessionId(ctx); if (!sessionId) { - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-embed", text: "## /ctx-embed\n\nNo active Pi session is available.", level: "error", @@ -185,7 +186,7 @@ export function registerCtxEmbedCommand( project.projectIdentity, sessionId, ); - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-embed", text: `## /ctx-embed\n\nPaused at ${cov.session.embedded}/${cov.session.total} compartments embedded.`, level: "info", @@ -194,7 +195,7 @@ export function registerCtxEmbedCommand( } if (memoryEnabled === false) { - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-embed", text: "## /ctx-embed\n\nMemory is disabled for this project, so there is no semantic embedding to backfill.", level: "info", @@ -211,18 +212,18 @@ export function registerCtxEmbedCommand( sessionId, { onStatus: (status) => - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-embed", ...status, }), }, ); - sendCtxStatusMessage(pi, { title: "/ctx-embed", text, level }); + sendStatus({ title: "/ctx-embed", text, level }); return; } if (sub !== "") { - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-embed", text: "## /ctx-embed\n\nUsage: `/ctx-embed` (status), `/ctx-embed start`, or `/ctx-embed pause`.", level: "info", @@ -236,7 +237,7 @@ export function registerCtxEmbedCommand( sessionId, ); const statusText = formatEmbedStatusText(coverage, { status: "idle" }); - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-embed", text: `## Embedding Status\n\n${statusText}`, level: "info", diff --git a/packages/pi-plugin/src/commands/ctx-flush.ts b/packages/pi-plugin/src/commands/ctx-flush.ts index d5cbb18ee..f9f6b5b6a 100644 --- a/packages/pi-plugin/src/commands/ctx-flush.ts +++ b/packages/pi-plugin/src/commands/ctx-flush.ts @@ -8,7 +8,7 @@ import { signalPiPendingMaterialization, signalPiSystemPromptRefresh, } from "../context-handler"; -import { resolveSessionId, sendCtxStatusMessage } from "./pi-command-utils"; +import { createCtxStatusSender, resolveSessionId } from "./pi-command-utils"; export function registerCtxFlushCommand( pi: ExtensionAPI, @@ -18,9 +18,10 @@ export function registerCtxFlushCommand( description: "Force pending Magic Context drops to materialize on the next provider call", handler: async (_args, ctx) => { + const sendStatus = createCtxStatusSender(pi, ctx); const sessionId = resolveSessionId(ctx); if (!sessionId) { - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-flush", text: "## /ctx-flush\n\nNo active Pi session is available.", level: "error", @@ -28,7 +29,7 @@ export function registerCtxFlushCommand( return; } if (deps.compactionOff) { - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-flush", text: COMPACTION_OFF_COMMAND_UNAVAILABLE, level: "warning", @@ -66,8 +67,7 @@ export function registerCtxFlushCommand( pendingBefore > 0 ? `## /ctx-flush\n\nFlushed ${pendingBefore} pending ops; next provider call will materialize.\n\n${result}` : `## /ctx-flush\n\n${result}`; - sendCtxStatusMessage( - pi, + sendStatus( { title: "/ctx-flush", text, diff --git a/packages/pi-plugin/src/commands/ctx-recomp.ts b/packages/pi-plugin/src/commands/ctx-recomp.ts index d467df3b7..f67a09beb 100644 --- a/packages/pi-plugin/src/commands/ctx-recomp.ts +++ b/packages/pi-plugin/src/commands/ctx-recomp.ts @@ -28,7 +28,7 @@ import { stagePiRecompMarker } from "../pi-recomp-marker"; import { isPiRecompInFlight, spawnPiRecompRun } from "../pi-recomp-runner"; import { readPiSessionMessages } from "../read-session-pi"; import { updateStatusLine } from "../status-line"; -import { resolveSessionId, sendCtxStatusMessage } from "./pi-command-utils"; +import { createCtxStatusSender, resolveSessionId } from "./pi-command-utils"; interface RecompConfirmation { timestamp: number; @@ -70,9 +70,10 @@ export function registerCtxRecompCommand( description: "Rebuild Magic Context compartments from raw Pi session history", handler: async (args, ctx) => { + const sendStatus = createCtxStatusSender(pi, ctx); const sessionId = resolveSessionId(ctx); if (!sessionId) { - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-recomp", text: "## Magic Recomp\n\nNo active Pi session is available.", level: "error", @@ -81,7 +82,7 @@ export function registerCtxRecompCommand( } const currentDeps = deps.resolveRuntimeDeps?.(ctx) ?? deps; if (currentDeps.compactionOff) { - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-recomp", text: COMPACTION_OFF_COMMAND_UNAVAILABLE, level: "warning", @@ -91,7 +92,7 @@ export function registerCtxRecompCommand( const parsed = parseRecompArgs(args); if (parsed.kind === "error") { - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-recomp", text: `## Magic Recomp — Invalid Arguments\n\n${parsed.message}`, level: "error", @@ -100,7 +101,7 @@ export function registerCtxRecompCommand( } if (parsed.kind === "upgrade") { - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-recomp", text: executeRecompUpgradeStub(currentDeps.db, sessionId), level: "info", @@ -110,7 +111,7 @@ export function registerCtxRecompCommand( } if (!currentDeps.historianModel) { - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-recomp", text: "## Magic Recomp\n\n/ctx-recomp is unavailable because `historian.model` is not configured.", level: "error", @@ -137,7 +138,7 @@ export function registerCtxRecompCommand( ); if (!warning.confirmable) confirmationBySession.delete(sessionId); else confirmationBySession.set(sessionId, { timestamp: now, argsKey }); - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-recomp", text: warning.text, level: warning.confirmable ? "warning" : "error", @@ -146,7 +147,7 @@ export function registerCtxRecompCommand( } if (isWrapupInProgress(currentDeps.db, sessionId)) { - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-recomp", text: "## Magic Recomp\n\n/ctx-wrapup is already compacting this session. Wait for it to finish, then try `/ctx-recomp` again.", level: "warning", @@ -155,7 +156,7 @@ export function registerCtxRecompCommand( } if (isPiRecompInFlight(sessionId)) { - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-recomp", text: "## Magic Recomp\n\nA recomp or upgrade is already running for this session in the background. Wait for it to finish, then try again.", level: "warning", @@ -164,7 +165,7 @@ export function registerCtxRecompCommand( } confirmationBySession.delete(sessionId); - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-recomp", text: parsed.kind === "partial" @@ -207,7 +208,7 @@ export function registerCtxRecompCommand( directory: ctx.cwd, accountingSessionId: sessionId, notify: (text) => { - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-recomp", text, level: inferLevel(text), @@ -269,7 +270,7 @@ export function registerCtxRecompCommand( signalPiDeferredHistoryRefresh(sessionId); signalPiDeferredMaterialization(sessionId); } - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-recomp", text: result.message, level: inferLevel(result.message), diff --git a/packages/pi-plugin/src/commands/ctx-session-upgrade.ts b/packages/pi-plugin/src/commands/ctx-session-upgrade.ts index 7621a0f0a..bcc40e21a 100644 --- a/packages/pi-plugin/src/commands/ctx-session-upgrade.ts +++ b/packages/pi-plugin/src/commands/ctx-session-upgrade.ts @@ -30,7 +30,7 @@ import { stagePiRecompMarker } from "../pi-recomp-marker"; import { isPiRecompInFlight, spawnPiRecompRun } from "../pi-recomp-runner"; import { readPiSessionMessages } from "../read-session-pi"; import { updateStatusLine } from "../status-line"; -import { resolveSessionId, sendCtxStatusMessage } from "./pi-command-utils"; +import { createCtxStatusSender, resolveSessionId } from "./pi-command-utils"; export interface CtxSessionUpgradeRuntimeDeps { db: ContextDatabase; @@ -74,9 +74,10 @@ export function registerCtxSessionUpgradeCommand( description: "Upgrade this session to the current Magic Context history format and re-organize project memories", handler: async (_args, ctx) => { + const sendStatus = createCtxStatusSender(pi, ctx); const sessionId = resolveSessionId(ctx); if (!sessionId) { - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-session-upgrade", text: "## Session Upgrade\n\nNo active Pi session is available.", level: "error", @@ -85,7 +86,7 @@ export function registerCtxSessionUpgradeCommand( } const currentDeps = deps.resolveRuntimeDeps?.(ctx) ?? deps; if (currentDeps.compactionOff) { - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-session-upgrade", text: COMPACTION_OFF_COMMAND_UNAVAILABLE, level: "warning", @@ -93,7 +94,7 @@ export function registerCtxSessionUpgradeCommand( return; } if (!currentDeps.historianModel) { - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-session-upgrade", text: "## Session Upgrade\n\nUnavailable because `historian.model` is not configured.", level: "error", @@ -102,7 +103,7 @@ export function registerCtxSessionUpgradeCommand( } if (isWrapupInProgress(currentDeps.db, sessionId)) { - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-session-upgrade", text: "## Session Upgrade\n\n/ctx-wrapup is already compacting this session. Wait for it to finish, then try again.", level: "warning", @@ -111,7 +112,7 @@ export function registerCtxSessionUpgradeCommand( } if (isPiRecompInFlight(sessionId)) { - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-session-upgrade", text: "## Session Upgrade\n\nAn upgrade or recomp is already running for this session in the background. Wait for it to finish, then try again.", level: "warning", @@ -188,7 +189,7 @@ export function registerCtxSessionUpgradeCommand( migrationEnabled && !isMemoryMigrationDone(currentDeps.db, projectPath); if (!migrationPending) { - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-session-upgrade", text: [ "## Session Upgrade — Already Up To Date", @@ -205,7 +206,7 @@ export function registerCtxSessionUpgradeCommand( // Compartments current but project memories never migrated — run // migration only. Detached so the single migration LLM call doesn't // block the Pi REPL either (parity with the full-recomp path below). - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-session-upgrade", text: "## Session Upgrade\n\nCompartments are already current. Re-organizing project memories. This may take a while.", level: "info", @@ -222,7 +223,7 @@ export function registerCtxSessionUpgradeCommand( }), work: async () => { const summary = await runMigration(); - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-session-upgrade", text: ["## Session Upgrade — Complete", "", summary].join("\n"), level: "info", @@ -233,7 +234,7 @@ export function registerCtxSessionUpgradeCommand( return; } - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-session-upgrade", text: "## Session Upgrade\n\nRebuilding compartments into the v2 format and re-organizing project memories. This may take a while.", level: "info", @@ -275,7 +276,7 @@ export function registerCtxSessionUpgradeCommand( { preserveUserQuotes: true }, ), notify: (text) => - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-session-upgrade", text, level: "info", @@ -326,7 +327,7 @@ export function registerCtxSessionUpgradeCommand( ? extractRecompReason(recompResult.message) : `Compartments were not fully rebuilt: ${extractRecompReason(recompResult.message)}`, ); - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-session-upgrade", text: `## Session Upgrade — Incomplete\n\n${reason}`, level: "error", @@ -364,7 +365,7 @@ export function registerCtxSessionUpgradeCommand( // Step 2 — memory migration (once per project, idempotent). const migrationSummary = await runMigration(); - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-session-upgrade", text: [ "## Session Upgrade — Complete", diff --git a/packages/pi-plugin/src/commands/ctx-status.ts b/packages/pi-plugin/src/commands/ctx-status.ts index 394f412cc..cfb9ce420 100644 --- a/packages/pi-plugin/src/commands/ctx-status.ts +++ b/packages/pi-plugin/src/commands/ctx-status.ts @@ -16,7 +16,7 @@ import { resolveTailHygieneStatus } from "@magic-context/core/shared/tail-hygien import { getPiChannel1Baseline } from "../ctx-reduce-nudge-pi"; import { showStatusDialog } from "../dialogs/status-dialog"; import { resolvePiWindowGeometry } from "../pi-context-limit"; -import { resolveSessionId, sendCtxStatusMessage } from "./pi-command-utils"; +import { createCtxStatusSender, resolveSessionId } from "./pi-command-utils"; export interface RegisterCtxStatusDeps { db: ContextDatabase; @@ -79,6 +79,7 @@ export function registerCtxStatusCommand( pi.registerCommand("ctx-status", { description: "Show Magic Context status for the current Pi session", handler: async (_args, ctx) => { + const sendStatus = createCtxStatusSender(pi, ctx); const runtimeDeps = deps.resolveStatusDeps?.(ctx) ?? deps; const projectIdentity = runtimeDeps.resolveProject?.(ctx).projectIdentity ?? @@ -86,7 +87,7 @@ export function registerCtxStatusCommand( const currentDeps = { ...runtimeDeps, projectIdentity }; const sessionId = resolveSessionId(ctx); if (!sessionId) { - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-status", text: "## Magic Status\n\nNo active Pi session is available.", level: "error", @@ -144,8 +145,7 @@ export function registerCtxStatusCommand( resolveTailHygieneStatus(getPiChannel1Baseline(sessionId)), ); const details = buildStatusDetails(currentDeps, sessionId); - sendCtxStatusMessage( - pi, + sendStatus( { title: "/ctx-status", text: statusText, @@ -155,7 +155,7 @@ export function registerCtxStatusCommand( details, ); } catch (error) { - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-status", text: `## Magic Status — Failed\n\n${describeError(error).brief}`, level: "error", diff --git a/packages/pi-plugin/src/commands/ctx-wrapup.ts b/packages/pi-plugin/src/commands/ctx-wrapup.ts index c9491dc55..ce186c2ec 100644 --- a/packages/pi-plugin/src/commands/ctx-wrapup.ts +++ b/packages/pi-plugin/src/commands/ctx-wrapup.ts @@ -42,7 +42,7 @@ import { runPiHistorian } from "../pi-historian-runner"; import { isPiRecompInFlight } from "../pi-recomp-runner"; import { readPiSessionMessages } from "../read-session-pi"; import { updateStatusLine } from "../status-line"; -import { resolveSessionId, sendCtxStatusMessage } from "./pi-command-utils"; +import { createCtxStatusSender, resolveSessionId } from "./pi-command-utils"; export interface RegisterCtxWrapupDeps { db: ContextDatabase; @@ -120,9 +120,10 @@ export function registerCtxWrapupCommand( description: "Compact older Magic Context history while keeping the newest messages raw", handler: async (args, ctx) => { + const sendStatus = createCtxStatusSender(pi, ctx); const sessionId = resolveSessionId(ctx); if (!sessionId) { - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-wrapup", text: "## Magic Wrapup\n\nNo active Pi session is available.", level: "error", @@ -131,7 +132,7 @@ export function registerCtxWrapupCommand( } const currentDeps = deps.resolveRuntimeDeps?.(ctx) ?? deps; if (currentDeps.compactionOff) { - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-wrapup", text: COMPACTION_OFF_COMMAND_UNAVAILABLE, level: "warning", @@ -141,7 +142,7 @@ export function registerCtxWrapupCommand( const sessionMeta = getOrCreateSessionMeta(currentDeps.db, sessionId); if (sessionMeta.isSubagent) { - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-wrapup", text: "## Magic Wrapup — Skipped\n\n/ctx-wrapup is only available in primary sessions.", level: "warning", @@ -150,7 +151,7 @@ export function registerCtxWrapupCommand( } if (!currentDeps.historianModel) { - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-wrapup", text: "## Magic Wrapup\n\n/ctx-wrapup is unavailable because `historian.model` is not configured.", level: "error", @@ -160,7 +161,7 @@ export function registerCtxWrapupCommand( const parsed = parseWrapupArgs(args); if (!parsed.ok) { - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-wrapup", text: `## Magic Wrapup — Invalid Arguments\n\n${parsed.message}`, level: "error", @@ -175,7 +176,7 @@ export function registerCtxWrapupCommand( sessionId, parsed.messagesToKeep, ); - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-wrapup", text: result, level: @@ -194,6 +195,7 @@ export async function runPiWrapup( sessionId: string, messagesToKeep: number, ): Promise { + const sendStatus = createCtxStatusSender(pi, ctx); if (getOrCreateSessionMeta(deps.db, sessionId).isSubagent) { return "## Magic Wrapup — Skipped\n\n/ctx-wrapup is only available in primary sessions."; } @@ -292,7 +294,7 @@ export async function runPiWrapup( } }, 60_000); try { - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-wrapup", text: `## Magic Wrapup\n\nEligible history is about ${initialPlan.snapshot.trueRawEligibleTokens.toLocaleString()} tokens across approximately ${estimateChunks(initialPlan.snapshot.trueRawEligibleTokens, deps.historianChunkTokens)} historian chunk(s).`, level: "info", @@ -365,7 +367,7 @@ export async function runPiWrapup( failure = `${ownershipLostReason}; wrapped up through message ${lastEnd}. Run /ctx-wrapup again to continue.`; break; } - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-wrapup", text: `## Magic Wrapup\n\nChunk ${chunkIndex}: wrapping messages ${plan.snapshot.offset}-${plan.snapshot.eligibleEndOrdinal - 1} (~${plan.snapshot.trueRawEligibleTokens.toLocaleString()} eligible tokens remain).`, level: "info", @@ -433,7 +435,7 @@ export async function runPiWrapup( compartmentLeaseHolderId: leaseHolder, readBranchEntries: () => readBranchEntries(ctx), notifyIssue: (text) => - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-wrapup", text, level: "warning", diff --git a/packages/pi-plugin/src/commands/pi-command-utils.test.ts b/packages/pi-plugin/src/commands/pi-command-utils.test.ts index 085b13d3f..03f5732f9 100644 --- a/packages/pi-plugin/src/commands/pi-command-utils.test.ts +++ b/packages/pi-plugin/src/commands/pi-command-utils.test.ts @@ -6,7 +6,6 @@ import { type PiMessageSender, registerCtxStatusEntryRenderer, sendCtxStatusMessage, - setCtxStatusPresenter, shouldShowCtxStatusDialog, showCtxStatusDialog, } from "./pi-command-utils"; @@ -14,7 +13,6 @@ import { describe("ctx-status entries", () => { it("appends model-invisible entry data instead of sending a message", () => { const appended: Array<{ customType: string; data: unknown }> = []; - const presented: CtxStatusEntryData[] = []; let sent = 0; const pi = { registerEntryRenderer() {}, @@ -27,7 +25,6 @@ describe("ctx-status entries", () => { } as unknown as PiMessageSender; registerCtxStatusEntryRenderer(pi); - setCtxStatusPresenter(pi, (content) => presented.push(content)); sendCtxStatusMessage( pi, { title: "Magic Embed", text: "Embedding history…", level: "info" }, @@ -46,7 +43,6 @@ describe("ctx-status entries", () => { }, ]); expect(sent).toBe(0); - expect(presented).toEqual([appended[0]?.data]); }); it("registers one ctx-status entry renderer and ignores malformed data", () => { diff --git a/packages/pi-plugin/src/commands/pi-command-utils.ts b/packages/pi-plugin/src/commands/pi-command-utils.ts index a34c9f328..a003b3f33 100644 --- a/packages/pi-plugin/src/commands/pi-command-utils.ts +++ b/packages/pi-plugin/src/commands/pi-command-utils.ts @@ -41,17 +41,6 @@ type PiEntryRendererRegistration = { export type PiMessageSender = Pick & PiEntryRendererRegistration; -type CtxStatusPresenter = (content: CtxStatusMessageContent) => void; - -const ctxStatusPresenters = new WeakMap(); - -export function setCtxStatusPresenter( - pi: PiMessageSender, - present: CtxStatusPresenter, -): void { - ctxStatusPresenters.set(pi, present); -} - export function shouldShowCtxStatusDialog( content: CtxStatusMessageContent, ): boolean { @@ -64,7 +53,7 @@ export function shouldShowCtxStatusDialog( } export function resolveSessionId( - ctx: ExtensionCommandContext, + ctx: Pick, ): string | undefined { const sm = ctx.sessionManager; const getSessionId = (sm as { getSessionId?: () => string | undefined }) @@ -179,10 +168,40 @@ export function registerCtxStatusEntryRenderer(pi: PiMessageSender): boolean { } } +function presentCtxStatusMessage( + ctx: Pick, + content: CtxStatusMessageContent, +): void { + if (ctx.mode !== "rpc") return; + const type = + content.level === "error" || content.level === "warning" + ? content.level + : "info"; + if (!shouldShowCtxStatusDialog(content)) { + ctx.ui.notify(content.text, type); + return; + } + void showCtxStatusDialog(ctx, content).catch((err) => { + sessionLog( + "pi-status", + `ctx status dialog failed: ${err instanceof Error ? err.message : String(err)}`, + ); + ctx.ui.notify(content.text, type); + }); +} + +export function createCtxStatusSender( + pi: PiMessageSender, + ctx: ExtensionCommandContext, +): (content: CtxStatusMessageContent, details?: unknown) => void { + return (content, details) => sendCtxStatusMessage(pi, content, details, ctx); +} + export function sendCtxStatusMessage( pi: PiMessageSender, content: CtxStatusMessageContent, details?: unknown, + ctx?: ExtensionCommandContext, ): void { const data: CtxStatusEntryData = { ...content, @@ -194,7 +213,7 @@ export function sendCtxStatusMessage( if (typeof pi.appendEntry === "function") { pi.appendEntry(CTX_STATUS_CUSTOM_TYPE, data); } - ctxStatusPresenters.get(pi)?.(data); + if (ctx) presentCtxStatusMessage(ctx, data); // Minimal non-interactive API shims may omit appendEntry; logging remains the // safe fallback and status text must never be routed through sendMessage. diff --git a/packages/pi-plugin/src/context-handler.test.ts b/packages/pi-plugin/src/context-handler.test.ts index 6eaae964e..2902aba5a 100644 --- a/packages/pi-plugin/src/context-handler.test.ts +++ b/packages/pi-plugin/src/context-handler.test.ts @@ -1074,6 +1074,44 @@ describe("registerPiContextHandler", () => { clearAutoSearchForPiSession("ses-sticky-context"); }); + it("awaits only the requested session's in-flight historian", async () => { + let resolveA!: () => void; + let resolveB!: () => void; + const historianA = new Promise((resolve) => { + resolveA = resolve; + }); + const historianB = new Promise((resolve) => { + resolveB = resolve; + }); + const restoreA = contextHandlerInternals.setInFlightHistorianForTests( + "ses-drain-a", + historianA, + ); + const restoreB = contextHandlerInternals.setInFlightHistorianForTests( + "ses-drain-b", + historianB, + ); + let sessionADrained = false; + const drainA = awaitInFlightHistorians("ses-drain-a").then(() => { + sessionADrained = true; + }); + + try { + resolveB(); + await awaitInFlightHistorians("ses-drain-b"); + expect(sessionADrained).toBe(false); + + resolveA(); + await drainA; + expect(sessionADrained).toBe(true); + } finally { + resolveA(); + resolveB(); + restoreA(); + restoreB(); + } + }); + it("does not reset Pi model-specific state when canonical and native alias spellings flip", async () => { const db = createTestDb(); const sessionId = "ses-pi-model-alias-switch"; diff --git a/packages/pi-plugin/src/context-handler.ts b/packages/pi-plugin/src/context-handler.ts index 51efae28c..06e0ff91f 100644 --- a/packages/pi-plugin/src/context-handler.ts +++ b/packages/pi-plugin/src/context-handler.ts @@ -3346,14 +3346,22 @@ export function registerPiContextHandler( const inFlightHistorian = new Map>(); /** - * Wait for all in-flight historian runs to complete. Called from the - * Pi `session_shutdown` event handler so historian can finish writing - * compartments before the process exits. Returns immediately if no - * runs are in-flight. + * Wait for one session's in-flight historian run to complete. Called from the + * Pi `session_shutdown` event handler so its historian can finish writing + * compartments before that session shuts down. Omitting the session id waits + * for all runs and remains available for process-exit callers and tests. Returns + * immediately if no matching runs are in-flight. */ -export async function awaitInFlightHistorians(): Promise { - if (inFlightHistorian.size === 0) return; - await Promise.allSettled(Array.from(inFlightHistorian.values())); +export async function awaitInFlightHistorians( + sessionId?: string, +): Promise { + const runs = sessionId + ? [inFlightHistorian.get(sessionId)].filter( + (run): run is Promise => run !== undefined, + ) + : [...inFlightHistorian.values()]; + if (runs.length === 0) return; + await Promise.allSettled(runs); } export function resolvePiHistorianTriggerInputs(args: { diff --git a/packages/pi-plugin/src/dreamer/index.test.ts b/packages/pi-plugin/src/dreamer/index.test.ts index d73f2d75b..2119ad4a9 100644 --- a/packages/pi-plugin/src/dreamer/index.test.ts +++ b/packages/pi-plugin/src/dreamer/index.test.ts @@ -138,6 +138,39 @@ describe("Pi dreamer wiring", () => { expect(__test.registeredProjectCount()).toBe(1); }); + test("shares registrations across jiti-style module instances", async () => { + db = createDb(); + let timerStarts = 0; + __test.setStartDreamScheduleTimerFactory(async () => { + timerStarts += 1; + return mock(() => {}); + }); + const opts = dreamerOptions({ + database: db, + projectDir: "/tmp/pi-shared-module", + projectIdentity: "git:pi-shared-module", + }); + registerPiDreamerProject(opts); + await flushMicrotasks(); + + const secondInstance = await import( + `./index.ts?registry-instance=${Date.now()}` + ); + secondInstance.__test.setStartDreamScheduleTimerFactory(async () => { + timerStarts += 1; + return mock(() => {}); + }); + secondInstance.registerPiDreamerProject({ + ...opts, + registrationOwner: {}, + }); + await flushMicrotasks(); + + expect(timerStarts).toBe(1); + expect(secondInstance.__test.registeredProjectCount()).toBe(1); + secondInstance.__test.reset(); + }); + test("threads language into scheduled dreamer registration", async () => { db = createDb(); let language: string | undefined; @@ -451,6 +484,57 @@ describe("Pi dreamer wiring", () => { expect(__test.registeredProjectCount()).toBe(0); }); + test("re-registration refreshes owner recency before active-owner handoff", async () => { + db = createDb(); + const dirs: string[] = []; + __test.setStartDreamScheduleTimerFactory(async (registration) => { + dirs.push((registration as { directory: string }).directory); + return mock(() => {}); + }); + + const ownerA = {}; + const firstA = dreamerOptions({ + database: db, + projectDir: "/tmp/worktree-A", + projectIdentity: "git:pi-owner-recency", + registrationOwner: ownerA, + }); + const ownerB = dreamerOptions({ + database: db, + projectDir: "/tmp/worktree-B", + projectIdentity: "git:pi-owner-recency", + }); + const refreshedA = dreamerOptions({ + database: db, + projectDir: "/tmp/worktree-A", + projectIdentity: "git:pi-owner-recency", + registrationOwner: ownerA, + }); + const activeC = dreamerOptions({ + database: db, + projectDir: "/tmp/worktree-C", + projectIdentity: "git:pi-owner-recency", + }); + for (const owner of [firstA, ownerB, refreshedA, activeC]) { + registerPiDreamerProject(owner); + await flushMicrotasks(); + } + + unregisterPiDreamerProject({ + projectIdentity: "git:pi-owner-recency", + registrationOwner: activeC.registrationOwner, + }); + await flushMicrotasks(); + + expect(dirs).toEqual([ + "/tmp/worktree-A", + "/tmp/worktree-B", + "/tmp/worktree-A", + "/tmp/worktree-C", + "/tmp/worktree-A", + ]); + }); + test("unregister removes the project", () => { db = createDb(); const opts = dreamerOptions({ @@ -472,6 +556,69 @@ describe("Pi dreamer wiring", () => { await expect(awaitInFlightDreamers()).resolves.toBeUndefined(); }); + test("awaitInFlightDreamers waits only for the requested owner", async () => { + db = createDb(); + const ownerA = {}; + const ownerB = {}; + const gates = [ + deferred<{ ok: true; assistantText: string }>(), + deferred<{ ok: true; assistantText: string }>(), + ]; + let nextRunner = 0; + const clients: CapturedDreamClient[] = []; + __test.setPiSubagentRunnerFactory(() => { + const gate = gates[nextRunner++]; + return { run: mock(() => gate.promise) } as never; + }); + __test.setStartDreamScheduleTimerFactory(async (registration) => { + clients.push(registration.client as unknown as CapturedDreamClient); + return mock(() => {}); + }); + + registerPiDreamerProject( + dreamerOptions({ + database: db, + projectIdentity: "git:pi-owner-a", + registrationOwner: ownerA, + }), + ); + registerPiDreamerProject( + dreamerOptions({ + database: db, + projectIdentity: "git:pi-owner-b", + registrationOwner: ownerB, + }), + ); + await flushMicrotasks(); + + const sessions = await Promise.all( + clients.map( + (client) => client.session.create({}) as Promise<{ id: string }>, + ), + ); + const prompts = clients.map((client, index) => + client.session.prompt({ + path: { id: sessions[index]?.id }, + body: { system: "system", parts: [{ text: "run dreamer" }] }, + }), + ); + await flushMicrotasks(); + + let ownerADrained = false; + const ownerADrain = awaitInFlightDreamers(ownerA).then(() => { + ownerADrained = true; + }); + gates[1]?.resolve({ ok: true, assistantText: "owner B done" }); + await prompts[1]; + await flushMicrotasks(); + expect(ownerADrained).toBe(false); + + gates[0]?.resolve({ ok: true, assistantText: "owner A done" }); + await ownerADrain; + await prompts[0]; + expect(ownerADrained).toBe(true); + }); + test("fires onAdjunctsRefreshNeeded after successful dreamer prompt", async () => { db = createDb(); let capturedClient: CapturedDreamClient | null = null; diff --git a/packages/pi-plugin/src/dreamer/index.ts b/packages/pi-plugin/src/dreamer/index.ts index 1a5160400..4c60678ba 100644 --- a/packages/pi-plugin/src/dreamer/index.ts +++ b/packages/pi-plugin/src/dreamer/index.ts @@ -90,7 +90,10 @@ interface ProjectRegistration { /** Run dream tasks for this project IMMEDIATELY (Dreamer v2 manual path). * `task` forces one task ignoring its gate; omitted runs all enabled. The * registered dreamer timer also runs due tasks on its own schedule. */ - runManual: (task?: DreamTaskName) => Promise; + runManual: ( + task?: DreamTaskName, + registrationOwner?: object, + ) => Promise; /** The directory this registration was built for. `resolveProjectIdentity` * is intentionally identical across worktrees/clones of one repo, so a * `/cd` into a different checkout of the SAME repo keeps the same identity @@ -109,9 +112,34 @@ interface PiDreamerSession { messages: unknown[]; } -const registeredProjects = new Map(); +const PI_DREAMER_PROJECTS = Symbol.for( + "magic-context.pi.dreamer-registered-projects", +); + +function getRegisteredProjects(): Map { + const globals = globalThis as Record; + const existing = globals[PI_DREAMER_PROJECTS]; + if (existing instanceof Map) { + return existing as Map; + } + const projects = new Map(); + globals[PI_DREAMER_PROJECTS] = projects; + return projects; +} + +const registeredProjects = getRegisteredProjects(); const sessionsById = new Map(); -const inFlightDreams = new Set>(); +const PI_DREAMER_IN_FLIGHT = Symbol.for("magic-context.pi.dreamer-in-flight"); +const inFlightDreams = (() => { + const globals = globalThis as Record; + const existing = globals[PI_DREAMER_IN_FLIGHT]; + if (existing instanceof Map) { + return existing as Map, object>; + } + const dreams = new Map, object>(); + globals[PI_DREAMER_IN_FLIGHT] = dreams; + return dreams; +})(); let sessionCounter = 0; let piSubagentRunnerFactory: PiSubagentRunnerFactory = () => new PiSubagentRunner(); @@ -127,6 +155,7 @@ export function registerPiDreamerProject(opts: PiDreamerOptions): void { const existing = registeredProjects.get(opts.projectIdentity); const owners = existing?.owners ?? new Map(); + owners.delete(opts.registrationOwner); owners.set(opts.registrationOwner, opts); if (existing) { // Same identity, same directory → genuinely already registered, no-op. @@ -146,8 +175,8 @@ export function registerPiDreamerProject(opts: PiDreamerOptions): void { registeredProjects.delete(opts.projectIdentity); } - // Build the dreamer client ONCE so both the timer and the immediate - // /ctx-dream path share the same `inFlightDreams` accounting + the + // Build the scheduled client once. Manual runs build owner-bound clients + // below; both paths share the same `inFlightDreams` accounting and the // same module-private `sessionsById` table. const client = createPiDreamerClient(opts); @@ -185,36 +214,46 @@ export function registerPiDreamerProject(opts: PiDreamerOptions): void { }); // Manual /ctx-dream (Dreamer v2): run dream tasks NOW via the per-task - // scheduler, using the same DreamTimerClient facade the timer uses (cast at - // the boundary — it implements the session.{create,prompt,messages,delete} + // scheduler, using an owner-bound DreamTimerClient facade (cast at the + // boundary — it implements the session.{create,prompt,messages,delete} // surface the executor consumes; TS can't see structural compatibility // through the wrapper). Project-scoped: only this project's tasks run. - const runManual = async (task?: DreamTaskName): Promise => - runManualDream({ - db: opts.db, - projectIdentity: opts.projectIdentity, + // Scheduled runs keep using the timer's client above; binding manual runs + // to their owner lets session_shutdown wait only for that instance's work. + const runManual = async ( + task?: DreamTaskName, + registrationOwner = opts.registrationOwner, + ): Promise => { + const manualOpts = owners.get(registrationOwner) ?? opts; + const manualClient = createPiDreamerClient(manualOpts); + return runManualDream({ + db: manualOpts.db, + projectIdentity: manualOpts.projectIdentity, tasks: buildDreamTaskRuntimeConfigs( - opts.config, + manualOpts.config, "pi", - opts.language, - opts.mural?.model, + manualOpts.language, + manualOpts.mural?.model, ), executor: createDreamTaskExecutor({ - client: client as never, - sessionDirectory: opts.projectDir, + client: manualClient as never, + sessionDirectory: manualOpts.projectDir, openOpenCodeDb, retrospectiveRawProvider: new PiRetrospectiveRawProvider({ - projectCwd: opts.projectDir, + projectCwd: manualOpts.projectDir, }), primerRawProviderFactory: createPiPrimerRawProviderFactory(), - userMemoryCollectionEnabled: userMemoryCollectionEnabled(opts.config), + userMemoryCollectionEnabled: userMemoryCollectionEnabled( + manualOpts.config, + ), ensureProjectRegistered: ensureProjectRegisteredFromPiDirectory, - language: opts.language, - retinaHandoff: opts.retinaHandoff, - mural: opts.mural, + language: manualOpts.language, + retinaHandoff: manualOpts.retinaHandoff, + mural: manualOpts.mural, }), task, }); + }; registeredProjects.set(opts.projectIdentity, { activeOwner: opts.registrationOwner, @@ -244,6 +283,7 @@ export function registerPiDreamerProject(opts: PiDreamerOptions): void { export async function runPiDreamForProject( projectIdentity: string, task?: DreamTaskName, + registrationOwner?: object, ): Promise { const registration = registeredProjects.get(projectIdentity); if (!registration) { @@ -251,7 +291,7 @@ export async function runPiDreamForProject( `Pi dreamer not registered for project ${projectIdentity}; call registerPiDreamerProject() first`, ); } - return registration.runManual(task); + return registration.runManual(task, registrationOwner); } /** Cleanup hook — call from session_shutdown to release this session's ownership. */ @@ -292,15 +332,21 @@ export function unregisterPiDreamerProject(opts: { registeredProjects.delete(opts.projectIdentity); } -/** Wait for any currently-running dreamer task to finish gracefully. Used - * in agent_end / session_shutdown so Pi doesn't kill an in-flight dream - * in `--print` mode. Same pattern as `awaitInFlightHistorians()`. */ -export async function awaitInFlightDreamers(): Promise { - if (inFlightDreams.size === 0) { - return; - } - - await Promise.allSettled(Array.from(inFlightDreams)); +/** Wait for any currently-running dreamer task owned by this extension + * instance to finish gracefully. Used in `session_shutdown`; omitting the owner + * waits for all tasks and remains available for process-exit callers and tests. + * Same pattern as `awaitInFlightHistorians()`. */ +export async function awaitInFlightDreamers( + registrationOwner?: object, +): Promise { + const runs = + registrationOwner === undefined + ? [...inFlightDreams.keys()] + : [...inFlightDreams.entries()] + .filter(([, owner]) => owner === registrationOwner) + .map(([run]) => run); + if (runs.length === 0) return; + await Promise.allSettled(runs); } function createPiDreamerClient(opts: PiDreamerOptions): DreamTimerClient { @@ -353,7 +399,7 @@ function createPiDreamerClient(opts: PiDreamerOptions): DreamTimerClient { // `--thinking` without letting a primary level leak to fallbacks. thinkingLevel: extractBodyVariant(args), }); - inFlightDreams.add(runPromise); + inFlightDreams.set(runPromise, opts.registrationOwner); try { const result = await runPromise; if (!result.ok) { diff --git a/packages/pi-plugin/src/index-in-process-latch.test.ts b/packages/pi-plugin/src/index-in-process-latch.test.ts index 24f133ea1..c4ee6669a 100644 --- a/packages/pi-plugin/src/index-in-process-latch.test.ts +++ b/packages/pi-plugin/src/index-in-process-latch.test.ts @@ -189,6 +189,44 @@ describe("Pi in-process child guard (#247)", () => { expect(sibling.commands).toContain("ctx-status"); }, 15_000); + it("does not suppress an independent session while a child marker is active", async () => { + isolateXdgEnv(); + delete process.env[MAGIC_CONTEXT_PI_SUBAGENT_ENV]; + + const parent = createCountingPi(); + await magicContextPiExtension(parent.pi); + + let childMarked!: () => void; + const marked = new Promise((resolve) => { + childMarked = resolve; + }); + let releaseChild!: () => void; + const release = new Promise((resolve) => { + releaseChild = resolve; + }); + const child = createCountingPi(); + const childBranch = Promise.resolve().then(async () => { + parent.emitEvent("subagents:child:session-created"); + await magicContextPiExtension(child.pi); + childMarked(); + await release; + parent.emitEvent("subagents:child:disposed"); + }); + + await marked; + try { + expect(child.tools).toEqual([]); + + const independent = createCountingPi(); + await magicContextPiExtension(independent.pi); + expect(independent.tools).toContain("ctx_search"); + expect(independent.commands).toContain("ctx-status"); + } finally { + releaseChild(); + await childBranch; + } + }, 15_000); + it("keeps the spawned-child environment guard", async () => { isolateXdgEnv(); process.env[MAGIC_CONTEXT_PI_SUBAGENT_ENV] = "1"; diff --git a/packages/pi-plugin/src/index.ts b/packages/pi-plugin/src/index.ts index 2e09c381c..971e80716 100644 --- a/packages/pi-plugin/src/index.ts +++ b/packages/pi-plugin/src/index.ts @@ -120,10 +120,8 @@ import { registerCtxStatusCommand } from "./commands/ctx-status"; import { registerCtxWrapupCommand } from "./commands/ctx-wrapup"; import { registerCtxStatusEntryRenderer, + resolveSessionId, sendCtxStatusMessage, - setCtxStatusPresenter, - shouldShowCtxStatusDialog, - showCtxStatusDialog, } from "./commands/pi-command-utils"; import { loadPiConfig } from "./config"; import { @@ -1281,23 +1279,6 @@ async function startPiMagicContextRuntime( ); pi.on("session_start", async (event, ctx) => { - if (ctx.mode === "rpc") { - setCtxStatusPresenter(pi, (content) => { - const type = - content.level === "error" || content.level === "warning" - ? content.level - : "info"; - const useDialog = shouldShowCtxStatusDialog(content); - if (!useDialog) { - ctx.ui.notify(content.text, type); - return; - } - void showCtxStatusDialog(ctx, content).catch((err) => { - warn("ctx status dialog failed:", err); - ctx.ui.notify(content.text, type); - }); - }); - } await handlePiCloneSessionStart(event, ctx, { db, signalPendingMarker: signalPiDeferredCompactionMarkerDrain, @@ -1550,6 +1531,7 @@ async function startPiMagicContextRuntime( onProjectSeen: (identity) => seenDreamerProjectIdentities.add(identity), ensureRegistered: (ctx) => syncDreamerProjectRegistration(resolveCurrentProjectDeps(ctx)), + registrationOwner: dreamerRegistrationOwner, }); info("registered /ctx-dream"); @@ -2321,9 +2303,9 @@ async function startPiMagicContextRuntime( } }); - // Unregister project from dreamer timer on session shutdown. Pi's - // `/reload` command tears down extensions and re-runs this default - // export — without unregistering, the dreamer timer would hold a + // Release this extension instance's dreamer registrations on session shutdown. + // Pi's `/reload` command tears down extensions and re-runs this default + // export — without releasing them, the dreamer timer would hold a // stale reference to the previous extension instance. // // IMPORTANT: We do NOT close the SQLite handle here. `openDatabase()` @@ -2342,26 +2324,35 @@ async function startPiMagicContextRuntime( // session_shutdown only fires when the user is actually leaving // the session, so a brief wait is acceptable — and lets the // JSONL session state reach a consistent compartment boundary - // before the process exits. + // before that session's cleanup completes. // - // 5-second cap protects interactive shutdown from a hung + // 5-second cap per drain protects interactive shutdown from a hung // subagent. In `pi --print` mode the process exits after // agent_end before this handler fires anyway, so the cap // doesn't help that mode (and we don't pretend it does — see // the comment block on the agent_end handler above). const SHUTDOWN_DRAIN_MS = 5_000; - try { - await withTimeout(awaitInFlightHistorians(), SHUTDOWN_DRAIN_MS); - } catch (err) { - warn("shutdown: historian drain threw:", err); - } - try { - await withTimeout(awaitInFlightRecomps(), SHUTDOWN_DRAIN_MS); - } catch (err) { - warn("shutdown: recomp drain threw:", err); + const sessionId = resolveSessionId(ctx); + if (sessionId) { + try { + await withTimeout( + awaitInFlightHistorians(sessionId), + SHUTDOWN_DRAIN_MS, + ); + } catch (err) { + warn("shutdown: historian drain threw:", err); + } + try { + await withTimeout(awaitInFlightRecomps(sessionId), SHUTDOWN_DRAIN_MS); + } catch (err) { + warn("shutdown: recomp drain threw:", err); + } } try { - await withTimeout(awaitInFlightDreamers(), SHUTDOWN_DRAIN_MS); + await withTimeout( + awaitInFlightDreamers(dreamerRegistrationOwner), + SHUTDOWN_DRAIN_MS, + ); } catch (err) { warn("shutdown: dreamer drain threw:", err); } diff --git a/packages/pi-plugin/src/pi-recomp-runner.test.ts b/packages/pi-plugin/src/pi-recomp-runner.test.ts new file mode 100644 index 000000000..2c30b466f --- /dev/null +++ b/packages/pi-plugin/src/pi-recomp-runner.test.ts @@ -0,0 +1,37 @@ +import { expect, test } from "bun:test"; +import { awaitInFlightRecomps, spawnPiRecompRun } from "./pi-recomp-runner"; + +function deferred() { + let resolve!: () => void; + const promise = new Promise((done) => { + resolve = done; + }); + return { promise, resolve }; +} + +test("awaitInFlightRecomps waits only for the requested session", async () => { + const sessionA = deferred(); + const sessionB = deferred(); + const spawn = (sessionId: string, work: Promise) => + spawnPiRecompRun({ + sessionId, + provider: { readMessages: async () => [] } as never, + onStatusChange: () => {}, + work: () => work, + }); + + spawn("session-a", sessionA.promise); + spawn("session-b", sessionB.promise); + let sessionADrained = false; + const drainA = awaitInFlightRecomps("session-a").then(() => { + sessionADrained = true; + }); + + sessionB.resolve(); + await awaitInFlightRecomps("session-b"); + expect(sessionADrained).toBe(false); + + sessionA.resolve(); + await drainA; + expect(sessionADrained).toBe(true); +}); diff --git a/packages/pi-plugin/src/pi-recomp-runner.ts b/packages/pi-plugin/src/pi-recomp-runner.ts index fa840174d..d8dfe50d1 100644 --- a/packages/pi-plugin/src/pi-recomp-runner.ts +++ b/packages/pi-plugin/src/pi-recomp-runner.ts @@ -7,7 +7,7 @@ import { setMagicContextRecompActive } from "./status-line"; /** * In-flight detached recomp / upgrade runs, keyed by session, so the - * `session_shutdown` handler can await them before Pi exits — mirrors + * `session_shutdown` handler can await only that session's work — mirrors * `inFlightHistorian` in context-handler.ts. * * Why detached: Pi's command handler IS the REPL turn (single process). Awaiting @@ -26,13 +26,19 @@ export function isPiRecompInFlight(sessionId: string): boolean { } /** - * Await all in-flight recomp/upgrade runs. Called from `session_shutdown` - * (bounded by a timeout there) so a background recomp can finish publishing - * before Pi tears the session down. + * Await one session's in-flight recomp/upgrade run. Called from + * `session_shutdown` (bounded by a timeout there) so a background recomp can + * finish publishing before Pi tears that session down. Omitting the session id + * waits for all runs and remains available for process-exit callers and tests. */ -export async function awaitInFlightRecomps(): Promise { - if (inFlightRecomp.size === 0) return; - await Promise.allSettled(Array.from(inFlightRecomp.values())); +export async function awaitInFlightRecomps(sessionId?: string): Promise { + const runs = sessionId + ? [inFlightRecomp.get(sessionId)].filter( + (run): run is Promise => run !== undefined, + ) + : [...inFlightRecomp.values()]; + if (runs.length === 0) return; + await Promise.allSettled(runs); } /** diff --git a/packages/pi-plugin/src/subagent-runner.test.ts b/packages/pi-plugin/src/subagent-runner.test.ts index f12df46f1..f0ba60ef9 100644 --- a/packages/pi-plugin/src/subagent-runner.test.ts +++ b/packages/pi-plugin/src/subagent-runner.test.ts @@ -10,6 +10,7 @@ import { import { EventEmitter } from "node:events"; import { existsSync, + mkdirSync, mkdtempSync, readFileSync, rmSync, @@ -1067,6 +1068,57 @@ describe("PiSubagentRunner spawn lifecycle", () => { }); }); + it("with no piBinary override, spawns the host runtime + cli.js (Windows-safe, #177)", async () => { + // Default resolution must NOT spawn a bare "pi" (which ENOENTs on Windows + // because npm installs a pi.cmd shim, not a literal pi). It re-invokes the + // exact host CLI: process.execPath + process.argv[1], with no shell. + const root = mkdtempSync(join(tmpdir(), "mc-pi-cli-")); + const distDir = join( + root, + "node_modules", + "@earendil-works", + "pi-coding-agent", + "dist", + ); + const cliPath = join(distDir, "cli.js"); + mkdirSync(distDir, { recursive: true }); + writeFileSync(cliPath, ""); + const previousScript = process.argv[1]; + process.argv[1] = cliPath; + try { + const child = createMockChild(); + const spawnImpl = mock(() => child as never); + const runner = new PiSubagentRunner({ spawnImpl: spawnImpl as never }); + + const resultPromise = runner.run(baseOptions); + child.writeStdoutLine({ type: "session", id: "s1" }); + child.writeStdoutLine( + agentEnd([ + { + role: "assistant", + content: [{ type: "text", text: "ok" }], + stopReason: "stop", + }, + ]), + ); + child.emitClose(0); + await resultPromise; + + const [command, spawnArgs, opts] = ( + spawnImpl.mock.calls as unknown[][] + )[0] as [string, string[], { shell?: boolean }]; + expect(command).toBe(process.execPath); + expect(spawnArgs[0]).toBe(cliPath); + // Crucially, never a bare "pi". + expect(command).not.toBe("pi"); + expect(opts.shell).toBeFalsy(); + } finally { + if (previousScript === undefined) delete process.argv[1]; + else process.argv[1] = previousScript; + rmSync(root, { recursive: true, force: true }); + } + }); + it("with no piBinary override, does not re-run an embedded host", async () => { // The Bun test file stands in for pi-web's Next.js argv[1]. const child = createMockChild(); From 6b490fcd7fa629fcfd9e21a477848e944249615d Mon Sep 17 00:00:00 2001 From: elrond <171308700+elrond298@users.noreply.github.com> Date: Fri, 21 Aug 2026 20:26:41 +0800 Subject: [PATCH 4/5] fix(pi-plugin): prevent stale Dreamer runs across worktrees - reject Dreamer prompts after their registration owner is removed or switches worktrees - track complete manual runs, including domain lease waits, during shutdown - unregister the session owner before draining its Dreamer work - notify every registered owner after successful adjunct updates - add regression coverage for shutdown and multi-worktree races --- .../pi-plugin/src/agent-end-handler.test.ts | 14 ++ packages/pi-plugin/src/dreamer/index.test.ts | 174 ++++++++++++++++++ packages/pi-plugin/src/dreamer/index.ts | 56 +++++- packages/pi-plugin/src/index.ts | 26 ++- 4 files changed, 253 insertions(+), 17 deletions(-) diff --git a/packages/pi-plugin/src/agent-end-handler.test.ts b/packages/pi-plugin/src/agent-end-handler.test.ts index 16af9aefe..99213908e 100644 --- a/packages/pi-plugin/src/agent-end-handler.test.ts +++ b/packages/pi-plugin/src/agent-end-handler.test.ts @@ -125,6 +125,20 @@ describe("session_shutdown handler (drain location)", () => { ); }); + test("stops Dreamer registration before draining its work", () => { + const shutdownAt = body.indexOf("sessionShuttingDown = true"); + const unregisterAt = body.indexOf("unregisterPiDreamerProject"); + const drainAt = body.indexOf("awaitInFlightDreamers"); + expect(shutdownAt).toBeGreaterThanOrEqual(0); + expect(unregisterAt).toBeGreaterThanOrEqual(0); + expect(drainAt).toBeGreaterThanOrEqual(0); + expect(shutdownAt).toBeLessThan(unregisterAt); + expect(unregisterAt).toBeLessThan(drainAt); + expect(INDEX_SRC).toMatch( + /function syncDreamerProjectRegistration[\s\S]*?if \(sessionShuttingDown\) return;/, + ); + }); + test("drain timeout uses unref/clear helper", () => { // The bounded wait goes through withTimeout(), whose implementation // calls unref() and clearTimeout() so an early drain cannot pin exit. diff --git a/packages/pi-plugin/src/dreamer/index.test.ts b/packages/pi-plugin/src/dreamer/index.test.ts index 2119ad4a9..10bcbda39 100644 --- a/packages/pi-plugin/src/dreamer/index.test.ts +++ b/packages/pi-plugin/src/dreamer/index.test.ts @@ -3,7 +3,12 @@ import { type DreamerConfig, DreamerConfigSchema, } from "@magic-context/core/config/schema/magic-context"; +import { + acquireLease, + releaseLease, +} from "@magic-context/core/features/magic-context/dreamer/lease"; import { getTaskScheduleState } from "@magic-context/core/features/magic-context/dreamer/storage-task-schedule"; +import { leaseKeyFor } from "@magic-context/core/features/magic-context/dreamer/task-registry"; import { insertMemory } from "@magic-context/core/features/magic-context/memory"; import { runMigrations } from "@magic-context/core/features/magic-context/migrations"; import { initializeDatabase } from "@magic-context/core/features/magic-context/storage-db"; @@ -419,6 +424,51 @@ describe("Pi dreamer wiring", () => { ]); }); + test("old timer client cannot prompt after its owner switches worktrees", async () => { + db = createDb(); + const clients: CapturedDreamClient[] = []; + const run = mock(async () => ({ + ok: true as const, + assistantText: "done", + })); + __test.setPiSubagentRunnerFactory(() => ({ run }) as never); + __test.setStartDreamScheduleTimerFactory(async (registration) => { + clients.push(registration.client as unknown as CapturedDreamClient); + return mock(() => {}); + }); + const projectIdentity = "git:pi-stale-worktree-client"; + const owner = {}; + registerPiDreamerProject( + dreamerOptions({ + database: db, + projectDir: "/tmp/worktree-A", + projectIdentity, + registrationOwner: owner, + }), + ); + await flushMicrotasks(); + const oldClient = clients[0]; + expect(oldClient).toBeDefined(); + if (!oldClient) throw new Error("first dreamer client was not captured"); + const created = (await oldClient.session.create({})) as { id: string }; + + registerPiDreamerProject( + dreamerOptions({ + database: db, + projectDir: "/tmp/worktree-B", + projectIdentity, + registrationOwner: owner, + }), + ); + await expect( + oldClient.session.prompt({ + path: { id: created.id }, + body: { system: "system", parts: [{ text: "run dreamer" }] }, + }), + ).rejects.toThrow("registration is no longer active"); + expect(run).not.toHaveBeenCalled(); + }); + test("active-owner handoff starts one timer when remaining worktree dirs repeat", async () => { db = createDb(); const dirs: string[] = []; @@ -535,6 +585,86 @@ describe("Pi dreamer wiring", () => { ]); }); + test("rejects a manual run after its explicit owner unregisters", async () => { + db = createDb(); + __test.setStartDreamScheduleTimerFactory(async () => mock(() => {})); + const projectIdentity = "git:pi-stale-manual-owner"; + const ownerA = dreamerOptions({ + database: db, + projectDir: "/tmp/worktree-A", + projectIdentity, + }); + const ownerB = dreamerOptions({ + database: db, + projectDir: "/tmp/worktree-B", + projectIdentity, + }); + registerPiDreamerProject(ownerA); + await flushMicrotasks(); + registerPiDreamerProject(ownerB); + await flushMicrotasks(); + unregisterPiDreamerProject({ + projectIdentity, + registrationOwner: ownerA.registrationOwner, + }); + __test.setPiSubagentRunnerFactory(() => { + throw new Error("manual client should not be created"); + }); + + await expect( + runPiDreamForProject( + projectIdentity, + undefined, + ownerA.registrationOwner, + ), + ).rejects.toThrow( + `Pi dreamer registration owner is no longer active for project ${projectIdentity}`, + ); + }); + + test("owner drain covers a lease wait and stale owner cannot start a prompt", async () => { + db = createDb(); + __test.setStartDreamScheduleTimerFactory(async () => mock(() => {})); + const run = mock(async () => ({ + ok: true as const, + assistantText: "", + })); + __test.setPiSubagentRunnerFactory(() => ({ run }) as never); + const projectIdentity = "git:pi-manual-lease-wait"; + const owner = {}; + const leaseKey = leaseKeyFor("curate", projectIdentity); + const blocker = "manual-lease-blocker"; + expect(acquireLease(db, blocker, leaseKey)).toBe(true); + registerPiDreamerProject( + dreamerOptions({ + database: db, + projectIdentity, + registrationOwner: owner, + config: DreamerConfigSchema.parse({ + model: "test/model", + tasks: { curate: { schedule: "0 4 * * *" } }, + }), + }), + ); + + const manualRun = runPiDreamForProject(projectIdentity, "curate", owner); + await flushMicrotasks(); + let drained = false; + const drain = awaitInFlightDreamers(owner).then(() => { + drained = true; + }); + await flushMicrotasks(); + expect(drained).toBe(false); + + unregisterPiDreamerProject({ projectIdentity, registrationOwner: owner }); + releaseLease(db, blocker, leaseKey); + const result = await manualRun; + await drain; + expect(drained).toBe(true); + expect(run).not.toHaveBeenCalled(); + expect(result.failed).toEqual(["curate"]); + }); + test("unregister removes the project", () => { db = createDb(); const opts = dreamerOptions({ @@ -655,6 +785,50 @@ describe("Pi dreamer wiring", () => { expect(onAdjunctsRefreshNeeded).toHaveBeenCalledWith("git:pi-g5-success"); }); + test("notifies every registered worktree after a successful dreamer prompt", async () => { + db = createDb(); + let capturedClient: CapturedDreamClient | null = null; + __test.setStartDreamScheduleTimerFactory(async (registration) => { + capturedClient = registration.client as unknown as CapturedDreamClient; + return mock(() => {}); + }); + __test.setPiSubagentRunnerFactory( + () => + ({ + run: mock(async () => ({ ok: true, assistantText: "done" })), + }) as never, + ); + const projectIdentity = "git:pi-g5-worktrees"; + const refreshA = mock(() => {}); + const refreshB = mock(() => {}); + registerPiDreamerProject( + dreamerOptions({ + database: db, + projectDir: "/tmp/worktree-A", + projectIdentity, + onAdjunctsRefreshNeeded: refreshA, + }), + ); + registerPiDreamerProject( + dreamerOptions({ + database: db, + projectDir: "/tmp/worktree-B", + projectIdentity, + onAdjunctsRefreshNeeded: refreshB, + }), + ); + + const client = requireCapturedClient(capturedClient); + const created = (await client.session.create({})) as { id: string }; + await client.session.prompt({ + path: { id: created.id }, + body: { system: "system", parts: [{ text: "run dreamer" }] }, + }); + + expect(refreshA).toHaveBeenCalledWith(projectIdentity); + expect(refreshB).toHaveBeenCalledWith(projectIdentity); + }); + test("undefined onAdjunctsRefreshNeeded is a no-op after successful dreamer prompt", async () => { db = createDb(); let capturedClient: CapturedDreamClient | null = null; diff --git a/packages/pi-plugin/src/dreamer/index.ts b/packages/pi-plugin/src/dreamer/index.ts index 4c60678ba..dda2234f1 100644 --- a/packages/pi-plugin/src/dreamer/index.ts +++ b/packages/pi-plugin/src/dreamer/index.ts @@ -157,6 +157,14 @@ export function registerPiDreamerProject(opts: PiDreamerOptions): void { const owners = existing?.owners ?? new Map(); owners.delete(opts.registrationOwner); owners.set(opts.registrationOwner, opts); + const notifyOwnersOfAdjunctRefresh = (projectIdentity: string): void => { + const callbacks = new Set( + [...owners.values()] + .map((owner) => owner.onAdjunctsRefreshNeeded) + .filter((callback) => callback !== undefined), + ); + for (const callback of callbacks) callback(projectIdentity); + }; if (existing) { // Same identity, same directory → genuinely already registered, no-op. // Keep this extension instance as an owner so another session cannot @@ -178,7 +186,11 @@ export function registerPiDreamerProject(opts: PiDreamerOptions): void { // Build the scheduled client once. Manual runs build owner-bound clients // below; both paths share the same `inFlightDreams` accounting and the // same module-private `sessionsById` table. - const client = createPiDreamerClient(opts); + const client = createPiDreamerClient( + opts, + notifyOwnersOfAdjunctRefresh, + () => owners.get(opts.registrationOwner)?.projectDir === opts.projectDir, + ); let cleanup: (() => void) | undefined; let cancelled = false; @@ -222,11 +234,23 @@ export function registerPiDreamerProject(opts: PiDreamerOptions): void { // to their owner lets session_shutdown wait only for that instance's work. const runManual = async ( task?: DreamTaskName, - registrationOwner = opts.registrationOwner, + registrationOwner?: object, ): Promise => { - const manualOpts = owners.get(registrationOwner) ?? opts; - const manualClient = createPiDreamerClient(manualOpts); - return runManualDream({ + const manualOpts = + registrationOwner === undefined ? opts : owners.get(registrationOwner); + if (!manualOpts) { + throw new Error( + `Pi dreamer registration owner is no longer active for project ${opts.projectIdentity}`, + ); + } + const manualClient = createPiDreamerClient( + manualOpts, + notifyOwnersOfAdjunctRefresh, + () => + owners.get(manualOpts.registrationOwner)?.projectDir === + manualOpts.projectDir, + ); + const manualRun = runManualDream({ db: manualOpts.db, projectIdentity: manualOpts.projectIdentity, tasks: buildDreamTaskRuntimeConfigs( @@ -253,6 +277,14 @@ export function registerPiDreamerProject(opts: PiDreamerOptions): void { }), task, }); + // Track the whole manual run, including lease waits before its first + // subagent prompt, so owner-scoped shutdown cannot miss it. + inFlightDreams.set(manualRun, manualOpts.registrationOwner); + try { + return await manualRun; + } finally { + inFlightDreams.delete(manualRun); + } }; registeredProjects.set(opts.projectIdentity, { @@ -349,7 +381,11 @@ export async function awaitInFlightDreamers( await Promise.allSettled(runs); } -function createPiDreamerClient(opts: PiDreamerOptions): DreamTimerClient { +function createPiDreamerClient( + opts: PiDreamerOptions, + onAdjunctsRefreshNeeded = opts.onAdjunctsRefreshNeeded, + isRegistrationOwnerActive: () => boolean = () => true, +): DreamTimerClient { const runner = piSubagentRunnerFactory(); const session = { @@ -371,6 +407,12 @@ function createPiDreamerClient(opts: PiDreamerOptions): DreamTimerClient { throw new Error(`Pi dreamer session not found: ${sessionId}`); } + if (!isRegistrationOwnerActive()) { + throw new Error( + `Pi dreamer registration is no longer active for project ${opts.projectIdentity}`, + ); + } + const userMessage = extractUserMessage(args); const systemPrompt = extractSystemPrompt(args); // Per-task model override (Dreamer v2): the SHARED executor @@ -427,7 +469,7 @@ function createPiDreamerClient(opts: PiDreamerOptions): DreamTimerClient { // can update , , or . The cost // of one extra disk read per session next turn is tiny compared to // stale adjuncts surviving until restart. - opts.onAdjunctsRefreshNeeded?.(opts.projectIdentity); + onAdjunctsRefreshNeeded?.(opts.projectIdentity); } finally { inFlightDreams.delete(runPromise); } diff --git a/packages/pi-plugin/src/index.ts b/packages/pi-plugin/src/index.ts index 971e80716..7a55dbbe1 100644 --- a/packages/pi-plugin/src/index.ts +++ b/packages/pi-plugin/src/index.ts @@ -955,6 +955,7 @@ async function startPiMagicContextRuntime( const projectDir = process.cwd(); const seenDreamerProjectIdentities = new Set(); const dreamerRegistrationOwner = {}; + let sessionShuttingDown = false; // Step 5b: load the user's full magic-context.jsonc config. The loader // reads the shared CortexKit project/user paths, validates them through the // shared Zod schema, falls back to Pi-owned legacy files only while migration @@ -1193,6 +1194,7 @@ async function startPiMagicContextRuntime( function syncDreamerProjectRegistration( current: ResolvedPiProjectDeps, ): void { + if (sessionShuttingDown) return; seenDreamerProjectIdentities.add(current.projectIdentity); if (!current.dreamerConfig) { unregisterPiDreamerProject({ @@ -2317,6 +2319,7 @@ async function startPiMagicContextRuntime( // re-runs the extension code but keeps the host process alive, so // the cached handle is still valid across reload boundaries. pi.on("session_shutdown", async (_event, ctx) => { + sessionShuttingDown = true; // Bounded drain of in-flight historian / dreamer runs that were // kicked off by recent turns. We moved the drain here from // `agent_end` because Pi awaits agent_end handlers and was @@ -2333,6 +2336,19 @@ async function startPiMagicContextRuntime( // the comment block on the agent_end handler above). const SHUTDOWN_DRAIN_MS = 5_000; const sessionId = resolveSessionId(ctx); + // Stop this owner from admitting new Dreamer runs before snapshotting its + // in-flight work. A command that already started remains owner-tracked and + // is drained below; a later command fails instead of outliving shutdown. + try { + for (const identity of seenDreamerProjectIdentities) { + unregisterPiDreamerProject({ + projectIdentity: identity, + registrationOwner: dreamerRegistrationOwner, + }); + } + } catch (err) { + warn("shutdown: unregisterPiDreamerProject threw:", err); + } if (sessionId) { try { await withTimeout( @@ -2356,16 +2372,6 @@ async function startPiMagicContextRuntime( } catch (err) { warn("shutdown: dreamer drain threw:", err); } - try { - for (const identity of seenDreamerProjectIdentities) { - unregisterPiDreamerProject({ - projectIdentity: identity, - registrationOwner: dreamerRegistrationOwner, - }); - } - } catch (err) { - warn("shutdown: unregisterPiDreamerProject threw:", err); - } // Clear per-session system-prompt adjunct caches (sticky date, // project docs, user profile, key files). Pi's // `_extensionRunner.invalidate` resets module state on session From abc835fed3775e93dc1f1061ab788dd948db2d3a Mon Sep 17 00:00:00 2001 From: elrond <171308700+elrond298@users.noreply.github.com> Date: Fri, 21 Aug 2026 21:33:47 +0800 Subject: [PATCH 5/5] fix(pi-plugin): require owners for manual Dreamer runs - resolve manual runs through the active registration owner's options - reject ownerless and deregistered-owner requests - preserve process-shared argument order across extension reloads - cover refreshed owners, stale owners, and lease-wait draining --- .../src/commands/ctx-commands.test.ts | 4 + packages/pi-plugin/src/commands/ctx-dream.ts | 2 +- packages/pi-plugin/src/dreamer/index.test.ts | 85 ++++++++++--------- packages/pi-plugin/src/dreamer/index.ts | 27 +++--- 4 files changed, 65 insertions(+), 53 deletions(-) diff --git a/packages/pi-plugin/src/commands/ctx-commands.test.ts b/packages/pi-plugin/src/commands/ctx-commands.test.ts index 93330c8e4..9dcb064bb 100644 --- a/packages/pi-plugin/src/commands/ctx-commands.test.ts +++ b/packages/pi-plugin/src/commands/ctx-commands.test.ts @@ -295,6 +295,7 @@ describe("Pi Magic Context commands", () => { db, projectDir: "/tmp/project", projectIdentity: "/tmp/project", + registrationOwner: {}, ensureRegistered: (ctx) => { registrationCwds.push(ctx.cwd); }, @@ -318,6 +319,7 @@ describe("Pi Magic Context commands", () => { db, projectDir: "/tmp/project", projectIdentity: "/tmp/project", + registrationOwner: {}, }); await handlers.get("ctx-dream")?.("verify", createCtx()); @@ -347,6 +349,7 @@ describe("Pi Magic Context commands", () => { db, projectDir: "/tmp/project", projectIdentity: "/tmp/project", + registrationOwner: {}, dreamerEnabled: false, }); await handlers.get("ctx-dream")?.("", createCtx()); @@ -362,6 +365,7 @@ describe("Pi Magic Context commands", () => { db, projectDir: "/tmp/project-a", projectIdentity: "/tmp/project-a", + registrationOwner: {}, resolveProject: (ctx) => ({ projectDir: ctx.cwd, projectIdentity: ctx.cwd, diff --git a/packages/pi-plugin/src/commands/ctx-dream.ts b/packages/pi-plugin/src/commands/ctx-dream.ts index 67f886425..ba7c6c7f6 100644 --- a/packages/pi-plugin/src/commands/ctx-dream.ts +++ b/packages/pi-plugin/src/commands/ctx-dream.ts @@ -25,7 +25,7 @@ export function registerCtxDreamCommand( resolveDreamerEnabled?: (ctx: { cwd: string }) => boolean | undefined; onProjectSeen?: (projectIdentity: string) => void; ensureRegistered?: (ctx: { cwd: string }) => void | Promise; - registrationOwner?: object; + registrationOwner: object; }, ): void { pi.registerCommand("ctx-dream", { diff --git a/packages/pi-plugin/src/dreamer/index.test.ts b/packages/pi-plugin/src/dreamer/index.test.ts index 10bcbda39..32341b493 100644 --- a/packages/pi-plugin/src/dreamer/index.test.ts +++ b/packages/pi-plugin/src/dreamer/index.test.ts @@ -196,7 +196,7 @@ describe("Pi dreamer wiring", () => { expect(language).toBe("es"); }); - test("manual dreamer passes a directive-bearing system prompt when language is set", async () => { + test("manual dreamer uses refreshed options for its explicit owner", async () => { db = createDb(); let capturedSystem = ""; __test.setStartDreamScheduleTimerFactory(async () => mock(() => {})); @@ -215,22 +215,23 @@ describe("Pi dreamer wiring", () => { content: "The Pi harness runs dreamer prompts through a subprocess.", }); - registerPiDreamerProject( - dreamerOptions({ - database: db, - projectDir: process.cwd(), - projectIdentity: "git:pi-manual-language", - config: DreamerConfigSchema.parse({ - model: "test/model", - tasks: { curate: { schedule: "0 4 * * *" } }, - }), - language: "es", + const opts = dreamerOptions({ + database: db, + projectDir: process.cwd(), + projectIdentity: "git:pi-manual-language", + config: DreamerConfigSchema.parse({ + model: "test/model", + tasks: { curate: { schedule: "0 4 * * *" } }, }), - ); + language: "en", + }); + registerPiDreamerProject(opts); + registerPiDreamerProject({ ...opts, language: "es" }); const result = await runPiDreamForProject( "git:pi-manual-language", "curate", + opts.registrationOwner, ); expect( getTaskScheduleState(db, "git:pi-manual-language", "curate")?.lastError, @@ -279,34 +280,34 @@ describe("Pi dreamer wiring", () => { content: "Use the shared release checklist before publishing.", }); - registerPiDreamerProject( - dreamerOptions({ - database: db, - projectDir: process.cwd(), - projectIdentity: "git:pi-curate-pseudo-tool-call", - // Model resolution is harness-scoped: scheduling remains at - // dreamer.tasks, while Pi's attempts live under dreamer.pi. - config: { - ...DreamerConfigSchema.parse({ - tasks: { curate: { schedule: "0 4 * * *" } }, - }), - pi: { - model: { model: "primary/curator", thinking_level: "high" }, - tasks: { - curate: { - fallback_models: [ - { model: "fallback/curator", thinking_level: "low" }, - ], - }, + const opts = dreamerOptions({ + database: db, + projectDir: process.cwd(), + projectIdentity: "git:pi-curate-pseudo-tool-call", + // Model resolution is harness-scoped: scheduling remains at + // dreamer.tasks, while Pi's attempts live under dreamer.pi. + config: { + ...DreamerConfigSchema.parse({ + tasks: { curate: { schedule: "0 4 * * *" } }, + }), + pi: { + model: { model: "primary/curator", thinking_level: "high" }, + tasks: { + curate: { + fallback_models: [ + { model: "fallback/curator", thinking_level: "low" }, + ], }, }, - } as never, - }), - ); + }, + } as never, + }); + registerPiDreamerProject(opts); const result = await runPiDreamForProject( "git:pi-curate-pseudo-tool-call", "curate", + opts.registrationOwner, ); expect(attemptedModels).toEqual(["primary/curator", "fallback/curator"]); @@ -585,7 +586,7 @@ describe("Pi dreamer wiring", () => { ]); }); - test("rejects a manual run after its explicit owner unregisters", async () => { + test("rejects ownerless and unregistered-owner manual runs", async () => { db = createDb(); __test.setStartDreamScheduleTimerFactory(async () => mock(() => {})); const projectIdentity = "git:pi-stale-manual-owner"; @@ -603,14 +604,20 @@ describe("Pi dreamer wiring", () => { await flushMicrotasks(); registerPiDreamerProject(ownerB); await flushMicrotasks(); - unregisterPiDreamerProject({ - projectIdentity, - registrationOwner: ownerA.registrationOwner, - }); __test.setPiSubagentRunnerFactory(() => { throw new Error("manual client should not be created"); }); + await expect( + runPiDreamForProject(projectIdentity, undefined, undefined as never), + ).rejects.toThrow( + `Pi dreamer registration owner is no longer active for project ${projectIdentity}`, + ); + + unregisterPiDreamerProject({ + projectIdentity, + registrationOwner: ownerA.registrationOwner, + }); await expect( runPiDreamForProject( projectIdentity, diff --git a/packages/pi-plugin/src/dreamer/index.ts b/packages/pi-plugin/src/dreamer/index.ts index dda2234f1..c262136fd 100644 --- a/packages/pi-plugin/src/dreamer/index.ts +++ b/packages/pi-plugin/src/dreamer/index.ts @@ -88,11 +88,12 @@ interface ProjectRegistration { activeOwner: object; owners: Map; /** Run dream tasks for this project IMMEDIATELY (Dreamer v2 manual path). - * `task` forces one task ignoring its gate; omitted runs all enabled. The - * registered dreamer timer also runs due tasks on its own schedule. */ + * `task` forces one task ignoring its gate; `undefined` runs all enabled. The + * registered dreamer timer also runs due tasks on its own schedule. + * Keep this parameter order stable: registrations are shared across reloads. */ runManual: ( - task?: DreamTaskName, - registrationOwner?: object, + task: DreamTaskName | undefined, + registrationOwner: object, ) => Promise; /** The directory this registration was built for. `resolveProjectIdentity` * is intentionally identical across worktrees/clones of one repo, so a @@ -233,11 +234,10 @@ export function registerPiDreamerProject(opts: PiDreamerOptions): void { // Scheduled runs keep using the timer's client above; binding manual runs // to their owner lets session_shutdown wait only for that instance's work. const runManual = async ( - task?: DreamTaskName, - registrationOwner?: object, + task: DreamTaskName | undefined, + registrationOwner: object, ): Promise => { - const manualOpts = - registrationOwner === undefined ? opts : owners.get(registrationOwner); + const manualOpts = owners.get(registrationOwner); if (!manualOpts) { throw new Error( `Pi dreamer registration owner is no longer active for project ${opts.projectIdentity}`, @@ -303,19 +303,20 @@ export function registerPiDreamerProject(opts: PiDreamerOptions): void { * Run one dream cycle IMMEDIATELY for the given project, mirroring * OpenCode's `/ctx-dream` behavior. Returns the run result, or `null` * if there's nothing to dequeue (queue empty or another worker holds - * the lease — see `processDreamQueue` semantics). Throws if the project - * isn't registered (call `registerPiDreamerProject` first). + * the lease — see `processDreamQueue` semantics). Throws if the project or + * owner isn't registered (call `registerPiDreamerProject` first). * * The user-visible reason this exists: without it, the user types * `/ctx-dream` and gets "queued, the timer will run it eventually" — * which makes the command feel broken even though the queue entry is * really there. Mirroring OpenCode's behavior lets us actually drain - * it on the same turn. + * it on the same turn. The owner is required so same-directory + * re-registration always resolves the current owner options. */ export async function runPiDreamForProject( projectIdentity: string, - task?: DreamTaskName, - registrationOwner?: object, + task: DreamTaskName | undefined, + registrationOwner: object, ): Promise { const registration = registeredProjects.get(projectIdentity); if (!registration) {