diff --git a/config/packaged-runtime-node-modules.cjs b/config/packaged-runtime-node-modules.cjs index d3f81817cee..6adf10b77e5 100644 --- a/config/packaged-runtime-node-modules.cjs +++ b/config/packaged-runtime-node-modules.cjs @@ -22,6 +22,7 @@ const PACKAGED_RUNTIME_PACKAGE_ROOTS = [ 'jsonc-parser', 'node-pty', 'posthog-node', + 'proper-lockfile', // serve-sim (for CLI JS entry + closure + state/middleware + to make packaged require('serve-sim') + its internal relatives work; mirrors other runtime JS like ws/yaml/zod. Natives/dylibs still via extraResources + the node_modules/serve-sim copy in resources from builder. Client if added too. 'serve-sim', 'qrcode', diff --git a/config/scripts/electron-builder-config.test.mjs b/config/scripts/electron-builder-config.test.mjs index faf29fa244f..95a1d6cc53e 100644 --- a/config/scripts/electron-builder-config.test.mjs +++ b/config/scripts/electron-builder-config.test.mjs @@ -249,6 +249,14 @@ describe('electron-builder config', () => { ).toBe(true) }) + it('includes proper-lockfile in the packaged runtime closure', () => { + const packagedTargets = createPackagedRuntimeNodeModuleResources().map( + (resource) => resource.to + ) + + expect(packagedTargets).toContain(join('node_modules', 'proper-lockfile')) + }) + it('prunes non-target @parcel/watcher platform subpackages from packaged runtime resources', async () => { const resourcesDir = await mkdtemp(join(tmpdir(), 'orca-parcel-watcher-prune-')) try { diff --git a/package.json b/package.json index a062d03ca1f..f088d212d01 100644 --- a/package.json +++ b/package.json @@ -112,6 +112,7 @@ "jsonc-parser": "^3.3.1", "node-pty": "^1.1.0", "posthog-node": "^5.33.3", + "proper-lockfile": "4.1.2", "qrcode": "^1.5.4", "react-i18next": "^17.0.8", "serve-sim": "^0.1.40", diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index e12b10d7464..84c1f191ad1 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -64,6 +64,9 @@ importers: posthog-node: specifier: ^5.33.3 version: 5.33.3 + proper-lockfile: + specifier: 4.1.2 + version: 4.1.2 qrcode: specifier: ^1.5.4 version: 1.5.4 diff --git a/src/main/rate-limits/kimi-credential-location.ts b/src/main/rate-limits/kimi-credential-location.ts new file mode 100644 index 00000000000..722e0d8a925 --- /dev/null +++ b/src/main/rate-limits/kimi-credential-location.ts @@ -0,0 +1,51 @@ +import { createHash } from 'node:crypto' +import { homedir } from 'node:os' +import { join } from 'node:path' + +const DEFAULT_OAUTH_HOST = 'https://auth.kimi.com' +const DEFAULT_BASE_URL = 'https://api.kimi.com/coding/v1' + +export type KimiCredentialLocation = { + home: string + oauthHost: string + baseUrl: string + storageName: string + credentialsPath: string + lockTarget: string + tokenUrl: string + usageUrl: string +} + +type KimiEnvironment = Partial> + +export function resolveKimiCredentialLocation( + env: KimiEnvironment = process.env, + fallbackHome = homedir() +): KimiCredentialLocation { + const home = env.KIMI_CODE_HOME || join(fallbackHome, '.kimi-code') + const oauthHost = (env.KIMI_CODE_OAUTH_HOST ?? env.KIMI_OAUTH_HOST ?? DEFAULT_OAUTH_HOST) + .trim() + .replace(/\/+$/, '') + // Why: match oauthHost trimming so whitespace in KIMI_CODE_BASE_URL cannot + // skew storageName hashing or produce a malformed usageUrl. + const baseUrl = (env.KIMI_CODE_BASE_URL ?? DEFAULT_BASE_URL).trim().replace(/\/+$/, '') + const isDefault = oauthHost === DEFAULT_OAUTH_HOST && baseUrl === DEFAULT_BASE_URL + const suffix = isDefault + ? '' + : `-env-${createHash('sha256') + .update(JSON.stringify({ oauthHost, baseUrl })) + .digest('hex') + .slice(0, 16)}` + const storageName = `kimi-code${suffix}` + + return { + home, + oauthHost, + baseUrl, + storageName, + credentialsPath: join(home, 'credentials', `${storageName}.json`), + lockTarget: join(home, 'oauth', storageName), + tokenUrl: `${oauthHost}/api/oauth/token`, + usageUrl: `${baseUrl}/usages` + } +} diff --git a/src/main/rate-limits/kimi-fetcher.test.ts b/src/main/rate-limits/kimi-fetcher.test.ts index a6274097f5f..f4d0c66d52e 100644 --- a/src/main/rate-limits/kimi-fetcher.test.ts +++ b/src/main/rate-limits/kimi-fetcher.test.ts @@ -1,9 +1,21 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { join } from 'node:path' const netFetchMock = vi.hoisted(() => vi.fn()) -const fsState = vi.hoisted<{ credentials: string | null; readError: Error | null }>(() => ({ +const fsState = vi.hoisted<{ + credentials: string | null + readError: Error | null + writes: Map + writeOptions: Map + renames: { from: string; to: string }[] + readPaths: string[] +}>(() => ({ credentials: null, - readError: null + readError: null, + writes: new Map(), + writeOptions: new Map(), + renames: [], + readPaths: [] })) vi.mock('electron', () => ({ @@ -12,7 +24,8 @@ vi.mock('electron', () => ({ vi.mock('node:fs', () => ({ existsSync: () => fsState.credentials !== null, - readFileSync: () => { + readFileSync: (path: string) => { + fsState.readPaths.push(path) if (fsState.readError) { throw fsState.readError } @@ -21,13 +34,35 @@ vi.mock('node:fs', () => ({ } return fsState.credentials }, - writeFileSync: () => {}, - renameSync: () => {} + writeFileSync: (path: string, value: string, options?: unknown) => { + fsState.writes.set(path, value) + fsState.writeOptions.set(path, options) + }, + renameSync: (from: string, to: string) => { + fsState.renames.push({ from, to }) + const value = fsState.writes.get(from) + if (value !== undefined) { + fsState.credentials = value + fsState.writes.delete(from) + } + } })) vi.mock('node:os', () => ({ homedir: () => '/home/test' })) +vi.mock('../../shared/secure-file', () => ({ + writeSecureJsonFile: (path: string, value: unknown) => { + const tmpPath = `${path}.test.tmp` + fsState.writes.set(tmpPath, `${JSON.stringify(value, null, 2)}\n`) + fsState.writeOptions.set(tmpPath, { encoding: 'utf-8', mode: 0o600 }) + fsState.renames.push({ from: tmpPath, to: path }) + fsState.credentials = fsState.writes.get(tmpPath) ?? null + fsState.writes.delete(tmpPath) + } +})) + import { fetchKimiRateLimits } from './kimi-fetcher' +import { resolveKimiCredentialLocation } from './kimi-credential-location' function jsonResponse(body: unknown, status = 200): Response { return { @@ -57,9 +92,14 @@ function freshCredentials(): string { describe('fetchKimiRateLimits', () => { beforeEach(() => { + vi.stubEnv('KIMI_DISABLE_OAUTH_LOCK', '1') netFetchMock.mockReset() fsState.credentials = null fsState.readError = null + fsState.writes.clear() + fsState.writeOptions.clear() + fsState.renames = [] + fsState.readPaths = [] }) afterEach(() => { @@ -133,8 +173,96 @@ describe('fetchKimiRateLimits', () => { expect(result.weekly).toBeNull() }) - it('does NOT refresh or call the API when the token is expired (read-only)', async () => { - // expires_at in the past → token stale; fetcher must not hit the network. + it('refreshes an expired access token before fetching usage', async () => { + fsState.credentials = JSON.stringify({ + access_token: 'tok-old', + refresh_token: 'refresh-old', + expires_at: 1, + token_type: 'Bearer', + scope: 'kimi-code' + }) + netFetchMock + .mockResolvedValueOnce( + jsonResponse({ + access_token: 'tok-new', + refresh_token: 'refresh-new', + expires_in: 900, + token_type: 'Bearer', + scope: 'kimi-code' + }) + ) + .mockResolvedValueOnce(jsonResponse(USAGE_RESPONSE)) + + const result = await fetchKimiRateLimits() + + expect(result.status).toBe('ok') + expect(netFetchMock).toHaveBeenCalledTimes(2) + const [refreshUrl, refreshInit] = netFetchMock.mock.calls[0] + expect(refreshUrl).toBe('https://auth.kimi.com/api/oauth/token') + const refreshBody = new URLSearchParams(refreshInit.body as string) + expect(refreshBody.get('client_id')).toBe('17e5f671-d194-4dfb-9706-5516cb48c098') + expect(refreshBody.get('grant_type')).toBe('refresh_token') + expect(refreshBody.get('refresh_token')).toBe('refresh-old') + const [, usageInit] = netFetchMock.mock.calls[1] + expect((usageInit.headers as Record).Authorization).toBe('Bearer tok-new') + expect(fsState.renames).toHaveLength(1) + const [{ from: tmpPath }] = fsState.renames + expect(fsState.writeOptions.get(tmpPath)).toEqual({ encoding: 'utf-8', mode: 0o600 }) + expect(JSON.parse(fsState.credentials ?? '{}')).toMatchObject({ + access_token: 'tok-new', + refresh_token: 'refresh-new', + expires_in: 900, + token_type: 'Bearer', + scope: 'kimi-code' + }) + expect(JSON.parse(fsState.credentials ?? '{}').expires_at).toBeGreaterThan(1) + }) + + it('uses credentials refreshed by the Kimi CLI when its refresh wins the race', async () => { + fsState.credentials = JSON.stringify({ + access_token: 'tok-old', + refresh_token: 'refresh-old', + expires_at: 1 + }) + netFetchMock + .mockImplementationOnce(async () => { + fsState.credentials = JSON.stringify({ + access_token: 'tok-cli', + refresh_token: 'refresh-cli', + expires_at: 99_999_999_999 + }) + return jsonResponse({ error: 'invalid_grant' }, 400) + }) + .mockResolvedValueOnce(jsonResponse(USAGE_RESPONSE)) + + const result = await fetchKimiRateLimits() + + expect(result.status).toBe('ok') + expect(netFetchMock).toHaveBeenCalledTimes(2) + const [, usageInit] = netFetchMock.mock.calls[1] + expect((usageInit.headers as Record).Authorization).toBe('Bearer tok-cli') + expect(fsState.renames).toHaveLength(0) + }) + + it('returns user-safe unauthorized refresh copy after persisting revocation', async () => { + fsState.credentials = JSON.stringify({ + access_token: 'tok-old', + refresh_token: 'refresh-old', + expires_at: 1 + }) + netFetchMock.mockResolvedValueOnce(jsonResponse({ error: 'invalid_grant' }, 401)) + + const result = await fetchKimiRateLimits() + + expect(result.status).toBe('error') + expect(result.error).toBe('Kimi token refresh unauthorized - run kimi login') + expect(JSON.parse(fsState.credentials ?? '{}')).toMatchObject({ + access_token: '', + refresh_token: '' + }) + }) + + it('does NOT refresh or call the API when the expired token has no refresh token', async () => { fsState.credentials = JSON.stringify({ access_token: 'tok-old', expires_at: 1 }) const result = await fetchKimiRateLimits() @@ -142,4 +270,88 @@ describe('fetchKimiRateLimits', () => { expect(result.error).toMatch(/expired/i) expect(netFetchMock).not.toHaveBeenCalled() }) + + it('uses only the endpoint-scoped credential for refresh and usage', async () => { + vi.stubEnv('KIMI_CODE_HOME', '/kimi-home') + vi.stubEnv('KIMI_CODE_OAUTH_HOST', 'https://auth.example.com') + vi.stubEnv('KIMI_CODE_BASE_URL', 'https://api.example.com/coding/v1') + fsState.credentials = JSON.stringify({ + access_token: 'scoped-old', + refresh_token: 'scoped-refresh', + expires_at: 1 + }) + netFetchMock + .mockResolvedValueOnce( + jsonResponse({ + access_token: 'scoped-new', + refresh_token: 'scoped-rotated', + expires_in: 900 + }) + ) + .mockResolvedValueOnce(jsonResponse(USAGE_RESPONSE)) + + const result = await fetchKimiRateLimits() + + const scopedPath = join('/kimi-home', 'credentials', 'kimi-code-env-3e58296a497444bf.json') + expect(result.status).toBe('ok') + expect(fsState.readPaths).toEqual([scopedPath, scopedPath]) + expect(fsState.readPaths).not.toContain('/kimi-home/credentials/kimi-code.json') + expect(netFetchMock.mock.calls[0][0]).toBe('https://auth.example.com/api/oauth/token') + expect(netFetchMock.mock.calls[1][0]).toBe('https://api.example.com/coding/v1/usages') + expect(fsState.renames[0].to).toBe(scopedPath) + }) +}) + +describe('resolveKimiCredentialLocation', () => { + it.each([ + ['https://auth.kimi.com', 'https://api.kimi.com/coding/v1', 'kimi-code.json'], + [ + 'https://auth.example.com', + 'https://api.kimi.com/coding/v1', + 'kimi-code-env-d5c8af65e27873b9.json' + ], + [ + 'https://auth.kimi.com', + 'https://api.example.com/coding/v1', + 'kimi-code-env-bf2242cfc5d51611.json' + ], + [ + 'https://auth.example.com', + 'https://api.example.com/coding/v1', + 'kimi-code-env-3e58296a497444bf.json' + ] + ])('matches the Kimi credential vector for %s and %s', (oauthHost, baseUrl, filename) => { + const location = resolveKimiCredentialLocation( + { + KIMI_CODE_HOME: '/kimi-home', + KIMI_CODE_OAUTH_HOST: oauthHost, + KIMI_CODE_BASE_URL: baseUrl + }, + '/fallback-home' + ) + + expect(location.credentialsPath).toBe(join('/kimi-home', 'credentials', filename)) + }) + + it('scopes credentials, lock, and request URLs to the normalized endpoint pair', () => { + const location = resolveKimiCredentialLocation( + { + KIMI_CODE_HOME: '/kimi-home', + KIMI_CODE_OAUTH_HOST: ' https://auth.example.com/// ', + KIMI_CODE_BASE_URL: ' https://api.example.com/coding/v1/// ' + }, + '/fallback-home' + ) + + expect(location).toEqual({ + home: '/kimi-home', + oauthHost: 'https://auth.example.com', + baseUrl: 'https://api.example.com/coding/v1', + storageName: 'kimi-code-env-3e58296a497444bf', + credentialsPath: join('/kimi-home', 'credentials', 'kimi-code-env-3e58296a497444bf.json'), + lockTarget: join('/kimi-home', 'oauth', 'kimi-code-env-3e58296a497444bf'), + tokenUrl: 'https://auth.example.com/api/oauth/token', + usageUrl: 'https://api.example.com/coding/v1/usages' + }) + }) }) diff --git a/src/main/rate-limits/kimi-fetcher.ts b/src/main/rate-limits/kimi-fetcher.ts index f957e171d88..2e48f1a33d1 100644 --- a/src/main/rate-limits/kimi-fetcher.ts +++ b/src/main/rate-limits/kimi-fetcher.ts @@ -1,34 +1,22 @@ import { existsSync, readFileSync } from 'node:fs' -import { homedir } from 'node:os' -import { join } from 'node:path' import { net } from 'electron' import type { ProviderRateLimits, RateLimitWindow } from '../../shared/rate-limit-types' +import { resolveKimiCredentialLocation } from './kimi-credential-location' +import { + KimiRefreshError, + refreshKimiCredentials, + type KimiCredentials +} from './kimi-oauth-refresh' // Why: Kimi Code's managed coding plan exposes subscription usage at // `${base}/usages` (see packages/oauth/src/managed-usage.ts in the CLI bundle). // The base URL is overridable via the same env var the CLI honours so Orca // stays aligned with a user's self-hosted/staging config. -const KIMI_BASE_URL = process.env.KIMI_CODE_BASE_URL ?? 'https://api.kimi.com/coding/v1' const API_TIMEOUT_MS = 10_000 const SESSION_WINDOW_MINUTES = 300 // 5h const WEEKLY_WINDOW_MINUTES = 10080 // 7d -function getKimiHome(): string { - // Why: match the CLI's `KIMI_CODE_HOME ?? ~/.kimi-code` resolution so we read - // the same OAuth credentials the running Kimi CLI writes. - return process.env.KIMI_CODE_HOME ?? join(homedir(), '.kimi-code') -} - -function getCredentialsPath(): string { - return join(getKimiHome(), 'credentials', 'kimi-code.json') -} - -type KimiCredentials = { - access_token?: string - expires_at?: number -} - type CredentialsReadResult = | { status: 'missing' } | { status: 'error'; error: string } @@ -38,18 +26,29 @@ function parseCredentials(value: unknown): KimiCredentials | null { if (typeof value !== 'object' || value === null) { return null } - const credentials: KimiCredentials = {} - if ('access_token' in value && typeof value.access_token === 'string') { - credentials.access_token = value.access_token + const credentials: KimiCredentials = { ...(value as Record) } + if (typeof credentials.access_token !== 'string') { + delete credentials.access_token + } + if (typeof credentials.refresh_token !== 'string') { + delete credentials.refresh_token } - if ('expires_at' in value && typeof value.expires_at === 'number') { - credentials.expires_at = value.expires_at + if (typeof credentials.expires_at !== 'number') { + delete credentials.expires_at + } + if (typeof credentials.expires_in !== 'number') { + delete credentials.expires_in + } + if (typeof credentials.scope !== 'string') { + delete credentials.scope + } + if (typeof credentials.token_type !== 'string') { + delete credentials.token_type } return credentials } -function readCredentials(): CredentialsReadResult { - const path = getCredentialsPath() +function readCredentials(path: string): CredentialsReadResult { if (!existsSync(path)) { return { status: 'missing' } } @@ -213,18 +212,19 @@ function result(status: ProviderRateLimits['status'], error: string | null): Pro } /** - * Read-only subscription usage for Kimi Code. + * Subscription usage for Kimi Code. * - * Why read-only: the access token lives in `~/.kimi-code/credentials/kimi-code.json` - * and is refreshed by the Kimi CLI itself (15-min TTL, refresh-token rotation). - * Orca must NEVER refresh or rewrite that file — a rotated refresh token would - * log out a live `kimi` session. We only read the current token and call the - * same `GET /usages` endpoint, with the same headers, that the CLI's own - * `/usage` command uses. The completion endpoint (the one Moonshot gates to - * approved coding agents) is never touched here. + * Why refresh here: Kimi Code 0.9 stores a 15-minute access token plus a + * refresh token in `~/.kimi-code/credentials/kimi-code.json`. If the CLI is not + * running, nobody rotates that access token, so background usage polling goes + * stale. Orca only refreshes this usage token when the cached access token has + * already expired, persists any rotated refresh token atomically, then calls + * the CLI's same `GET /usages` endpoint. The completion endpoint (the one + * Moonshot gates to approved coding agents) is never touched here. */ export async function fetchKimiRateLimits(): Promise { - const readResult = readCredentials() + const location = resolveKimiCredentialLocation() + const readResult = readCredentials(location.credentialsPath) if (readResult.status === 'missing') { return result('unavailable', 'Not signed in to Kimi Code') } @@ -236,14 +236,28 @@ export async function fetchKimiRateLimits(): Promise { return result('error', 'Kimi credentials file is missing an access token') } if (!isAccessTokenFresh(creds)) { - // Why: don't refresh — the CLI owns the token lifecycle. Report a transient - // error so the rate-limit service keeps the last good snapshot (stale - // policy) until the user next runs Kimi and the CLI refreshes the file. - return result('error', 'Kimi token expired — open Kimi to refresh') + if (typeof creds.refresh_token !== 'string' || creds.refresh_token.length === 0) { + return result('error', 'Kimi token expired — open Kimi to refresh') + } + let refreshed: KimiCredentials | null + try { + refreshed = await refreshKimiCredentials(creds, location) + } catch (error) { + return result( + 'error', + error instanceof KimiRefreshError && error.kind === 'unauthorized' + ? 'Kimi token refresh unauthorized - run kimi login' + : 'Kimi token refresh failed - run kimi login' + ) + } + if (!refreshed) { + return result('error', 'Kimi token refresh failed - run kimi login') + } + creds.access_token = refreshed.access_token } try { - const res = await net.fetch(`${KIMI_BASE_URL.replace(/\/$/, '')}/usages`, { + const res = await net.fetch(location.usageUrl, { // Why: identical to the CLI's fetchManagedUsage — bearer token + Accept. // No extra User-Agent: the usages endpoint authenticates by token only. headers: { Authorization: `Bearer ${creds.access_token}`, Accept: 'application/json' }, diff --git a/src/main/rate-limits/kimi-oauth-refresh.test.ts b/src/main/rate-limits/kimi-oauth-refresh.test.ts new file mode 100644 index 00000000000..5c2807ad0c2 --- /dev/null +++ b/src/main/rate-limits/kimi-oauth-refresh.test.ts @@ -0,0 +1,254 @@ +import { mkdtempSync, rmSync } from 'node:fs' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import type { lock as ProperLock } from 'proper-lockfile' +import { afterEach, describe, expect, it, vi } from 'vitest' + +const lockMock = vi.hoisted(() => vi.fn()) + +vi.mock('proper-lockfile', () => ({ lock: lockMock })) +import type { KimiCredentialLocation } from './kimi-credential-location' +import { + acquireKimiRefreshLock, + refreshKimiCredentials, + type KimiCredentials, + type KimiRefreshError, + type KimiRefreshDependencies +} from './kimi-oauth-refresh' + +const location: KimiCredentialLocation = { + home: '/kimi-home', + oauthHost: 'https://auth.example.com', + baseUrl: 'https://api.example.com/coding/v1', + storageName: 'kimi-code-env-test', + credentialsPath: '/kimi-home/credentials/kimi-code-env-test.json', + lockTarget: '/kimi-home/oauth/kimi-code-env-test', + tokenUrl: 'https://auth.example.com/api/oauth/token', + usageUrl: 'https://api.example.com/coding/v1/usages' +} + +function expired(refreshToken: string): KimiCredentials { + return { access_token: 'expired-access', refresh_token: refreshToken, expires_at: 1 } +} + +describe('refreshKimiCredentials', () => { + afterEach(() => { + vi.restoreAllMocks() + vi.unstubAllEnvs() + lockMock.mockReset() + }) + + it.skipIf(process.platform === 'win32')( + 'uses Kimi proper-lockfile options and releases after success', + async () => { + const root = mkdtempSync(join(tmpdir(), 'orca-kimi-lock-')) + const target = join(root, 'oauth', 'kimi-code') + const release = vi.fn().mockResolvedValue(undefined) + lockMock.mockResolvedValue(release) + + await expect(acquireKimiRefreshLock(target, async () => 'done')).resolves.toBe('done') + + expect(lockMock).toHaveBeenCalledWith(target, { + retries: { retries: 120, factor: 1, minTimeout: 500, maxTimeout: 1000 }, + stale: 5000, + realpath: false + }) + expect(release).toHaveBeenCalledOnce() + rmSync(root, { recursive: true, force: true }) + } + ) + + it.skipIf(process.platform === 'win32')( + 'serializes overlapping actions through the production proper-lockfile path', + async () => { + const { lock: actualLock } = await vi.importActual<{ lock: typeof ProperLock }>( + 'proper-lockfile' + ) + lockMock.mockImplementation(actualLock) + const root = mkdtempSync(join(tmpdir(), 'orca-kimi-lock-')) + const target = join(root, 'oauth', 'kimi-code') + const events: string[] = [] + let markFirstStarted!: () => void + const firstStarted = new Promise((resolve) => { + markFirstStarted = resolve + }) + + const first = acquireKimiRefreshLock(target, async () => { + events.push('first:start') + markFirstStarted() + await new Promise((resolve) => setTimeout(resolve, 20)) + events.push('first:end') + }) + await firstStarted + const second = acquireKimiRefreshLock(target, async () => { + events.push('second:start') + events.push('second:end') + }) + await Promise.all([first, second]) + + expect(events).toEqual(['first:start', 'first:end', 'second:start', 'second:end']) + rmSync(root, { recursive: true, force: true }) + } + ) + + it.skipIf(process.platform === 'win32')( + 'fails closed when production lock acquisition fails', + async () => { + const action = vi.fn() + lockMock.mockRejectedValue(new Error('lock unavailable')) + + await expect(acquireKimiRefreshLock('/tmp/kimi-lock-target', action)).rejects.toThrow( + 'lock unavailable' + ) + expect(action).not.toHaveBeenCalled() + } + ) + + it.skipIf(process.platform === 'win32')( + 'releases the production lock after the protected action throws', + async () => { + const root = mkdtempSync(join(tmpdir(), 'orca-kimi-lock-')) + const target = join(root, 'oauth', 'kimi-code') + const release = vi.fn().mockResolvedValue(undefined) + lockMock.mockResolvedValue(release) + + await expect( + acquireKimiRefreshLock(target, async () => { + throw new Error('refresh failed') + }) + ).rejects.toThrow('refresh failed') + expect(release).toHaveBeenCalledOnce() + rmSync(root, { recursive: true, force: true }) + } + ) + + it.each([ + ['the disable flag', () => vi.stubEnv('KIMI_DISABLE_OAUTH_LOCK', '1')], + ['Windows', () => vi.spyOn(process, 'platform', 'get').mockReturnValue('win32')] + ])('does not acquire a filesystem lock on %s', async (_name, configure) => { + configure() + + await expect(acquireKimiRefreshLock('/unwritable/kimi', async () => 'done')).resolves.toBe( + 'done' + ) + expect(lockMock).not.toHaveBeenCalled() + }) + + it('serializes overlapping refreshes and rereads the rotated winner after locking', async () => { + let stored = expired('initial-refresh') + let lockTail = Promise.resolve() + const fetchToken = vi.fn(async (_url: string, init: RequestInit) => { + const body = new URLSearchParams(init.body as string) + expect(body.get('refresh_token')).toBe('initial-refresh') + await Promise.resolve() + return { + ok: true, + status: 200, + json: async () => ({ + access_token: 'winner-access', + refresh_token: 'winner-refresh', + expires_in: 900 + }) + } as Response + }) + const dependencies: KimiRefreshDependencies = { + acquireLock: async (_target, action) => { + const previous = lockTail + let release!: () => void + lockTail = new Promise((resolve) => { + release = resolve + }) + await previous + try { + return await action() + } finally { + release() + } + }, + readCredentials: () => stored, + saveCredentials: (_path, credentials) => { + stored = credentials + }, + fetchToken, + nowSeconds: () => 100 + } + + const [first, second] = await Promise.all([ + refreshKimiCredentials(expired('initial-refresh'), location, dependencies), + refreshKimiCredentials(expired('initial-refresh'), location, dependencies) + ]) + + expect(first?.access_token).toBe('winner-access') + expect(second?.access_token).toBe('winner-access') + expect(fetchToken).toHaveBeenCalledTimes(1) + expect(stored.refresh_token).toBe('winner-refresh') + }) + + it('accepts only a changed-token fresh winner after unauthorized refresh', async () => { + const initial = expired('initial-refresh') + const winner = { + access_token: 'winner-access', + refresh_token: 'winner-refresh', + expires_at: 1000 + } + const reads = [initial, winner] + const saveCredentials = vi.fn() + const delay = vi.fn().mockResolvedValue(undefined) + const dependencies: KimiRefreshDependencies = { + acquireLock: async (_target, action) => action(), + readCredentials: () => reads.shift() ?? winner, + saveCredentials, + fetchToken: vi.fn().mockResolvedValue({ ok: false, status: 401 } as Response), + nowSeconds: () => 100, + delay + } + + await expect(refreshKimiCredentials(initial, location, dependencies)).resolves.toBe(winner) + expect(delay).toHaveBeenCalledWith(100) + expect(saveCredentials).not.toHaveBeenCalled() + }) + + it('persists a revoked tombstone and rethrows unauthorized refresh without a winner', async () => { + const initial = expired('initial-refresh') + const saveCredentials = vi.fn() + const dependencies: KimiRefreshDependencies = { + acquireLock: async (_target, action) => action(), + readCredentials: () => initial, + saveCredentials, + fetchToken: vi.fn().mockResolvedValue({ ok: false, status: 403 } as Response), + nowSeconds: () => 100, + delay: vi.fn().mockResolvedValue(undefined) + } + + await expect(refreshKimiCredentials(initial, location, dependencies)).rejects.toMatchObject({ + name: 'KimiRefreshError', + kind: 'unauthorized', + status: 403 + } satisfies Partial) + expect(saveCredentials).toHaveBeenCalledWith(location.credentialsPath, { + ...initial, + access_token: '', + refresh_token: '' + }) + }) + + it('rethrows non-auth refresh failures without overwriting credentials', async () => { + const initial = expired('initial-refresh') + const saveCredentials = vi.fn() + const dependencies: KimiRefreshDependencies = { + acquireLock: async (_target, action) => action(), + readCredentials: () => initial, + saveCredentials, + fetchToken: vi.fn().mockResolvedValue({ ok: false, status: 500 } as Response), + nowSeconds: () => 100, + delay: vi.fn().mockResolvedValue(undefined) + } + + await expect(refreshKimiCredentials(initial, location, dependencies)).rejects.toMatchObject({ + name: 'KimiRefreshError', + kind: 'request', + status: 500 + } satisfies Partial) + expect(saveCredentials).not.toHaveBeenCalled() + }) +}) diff --git a/src/main/rate-limits/kimi-oauth-refresh.ts b/src/main/rate-limits/kimi-oauth-refresh.ts new file mode 100644 index 00000000000..d00b3728ba8 --- /dev/null +++ b/src/main/rate-limits/kimi-oauth-refresh.ts @@ -0,0 +1,223 @@ +import { appendFileSync, existsSync, mkdirSync, readFileSync } from 'node:fs' +import { dirname } from 'node:path' +import { net } from 'electron' +import { lock } from 'proper-lockfile' +import { writeSecureJsonFile } from '../../shared/secure-file' +import type { KimiCredentialLocation } from './kimi-credential-location' + +const KIMI_OAUTH_CLIENT_ID = '17e5f671-d194-4dfb-9706-5516cb48c098' +const API_TIMEOUT_MS = 10_000 + +export type KimiCredentials = { + access_token?: string + refresh_token?: string + expires_at?: number + expires_in?: number + scope?: string + token_type?: string + [key: string]: unknown +} + +type KimiTokenEndpointResponse = { + access_token?: unknown + refresh_token?: unknown + expires_in?: unknown + scope?: unknown + token_type?: unknown +} + +export type KimiRefreshDependencies = { + acquireLock: (target: string, action: () => Promise) => Promise + readCredentials: (path: string) => KimiCredentials | null + saveCredentials: (path: string, credentials: KimiCredentials) => void + fetchToken: (url: string, init: RequestInit) => Promise + nowSeconds: () => number + delay?: (milliseconds: number) => Promise +} + +export class KimiRefreshError extends Error { + readonly kind: 'unauthorized' | 'request' + readonly status: number + + constructor(kind: 'unauthorized' | 'request', status: number) { + super(kind === 'unauthorized' ? 'Kimi refresh unauthorized' : 'Kimi refresh request failed') + this.name = 'KimiRefreshError' + this.kind = kind + this.status = status + } +} + +function isFresh(credentials: KimiCredentials, nowSeconds: number): boolean { + return ( + typeof credentials.access_token === 'string' && + credentials.access_token.length > 0 && + typeof credentials.expires_at === 'number' && + credentials.expires_at - nowSeconds > 5 + ) +} + +function readCredentials(path: string): KimiCredentials | null { + try { + const parsed: unknown = JSON.parse(readFileSync(path, 'utf-8')) + return typeof parsed === 'object' && parsed !== null ? (parsed as KimiCredentials) : null + } catch { + return null + } +} + +function applyRefreshedCredentials( + credentials: KimiCredentials, + response: KimiTokenEndpointResponse, + nowSeconds: number +): KimiCredentials | null { + if (typeof response.access_token !== 'string' || response.access_token.length === 0) { + return null + } + if ( + typeof response.expires_in !== 'number' || + !Number.isFinite(response.expires_in) || + response.expires_in <= 5 + ) { + return null + } + const refreshToken = + typeof response.refresh_token === 'string' && response.refresh_token.length > 0 + ? response.refresh_token + : credentials.refresh_token + if (typeof refreshToken !== 'string' || refreshToken.length === 0) { + return null + } + return { + ...credentials, + access_token: response.access_token, + refresh_token: refreshToken, + expires_at: nowSeconds + response.expires_in, + expires_in: response.expires_in, + token_type: + typeof response.token_type === 'string' && response.token_type.length > 0 + ? response.token_type + : credentials.token_type, + scope: + typeof response.scope === 'string' && response.scope.length > 0 + ? response.scope + : credentials.scope + } +} + +function sleep(milliseconds: number): Promise { + return new Promise((resolve) => setTimeout(resolve, milliseconds)) +} + +async function isUnauthorizedRefreshResponse(response: Response): Promise { + if (response.status === 401 || response.status === 403) { + return true + } + try { + const body: unknown = await response.json() + return ( + typeof body === 'object' && body !== null && 'error' in body && body.error === 'invalid_grant' + ) + } catch { + return false + } +} + +export async function acquireKimiRefreshLock( + target: string, + action: () => Promise +): Promise { + if (process.platform === 'win32' || process.env.KIMI_DISABLE_OAUTH_LOCK === '1') { + return action() + } + mkdirSync(dirname(target), { recursive: true, mode: 0o700 }) + if (!existsSync(target)) { + appendFileSync(target, '', { mode: 0o600 }) + } + // Why: these are Kimi 0.23.3's exact lock settings, so Orca coordinates with + // the CLI's rotating refresh-token critical section across processes. + const release = await lock(target, { + retries: { retries: 120, factor: 1, minTimeout: 500, maxTimeout: 1000 }, + stale: 5000, + realpath: false + }) + try { + return await action() + } finally { + try { + await release() + } catch { + // Kimi ignores release errors as well. + } + } +} + +const defaultDependencies: KimiRefreshDependencies = { + acquireLock: acquireKimiRefreshLock, + readCredentials, + saveCredentials: writeSecureJsonFile, + fetchToken: (url, init) => net.fetch(url, init), + nowSeconds: () => Math.floor(Date.now() / 1000) +} + +async function refreshUnderLock( + initial: KimiCredentials, + location: KimiCredentialLocation, + dependencies: KimiRefreshDependencies +): Promise { + const latest = dependencies.readCredentials(location.credentialsPath) ?? initial + if (isFresh(latest, dependencies.nowSeconds())) { + return latest + } + if (typeof latest.refresh_token !== 'string' || latest.refresh_token.length === 0) { + return null + } + const response = await dependencies.fetchToken(location.tokenUrl, { + method: 'POST', + headers: { 'Content-Type': 'application/x-www-form-urlencoded', Accept: 'application/json' }, + body: new URLSearchParams({ + client_id: KIMI_OAUTH_CLIENT_ID, + grant_type: 'refresh_token', + refresh_token: latest.refresh_token + }).toString(), + signal: AbortSignal.timeout(API_TIMEOUT_MS) + }) + if (!response.ok) { + const unauthorized = await isUnauthorizedRefreshResponse(response) + if (!unauthorized) { + throw new KimiRefreshError('request', response.status) + } + await (dependencies.delay ?? sleep)(100) + const winner = dependencies.readCredentials(location.credentialsPath) + if ( + winner && + winner.refresh_token !== latest.refresh_token && + isFresh(winner, dependencies.nowSeconds()) + ) { + return winner + } + // Why: Kimi revokes the stored token after a rejected refresh unless a + // concurrently rotated fresh winner proves another process succeeded. + dependencies.saveCredentials(location.credentialsPath, { + ...latest, + access_token: '', + refresh_token: '' + }) + throw new KimiRefreshError('unauthorized', response.status) + } + const data = (await response.json()) as KimiTokenEndpointResponse + const refreshed = applyRefreshedCredentials(latest, data, dependencies.nowSeconds()) + if (refreshed) { + dependencies.saveCredentials(location.credentialsPath, refreshed) + } + return refreshed +} + +export async function refreshKimiCredentials( + credentials: KimiCredentials, + location: KimiCredentialLocation, + dependencies: KimiRefreshDependencies = defaultDependencies +): Promise { + return dependencies.acquireLock(location.lockTarget, () => + refreshUnderLock(credentials, location, dependencies) + ) +}