From 158d19c9bed14e4bddf81de340c506136c787671 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?=
 <zhangyanghaha0407@outlook.com>
Date: Wed, 29 Apr 2026 20:40:05 -0400
Subject: [PATCH] fix(budget): add cost fallback for all zero-pricing providers
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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.
---
 src/main/agent-runtime/cost-writer.ts | 16 +++-------
 src/main/model-resolver.ts            | 43 +++++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 12 deletions(-)

diff --git a/src/main/agent-runtime/cost-writer.ts b/src/main/agent-runtime/cost-writer.ts
index 1e8407c4..21771062 100644
--- a/src/main/agent-runtime/cost-writer.ts
+++ b/src/main/agent-runtime/cost-writer.ts
@@ -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}`,
diff --git a/src/main/model-resolver.ts b/src/main/model-resolver.ts
index be873cdb..8daa88b3 100644
--- a/src/main/model-resolver.ts
+++ b/src/main/model-resolver.ts
@@ -94,6 +94,49 @@ export function resolveModel(provider: string, modelId: string): Model<Api> | nu
   }
 }
 
+const CROSS_PROVIDER_COST_LOOKUP: Record<string, string[]> = {
+  "github-copilot": ["anthropic", "openai", "google", "xai"],
+  "google-gemini-cli": ["google"],
+};
+
+const PROVIDER_FALLBACK_COST: Record<string, { input: number; output: number }> = {
+  "zai": { input: 1, output: 3.2 },
+  "kimi-coding": { input: 0.6, output: 3 },
+};
+
+function normalizeModelId(id: string): string {
+  return id.replace(/\./g, "-");
+}
+
+export function resolveModelCostFallback(model: Model<Api>, provider: string): Model<Api> {
+  if (model.cost.input > 0 || model.cost.output > 0) return model;
+
+  const crossProviders = CROSS_PROVIDER_COST_LOOKUP[provider];
+  if (crossProviders) {
+    const normalized = normalizeModelId(model.id);
+    for (const cp of crossProviders) {
+      try {
+        const sourceModels = getModels(cp as Parameters<typeof getModels>[0]);
+        const match = sourceModels.find((m) => m.id === model.id)
+          ?? sourceModels.find((m) => normalizeModelId(m.id) === normalized)
+          ?? sourceModels.find((m) => normalized.includes(normalizeModelId(m.id)) || normalizeModelId(m.id).includes(normalized));
+        if (match && (match.cost.input > 0 || match.cost.output > 0)) {
+          return { ...model, cost: { ...match.cost } };
+        }
+      } catch {
+        continue;
+      }
+    }
+  }
+
+  const fallback = PROVIDER_FALLBACK_COST[provider];
+  if (fallback) {
+    return { ...model, cost: { ...model.cost, input: fallback.input, output: fallback.output } };
+  }
+
+  return model;
+}
+
 export function resolveModels(provider: string): Array<{ id: string; name: string }> {
   if (provider === "fireworks-firepass") {
     return [{ id: FIREWORKS_FIREPASS_MODEL_ID, name: FIREWORKS_FIREPASS_MODEL_NAME }];