diff --git a/packages/agent-core-v2/src/app/auth/authService.ts b/packages/agent-core-v2/src/app/auth/authService.ts index e52b69e48..01a2b6d5d 100644 --- a/packages/agent-core-v2/src/app/auth/authService.ts +++ b/packages/agent-core-v2/src/app/auth/authService.ts @@ -97,6 +97,14 @@ export class OAuthService extends Disposable implements IOAuthService { declare readonly _serviceBrand: undefined; private readonly flows = new Map(); + /** + * Serializes managed-provider model refreshes so a refresh triggered by + * login completion and a manual `:refresh_oauth` (or two overlapping manual + * ones) never race on reading/patching the persisted config. Mirrors v1's + * `_refreshChain`. + */ + private refreshChain: Promise = Promise.resolve(); + constructor( @IOAuthToolkit private readonly toolkit: IOAuthToolkit, @IProviderService private readonly providerService: IProviderService, @@ -249,7 +257,16 @@ export class OAuthService extends Disposable implements IOAuthService { return this.toolkit.getCachedAccessToken(provider, this.resolveRuntimeOAuthRef(provider, oauthRef)); } - async refreshOAuthProviderModels(): Promise { + refreshOAuthProviderModels(): Promise { + const run = this.refreshChain.then(() => this.doRefreshOAuthProviderModels()); + this.refreshChain = run.then( + () => undefined, + () => undefined, + ); + return run; + } + + private async doRefreshOAuthProviderModels(): Promise { const changed: RefreshOAuthProviderModelsResponse['changed'] = []; const unchanged: string[] = []; const failed: RefreshOAuthProviderModelsResponse['failed'] = []; diff --git a/packages/agent-core-v2/src/app/modelCatalog/modelCatalogService.ts b/packages/agent-core-v2/src/app/modelCatalog/modelCatalogService.ts index 9673bdd67..694d9a29c 100644 --- a/packages/agent-core-v2/src/app/modelCatalog/modelCatalogService.ts +++ b/packages/agent-core-v2/src/app/modelCatalog/modelCatalogService.ts @@ -53,6 +53,13 @@ const THINKING_SECTION = 'thinking'; export class ModelCatalogService implements IModelCatalogService { declare readonly _serviceBrand: undefined; + /** + * Serializes refresh runs so a scheduled refresh and a manual one (or two + * manual ones with different options) never race on reading/patching the + * persisted config. Mirrors v1's `_refreshChain`. + */ + private refreshChain: Promise = Promise.resolve(); + constructor( @IModelService private readonly modelService: IModelService, @IProviderService private readonly providerService: IProviderService, @@ -100,8 +107,19 @@ export class ModelCatalogService implements IModelCatalogService { }; } - async refreshProviderModels( + refreshProviderModels( options: RefreshProviderModelsOptions = {}, + ): Promise { + const run = this.refreshChain.then(() => this.doRefreshProviderModels(options)); + this.refreshChain = run.then( + () => undefined, + () => undefined, + ); + return run; + } + + private async doRefreshProviderModels( + options: RefreshProviderModelsOptions, ): Promise { await this.config.reload(); if (options.providerId !== undefined) { diff --git a/packages/agent-core-v2/test/auth/auth.test.ts b/packages/agent-core-v2/test/auth/auth.test.ts index 27c69c766..43eb7e786 100644 --- a/packages/agent-core-v2/test/auth/auth.test.ts +++ b/packages/agent-core-v2/test/auth/auth.test.ts @@ -588,6 +588,39 @@ describe('OAuthService', () => { }, ]); }); + + it('serializes concurrent refreshOAuthProviderModels runs so they never overlap', async () => { + let inFlight = 0; + let maxInFlight = 0; + const fetchMock = vi.fn().mockImplementation(async () => { + inFlight++; + maxInFlight = Math.max(maxInFlight, inFlight); + await new Promise((resolve) => setTimeout(resolve, 20)); + inFlight--; + return { + ok: true, + json: async () => ({ + data: [ + { + id: 'kimi-k2', + context_length: 131072, + supports_reasoning: true, + display_name: 'Kimi K2', + }, + ], + }), + }; + }); + vi.stubGlobal('fetch', fetchMock); + const svc = createService(); + + await Promise.all([svc.refreshOAuthProviderModels(), svc.refreshOAuthProviderModels()]); + + // Without the refresh chain both remote fetches would overlap (peak 2); the + // chain holds the second run until the first finishes, so the peak stays 1. + expect(maxInFlight).toBe(1); + expect(fetchMock).toHaveBeenCalledTimes(2); + }); }); describe('WebSearchProviderService', () => { diff --git a/packages/agent-core-v2/test/modelCatalog/modelCatalog.test.ts b/packages/agent-core-v2/test/modelCatalog/modelCatalog.test.ts index 8989dd7e9..4b612e0aa 100644 --- a/packages/agent-core-v2/test/modelCatalog/modelCatalog.test.ts +++ b/packages/agent-core-v2/test/modelCatalog/modelCatalog.test.ts @@ -8,6 +8,7 @@ * refresh is covered in `auth/auth.test.ts`. */ +import { KIMI_CODE_PROVIDER_NAME } from '@moonshot-ai/kimi-code-oauth'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { DisposableStore } from '#/_base/di/lifecycle'; @@ -69,6 +70,7 @@ describe('ModelCatalogService', () => { let configSet: ReturnType; let configReplace: ReturnType; let getCachedAccessToken: ReturnType; + let resolveTokenProvider: ReturnType; let publishEvent: ReturnType; beforeEach(() => { @@ -82,6 +84,7 @@ describe('ModelCatalogService', () => { : patch; }); getCachedAccessToken = vi.fn().mockResolvedValue(undefined); + resolveTokenProvider = vi.fn().mockReturnValue(undefined); configReplace = vi.fn().mockImplementation(async (domain: string, value: unknown) => { (backing as unknown as Record)[domain] = value; }); @@ -106,9 +109,8 @@ describe('ModelCatalogService', () => { }); reg.definePartialInstance(IOAuthService, { getCachedAccessToken: getCachedAccessToken as unknown as IOAuthService['getCachedAccessToken'], - resolveTokenProvider: vi - .fn() - .mockReturnValue(undefined) as unknown as IOAuthService['resolveTokenProvider'], + resolveTokenProvider: + resolveTokenProvider as unknown as IOAuthService['resolveTokenProvider'], }); reg.definePartialInstance(IEventService, { publish: publishEvent as unknown as IEventService['publish'], @@ -265,4 +267,51 @@ describe('ModelCatalogService', () => { expect(result).toEqual({ changed: [], unchanged: [], failed: [] }); expect(publishEvent).not.toHaveBeenCalled(); }); + + it('serializes concurrent refreshProviderModels runs so they never overlap', async () => { + // Seed the managed OAuth provider so the orchestrator actually refreshes it + // (a plain api-key provider is a no-op and would not exercise the chain). + backing.providers = { + [KIMI_CODE_PROVIDER_NAME]: { + type: 'kimi', + baseUrl: 'https://api.example.test/v1', + oauth: { storage: 'file', key: 'oauth/kimi-code' }, + }, + }; + backing.models = {}; + resolveTokenProvider.mockReturnValue({ getAccessToken: async () => 'access-token' }); + + let inFlight = 0; + let maxInFlight = 0; + const fetchMock = vi.fn().mockImplementation(async () => { + inFlight++; + maxInFlight = Math.max(maxInFlight, inFlight); + await new Promise((resolve) => setTimeout(resolve, 20)); + inFlight--; + return { + ok: true, + json: async () => ({ + data: [ + { + id: 'kimi-k2', + context_length: 131072, + supports_reasoning: true, + display_name: 'Kimi K2', + }, + ], + }), + }; + }); + vi.stubGlobal('fetch', fetchMock); + + await Promise.all([ + catalog().refreshProviderModels({ scope: 'all' }), + catalog().refreshProviderModels({ scope: 'all' }), + ]); + + // Without the refresh chain both remote fetches would overlap (peak 2); the + // chain holds the second run until the first finishes, so the peak stays 1. + expect(maxInFlight).toBe(1); + expect(fetchMock).toHaveBeenCalledTimes(2); + }); });