{"version":3,"file":"budget.mjs","names":[],"sources":["../../../../../../../../ai/src/middleware/builtins/budget.ts"],"sourcesContent":["import type {\n  AgentMiddleware,\n  MiddlewareExecuteContext,\n} from \"../../contracts/middleware\";\nimport { BudgetExceededError, type BudgetUnit } from \"../../errors\";\nimport { namespacedState } from \"../utils\";\nimport type {\n  BudgetContract,\n  BudgetContractDimension,\n  BudgetContractViolation,\n} from \"./budget-contract.type\";\n\nexport type {\n  BudgetContract,\n  BudgetContractDimension,\n  BudgetContractFallback,\n  BudgetContractViolation,\n  BudgetContractViolationMode,\n} from \"./budget-contract.type\";\n\n/**\n * Per-model pricing used to compute USD cost from token counts.\n * Caller-supplied — no bundled table. Keys are model names (the\n * `ModelContract.name` value); values are input / output token\n * prices expressed as **USD per 1K tokens** to match every major\n * provider's published pricing sheet.\n */\nexport type BudgetPricing = Record<\n  string,\n  {\n    /** USD per 1,000 input tokens. */\n    inputPer1K: number;\n    /** USD per 1,000 output tokens. */\n    outputPer1K: number;\n  }\n>;\n\n/**\n * Configuration for `budget()`. At least one of `maxTokens` or\n * `maxCostUSD` must be supplied — a budget with no cap is a no-op.\n */\nexport type BudgetOptions = {\n  /**\n   * Hard cap on cumulative total tokens (input + output, summed\n   * across every trip of the run). Inclusive — exceeding triggers\n   * the configured `onExceeded`.\n   */\n  maxTokens?: number;\n  /**\n   * Hard cap on cumulative USD cost. Requires `pricing` for the\n   * agent's configured model — without a pricing entry the USD check\n   * silently skips (tokens-only enforcement still applies).\n   */\n  maxCostUSD?: number;\n  /**\n   * Per-model pricing table used to compute USD cost. Only consulted\n   * when `maxCostUSD` is set. Model names must match the running\n   * agent's `ModelContract.name` exactly.\n   */\n  pricing?: BudgetPricing;\n  /**\n   * Behavior when a cap is breached. `\"abort\"` throws\n   * `BudgetExceededError` — surfaces on `result.error`, stops the\n   * run at the next trip boundary. `\"warn\"` logs a warning and\n   * lets the run continue (useful for observability-first rollouts\n   * before flipping the switch to abort). Default `\"abort\"`.\n   */\n  onExceeded?: \"abort\" | \"warn\";\n  /**\n   * Override the middleware name. Useful when two budgets coexist\n   * (e.g. a per-request cap plus a session-wide cap via different\n   * instances). Default `\"budget\"`.\n   */\n  name?: string;\n  /**\n   * Declarative SLO / cost contract enforced on top of (and\n   * independently of) the legacy `maxTokens` / `maxCostUSD` caps.\n   * Adds a wall-clock `maxLatencyMs` dimension and a per-contract\n   * `onViolation` reaction (`\"abort\"` hard-stops, `\"fallback\"` records\n   * a signal + fires `fallback` and lets the run continue). Omit to\n   * keep the classic budget behavior unchanged.\n   *\n   * Read a recorded fallback signal back with\n   * {@link readBudgetFallbackSignal}.\n   */\n  contract?: BudgetContract;\n};\n\ntype BudgetCounters = {\n  tokens: number;\n  costUSD: number;\n  warned: boolean;\n  /**\n   * Wall-clock epoch ms captured at `execute.before`. Used to derive\n   * cumulative run latency for the contract's `maxLatencyMs` clause.\n   */\n  startedAt: number;\n  /**\n   * Set once a `\"fallback\"` contract clause has fired, so the signal +\n   * callback are emitted at most once per run even if later trips keep\n   * breaching.\n   */\n  fallbackFired: boolean;\n};\n\n/**\n * Recorded contract fallback signal, stashed under the `<name>.fallback`\n * state key when a `\"fallback\"` clause trips. A fallback orchestrator\n * reads it via {@link readBudgetFallbackSignal} to decide how to degrade.\n */\nexport type BudgetFallbackSignal = BudgetContractViolation;\n\n/**\n * The `BudgetUnit` to stamp on the thrown error per contract dimension.\n * Latency has no native unit — it borrows `\"requests\"` as the least-wrong\n * operational measure, while the authoritative detail rides on the\n * error's `context.dimension`.\n */\nconst DIMENSION_UNIT: Record<BudgetContractDimension, BudgetUnit> = {\n  tokens: \"tokens\",\n  cost: \"usd\",\n  latency: \"requests\",\n};\n\nfunction breach(\n  limit: number,\n  actual: number,\n  unit: BudgetUnit,\n  name: string,\n): never {\n  throw new BudgetExceededError(\n    `budget \"${name}\" exceeded — ${actual} ${unit} (cap: ${limit})`,\n    { limit, actual, unit },\n  );\n}\n\nfunction breachContract(\n  name: string,\n  dimension: BudgetContractDimension,\n  limit: number,\n  actual: number,\n): never {\n  throw new BudgetExceededError(\n    `budget \"${name}\" contract exceeded — ${dimension} ${actual} (cap: ${limit})`,\n    {\n      limit,\n      actual,\n      unit: DIMENSION_UNIT[dimension],\n      context: { dimension, limit, actual, source: \"contract\" },\n    },\n  );\n}\n\n/**\n * Read the contract fallback signal recorded by a `budget()` middleware\n * running under `contract.onViolation: \"fallback\"`. Returns `undefined`\n * when no clause was breached.\n *\n * **Role.** The middleware cannot itself switch models on a soft breach,\n * so it records a typed {@link BudgetFallbackSignal} in the shared state\n * bag and lets the run continue. A fallback orchestrator (or the\n * `execute.after` hook of an outer middleware) reads it back here and\n * decides how to degrade the next run — cheaper model, cached answer,\n * truncated context.\n *\n * @param state - The middleware state bag (`ctx.state`).\n * @param name - The budget middleware's name. Default `\"budget\"`,\n *   matching `BudgetOptions.name`'s default.\n *\n * @example\n * const guard = budget({ contract: { maxCostUSD: 0.05, onViolation: \"fallback\" } });\n *\n * // In an outer middleware's execute.after, after the run:\n * const signal = readBudgetFallbackSignal(ctx.state);\n * if (signal?.dimension === \"cost\") {\n *   await rerunOnCheaperModel();\n * }\n */\nexport function readBudgetFallbackSignal(\n  state: MiddlewareExecuteContext[\"state\"],\n  name = \"budget\",\n): BudgetFallbackSignal | undefined {\n  return namespacedState<BudgetFallbackSignal>(\n    { state },\n    `${name}.fallback`,\n  ).get();\n}\n\n/**\n * Enforced token and / or USD budget for an agent run.\n *\n * **Role.** Guards against runaway tool loops, misconfigured\n * prompts, and unexpected provider price swings by capping\n * cumulative usage across every LLM trip of a single execution.\n * Aborts the run with a typed `BudgetExceededError` the moment a cap\n * is breached, rather than letting the damage grow trip by trip.\n *\n * **Scope.** Per-execution. A fresh counter is created at\n * `execute.before` and lives in the middleware state bag until the\n * run ends. Two concurrent `agent.execute()` calls on the same\n * agent therefore enforce the cap independently.\n *\n * **Token accounting.** After each successful trip, the middleware\n * adds `response.usage.total` to its running total and checks\n * against `maxTokens`. Synthetic trips (cache hits) contribute\n * `usage.total` as returned by the cache — cache middleware is\n * expected to surface zero usage on a hit, which naturally excludes\n * those trips from the budget.\n *\n * **USD accounting.** When `maxCostUSD` + `pricing[modelName]` are\n * both present, the middleware converts per-trip input / output\n * tokens to USD and accumulates. Missing pricing silently degrades\n * to tokens-only — explicit rather than guessing.\n *\n * **Warn mode.** `onExceeded: \"warn\"` logs a single warning the first\n * time a cap is breached and lets the run continue. Useful for\n * measuring real-world traffic against a proposed cap before flipping\n * to `\"abort\"` in production.\n *\n * **Contract / SLO mode.** Pass `contract` to enforce a declarative\n * service-level objective — `maxCostUSD`, `maxLatencyMs`, `maxTokens` —\n * on top of the legacy caps, with a single `onViolation` reaction:\n * `\"abort\"` hard-stops with `BudgetExceededError`; `\"fallback\"` records\n * a typed signal (read it via {@link readBudgetFallbackSignal}), fires\n * the optional `fallback` callback, and lets the run continue so an\n * outer layer can degrade gracefully. The contract's clauses are\n * evaluated independently of — and after — the top-level caps; the\n * top-level caps stay fully functional with or without a contract.\n *\n * @example\n * const budgetMiddleware = budget({ maxTokens: 50_000 });\n *\n * const myAgent = agent({\n *   model,\n *   middleware: [budgetMiddleware],\n * });\n *\n * @example\n * // With USD cap and custom pricing\n * const guard = budget({\n *   maxCostUSD: 0.5,\n *   pricing: {\n *     \"gpt-4o\": { inputPer1K: 0.005, outputPer1K: 0.015 },\n *   },\n * });\n *\n * @example\n * // SLO contract — soft-fallback on any breach\n * const sloGuard = budget({\n *   pricing: { \"gpt-4o\": { inputPer1K: 0.005, outputPer1K: 0.015 } },\n *   contract: {\n *     maxCostUSD: 0.05,\n *     maxLatencyMs: 8_000,\n *     maxTokens: 40_000,\n *     onViolation: \"fallback\",\n *     fallback: (violation) => routeToCheaperModel(violation.dimension),\n *   },\n * });\n */\nexport function budget(options: BudgetOptions): AgentMiddleware {\n  const name = options.name ?? \"budget\";\n  const onExceeded = options.onExceeded ?? \"abort\";\n  const hasTokenCap = typeof options.maxTokens === \"number\";\n  const hasCostCap = typeof options.maxCostUSD === \"number\";\n\n  const contract = options.contract;\n  const contractMode = contract?.onViolation ?? \"abort\";\n  const hasContractTokenCap = typeof contract?.maxTokens === \"number\";\n  const hasContractCostCap = typeof contract?.maxCostUSD === \"number\";\n  const hasContractLatencyCap = typeof contract?.maxLatencyMs === \"number\";\n  const contractNeedsCost = hasCostCap || hasContractCostCap;\n  // Warn once per model when a cost cap is configured but the running model\n  // has no pricing entry — without this the USD cap silently never enforces\n  // (costUSD stays 0), a fail-open the JSDoc on `maxCostUSD` documents.\n  const warnedUnpricedModels = new Set<string>();\n\n  return {\n    name,\n    execute: {\n      before(context) {\n        const counters = namespacedState<BudgetCounters>(context, name);\n        counters.set({\n          tokens: 0,\n          costUSD: 0,\n          warned: false,\n          startedAt: Date.now(),\n          fallbackFired: false,\n        });\n      },\n    },\n    trip: {\n      async after(context, response) {\n        const counters = namespacedState<BudgetCounters>(context, name).get();\n\n        if (!counters) {\n          return;\n        }\n\n        counters.tokens += response.usage.total;\n\n        if (contractNeedsCost) {\n          const pricing = options.pricing?.[context.model.name];\n\n          if (pricing) {\n            const tripCost =\n              (response.usage.input / 1000) * pricing.inputPer1K +\n              (response.usage.output / 1000) * pricing.outputPer1K;\n            counters.costUSD += tripCost;\n          } else if (!warnedUnpricedModels.has(context.model.name)) {\n            // A cost cap is set but no pricing matched the running model, so\n            // costUSD can never grow and the USD cap silently never fires.\n            // Surface the fail-open once per model instead of swallowing it.\n            warnedUnpricedModels.add(context.model.name);\n            console.warn(\n              `ai.middleware.budget(\"${name}\"): a USD cost cap is set but no pricing entry ` +\n                `matches the running model \"${context.model.name}\" — the cap cannot be enforced ` +\n                `for it. Add a pricing entry for \"${context.model.name}\" to options.pricing.`,\n            );\n          }\n        }\n\n        if (hasTokenCap && counters.tokens > options.maxTokens!) {\n          if (onExceeded === \"abort\") {\n            breach(options.maxTokens!, counters.tokens, \"tokens\", name);\n          }\n\n          if (!counters.warned) {\n            counters.warned = true;\n          }\n        }\n\n        if (hasCostCap && counters.costUSD > options.maxCostUSD!) {\n          if (onExceeded === \"abort\") {\n            breach(options.maxCostUSD!, counters.costUSD, \"usd\", name);\n          }\n\n          if (!counters.warned) {\n            counters.warned = true;\n          }\n        }\n\n        if (!contract) {\n          return;\n        }\n\n        if (hasContractTokenCap && counters.tokens > contract.maxTokens!) {\n          await enforceContract(\n            context,\n            counters,\n            name,\n            contractMode,\n            contract,\n            \"tokens\",\n            contract.maxTokens!,\n            counters.tokens,\n          );\n        }\n\n        if (hasContractCostCap && counters.costUSD > contract.maxCostUSD!) {\n          await enforceContract(\n            context,\n            counters,\n            name,\n            contractMode,\n            contract,\n            \"cost\",\n            contract.maxCostUSD!,\n            counters.costUSD,\n          );\n        }\n\n        if (hasContractLatencyCap) {\n          const elapsedMs = Date.now() - counters.startedAt;\n\n          if (elapsedMs > contract.maxLatencyMs!) {\n            await enforceContract(\n              context,\n              counters,\n              name,\n              contractMode,\n              contract,\n              \"latency\",\n              contract.maxLatencyMs!,\n              elapsedMs,\n            );\n          }\n        }\n      },\n    },\n  };\n}\n\n/**\n * Apply the contract's reaction to a single breached clause. `\"abort\"`\n * throws `BudgetExceededError` (stops the run); `\"fallback\"` records the\n * signal once, fires the callback, and returns so the run continues.\n *\n * The callback is invoked at most once per run (guarded by\n * `counters.fallbackFired`) and its rejections are swallowed — a buggy\n * fallback hook must never crash the agent.\n */\nasync function enforceContract(\n  context: MiddlewareExecuteContext,\n  counters: BudgetCounters,\n  name: string,\n  mode: NonNullable<BudgetContract[\"onViolation\"]>,\n  contract: BudgetContract,\n  dimension: BudgetContractDimension,\n  limit: number,\n  actual: number,\n): Promise<void> {\n  if (mode === \"abort\") {\n    breachContract(name, dimension, limit, actual);\n  }\n\n  if (counters.fallbackFired) {\n    return;\n  }\n\n  counters.fallbackFired = true;\n\n  const violation: BudgetContractViolation = {\n    dimension,\n    limit,\n    actual,\n    mode,\n  };\n\n  namespacedState<BudgetContractViolation>(context, `${name}.fallback`).set(\n    violation,\n  );\n\n  if (!contract.fallback) {\n    return;\n  }\n\n  try {\n    await contract.fallback(violation, context);\n  } catch {\n    // A fallback callback is a notification hook — its failure must\n    // never crash the run. Swallow deliberately.\n  }\n}\n"],"mappings":";;;;;;;;;;;AAsHA,MAAM,iBAA8D;CAClE,QAAQ;CACR,MAAM;CACN,SAAS;AACX;AAEA,SAAS,OACP,OACA,QACA,MACA,MACO;CACP,MAAM,IAAI,oBACR,WAAW,KAAK,eAAe,OAAO,GAAG,KAAK,SAAS,MAAM,IAC7D;EAAE;EAAO;EAAQ;CAAK,CACxB;AACF;AAEA,SAAS,eACP,MACA,WACA,OACA,QACO;CACP,MAAM,IAAI,oBACR,WAAW,KAAK,wBAAwB,UAAU,GAAG,OAAO,SAAS,MAAM,IAC3E;EACE;EACA;EACA,MAAM,eAAe;EACrB,SAAS;GAAE;GAAW;GAAO;GAAQ,QAAQ;EAAW;CAC1D,CACF;AACF;;;;;;;;;;;;;;;;;;;;;;;;;;AA2BA,SAAgB,yBACd,OACA,OAAO,UAC2B;CAClC,OAAO,gBACL,EAAE,MAAM,GACR,GAAG,KAAK,UACV,CAAC,CAAC,IAAI;AACR;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;AAyEA,SAAgB,OAAO,SAAyC;CAC9D,MAAM,OAAO,QAAQ,QAAQ;CAC7B,MAAM,aAAa,QAAQ,cAAc;CACzC,MAAM,cAAc,OAAO,QAAQ,cAAc;CACjD,MAAM,aAAa,OAAO,QAAQ,eAAe;CAEjD,MAAM,WAAW,QAAQ;CACzB,MAAM,eAAe,UAAU,eAAe;CAC9C,MAAM,sBAAsB,OAAO,UAAU,cAAc;CAC3D,MAAM,qBAAqB,OAAO,UAAU,eAAe;CAC3D,MAAM,wBAAwB,OAAO,UAAU,iBAAiB;CAChE,MAAM,oBAAoB,cAAc;CAIxC,MAAM,uCAAuB,IAAI,IAAY;CAE7C,OAAO;EACL;EACA,SAAS,EACP,OAAO,SAAS;GAEd,AADiB,gBAAgC,SAAS,IACnD,CAAC,CAAC,IAAI;IACX,QAAQ;IACR,SAAS;IACT,QAAQ;IACR,WAAW,KAAK,IAAI;IACpB,eAAe;GACjB,CAAC;EACH,EACF;EACA,MAAM,EACJ,MAAM,MAAM,SAAS,UAAU;GAC7B,MAAM,WAAW,gBAAgC,SAAS,IAAI,CAAC,CAAC,IAAI;GAEpE,IAAI,CAAC,UACH;GAGF,SAAS,UAAU,SAAS,MAAM;GAElC,IAAI,mBAAmB;IACrB,MAAM,UAAU,QAAQ,UAAU,QAAQ,MAAM;IAEhD,IAAI,SAAS;KACX,MAAM,WACH,SAAS,MAAM,QAAQ,MAAQ,QAAQ,aACvC,SAAS,MAAM,SAAS,MAAQ,QAAQ;KAC3C,SAAS,WAAW;IACtB,OAAO,IAAI,CAAC,qBAAqB,IAAI,QAAQ,MAAM,IAAI,GAAG;KAIxD,qBAAqB,IAAI,QAAQ,MAAM,IAAI;KAC3C,QAAQ,KACN,yBAAyB,KAAK,4EACE,QAAQ,MAAM,KAAK,kEACb,QAAQ,MAAM,KAAK,sBAC3D;IACF;GACF;GAEA,IAAI,eAAe,SAAS,SAAS,QAAQ,WAAY;IACvD,IAAI,eAAe,SACjB,OAAO,QAAQ,WAAY,SAAS,QAAQ,UAAU,IAAI;IAG5D,IAAI,CAAC,SAAS,QACZ,SAAS,SAAS;GAEtB;GAEA,IAAI,cAAc,SAAS,UAAU,QAAQ,YAAa;IACxD,IAAI,eAAe,SACjB,OAAO,QAAQ,YAAa,SAAS,SAAS,OAAO,IAAI;IAG3D,IAAI,CAAC,SAAS,QACZ,SAAS,SAAS;GAEtB;GAEA,IAAI,CAAC,UACH;GAGF,IAAI,uBAAuB,SAAS,SAAS,SAAS,WACpD,MAAM,gBACJ,SACA,UACA,MACA,cACA,UACA,UACA,SAAS,WACT,SAAS,MACX;GAGF,IAAI,sBAAsB,SAAS,UAAU,SAAS,YACpD,MAAM,gBACJ,SACA,UACA,MACA,cACA,UACA,QACA,SAAS,YACT,SAAS,OACX;GAGF,IAAI,uBAAuB;IACzB,MAAM,YAAY,KAAK,IAAI,IAAI,SAAS;IAExC,IAAI,YAAY,SAAS,cACvB,MAAM,gBACJ,SACA,UACA,MACA,cACA,UACA,WACA,SAAS,cACT,SACF;GAEJ;EACF,EACF;CACF;AACF;;;;;;;;;;AAWA,eAAe,gBACb,SACA,UACA,MACA,MACA,UACA,WACA,OACA,QACe;CACf,IAAI,SAAS,SACX,eAAe,MAAM,WAAW,OAAO,MAAM;CAG/C,IAAI,SAAS,eACX;CAGF,SAAS,gBAAgB;CAEzB,MAAM,YAAqC;EACzC;EACA;EACA;EACA;CACF;CAEA,gBAAyC,SAAS,GAAG,KAAK,UAAU,CAAC,CAAC,IACpE,SACF;CAEA,IAAI,CAAC,SAAS,UACZ;CAGF,IAAI;EACF,MAAM,SAAS,SAAS,WAAW,OAAO;CAC5C,QAAQ,CAGR;AACF"}