From bb509f41ab89326320377fa06139595b5fa96453 Mon Sep 17 00:00:00 2001
From: Mr-Neutr0n <64578610+Mr-Neutr0n@users.noreply.github.com>
Date: Sun, 23 Aug 2026 12:16:00 +0000
Subject: [PATCH] fix(budget): bill metered LLM usage instead of a flat
 input-only estimate

The main turn path recorded one estimateCostUsd(inputTokens) per turn,
pricing every model at DEFAULT_AGENT_INPUT_USD_PER_MTOK and ignoring
output/cache tokens entirely, while the per-step metered usage that each
harness already reports through recordLlmRequest (input, output,
cacheRead, cacheWrite, provider-computed costUsd) never reached the
budget tracker.

Budget accounting for turns now comes from that metered usage via
costFromUsage(): it prefers the provider-reported costUsd and falls back
to a total-token estimate at the default rate when a step reports tokens
but no cost. Steps with no reportable usage bill zero rather than a
guessed figure.

Fixes #586

This change was prepared with AI assistance under human direction and review.
---
 src/core/orchestrator.ts |  4 ++--
 src/ratelimit/budget.ts  | 19 +++++++++++++++++++
 test/budget.test.ts      | 27 ++++++++++++++++++++++++++-
 3 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/src/core/orchestrator.ts b/src/core/orchestrator.ts
index d6071d572..3e33d389d 100644
--- a/src/core/orchestrator.ts
+++ b/src/core/orchestrator.ts
@@ -44,7 +44,7 @@ import {
 } from "../credentials/connector-status.ts";
 import { renderComputerBlock, renderResidentLoginsBlock, renderConnectedAppsBlock } from "./environment-facts.ts";
 import { PROVIDERS } from "../connectors/oauth.ts";
-import { estimateCostUsd } from "../ratelimit/budget.ts";
+import { costFromUsage, estimateCostUsd } from "../ratelimit/budget.ts";
 import {
   mintCapabilityToken,
   CAPABILITY_TTL_MS,
@@ -2528,7 +2528,6 @@ export function createOrchestrator(deps: OrchestratorDeps): Orchestrator {
             },
             recordModelCall: (rec) => {
               deps.modelGateway.recordCall({ at: Date.now(), scopeLabel: scopeId, ...rec });
-              void deps.budget?.record(actor.id, estimateCostUsd(rec.inputTokens));
             },
             recordLlmRequest: async (rec) => {
               try {
@@ -2536,6 +2535,7 @@ export function createOrchestrator(deps: OrchestratorDeps): Orchestrator {
               } catch (err) {
                 console.error("[orchestrator] failed to persist LLM request snapshot:", errMessage(err));
               }
+              void deps.budget?.record(actor.id, costFromUsage(rec.usage));
             },
           });
         };
diff --git a/src/ratelimit/budget.ts b/src/ratelimit/budget.ts
index 0825da8c4..46b10136a 100644
--- a/src/ratelimit/budget.ts
+++ b/src/ratelimit/budget.ts
@@ -16,6 +16,25 @@ export function estimateCostUsd(inputTokens: number, usdPerMTok = DEFAULT_AGENT_
   return (inputTokens / 1_000_000) * usdPerMTok;
 }
 
+export function costFromUsage(
+  usage:
+    | {
+        input: number;
+        output: number;
+        cacheRead: number;
+        cacheWrite: number;
+        costUsd: number;
+      }
+    | null
+    | undefined,
+): number {
+  if (!usage) return 0;
+  if (Number.isFinite(usage.costUsd) && usage.costUsd > 0) return usage.costUsd;
+  const tokens = usage.input + usage.output + usage.cacheRead + usage.cacheWrite;
+  if (!Number.isFinite(tokens) || tokens <= 0) return 0;
+  return estimateCostUsd(tokens);
+}
+
 export function createBudgetTracker(
   opts: { limitUsd?: number; orgLimitUsd?: number; windowMs?: number } = {},
 ): BudgetTracker {
diff --git a/test/budget.test.ts b/test/budget.test.ts
index 7eeeb81ed..823ec69dd 100644
--- a/test/budget.test.ts
+++ b/test/budget.test.ts
@@ -5,7 +5,7 @@ import assert from "node:assert/strict";
 import { mkdtempSync } from "node:fs";
 import { tmpdir } from "node:os";
 import { join } from "node:path";
-import { createBudgetTracker, estimateCostUsd } from "../src/ratelimit/budget.ts";
+import { costFromUsage, createBudgetTracker, estimateCostUsd } from "../src/ratelimit/budget.ts";
 import { DEFAULT_AGENT_INPUT_USD_PER_MTOK } from "../src/model/pi-models.ts";
 import { buildApp } from "../src/wiring.ts";
 import type { TurnRequest } from "../src/types.ts";
@@ -59,3 +59,28 @@ test("a principal over budget is refused by the app", async () => {
   assert.equal(second.status, "refused");
   assert.match(second.reason ?? "", /budget exceeded/);
 });
+
+test("costFromUsage prefers the provider-reported cost when it is usable", () => {
+  assert.equal(costFromUsage({ input: 1_000_000, output: 0, cacheRead: 0, cacheWrite: 0, costUsd: 0.03 }), 0.03);
+});
+
+test("costFromUsage falls back to a token-total estimate without a usable cost", () => {
+  const fallback = { input: 500_000, output: 300_000, cacheRead: 100_000, cacheWrite: 100_000, costUsd: 0 };
+  assert.equal(costFromUsage(fallback), DEFAULT_AGENT_INPUT_USD_PER_MTOK);
+  assert.equal(
+    costFromUsage({ input: 0, output: 0, cacheRead: 0, cacheWrite: 0, costUsd: Number.NaN }),
+    0,
+    "no tokens and no cost means nothing to bill",
+  );
+});
+
+test("costFromUsage treats missing metering as zero cost", () => {
+  assert.equal(costFromUsage(null), 0);
+  assert.equal(costFromUsage(undefined), 0);
+});
+
+test("metered output tokens count against the budget even with no input growth", async () => {
+  const b = createBudgetTracker({ limitUsd: DEFAULT_AGENT_INPUT_USD_PER_MTOK, windowMs: 60_000 });
+  await b.record("U1", costFromUsage({ input: 0, output: 2_000_000, cacheRead: 0, cacheWrite: 0, costUsd: 0 }), 1000);
+  assert.equal((await b.check("U1", 1000)).allowed, false);
+});