From f6ce001add30f488d213a12542731e9dfd08de3a Mon Sep 17 00:00:00 2001 From: Hunter <82844928+UnbrokenHunter@users.noreply.github.com> Date: Thu, 30 Jul 2026 23:01:56 -0400 Subject: [PATCH] Refresh expired Kimi usage tokens --- .../services/quota-fetcher/providers/kimi.ts | 171 +++++++++++++++--- 1 file changed, 147 insertions(+), 24 deletions(-) diff --git a/packages/server/src/services/quota-fetcher/providers/kimi.ts b/packages/server/src/services/quota-fetcher/providers/kimi.ts index af3e6b5602..d1416990bd 100644 --- a/packages/server/src/services/quota-fetcher/providers/kimi.ts +++ b/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; +type KimiTokenRefresh = z.infer; + +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 { - 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 { + private async callUsageApi(token: string): Promise { + return fetchProviderApi(this.fetchApi, KIMI_USAGE_URL, { + headers: { + Authorization: `Bearer ${token}`, + Accept: "application/json", + }, + }); + } + + private async refreshToken(refreshToken: string): Promise { + 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 { + 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 { + 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 { + if (!existsSync(path)) return null; + try { + return KimiAuthSchema.parse(JSON.parse(await fs.readFile(path, "utf8"))); + } catch { + return null; } - return null; } }