Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
63 changes: 62 additions & 1 deletion packages/cli/src/ui/AppContainer.test.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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();
});
});
});
23 changes: 21 additions & 2 deletions packages/cli/src/ui/AppContainer.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -1321,6 +1330,7 @@ Logging in with Google... Restarting Gemini CLI to continue.
streamingState,
submitQuery,
isMcpReady,
isCompressing,
});

cancelHandlerRef.current = useCallback(
Expand Down Expand Up @@ -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)
) {
Comment thread
cocosheng-g marked this conversation as resolved.
Comment thread
cocosheng-g marked this conversation as resolved.
if (!isSlash) {
const permissions = await checkPermissions(submittedValue, config);
if (permissions.length > 0) {
Expand All @@ -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
Expand All @@ -1458,6 +1476,7 @@ Logging in with Google... Restarting Gemini CLI to continue.
slashCommands,
isMcpReady,
streamingState,
isCompressing,
messageQueue.length,
pendingHistoryItems,
config,
Expand Down
5 changes: 5 additions & 0 deletions packages/cli/src/ui/commands/compressCommand.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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({
Expand All @@ -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({
Expand All @@ -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);
});

Expand Down
73 changes: 38 additions & 35 deletions packages/cli/src/ui/commands/compressCommand.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
})();
},
};
49 changes: 49 additions & 0 deletions packages/cli/src/ui/hooks/useMessageQueue.test.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down Expand Up @@ -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([]);
});
});
});
});
4 changes: 4 additions & 0 deletions packages/cli/src/ui/hooks/useMessageQueue.ts
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ export interface UseMessageQueueOptions {
streamingState: StreamingState;
submitQuery: (query: string) => void;
isMcpReady: boolean;
isCompressing?: boolean;
}

export interface UseMessageQueueReturn {
Expand All @@ -32,6 +33,7 @@ export function useMessageQueue({
streamingState,
submitQuery,
isMcpReady,
isCompressing = false,
}: UseMessageQueueOptions): UseMessageQueueReturn {
const [messageQueue, setMessageQueue] = useState<string[]>([]);

Expand Down Expand Up @@ -69,6 +71,7 @@ export function useMessageQueue({
if (
isConfigInitialized &&
streamingState === StreamingState.Idle &&
!isCompressing &&
isMcpReady &&
messageQueue.length > 0
) {
Expand All @@ -84,6 +87,7 @@ export function useMessageQueue({
isMcpReady,
messageQueue,
submitQuery,
isCompressing,
]);

return {
Expand Down
Loading