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..2f1a18e15f 100644 --- a/apps/kimi-code/src/tui/utils/refresh-providers.ts +++ b/apps/kimi-code/src/tui/utils/refresh-providers.ts @@ -12,6 +12,7 @@ import { isOpenPlatformId, removeCustomRegistryProvider, removeOpenPlatformConfig, + resolveKimiCodeRuntimeAuth, type CustomRegistrySource, type ManagedKimiConfigShape, } from '@moonshot-ai/kimi-code-oauth'; @@ -120,13 +121,14 @@ export async function refreshAllProviderModels(host: RefreshProviderHost): Promi managedProvider.oauth !== undefined ) { try { - const accessToken = await host.resolveOAuthToken( - KIMI_CODE_PROVIDER_NAME, - managedProvider.oauth, - ); + const auth = resolveKimiCodeRuntimeAuth({ + configuredBaseUrl: managedProvider.baseUrl, + configuredOAuthRef: managedProvider.oauth, + }); + const accessToken = await host.resolveOAuthToken(KIMI_CODE_PROVIDER_NAME, auth.oauthRef); const models = await fetchManagedKimiCodeModels({ accessToken, - baseUrl: managedProvider.baseUrl, + baseUrl: auth.baseUrl, }); if (models.length > 0) { const beforeIds = collectModelIdsForProvider(config, KIMI_CODE_PROVIDER_NAME); @@ -140,7 +142,9 @@ export async function refreshAllProviderModels(host: RefreshProviderHost): Promi clearManagedKimiCodeConfig(asManaged(config)); applyManagedKimiCodeConfig(asManaged(config), { models, - baseUrl: managedProvider.baseUrl, + baseUrl: auth.baseUrl, + oauthKey: auth.oauthRef.key, + oauthHost: auth.oauthRef.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..5443464a2e 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, + resolveKimiCodeLoginAuth, + resolveKimiCodeRuntimeAuth, type AuthManagedUsageResult, type AuthStatus, type BearerTokenProvider, @@ -68,14 +70,27 @@ 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 loginAuth = resolveKimiCodeLoginAuth({ + configuredBaseUrl: auth.baseUrl, + configuredOAuthRef: auth.oauthRef, + requestedBaseUrl: options.baseUrl, + requestedOAuthHost: options.oauthHost, + }); + const result = await this.toolkit.login(providerName, { + ...options, + baseUrl: loginAuth.baseUrl, + oauthHost: loginAuth.oauthHost, + oauthRef: options.oauthRef ?? loginAuth.oauthRef, + provisionConfig: true, + }); if (result.provision === undefined) { throw new Error('Kimi auth login did not provision model config.'); } @@ -91,7 +106,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 +119,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 +140,63 @@ 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 resolveRuntimeManagedAuth(providerName?: string | undefined): { + readonly oauthRef: OAuthRef; + readonly baseUrl?: string | undefined; + } { + const auth = this.resolveManagedAuth(providerName); + return resolveKimiCodeRuntimeAuth({ + configuredBaseUrl: auth.baseUrl, + configuredOAuthRef: auth.oauthRef, + }); + } + + 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); + return resolveKimiCodeRuntimeAuth({ + configuredBaseUrl: auth.baseUrl, + configuredOAuthRef: oauthRef ?? auth.oauthRef, + }).oauthRef; + } } 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..02b9e2d873 100644 --- a/packages/oauth/src/index.ts +++ b/packages/oauth/src/index.ts @@ -42,19 +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 caca34c3bc..ee612ea631 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; @@ -44,8 +49,50 @@ export interface ManagedKimiCodeCleanupResult { } export interface ManagedKimiOAuthRef { - readonly storage: 'file'; - readonly key: typeof KIMI_CODE_OAUTH_KEY; + 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 { @@ -94,6 +141,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 +154,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 +182,160 @@ 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(/\/+$/, ''); +} + +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; + readonly storage?: 'file' | 'keyring' | undefined; +}): ManagedKimiOAuthRef { + const oauthHost = persistedOAuthHost(options); + return { + 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; +}): 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, + }); +} + +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; @@ -168,12 +373,18 @@ 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 ManagedKimiCodeModelsAuthError({ + status: response.status, + baseUrl, + message, + }); + } + throw new Error(message); } const payload: unknown = await response.json(); if (!isRecord(payload) || !Array.isArray(payload['data'])) { @@ -189,6 +400,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 +413,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 +426,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 +452,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 +605,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..ea1d4f1def 100644 --- a/packages/oauth/test/managed-kimi-code.test.ts +++ b/packages/oauth/test/managed-kimi-code.test.ts @@ -5,10 +5,17 @@ import { applyManagedKimiCodeConfig, clearManagedKimiCodeConfig, 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'; function makeModelsResponse(): Response { return new Response( @@ -37,6 +44,149 @@ 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('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: { @@ -122,6 +272,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 +704,78 @@ 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; + + 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.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', () => { 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 () => {