@cortexkit/opencode-magic-context 0.44.0 → 0.44.2

package/src/config/live-run-config.test.ts

@@ -7,12 +7,13 @@ import {
     buildDreamTaskRuntimeConfigs,
     userMemoryCollectionEnabled,
 } from "../features/magic-context/dreamer/task-config";
+import { producerInputTokenLimit } from "../hooks/magic-context/producer-window-guard";
 import { resolveHistorianModel } from "../shared/model-resolution";
 import { loadPluginConfigDetailed } from "./index";
 import { dreamerRunConfig, historianRunConfig } from "./live-run-config";
 import { LiveConfigReader } from "./live-snapshot";
 
-for (const host of ["OC1", "OC2", "Pi"] as const) {
+for (const host of ["OC1", "OC2", "Pi", "OMP"] as const) {
     test(`${host} historian samples the next generation while an existing run keeps its model and fallback`, () => {
         const root = mkdtempSync(join(tmpdir(), "mc-live-historian-"));
         const previous = { home: process.env.HOME, config: process.env.XDG_CONFIG_HOME };
@@ -23,22 +24,23 @@ for (const host of ["OC1", "OC2", "Pi"] as const) {
         mkdirSync(join(root, "config", "cortexkit"), { recursive: true });
         mkdirSync(join(directory, ".cortexkit"), { recursive: true });
         try {
-            const block = host === "Pi" ? "pi" : "opencode";
+            const block = host === "Pi" ? "pi" : host === "OMP" ? "omp" : "opencode";
             const write = (
                 model: string,
                 fallback: string,
                 autoPromote: boolean,
                 minClusters: number,
+                maxTokens: number,
             ) =>
                 writeFileSync(
                     file,
                     JSON.stringify({
-                        historian: { [block]: { model, fallback_models: [fallback] } },
+                        historian: { maxTokens, [block]: { model, fallback_models: [fallback] } },
                         memory: { auto_promote: autoPromote },
                         commit_cluster_trigger: { enabled: true, min_clusters: minClusters },
                     }),
                 );
-            write("anthropic/old-model", "anthropic/old-fallback", true, 3);
+            write("anthropic/old-model", "anthropic/old-fallback", true, 3, 4096);
             const load = () => loadPluginConfigDetailed(directory, false).config;
             const boot = load();
             const reader = new LiveConfigReader(directory, boot, load, () => {});
@@ -49,9 +51,10 @@ for (const host of ["OC1", "OC2", "Pi"] as const) {
                 "anthropic/new-fallback-with-longer-name",
                 false,
                 5,
+                8192,
             );
             const runTwo = historianRunConfig(boot, reader.poll().effective);
-            const harness = host === "Pi" ? "pi" : "opencode";
+            const harness = host === "Pi" ? "pi" : host === "OMP" ? "omp" : "opencode";
             expect(resolveHistorianModel(runOne, harness).primary?.model).toBe(
                 "anthropic/old-model",
             );
@@ -68,6 +71,9 @@ for (const host of ["OC1", "OC2", "Pi"] as const) {
             expect(runTwo.memory.auto_promote).toBe(false);
             expect(runOne.commit_cluster_trigger.min_clusters).toBe(3);
             expect(runTwo.commit_cluster_trigger.min_clusters).toBe(5);
+            expect(runOne.historian?.maxTokens).toBe(4096);
+            expect(runTwo.historian?.maxTokens).toBe(8192);
+            expect(producerInputTokenLimit(32_000, runTwo.historian?.maxTokens ?? 0)).toBe(23_093);
             expect(reader.current().generation).toBe(2);
         } finally {
             if (previous.home === undefined) delete process.env.HOME;
@@ -79,7 +85,7 @@ for (const host of ["OC1", "OC2", "Pi"] as const) {
     });
 }
 
-for (const host of ["OC1", "OC2", "Pi"] as const) {
+for (const host of ["OC1", "OC2", "Pi", "OMP"] as const) {
     test(`${host} dreamer samples the new schedule and model chain without changing a running task`, () => {
         const root = mkdtempSync(join(tmpdir(), "mc-live-dreamer-"));
         const previous = { home: process.env.HOME, config: process.env.XDG_CONFIG_HOME };
@@ -91,11 +97,14 @@ for (const host of ["OC1", "OC2", "Pi"] as const) {
         mkdirSync(join(root, "config", "cortexkit"), { recursive: true });
         mkdirSync(join(directory, ".cortexkit"), { recursive: true });
         try {
-            const block = host === "Pi" ? "pi" : "opencode";
+            const block = host === "Pi" ? "pi" : host === "OMP" ? "omp" : "opencode";
             writeFileSync(
                 userFile,
                 JSON.stringify({
-                    dreamer: { [block]: { model: "old/model", fallback_models: ["old/fallback"] } },
+                    dreamer: {
+                        maxTokens: 4096,
+                        [block]: { model: "old/model", fallback_models: ["old/fallback"] },
+                    },
                 }),
             );
             const writeSchedule = (schedule: string) =>
@@ -122,6 +131,7 @@ for (const host of ["OC1", "OC2", "Pi"] as const) {
                 userFile,
                 JSON.stringify({
                     dreamer: {
+                        maxTokens: 8192,
                         [block]: {
                             model: "new/model-long",
                             fallback_models: ["new/fallback-long"],
@@ -132,14 +142,17 @@ for (const host of ["OC1", "OC2", "Pi"] as const) {
             writeSchedule("15 4 * * *");
             const runTwo = dreamerRunConfig(boot, reader.poll().effective);
             const task = (cfg: typeof runOne) =>
-                buildDreamTaskRuntimeConfigs(cfg.dreamer, host === "Pi" ? "pi" : "opencode").find(
-                    (entry) => entry.task === "verify",
-                )!;
+                buildDreamTaskRuntimeConfigs(
+                    cfg.dreamer,
+                    host === "Pi" ? "pi" : host === "OMP" ? "omp" : "opencode",
+                ).find((entry) => entry.task === "verify")!;
             expect(task(runOne).schedule).toBe("0 3 * * *");
             expect(userMemoryCollectionEnabled(runOne.dreamer)).toBe(false);
             expect(task(runOne).model?.model).toBe("old/model");
             expect(task(runOne).fallbackModels[0]?.model).toBe("old/fallback");
             expect(task(runTwo).schedule).toBe("15 4 * * *");
+            expect(runOne.dreamer?.maxTokens).toBe(4096);
+            expect(runTwo.dreamer?.maxTokens).toBe(8192);
             expect(
                 userMemoryCollectionEnabled(
                     historianRunConfig(boot, reader.current().effective).dreamer,

package/src/config/schema/magic-context.test.ts

@@ -38,7 +38,7 @@ describe("MagicContextConfigSchema", () => {
                     retrieval_count_promotion_threshold: 3,
                 },
                 todowrite: {
-                    enabled: true,
+                    enabled: false,
                     overlay: true,
                 },
             });
@@ -50,6 +50,21 @@ describe("MagicContextConfigSchema", () => {
         });
     });
 
+    it("defaults todowrite off but honors explicit enablement and disablement", () => {
+        expect(MagicContextConfigSchema.parse({}).todowrite).toEqual({
+            enabled: false,
+            overlay: true,
+        });
+        expect(MagicContextConfigSchema.parse({ todowrite: {} }).todowrite.enabled).toBe(false);
+        expect(MagicContextConfigSchema.parse({ todowrite: { enabled: true } }).todowrite).toEqual({
+            enabled: true,
+            overlay: true,
+        });
+        expect(
+            MagicContextConfigSchema.parse({ todowrite: { enabled: false } }).todowrite.enabled,
+        ).toBe(false);
+    });
+
     describe("budget configuration", () => {
         it("accepts 90% execute thresholds and rejects 91%", () => {
             expect(

package/src/config/schema/magic-context.ts

@@ -484,14 +484,27 @@ const CronScheduleSchema = z
 const DreamTaskBaseConfigSchema = z
     .object({
         schedule: CronScheduleSchema.default(""),
+        token_budget: z
+            .number()
+            .int()
+            .positive()
+            .optional()
+            .describe(
+                "Maximum cumulative prompt tokens (input + cache read + cache write) for one tool-loop child. Defaults vary by task.",
+            ),
     })
     .strict();
 
+export const DREAM_TASK_PROMOTION_DEFAULTS = {
+    "review-user-memories": 3,
+    "promote-primers": 2,
+} as const;
+
 const PromotionThresholdSchema = z
     .number()
     .min(2)
     .max(20)
-    .optional()
+    .default(DREAM_TASK_PROMOTION_DEFAULTS["review-user-memories"])
     .describe(
         "review-user-memories: min candidate observations before promotion is considered (default: 3)",
     );
@@ -499,7 +512,7 @@ const PrimerPromotionThresholdSchema = z
     .number()
     .min(2)
     .max(20)
-    .optional()
+    .default(DREAM_TASK_PROMOTION_DEFAULTS["promote-primers"])
     .describe(
         "promote-primers: min recurring source days before promotion is considered (default: 2)",
     );
@@ -548,8 +561,10 @@ const DEFAULT_TASK_SCHEDULES: Record<DreamTaskName, string> = {
 
 function defaultTaskConfig(task: DreamTaskName): z.input<typeof DreamTaskConfigSchema> {
     const base: z.input<typeof DreamTaskConfigSchema> = { schedule: DEFAULT_TASK_SCHEDULES[task] };
-    if (task === "review-user-memories") base.promotion_threshold = 3;
-    if (task === "promote-primers") base.promotion_threshold = 2;
+    if (task === "review-user-memories")
+        base.promotion_threshold = DREAM_TASK_PROMOTION_DEFAULTS["review-user-memories"];
+    if (task === "promote-primers")
+        base.promotion_threshold = DREAM_TASK_PROMOTION_DEFAULTS["promote-primers"];
     return base;
 }
 
@@ -888,7 +903,7 @@ export interface MuralConfig {
 
 export interface MagicContextConfig {
     enabled: boolean;
-    /** User-level setting that lets a session started exactly in the canonical home directory use a deterministic directory identity. */
+    /** User-level setting that lets a session in the canonical home directory use project memory. */
     allow_home_project: boolean;
     mural: MuralConfig;
     /** Selects the runtime implementation for this project. Rust mode is experimental and requires user-level subc configuration. */
@@ -1074,7 +1089,7 @@ export const MagicContextConfigSchema = z
             .boolean()
             .default(false)
             .describe(
-                "Allow Magic Context sessions launched from the exact canonical home directory. The home session uses its deterministic dir: identity so pre-gate memories reconnect. USER-LEVEL ONLY: project config is ignored. The home identity is excluded from registry seed exports, never resolves descendants by containment, and cannot join a workspace.",
+                "Allow Magic Context sessions launched from the exact canonical home directory. A non-git home uses its deterministic dir: identity; a home repository uses its git: identity. USER-LEVEL ONLY: project config is ignored. The home identity is excluded from registry seed exports, never resolves descendants by containment, and cannot join a workspace.",
             ),
         mural: z
             .object({
@@ -1156,7 +1171,7 @@ export const MagicContextConfigSchema = z
             .union([z.string(), z.object({ default: z.string() }).catchall(z.string())])
             .default("5m")
             .describe(
-                'How long Magic Context assumes the provider\'s cached prefix stays valid. This is MC\'s own deferral gate — it does not change the provider\'s actual cache lifetime. String (e.g. "5m", "1h", "30s") or per-model object ({ default: "5m", "provider/model": "1h", "provider/*": "never" }); keys resolve most-specific first (exact provider/model, bare model ID, shorter dash-prefixes, then the provider/* wildcard, then default). Set to "never" to mean MC never assumes expiry (for lanes kept warm externally by a cache-keep tool) — disables the idle-TTL heuristic so MC never initiates a rebuild based on elapsed time. Provider-side extended TTL is a separate request-level concern (cache_control: { ttl } in the request body).',
+                'How long Magic Context assumes the provider\'s cached prefix stays valid. This is MC\'s own deferral gate — it does not change the provider\'s actual cache lifetime. String (e.g. "5m", "1h", "30s") or per-model object ({ default: "5m", "provider/model": "1h", "provider/*": "never" }); keys resolve most-specific first (exact provider/model, bare model ID, shorter dash-prefixes, then the provider/* wildcard). Explicit per-model entries win; otherwise GPT-5.6 and later (including gpt-6*, through any provider prefix) use a built-in 30m lifetime before the object default or 5m fallback. An unset or global "5m" opts into built-in defaults; any other global string is an explicit policy and wins. Policy is frozen per session, including across restarts; a model switch resolves against that frozen policy. /ctx-status shows the effective value and source. OpenAI documents at least 30 minutes since the latest write or reuse: https://developers.openai.com/api/docs/guides/prompt-caching (Cache lifetime and Summary of model differences). Set to "never" to mean MC never assumes expiry (for lanes kept warm externally by a cache-keep tool) — disables the idle-TTL heuristic so MC never initiates a rebuild based on elapsed time. Provider-side extended TTL is a separate request-level concern (cache_control: { ttl } in the request body).',
             ),
         prompt_surface: PromptSurfaceConfigSchema.default({ default: "full" }).describe(
             "Prompt-surface presets: default is full; models use bare model IDs, provider/model, or provider/* routing keys. Guidance and tool-description overrides are user-level only. OpenCode 1.x, Pi, and OMP register tool descriptions once per process (they follow the default preset). OpenCode 2 rewrites the five ctx_* descriptions per request from the draft model.",
@@ -1388,9 +1403,9 @@ export const MagicContextConfigSchema = z
             .object({
                 enabled: z
                     .boolean()
-                    .default(true)
+                    .default(false)
                     .describe(
-                        "Pi only: register Magic Context's todowrite task-list tool. Disable if you use your own todo extension. OpenCode ships its own built-in todowrite; this setting has no effect there.",
+                        "Pi only: off by default. Set todowrite.enabled=true to register Magic Context's todowrite task-list tool and /todos command. OpenCode ships its own built-in todowrite; this setting has no effect there.",
                     ),
                 overlay: z
                     .boolean()
@@ -1399,7 +1414,7 @@ export const MagicContextConfigSchema = z
                         "Pi only: show the persistent todo overlay above the editor while tasks are active.",
                     ),
             })
-            .default({ enabled: true, overlay: true })
+            .default({ enabled: false, overlay: true })
             .describe(
                 "Pi-only todowrite tool and overlay controls. Pi registers tools and widgets at extension boot, so changing this after /cd requires /reload or restart.",
             ),
@@ -1552,10 +1567,12 @@ export const LIVE_RELOAD_CONFIG_PATHS = [
     "historian.omp.fallback_models",
     "historian.omp.thinking_level",
     "historian.two_pass",
+    "historian.maxTokens",
     "historian_timeout_ms",
     "commit_cluster_trigger.enabled",
     "commit_cluster_trigger.min_clusters",
     "memory.auto_promote",
+    "dreamer.maxTokens",
     "dreamer.opencode.model",
     "dreamer.opencode.fallback_models",
     "dreamer.opencode.variant",

package/src/features/magic-context/dreamer/task-registry.ts

@@ -29,6 +29,19 @@ export const CANONICAL_DREAM_TASKS = [
 
 export type DreamTaskName = (typeof CANONICAL_DREAM_TASKS)[number];
 
+/** Prompt-token ceilings informed by docs/reports/dreamer-token-usage.md.
+ * Multi-memory batches need more room than a single primer investigation;
+ * curate's 150-step tool loop warrants a larger allowance too. */
+export const DREAM_TOOL_LOOP_TOKEN_BUDGETS = {
+    "map-memories": 1_500_000,
+    verify: 1_700_000,
+    "verify-broad": 1_700_000,
+    curate: 1_500_000,
+    retrospective: 300_000,
+    "maintain-docs": 1_600_000,
+    "refresh-primers": 350_000,
+} as const;
+
 /**
  * How a Dreamer task reaches its result, which is what decides whether a host
  * without a tool loop can run it:

package/src/shared/cache-ttl-display.ts

@@ -1,7 +1,7 @@
 import type { MagicContextConfig } from "../config/schema/magic-context";
-import { resolveModelConfigValue } from "./prompt-surface";
+import { type ResolvedCacheTtl, resolveModelCacheTtl } from "./model-cache-ttl";
 
-export type CacheTtlDisplaySource = "config" | "session" | "default";
+export type CacheTtlDisplaySource = ResolvedCacheTtl["source"] | "session";
 
 export interface CacheTtlDisplay {
     value: string;
@@ -10,6 +10,7 @@ export interface CacheTtlDisplay {
 }
 
 export interface ResolveCacheTtlDisplayArgs {
+    frozen?: ResolvedCacheTtl;
     configured: MagicContextConfig["cache_ttl"];
     configuredExplicitly: boolean;
     modelKey: string | undefined;
@@ -24,6 +25,7 @@ export interface ResolveCacheTtlDisplayArgs {
  * displayed; otherwise use live config because the session value may be for another model.
  */
 export function resolveCacheTtlDisplay(args: ResolveCacheTtlDisplayArgs): CacheTtlDisplay {
+    if (args.frozen) return args.frozen;
     if (
         (args.sessionModelKey && (!args.modelKey || args.sessionModelKey === args.modelKey)) ||
         (!args.modelKey && !args.sessionModelKey && args.sessionValue !== "5m")
@@ -35,27 +37,19 @@ export function resolveCacheTtlDisplay(args: ResolveCacheTtlDisplayArgs): CacheT
         };
     }
 
-    if (typeof args.configured === "string") {
-        return {
-            value: args.configured,
-            source: args.configuredExplicitly ? "config" : "default",
-            modelKey: args.modelKey,
-        };
-    }
-
-    const matched = resolveModelConfigValue(args.configured, args.modelKey);
-    if (matched) {
-        return { value: matched.value, source: "config", modelKey: args.modelKey };
-    }
-
-    return {
-        value: args.configured.default ?? "5m",
-        source: "default",
-        modelKey: args.modelKey,
-    };
+    const resolved = resolveModelCacheTtl(args.configured, args.modelKey);
+    if (
+        resolved.source === "default" &&
+        typeof args.configured === "string" &&
+        args.configuredExplicitly
+    )
+        return { ...resolved, source: "config" };
+    return resolved;
 }
 
 export function formatCacheTtlDisplay(display: CacheTtlDisplay): string {
+    if (display.source === "OpenAI GPT-5.6+ default")
+        return `Cache TTL: ${display.value} (${display.source})`;
     if (display.source === "session") return `Cache TTL: ${display.value} (session)`;
     if (display.source === "config") {
         return `Cache TTL: ${display.value} (config for ${display.modelKey ?? "current model"})`;

package/src/shared/harness.ts

@@ -52,6 +52,20 @@ export function getHarness(): HarnessId {
     return currentHarness;
 }
 
+/**
+ * Whether this process may read OpenCode's own session store (opencode.db).
+ *
+ * Only the OpenCode plugins own that store. A Pi or OMP process keeps its
+ * history in Pi JSONL sessions and shares only Magic Context's context.db with
+ * OpenCode. Any history read in a Pi process for a session with no Pi provider
+ * used to fall through to opencode.db, which let background work load another
+ * project's OpenCode sessions into Pi's heap. Pi therefore treats the OpenCode
+ * store as absent, the same as on a Pi-only install.
+ */
+export function harnessOwnsOpenCodeStore(harness: HarnessId = currentHarness): boolean {
+    return harness === "opencode" || harness === "opencode2";
+}
+
 /**
  * Test-only helper to reset harness state between test cases. Do NOT call
  * from production code paths.

package/src/shared/logger.test.ts

@@ -189,6 +189,7 @@ async function runLoggerScenario<T>(
     const root = mkdtempSync(path.join(os.tmpdir(), "magic-context-logger-test-"));
     scenarioRoots.push(root);
     const child = Bun.spawn({
+        windowsHide: true,
         cmd: ["bun", "--eval", loggerScenario],
         cwd: import.meta.dir,
         env: {

package/src/shared/model-cache-ttl.test.ts

@@ -0,0 +1,40 @@
+import { describe, expect, it } from "bun:test";
+import { MagicContextConfigSchema } from "../config/schema/magic-context";
+import { resolveCacheTtl } from "../hooks/magic-context/event-resolvers";
+
+describe("known model cache TTL", () => {
+    it("matches canonical model versions through provider prefixes", () => {
+        for (const prefix of ["", "openai/", "openai-codex/", "openrouter/openai/", "azure/"]) {
+            for (const model of [
+                "gpt-5.6",
+                "gpt-5.6-sol",
+                "gpt-5.7",
+                "gpt-5.10",
+                "gpt-6",
+                "gpt-6.1",
+                "gpt-6-mini",
+                "gpt-7",
+            ]) {
+                expect(resolveCacheTtl("5m", prefix + model), prefix + model).toBe("30m");
+            }
+            for (const model of ["gpt-5.5", "gpt-5.1", "gpt-5", "gpt-4o", "unknown", "not-gpt-6"]) {
+                expect(resolveCacheTtl("5m", prefix + model), prefix + model).toBe("5m");
+            }
+        }
+    });
+    it("uses built-ins for unset and global 5m, but honors global 10m", () => {
+        expect(resolveCacheTtl(MagicContextConfigSchema.parse({}).cache_ttl, "openai/gpt-6")).toBe(
+            "30m",
+        );
+        expect(resolveCacheTtl("5m", "openai/gpt-6")).toBe("30m");
+        expect(resolveCacheTtl("10m", "openai/gpt-6")).toBe("10m");
+    });
+    it("honors per-model entries before built-ins and built-ins before object defaults", () => {
+        expect(resolveCacheTtl({ default: "10m" }, "openai/gpt-6")).toBe("30m");
+        expect(resolveCacheTtl({ default: "10m", "gpt-6": "5m" }, "openai/gpt-6")).toBe("5m");
+        expect(resolveCacheTtl({ default: "10m", "openai/*": "never" }, "openai/gpt-6")).toBe(
+            "never",
+        );
+        expect(resolveCacheTtl({ default: "10m" }, "other/model")).toBe("10m");
+    });
+});

package/src/shared/model-cache-ttl.ts

@@ -0,0 +1,58 @@
+import { canonicalModelIdentity } from "./harness-provider-map";
+import { resolveModelConfigValue } from "./prompt-surface";
+
+export type CacheTtlConfig = string | Record<string, string>;
+export type CacheTtlSource = "config" | "default" | "OpenAI GPT-5.6+ default";
+export interface ResolvedCacheTtl {
+    value: string;
+    source: CacheTtlSource;
+    modelKey: string | undefined;
+}
+
+// Provider documentation: https://developers.openai.com/api/docs/guides/prompt-caching
+// Cache lifetime / Summary of model differences: at least 30 minutes since write or reuse.
+const MODEL_CACHE_LIFETIMES = [
+    {
+        source: "OpenAI GPT-5.6+ default" as const,
+        value: "30m",
+        matches(model: string): boolean {
+            const version = /^gpt-(\d+)(?:\.(\d+))?(?:$|[^\d.])/.exec(model);
+            if (!version) return false;
+            const major = Number(version[1]);
+            const minor = Number(version[2] ?? 0);
+            return major > 5 || (major === 5 && minor >= 6);
+        },
+    },
+];
+
+export function resolveModelCacheTtl(
+    config: CacheTtlConfig | undefined,
+    modelKey: string | undefined,
+): ResolvedCacheTtl {
+    if (config && typeof config !== "string") {
+        const match =
+            modelKey &&
+            !modelKey.includes("/") &&
+            Object.hasOwn(config, modelKey) &&
+            modelKey !== "default"
+                ? config[modelKey]
+                : resolveModelConfigValue(config, modelKey)?.value;
+        if (match !== undefined) return { value: match, source: "config", modelKey };
+    }
+    // A non-5m global string is an explicit policy. The generic 5m default is
+    // not evidence of provider eviction: documented lifetimes avoid paid rewrites.
+    if (typeof config === "string" && config !== "5m")
+        return { value: config, source: "config", modelKey };
+    const model =
+        canonicalModelIdentity(modelKey ?? "")
+            .toLowerCase()
+            .split("/")
+            .at(-1) ?? "";
+    const known = MODEL_CACHE_LIFETIMES.find((entry) => entry.matches(model));
+    if (known) return { value: known.value, source: known.source, modelKey };
+    return {
+        value: (typeof config === "object" ? config.default : config) ?? "5m",
+        source: "default",
+        modelKey,
+    };
+}

package/src/shared/model-suggestion-retry.ts

@@ -1,5 +1,5 @@
 import type { createOpencodeClient } from "@opencode-ai/sdk";
-
+import { DreamTokenBudgetExceeded } from "../features/magic-context/dreamer/token-budget";
 import { detectOverflow } from "../features/magic-context/overflow-detection";
 import { HiddenAgentStepLimit } from "../v2/hooks/hidden-child";
 import {
@@ -115,6 +115,7 @@ export type PromptFailureClass =
     | "provider_timeout"
     | "provider_error"
     | "step_limit"
+    | "token_budget"
     | "empty_completion"
     | "no_models"
     | "child_aborted"
@@ -351,7 +352,12 @@ async function abortChildRun(client: Client, sessionId: string): Promise<void> {
  * different model.
  */
 function isNonRetryable(error: unknown, externalSignal?: AbortSignal): boolean {
-    if (externalSignal?.aborted || error instanceof HiddenAgentStepLimit) return true;
+    if (
+        externalSignal?.aborted ||
+        error instanceof HiddenAgentStepLimit ||
+        error instanceof DreamTokenBudgetExceeded
+    )
+        return true;
 
     if (error instanceof Error) {
         if (error.name === "AbortError") return true;
@@ -393,6 +399,7 @@ function classifyPromptFailure(
 ): PromptFailureClass {
     const message = extractMessage(error);
     if (error instanceof HiddenAgentStepLimit) return "step_limit";
+    if (error instanceof DreamTokenBudgetExceeded) return "token_budget";
     if (externalSignal?.aborted || message === "prompt aborted by external signal") {
         return "child_aborted";
     }