diff --git a/plugins/kimi-usage-source/server/usage.test.ts b/plugins/kimi-usage-source/server/usage.test.ts index a67e416272..5c5783a4d5 100644 --- a/plugins/kimi-usage-source/server/usage.test.ts +++ b/plugins/kimi-usage-source/server/usage.test.ts @@ -5,14 +5,22 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { discover, fetchUsage } from "./usage.js"; import type { UsageReport } from "@getpaseo/plugin/server/usage"; -function kimiCredentialPath(dir: string): string { - return join(dir, "credentials", "kimi-code.json"); +// The slot the Kimi CLI writes for a login to the global (.ai) deployment. +const GLOBAL_SLOT = "kimi-code-env-0e4f99c69cc27850"; + +function kimiCredentialPath(dir: string, slot = "kimi-code"): string { + return join(dir, "credentials", `${slot}.json`); } -function writeKimiCredentials(dir: string, accessToken: string, overrides: object = {}): void { +function writeKimiCredentials( + dir: string, + accessToken: string, + overrides: object = {}, + slot = "kimi-code", +): void { mkdirSync(join(dir, "credentials"), { recursive: true }); writeFileSync( - kimiCredentialPath(dir), + kimiCredentialPath(dir, slot), JSON.stringify({ access_token: accessToken, refresh_token: "rt_kimi", @@ -173,6 +181,56 @@ describe("kimi usage source", () => { }); }); + it("fetches usage from the global deployment for a kimi.ai login", async () => { + writeKimiCredentials(join(homeDir, ".kimi-code"), "kimi_global_token", {}, GLOBAL_SLOT); + let requestedUrl: string | null = null; + let authorization: string | null = null; + fetchApi = (async (url: RequestInfo | URL, init?: RequestInit) => { + requestedUrl = url.toString(); + authorization = (init?.headers as Record | undefined)?.Authorization ?? null; + return jsonResponse({ usage: { limit: "100", remaining: "60" } }); + }) as unknown as typeof fetch; + + const kimi = findProvider(await service().listUsage(), "kimi"); + + expect(requestedUrl).toBe("https://api.kimi.ai/coding/v1/usages"); + expect(authorization).toBe("Bearer kimi_global_token"); + expect(kimi.status).toBe("available"); + }); + + it("prefers the login the CLI still refreshes over a stale slot", async () => { + const kimiHome = join(homeDir, ".kimi-code"); + writeKimiCredentials(kimiHome, "kimi_stale_token", { expires_at: 1_791_426_752 }); + writeKimiCredentials(kimiHome, "kimi_live_token", { expires_at: 1_793_292_337 }, GLOBAL_SLOT); + + expect(await discover()).toEqual([ + { + key: "default", + input: { + store: "file", + locator: kimiCredentialPath(kimiHome, GLOBAL_SLOT), + baseUrl: "https://api.kimi.ai/coding/v1", + }, + }, + ]); + }); + + it("asks the user to run kimi when the CLI token has expired", async () => { + writeKimiCredentials(join(homeDir, ".kimi-code"), "kimi_cli_token", { + expires_at: 1_700_000_000, + }); + + const kimi = findProvider(await service().listUsage(), "kimi"); + + expect(kimi).toEqual({ + providerId: "kimi", + status: "unavailable", + problem: { kind: "expired", expiresAt: "2023-11-14T22:13:20.000Z", refreshedBy: "kimi" }, + error: null, + planLabel: null, + }); + }); + it("reads Kimi credentials whose optional fields are null", async () => { writeKimiCredentials(join(homeDir, ".kimi-code"), "kimi_cli_token", { expires_at: null, diff --git a/plugins/kimi-usage-source/server/usage.ts b/plugins/kimi-usage-source/server/usage.ts index 556008990f..f002d032a9 100644 --- a/plugins/kimi-usage-source/server/usage.ts +++ b/plugins/kimi-usage-source/server/usage.ts @@ -1,4 +1,5 @@ import type { UsageInput } from "../shared/input.js"; +import { createHash } from "node:crypto"; import { promises as fs } from "node:fs"; import { homedir } from "node:os"; import { join } from "node:path"; @@ -18,7 +19,21 @@ const ApiOptionalStringSchema = z.preprocess( z.coerce.string().optional(), ); -const KIMI_USAGE_URL = "https://api.kimi.com/coding/v1/usages"; +const KIMI_CODE_BASE_URL = "https://api.kimi.com/coding/v1"; + +// The Kimi CLI's official deployments. A login writes its token to a credential +// slot derived from the deployment, so a global (.ai) login never touches +// kimi-code.json. Mirrors resolveKimiCodeOAuthKey in the Kimi CLI. +const KIMI_CODE_DEPLOYMENTS = [ + { oauthHost: "https://auth.kimi.com", baseUrl: KIMI_CODE_BASE_URL }, + { oauthHost: "https://auth.kimi.ai", baseUrl: "https://api.kimi.ai/coding/v1" }, +]; + +function credentialSlot(deployment: { oauthHost: string; baseUrl: string }): string { + if (deployment.baseUrl === KIMI_CODE_BASE_URL) return "kimi-code"; + const hash = createHash("sha256").update(JSON.stringify(deployment)).digest("hex"); + return `kimi-code-env-${hash.slice(0, 16)}`; +} const KimiUsageFieldsSchema = z.object({ limit: ApiOptionalStringSchema, @@ -257,7 +272,7 @@ export async function fetchUsage( fetchApi: typeof fetch = fetch, ): Promise { async function callUsageApi(token: string): Promise { - return fetchApi(KIMI_USAGE_URL, { + return fetchApi(`${input.baseUrl ?? KIMI_CODE_BASE_URL}/usages`, { signal: AbortSignal.timeout(15_000), headers: { Authorization: `Bearer ${token}`, @@ -272,6 +287,7 @@ export async function fetchUsage( return unavailable({ kind: "expired", expiresAt: new Date(credentials.expires_at * 1000).toISOString(), + refreshedBy: "kimi", }); const res = await callUsageApi(credentials.access_token); @@ -306,21 +322,29 @@ async function readCredentials(input: UsageInput): Promise { - const candidates: UsageInput[] = ["KIMI_TOKEN", "KIMI_API_KEY"].map((locator) => ({ - store: "env", - locator, - })); - candidates.push( - ...[ - join( - process.env.KIMI_CODE_HOME || join(homedir(), ".kimi-code"), - "credentials", - "kimi-code.json", - ), - join(homedir(), ".kimi", "credentials", "kimi-code.json"), - ].map((locator) => ({ store: "file" as const, locator })), - ); - for (const input of candidates) + for (const locator of ["KIMI_TOKEN", "KIMI_API_KEY"]) { + const input: UsageInput = { store: "env", locator }; if (await readCredentials(input)) return [{ key: "default", input }]; - return []; + } + const homes = [ + process.env.KIMI_CODE_HOME || join(homedir(), ".kimi-code"), + join(homedir(), ".kimi"), + ]; + // A login to another deployment leaves the old slot behind with a dead token. + // The CLI keeps refreshing the slot it uses, so the latest expiry wins. + let found: { input: UsageInput; expiresAt: number } | undefined; + for (const home of homes) { + for (const deployment of KIMI_CODE_DEPLOYMENTS) { + const input: UsageInput = { + store: "file", + locator: join(home, "credentials", `${credentialSlot(deployment)}.json`), + baseUrl: deployment.baseUrl, + }; + const credentials = await readCredentials(input); + if (!credentials) continue; + const expiresAt = credentials.expires_at ?? Number.POSITIVE_INFINITY; + if (!found || expiresAt > found.expiresAt) found = { input, expiresAt }; + } + } + return found ? [{ key: "default", input: found.input }] : []; } diff --git a/plugins/kimi-usage-source/shared/input.ts b/plugins/kimi-usage-source/shared/input.ts index 4be5cce316..edcb55302e 100644 --- a/plugins/kimi-usage-source/shared/input.ts +++ b/plugins/kimi-usage-source/shared/input.ts @@ -3,6 +3,7 @@ export const inputSchema = z .object({ store: z.enum(["env", "file"]), locator: z.string().min(1), + baseUrl: z.url().optional(), }) .strict(); export type UsageInput = z.infer;