@cortexkit/opencode-magic-context 0.44.1 → 0.44.3

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

@@ -1,5 +1,6 @@
 import { describe, expect, mock, test } from "bun:test";
-
+import { refuseBudgetedToolCall } from "../features/magic-context/dreamer/token-budget";
+import { createDroppedInputToolExecuteBeforeHook } from "../hooks/magic-context/dropped-input-guard";
 import { promptSyncWithValidatedOutputRetry } from "./model-suggestion-retry";
 import {
     createPromptAsyncTransport,
@@ -188,6 +189,258 @@ describe("promptAsyncAndWaitForIdle", () => {
     });
 });
 
+describe("dreamer budget on an async child", () => {
+    test("sums prompt usage across validation retries in the same child", async () => {
+        const messages: unknown[] = [];
+        const states: Array<{ spent: number }> = [];
+        let sends = 0;
+        let aborted = false;
+        const client = {
+            session: {
+                messages: async () => ({ data: [...messages] }),
+                promptAsync: async (req: { body: { parts: unknown[] } }) => {
+                    sends++;
+                    messages.push({
+                        info: { id: `user-${sends}`, role: "user" },
+                        parts: req.body.parts,
+                    });
+                    return { data: undefined };
+                },
+                abort: async () => {
+                    aborted = true;
+                    return { data: true };
+                },
+                status: async () => {
+                    const latest = messages.at(-1) as { info?: { id?: string } } | undefined;
+                    if (sends === 1 && latest?.info?.id === "user-1")
+                        messages.push({
+                            info: {
+                                id: "a1",
+                                role: "assistant",
+                                finish: "stop",
+                                tokens: { input: 60 },
+                                time: { created: 1, completed: 2 },
+                            },
+                            parts: [{ type: "text", text: "invalid manifest" }],
+                        });
+                    if (sends === 2 && latest?.info?.id === "user-2")
+                        messages.push({
+                            info: {
+                                id: "a2",
+                                role: "assistant",
+                                finish: "tool-calls",
+                                tokens: { input: 25 },
+                                time: { created: 3, completed: 4 },
+                            },
+                            parts: [{ type: "tool" }],
+                        });
+                    if (sends === 3 && latest?.info?.id === "user-3")
+                        messages.push({
+                            info: {
+                                id: "a3",
+                                role: "assistant",
+                                finish: "stop",
+                                tokens: { input: 1 },
+                                time: { created: 5, completed: 6 },
+                            },
+                            parts: [{ type: "text", text: "<mappings/>" }],
+                        });
+                    return {
+                        data: sends === 2 && !aborted ? { "ses-child": { type: "busy" } } : {},
+                    };
+                },
+            },
+        } as never;
+        const transport = createPromptAsyncTransport(client, "ses-child", {
+            pollIntervalMs: 1,
+            tokenBudget: 100,
+            onBudgetUpdate: (state) => states.push(state),
+        });
+        if (!transport) throw new Error("missing async transport");
+        await transport(request());
+        await transport(request());
+        expect(sends).toBe(3);
+        expect(states.at(-1)?.spent).toBe(86);
+    });
+    test("does not send a finalize turn when the child finishes below budget", async () => {
+        const host = asyncHost({ plan: ["busy", "settle"] });
+        const originalMessages = host.client.session.messages;
+        host.client.session.messages = mock(
+            async (...args: Parameters<typeof originalMessages>) => {
+                const response = await originalMessages(...args);
+                for (const message of response.data) {
+                    if (message.info.role === "assistant") message.info.tokens = { input: 39 };
+                }
+                return response;
+            },
+        );
+        await promptAsyncAndWaitForIdle(host.client as never, request(), {
+            pollIntervalMs: 1,
+            tokenBudget: 100,
+        });
+        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" }],
+            },
+        };
+        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: {
+                id: string;
+                role: string;
+                tokens?: unknown;
+                time?: { completed: number };
+                finish?: string;
+            };
+            parts: unknown[];
+        }> = [];
+        const toolsByAgent: Record<string, unknown[]> = {
+            "dreamer-memory-mapper": [{ name: "read", parameters: { type: "object" } }],
+        };
+        const capturedTools: string[] = [];
+        const prompts: unknown[] = [];
+        let stage = 0;
+        const client = {
+            session: {
+                messages: async () => ({ data: [...messages] }),
+                promptAsync: async (req: unknown) => {
+                    prompts.push(req);
+                    const agent = (req as { body: { agent: string } }).body.agent;
+                    capturedTools.push(JSON.stringify(toolsByAgent[agent] ?? []));
+                    if (prompts.length === 1)
+                        messages.push({ info: { id: "u1", role: "user" }, parts: [] });
+                    else messages.push({ info: { id: "u2", role: "user" }, parts: [] });
+                    return { data: undefined };
+                },
+                abort: mock(async () => ({ data: true })),
+                status: async () => {
+                    stage++;
+                    if (stage === 1)
+                        messages.push({
+                            info: {
+                                id: "a1",
+                                role: "assistant",
+                                tokens: { input: 10, cache: { read: 71, write: 0 } },
+                                time: { completed: 1 },
+                                finish: "tool-calls",
+                            },
+                            parts: [],
+                        });
+                    if (prompts.length === 2 && !messages.some((m) => m.info.id === "a2"))
+                        messages.push({
+                            info: {
+                                id: "a2",
+                                role: "assistant",
+                                tokens: { input: 1, cache: { read: 1, write: 0 } },
+                                time: { completed: 2 },
+                                finish: "stop",
+                            },
+                            parts: [{ type: "text", text: "<manifest/>" }],
+                        });
+                    return { data: stage === 1 ? { "ses-child": { type: "busy" } } : {} };
+                },
+            },
+        } as never;
+        const run = promptAsyncAndWaitForIdle(
+            client,
+            {
+                ...request(),
+                body: { agent: "dreamer-memory-mapper", parts: [{ type: "text", text: "go" }] },
+            },
+            { pollIntervalMs: 1, tokenBudget: 100 },
+        );
+        // Wait for the follow-up to be enqueued before simulating the tool hook.
+        while (prompts.length < 2) await new Promise((resolve) => setTimeout(resolve, 1));
+        const before = createDroppedInputToolExecuteBeforeHook();
+        await expect(
+            before({ sessionID: "ses-child", tool: "read" }, { args: { path: "a" } }),
+        ).rejects.toThrow("Out of token budget");
+        await expect(
+            before({ sessionID: "ses-child", tool: "read" }, { args: { path: "b" } }),
+        ).rejects.toThrow("Out of token budget");
+        await expect(run).rejects.toThrow("token_budget");
+        expect(prompts).toHaveLength(2);
+        expect(capturedTools[0]).not.toBe("[]");
+        expect(capturedTools[1]).toBe(capturedTools[0]);
+        expect(
+            (prompts[1] as { body: { parts: Array<{ text: string }> } }).body.parts[0]?.text,
+        ).toContain("no more tool calls");
+        expect(refuseBudgetedToolCall("ses-child")).toBeNull();
+    });
+});
+
 describe("prompt_async transport under the retry chain", () => {
     test("our slice stays the only timer and aborts the busy child when it expires", async () => {
         const host = asyncHost({ plan: ["busy"] });

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);
     }
 }