Allow prompt queueing during MCP initialization (#17395)

This commit is contained in:
Adib234
2026-01-23 18:32:06 -05:00
committed by GitHub
parent 93da9817b6
commit 12a5490bcf
8 changed files with 234 additions and 155 deletions

View File

@@ -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(() => {

View File

@@ -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();
});
});
});

View File

@@ -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 (

View File

@@ -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<typeof useMcpStatus>;
function TestComponent({ config }: { config: Config }) {
hookResult = useMcpStatus(config);
return null;
}
render(<TestComponent config={config} />);
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);
});
});

View File

@@ -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<MCPDiscoveryState>(
() =>
config.getMcpClientManager()?.getDiscoveryState() ??
MCPDiscoveryState.NOT_STARTED,
);
const [mcpServerCount, setMcpServerCount] = useState<number>(
() => 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,
};
}

View File

@@ -28,6 +28,7 @@ describe('useMessageQueue', () => {
isConfigInitialized: boolean;
streamingState: StreamingState;
submitQuery: (query: string) => void;
isMcpReady: boolean;
}) => {
let hookResult: ReturnType<typeof useMessageQueue>;
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

View File

@@ -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<string[]>([]);
@@ -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,

View File

@@ -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;