From 274318aebfde8bb8f387e053f3c6b4899e4b6e7f Mon Sep 17 00:00:00 2001 From: gsxdsm Date: Thu, 16 Jul 2026 01:52:10 -0700 Subject: [PATCH] FN-8056: enforce task token budgets Enforce configured task token budgets whenever session usage is persisted. - Apply soft alerts and hard pauses atomically from all executor persistence paths. - Exclude cache-read tokens from budget usage and dispatch budget notifications once. - Document budget semantics and add regression coverage. Files changed: .changeset/fn-8056-token-budget-enforcement.md | 7 ++ docs/settings-reference.md | 2 + packages/core/src/types.ts | 4 +- .../src/__tests__/session-token-usage.test.ts | 101 ++++++++++++++++++++- .../src/__tests__/token-budget-enforcer.test.ts | 81 ++++++++++------- packages/engine/src/executor.ts | 22 ++++- packages/engine/src/session-token-usage.ts | 8 +- packages/engine/src/token-budget-enforcer.ts | 98 +++++++++++++++++--- 8 files changed, 262 insertions(+), 61 deletions(-) Fusion-Task-Id: FN-8056 Fusion-Task-Lineage: 5f5ed522-f950-42ce-b4fd-e0b1d45b5815 Co-authored-by: Fusion (runfusion.ai) --- .../fn-8056-token-budget-enforcement.md | 7 ++ docs/settings-reference.md | 2 + packages/core/src/types.ts | 4 +- .../src/__tests__/session-token-usage.test.ts | 101 +++++++++++++++++- .../__tests__/token-budget-enforcer.test.ts | 83 ++++++++------ packages/engine/src/executor.ts | 22 +++- packages/engine/src/session-token-usage.ts | 8 +- packages/engine/src/token-budget-enforcer.ts | 98 ++++++++++++++--- 8 files changed, 263 insertions(+), 62 deletions(-) create mode 100644 .changeset/fn-8056-token-budget-enforcement.md diff --git a/.changeset/fn-8056-token-budget-enforcement.md b/.changeset/fn-8056-token-budget-enforcement.md new file mode 100644 index 0000000000..aac9544a20 --- /dev/null +++ b/.changeset/fn-8056-token-budget-enforcement.md @@ -0,0 +1,7 @@ +--- +"@runfusion/fusion": minor +--- + +summary: Per-task token budgets now enforce — soft caps alert once and hard caps pause the task. +category: fix +dev: Wires persist-time enforcement and token-budget notifications; budgets exclude cache-read tokens. diff --git a/docs/settings-reference.md b/docs/settings-reference.md index 345599f138..33c9ce9f46 100644 --- a/docs/settings-reference.md +++ b/docs/settings-reference.md @@ -762,6 +762,8 @@ Backlog health is the alert family for scheduler/backlog imbalance, dependency-b 4. Global per-size (`global.taskTokenBudget.perSize[task.size]`) 5. Global base (`global.taskTokenBudget.soft/hard`) +Budgets measure **input + output + cache-write tokens**. Cache-read tokens are deliberately excluded: reading a cached prompt can be very large without representing newly processed model work. A soft cap records one alert timestamp and dispatches one `token-budget` notification; a hard cap records its timestamp, pauses the task with `pausedReason: "token_budget_exceeded"`, then dispatches one notification. These transitions are atomic and are not repeated on later token persists. + Example: ```json diff --git a/packages/core/src/types.ts b/packages/core/src/types.ts index bb514eb90f..4e072aa5f1 100644 --- a/packages/core/src/types.ts +++ b/packages/core/src/types.ts @@ -1111,9 +1111,9 @@ export interface TaskTokenUsage { } export interface TaskTokenBudget { - /** Total-token soft cap. When reached, emits one notification and continues. */ + /** Input, output, and cache-write token soft cap (cache reads excluded). When reached, emits one notification and continues. */ soft?: number; - /** Total-token hard cap. When reached, pauses the task with pausedReason="token_budget_exceeded". */ + /** Input, output, and cache-write token hard cap (cache reads excluded). When reached, pauses the task with pausedReason="token_budget_exceeded". */ hard?: number; /** Optional per-size overrides keyed by Task.size (S/M/L). Falls back to soft/hard when absent. */ perSize?: { S?: { soft?: number; hard?: number }; M?: { soft?: number; hard?: number }; L?: { soft?: number; hard?: number } }; diff --git a/packages/engine/src/__tests__/session-token-usage.test.ts b/packages/engine/src/__tests__/session-token-usage.test.ts index 523eabe459..cad4e79652 100644 --- a/packages/engine/src/__tests__/session-token-usage.test.ts +++ b/packages/engine/src/__tests__/session-token-usage.test.ts @@ -1,8 +1,12 @@ import { describe, it, expect, vi, beforeEach } from "vitest"; import type { Task, TaskStore } from "@fusion/core"; import { accumulateSessionTokenUsage, computeCacheHitRatio } from "../session-token-usage.js"; +import { enforceTaskTokenBudgetForPersist } from "../token-budget-enforcer.js"; import { TaskExecutor } from "../executor.js"; +const { notificationService } = vi.hoisted(() => ({ notificationService: { dispatch: vi.fn() } })); +vi.mock("../notifier.js", () => ({ getActiveNotificationService: () => notificationService })); + interface MockSessionStats { tokens?: { input?: number; output?: number; cacheRead?: number; cacheWrite?: number; total?: number }; } @@ -14,17 +18,32 @@ function createSession( return { getSessionStats: vi.fn(() => stats), ...(model ? { model } : {}) } as unknown as Parameters[2]; } -function createStore(initial: Task["tokenUsage"]): TaskStore & { _task: Task; updateTask: ReturnType } { - const task = { id: "FN-1", tokenUsage: initial } as Task; +function createStore(initial: Task["tokenUsage"], budget?: { soft?: number; hard?: number }): TaskStore & { _task: Task; updateTask: ReturnType; pauseTask: ReturnType } { + const task = { id: "FN-1", title: "Budget task", tokenUsage: initial } as Task; const updateTask = vi.fn(async (_id: string, updates: Partial) => { - if (updates.tokenUsage !== undefined) task.tokenUsage = updates.tokenUsage as Task["tokenUsage"]; + Object.assign(task, updates); return task; }); + const pauseTask = vi.fn(async () => task); + // Model TaskStore's atomic transaction boundary: each updater observes the latest row. + let atomicQueue: Promise = Promise.resolve(); + const updateTaskAtomic = vi.fn((_id: string, updater: (current: Task) => Partial | null) => { + const result = atomicQueue.then(async () => { + const patch = await updater(task); + if (patch) Object.assign(task, patch); + return task; + }); + atomicQueue = result.then(() => undefined, () => undefined); + return result; + }); const store = { _task: task, getTask: vi.fn(async () => task), + getSettingsByScope: vi.fn(async () => ({ project: budget ? { taskTokenBudget: budget } : {}, global: {} })), updateTask, - } as unknown as TaskStore & { _task: Task; updateTask: ReturnType }; + updateTaskAtomic, + pauseTask, + } as unknown as TaskStore & { _task: Task; updateTask: ReturnType; pauseTask: ReturnType }; return store; } @@ -32,6 +51,7 @@ describe("accumulateSessionTokenUsage", () => { beforeEach(() => { vi.useRealTimers(); vi.restoreAllMocks(); + notificationService.dispatch.mockReset(); }); it("writes initial token usage and emits cache metrics log", async () => { @@ -155,6 +175,7 @@ describe("accumulateSessionTokenUsage", () => { executor.store = store; executor.tokenUsageBaselines = new Map(); executor.activeSessions = new Map(); + executor.currentRunContexts = new Map(); await executor.persistTokenUsage("FN-1", { getSessionStats: () => ({ tokens: { input: 3, output: 2, cacheRead: 1, cacheWrite: 0, total: 6 } }), @@ -167,6 +188,78 @@ describe("accumulateSessionTokenUsage", () => { expect(call.tokenUsage).toMatchObject({ modelProvider: "mock", modelId: "scripted" }); }); + it("enforces soft and hard budgets through the real persist helper exactly once", async () => { + const store = createStore(undefined, { soft: 10, hard: 20 }); + const session = createSession({ tokens: { input: 12, output: 0, cacheRead: 0, cacheWrite: 0 } }); + + await accumulateSessionTokenUsage(store, "FN-1", session); + session.getSessionStats.mockReturnValue({ tokens: { input: 25, output: 0, cacheRead: 0, cacheWrite: 0 } }); + await accumulateSessionTokenUsage(store, "FN-1", session); + session.getSessionStats.mockReturnValue({ tokens: { input: 30, output: 0, cacheRead: 0, cacheWrite: 0 } }); + await accumulateSessionTokenUsage(store, "FN-1", session); + + expect(store._task.tokenBudgetSoftAlertedAt).toBeTruthy(); + expect(store._task.tokenBudgetHardAlertedAt).toBeTruthy(); + expect(store.pauseTask).toHaveBeenCalledOnce(); + expect(store.pauseTask).toHaveBeenCalledWith("FN-1", true, undefined, { pausedReason: "token_budget_exceeded" }); + expect(notificationService.dispatch).toHaveBeenCalledTimes(2); + expect(notificationService.dispatch).toHaveBeenNthCalledWith(1, "token-budget", expect.objectContaining({ metadata: expect.objectContaining({ kind: "soft" }) })); + expect(notificationService.dispatch).toHaveBeenNthCalledWith(2, "token-budget", expect.objectContaining({ metadata: expect.objectContaining({ kind: "hard" }) })); + }); + + it("atomically claims concurrent soft and hard enforcement once", async () => { + const store = createStore({ inputTokens: 25, outputTokens: 0, cacheWriteTokens: 0, totalTokens: 25 }, { soft: 10, hard: 20 }); + + await Promise.all(Array.from({ length: 8 }, () => enforceTaskTokenBudgetForPersist(store, "FN-1"))); + + expect(store.pauseTask).toHaveBeenCalledOnce(); + expect(notificationService.dispatch).toHaveBeenCalledTimes(2); + expect(store._task.tokenBudgetSoftAlertedAt).toBeTruthy(); + expect(store._task.tokenBudgetHardAlertedAt).toBeTruthy(); + }); + + it("enforces direct executor token persistence through its shared seam", async () => { + const store = createStore(undefined, { hard: 10 }); + const executor = Object.create(TaskExecutor.prototype) as any; + executor.store = store; + executor.currentRunContexts = new Map(); + + await executor.persistTaskTokenUsage("FN-1", { inputTokens: 20, outputTokens: 0, cachedTokens: 0, cacheWriteTokens: 0, totalTokens: 20 }); + + expect(store.pauseTask).toHaveBeenCalledOnce(); + expect(notificationService.dispatch).toHaveBeenCalledWith("token-budget", expect.objectContaining({ metadata: expect.objectContaining({ kind: "hard" }) })); + }); + + it("retries a hard pause after a failed persisted enforcement attempt", async () => { + const store = createStore({ inputTokens: 20, outputTokens: 0, cacheWriteTokens: 0, totalTokens: 20 }, { hard: 10 }); + store.pauseTask.mockRejectedValueOnce(new Error("temporary pause failure")); + + await expect(enforceTaskTokenBudgetForPersist(store, "FN-1")).resolves.toBeUndefined(); + expect(store._task.tokenBudgetHardAlertedAt).toBeNull(); + await enforceTaskTokenBudgetForPersist(store, "FN-1"); + + expect(store.pauseTask).toHaveBeenCalledTimes(2); + expect(store._task.tokenBudgetHardAlertedAt).toBeTruthy(); + expect(notificationService.dispatch).toHaveBeenCalledOnce(); + }); + + it("keeps a successful hard pause durable when notification dispatch fails", async () => { + notificationService.dispatch.mockRejectedValueOnce(new Error("notification unavailable")); + const store = createStore(undefined, { hard: 10 }); + + await expect(accumulateSessionTokenUsage(store, "FN-1", createSession({ tokens: { input: 20, output: 0, cacheRead: 0, cacheWrite: 0 } }))).resolves.toBeUndefined(); + + expect(store.pauseTask).toHaveBeenCalledOnce(); + expect(store._task.tokenBudgetHardAlertedAt).toBeTruthy(); + }); + + it("does not enforce when no budget is configured", async () => { + const store = createStore(undefined); + await accumulateSessionTokenUsage(store, "FN-1", createSession({ tokens: { input: 100, output: 0, cacheRead: 0, cacheWrite: 0 } })); + expect(store.pauseTask).not.toHaveBeenCalled(); + expect(notificationService.dispatch).not.toHaveBeenCalled(); + }); + it("swallows store errors instead of throwing", async () => { const store = createStore(undefined); (store.updateTask as ReturnType).mockRejectedValueOnce(new Error("db down")); diff --git a/packages/engine/src/__tests__/token-budget-enforcer.test.ts b/packages/engine/src/__tests__/token-budget-enforcer.test.ts index 82092fe395..90a4fe151d 100644 --- a/packages/engine/src/__tests__/token-budget-enforcer.test.ts +++ b/packages/engine/src/__tests__/token-budget-enforcer.test.ts @@ -1,44 +1,61 @@ import { describe, expect, it, vi } from "vitest"; -import { enforceTaskTokenBudget, resolveTaskTokenBudget } from "../token-budget-enforcer.js"; +import { enforceTaskTokenBudget, getTokenBudgetUsage, resolveTaskTokenBudget } from "../token-budget-enforcer.js"; + +const task = (patch: Record = {}) => ({ + id: "FN-1", description: "x", column: "todo", dependencies: [], steps: [], currentStep: 0, createdAt: "", updatedAt: "", ...patch, +}) as any; describe("resolveTaskTokenBudget", () => { - it("prefers task override", () => { - const result = resolveTaskTokenBudget( - { id: "FN-1", description: "x", column: "todo", dependencies: [], steps: [], currentStep: 0, createdAt: "", updatedAt: "", tokenBudgetOverride: { soft: 10, hard: 20 } } as any, - { taskTokenBudget: { soft: 100, hard: 200 } } as any, - { taskTokenBudget: { soft: 1000, hard: 2000 } } as any, - ); - expect(result).toEqual({ soft: 10, hard: 20, source: "task-override" }); + it.each([ + ["task override", task({ size: "M", tokenBudgetOverride: { soft: 1, hard: 2 } }), { taskTokenBudget: { soft: 10, hard: 20, perSize: { M: { soft: 11 } } } }, { taskTokenBudget: { soft: 100 } }, { soft: 1, hard: 2, source: "task-override" }], + ["project per-size", task({ size: "M" }), { taskTokenBudget: { soft: 10, hard: 20, perSize: { M: { soft: 11 } } } }, { taskTokenBudget: { soft: 100 } }, { soft: 11, hard: 20, source: "project-per-size" }], + ["project base", task(), { taskTokenBudget: { soft: 10, hard: 20 } }, { taskTokenBudget: { soft: 100 } }, { soft: 10, hard: 20, source: "project" }], + ["global per-size", task({ size: "M" }), {}, { taskTokenBudget: { soft: 100, hard: 200, perSize: { M: { hard: 201 } } } }, { soft: 100, hard: 201, source: "global-per-size" }], + ["global base", task(), {}, { taskTokenBudget: { soft: 100, hard: 200 } }, { soft: 100, hard: 200, source: "global" }], + ["none", task(), {}, {}, { source: "none" }], + ])("uses %s precedence", (_name, subject, project, global, expected) => { + expect(resolveTaskTokenBudget(subject, project as any, global as any)).toEqual(expected); }); }); describe("enforceTaskTokenBudget", () => { - it("fires soft once and hard pause once", async () => { - const updateTask = vi.fn(async () => undefined); - const pauseTask = vi.fn(async () => undefined); - const notify = vi.fn(async () => undefined); - const task = { - id: "FN-1", - description: "x", - column: "in-progress", - dependencies: [], - steps: [], - currentStep: 0, - createdAt: "", - updatedAt: "", - tokenUsage: { totalTokens: 150 }, - } as any; - - await enforceTaskTokenBudget({ - store: { updateTask, pauseTask }, - task, - projectSettings: { taskTokenBudget: { soft: 100, hard: 140 } } as any, - globalSettings: {} as any, - notify, + it("claims soft and hard caps once and pauses through pauseTask", async () => { + const current = task({ tokenUsage: { inputTokens: 150, outputTokens: 0, cacheWriteTokens: 0 } }); + const updateTaskAtomic = vi.fn(async (_id, updater) => { + const patch = await updater(current); + if (patch) Object.assign(current, patch); + return current; }); + const pauseTask = vi.fn(async () => current); + const notify = vi.fn(async () => undefined); - expect(updateTask).toHaveBeenCalled(); - expect(pauseTask).toHaveBeenCalledWith("FN-1", true, undefined); - expect(notify).toHaveBeenCalledWith(expect.objectContaining({ kind: "hard" })); + await enforceTaskTokenBudget({ store: { updateTaskAtomic, pauseTask } as any, task: current, projectSettings: { taskTokenBudget: { soft: 100, hard: 140 } } as any, globalSettings: {} as any, notify }); + await enforceTaskTokenBudget({ store: { updateTaskAtomic, pauseTask } as any, task: current, projectSettings: { taskTokenBudget: { soft: 100, hard: 140 } } as any, globalSettings: {} as any, notify }); + + expect(pauseTask).toHaveBeenCalledOnce(); + expect(pauseTask).toHaveBeenCalledWith("FN-1", true, undefined, { pausedReason: "token_budget_exceeded" }); + expect(notify).toHaveBeenCalledTimes(2); + }); + + it("releases a hard claim when pause fails so a later persist retries", async () => { + const current = task({ tokenUsage: { inputTokens: 25 } }); + const updateTaskAtomic = vi.fn(async (_id, updater) => { + const patch = await updater(current); + if (patch) Object.assign(current, patch); + return current; + }); + const pauseTask = vi.fn().mockRejectedValueOnce(new Error("transient")).mockResolvedValueOnce(current); + const notify = vi.fn(async () => undefined); + const params = { store: { updateTaskAtomic, pauseTask } as any, task: current, projectSettings: { taskTokenBudget: { hard: 20 } } as any, globalSettings: {} as any, notify }; + + await expect(enforceTaskTokenBudget(params)).rejects.toThrow("transient"); + expect(current.tokenBudgetHardAlertedAt).toBeNull(); + await enforceTaskTokenBudget(params); + expect(pauseTask).toHaveBeenCalledTimes(2); + expect(notify).toHaveBeenCalledOnce(); + }); + + it("excludes cache reads from the budget basis", () => { + expect(getTokenBudgetUsage({ inputTokens: 10, outputTokens: 20, cacheWriteTokens: 30, cachedTokens: 9_999, totalTokens: 10_059 } as any)).toBe(60); }); }); diff --git a/packages/engine/src/executor.ts b/packages/engine/src/executor.ts index 126ad285ed..371557ffbf 100644 --- a/packages/engine/src/executor.ts +++ b/packages/engine/src/executor.ts @@ -90,6 +90,7 @@ import { Type, type Static } from "@earendil-works/pi-ai"; import { describeModel, formatModelMarkerDetails, promptWithFallback, compactSessionContext } from "./pi.js"; import { buildAgentGatedActionSummary } from "./permanent-agent-gating.js"; import { accumulateSessionTokenUsage, mergeTokenUsagePerModel } from "./session-token-usage.js"; +import { enforceTaskTokenBudgetForPersist } from "./token-budget-enforcer.js"; import { createResolvedAgentSession, extractRuntimeHint, @@ -4207,6 +4208,17 @@ export class TaskExecutor { } } + /** + * FNXC:TokenBudget 2026-07-16-00:00: + * Step-session token usage bypasses the shared session helper, so all executor + * writes use this seam to retain the required persist-time budget enforcement. + */ + private async persistTaskTokenUsage(taskId: string, tokenUsage: TaskTokenUsage): Promise { + const runContext = this.getRunContextFor(taskId); + await this.store.updateTask(taskId, { tokenUsage }, runContext); + await enforceTaskTokenBudgetForPersist(this.store, taskId, runContext); + } + private async persistTokenUsage(taskId: string, session?: AgentSession): Promise { const activeSession = session ?? this.activeSessions.get(taskId)?.session; const currentUsage = await this.extractSessionTokenUsage(activeSession); @@ -4250,7 +4262,7 @@ export class TaskExecutor { hitRatio: tokenUsage.inputTokens + tokenUsage.cachedTokens > 0 ? tokenUsage.cachedTokens / (tokenUsage.inputTokens + tokenUsage.cachedTokens) : 0, })); - await this.store.updateTask(taskId, { tokenUsage }); + await this.persistTaskTokenUsage(taskId, tokenUsage); } /** @@ -10482,7 +10494,7 @@ export class TaskExecutor { return; } - this.store.updateTask(task.id, { tokenUsage: accumulatedStepTokenUsage }).catch((err) => { + this.persistTaskTokenUsage(task.id, accumulatedStepTokenUsage).catch((err) => { executorLog.warn(`${task.id}: failed to persist token usage on step ${stepIndex} complete: ${err}`); }); }, @@ -10531,7 +10543,7 @@ export class TaskExecutor { } if (accumulatedStepTokenUsage) { - await this.store.updateTask(task.id, { tokenUsage: accumulatedStepTokenUsage }); + await this.persistTaskTokenUsage(task.id, accumulatedStepTokenUsage); } const allSuccess = results.every(r => r.success); @@ -10850,13 +10862,13 @@ export class TaskExecutor { nextRecoveryAt: null, }); if (accumulatedStepTokenUsage) { - await this.store.updateTask(task.id, { tokenUsage: accumulatedStepTokenUsage }); + await this.persistTaskTokenUsage(task.id, accumulatedStepTokenUsage); } executorLog.log(`✗ ${task.id} transient retries exhausted — failed in execution`); this.options.onError?.(task, err instanceof Error ? err : new Error(errorMessage)); } else { if (accumulatedStepTokenUsage) { - await this.store.updateTask(task.id, { tokenUsage: accumulatedStepTokenUsage }); + await this.persistTaskTokenUsage(task.id, accumulatedStepTokenUsage); } if (await this.handleNonContinuableSessionError(task, false, errorMessage)) { return; diff --git a/packages/engine/src/session-token-usage.ts b/packages/engine/src/session-token-usage.ts index fa8c9410a0..30fcdad192 100644 --- a/packages/engine/src/session-token-usage.ts +++ b/packages/engine/src/session-token-usage.ts @@ -1,6 +1,7 @@ -import type { AgentRole, TaskStore, TaskTokenUsage, TaskTokenUsagePerModel } from "@fusion/core"; +import type { AgentRole, RunMutationContext, TaskStore, TaskTokenUsage, TaskTokenUsagePerModel } from "@fusion/core"; import type { AgentSession } from "@earendil-works/pi-coding-agent"; import { createLogger } from "./logger.js"; +import { enforceTaskTokenBudgetForPersist } from "./token-budget-enforcer.js"; const log = createLogger("session-token-usage"); const cacheMetricsLog = createLogger("token-cache-metrics"); @@ -83,7 +84,7 @@ export async function accumulateSessionTokenUsage( store: TaskStore, taskId: string, session: AgentSession, - options?: { agentId?: string; role?: AgentRole }, + options?: { agentId?: string; role?: AgentRole; runContext?: RunMutationContext }, ): Promise { try { const stats = readSessionStats(session); @@ -156,7 +157,8 @@ export async function accumulateSessionTokenUsage( hitRatio: computeCacheHitRatio(tokenUsage.inputTokens, tokenUsage.cachedTokens), })); - await store.updateTask(taskId, { tokenUsage }); + await store.updateTask(taskId, { tokenUsage }, options?.runContext); + await enforceTaskTokenBudgetForPersist(store, taskId, options?.runContext); } catch (err) { const message = err instanceof Error ? err.message : String(err); log.warn(`${taskId}: session token usage accumulate failed: ${message}`); diff --git a/packages/engine/src/token-budget-enforcer.ts b/packages/engine/src/token-budget-enforcer.ts index 887024e740..43f5f5452c 100644 --- a/packages/engine/src/token-budget-enforcer.ts +++ b/packages/engine/src/token-budget-enforcer.ts @@ -1,10 +1,13 @@ -import type { GlobalSettings, ProjectSettings, RunMutationContext, Task } from "@fusion/core"; +import type { GlobalSettings, ProjectSettings, RunMutationContext, Task, TaskStore } from "@fusion/core"; import { createLogger } from "./logger.js"; +import { getActiveNotificationService } from "./notifier.js"; const log = createLogger("token-budget-enforcer"); type BudgetSource = "task-override" | "project-per-size" | "project" | "global-per-size" | "global" | "none"; +type TokenBudgetStore = Pick; + export interface ResolvedTaskTokenBudget { soft?: number; hard?: number; @@ -62,27 +65,92 @@ export function resolveTaskTokenBudget( return { source: "none" }; } +/** + * FNXC:TokenBudget 2026-07-16-00:00: + * Budget limits measure newly processed input/output plus cache writes, not cache reads. + * Cache reads can dominate a session without representing new model work, so using them + * would unexpectedly pause existing tasks calibrated for actual work tokens. + */ +export function getTokenBudgetUsage(tokenUsage: Task["tokenUsage"]): number { + return (tokenUsage?.inputTokens ?? 0) + (tokenUsage?.outputTokens ?? 0) + (tokenUsage?.cacheWriteTokens ?? 0); +} + export async function enforceTaskTokenBudget( - params: { store: { updateTask: (id: string, updates: Record, runContext?: RunMutationContext) => Promise; pauseTask: (id: string, paused: boolean, runContext?: RunMutationContext) => Promise }; task: Task } & EnforcementContext, + params: { store: Pick; task: Task } & EnforcementContext, ): Promise { const { store, task, projectSettings, globalSettings, runContext, notify } = params; - const total = task.tokenUsage?.totalTokens ?? 0; - const resolved = resolveTaskTokenBudget(task, projectSettings, globalSettings); - const { soft, hard } = resolved; + const total = getTokenBudgetUsage(task.tokenUsage); + const { soft, hard } = resolveTaskTokenBudget(task, projectSettings, globalSettings); - if (soft !== undefined && total >= soft && !task.tokenBudgetSoftAlertedAt) { + if (soft !== undefined && total >= soft) { const now = new Date().toISOString(); - await store.updateTask(task.id, { tokenBudgetSoftAlertedAt: now }, runContext); - log.warn(`${task.id}: soft token budget reached (${total}/${soft})`); - await notify({ kind: "soft", task, total, soft, hard }); + let claimedSoft = false; + await store.updateTaskAtomic(task.id, (current) => { + if (current.tokenBudgetSoftAlertedAt) return null; + claimedSoft = true; + return { tokenBudgetSoftAlertedAt: now }; + }, runContext); + if (claimedSoft) { + log.warn(`${task.id}: soft token budget reached (${total}/${soft})`); + await notify({ kind: "soft", task, total, soft, hard }); + } } - if (hard !== undefined && total >= hard && !task.tokenBudgetHardAlertedAt) { + if (hard !== undefined && total >= hard) { const now = new Date().toISOString(); - await store.updateTask(task.id, { tokenBudgetHardAlertedAt: now }, runContext); - await store.pauseTask(task.id, true, runContext); - await store.updateTask(task.id, { pausedReason: "token_budget_exceeded" }, runContext); - log.error(`${task.id}: hard token budget reached (${total}/${hard}), task paused`); - await notify({ kind: "hard", task, total, soft, hard }); + let claimedHard = false; + await store.updateTaskAtomic(task.id, (current) => { + if (current.tokenBudgetHardAlertedAt) return null; + claimedHard = true; + return { tokenBudgetHardAlertedAt: now }; + }, runContext); + if (claimedHard) { + try { + await store.pauseTask(task.id, true, runContext, { pausedReason: "token_budget_exceeded" }); + } catch (err) { + await store.updateTaskAtomic(task.id, (current) => + current.tokenBudgetHardAlertedAt === now ? { tokenBudgetHardAlertedAt: null } : null, + runContext).catch(() => {}); + throw err; + } + log.error(`${task.id}: hard token budget reached (${total}/${hard}), task paused`); + await notify({ kind: "hard", task, total, soft, hard }); + } + } +} + +/** + * FNXC:TokenBudget 2026-07-16-00:00: + * FN-8056 found the enforcer was dead code; every persisted task.tokenUsage must enter + * this best-effort helper so documented soft alerts and hard pauses remain live. + */ +export async function enforceTaskTokenBudgetForPersist( + store: TokenBudgetStore, + taskId: string, + runContext?: RunMutationContext, +): Promise { + try { + const [task, settings] = await Promise.all([store.getTask(taskId), store.getSettingsByScope()]); + await enforceTaskTokenBudget({ + store, + task, + projectSettings: settings.project as ProjectSettings, + globalSettings: settings.global, + runContext, + notify: async ({ kind, task: notifiedTask, total, soft, hard }) => { + const notificationService = getActiveNotificationService(); + if (!notificationService) return; + await notificationService.dispatch("token-budget", { + taskId: notifiedTask.id, + taskTitle: notifiedTask.title, + event: "token-budget", + timestamp: new Date().toISOString(), + metadata: { kind, total, soft, hard }, + }); + }, + }); + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + log.warn(`${taskId}: token budget enforcement failed: ${message}`); } }