@cortexkit/opencode-magic-context 0.44.3 → 0.45.0

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

@@ -5,6 +5,7 @@ import { promptSyncWithValidatedOutputRetry } from "./model-suggestion-retry";
 import {
     createPromptAsyncTransport,
     promptAsyncAndWaitForIdle,
+    recordPromptSessionError,
     supportsPromptAsync,
 } from "./prompt-async-transport";
 
@@ -190,6 +191,96 @@ describe("promptAsyncAndWaitForIdle", () => {
 });
 
 describe("dreamer budget on an async child", () => {
+    test("forwards Pi's finalized per-child budget before manifest application", async () => {
+        let ready = false;
+        const states: unknown[] = [];
+        const client = {
+            session: {
+                prompt: async () => {
+                    ready = true;
+                },
+            },
+            readChildTokenBudget: (id: string) => {
+                expect(id).toBe("ses-child");
+                expect(ready).toBe(true);
+                return { spent: 58, finalizeFired: true };
+            },
+        };
+        const transport = createPromptAsyncTransport(client as never, "ses-child", {
+            onBudgetUpdate: (state) => states.push(state),
+        });
+        expect(transport).toBeDefined();
+        await transport?.(request());
+        expect(states).toEqual([{ spent: 58, finalizeFired: true, sessionId: "ses-child" }]);
+    });
+    test("finalizes a zero-usage OpenCode mapper at 58 steps, before the hard ceiling", async () => {
+        const messages: unknown[] = [];
+        let sends = 0;
+        let aborted = false;
+        const states: unknown[] = [];
+        const client = {
+            session: {
+                messages: async () => ({ data: [...messages] }),
+                promptAsync: async (req: { body: { parts: unknown[] } }) => {
+                    sends++;
+                    if (sends === 2)
+                        expect(JSON.stringify(req.body.parts)).toContain("no more tool calls");
+                    messages.push({
+                        info: { id: `u${sends}`, role: "user" },
+                        parts: req.body.parts,
+                    });
+                },
+                abort: async () => {
+                    aborted = true;
+                },
+                status: async () => {
+                    if (sends === 1 && messages.length === 1) {
+                        for (let step = 1; step <= 58; step++)
+                            messages.push({
+                                info: {
+                                    id: `a${step}`,
+                                    role: "assistant",
+                                    finish: "tool-calls",
+                                    tokens: { input: 0 },
+                                    time: { completed: 1 },
+                                },
+                                parts: [{ type: "tool" }],
+                            });
+                    }
+                    if (sends === 2 && messages.length === 60)
+                        messages.push({
+                            info: {
+                                id: "final",
+                                role: "assistant",
+                                finish: "stop",
+                                tokens: { input: 1 },
+                                time: { completed: 2 },
+                            },
+                            parts: [{ type: "text", text: "<verify/>" }],
+                        });
+                    return {
+                        data: sends === 1 && !aborted ? { "ses-child": { type: "busy" } } : {},
+                    };
+                },
+            },
+        };
+        await promptAsyncAndWaitForIdle(
+            client as never,
+            {
+                ...request(),
+                signal: AbortSignal.timeout(200),
+                body: { ...request().body, agent: "dreamer-memory-mapper" },
+            },
+            {
+                tokenBudget: 3_000_000,
+                pollIntervalMs: 1,
+                startGraceMs: 100,
+                onBudgetUpdate: (state) => states.push(state),
+            },
+        );
+        expect(sends).toBe(2);
+        expect(states).toContainEqual(expect.objectContaining({ finalizeFired: true, spent: 0 }));
+    });
     test("sums prompt usage across validation retries in the same child", async () => {
         const messages: unknown[] = [];
         const states: Array<{ spent: number }> = [];
@@ -281,80 +372,121 @@ describe("dreamer budget on an async child", () => {
         expect(host.promptAsync).toHaveBeenCalledTimes(1);
         expect(host.abort).not.toHaveBeenCalled();
     });
-    test("returns the final manifest after a soft nudge without changing the child agent", async () => {
-        const host = asyncHost({ plan: ["busy", "settle"] });
-        const messages: Array<{
-            info: {
-                id: string;
-                role: string;
-                tokens?: unknown;
-                time?: { completed: number };
-                finish?: string;
-            };
-            parts: unknown[];
-        }> = [];
-        let sends = 0;
-        const client = {
-            session: {
-                messages: async () => ({ data: [...messages] }),
-                promptAsync: mock(async (_req: unknown) => {
-                    sends++;
-                    messages.push({ info: { id: `u${sends}`, role: "user" }, parts: [] });
-                    return { data: undefined };
-                }),
-                abort: host.abort,
-                status: async () => {
-                    if (sends === 1 && !messages.some((m) => m.info.id === "a1")) {
-                        messages.push({
-                            info: {
-                                id: "a1",
-                                role: "assistant",
-                                tokens: { input: 81 },
-                                time: { completed: 1 },
-                                finish: "tool-calls",
-                            },
-                            parts: [],
-                        });
-                    }
-                    if (sends === 2 && !messages.some((m) => m.info.id === "a2")) {
-                        messages.push({
-                            info: {
-                                id: "a2",
-                                role: "assistant",
-                                tokens: { input: 1 },
-                                time: { completed: 2 },
-                                finish: "stop",
-                            },
-                            parts: [{ type: "text", text: "<mappings/>" }],
-                        });
-                    }
-                    return {
-                        data:
-                            sends === 1 && host.abort.mock.calls.length === 0
-                                ? { "ses-child": { type: "busy" } }
-                                : {},
-                    };
-                },
-            },
-        } as never;
-        const body = {
-            ...request(),
-            body: {
-                agent: "dreamer-memory-mapper",
-                system: "same system",
-                parts: [{ type: "text", text: "work" }],
+    test("keeps a completed answer when a poll receives over-budget tool usage and completion together", async () => {
+        const host = asyncHost({ plan: ["settle"] });
+        const originalMessages = host.client.session.messages;
+        host.client.session.messages = mock(
+            async (...args: Parameters<typeof originalMessages>) => {
+                const response = await originalMessages(...args);
+                if (response.data.some((message) => message.info.role === "assistant")) {
+                    response.data.unshift({
+                        info: {
+                            id: "coarse-tool",
+                            role: "assistant",
+                            finish: "tool-calls",
+                            tokens: { input: 120 },
+                            time: { created: 1, completed: 2 },
+                        },
+                        parts: [],
+                    } as never);
+                }
+                return response;
             },
-        };
-        await promptAsyncAndWaitForIdle(client, body, { pollIntervalMs: 1, tokenBudget: 100 });
-        const prompts = (client as { session: { promptAsync: ReturnType<typeof mock> } }).session
-            .promptAsync.mock.calls;
-        expect(prompts).toHaveLength(2);
-        expect((prompts[1]?.[0] as typeof body).body.agent).toBe(
-            (prompts[0]?.[0] as typeof body).body.agent,
         );
-        expect(messages.at(-1)?.parts).toEqual([{ type: "text", text: "<mappings/>" }]);
-        expect(host.abort).toHaveBeenCalledTimes(1);
+        await promptAsyncAndWaitForIdle(host.client as never, request(), {
+            pollIntervalMs: 1,
+            tokenBudget: 100,
+        });
+        expect(host.promptAsync).toHaveBeenCalledTimes(1);
+        expect(host.abort).not.toHaveBeenCalled();
     });
+
+    for (const finalUsage of [1, 25]) {
+        test(`returns the final manifest after a soft nudge with ${finalUsage} final tokens`, async () => {
+            const host = asyncHost({ plan: ["busy", "settle"] });
+            const messages: Array<{
+                info: {
+                    id: string;
+                    role: string;
+                    tokens?: unknown;
+                    time?: { completed: number };
+                    finish?: string;
+                };
+                parts: unknown[];
+            }> = [];
+            let sends = 0;
+            const client = {
+                session: {
+                    messages: async () => ({ data: [...messages] }),
+                    promptAsync: mock(async (_req: unknown) => {
+                        sends++;
+                        messages.push({ info: { id: `u${sends}`, role: "user" }, parts: [] });
+                        return { data: undefined };
+                    }),
+                    abort: async () => {
+                        recordPromptSessionError("ses-child", {
+                            name: "MessageAbortedError",
+                            data: { message: "Aborted" },
+                        });
+                        return host.abort();
+                    },
+                    status: async () => {
+                        if (sends === 1 && !messages.some((m) => m.info.id === "a1")) {
+                            messages.push({
+                                info: {
+                                    id: "a1",
+                                    role: "assistant",
+                                    tokens: { input: 81 },
+                                    time: { completed: 1 },
+                                    finish: "tool-calls",
+                                },
+                                parts: [],
+                            });
+                        }
+                        if (sends === 2 && !messages.some((m) => m.info.id === "a2")) {
+                            if (finalUsage === 25) {
+                                refuseBudgetedToolCall("ses-child");
+                                refuseBudgetedToolCall("ses-child");
+                            }
+                            messages.push({
+                                info: {
+                                    id: "a2",
+                                    role: "assistant",
+                                    tokens: { input: finalUsage },
+                                    time: { completed: 2 },
+                                    finish: "stop",
+                                },
+                                parts: [{ type: "text", text: "<mappings/>" }],
+                            });
+                        }
+                        return {
+                            data:
+                                sends === 1 && host.abort.mock.calls.length === 0
+                                    ? { "ses-child": { type: "busy" } }
+                                    : {},
+                        };
+                    },
+                },
+            } as never;
+            const body = {
+                ...request(),
+                body: {
+                    agent: "dreamer-memory-mapper",
+                    system: "same system",
+                    parts: [{ type: "text", text: "work" }],
+                },
+            };
+            await promptAsyncAndWaitForIdle(client, body, { pollIntervalMs: 1, tokenBudget: 100 });
+            const prompts = (client as { session: { promptAsync: ReturnType<typeof mock> } })
+                .session.promptAsync.mock.calls;
+            expect(prompts).toHaveLength(2);
+            expect((prompts[1]?.[0] as typeof body).body.agent).toBe(
+                (prompts[0]?.[0] as typeof body).body.agent,
+            );
+            expect(messages.at(-1)?.parts).toEqual([{ type: "text", text: "<mappings/>" }]);
+            expect(host.abort).toHaveBeenCalledTimes(1);
+        });
+    }
     test("finalizes once, refuses two tool calls, and keeps agent and tools unchanged", async () => {
         const messages: Array<{
             info: {
@@ -405,7 +537,7 @@ describe("dreamer budget on an async child", () => {
                                 role: "assistant",
                                 tokens: { input: 1, cache: { read: 1, write: 0 } },
                                 time: { completed: 2 },
-                                finish: "stop",
+                                finish: "tool-calls",
                             },
                             parts: [{ type: "text", text: "<manifest/>" }],
                         });
@@ -468,3 +600,48 @@ describe("prompt_async transport under the retry chain", () => {
         expect(host.abort).toHaveBeenCalledWith({ path: { id: "ses-child" } });
     });
 });
+
+test("an idle accepted user row is not a no-output completion after 30 seconds", async () => {
+    const host = asyncHost({ plan: ["idle-pending"] });
+    await expect(
+        promptAsyncAndWaitForIdle(host.client, request(), {
+            pollIntervalMs: 1,
+            startGraceMs: 20,
+        }),
+    ).rejects.toThrow("settled assistant");
+});
+
+test("an early child session error preserves DeepSeek's unsupported-model 400 without waiting for an assistant", async () => {
+    const host = asyncHost({ plan: ["idle-pending"] });
+    const message =
+        "The supported API model names are deepseek-flash, deepseek-v4-pro, but you passed deepseek-v4-flash";
+    host.promptAsync.mockImplementation(async () => {
+        recordPromptSessionError("ses-child", {
+            name: "APIError",
+            data: { statusCode: 400, message },
+        });
+        recordPromptSessionError("ses-child", {
+            name: "UnknownError",
+            data: { message: "duplicate stack" },
+        });
+        return { data: undefined };
+    });
+    await expect(
+        promptAsyncAndWaitForIdle(host.client, request(), {
+            pollIntervalMs: 1,
+            startGraceMs: 60000,
+        }),
+    ).rejects.toThrow(`${message} (status=400)`);
+    expect(host.polls()).toBe(0);
+});
+
+test("session errors belonging to another child do not fail the run", async () => {
+    const host = asyncHost({ plan: ["settle"] });
+    const send = host.promptAsync.getMockImplementation()!;
+    host.promptAsync.mockImplementation(async (req) => {
+        recordPromptSessionError("other-child", "unrelated");
+        return send(req);
+    });
+    await promptAsyncAndWaitForIdle(host.client, request(), { pollIntervalMs: 1 });
+    expect(host.polls()).toBe(1);
+});

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

@@ -1,4 +1,5 @@
 import type { createOpencodeClient } from "@opencode-ai/sdk";
+import { buildHiddenAgentRegistrations } from "../agents/hidden-agent-registrations";
 import {
     createDreamTokenBudget,
     DreamTokenBudgetExceeded,
@@ -7,10 +8,28 @@ import {
     TOKEN_BUDGET_FINALIZE_MESSAGE,
 } from "../features/magic-context/dreamer/token-budget";
 import { sumTokensFromChildMessages } from "../features/magic-context/subagent-token-capture";
+import { describeAssistantError } from "./assistant-message-extractor";
+import { sessionLog } from "./logger";
 import type { PromptArgs, PromptTransport } from "./model-suggestion-retry";
 
 type Client = ReturnType<typeof createOpencodeClient> | undefined;
 
+const activeSessionErrors = new Map<string, unknown>();
+const budgetAbortedSessions = new Set<string>();
+
+/** Feed host errors only to active background waits, not ordinary user sessions. */
+export function recordPromptSessionError(sessionId: string, error: unknown): void {
+    // The finalize nudge intentionally aborts the previous generation. Its late
+    // event must not poison the replacement turn; other provider errors still count.
+    if (
+        budgetAbortedSessions.has(sessionId) &&
+        describeAssistantError(error).startsWith("MessageAbortedError")
+    )
+        return;
+    if (activeSessionErrors.has(sessionId) && activeSessionErrors.get(sessionId) === undefined)
+        activeSessionErrors.set(sessionId, error);
+}
+
 /**
  * Why this exists: a synchronous `session.prompt` keeps one HTTP request open
  * for a child's whole agent loop, which for a dreamer batch is many minutes.
@@ -29,6 +48,12 @@ 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 = 100;
+const MAPPER_STEP_CAP = buildHiddenAgentRegistrations({
+    dreamerPrompt: undefined,
+    historianPrompt: undefined,
+    historianEditorPrompt: undefined,
+    historianDisallowed: [],
+}).find((agent) => agent.id === "dreamer-memory-mapper")?.maxSteps;
 
 export interface PromptAsyncWaitOptions {
     pollIntervalMs?: number;
@@ -44,6 +69,7 @@ export interface PromptAsyncWaitOptions {
 }
 
 type SessionApi = {
+    prompt?: (options: unknown) => Promise<unknown>;
     promptAsync?: (options: unknown) => Promise<unknown>;
     status?: (options?: unknown) => Promise<unknown>;
     messages?: (options: unknown) => Promise<unknown>;
@@ -54,6 +80,18 @@ function sessionApi(client: Client): SessionApi | undefined {
     return (client as { session?: SessionApi } | undefined)?.session;
 }
 
+/** Marks the error the wait throws when the host reported a session error for the child. */
+const HOST_SESSION_ERROR = Symbol("promptAsyncHostSessionError");
+
+/** True for the error the wait throws when the host recorded a session error for the child. */
+export function isHostSessionError(error: unknown): boolean {
+    return (
+        error !== null &&
+        typeof error === "object" &&
+        (error as { [HOST_SESSION_ERROR]?: unknown })[HOST_SESSION_ERROR] === true
+    );
+}
+
 /** True when the host client can send without waiting and report session status. */
 export function supportsPromptAsync(client: Client): boolean {
     const session = sessionApi(client);
@@ -74,7 +112,31 @@ export function createPromptAsyncTransport(
     childSessionId: string,
     options: PromptAsyncWaitOptions = {},
 ): PromptTransport | undefined {
-    if (!supportsPromptAsync(client)) return undefined;
+    if (!supportsPromptAsync(client)) {
+        // Pi's synchronous facade owns the subprocess budget. Forward its final
+        // per-child state before manifest application, not just into the run ledger.
+        const readChildTokenBudget = (
+            client as unknown as
+                | {
+                      readChildTokenBudget?: (
+                          id: string,
+                      ) =>
+                          | ReturnType<ReturnType<typeof createDreamTokenBudget>["snapshot"]>
+                          | undefined;
+                  }
+                | undefined
+        )?.readChildTokenBudget;
+        const session = sessionApi(client);
+        if (!readChildTokenBudget || !session?.prompt) return undefined;
+        return Object.assign(
+            async (request: PromptArgs) => {
+                await session.prompt?.(request);
+                const state = readChildTokenBudget(childSessionId);
+                if (state) options.onBudgetUpdate?.({ ...state, sessionId: childSessionId });
+            },
+            { childSessionId },
+        );
+    }
     const budgetGuard = options.tokenBudget
         ? createDreamTokenBudget(options.tokenBudget)
         : undefined;
@@ -224,9 +286,11 @@ export async function promptAsyncAndWaitForIdle(
     const budget =
         options.budgetGuard ??
         (options.tokenBudget ? createDreamTokenBudget(options.tokenBudget) : undefined);
-    if (budget?.snapshot().finalizeFired)
+    if (budget && (budget.snapshot().finalizeFired || budget.snapshot().hardStopped))
         throw new DreamTokenBudgetExceeded(sessionId, budget.snapshot().spent);
     const seenUsage = new Set<string>();
+    const seenSteps = new Set<string>();
+    const mapper = request.body.agent === "dreamer-memory-mapper";
     let finalizing = false;
     let completionBaseline = new Set(
         (await readMessages(session, sessionId, dir, signal))
@@ -234,35 +298,74 @@ export async function promptAsyncAndWaitForIdle(
             .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)}`);
-    }
-    let sentAt = Date.now();
-    let sawBusy = false;
-
+    // Dispatch can fail before an assistant row exists. The plugin event hook
+    // retains that provider/model error while this child is actively awaited.
+    activeSessionErrors.set(sessionId, undefined);
     try {
+        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)}`);
+            }
+        }
+        let sentAt = Date.now();
+        let sawBusy = false;
+
         for (;;) {
             await sleep(pollIntervalMs, signal);
+            const hostError = activeSessionErrors.get(sessionId);
+            if (hostError !== undefined) {
+                const error = new Error(
+                    `Host recorded session error: ${describeAssistantError(hostError)}`,
+                );
+                Object.assign(error, {
+                    name: "DreamerProviderOutputFailureError",
+                    transient: true,
+                    [HOST_SESSION_ERROR]: true,
+                });
+                throw error;
+            }
             const status = await readStatus(session, sessionId, dir, signal);
             const messages = await readMessages(session, sessionId, dir, signal);
             if (budget) {
+                const latest = messages.at(-1);
+                const latestId = messageId(latest);
+                const completedAnswer =
+                    latestId !== null &&
+                    !completionBaseline.has(latestId) &&
+                    isTerminalAssistant(latest) &&
+                    infoOf(latest).error == null;
                 for (const message of messages) {
                     const id = messageId(message);
-                    if (!id || seenUsage.has(id) || infoOf(message).role !== "assistant") continue;
+                    if (!id || infoOf(message).role !== "assistant") continue;
                     // The initial baseline belongs to a previous child attempt.
                     if (initialBaseline.has(id)) continue;
+                    if (isSettledAssistant(message)) seenSteps.add(id);
+                    if (seenUsage.has(id)) continue;
                     const tokens = sumTokensFromChildMessages([message]);
-                    if (tokens.input + tokens.cacheRead + tokens.cacheWrite === 0) continue;
-                    seenUsage.add(id);
-                    const decision = budget.charge(
+                    // A streaming row may not have usage yet; let a later poll
+                    // charge its final report rather than remembering zero tokens.
+                    if (tokens.input + tokens.cacheRead + tokens.cacheWrite > 0) seenUsage.add(id);
+                    let decision = budget.charge(
                         tokens.input,
                         tokens.cacheRead,
                         tokens.cacheWrite,
-                        isTerminalAssistant(message),
+                        // A poll can receive several usage rows together. Charge
+                        // all of them, but never abort an answer already completed.
+                        completedAnswer ||
+                            (isTerminalAssistant(message) && infoOf(message).error == null),
                     );
+                    // Polling can see an in-flight turn after the newest usage row;
+                    // reserve it and a final answer before OpenCode's hard step cap.
+                    if (
+                        mapper &&
+                        MAPPER_STEP_CAP !== undefined &&
+                        !completedAnswer &&
+                        decision === "continue" &&
+                        seenSteps.size >= MAPPER_STEP_CAP - 2
+                    )
+                        decision = budget.finalize();
                     options.onBudgetUpdate?.({ ...budget.snapshot(), sessionId });
                     if (decision === "stop") {
                         await session.abort?.({ path: { id: sessionId } });
@@ -270,6 +373,12 @@ export async function promptAsyncAndWaitForIdle(
                     }
                     if (decision === "finalize") {
                         finalizing = true;
+                        sessionLog(
+                            sessionId,
+                            "dreamer token budget: finalize fired",
+                            budget.snapshot(),
+                        );
+                        budgetAbortedSessions.add(sessionId);
                         registerBudgetFinalizeChild(sessionId, budget);
                         if (!session.abort)
                             throw new DreamTokenBudgetExceeded(sessionId, budget.snapshot().spent);
@@ -311,7 +420,7 @@ export async function promptAsyncAndWaitForIdle(
                         break;
                     }
                 }
-                if (budget.snapshot().hardStopped) {
+                if (budget.snapshot().hardStopped && !completedAnswer) {
                     await session.abort?.({ path: { id: sessionId } });
                     throw new DreamTokenBudgetExceeded(sessionId, budget.snapshot().spent);
                 }
@@ -333,19 +442,17 @@ export async function promptAsyncAndWaitForIdle(
                 isSettledAssistant(last)
             )
                 return;
-            const fresh = current.some((message) => {
-                const id = messageId(message);
-                return id !== null && !completionBaseline.has(id);
-            });
-            if (sawBusy && fresh) return;
+            // An accepted user message is not a completion. A host-side rejection
+            // can leave only that row behind while the session remains idle.
             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`,
+                    `prompt_async did not start a run with a settled assistant in child session ${sessionId} within ${startGraceMs}ms (observed_busy=${sawBusy})`,
                 );
             }
         }
     } finally {
+        activeSessionErrors.delete(sessionId);
+        budgetAbortedSessions.delete(sessionId);
         if (budget) releaseBudgetFinalizeChild(sessionId);
     }
 }