packages/forkhorse-mcp/src/usage.ts

1import { z } from "zod";
2import { PublicError } from "./core";
3
4export const tokenCount = z.number().int().min(0).max(1_000_000_000);
5const record = z.object({
6  inputTokens: tokenCount,
7  outputTokens: tokenCount,
8  model: z.string().optional(),
9  source: z.enum(["reported", "metered"]),
10  at: z.string(),
11});
12export const usageSourcesSchema = z.object({
13  reported: record.extend({ source: z.literal("reported") }).nullable(),
14  metered: record.extend({ source: z.literal("metered") }).nullable(),
15});
16export const budgetStateSchema = z.enum(["unchecked", "within", "near", "over"]);
17const binding = z.object({
18  agent: z.string(),
19  repoName: z.string(),
20  // Older core records have no usage fields. Missing means unreported, never zero.
21  usage: usageSourcesSchema.optional(),
22});
23const usageState = z.object({
24  claims: z.array(binding.extend({ packet: z.string() })),
25  races: z
26    .array(
27      z.object({
28        packet: z.string(),
29        tokenBudget: z.number().int().positive().nullable().optional(),
30        entrants: z.array(binding.extend({ budget: budgetStateSchema.optional() })),
31      }),
32    )
33    .default([]),
34});
35
36/** FH-060 state projection only: core computes race budgets and enforces live bindings. */
37export function callerUsage(snapshot: unknown, agent: string) {
38  const state = usageState.safeParse(snapshot);
39  if (!state.success) throw new PublicError(502, "Core usage state unavailable; refresh status.");
40  const empty = { reported: null, metered: null };
41  return [
42    ...state.data.claims
43      .filter((claim) => claim.agent === agent)
44      .map((claim) => ({
45        packet: claim.packet,
46        repoName: claim.repoName,
47        usage: claim.usage ?? empty,
48        tokenBudget: null,
49        budget: "unchecked" as const,
50      })),
51    ...state.data.races.flatMap((race) =>
52      race.entrants
53        .filter((entrant) => entrant.agent === agent)
54        .map((entrant) => {
55          // A budget without core's computed state must never look safe to an agent.
56          if (race.tokenBudget != null && entrant.budget === undefined)
57            throw new PublicError(502, "Core race budget state unavailable; FH-060 is required.");
58          return {
59            packet: race.packet,
60            repoName: entrant.repoName,
61            usage: entrant.usage ?? empty,
62            tokenBudget: race.tokenBudget ?? null,
63            budget: entrant.budget ?? "unchecked",
64          };
65        }),
66    ),
67  ];
68}
69