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 | nu } } +const CROSS_PROVIDER_COST_LOOKUP: Record = { + "github-copilot": ["anthropic", "openai", "google", "xai"], + "google-gemini-cli": ["google"], +}; + +const PROVIDER_FALLBACK_COST: Record = { + "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, provider: string): Model { + 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[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 }];