diff --git a/src/infra/provider-usage.auth.ts b/src/infra/provider-usage.auth.ts index c503779b6f5..3d9e30fa15f 100644 --- a/src/infra/provider-usage.auth.ts +++ b/src/infra/provider-usage.auth.ts @@ -220,6 +220,14 @@ async function resolveProviderUsageAuthFallback(params: { }); return apiKey ? { provider: "xiaomi", token: apiKey } : null; } + case "kilocode": { + const apiKey = resolveProviderApiKeyFromConfigAndStore({ + state: params.state, + providerIds: ["kilocode"], + envDirect: [params.state.env.KILOCODE_API_KEY], + }); + return apiKey ? { provider: "kilocode", token: apiKey } : null; + } default: return null; } diff --git a/src/infra/provider-usage.fetch.kilocode.test.ts b/src/infra/provider-usage.fetch.kilocode.test.ts new file mode 100644 index 00000000000..bd1264fba89 --- /dev/null +++ b/src/infra/provider-usage.fetch.kilocode.test.ts @@ -0,0 +1,105 @@ +import { describe, expect, it } from "vitest"; +import { createProviderUsageFetch, makeResponse } from "../test-utils/provider-usage-fetch.js"; +import { fetchKilocodeUsage } from "./provider-usage.fetch.kilocode.js"; + +describe("fetchKilocodeUsage", () => { + it("returns HTTP error snapshot for failed requests", async () => { + const mockFetch = createProviderUsageFetch(async () => makeResponse(401, "unauthorized")); + const result = await fetchKilocodeUsage("key", 5000, mockFetch); + + expect(result.error).toBe("HTTP 401"); + expect(result.windows).toHaveLength(0); + }); + + it("returns HTTP error snapshot for server errors", async () => { + const mockFetch = createProviderUsageFetch(async () => makeResponse(503, "unavailable")); + const result = await fetchKilocodeUsage("key", 5000, mockFetch); + + expect(result.error).toBe("HTTP 503"); + expect(result.windows).toHaveLength(0); + }); + + it("returns balance as plan label when account is active", async () => { + const mockFetch = createProviderUsageFetch(async () => + makeResponse(200, { balance: 19.604203, isDepleted: false }), + ); + + const result = await fetchKilocodeUsage("key", 5000, mockFetch); + + expect(result.provider).toBe("kilocode"); + expect(result.displayName).toBe("Kilo"); + expect(result.plan).toBe("$19.60"); + expect(result.windows).toHaveLength(0); + expect(result.error).toBeUndefined(); + }); + + it("shows depleted label in plan when isDepleted is true", async () => { + const mockFetch = createProviderUsageFetch(async () => + makeResponse(200, { balance: 0, isDepleted: true }), + ); + + const result = await fetchKilocodeUsage("key", 5000, mockFetch); + + expect(result.provider).toBe("kilocode"); + expect(result.plan).toBe("$0.00 (depleted)"); + expect(result.error).toBeUndefined(); + expect(result.windows).toHaveLength(0); + }); + + it("shows depleted label without balance when isDepleted and no balance", async () => { + const mockFetch = createProviderUsageFetch(async () => makeResponse(200, { isDepleted: true })); + + const result = await fetchKilocodeUsage("key", 5000, mockFetch); + + expect(result.plan).toBe("depleted"); + expect(result.error).toBeUndefined(); + expect(result.windows).toHaveLength(0); + }); + + it("handles missing balance field gracefully", async () => { + const mockFetch = createProviderUsageFetch(async () => + makeResponse(200, { isDepleted: false }), + ); + + const result = await fetchKilocodeUsage("key", 5000, mockFetch); + + expect(result.plan).toBeUndefined(); + expect(result.windows).toHaveLength(0); + expect(result.error).toBeUndefined(); + }); + + it("rounds balance to two decimal places", async () => { + const mockFetch = createProviderUsageFetch(async () => + makeResponse(200, { balance: 100, isDepleted: false }), + ); + + const result = await fetchKilocodeUsage("key", 5000, mockFetch); + expect(result.plan).toBe("$100.00"); + }); + + it("returns error snapshot for invalid JSON response", async () => { + const mockFetch = createProviderUsageFetch(async () => { + const res = new Response("not-json", { status: 200 }); + return res; + }); + + const result = await fetchKilocodeUsage("key", 5000, mockFetch); + + expect(result.error).toBe("Invalid JSON"); + expect(result.windows).toHaveLength(0); + }); + + it("sends Authorization header with Bearer token", async () => { + let capturedInit: RequestInit | undefined; + const mockFetch = createProviderUsageFetch(async (_url, init) => { + capturedInit = init; + return makeResponse(200, { balance: 5.0, isDepleted: false }); + }); + + await fetchKilocodeUsage("test-api-key", 5000, mockFetch); + + expect((capturedInit?.headers as Record)?.["Authorization"]).toBe( + "Bearer test-api-key", + ); + }); +}); diff --git a/src/infra/provider-usage.fetch.kilocode.ts b/src/infra/provider-usage.fetch.kilocode.ts new file mode 100644 index 00000000000..6fb541a1916 --- /dev/null +++ b/src/infra/provider-usage.fetch.kilocode.ts @@ -0,0 +1,60 @@ +import { + buildUsageErrorSnapshot, + buildUsageHttpErrorSnapshot, + fetchJson, +} from "./provider-usage.fetch.shared.js"; +import { PROVIDER_LABELS } from "./provider-usage.shared.js"; +import type { ProviderUsageSnapshot } from "./provider-usage.types.js"; + +const KILOCODE_BALANCE_URL = "https://api.kilo.ai/api/profile/balance"; + +type KilocodeBalanceResponse = { + balance?: number; + isDepleted?: boolean; +}; + +export async function fetchKilocodeUsage( + apiKey: string, + timeoutMs: number, + fetchFn: typeof fetch, +): Promise { + const res = await fetchJson( + KILOCODE_BALANCE_URL, + { + method: "GET", + headers: { + Authorization: `Bearer ${apiKey}`, + Accept: "application/json", + }, + }, + timeoutMs, + fetchFn, + ); + + if (!res.ok) { + return buildUsageHttpErrorSnapshot({ provider: "kilocode", status: res.status }); + } + + const data = (await res.json().catch(() => null)) as KilocodeBalanceResponse | null; + if (!data) { + return buildUsageErrorSnapshot("kilocode", "Invalid JSON"); + } + const balance = typeof data.balance === "number" ? data.balance : null; + // Show credit balance as a plan label since Kilo's endpoint returns a dollar + // amount rather than quota windows with usage percentages. Fold the depleted + // flag into the label so status renderers display it rather than dropping the + // snapshot (formatUsageWindowSummary silently drops snapshots with error set). + const balanceLabel = balance !== null ? `$${balance.toFixed(2)}` : undefined; + const plan = data.isDepleted + ? balanceLabel + ? `${balanceLabel} (depleted)` + : "depleted" + : balanceLabel; + + return { + provider: "kilocode", + displayName: PROVIDER_LABELS.kilocode, + windows: [], + plan, + }; +} diff --git a/src/infra/provider-usage.fetch.ts b/src/infra/provider-usage.fetch.ts index 87f216eef24..e5c781aafc0 100644 --- a/src/infra/provider-usage.fetch.ts +++ b/src/infra/provider-usage.fetch.ts @@ -1,5 +1,6 @@ export { fetchClaudeUsage } from "./provider-usage.fetch.claude.js"; export { fetchCodexUsage } from "./provider-usage.fetch.codex.js"; export { fetchGeminiUsage } from "./provider-usage.fetch.gemini.js"; +export { fetchKilocodeUsage } from "./provider-usage.fetch.kilocode.js"; export { fetchMinimaxUsage } from "./provider-usage.fetch.minimax.js"; export { fetchZaiUsage } from "./provider-usage.fetch.zai.js"; diff --git a/src/infra/provider-usage.load.ts b/src/infra/provider-usage.load.ts index ec870aa27ee..ea6463bb816 100644 --- a/src/infra/provider-usage.load.ts +++ b/src/infra/provider-usage.load.ts @@ -6,6 +6,7 @@ import { fetchClaudeUsage, fetchCodexUsage, fetchGeminiUsage, + fetchKilocodeUsage, fetchMinimaxUsage, fetchZaiUsage, } from "./provider-usage.fetch.js"; @@ -95,6 +96,8 @@ async function fetchProviderUsageSnapshotFallback(params: { params.timeoutMs, params.fetchFn, ); + case "kilocode": + return await fetchKilocodeUsage(params.auth.token, params.timeoutMs, params.fetchFn); case "zai": return await fetchZaiUsage(params.auth.token, params.timeoutMs, params.fetchFn); case "minimax": diff --git a/src/infra/provider-usage.shared.ts b/src/infra/provider-usage.shared.ts index b801da4824c..690c5971c85 100644 --- a/src/infra/provider-usage.shared.ts +++ b/src/infra/provider-usage.shared.ts @@ -11,6 +11,7 @@ export const PROVIDER_LABELS: Record = { anthropic: "Claude", "github-copilot": "Copilot", "google-gemini-cli": "Gemini", + kilocode: "Kilo", minimax: "MiniMax", "openai-codex": "Codex", xiaomi: "Xiaomi", @@ -21,6 +22,7 @@ export const usageProviders: UsageProviderId[] = [ "anthropic", "github-copilot", "google-gemini-cli", + "kilocode", "minimax", "openai-codex", "xiaomi", diff --git a/src/infra/provider-usage.types.ts b/src/infra/provider-usage.types.ts index af5e2e93c8b..506f56c8ce0 100644 --- a/src/infra/provider-usage.types.ts +++ b/src/infra/provider-usage.types.ts @@ -21,6 +21,7 @@ export type UsageProviderId = | "anthropic" | "github-copilot" | "google-gemini-cli" + | "kilocode" | "minimax" | "openai-codex" | "xiaomi"