Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
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
46 changes: 45 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,48 @@ describe('AppContainer State Management', () => {
unmount();
});
});

describe('Compression Queuing', () => {
it('queues messages during compression instead of handling as steering hints', 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,
},
},
],
}));

const { unmount } = await act(async () => renderAppContainer());

// 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();
});
});
});
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) ||
(isIdle && !isCompressing && 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
75 changes: 39 additions & 36 deletions packages/cli/src/ui/commands/compressCommand.ts
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ export const compressCommand: SlashCommand = {
description: 'Compresses the context by replacing it with a summary',
kind: CommandKind.BUILT_IN,
autoExecute: true,
action: async (context) => {
action: (context) => {
Comment thread
cocosheng-g marked this conversation as resolved.
Outdated
const { ui } = context;
if (ui.pendingItem) {
ui.addItem(
Expand All @@ -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);
Comment thread
cocosheng-g marked this conversation as resolved.
}
} 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);
}
})();
},
};
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