fix(auth): keep OAuth login alive when its own provisioning writes the provider

.changeset/fix-oauth-login-self-cancel.md

@@ -0,0 +1,5 @@
+---
+"@moonshot-ai/kimi-code": patch
+---
+
+Fix sign-in briefly showing a device-code-expired error after a successful authorization.

packages/agent-core-v2/src/app/auth/authService.ts

@@ -86,6 +86,7 @@ interface FlowState {
   readonly loginBaseUrl: string | undefined;
   device: DeviceAuthorization | undefined;
   status: OAuthFlowStatus;
+  tokenGranted: boolean;
   expiresAt: number;
   gcTimer: ReturnType<typeof setTimeout> | undefined;
   errorMessage: string | undefined;
@@ -135,6 +136,7 @@ export class OAuthService extends Disposable implements IOAuthService {
       loginBaseUrl: loginAuth.baseUrl,
       device: undefined,
       status: 'pending',
+      tokenGranted: false,
       expiresAt: Date.now() + DEFAULT_DEVICE_EXPIRES_IN_SEC * 1000,
       gcTimer: undefined,
       errorMessage: undefined,
@@ -469,6 +471,7 @@ export class OAuthService extends Disposable implements IOAuthService {
     for (const state of this.flows.values()) {
       if (!affected.has(state.provider)) continue;
       if (state.status !== 'pending') continue;
+      if (state.tokenGranted) continue;
       state.controller.abort();
       state.errorMessage = 'Provider configuration changed during login.';
       this.setTerminal(state, 'cancelled');
@@ -477,20 +480,22 @@ export class OAuthService extends Disposable implements IOAuthService {
 
   private handleSuccess(state: FlowState): void {
     if (state.status !== 'pending') return;
-    void this.finalizeAuthentication(state);
+    state.tokenGranted = true;
+    void this.provisionAfterSuccess(state);
   }
 
   private async completeAlreadyAuthenticatedLogin(state: FlowState): Promise<void> {
-    await this.finalizeAuthentication(state);
+    if (state.status !== 'pending') return;
+    state.tokenGranted = true;
+    await this.provisionAfterSuccess(state);
   }
 
-  private async finalizeAuthentication(state: FlowState): Promise<void> {
+  private async provisionAfterSuccess(state: FlowState): Promise<void> {
     try {
       await this.provisionProvider(state.provider, state.oauthRef, state.loginBaseUrl);
-      if (state.status !== 'pending') return;
+      if (this.flows.get(state.provider) !== state) return;
       if (state.provider === KIMI_CODE_PROVIDER_NAME) {
         await this.refreshOAuthProviderModelsBestEffort(state.provider);
-        if (state.status !== 'pending') return;
       }
     } catch (error) {
       this.log.warn('oauth provider provisioning failed', {

packages/agent-core-v2/test/app/auth/auth.test.ts

@@ -98,7 +98,7 @@ describe('OAuthService', () => {
   let defaultModel: string | undefined;
   let thinking: { enabled?: boolean; effort?: string } | undefined;
   let toolkit: FakeToolkit;
-  let providerSet: ReturnType<typeof vi.fn>;
+  let providerSet: ReturnType<typeof vi.fn<(name: string, config: ProviderConfig) => Promise<void>>>;
   let configSet: ReturnType<typeof vi.fn>;
   let configReplace: ReturnType<typeof vi.fn>;
   let events: Event2[];
@@ -661,6 +661,50 @@ describe('OAuthService', () => {
     await vi.waitFor(() => expect(svc.getFlow(OAUTH_PROVIDER)?.status).toBe('cancelled'));
   });
 
+  it('reports pending until provisioning finishes after the grant settles', async () => {
+    stubManagedModelsFetch();
+    toolkit.login.mockImplementation((_provider, options) => {
+      options.onDeviceCode(deviceAuth);
+      return Promise.resolve({ providerName: OAUTH_PROVIDER, ok: true });
+    });
+    let resolveProvision!: () => void;
+    providerSet.mockImplementation((name: string, config: ProviderConfig) => {
+      providers = { ...providers, [name]: config };
+      return new Promise<void>((resolve) => {
+        resolveProvision = resolve;
+      });
+    });
+    const svc = createService();
+    await svc.startLogin(OAUTH_PROVIDER);
+
+    await vi.waitFor(() => {
+      expect(providerSet).toHaveBeenCalled();
+    });
+    expect(svc.getFlow(OAUTH_PROVIDER)?.status).toBe('pending');
+
+    resolveProvision();
+    await vi.waitFor(() => {
+      expect(svc.getFlow(OAUTH_PROVIDER)?.status).toBe('authenticated');
+    });
+  });
+
+  it('keeps the login authenticated when its own provisioning fires a provider change', async () => {
+    stubManagedModelsFetch();
+    toolkit.login.mockImplementation((_provider, options) => {
+      options.onDeviceCode(deviceAuth);
+      return Promise.resolve({ providerName: OAUTH_PROVIDER, ok: true });
+    });
+    providerSet.mockImplementation((name: string, config: ProviderConfig) => {
+      providers = { ...providers, [name]: config };
+      providerChangedEmitter.fire({ added: [], removed: [], changed: [name] });
+      return Promise.resolve();
+    });
+    const svc = createService();
+    await svc.startLogin(OAUTH_PROVIDER);
+
+    await vi.waitFor(() => expect(svc.getFlow(OAUTH_PROVIDER)?.status).toBe('authenticated'));
+  });
+
   it('cancelLogin aborts a pending flow and marks it cancelled', async () => {
     let capturedSignal: AbortSignal | undefined;
     toolkit.login.mockImplementation((_provider, options) => {