From 12a5490bcf25d046a789f8ae8267c5f40a73af11 Mon Sep 17 00:00:00 2001 From: Adib234 <30782825+Adib234@users.noreply.github.com> Date: Fri, 23 Jan 2026 18:32:06 -0500 Subject: [PATCH] Allow prompt queueing during MCP initialization (#17395) --- packages/cli/src/ui/AppContainer.tsx | 30 +++- .../cli/src/ui/hooks/useGeminiStream.test.tsx | 131 ------------------ packages/cli/src/ui/hooks/useGeminiStream.ts | 20 --- .../cli/src/ui/hooks/useMcpStatus.test.tsx | 97 +++++++++++++ packages/cli/src/ui/hooks/useMcpStatus.ts | 51 +++++++ .../cli/src/ui/hooks/useMessageQueue.test.tsx | 42 +++++- packages/cli/src/ui/hooks/useMessageQueue.ts | 11 +- packages/core/src/tools/mcp-client-manager.ts | 7 + 8 files changed, 234 insertions(+), 155 deletions(-) create mode 100644 packages/cli/src/ui/hooks/useMcpStatus.test.tsx create mode 100644 packages/cli/src/ui/hooks/useMcpStatus.ts diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index 2d77bc4910..70cf7b6cb1 100644 --- a/packages/cli/src/ui/AppContainer.tsx +++ b/packages/cli/src/ui/AppContainer.tsx @@ -106,6 +106,7 @@ import { registerCleanup, runExitCleanup } from '../utils/cleanup.js'; import { RELAUNCH_EXIT_CODE } from '../utils/processUtils.js'; import type { SessionInfo } from '../utils/sessionUtils.js'; import { useMessageQueue } from './hooks/useMessageQueue.js'; +import { useMcpStatus } from './hooks/useMcpStatus.js'; import { useApprovalModeIndicator } from './hooks/useApprovalModeIndicator.js'; import { useSessionStats } from './contexts/SessionContext.js'; import { useGitBranchName } from './hooks/useGitBranchName.js'; @@ -131,6 +132,7 @@ import { QUEUE_ERROR_DISPLAY_DURATION_MS, } from './constants.js'; import { LoginWithGoogleRestartDialog } from './auth/LoginWithGoogleRestartDialog.js'; +import { isSlashCommand } from './utils/commandUtils.js'; function isToolExecuting(pendingHistoryItems: HistoryItemWithoutId[]) { return pendingHistoryItems.some((item) => { @@ -910,6 +912,8 @@ Logging in with Google... Restarting Gemini CLI to continue. isActive: !embeddedShellFocused, }); + const { isMcpReady } = useMcpStatus(config); + const { messageQueue, addMessage, @@ -920,6 +924,7 @@ Logging in with Google... Restarting Gemini CLI to continue. isConfigInitialized, streamingState, submitQuery, + isMcpReady, }); cancelHandlerRef.current = useCallback( @@ -961,10 +966,31 @@ Logging in with Google... Restarting Gemini CLI to continue. const handleFinalSubmit = useCallback( (submittedValue: string) => { - addMessage(submittedValue); + const isSlash = isSlashCommand(submittedValue.trim()); + const isIdle = streamingState === StreamingState.Idle; + + if (isSlash || (isIdle && isMcpReady)) { + void submitQuery(submittedValue); + } else { + // Check messageQueue.length === 0 to only notify on the first queued item + if (isIdle && !isMcpReady && messageQueue.length === 0) { + coreEvents.emitFeedback( + 'info', + 'Waiting for MCP servers to initialize... Slash commands are still available and prompts will be queued.', + ); + } + addMessage(submittedValue); + } addInput(submittedValue); // Track input for up-arrow history }, - [addMessage, addInput], + [ + addMessage, + addInput, + submitQuery, + isMcpReady, + streamingState, + messageQueue.length, + ], ); const handleClearScreen = useCallback(() => { diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx index a0f6fdfa68..730ea32e5e 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx +++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx @@ -1962,73 +1962,6 @@ describe('useGeminiStream', () => { }); }); - describe('MCP Discovery State', () => { - it('should block non-slash command queries when discovery is in progress and servers exist', async () => { - const mockMcpClientManager = { - getDiscoveryState: vi - .fn() - .mockReturnValue(MCPDiscoveryState.IN_PROGRESS), - getMcpServerCount: vi.fn().mockReturnValue(1), - }; - mockConfig.getMcpClientManager = () => mockMcpClientManager as any; - - const { result } = renderTestHook(); - - await act(async () => { - await result.current.submitQuery('test query'); - }); - - expect(coreEvents.emitFeedback).toHaveBeenCalledWith( - 'info', - 'Waiting for MCP servers to initialize... Slash commands are still available.', - ); - expect(mockSendMessageStream).not.toHaveBeenCalled(); - }); - - it('should NOT block queries when discovery is NOT_STARTED but there are no servers', async () => { - const mockMcpClientManager = { - getDiscoveryState: vi - .fn() - .mockReturnValue(MCPDiscoveryState.NOT_STARTED), - getMcpServerCount: vi.fn().mockReturnValue(0), - }; - mockConfig.getMcpClientManager = () => mockMcpClientManager as any; - - const { result } = renderTestHook(); - - await act(async () => { - await result.current.submitQuery('test query'); - }); - - expect(coreEvents.emitFeedback).not.toHaveBeenCalledWith( - 'info', - 'Waiting for MCP servers to initialize... Slash commands are still available.', - ); - expect(mockSendMessageStream).toHaveBeenCalled(); - }); - - it('should NOT block slash commands even when discovery is in progress', async () => { - const mockMcpClientManager = { - getDiscoveryState: vi - .fn() - .mockReturnValue(MCPDiscoveryState.IN_PROGRESS), - getMcpServerCount: vi.fn().mockReturnValue(1), - }; - mockConfig.getMcpClientManager = () => mockMcpClientManager as any; - - const { result } = renderTestHook(); - - await act(async () => { - await result.current.submitQuery('/help'); - }); - - expect(coreEvents.emitFeedback).not.toHaveBeenCalledWith( - 'info', - 'Waiting for MCP servers to initialize... Slash commands are still available.', - ); - }); - }); - describe('handleFinishedEvent', () => { it('should add info message for MAX_TOKENS finish reason', async () => { // Setup mock to return a stream with MAX_TOKENS finish reason @@ -3270,68 +3203,4 @@ describe('useGeminiStream', () => { }); }); }); - - describe('MCP Server Initialization', () => { - it('should allow slash commands to run while MCP servers are initializing', async () => { - const mockMcpClientManager = { - getDiscoveryState: vi - .fn() - .mockReturnValue(MCPDiscoveryState.IN_PROGRESS), - getMcpServerCount: vi.fn().mockReturnValue(1), - }; - mockConfig.getMcpClientManager = () => mockMcpClientManager as any; - - const { result } = renderTestHook(); - - await act(async () => { - await result.current.submitQuery('/help'); - }); - - // Slash command should be handled, and no Gemini call should be made. - expect(mockHandleSlashCommand).toHaveBeenCalledWith('/help'); - expect(coreEvents.emitFeedback).not.toHaveBeenCalled(); - }); - - it('should block normal prompts and provide feedback while MCP servers are initializing', async () => { - const mockMcpClientManager = { - getDiscoveryState: vi - .fn() - .mockReturnValue(MCPDiscoveryState.IN_PROGRESS), - getMcpServerCount: vi.fn().mockReturnValue(1), - }; - mockConfig.getMcpClientManager = () => mockMcpClientManager as any; - const { result } = renderTestHook(); - - await act(async () => { - await result.current.submitQuery('a normal prompt'); - }); - - // No slash command, no Gemini call, but feedback should be emitted. - expect(mockHandleSlashCommand).not.toHaveBeenCalled(); - expect(mockSendMessageStream).not.toHaveBeenCalled(); - expect(coreEvents.emitFeedback).toHaveBeenCalledWith( - 'info', - 'Waiting for MCP servers to initialize... Slash commands are still available.', - ); - }); - - it('should allow normal prompts to run when MCP servers are finished initializing', async () => { - const mockMcpClientManager = { - getDiscoveryState: vi.fn().mockReturnValue(MCPDiscoveryState.COMPLETED), - getMcpServerCount: vi.fn().mockReturnValue(1), - }; - mockConfig.getMcpClientManager = () => mockMcpClientManager as any; - - const { result } = renderTestHook(); - - await act(async () => { - await result.current.submitQuery('a normal prompt'); - }); - - // Prompt should be sent to Gemini. - expect(mockHandleSlashCommand).not.toHaveBeenCalled(); - expect(mockSendMessageStream).toHaveBeenCalled(); - expect(coreEvents.emitFeedback).not.toHaveBeenCalled(); - }); - }); }); diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index 9c44b0ee11..b6da0a84d5 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -31,7 +31,6 @@ import { ValidationRequiredError, coreEvents, CoreEvent, - MCPDiscoveryState, } from '@google/gemini-cli-core'; import type { Config, @@ -991,25 +990,6 @@ export const useGeminiStream = ( async ({ metadata: spanMetadata }) => { spanMetadata.input = query; - const discoveryState = config - .getMcpClientManager() - ?.getDiscoveryState(); - const mcpServerCount = - config.getMcpClientManager()?.getMcpServerCount() ?? 0; - if ( - !options?.isContinuation && - typeof query === 'string' && - !isSlashCommand(query.trim()) && - mcpServerCount > 0 && - discoveryState !== MCPDiscoveryState.COMPLETED - ) { - coreEvents.emitFeedback( - 'info', - 'Waiting for MCP servers to initialize... Slash commands are still available.', - ); - return; - } - const queryId = `${Date.now()}-${Math.random()}`; activeQueryIdRef.current = queryId; if ( diff --git a/packages/cli/src/ui/hooks/useMcpStatus.test.tsx b/packages/cli/src/ui/hooks/useMcpStatus.test.tsx new file mode 100644 index 0000000000..0311f03c63 --- /dev/null +++ b/packages/cli/src/ui/hooks/useMcpStatus.test.tsx @@ -0,0 +1,97 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi, beforeEach, type Mock } from 'vitest'; +import { act } from 'react'; +import { render } from '../../test-utils/render.js'; +import { useMcpStatus } from './useMcpStatus.js'; +import { + MCPDiscoveryState, + type Config, + CoreEvent, + coreEvents, +} from '@google/gemini-cli-core'; + +describe('useMcpStatus', () => { + let mockConfig: Config; + let mockMcpClientManager: { + getDiscoveryState: Mock<() => MCPDiscoveryState>; + getMcpServerCount: Mock<() => number>; + }; + + beforeEach(() => { + mockMcpClientManager = { + getDiscoveryState: vi.fn().mockReturnValue(MCPDiscoveryState.NOT_STARTED), + getMcpServerCount: vi.fn().mockReturnValue(0), + }; + + mockConfig = { + getMcpClientManager: vi.fn().mockReturnValue(mockMcpClientManager), + } as unknown as Config; + }); + + const renderMcpStatusHook = (config: Config) => { + let hookResult: ReturnType; + function TestComponent({ config }: { config: Config }) { + hookResult = useMcpStatus(config); + return null; + } + render(); + return { + result: { + get current() { + return hookResult; + }, + }, + }; + }; + + it('should initialize with correct values (no servers)', () => { + const { result } = renderMcpStatusHook(mockConfig); + + expect(result.current.discoveryState).toBe(MCPDiscoveryState.NOT_STARTED); + expect(result.current.mcpServerCount).toBe(0); + expect(result.current.isMcpReady).toBe(true); + }); + + it('should initialize with correct values (with servers, not started)', () => { + mockMcpClientManager.getMcpServerCount.mockReturnValue(1); + const { result } = renderMcpStatusHook(mockConfig); + + expect(result.current.isMcpReady).toBe(false); + }); + + it('should not be ready while in progress', () => { + mockMcpClientManager.getDiscoveryState.mockReturnValue( + MCPDiscoveryState.IN_PROGRESS, + ); + mockMcpClientManager.getMcpServerCount.mockReturnValue(1); + const { result } = renderMcpStatusHook(mockConfig); + + expect(result.current.isMcpReady).toBe(false); + }); + + it('should update state when McpClientUpdate is emitted', () => { + mockMcpClientManager.getMcpServerCount.mockReturnValue(1); + mockMcpClientManager.getDiscoveryState.mockReturnValue( + MCPDiscoveryState.IN_PROGRESS, + ); + const { result } = renderMcpStatusHook(mockConfig); + + expect(result.current.isMcpReady).toBe(false); + + mockMcpClientManager.getDiscoveryState.mockReturnValue( + MCPDiscoveryState.COMPLETED, + ); + + act(() => { + coreEvents.emit(CoreEvent.McpClientUpdate, new Map()); + }); + + expect(result.current.discoveryState).toBe(MCPDiscoveryState.COMPLETED); + expect(result.current.isMcpReady).toBe(true); + }); +}); diff --git a/packages/cli/src/ui/hooks/useMcpStatus.ts b/packages/cli/src/ui/hooks/useMcpStatus.ts new file mode 100644 index 0000000000..cc4d325cd7 --- /dev/null +++ b/packages/cli/src/ui/hooks/useMcpStatus.ts @@ -0,0 +1,51 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { useEffect, useState } from 'react'; +import { + type Config, + coreEvents, + MCPDiscoveryState, + CoreEvent, +} from '@google/gemini-cli-core'; + +export function useMcpStatus(config: Config) { + const [discoveryState, setDiscoveryState] = useState( + () => + config.getMcpClientManager()?.getDiscoveryState() ?? + MCPDiscoveryState.NOT_STARTED, + ); + + const [mcpServerCount, setMcpServerCount] = useState( + () => config.getMcpClientManager()?.getMcpServerCount() ?? 0, + ); + + useEffect(() => { + const onChange = () => { + const manager = config.getMcpClientManager(); + if (manager) { + setDiscoveryState(manager.getDiscoveryState()); + setMcpServerCount(manager.getMcpServerCount()); + } + }; + + coreEvents.on(CoreEvent.McpClientUpdate, onChange); + return () => { + coreEvents.off(CoreEvent.McpClientUpdate, onChange); + }; + }, [config]); + + // We are ready if discovery has completed, OR if it hasn't even started and there are no servers. + const isMcpReady = + discoveryState === MCPDiscoveryState.COMPLETED || + (discoveryState === MCPDiscoveryState.NOT_STARTED && mcpServerCount === 0); + + return { + discoveryState, + mcpServerCount, + isMcpReady, + }; +} diff --git a/packages/cli/src/ui/hooks/useMessageQueue.test.tsx b/packages/cli/src/ui/hooks/useMessageQueue.test.tsx index b3464af635..5b05d2a9f1 100644 --- a/packages/cli/src/ui/hooks/useMessageQueue.test.tsx +++ b/packages/cli/src/ui/hooks/useMessageQueue.test.tsx @@ -28,6 +28,7 @@ describe('useMessageQueue', () => { isConfigInitialized: boolean; streamingState: StreamingState; submitQuery: (query: string) => void; + isMcpReady: boolean; }) => { let hookResult: ReturnType; function TestComponent(props: typeof initialProps) { @@ -51,6 +52,7 @@ describe('useMessageQueue', () => { isConfigInitialized: true, streamingState: StreamingState.Idle, submitQuery: mockSubmitQuery, + isMcpReady: true, }); expect(result.current.messageQueue).toEqual([]); @@ -62,6 +64,7 @@ describe('useMessageQueue', () => { isConfigInitialized: true, streamingState: StreamingState.Responding, submitQuery: mockSubmitQuery, + isMcpReady: true, }); act(() => { @@ -80,6 +83,7 @@ describe('useMessageQueue', () => { isConfigInitialized: true, streamingState: StreamingState.Responding, submitQuery: mockSubmitQuery, + isMcpReady: true, }); act(() => { @@ -100,6 +104,7 @@ describe('useMessageQueue', () => { isConfigInitialized: true, streamingState: StreamingState.Responding, submitQuery: mockSubmitQuery, + isMcpReady: true, }); act(() => { @@ -120,6 +125,7 @@ describe('useMessageQueue', () => { isConfigInitialized: true, streamingState: StreamingState.Responding, submitQuery: mockSubmitQuery, + isMcpReady: true, }); act(() => { @@ -133,11 +139,12 @@ describe('useMessageQueue', () => { ); }); - it('should auto-submit queued messages when transitioning to Idle', async () => { + it('should auto-submit queued messages when transitioning to Idle and MCP is ready', async () => { const { result, rerender } = renderMessageQueueHook({ isConfigInitialized: true, streamingState: StreamingState.Responding, submitQuery: mockSubmitQuery, + isMcpReady: true, }); // Add some messages @@ -157,11 +164,37 @@ describe('useMessageQueue', () => { }); }); + it('should wait for MCP readiness before auto-submitting', async () => { + const { result, rerender } = renderMessageQueueHook({ + isConfigInitialized: true, + streamingState: StreamingState.Idle, + submitQuery: mockSubmitQuery, + isMcpReady: false, + }); + + // Add some messages while Idle but MCP not ready + act(() => { + result.current.addMessage('Delayed message'); + }); + + expect(result.current.messageQueue).toEqual(['Delayed message']); + expect(mockSubmitQuery).not.toHaveBeenCalled(); + + // Transition MCP to ready + rerender({ isMcpReady: true }); + + await waitFor(() => { + expect(mockSubmitQuery).toHaveBeenCalledWith('Delayed message'); + expect(result.current.messageQueue).toEqual([]); + }); + }); + it('should not auto-submit when queue is empty', () => { const { rerender } = renderMessageQueueHook({ isConfigInitialized: true, streamingState: StreamingState.Responding, submitQuery: mockSubmitQuery, + isMcpReady: true, }); // Transition to Idle with empty queue @@ -175,6 +208,7 @@ describe('useMessageQueue', () => { isConfigInitialized: true, streamingState: StreamingState.Responding, submitQuery: mockSubmitQuery, + isMcpReady: true, }); // Add messages @@ -194,6 +228,7 @@ describe('useMessageQueue', () => { isConfigInitialized: true, streamingState: StreamingState.Idle, submitQuery: mockSubmitQuery, + isMcpReady: true, }); // Start responding @@ -235,6 +270,7 @@ describe('useMessageQueue', () => { isConfigInitialized: true, streamingState: StreamingState.Responding, submitQuery: mockSubmitQuery, + isMcpReady: true, }); // Add multiple messages @@ -265,6 +301,7 @@ describe('useMessageQueue', () => { isConfigInitialized: true, streamingState: StreamingState.Responding, submitQuery: mockSubmitQuery, + isMcpReady: true, }); let poppedMessages: string | undefined = 'not-undefined'; @@ -281,6 +318,7 @@ describe('useMessageQueue', () => { isConfigInitialized: true, streamingState: StreamingState.Responding, submitQuery: mockSubmitQuery, + isMcpReady: false, }); act(() => { @@ -301,6 +339,7 @@ describe('useMessageQueue', () => { isConfigInitialized: true, streamingState: StreamingState.Responding, submitQuery: mockSubmitQuery, + isMcpReady: false, }); act(() => { @@ -330,6 +369,7 @@ describe('useMessageQueue', () => { isConfigInitialized: true, streamingState: StreamingState.Responding, submitQuery: mockSubmitQuery, + isMcpReady: false, }); // Add messages diff --git a/packages/cli/src/ui/hooks/useMessageQueue.ts b/packages/cli/src/ui/hooks/useMessageQueue.ts index 58a1e890f3..93bb0ab7a9 100644 --- a/packages/cli/src/ui/hooks/useMessageQueue.ts +++ b/packages/cli/src/ui/hooks/useMessageQueue.ts @@ -11,6 +11,7 @@ export interface UseMessageQueueOptions { isConfigInitialized: boolean; streamingState: StreamingState; submitQuery: (query: string) => void; + isMcpReady: boolean; } export interface UseMessageQueueReturn { @@ -30,6 +31,7 @@ export function useMessageQueue({ isConfigInitialized, streamingState, submitQuery, + isMcpReady, }: UseMessageQueueOptions): UseMessageQueueReturn { const [messageQueue, setMessageQueue] = useState([]); @@ -67,6 +69,7 @@ export function useMessageQueue({ if ( isConfigInitialized && streamingState === StreamingState.Idle && + isMcpReady && messageQueue.length > 0 ) { // Combine all messages with double newlines for clarity @@ -75,7 +78,13 @@ export function useMessageQueue({ setMessageQueue([]); submitQuery(combinedMessage); } - }, [isConfigInitialized, streamingState, messageQueue, submitQuery]); + }, [ + isConfigInitialized, + streamingState, + isMcpReady, + messageQueue, + submitQuery, + ]); return { messageQueue, diff --git a/packages/core/src/tools/mcp-client-manager.ts b/packages/core/src/tools/mcp-client-manager.ts index 657699ca1c..743d7adb47 100644 --- a/packages/core/src/tools/mcp-client-manager.ts +++ b/packages/core/src/tools/mcp-client-manager.ts @@ -279,6 +279,7 @@ export class McpClientManager { if (currentPromise === this.discoveryPromise) { this.discoveryPromise = undefined; this.discoveryState = MCPDiscoveryState.COMPLETED; + this.eventEmitter?.emit('mcp-client-update', this.clients); } }) .catch(() => {}); // Prevents unhandled rejection from the .finally branch @@ -307,6 +308,12 @@ export class McpClientManager { this.cliConfig.getMcpServerCommand(), ); + if (Object.keys(servers).length === 0) { + this.discoveryState = MCPDiscoveryState.COMPLETED; + this.eventEmitter?.emit('mcp-client-update', this.clients); + return; + } + // Set state synchronously before any await yields control if (!this.discoveryPromise) { this.discoveryState = MCPDiscoveryState.IN_PROGRESS;