diff --git a/integration-tests/terminal-capture/scenario-runner.ts b/integration-tests/terminal-capture/scenario-runner.ts index 93640694b..ff4920aa7 100644 --- a/integration-tests/terminal-capture/scenario-runner.ts +++ b/integration-tests/terminal-capture/scenario-runner.ts @@ -31,6 +31,24 @@ export interface FlowStep { capture?: string; /** Explicit screenshot: full scrollback buffer long image (standalone capture when no type) */ captureFull?: string; + /** + * Explicit sleep before executing this step (milliseconds). + * + * The runner's built-in idle detection (`idle(2000, 60000)`) works well for + * synchronous streaming, but cannot anticipate async responses that arrive + * after output has already stabilized (e.g., a /btw side-question whose API + * response is serialized behind a main streaming task). In such cases, the + * idle detector triggers too early and the async response is missed. + * + * Use `sleep` to bridge that gap — it inserts a fixed delay before the step + * runs, giving async operations time to complete. Optional; omitting it (or + * setting it to 0) has no effect on existing scenarios. + * + * @example + * // Wait 20s for a /btw response before capturing the result + * { sleep: 20000, capture: 'btw-answered.png' } + */ + sleep?: number; /** * Streaming capture: capture multiple screenshots during execution at intervals. * Useful for demonstrating real-time output like progress bars. @@ -159,6 +177,11 @@ export async function runScenario( const step = config.flow[i]; const label = `[${i + 1}/${config.flow.length}]`; + if (step.sleep && step.sleep > 0) { + console.log(` ${label} 💤 sleep: ${step.sleep}ms`); + await sleep(step.sleep); + } + if (step.type) { const display = step.type.length > 60 ? step.type.slice(0, 60) + '...' : step.type; diff --git a/packages/cli/src/nonInteractiveCliCommands.test.ts b/packages/cli/src/nonInteractiveCliCommands.test.ts index 76b29f3e0..c1c47c678 100644 --- a/packages/cli/src/nonInteractiveCliCommands.test.ts +++ b/packages/cli/src/nonInteractiveCliCommands.test.ts @@ -149,6 +149,33 @@ describe('handleSlashCommand', () => { } }); + it('should execute /btw when using the default allowed list', async () => { + const mockBtwCommand = { + name: 'btw', + description: 'Ask a side question', + kind: CommandKind.BUILT_IN, + action: vi.fn().mockResolvedValue({ + type: 'message', + messageType: 'info', + content: 'btw> question\nanswer', + }), + }; + mockGetCommands.mockReturnValue([mockBtwCommand]); + + const result = await handleSlashCommand( + '/btw question', + abortController, + mockConfig, + mockSettings, + ); + + expect(mockBtwCommand.action).toHaveBeenCalled(); + expect(result.type).toBe('message'); + if (result.type === 'message') { + expect(result.content).toBe('btw> question\nanswer'); + } + }); + it('should execute file commands regardless of allowed list', async () => { const mockFileCommand = { name: 'custom', diff --git a/packages/cli/src/nonInteractiveCliCommands.ts b/packages/cli/src/nonInteractiveCliCommands.ts index b089fa6c2..e6344f5d0 100644 --- a/packages/cli/src/nonInteractiveCliCommands.ts +++ b/packages/cli/src/nonInteractiveCliCommands.ts @@ -42,6 +42,7 @@ export const ALLOWED_BUILTIN_COMMANDS_NON_INTERACTIVE = [ 'init', 'summary', 'compress', + 'btw', 'bug', ] as const; diff --git a/packages/cli/src/services/BuiltinCommandLoader.ts b/packages/cli/src/services/BuiltinCommandLoader.ts index 73c233209..f379a39de 100644 --- a/packages/cli/src/services/BuiltinCommandLoader.ts +++ b/packages/cli/src/services/BuiltinCommandLoader.ts @@ -12,6 +12,7 @@ import { agentsCommand } from '../ui/commands/agentsCommand.js'; import { arenaCommand } from '../ui/commands/arenaCommand.js'; import { approvalModeCommand } from '../ui/commands/approvalModeCommand.js'; import { authCommand } from '../ui/commands/authCommand.js'; +import { btwCommand } from '../ui/commands/btwCommand.js'; import { bugCommand } from '../ui/commands/bugCommand.js'; import { clearCommand } from '../ui/commands/clearCommand.js'; import { compressCommand } from '../ui/commands/compressCommand.js'; @@ -67,6 +68,7 @@ export class BuiltinCommandLoader implements ICommandLoader { arenaCommand, approvalModeCommand, authCommand, + btwCommand, bugCommand, clearCommand, compressCommand, diff --git a/packages/cli/src/test-utils/mockCommandContext.ts b/packages/cli/src/test-utils/mockCommandContext.ts index fd825b9df..d6a6c3e6d 100644 --- a/packages/cli/src/test-utils/mockCommandContext.ts +++ b/packages/cli/src/test-utils/mockCommandContext.ts @@ -55,6 +55,10 @@ export const createMockCommandContext = ( setDebugMessage: vi.fn(), pendingItem: null, setPendingItem: vi.fn(), + btwItem: null, + setBtwItem: vi.fn(), + cancelBtw: vi.fn(), + btwAbortControllerRef: { current: null }, loadHistory: vi.fn(), toggleVimEnabled: vi.fn(), extensionsUpdateState: new Map(), diff --git a/packages/cli/src/ui/AppContainer.test.tsx b/packages/cli/src/ui/AppContainer.test.tsx index 91b8ae644..4e8091378 100644 --- a/packages/cli/src/ui/AppContainer.test.tsx +++ b/packages/cli/src/ui/AppContainer.test.tsx @@ -434,6 +434,41 @@ describe('AppContainer State Management', () => { ); }).not.toThrow(); }); + + it('submits /btw immediately instead of queueing while responding', () => { + const mockSubmitQuery = vi.fn(); + const mockQueueMessage = vi.fn(); + + mockedUseGeminiStream.mockReturnValue({ + streamingState: 'responding', + 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(''), + }); + + render( + , + ); + + capturedUIActions.handleFinalSubmit('/btw quick side question'); + + expect(mockSubmitQuery).toHaveBeenCalledWith('/btw quick side question'); + expect(mockQueueMessage).not.toHaveBeenCalled(); + }); }); describe('Settings Integration', () => { diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index 2574f5bf0..b1918ebaa 100644 --- a/packages/cli/src/ui/AppContainer.tsx +++ b/packages/cli/src/ui/AppContainer.tsx @@ -71,6 +71,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 { type LoadedSettings, SettingScope } from '../config/settings.js'; import { type InitializationResult } from '../core/initializer.js'; import { useFocus } from './hooks/useFocus.js'; @@ -599,6 +600,9 @@ export const AppContainer = (props: AppContainerProps) => { handleSlashCommand, slashCommands, pendingHistoryItems: pendingSlashCommandHistoryItems, + btwItem, + setBtwItem, + cancelBtw, commandContext, shellConfirmationRequest, confirmationRequest, @@ -747,9 +751,16 @@ export const AppContainer = (props: AppContainerProps) => { return; } } + if ( + streamingState === StreamingState.Responding && + isBtwCommand(submittedValue) + ) { + void submitQuery(submittedValue); + return; + } addMessage(submittedValue); }, - [addMessage, agentViewState], + [addMessage, agentViewState, streamingState, submitQuery], ); const handleArenaModelsSelected = useCallback( @@ -947,6 +958,7 @@ export const AppContainer = (props: AppContainerProps) => { const ctrlDTimerRef = useRef(null); const [escapePressedOnce, setEscapePressedOnce] = useState(false); const escapeTimerRef = useRef(null); + const dialogsVisibleRef = useRef(false); const [constrainHeight, setConstrainHeight] = useState(true); const [ideContextState, setIdeContextState] = useState< IdeContext | undefined @@ -1233,7 +1245,13 @@ export const AppContainer = (props: AppContainerProps) => { handleExit(ctrlDPressedOnce, setCtrlDPressedOnce, ctrlDTimerRef); return; } else if (keyMatchers[Command.ESCAPE](key)) { - // Escape key handling + // Dismiss or cancel btw side-question on Escape, + // but only when btw is actually visible (not hidden behind a dialog). + if (btwItem && !dialogsVisibleRef.current) { + cancelBtw(); + return; + } + // Skip if shell is focused (to allow shell's own escape handling) if (embeddedShellFocused) { return; @@ -1275,6 +1293,20 @@ export const AppContainer = (props: AppContainerProps) => { return; } + // Dismiss completed btw side-question on Space or Enter, + // but only when btw is visible and the input buffer is empty. + if ( + btwItem && + !btwItem.btw.isPending && + !dialogsVisibleRef.current && + buffer.text.length === 0 + ) { + if (key.name === 'return' || key.sequence === ' ') { + setBtwItem(null); + return; + } + } + let enteringConstrainHeightMode = false; if (!constrainHeight) { enteringConstrainHeightMode = true; @@ -1329,6 +1361,9 @@ export const AppContainer = (props: AppContainerProps) => { handleSlashCommand, activePtyId, embeddedShellFocused, + btwItem, + setBtwItem, + cancelBtw, settings.merged.general?.debugKeystrokeLogging, isAuthenticating, ], @@ -1402,6 +1437,7 @@ export const AppContainer = (props: AppContainerProps) => { isApprovalModeDialogOpen || isResumeDialogOpen || isExtensionsManagerDialogOpen; + dialogsVisibleRef.current = dialogsVisible; const { isFeedbackDialogOpen, @@ -1492,6 +1528,9 @@ export const AppContainer = (props: AppContainerProps) => { staticExtraHeight, dialogsVisible, pendingHistoryItems, + btwItem, + setBtwItem, + cancelBtw, nightly, branchName, sessionStats, @@ -1588,6 +1627,9 @@ export const AppContainer = (props: AppContainerProps) => { staticExtraHeight, dialogsVisible, pendingHistoryItems, + btwItem, + setBtwItem, + cancelBtw, nightly, branchName, sessionStats, diff --git a/packages/cli/src/ui/commands/btwCommand.test.ts b/packages/cli/src/ui/commands/btwCommand.test.ts new file mode 100644 index 000000000..99dfa40d3 --- /dev/null +++ b/packages/cli/src/ui/commands/btwCommand.test.ts @@ -0,0 +1,464 @@ +/** + * @license + * Copyright 2025 Qwen Code + * SPDX-License-Identifier: Apache-2.0 + */ + +import { vi, describe, it, expect, beforeEach } from 'vitest'; +import { btwCommand } from './btwCommand.js'; +import { type CommandContext } from './types.js'; +import { createMockCommandContext } from '../../test-utils/mockCommandContext.js'; +import { CommandKind } from './types.js'; +import { MessageType } from '../types.js'; + +vi.mock('../../i18n/index.js', () => ({ + t: (key: string, params?: Record) => { + if (params) { + return Object.entries(params).reduce( + (str, [k, v]) => str.replace(`{{${k}}}`, v), + key, + ); + } + return key; + }, +})); + +describe('btwCommand', () => { + let mockContext: CommandContext; + let mockGenerateContent: ReturnType; + let mockGetHistory: ReturnType; + const createConfig = (overrides: Record = {}) => ({ + getGeminiClient: () => ({ + getHistory: mockGetHistory, + generateContent: mockGenerateContent, + }), + getModel: () => 'test-model', + getSessionId: () => 'test-session-id', + ...overrides, + }); + + beforeEach(() => { + vi.clearAllMocks(); + + mockGenerateContent = vi.fn(); + mockGetHistory = vi.fn().mockReturnValue([]); + + mockContext = createMockCommandContext({ + services: { + config: createConfig(), + }, + }); + }); + + it('should have correct metadata', () => { + expect(btwCommand.name).toBe('btw'); + expect(btwCommand.kind).toBe(CommandKind.BUILT_IN); + expect(btwCommand.description).toBeTruthy(); + }); + + it('should return error when no question is provided', async () => { + const result = await btwCommand.action!(mockContext, ''); + + expect(result).toEqual({ + type: 'message', + messageType: 'error', + content: 'Please provide a question. Usage: /btw ', + }); + }); + + it('should return error when only whitespace is provided', async () => { + const result = await btwCommand.action!(mockContext, ' '); + + expect(result).toEqual({ + type: 'message', + messageType: 'error', + content: 'Please provide a question. Usage: /btw ', + }); + }); + + it('should return error when config is not loaded', async () => { + const noConfigContext = createMockCommandContext({ + services: { config: null }, + }); + + const result = await btwCommand.action!(noConfigContext, 'test question'); + + expect(result).toEqual({ + type: 'message', + messageType: 'error', + content: 'Config not loaded.', + }); + }); + + it('should return error when model is not configured', async () => { + const noModelContext = createMockCommandContext({ + services: { + config: createConfig({ + getModel: () => '', + }), + }, + }); + + const result = await btwCommand.action!(noModelContext, 'test question'); + + expect(result).toEqual({ + type: 'message', + messageType: 'error', + content: 'No model configured.', + }); + }); + + describe('interactive mode', () => { + const flushPromises = () => + new Promise((resolve) => setTimeout(resolve, 0)); + + it('should set btwItem and update it on success', async () => { + mockGenerateContent.mockResolvedValue({ + candidates: [ + { + content: { + parts: [{ text: 'The answer is 42.' }], + }, + }, + ], + }); + + await btwCommand.action!(mockContext, 'what is the meaning of life?'); + + // Action returns immediately; btwItem is set synchronously + expect(mockContext.ui.setBtwItem).toHaveBeenCalledWith({ + type: MessageType.BTW, + btw: { + question: 'what is the meaning of life?', + answer: '', + isPending: true, + }, + }); + + // pendingItem should NOT be used + expect(mockContext.ui.setPendingItem).not.toHaveBeenCalled(); + + await flushPromises(); + + // On success, setBtwItem is called with the completed answer + expect(mockContext.ui.setBtwItem).toHaveBeenCalledWith({ + type: MessageType.BTW, + btw: { + question: 'what is the meaning of life?', + answer: 'The answer is 42.', + isPending: false, + }, + }); + + // addItem should NOT be called (btw stays in fixed area, not in history) + expect(mockContext.ui.addItem).not.toHaveBeenCalled(); + }); + + it('should pass conversation history to generateContent', async () => { + const history = [ + { role: 'user', parts: [{ text: 'Hello' }] }, + { role: 'model', parts: [{ text: 'Hi!' }] }, + ]; + mockGetHistory.mockReturnValue(history); + mockGenerateContent.mockResolvedValue({ + candidates: [{ content: { parts: [{ text: 'answer' }] } }], + }); + + await btwCommand.action!(mockContext, 'my question'); + await flushPromises(); + + expect(mockGenerateContent).toHaveBeenCalledWith( + [ + ...history, + { + role: 'user', + parts: [ + { + text: expect.stringContaining('my question'), + }, + ], + }, + ], + {}, + expect.any(AbortSignal), + 'test-model', + expect.stringMatching(/^test-session-id########btw-/), + ); + }); + + it('should add error item on failure and clear btwItem', async () => { + mockGenerateContent.mockRejectedValue(new Error('API error')); + + await btwCommand.action!(mockContext, 'test question'); + await flushPromises(); + + // btwItem should be cleared on error + expect(mockContext.ui.setBtwItem).toHaveBeenCalledWith(null); + + // Error goes to history + expect(mockContext.ui.addItem).toHaveBeenCalledWith( + { + type: MessageType.ERROR, + text: 'Failed to answer btw question: API error', + }, + expect.any(Number), + ); + }); + + it('should handle non-Error exceptions', async () => { + mockGenerateContent.mockRejectedValue('string error'); + + await btwCommand.action!(mockContext, 'test question'); + await flushPromises(); + + expect(mockContext.ui.addItem).toHaveBeenCalledWith( + { + type: MessageType.ERROR, + text: 'Failed to answer btw question: string error', + }, + expect.any(Number), + ); + }); + + it('should not block when another pendingItem exists', async () => { + const busyContext = createMockCommandContext({ + services: { + config: createConfig(), + }, + ui: { + pendingItem: { type: 'info' }, + }, + }); + + mockGenerateContent.mockResolvedValue({ + candidates: [{ content: { parts: [{ text: 'answer' }] } }], + }); + + // btw should NOT be blocked by pendingItem anymore + const result = await btwCommand.action!(busyContext, 'test question'); + expect(result).toBeUndefined(); + expect(busyContext.ui.setBtwItem).toHaveBeenCalled(); + }); + + it('should not update btwItem when cancelled via btwAbortControllerRef', async () => { + mockGenerateContent.mockImplementation( + () => + new Promise((resolve) => + setTimeout( + () => + resolve({ + candidates: [ + { content: { parts: [{ text: 'late answer' }] } }, + ], + }), + 50, + ), + ), + ); + + await btwCommand.action!(mockContext, 'test question'); + + // The btw command should have registered its AbortController + expect(mockContext.ui.btwAbortControllerRef.current).toBeInstanceOf( + AbortController, + ); + + // Simulate user pressing ESC: cancel the in-flight btw + mockContext.ui.btwAbortControllerRef.current!.abort(); + + await flushPromises(); + + // setBtwItem should only have the initial pending call (no completion) + expect(mockContext.ui.setBtwItem).toHaveBeenCalledTimes(1); + expect(mockContext.ui.addItem).not.toHaveBeenCalled(); + }); + + it('should clear btwAbortControllerRef after successful completion', async () => { + mockGenerateContent.mockResolvedValue({ + candidates: [{ content: { parts: [{ text: 'answer' }] } }], + }); + + await btwCommand.action!(mockContext, 'test question'); + + // Ref is set during the call + expect(mockContext.ui.btwAbortControllerRef.current).toBeInstanceOf( + AbortController, + ); + + await flushPromises(); + + // After completion, ref should be cleaned up + expect(mockContext.ui.btwAbortControllerRef.current).toBeNull(); + }); + + it('should clear btwAbortControllerRef after error', async () => { + mockGenerateContent.mockRejectedValue(new Error('API error')); + + await btwCommand.action!(mockContext, 'test question'); + + expect(mockContext.ui.btwAbortControllerRef.current).toBeInstanceOf( + AbortController, + ); + + await flushPromises(); + + expect(mockContext.ui.btwAbortControllerRef.current).toBeNull(); + }); + + it('should cancel previous btw when starting a new one', async () => { + mockGenerateContent.mockResolvedValue({ + candidates: [{ content: { parts: [{ text: 'answer' }] } }], + }); + + await btwCommand.action!(mockContext, 'first question'); + + // cancelBtw should have been called to clean up any previous btw + expect(mockContext.ui.cancelBtw).toHaveBeenCalledTimes(1); + + // Second btw call + await btwCommand.action!(mockContext, 'second question'); + + // cancelBtw called again for the second invocation + expect(mockContext.ui.cancelBtw).toHaveBeenCalledTimes(2); + }); + + it('should return fallback text when response has no parts', async () => { + mockGenerateContent.mockResolvedValue({ + candidates: [{ content: { parts: [] } }], + }); + + await btwCommand.action!(mockContext, 'test question'); + await flushPromises(); + + expect(mockContext.ui.setBtwItem).toHaveBeenCalledWith({ + type: MessageType.BTW, + btw: { + question: 'test question', + answer: 'No response received.', + isPending: false, + }, + }); + }); + + it('should return void immediately without blocking', async () => { + mockGenerateContent.mockResolvedValue({ + candidates: [{ content: { parts: [{ text: 'answer' }] } }], + }); + + const result = await btwCommand.action!(mockContext, 'test question'); + + expect(result).toBeUndefined(); + + // Only the pending setBtwItem called so far + expect(mockContext.ui.setBtwItem).toHaveBeenCalledTimes(1); + + await flushPromises(); + + // Now the completed setBtwItem has been called + expect(mockContext.ui.setBtwItem).toHaveBeenCalledTimes(2); + }); + }); + + describe('non-interactive mode', () => { + let nonInteractiveContext: CommandContext; + + beforeEach(() => { + nonInteractiveContext = createMockCommandContext({ + executionMode: 'non_interactive', + services: { + config: createConfig(), + }, + }); + }); + + it('should return info message on success', async () => { + mockGenerateContent.mockResolvedValue({ + candidates: [{ content: { parts: [{ text: 'the answer' }] } }], + }); + + const result = await btwCommand.action!( + nonInteractiveContext, + 'my question', + ); + + expect(result).toEqual({ + type: 'message', + messageType: 'info', + content: 'btw> my question\nthe answer', + }); + }); + + it('should return error message on failure', async () => { + mockGenerateContent.mockRejectedValue(new Error('network error')); + + const result = await btwCommand.action!( + nonInteractiveContext, + 'my question', + ); + + expect(result).toEqual({ + type: 'message', + messageType: 'error', + content: 'Failed to answer btw question: network error', + }); + }); + }); + + describe('acp mode', () => { + let acpContext: CommandContext; + + beforeEach(() => { + acpContext = createMockCommandContext({ + executionMode: 'acp', + services: { + config: createConfig(), + }, + }); + }); + + it('should return stream_messages generator on success', async () => { + mockGenerateContent.mockResolvedValue({ + candidates: [{ content: { parts: [{ text: 'streamed answer' }] } }], + }); + + const result = (await btwCommand.action!(acpContext, 'my question')) as { + type: string; + messages: AsyncGenerator; + }; + + expect(result.type).toBe('stream_messages'); + + const messages = []; + for await (const msg of result.messages) { + messages.push(msg); + } + + expect(messages).toEqual([ + { messageType: 'info', content: 'Thinking...' }, + { messageType: 'info', content: 'btw> my question\nstreamed answer' }, + ]); + }); + + it('should yield error message on failure', async () => { + mockGenerateContent.mockRejectedValue(new Error('api failure')); + + const result = (await btwCommand.action!(acpContext, 'my question')) as { + type: string; + messages: AsyncGenerator; + }; + + const messages = []; + for await (const msg of result.messages) { + messages.push(msg); + } + + expect(messages).toEqual([ + { messageType: 'info', content: 'Thinking...' }, + { + messageType: 'error', + content: 'Failed to answer btw question: api failure', + }, + ]); + }); + }); +}); diff --git a/packages/cli/src/ui/commands/btwCommand.ts b/packages/cli/src/ui/commands/btwCommand.ts new file mode 100644 index 000000000..60a3ab8dd --- /dev/null +++ b/packages/cli/src/ui/commands/btwCommand.ts @@ -0,0 +1,226 @@ +/** + * @license + * Copyright 2025 Qwen Code + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { + CommandContext, + SlashCommand, + SlashCommandActionReturn, +} from './types.js'; +import { CommandKind } from './types.js'; +import { MessageType } from '../types.js'; +import type { HistoryItemBtw } from '../types.js'; +import { t } from '../../i18n/index.js'; +import type { GeminiClient } from '@qwen-code/qwen-code-core'; + +function makeBtwPromptId(sessionId: string): string { + return `${sessionId}########btw-${Date.now()}`; +} + +function formatBtwError(error: unknown): string { + return t('Failed to answer btw question: {{error}}', { + error: + error instanceof Error ? error.message : String(error || 'Unknown error'), + }); +} + +/** + * Helper to make the ephemeral generateContent call and extract the answer. + * Uses a snapshot of the current conversation history as context. + */ +async function askBtw( + geminiClient: GeminiClient, + model: string, + question: string, + abortSignal: AbortSignal, + promptId: string, +): Promise { + const history = geminiClient.getHistory(); + + const response = await geminiClient.generateContent( + [ + ...history, + { + role: 'user', + parts: [ + { + text: `[Side question - answer briefly and concisely, this is a "by the way" question that doesn't need to be part of our main conversation]\n\n${question}`, + }, + ], + }, + ], + {}, + abortSignal, + model, + promptId, + ); + + const parts = response.candidates?.[0]?.content?.parts; + return ( + parts + ?.map((part) => part.text) + .filter((text): text is string => typeof text === 'string') + .join('') || t('No response received.') + ); +} + +export const btwCommand: SlashCommand = { + name: 'btw', + get description() { + return t( + 'Ask a quick side question without affecting the main conversation', + ); + }, + kind: CommandKind.BUILT_IN, + action: async ( + context: CommandContext, + args: string, + ): Promise => { + const question = args.trim(); + const executionMode = context.executionMode ?? 'interactive'; + const abortSignal = context.abortSignal ?? new AbortController().signal; + + if (!question) { + return { + type: 'message', + messageType: 'error', + content: t('Please provide a question. Usage: /btw '), + }; + } + + const { config } = context.services; + const { ui } = context; + + if (!config) { + return { + type: 'message', + messageType: 'error', + content: t('Config not loaded.'), + }; + } + + const geminiClient = config.getGeminiClient(); + const model = config.getModel(); + const sessionId = config.getSessionId(); + + if (!model) { + return { + type: 'message', + messageType: 'error', + content: t('No model configured.'), + }; + } + + // ACP mode: return a stream_messages async generator + if (executionMode === 'acp') { + const btwPromptId = makeBtwPromptId(sessionId); + const messages = async function* () { + try { + yield { + messageType: 'info' as const, + content: t('Thinking...'), + }; + + const answer = await askBtw( + geminiClient, + model, + question, + abortSignal, + btwPromptId, + ); + + yield { + messageType: 'info' as const, + content: `btw> ${question}\n${answer}`, + }; + } catch (error) { + yield { + messageType: 'error' as const, + content: formatBtwError(error), + }; + } + }; + + return { type: 'stream_messages', messages: messages() }; + } + + // Non-interactive mode: return a simple message result + if (executionMode === 'non_interactive') { + try { + const btwPromptId = makeBtwPromptId(sessionId); + const answer = await askBtw( + geminiClient, + model, + question, + abortSignal, + btwPromptId, + ); + return { + type: 'message', + messageType: 'info', + content: `btw> ${question}\n${answer}`, + }; + } catch (error) { + return { + type: 'message', + messageType: 'error', + content: formatBtwError(error), + }; + } + } + + // Interactive mode: use dedicated btwItem state for the fixed bottom area. + // This does NOT occupy pendingItem, so the main conversation is never blocked. + + // Cancel any previous in-flight btw before starting a new one. + ui.cancelBtw(); + + const btwAbortController = new AbortController(); + const btwSignal = btwAbortController.signal; + ui.btwAbortControllerRef.current = btwAbortController; + + const pendingItem: HistoryItemBtw = { + type: MessageType.BTW, + btw: { + question, + answer: '', + isPending: true, + }, + }; + ui.setBtwItem(pendingItem); + + // Fire-and-forget: run the API call in the background so the main + // conversation is not blocked while waiting for the btw answer. + const btwPromptId = makeBtwPromptId(sessionId); + void askBtw(geminiClient, model, question, btwSignal, btwPromptId) + .then((answer) => { + if (btwSignal.aborted) return; + + ui.btwAbortControllerRef.current = null; + const completedItem: HistoryItemBtw = { + type: MessageType.BTW, + btw: { + question, + answer, + isPending: false, + }, + }; + ui.setBtwItem(completedItem); + }) + .catch((error) => { + if (btwSignal.aborted) return; + + ui.btwAbortControllerRef.current = null; + ui.setBtwItem(null); + ui.addItem( + { + type: MessageType.ERROR, + text: formatBtwError(error), + }, + Date.now(), + ); + }); + }, +}; diff --git a/packages/cli/src/ui/commands/types.ts b/packages/cli/src/ui/commands/types.ts index 49f937027..d74f3e393 100644 --- a/packages/cli/src/ui/commands/types.ts +++ b/packages/cli/src/ui/commands/types.ts @@ -4,12 +4,13 @@ * SPDX-License-Identifier: Apache-2.0 */ -import type { ReactNode } from 'react'; +import type { MutableRefObject, ReactNode } from 'react'; import type { Content, PartListUnion } from '@google/genai'; import type { Config, GitService, Logger } from '@qwen-code/qwen-code-core'; import type { HistoryItemWithoutId, HistoryItem, + HistoryItemBtw, ConfirmationRequest, } from '../types.js'; import type { LoadedSettings } from '../../config/settings.js'; @@ -66,6 +67,14 @@ export interface CommandContext { * @param item The history item to display as pending, or `null` to clear. */ setPendingItem: (item: HistoryItemWithoutId | null) => void; + /** The current btw side-question item rendered in the fixed bottom area. */ + btwItem: HistoryItemBtw | null; + /** Sets the btw item independently of the main pendingItem. */ + setBtwItem: (item: HistoryItemBtw | null) => void; + /** Cancels a pending btw (aborts the in-flight API call and clears the btw area). */ + cancelBtw: () => void; + /** Ref to the btw AbortController, set by btwCommand so cancelBtw can abort it. */ + btwAbortControllerRef: MutableRefObject; /** * Loads a new set of history items, replacing the current history. * diff --git a/packages/cli/src/ui/components/HistoryItemDisplay.tsx b/packages/cli/src/ui/components/HistoryItemDisplay.tsx index b52a2b9bf..12a46380e 100644 --- a/packages/cli/src/ui/components/HistoryItemDisplay.tsx +++ b/packages/cli/src/ui/components/HistoryItemDisplay.tsx @@ -42,6 +42,7 @@ import { McpStatus } from './views/McpStatus.js'; import { ContextUsage } from './views/ContextUsage.js'; import { ArenaAgentCard, ArenaSessionCard } from './arena/ArenaCards.js'; import { InsightProgressMessage } from './messages/InsightProgressMessage.js'; +import { BtwMessage } from './messages/BtwMessage.js'; interface HistoryItemDisplayProps { item: HistoryItem; @@ -226,6 +227,9 @@ const HistoryItemDisplayComponent: React.FC = ({ {itemForDisplay.type === 'insight_progress' && ( )} + {itemForDisplay.type === 'btw' && itemForDisplay.btw && ( + + )} ); }; diff --git a/packages/cli/src/ui/components/messages/BtwMessage.test.tsx b/packages/cli/src/ui/components/messages/BtwMessage.test.tsx new file mode 100644 index 000000000..da784dc0d --- /dev/null +++ b/packages/cli/src/ui/components/messages/BtwMessage.test.tsx @@ -0,0 +1,34 @@ +/** + * @license + * Copyright 2025 Qwen Code + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, expect, it } from 'vitest'; +import { render } from 'ink-testing-library'; +import { BtwMessage } from './BtwMessage.js'; + +describe('BtwMessage', () => { + it('is wrapped in React.memo to avoid unnecessary layout rerenders', () => { + expect((BtwMessage as unknown as { $$typeof?: symbol }).$$typeof).toBe( + Symbol.for('react.memo'), + ); + }); + + it('renders the side question and answer', () => { + const { lastFrame } = render( + , + ); + + const output = lastFrame() ?? ''; + expect(output).toContain('/btw'); + expect(output).toContain('side question'); + expect(output).toContain('side answer'); + }); +}); diff --git a/packages/cli/src/ui/components/messages/BtwMessage.tsx b/packages/cli/src/ui/components/messages/BtwMessage.tsx new file mode 100644 index 000000000..9b28ecc49 --- /dev/null +++ b/packages/cli/src/ui/components/messages/BtwMessage.tsx @@ -0,0 +1,54 @@ +/** + * @license + * Copyright 2025 Qwen Code + * SPDX-License-Identifier: Apache-2.0 + */ + +import React from 'react'; +import { Box, Text } from 'ink'; +import type { BtwProps } from '../../types.js'; +import { Colors } from '../../colors.js'; +import { t } from '../../../i18n/index.js'; + +export interface BtwDisplayProps { + btw: BtwProps; +} + +const BtwMessageInternal: React.FC = ({ btw }) => ( + + + + {'/btw '} + + + {btw.question} + + + {btw.isPending ? ( + + + {'+ '} + {t('Answering...')} + + + {t('Press Escape to cancel')} + + + ) : ( + + {btw.answer} + + {t('Press Space, Enter, or Escape to dismiss')} + + + )} + +); + +export const BtwMessage = React.memo(BtwMessageInternal); diff --git a/packages/cli/src/ui/contexts/UIStateContext.tsx b/packages/cli/src/ui/contexts/UIStateContext.tsx index 986b07899..03bda1e58 100644 --- a/packages/cli/src/ui/contexts/UIStateContext.tsx +++ b/packages/cli/src/ui/contexts/UIStateContext.tsx @@ -7,6 +7,7 @@ import { createContext, useContext } from 'react'; import type { HistoryItem, + HistoryItemBtw, ThoughtSummary, ShellConfirmationRequest, ConfirmationRequest, @@ -104,6 +105,9 @@ export interface UIState { staticExtraHeight: number; dialogsVisible: boolean; pendingHistoryItems: HistoryItemWithoutId[]; + btwItem: HistoryItemBtw | null; + setBtwItem: (item: HistoryItemBtw | null) => void; + cancelBtw: () => void; nightly: boolean; branchName: string | undefined; sessionStats: SessionStatsState; diff --git a/packages/cli/src/ui/hooks/slashCommandProcessor.ts b/packages/cli/src/ui/hooks/slashCommandProcessor.ts index d799a402d..2d61409f4 100644 --- a/packages/cli/src/ui/hooks/slashCommandProcessor.ts +++ b/packages/cli/src/ui/hooks/slashCommandProcessor.ts @@ -23,6 +23,7 @@ import { useSessionStats } from '../contexts/SessionContext.js'; import type { Message, HistoryItemWithoutId, + HistoryItemBtw, SlashCommandProcessorResult, HistoryItem, ConfirmationRequest, @@ -36,6 +37,7 @@ import { BundledSkillLoader } from '../../services/BundledSkillLoader.js'; import { FileCommandLoader } from '../../services/FileCommandLoader.js'; import { McpPromptLoader } from '../../services/McpPromptLoader.js'; import { parseSlashCommand } from '../../utils/commands.js'; +import { isBtwCommand } from '../utils/commandUtils.js'; import { clearScreen } from '../../utils/stdioHelpers.js'; import { useKeypress } from './useKeypress.js'; import { @@ -63,6 +65,7 @@ const SLASH_COMMANDS_SKIP_RECORDING = new Set([ 'reset', 'new', 'resume', + 'btw', ]); interface SlashCommandProcessorActions { @@ -139,10 +142,20 @@ export const useSlashCommandProcessor = ( null, ); + const [btwItem, setBtwItem] = useState(null); + const btwAbortControllerRef = useRef(null); + + const cancelBtw = useCallback(() => { + btwAbortControllerRef.current?.abort(); + btwAbortControllerRef.current = null; + setBtwItem(null); + }, []); + // AbortController for cancelling async slash commands via ESC const abortControllerRef = useRef(null); const cancelSlashCommand = useCallback(() => { + cancelBtw(); if (!abortControllerRef.current) { return; } @@ -156,7 +169,7 @@ export const useSlashCommandProcessor = ( ); setPendingItem(null); setIsProcessing(false); - }, [addItem, setIsProcessing]); + }, [addItem, setIsProcessing, cancelBtw]); useKeypress( (key) => { @@ -251,6 +264,10 @@ export const useSlashCommandProcessor = ( setDebugMessage: actions.setDebugMessage, pendingItem, setPendingItem, + btwItem, + setBtwItem, + cancelBtw, + btwAbortControllerRef, toggleVimEnabled, setGeminiMdFileCount, reloadCommands, @@ -279,6 +296,9 @@ export const useSlashCommandProcessor = ( actions, pendingItem, setPendingItem, + btwItem, + setBtwItem, + cancelBtw, toggleVimEnabled, sessionShellAllowlist, setGeminiMdFileCount, @@ -366,10 +386,12 @@ export const useSlashCommandProcessor = ( abortControllerRef.current = abortController; const userMessageTimestamp = Date.now(); - addItemWithRecording( - { type: MessageType.USER, text: trimmed }, - userMessageTimestamp, - ); + if (!isBtwCommand(trimmed)) { + addItemWithRecording( + { type: MessageType.USER, text: trimmed }, + userMessageTimestamp, + ); + } let hasError = false; const { @@ -727,6 +749,9 @@ export const useSlashCommandProcessor = ( handleSlashCommand, slashCommands: commands, pendingHistoryItems, + btwItem, + setBtwItem, + cancelBtw, commandContext, shellConfirmationRequest, confirmationRequest, diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx index 4330ba7a5..2234db6bd 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx +++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx @@ -834,7 +834,7 @@ describe('useGeminiStream', () => { // Wait for the first part of the response await waitFor(() => { - expect(result.current.streamingState).toBe(StreamingState.Responding); + expect(mockSendMessageStream).toHaveBeenCalledTimes(1); }); // Call cancelOngoingRequest directly @@ -983,7 +983,7 @@ describe('useGeminiStream', () => { }); await waitFor(() => { - expect(result.current.streamingState).toBe(StreamingState.Responding); + expect(mockSendMessageStream).toHaveBeenCalledTimes(1); }); // Cancel the request @@ -2709,6 +2709,109 @@ describe('useGeminiStream', () => { }); describe('Concurrent Execution Prevention', () => { + it('should allow /btw slash commands while a main response is in progress', async () => { + let resolveFirstCall!: () => void; + + const firstCallPromise = new Promise((resolve) => { + resolveFirstCall = resolve; + }); + + const firstStream = (async function* () { + yield { + type: ServerGeminiEventType.Content, + value: 'First call content', + }; + await firstCallPromise; + })(); + + mockSendMessageStream.mockImplementation(() => firstStream); + mockHandleSlashCommand.mockImplementation(async (command) => { + if (command === '/btw quick side question') { + return { type: 'handled' }; + } + return false; + }); + + const { result } = renderTestHook(); + + let mainRequest!: Promise; + await act(async () => { + mainRequest = result.current.submitQuery('First query'); + }); + + try { + await waitFor(() => { + expect(mockSendMessageStream).toHaveBeenCalledTimes(1); + expect(result.current.streamingState).toBe(StreamingState.Responding); + }); + + await act(async () => { + await result.current.submitQuery('/btw quick side question'); + }); + + expect(mockHandleSlashCommand).toHaveBeenCalledWith( + '/btw quick side question', + ); + expect(mockSendMessageStream).toHaveBeenCalledTimes(1); + } finally { + resolveFirstCall(); + await mainRequest; + } + }); + + it('should keep the main request cancellable after submitting /btw in parallel', async () => { + let resolveFirstCall!: () => void; + let mainAbortSignal: AbortSignal | undefined; + + const firstCallPromise = new Promise((resolve) => { + resolveFirstCall = resolve; + }); + + mockSendMessageStream.mockImplementation((_query, signal) => { + mainAbortSignal = signal; + return (async function* () { + yield { + type: ServerGeminiEventType.Content, + value: 'First call content', + }; + await firstCallPromise; + })(); + }); + mockHandleSlashCommand.mockImplementation(async (command) => { + if (command === '/btw quick side question') { + return { type: 'handled' }; + } + return false; + }); + + const { result } = renderTestHook(); + + let mainRequest!: Promise; + await act(async () => { + mainRequest = result.current.submitQuery('First query'); + }); + + try { + await waitFor(() => { + expect(mainAbortSignal).toBeDefined(); + expect(result.current.streamingState).toBe(StreamingState.Responding); + }); + + await act(async () => { + await result.current.submitQuery('/btw quick side question'); + }); + + act(() => { + result.current.cancelOngoingRequest(); + }); + + expect(mainAbortSignal?.aborted).toBe(true); + } finally { + resolveFirstCall(); + await mainRequest; + } + }); + it('should prevent concurrent submitQuery calls', async () => { let resolveFirstCall!: () => void; let resolveSecondCall!: () => void; diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index 108b2fe83..5d39654b1 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -49,7 +49,11 @@ import type { SlashCommandProcessorResult, } from '../types.js'; import { StreamingState, MessageType, ToolCallStatus } from '../types.js'; -import { isAtCommand, isSlashCommand } from '../utils/commandUtils.js'; +import { + isAtCommand, + isBtwCommand, + isSlashCommand, +} from '../utils/commandUtils.js'; import { useShellCommandProcessor } from './shellCommandProcessor.js'; import { handleAtCommand } from './atCommandProcessor.js'; import { findLastSafeSplitPoint } from '../utils/markdownUtilities.js'; @@ -1094,11 +1098,18 @@ export const useGeminiStream = ( submitType: SendMessageType = SendMessageType.UserQuery, prompt_id?: string, ) => { + const allowConcurrentBtwDuringResponse = + submitType === SendMessageType.UserQuery && + streamingState === StreamingState.Responding && + typeof query === 'string' && + isBtwCommand(query); + // Prevent concurrent executions of submitQuery, but allow continuations // which are part of the same logical flow (tool responses) if ( isSubmittingQueryRef.current && - submitType !== SendMessageType.ToolResult + submitType !== SendMessageType.ToolResult && + !allowConcurrentBtwDuringResponse ) { return; } @@ -1106,7 +1117,8 @@ export const useGeminiStream = ( if ( (streamingState === StreamingState.Responding || streamingState === StreamingState.WaitingForConfirmation) && - submitType !== SendMessageType.ToolResult + submitType !== SendMessageType.ToolResult && + !allowConcurrentBtwDuringResponse ) return; @@ -1116,7 +1128,10 @@ export const useGeminiStream = ( const userMessageTimestamp = Date.now(); // Reset quota error flag when starting a new query (not a continuation) - if (submitType !== SendMessageType.ToolResult) { + if ( + submitType !== SendMessageType.ToolResult && + !allowConcurrentBtwDuringResponse + ) { setModelSwitchedFromQuotaError(false); // Commit any pending retry error to history (without hint) since the // user is starting a new conversation turn. @@ -1130,9 +1145,15 @@ export const useGeminiStream = ( } } - abortControllerRef.current = new AbortController(); - const abortSignal = abortControllerRef.current.signal; - turnCancelledRef.current = false; + const abortController = new AbortController(); + const abortSignal = abortController.signal; + + // Keep the main stream's cancellation state intact while /btw is handled + // in parallel. The side-question can use its own local abort signal. + if (!allowConcurrentBtwDuringResponse) { + abortControllerRef.current = abortController; + turnCancelledRef.current = false; + } if (!prompt_id) { prompt_id = config.getSessionId() + '########' + getPromptCount(); diff --git a/packages/cli/src/ui/layouts/DefaultAppLayout.tsx b/packages/cli/src/ui/layouts/DefaultAppLayout.tsx index ddb3f2df0..479730cb4 100644 --- a/packages/cli/src/ui/layouts/DefaultAppLayout.tsx +++ b/packages/cli/src/ui/layouts/DefaultAppLayout.tsx @@ -11,6 +11,7 @@ import { MainContent } from '../components/MainContent.js'; import { DialogManager } from '../components/DialogManager.js'; import { Composer } from '../components/Composer.js'; import { ExitWarning } from '../components/ExitWarning.js'; +import { BtwMessage } from '../components/messages/BtwMessage.js'; import { AgentTabBar } from '../components/agent-view/AgentTabBar.js'; import { AgentChatView } from '../components/agent-view/AgentChatView.js'; import { AgentComposer } from '../components/agent-view/AgentComposer.js'; @@ -66,6 +67,10 @@ export const DefaultAppLayout: React.FC = () => { addItem={uiState.historyManager.addItem} /> + ) : uiState.btwItem ? ( + + + ) : ( )} diff --git a/packages/cli/src/ui/layouts/ScreenReaderAppLayout.tsx b/packages/cli/src/ui/layouts/ScreenReaderAppLayout.tsx index b4967a5f4..f9e876a48 100644 --- a/packages/cli/src/ui/layouts/ScreenReaderAppLayout.tsx +++ b/packages/cli/src/ui/layouts/ScreenReaderAppLayout.tsx @@ -12,6 +12,7 @@ import { DialogManager } from '../components/DialogManager.js'; import { Composer } from '../components/Composer.js'; import { Footer } from '../components/Footer.js'; import { ExitWarning } from '../components/ExitWarning.js'; +import { BtwMessage } from '../components/messages/BtwMessage.js'; import { useUIState } from '../contexts/UIStateContext.js'; export const ScreenReaderAppLayout: React.FC = () => { @@ -24,6 +25,7 @@ export const ScreenReaderAppLayout: React.FC = () => { + {uiState.dialogsVisible ? ( { addItem={uiState.historyManager.addItem} /> + ) : uiState.btwItem ? ( + + + ) : ( )} diff --git a/packages/cli/src/ui/noninteractive/nonInteractiveUi.ts b/packages/cli/src/ui/noninteractive/nonInteractiveUi.ts index 779293330..dbdf4e2e3 100644 --- a/packages/cli/src/ui/noninteractive/nonInteractiveUi.ts +++ b/packages/cli/src/ui/noninteractive/nonInteractiveUi.ts @@ -20,6 +20,10 @@ export function createNonInteractiveUI(): CommandContext['ui'] { loadHistory: (_newHistory) => {}, pendingItem: null, setPendingItem: (_item) => {}, + btwItem: null, + setBtwItem: (_item) => {}, + cancelBtw: () => {}, + btwAbortControllerRef: { current: null }, toggleVimEnabled: async () => false, setGeminiMdFileCount: (_count) => {}, reloadCommands: () => {}, diff --git a/packages/cli/src/ui/types.ts b/packages/cli/src/ui/types.ts index 64353066e..7f9b4c176 100644 --- a/packages/cli/src/ui/types.ts +++ b/packages/cli/src/ui/types.ts @@ -350,6 +350,17 @@ export type HistoryItemInsightProgress = HistoryItemBase & { progress: InsightProgressProps; }; +export interface BtwProps { + question: string; + answer: string; + isPending: boolean; +} + +export type HistoryItemBtw = HistoryItemBase & { + type: 'btw'; + btw: BtwProps; +}; + // Using Omit seems to have some issues with typescript's // type inference e.g. historyItem.type === 'tool_group' isn't auto-inferring that // 'tools' in historyItem. @@ -383,7 +394,8 @@ export type HistoryItemWithoutId = | HistoryItemContextUsage | HistoryItemArenaAgentComplete | HistoryItemArenaSessionComplete - | HistoryItemInsightProgress; + | HistoryItemInsightProgress + | HistoryItemBtw; export type HistoryItem = HistoryItemWithoutId & { id: number }; @@ -411,6 +423,7 @@ export enum MessageType { ARENA_AGENT_COMPLETE = 'arena_agent_complete', ARENA_SESSION_COMPLETE = 'arena_session_complete', INSIGHT_PROGRESS = 'insight_progress', + BTW = 'btw', } export interface InsightProgressProps { diff --git a/packages/cli/src/ui/utils/commandUtils.ts b/packages/cli/src/ui/utils/commandUtils.ts index 802107f6b..9436447f7 100644 --- a/packages/cli/src/ui/utils/commandUtils.ts +++ b/packages/cli/src/ui/utils/commandUtils.ts @@ -62,6 +62,17 @@ export const isSlashCommand = (query: string): boolean => { return true; }; +const BTW_COMMAND_RE = /^[/?]btw(?:\s|$)/; + +/** + * Checks if a query is a /btw side-question invocation. + * Accepts both "/btw" and "?btw" prefixes. + */ +export const isBtwCommand = (query: string): boolean => { + const trimmed = query.trim(); + return trimmed.length > 0 && BTW_COMMAND_RE.test(trimmed); +}; + const debugLogger = createDebugLogger('COMMAND_UTILS'); // Copies a string snippet to the clipboard for different platforms diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index 01d97ffbd..9527ef071 100644 --- a/packages/core/src/core/client.test.ts +++ b/packages/core/src/core/client.test.ts @@ -34,6 +34,7 @@ import { import { getCoreSystemPrompt, getCustomSystemPrompt } from './prompts.js'; import { DEFAULT_QWEN_FLASH_MODEL } from '../config/models.js'; import { FileDiscoveryService } from '../services/fileDiscoveryService.js'; +import { promptIdContext } from '../utils/promptIdContext.js'; import { setSimulate429 } from '../utils/testUtils.js'; import { ideContextStore } from '../ide/ideContext.js'; import { uiTelemetryService } from '../telemetry/uiTelemetry.js'; @@ -2441,6 +2442,55 @@ Other open files: ); }); + it('should prefer the current prompt id context for stateless requests', async () => { + const contents = [{ role: 'user', parts: [{ text: 'hello' }] }]; + const abortSignal = new AbortController().signal; + + await promptIdContext.run('btw-prompt-id', async () => { + await client.generateContent( + contents, + {}, + abortSignal, + DEFAULT_QWEN_FLASH_MODEL, + ); + }); + + expect(mockContentGenerator.generateContent).toHaveBeenCalledWith( + expect.objectContaining({ + model: DEFAULT_QWEN_FLASH_MODEL, + contents, + }), + 'btw-prompt-id', + ); + }); + + it('should prefer an explicit prompt id override over the current context', async () => { + const contents = [{ role: 'user', parts: [{ text: 'hello' }] }]; + const abortSignal = new AbortController().signal; + + await promptIdContext.run('context-prompt-id', async () => { + await ( + client.generateContent as unknown as ( + ...args: unknown[] + ) => Promise + )( + contents, + {}, + abortSignal, + DEFAULT_QWEN_FLASH_MODEL, + 'override-prompt-id', + ); + }); + + expect(mockContentGenerator.generateContent).toHaveBeenCalledWith( + expect.objectContaining({ + model: DEFAULT_QWEN_FLASH_MODEL, + contents, + }), + 'override-prompt-id', + ); + }); + it('should use config system prompt override when provided', async () => { const contents = [{ role: 'user', parts: [{ text: 'hello' }] }]; const abortSignal = new AbortController().signal; diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index c7a04d2fe..4a0de9746 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -68,6 +68,7 @@ import { reportError } from '../utils/errorReporting.js'; import { getErrorMessage } from '../utils/errors.js'; import { checkNextSpeaker } from '../utils/nextSpeakerChecker.js'; import { flatMapTextParts } from '../utils/partUtils.js'; +import { promptIdContext } from '../utils/promptIdContext.js'; import { retryWithBackoff } from '../utils/retry.js'; // Hook types and utilities @@ -786,8 +787,11 @@ export class GeminiClient { generationConfig: GenerateContentConfig, abortSignal: AbortSignal, model: string, + promptIdOverride?: string, ): Promise { let currentAttemptModel: string = model; + const promptId = + promptIdOverride ?? promptIdContext.getStore() ?? this.lastPromptId!; try { const userMemory = this.config.getUserMemory(); @@ -810,7 +814,7 @@ export class GeminiClient { config: requestConfig, contents, }, - this.lastPromptId!, + promptId, ); }; const result = await retryWithBackoff(apiCall, {