diff --git a/open-sse/services/usage.ts b/open-sse/services/usage.ts index 0493f7e65d..171f68dfe4 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": @@ -3103,6 +3108,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, @@ -3331,6 +3383,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 2ae6da88d2..b14997ed04 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 514fe4b8a2..6c10405ec6 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 38e8d672c8..184ffbbef8 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 0000000000..4a9bb5e13b --- /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"); + }); +});