diff --git a/packages/agent-core-v2/src/app/mcpConfig/configStore.ts b/packages/agent-core-v2/src/app/mcpConfig/configStore.ts index 0b56c1e6b..b0f94a169 100644 --- a/packages/agent-core-v2/src/app/mcpConfig/configStore.ts +++ b/packages/agent-core-v2/src/app/mcpConfig/configStore.ts @@ -3,7 +3,7 @@ import { join } from 'pathe'; import { createDecorator, type ServiceIdentifier } from '#/_base/di/instantiation'; import { Disposable } from '#/_base/di/lifecycle'; import { ScopeActivation, registerScopedService } from '#/_base/di/scope'; -import { Emitter, type Event } from '#/_base/event'; +import { AsyncEmitter, type Event, type IWaitUntil } from '#/_base/event'; import { IBootstrapService } from '#/app/bootstrap/bootstrap'; import { LifecycleScope } from '#/app/scopes'; import { ErrorCodes, Error2 } from '#/errors'; @@ -12,10 +12,12 @@ import { IFileSystemStorageService } from '#/persistence/interface/storage'; export type GlobalMcpServerConfig = McpServerConfig & { readonly name: string }; +export type McpConfigWriteEvent = IWaitUntil; + export interface IMcpConfigStore { readonly _serviceBrand: undefined; readonly path: string; - readonly onDidWrite: Event; + readonly onDidWrite: Event; list(): Promise; get(name: string): Promise; add(server: GlobalMcpServerConfig): Promise; @@ -43,8 +45,8 @@ export class McpConfigStore extends Disposable implements IMcpConfigStore { readonly path: string; - private readonly writeEmitter = this._register(new Emitter()); - readonly onDidWrite: Event = this.writeEmitter.event; + private readonly writeEmitter = this._register(new AsyncEmitter()); + readonly onDidWrite: Event = this.writeEmitter.event; private mutationTail: Promise = Promise.resolve(); constructor( @@ -160,10 +162,12 @@ export class McpConfigStore extends Disposable implements IMcpConfigStore { await this.storage.write(CONFIG_SCOPE, MCP_CONFIG_KEY, textEncoder.encode(text), { atomic: true, }); - this.writeEmitter.fire(); + await this.writeEmitter.fireAsync({}, NO_ABORT); } } +const NO_ABORT = new AbortController().signal; + function parseServerInput(server: GlobalMcpServerConfig): GlobalMcpServerConfig { return parseServer(normalizeServerName(server.name), server); } diff --git a/packages/agent-core-v2/src/workspace/workspaceMcp/workspaceMcpService.ts b/packages/agent-core-v2/src/workspace/workspaceMcp/workspaceMcpService.ts index a6e889432..fb10f0d0b 100644 --- a/packages/agent-core-v2/src/workspace/workspaceMcp/workspaceMcpService.ts +++ b/packages/agent-core-v2/src/workspace/workspaceMcp/workspaceMcpService.ts @@ -67,7 +67,7 @@ export class WorkspaceMcpService extends Disposable implements IWorkspaceMcpServ this._register({ dispose: () => void this.manager.shutdown() }); this._register( this.mcpConfig.onDidChange((change) => { - this.scheduleApply(change); + change.waitUntil(this.scheduleApply(change)); }), ); this._register({ dispose: this.oauthEventSubscription(this.manager) }); @@ -252,8 +252,8 @@ export class WorkspaceMcpService extends Disposable implements IWorkspaceMcpServ this.trackMcpInitialLoad(); } - private scheduleApply(change: McpServersChange): void { - void this.ready + private scheduleApply(change: McpServersChange): Promise { + return this.ready .then(() => this.mutate(() => this.apply(change))) .catch((error) => { this.log.warn(`mcp server change apply failed: ${String(error)}`); diff --git a/packages/agent-core-v2/src/workspace/workspaceMcpConfig/workspaceMcpConfig.ts b/packages/agent-core-v2/src/workspace/workspaceMcpConfig/workspaceMcpConfig.ts index 6ae1194a2..783a4c7f3 100644 --- a/packages/agent-core-v2/src/workspace/workspaceMcpConfig/workspaceMcpConfig.ts +++ b/packages/agent-core-v2/src/workspace/workspaceMcpConfig/workspaceMcpConfig.ts @@ -1,5 +1,5 @@ import { createDecorator, type ServiceIdentifier } from '#/_base/di/instantiation'; -import type { Event } from '#/_base/event'; +import type { Event, IWaitUntil } from '#/_base/event'; import type { McpServerConfig } from '#/mcpCore/config-schema'; export interface McpServersChange { @@ -7,6 +7,8 @@ export interface McpServersChange { readonly remove: readonly string[]; } +export type McpServersChangeEvent = McpServersChange & IWaitUntil; + export interface McpTunables { readonly startupTimeoutMs?: number; readonly toolTimeoutMs?: number; @@ -21,7 +23,7 @@ export interface IWorkspaceMcpConfigService { tunables(): McpTunables; - readonly onDidChange: Event; + readonly onDidChange: Event; } export const IWorkspaceMcpConfigService: ServiceIdentifier = diff --git a/packages/agent-core-v2/src/workspace/workspaceMcpConfig/workspaceMcpConfigService.ts b/packages/agent-core-v2/src/workspace/workspaceMcpConfig/workspaceMcpConfigService.ts index 85f86b7d5..9be8a528a 100644 --- a/packages/agent-core-v2/src/workspace/workspaceMcpConfig/workspaceMcpConfigService.ts +++ b/packages/agent-core-v2/src/workspace/workspaceMcpConfig/workspaceMcpConfigService.ts @@ -1,7 +1,7 @@ import { dirname } from 'pathe'; import { Disposable } from '#/_base/di/lifecycle'; -import { Emitter } from '#/_base/event'; +import { AsyncEmitter } from '#/_base/event'; import { ILogService } from '#/_base/log/log'; import { subtreeWatchFilter } from '#/_base/utils/paths'; import { TimeoutTimer } from '#/_base/utils/timer'; @@ -19,7 +19,7 @@ import { IWorkspaceTrust } from '#/workspace/workspaceTrust/workspaceTrust'; import { IWorkspaceMcpConfigService, - type McpServersChange, + type McpServersChangeEvent, type McpTunables, } from './workspaceMcpConfig'; @@ -34,7 +34,7 @@ export class WorkspaceMcpConfigService extends Disposable implements IWorkspaceM private pluginServers = new Map(); private current: Readonly> = {}; private readonly watchDebounce = this._register(new TimeoutTimer()); - private readonly changeEmitter = this._register(new Emitter()); + private readonly changeEmitter = this._register(new AsyncEmitter()); readonly onDidChange = this.changeEmitter.event; constructor( @@ -67,10 +67,12 @@ export class WorkspaceMcpConfigService extends Disposable implements IWorkspaceM }), ); this._register( - mcpConfigStore.onDidWrite(() => { - void this.reloadFileServers().catch((error) => { - this.log.warn(`mcp config reload after management write failed: ${String(error)}`); - }); + mcpConfigStore.onDidWrite((event) => { + event.waitUntil( + this.reloadFileServers().catch((error) => { + this.log.warn(`mcp config reload after management write failed: ${String(error)}`); + }), + ); }), ); void this.watchConfigFiles(); @@ -164,7 +166,7 @@ export class WorkspaceMcpConfigService extends Disposable implements IWorkspaceM includeProject: this.trust.isTrusted(), }); this.fileServers = new Map(Object.entries(fresh)); - this.publishIfChanged(); + await this.publishIfChanged(); }); } @@ -173,11 +175,11 @@ export class WorkspaceMcpConfigService extends Disposable implements IWorkspaceM await this.mutate(async () => { const fresh = await this.plugins.enabledMcpServers(); this.pluginServers = new Map(Object.entries(fresh)); - this.publishIfChanged(); + await this.publishIfChanged(); }); } - private publishIfChanged(): void { + private async publishIfChanged(): Promise { const next = this.merged(); const upsert: Record = Object.create(null); const remove: string[] = []; @@ -192,10 +194,12 @@ export class WorkspaceMcpConfigService extends Disposable implements IWorkspaceM } this.current = next; if (Object.keys(upsert).length === 0 && remove.length === 0) return; - this.changeEmitter.fire({ upsert, remove }); + await this.changeEmitter.fireAsync({ upsert, remove }, NO_ABORT); } } +const NO_ABORT = new AbortController().signal; + function fingerprintConfig(config: McpServerConfig): string { return JSON.stringify(sortKeysDeep(config)); } diff --git a/packages/agent-core-v2/test/app/mcpConfig/configStore.test.ts b/packages/agent-core-v2/test/app/mcpConfig/configStore.test.ts index 394a647ad..19948f217 100644 --- a/packages/agent-core-v2/test/app/mcpConfig/configStore.test.ts +++ b/packages/agent-core-v2/test/app/mcpConfig/configStore.test.ts @@ -290,5 +290,32 @@ describe('McpConfigStore', () => { await store.remove('ghost'); expect(fired).toBe(0); }); + + it('waits for asynchronous listeners before resolving a mutation', async () => { + let resolveStarted!: () => void; + const started = new Promise((resolve) => { + resolveStarted = resolve; + }); + let release!: () => void; + const gate = new Promise((resolve) => { + release = resolve; + }); + store.onDidWrite((event) => { + resolveStarted(); + event.waitUntil(gate); + }); + + let completed = false; + const mutation = store.add(stdioServer('alpha')).then(() => { + completed = true; + }); + await started; + await Promise.resolve(); + expect(completed).toBe(false); + + release(); + await mutation; + expect(completed).toBe(true); + }); }); }); diff --git a/packages/agent-core-v2/test/app/mcpManagement/mcpManagement.test.ts b/packages/agent-core-v2/test/app/mcpManagement/mcpManagement.test.ts index 73380ad91..b881b06c7 100644 --- a/packages/agent-core-v2/test/app/mcpManagement/mcpManagement.test.ts +++ b/packages/agent-core-v2/test/app/mcpManagement/mcpManagement.test.ts @@ -344,6 +344,33 @@ describe('McpManagementService', () => { await expect(store.list()).resolves.toEqual([]); }); + it('waits for live config reconciliation listeners before returning', async () => { + let resolveStarted!: () => void; + const started = new Promise((resolve) => { + resolveStarted = resolve; + }); + let release!: () => void; + const gate = new Promise((resolve) => { + release = resolve; + }); + store.onDidWrite((event) => { + resolveStarted(); + event.waitUntil(gate); + }); + + let completed = false; + const mutation = management.addServer(stdioServer('alpha')).then(() => { + completed = true; + }); + await started; + await Promise.resolve(); + expect(completed).toBe(false); + + release(); + await mutation; + expect(completed).toBe(true); + }); + it('normalizes server names so the guard, the persisted key, and the list agree', async () => { const added = await management.addServer(stdioServer(' alpha ')); diff --git a/packages/agent-core-v2/test/workspace/workspaceMcp/initialization.test.ts b/packages/agent-core-v2/test/workspace/workspaceMcp/initialization.test.ts index 634217e46..28aafe4be 100644 --- a/packages/agent-core-v2/test/workspace/workspaceMcp/initialization.test.ts +++ b/packages/agent-core-v2/test/workspace/workspaceMcp/initialization.test.ts @@ -13,7 +13,7 @@ import { ILogService } from '#/_base/log/log'; import { McpConnectionManager } from '#/mcpCore/connection-manager'; import { MCP_SECTION, type McpSection } from '#/app/mcpConfig/configSection'; import { IMcpOAuthService } from '#/app/mcpConfig/oauthService'; -import { IMcpConfigStore } from '#/app/mcpConfig/configStore'; +import { IMcpConfigStore, type McpConfigWriteEvent } from '#/app/mcpConfig/configStore'; import { McpOAuthService } from '#/mcpCore/oauth/service'; import { IBootstrapService } from '#/app/bootstrap/bootstrap'; import { IConfigService } from '#/app/config/config'; @@ -81,7 +81,9 @@ describe('Workspace MCP initialization', () => { IMcpOAuthService, new McpOAuthService({ store: createMemoryMcpOAuthStore() }), ); - reg.definePartialInstance(IMcpConfigStore, { onDidWrite: Event.None as Event }); + reg.definePartialInstance(IMcpConfigStore, { + onDidWrite: Event.None as Event, + }); reg.defineInstance(ILogService, stubLog()); reg.defineInstance(ITelemetryService, noopTelemetryService); const runtime = Object.assign( diff --git a/packages/agent-core-v2/test/workspace/workspaceMcp/workspaceMcp.test.ts b/packages/agent-core-v2/test/workspace/workspaceMcp/workspaceMcp.test.ts index 62ca5f458..c4736884f 100644 --- a/packages/agent-core-v2/test/workspace/workspaceMcp/workspaceMcp.test.ts +++ b/packages/agent-core-v2/test/workspace/workspaceMcp/workspaceMcp.test.ts @@ -12,7 +12,7 @@ import type { OAuthClientInformationFull } from '@modelcontextprotocol/sdk/share import type { ServiceIdentifier } from '#/_base/di/instantiation'; import { DisposableStore } from '#/_base/di/lifecycle'; import { createServices } from '#/_base/di/test'; -import { Emitter } from '#/_base/event'; +import { AsyncEmitter, Emitter } from '#/_base/event'; import { ILogService } from '#/_base/log/log'; import { IMcpOAuthService } from '#/app/mcpConfig/oauthService'; import { ISessionManager } from '#/app/sessionManager/sessionManager'; @@ -40,7 +40,7 @@ import { import { WorkspaceMcpService } from '#/workspace/workspaceMcp/workspaceMcpService'; import { IWorkspaceMcpConfigService, - type McpServersChange, + type McpServersChangeEvent, type McpTunables, } from '#/workspace/workspaceMcpConfig/workspaceMcpConfig'; @@ -68,7 +68,7 @@ describe('WorkspaceMcpService', () => { let current: Record; let tunablesValue: McpTunables; let tunablesFn: Mock<() => McpTunables>; - let configChanges: Emitter; + let configChanges: AsyncEmitter; let assemblyEvents: Emitter; let oauthService: McpOAuthService; let oauthScheduler: ManualMcpOAuthScheduler; @@ -80,7 +80,7 @@ describe('WorkspaceMcpService', () => { current = {}; tunablesValue = {}; tunablesFn = vi.fn(() => tunablesValue); - configChanges = new Emitter(); + configChanges = disposables.add(new AsyncEmitter()); assemblyEvents = disposables.add(new Emitter()); oauthScheduler = new ManualMcpOAuthScheduler(); oauthService = new McpOAuthService({ @@ -172,15 +172,13 @@ describe('WorkspaceMcpService', () => { await service.ready; expect(manager.get('alpha')?.status).toBe('connected'); - configChanges.fire({ upsert: { beta: stdioServer() }, remove: ['alpha'] }); - - await vi.waitFor( - () => { - expect(manager?.get('alpha')?.status).toBe('removed'); - expect(manager?.get('beta')?.status).toBe('connected'); - }, - { timeout: 10000, interval: 50 }, + await configChanges.fireAsync( + { upsert: { beta: stdioServer() }, remove: ['alpha'] }, + new AbortController().signal, ); + + expect(manager.get('alpha')?.status).toBe('removed'); + expect(manager.get('beta')?.status).toBe('connected'); }, 20000); it('queues change events until the initial connect settles', async () => { @@ -207,7 +205,10 @@ describe('WorkspaceMcpService', () => { const service = createService(); manager = service.connectionManager(); - configChanges.fire({ upsert: { beta: stdioServer() }, remove: ['alpha'] }); + void configChanges.fireAsync( + { upsert: { beta: stdioServer() }, remove: ['alpha'] }, + new AbortController().signal, + ); await connectAllStarted; expect(connect).not.toHaveBeenCalled(); expect(markRemoved).not.toHaveBeenCalled(); diff --git a/packages/agent-core-v2/test/workspace/workspaceMcpConfig/workspaceMcpConfig.test.ts b/packages/agent-core-v2/test/workspace/workspaceMcpConfig/workspaceMcpConfig.test.ts index c70a634a5..9076e0ce9 100644 --- a/packages/agent-core-v2/test/workspace/workspaceMcpConfig/workspaceMcpConfig.test.ts +++ b/packages/agent-core-v2/test/workspace/workspaceMcpConfig/workspaceMcpConfig.test.ts @@ -7,12 +7,15 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { DisposableStore } from '#/_base/di/lifecycle'; import { createServices } from '#/_base/di/test'; -import { Emitter } from '#/_base/event'; +import { AsyncEmitter, Emitter } from '#/_base/event'; import { ILogService } from '#/_base/log/log'; import { IBootstrapService } from '#/app/bootstrap/bootstrap'; import { IConfigService } from '#/app/config/config'; import { MCP_SECTION, type McpSection } from '#/app/mcpConfig/configSection'; -import { IMcpConfigStore } from '#/app/mcpConfig/configStore'; +import { + IMcpConfigStore, + type McpConfigWriteEvent, +} from '#/app/mcpConfig/configStore'; import { IPluginService } from '#/app/plugin/plugin'; import type { ReloadSummary } from '#/app/plugin/types'; import type { McpServerConfig } from '#/mcpCore/config-schema'; @@ -47,7 +50,7 @@ describe('WorkspaceMcpConfigService', () => { let watchFires: Map>; let pluginServers: Record; let pluginReloads: Emitter; - let storeWrites: Emitter; + let storeWrites: AsyncEmitter; let trusted: boolean; let trustFlips: Emitter; let changes: McpServersChange[]; @@ -59,7 +62,7 @@ describe('WorkspaceMcpConfigService', () => { watchFires = new Map(); pluginServers = {}; pluginReloads = new Emitter(); - storeWrites = new Emitter(); + storeWrites = disposables.add(new AsyncEmitter()); trusted = true; trustFlips = new Emitter(); changes = []; @@ -116,7 +119,7 @@ describe('WorkspaceMcpConfigService', () => { }, }); const service = ix.get(IWorkspaceMcpConfigService); - service.onDidChange((change) => changes.push(change)); + service.onDidChange(({ upsert, remove }) => changes.push({ upsert, remove })); return service; } @@ -231,7 +234,7 @@ describe('WorkspaceMcpConfigService', () => { expect(service.servers()).toEqual({ shared: stdioConfig('plugin-version') }); await writeProjectConfig({}); - storeWrites.fire(); + await storeWrites.fireAsync({}, new AbortController().signal); pluginServers = { shared: stdioConfig('plugin-version'), pluginOnly: stdioConfig('plugin'), @@ -313,7 +316,7 @@ describe('WorkspaceMcpConfigService', () => { JSON.stringify({ mcpServers: { added: stdioConfig('added') } }), 'utf8', ); - storeWrites.fire(); + await storeWrites.fireAsync({}, new AbortController().signal); await vi.waitFor( () => {