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
117 changes: 116 additions & 1 deletion packages/cli/src/ui/AppContainer.test.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ import {
UIActionsContext,
type UIActions,
} from './contexts/UIActionsContext.js';
import { ToolCallStatus } from './types.js';
import { useContext } from 'react';

// Mock useStdout to capture terminal title writes
Expand Down Expand Up @@ -245,6 +246,7 @@ describe('AppContainer State Management', () => {
getQueuedMessagesText: vi.fn().mockReturnValue(''),
popAllMessages: vi.fn().mockReturnValue(null),
drainQueue: vi.fn().mockReturnValue([]),
popNextSegment: vi.fn().mockReturnValue(null),
});
mockedUseAutoAcceptIndicator.mockReturnValue(false);
mockedUseGitBranchName.mockReturnValue('main');
Expand Down Expand Up @@ -459,6 +461,7 @@ describe('AppContainer State Management', () => {
getQueuedMessagesText: vi.fn().mockReturnValue(''),
popAllMessages: vi.fn().mockReturnValue(null),
drainQueue: vi.fn().mockReturnValue([]),
popNextSegment: vi.fn().mockReturnValue(null),
});

render(
Expand All @@ -476,6 +479,44 @@ describe('AppContainer State Management', () => {
expect(mockQueueMessage).not.toHaveBeenCalled();
});

it('submits slash commands immediately instead of queueing while idle', () => {
const mockSubmitQuery = vi.fn();
const mockQueueMessage = vi.fn();

mockedUseGeminiStream.mockReturnValue({
streamingState: 'idle',
submitQuery: mockSubmitQuery,
initError: null,
pendingHistoryItems: [],
thought: null,
cancelOngoingRequest: vi.fn(),
retryLastPrompt: vi.fn(),
});
mockedUseMessageQueue.mockReturnValue({
messageQueue: [],
addMessage: mockQueueMessage,
clearQueue: vi.fn(),
getQueuedMessagesText: vi.fn().mockReturnValue(''),
popAllMessages: vi.fn().mockReturnValue(null),
drainQueue: vi.fn().mockReturnValue([]),
popNextSegment: vi.fn().mockReturnValue(null),
});

render(
<AppContainer
config={mockConfig}
settings={mockSettings}
version="1.0.0"
initializationResult={mockInitResult}
/>,
);

capturedUIActions.handleFinalSubmit('/model');

expect(mockSubmitQuery).toHaveBeenCalledWith('/model');
expect(mockQueueMessage).not.toHaveBeenCalled();
});

it.each(['exit', 'quit', ':q', ':q!', ':wq', ':wq!'])(
'routes bare "%s" to /quit instead of sending as a message',
(command) => {
Expand All @@ -497,6 +538,7 @@ describe('AppContainer State Management', () => {
getQueuedMessagesText: vi.fn().mockReturnValue(''),
popAllMessages: vi.fn().mockReturnValue(null),
drainQueue: vi.fn().mockReturnValue([]),
popNextSegment: vi.fn().mockReturnValue(null),
});

render(
Expand Down Expand Up @@ -577,6 +619,7 @@ describe('AppContainer State Management', () => {
getQueuedMessagesText: vi.fn().mockReturnValue(''),
popAllMessages: vi.fn().mockReturnValue(null),
drainQueue: vi.fn().mockReturnValue([]),
popNextSegment: vi.fn().mockReturnValue(null),
});

render(
Expand Down Expand Up @@ -605,6 +648,7 @@ describe('AppContainer State Management', () => {
it('moves queued follow-up messages into an empty buffer on cancel', async () => {
const mockSetText = vi.fn();
const mockPopAllMessages = vi.fn().mockReturnValue('queued follow-up');
const mockClearQueue = vi.fn();
mockedUseTextBuffer.mockReturnValue({
text: '',
setText: mockSetText,
Expand All @@ -626,10 +670,11 @@ describe('AppContainer State Management', () => {
mockedUseMessageQueue.mockReturnValue({
messageQueue: ['queued follow-up'],
addMessage: vi.fn(),
clearQueue: vi.fn(),
clearQueue: mockClearQueue,
getQueuedMessagesText: vi.fn().mockReturnValue('queued follow-up'),
popAllMessages: mockPopAllMessages,
drainQueue: vi.fn().mockReturnValue(['queued follow-up']),
popNextSegment: vi.fn().mockReturnValue('queued follow-up'),
});

render(
Expand All @@ -653,6 +698,75 @@ describe('AppContainer State Management', () => {
expect.stringContaining('the previous prompt'),
);
expect(mockPopAllMessages).toHaveBeenCalled();
// popAllForEdit drains the queue internally, so the cancel handler
// does not need to call clearQueue separately on this path.
expect(mockClearQueue).not.toHaveBeenCalled();
});

it('drops the queue when cancelling during tool execution', async () => {
// Simulates: user asks for a shell tool (e.g. sleep 30), queues
// `/model` and `hi` while the tool is running, then hits Ctrl+C.
// The cancel must clear BOTH the buffer and the queue so that
// `hi` does not auto-fire once the tool settles and the app
// returns to idle.
const mockSetText = vi.fn();
const mockClearQueue = vi.fn();
mockedUseTextBuffer.mockReturnValue({
text: '',
setText: mockSetText,
});
installCancelCapture({
streamingState: 'responding',
submitQuery: vi.fn(),
initError: null,
pendingHistoryItems: [
{
type: 'tool_group',
tools: [
{
callId: 'call-1',
name: 'run_shell_command',
description: 'sleep 30',
status: ToolCallStatus.Executing,
resultDisplay: undefined,
confirmationDetails: undefined,
renderOutputAsMarkdown: false,
},
],
},
],
thought: null,
cancelOngoingRequest: vi.fn(),
retryLastPrompt: vi.fn(),
});
mockedUseMessageQueue.mockReturnValue({
messageQueue: ['/model', 'hi'],
addMessage: vi.fn(),
clearQueue: mockClearQueue,
getQueuedMessagesText: vi.fn().mockReturnValue('/model\n\nhi'),
popAllMessages: vi.fn().mockReturnValue('/model'),
drainQueue: vi.fn().mockReturnValue([]),
popNextSegment: vi.fn().mockReturnValue('/model'),
});

render(
<AppContainer
config={mockConfig}
settings={mockSettings}
version="1.0.0"
initializationResult={mockInitResult}
/>,
);

await Promise.resolve();
await Promise.resolve();

triggerCancel();

// Buffer cleared and queue dropped — same "abort and redirect"
// contract as the non-tool cancel path.
expect(mockSetText).toHaveBeenCalledWith('');
expect(mockClearQueue).toHaveBeenCalled();
});

it('preserves an in-progress draft when restoring queued messages on cancel', async () => {
Expand Down Expand Up @@ -680,6 +794,7 @@ describe('AppContainer State Management', () => {
getQueuedMessagesText: vi.fn().mockReturnValue('queued follow-up'),
popAllMessages: vi.fn().mockReturnValue('queued follow-up'),
drainQueue: vi.fn().mockReturnValue(['queued follow-up']),
popNextSegment: vi.fn().mockReturnValue('queued follow-up'),
});

render(
Expand Down
75 changes: 56 additions & 19 deletions packages/cli/src/ui/AppContainer.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -88,7 +88,7 @@ import { useTextBuffer } from './components/shared/text-buffer.js';
import { useLogger } from './hooks/useLogger.js';
import { useGeminiStream } from './hooks/useGeminiStream.js';
import { useVim } from './hooks/vim.js';
import { isBtwCommand } from './utils/commandUtils.js';
import { isBtwCommand, isSlashCommand } from './utils/commandUtils.js';
import { type LoadedSettings, SettingScope } from '../config/settings.js';
import { type InitializationResult } from '../core/initializer.js';
import { useFocus } from './hooks/useFocus.js';
Expand Down Expand Up @@ -873,16 +873,18 @@ export const AppContainer = (props: AppContainerProps) => {
disabled: agentViewState.activeView !== 'main',
});

const { messageQueue, addMessage, popAllMessages, drainQueue } =
useMessageQueue({
isConfigInitialized,
streamingState,
submitQuery,
});
const {
messageQueue,
addMessage,
clearQueue,
popAllMessages,
drainQueue,
popNextSegment,
} = useMessageQueue();

// Bridge message queue to mid-turn drain via ref.
// drainQueue reads the synchronous queueRef inside the hook, so it
// stays consistent with popAllMessages even before React re-renders.
// stays consistent with popNextSegment even before React re-renders.
midTurnDrainRef.current = drainQueue;

// Connect remote input watcher to submitQuery for bidirectional sync.
Expand Down Expand Up @@ -1158,6 +1160,14 @@ export const AppContainer = (props: AppContainerProps) => {
speculationRef.current = IDLE_SPECULATION;
}

if (
streamingState === StreamingState.Idle &&
isSlashCommand(submittedValue)
) {
void submitQuery(submittedValue);
return;
}

addMessage(submittedValue);
},
[
Expand Down Expand Up @@ -1204,28 +1214,22 @@ export const AppContainer = (props: AppContainerProps) => {
...pendingGeminiHistoryItems,
];
if (isToolExecuting(pendingHistoryItems)) {
buffer.setText(''); // Just clear the prompt
// Tool-cancel: drop both buffer and queue so nothing auto-fires later.
buffer.setText('');
clearQueue();
return;
}

// Move any queued follow-up messages back into the buffer so the user
// can edit or resubmit them. Otherwise leave the buffer alone — in
// particular, do NOT repopulate it with the previous prompt; the user
// can still recall it via history navigation (Up/Ctrl+P).
//
// popAllMessages is atomic via the queue's synchronous ref, matching
// the drain behavior used during tool completion.
// Restore queued input joined into the buffer for editing.
const popped = popAllMessages();
if (popped) {
const currentText = buffer.text;
// Preserve any in-progress draft the user typed since submitting (this
// is reachable via Ctrl+C cancel, which fires regardless of buffer
// content). Mirrors the popQueueIntoInput convention in InputPrompt.
buffer.setText(currentText ? `${popped}\n${currentText}` : popped);
}
}, [
buffer,
popAllMessages,
clearQueue,
pendingSlashCommandHistoryItems,
pendingGeminiHistoryItems,
]);
Expand Down Expand Up @@ -2029,6 +2033,39 @@ export const AppContainer = (props: AppContainerProps) => {
isExtensionsManagerDialogOpen;
dialogsVisibleRef.current = dialogsVisible;

// Drain queued messages when idle. `queueDrainNonce` re-fires the effect
// after each submission settles so multi-step queues drain end-to-end.
const queueDrainingRef = useRef(false);
const [queueDrainNonce, setQueueDrainNonce] = useState(0);
useEffect(() => {
if (queueDrainingRef.current) return;
if (!isConfigInitialized) return;
if (streamingState !== StreamingState.Idle) return;
if (dialogsVisible) return;
if (messageQueue.length === 0) return;

// Two-phase: batch plain prompts as one turn, else pop next slash command.
const plainPrompts = drainQueue();
const submission =
plainPrompts.length > 0 ? plainPrompts.join('\n\n') : popNextSegment();
if (submission === null) return;

queueDrainingRef.current = true;
Promise.resolve(submitQuery(submission)).finally(() => {
queueDrainingRef.current = false;
setQueueDrainNonce((n) => n + 1);
});
}, [
isConfigInitialized,
streamingState,
dialogsVisible,
messageQueue,
drainQueue,
popNextSegment,
submitQuery,
queueDrainNonce,
]);

const {
isFeedbackDialogOpen,
openFeedbackDialog,
Expand Down
Loading
Loading