diff --git a/.changeset/v2-login-env-aware-auth.md b/.changeset/v2-login-env-aware-auth.md new file mode 100644 index 0000000000..fc1297a3e2 --- /dev/null +++ b/.changeset/v2-login-env-aware-auth.md @@ -0,0 +1,5 @@ +--- +"@moonshot-ai/agent-core-v2": patch +--- + +Fix v2 managed OAuth login ignoring `KIMI_CODE_BASE_URL` / `KIMI_CODE_OAUTH_HOST`: the login environment is now resolved env-aware (v1 parity), so the credential slot a token is written to always matches the slot the runtime reads — no more "login succeeds but every call 401s" against non-default environments. The provisioned provider entry records the login environment and credential slot explicitly, and logout deletes from the runtime (env-aware) slot. diff --git a/packages/agent-core-v2/src/app/auth/authService.ts b/packages/agent-core-v2/src/app/auth/authService.ts index 01a2b6d5d5..9c61f1f5d2 100644 --- a/packages/agent-core-v2/src/app/auth/authService.ts +++ b/packages/agent-core-v2/src/app/auth/authService.ts @@ -24,6 +24,7 @@ import { applyManagedKimiCodeConfig, clearManagedKimiCodeConfig, fetchManagedKimiCodeModels, + resolveKimiCodeLoginAuth, resolveKimiCodeOAuthRef, resolveKimiCodeRuntimeAuth, type BearerTokenProvider, @@ -85,6 +86,8 @@ interface FlowState { readonly provider: string; readonly controller: AbortController; readonly oauthRef: OAuthRef | undefined; + /** Base URL of the environment the login targeted (env-aware); drives the provisioned provider entry. */ + readonly loginBaseUrl: string | undefined; device: DeviceAuthorization | undefined; status: OAuthFlowStatus; expiresAt: number; @@ -121,10 +124,12 @@ export class OAuthService extends Disposable implements IOAuthService { async startLogin(provider = KIMI_CODE_PROVIDER_NAME): Promise { this.log.info('oauth startLogin: enter', { provider }); - const oauthRef = this.resolveOAuthRef(provider); - this.log.info('oauth startLogin: resolved oauthRef', { + const loginAuth = this.resolveLoginAuth(provider); + this.log.info('oauth startLogin: resolved login auth', { provider, - hasOAuthRef: oauthRef !== undefined, + hasOAuthRef: loginAuth.oauthRef !== undefined, + hasBaseUrl: loginAuth.baseUrl !== undefined, + hasOAuthHost: loginAuth.oauthHost !== undefined, }); this.abortExisting(provider); @@ -132,7 +137,8 @@ export class OAuthService extends Disposable implements IOAuthService { flowId: `oauth_${randomUUID()}`, provider, controller: new AbortController(), - oauthRef, + oauthRef: loginAuth.oauthRef, + loginBaseUrl: loginAuth.baseUrl, device: undefined, status: 'pending', expiresAt: Date.now() + DEFAULT_DEVICE_EXPIRES_IN_SEC * 1000, @@ -152,7 +158,9 @@ export class OAuthService extends Disposable implements IOAuthService { this.log.info('oauth startLogin: calling toolkit.login', { provider }); const loginPromise = this.toolkit.login(provider, { signal: state.controller.signal, - oauthRef, + oauthRef: loginAuth.oauthRef, + baseUrl: loginAuth.baseUrl, + oauthHost: loginAuth.oauthHost, onDeviceCode: (auth) => { this.log.info('oauth startLogin: onDeviceCode fired', { provider }); state.device = auth; @@ -226,7 +234,12 @@ export class OAuthService extends Disposable implements IOAuthService { } async logout(provider = KIMI_CODE_PROVIDER_NAME): Promise { - const oauthRef = this.readOAuthRefOptional(provider); + // Delete the token from the slot the runtime actually reads (v1 parity): + // env-aware for kimi-code, so an env-scoped login's token is removed too. + const oauthRef = + provider === KIMI_CODE_PROVIDER_NAME + ? this.resolveRuntimeOAuthRef(provider) + : this.readOAuthRefOptional(provider); const result = await this.toolkit.logout(provider, oauthRef); this.abortExisting(provider); await this.deprovisionProvider(provider); @@ -367,11 +380,44 @@ export class OAuthService extends Disposable implements IOAuthService { }; } - private resolveOAuthRef(provider: string): OAuthRef | undefined { + /** + * Resolve the environment the login should target (v1's + * `managedAuth.login`): env-aware via `resolveKimiCodeLoginAuth`, so + * `KIMI_CODE_BASE_URL` / `KIMI_CODE_OAUTH_HOST` steer the credential slot + * the token is written to the same way they steer the runtime token reads + * (`resolveKimiCodeRuntimeAuth`). A mismatched slot is the "login succeeds + * but every call 401s" bug. + */ + private resolveLoginAuth(provider: string): { + readonly oauthRef: OAuthRef | undefined; + readonly baseUrl: string | undefined; + readonly oauthHost: string | undefined; + } { const config = this.providerService.get(provider); - if (config?.oauth !== undefined) return config.oauth; - if (provider !== KIMI_CODE_PROVIDER_NAME) return undefined; - return resolveKimiCodeOAuthRef({ baseUrl: config?.baseUrl }); + if (provider !== KIMI_CODE_PROVIDER_NAME) { + return { oauthRef: config?.oauth, baseUrl: undefined, oauthHost: undefined }; + } + const loginAuth = resolveKimiCodeLoginAuth({ + configuredBaseUrl: config?.baseUrl, + configuredOAuthRef: config?.oauth, + }); + // Always resolve to a concrete ref for kimi-code: when the login env + // overrides the configured one, the provisioned entry must record the + // env-scoped slot explicitly (v1's `provisionManagedKimiCodeConfig` + // writes the login (oauthKey, oauthHost)) — not only so the runtime can + // find it, but so `isKimiOAuthProvider` still holds and the post-login + // model refresh runs. + const oauthRef = + loginAuth.oauthRef ?? + resolveKimiCodeOAuthRef({ + oauthHost: loginAuth.oauthHost, + baseUrl: loginAuth.baseUrl, + }); + return { + oauthRef, + baseUrl: loginAuth.baseUrl, + oauthHost: loginAuth.oauthHost, + }; } private readOAuthRefOptional(provider: string): OAuthRef | undefined { @@ -425,7 +471,7 @@ export class OAuthService extends Disposable implements IOAuthService { private async finalizeAuthentication(state: FlowState): Promise { try { - await this.provisionProvider(state.provider, state.oauthRef); + await this.provisionProvider(state.provider, state.oauthRef, state.loginBaseUrl); if (state.status !== 'pending') return; if (state.provider === KIMI_CODE_PROVIDER_NAME) { await this.refreshOAuthProviderModelsBestEffort(state.provider); @@ -443,9 +489,18 @@ export class OAuthService extends Disposable implements IOAuthService { } } - private async provisionProvider(provider: string, oauthRef: OAuthRef | undefined): Promise { - if (oauthRef === undefined) return; - const baseUrl = this.providerService.get(provider)?.baseUrl ?? kimiCodeBaseUrl(); + private async provisionProvider( + provider: string, + oauthRef: OAuthRef | undefined, + loginBaseUrl: string | undefined, + ): Promise { + // `baseUrl` comes from the login environment (env-aware), not a stale + // configured one, and `oauth` records the login credential slot — v1 + // parity: `provisionManagedKimiCodeConfig` rewrites both from the login + // auth. Non-kimi providers without a ref keep the old skip. + if (oauthRef === undefined && provider !== KIMI_CODE_PROVIDER_NAME) return; + const baseUrl = + loginBaseUrl ?? this.providerService.get(provider)?.baseUrl ?? kimiCodeBaseUrl(); await this.providerService.set(provider, { type: 'kimi', baseUrl, diff --git a/packages/agent-core-v2/test/app/auth/auth.test.ts b/packages/agent-core-v2/test/app/auth/auth.test.ts index 6d7872fced..2494f58499 100644 --- a/packages/agent-core-v2/test/app/auth/auth.test.ts +++ b/packages/agent-core-v2/test/app/auth/auth.test.ts @@ -8,6 +8,7 @@ import { afterEach, beforeEach, describe, expect, it, vi, type Mock } from 'vitest'; import { clearManagedKimiCodeConfig, + resolveKimiCodeOAuthKey, resolveKimiCodeRuntimeAuth, } from '@moonshot-ai/kimi-code-oauth'; @@ -53,6 +54,27 @@ const deviceAuth = { const flush = (): Promise => new Promise((resolve) => setTimeout(resolve, 0)); +/** + * Scoped credential ref derived for the `https://api.example.com` fixture + * environment (default OAuth host) — what login resolves when the configured + * ref does not match its (host, baseUrl) environment. + */ +const EXAMPLE_COM_SCOPED_REF = { + storage: 'file', + key: resolveKimiCodeOAuthKey({ baseUrl: 'https://api.example.com' }), + oauthHost: 'https://auth.kimi.com', +} as const; + +/** Scoped credential ref for the env-override fixture environment. */ +const ENV_SCOPED_REF = { + storage: 'file', + key: resolveKimiCodeOAuthKey({ + oauthHost: 'https://env-auth.example.com', + baseUrl: 'https://env-api.example.com/coding/v1', + }), + oauthHost: 'https://env-auth.example.com', +} as const; + interface FakeToolkit { readonly login: Mock<(...args: any[]) => any>; readonly logout: ReturnType; @@ -167,6 +189,7 @@ describe('OAuthService', () => { afterEach(() => { disposables.dispose(); vi.unstubAllGlobals(); + vi.unstubAllEnvs(); }); function createService(): IOAuthService { @@ -214,7 +237,14 @@ describe('OAuthService', () => { }); expect(toolkit.login).toHaveBeenCalledWith( OAUTH_PROVIDER, - expect.objectContaining({ oauthRef: { storage: 'file', key: 'oauth/kimi-code' } }), + expect.objectContaining({ + // The fixture's configured key does not match its (host, baseUrl) + // environment, so login re-derives the slot from the environment + // (v1 parity) instead of trusting the stale ref. + oauthRef: EXAMPLE_COM_SCOPED_REF, + baseUrl: 'https://api.example.com', + oauthHost: undefined, + }), ); await vi.waitFor(() => expect(svc.getFlow(OAUTH_PROVIDER)?.status).toBe('authenticated')); @@ -236,12 +266,14 @@ describe('OAuthService', () => { type: 'kimi', baseUrl: 'https://api.example.com', apiKey: '', - oauth: { storage: 'file', key: 'oauth/kimi-code' }, + // The provisioned entry records the env-scoped slot explicitly, so + // the runtime reads the same slot login wrote (v1 parity). + oauth: EXAMPLE_COM_SCOPED_REF, }), ); }); - it('startLogin resolves a default oauth ref for the managed provider without oauth config', async () => { + it('startLogin resolves an env-scoped oauth ref for the managed provider without oauth config', async () => { providers[OAUTH_PROVIDER] = { type: 'kimi', baseUrl: 'https://api.example.com' }; stubManagedModelsFetch(); toolkit.login.mockImplementation((_provider, options) => { @@ -254,7 +286,61 @@ describe('OAuthService', () => { expect(toolkit.login).toHaveBeenCalledWith( OAUTH_PROVIDER, expect.objectContaining({ - oauthRef: expect.objectContaining({ storage: 'file', key: expect.any(String) }), + oauthRef: EXAMPLE_COM_SCOPED_REF, + baseUrl: 'https://api.example.com', + }), + ); + await flush(); + expect(providerSet).toHaveBeenCalledWith( + OAUTH_PROVIDER, + expect.objectContaining({ + type: 'kimi', + baseUrl: 'https://api.example.com', + oauth: EXAMPLE_COM_SCOPED_REF, + }), + ); + }); + + it('startLogin reuses the configured oauth ref when it matches the login environment', async () => { + providers[OAUTH_PROVIDER] = { + type: 'kimi', + baseUrl: 'https://api.kimi.com/coding/v1', + oauth: { storage: 'file', key: 'oauth/kimi-code' }, + }; + stubManagedModelsFetch(); + toolkit.login.mockImplementation((_provider, options) => { + options.onDeviceCode(deviceAuth); + return Promise.resolve({ providerName: OAUTH_PROVIDER, ok: true }); + }); + const svc = createService(); + await svc.startLogin(OAUTH_PROVIDER); + + expect(toolkit.login).toHaveBeenCalledWith( + OAUTH_PROVIDER, + expect.objectContaining({ + oauthRef: { storage: 'file', key: 'oauth/kimi-code' }, + baseUrl: 'https://api.kimi.com/coding/v1', + }), + ); + }); + + it('startLogin honors KIMI_CODE_BASE_URL / KIMI_CODE_OAUTH_HOST for the login environment', async () => { + vi.stubEnv('KIMI_CODE_BASE_URL', 'https://env-api.example.com/coding/v1'); + vi.stubEnv('KIMI_CODE_OAUTH_HOST', 'https://env-auth.example.com'); + stubManagedModelsFetch(); + toolkit.login.mockImplementation((_provider, options) => { + options.onDeviceCode(deviceAuth); + return Promise.resolve({ providerName: OAUTH_PROVIDER, ok: true }); + }); + const svc = createService(); + await svc.startLogin(OAUTH_PROVIDER); + + expect(toolkit.login).toHaveBeenCalledWith( + OAUTH_PROVIDER, + expect.objectContaining({ + oauthRef: ENV_SCOPED_REF, + baseUrl: 'https://env-api.example.com/coding/v1', + oauthHost: 'https://env-auth.example.com', }), ); await flush(); @@ -262,7 +348,37 @@ describe('OAuthService', () => { OAUTH_PROVIDER, expect.objectContaining({ type: 'kimi', - oauth: expect.objectContaining({ storage: 'file', key: expect.any(String) }), + // The provisioned entry targets the env environment, not the stale + // configured one — so runtime reads hit the same credential slot. + baseUrl: 'https://env-api.example.com/coding/v1', + oauth: ENV_SCOPED_REF, + }), + ); + }); + + it('resolves the runtime credential slot to the env environment after an env-scoped login', async () => { + vi.stubEnv('KIMI_CODE_BASE_URL', 'https://env-api.example.com/coding/v1'); + vi.stubEnv('KIMI_CODE_OAUTH_HOST', 'https://env-auth.example.com'); + stubManagedModelsFetch(); + toolkit.login.mockImplementation((_provider, options) => { + options.onDeviceCode(deviceAuth); + return Promise.resolve({ providerName: OAUTH_PROVIDER, ok: true }); + }); + const svc = createService(); + await svc.startLogin(OAUTH_PROVIDER); + await vi.waitFor(() => expect(svc.getFlow(OAUTH_PROVIDER)?.status).toBe('authenticated')); + + // The slot login targeted and the slot the runtime reads must be the + // same env-scoped key — the mismatch was "login succeeds but every + // call 401s". + await svc.status(OAUTH_PROVIDER); + expect(toolkit.getCachedAccessToken).toHaveBeenCalledWith( + OAUTH_PROVIDER, + expect.objectContaining({ + key: resolveKimiCodeOAuthKey({ + oauthHost: 'https://env-auth.example.com', + baseUrl: 'https://env-api.example.com/coding/v1', + }), }), ); }); @@ -291,7 +407,7 @@ describe('OAuthService', () => { expect.objectContaining({ type: 'kimi', baseUrl: 'https://api.example.com', - oauth: { storage: 'file', key: 'oauth/kimi-code' }, + oauth: EXAMPLE_COM_SCOPED_REF, }), ); expect(fetchMock).toHaveBeenCalledTimes(1); @@ -315,7 +431,7 @@ describe('OAuthService', () => { expect.objectContaining({ type: 'kimi', baseUrl: 'https://api.example.com', - oauth: { storage: 'file', key: 'oauth/kimi-code' }, + oauth: EXAMPLE_COM_SCOPED_REF, }), ); expect(configSet).not.toHaveBeenCalledWith('defaultModel', expect.any(String)); @@ -355,7 +471,7 @@ describe('OAuthService', () => { OAUTH_PROVIDER, expect.objectContaining({ type: 'kimi', - oauth: { storage: 'file', key: 'oauth/kimi-code' }, + oauth: EXAMPLE_COM_SCOPED_REF, }), ); expect(configReplace).toHaveBeenCalledWith( @@ -421,10 +537,10 @@ describe('OAuthService', () => { const result = await svc.logout(OAUTH_PROVIDER); expect(result).toEqual({ logged_out: true, provider: OAUTH_PROVIDER }); - expect(toolkit.logout).toHaveBeenCalledWith(OAUTH_PROVIDER, { - storage: 'file', - key: 'oauth/kimi-code', - }); + // Logout deletes from the slot the runtime reads: the fixture's configured + // key does not match its (host, baseUrl) environment, so the env-derived + // scoped slot is the one cleared (v1 parity). + expect(toolkit.logout).toHaveBeenCalledWith(OAUTH_PROVIDER, EXAMPLE_COM_SCOPED_REF); expect(configReplace).toHaveBeenCalledWith('providers', { [NON_OAUTH_PROVIDER]: { type: 'openai', apiKey: 'sk-test' }, }); @@ -500,10 +616,7 @@ describe('OAuthService', () => { const svc = createService(); await expect(svc.logout(OAUTH_PROVIDER)).rejects.toThrow('config write failed'); - expect(toolkit.logout).toHaveBeenCalledWith(OAUTH_PROVIDER, { - storage: 'file', - key: 'oauth/kimi-code', - }); + expect(toolkit.logout).toHaveBeenCalledWith(OAUTH_PROVIDER, EXAMPLE_COM_SCOPED_REF); }); it('status reports loggedIn based on the cached access token', async () => {