fix(usage): preserve the selected Kimi credential home

Compare deployment credentials within each home so a newer legacy-home
login cannot override KIMI_CODE_HOME. Cover the expired-default plus live
global workflow through discovery and usage fetching.

Co-authored-by: benthecarman <benthecarman@live.com>

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,102 @@ 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<string, string> | 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("fetches global usage when an expired default-deployment login remains", async () => {
+    const kimiHome = join(homeDir, ".kimi-code");
+    writeKimiCredentials(kimiHome, "expired_default", { expires_at: 1_700_000_000 });
+    writeKimiCredentials(kimiHome, "live_global", { expires_at: 2_000_000_000 }, GLOBAL_SLOT);
+    const [account] = await discover();
+    if (!account) throw new Error("No configured Kimi account");
+    const requests: string[] = [];
+
+    const report = await fetchUsage(
+      account.input as Parameters<typeof fetchUsage>[0],
+      async (url, init) => {
+        requests.push(url.toString());
+        expect(new Headers(init?.headers).get("Authorization")).toBe("Bearer live_global");
+        return jsonResponse({ usage: { limit: "100", remaining: "60" } });
+      },
+    );
+
+    expect(requests).toEqual(["https://api.kimi.ai/coding/v1/usages"]);
+    expect(report).toMatchObject({ status: "available", windows: [{ usedPct: 40 }] });
+  });
+
+  it("keeps KIMI_CODE_HOME ahead of a newer legacy-home login", async () => {
+    const selectedHome = join(homeDir, "selected-kimi");
+    process.env.KIMI_CODE_HOME = selectedHome;
+    writeKimiCredentials(
+      selectedHome,
+      "selected_global_token",
+      { expires_at: 2_000_000_000 },
+      GLOBAL_SLOT,
+    );
+    writeKimiCredentials(join(homeDir, ".kimi"), "legacy_token", { expires_at: 2_000_000_100 });
+
+    const accounts = await discover();
+
+    expect(accounts).toEqual([
+      {
+        key: "default",
+        input: {
+          store: "file",
+          locator: kimiCredentialPath(selectedHome, 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,

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<UsageReport> {
   async function callUsageApi(token: string): Promise<Response> {
-    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,31 @@ async function readCredentials(input: UsageInput): Promise<KimiCredentials | und
   }
 }
 export async function discover(): Promise<UsageAccount[]> {
-  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 }];
+  }
+  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.
+  // Within each home, prefer the latest expiry. Keep the configured home ahead
+  // of the legacy fallback, even if that fallback holds a newer token.
+  for (const home of homes) {
+    let found: { input: UsageInput; expiresAt: number } | undefined;
+    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 };
+    }
+    if (found) return [{ key: "default", input: found.input }];
+  }
   return [];
 }

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<typeof inputSchema>;