diff --git a/packages/pi-plugin/PARITY.md b/packages/pi-plugin/PARITY.md index 25a14ae10..a52958954 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,14 @@ 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. RPC hosts that execute the `ctx.ui.custom` component factory (such as +pi-web) present detailed results as dialogs; hosts where `custom` resolves without +executing the factory receive the same details through a notification fallback. +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 +168,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 +398,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..47b9f174c 100644 --- a/packages/pi-plugin/src/agent-end-handler.test.ts +++ b/packages/pi-plugin/src/agent-end-handler.test.ts @@ -107,15 +107,42 @@ 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("aborts a recomp that outlives the graceful drain before shutdown returns", () => { + const drainAt = body.indexOf("awaitInFlightRecomps(sessionId)"); + const abortAt = body.indexOf("abortInFlightRecomps(sessionId)"); + expect(abortAt).toBeGreaterThan(drainAt); + }); + + 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..6e87ea0a1 100644 --- a/packages/pi-plugin/src/commands/ctx-commands.test.ts +++ b/packages/pi-plugin/src/commands/ctx-commands.test.ts @@ -1,12 +1,28 @@ import { describe, expect, it } from "bun:test"; import { replaceAllCompartmentState } from "@magic-context/core/features/magic-context/compartment-storage"; +import { isMemoryMigrationDone } from "@magic-context/core/features/magic-context/memory/memory-migration"; +import { resolveProjectIdentity } from "@magic-context/core/features/magic-context/memory/project-identity"; +import { + getMemoriesByProject, + insertMemory, +} from "@magic-context/core/features/magic-context/memory/storage-memory"; import { runMigrations } from "@magic-context/core/features/magic-context/migrations"; +import { + getPendingPiCompactionMarkerState, + insertTag, +} from "@magic-context/core/features/magic-context/storage"; import { initializeDatabase } from "@magic-context/core/features/magic-context/storage-db"; import { queuePendingOp } from "@magic-context/core/features/magic-context/storage-ops"; -import { insertTag } from "@magic-context/core/features/magic-context/storage-tags"; import { Database } from "@magic-context/core/shared/sqlite"; -import { awaitInFlightRecomps } from "../pi-recomp-runner"; +import { + consumeDeferredHistoryRefresh, + consumeDeferredMaterialization, +} from "../context-handler"; +import { + abortInFlightRecomps, + awaitInFlightRecomps, +} from "../pi-recomp-runner"; import { registerCtxDreamCommand } from "./ctx-dream"; import { registerCtxFlushCommand } from "./ctx-flush"; import { registerCtxRecompCommand } from "./ctx-recomp"; @@ -29,8 +45,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?: { @@ -85,7 +103,10 @@ function createCtx(sessionId = "ses-1"): MockCommandContext { type: "message", message: { role: index % 2 === 0 ? "user" : "assistant", - content: `message ${index + 1}`, + content: + index % 2 === 0 + ? `message ${index + 1}` + : [{ type: "text", text: `message ${index + 1}` }], }, })); return { @@ -108,6 +129,76 @@ function createCtx(sessionId = "ses-1"): MockCommandContext { }; } +function deferred() { + let resolve!: () => void; + const promise = new Promise((done) => { + resolve = done; + }); + return { promise, resolve }; +} + +function validCompartmentForPrompt(prompt: string): string { + const ordinals = [...prompt.matchAll(/^\[(\d+)\] [UAT]:/gm)].map((match) => + Number(match[1]), + ); + const start = ordinals[0]; + const end = ordinals.at(-1); + if (start === undefined || end === undefined) { + throw new Error("expected tagged transcript ordinals in historian prompt"); + } + return `Covered messages ${start}-${end}.`; +} + +function probeLiveCommandContext(ctx: MockCommandContext): { + endLifecycle(): void; + lateAccesses(): number; +} { + let live = true; + let cwd = ctx.cwd; + let model = ctx.model; + let lateAccessCount = 0; + const guard = () => { + if (!live) lateAccessCount += 1; + }; + Object.defineProperty(ctx, "cwd", { + configurable: true, + get: () => { + guard(); + return cwd; + }, + }); + Object.defineProperty(ctx, "model", { + configurable: true, + get: () => { + guard(); + return model; + }, + }); + const custom = ctx.ui.custom; + ctx.ui.custom = (...args) => { + guard(); + return custom(...args); + }; + const notify = ctx.ui.notify; + ctx.ui.notify = (...args) => { + guard(); + notify?.(...args); + }; + const setStatus = ctx.ui.setStatus; + ctx.ui.setStatus = (...args) => { + guard(); + setStatus?.(...args); + }; + return { + endLifecycle() { + live = false; + cwd = "/tmp/reused-session-context"; + model = { provider: "other", id: "replacement" }; + }, + lateAccesses: () => lateAccessCount, + }; +} + describe("Pi Magic Context commands", () => { it("registers /ctx-status and opens a UI overlay when UI is available", async () => { const db = createDb(); @@ -152,6 +243,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 +378,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 +408,7 @@ describe("Pi Magic Context commands", () => { db, projectDir: "/tmp/project", projectIdentity: "/tmp/project", + registrationOwner: {}, }); await handlers.get("ctx-dream")?.("verify", createCtx()); @@ -311,6 +438,7 @@ describe("Pi Magic Context commands", () => { db, projectDir: "/tmp/project", projectIdentity: "/tmp/project", + registrationOwner: {}, dreamerEnabled: false, }); await handlers.get("ctx-dream")?.("", createCtx()); @@ -326,6 +454,7 @@ describe("Pi Magic Context commands", () => { db, projectDir: "/tmp/project-a", projectIdentity: "/tmp/project-a", + registrationOwner: {}, resolveProject: (ctx) => ({ projectDir: ctx.cwd, projectIdentity: ctx.cwd, @@ -531,6 +660,13 @@ describe("Pi Magic Context commands", () => { }); const ctx = createCtx("ses-budget"); + const branch = ctx.sessionManager.getBranch?.() as Array<{ + message: { content: unknown }; + }>; + const firstMessage = branch[0]; + if (!firstMessage) throw new Error("expected a first message fixture"); + firstMessage.message.content = `message 1 ${"word ".repeat(100)}`; + ctx.sessionManager.getBranch = () => branch; // First call shows the confirmation warning; second call confirms and // spawns the DETACHED recomp. The recomp now runs in the background // (parity with OpenCode), so await the in-flight run before asserting @@ -542,4 +678,250 @@ describe("Pi Magic Context commands", () => { expect(promptText).toContain("[1] U: message 1"); expect(promptText).not.toContain("[2] A: message 2"); }); + + it("control: a valid /ctx-recomp publishes and stages deferred effects", async () => { + const sessionId = "ses-recomp-control"; + const db = createDb(); + const { pi, handlers, sent } = createMockPi(); + registerCtxRecompCommand(pi as never, { + db, + runner: { + run: async (args) => ({ + ok: true as const, + assistantText: validCompartmentForPrompt(args.userMessage), + cost: 0, + durationMs: 1, + }), + }, + historianModel: "anthropic/claude", + historianChunkTokens: 100_000, + memoryEnabled: false, + autoPromote: false, + }); + + const ctx = createCtx(sessionId); + await handlers.get("ctx-recomp")?.("", ctx); + await handlers.get("ctx-recomp")?.("", ctx); + await awaitInFlightRecomps(sessionId); + + expect(getPendingPiCompactionMarkerState(db, sessionId)).not.toBeNull(); + expect(consumeDeferredHistoryRefresh(sessionId)).toBe(true); + expect(consumeDeferredMaterialization(sessionId)).toBe(true); + expect( + sent.some((entry) => entry.data.text.includes("Magic Recomp — Complete")), + ).toBe(true); + }); + + it("fences /ctx-recomp side effects when an aborted runner settles late", async () => { + const sessionId = "ses-recomp-late"; + const db = createDb(); + const { pi, handlers, sent } = createMockPi(); + const runStarted = deferred(); + const releaseRun = deferred(); + let observedSignal: AbortSignal | undefined; + let observedDirectory: string | undefined; + registerCtxRecompCommand(pi as never, { + db, + runner: { + run: async (args) => { + observedSignal = args.signal; + observedDirectory = args.cwd; + runStarted.resolve(); + await releaseRun.promise; + return { + ok: true as const, + assistantText: validCompartmentForPrompt(args.userMessage), + cost: 0, + durationMs: 1, + }; + }, + }, + historianModel: "anthropic/claude", + historianChunkTokens: 20, + memoryEnabled: false, + autoPromote: false, + }); + + const ctx = createCtx(sessionId); + const contextProbe = probeLiveCommandContext(ctx); + let branchReads = 0; + const getBranch = ctx.sessionManager.getBranch; + ctx.sessionManager.getBranch = () => { + branchReads += 1; + if (branchReads > 1) throw new Error("late session access"); + return getBranch?.() ?? []; + }; + await handlers.get("ctx-recomp")?.("", ctx); + await handlers.get("ctx-recomp")?.("", ctx); + await runStarted.promise; + const sentBeforeAbort = sent.length; + + abortInFlightRecomps(sessionId); + contextProbe.endLifecycle(); + releaseRun.resolve(); + await awaitInFlightRecomps(sessionId); + + expect(observedSignal?.aborted).toBe(true); + expect(observedDirectory).toBe("/tmp/project"); + expect(contextProbe.lateAccesses()).toBe(0); + expect(branchReads).toBe(1); + expect(sent).toHaveLength(sentBeforeAbort); + expect(getPendingPiCompactionMarkerState(db, sessionId)).toBeNull(); + expect(consumeDeferredHistoryRefresh(sessionId)).toBe(false); + expect(consumeDeferredMaterialization(sessionId)).toBe(false); + }); + + it("fences /ctx-session-upgrade side effects when an aborted runner settles late", async () => { + const sessionId = "ses-upgrade-late"; + const db = createDb(); + replaceAllCompartmentState( + db, + sessionId, + [ + { + sequence: 1, + startMessage: 1, + endMessage: 2, + startMessageId: "m1", + endMessageId: "m2", + title: "legacy", + content: "legacy content", + }, + ], + [], + ); + const { pi, handlers, sent } = createMockPi(); + const runStarted = deferred(); + const releaseRun = deferred(); + let observedSignal: AbortSignal | undefined; + let observedDirectory: string | undefined; + let runnerCalls = 0; + + registerCtxSessionUpgradeCommand(pi as never, { + db, + runner: { + run: async (args) => { + runnerCalls += 1; + observedSignal = args.signal; + observedDirectory = args.cwd; + runStarted.resolve(); + await releaseRun.promise; + return { + ok: true as const, + assistantText: validCompartmentForPrompt(args.userMessage), + cost: 0, + durationMs: 1, + }; + }, + }, + historianModel: "anthropic/claude", + historianChunkTokens: 20, + memoryEnabled: true, + autoPromote: false, + }); + + const ctx = createCtx(sessionId); + const contextProbe = probeLiveCommandContext(ctx); + let branchReads = 0; + const getBranch = ctx.sessionManager.getBranch; + ctx.sessionManager.getBranch = () => { + branchReads += 1; + if (branchReads > 1) throw new Error("late session access"); + return getBranch?.() ?? []; + }; + await handlers.get("ctx-session-upgrade")?.("", ctx); + await runStarted.promise; + const sentBeforeAbort = sent.length; + + abortInFlightRecomps(sessionId); + contextProbe.endLifecycle(); + releaseRun.resolve(); + await awaitInFlightRecomps(sessionId); + + expect(observedSignal?.aborted).toBe(true); + expect(observedDirectory).toBe("/tmp/project"); + expect(contextProbe.lateAccesses()).toBe(0); + expect(runnerCalls).toBe(1); + expect(branchReads).toBe(1); + expect(sent).toHaveLength(sentBeforeAbort); + expect(getPendingPiCompactionMarkerState(db, sessionId)).toBeNull(); + expect(consumeDeferredHistoryRefresh(sessionId)).toBe(false); + expect(consumeDeferredMaterialization(sessionId)).toBe(false); + }); + + it("keeps migration-only upgrade state unchanged after a late cancelled result", async () => { + const sessionId = "ses-upgrade-migration-late"; + const db = createDb(); + const { pi, handlers, sent } = createMockPi(); + const ctx = createCtx(sessionId); + const projectPath = resolveProjectIdentity(ctx.cwd); + insertMemory(db, { + projectPath, + category: "ARCHITECTURE_DECISIONS", + content: "Legacy migration fixture.", + }); + const before = getMemoriesByProject(db, projectPath).map((memory) => ({ + id: memory.id, + category: memory.category, + content: memory.content, + })); + const runStarted = deferred(); + const releaseRun = deferred(); + let observedSignal: AbortSignal | undefined; + registerCtxSessionUpgradeCommand(pi as never, { + db, + runner: { + run: async (args) => { + observedSignal = args.signal; + runStarted.resolve(); + await releaseRun.promise; + return { + ok: true as const, + assistantText: [ + "", + "", + "* Migrated replacement.", + "", + "", + ].join("\n"), + cost: 0, + durationMs: 1, + }; + }, + }, + historianModel: "anthropic/claude", + historianChunkTokens: 100_000, + memoryEnabled: true, + autoPromote: false, + }); + const contextProbe = probeLiveCommandContext(ctx); + let branchReads = 0; + const getBranch = ctx.sessionManager.getBranch; + ctx.sessionManager.getBranch = () => { + branchReads += 1; + if (branchReads > 1) throw new Error("late session access"); + return getBranch?.() ?? []; + }; + + await handlers.get("ctx-session-upgrade")?.("", ctx); + await runStarted.promise; + const sentBeforeAbort = sent.length; + abortInFlightRecomps(sessionId); + contextProbe.endLifecycle(); + releaseRun.resolve(); + await awaitInFlightRecomps(sessionId); + + expect(observedSignal?.aborted).toBe(true); + expect(branchReads).toBe(1); + expect(contextProbe.lateAccesses()).toBe(0); + expect(sent).toHaveLength(sentBeforeAbort); + expect( + getMemoriesByProject(db, projectPath).map((memory) => ({ + id: memory.id, + category: memory.category, + content: memory.content, + })), + ).toEqual(before); + expect(isMemoryMigrationDone(db, projectPath)).toBe(false); + }); }); 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-signals.test.ts b/packages/pi-plugin/src/commands/ctx-recomp-signals.test.ts index 295011706..67666eb6e 100644 --- a/packages/pi-plugin/src/commands/ctx-recomp-signals.test.ts +++ b/packages/pi-plugin/src/commands/ctx-recomp-signals.test.ts @@ -73,4 +73,13 @@ describe("/ctx-recomp post-completion signal contract", () => { ); expect(clearCall).toBeGreaterThan(publishedGate); }); + + test("captures session data before detached work and threads its abort signal", () => { + expect(codeOnly).toContain("const snapshot = readPiSessionSnapshot(ctx)"); + expect(codeOnly).toContain("readMessages: () => snapshot.rawMessages"); + expect(codeOnly).toContain("branchEntries: snapshot.branchEntries"); + expect(codeOnly).toContain("work: async (signal)"); + expect(codeOnly).toContain("signal,"); + expect(codeOnly).not.toContain("readPiSessionMessages(ctx)"); + }); }); diff --git a/packages/pi-plugin/src/commands/ctx-recomp.ts b/packages/pi-plugin/src/commands/ctx-recomp.ts index 3f81b2bc1..f04b30cc2 100644 --- a/packages/pi-plugin/src/commands/ctx-recomp.ts +++ b/packages/pi-plugin/src/commands/ctx-recomp.ts @@ -26,9 +26,9 @@ import { ensureProjectRegisteredFromPiDirectory } from "../embedding-bootstrap"; import { createPiHistorianClient } from "../pi-recomp-client-shared"; import { stagePiRecompMarker } from "../pi-recomp-marker"; import { isPiRecompInFlight, spawnPiRecompRun } from "../pi-recomp-runner"; -import { readPiSessionMessages } from "../read-session-pi"; +import { readPiSessionSnapshot } 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" @@ -172,8 +174,13 @@ export function registerCtxRecompCommand( level: "info", }); + const cwd = ctx.cwd; + const fallbackModelId = ctx.model + ? `${ctx.model.provider}/${ctx.model.id}` + : undefined; + const snapshot = readPiSessionSnapshot(ctx); const provider = { - readMessages: () => readPiSessionMessages(ctx), + readMessages: () => snapshot.rawMessages, } satisfies RawMessageProvider; // Detached: the recomp runs in the background so the Pi REPL stays @@ -187,9 +194,10 @@ export function registerCtxRecompCommand( onStatusChange: () => updateStatusLine(ctx, { db: currentDeps.db, - projectIdentity: ctx.cwd, + projectIdentity: cwd, }), - work: async () => { + work: async (signal) => { + const detachedSendStatus = createCtxStatusSender(pi, ctx, signal); const result = await executeContextRecompWithResult( { client: createPiHistorianClient({ @@ -203,10 +211,11 @@ export function registerCtxRecompCommand( fallbackModels: currentDeps.historianFallbacks, timeoutMs: currentDeps.historianTimeoutMs, thinkingLevel: currentDeps.historianThinkingLevel, - directory: ctx.cwd, + directory: cwd, accountingSessionId: sessionId, + signal, notify: (text) => { - sendCtxStatusMessage(pi, { + detachedSendStatus({ title: "/ctx-recomp", text, level: inferLevel(text), @@ -216,7 +225,7 @@ export function registerCtxRecompCommand( db: currentDeps.db, sessionId, historianChunkTokens: currentDeps.historianChunkTokens, - directory: ctx.cwd, + directory: cwd, historianTimeoutMs: currentDeps.historianTimeoutMs, memoryEnabled: currentDeps.memoryEnabled, autoPromote: currentDeps.autoPromote, @@ -228,12 +237,11 @@ export function registerCtxRecompCommand( // fallbacks + the session's own model as last-ditch retry. fallbackModels: currentDeps.historianFallbacks, language: currentDeps.language, - fallbackModelId: ctx.model - ? `${ctx.model.provider}/${ctx.model.id}` - : undefined, + fallbackModelId, }, parsed.kind === "partial" ? { range: parsed.range } : {}, ); + if (signal.aborted) return; if (result.published) { // A successful recomp resolves the overflow that may have armed // needs_emergency_recovery — clear it so the flag stops force- @@ -258,7 +266,11 @@ export function registerCtxRecompCommand( // mid-turn, busting the cache. Mirrors the background // historian's onPublished (signalPiDeferred*). try { - stagePiRecompMarker({ db: currentDeps.db, sessionId, ctx }); + stagePiRecompMarker({ + db: currentDeps.db, + sessionId, + branchEntries: snapshot.branchEntries, + }); } catch (markerError) { sessionLog( sessionId, @@ -268,7 +280,7 @@ export function registerCtxRecompCommand( signalPiDeferredHistoryRefresh(sessionId); signalPiDeferredMaterialization(sessionId); } - sendCtxStatusMessage(pi, { + detachedSendStatus({ title: "/ctx-recomp", text: result.message, level: inferLevel(result.message), diff --git a/packages/pi-plugin/src/commands/ctx-session-upgrade-signals.test.ts b/packages/pi-plugin/src/commands/ctx-session-upgrade-signals.test.ts index 7c76ad2f0..a21dad910 100644 --- a/packages/pi-plugin/src/commands/ctx-session-upgrade-signals.test.ts +++ b/packages/pi-plugin/src/commands/ctx-session-upgrade-signals.test.ts @@ -52,4 +52,13 @@ describe("/ctx-session-upgrade detached execution contract", () => { expect(codeOnly).toContain("isRecompComplete(recompResult.message)"); expect(codeOnly).toContain("!recompResult.published"); }); + + test("captures snapshots and threads cancellation through recomp and migration", () => { + expect(codeOnly).toContain("readPiSessionSnapshot(ctx)"); + expect(codeOnly).toContain("readMessages: () => snapshot.rawMessages"); + expect(codeOnly).toContain("branchEntries: snapshot.branchEntries"); + expect(codeOnly).toContain("runMigration(signal)"); + expect(codeOnly).toContain("work: async (signal)"); + expect(codeOnly).not.toContain("readPiSessionMessages(ctx)"); + }); }); diff --git a/packages/pi-plugin/src/commands/ctx-session-upgrade.ts b/packages/pi-plugin/src/commands/ctx-session-upgrade.ts index 71936fa3c..0d9877833 100644 --- a/packages/pi-plugin/src/commands/ctx-session-upgrade.ts +++ b/packages/pi-plugin/src/commands/ctx-session-upgrade.ts @@ -28,9 +28,9 @@ import { runPiMemoryMigration } from "../pi-memory-migration"; import { createPiHistorianClient } from "../pi-recomp-client-shared"; import { stagePiRecompMarker } from "../pi-recomp-marker"; import { isPiRecompInFlight, spawnPiRecompRun } from "../pi-recomp-runner"; -import { readPiSessionMessages } from "../read-session-pi"; +import { readPiSessionSnapshot } 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", @@ -133,6 +134,7 @@ export function registerCtxSessionUpgradeCommand( // OpenCode's primaryModelId): a quality-sensitive consolidation should // run on the user's working model, not the (possibly misconfigured) // historian model. Historian model + fallbacks remain the safety net. + const cwd = ctx.cwd; const sessionMainModel = ctx.model ? `${ctx.model.provider}/${ctx.model.id}` : undefined; @@ -145,7 +147,7 @@ export function registerCtxSessionUpgradeCommand( // could touch a pool the user opted out of at worst. const migrationEnabled = currentDeps.memoryEnabled; - const runMigration = async (): Promise => { + const runMigration = async (signal: AbortSignal): Promise => { if (!migrationEnabled) { return "Memory migration skipped (memory disabled)."; } @@ -160,14 +162,16 @@ export function registerCtxSessionUpgradeCommand( fallbackModels: currentDeps.historianFallbacks, timeoutMs: currentDeps.historianTimeoutMs, thinkingLevel: currentDeps.historianThinkingLevel, - directory: ctx.cwd, + directory: cwd, allowHomeProject: currentDeps.allowHomeProject, sessionId, userMemoriesEnabled: currentDeps.userMemoriesEnabled, language: currentDeps.language, + signal, }); return outcome.summary; } catch (error) { + if (signal.aborted) return "Memory migration cancelled."; return `Memory migration skipped (error): ${describeError(error).brief}`; } }; @@ -178,7 +182,7 @@ export function registerCtxSessionUpgradeCommand( // • none + migration still pending → migration only (skip recomp) if (upgradableCount === 0) { const projectPath = resolveProjectIdentityForSession( - ctx.cwd, + cwd, currentDeps.allowHomeProject, ); if (!projectPath) return; @@ -188,7 +192,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,47 +202,53 @@ 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", }); + const snapshot = readPiSessionSnapshot(ctx); spawnPiRecompRun({ sessionId, provider: { - readMessages: () => readPiSessionMessages(ctx), + readMessages: () => snapshot.rawMessages, } satisfies RawMessageProvider, onStatusChange: () => updateStatusLine(ctx, { db: currentDeps.db, - projectIdentity: ctx.cwd, + projectIdentity: cwd, }), - work: async () => { - const summary = await runMigration(); - sendCtxStatusMessage(pi, { + work: async (signal) => { + const detachedSendStatus = createCtxStatusSender(pi, ctx, signal); + const summary = await runMigration(signal); + if (signal.aborted) return; + detachedSendStatus({ 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", }); + const snapshot = readPiSessionSnapshot(ctx); const provider = { - readMessages: () => readPiSessionMessages(ctx), + readMessages: () => snapshot.rawMessages, } satisfies RawMessageProvider; // Detached: the upgrade (multi-pass recomp + memory migration) runs in @@ -253,9 +263,10 @@ export function registerCtxSessionUpgradeCommand( onStatusChange: () => updateStatusLine(ctx, { db: currentDeps.db, - projectIdentity: ctx.cwd, + projectIdentity: cwd, }), - work: async () => { + work: async (signal) => { + const detachedSendStatus = createCtxStatusSender(pi, ctx, signal); // Step 1 — compartment upgrade via full recomp. const recompResult = await executeContextRecompWithResult( { @@ -265,15 +276,16 @@ export function registerCtxSessionUpgradeCommand( fallbackModels: currentDeps.historianFallbacks, timeoutMs: currentDeps.historianTimeoutMs, thinkingLevel: currentDeps.historianThinkingLevel, - directory: ctx.cwd, + directory: cwd, accountingSessionId: sessionId, + signal, systemPrompt: withContentLanguageDirective( COMPARTMENT_STRUCTURAL_SYSTEM_PROMPT, currentDeps.language, { preserveUserQuotes: true }, ), notify: (text) => - sendCtxStatusMessage(pi, { + detachedSendStatus({ title: "/ctx-session-upgrade", text, level: "info", @@ -282,7 +294,7 @@ export function registerCtxSessionUpgradeCommand( db: currentDeps.db, sessionId, historianChunkTokens: currentDeps.historianChunkTokens, - directory: ctx.cwd, + directory: cwd, historianTimeoutMs: currentDeps.historianTimeoutMs, memoryEnabled: currentDeps.memoryEnabled, autoPromote: currentDeps.autoPromote, @@ -300,6 +312,7 @@ export function registerCtxSessionUpgradeCommand( }, {}, ); + if (signal.aborted) return; // Gate migration + "Complete" on `published` — the GROUND TRUTH // that recomp actually rebuilt compartments (parity with OpenCode @@ -324,7 +337,7 @@ export function registerCtxSessionUpgradeCommand( ? extractRecompReason(recompResult.message) : `Compartments were not fully rebuilt: ${extractRecompReason(recompResult.message)}`, ); - sendCtxStatusMessage(pi, { + detachedSendStatus({ title: "/ctx-session-upgrade", text: `## Session Upgrade — Incomplete\n\n${reason}`, level: "error", @@ -348,7 +361,11 @@ export function registerCtxSessionUpgradeCommand( // migration, or the "Complete" message below — recomp already // published. try { - stagePiRecompMarker({ db: currentDeps.db, sessionId, ctx }); + stagePiRecompMarker({ + db: currentDeps.db, + sessionId, + branchEntries: snapshot.branchEntries, + }); } catch (markerError) { sessionLog( sessionId, @@ -360,9 +377,10 @@ export function registerCtxSessionUpgradeCommand( signalPiDeferredMaterialization(sessionId); // Step 2 — memory migration (once per project, idempotent). - const migrationSummary = await runMigration(); + const migrationSummary = await runMigration(signal); + if (signal.aborted) return; - sendCtxStatusMessage(pi, { + detachedSendStatus({ title: "/ctx-session-upgrade", text: [ "## Session Upgrade — Complete", @@ -375,6 +393,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..3a6780245 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,217 @@ 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"); + }, + }, + }; + + const shown = await showCtxStatusDialog(ctx as never, { + title: "/ctx-flush", + text: "## /ctx-flush\n\nDetailed result", + level: "success", + }); + expect(shown).toBe(true); + + expect(rendered.join("\n")).toContain("Detailed result"); + expect(closed).toBe(true); + expect(options).toEqual({ + overlay: true, + overlayOptions: { anchor: "center", width: 92 }, + }); + }); + + it("falls back when Pi RPC resolves custom without invoking its factory", async () => { + const notifications: string[] = []; + const ctx = { + mode: "rpc", + ui: { + custom: async () => undefined, + notify: (text: string) => notifications.push(text), + }, + }; + expect( + await showCtxStatusDialog(ctx as never, { + title: "/ctx-status", + text: "Detailed status", + rpcDisplay: "dialog", + }), + ).toBe(false); + + sendCtxStatusMessage( + { appendEntry() {} } as never, + { + title: "/ctx-status", + text: "Detailed status", + rpcDisplay: "dialog", + }, + undefined, + ctx as never, + ); + await new Promise((resolve) => setTimeout(resolve, 0)); + expect(notifications).toEqual(["Detailed status"]); + }); + + it("does not duplicate a dialog as a notification on a custom-capable host", async () => { + let factories = 0; + let notifications = 0; + const ctx = { + mode: "rpc", + ui: { + custom: async (factory: (...args: unknown[]) => unknown) => { + factories += 1; + factory( + {}, + { + fg: (_name: string, text: string) => text, + bold: (text: string) => text, + }, + {}, + () => {}, + ); + }, + notify: () => { + notifications += 1; + }, + }, + }; + sendCtxStatusMessage( + { appendEntry() {} } as never, + { title: "/ctx-status", text: "Detailed", rpcDisplay: "dialog" }, + undefined, + ctx as never, + ); + await new Promise((resolve) => setTimeout(resolve, 0)); + expect(factories).toBe(1); + expect(notifications).toBe(0); + }); + + it("falls back to a notification when custom rejects", async () => { + const notifications: string[] = []; + const ctx = { + mode: "rpc", + ui: { + custom: async () => { + throw new Error("unsupported"); + }, + notify: (text: string) => notifications.push(text), + }, + }; + sendCtxStatusMessage( + { appendEntry() {} } as never, + { title: "/ctx-status", text: "Detailed", rpcDisplay: "dialog" }, + undefined, + ctx as never, + ); + await new Promise((resolve) => setTimeout(resolve, 0)); + expect(notifications).toEqual(["Detailed"]); + }); + + it("suppresses a late custom failure fallback after lifecycle abort", async () => { + let rejectCustom!: (error: Error) => void; + let notifications = 0; + const controller = new AbortController(); + const ctx = { + mode: "rpc", + ui: { + custom: () => + new Promise((_resolve, reject) => { + rejectCustom = reject; + }), + notify: () => { + notifications += 1; + }, + }, + }; + sendCtxStatusMessage( + { appendEntry() {} } as never, + { title: "/ctx-status", text: "Detailed", rpcDisplay: "dialog" }, + undefined, + ctx as never, + controller.signal, + ); + controller.abort(); + rejectCustom(new Error("host disposed")); + await new Promise((resolve) => setTimeout(resolve, 0)); + expect(notifications).toBe(0); + }); + + it("suppresses a late no-op custom fallback after lifecycle abort", async () => { + let resolveCustom!: (value: undefined) => void; + let notifications = 0; + const controller = new AbortController(); + const ctx = { + mode: "rpc", + ui: { + custom: () => + new Promise((resolve) => { + resolveCustom = resolve; + }), + notify: () => { + notifications += 1; + }, + }, + }; + sendCtxStatusMessage( + { appendEntry() {} } as never, + { title: "/ctx-status", text: "Detailed", rpcDisplay: "dialog" }, + undefined, + ctx as never, + controller.signal, + ); + controller.abort(); + resolveCustom(undefined); + await new Promise((resolve) => setTimeout(resolve, 0)); + expect(notifications).toBe(0); + }); }); diff --git a/packages/pi-plugin/src/commands/pi-command-utils.ts b/packages/pi-plugin/src/commands/pi-command-utils.ts index ee22b8471..08a416324 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,28 @@ type PiEntryRendererRegistration = { export type PiMessageSender = Pick & PiEntryRendererRegistration; +const statusLifecycleSignals = new WeakMap(); + +export function registerCtxStatusLifecycleSignal( + pi: PiMessageSender, + signal: AbortSignal, +): void { + statusLifecycleSignals.set(pi, signal); +} + +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 +76,56 @@ export function resolveSessionId( } } +export async function showCtxStatusDialog( + ctx: Pick, + content: CtxStatusMessageContent, +): Promise { + let factoryInvoked = false; + await ctx.ui.custom( + (_tui, theme, _keybindings, done) => { + factoryInvoked = true; + return new CtxStatusDialog(content, theme, done); + }, + { overlay: true, overlayOptions: { anchor: "center", width: 92 } }, + ); + return factoryInvoked; +} + +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,11 +181,57 @@ export function registerCtxStatusEntryRenderer(pi: PiMessageSender): boolean { } } +async function presentCtxStatusMessage( + ctx: Pick, + content: CtxStatusMessageContent, + signal?: AbortSignal, +): Promise { + if (ctx.mode !== "rpc" || signal?.aborted) return; + const type = + content.level === "error" || content.level === "warning" + ? content.level + : "info"; + if (!shouldShowCtxStatusDialog(content)) { + if (!signal?.aborted) ctx.ui.notify(content.text, type); + return; + } + try { + const shown = await showCtxStatusDialog(ctx, content); + if (!shown && !signal?.aborted) ctx.ui.notify(content.text, type); + } catch (err) { + if (signal?.aborted) return; + 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, + signal?: AbortSignal, +): (content: CtxStatusMessageContent, details?: unknown) => void { + const lifecycleSignal = statusLifecycleSignals.get(pi); + const effectiveSignal = + signal && lifecycleSignal + ? AbortSignal.any([signal, lifecycleSignal]) + : (signal ?? lifecycleSignal); + return (content, details) => { + if (!effectiveSignal?.aborted) + sendCtxStatusMessage(pi, content, details, ctx, effectiveSignal); + }; +} + export function sendCtxStatusMessage( pi: PiMessageSender, content: CtxStatusMessageContent, details?: unknown, + ctx?: ExtensionCommandContext, + signal?: AbortSignal, ): void { + if (signal?.aborted) return; const data: CtxStatusEntryData = { ...content, details: details ?? content.details, @@ -125,6 +242,8 @@ export function sendCtxStatusMessage( if (typeof pi.appendEntry === "function") { pi.appendEntry(CTX_STATUS_CUSTOM_TYPE, data); } + if (ctx) void presentCtxStatusMessage(ctx, data, signal); + // 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 217fe6089..d4c567ae8 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 805d9dc03..c422ae6e9 100644 --- a/packages/pi-plugin/src/context-handler.ts +++ b/packages/pi-plugin/src/context-handler.ts @@ -3363,14 +3363,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..462287ae1 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"; @@ -23,6 +28,7 @@ type CapturedDreamClient = { session: { create: (args: unknown) => Promise; prompt: (args: unknown) => Promise; + messages: (args: unknown) => Promise; }; }; @@ -56,6 +62,7 @@ function dreamerOptions(args: { database: Database; projectIdentity: string; projectDir?: string; + registrationOwner?: object; config?: DreamerConfig; language?: string; onAdjunctsRefreshNeeded?: (projectIdentity: string) => void; @@ -66,6 +73,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 +144,101 @@ 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("shares manual-run draining across jiti-style module instances", async () => { + db = createDb(); + const gate = deferred<{ ok: true; assistantText: string }>(); + const runStarted = deferred(); + __test.setStartDreamScheduleTimerFactory(async () => mock(() => {})); + __test.setPiSubagentRunnerFactory( + () => + ({ + run: mock(() => { + runStarted.resolve(); + return gate.promise; + }), + }) as never, + ); + const projectIdentity = "git:pi-shared-manual-drain"; + const ownerA = {}; + const ownerB = {}; + const config = DreamerConfigSchema.parse({ + model: "test/model", + tasks: { curate: { schedule: "0 4 * * *" } }, + }); + insertMemory(db, { + projectPath: projectIdentity, + category: "PROJECT_RULES", + content: "Keep reload-safe Dreamer lifecycle accounting process-shared.", + }); + const opts = dreamerOptions({ + database: db, + projectDir: process.cwd(), + projectIdentity, + registrationOwner: ownerA, + config, + }); + registerPiDreamerProject(opts); + + const secondInstance = await import( + `./index.ts?manual-drain-instance=${Date.now()}` + ); + secondInstance.registerPiDreamerProject({ + ...opts, + registrationOwner: ownerB, + }); + const manualRun = secondInstance.runPiDreamForProject( + projectIdentity, + "curate", + ownerB, + ); + await runStarted.promise; + let drained = false; + const drain = secondInstance.awaitInFlightDreamers(ownerB).then(() => { + drained = true; + }); + await flushMicrotasks(); + expect(drained).toBe(false); + + gate.resolve({ ok: true, assistantText: "curation complete" }); + await manualRun; + await drain; + expect(drained).toBe(true); + secondInstance.__test.reset(); + }); + test("threads language into scheduled dreamer registration", async () => { db = createDb(); let language: string | undefined; @@ -156,7 +259,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 +278,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 +343,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 +399,75 @@ describe("Pi dreamer wiring", () => { expect(timerCleanup).not.toHaveBeenCalled(); }); + test("legacy same-dir registration rebuilds once", async () => { + db = createDb(); + const firstCleanup = mock(() => {}); + const secondCleanup = mock(() => {}); + const cleanups = [firstCleanup, secondCleanup]; + let timerStarts = 0; + __test.setStartDreamScheduleTimerFactory(async () => { + timerStarts += 1; + return cleanups.shift() ?? mock(() => {}); + }); + const projectIdentity = "git:pi-legacy-registration"; + const opts = dreamerOptions({ + database: db, + projectDir: "/tmp/pi-legacy", + projectIdentity, + }); + + registerPiDreamerProject(opts); + await flushMicrotasks(); + __test.clearRegistrationGeneration(projectIdentity); + registerPiDreamerProject(opts); + await flushMicrotasks(); + registerPiDreamerProject(opts); + await flushMicrotasks(); + + expect(timerStarts).toBe(2); + expect(firstCleanup).toHaveBeenCalledTimes(1); + expect(secondCleanup).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 +480,512 @@ 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("keeps only the final timer when A-B-A registrations start concurrently", async () => { + db = createDb(); + const gates = [ + deferred<() => void>(), + deferred<() => void>(), + deferred<() => void>(), + ]; + const cleanups = [mock(() => {}), mock(() => {}), mock(() => {})]; + const clients: CapturedDreamClient[] = []; + const timerRegistrations = new Map(); + let timerIndex = 0; + __test.setPiSubagentRunnerFactory( + () => + ({ + run: mock(async () => ({ ok: true, assistantText: "done" })), + }) as never, + ); + __test.setStartDreamScheduleTimerFactory(async (registration) => { + const index = timerIndex++; + clients.push(registration.client as unknown as CapturedDreamClient); + timerRegistrations.set(registration.directory, index); + const cleanup = await (gates[index]?.promise ?? + Promise.resolve(mock(() => {}))); + return () => { + cleanup(); + timerRegistrations.delete(registration.directory); + }; + }); + const projectIdentity = "git:pi-overlapping-handoff"; + const ownerA = {}; + const ownerB = {}; + registerPiDreamerProject( + dreamerOptions({ + database: db, + projectDir: "/tmp/worktree-A", + projectIdentity, + registrationOwner: ownerA, + }), + ); + registerPiDreamerProject( + dreamerOptions({ + database: db, + projectDir: "/tmp/worktree-B", + projectIdentity, + registrationOwner: ownerB, + }), + ); registerPiDreamerProject( dreamerOptions({ database: db, projectDir: "/tmp/worktree-A", - projectIdentity: "git:pi-worktree", + projectIdentity, + registrationOwner: ownerA, }), ); + expect(timerIndex).toBe(3); + + gates[1]?.resolve(cleanups[1] as () => void); await flushMicrotasks(); - // Worktree B of the SAME repo (same identity, different dir). + gates[0]?.resolve(cleanups[0] as () => void); + await flushMicrotasks(); + gates[2]?.resolve(cleanups[2] as () => void); + await flushMicrotasks(); + + expect(cleanups[0]).not.toHaveBeenCalled(); + expect(cleanups[1]).toHaveBeenCalledTimes(1); + expect(cleanups[2]).not.toHaveBeenCalled(); + expect(timerRegistrations).toEqual(new Map([["/tmp/worktree-A", 2]])); + await expect(clients[0]?.session.create({})).rejects.toThrow( + "registration is no longer active", + ); + await expect(clients[1]?.session.create({})).rejects.toThrow( + "registration is no longer active", + ); + const activeClient = requireCapturedClient(clients[2] ?? null); + const session = (await activeClient.session.create({})) as { id: string }; + await activeClient.session.prompt({ + path: { id: session.id }, + body: { system: "system", parts: [{ text: "run dreamer" }] }, + }); + }); + + test("stale sibling timer clients stay invalid across a worktree handoff", 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 ownerA = {}; + const ownerB = {}; + registerPiDreamerProject( + dreamerOptions({ + database: db, + projectDir: "/tmp/worktree-A", + projectIdentity, + registrationOwner: ownerA, + }), + ); + await flushMicrotasks(); + const oldClient = requireCapturedClient(clients[0] ?? null); + const created = (await oldClient.session.create({})) as { id: string }; + registerPiDreamerProject( dreamerOptions({ database: db, projectDir: "/tmp/worktree-B", - projectIdentity: "git:pi-worktree", + projectIdentity, + registrationOwner: ownerB, }), ); await flushMicrotasks(); + await expect(oldClient.session.create({})).rejects.toThrow( + "registration is no longer active", + ); + await expect( + oldClient.session.prompt({ + path: { id: created.id }, + body: { system: "system", parts: [{ text: "run dreamer" }] }, + }), + ).rejects.toThrow("registration is no longer active"); - // 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, + registrationOwner: ownerB, + }); + await flushMicrotasks(); + const replacementClient = requireCapturedClient(clients[2] ?? null); + await expect(oldClient.session.create({})).rejects.toThrow( + "registration is no longer active", + ); + const replacementSession = (await replacementClient.session.create({})) as { + id: string; + }; + await replacementClient.session.prompt({ + path: { id: replacementSession.id }, + body: { system: "system", parts: [{ text: "run dreamer" }] }, + }); + expect(run).toHaveBeenCalledTimes(1); }); - test("unregister removes the project", () => { + test("discards a stale timer result that settles after a worktree handoff", async () => { + db = createDb(); + const gate = deferred<{ ok: true; assistantText: string }>(); + const clients: CapturedDreamClient[] = []; + const refresh = mock(() => {}); + __test.setPiSubagentRunnerFactory( + () => ({ run: mock(() => gate.promise) }) as never, + ); + __test.setStartDreamScheduleTimerFactory(async (registration) => { + clients.push(registration.client as unknown as CapturedDreamClient); + return mock(() => {}); + }); + const projectIdentity = "git:pi-late-stale-result"; + registerPiDreamerProject( + dreamerOptions({ + database: db, + projectDir: "/tmp/worktree-A", + projectIdentity, + registrationOwner: {}, + onAdjunctsRefreshNeeded: refresh, + }), + ); + await flushMicrotasks(); + const oldClient = requireCapturedClient(clients[0] ?? null); + const created = (await oldClient.session.create({})) as { id: string }; + const prompt = oldClient.session.prompt({ + path: { id: created.id }, + body: { system: "system", parts: [{ text: "run dreamer" }] }, + }); + await flushMicrotasks(); + + registerPiDreamerProject( + dreamerOptions({ + database: db, + projectDir: "/tmp/worktree-B", + projectIdentity, + registrationOwner: {}, + }), + ); + await flushMicrotasks(); + gate.resolve({ ok: true, assistantText: "stale result" }); + + await expect(prompt).rejects.toThrow("registration is no longer active"); + await expect( + oldClient.session.messages({ path: { id: created.id } }), + ).rejects.toThrow("registration is no longer active"); + expect(refresh).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"]); + + unregisterPiDreamerProject({ + projectIdentity: "git:pi-handoff", + registrationOwner: firstA.registrationOwner, + }); + expect(__test.registeredProjectCount()).toBe(0); + }); + + test("re-registration refreshes owner recency before active-owner handoff", async () => { + db = createDb(); + const dirs: string[] = []; + __test.setStartDreamScheduleTimerFactory(async (registration) => { + dirs.push((registration as { directory: string }).directory); + return mock(() => {}); + }); + + const ownerA = {}; + const firstA = dreamerOptions({ + database: db, + projectDir: "/tmp/worktree-A", + projectIdentity: "git:pi-owner-recency", + registrationOwner: ownerA, + }); + const ownerB = dreamerOptions({ + database: db, + projectDir: "/tmp/worktree-B", + projectIdentity: "git:pi-owner-recency", + }); + const refreshedA = dreamerOptions({ + database: db, + projectDir: "/tmp/worktree-A", + projectIdentity: "git:pi-owner-recency", + registrationOwner: ownerA, + }); + const activeC = dreamerOptions({ + database: db, + projectDir: "/tmp/worktree-C", + projectIdentity: "git:pi-owner-recency", + }); + for (const owner of [firstA, ownerB, refreshedA, activeC]) { + registerPiDreamerProject(owner); + await flushMicrotasks(); + } + + unregisterPiDreamerProject({ + projectIdentity: "git:pi-owner-recency", + registrationOwner: activeC.registrationOwner, + }); + await flushMicrotasks(); + + expect(dirs).toEqual([ + "/tmp/worktree-A", + "/tmp/worktree-B", + "/tmp/worktree-A", + "/tmp/worktree-C", + "/tmp/worktree-A", + ]); + }); + + test("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, + projectIdentity, + registrationOwner: owner, + config: DreamerConfigSchema.parse({ + model: "test/model", + tasks: { curate: { schedule: "0 4 * * *" } }, + }), + }), + ); + + const manualRun = runPiDreamForProject(projectIdentity, "curate", owner); + await flushMicrotasks(); + let drained = false; + const drain = awaitInFlightDreamers(owner).then(() => { + drained = true; + }); + await flushMicrotasks(); + expect(drained).toBe(false); + + unregisterPiDreamerProject({ projectIdentity, registrationOwner: owner }); + releaseLease(db, blocker, leaseKey); + const result = await manualRun; + await drain; + expect(drained).toBe(true); + expect(run).not.toHaveBeenCalled(); + expect(result.failed).toEqual(["curate"]); + }); + + test("drains a manual run whose successful result arrives after unregister", async () => { db = createDb(); + const gate = deferred<{ ok: true; assistantText: string }>(); + const runStarted = deferred(); + const refresh = mock(() => {}); + __test.setStartDreamScheduleTimerFactory(async () => mock(() => {})); + __test.setPiSubagentRunnerFactory( + () => + ({ + run: mock(() => { + runStarted.resolve(); + return gate.promise; + }), + }) as never, + ); + const projectIdentity = "git:pi-manual-late-unregister"; + const owner = {}; + insertMemory(db, { + projectPath: projectIdentity, + category: "PROJECT_RULES", + content: "Ignore Dreamer results after their registration owner exits.", + }); registerPiDreamerProject( dreamerOptions({ database: db, - projectDir: "/tmp/pi-project-unregister", - projectIdentity: "git:pi-unregister", + projectDir: process.cwd(), + projectIdentity, + registrationOwner: owner, + config: DreamerConfigSchema.parse({ + model: "test/model", + tasks: { curate: { schedule: "0 4 * * *" } }, + }), + onAdjunctsRefreshNeeded: refresh, }), ); - unregisterPiDreamerProject({ projectIdentity: "git:pi-unregister" }); + const manualRun = runPiDreamForProject(projectIdentity, "curate", owner); + await runStarted.promise; + unregisterPiDreamerProject({ + projectIdentity, + registrationOwner: owner, + }); + let drained = false; + const drain = awaitInFlightDreamers(owner).then(() => { + drained = true; + }); + await flushMicrotasks(); + expect(drained).toBe(false); + + gate.resolve({ ok: true, assistantText: "curation complete" }); + const result = await manualRun; + await drain; + expect(drained).toBe(true); + expect(result.failed).toEqual(["curate"]); + expect(refresh).not.toHaveBeenCalled(); + }); + + 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 +994,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 +1093,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 +1248,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 +1271,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..bf3bb4cfb 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; /** @@ -82,11 +84,19 @@ interface SessionPromptArgs extends SessionMessagesArgs { type SessionDeleteArgs = SessionMessagesArgs; interface ProjectRegistration { + /** Unique per timer build. Optional only for registrations retained across reload from older code. */ + generation?: object; 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 +115,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,26 +157,45 @@ 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. - if (existing.projectDir === opts.projectDir) { + // Same identity and directory genuinely reuses the timer. Registrations + // retained by an older module have no generation and must rebuild once. + if (existing.generation && existing.projectDir === opts.projectDir) { return; } - // Same identity, DIFFERENT directory: a worktree/clone switch in the same - // process. The existing registration's timer + client closure are pinned - // to the OLD checkout and its boot-time dreamerConfig. Tear it down and - // rebuild against the new directory + freshly-resolved config below, so - // the dreamer runs in the right checkout (and honors a `dreamer.disable` - // that may differ between checkouts — handled by the disable early-return - // above, which fires before this). + // A different checkout or legacy registration has a timer + client closure + // pinned to stale options. Tear it down before rebuilding below. existing.cleanup(); 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 + const generation = {}; + // 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, + () => { + const current = registeredProjects.get(opts.projectIdentity); + return ( + current?.generation === generation && + current.owners.get(opts.registrationOwner)?.projectDir === + opts.projectDir + ); + }, + ); let cleanup: (() => void) | undefined; let cancelled = false; @@ -168,47 +222,85 @@ export function registerPiDreamerProject(opts: PiDreamerOptions): void { primerRawProviderFactory: createPiPrimerRawProviderFactory(), }).then((timerCleanup) => { if (cancelled) { - // Registration was cancelled before timer setup completed — - // immediately invoke cleanup to prevent leaked timer registration. - timerCleanup?.(); + // The shared timer registry is keyed by directory. A newer A registration + // may already have replaced this stale A entry during an async A→B→A handoff, + // so its directory must not be deleted by the stale cleanup. + if ( + registeredProjects.get(opts.projectIdentity)?.projectDir !== + opts.projectDir + ) { + timerCleanup?.(); + } return; } cleanup = timerCleanup; }); // 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, { + generation, + activeOwner: opts.registrationOwner, + owners, cleanup: () => { cancelled = true; cleanup?.(); @@ -222,18 +314,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 +335,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,22 +376,40 @@ 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 assertRegistrationOwnerActive = (): void => { + if (!isRegistrationOwnerActive()) { + throw new Error( + `Pi dreamer registration is no longer active for project ${opts.projectIdentity}`, + ); + } + }; const session = { create: async (args: SessionCreateArgs) => { + assertRegistrationOwnerActive(); const sessionId = `magic-context-pi-dream-${++sessionCounter}`; sessionsById.set(sessionId, { id: sessionId, @@ -290,6 +427,8 @@ function createPiDreamerClient(opts: PiDreamerOptions): DreamTimerClient { throw new Error(`Pi dreamer session not found: ${sessionId}`); } + assertRegistrationOwnerActive(); + const userMessage = extractUserMessage(args); const systemPrompt = extractSystemPrompt(args); // Per-task model override (Dreamer v2): the SHARED executor @@ -318,9 +457,10 @@ 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; + assertRegistrationOwnerActive(); if (!result.ok) { const error = new Error( `Pi dreamer subagent failed (${result.reason}): ${result.error}`, @@ -346,12 +486,13 @@ 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); } }, messages: async (args: SessionMessagesArgs) => { + assertRegistrationOwnerActive(); const dreamSession = sessionsById.get(args.path.id); return { data: dreamSession?.messages ?? [] }; }, @@ -485,6 +626,10 @@ function makeMessage( export const __test = { registeredProjectCount: () => registeredProjects.size, + clearRegistrationGeneration: (projectIdentity: string) => { + const registration = registeredProjects.get(projectIdentity); + if (registration) delete registration.generation; + }, setPiSubagentRunnerFactory: (factory: PiSubagentRunnerFactory) => { piSubagentRunnerFactory = factory; }, 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..72a4abc1a 100644 --- a/packages/pi-plugin/src/index-in-process-latch.test.ts +++ b/packages/pi-plugin/src/index-in-process-latch.test.ts @@ -1,9 +1,10 @@ -import { afterEach, describe, expect, it, mock } from "bun:test"; +import { afterEach, describe, expect, it, mock, spyOn } from "bun:test"; import { mkdtempSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import type { ExtensionAPI } from "@earendil-works/pi-coding-agent"; import magicContextPiExtension, { __test } from "./index"; +import { awaitInFlightRecomps, spawnPiRecompRun } from "./pi-recomp-runner"; import { MAGIC_CONTEXT_PI_SUBAGENT_ENV } from "./subagent-runner"; const originalEnv = { @@ -22,35 +23,63 @@ 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[] = []; const tools: string[] = []; const flags: string[] = []; const commands: string[] = []; + const commandHandlers = new Map< + string, + (args: string, ctx: unknown) => unknown + >(); 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 ?? ""); }), registerFlag: mock((name: string) => { flags.push(name); }), - registerCommand: mock((name: string) => { - commands.push(name); - }), + registerCommand: mock( + ( + name: string, + command: { handler: (args: string, ctx: unknown) => unknown }, + ) => { + commands.push(name); + commandHandlers.set(name, command.handler); + }, + ), registerEntryRenderer: mock((customType: string) => { entryRenderers.push(customType); }), @@ -58,73 +87,369 @@ 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, + runCommand(name: string, args: string, ctx: unknown) { + const handler = commandHandlers.get(name); + if (!handler) throw new Error(`Command not registered: ${name}`); + return handler(args, ctx); + }, + 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("fences a registered command's late RPC fallback on shutdown", async () => { + isolateXdgEnv(); + delete process.env[MAGIC_CONTEXT_PI_SUBAGENT_ENV]; + const runtime = createCountingPi(); + await magicContextPiExtension(runtime.pi); + + let resolveCustom!: (value: undefined) => void; + let notifications = 0; + const ctx = { + mode: "rpc", + hasUI: false, + cwd: process.cwd(), + model: undefined, + sessionManager: { getSessionId: () => undefined }, + ui: { + custom: () => + new Promise((resolve) => { + resolveCustom = resolve; + }), + notify: () => { + notifications += 1; + }, + setStatus: () => undefined, + }, + }; + await runtime.runCommand("ctx-status", "", ctx); + expect(resolveCustom).toBeDefined(); + + await runtime.emitPiEvent("session_shutdown", {}, ctx); + resolveCustom(undefined); + await Promise.resolve(); + await Promise.resolve(); + expect(notifications).toBe(0); + }); + + it("aborts a recomp only after the five-second shutdown drain expires", async () => { + isolateXdgEnv(); + delete process.env[MAGIC_CONTEXT_PI_SUBAGENT_ENV]; + const runtime = createCountingPi(); + await magicContextPiExtension(runtime.pi); + + let observedSignal: AbortSignal | undefined; + let releaseRun!: () => void; + const runGate = new Promise((resolve) => { + releaseRun = resolve; + }); + spawnPiRecompRun({ + sessionId: "ses-shutdown-timeout", + provider: { readMessages: () => [] }, + onStatusChange() {}, + work: async (signal) => { + observedSignal = signal; + await runGate; + }, + }); + await Promise.resolve(); + + const timers: Array<{ + active: boolean; + callback: () => void; + delay: number; + handle: ReturnType; + }> = []; + const setTimeoutSpy = spyOn(globalThis, "setTimeout").mockImplementation((( + callback: (...args: unknown[]) => void, + delay?: number, + ) => { + const timer = { + active: true, + callback: () => callback(), + delay: delay ?? 0, + handle: { unref() {} } as ReturnType, + }; + timers.push(timer); + return timer.handle; + }) as typeof setTimeout); + const clearTimeoutSpy = spyOn( + globalThis, + "clearTimeout", + ).mockImplementation(((handle: ReturnType) => { + const timer = timers.find((candidate) => candidate.handle === handle); + if (timer) timer.active = false; + }) as typeof clearTimeout); + + try { + const shutdown = runtime.emitPiEvent( + "session_shutdown", + {}, + { + sessionManager: { getSessionId: () => "ses-shutdown-timeout" }, + ui: { setStatus: () => undefined }, + }, + ); + for (let attempt = 0; attempt < 20 && timers.length < 2; attempt += 1) { + await Promise.resolve(); + } + const timeout = timers.findLast((timer) => timer.active); + expect(timeout?.delay).toBe(5_000); + expect(observedSignal?.aborted).toBe(false); + + timeout?.callback(); + await shutdown; + expect(observedSignal?.aborted).toBe(true); + } finally { + setTimeoutSpy.mockRestore(); + clearTimeoutSpy.mockRestore(); + releaseRun(); + await awaitInFlightRecomps("ses-shutdown-timeout"); + } + }, 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("suppresses four overlapping child initializations without leaking ALS state", async () => { + isolateXdgEnv(); + delete process.env[MAGIC_CONTEXT_PI_SUBAGENT_ENV]; + + const parent = createCountingPi(); + await magicContextPiExtension(parent.pi); + __test.clearPiStartupMaintenanceClaim(); + + const children = Array.from({ length: 4 }, () => createCountingPi()); + let markedCount = 0; + let releaseBarrier!: () => void; + const allMarked = new Promise((resolve) => { + releaseBarrier = resolve; + }); + const branches = children.map((child) => + Promise.resolve().then(async () => { + parent.emitEvent("subagents:child:session-created"); + markedCount += 1; + if (markedCount === children.length) releaseBarrier(); + await allMarked; + await magicContextPiExtension(child.pi); + parent.emitEvent("subagents:child:disposed"); + }), + ); + await allMarked; + await Promise.all(branches); + + for (const child of children) { + expect(child.events).toEqual([]); + expect(child.tools).toEqual([]); + expect(child.flags).toEqual([]); + expect(child.commands).toEqual([]); + expect(child.entryRenderers).toEqual([]); + expect( + child.eventBusHandlerCount("subagents:child:session-created"), + ).toBe(0); + expect(child.eventBusHandlerCount("subagents:child:disposed")).toBe(0); + } + // Child entry must return before claiming process-wide startup maintenance. + expect(__test.claimPiStartupMaintenance()).toBe(true); + __test.clearPiStartupMaintenanceClaim(); + + const independent = createCountingPi(); + await magicContextPiExtension(independent.pi); + expect(independent.tools).toContain("ctx_search"); + expect(independent.commands).toContain("ctx-status"); + expect(independent.entryRenderers).toEqual(["ctx-status"]); + expect(__test.claimPiStartupMaintenance()).toBe(false); + }, 20_000); + + it("keeps sibling child markers isolated after one child disposes early", async () => { + isolateXdgEnv(); + delete process.env[MAGIC_CONTEXT_PI_SUBAGENT_ENV]; + + const parent = createCountingPi(); + await magicContextPiExtension(parent.pi); + __test.clearPiStartupMaintenanceClaim(); + const children = Array.from({ length: 4 }, () => createCountingPi()); + let markedCount = 0; + let releaseBarrier!: () => void; + const allMarked = new Promise((resolve) => { + releaseBarrier = resolve; + }); + let firstDisposed!: () => void; + const firstChildDisposed = new Promise((resolve) => { + firstDisposed = resolve; + }); + + const branches = children.map((child, index) => + Promise.resolve().then(async () => { + parent.emitEvent("subagents:child:session-created"); + markedCount += 1; + if (markedCount === children.length) releaseBarrier(); + await allMarked; + if (index !== 0) await firstChildDisposed; + await magicContextPiExtension(child.pi); + parent.emitEvent("subagents:child:disposed"); + if (index === 0) firstDisposed(); + }), + ); + await Promise.all(branches); + + for (const child of children) { + expect(child.events).toEqual([]); + expect(child.tools).toEqual([]); + expect(child.flags).toEqual([]); + expect(child.commands).toEqual([]); + expect(child.entryRenderers).toEqual([]); + } + expect(__test.claimPiStartupMaintenance()).toBe(true); + __test.clearPiStartupMaintenanceClaim(); + }, 20_000); + + 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 +459,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..61b9e8e7b 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,8 @@ import { registerCtxStatusCommand } from "./commands/ctx-status"; import { registerCtxWrapupCommand } from "./commands/ctx-wrapup"; import { registerCtxStatusEntryRenderer, + registerCtxStatusLifecycleSignal, + resolveSessionId, sendCtxStatusMessage, } from "./commands/pi-command-utils"; import { loadPiConfig } from "./config"; @@ -156,7 +159,7 @@ import { ensureProjectRegisteredFromPiDirectory } from "./embedding-bootstrap"; import { registerPiFailClosedSurface } from "./fail-closed-pi"; import { resolvePiUsableContextLimit } from "./pi-context-limit"; import { computePiPressure, extractAssistantUsage } from "./pi-pressure"; -import { awaitInFlightRecomps } from "./pi-recomp-runner"; +import { abortInFlightRecomps, awaitInFlightRecomps } from "./pi-recomp-runner"; import { readPiSessionMessages } from "./read-session-pi"; import { registerStatusLine, updateStatusLine } from "./status-line"; import { stripTagPrefixFromAssistantMessage } from "./strip-tag-prefix"; @@ -184,7 +187,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 +200,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 +497,9 @@ export const __test = { resetLoggedPiConfigDirs(): void { loggedPiConfigDirs.clear(); }, - isPiMagicContextActiveInProcess, - markPiMagicContextActive, - clearPiMagicContextActive, + clearPiInProcessSubagentInitContext, + claimPiStartupMaintenance, + clearPiStartupMaintenanceClaim, }; function formatTokens(value: number): string { @@ -732,18 +764,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 +893,72 @@ 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 = {}; + const commandLifecycleController = new AbortController(); + registerCtxStatusLifecycleSignal(pi, commandLifecycleController.signal); + 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 +1193,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 +1534,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 +1563,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 +1638,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 +1652,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 +2308,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 +2322,8 @@ 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; + commandLifecycleController.abort(); // 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 +2331,54 @@ 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); + } + // Timeout only stops waiting. Fence and cancel any run still alive before + // this handler returns and Pi disposes its command context. + abortInFlightRecomps(sessionId); } 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 +2405,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 +2445,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-historian-runner.ts b/packages/pi-plugin/src/pi-historian-runner.ts index 5e246cd45..c6f5f546b 100644 --- a/packages/pi-plugin/src/pi-historian-runner.ts +++ b/packages/pi-plugin/src/pi-historian-runner.ts @@ -1680,7 +1680,7 @@ export function buildPiCompactionSummary( * cut at, so defer there too. */ export function findFirstKeptEntryId( - entries: unknown[], + entries: readonly unknown[], lastCompactedOrdinal: number, ): string | null { const rawMessages = convertEntriesToRawMessages(entries); diff --git a/packages/pi-plugin/src/pi-memory-migration.test.ts b/packages/pi-plugin/src/pi-memory-migration.test.ts index 8dad61fd9..a8d1200df 100644 --- a/packages/pi-plugin/src/pi-memory-migration.test.ts +++ b/packages/pi-plugin/src/pi-memory-migration.test.ts @@ -253,4 +253,39 @@ describe("Pi memory migration (E6c)", () => { // `primaryModelId ?? undefined` falling through to the agent default). expect(models[0]).toBe("historian/model"); }); + + it("cancellation stops the fallback chain before applying memory changes", async () => { + const db = openDatabase(); + const projectPath = resolveProjectIdentity(process.cwd()); + insertMemory(db, { + projectPath, + category: "ARCHITECTURE_DECISIONS", + content: "keep me", + }); + const controller = new AbortController(); + let calls = 0; + let observedSignal: AbortSignal | undefined; + const outcome = await runPiMemoryMigration({ + db, + runner: { + run: async (options: { signal?: AbortSignal }) => { + calls += 1; + observedSignal = options.signal; + controller.abort(); + return { ok: true, assistantText: "no migrated block" }; + }, + } as never, + model: "historian/model", + fallbackModels: ["fallback/model"], + directory: process.cwd(), + sessionId: "ses-pi-mig-cancel", + signal: controller.signal, + }); + + expect(outcome.summary).toContain("cancelled"); + expect(observedSignal).toBe(controller.signal); + expect(calls).toBe(1); + expect(getMemoriesByProject(db, projectPath)).toHaveLength(1); + expect(isMemoryMigrationDone(db, projectPath)).toBe(false); + }); }); diff --git a/packages/pi-plugin/src/pi-memory-migration.ts b/packages/pi-plugin/src/pi-memory-migration.ts index 32469bc86..fbed3d981 100644 --- a/packages/pi-plugin/src/pi-memory-migration.ts +++ b/packages/pi-plugin/src/pi-memory-migration.ts @@ -55,6 +55,7 @@ export interface PiMemoryMigrationDeps { /** Route user_observations to the user-memory candidate pool when enabled. */ userMemoriesEnabled?: boolean; language?: string; + signal?: AbortSignal; } export interface PiMemoryMigrationOutcome { @@ -62,9 +63,15 @@ export interface PiMemoryMigrationOutcome { summary: string; } +const CANCELLED_OUTCOME: PiMemoryMigrationOutcome = { + ran: false, + summary: "Memory migration cancelled during session shutdown.", +}; + export async function runPiMemoryMigration( deps: PiMemoryMigrationDeps, ): Promise { + if (deps.signal?.aborted) return CANCELLED_OUTCOME; const projectPath = resolveProjectIdentityForSession( deps.directory, deps.allowHomeProject, @@ -139,6 +146,7 @@ export async function runPiMemoryMigration( let parsed: ReturnType | null = null; let lastFailReason = "no output"; for (let i = 0; i < modelChain.length; i += 1) { + if (deps.signal?.aborted) return CANCELLED_OUTCOME; const model = modelChain[i]; if (i > 0) { sessionLog( @@ -164,7 +172,9 @@ export async function runPiMemoryMigration( // Reuse the "recomp" accounting bucket — memory migration is part of the // session-upgrade flow and there is no dedicated subagent tag for it. accountingSubagent: "recomp", + signal: deps.signal, }); + if (deps.signal?.aborted) return CANCELLED_OUTCOME; if (!result.ok) { lastFailReason = `historian ${result.reason}`; @@ -210,6 +220,7 @@ export async function runPiMemoryMigration( }; } + if (deps.signal?.aborted) return CANCELLED_OUTCOME; // Persist observations BEFORE the destructive apply. let routed = 0; if (deps.userMemoriesEnabled && parsed.userObservations.length > 0) { diff --git a/packages/pi-plugin/src/pi-recomp-client-shared.ts b/packages/pi-plugin/src/pi-recomp-client-shared.ts index e089b6e73..c2a3cb75e 100644 --- a/packages/pi-plugin/src/pi-recomp-client-shared.ts +++ b/packages/pi-plugin/src/pi-recomp-client-shared.ts @@ -22,11 +22,14 @@ export function createPiHistorianClient(args: { thinkingLevel?: string; directory: string; accountingSessionId: string; + signal?: AbortSignal; notify: (text: string) => void; }) { const sessions = new Map(); let counter = 0; async function prompt(input: unknown): Promise> { + if (args.signal?.aborted) + throw new Error("prompt aborted by external signal"); const body = readBody(input); const sessionId = readPathId(input); if (body.noReply) { @@ -60,7 +63,10 @@ export function createPiHistorianClient(args: { : args.thinkingLevel, accountingSessionId: args.accountingSessionId, accountingSubagent: "recomp", + signal: args.signal, }); + if (args.signal?.aborted) + throw new Error("prompt aborted by external signal"); if (!result.ok) { throw new Error( `Pi recomp historian failed (${result.reason}): ${result.error}`, @@ -73,6 +79,8 @@ export function createPiHistorianClient(args: { session: { get: async () => ({ directory: args.directory }), create: async () => { + if (args.signal?.aborted) + throw new Error("prompt aborted by external signal"); const id = `magic-context-pi-recomp-${++counter}`; sessions.set(id, []); return { id }; diff --git a/packages/pi-plugin/src/pi-recomp-marker.ts b/packages/pi-plugin/src/pi-recomp-marker.ts index 3a1a1bf88..dcd528300 100644 --- a/packages/pi-plugin/src/pi-recomp-marker.ts +++ b/packages/pi-plugin/src/pi-recomp-marker.ts @@ -46,11 +46,8 @@ import { export function stagePiRecompMarker(args: { db: ContextDatabase; sessionId: string; - ctx: unknown; + branchEntries: readonly unknown[]; }): void { - const readBranchEntries = resolvePiReadBranchEntries(args.ctx); - if (!readBranchEntries) return; - const compartments = getCompartments(args.db, args.sessionId); const last = compartments[compartments.length - 1]; if (!last) return; @@ -58,7 +55,7 @@ export function stagePiRecompMarker(args: { let firstKeptEntryId: string | null = null; try { firstKeptEntryId = findFirstKeptEntryId( - readBranchEntries(), + args.branchEntries, last.endMessage, ); } catch { 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..73d735bdf --- /dev/null +++ b/packages/pi-plugin/src/pi-recomp-runner.test.ts @@ -0,0 +1,106 @@ +import { expect, test } from "bun:test"; +import { createPiHistorianClient } from "./pi-recomp-client-shared"; +import { + abortInFlightRecomps, + 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); +}); + +test("abort fences a detached run before a late settle", async () => { + const gate = deferred(); + let observedSignal: AbortSignal | undefined; + let statusChanges = 0; + let settled = false; + + spawnPiRecompRun({ + sessionId: "session-abort", + provider: { readMessages: async () => [] } as never, + onStatusChange: () => { + statusChanges += 1; + }, + work: async (signal) => { + observedSignal = signal; + await gate.promise; + settled = true; + }, + }); + await Promise.resolve(); + + abortInFlightRecomps("session-abort"); + expect(observedSignal?.aborted).toBe(true); + gate.resolve(); + await awaitInFlightRecomps("session-abort"); + + expect(settled).toBe(true); + expect(statusChanges).toBe(1); +}); + +test("historian client forwards cancellation and rejects a late result", async () => { + const gate = deferred(); + const controller = new AbortController(); + let runnerSignal: AbortSignal | undefined; + let runCalls = 0; + const client = createPiHistorianClient({ + runner: { + run: async (options: { signal?: AbortSignal }) => { + runCalls += 1; + runnerSignal = options.signal; + await gate.promise; + return { ok: true, assistantText: "done" }; + }, + } as never, + model: "provider/model", + systemPrompt: "system", + directory: "/project", + accountingSessionId: "session-client-abort", + signal: controller.signal, + notify: () => {}, + }); + const session = await client.session.create(); + const prompt = client.session.prompt({ + path: { id: session.id }, + body: { parts: [{ text: "rebuild" }] }, + }); + await Promise.resolve(); + + controller.abort(); + gate.resolve(); + await expect(prompt).rejects.toThrow("prompt aborted by external signal"); + expect(runnerSignal).toBe(controller.signal); + expect(runCalls).toBe(1); +}); diff --git a/packages/pi-plugin/src/pi-recomp-runner.ts b/packages/pi-plugin/src/pi-recomp-runner.ts index fa840174d..dfbc32914 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 @@ -18,7 +18,12 @@ import { setMagicContextRecompActive } from "./status-line"; * equivalent fire-and-forget so the REPL stays responsive while the historian * passes run in the background — the same pattern as `spawnPiHistorianRun`. */ -const inFlightRecomp = new Map>(); +interface InFlightRecompRun { + promise: Promise; + controller: AbortController; +} + +const inFlightRecomp = new Map(); /** True when a detached recomp/upgrade is already running for this session. */ export function isPiRecompInFlight(sessionId: string): boolean { @@ -26,13 +31,24 @@ 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)?.promise].filter( + (run): run is Promise => run !== undefined, + ) + : [...inFlightRecomp.values()].map((run) => run.promise); + if (runs.length === 0) return; + await Promise.allSettled(runs); +} + +/** Fence and cancel one session's detached recomp/upgrade, if still running. */ +export function abortInFlightRecomps(sessionId: string): void { + inFlightRecomp.get(sessionId)?.controller.abort(); } /** @@ -54,26 +70,36 @@ export function spawnPiRecompRun(args: { sessionId: string; provider: RawMessageProvider; onStatusChange: () => void; - work: () => Promise; + work: (signal: AbortSignal) => Promise; }): void { const { sessionId, provider, onStatusChange, work } = args; + const controller = new AbortController(); const unregister = setRawMessageProvider(sessionId, provider); setMagicContextRecompActive(sessionId, true); + + let run: InFlightRecompRun; + const runPromise = Promise.resolve() + .then(async () => { + try { + await work(controller.signal); + } catch (err) { + if (!controller.signal.aborted) { + sessionLog( + sessionId, + `pi recomp run failed (detached): ${err instanceof Error ? err.message : String(err)}`, + ); + } + } + }) + .finally(() => { + if (inFlightRecomp.get(sessionId) === run) { + inFlightRecomp.delete(sessionId); + } + setMagicContextRecompActive(sessionId, false); + unregister(); + if (!controller.signal.aborted) onStatusChange(); + }); + run = { promise: runPromise, controller }; + inFlightRecomp.set(sessionId, run); onStatusChange(); - const runPromise = (async () => { - try { - await work(); - } catch (err) { - sessionLog( - sessionId, - `pi recomp run failed (detached): ${err instanceof Error ? err.message : String(err)}`, - ); - } - })().finally(() => { - inFlightRecomp.delete(sessionId); - setMagicContextRecompActive(sessionId, false); - unregister(); - onStatusChange(); - }); - inFlightRecomp.set(sessionId, runPromise); } diff --git a/packages/pi-plugin/src/read-session-pi.ts b/packages/pi-plugin/src/read-session-pi.ts index cb5af6e91..975271f30 100644 --- a/packages/pi-plugin/src/read-session-pi.ts +++ b/packages/pi-plugin/src/read-session-pi.ts @@ -246,22 +246,37 @@ function getToolCallIds(content: unknown): Set { * repeated calls inside a single trigger evaluation don't re-walk the * branch. */ -export function readPiSessionMessages(ctx: ExtensionContext): RawMessage[] { +export interface PiSessionSnapshot { + branchEntries: readonly unknown[]; + rawMessages: RawMessage[]; +} + +export function readPiSessionSnapshot( + ctx: ExtensionContext, +): PiSessionSnapshot { const sm = ctx.sessionManager; - if (sm === undefined) return []; + if (sm === undefined) return { branchEntries: [], rawMessages: [] }; const getBranch = (sm as { getBranch?: (fromId?: string) => unknown[] }) .getBranch; - if (typeof getBranch !== "function") return []; + if (typeof getBranch !== "function") + return { branchEntries: [], rawMessages: [] }; let entries: unknown[]; try { entries = getBranch.call(sm); } catch { - return []; + return { branchEntries: [], rawMessages: [] }; } - if (!Array.isArray(entries)) return []; + if (!Array.isArray(entries)) return { branchEntries: [], rawMessages: [] }; - return convertEntriesToRawMessages(entries); + return { + branchEntries: entries, + rawMessages: convertEntriesToRawMessages(entries), + }; +} + +export function readPiSessionMessages(ctx: ExtensionContext): RawMessage[] { + return readPiSessionSnapshot(ctx).rawMessages; } /** @@ -340,7 +355,9 @@ function attachPiPartVersion( * Pure conversion exposed for unit testing — call sites in production * always go through `readPiSessionMessages`. */ -export function convertEntriesToRawMessages(entries: unknown[]): RawMessage[] { +export function convertEntriesToRawMessages( + entries: readonly unknown[], +): RawMessage[] { const result: RawMessage[] = []; let nextOrdinal = 1; 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");