From 29d14993282334cffa0f9a2262572502535bb0a3 Mon Sep 17 00:00:00 2001 From: artickc Date: Fri, 12 Jun 2026 16:15:33 +0300 Subject: [PATCH] feat(vertex): self-tracked USD spend since account added Vertex AI exposes no usage/quota API for an API key or Service Account (credit balance lives behind the Cloud Billing API the proxy credential can't reach), so the Limits page showed nothing for Vertex. Add a vertex/vertex-partner usage fetcher that reports the USD OmniRoute has spent through the connection since it was added: getConnectionSpendUsdSinceAdded sums per-model tokens from usage_history (scoped to the provider and success=1 so failed/cross-provider rows don't inflate the total) and prices them via the backend pricing table (calculateCost). Returns a '\ used since this account was added' message plus a 'spend' quota (quotaSource localUsageHistory) so the limits cache persists it. Registered in USAGE_FETCHER_PROVIDERS, USAGE_SUPPORTED_PROVIDERS and PROVIDER_LIMITS_APIKEY_PROVIDERS. Tests included. --- open-sse/services/usage.ts | 53 +++++++++++++ src/lib/usage/providerLimits.ts | 2 + src/lib/usage/usageStats.ts | 59 +++++++++++++++ src/shared/constants/providers.ts | 2 + tests/unit/vertex-spend-usage.test.ts | 105 ++++++++++++++++++++++++++ 5 files changed, 221 insertions(+) create mode 100644 tests/unit/vertex-spend-usage.test.ts diff --git a/open-sse/services/usage.ts b/open-sse/services/usage.ts index d33cbcea4ab6..f198ddcfe429 100644 --- a/open-sse/services/usage.ts +++ b/open-sse/services/usage.ts @@ -1477,6 +1477,8 @@ export const USAGE_FETCHER_PROVIDERS = [ "opencode", "opencode-zen", "xiaomi-mimo", + "vertex", + "vertex-partner", ] as const; export type UsageFetcherProvider = (typeof USAGE_FETCHER_PROVIDERS)[number]; @@ -1516,6 +1518,9 @@ export async function getUsageForProvider( case "kiro": case "amazon-q": return await getKiroUsage(accessToken, providerSpecificData); + case "vertex": + case "vertex-partner": + return await getVertexUsage(id || "", provider); case "kimi-coding": return await getKimiUsage(accessToken); case "qwen": @@ -3051,6 +3056,53 @@ async function getKiroUsage(accessToken?: string, providerSpecificData?: JsonRec } } +/** + * Vertex AI — SELF-TRACKED spend. + * + * Vertex AI exposes no usage/quota API for an API key or Service Account (billing/credit balance + * lives behind the Cloud Billing API, which the proxy credential can't reach). Instead we report + * the USD that OmniRoute has spent through this connection since the account was added — summed + * from `usage_history` and priced via the backend pricing table. Returns a `message` (with the $ + * figure) plus a `spend` quota entry so the limits cache persists it (a message-only result is + * treated as a transient error and not cached). + */ +async function getVertexUsage(connectionId: string, provider: string) { + if (!connectionId) { + return { message: "Vertex connected. Connection id unavailable for usage tracking." }; + } + try { + const { getConnectionSpendUsdSinceAdded } = await import("@/lib/usage/usageStats"); + const { costUsd, requests } = await getConnectionSpendUsdSinceAdded(provider, connectionId); + + const spend: JsonRecord = { + used: Number(costUsd.toFixed(6)), + displayName: "Spend (USD)", + quotaSource: "localUsageHistory", + resetAt: null, + unlimited: false, + }; + + if (requests === 0) { + return { + plan: "Vertex AI", + message: "Vertex connected. No usage recorded through OmniRoute yet for this account.", + quotas: { spend }, + }; + } + + const costStr = costUsd >= 1 ? costUsd.toFixed(2) : costUsd.toFixed(4); + return { + plan: "Vertex AI", + message: `$${costStr} used since this account was added \u00b7 ${requests} request${ + requests === 1 ? "" : "s" + }`, + quotas: { spend }, + }; + } catch (error) { + return { message: `Vertex usage tracking error: ${(error as Error).message}` }; + } +} + /** * Map Kimi membership level to display name * LEVEL_BASIC = Moderato, LEVEL_INTERMEDIATE = Allegretto, @@ -3279,6 +3331,7 @@ export const __testing = { getMiniMaxRemainingPercent, getMiniMaxUsage, getXiaomiMimoUsage, + getVertexUsage, getMiniMaxAuthErrorMessage, getMiniMaxErrorSummary, mapCodeAssistSubscriptionToPlanLabel, diff --git a/src/lib/usage/providerLimits.ts b/src/lib/usage/providerLimits.ts index 2ae6da88d23e..b14997ed04bb 100644 --- a/src/lib/usage/providerLimits.ts +++ b/src/lib/usage/providerLimits.ts @@ -70,6 +70,8 @@ const PROVIDER_LIMITS_APIKEY_PROVIDERS = new Set([ "nanogpt", "deepseek", "xiaomi-mimo", + "vertex", + "vertex-partner", ]); const DEFAULT_PROVIDER_LIMITS_SYNC_INTERVAL_MINUTES = 70; const PROVIDER_LIMITS_AUTO_SYNC_SETTING_KEY = "provider_limits_auto_sync_last_run"; diff --git a/src/lib/usage/usageStats.ts b/src/lib/usage/usageStats.ts index 514fe4b8a2bf..6c10405ec6a2 100644 --- a/src/lib/usage/usageStats.ts +++ b/src/lib/usage/usageStats.ts @@ -219,6 +219,65 @@ export function getMonthlyProviderTokensForConnection( return Math.max(0, Number(row?.total ?? 0)); } +/** + * Total USD spend OmniRoute has recorded for a single provider connection, across all time + * (i.e. "since the account was added" — usage_history rows only exist from first use onward). + * + * Sums per-model token usage from `usage_history` for the connection and prices each model via the + * backend pricing table (`calculateCost`). Scoped to the given `provider` and to **successful** + * requests (`success = 1`) so failed/errored calls and any cross-provider rows can't inflate the + * total. Only reflects traffic that went THROUGH OmniRoute, not the provider's own dashboard. Used + * to surface a "$X used since added" figure for providers that expose no native usage/quota API + * (e.g. Vertex AI). + */ +export async function getConnectionSpendUsdSinceAdded( + provider: string, + connectionId: string +): Promise<{ costUsd: number; requests: number }> { + if (!provider || !connectionId) return { costUsd: 0, requests: 0 }; + + const db = getDbInstance(); + const rows = db + .prepare( + `SELECT model, + COALESCE(SUM(tokens_input), 0) AS input, + COALESCE(SUM(tokens_output), 0) AS output, + COALESCE(SUM(tokens_cache_read), 0) AS cacheRead, + COALESCE(SUM(tokens_cache_creation), 0) AS cacheCreation, + COALESCE(SUM(tokens_reasoning), 0) AS reasoning, + COUNT(*) AS requests + FROM usage_history + WHERE connection_id = ? AND provider = ? AND success = 1 + GROUP BY model` + ) + .all(connectionId, provider) as Array<{ + model?: string; + input?: number; + output?: number; + cacheRead?: number; + cacheCreation?: number; + reasoning?: number; + requests?: number; + }>; + + let costUsd = 0; + let requests = 0; + for (const row of rows) { + requests += Math.max(0, Number(row.requests ?? 0)); + const model = typeof row.model === "string" ? row.model : ""; + const tokens = { + input: Number(row.input ?? 0), + output: Number(row.output ?? 0), + cacheRead: Number(row.cacheRead ?? 0), + cacheCreation: Number(row.cacheCreation ?? 0), + reasoning: Number(row.reasoning ?? 0), + }; + costUsd += await calculateCost(provider, model, tokens, { provider, model }); + } + + return { costUsd: Math.max(0, costUsd), requests }; +} + /** * Get aggregated usage stats. * Uses UNION of recent raw data and older aggregated data when aggregation is enabled. diff --git a/src/shared/constants/providers.ts b/src/shared/constants/providers.ts index 38e8d672c897..184ffbbef811 100644 --- a/src/shared/constants/providers.ts +++ b/src/shared/constants/providers.ts @@ -3127,6 +3127,8 @@ export const USAGE_SUPPORTED_PROVIDERS = [ "nanogpt", "deepseek", "xiaomi-mimo", + "vertex", + "vertex-partner", ]; // ── Zod validation at module load (Phase 7.2) ── diff --git a/tests/unit/vertex-spend-usage.test.ts b/tests/unit/vertex-spend-usage.test.ts new file mode 100644 index 000000000000..4a9bb5e13be7 --- /dev/null +++ b/tests/unit/vertex-spend-usage.test.ts @@ -0,0 +1,105 @@ +/** + * tests/unit/vertex-spend-usage.test.ts + * + * Vertex AI exposes no native usage/quota API for an API key or Service Account, so OmniRoute + * SELF-TRACKS spend: it sums the tokens it routed to the connection (usage_history) and prices + * them via the backend pricing table, surfacing a "$X used since this account was added" figure. + * These tests cover the aggregation helper + the fetcher response shape with a real temp DB. + */ + +import { describe, it, before, after } from "node:test"; +import assert from "node:assert/strict"; +import os from "node:os"; +import path from "node:path"; +import fs from "node:fs"; + +// DATA_DIR must be set before any module that opens the DB is imported. +const TMP = fs.mkdtempSync(path.join(os.tmpdir(), "omni-vertex-")); +process.env.DATA_DIR = TMP; + +const core = await import("../../src/lib/db/core.ts"); +const { getConnectionSpendUsdSinceAdded } = await import("../../src/lib/usage/usageStats.ts"); +const { __testing } = await import("../../open-sse/services/usage.ts"); +const { getVertexUsage } = __testing; + +function insertUsage( + connectionId: string, + provider: string, + model: string, + tokensIn: number, + tokensOut: number, + success = 1 +) { + const db = core.getDbInstance(); + db.prepare( + `INSERT INTO usage_history (provider, model, connection_id, tokens_input, tokens_output, success, timestamp) + VALUES (?, ?, ?, ?, ?, ?, ?)` + ).run(provider, model, connectionId, tokensIn, tokensOut, success, new Date().toISOString()); +} + +describe("vertex self-tracked spend", () => { + before(() => { + core.getDbInstance(); // trigger migrations + // conn-v: two SUCCESSFUL priced requests across two models. + insertUsage("conn-v", "vertex", "gemini-2.5-flash", 1_000_000, 500_000, 1); + insertUsage("conn-v", "vertex", "gemini-3-pro-image-preview", 200_000, 100_000, 1); + // a FAILED request on the same connection must NOT count toward spend. + insertUsage("conn-v", "vertex", "gemini-2.5-flash", 5_000_000, 5_000_000, 0); + // a row with a different provider on the same connection id must NOT bleed in. + insertUsage("conn-v", "vertex-partner", "claude-opus-4-7", 9_000_000, 9_000_000, 1); + // a different connection must not bleed in + insertUsage("conn-other", "vertex", "gemini-2.5-flash", 9_000_000, 9_000_000, 1); + }); + + after(() => { + core.resetDbInstance(); + try { + fs.rmSync(TMP, { recursive: true, force: true }); + } catch { + // best-effort temp cleanup + } + }); + + it("counts only the connection's successful, same-provider requests", async () => { + const { costUsd, requests } = await getConnectionSpendUsdSinceAdded("vertex", "conn-v"); + assert.equal( + requests, + 2, + "only the two successful vertex rows count (failed + vertex-partner + other-conn excluded)" + ); + assert.ok(Number.isFinite(costUsd) && costUsd >= 0, "cost is a finite, non-negative number"); + }); + + it("returns 0/0 for an unknown connection (no bleed)", async () => { + const { costUsd, requests } = await getConnectionSpendUsdSinceAdded("vertex", "conn-none"); + assert.equal(requests, 0); + assert.equal(costUsd, 0); + }); + + it("getVertexUsage returns a spend quota + $ message for a used connection", async () => { + const r = (await getVertexUsage("conn-v", "vertex")) as { + plan?: string; + message?: string; + quotas?: Record; + }; + assert.ok(r.quotas?.spend, "spend quota present (so the limits cache persists it)"); + assert.equal(r.quotas!.spend.quotaSource, "localUsageHistory"); + assert.ok(typeof r.quotas!.spend.used === "number" && r.quotas!.spend.used >= 0); + assert.ok(r.message && r.message.includes("$"), "message carries the dollar figure"); + assert.ok(r.message!.includes("2 requests"), "message reports the request count"); + }); + + it("getVertexUsage reports no-usage cleanly when nothing was routed", async () => { + const r = (await getVertexUsage("conn-empty", "vertex")) as { + message?: string; + quotas?: Record; + }; + assert.ok(r.message && /no usage/i.test(r.message), "informative no-usage message"); + assert.equal(r.quotas?.spend.used, 0); + }); + + it("getVertexUsage returns a message when connection id is missing", async () => { + const r = (await getVertexUsage("", "vertex")) as { message?: string; quotas?: unknown }; + assert.ok(r.message && !r.quotas, "no spend quota without a connection id"); + }); +});