packages/forkhorse-core/src/usage.ts

1import { HttpError } from "./http";
2import type { TokenBudgetState, UsageSources } from "./workflow-types";
3
4/** Normalize legacy records without backfilling or inventing measurements. */
5export const usageSources = (usage?: UsageSources): UsageSources => ({
6  reported: usage?.reported ?? null, metered: usage?.metered ?? null,
7});
8
9export const tokenCount = (value: unknown): value is number =>
10  typeof value === "number" && Number.isInteger(value) && value >= 0 && value <= 1_000_000_000;
11
12export function tokenBudgetState(tokenBudget: number | undefined, usage?: UsageSources): TokenBudgetState {
13  const measured = usage?.metered ?? usage?.reported;
14  if (tokenBudget === undefined || !measured) return "unchecked";
15  const total = measured.inputTokens + measured.outputTokens;
16  return total > tokenBudget ? "over" : total >= tokenBudget * 0.8 ? "near" : "within";
17}
18
19export function requireTokenBudget(tokenBudget: number | undefined, usage?: UsageSources): void {
20  if (tokenBudget === undefined) return;
21  const status = tokenBudgetState(tokenBudget, usage);
22  if (status === "unchecked")
23    throw new HttpError(409, "Report usage first: this race has a token budget.");
24  if (status === "over")
25    throw new HttpError(409, `Over its token budget (${usage?.metered ? "metered" : "reported"}).`);
26}
27