fix(budget): add cost fallback for all zero-pricing providers

Pi-ai ships $0 pricing for github-copilot (25 models), google-gemini-cli
(6 models), zai (13 models), and kimi-coding (2 models). This made cost
tracking produce $0 for every run through these providers, rendering
budget enforcement useless.

Added resolveModelCostFallback() in model-resolver.ts with two strategies:
1. Cross-provider lookup: subscription providers (github-copilot,
   google-gemini-cli) use the same models as their API counterparts.
   Look up the model in anthropic/openai/google/xai catalogs using
   normalized ID matching (dots→hyphens) to find real pricing.
2. Per-provider fallback: API-key providers with no cross-provider
   equivalent (zai, kimi-coding) get a static per-provider cost.

CostWriter now calls resolveModelCostFallback after resolveModel, so
all consumers get patched pricing without touching the pi-ai library.
This commit is contained in:
ZacharyZhang-NY committed 2026-04-29 20:40:05 -04:00
1 parent 7a78a34558
commit 158d19c9be
2 files changed
+47 -12

No files matched your search

+4 -12
View File
@@ -1,7 +1,7 @@
import { calculateCost, type Usage } from "@mariozechner/pi-ai";
import { calculateCost, type Usage, type Api, type Model } from "@mariozechner/pi-ai";
import type { AppDatabase } from "@main/database";
import { getConnectorDefinition } from "@main/connectors";
import { resolveModel } from "@main/model-resolver";
import { resolveModel, resolveModelCostFallback } from "@main/model-resolver";
export interface CostWriterDeps {
db: AppDatabase;
@@ -17,19 +17,10 @@ interface ObserveInput {
usage: Usage;
}
/**
* Writes a `cost_entries` row for each pi `message_end`. finishRun reads
* SUM(cost_entries.amount_usd) for the run-level aggregate, so this is
* the single source of truth for cost attribution.
*/
export class CostWriter {
constructor(private readonly deps: CostWriterDeps) {}
observeAssistantMessage(input: ObserveInput): void {
// Fail closed on cost calculation — budget policies read `cost_entries`
// to enforce spending gates; silently emitting an unattributed row
// would let an agent run unbounded. If pricing cannot be resolved,
// throw so the caller logs and the missing cost is visible.
const definition = getConnectorDefinition(input.connectorId);
if (!definition.provider) {
throw new Error(`[cost-writer] connector ${input.connectorId} has no provider mapping`);
@@ -38,7 +29,8 @@ export class CostWriter {
if (!model) {
throw new Error(`[cost-writer] model ${input.modelId} not found for provider ${definition.provider}`);
}
const cost = calculateCost(model, input.usage);
const patchedModel = resolveModelCostFallback(model, definition.provider);
const cost = calculateCost(patchedModel, input.usage);
if (!Number.isFinite(cost.total)) {
throw new Error(
`[cost-writer] cost calculation returned non-finite total for ${input.connectorId}/${input.modelId}`,