From 4bc7e2554f24acd91a9149c3a682b718be95b7c5 Mon Sep 17 00:00:00 2001 From: Mahima Shanware Date: Fri, 27 Mar 2026 16:42:54 +0000 Subject: [PATCH] feat(cli): add useBtw hook and slash command processing --- packages/cli/src/ui/commands/btwCommand.ts | 12 +- packages/cli/src/ui/commands/types.ts | 6 +- .../cli/src/ui/hooks/slashCommandProcessor.ts | 3 + packages/cli/src/ui/hooks/useBtw.test.ts | 120 +++++++++++ packages/cli/src/ui/hooks/useBtw.ts | 187 ++++++++++++++++++ packages/cli/src/ui/hooks/useGeminiStream.ts | 3 + packages/cli/src/ui/types.ts | 8 +- 7 files changed, 326 insertions(+), 13 deletions(-) create mode 100644 packages/cli/src/ui/hooks/useBtw.test.ts create mode 100644 packages/cli/src/ui/hooks/useBtw.ts diff --git a/packages/cli/src/ui/commands/btwCommand.ts b/packages/cli/src/ui/commands/btwCommand.ts index c39904cb2f..696a9ed412 100644 --- a/packages/cli/src/ui/commands/btwCommand.ts +++ b/packages/cli/src/ui/commands/btwCommand.ts @@ -12,16 +12,14 @@ export const btwCommand: SlashCommand = { kind: CommandKind.BUILT_IN, autoExecute: true, isSafeConcurrent: true, - action: (context, args) => { + action: (_context, args) => { const query = args.trim(); if (!query) { return { - type: 'command_output', - value: [ - { - text: 'Please provide a question, e.g. /btw what is this regex doing?', - }, - ], + type: 'message', + messageType: 'error', + content: + 'Please provide a question, e.g. /btw what is this regex doing?', }; } diff --git a/packages/cli/src/ui/commands/types.ts b/packages/cli/src/ui/commands/types.ts index 2a9ee5914e..7d97710270 100644 --- a/packages/cli/src/ui/commands/types.ts +++ b/packages/cli/src/ui/commands/types.ts @@ -9,6 +9,7 @@ import type { HistoryItemWithoutId, HistoryItem, ConfirmationRequest, + BtwActionReturn, } from '../types.js'; import type { GitService, @@ -166,11 +167,6 @@ export interface LogoutActionReturn { type: 'logout'; } -export interface BtwActionReturn { - type: 'btw'; - query: string; -} - export type SlashCommandActionReturn = | CommandActionReturn | QuitActionReturn diff --git a/packages/cli/src/ui/hooks/slashCommandProcessor.ts b/packages/cli/src/ui/hooks/slashCommandProcessor.ts index f55503ad25..8cdbff1781 100644 --- a/packages/cli/src/ui/hooks/slashCommandProcessor.ts +++ b/packages/cli/src/ui/hooks/slashCommandProcessor.ts @@ -662,6 +662,9 @@ export const useSlashCommandProcessor = ( setCustomDialog(result.component); return { type: 'handled' }; } + case 'btw': { + return result; + } default: { const unhandled: never = result; throw new Error( diff --git a/packages/cli/src/ui/hooks/useBtw.test.ts b/packages/cli/src/ui/hooks/useBtw.test.ts new file mode 100644 index 0000000000..45171ae6f2 --- /dev/null +++ b/packages/cli/src/ui/hooks/useBtw.test.ts @@ -0,0 +1,120 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { act } from 'react'; +import { renderHook } from '../../test-utils/render.js'; +import { useBtw } from './useBtw.js'; +import { type GeminiClient, GeminiEventType } from '@google/gemini-cli-core'; + +describe('useBtw', () => { + let mockGeminiClient: { + sendBtwStream: ReturnType; + }; + + beforeEach(() => { + mockGeminiClient = { + sendBtwStream: vi.fn(), + }; + }); + + it('should initialize with inactive state', async () => { + const { result } = await renderHook(() => + useBtw(mockGeminiClient as unknown as GeminiClient), + ); + expect(result.current.isActive).toBe(false); + expect(result.current.query).toBe(''); + expect(result.current.response).toBe(''); + expect(result.current.isStreaming).toBe(false); + }); + + it('should update state during streaming', async () => { + let resolveStream: (value: void) => void; + const streamGate = new Promise((resolve) => { + resolveStream = resolve; + }); + + const mockStream = (async function* () { + yield { type: GeminiEventType.Content, value: 'Hello' }; + await streamGate; + yield { type: GeminiEventType.Content, value: ' world' }; + yield { type: GeminiEventType.Finished, value: { reason: 'STOP' } }; + })(); + mockGeminiClient.sendBtwStream.mockReturnValue(mockStream); + + const { result } = await renderHook(() => + useBtw(mockGeminiClient as unknown as GeminiClient), + ); + + let submitPromise!: Promise; + await act(async () => { + submitPromise = result.current.submitBtw('test query'); + }); + + // Check immediate state (it should be streaming because of streamGate) + expect(result.current.isActive).toBe(true); + expect(result.current.query).toBe('test query'); + expect(result.current.isStreaming).toBe(true); + expect(result.current.response).toBe('Hello'); + + await act(async () => { + resolveStream(); + await submitPromise; + }); + + // Check final state + expect(result.current.response).toBe('Hello world'); + expect(result.current.isStreaming).toBe(false); + }); + + it('should handle errors', async () => { + const mockStream = (async function* () { + yield { + type: GeminiEventType.Error, + value: { error: { message: 'API Error' } }, + }; + })(); + mockGeminiClient.sendBtwStream.mockReturnValue(mockStream); + + const { result } = await renderHook(() => + useBtw(mockGeminiClient as unknown as GeminiClient), + ); + + await act(async () => { + await result.current.submitBtw('test query'); + }); + + expect(result.current.error).toBe('API Error'); + expect(result.current.isStreaming).toBe(false); + }); + + it('should reset state on dismiss', async () => { + const mockStream = (async function* () { + yield { type: GeminiEventType.Content, value: 'partial' }; + // Hang + await new Promise(() => {}); + })(); + mockGeminiClient.sendBtwStream.mockReturnValue(mockStream); + + const { result } = await renderHook(() => + useBtw(mockGeminiClient as unknown as GeminiClient), + ); + + act(() => { + void result.current.submitBtw('test query'); + }); + + expect(result.current.isActive).toBe(true); + + act(() => { + result.current.dismissBtw(); + }); + + expect(result.current.isActive).toBe(false); + expect(result.current.query).toBe(''); + expect(result.current.response).toBe(''); + }); +}); diff --git a/packages/cli/src/ui/hooks/useBtw.ts b/packages/cli/src/ui/hooks/useBtw.ts new file mode 100644 index 0000000000..5035045416 --- /dev/null +++ b/packages/cli/src/ui/hooks/useBtw.ts @@ -0,0 +1,187 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { useCallback, useRef, useEffect, useReducer } from 'react'; +import { type GeminiClient, GeminiEventType } from '@google/gemini-cli-core'; + +export interface UseBtwReturn { + isActive: boolean; + query: string; + response: string; + isStreaming: boolean; + error: string | null; + submitBtw: (query: string) => Promise; + dismissBtw: () => void; +} + +interface BtwState { + isActive: boolean; + query: string; + response: string; + isStreaming: boolean; + error: string | null; +} + +type BtwAction = + | { type: 'SUBMIT'; query: string } + | { type: 'APPEND_CONTENT'; content: string } + | { type: 'ERROR'; error: string } + | { type: 'FINISHED' } + | { type: 'DISMISS' }; + +const initialState: BtwState = { + isActive: false, + query: '', + response: '', + isStreaming: false, + error: null, +}; + +const btwReducer = (state: BtwState, action: BtwAction): BtwState => { + switch (action.type) { + case 'SUBMIT': + return { + ...state, + isActive: true, + query: action.query, + response: '', + isStreaming: true, + error: null, + }; + case 'APPEND_CONTENT': + return { + ...state, + response: state.response + action.content, + }; + case 'ERROR': + return { + ...state, + error: action.error, + isStreaming: false, + }; + case 'FINISHED': + return { + ...state, + isStreaming: false, + }; + case 'DISMISS': + return initialState; + default: + return state; + } +}; + +export const useBtw = ( + geminiClient: GeminiClient | undefined, +): UseBtwReturn => { + const [state, dispatch] = useReducer(btwReducer, initialState); + + const abortControllerRef = useRef(null); + + const dismissBtw = useCallback(() => { + if (abortControllerRef.current) { + abortControllerRef.current.abort(); + abortControllerRef.current = null; + } + dispatch({ type: 'DISMISS' }); + }, []); + + const submitBtw = useCallback( + async (newQuery: string) => { + if (!geminiClient) return; + + // Abort any ongoing BTW stream + if (abortControllerRef.current) { + abortControllerRef.current.abort(); + } + + const abortController = new AbortController(); + abortControllerRef.current = abortController; + + dispatch({ type: 'SUBMIT', query: newQuery }); + + try { + const stream = geminiClient.sendBtwStream( + [{ text: newQuery }], + abortController.signal, + `btw-${Date.now()}`, + ); + + for await (const event of stream) { + if (abortController.signal.aborted) break; + + switch (event.type) { + case GeminiEventType.Content: + dispatch({ type: 'APPEND_CONTENT', content: event.value ?? '' }); + break; + case GeminiEventType.Error: { + const value = event.value; + let errorMessage = 'Unknown error'; + if ( + typeof value === 'object' && + value !== null && + 'error' in value && + typeof value.error === 'object' && + value.error !== null + ) { + const errorObj = value.error; + errorMessage = + 'message' in errorObj + ? String(errorObj.message) + : String(errorObj); + } else { + errorMessage = String(value); + } + dispatch({ + type: 'ERROR', + error: errorMessage, + }); + break; + } + case GeminiEventType.Finished: + dispatch({ type: 'FINISHED' }); + break; + case GeminiEventType.UserCancelled: + dispatch({ type: 'FINISHED' }); + break; + default: + break; + } + } + } catch (err) { + if (err instanceof Error && err.name === 'AbortError') { + // Ignore aborts + } else { + dispatch({ + type: 'ERROR', + error: err instanceof Error ? err.message : String(err), + }); + } + } finally { + if (abortControllerRef.current === abortController) { + dispatch({ type: 'FINISHED' }); + } + } + }, + [geminiClient], + ); + + // Cleanup on unmount + useEffect( + () => () => { + if (abortControllerRef.current) { + abortControllerRef.current.abort(); + } + }, + [], + ); + + return { + ...state, + submitBtw, + dismissBtw, + }; +}; diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index a2621c4546..bd14dfc646 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -990,6 +990,9 @@ export const useGeminiStream = ( case 'handled': { return { queryToSend: null, shouldProceed: false }; } + case 'btw': { + return { queryToSend: null, shouldProceed: false }; + } default: { const unreachable: never = slashCommandResult; throw new Error( diff --git a/packages/cli/src/ui/types.ts b/packages/cli/src/ui/types.ts index 6fbc3151d8..f106a8837b 100644 --- a/packages/cli/src/ui/types.ts +++ b/packages/cli/src/ui/types.ts @@ -500,6 +500,11 @@ export interface SubmitPromptResult { content: PartListUnion; } +export interface BtwActionReturn { + type: 'btw'; + query: string; +} + /** * Defines the result of the slash command processor for its consumer (useGeminiStream). */ @@ -513,7 +518,8 @@ export type SlashCommandProcessorResult = | { type: 'handled'; // Indicates the command was processed and no further action is needed. } - | SubmitPromptResult; + | SubmitPromptResult + | BtwActionReturn; export interface ConfirmationRequest { prompt: ReactNode;