feat: allow queuing messages during compression (#24071) (#26506)

This commit is contained in:
Coco Sheng
2026-05-05 13:52:08 -04:00
committed by GitHub
parent 7cc19c2a1b
commit e80d7cc083
6 changed files with 179 additions and 38 deletions
+62 -1
View File
@@ -100,7 +100,7 @@ import { type LoadedSettings } from '../config/settings.js';
import { createMockSettings } from '../test-utils/settings.js';
import type { InitializationResult } from '../core/initializer.js';
import { useQuotaAndFallback } from './hooks/useQuotaAndFallback.js';
import { StreamingState } from './types.js';
import { StreamingState, MessageType } from './types.js';
import { UIStateContext, type UIState } from './contexts/UIStateContext.js';
import {
UIActionsContext,
@@ -3576,4 +3576,65 @@ describe('AppContainer State Management', () => {
unmount();
});
});
describe('Compression Queuing', () => {
beforeEach(async () => {
const { checkPermissions } = await import(
'./hooks/atCommandProcessor.js'
);
vi.mocked(checkPermissions).mockResolvedValue([]);
vi.spyOn(mockConfig, 'isModelSteeringEnabled').mockReturnValue(true);
const actual = await vi.importActual('./hooks/useMessageQueue.js');
// eslint-disable-next-line @typescript-eslint/no-explicit-any
const { useMessageQueue: realUseMessageQueue } = actual as any;
mockedUseMessageQueue.mockImplementation(realUseMessageQueue);
// Start compression by mocking pendingHistoryItems to include a pending compression
mockedUseGeminiStream.mockImplementation(() => ({
...DEFAULT_GEMINI_STREAM_MOCK,
pendingHistoryItems: [
{
type: MessageType.COMPRESSION,
compression: {
isPending: true,
originalTokenCount: null,
newTokenCount: null,
compressionStatus: null,
},
},
],
}));
});
it('queues messages during compression instead of handling as steering hints', async () => {
const { unmount } = await act(async () => renderAppContainer());
// Verify state isolation
expect(capturedUIState.streamingState).toBe(StreamingState.Idle);
// Submit a message
await act(async () =>
capturedUIActions.handleFinalSubmit('follow up message'),
);
// Verify it was queued, not submitted as steering hint
expect(capturedUIState.messageQueue).toContain('follow up message');
unmount();
});
it('executes slash commands immediately during compression', async () => {
const { unmount } = await act(async () => renderAppContainer());
// Submit a slash command
await act(async () => capturedUIActions.handleFinalSubmit('/help'));
// Verify it was NOT queued
expect(capturedUIState.messageQueue).not.toContain('/help');
unmount();
});
});
});
+21 -2
View File
@@ -1310,6 +1310,15 @@ Logging in with Google... Restarting Gemini CLI to continue.
const { isMcpReady } = useMcpStatus(config);
const isCompressing = useMemo(
() =>
pendingHistoryItems.some(
(item) =>
item.type === MessageType.COMPRESSION && item.compression.isPending,
),
[pendingHistoryItems],
);
const {
messageQueue,
addMessage,
@@ -1321,6 +1330,7 @@ Logging in with Google... Restarting Gemini CLI to continue.
streamingState,
submitQuery,
isMcpReady,
isCompressing,
});
cancelHandlerRef.current = useCallback(
@@ -1415,7 +1425,10 @@ Logging in with Google... Restarting Gemini CLI to continue.
}
const isMcpOrConfigReady = isConfigInitialized && isMcpReady;
if ((isSlash && isConfigInitialized) || (isIdle && isMcpOrConfigReady)) {
if (
(isSlash && isConfigInitialized) ||
(!isCompressing && isIdle && isMcpOrConfigReady)
) {
if (!isSlash) {
const permissions = await checkPermissions(submittedValue, config);
if (permissions.length > 0) {
@@ -1438,7 +1451,12 @@ Logging in with Google... Restarting Gemini CLI to continue.
void submitQuery(submittedValue);
} else {
// Check messageQueue.length === 0 to only notify on the first queued item
if (isIdle && !isMcpOrConfigReady && messageQueue.length === 0) {
if (
isIdle &&
!isCompressing &&
!isMcpOrConfigReady &&
messageQueue.length === 0
) {
coreEvents.emitFeedback(
'info',
!isConfigInitialized
@@ -1458,6 +1476,7 @@ Logging in with Google... Restarting Gemini CLI to continue.
slashCommands,
isMcpReady,
streamingState,
isCompressing,
messageQueue.length,
pendingHistoryItems,
config,
@@ -42,6 +42,7 @@ describe('compressCommand', () => {
},
};
await compressCommand.action!(context, '');
await new Promise((r) => setTimeout(r, 0));
expect(context.ui.addItem).toHaveBeenCalledWith(
expect.objectContaining({
type: MessageType.ERROR,
@@ -62,6 +63,7 @@ describe('compressCommand', () => {
mockTryCompressChat.mockResolvedValue(compressedResult);
await compressCommand.action!(context, '');
await new Promise((r) => setTimeout(r, 0));
expect(context.ui.setPendingItem).toHaveBeenNthCalledWith(1, {
type: MessageType.COMPRESSION,
@@ -98,6 +100,7 @@ describe('compressCommand', () => {
mockTryCompressChat.mockResolvedValue(null);
await compressCommand.action!(context, '');
await new Promise((r) => setTimeout(r, 0));
expect(context.ui.addItem).toHaveBeenCalledWith(
expect.objectContaining({
@@ -114,6 +117,7 @@ describe('compressCommand', () => {
mockTryCompressChat.mockRejectedValue(error);
await compressCommand.action!(context, '');
await new Promise((r) => setTimeout(r, 0));
expect(context.ui.addItem).toHaveBeenCalledWith(
expect.objectContaining({
@@ -128,6 +132,7 @@ describe('compressCommand', () => {
it('should clear the pending item in a finally block', async () => {
mockTryCompressChat.mockRejectedValue(new Error('some error'));
await compressCommand.action!(context, '');
await new Promise((r) => setTimeout(r, 0));
expect(context.ui.setPendingItem).toHaveBeenCalledWith(null);
});
+38 -35
View File
@@ -36,48 +36,51 @@ export const compressCommand: SlashCommand = {
},
};
try {
ui.setPendingItem(pendingMessage);
const promptId = `compress-${Date.now()}`;
const compressed =
await context.services.agentContext?.geminiClient?.tryCompressChat(
promptId,
true,
);
if (compressed) {
ui.addItem(
{
type: MessageType.COMPRESSION,
compression: {
isPending: false,
originalTokenCount: compressed.originalTokenCount,
newTokenCount: compressed.newTokenCount,
compressionStatus: compressed.compressionStatus,
ui.setPendingItem(pendingMessage);
void (async () => {
try {
const promptId = `compress-${Date.now()}`;
const compressed =
await context.services.agentContext?.geminiClient?.tryCompressChat(
promptId,
true,
);
if (compressed) {
ui.addItem(
{
type: MessageType.COMPRESSION,
compression: {
isPending: false,
originalTokenCount: compressed.originalTokenCount,
newTokenCount: compressed.newTokenCount,
compressionStatus: compressed.compressionStatus,
},
} as HistoryItemCompression,
Date.now(),
);
} else {
ui.addItem(
{
type: MessageType.ERROR,
text: 'Failed to compress chat history.',
},
} as HistoryItemCompression,
Date.now(),
);
} else {
Date.now(),
);
}
} catch (e) {
ui.addItem(
{
type: MessageType.ERROR,
text: 'Failed to compress chat history.',
text: `Failed to compress chat history: ${
e instanceof Error ? e.message : String(e)
}`,
},
Date.now(),
);
} finally {
ui.setPendingItem(null);
}
} catch (e) {
ui.addItem(
{
type: MessageType.ERROR,
text: `Failed to compress chat history: ${
e instanceof Error ? e.message : String(e)
}`,
},
Date.now(),
);
} finally {
ui.setPendingItem(null);
}
})();
},
};
@@ -29,6 +29,7 @@ describe('useMessageQueue', () => {
streamingState: StreamingState;
submitQuery: (query: string) => void;
isMcpReady: boolean;
isCompressing?: boolean;
}) => {
let hookResult: ReturnType<typeof useMessageQueue>;
function TestComponent(props: typeof initialProps) {
@@ -402,4 +403,52 @@ describe('useMessageQueue', () => {
expect(result.current.messageQueue).toEqual([]);
});
});
describe('isCompressing logic', () => {
it('should not auto-submit when isCompressing is true, even if streamingState is Idle', async () => {
const { result } = await renderMessageQueueHook({
isConfigInitialized: true,
streamingState: StreamingState.Idle,
submitQuery: mockSubmitQuery,
isMcpReady: true,
isCompressing: true,
});
// Add messages
act(() => {
result.current.addMessage('Compression message');
});
expect(mockSubmitQuery).not.toHaveBeenCalled();
expect(result.current.messageQueue).toEqual(['Compression message']);
});
it('should auto-submit queued messages when isCompressing becomes false', async () => {
const { result, rerender } = await renderMessageQueueHook({
isConfigInitialized: true,
streamingState: StreamingState.Idle,
submitQuery: mockSubmitQuery,
isMcpReady: true,
isCompressing: true,
});
// Add messages
act(() => {
result.current.addMessage('Pending compression message 1');
result.current.addMessage('Pending compression message 2');
});
expect(mockSubmitQuery).not.toHaveBeenCalled();
// Transition isCompressing to false
rerender({ isCompressing: false });
await waitFor(() => {
expect(mockSubmitQuery).toHaveBeenCalledWith(
'Pending compression message 1\n\nPending compression message 2',
);
expect(result.current.messageQueue).toEqual([]);
});
});
});
});
@@ -12,6 +12,7 @@ export interface UseMessageQueueOptions {
streamingState: StreamingState;
submitQuery: (query: string) => void;
isMcpReady: boolean;
isCompressing?: boolean;
}
export interface UseMessageQueueReturn {
@@ -32,6 +33,7 @@ export function useMessageQueue({
streamingState,
submitQuery,
isMcpReady,
isCompressing = false,
}: UseMessageQueueOptions): UseMessageQueueReturn {
const [messageQueue, setMessageQueue] = useState<string[]>([]);
@@ -69,6 +71,7 @@ export function useMessageQueue({
if (
isConfigInitialized &&
streamingState === StreamingState.Idle &&
!isCompressing &&
isMcpReady &&
messageQueue.length > 0
) {
@@ -84,6 +87,7 @@ export function useMessageQueue({
isMcpReady,
messageQueue,
submitQuery,
isCompressing,
]);
return {