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..99213908e 100644 --- a/packages/pi-plugin/src/agent-end-handler.test.ts +++ b/packages/pi-plugin/src/agent-end-handler.test.ts @@ -107,15 +107,36 @@ 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("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", () => { diff --git a/packages/pi-plugin/src/commands/ctx-commands.test.ts b/packages/pi-plugin/src/commands/ctx-commands.test.ts index 31381b87e..9dcb064bb 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"; @@ -259,17 +289,24 @@ 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", + registrationOwner: {}, + 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. + // The injected registration sync runs immediately before runManual. 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"); }); @@ -282,6 +319,7 @@ describe("Pi Magic Context commands", () => { db, projectDir: "/tmp/project", projectIdentity: "/tmp/project", + registrationOwner: {}, }); await handlers.get("ctx-dream")?.("verify", createCtx()); @@ -311,6 +349,7 @@ describe("Pi Magic Context commands", () => { db, projectDir: "/tmp/project", projectIdentity: "/tmp/project", + registrationOwner: {}, dreamerEnabled: false, }); await handlers.get("ctx-dream")?.("", createCtx()); @@ -326,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 27f653eca..ba7c6c7f6 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, @@ -24,11 +24,14 @@ export function registerCtxDreamCommand( dreamerEnabled?: boolean; 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, @@ -43,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}".`, @@ -60,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`).", @@ -83,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: [ @@ -108,9 +108,11 @@ 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, + deps.registrationOwner, ); const lines: string[] = []; if (result.ran.length > 0) lines.push(`Ran: ${result.ran.join(", ")}`); @@ -140,12 +142,12 @@ 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"), level: result.ran.length > 0 ? "success" : "info", + rpcDisplay: "dialog", }, { projectDir: project.projectDir, @@ -155,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 c7d5e7911..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,10 +237,11 @@ export function registerCtxEmbedCommand( sessionId, ); const statusText = formatEmbedStatusText(coverage, { status: "idle" }); - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-embed", text: `## Embedding Status\n\n${statusText}`, level: "info", + rpcDisplay: "dialog", }); }, }); 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 3f81b2bc1..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,16 +101,17 @@ export function registerCtxRecompCommand( } if (parsed.kind === "upgrade") { - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-recomp", text: executeRecompUpgradeStub(currentDeps.db, sessionId), level: "info", + rpcDisplay: "dialog", }); return; } 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", @@ -136,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", @@ -145,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", @@ -154,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", @@ -163,7 +165,7 @@ export function registerCtxRecompCommand( } confirmationBySession.delete(sessionId); - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-recomp", text: parsed.kind === "partial" @@ -206,7 +208,7 @@ export function registerCtxRecompCommand( directory: ctx.cwd, accountingSessionId: sessionId, notify: (text) => { - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-recomp", text, level: inferLevel(text), @@ -268,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 71936fa3c..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", @@ -198,13 +199,14 @@ export function registerCtxSessionUpgradeCommand( : "This session's compartments are already in the current format.", ].join("\n"), level: "info", + rpcDisplay: "dialog", }); return; } // 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", @@ -221,17 +223,18 @@ 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", + rpcDisplay: "dialog", }); }, }); 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", @@ -273,7 +276,7 @@ export function registerCtxSessionUpgradeCommand( { preserveUserQuotes: true }, ), notify: (text) => - sendCtxStatusMessage(pi, { + sendStatus({ title: "/ctx-session-upgrade", text, level: "info", @@ -324,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", @@ -362,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", @@ -375,6 +378,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..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,13 +145,17 @@ export function registerCtxStatusCommand( resolveTailHygieneStatus(getPiChannel1Baseline(sessionId)), ); const details = buildStatusDetails(currentDeps, sessionId); - sendCtxStatusMessage( - pi, - { title: "/ctx-status", text: statusText, level: "info" }, + sendStatus( + { + title: "/ctx-status", + text: statusText, + level: "info", + rpcDisplay: "dialog", + }, 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 e627a14ef..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,6 +6,8 @@ import { type PiMessageSender, registerCtxStatusEntryRenderer, sendCtxStatusMessage, + shouldShowCtxStatusDialog, + showCtxStatusDialog, } from "./pi-command-utils"; describe("ctx-status entries", () => { @@ -91,4 +93,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..a003b3f33 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,8 +41,19 @@ type PiEntryRendererRegistration = { export type PiMessageSender = Pick & PiEntryRendererRegistration; +export function shouldShowCtxStatusDialog( + content: CtxStatusMessageContent, +): boolean { + return ( + content.rpcDisplay === "dialog" || + (content.rpcDisplay !== "notification" && + content.level !== undefined && + content.level !== "info") + ); +} + export function resolveSessionId( - ctx: ExtensionCommandContext, + ctx: Pick, ): string | undefined { const sm = ctx.sessionManager; const getSessionId = (sm as { getSessionId?: () => string | undefined }) @@ -55,6 +67,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": @@ -110,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, @@ -125,6 +213,8 @@ export function sendCtxStatusMessage( if (typeof pi.appendEntry === "function") { pi.appendEntry(CTX_STATUS_CUSTOM_TYPE, 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. sessionLog("pi-status", `${content.title}: ${content.text}`); 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 120f32745..32341b493 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"; @@ -56,6 +61,7 @@ function dreamerOptions(args: { database: Database; projectIdentity: string; projectDir?: string; + registrationOwner?: object; config?: DreamerConfig; language?: string; onAdjunctsRefreshNeeded?: (projectIdentity: string) => void; @@ -66,6 +72,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, @@ -136,6 +143,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; @@ -156,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(() => {})); @@ -175,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, @@ -239,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"]); @@ -295,6 +336,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,42 +387,304 @@ describe("Pi dreamer wiring", () => { }); // Worktree A of the same repo → identity X. + 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). + 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 + // timer started against worktree B. + 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("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: "git:pi-worktree", + projectIdentity, + registrationOwner: owner, }), ); await flushMicrotasks(); - // Worktree B of the SAME repo (same identity, different dir). + 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: "git:pi-worktree", + 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[] = []; + __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"]); - // Still one registration, but rebuilt: first timer torn down, second - // timer started against worktree B. - expect(__test.registeredProjectCount()).toBe(1); - expect(firstCleanup).toHaveBeenCalledTimes(1); - expect(dirs).toEqual(["/tmp/worktree-A", "/tmp/worktree-B"]); + unregisterPiDreamerProject({ + projectIdentity: "git:pi-handoff", + registrationOwner: firstA.registrationOwner, + }); + expect(__test.registeredProjectCount()).toBe(0); }); - test("unregister removes the project", () => { + 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("rejects ownerless and unregistered-owner manual runs", 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(); + __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, + 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, - projectDir: "/tmp/pi-project-unregister", - projectIdentity: "git:pi-unregister", + projectIdentity, + registrationOwner: owner, + config: DreamerConfigSchema.parse({ + model: "test/model", + tasks: { curate: { schedule: "0 4 * * *" } }, + }), }), ); - unregisterPiDreamerProject({ projectIdentity: "git:pi-unregister" }); + 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({ + database: db, + projectDir: "/tmp/pi-project-unregister", + projectIdentity: "git:pi-unregister", + }); + registerPiDreamerProject(opts); + + unregisterPiDreamerProject({ + projectIdentity: "git:pi-unregister", + registrationOwner: opts.registrationOwner, + }); expect(__test.registeredProjectCount()).toBe(0); }); @@ -351,6 +693,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; @@ -387,6 +792,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; @@ -498,10 +947,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 +970,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..c262136fd 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,10 +85,16 @@ 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. */ - runManual: (task?: DreamTaskName) => Promise; + * `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 | undefined, + 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 @@ -105,9 +113,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(); @@ -122,8 +155,21 @@ 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); + 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 + // deregister the shared timer while it is still active. if (existing.projectDir === opts.projectDir) { return; } @@ -138,10 +184,14 @@ 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); + const client = createPiDreamerClient( + opts, + notifyOwnersOfAdjunctRefresh, + () => owners.get(opts.registrationOwner)?.projectDir === opts.projectDir, + ); let cleanup: (() => void) | undefined; let cancelled = false; @@ -177,38 +227,69 @@ 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 | undefined, + registrationOwner: object, + ): Promise => { + const manualOpts = 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( - 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, }); + // 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, { + activeOwner: opts.registrationOwner, + owners, cleanup: () => { cancelled = true; cleanup?.(); @@ -222,18 +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, + task: DreamTaskName | undefined, + registrationOwner: object, ): Promise { const registration = registeredProjects.get(projectIdentity); if (!registration) { @@ -241,15 +324,40 @@ 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 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 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); + registerPiDreamerProject(replacementOptions); + const replacement = registeredProjects.get(opts.projectIdentity); + if (replacement) { + for (const remainingOptions of remaining) { + replacement.owners.set( + remainingOptions.registrationOwner, + remainingOptions, + ); + } + } return; } @@ -257,18 +365,28 @@ 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 { +function createPiDreamerClient( + opts: PiDreamerOptions, + onAdjunctsRefreshNeeded = opts.onAdjunctsRefreshNeeded, + isRegistrationOwnerActive: () => boolean = () => true, +): DreamTimerClient { const runner = piSubagentRunnerFactory(); const session = { @@ -290,6 +408,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 @@ -318,7 +442,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) { @@ -346,7 +470,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-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..c4ee6669a 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,160 @@ 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); + const second = createCountingPi(); + await magicContextPiExtension(second.pi); + 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("unsubscribes child lifecycle listeners on session shutdown", async () => { + isolateXdgEnv(); + delete process.env[MAGIC_CONTEXT_PI_SUBAGENT_ENV]; + + 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); + + it("skips only the marked in-process child", async () => { + isolateXdgEnv(); + delete process.env[MAGIC_CONTEXT_PI_SUBAGENT_ENV]; + + const parent = createCountingPi(); + await magicContextPiExtension(parent.pi); + parent.emitEvent("subagents:child:spawning"); + parent.emitEvent("subagents:child:session-created"); // 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); + const child = createCountingPi(); + await magicContextPiExtension(child.pi); + expect(child.events).toEqual([]); + expect(child.tools).toEqual([]); + expect(child.commands).toEqual([]); - expect(second.events).toEqual([]); - expect(second.tools).toEqual([]); - expect(second.flags).toEqual([]); - expect(second.commands).toEqual([]); - expect(second.entryRenderers).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("clearing the latch (dispose) allows a full re-init", async () => { + it("does not suppress an independent session while a child marker is active", 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 parent = createCountingPi(); + await magicContextPiExtension(parent.pi); - // Simulate the session_shutdown dispose path clearing the latch. - __test.clearPiMagicContextActive(); - expect(__test.isPiMagicContextActiveInProcess()).toBe(false); + 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"); + }); - // A subsequent init re-registers the full runtime. - const second = createCountingPi(); - await magicContextPiExtension(second.pi); + await marked; + try { + expect(child.tools).toEqual([]); - expect(second.events.length).toBeGreaterThan(0); - expect(second.tools.length).toBeGreaterThan(0); - expect(second.commands.length).toBeGreaterThan(0); - expect(second.entryRenderers).toEqual(["ctx-status"]); + 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("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 +239,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..7a55dbbe1 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"; @@ -119,6 +120,7 @@ import { registerCtxStatusCommand } from "./commands/ctx-status"; import { registerCtxWrapupCommand } from "./commands/ctx-wrapup"; import { registerCtxStatusEntryRenderer, + resolveSessionId, sendCtxStatusMessage, } from "./commands/pi-command-utils"; import { loadPiConfig } from "./config"; @@ -184,7 +186,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 +199,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 +496,9 @@ export const __test = { resetLoggedPiConfigDirs(): void { loggedPiConfigDirs.clear(); }, - isPiMagicContextActiveInProcess, - markPiMagicContextActive, - clearPiMagicContextActive, + clearPiInProcessSubagentInitContext, + claimPiStartupMaintenance, + clearPiStartupMaintenanceClaim, }; function formatTokens(value: number): string { @@ -732,18 +763,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 +892,70 @@ 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 = {}; + 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 @@ -1129,6 +1190,38 @@ async function startPiMagicContextRuntime( const bootProjectDeps = buildProjectDeps(projectDir, projectIdentity, config); projectDepsByDir.set(projectDir, bootProjectDeps); + + function syncDreamerProjectRegistration( + current: ResolvedPiProjectDeps, + ): void { + if (sessionShuttingDown) return; + 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; @@ -1438,6 +1531,9 @@ async function startPiMagicContextRuntime( resolveDreamerEnabled: (ctx) => resolveCurrentProjectDeps(ctx).dreamerEnabled, onProjectSeen: (identity) => seenDreamerProjectIdentities.add(identity), + ensureRegistered: (ctx) => + syncDreamerProjectRegistration(resolveCurrentProjectDeps(ctx)), + registrationOwner: dreamerRegistrationOwner, }); info("registered /ctx-dream"); @@ -1464,23 +1560,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 +1635,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 +1649,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 @@ -2248,9 +2305,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()` @@ -2262,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 @@ -2269,36 +2327,51 @@ 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; + 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 { - await withTimeout(awaitInFlightHistorians(), SHUTDOWN_DRAIN_MS); + for (const identity of seenDreamerProjectIdentities) { + unregisterPiDreamerProject({ + projectIdentity: identity, + registrationOwner: dreamerRegistrationOwner, + }); + } } catch (err) { - warn("shutdown: historian drain threw:", err); + warn("shutdown: unregisterPiDreamerProject threw:", err); } - try { - await withTimeout(awaitInFlightRecomps(), SHUTDOWN_DRAIN_MS); - } catch (err) { - warn("shutdown: recomp drain threw:", err); + 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); } - try { - for (const identity of seenDreamerProjectIdentities) { - unregisterPiDreamerProject({ projectIdentity: identity }); - } - } 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 @@ -2325,16 +2398,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 +2438,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/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/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..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, @@ -202,6 +203,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 +236,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({ @@ -1061,6 +1072,55 @@ describe("PiSubagentRunner spawn lifecycle", () => { // 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(); const spawnImpl = mock(() => child as never); const { PiSubagentRunner } = await import("./subagent-runner"); @@ -1084,17 +1144,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 +2534,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 +2561,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");