diff --git a/packages/a2a-server/src/config/config.test.ts b/packages/a2a-server/src/config/config.test.ts index e68ebc4431..c676e46289 100644 --- a/packages/a2a-server/src/config/config.test.ts +++ b/packages/a2a-server/src/config/config.test.ts @@ -28,6 +28,7 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => { const mockConfig = { ...params, initialize: vi.fn(), + waitForMcpInit: vi.fn(), refreshAuth: vi.fn(), getExperiments: vi.fn().mockReturnValue({ flags: { @@ -94,6 +95,7 @@ describe('loadConfig', () => { const mockConfig = { ...(params as object), initialize: vi.fn(), + waitForMcpInit: vi.fn(), refreshAuth: vi.fn(), getExperiments: vi.fn().mockReturnValue({ flags: { diff --git a/packages/a2a-server/src/config/config.ts b/packages/a2a-server/src/config/config.ts index 0873f43d98..f3100bce4d 100644 --- a/packages/a2a-server/src/config/config.ts +++ b/packages/a2a-server/src/config/config.ts @@ -166,6 +166,8 @@ export async function loadConfig( // Needed to initialize ToolRegistry, and git checkpointing if enabled await config.initialize(); + + await config.waitForMcpInit(); startupProfiler.flush(config); await refreshAuthentication(config, adcFilePath, 'Config'); diff --git a/packages/cli/src/zed-integration/zedIntegration.test.ts b/packages/cli/src/zed-integration/zedIntegration.test.ts index f3586b17f8..e8e5355dc0 100644 --- a/packages/cli/src/zed-integration/zedIntegration.test.ts +++ b/packages/cli/src/zed-integration/zedIntegration.test.ts @@ -129,6 +129,7 @@ describe('GeminiAgent', () => { mockConfig = { refreshAuth: vi.fn(), initialize: vi.fn(), + waitForMcpInit: vi.fn(), getFileSystemService: vi.fn(), setFileSystemService: vi.fn(), getContentGeneratorConfig: vi.fn(), @@ -486,6 +487,7 @@ describe('Session', () => { getMessageBus: vi.fn().mockReturnValue(mockMessageBus), setApprovalMode: vi.fn(), isPlanEnabled: vi.fn().mockReturnValue(false), + waitForMcpInit: vi.fn(), } as unknown as Mocked; mockConnection = { sessionUpdate: vi.fn(), @@ -500,6 +502,28 @@ describe('Session', () => { vi.clearAllMocks(); }); + it('should await MCP initialization before processing a prompt', async () => { + const stream = createMockStream([ + { + type: StreamEventType.CHUNK, + value: { candidates: [{ content: { parts: [{ text: 'Hi' }] } }] }, + }, + ]); + mockChat.sendMessageStream.mockResolvedValue(stream); + + await session.prompt({ + sessionId: 'session-1', + prompt: [{ type: 'text', text: 'test' }], + }); + + expect(mockConfig.waitForMcpInit).toHaveBeenCalledOnce(); + const waitOrder = (mockConfig.waitForMcpInit as Mock).mock + .invocationCallOrder[0]; + const sendOrder = (mockChat.sendMessageStream as Mock).mock + .invocationCallOrder[0]; + expect(waitOrder).toBeLessThan(sendOrder); + }); + it('should handle prompt with text response', async () => { const stream = createMockStream([ { diff --git a/packages/cli/src/zed-integration/zedIntegration.ts b/packages/cli/src/zed-integration/zedIntegration.ts index 1cc9f1f3b5..98c9efdc75 100644 --- a/packages/cli/src/zed-integration/zedIntegration.ts +++ b/packages/cli/src/zed-integration/zedIntegration.ts @@ -521,6 +521,8 @@ export class Session { const pendingSend = new AbortController(); this.pendingPrompt = pendingSend; + await this.config.waitForMcpInit(); + const promptId = Math.random().toString(16).slice(2); const chat = this.chat; diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 6262956c11..87633d35b6 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -716,6 +716,7 @@ export class Config implements McpContext { private compressionTruncationCounter = 0; private initialized = false; private initPromise: Promise | undefined; + private mcpInitializationPromise: Promise | null = null; readonly storage: Storage; private readonly fileExclusions: FileExclusions; private readonly eventEmitter?: EventEmitter; @@ -1124,7 +1125,7 @@ export class Config implements McpContext { ); // We do not await this promise so that the CLI can start up even if // MCP servers are slow to connect. - const mcpInitialization = Promise.allSettled([ + this.mcpInitializationPromise = Promise.allSettled([ this.mcpClientManager.startConfiguredMcpServers(), this.getExtensionLoader().start(this), ]).then((results) => { @@ -1136,7 +1137,7 @@ export class Config implements McpContext { }); if (!this.interactive || this.experimentalZedIntegration) { - await mcpInitialization; + await this.mcpInitializationPromise; } if (this.skillsSupport) { @@ -2234,6 +2235,12 @@ export class Config implements McpContext { return this.experimentalZedIntegration; } + async waitForMcpInit(): Promise { + if (this.mcpInitializationPromise) { + await this.mcpInitializationPromise; + } + } + getListExtensions(): boolean { return this.listExtensions; }