@cortexkit/opencode-magic-context 0.44.1 → 0.44.2

package/src/shared/prompt-async-transport.ts

@@ -1,5 +1,12 @@
 import type { createOpencodeClient } from "@opencode-ai/sdk";
-
+import {
+    createDreamTokenBudget,
+    DreamTokenBudgetExceeded,
+    registerBudgetFinalizeChild,
+    releaseBudgetFinalizeChild,
+    TOKEN_BUDGET_FINALIZE_MESSAGE,
+} from "../features/magic-context/dreamer/token-budget";
+import { sumTokensFromChildMessages } from "../features/magic-context/subagent-token-capture";
 import type { PromptArgs, PromptTransport } from "./model-suggestion-retry";
 
 type Client = ReturnType<typeof createOpencodeClient> | undefined;
@@ -21,17 +28,26 @@ const DEFAULT_POLL_INTERVAL_MS = 1_000;
 const DEFAULT_START_GRACE_MS = 30_000;
 /** The wait only looks at the newest messages: the ones this send appended and
  * the ones just before it. The host returns the most recent `limit` messages. */
-const MESSAGE_WINDOW = 20;
+const MESSAGE_WINDOW = 100;
 
 export interface PromptAsyncWaitOptions {
     pollIntervalMs?: number;
     startGraceMs?: number;
+    tokenBudget?: number;
+    /** Shared across retries of one child session. */
+    budgetGuard?: ReturnType<typeof createDreamTokenBudget>;
+    onBudgetUpdate?: (
+        state: ReturnType<ReturnType<typeof createDreamTokenBudget>["snapshot"]> & {
+            sessionId: string;
+        },
+    ) => void;
 }
 
 type SessionApi = {
     promptAsync?: (options: unknown) => Promise<unknown>;
     status?: (options?: unknown) => Promise<unknown>;
     messages?: (options: unknown) => Promise<unknown>;
+    abort?: (options: unknown) => Promise<unknown>;
 };
 
 function sessionApi(client: Client): SessionApi | undefined {
@@ -59,8 +75,12 @@ export function createPromptAsyncTransport(
     options: PromptAsyncWaitOptions = {},
 ): PromptTransport | undefined {
     if (!supportsPromptAsync(client)) return undefined;
+    const budgetGuard = options.tokenBudget
+        ? createDreamTokenBudget(options.tokenBudget)
+        : undefined;
     return Object.assign(
-        (request: PromptArgs) => promptAsyncAndWaitForIdle(client, request, options),
+        (request: PromptArgs) =>
+            promptAsyncAndWaitForIdle(client, request, { ...options, budgetGuard }),
         { childSessionId },
     );
 }
@@ -90,6 +110,23 @@ function isSettledAssistant(message: unknown): boolean {
     return info.time?.completed != null || info.error != null || info.finish != null;
 }
 
+function isTerminalAssistant(message: unknown): boolean {
+    const info = infoOf(message);
+    return (
+        isSettledAssistant(message) &&
+        info.finish !== "tool-calls" &&
+        info.finish !== "tool_use" &&
+        !(
+            message &&
+            typeof message === "object" &&
+            Array.isArray((message as { parts?: unknown[] }).parts) &&
+            (message as { parts: Array<{ type?: string }> }).parts.some(
+                (part) => part.type === "tool",
+            )
+        )
+    );
+}
+
 function unwrapData(response: unknown): unknown {
     if (response && typeof response === "object" && "data" in response) {
         return (response as { data?: unknown }).data;
@@ -169,13 +206,6 @@ async function readStatus(
     }
 }
 
-/**
- * Send one prompt with `prompt_async`, then poll until the child session is idle
- * and its newest message is a finished assistant reply from this send. Returns
- * without judging that reply: the caller's output validation reads it. Rejects
- * when the caller's signal aborts, and when the host accepted the send but never
- * started a run within the start grace.
- */
 export async function promptAsyncAndWaitForIdle(
     client: Client,
     request: PromptArgs,
@@ -191,52 +221,131 @@ export async function promptAsyncAndWaitForIdle(
     const signal = request.signal;
     const pollIntervalMs = options.pollIntervalMs ?? DEFAULT_POLL_INTERVAL_MS;
     const startGraceMs = options.startGraceMs ?? DEFAULT_START_GRACE_MS;
-
-    // A fallback attempt reuses the child session, so an earlier attempt's final
-    // reply must not be mistaken for this send's reply.
-    const baseline = new Set(
+    const budget =
+        options.budgetGuard ??
+        (options.tokenBudget ? createDreamTokenBudget(options.tokenBudget) : undefined);
+    if (budget?.snapshot().finalizeFired)
+        throw new DreamTokenBudgetExceeded(sessionId, budget.snapshot().spent);
+    const seenUsage = new Set<string>();
+    let finalizing = false;
+    let completionBaseline = new Set(
         (await readMessages(session, sessionId, dir, signal))
             .map(messageId)
             .filter((id): id is string => id !== null),
     );
-
+    const initialBaseline = new Set(completionBaseline);
     const sent = await session.promptAsync?.(request);
     if (sent && typeof sent === "object" && "error" in sent) {
         const rejection = (sent as { error?: unknown }).error;
         if (rejection)
             throw new Error(`prompt_async was rejected: ${describeRejection(rejection)}`);
     }
-    const sentAt = Date.now();
+    let sentAt = Date.now();
     let sawBusy = false;
 
-    for (;;) {
-        await sleep(pollIntervalMs, signal);
-        const status = await readStatus(session, sessionId, dir, signal);
-        if (status === "busy" || status === "retry") {
-            sawBusy = true;
-            continue;
-        }
-        // An unreadable status says nothing about the run; keep waiting under the
-        // caller's timer rather than guessing.
-        if (status === null) continue;
-
-        const messages = await readMessages(session, sessionId, dir, signal);
-        const last = messages.at(-1);
-        const lastId = messageId(last);
-        if (last && lastId !== null && !baseline.has(lastId) && isSettledAssistant(last)) return;
-
-        const fresh = messages.some((message) => {
-            const id = messageId(message);
-            return id !== null && !baseline.has(id);
-        });
-        // The loop ran and ended without a final reply (for example it failed
-        // before writing one). Output validation reports what is missing.
-        if (sawBusy && fresh) return;
-        if (Date.now() - sentAt >= startGraceMs) {
-            if (fresh) return;
-            throw new Error(
-                `prompt_async did not start a run in child session ${sessionId} within ${startGraceMs}ms`,
-            );
+    try {
+        for (;;) {
+            await sleep(pollIntervalMs, signal);
+            const status = await readStatus(session, sessionId, dir, signal);
+            const messages = await readMessages(session, sessionId, dir, signal);
+            if (budget) {
+                for (const message of messages) {
+                    const id = messageId(message);
+                    if (!id || seenUsage.has(id) || infoOf(message).role !== "assistant") continue;
+                    // The initial baseline belongs to a previous child attempt.
+                    if (initialBaseline.has(id)) continue;
+                    const tokens = sumTokensFromChildMessages([message]);
+                    if (tokens.input + tokens.cacheRead + tokens.cacheWrite === 0) continue;
+                    seenUsage.add(id);
+                    const decision = budget.charge(
+                        tokens.input,
+                        tokens.cacheRead,
+                        tokens.cacheWrite,
+                        isTerminalAssistant(message),
+                    );
+                    options.onBudgetUpdate?.({ ...budget.snapshot(), sessionId });
+                    if (decision === "stop") {
+                        await session.abort?.({ path: { id: sessionId } });
+                        throw new DreamTokenBudgetExceeded(sessionId, budget.snapshot().spent);
+                    }
+                    if (decision === "finalize") {
+                        finalizing = true;
+                        registerBudgetFinalizeChild(sessionId, budget);
+                        if (!session.abort)
+                            throw new DreamTokenBudgetExceeded(sessionId, budget.snapshot().spent);
+                        await session.abort({ path: { id: sessionId } });
+                        // Wait for the cancelled generation to become idle before
+                        // submitting a new user turn in the same session.
+                        const abortDeadline = Date.now() + startGraceMs;
+                        while (Date.now() < abortDeadline) {
+                            const state = await readStatus(session, sessionId, dir, signal);
+                            if (state === "idle") break;
+                            await sleep(pollIntervalMs, signal);
+                        }
+                        if ((await readStatus(session, sessionId, dir, signal)) !== "idle") {
+                            throw new DreamTokenBudgetExceeded(sessionId, budget.snapshot().spent);
+                        }
+                        // An aborted assistant is not the answer to the new user turn.
+                        completionBaseline = new Set(
+                            (await readMessages(session, sessionId, dir, signal))
+                                .map(messageId)
+                                .filter((value): value is string => value !== null),
+                        );
+                        const finalizeResponse = await session.promptAsync?.({
+                            ...request,
+                            body: {
+                                ...request.body,
+                                parts: [{ type: "text", text: TOKEN_BUDGET_FINALIZE_MESSAGE }],
+                            },
+                        });
+                        if (
+                            finalizeResponse &&
+                            typeof finalizeResponse === "object" &&
+                            "error" in finalizeResponse &&
+                            finalizeResponse.error
+                        ) {
+                            throw new DreamTokenBudgetExceeded(sessionId, budget.snapshot().spent);
+                        }
+                        sentAt = Date.now();
+                        sawBusy = false;
+                        break;
+                    }
+                }
+                if (budget.snapshot().hardStopped) {
+                    await session.abort?.({ path: { id: sessionId } });
+                    throw new DreamTokenBudgetExceeded(sessionId, budget.snapshot().spent);
+                }
+            }
+            if (status === "busy" || status === "retry") {
+                sawBusy = true;
+                continue;
+            }
+            if (status === null) continue;
+            const current = finalizing
+                ? await readMessages(session, sessionId, dir, signal)
+                : messages;
+            const last = current.at(-1);
+            const lastId = messageId(last);
+            if (
+                last &&
+                lastId !== null &&
+                !completionBaseline.has(lastId) &&
+                isSettledAssistant(last)
+            )
+                return;
+            const fresh = current.some((message) => {
+                const id = messageId(message);
+                return id !== null && !completionBaseline.has(id);
+            });
+            if (sawBusy && fresh) return;
+            if (Date.now() - sentAt >= startGraceMs) {
+                if (fresh) return;
+                throw new Error(
+                    `prompt_async did not start a run in child session ${sessionId} within ${startGraceMs}ms`,
+                );
+            }
         }
+    } finally {
+        if (budget) releaseBudgetFinalizeChild(sessionId);
     }
 }

package/src/shared/rpc-async-probes.test.ts

@@ -0,0 +1,230 @@
+import { afterEach, expect, test } from "bun:test";
+import type { execFileSync } from "node:child_process";
+import { mkdirSync, mkdtempSync, rmSync, writeFileSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import {
+    __resetSchemaFenceStateForTests,
+    getMigrationOnOpenRefusal,
+    openDatabaseAsync,
+} from "../features/magic-context/storage-db";
+import { createV2StorageGate, probeV2StorageAtBoot } from "../v2/hooks/storage-gate";
+import {
+    __resetRpcIdentityTestHooks,
+    __setAsyncProcessProbeForTests,
+    __setRpcIdentityTestHooks,
+    inspectProcessesAsync,
+} from "./rpc-utils";
+import { Database } from "./sqlite";
+
+afterEach(() => {
+    __resetSchemaFenceStateForTests();
+    __resetRpcIdentityTestHooks();
+    __setAsyncProcessProbeForTests();
+});
+
+test("async Windows inspection bypasses slow synchronous probes and shares a bounded snapshot", async () => {
+    let synchronousCalls = 0;
+    const slow = (() => {
+        synchronousCalls++;
+        Atomics.wait(new Int32Array(new SharedArrayBuffer(4)), 0, 0, 5000);
+        return "";
+    }) as typeof execFileSync;
+    __setRpcIdentityTestHooks({
+        platform: "win32",
+        execFileSync: slow,
+        processListExecFileSync: slow,
+    });
+    let release!: (value: string) => void;
+    let calls = 0;
+    __setAsyncProcessProbeForTests(async (file, _args, timeout) => {
+        expect(file).toBe("powershell");
+        expect(timeout).toBe(5000);
+        calls++;
+        return new Promise<string>((resolve) => {
+            release = resolve;
+        });
+    });
+    const first = inspectProcessesAsync();
+    const second = inspectProcessesAsync();
+    expect(second).toBe(first);
+    const heartbeat = await Promise.race([
+        first.then(() => "probe"),
+        new Promise<string>((resolve) => setTimeout(() => resolve("heartbeat"), 10)),
+    ]);
+    expect(heartbeat).toBe("heartbeat");
+    release(
+        JSON.stringify([
+            {
+                ProcessId: 12345,
+                ParentProcessId: 1,
+                CommandLine: "opencode serve",
+                CreationDate: "2026-09-28T00:00:00Z",
+            },
+        ]),
+    );
+    const result = await first;
+    expect(result.liveness(12345)).toBe("alive");
+    expect(result.evidence(12345).commandLine).toBe("opencode serve");
+    expect(result.liveness(12346)).toBe("dead");
+    expect(await inspectProcessesAsync()).toBe(result);
+    expect(calls).toBe(1);
+    expect(synchronousCalls).toBe(0);
+});
+
+test("offline process inspection can bypass a cached snapshot", async () => {
+    __setRpcIdentityTestHooks({ platform: "win32" });
+    let calls = 0;
+    __setAsyncProcessProbeForTests(async () => {
+        calls++;
+        return JSON.stringify([
+            {
+                ProcessId: calls + 100,
+                ParentProcessId: 1,
+                Name: "opencode.exe",
+                CommandLine: null,
+                CreationDate: "2026-09-28T00:00:00Z",
+            },
+        ]);
+    });
+    const cached = await inspectProcessesAsync();
+    expect((await inspectProcessesAsync()).processSnapshot).toBe(cached.processSnapshot);
+    const fresh = await inspectProcessesAsync(true);
+    expect(fresh.processSnapshot?.source).toBe("cim");
+    expect(fresh.processSnapshot?.facts).toContainEqual(
+        expect.objectContaining({ pid: 102, imageName: "opencode.exe", commandLine: null }),
+    );
+    expect(calls).toBe(2);
+});
+
+test("async Windows inspection bounds CIM and tasklist fallback and caches failure", async () => {
+    __setRpcIdentityTestHooks({ platform: "win32" });
+    const calls: Array<[string, number]> = [];
+    __setAsyncProcessProbeForTests(async (file, _args, timeout) => {
+        calls.push([file, timeout]);
+        throw new Error("probe timed out");
+    });
+    const result = await inspectProcessesAsync();
+    expect(result.pi.state).toBe("unreadable");
+    expect(result.liveness(123)).toBe("inconclusive");
+    expect(result.evidence(123)).toEqual({ startTime: null, commandLine: null });
+    await inspectProcessesAsync();
+    expect(calls).toEqual([
+        ["powershell", 5000],
+        ["tasklist", 1000],
+    ]);
+});
+
+test("boot storage wait stays responsive with a locked v90 store and slow Windows probes", async () => {
+    const root = mkdtempSync(join(tmpdir(), "async-storage-guard-"));
+    const dbPath = join(root, "context.db");
+    const seeded = new Database(dbPath);
+    seeded.exec(
+        "CREATE TABLE schema_migrations(version INTEGER PRIMARY KEY); INSERT INTO schema_migrations VALUES(90)",
+    );
+    seeded.exec("BEGIN IMMEDIATE");
+    mkdirSync(join(root, "rpc", "blocker"), { recursive: true });
+    writeFileSync(
+        join(root, "rpc", "blocker", "port-12345.json"),
+        JSON.stringify({
+            pid: 12345,
+            port: 43123,
+            started_at: Date.now(),
+            kind: "OpenCode server",
+        }),
+    );
+    let syncCalls = 0;
+    const slow = (() => {
+        syncCalls++;
+        Atomics.wait(new Int32Array(new SharedArrayBuffer(4)), 0, 0, 5000);
+        return "";
+    }) as typeof execFileSync;
+    __setRpcIdentityTestHooks({
+        platform: "win32",
+        execFileSync: slow,
+        processListExecFileSync: slow,
+    });
+    let release!: (value: string) => void;
+    __setAsyncProcessProbeForTests(
+        () =>
+            new Promise<string>((resolve) => {
+                release = resolve;
+            }),
+    );
+    try {
+        const gate = createV2StorageGate({
+            open: () => openDatabaseAsync({ dbPath, busyTimeoutMs: 0 }),
+        });
+        const started = performance.now();
+        const boot = probeV2StorageAtBoot(gate);
+        const opening = gate.probe();
+        expect(
+            await Promise.race([
+                opening.then(() => "opened"),
+                new Promise<string>((resolve) => setTimeout(() => resolve("responsive"), 10)),
+            ]),
+        ).toBe("responsive");
+        expect(performance.now() - started).toBeLessThan(300);
+        expect(gate.current()).toBeUndefined();
+        release(
+            JSON.stringify([
+                {
+                    ProcessId: 12345,
+                    ParentProcessId: 1,
+                    CommandLine: "opencode serve",
+                    CreationDate: "2026-01-01T00:00:00Z",
+                },
+            ]),
+        );
+        expect(await opening).toBeUndefined();
+        expect(await boot).toBeUndefined();
+        expect(gate.reason()?.kind).toBe("migration_guard");
+        expect(getMigrationOnOpenRefusal()?.serverPids).toEqual([12345]);
+        expect(syncCalls).toBe(0);
+        const checked = new Database(dbPath);
+        expect(
+            checked.prepare("SELECT MAX(version) AS version FROM schema_migrations").get(),
+        ).toEqual({ version: 90 });
+        checked.close();
+    } finally {
+        seeded.close();
+        rmSync(root, { recursive: true, force: true });
+    }
+});
+
+test("a timed-out refresh retains confirmed blockers until a successful process scan", async () => {
+    let now = 1000;
+    __setRpcIdentityTestHooks({ platform: "win32", nowMs: () => now });
+    let fail = false;
+    let calls = 0;
+    let pid = 12345;
+    __setAsyncProcessProbeForTests(async () => {
+        calls++;
+        if (fail) throw new Error("timeout");
+        return JSON.stringify([
+            {
+                ProcessId: pid,
+                ParentProcessId: 1,
+                CommandLine: "opencode serve",
+                CreationDate: "2026-01-01T00:00:00Z",
+            },
+        ]);
+    });
+    const first = await inspectProcessesAsync();
+    now += 1999;
+    expect(await inspectProcessesAsync()).toBe(first);
+    expect(calls).toBe(1);
+    now++;
+    fail = true;
+    const failed = await inspectProcessesAsync();
+    expect(calls).toBe(3);
+    expect(failed.liveness(12345)).toBe("alive");
+    expect(failed.liveness(54321)).toBe("inconclusive");
+    expect(failed.evidence(12345)).toEqual(first.evidence(12345));
+    now += 2000;
+    fail = false;
+    pid = 54321;
+    const refreshed = await inspectProcessesAsync();
+    expect(calls).toBe(4);
+    expect(refreshed.liveness(12345)).toBe("dead");
+});

package/src/shared/rpc-server.ts

@@ -19,7 +19,7 @@ import {
     registerNotificationSink,
 } from "./rpc-notifications";
 import {
-    isPidAlive,
+    inspectProcessesAsync,
     parseRpcPortFile,
     registerOwnRpcServerInstance,
     rpcPortDir,
@@ -176,7 +176,7 @@ export class MagicContextRpcServer {
         // supported: TUI discovery scans all live files and picks the most
         // recent instead of cross-wiring via one shared project file.
         try {
-            this.warnIfOtherLiveInstance();
+            void this.warnIfOtherLiveInstance();
             const dir = dirname(this.portFilePath);
             // The port file carries the RPC bearer token. The normal policy keeps
             // it owner-only; a trusted-group deployment explicitly delegates every
@@ -263,12 +263,14 @@ export class MagicContextRpcServer {
         this.unregisterOwnInstance = null;
     }
 
-    private warnIfOtherLiveInstance(): void {
+    private async warnIfOtherLiveInstance(): Promise<void> {
         try {
             for (const entry of readdirSync(this.portDir)) {
                 if (!entry.startsWith("port-") || !entry.endsWith(".json")) continue;
                 const record = parseRpcPortFile(readFileSync(`${this.portDir}/${entry}`, "utf-8"));
-                if (!record || record.pid === process.pid || !isPidAlive(record.pid)) continue;
+                if (!record || record.pid === process.pid) continue;
+                const processes = await inspectProcessesAsync();
+                if (processes.liveness(record.pid) !== "alive") continue;
                 log(
                     `[rpc] another Magic Context RPC server is active for this project (pid ${record.pid}, port ${record.port}); starting separate instance on a new port`,
                 );

package/src/shared/rpc-utils.test.ts

@@ -183,6 +183,32 @@ describe("discoverLivePiProcessIds", () => {
 });
 
 describe("#411 Windows OMP/Pi live-process scan", () => {
+    test("reuses a Windows snapshot briefly, then refreshes process discovery", () => {
+        let now = PLUGIN_BUILD_MS;
+        const calls: string[] = [];
+        __setRpcIdentityTestHooks({
+            platform: "win32",
+            nowMs: () => now,
+            processListExecFileSync: windowsProcessListExec((file) => {
+                calls.push(file);
+                return cimOutput([
+                    {
+                        ProcessId: WINDOWS_FOREIGN_PID,
+                        ParentProcessId: 4,
+                        CommandLine: OMP_PI_ARC_COMMAND,
+                        CreationDate: wmiCreationDate(OLDER_THAN_PLUGIN_MS),
+                    },
+                ]);
+            }),
+        });
+        expect(inspectLivePiProcesses().processIds).toEqual([WINDOWS_FOREIGN_PID]);
+        now += 1_999;
+        expect(inspectLivePiProcesses().processIds).toEqual([WINDOWS_FOREIGN_PID]);
+        expect(calls).toEqual(["powershell"]);
+        now += 1;
+        inspectLivePiProcesses();
+        expect(calls).toEqual(["powershell", "powershell"]);
+    });
     test("(a) tasklist-only fallback skips a parent omp.exe as an ancestor", () => {
         const calls: string[] = [];
         __setRpcIdentityTestHooks({
@@ -607,7 +633,7 @@ describe("isPidIdentityPlausible", () => {
                 args: [
                     "-NoProfile",
                     "-Command",
-                    "Get-CimInstance Win32_Process | Select-Object ProcessId,ParentProcessId,CommandLine,CreationDate | ConvertTo-Json -Compress",
+                    "Get-CimInstance Win32_Process | Select-Object ProcessId,ParentProcessId,Name,CommandLine,CreationDate | ConvertTo-Json -Compress",
                 ],
             },
         ]);