fix(agent-core-v2): make managed OAuth login env-aware

.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.

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,18 +124,21 @@ export class OAuthService extends Disposable implements IOAuthService {
 
   async startLogin(provider = KIMI_CODE_PROVIDER_NAME): Promise<OAuthFlowStart> {
     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);
 
     const state: FlowState = {
       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<OAuthLogoutResponse> {
-    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<void> {
     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<void> {
-    if (oauthRef === undefined) return;
-    const baseUrl = this.providerService.get(provider)?.baseUrl ?? kimiCodeBaseUrl();
+  private async provisionProvider(
+    provider: string,
+    oauthRef: OAuthRef | undefined,
+    loginBaseUrl: string | undefined,
+  ): Promise<void> {
+    // `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,

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<void> => 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<typeof vi.fn>;
@@ -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,15 +286,99 @@ 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();
     expect(providerSet).toHaveBeenCalledWith(
       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 () => {