Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 53 additions & 0 deletions open-sse/services/usage.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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];
Expand Down Expand Up @@ -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":
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -3279,6 +3331,7 @@ export const __testing = {
getMiniMaxRemainingPercent,
getMiniMaxUsage,
getXiaomiMimoUsage,
getVertexUsage,
getMiniMaxAuthErrorMessage,
getMiniMaxErrorSummary,
mapCodeAssistSubscriptionToPlanLabel,
Expand Down
2 changes: 2 additions & 0 deletions src/lib/usage/providerLimits.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down
59 changes: 59 additions & 0 deletions src/lib/usage/usageStats.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
2 changes: 2 additions & 0 deletions src/shared/constants/providers.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3127,6 +3127,8 @@ export const USAGE_SUPPORTED_PROVIDERS = [
"nanogpt",
"deepseek",
"xiaomi-mimo",
"vertex",
"vertex-partner",
];

// ── Zod validation at module load (Phase 7.2) ──
Expand Down
105 changes: 105 additions & 0 deletions tests/unit/vertex-spend-usage.test.ts
Original file line number Diff line number Diff line change
@@ -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<string, { used: number; quotaSource?: string; displayName?: string }>;
};
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<string, { used: number }>;
};
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");
});
});