From 9ce662b7bbfe590154d1f2dce7de7e95f4b0cae0 Mon Sep 17 00:00:00 2001 From: 7Sageer <7sageer@djwcb.cn> Date: Thu, 4 Jun 2026 11:33:49 +0800 Subject: [PATCH 1/2] fix: scope managed OAuth credentials --- .changeset/scope-managed-oauth-login.md | 8 + .../src/tui/utils/refresh-providers.ts | 27 +- .../test/tui/utils/refresh-providers.test.ts | 96 ++++ packages/agent-core/src/config/schema.ts | 1 + packages/agent-core/src/config/toml.ts | 2 +- .../agent-core/test/config/configs.test.ts | 32 ++ packages/node-sdk/src/auth.ts | 140 +++++- .../node-sdk/src/kimi-code-model-provider.ts | 13 +- packages/node-sdk/test/auth-facade.test.ts | 474 +++++++++++++++++- packages/oauth/src/constants.ts | 4 +- packages/oauth/src/index.ts | 3 + packages/oauth/src/managed-kimi-code.ts | 113 ++++- packages/oauth/src/managed-usage.ts | 3 +- packages/oauth/src/toolkit.ts | 190 +++++-- packages/oauth/test/managed-kimi-code.test.ts | 193 +++++++ packages/oauth/test/toolkit.test.ts | 320 +++++++++++- 16 files changed, 1532 insertions(+), 87 deletions(-) create mode 100644 .changeset/scope-managed-oauth-login.md create mode 100644 apps/kimi-code/test/tui/utils/refresh-providers.test.ts diff --git a/.changeset/scope-managed-oauth-login.md b/.changeset/scope-managed-oauth-login.md new file mode 100644 index 0000000000..ba75c73829 --- /dev/null +++ b/.changeset/scope-managed-oauth-login.md @@ -0,0 +1,8 @@ +--- +"@moonshot-ai/agent-core": patch +"@moonshot-ai/kimi-code-oauth": patch +"@moonshot-ai/kimi-code-sdk": patch +"@moonshot-ai/kimi-code": patch +--- + +Keep managed OAuth credentials scoped to their configured authentication and API endpoints. diff --git a/apps/kimi-code/src/tui/utils/refresh-providers.ts b/apps/kimi-code/src/tui/utils/refresh-providers.ts index 3bbda132ab..35292ac281 100644 --- a/apps/kimi-code/src/tui/utils/refresh-providers.ts +++ b/apps/kimi-code/src/tui/utils/refresh-providers.ts @@ -10,8 +10,10 @@ import { filterModelsByPrefix, getOpenPlatformById, isOpenPlatformId, + kimiCodeBaseUrl, removeCustomRegistryProvider, removeOpenPlatformConfig, + resolveKimiCodeOAuthRef, type CustomRegistrySource, type ManagedKimiConfigShape, } from '@moonshot-ai/kimi-code-oauth'; @@ -56,6 +58,15 @@ function asManaged(config: KimiConfig): ManagedKimiConfigShape { return config as unknown as ManagedKimiConfigShape; } +function managedRuntimeBaseUrl(configuredBaseUrl: string | undefined): string | undefined { + if (process.env['KIMI_CODE_BASE_URL'] !== undefined) return kimiCodeBaseUrl(); + return configuredBaseUrl; +} + +function envOAuthHost(): string | undefined { + return process.env['KIMI_CODE_OAUTH_HOST'] ?? process.env['KIMI_OAUTH_HOST']; +} + function collectModelIdsForProvider(config: KimiConfig, providerId: string): Set { const ids = new Set(); for (const alias of Object.values(config.models ?? {})) { @@ -120,13 +131,15 @@ export async function refreshAllProviderModels(host: RefreshProviderHost): Promi managedProvider.oauth !== undefined ) { try { - const accessToken = await host.resolveOAuthToken( - KIMI_CODE_PROVIDER_NAME, - managedProvider.oauth, - ); + const baseUrl = managedRuntimeBaseUrl(managedProvider.baseUrl); + const oauth = resolveKimiCodeOAuthRef({ + oauthHost: envOAuthHost() ?? managedProvider.oauth.oauthHost, + baseUrl, + }); + const accessToken = await host.resolveOAuthToken(KIMI_CODE_PROVIDER_NAME, oauth); const models = await fetchManagedKimiCodeModels({ accessToken, - baseUrl: managedProvider.baseUrl, + baseUrl, }); if (models.length > 0) { const beforeIds = collectModelIdsForProvider(config, KIMI_CODE_PROVIDER_NAME); @@ -140,7 +153,9 @@ export async function refreshAllProviderModels(host: RefreshProviderHost): Promi clearManagedKimiCodeConfig(asManaged(config)); applyManagedKimiCodeConfig(asManaged(config), { models, - baseUrl: managedProvider.baseUrl, + baseUrl, + oauthKey: oauth.key, + oauthHost: oauth.oauthHost, preserveDefaultModel: true, }); await host.setConfig({ diff --git a/apps/kimi-code/test/tui/utils/refresh-providers.test.ts b/apps/kimi-code/test/tui/utils/refresh-providers.test.ts new file mode 100644 index 0000000000..712e226f93 --- /dev/null +++ b/apps/kimi-code/test/tui/utils/refresh-providers.test.ts @@ -0,0 +1,96 @@ +import { + KIMI_CODE_PROVIDER_NAME, + resolveKimiCodeOAuthKey, + resolveKimiCodeOAuthRef, +} from '@moonshot-ai/kimi-code-oauth'; +import { afterEach, describe, expect, it, vi } from 'vitest'; + +import { refreshAllProviderModels } from '../../../src/tui/utils/refresh-providers'; +import type { KimiConfig } from '@moonshot-ai/kimi-code-sdk'; + +type FetchMock = ( + input: Parameters[0], + init?: Parameters[1], +) => Promise; + +function fetchInputUrl(input: Parameters[0]): string { + if (typeof input === 'string') return input; + if (input instanceof URL) return input.href; + return input.url; +} + +describe('refreshAllProviderModels', () => { + afterEach(() => { + vi.unstubAllEnvs(); + vi.unstubAllGlobals(); + }); + + it('refreshes managed Kimi Code against environment endpoints over persisted config', async () => { + const configuredBaseUrl = 'https://api.configured.example.test/coding/v1'; + const envBaseUrl = 'https://api.env.example.test/coding/v1'; + const envOauthHost = 'https://auth.env.example.test'; + const configuredOauthKey = resolveKimiCodeOAuthKey({ baseUrl: configuredBaseUrl }); + const envOauthRef = resolveKimiCodeOAuthRef({ + oauthHost: envOauthHost, + baseUrl: envBaseUrl, + }); + const config: KimiConfig = { + providers: { + [KIMI_CODE_PROVIDER_NAME]: { + type: 'kimi', + baseUrl: configuredBaseUrl, + apiKey: '', + oauth: { + storage: 'file', + key: configuredOauthKey, + oauthHost: 'https://auth.kimi.com', + }, + }, + }, + models: { + 'kimi-code/kimi-for-coding': { + provider: KIMI_CODE_PROVIDER_NAME, + model: 'kimi-for-coding', + maxContextSize: 262144, + }, + }, + defaultModel: 'kimi-code/kimi-for-coding', + telemetry: true, + }; + vi.stubEnv('KIMI_CODE_BASE_URL', envBaseUrl); + vi.stubEnv('KIMI_CODE_OAUTH_HOST', envOauthHost); + const resolveOAuthToken = vi.fn(async (_providerName, oauthRef) => { + expect(oauthRef).toEqual(envOauthRef); + return 'env-access-token'; + }); + const fetchMock = vi.fn(async (input, init) => { + expect(fetchInputUrl(input)).toBe(`${envBaseUrl}/models`); + expect(new Headers(init?.headers).get('authorization')).toBe('Bearer env-access-token'); + return new Response( + JSON.stringify({ + data: [ + { + id: 'kimi-for-coding', + context_length: 262144, + supports_reasoning: true, + }, + ], + }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); + }); + vi.stubGlobal('fetch', fetchMock); + + const result = await refreshAllProviderModels({ + getConfig: async () => config, + removeProvider: vi.fn(), + setConfig: vi.fn(), + resolveOAuthToken, + }); + + expect(result.failed).toEqual([]); + expect(result.unchanged).toEqual([KIMI_CODE_PROVIDER_NAME]); + expect(fetchMock).toHaveBeenCalledTimes(1); + expect(resolveOAuthToken).toHaveBeenCalledWith(KIMI_CODE_PROVIDER_NAME, envOauthRef); + }); +}); diff --git a/packages/agent-core/src/config/schema.ts b/packages/agent-core/src/config/schema.ts index d894ef1d45..ae10f5c8f3 100644 --- a/packages/agent-core/src/config/schema.ts +++ b/packages/agent-core/src/config/schema.ts @@ -17,6 +17,7 @@ export type ProviderType = z.infer; export const OAuthRefSchema = z.object({ storage: z.enum(['file', 'keyring']), key: z.string().min(1), + oauthHost: z.string().min(1).optional(), }); export type OAuthRef = z.infer; diff --git a/packages/agent-core/src/config/toml.ts b/packages/agent-core/src/config/toml.ts index c8adb37946..36191f6037 100644 --- a/packages/agent-core/src/config/toml.ts +++ b/packages/agent-core/src/config/toml.ts @@ -482,7 +482,7 @@ function hookToToml(hook: HookDefConfig): Record { function oauthToToml(oauth: OAuthRef): Record { const out: Record = {}; for (const [key, value] of Object.entries(oauth)) { - out[camelToSnake(key)] = value; + setDefined(out, camelToSnake(key), value); } return out; } diff --git a/packages/agent-core/test/config/configs.test.ts b/packages/agent-core/test/config/configs.test.ts index 947469e379..dcade8cfd1 100644 --- a/packages/agent-core/test/config/configs.test.ts +++ b/packages/agent-core/test/config/configs.test.ts @@ -222,6 +222,38 @@ source = { kind = "apiJson", url = "https://registry.example/api.json", apiKey = }); }); + it('round-trips OAuth refs with scoped OAuth hosts', async () => { + const dir = makeTempDir(); + const configPath = join(dir, 'oauth-ref.toml'); + const toml = ` +[providers."managed:kimi-code"] +type = "kimi" +base_url = "https://coding.deva.msh.team/coding/v1" +api_key = "" +oauth = { storage = "file", key = "oauth/kimi-code-env-1234", oauth_host = "https://auth.dev.kimi.team" } + +[services.moonshot_search] +base_url = "https://coding.deva.msh.team/coding/v1/search" +api_key = "" +oauth = { storage = "file", key = "oauth/kimi-code-env-1234", oauth_host = "https://auth.dev.kimi.team" } +`; + const config = parseConfigString(toml, configPath); + expect(config.providers['managed:kimi-code']?.oauth).toEqual({ + storage: 'file', + key: 'oauth/kimi-code-env-1234', + oauthHost: 'https://auth.dev.kimi.team', + }); + expect(config.services?.moonshotSearch?.oauth?.oauthHost).toBe('https://auth.dev.kimi.team'); + + await writeConfigFile(configPath, config); + const text = await readFile(configPath, 'utf-8'); + expect(text).toContain('oauth_host = "https://auth.dev.kimi.team"'); + const roundTripped = parseConfigString(text, configPath); + expect(roundTripped.providers['managed:kimi-code']?.oauth?.oauthHost).toBe( + 'https://auth.dev.kimi.team', + ); + }); + it('loads defaults for absent files and writes typed fields without dropping raw sections', async () => { const dir = makeTempDir(); const configPath = join(dir, 'config.toml'); diff --git a/packages/node-sdk/src/auth.ts b/packages/node-sdk/src/auth.ts index 15a01eff8d..4a6ce30eb0 100644 --- a/packages/node-sdk/src/auth.ts +++ b/packages/node-sdk/src/auth.ts @@ -4,6 +4,8 @@ import { applyManagedKimiCodeLogoutConfig, KIMI_CODE_PROVIDER_NAME, KimiOAuthToolkit, + resolveKimiCodeOAuthKey, + resolveKimiCodeOAuthRef, type AuthManagedUsageResult, type AuthStatus, type BearerTokenProvider, @@ -68,14 +70,23 @@ export class KimiAuthFacade { } async status(providerName?: string | undefined): Promise { - return this.toolkit.status(providerName); + return this.toolkit.status(providerName, this.resolveRuntimeManagedAuth(providerName).oauthRef); } async login( providerName: string | undefined = KIMI_CODE_PROVIDER_NAME, options: KimiAuthLoginOptions = {}, ): Promise { - const result = await this.toolkit.login(providerName, { ...options, provisionConfig: true }); + const auth = this.resolveManagedAuth(providerName); + const baseUrl = options.baseUrl ?? this.configBaseUrlForLogin(auth.baseUrl); + const oauthHost = options.oauthHost ?? this.envOAuthHost(); + const result = await this.toolkit.login(providerName, { + ...options, + baseUrl, + oauthHost, + oauthRef: options.oauthRef ?? this.configOAuthRefForLogin(auth.oauthRef, options, baseUrl), + provisionConfig: true, + }); if (result.provision === undefined) { throw new Error('Kimi auth login did not provision model config.'); } @@ -91,7 +102,10 @@ export class KimiAuthFacade { } async logout(providerName?: string | undefined): Promise { - const result = await this.toolkit.logout(providerName); + const result = await this.toolkit.logout( + providerName, + this.resolveRuntimeManagedAuth(providerName).oauthRef, + ); const updated = readConfigFile(this.options.configPath); this.options.onConfigUpdated?.(updated); return { @@ -101,13 +115,18 @@ export class KimiAuthFacade { } async getManagedUsage(providerName?: string | undefined): Promise { - return this.toolkit.getManagedUsage(providerName); + const auth = this.resolveRuntimeManagedAuth(providerName); + return this.toolkit.getManagedUsage(providerName, { + oauthRef: auth.oauthRef, + baseUrl: auth.baseUrl, + }); } async submitFeedback( input: KimiAuthSubmitFeedbackInput, providerName?: string | undefined, ): Promise { + const auth = this.resolveRuntimeManagedAuth(providerName); return this.toolkit.submitFeedback( { session_id: input.sessionId, @@ -117,17 +136,124 @@ export class KimiAuthFacade { model: input.model, }, providerName, + { + oauthRef: auth.oauthRef, + baseUrl: auth.baseUrl, + }, ); } - async getCachedAccessToken(providerName?: string): Promise { - return this.toolkit.getCachedAccessToken(providerName); + async getCachedAccessToken( + providerName?: string, + oauthRef?: OAuthRef | undefined, + ): Promise { + return this.toolkit.getCachedAccessToken( + providerName, + this.runtimeOAuthRef(providerName, oauthRef), + ); } readonly resolveOAuthTokenProvider = ( providerName: string, oauthRef?: OAuthRef | undefined, ): BearerTokenProvider => { - return this.toolkit.tokenProvider(providerName, oauthRef); + return this.toolkit.tokenProvider(providerName, this.runtimeOAuthRef(providerName, oauthRef)); }; + + private resolveManagedAuth(providerName?: string | undefined): { + readonly oauthRef?: OAuthRef | undefined; + readonly baseUrl?: string | undefined; + } { + const name = providerName ?? KIMI_CODE_PROVIDER_NAME; + const config = readConfigFile(this.options.configPath); + const provider = config.providers[name]; + return { + oauthRef: provider?.oauth, + baseUrl: provider?.baseUrl, + }; + } + + private configBaseUrlForLogin(baseUrl?: string | undefined): string | undefined { + return process.env['KIMI_CODE_BASE_URL'] ?? baseUrl; + } + + private configOAuthRefForLogin( + oauthRef: OAuthRef | undefined, + options: KimiAuthLoginOptions, + baseUrl?: string | undefined, + ): OAuthRef | undefined { + if ( + options.baseUrl !== undefined || + options.oauthHost !== undefined || + process.env['KIMI_CODE_BASE_URL'] !== undefined || + process.env['KIMI_CODE_OAUTH_HOST'] !== undefined || + process.env['KIMI_OAUTH_HOST'] !== undefined + ) { + return undefined; + } + if (oauthRef !== undefined && oauthRef.key !== this.expectedOAuthKey(oauthRef, baseUrl)) { + return undefined; + } + return oauthRef; + } + + private configBaseUrlForRuntime(baseUrl?: string | undefined): string | undefined { + return process.env['KIMI_CODE_BASE_URL'] ?? baseUrl; + } + + private configOAuthRefForRuntime( + oauthRef: OAuthRef | undefined, + baseUrl?: string | undefined, + ): OAuthRef { + const expected = this.expectedOAuthRef(oauthRef, baseUrl); + if (oauthRef === undefined) return expected; + if (process.env['KIMI_CODE_BASE_URL'] !== undefined || this.envOAuthHost() !== undefined) { + return expected; + } + if (oauthRef.key !== expected.key) return expected; + return oauthRef; + } + + private resolveRuntimeManagedAuth(providerName?: string | undefined): { + readonly oauthRef: OAuthRef; + readonly baseUrl?: string | undefined; + } { + const auth = this.resolveManagedAuth(providerName); + const baseUrl = this.configBaseUrlForRuntime(auth.baseUrl); + return { + oauthRef: this.configOAuthRefForRuntime(auth.oauthRef, baseUrl), + baseUrl, + }; + } + + private runtimeOAuthRef( + providerName: string | undefined, + oauthRef?: OAuthRef | undefined, + ): OAuthRef | undefined { + if ((providerName ?? KIMI_CODE_PROVIDER_NAME) !== KIMI_CODE_PROVIDER_NAME) return oauthRef; + const auth = this.resolveManagedAuth(providerName); + const baseUrl = this.configBaseUrlForRuntime(auth.baseUrl); + return this.configOAuthRefForRuntime(oauthRef ?? auth.oauthRef, baseUrl); + } + + private expectedOAuthKey(oauthRef: OAuthRef, baseUrl?: string | undefined): string { + return resolveKimiCodeOAuthKey({ + oauthHost: oauthRef.oauthHost, + baseUrl, + }); + } + + private expectedOAuthRef(oauthRef: OAuthRef | undefined, baseUrl?: string | undefined): OAuthRef { + const envOAuthHost = this.envOAuthHost(); + const hasEnvOverride = + process.env['KIMI_CODE_BASE_URL'] !== undefined || envOAuthHost !== undefined; + return resolveKimiCodeOAuthRef({ + oauthHost: hasEnvOverride ? envOAuthHost : oauthRef?.oauthHost, + baseUrl, + }); + } + + private envOAuthHost(): string | undefined { + return process.env['KIMI_CODE_OAUTH_HOST'] ?? process.env['KIMI_OAUTH_HOST']; + } } diff --git a/packages/node-sdk/src/kimi-code-model-provider.ts b/packages/node-sdk/src/kimi-code-model-provider.ts index 0e03944e6f..b9b3e26316 100644 --- a/packages/node-sdk/src/kimi-code-model-provider.ts +++ b/packages/node-sdk/src/kimi-code-model-provider.ts @@ -8,10 +8,13 @@ import { } from '@moonshot-ai/agent-core'; import { createKimiDefaultHeaders, + KIMI_CODE_FLOW_CONFIG, KIMI_CODE_PROVIDER_NAME, KimiOAuthToolkit, kimiCodeBaseUrl, + resolveKimiCodeOAuthRef, type KimiHostIdentity, + type ManagedKimiOAuthRef, } from '@moonshot-ai/kimi-code-oauth'; import type { ProviderConfig as KosongProviderConfig, @@ -35,6 +38,7 @@ export class KimiForCodingProvider implements ModelProvider { private readonly toolkit: KimiOAuthToolkit; private readonly homeDir: string; private readonly identity: KimiHostIdentity; + private readonly oauthRef: ManagedKimiOAuthRef; constructor(options: KimiForCodingProviderOptions) { this.model = options.model ?? 'kimi-for-coding'; @@ -47,6 +51,10 @@ export class KimiForCodingProvider implements ModelProvider { version: options.version, userAgentSuffix: options.userAgentSuffix, }; + this.oauthRef = resolveKimiCodeOAuthRef({ + oauthHost: KIMI_CODE_FLOW_CONFIG.oauthHost, + baseUrl: this.baseUrl, + }); this.toolkit = new KimiOAuthToolkit({ homeDir: this.homeDir, identity: this.identity, @@ -111,7 +119,10 @@ export class KimiForCodingProvider implements ModelProvider { } private async buildAuth(force: boolean): Promise { - const apiKey = await this.toolkit.ensureFresh(KIMI_CODE_PROVIDER_NAME, { force }); + const apiKey = await this.toolkit.ensureFresh(KIMI_CODE_PROVIDER_NAME, { + force, + oauthRef: this.oauthRef, + }); return { apiKey }; } } diff --git a/packages/node-sdk/test/auth-facade.test.ts b/packages/node-sdk/test/auth-facade.test.ts index 2b493a755a..09e8b71fb7 100644 --- a/packages/node-sdk/test/auth-facade.test.ts +++ b/packages/node-sdk/test/auth-facade.test.ts @@ -2,7 +2,13 @@ import { mkdtemp, readFile, rm, writeFile } from 'node:fs/promises'; import { tmpdir } from 'node:os'; import { join } from 'node:path'; -import { FileTokenStorage, KIMI_CODE_PROVIDER_NAME, type TokenInfo } from '@moonshot-ai/kimi-code-oauth'; +import { + FileTokenStorage, + KIMI_CODE_PROVIDER_NAME, + resolveKimiCodeOAuthKey, + resolveKimiTokenStorageName, + type TokenInfo, +} from '@moonshot-ai/kimi-code-oauth'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { createKimiHarness } from '#/index'; @@ -17,6 +23,12 @@ type FetchMock = ( init?: Parameters[1], ) => Promise; +function fetchInputUrl(input: Parameters[0]): string { + if (typeof input === 'string') return input; + if (input instanceof URL) return input.href; + return input.url; +} + function freshToken(): TokenInfo { return { accessToken: 'oauth-access-token', @@ -34,6 +46,7 @@ beforeEach(async () => { afterEach(async () => { vi.unstubAllGlobals(); + vi.unstubAllEnvs(); await rm(homeDir, { recursive: true, force: true }); }); @@ -49,6 +62,56 @@ describe('KimiHarness.auth', () => { await expect(harness.auth.getCachedAccessToken()).resolves.toBe('oauth-access-token'); }); + it('resolves cached access tokens from the configured scoped OAuth ref', async () => { + const oauthKey = resolveKimiCodeOAuthKey({ + oauthHost: 'https://auth.dev.kimi.team', + baseUrl: 'https://coding.deva.msh.team/coding/v1', + }); + const storageName = resolveKimiTokenStorageName({ oauthKey }); + const storage = new FileTokenStorage(join(homeDir, 'credentials')); + await storage.save('kimi-code', freshToken()); + await storage.save(storageName, { ...freshToken(), accessToken: 'dev-access-token' }); + await writeFile( + join(homeDir, 'config.toml'), + ` +[providers."managed:kimi-code"] +type = "kimi" +base_url = "https://coding.deva.msh.team/coding/v1" +api_key = "" +oauth = { storage = "file", key = "${oauthKey}", oauth_host = "https://auth.dev.kimi.team" } +`, + ); + const harness = createKimiHarness({ homeDir, identity: TEST_IDENTITY }); + + await expect(harness.auth.getCachedAccessToken()).resolves.toBe('dev-access-token'); + }); + + it('reports auth status from the configured scoped OAuth ref', async () => { + const oauthKey = resolveKimiCodeOAuthKey({ + oauthHost: 'https://auth.dev.kimi.team', + baseUrl: 'https://coding.deva.msh.team/coding/v1', + }); + await new FileTokenStorage(join(homeDir, 'credentials')).save( + resolveKimiTokenStorageName({ oauthKey }), + { ...freshToken(), accessToken: 'dev-access-token' }, + ); + await writeFile( + join(homeDir, 'config.toml'), + ` +[providers."managed:kimi-code"] +type = "kimi" +base_url = "https://coding.deva.msh.team/coding/v1" +api_key = "" +oauth = { storage = "file", key = "${oauthKey}", oauth_host = "https://auth.dev.kimi.team" } +`, + ); + const harness = createKimiHarness({ homeDir, identity: TEST_IDENTITY }); + + await expect(harness.auth.status()).resolves.toEqual({ + providers: [{ providerName: KIMI_CODE_PROVIDER_NAME, hasToken: true }], + }); + }); + it('provisions SDK config using an existing Kimi OAuth token', async () => { await new FileTokenStorage(join(homeDir, 'credentials')).save('kimi-code', freshToken()); const fetchMock = vi.fn( @@ -110,6 +173,200 @@ describe('KimiHarness.auth', () => { }); }); + it('logs in against the configured scoped OAuth host and base URL when env is absent', async () => { + const baseUrl = 'https://coding.deva.msh.team/coding/v1'; + const oauthHost = 'https://auth.dev.kimi.team'; + const oauthKey = resolveKimiCodeOAuthKey({ oauthHost, baseUrl }); + const storageName = resolveKimiTokenStorageName({ oauthKey }); + const storage = new FileTokenStorage(join(homeDir, 'credentials')); + await storage.save(storageName, { + ...freshToken(), + accessToken: 'expired-dev-access-token', + refreshToken: 'dev-refresh-token', + expiresAt: 1, + }); + await writeFile( + join(homeDir, 'config.toml'), + ` +[providers."managed:kimi-code"] +type = "kimi" +base_url = "${baseUrl}" +api_key = "" +oauth = { storage = "file", key = "${oauthKey}", oauth_host = "${oauthHost}" } +`, + ); + const fetchMock = vi.fn(async (input, init) => { + const url = fetchInputUrl(input); + if (url === `${oauthHost}/api/oauth/token`) { + if (typeof init?.body !== 'string') throw new TypeError('expected form body'); + const body = new URLSearchParams(init.body); + expect(body.get('grant_type')).toBe('refresh_token'); + expect(body.get('refresh_token')).toBe('dev-refresh-token'); + return new Response( + JSON.stringify({ + access_token: 'rotated-dev-access-token', + refresh_token: 'rotated-dev-refresh-token', + expires_in: 3600, + scope: '', + token_type: 'Bearer', + }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); + } + if (url === `${baseUrl}/models`) { + expect(new Headers(init?.headers).get('authorization')).toBe( + 'Bearer rotated-dev-access-token', + ); + return new Response( + JSON.stringify({ + data: [ + { + id: 'kimi-for-coding', + context_length: 262144, + supports_reasoning: true, + }, + ], + }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); + } + throw new Error(`unexpected request: ${url}`); + }); + vi.stubGlobal('fetch', fetchMock); + const harness = createKimiHarness({ homeDir, identity: TEST_IDENTITY }); + + await expect(harness.auth.login()).resolves.toMatchObject({ + providerName: KIMI_CODE_PROVIDER_NAME, + ok: true, + defaultModel: 'kimi-code/kimi-for-coding', + }); + await expect(storage.load(storageName)).resolves.toMatchObject({ + accessToken: 'rotated-dev-access-token', + }); + const config = await harness.getConfig({ reload: true }); + expect(config.providers[KIMI_CODE_PROVIDER_NAME]).toMatchObject({ + baseUrl, + oauth: { storage: 'file', key: oauthKey, oauthHost }, + }); + expect(fetchMock.mock.calls.map((call) => fetchInputUrl(call[0]))).toEqual([ + `${oauthHost}/api/oauth/token`, + `${baseUrl}/models`, + ]); + }); + + it('recomputes legacy managed OAuth refs during login for non-default base URLs', async () => { + const baseUrl = 'https://api.example.test/coding/v1'; + const oauthKey = resolveKimiCodeOAuthKey({ baseUrl }); + const scopedStorageName = resolveKimiTokenStorageName({ oauthKey }); + const storage = new FileTokenStorage(join(homeDir, 'credentials')); + await storage.save('kimi-code', { ...freshToken(), accessToken: 'legacy-access-token' }); + await storage.save(scopedStorageName, { + ...freshToken(), + accessToken: 'scoped-access-token', + }); + await writeFile( + join(homeDir, 'config.toml'), + ` +[providers."managed:kimi-code"] +type = "kimi" +base_url = "${baseUrl}" +api_key = "" +oauth = { storage = "file", key = "oauth/kimi-code" } +`, + ); + const fetchMock = vi.fn(async (input, init) => { + expect(fetchInputUrl(input)).toBe(`${baseUrl}/models`); + expect(new Headers(init?.headers).get('authorization')).toBe('Bearer scoped-access-token'); + return new Response( + JSON.stringify({ + data: [ + { + id: 'kimi-for-coding', + context_length: 262144, + supports_reasoning: true, + }, + ], + }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); + }); + vi.stubGlobal('fetch', fetchMock); + const harness = createKimiHarness({ homeDir, identity: TEST_IDENTITY }); + + await expect(harness.auth.login()).resolves.toMatchObject({ + providerName: KIMI_CODE_PROVIDER_NAME, + ok: true, + defaultModel: 'kimi-code/kimi-for-coding', + }); + + expect(fetchMock).toHaveBeenCalledTimes(1); + const config = await harness.getConfig({ reload: true }); + expect(config.providers[KIMI_CODE_PROVIDER_NAME]).toMatchObject({ + baseUrl, + oauth: { storage: 'file', key: oauthKey, oauthHost: 'https://auth.kimi.com' }, + }); + }); + + it('logs in against environment OAuth host and base URL over persisted config', async () => { + const configuredBaseUrl = 'https://api.configured.example.test/coding/v1'; + const envBaseUrl = 'https://api.env.example.test/coding/v1'; + const envOauthHost = 'https://auth.env.example.test'; + const configuredOauthKey = resolveKimiCodeOAuthKey({ baseUrl: configuredBaseUrl }); + const envOauthKey = resolveKimiCodeOAuthKey({ oauthHost: envOauthHost, baseUrl: envBaseUrl }); + const storage = new FileTokenStorage(join(homeDir, 'credentials')); + await storage.save(resolveKimiTokenStorageName({ oauthKey: configuredOauthKey }), { + ...freshToken(), + accessToken: 'configured-access-token', + }); + await storage.save(resolveKimiTokenStorageName({ oauthKey: envOauthKey }), { + ...freshToken(), + accessToken: 'env-access-token', + }); + await writeFile( + join(homeDir, 'config.toml'), + ` +[providers."managed:kimi-code"] +type = "kimi" +base_url = "${configuredBaseUrl}" +api_key = "" +oauth = { storage = "file", key = "${configuredOauthKey}", oauth_host = "https://auth.kimi.com" } +`, + ); + vi.stubEnv('KIMI_CODE_BASE_URL', envBaseUrl); + vi.stubEnv('KIMI_CODE_OAUTH_HOST', envOauthHost); + const fetchMock = vi.fn(async (input, init) => { + expect(fetchInputUrl(input)).toBe(`${envBaseUrl}/models`); + expect(new Headers(init?.headers).get('authorization')).toBe('Bearer env-access-token'); + return new Response( + JSON.stringify({ + data: [ + { + id: 'kimi-for-coding', + context_length: 262144, + supports_reasoning: true, + }, + ], + }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); + }); + vi.stubGlobal('fetch', fetchMock); + const harness = createKimiHarness({ homeDir, identity: TEST_IDENTITY }); + + await expect(harness.auth.login()).resolves.toMatchObject({ + providerName: KIMI_CODE_PROVIDER_NAME, + ok: true, + defaultModel: 'kimi-code/kimi-for-coding', + }); + + expect(fetchMock).toHaveBeenCalledTimes(1); + const config = await harness.getConfig({ reload: true }); + expect(config.providers[KIMI_CODE_PROVIDER_NAME]).toMatchObject({ + baseUrl: envBaseUrl, + oauth: { storage: 'file', key: envOauthKey, oauthHost: envOauthHost }, + }); + }); + it('fails clearly when a configured model alias does not have max_context_size', async () => { await new FileTokenStorage(join(homeDir, 'credentials')).save('kimi-code', freshToken()); await writeFile( @@ -215,6 +472,85 @@ oauth = { storage = "file", key = "oauth/kimi-code" } expect(text).not.toContain('moonshot_search'); }); + it('removes the configured scoped OAuth token on logout without touching the production token', async () => { + const oauthKey = resolveKimiCodeOAuthKey({ + oauthHost: 'https://auth.dev.kimi.team', + baseUrl: 'https://coding.deva.msh.team/coding/v1', + }); + const storageName = resolveKimiTokenStorageName({ oauthKey }); + const storage = new FileTokenStorage(join(homeDir, 'credentials')); + await storage.save('kimi-code', freshToken()); + await storage.save(storageName, { ...freshToken(), accessToken: 'dev-access-token' }); + await writeFile( + join(homeDir, 'config.toml'), + ` +default_model = "kimi-code/kimi-for-coding" + +[providers."managed:kimi-code"] +type = "kimi" +base_url = "https://coding.deva.msh.team/coding/v1" +api_key = "" +oauth = { storage = "file", key = "${oauthKey}", oauth_host = "https://auth.dev.kimi.team" } + +[models."kimi-code/kimi-for-coding"] +provider = "managed:kimi-code" +model = "kimi-for-coding" +max_context_size = 262144 +`, + ); + const harness = createKimiHarness({ homeDir, identity: TEST_IDENTITY }); + + await expect(harness.auth.logout()).resolves.toMatchObject({ + providerName: KIMI_CODE_PROVIDER_NAME, + ok: true, + }); + + await expect(storage.load(storageName)).resolves.toBeUndefined(); + await expect(storage.load('kimi-code')).resolves.toMatchObject({ + accessToken: 'oauth-access-token', + }); + }); + + it('recomputes legacy managed OAuth refs during logout for non-default base URLs', async () => { + const baseUrl = 'https://api.example.test/coding/v1'; + const oauthKey = resolveKimiCodeOAuthKey({ baseUrl }); + const scopedStorageName = resolveKimiTokenStorageName({ oauthKey }); + const storage = new FileTokenStorage(join(homeDir, 'credentials')); + await storage.save('kimi-code', freshToken()); + await storage.save(scopedStorageName, { + ...freshToken(), + accessToken: 'scoped-access-token', + }); + await writeFile( + join(homeDir, 'config.toml'), + ` +default_model = "kimi-code/kimi-for-coding" + +[providers."managed:kimi-code"] +type = "kimi" +base_url = "${baseUrl}" +api_key = "" +oauth = { storage = "file", key = "oauth/kimi-code" } + +[models."kimi-code/kimi-for-coding"] +provider = "managed:kimi-code" +model = "kimi-for-coding" +max_context_size = 262144 +`, + ); + const harness = createKimiHarness({ homeDir, identity: TEST_IDENTITY }); + + await expect(harness.auth.logout()).resolves.toMatchObject({ + providerName: KIMI_CODE_PROVIDER_NAME, + ok: true, + }); + + await expect(storage.load(scopedStorageName)).resolves.toBeUndefined(); + await expect(storage.load('kimi-code')).resolves.toMatchObject({ + accessToken: 'oauth-access-token', + }); + }); + it('gets managed usage without host identity and sends only auth headers', async () => { await new FileTokenStorage(join(homeDir, 'credentials')).save('kimi-code', freshToken()); const fetchMock = vi.fn( @@ -243,6 +579,142 @@ oauth = { storage = "file", key = "oauth/kimi-code" } expect(headers.get('x-msh-platform')).toBeNull(); }); + it('uses configured scoped OAuth refs and base URLs for managed usage and feedback', async () => { + const baseUrl = 'https://coding.deva.msh.team/coding/v1'; + const oauthKey = resolveKimiCodeOAuthKey({ + oauthHost: 'https://auth.dev.kimi.team', + baseUrl, + }); + const storageName = resolveKimiTokenStorageName({ oauthKey }); + await new FileTokenStorage(join(homeDir, 'credentials')).save(storageName, { + ...freshToken(), + accessToken: 'dev-access-token', + }); + await writeFile( + join(homeDir, 'config.toml'), + ` +[providers."managed:kimi-code"] +type = "kimi" +base_url = "${baseUrl}" +api_key = "" +oauth = { storage = "file", key = "${oauthKey}", oauth_host = "https://auth.dev.kimi.team" } +`, + ); + const fetchMock = vi.fn(async (input) => { + const url = fetchInputUrl(input); + if (url.endsWith('/usages')) { + return new Response( + JSON.stringify({ usage: { used: 2, limit: 10, name: 'Dev limit' } }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); + } + return new Response('', { status: 200 }); + }); + vi.stubGlobal('fetch', fetchMock); + const harness = createKimiHarness({ homeDir }); + + await expect(harness.auth.getManagedUsage()).resolves.toMatchObject({ + kind: 'ok', + summary: { label: 'Dev limit', used: 2, limit: 10 }, + }); + await expect( + harness.auth.submitFeedback({ + content: 'dev feedback', + sessionId: 'sess-dev', + version: 'kimi-code-0.1.1', + os: 'Darwin 25.3.0', + model: 'kimi-code/kimi-for-coding', + }), + ).resolves.toEqual({ kind: 'ok' }); + + expect(fetchMock.mock.calls[0]?.[0]).toBe(`${baseUrl}/usages`); + expect(fetchMock.mock.calls[1]?.[0]).toBe(`${baseUrl}/feedback`); + for (const call of fetchMock.mock.calls) { + const init = call[1]; + expect(new Headers(init?.headers).get('authorization')).toBe('Bearer dev-access-token'); + } + }); + + it('uses environment managed endpoints for usage and feedback over persisted config', async () => { + const configuredBaseUrl = 'https://api.configured.example.test/coding/v1'; + const envBaseUrl = 'https://api.env.example.test/coding/v1'; + const envOauthHost = 'https://auth.env.example.test'; + const configuredOauthKey = resolveKimiCodeOAuthKey({ baseUrl: configuredBaseUrl }); + const envOauthKey = resolveKimiCodeOAuthKey({ + oauthHost: envOauthHost, + baseUrl: envBaseUrl, + }); + const storage = new FileTokenStorage(join(homeDir, 'credentials')); + await storage.save(resolveKimiTokenStorageName({ oauthKey: configuredOauthKey }), { + ...freshToken(), + accessToken: 'configured-access-token', + }); + await storage.save(resolveKimiTokenStorageName({ oauthKey: envOauthKey }), { + ...freshToken(), + accessToken: 'env-access-token', + }); + await writeFile( + join(homeDir, 'config.toml'), + ` +[providers."managed:kimi-code"] +type = "kimi" +base_url = "${configuredBaseUrl}" +api_key = "" +oauth = { storage = "file", key = "${configuredOauthKey}", oauth_host = "https://auth.kimi.com" } +`, + ); + vi.stubEnv('KIMI_CODE_BASE_URL', envBaseUrl); + vi.stubEnv('KIMI_CODE_OAUTH_HOST', envOauthHost); + const fetchMock = vi.fn(async (input) => { + const url = fetchInputUrl(input); + if (url.endsWith('/usages')) { + return new Response( + JSON.stringify({ usage: { used: 3, limit: 10, name: 'Env limit' } }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); + } + return new Response('', { status: 200 }); + }); + vi.stubGlobal('fetch', fetchMock); + const harness = createKimiHarness({ homeDir }); + + await expect(harness.auth.status()).resolves.toEqual({ + providers: [{ providerName: KIMI_CODE_PROVIDER_NAME, hasToken: true }], + }); + await expect(harness.auth.getCachedAccessToken()).resolves.toBe('env-access-token'); + await expect( + harness.auth.resolveOAuthTokenProvider(KIMI_CODE_PROVIDER_NAME).getAccessToken(), + ).resolves.toBe('env-access-token'); + await expect( + harness.auth + .resolveOAuthTokenProvider(KIMI_CODE_PROVIDER_NAME, { + storage: 'file', + key: configuredOauthKey, + oauthHost: 'https://auth.kimi.com', + }) + .getAccessToken(), + ).resolves.toBe('env-access-token'); + await expect(harness.auth.getManagedUsage()).resolves.toMatchObject({ + kind: 'ok', + summary: { label: 'Env limit', used: 3, limit: 10 }, + }); + await expect( + harness.auth.submitFeedback({ + content: 'env feedback', + sessionId: 'sess-env', + version: 'kimi-code-0.1.1', + os: 'Darwin 25.3.0', + model: 'kimi-code/kimi-for-coding', + }), + ).resolves.toEqual({ kind: 'ok' }); + + expect(fetchMock.mock.calls[0]?.[0]).toBe(`${envBaseUrl}/usages`); + expect(fetchMock.mock.calls[1]?.[0]).toBe(`${envBaseUrl}/feedback`); + for (const call of fetchMock.mock.calls) { + expect(new Headers(call[1]?.headers).get('authorization')).toBe('Bearer env-access-token'); + } + }); + it('submitFeedback maps camelCase input to snake_case body and posts with bearer auth', async () => { await new FileTokenStorage(join(homeDir, 'credentials')).save('kimi-code', freshToken()); const fetchMock = vi.fn(async () => new Response('', { status: 200 })); diff --git a/packages/oauth/src/constants.ts b/packages/oauth/src/constants.ts index 3c338fb6eb..58ad3f2596 100644 --- a/packages/oauth/src/constants.ts +++ b/packages/oauth/src/constants.ts @@ -1,10 +1,12 @@ import type { OAuthFlowConfig } from './types'; +export const DEFAULT_KIMI_CODE_OAUTH_HOST = 'https://auth.kimi.com'; + export const KIMI_CODE_FLOW_CONFIG: OAuthFlowConfig = { name: 'kimi-code', oauthHost: process.env['KIMI_CODE_OAUTH_HOST'] ?? process.env['KIMI_OAUTH_HOST'] ?? - 'https://auth.kimi.com', + DEFAULT_KIMI_CODE_OAUTH_HOST, clientId: '17e5f671-d194-4dfb-9706-5516cb48c098', }; diff --git a/packages/oauth/src/index.ts b/packages/oauth/src/index.ts index 4aeda80e3c..84b95ae6d0 100644 --- a/packages/oauth/src/index.ts +++ b/packages/oauth/src/index.ts @@ -46,6 +46,8 @@ export { KIMI_CODE_PLATFORM_ID, KIMI_CODE_PROVIDER_NAME, provisionManagedKimiCodeConfig, + resolveKimiCodeOAuthKey, + resolveKimiCodeOAuthRef, } from './managed-kimi-code'; export type { FetchManagedKimiCodeModelsOptions, @@ -55,6 +57,7 @@ export type { ManagedKimiCodeProvisionResult, ManagedKimiConfigAdapter, ManagedKimiConfigShape, + ManagedKimiOAuthRef, ProvisionManagedKimiCodeConfigOptions, } from './managed-kimi-code'; diff --git a/packages/oauth/src/managed-kimi-code.ts b/packages/oauth/src/managed-kimi-code.ts index caca34c3bc..3b7a3a0b61 100644 --- a/packages/oauth/src/managed-kimi-code.ts +++ b/packages/oauth/src/managed-kimi-code.ts @@ -1,10 +1,15 @@ +import { createHash } from 'node:crypto'; + import { readApiErrorMessage } from './api-error'; -import { kimiCodeBaseUrl } from './managed-usage'; +import { DEFAULT_KIMI_CODE_OAUTH_HOST } from './constants'; +import { OAuthUnauthorizedError } from './errors'; +import { DEFAULT_KIMI_CODE_BASE_URL, kimiCodeBaseUrl } from './managed-usage'; import { isRecord } from './utils'; export const KIMI_CODE_PLATFORM_ID = 'kimi-code'; export const KIMI_CODE_PROVIDER_NAME = 'managed:kimi-code'; export const KIMI_CODE_OAUTH_KEY = 'oauth/kimi-code'; +const KIMI_CODE_SCOPED_OAUTH_KEY_PREFIX = 'oauth/kimi-code-env-'; export interface ManagedKimiCodeModelInfo { readonly id: string; @@ -45,7 +50,8 @@ export interface ManagedKimiCodeCleanupResult { export interface ManagedKimiOAuthRef { readonly storage: 'file'; - readonly key: typeof KIMI_CODE_OAUTH_KEY; + readonly key: string; + readonly oauthHost?: string | undefined; } export interface ManagedKimiProviderConfig { @@ -94,6 +100,8 @@ export interface ManagedKimiConfigAdapter { input: { readonly models: readonly ManagedKimiCodeModelInfo[]; readonly baseUrl?: string | undefined; + readonly oauthKey?: string | undefined; + readonly oauthHost?: string | undefined; readonly preserveDefaultModel?: boolean | undefined; }, ): ManagedKimiCodeApplyResult; @@ -105,6 +113,8 @@ export interface ProvisionManagedKimiCodeConfigOptions { readonly adapter: ManagedKimiConfigAdapter; readonly accessToken: string; readonly baseUrl?: string | undefined; + readonly oauthKey?: string | undefined; + readonly oauthHost?: string | undefined; readonly preserveDefaultModel?: boolean | undefined; readonly fetchImpl?: typeof fetch | undefined; } @@ -131,6 +141,77 @@ function defaultBaseUrl(baseUrl: string | undefined): string { return (baseUrl ?? kimiCodeBaseUrl()).replace(/\/+$/, ''); } +function normalizeEndpoint(value: string): string { + return value.trim().replace(/\/+$/, ''); +} + +function persistedOAuthHost(options: { + readonly key: string; + readonly oauthHost?: string | undefined; +}): string | undefined { + const oauthHost = options.oauthHost; + const normalized = normalizeEndpoint(oauthHost ?? DEFAULT_KIMI_CODE_OAUTH_HOST); + if ( + options.key === KIMI_CODE_OAUTH_KEY && + normalized === normalizeEndpoint(DEFAULT_KIMI_CODE_OAUTH_HOST) + ) { + return undefined; + } + return normalized; +} + +function managedOAuthRef(options: { + readonly key: string; + readonly oauthHost?: string | undefined; +}): ManagedKimiOAuthRef { + const oauthHost = persistedOAuthHost(options); + return { + storage: 'file', + key: options.key, + oauthHost, + }; +} + +export function resolveKimiCodeOAuthKey(options: { + readonly oauthHost?: string | undefined; + readonly baseUrl?: string | undefined; +}): string { + const oauthHost = normalizeEndpoint(options.oauthHost ?? DEFAULT_KIMI_CODE_OAUTH_HOST); + const baseUrl = defaultBaseUrl(options.baseUrl); + const defaultOauthHost = normalizeEndpoint(DEFAULT_KIMI_CODE_OAUTH_HOST); + const defaultApiBaseUrl = normalizeEndpoint(DEFAULT_KIMI_CODE_BASE_URL); + + if (oauthHost === defaultOauthHost && baseUrl === defaultApiBaseUrl) { + return KIMI_CODE_OAUTH_KEY; + } + + const digest = createHash('sha256') + .update(JSON.stringify({ oauthHost, baseUrl })) + .digest('hex') + .slice(0, 16); + return `${KIMI_CODE_SCOPED_OAUTH_KEY_PREFIX}${digest}`; +} + +/** + * Resolve the full managed-Kimi-Code OAuth ref (credential storage key + + * persisted host) for an (oauthHost, baseUrl) environment. + * + * Single source of truth for "which credential slot does this environment map + * to". Login, provisioning, and the runtime provider all derive their ref + * through here, so the slot a token is written to always matches the slot it + * is later read from — preventing the env-mismatch credential mix-ups this + * scoping is meant to fix. + */ +export function resolveKimiCodeOAuthRef(options: { + readonly oauthHost?: string | undefined; + readonly baseUrl?: string | undefined; +}): ManagedKimiOAuthRef { + return managedOAuthRef({ + key: resolveKimiCodeOAuthKey(options), + oauthHost: options.oauthHost, + }); +} + function toModelInfo(item: unknown): ManagedKimiCodeModelInfo | undefined { if (!isRecord(item) || typeof item['id'] !== 'string' || item['id'].length === 0) { return undefined; @@ -168,12 +249,16 @@ export async function fetchManagedKimiCodeModels( }, }); if (!response.ok) { - throw new Error( - await readApiErrorMessage( - response, - `Failed to list Kimi Code models (HTTP ${response.status}).`, - ), + const message = await readApiErrorMessage( + response, + `Failed to list Kimi Code models (HTTP ${response.status}).`, ); + if (response.status === 401 || response.status === 402 || response.status === 403) { + throw new OAuthUnauthorizedError( + `Kimi Code models endpoint ${baseUrl} rejected OAuth credentials: ${message}`, + ); + } + throw new Error(message); } const payload: unknown = await response.json(); if (!isRecord(payload) || !Array.isArray(payload['data'])) { @@ -189,6 +274,8 @@ export function applyManagedKimiCodeConfig( options: { readonly models: readonly ManagedKimiCodeModelInfo[]; readonly baseUrl?: string | undefined; + readonly oauthKey?: string | undefined; + readonly oauthHost?: string | undefined; readonly preserveDefaultModel?: boolean | undefined; }, ): ManagedKimiCodeApplyResult { @@ -200,6 +287,10 @@ export function applyManagedKimiCodeConfig( } const baseUrl = defaultBaseUrl(options.baseUrl); + const oauth = + options.oauthKey !== undefined + ? managedOAuthRef({ key: options.oauthKey, oauthHost: options.oauthHost }) + : resolveKimiCodeOAuthRef({ baseUrl, oauthHost: options.oauthHost }); const existingModels = config.models ?? {}; const selectedDefault = selectDefaultModel(config, options.models, { preserveExisting: options.preserveDefaultModel === true, @@ -209,7 +300,7 @@ export function applyManagedKimiCodeConfig( type: 'kimi', baseUrl, apiKey: '', - oauth: { storage: 'file', key: KIMI_CODE_OAUTH_KEY }, + oauth, }; for (const [key, model] of Object.entries(existingModels)) { @@ -235,12 +326,12 @@ export function applyManagedKimiCodeConfig( moonshotSearch: { baseUrl: `${baseUrl}/search`, apiKey: '', - oauth: { storage: 'file', key: KIMI_CODE_OAUTH_KEY }, + oauth, }, moonshotFetch: { baseUrl: `${baseUrl}/fetch`, apiKey: '', - oauth: { storage: 'file', key: KIMI_CODE_OAUTH_KEY }, + oauth, }, }; @@ -388,6 +479,8 @@ export async function provisionManagedKimiCodeConfig( const applied = options.adapter.apply(config, { models, baseUrl: options.baseUrl, + oauthKey: options.oauthKey, + oauthHost: options.oauthHost, preserveDefaultModel: options.preserveDefaultModel, }); await options.adapter.write(config); diff --git a/packages/oauth/src/managed-usage.ts b/packages/oauth/src/managed-usage.ts index 771223395e..534104a20e 100644 --- a/packages/oauth/src/managed-usage.ts +++ b/packages/oauth/src/managed-usage.ts @@ -22,6 +22,7 @@ import { isRecord } from './utils'; const MANAGED_PREFIX = 'managed:'; const KIMI_CODE_PLATFORM_ID = 'kimi-code'; +export const DEFAULT_KIMI_CODE_BASE_URL = 'https://api.kimi.com/coding/v1'; export function isManagedKimiCode(providerKey?: string | null): boolean { if (!providerKey) return false; @@ -30,7 +31,7 @@ export function isManagedKimiCode(providerKey?: string | null): boolean { } export function kimiCodeBaseUrl(): string { - return process.env['KIMI_CODE_BASE_URL'] ?? 'https://api.kimi.com/coding/v1'; + return process.env['KIMI_CODE_BASE_URL'] ?? DEFAULT_KIMI_CODE_BASE_URL; } export function kimiCodeUsageUrl(): string { diff --git a/packages/oauth/src/toolkit.ts b/packages/oauth/src/toolkit.ts index c09d3b59f3..e5726b677a 100644 --- a/packages/oauth/src/toolkit.ts +++ b/packages/oauth/src/toolkit.ts @@ -14,6 +14,7 @@ import { KIMI_CODE_OAUTH_KEY, KIMI_CODE_PROVIDER_NAME, provisionManagedKimiCodeConfig, + resolveKimiCodeOAuthKey, type ManagedKimiCodeProvisionResult, type ManagedKimiConfigAdapter, } from './managed-kimi-code'; @@ -58,10 +59,13 @@ export interface KimiOAuthToolkitOptions { export interface KimiOAuthLoginOptions extends LoginOptions { readonly provisionConfig?: boolean | undefined; readonly baseUrl?: string | undefined; + readonly oauthRef?: KimiOAuthTokenRef | undefined; + readonly oauthHost?: string | undefined; } export interface KimiOAuthTokenRef { readonly key?: string | undefined; + readonly oauthHost?: string | undefined; } export interface KimiOAuthLoginResult { @@ -114,13 +118,18 @@ export class KimiOAuthToolkit { }; } - async status(providerName?: string | undefined): Promise { + async status( + providerName?: string | undefined, + oauthRef?: KimiOAuthTokenRef | undefined, + ): Promise { const name = providerName ?? KIMI_CODE_PROVIDER_NAME; + const oauthHost = this.oauthHostFor(oauthRef); + const oauthKey = oauthRef?.key ?? this.defaultOAuthKey(undefined, oauthHost); return { providers: [ { providerName: name, - hasToken: await this.managerFor(name).hasToken(), + hasToken: await this.managerFor(name, oauthKey, oauthHost).hasToken(), }, ], }; @@ -131,48 +140,81 @@ export class KimiOAuthToolkit { options: KimiOAuthLoginOptions = {}, ): Promise { const name = providerName ?? KIMI_CODE_PROVIDER_NAME; - const manager = this.managerFor(name); + const oauthHost = this.oauthHostFor(options.oauthRef, options.oauthHost); + const oauthKey = options.oauthRef?.key ?? this.defaultOAuthKey(options.baseUrl, oauthHost); + const manager = this.managerFor(name, oauthKey, oauthHost); const hadToken = await manager.hasToken(); + let usedDeviceLogin = false; + const loginWithDevice = async (): Promise => { + usedDeviceLogin = true; + return ( + await manager.login({ + signal: options.signal, + onDeviceCode: options.onDeviceCode, + }) + ).accessToken; + }; let accessToken: string; if (hadToken) { try { accessToken = await manager.ensureFresh(); } catch (error) { if (!(error instanceof OAuthUnauthorizedError)) throw error; - accessToken = ( - await manager.login({ - signal: options.signal, - onDeviceCode: options.onDeviceCode, - }) - ).accessToken; + accessToken = await loginWithDevice(); } } else { - accessToken = ( - await manager.login({ - signal: options.signal, - onDeviceCode: options.onDeviceCode, - }) - ).accessToken; + accessToken = await loginWithDevice(); } const shouldProvision = options.provisionConfig ?? this.configAdapter !== undefined; - const provision = - shouldProvision && this.configAdapter !== undefined - ? await provisionManagedKimiCodeConfig({ - accessToken, - adapter: this.configAdapter, - baseUrl: options.baseUrl, - preserveDefaultModel: hadToken, - fetchImpl: this.fetchImpl, - }) - : undefined; + const configAdapter = this.configAdapter; + let provision: ManagedKimiCodeProvisionResult | undefined; + if (shouldProvision && configAdapter !== undefined) { + const provisionWithToken = (token: string): Promise => + provisionManagedKimiCodeConfig({ + accessToken: token, + adapter: configAdapter, + baseUrl: options.baseUrl, + oauthKey, + oauthHost, + preserveDefaultModel: hadToken, + fetchImpl: this.fetchImpl, + }); + try { + provision = await provisionWithToken(accessToken); + } catch (error) { + if (!(error instanceof OAuthUnauthorizedError) || !hadToken || usedDeviceLogin) { + throw error; + } + let retryToken: string; + try { + retryToken = await manager.ensureFresh({ force: true }); + } catch (refreshError) { + if (!(refreshError instanceof OAuthUnauthorizedError)) throw refreshError; + retryToken = await loginWithDevice(); + } + try { + provision = await provisionWithToken(retryToken); + } catch (retryError) { + if (!(retryError instanceof OAuthUnauthorizedError) || usedDeviceLogin) { + throw retryError; + } + provision = await provisionWithToken(await loginWithDevice()); + } + } + } return { providerName: name, ok: true, provision }; } - async logout(providerName?: string | undefined): Promise { + async logout( + providerName?: string | undefined, + oauthRef?: KimiOAuthTokenRef | undefined, + ): Promise { const name = providerName ?? KIMI_CODE_PROVIDER_NAME; - await this.managerFor(name).logout(); + const oauthHost = this.oauthHostFor(oauthRef); + const oauthKey = oauthRef?.key ?? this.defaultOAuthKey(undefined, oauthHost); + await this.managerFor(name, oauthKey, oauthHost).logout(); if (this.configAdapter?.remove !== undefined && name === KIMI_CODE_PROVIDER_NAME) { const config = await this.configAdapter.read(); this.configAdapter.remove(config); @@ -183,10 +225,15 @@ export class KimiOAuthToolkit { async ensureFresh( providerName?: string | undefined, - options: { readonly force?: boolean | undefined } = {}, + options: { + readonly force?: boolean | undefined; + readonly oauthRef?: KimiOAuthTokenRef | undefined; + } = {}, ): Promise { const name = providerName ?? KIMI_CODE_PROVIDER_NAME; - return this.managerFor(name).ensureFresh(options); + const oauthHost = this.oauthHostFor(options.oauthRef); + const oauthKey = options.oauthRef?.key ?? this.defaultOAuthKey(undefined, oauthHost); + return this.managerFor(name, oauthKey, oauthHost).ensureFresh(options); } async getCachedAccessToken( @@ -194,8 +241,9 @@ export class KimiOAuthToolkit { oauthRef?: KimiOAuthTokenRef, ): Promise { const name = providerName ?? KIMI_CODE_PROVIDER_NAME; - const oauthKey = oauthRef?.key ?? KIMI_CODE_OAUTH_KEY; - return this.managerFor(name, oauthKey).getCachedAccessToken(); + const oauthHost = this.oauthHostFor(oauthRef); + const oauthKey = oauthRef?.key ?? this.defaultOAuthKey(undefined, oauthHost); + return this.managerFor(name, oauthKey, oauthHost).getCachedAccessToken(); } tokenProvider( @@ -203,17 +251,26 @@ export class KimiOAuthToolkit { oauthRef?: KimiOAuthTokenRef | undefined, ): BearerTokenProvider { const name = providerName ?? KIMI_CODE_PROVIDER_NAME; - const oauthKey = oauthRef?.key ?? KIMI_CODE_OAUTH_KEY; + const oauthHost = this.oauthHostFor(oauthRef); + const oauthKey = oauthRef?.key ?? this.defaultOAuthKey(undefined, oauthHost); return { - getAccessToken: (options) => this.managerFor(name, oauthKey).ensureFresh(options), + getAccessToken: (options) => this.managerFor(name, oauthKey, oauthHost).ensureFresh(options), }; } - async getManagedUsage(providerName?: string | undefined): Promise { + async getManagedUsage( + providerName?: string | undefined, + options: { + readonly oauthRef?: KimiOAuthTokenRef | undefined; + readonly baseUrl?: string | undefined; + } = {}, + ): Promise { const name = providerName ?? KIMI_CODE_PROVIDER_NAME; try { - const accessToken = await this.ensureFresh(name); - const result = await fetchManagedUsage(kimiCodeUsageUrl(), accessToken); + const accessToken = await this.ensureFresh(name, { + oauthRef: options.oauthRef ?? this.defaultOAuthRef(options.baseUrl), + }); + const result = await fetchManagedUsage(managedUsageUrl(options.baseUrl), accessToken); if (result.kind === 'error') return result; return { kind: 'ok', @@ -231,11 +288,17 @@ export class KimiOAuthToolkit { async submitFeedback( body: SubmitFeedbackBody, providerName?: string | undefined, + options: { + readonly oauthRef?: KimiOAuthTokenRef | undefined; + readonly baseUrl?: string | undefined; + } = {}, ): Promise { const name = providerName ?? KIMI_CODE_PROVIDER_NAME; try { - const accessToken = await this.ensureFresh(name); - return await fetchSubmitFeedback(kimiCodeFeedbackUrl(), accessToken, body); + const accessToken = await this.ensureFresh(name, { + oauthRef: options.oauthRef ?? this.defaultOAuthRef(options.baseUrl), + }); + return await fetchSubmitFeedback(managedFeedbackUrl(options.baseUrl), accessToken, body); } catch (error) { return { kind: 'error', @@ -244,15 +307,22 @@ export class KimiOAuthToolkit { } } - managerFor(providerName: string, oauthKey = KIMI_CODE_OAUTH_KEY): OAuthManager { + managerFor( + providerName: string, + oauthKey = KIMI_CODE_OAUTH_KEY, + oauthHost?: string | undefined, + ): OAuthManager { const storageName = resolveKimiTokenStorageName({ providerName, oauthKey }); - let manager = this.managers.get(storageName); + const effectiveOAuthHost = oauthHost ?? this.flowConfig.oauthHost; + const managerKey = `${storageName}\0${normalizeOAuthHost(effectiveOAuthHost)}`; + let manager = this.managers.get(managerKey); if (manager !== undefined) return manager; const identity = this.identity; manager = new OAuthManager({ config: { ...this.flowConfig, + oauthHost: effectiveOAuthHost, name: storageName, }, storage: this.storage, @@ -267,9 +337,33 @@ export class KimiOAuthToolkit { }), ...this.managerOptions, }); - this.managers.set(storageName, manager); + this.managers.set(managerKey, manager); return manager; } + + private defaultOAuthKey( + baseUrl?: string | undefined, + oauthHost?: string | undefined, + ): string { + return resolveKimiCodeOAuthKey({ + oauthHost: oauthHost ?? this.flowConfig.oauthHost, + baseUrl, + }); + } + + private defaultOAuthRef(baseUrl?: string | undefined): KimiOAuthTokenRef { + return { + key: this.defaultOAuthKey(baseUrl, this.flowConfig.oauthHost), + oauthHost: this.flowConfig.oauthHost, + }; + } + + private oauthHostFor( + oauthRef?: KimiOAuthTokenRef | undefined, + oauthHost?: string | undefined, + ): string { + return oauthRef?.oauthHost ?? oauthHost ?? this.flowConfig.oauthHost; + } } export function resolveKimiTokenStorageName(input: { @@ -298,3 +392,17 @@ function defaultKimiHome(): string { if (override !== undefined && override.length > 0) return override; return join(homedir(), '.kimi-code'); } + +function managedUsageUrl(baseUrl: string | undefined): string { + if (baseUrl === undefined) return kimiCodeUsageUrl(); + return `${baseUrl.replace(/\/+$/, '')}/usages`; +} + +function managedFeedbackUrl(baseUrl: string | undefined): string { + if (baseUrl === undefined) return kimiCodeFeedbackUrl(); + return `${baseUrl.replace(/\/+$/, '')}/feedback`; +} + +function normalizeOAuthHost(oauthHost: string): string { + return oauthHost.trim().replace(/\/+$/, ''); +} diff --git a/packages/oauth/test/managed-kimi-code.test.ts b/packages/oauth/test/managed-kimi-code.test.ts index 504775820d..f09f681c3d 100644 --- a/packages/oauth/test/managed-kimi-code.test.ts +++ b/packages/oauth/test/managed-kimi-code.test.ts @@ -5,10 +5,14 @@ import { applyManagedKimiCodeConfig, clearManagedKimiCodeConfig, fetchManagedKimiCodeModels, + KIMI_CODE_OAUTH_KEY, KIMI_CODE_PROVIDER_NAME, provisionManagedKimiCodeConfig, + resolveKimiCodeOAuthKey, + resolveKimiCodeOAuthRef, type ManagedKimiConfigShape, } from '../src/managed-kimi-code'; +import { OAuthUnauthorizedError } from '../src/errors'; function makeModelsResponse(): Response { return new Response( @@ -37,6 +41,69 @@ function makeModelsResponse(): Response { } describe('provisionManagedKimiCodeConfig', () => { + it('keeps the legacy credential key for the default production environment', () => { + expect( + resolveKimiCodeOAuthKey({ + oauthHost: 'https://auth.kimi.com/', + baseUrl: 'https://api.kimi.com/coding/v1/', + }), + ).toBe(KIMI_CODE_OAUTH_KEY); + }); + + it('scopes credential keys for non-default OAuth hosts and API base URLs', () => { + const devKey = resolveKimiCodeOAuthKey({ + oauthHost: 'https://auth.dev.kimi.team', + baseUrl: 'https://coding.deva.msh.team/coding/v1', + }); + + expect(devKey).not.toBe(KIMI_CODE_OAUTH_KEY); + expect(devKey).toMatch(/^oauth\/kimi-code-env-[a-f0-9]{16}$/); + expect( + resolveKimiCodeOAuthKey({ + oauthHost: 'https://auth.dev.kimi.team/', + baseUrl: 'https://coding.deva.msh.team/coding/v1/', + }), + ).toBe(devKey); + }); + + it('derives a full OAuth ref whose key and persisted host stay in sync', () => { + // Default environment collapses to the legacy ref (no persisted host), so + // existing production credentials keep resolving to `kimi-code.json`. + expect( + resolveKimiCodeOAuthRef({ + oauthHost: 'https://auth.kimi.com/', + baseUrl: 'https://api.kimi.com/coding/v1/', + }), + ).toEqual({ storage: 'file', key: KIMI_CODE_OAUTH_KEY, oauthHost: undefined }); + + const defaultAuthCustomApiRef = resolveKimiCodeOAuthRef({ + baseUrl: 'https://api.example.test/coding/v1', + }); + expect(defaultAuthCustomApiRef).toEqual({ + storage: 'file', + key: resolveKimiCodeOAuthKey({ + oauthHost: 'https://auth.kimi.com', + baseUrl: 'https://api.example.test/coding/v1', + }), + oauthHost: 'https://auth.kimi.com', + }); + + // A non-default environment yields a scoped key AND the normalized host, + // both derived from the same input — login and runtime cannot drift apart. + const devRef = resolveKimiCodeOAuthRef({ + oauthHost: 'https://auth.dev.kimi.team/', + baseUrl: 'https://coding.deva.msh.team/coding/v1', + }); + expect(devRef).toEqual({ + storage: 'file', + key: resolveKimiCodeOAuthKey({ + oauthHost: 'https://auth.dev.kimi.team', + baseUrl: 'https://coding.deva.msh.team/coding/v1', + }), + oauthHost: 'https://auth.dev.kimi.team', + }); + }); + it('writes the managed provider, models, services, and default model through an adapter', async () => { const config: ManagedKimiConfigShape = { providers: { @@ -122,6 +189,76 @@ describe('provisionManagedKimiCodeConfig', () => { expect(Object.keys(config.services ?? {})).toEqual(['moonshotSearch', 'moonshotFetch']); }); + it('writes scoped OAuth refs when provisioning against a non-default environment', async () => { + const config: ManagedKimiConfigShape = { + providers: {}, + }; + const oauthKey = resolveKimiCodeOAuthKey({ + oauthHost: 'https://auth.dev.kimi.team', + baseUrl: 'https://coding.deva.msh.team/coding/v1', + }); + + await provisionManagedKimiCodeConfig({ + accessToken: 'oauth-access-token', + baseUrl: 'https://coding.deva.msh.team/coding/v1', + oauthKey, + oauthHost: 'https://auth.dev.kimi.team', + fetchImpl: vi.fn(async () => makeModelsResponse()) as unknown as typeof fetch, + adapter: { + read: () => config, + write: vi.fn(), + apply: applyManagedKimiCodeConfig, + }, + }); + + expect(config.providers[KIMI_CODE_PROVIDER_NAME]).toMatchObject({ + baseUrl: 'https://coding.deva.msh.team/coding/v1', + oauth: { + storage: 'file', + key: oauthKey, + oauthHost: 'https://auth.dev.kimi.team', + }, + }); + expect(config.services?.moonshotSearch?.oauth).toEqual({ + storage: 'file', + key: oauthKey, + oauthHost: 'https://auth.dev.kimi.team', + }); + expect(config.services?.moonshotFetch?.oauth).toEqual({ + storage: 'file', + key: oauthKey, + oauthHost: 'https://auth.dev.kimi.team', + }); + }); + + it('persists the default OAuth host when only the API base URL is scoped', async () => { + const config: ManagedKimiConfigShape = { + providers: {}, + }; + const baseUrl = 'https://api.example.test/coding/v1'; + const oauthKey = resolveKimiCodeOAuthKey({ baseUrl }); + + await provisionManagedKimiCodeConfig({ + accessToken: 'oauth-access-token', + baseUrl, + fetchImpl: vi.fn(async () => makeModelsResponse()) as unknown as typeof fetch, + adapter: { + read: () => config, + write: vi.fn(), + apply: applyManagedKimiCodeConfig, + }, + }); + + expect(config.providers[KIMI_CODE_PROVIDER_NAME]).toMatchObject({ + baseUrl, + oauth: { + storage: 'file', + key: oauthKey, + oauthHost: 'https://auth.kimi.com', + }, + }); + }); + it('preserves an existing valid default model during refresh', async () => { const config: ManagedKimiConfigShape = { providers: { @@ -484,6 +621,62 @@ describe('provisionManagedKimiCodeConfig', () => { ).rejects.toThrow('quota exceeded'); }); + it('classifies model listing 401 responses as OAuth unauthorized', async () => { + const fetchImpl = vi.fn( + async () => + new Response( + JSON.stringify({ + error: { message: 'The API Key appears to be invalid or may have expired.' }, + }), + { + status: 401, + headers: { 'Content-Type': 'application/json' }, + }, + ), + ) as unknown as typeof fetch; + + await expect( + fetchManagedKimiCodeModels({ + accessToken: 'oauth-access-token', + fetchImpl, + }), + ).rejects.toBeInstanceOf(OAuthUnauthorizedError); + }); + + it('classifies membership-check 402 responses as OAuth unauthorized', async () => { + const fetchImpl = vi.fn( + async () => + new Response( + JSON.stringify({ + error: { + message: + "We're unable to verify your membership benefits at this time. Please ensure your membership is active.", + }, + }), + { + status: 402, + headers: { 'Content-Type': 'application/json' }, + }, + ), + ) as unknown as typeof fetch; + + await expect( + fetchManagedKimiCodeModels({ + accessToken: 'oauth-access-token', + baseUrl: 'https://coding.deva.msh.team/coding/v1', + fetchImpl, + }), + ).rejects.toThrow( + "Kimi Code models endpoint https://coding.deva.msh.team/coding/v1 rejected OAuth credentials: We're unable to verify your membership benefits at this time. Please ensure your membership is active.", + ); + await expect( + fetchManagedKimiCodeModels({ + accessToken: 'oauth-access-token', + fetchImpl, + }), + ).rejects.toBeInstanceOf(OAuthUnauthorizedError); + }); + it('clears managed provider, models, default model, and services on logout', () => { const config: ManagedKimiConfigShape = { providers: { diff --git a/packages/oauth/test/toolkit.test.ts b/packages/oauth/test/toolkit.test.ts index 9a3ff1ca28..a50a07e0c7 100644 --- a/packages/oauth/test/toolkit.test.ts +++ b/packages/oauth/test/toolkit.test.ts @@ -3,9 +3,12 @@ import { join } from 'node:path'; import { afterEach, describe, expect, it, vi } from 'vitest'; import { + applyManagedKimiCodeConfig, KIMI_CODE_PROVIDER_NAME, KimiOAuthToolkit, + resolveKimiCodeOAuthKey, resolveKimiTokenStorageName, + type ManagedKimiConfigShape, type TokenInfo, type TokenStorage, } from '../src'; @@ -48,8 +51,31 @@ const TEST_IDENTITY = { afterEach(() => { vi.restoreAllMocks(); + vi.unstubAllGlobals(); }); +function managedModelsResponse(): Response { + return new Response( + JSON.stringify({ + data: [ + { + id: 'kimi-for-coding', + context_length: 262144, + supports_reasoning: true, + }, + ], + }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); +} + +function fetchInputUrl(input: unknown): string { + if (typeof input === 'string') return input; + if (input instanceof URL) return input.href; + if (input instanceof Request) return input.url; + throw new TypeError('expected fetch input to be a string, URL, or Request'); +} + describe('resolveKimiTokenStorageName', () => { it('maps config oauth keys to the file storage token name', () => { expect( @@ -118,6 +144,108 @@ describe('KimiOAuthToolkit', () => { ).resolves.toBe('custom-access'); }); + it('refreshes configured bearer token refs against their OAuth host', async () => { + const storage = new MemoryTokenStorage(); + const oauthHost = 'https://auth.dev.kimi.team'; + const oauthKey = resolveKimiCodeOAuthKey({ + oauthHost, + baseUrl: 'https://coding.deva.msh.team/coding/v1', + }); + storage.tokens.set(resolveKimiTokenStorageName({ oauthKey }), { + ...token('expired-dev-access'), + expiresAt: 100, + }); + const fetchImpl = vi.fn(async (input: unknown, init?: RequestInit) => { + expect(fetchInputUrl(input)).toBe(`${oauthHost}/api/oauth/token`); + if (typeof init?.body !== 'string') throw new TypeError('expected form body'); + expect(new URLSearchParams(init.body).get('grant_type')).toBe('refresh_token'); + return new Response( + JSON.stringify({ + access_token: 'rotated-dev-access', + refresh_token: 'rotated-dev-refresh', + expires_in: 3600, + scope: '', + token_type: 'Bearer', + }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); + }); + vi.stubGlobal('fetch', fetchImpl); + const toolkit = new KimiOAuthToolkit({ + homeDir: join('/tmp', 'kimi-oauth-toolkit-test'), + identity: TEST_IDENTITY, + storage, + now: () => 1_000, + flowConfig: { + name: 'kimi-code', + oauthHost: 'https://auth.kimi.com', + clientId: 'test-client-id', + }, + }); + + await expect( + toolkit + .tokenProvider(KIMI_CODE_PROVIDER_NAME, { key: oauthKey, oauthHost }) + .getAccessToken(), + ).resolves.toBe('rotated-dev-access'); + }); + + it('does not reuse a cached OAuth manager across different hosts for the same token key', async () => { + const storage = new MemoryTokenStorage(); + storage.tokens.set('custom-kimi-code', { + ...token('expired-custom-access'), + expiresAt: 100, + }); + const requests: string[] = []; + const fetchImpl = vi.fn(async (input: unknown) => { + requests.push(fetchInputUrl(input)); + return new Response( + JSON.stringify({ + access_token: `rotated-${String(requests.length)}`, + refresh_token: `refresh-${String(requests.length)}`, + expires_in: 3600, + scope: '', + token_type: 'Bearer', + }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); + }); + vi.stubGlobal('fetch', fetchImpl); + const toolkit = new KimiOAuthToolkit({ + homeDir: join('/tmp', 'kimi-oauth-toolkit-test'), + identity: TEST_IDENTITY, + storage, + now: () => 1_000, + flowConfig: { + name: 'kimi-code', + oauthHost: 'https://auth.kimi.com', + clientId: 'test-client-id', + }, + }); + + await expect( + toolkit + .tokenProvider(KIMI_CODE_PROVIDER_NAME, { + key: 'oauth/custom-kimi-code', + oauthHost: 'https://auth.one.test/', + }) + .getAccessToken({ force: true }), + ).resolves.toBe('rotated-1'); + await expect( + toolkit + .tokenProvider(KIMI_CODE_PROVIDER_NAME, { + key: 'oauth/custom-kimi-code', + oauthHost: 'https://auth.two.test', + }) + .getAccessToken({ force: true }), + ).resolves.toBe('rotated-2'); + + expect(requests).toEqual([ + 'https://auth.one.test/api/oauth/token', + 'https://auth.two.test/api/oauth/token', + ]); + }); + it('returns the cached access token without refreshing it', async () => { const storage = new MemoryTokenStorage(); storage.tokens.set('kimi-code', { @@ -163,21 +291,7 @@ describe('KimiOAuthToolkit', () => { it('provisions managed config after login when an adapter is configured', async () => { const storage = new MemoryTokenStorage(); const write = vi.fn(); - const fetchImpl = vi.fn( - async () => - new Response( - JSON.stringify({ - data: [ - { - id: 'kimi-for-coding', - context_length: 262144, - supports_reasoning: true, - }, - ], - }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, - ), - ) as unknown as typeof fetch; + const fetchImpl = vi.fn(async () => managedModelsResponse()) as unknown as typeof fetch; const config = { providers: {} }; const toolkit = new KimiOAuthToolkit({ homeDir: join('/tmp', 'kimi-oauth-toolkit-test'), @@ -212,9 +326,179 @@ describe('KimiOAuthToolkit', () => { expect(write).toHaveBeenCalledWith(config); }); + it.each([401, 402])( + 'force-refreshes a stored token when managed model provisioning rejects cached auth with HTTP %i', + async (status) => { + const storage = new MemoryTokenStorage(); + const write = vi.fn(); + const onDeviceCode = vi.fn(); + const config = { providers: {} }; + const oauthHost = 'https://auth.test'; + const oauthKey = resolveKimiCodeOAuthKey({ oauthHost }); + storage.tokens.set(resolveKimiTokenStorageName({ oauthKey }), token('stale-access')); + const fetchMock = vi + .fn() + .mockResolvedValueOnce( + new Response( + JSON.stringify({ + error: { message: 'The API Key appears to be invalid or may have expired.' }, + }), + { status, headers: { 'Content-Type': 'application/json' } }, + ), + ) + .mockResolvedValueOnce(managedModelsResponse()); + const fetchImpl = fetchMock as unknown as typeof fetch; + const oauthFetch = vi.fn(async (_input: unknown, init?: RequestInit) => { + if (typeof init?.body !== 'string') throw new TypeError('expected form body'); + const body = new URLSearchParams(init.body); + if (body.get('grant_type') !== 'refresh_token') { + throw new Error(`unexpected OAuth grant: ${body.get('grant_type') ?? ''}`); + } + return new Response( + JSON.stringify({ + access_token: 'rotated-access', + refresh_token: 'rotated-refresh', + expires_in: 3600, + scope: '', + token_type: 'Bearer', + }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); + }); + vi.stubGlobal('fetch', oauthFetch); + const toolkit = new KimiOAuthToolkit({ + homeDir: join('/tmp', 'kimi-oauth-toolkit-test'), + identity: TEST_IDENTITY, + storage, + now: () => 100, + fetchImpl, + flowConfig: { + name: 'kimi-code', + oauthHost, + clientId: 'test-client-id', + }, + configAdapter: { + read: () => config, + write, + apply: (target, input) => { + target.providers[KIMI_CODE_PROVIDER_NAME] = { + type: 'kimi', + apiKey: '', + }; + return { + defaultModel: `kimi-code/${input.models[0]?.id ?? 'unknown'}`, + defaultThinking: true, + }; + }, + }, + }); + + await expect(toolkit.login(undefined, { onDeviceCode })).resolves.toMatchObject({ + providerName: KIMI_CODE_PROVIDER_NAME, + ok: true, + provision: { + defaultModel: 'kimi-code/kimi-for-coding', + }, + }); + expect(fetchMock).toHaveBeenCalledTimes(2); + const firstModelRequest = fetchMock.mock.calls[0]?.[1] as RequestInit | undefined; + const secondModelRequest = fetchMock.mock.calls[1]?.[1] as RequestInit | undefined; + expect(new Headers(firstModelRequest?.headers).get('authorization')).toBe( + 'Bearer stale-access', + ); + expect(new Headers(secondModelRequest?.headers).get('authorization')).toBe( + 'Bearer rotated-access', + ); + expect(oauthFetch).toHaveBeenCalledTimes(1); + expect(onDeviceCode).not.toHaveBeenCalled(); + expect(write).toHaveBeenCalledWith(config); + }, + ); + + it('uses a scoped credential slot for non-default OAuth login environments', async () => { + const storage = new MemoryTokenStorage(); + storage.tokens.set('kimi-code', token('prod-access')); + const config: ManagedKimiConfigShape = { providers: {} }; + const devBaseUrl = 'https://coding.deva.msh.team/coding/v1'; + const devOauthHost = 'https://auth.dev.kimi.team'; + const devOauthKey = resolveKimiCodeOAuthKey({ + oauthHost: devOauthHost, + baseUrl: devBaseUrl, + }); + const devStorageName = resolveKimiTokenStorageName({ oauthKey: devOauthKey }); + const write = vi.fn(); + const fetchMock = vi.fn(async (_input: unknown, _init?: RequestInit) => + managedModelsResponse(), + ); + const oauthFetch = vi.fn(async (_input: unknown, init?: RequestInit) => { + if (typeof init?.body !== 'string') throw new TypeError('expected form body'); + const body = new URLSearchParams(init.body); + if (body.get('grant_type') === 'urn:ietf:params:oauth:grant-type:device_code') { + return new Response( + JSON.stringify({ + access_token: 'dev-access', + refresh_token: 'dev-refresh', + expires_in: 3600, + scope: '', + token_type: 'Bearer', + }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); + } + return new Response( + JSON.stringify({ + user_code: 'WDJB-MJHT', + device_code: 'device-code', + verification_uri: `${devOauthHost}/verify`, + verification_uri_complete: `${devOauthHost}/verify?user_code=WDJB-MJHT`, + expires_in: 600, + interval: 1, + }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); + }); + vi.stubGlobal('fetch', oauthFetch); + const toolkit = new KimiOAuthToolkit({ + homeDir: join('/tmp', 'kimi-oauth-toolkit-test'), + identity: TEST_IDENTITY, + storage, + now: () => 100, + fetchImpl: fetchMock as unknown as typeof fetch, + flowConfig: { + name: 'kimi-code', + oauthHost: devOauthHost, + clientId: 'test-client-id', + }, + configAdapter: { + read: () => config, + write, + apply: applyManagedKimiCodeConfig, + }, + }); + + await expect(toolkit.login(undefined, { baseUrl: devBaseUrl })).resolves.toMatchObject({ + providerName: KIMI_CODE_PROVIDER_NAME, + ok: true, + }); + expect(oauthFetch).toHaveBeenCalledTimes(2); + expect(storage.tokens.get('kimi-code')?.accessToken).toBe('prod-access'); + expect(storage.tokens.get(devStorageName)?.accessToken).toBe('dev-access'); + expect(config.providers[KIMI_CODE_PROVIDER_NAME]?.oauth).toEqual({ + storage: 'file', + key: devOauthKey, + oauthHost: devOauthHost, + }); + const modelRequest = fetchMock.mock.calls[0]?.[1] as RequestInit | undefined; + expect(new Headers(modelRequest?.headers).get('authorization')).toBe('Bearer dev-access'); + expect(write).toHaveBeenCalledWith(config); + }); + it('starts a new device flow when the stored refresh token is invalid', async () => { const storage = new MemoryTokenStorage(); - storage.tokens.set('kimi-code', { + const oauthHost = 'https://auth.test'; + const oauthKey = resolveKimiCodeOAuthKey({ oauthHost }); + const storageName = resolveKimiTokenStorageName({ oauthKey }); + storage.tokens.set(storageName, { ...token('stale-access'), refreshToken: 'revoked-refresh', expiresAt: 101, @@ -264,7 +548,7 @@ describe('KimiOAuthToolkit', () => { now: () => 100, flowConfig: { name: 'kimi-code', - oauthHost: 'https://auth.test', + oauthHost, clientId: 'test-client-id', }, }); @@ -274,7 +558,7 @@ describe('KimiOAuthToolkit', () => { ok: true, }); expect(onDeviceCode).toHaveBeenCalledTimes(1); - expect((await storage.load('kimi-code'))?.accessToken).toBe('fresh-access'); + expect((await storage.load(storageName))?.accessToken).toBe('fresh-access'); }); it('removes managed config on logout when an adapter supports cleanup', async () => { From 78f5c1adca10173637feb06235fa238e91fa10ad Mon Sep 17 00:00:00 2001 From: 7Sageer <7sageer@djwcb.cn> Date: Thu, 4 Jun 2026 12:56:19 +0800 Subject: [PATCH 2/2] refactor: centralize managed OAuth auth resolution --- .../src/tui/utils/refresh-providers.ts | 29 ++-- packages/node-sdk/src/auth.ts | 95 +++--------- packages/oauth/src/index.ts | 9 ++ packages/oauth/src/managed-kimi-code.ts | 136 +++++++++++++++++- packages/oauth/test/managed-kimi-code.test.ts | 105 +++++++++++++- 5 files changed, 270 insertions(+), 104 deletions(-) diff --git a/apps/kimi-code/src/tui/utils/refresh-providers.ts b/apps/kimi-code/src/tui/utils/refresh-providers.ts index 35292ac281..2f1a18e15f 100644 --- a/apps/kimi-code/src/tui/utils/refresh-providers.ts +++ b/apps/kimi-code/src/tui/utils/refresh-providers.ts @@ -10,10 +10,9 @@ import { filterModelsByPrefix, getOpenPlatformById, isOpenPlatformId, - kimiCodeBaseUrl, removeCustomRegistryProvider, removeOpenPlatformConfig, - resolveKimiCodeOAuthRef, + resolveKimiCodeRuntimeAuth, type CustomRegistrySource, type ManagedKimiConfigShape, } from '@moonshot-ai/kimi-code-oauth'; @@ -58,15 +57,6 @@ function asManaged(config: KimiConfig): ManagedKimiConfigShape { return config as unknown as ManagedKimiConfigShape; } -function managedRuntimeBaseUrl(configuredBaseUrl: string | undefined): string | undefined { - if (process.env['KIMI_CODE_BASE_URL'] !== undefined) return kimiCodeBaseUrl(); - return configuredBaseUrl; -} - -function envOAuthHost(): string | undefined { - return process.env['KIMI_CODE_OAUTH_HOST'] ?? process.env['KIMI_OAUTH_HOST']; -} - function collectModelIdsForProvider(config: KimiConfig, providerId: string): Set { const ids = new Set(); for (const alias of Object.values(config.models ?? {})) { @@ -131,15 +121,14 @@ export async function refreshAllProviderModels(host: RefreshProviderHost): Promi managedProvider.oauth !== undefined ) { try { - const baseUrl = managedRuntimeBaseUrl(managedProvider.baseUrl); - const oauth = resolveKimiCodeOAuthRef({ - oauthHost: envOAuthHost() ?? managedProvider.oauth.oauthHost, - baseUrl, + const auth = resolveKimiCodeRuntimeAuth({ + configuredBaseUrl: managedProvider.baseUrl, + configuredOAuthRef: managedProvider.oauth, }); - const accessToken = await host.resolveOAuthToken(KIMI_CODE_PROVIDER_NAME, oauth); + const accessToken = await host.resolveOAuthToken(KIMI_CODE_PROVIDER_NAME, auth.oauthRef); const models = await fetchManagedKimiCodeModels({ accessToken, - baseUrl, + baseUrl: auth.baseUrl, }); if (models.length > 0) { const beforeIds = collectModelIdsForProvider(config, KIMI_CODE_PROVIDER_NAME); @@ -153,9 +142,9 @@ export async function refreshAllProviderModels(host: RefreshProviderHost): Promi clearManagedKimiCodeConfig(asManaged(config)); applyManagedKimiCodeConfig(asManaged(config), { models, - baseUrl, - oauthKey: oauth.key, - oauthHost: oauth.oauthHost, + baseUrl: auth.baseUrl, + oauthKey: auth.oauthRef.key, + oauthHost: auth.oauthRef.oauthHost, preserveDefaultModel: true, }); await host.setConfig({ diff --git a/packages/node-sdk/src/auth.ts b/packages/node-sdk/src/auth.ts index 4a6ce30eb0..5443464a2e 100644 --- a/packages/node-sdk/src/auth.ts +++ b/packages/node-sdk/src/auth.ts @@ -4,8 +4,8 @@ import { applyManagedKimiCodeLogoutConfig, KIMI_CODE_PROVIDER_NAME, KimiOAuthToolkit, - resolveKimiCodeOAuthKey, - resolveKimiCodeOAuthRef, + resolveKimiCodeLoginAuth, + resolveKimiCodeRuntimeAuth, type AuthManagedUsageResult, type AuthStatus, type BearerTokenProvider, @@ -78,13 +78,17 @@ export class KimiAuthFacade { options: KimiAuthLoginOptions = {}, ): Promise { const auth = this.resolveManagedAuth(providerName); - const baseUrl = options.baseUrl ?? this.configBaseUrlForLogin(auth.baseUrl); - const oauthHost = options.oauthHost ?? this.envOAuthHost(); + const loginAuth = resolveKimiCodeLoginAuth({ + configuredBaseUrl: auth.baseUrl, + configuredOAuthRef: auth.oauthRef, + requestedBaseUrl: options.baseUrl, + requestedOAuthHost: options.oauthHost, + }); const result = await this.toolkit.login(providerName, { ...options, - baseUrl, - oauthHost, - oauthRef: options.oauthRef ?? this.configOAuthRefForLogin(auth.oauthRef, options, baseUrl), + baseUrl: loginAuth.baseUrl, + oauthHost: loginAuth.oauthHost, + oauthRef: options.oauthRef ?? loginAuth.oauthRef, provisionConfig: true, }); if (result.provision === undefined) { @@ -173,57 +177,15 @@ export class KimiAuthFacade { }; } - private configBaseUrlForLogin(baseUrl?: string | undefined): string | undefined { - return process.env['KIMI_CODE_BASE_URL'] ?? baseUrl; - } - - private configOAuthRefForLogin( - oauthRef: OAuthRef | undefined, - options: KimiAuthLoginOptions, - baseUrl?: string | undefined, - ): OAuthRef | undefined { - if ( - options.baseUrl !== undefined || - options.oauthHost !== undefined || - process.env['KIMI_CODE_BASE_URL'] !== undefined || - process.env['KIMI_CODE_OAUTH_HOST'] !== undefined || - process.env['KIMI_OAUTH_HOST'] !== undefined - ) { - return undefined; - } - if (oauthRef !== undefined && oauthRef.key !== this.expectedOAuthKey(oauthRef, baseUrl)) { - return undefined; - } - return oauthRef; - } - - private configBaseUrlForRuntime(baseUrl?: string | undefined): string | undefined { - return process.env['KIMI_CODE_BASE_URL'] ?? baseUrl; - } - - private configOAuthRefForRuntime( - oauthRef: OAuthRef | undefined, - baseUrl?: string | undefined, - ): OAuthRef { - const expected = this.expectedOAuthRef(oauthRef, baseUrl); - if (oauthRef === undefined) return expected; - if (process.env['KIMI_CODE_BASE_URL'] !== undefined || this.envOAuthHost() !== undefined) { - return expected; - } - if (oauthRef.key !== expected.key) return expected; - return oauthRef; - } - private resolveRuntimeManagedAuth(providerName?: string | undefined): { readonly oauthRef: OAuthRef; readonly baseUrl?: string | undefined; } { const auth = this.resolveManagedAuth(providerName); - const baseUrl = this.configBaseUrlForRuntime(auth.baseUrl); - return { - oauthRef: this.configOAuthRefForRuntime(auth.oauthRef, baseUrl), - baseUrl, - }; + return resolveKimiCodeRuntimeAuth({ + configuredBaseUrl: auth.baseUrl, + configuredOAuthRef: auth.oauthRef, + }); } private runtimeOAuthRef( @@ -232,28 +194,9 @@ export class KimiAuthFacade { ): OAuthRef | undefined { if ((providerName ?? KIMI_CODE_PROVIDER_NAME) !== KIMI_CODE_PROVIDER_NAME) return oauthRef; const auth = this.resolveManagedAuth(providerName); - const baseUrl = this.configBaseUrlForRuntime(auth.baseUrl); - return this.configOAuthRefForRuntime(oauthRef ?? auth.oauthRef, baseUrl); - } - - private expectedOAuthKey(oauthRef: OAuthRef, baseUrl?: string | undefined): string { - return resolveKimiCodeOAuthKey({ - oauthHost: oauthRef.oauthHost, - baseUrl, - }); - } - - private expectedOAuthRef(oauthRef: OAuthRef | undefined, baseUrl?: string | undefined): OAuthRef { - const envOAuthHost = this.envOAuthHost(); - const hasEnvOverride = - process.env['KIMI_CODE_BASE_URL'] !== undefined || envOAuthHost !== undefined; - return resolveKimiCodeOAuthRef({ - oauthHost: hasEnvOverride ? envOAuthHost : oauthRef?.oauthHost, - baseUrl, - }); - } - - private envOAuthHost(): string | undefined { - return process.env['KIMI_CODE_OAUTH_HOST'] ?? process.env['KIMI_OAUTH_HOST']; + return resolveKimiCodeRuntimeAuth({ + configuredBaseUrl: auth.baseUrl, + configuredOAuthRef: oauthRef ?? auth.oauthRef, + }).oauthRef; } } diff --git a/packages/oauth/src/index.ts b/packages/oauth/src/index.ts index 84b95ae6d0..02b9e2d873 100644 --- a/packages/oauth/src/index.ts +++ b/packages/oauth/src/index.ts @@ -42,22 +42,31 @@ export { applyManagedKimiCodeConfig, clearManagedKimiCodeConfig, fetchManagedKimiCodeModels, + kimiCodeEnvBaseUrl, + kimiCodeEnvOAuthHost, KIMI_CODE_OAUTH_KEY, KIMI_CODE_PLATFORM_ID, KIMI_CODE_PROVIDER_NAME, + ManagedKimiCodeModelsAuthError, provisionManagedKimiCodeConfig, + resolveKimiCodeLoginAuth, resolveKimiCodeOAuthKey, resolveKimiCodeOAuthRef, + resolveKimiCodeRuntimeAuth, } from './managed-kimi-code'; export type { FetchManagedKimiCodeModelsOptions, ManagedKimiCodeApplyResult, ManagedKimiCodeCleanupResult, + ManagedKimiEnv, + ManagedKimiLoginAuth, ManagedKimiCodeModelInfo, ManagedKimiCodeProvisionResult, ManagedKimiConfigAdapter, ManagedKimiConfigShape, ManagedKimiOAuthRef, + ManagedKimiOAuthRefInput, + ManagedKimiRuntimeAuth, ProvisionManagedKimiCodeConfigOptions, } from './managed-kimi-code'; diff --git a/packages/oauth/src/managed-kimi-code.ts b/packages/oauth/src/managed-kimi-code.ts index 3b7a3a0b61..ee612ea631 100644 --- a/packages/oauth/src/managed-kimi-code.ts +++ b/packages/oauth/src/managed-kimi-code.ts @@ -49,11 +49,52 @@ export interface ManagedKimiCodeCleanupResult { } export interface ManagedKimiOAuthRef { - readonly storage: 'file'; + readonly storage: 'file' | 'keyring'; readonly key: string; readonly oauthHost?: string | undefined; } +export interface ManagedKimiOAuthRefInput { + readonly storage?: 'file' | 'keyring' | undefined; + readonly key?: string | undefined; + readonly oauthHost?: string | undefined; +} + +export interface ManagedKimiRuntimeAuth { + readonly baseUrl?: string | undefined; + readonly oauthRef: ManagedKimiOAuthRef; +} + +export interface ManagedKimiLoginAuth { + readonly baseUrl?: string | undefined; + readonly oauthHost?: string | undefined; + readonly oauthRef?: ManagedKimiOAuthRef | undefined; +} + +export interface ManagedKimiEnv { + readonly KIMI_CODE_BASE_URL?: string | undefined; + readonly KIMI_CODE_OAUTH_HOST?: string | undefined; + readonly KIMI_OAUTH_HOST?: string | undefined; +} + +export class ManagedKimiCodeModelsAuthError extends OAuthUnauthorizedError { + readonly status: number; + readonly baseUrl: string; + + constructor(options: { + readonly status: number; + readonly baseUrl: string; + readonly message: string; + }) { + super( + `Kimi Code models endpoint ${options.baseUrl} rejected OAuth credentials: ${options.message}`, + ); + this.name = 'ManagedKimiCodeModelsAuthError'; + this.status = options.status; + this.baseUrl = options.baseUrl; + } +} + export interface ManagedKimiProviderConfig { type: 'kimi'; baseUrl?: string | undefined; @@ -141,6 +182,10 @@ function defaultBaseUrl(baseUrl: string | undefined): string { return (baseUrl ?? kimiCodeBaseUrl()).replace(/\/+$/, ''); } +function normalizeBaseUrl(baseUrl: string): string { + return baseUrl.replace(/\/+$/, ''); +} + function normalizeEndpoint(value: string): string { return value.trim().replace(/\/+$/, ''); } @@ -163,15 +208,37 @@ function persistedOAuthHost(options: { function managedOAuthRef(options: { readonly key: string; readonly oauthHost?: string | undefined; + readonly storage?: 'file' | 'keyring' | undefined; }): ManagedKimiOAuthRef { const oauthHost = persistedOAuthHost(options); return { - storage: 'file', + storage: options.storage ?? 'file', key: options.key, oauthHost, }; } +function configuredOAuthRef( + oauthRef: ManagedKimiOAuthRefInput | undefined, +): ManagedKimiOAuthRef | undefined { + if (oauthRef === undefined) return undefined; + const key = oauthRef.key; + if (key === undefined) return undefined; + return managedOAuthRef({ + storage: oauthRef.storage, + key, + oauthHost: oauthRef.oauthHost, + }); +} + +export function kimiCodeEnvBaseUrl(env: ManagedKimiEnv = process.env): string | undefined { + return env.KIMI_CODE_BASE_URL; +} + +export function kimiCodeEnvOAuthHost(env: ManagedKimiEnv = process.env): string | undefined { + return env.KIMI_CODE_OAUTH_HOST ?? env.KIMI_OAUTH_HOST; +} + export function resolveKimiCodeOAuthKey(options: { readonly oauthHost?: string | undefined; readonly baseUrl?: string | undefined; @@ -212,6 +279,63 @@ export function resolveKimiCodeOAuthRef(options: { }); } +export function resolveKimiCodeRuntimeAuth(options: { + readonly configuredBaseUrl?: string | undefined; + readonly configuredOAuthRef?: ManagedKimiOAuthRefInput | undefined; + readonly env?: ManagedKimiEnv | undefined; +}): ManagedKimiRuntimeAuth { + const env = options.env ?? process.env; + const envBaseUrl = kimiCodeEnvBaseUrl(env); + const envOAuthHost = kimiCodeEnvOAuthHost(env); + const hasEnvOverride = envBaseUrl !== undefined || envOAuthHost !== undefined; + const baseUrl = + envBaseUrl !== undefined ? normalizeBaseUrl(envBaseUrl) : options.configuredBaseUrl; + const expected = resolveKimiCodeOAuthRef({ + oauthHost: hasEnvOverride ? envOAuthHost : options.configuredOAuthRef?.oauthHost, + baseUrl, + }); + const configured = configuredOAuthRef(options.configuredOAuthRef); + if (configured === undefined) return { baseUrl, oauthRef: expected }; + if (hasEnvOverride) return { baseUrl, oauthRef: expected }; + if (configured.key !== expected.key) return { baseUrl, oauthRef: expected }; + return { baseUrl, oauthRef: configured }; +} + +export function resolveKimiCodeLoginAuth(options: { + readonly configuredBaseUrl?: string | undefined; + readonly configuredOAuthRef?: ManagedKimiOAuthRefInput | undefined; + readonly requestedBaseUrl?: string | undefined; + readonly requestedOAuthHost?: string | undefined; + readonly env?: ManagedKimiEnv | undefined; +}): ManagedKimiLoginAuth { + const env = options.env ?? process.env; + const envBaseUrl = kimiCodeEnvBaseUrl(env); + const envOAuthHost = kimiCodeEnvOAuthHost(env); + const hasOverride = + options.requestedBaseUrl !== undefined || + options.requestedOAuthHost !== undefined || + envBaseUrl !== undefined || + envOAuthHost !== undefined; + const baseUrl = + options.requestedBaseUrl !== undefined + ? normalizeBaseUrl(options.requestedBaseUrl) + : envBaseUrl !== undefined + ? normalizeBaseUrl(envBaseUrl) + : options.configuredBaseUrl; + const oauthHost = options.requestedOAuthHost ?? envOAuthHost; + if (hasOverride) return { baseUrl, oauthHost }; + + const configured = configuredOAuthRef(options.configuredOAuthRef); + if (configured === undefined) return { baseUrl, oauthHost }; + const expectedKey = resolveKimiCodeOAuthKey({ + oauthHost: configured.oauthHost, + baseUrl, + }); + return configured.key === expectedKey + ? { baseUrl, oauthHost, oauthRef: configured } + : { baseUrl, oauthHost }; +} + function toModelInfo(item: unknown): ManagedKimiCodeModelInfo | undefined { if (!isRecord(item) || typeof item['id'] !== 'string' || item['id'].length === 0) { return undefined; @@ -254,9 +378,11 @@ export async function fetchManagedKimiCodeModels( `Failed to list Kimi Code models (HTTP ${response.status}).`, ); if (response.status === 401 || response.status === 402 || response.status === 403) { - throw new OAuthUnauthorizedError( - `Kimi Code models endpoint ${baseUrl} rejected OAuth credentials: ${message}`, - ); + throw new ManagedKimiCodeModelsAuthError({ + status: response.status, + baseUrl, + message, + }); } throw new Error(message); } diff --git a/packages/oauth/test/managed-kimi-code.test.ts b/packages/oauth/test/managed-kimi-code.test.ts index f09f681c3d..ea1d4f1def 100644 --- a/packages/oauth/test/managed-kimi-code.test.ts +++ b/packages/oauth/test/managed-kimi-code.test.ts @@ -7,9 +7,12 @@ import { fetchManagedKimiCodeModels, KIMI_CODE_OAUTH_KEY, KIMI_CODE_PROVIDER_NAME, + ManagedKimiCodeModelsAuthError, provisionManagedKimiCodeConfig, + resolveKimiCodeLoginAuth, resolveKimiCodeOAuthKey, resolveKimiCodeOAuthRef, + resolveKimiCodeRuntimeAuth, type ManagedKimiConfigShape, } from '../src/managed-kimi-code'; import { OAuthUnauthorizedError } from '../src/errors'; @@ -104,6 +107,86 @@ describe('provisionManagedKimiCodeConfig', () => { }); }); + it('resolves runtime auth from environment overrides over persisted config', () => { + const configuredBaseUrl = 'https://api.configured.example.test/coding/v1'; + const envBaseUrl = 'https://api.env.example.test/coding/v1/'; + const envOauthHost = 'https://auth.env.example.test/'; + const configuredOAuthRef = resolveKimiCodeOAuthRef({ + baseUrl: configuredBaseUrl, + }); + + const auth = resolveKimiCodeRuntimeAuth({ + configuredBaseUrl, + configuredOAuthRef, + env: { + KIMI_CODE_BASE_URL: envBaseUrl, + KIMI_CODE_OAUTH_HOST: envOauthHost, + }, + }); + + expect(auth.baseUrl).toBe('https://api.env.example.test/coding/v1'); + expect(auth.oauthRef).toEqual({ + storage: 'file', + key: resolveKimiCodeOAuthKey({ + oauthHost: 'https://auth.env.example.test', + baseUrl: 'https://api.env.example.test/coding/v1', + }), + oauthHost: 'https://auth.env.example.test', + }); + }); + + it('preserves a matching configured runtime OAuth ref when env is not overridden', () => { + const baseUrl = 'https://coding.deva.msh.team/coding/v1'; + const configuredOAuthRef = { + storage: 'keyring' as const, + key: resolveKimiCodeOAuthKey({ + oauthHost: 'https://auth.dev.kimi.team', + baseUrl, + }), + oauthHost: 'https://auth.dev.kimi.team', + }; + + expect( + resolveKimiCodeRuntimeAuth({ + configuredBaseUrl: baseUrl, + configuredOAuthRef, + env: {}, + }), + ).toEqual({ + baseUrl, + oauthRef: configuredOAuthRef, + }); + }); + + it('resolves login auth without reusing persisted refs under explicit or env overrides', () => { + const configuredBaseUrl = 'https://api.configured.example.test/coding/v1'; + const configuredOAuthRef = resolveKimiCodeOAuthRef({ baseUrl: configuredBaseUrl }); + + expect( + resolveKimiCodeLoginAuth({ + configuredBaseUrl, + configuredOAuthRef, + requestedBaseUrl: 'https://api.requested.example.test/coding/v1/', + env: {}, + }), + ).toEqual({ + baseUrl: 'https://api.requested.example.test/coding/v1', + oauthHost: undefined, + }); + + expect( + resolveKimiCodeLoginAuth({ + configuredBaseUrl, + configuredOAuthRef, + env: {}, + }), + ).toEqual({ + baseUrl: configuredBaseUrl, + oauthHost: undefined, + oauthRef: configuredOAuthRef, + }); + }); + it('writes the managed provider, models, services, and default model through an adapter', async () => { const config: ManagedKimiConfigShape = { providers: { @@ -660,21 +743,37 @@ describe('provisionManagedKimiCodeConfig', () => { ), ) as unknown as typeof fetch; + const promise = fetchManagedKimiCodeModels({ + accessToken: 'oauth-access-token', + baseUrl: 'https://coding.deva.msh.team/coding/v1', + fetchImpl, + }); + + await expect(promise).rejects.toThrow( + "Kimi Code models endpoint https://coding.deva.msh.team/coding/v1 rejected OAuth credentials: We're unable to verify your membership benefits at this time. Please ensure your membership is active.", + ); await expect( fetchManagedKimiCodeModels({ accessToken: 'oauth-access-token', baseUrl: 'https://coding.deva.msh.team/coding/v1', fetchImpl, }), - ).rejects.toThrow( - "Kimi Code models endpoint https://coding.deva.msh.team/coding/v1 rejected OAuth credentials: We're unable to verify your membership benefits at this time. Please ensure your membership is active.", - ); + ).rejects.toMatchObject({ + status: 402, + baseUrl: 'https://coding.deva.msh.team/coding/v1', + }); await expect( fetchManagedKimiCodeModels({ accessToken: 'oauth-access-token', fetchImpl, }), ).rejects.toBeInstanceOf(OAuthUnauthorizedError); + await expect( + fetchManagedKimiCodeModels({ + accessToken: 'oauth-access-token', + fetchImpl, + }), + ).rejects.toBeInstanceOf(ManagedKimiCodeModelsAuthError); }); it('clears managed provider, models, default model, and services on logout', () => {