Refresh expired Kimi usage tokens

packages/server/src/services/quota-fetcher/providers/kimi.ts

@@ -1,6 +1,6 @@
 import { existsSync, promises as fs } from "node:fs";
 import { homedir } from "node:os";
-import { join } from "node:path";
+import { dirname, join } from "node:path";
 import type { Logger } from "pino";
 import { z } from "zod";
 import type { ProviderUsage } from "../../../server/messages.js";
@@ -12,6 +12,10 @@ import {
   unavailableUsage,
 } from "../usage.js";
 
+const KIMI_CLIENT_ID = "17e5f671-d194-4dfb-9706-5516cb48c098";
+const KIMI_USAGE_URL = "https://api.kimi.com/coding/v1/usages";
+const KIMI_TOKEN_URL = "https://auth.kimi.com/api/oauth/token";
+
 const KimiUsageResponseSchema = z.object({
   usage: z
     .object({
@@ -22,10 +26,33 @@ const KimiUsageResponseSchema = z.object({
     .nullish(),
 });
 
-const KimiAuthSchema = z.object({
-  access_token: z.string().optional(),
+const KimiAuthSchema = z
+  .object({
+    access_token: z.string().optional(),
+    refresh_token: z.string().optional(),
+    expires_at: z.number().optional(),
+    expires_in: z.number().optional(),
+    scope: z.string().optional(),
+    token_type: z.string().optional(),
+  })
+  .passthrough();
+
+const KimiTokenRefreshSchema = z.object({
+  access_token: z.string(),
+  refresh_token: z.string().optional(),
+  expires_in: z.coerce.number().optional(),
+  scope: z.string().optional(),
+  token_type: z.string().optional(),
 });
 
+type KimiAuth = z.infer<typeof KimiAuthSchema>;
+type KimiTokenRefresh = z.infer<typeof KimiTokenRefreshSchema>;
+
+interface KimiCredentialRecord {
+  credentials: KimiAuth & { access_token: string };
+  filePath: string | null;
+}
+
 interface KimiQuotaProviderOptions {
   logger: Logger;
   fetch?: ProviderApiFetch;
@@ -41,23 +68,45 @@ export class KimiQuotaProvider implements ProviderUsageFetcher {
   private readonly homeDir?: string;
 
   constructor(options: KimiQuotaProviderOptions) {
-    this.logger = options.logger;
+    this.logger = options.logger.child({ module: "kimi-quota-provider" });
     this.fetchApi = options.fetch ?? fetch;
     this.homeDir = options.homeDir;
   }
 
   async fetchUsage(): Promise<ProviderUsage> {
-    const token =
-      process.env["KIMI_TOKEN"] || process.env["KIMI_API_KEY"] || (await this.readKimiToken());
+    const credentials = await this.readCredentials();
+    if (!credentials) return unavailableUsage(this);
 
-    if (!token) return unavailableUsage(this);
+    let token = credentials.credentials.access_token;
+    let res = await this.callUsageApi(token);
 
-    const res = await fetchProviderApi(this.fetchApi, "https://api.kimi.com/coding/v1/usages", {
-      headers: {
-        Authorization: `Bearer ${token}`,
-        Accept: "application/json",
-      },
-    });
+    if (res.status === 401 || res.status === 403) {
+      if (!credentials.filePath || !credentials.credentials.refresh_token) {
+        return unavailableUsage(this);
+      }
+
+      // Kimi Code may have refreshed the same credential file while this request was in
+      // flight. Prefer that newer token before making a second refresh request ourselves.
+      const latest = await this.readCredentialFile(credentials.filePath);
+      if (latest?.access_token && latest.access_token !== token) {
+        token = latest.access_token;
+        res = await this.callUsageApi(token);
+      }
+
+      if (res.status === 401 || res.status === 403) {
+        const refreshToken = latest?.refresh_token ?? credentials.credentials.refresh_token;
+        const refreshed = await this.refreshToken(refreshToken);
+        if (!refreshed) return unavailableUsage(this);
+
+        const merged = this.mergeRefreshedCredentials(
+          latest ?? credentials.credentials,
+          refreshed,
+        );
+        await this.saveCredentials(credentials.filePath, merged);
+
+        res = await this.callUsageApi(refreshed.access_token);
+      }
+    }
 
     if (!res.ok) {
       this.logger.debug({ status: res.status }, "Kimi usage fetch failed");
@@ -95,9 +144,85 @@ export class KimiQuotaProvider implements ProviderUsageFetcher {
     };
   }
 
-  private async readKimiToken(): Promise<string | null> {
+  private async callUsageApi(token: string): Promise<Response> {
+    return fetchProviderApi(this.fetchApi, KIMI_USAGE_URL, {
+      headers: {
+        Authorization: `Bearer ${token}`,
+        Accept: "application/json",
+      },
+    });
+  }
+
+  private async refreshToken(refreshToken: string): Promise<KimiTokenRefresh | null> {
+    const body = new URLSearchParams({
+      client_id: KIMI_CLIENT_ID,
+      grant_type: "refresh_token",
+      refresh_token: refreshToken,
+    });
+    const res = await fetchProviderApi(this.fetchApi, KIMI_TOKEN_URL, {
+      method: "POST",
+      headers: { "Content-Type": "application/x-www-form-urlencoded" },
+      body,
+    });
+    if (!res.ok) {
+      this.logger.debug({ status: res.status }, "Kimi token refresh failed");
+      return null;
+    }
+    const parsed = KimiTokenRefreshSchema.safeParse(await res.json());
+    return parsed.success ? parsed.data : null;
+  }
+
+  private mergeRefreshedCredentials(
+    existing: KimiAuth,
+    refreshed: KimiTokenRefresh,
+  ): KimiAuth & { access_token: string } {
+    const expiresIn = refreshed.expires_in ?? existing.expires_in;
+    return {
+      ...existing,
+      access_token: refreshed.access_token,
+      refresh_token: refreshed.refresh_token ?? existing.refresh_token,
+      expires_in: expiresIn,
+      expires_at: expiresIn === undefined ? existing.expires_at : Date.now() / 1000 + expiresIn,
+      scope: refreshed.scope ?? existing.scope,
+      token_type: refreshed.token_type ?? existing.token_type,
+    };
+  }
+
+  private async saveCredentials(filePath: string, credentials: KimiAuth): Promise<void> {
+    const tempPath = join(
+      dirname(filePath),
+      `.kimi-code.json.${process.pid}.${Date.now()}.tmp`,
+    );
+    try {
+      await fs.writeFile(tempPath, JSON.stringify(credentials), { mode: 0o600 });
+      await fs.rename(tempPath, filePath);
+    } catch (error) {
+      this.logger.debug({ err: error }, "Failed to persist refreshed Kimi credentials");
+      await fs.rm(tempPath, { force: true }).catch(() => undefined);
+    }
+  }
+
+  private async readCredentials(): Promise<KimiCredentialRecord | null> {
+    const environmentToken = process.env["KIMI_TOKEN"] || process.env["KIMI_API_KEY"];
+    if (environmentToken) {
+      return { credentials: { access_token: environmentToken }, filePath: null };
+    }
+
+    for (const path of this.credentialPaths()) {
+      const credentials = await this.readCredentialFile(path);
+      if (credentials?.access_token) {
+        return {
+          credentials: { ...credentials, access_token: credentials.access_token },
+          filePath: path,
+        };
+      }
+    }
+    return null;
+  }
+
+  private credentialPaths(): string[] {
     const homeDir = this.homeDir ?? homedir();
-    const paths = [
+    return [
       join(
         process.env["KIMI_CODE_HOME"] || join(homeDir, ".kimi-code"),
         "credentials",
@@ -105,16 +230,14 @@ export class KimiQuotaProvider implements ProviderUsageFetcher {
       ),
       join(homeDir, ".kimi", "credentials", "kimi-code.json"),
     ];
+  }
 
-    for (const path of paths) {
-      if (!existsSync(path)) continue;
-      try {
-        const credentials = KimiAuthSchema.parse(JSON.parse(await fs.readFile(path, "utf8")));
-        if (credentials.access_token) return credentials.access_token;
-      } catch {
-        continue;
-      }
+  private async readCredentialFile(path: string): Promise<KimiAuth | null> {
+    if (!existsSync(path)) return null;
+    try {
+      return KimiAuthSchema.parse(JSON.parse(await fs.readFile(path, "utf8")));
+    } catch {
+      return null;
     }
-    return null;
   }
 }