refactor(cli): Centralize history management via useHistoryManager hook (#261)

This commit is contained in:
Allen Hutchison 2025-05-06 16:20:28 -07:00 committed by GitHub
parent adeda6a5b3
commit 7d13f24288
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 293 additions and 455 deletions

View File

@ -27,6 +27,7 @@ import { HistoryItemDisplay } from './components/HistoryItemDisplay.js';
import { useCompletion } from './hooks/useCompletion.js';
import { SuggestionsDisplay } from './components/SuggestionsDisplay.js';
import { isAtCommand, isSlashCommand } from './utils/commandUtils.js';
import { useHistory } from './hooks/useHistoryManager.js';
interface AppProps {
config: Config;
@ -35,7 +36,7 @@ interface AppProps {
}
export const App = ({ config, settings, cliVersion }: AppProps) => {
const [history, setHistory] = useState<HistoryItem[]>([]);
const { history, addItem, updateItem, clearItems } = useHistory();
const [startupWarnings, setStartupWarnings] = useState<string[]>([]);
const [showHelp, setShowHelp] = useState<boolean>(false);
const {
@ -57,7 +58,9 @@ export const App = ({ config, settings, cliVersion }: AppProps) => {
debugMessage,
slashCommands,
} = useGeminiStream(
setHistory,
addItem,
updateItem,
clearItems,
refreshStatic,
setShowHelp,
config,

View File

@ -18,25 +18,15 @@ import {
IndividualToolCallDisplay,
ToolCallStatus,
} from '../types.js';
const addHistoryItem = (
setHistory: React.Dispatch<React.SetStateAction<HistoryItem[]>>,
itemData: Omit<HistoryItem, 'id'>,
id: number,
) => {
setHistory((prevHistory) => [
...prevHistory,
{ ...itemData, id } as HistoryItem,
]);
};
import { UseHistoryManagerReturn } from './useHistoryManager.js';
interface HandleAtCommandParams {
query: string;
config: Config;
setHistory: React.Dispatch<React.SetStateAction<HistoryItem[]>>;
addItem: UseHistoryManagerReturn['addItem'];
updateItem: UseHistoryManagerReturn['updateItem'];
setDebugMessage: React.Dispatch<React.SetStateAction<string>>;
getNextMessageId: (baseTimestamp: number) => number;
userMessageTimestamp: number;
messageId: number;
}
interface HandleAtCommandResult {
@ -53,7 +43,6 @@ function parseAtCommand(
): { textBefore: string; atPath: string; textAfter: string } | null {
let atIndex = -1;
for (let i = 0; i < query.length; i++) {
// Find the first '@' that is not preceded by a '\'
if (query[i] === '@' && (i === 0 || query[i - 1] !== '\\')) {
atIndex = i;
break;
@ -61,7 +50,7 @@ function parseAtCommand(
}
if (atIndex === -1) {
return null; // No '@' command found
return null;
}
const textBefore = query.substring(0, atIndex).trim();
@ -70,15 +59,11 @@ function parseAtCommand(
while (pathEndIndex < query.length) {
const char = query[pathEndIndex];
if (inEscape) {
// Current char is escaped, move past it
inEscape = false;
} else if (char === '\\') {
// Start of an escape sequence
inEscape = true;
} else if (/\s/.test(char)) {
// Unescaped whitespace marks the end of the path
break;
}
pathEndIndex++;
@ -86,7 +71,6 @@ function parseAtCommand(
const rawAtPath = query.substring(atIndex, pathEndIndex);
const textAfter = query.substring(pathEndIndex).trim();
const atPath = unescapePath(rawAtPath);
return { textBefore, atPath, textAfter };
@ -94,80 +78,40 @@ function parseAtCommand(
/**
* Processes user input potentially containing an '@<path>' command.
* It finds the first '@<path>', checks if the path is a file or directory,
* prepares the appropriate path specification for the read_many_files tool,
* updates the UI, and prepares the query for the LLM, incorporating the
* file content and surrounding text.
* If found, it attempts to read the specified file/directory using the
* 'read_many_files' tool, adds the user query and tool result/error to history,
* and prepares the content for the LLM.
*
* @returns An object containing the potentially modified query (or null)
* and a flag indicating if the main hook should proceed.
* @returns An object indicating whether the main hook should proceed with an
* LLM call and the processed query parts (including file content).
*/
export async function handleAtCommand({
query,
config,
setHistory,
addItem: addItem,
setDebugMessage,
getNextMessageId,
userMessageTimestamp,
messageId: userMessageTimestamp,
}: HandleAtCommandParams): Promise<HandleAtCommandResult> {
const trimmedQuery = query.trim();
const parsedCommand = parseAtCommand(trimmedQuery);
// If no @ command, add user query normally and proceed to LLM
if (!parsedCommand) {
// If no '@' was found, treat the whole query as user text and proceed
// This allows users to just type text without an @ command
addHistoryItem(
setHistory,
{ type: 'user', text: query },
userMessageTimestamp,
);
// Let the main hook decide what to do (likely send to LLM)
addItem({ type: 'user', text: query }, userMessageTimestamp);
return { processedQuery: [{ text: query }], shouldProceed: true };
// Or, if an @ command is *required* when the function is called:
/*
const errorTimestamp = getNextMessageId(userMessageTimestamp);
addHistoryItem(
setHistory,
{ type: 'error', text: 'Error: Could not find @ command.' },
errorTimestamp,
);
return { processedQuery: null, shouldProceed: false };
*/
}
const { textBefore, atPath, textAfter } = parsedCommand;
// Add the original user query to history *before* processing
addHistoryItem(
setHistory,
{ type: 'user', text: query },
userMessageTimestamp,
);
// Add the original user query to history first
addItem({ type: 'user', text: query }, userMessageTimestamp);
const pathPart = atPath.substring(1); // Remove the leading '@'
const pathPart = atPath.substring(1); // Remove leading '@'
if (!pathPart) {
const errorTimestamp = getNextMessageId(userMessageTimestamp);
addHistoryItem(
setHistory,
addItem(
{ type: 'error', text: 'Error: No path specified after @.' },
errorTimestamp,
);
return { processedQuery: null, shouldProceed: false };
}
addHistoryItem(
setHistory,
{ type: 'user', text: query },
userMessageTimestamp,
);
if (!pathPart) {
const errorTimestamp = getNextMessageId(userMessageTimestamp);
addHistoryItem(
setHistory,
{ type: 'error', text: 'Error: No path specified after @.' },
errorTimestamp,
userMessageTimestamp,
);
return { processedQuery: null, shouldProceed: false };
}
@ -176,39 +120,31 @@ export async function handleAtCommand({
const readManyFilesTool = toolRegistry.getTool('read_many_files');
if (!readManyFilesTool) {
const errorTimestamp = getNextMessageId(userMessageTimestamp);
addHistoryItem(
setHistory,
addItem(
{ type: 'error', text: 'Error: read_many_files tool not found.' },
errorTimestamp,
userMessageTimestamp,
);
return { processedQuery: null, shouldProceed: false };
}
// --- Path Handling for @ command ---
// Determine path spec (file or directory glob)
let pathSpec = pathPart;
const contentLabel = pathPart;
try {
// Resolve the path relative to the target directory
const absolutePath = path.resolve(config.getTargetDir(), pathPart);
const stats = await fs.stat(absolutePath);
if (stats.isDirectory()) {
// If it's a directory, ensure it ends with a globstar for recursive read
pathSpec = pathPart.endsWith('/') ? `${pathPart}**` : `${pathPart}/**`;
setDebugMessage(`Path resolved to directory, using glob: ${pathSpec}`);
} else {
// It's a file, use the original pathPart as pathSpec
setDebugMessage(`Path resolved to file: ${pathSpec}`);
}
} catch (error) {
// If stat fails (e.g., file/dir not found), proceed with the original pathPart.
// The read_many_files tool will handle the error if it's invalid.
// If stat fails (e.g., not found), proceed with original path.
// The tool itself will handle the error during execution.
if (isNodeError(error) && error.code === 'ENOENT') {
setDebugMessage(`Path not found, proceeding with original: ${pathSpec}`);
} else {
// Log other stat errors but still proceed
console.error(`Error stating path ${pathPart}:`, error);
setDebugMessage(
`Error stating path, proceeding with original: ${pathSpec}`,
@ -217,8 +153,6 @@ export async function handleAtCommand({
}
const toolArgs = { paths: [pathSpec] };
// --- End Path Handling ---
let toolCallDisplay: IndividualToolCallDisplay;
try {
@ -234,6 +168,7 @@ export async function handleAtCommand({
confirmationDetails: undefined,
};
// Prepare the query parts for the LLM
const processedQueryParts = [];
if (textBefore) {
processedQueryParts.push({ text: textBefore });
@ -244,21 +179,20 @@ export async function handleAtCommand({
if (textAfter) {
processedQueryParts.push({ text: textAfter });
}
const processedQuery: PartListUnion = processedQueryParts;
const toolGroupId = getNextMessageId(userMessageTimestamp);
addHistoryItem(
setHistory,
// Add the successful tool result to history
addItem(
{ type: 'tool_group', tools: [toolCallDisplay] } as Omit<
HistoryItem,
'id'
>,
toolGroupId,
userMessageTimestamp,
);
return { processedQuery, shouldProceed: true };
} catch (error) {
// Handle errors during tool execution
toolCallDisplay = {
callId: `client-read-${userMessageTimestamp}`,
name: readManyFilesTool.displayName,
@ -268,14 +202,13 @@ export async function handleAtCommand({
confirmationDetails: undefined,
};
const toolGroupId = getNextMessageId(userMessageTimestamp);
addHistoryItem(
setHistory,
// Add the error tool result to history
addItem(
{ type: 'tool_group', tools: [toolCallDisplay] } as Omit<
HistoryItem,
'id'
>,
toolGroupId,
userMessageTimestamp,
);
return { processedQuery: null, shouldProceed: false };

View File

@ -8,90 +8,79 @@ import { exec as _exec } from 'child_process';
import { useCallback } from 'react';
import { Config } from '@gemini-code/server';
import { type PartListUnion } from '@google/genai';
import { HistoryItem, StreamingState } from '../types.js';
import { StreamingState } from '../types.js';
import { getCommandFromQuery } from '../utils/commandUtils.js';
import { UseHistoryManagerReturn } from './useHistoryManager.js';
// Helper function (consider moving to a shared util if used elsewhere)
const addHistoryItem = (
setHistory: React.Dispatch<React.SetStateAction<HistoryItem[]>>,
itemData: Omit<HistoryItem, 'id'>,
id: number,
) => {
setHistory((prevHistory) => [
...prevHistory,
{ ...itemData, id } as HistoryItem,
]);
};
/**
* Hook to process shell commands (e.g., !ls, $pwd).
* Executes the command in the target directory and adds output/errors to history.
*/
export const useShellCommandProcessor = (
setHistory: React.Dispatch<React.SetStateAction<HistoryItem[]>>,
addItemToHistory: UseHistoryManagerReturn['addItem'],
setStreamingState: React.Dispatch<React.SetStateAction<StreamingState>>,
setDebugMessage: React.Dispatch<React.SetStateAction<string>>,
getNextMessageId: (baseTimestamp: number) => number,
config: Config,
) => {
/**
* Checks if the query is a shell command, executes it, and adds results to history.
* @returns True if the query was handled as a shell command, false otherwise.
*/
const handleShellCommand = useCallback(
(rawQuery: PartListUnion): boolean => {
if (typeof rawQuery !== 'string') {
return false; // Passthrough only works with string commands
return false;
}
const [symbol] = getCommandFromQuery(rawQuery);
if (symbol !== '!' && symbol !== '$') {
return false;
}
// Remove symbol from rawQuery
const trimmed = rawQuery.trim().slice(1).trimStart();
const commandToExecute = rawQuery.trim().slice(1).trimStart();
// Stop if command is empty
if (!trimmed) {
return false;
const userMessageTimestamp = Date.now();
addItemToHistory({ type: 'user', text: rawQuery }, userMessageTimestamp);
if (!commandToExecute) {
addItemToHistory(
{ type: 'error', text: 'Empty shell command.' },
userMessageTimestamp,
);
return true; // Handled (by showing error)
}
// Add user message *before* execution starts
const userMessageTimestamp = Date.now();
addHistoryItem(
setHistory,
{ type: 'user', text: rawQuery },
userMessageTimestamp,
);
// Execute and capture output
const targetDir = config.getTargetDir();
setDebugMessage(`Executing shell command in ${targetDir}: ${trimmed}`);
setDebugMessage(
`Executing shell command in ${targetDir}: ${commandToExecute}`,
);
const execOptions = {
cwd: targetDir,
};
// Set state to Responding while the command runs
setStreamingState(StreamingState.Responding);
_exec(trimmed, execOptions, (error, stdout, stderr) => {
const timestamp = getNextMessageId(userMessageTimestamp); // Use user message time as base
_exec(commandToExecute, execOptions, (error, stdout, stderr) => {
if (error) {
addHistoryItem(
setHistory,
addItemToHistory(
{ type: 'error', text: error.message },
timestamp,
userMessageTimestamp,
);
} else if (stderr) {
// Treat stderr as info for passthrough, as some tools use it for non-error output
addHistoryItem(setHistory, { type: 'info', text: stderr }, timestamp);
} else {
// Add stdout as an info message
addHistoryItem(
setHistory,
{ type: 'info', text: stdout || '(Command produced no output)' },
timestamp,
let output = '';
if (stdout) output += stdout;
if (stderr) output += (output ? '\n' : '') + stderr; // Include stderr as info
addItemToHistory(
{ type: 'info', text: output || '(Command produced no output)' },
userMessageTimestamp,
);
}
// Set state back to Idle *after* command finishes and output is added
setStreamingState(StreamingState.Idle);
});
return true; // Command was handled
return true; // Command was initiated
},
[config, setDebugMessage, setHistory, setStreamingState, getNextMessageId],
[config, setDebugMessage, addItemToHistory, setStreamingState],
);
return { handleShellCommand };

View File

@ -6,33 +6,25 @@
import { useCallback, useMemo } from 'react';
import { type PartListUnion } from '@google/genai';
import { HistoryItem } from '../types.js';
import { getCommandFromQuery } from '../utils/commandUtils.js';
import { UseHistoryManagerReturn } from './useHistoryManager.js';
export interface SlashCommand {
name: string; // slash command
altName?: string; // alternative name for the command
description: string; // flavor text in UI
name: string;
altName?: string;
description: string;
action: (value: PartListUnion) => void;
}
const addHistoryItem = (
setHistory: React.Dispatch<React.SetStateAction<HistoryItem[]>>,
itemData: Omit<HistoryItem, 'id'>,
id: number,
) => {
setHistory((prevHistory) => [
...prevHistory,
{ ...itemData, id } as HistoryItem,
]);
};
/**
* Hook to define and process slash commands (e.g., /help, /clear).
*/
export const useSlashCommandProcessor = (
setHistory: React.Dispatch<React.SetStateAction<HistoryItem[]>>,
addItem: UseHistoryManagerReturn['addItem'],
clearItems: UseHistoryManagerReturn['clearItems'],
refreshStatic: () => void,
setShowHelp: React.Dispatch<React.SetStateAction<boolean>>,
setDebugMessage: React.Dispatch<React.SetStateAction<string>>,
getNextMessageId: (baseTimestamp: number) => number,
openThemeDialog: () => void,
) => {
const slashCommands: SlashCommand[] = useMemo(
@ -50,9 +42,8 @@ export const useSlashCommandProcessor = (
name: 'clear',
description: 'clear the screen',
action: (_value: PartListUnion) => {
// This just clears the *UI* history, not the model history.
setDebugMessage('Clearing terminal.');
setHistory((_) => []);
clearItems();
refreshStatic();
},
},
@ -69,22 +60,19 @@ export const useSlashCommandProcessor = (
description: '',
action: (_value: PartListUnion) => {
setDebugMessage('Quitting. Good-bye.');
getNextMessageId(Date.now());
process.exit(0);
},
},
],
[
setDebugMessage,
setShowHelp,
setHistory,
refreshStatic,
openThemeDialog,
getNextMessageId,
],
[setDebugMessage, setShowHelp, refreshStatic, openThemeDialog, clearItems],
);
// Checks if the query is a slash command and executes the command if it is.
/**
* Checks if the query is a slash command and executes it if found.
* Adds user query and potential error messages to history.
* @returns True if the query was handled as a slash command (valid or invalid),
* false otherwise.
*/
const handleSlashCommand = useCallback(
(rawQuery: PartListUnion): boolean => {
if (typeof rawQuery !== 'string') {
@ -94,32 +82,33 @@ export const useSlashCommandProcessor = (
const trimmed = rawQuery.trim();
const [symbol, test] = getCommandFromQuery(trimmed);
// Skip non slash commands
if (symbol !== '/' && symbol !== '?') {
return false;
}
const userMessageTimestamp = Date.now();
addItem({ type: 'user', text: trimmed }, userMessageTimestamp);
for (const cmd of slashCommands) {
if (
test === cmd.name ||
test === cmd.altName ||
symbol === cmd.altName
) {
// Add user message *before* execution
const userMessageTimestamp = Date.now();
addHistoryItem(
setHistory,
{ type: 'user', text: trimmed },
userMessageTimestamp,
);
cmd.action(trimmed);
return true; // Command was handled
return true;
}
}
return false; // Not a recognized slash command
// Unknown command: Add error message
addItem(
{ type: 'error', text: `Unknown command: ${trimmed}` },
userMessageTimestamp, // Use same base timestamp for related error
);
return true; // Indicate command was processed (even though invalid)
},
[setHistory, slashCommands],
[addItem, slashCommands],
);
return { handleSlashCommand, slashCommands };

View File

@ -8,7 +8,7 @@ import { useState, useRef, useCallback, useEffect } from 'react';
import { useInput } from 'ink';
import {
GeminiClient,
GeminiEventType as ServerGeminiEventType, // Rename to avoid conflict
GeminiEventType as ServerGeminiEventType,
getErrorMessage,
isNodeError,
Config,
@ -32,21 +32,16 @@ import { useSlashCommandProcessor } from './slashCommandProcessor.js';
import { useShellCommandProcessor } from './shellCommandProcessor.js';
import { handleAtCommand } from './atCommandProcessor.js';
import { findSafeSplitPoint } from '../utils/markdownUtilities.js';
import { UseHistoryManagerReturn } from './useHistoryManager.js';
const addHistoryItem = (
setHistory: React.Dispatch<React.SetStateAction<HistoryItem[]>>,
itemData: Omit<HistoryItem, 'id'>,
id: number,
) => {
setHistory((prevHistory) => [
...prevHistory,
{ ...itemData, id } as HistoryItem,
]);
};
// Hook now accepts apiKey and model
/**
* Hook to manage the Gemini stream, handle user input, process commands,
* and interact with the Gemini API and history manager.
*/
export const useGeminiStream = (
setHistory: React.Dispatch<React.SetStateAction<HistoryItem[]>>,
addItem: UseHistoryManagerReturn['addItem'],
updateItem: UseHistoryManagerReturn['updateItem'],
clearItems: UseHistoryManagerReturn['clearItems'],
refreshStatic: () => void,
setShowHelp: React.Dispatch<React.SetStateAction<boolean>>,
config: Config,
@ -61,99 +56,56 @@ export const useGeminiStream = (
const abortControllerRef = useRef<AbortController | null>(null);
const chatSessionRef = useRef<Chat | null>(null);
const geminiClientRef = useRef<GeminiClient | null>(null);
const messageIdCounterRef = useRef(0);
const currentGeminiMessageIdRef = useRef<number | null>(null);
// ID Generation Callback
const getNextMessageId = useCallback((baseTimestamp: number): number => {
// Increment *before* adding to ensure uniqueness against the base timestamp
messageIdCounterRef.current += 1;
return baseTimestamp + messageIdCounterRef.current;
}, []);
// Instantiate command processors
const { handleSlashCommand, slashCommands } = useSlashCommandProcessor(
setHistory,
addItem,
clearItems,
refreshStatic,
setShowHelp,
setDebugMessage,
getNextMessageId,
openThemeDialog,
);
const { handleShellCommand } = useShellCommandProcessor(
setHistory,
addItem,
setStreamingState,
setDebugMessage,
getNextMessageId,
config,
);
// Initialize Client Effect - uses props now
useEffect(() => {
setInitError(null);
if (!geminiClientRef.current) {
try {
geminiClientRef.current = new GeminiClient(config);
} catch (error: unknown) {
setInitError(
`Failed to initialize client: ${getErrorMessage(error) || 'Unknown error'}`,
);
const errorMsg = `Failed to initialize client: ${getErrorMessage(error) || 'Unknown error'}`;
setInitError(errorMsg);
addItem({ type: 'error', text: errorMsg }, Date.now());
}
}
}, [config]);
}, [config, addItem]);
// Input Handling Effect (remains the same)
useInput((input, key) => {
useInput((_input, key) => {
if (streamingState === StreamingState.Responding && key.escape) {
abortControllerRef.current?.abort();
}
});
// Helper function to update Gemini message content
const updateGeminiMessage = useCallback(
(messageId: number, newContent: string) => {
setHistory((prevHistory) =>
prevHistory.map((item) =>
item.id === messageId && item.type === 'gemini'
? { ...item, text: newContent }
: item,
),
);
updateItem(messageId, { text: newContent });
},
[setHistory],
[updateItem],
);
// Helper function to update Gemini message content
const updateAndAddGeminiMessageContent = useCallback(
(
messageId: number,
previousContent: string,
nextId: number,
nextContent: string,
) => {
setHistory((prevHistory) => {
const beforeNextHistory = prevHistory.map((item) =>
item.id === messageId ? { ...item, text: previousContent } : item,
);
return [
...beforeNextHistory,
{ id: nextId, type: 'gemini_content', text: nextContent },
];
});
},
[setHistory],
);
// Improved submit query function
const submitQuery = useCallback(
async (query: PartListUnion) => {
if (streamingState === StreamingState.Responding) return;
if (typeof query === 'string' && query.trim().length === 0) return;
const userMessageTimestamp = Date.now();
messageIdCounterRef.current = 0; // Reset counter for this new submission
let queryToSendToGemini: PartListUnion | null = null;
setShowHelp(false);
@ -162,50 +114,33 @@ export const useGeminiStream = (
const trimmedQuery = query.trim();
setDebugMessage(`User query: '${trimmedQuery}'`);
// 1. Check for Slash Commands (/)
if (handleSlashCommand(trimmedQuery)) {
return;
}
// Handle UI-only commands first
if (handleSlashCommand(trimmedQuery)) return;
if (handleShellCommand(trimmedQuery)) return;
// 2. Check for Shell Commands (! or $)
if (handleShellCommand(trimmedQuery)) {
return;
}
// 3. Check for @ Commands using the utility function
// Handle @-commands (which might involve tool calls)
if (isAtCommand(trimmedQuery)) {
const atCommandResult = await handleAtCommand({
query: trimmedQuery,
config,
setHistory,
addItem,
updateItem,
setDebugMessage,
getNextMessageId,
userMessageTimestamp,
messageId: userMessageTimestamp,
});
if (!atCommandResult.shouldProceed) {
return; // @ command handled it (e.g., error) or decided not to proceed
}
if (!atCommandResult.shouldProceed) return;
queryToSendToGemini = atCommandResult.processedQuery;
// User message and tool UI were added by handleAtCommand
} else {
// 4. It's a normal query for Gemini
addHistoryItem(
setHistory,
{ type: 'user', text: trimmedQuery },
userMessageTimestamp,
);
// Normal query for Gemini
addItem({ type: 'user', text: trimmedQuery }, userMessageTimestamp);
queryToSendToGemini = trimmedQuery;
}
} else {
// 5. It's a function response (PartListUnion that isn't a string)
// Tool call/response UI handles history. Always proceed.
// It's a function response (PartListUnion that isn't a string)
queryToSendToGemini = query;
}
// --- Proceed to Gemini API call ---
if (queryToSendToGemini === null) {
// Should only happen if @ command failed and returned null query
setDebugMessage(
'Query processing resulted in null, not sending to Gemini.',
);
@ -214,7 +149,9 @@ export const useGeminiStream = (
const client = geminiClientRef.current;
if (!client) {
setInitError('Gemini client is not available.');
const errorMsg = 'Gemini client is not available.';
setInitError(errorMsg);
addItem({ type: 'error', text: errorMsg }, Date.now());
return;
}
@ -222,7 +159,9 @@ export const useGeminiStream = (
try {
chatSessionRef.current = await client.startChat();
} catch (err: unknown) {
setInitError(`Failed to start chat: ${getErrorMessage(err)}`);
const errorMsg = `Failed to start chat: ${getErrorMessage(err)}`;
setInitError(errorMsg);
addItem({ type: 'error', text: errorMsg }, Date.now());
setStreamingState(StreamingState.Idle);
return;
}
@ -231,51 +170,39 @@ export const useGeminiStream = (
setStreamingState(StreamingState.Responding);
setInitError(null);
const chat = chatSessionRef.current;
let currentToolGroupId: number | null = null;
let currentToolGroupMessageId: number | null = null;
try {
abortControllerRef.current = new AbortController();
const signal = abortControllerRef.current.signal;
// Use the determined query for the Gemini call
const stream = client.sendMessageStream(
chat,
queryToSendToGemini,
signal,
);
// Process the stream events from the server logic
let currentGeminiText = ''; // To accumulate message content
let currentGeminiText = '';
let hasInitialGeminiResponse = false;
for await (const event of stream) {
if (signal.aborted) break;
if (event.type === ServerGeminiEventType.Content) {
// For content events, accumulate the text and update an existing message or create a new one
currentGeminiText += event.value;
// Reset group because we're now adding a user message to the history. If we didn't reset the
// group here then any subsequent tool calls would get grouped before this message resulting in
// a misordering of history.
currentToolGroupId = null;
currentToolGroupMessageId = null; // Reset group on new text content
if (!hasInitialGeminiResponse) {
// Create a new Gemini message if this is the first content event
hasInitialGeminiResponse = true;
const eventTimestamp = getNextMessageId(userMessageTimestamp);
currentGeminiMessageIdRef.current = eventTimestamp;
addHistoryItem(
setHistory,
const eventId = addItem(
{ type: 'gemini', text: currentGeminiText },
eventTimestamp,
userMessageTimestamp,
);
currentGeminiMessageIdRef.current = eventId;
} else if (currentGeminiMessageIdRef.current !== null) {
// Split large messages for better rendering performance
const splitPoint = findSafeSplitPoint(currentGeminiText);
if (splitPoint === currentGeminiText.length) {
// Update the existing message with accumulated content
updateGeminiMessage(
currentGeminiMessageIdRef.current,
currentGeminiText,
@ -291,40 +218,33 @@ export const useGeminiStream = (
// broken up so that there are more "statically" rendered.
const originalMessageRef = currentGeminiMessageIdRef.current;
const beforeText = currentGeminiText.substring(0, splitPoint);
currentGeminiMessageIdRef.current =
getNextMessageId(userMessageTimestamp);
const afterText = currentGeminiText.substring(splitPoint);
currentGeminiText = afterText;
updateAndAddGeminiMessageContent(
originalMessageRef,
beforeText,
currentGeminiMessageIdRef.current,
afterText,
currentGeminiText = afterText; // Continue accumulating from split point
updateItem(originalMessageRef, { text: beforeText });
const nextId = addItem(
{ type: 'gemini_content', text: afterText },
userMessageTimestamp,
);
currentGeminiMessageIdRef.current = nextId;
}
}
} else if (event.type === ServerGeminiEventType.ToolCallRequest) {
// Reset the Gemini message tracking for the next response
currentGeminiText = '';
hasInitialGeminiResponse = false;
currentGeminiMessageIdRef.current = null;
const { callId, name, args } = event.value;
const cliTool = toolRegistry.getTool(name); // Get the full CLI tool
const cliTool = toolRegistry.getTool(name);
if (!cliTool) {
console.error(`CLI Tool "${name}" not found!`);
continue;
}
if (currentToolGroupId === null) {
currentToolGroupId = getNextMessageId(userMessageTimestamp);
// Add explicit cast to Omit<HistoryItem, 'id'>
addHistoryItem(
setHistory,
// Create a new tool group if needed
if (currentToolGroupMessageId === null) {
currentToolGroupMessageId = addItem(
{ type: 'tool_group', tools: [] } as Omit<HistoryItem, 'id'>,
currentToolGroupId,
userMessageTimestamp,
);
}
@ -335,7 +255,6 @@ export const useGeminiStream = (
description = `Error: Unable to get description: ${getErrorMessage(e)}`;
}
// Create the UI display object matching IndividualToolCallDisplay
const toolCallDisplay: IndividualToolCallDisplay = {
callId,
name: cliTool.displayName,
@ -345,25 +264,27 @@ export const useGeminiStream = (
confirmationDetails: undefined,
};
// Add pending tool call to the UI history group
setHistory((prevHistory) =>
prevHistory.map((item) => {
if (
item.id === currentToolGroupId &&
item.type === 'tool_group'
) {
// Ensure item.tools exists and is an array before spreading
const currentTools = Array.isArray(item.tools)
? item.tools
: [];
// Add the pending tool call to the current group
if (currentToolGroupMessageId !== null) {
updateItem(
currentToolGroupMessageId,
(
currentItem: HistoryItem,
): Partial<Omit<HistoryItem, 'id'>> => {
if (currentItem?.type !== 'tool_group') {
console.error(
`Attempted to update non-tool-group item ${currentItem?.id} as tool group.`,
);
return currentItem as Partial<Omit<HistoryItem, 'id'>>;
}
const currentTools = currentItem.tools;
return {
...item,
tools: [...currentTools, toolCallDisplay], // Add the complete display object
};
}
return item;
}),
);
...currentItem,
tools: [...currentTools, toolCallDisplay],
} as Partial<Omit<HistoryItem, 'id'>>;
},
);
}
} else if (event.type === ServerGeminiEventType.ToolCallResponse) {
const status = event.value.error
? ToolCallStatus.Error
@ -378,21 +299,20 @@ export const useGeminiStream = (
confirmationDetails,
);
setStreamingState(StreamingState.WaitingForConfirmation);
return;
return; // Wait for user confirmation
}
}
} // End stream loop
setStreamingState(StreamingState.Idle);
} catch (error: unknown) {
if (!isNodeError(error) || error.name !== 'AbortError') {
console.error('Error processing stream or executing tool:', error);
addHistoryItem(
setHistory,
addItem(
{
type: 'error',
text: `[Error: ${getErrorMessage(error)}]`,
text: `[Stream Error: ${getErrorMessage(error)}]`,
},
getNextMessageId(userMessageTimestamp),
userMessageTimestamp,
);
}
setStreamingState(StreamingState.Idle);
@ -400,28 +320,35 @@ export const useGeminiStream = (
abortControllerRef.current = null;
}
// --- Helper functions for updating tool UI ---
function updateConfirmingFunctionStatusUI(
callId: string,
confirmationDetails: ToolCallConfirmationDetails | undefined,
) {
setHistory((prevHistory) =>
prevHistory.map((item) => {
if (item.id === currentToolGroupId && item.type === 'tool_group') {
return {
...item,
tools: item.tools.map((tool) =>
tool.callId === callId
? {
...tool,
status: ToolCallStatus.Confirming,
confirmationDetails,
}
: tool,
),
};
if (currentToolGroupMessageId === null) return;
updateItem(
currentToolGroupMessageId,
(currentItem: HistoryItem): Partial<Omit<HistoryItem, 'id'>> => {
if (currentItem?.type !== 'tool_group') {
console.error(
`Attempted to update non-tool-group item ${currentItem?.id} status.`,
);
return currentItem as Partial<Omit<HistoryItem, 'id'>>;
}
return item;
}),
return {
...currentItem,
tools: (currentItem.tools || []).map((tool) =>
tool.callId === callId
? {
...tool,
status: ToolCallStatus.Confirming,
confirmationDetails,
}
: tool,
),
} as Partial<Omit<HistoryItem, 'id'>>;
},
);
}
@ -429,29 +356,35 @@ export const useGeminiStream = (
toolResponse: ToolCallResponseInfo,
status: ToolCallStatus,
) {
setHistory((prevHistory) =>
prevHistory.map((item) => {
if (item.id === currentToolGroupId && item.type === 'tool_group') {
return {
...item,
tools: item.tools.map((tool) => {
if (tool.callId === toolResponse.callId) {
return {
...tool,
status,
resultDisplay: toolResponse.resultDisplay,
};
} else {
return tool;
}
}),
};
if (currentToolGroupMessageId === null) return;
updateItem(
currentToolGroupMessageId,
(currentItem: HistoryItem): Partial<Omit<HistoryItem, 'id'>> => {
if (currentItem?.type !== 'tool_group') {
console.error(
`Attempted to update non-tool-group item ${currentItem?.id} response.`,
);
return currentItem as Partial<Omit<HistoryItem, 'id'>>;
}
return item;
}),
return {
...currentItem,
tools: (currentItem.tools || []).map((tool) => {
if (tool.callId === toolResponse.callId) {
return {
...tool,
status,
resultDisplay: toolResponse.resultDisplay,
};
} else {
return tool;
}
}),
} as Partial<Omit<HistoryItem, 'id'>>;
},
);
}
// Wires the server-side confirmation callback to UI updates and state changes
function wireConfirmationSubmission(
confirmationDetails: ServerToolCallConfirmationDetails,
): ToolCallConfirmationDetails {
@ -460,6 +393,7 @@ export const useGeminiStream = (
const resubmittingConfirm = async (
outcome: ToolConfirmationOutcome,
) => {
// Call the original server-side handler first
originalConfirmationDetails.onConfirm(outcome);
if (outcome === ToolConfirmationOutcome.Cancel) {
@ -480,41 +414,18 @@ export const useGeminiStream = (
response: { error: 'User rejected function call.' },
},
};
const responseInfo: ToolCallResponseInfo = {
callId: request.callId,
responsePart: functionResponse,
resultDisplay,
error: undefined,
error: new Error('User rejected function call.'),
};
// Update UI to show cancellation/error
updateFunctionResponseUI(responseInfo, ToolCallStatus.Error);
setStreamingState(StreamingState.Idle);
} else {
const tool = toolRegistry.getTool(request.name);
if (!tool) {
throw new Error(
`Tool "${request.name}" not found or is not registered.`,
);
}
const result = await tool.execute(request.args);
const functionResponse: Part = {
functionResponse: {
name: request.name,
id: request.callId,
response: { output: result.llmContent },
},
};
const responseInfo: ToolCallResponseInfo = {
callId: request.callId,
responsePart: functionResponse,
resultDisplay: result.returnDisplay,
error: undefined,
};
updateFunctionResponseUI(responseInfo, ToolCallStatus.Success);
setStreamingState(StreamingState.Idle);
await submitQuery(functionResponse);
// If accepted, set state back to Responding to wait for server execution/response
setStreamingState(StreamingState.Responding);
}
};
@ -524,21 +435,19 @@ export const useGeminiStream = (
};
}
},
// Dependencies need careful review
[
streamingState,
setHistory,
config,
getNextMessageId,
updateGeminiMessage,
handleSlashCommand,
handleShellCommand,
// handleAtCommand is implicitly included via its direct call
setDebugMessage, // Added dependency for handleAtCommand & passthrough
setStreamingState, // Added dependency for handlePassthroughCommand
updateAndAddGeminiMessageContent,
setDebugMessage,
setStreamingState,
addItem,
updateItem,
setShowHelp,
toolRegistry,
setInitError,
],
);

View File

@ -6,17 +6,17 @@
import { describe, it, expect } from 'vitest';
import { renderHook, act } from '@testing-library/react';
import { useHistoryManager } from './useHistoryManager.js';
import { useHistory } from './useHistoryManager.js';
import { HistoryItem } from '../types.js';
describe('useHistoryManager', () => {
it('should initialize with an empty history', () => {
const { result } = renderHook(() => useHistoryManager());
const { result } = renderHook(() => useHistory());
expect(result.current.history).toEqual([]);
});
it('should add an item to history with a unique ID', () => {
const { result } = renderHook(() => useHistoryManager());
const { result } = renderHook(() => useHistory());
const timestamp = Date.now();
const itemData: Omit<HistoryItem, 'id'> = {
type: 'user', // Replaced HistoryItemType.User
@ -24,7 +24,7 @@ describe('useHistoryManager', () => {
};
act(() => {
result.current.addItemToHistory(itemData, timestamp);
result.current.addItem(itemData, timestamp);
});
expect(result.current.history).toHaveLength(1);
@ -39,7 +39,7 @@ describe('useHistoryManager', () => {
});
it('should generate unique IDs for items added with the same base timestamp', () => {
const { result } = renderHook(() => useHistoryManager());
const { result } = renderHook(() => useHistory());
const timestamp = Date.now();
const itemData1: Omit<HistoryItem, 'id'> = {
type: 'user', // Replaced HistoryItemType.User
@ -54,8 +54,8 @@ describe('useHistoryManager', () => {
let id2!: number;
act(() => {
id1 = result.current.addItemToHistory(itemData1, timestamp);
id2 = result.current.addItemToHistory(itemData2, timestamp);
id1 = result.current.addItem(itemData1, timestamp);
id2 = result.current.addItem(itemData2, timestamp);
});
expect(result.current.history).toHaveLength(2);
@ -67,7 +67,7 @@ describe('useHistoryManager', () => {
});
it('should update an existing history item', () => {
const { result } = renderHook(() => useHistoryManager());
const { result } = renderHook(() => useHistory());
const timestamp = Date.now();
const initialItem: Omit<HistoryItem, 'id'> = {
type: 'gemini', // Replaced HistoryItemType.Gemini
@ -76,12 +76,12 @@ describe('useHistoryManager', () => {
let itemId!: number;
act(() => {
itemId = result.current.addItemToHistory(initialItem, timestamp);
itemId = result.current.addItem(initialItem, timestamp);
});
const updatedText = 'Updated content';
act(() => {
result.current.updateHistoryItem(itemId, { text: updatedText });
result.current.updateItem(itemId, { text: updatedText });
});
expect(result.current.history).toHaveLength(1);
@ -93,7 +93,7 @@ describe('useHistoryManager', () => {
});
it('should not change history if updateHistoryItem is called with a non-existent ID', () => {
const { result } = renderHook(() => useHistoryManager());
const { result } = renderHook(() => useHistory());
const timestamp = Date.now();
const itemData: Omit<HistoryItem, 'id'> = {
type: 'user', // Replaced HistoryItemType.User
@ -101,20 +101,20 @@ describe('useHistoryManager', () => {
};
act(() => {
result.current.addItemToHistory(itemData, timestamp);
result.current.addItem(itemData, timestamp);
});
const originalHistory = [...result.current.history]; // Clone before update attempt
act(() => {
result.current.updateHistoryItem(99999, { text: 'Should not apply' }); // Non-existent ID
result.current.updateItem(99999, { text: 'Should not apply' }); // Non-existent ID
});
expect(result.current.history).toEqual(originalHistory);
});
it('should clear the history', () => {
const { result } = renderHook(() => useHistoryManager());
const { result } = renderHook(() => useHistory());
const timestamp = Date.now();
const itemData1: Omit<HistoryItem, 'id'> = {
type: 'user', // Replaced HistoryItemType.User
@ -126,14 +126,14 @@ describe('useHistoryManager', () => {
};
act(() => {
result.current.addItemToHistory(itemData1, timestamp);
result.current.addItemToHistory(itemData2, timestamp);
result.current.addItem(itemData1, timestamp);
result.current.addItem(itemData2, timestamp);
});
expect(result.current.history).toHaveLength(2);
act(() => {
result.current.clearHistory();
result.current.clearItems();
});
expect(result.current.history).toEqual([]);

View File

@ -7,17 +7,19 @@
import { useState, useRef, useCallback } from 'react';
import { HistoryItem } from '../types.js';
// Type for the updater function passed to updateHistoryItem
type HistoryItemUpdater = (
prevItem: HistoryItem,
) => Partial<Omit<HistoryItem, 'id'>>;
export interface UseHistoryManagerReturn {
history: HistoryItem[];
addItemToHistory: (
itemData: Omit<HistoryItem, 'id'>,
baseTimestamp: number,
) => number; // Return the ID of the added item
updateHistoryItem: (
addItem: (itemData: Omit<HistoryItem, 'id'>, baseTimestamp: number) => number; // Returns the generated ID
updateItem: (
id: number,
updates: Partial<Omit<HistoryItem, 'id'>>,
updates: Partial<Omit<HistoryItem, 'id'>> | HistoryItemUpdater,
) => void;
clearHistory: () => void;
clearItems: () => void;
}
/**
@ -26,19 +28,18 @@ export interface UseHistoryManagerReturn {
* Encapsulates the history array, message ID generation, adding items,
* updating items, and clearing the history.
*/
export function useHistoryManager(): UseHistoryManagerReturn {
export function useHistory(): UseHistoryManagerReturn {
const [history, setHistory] = useState<HistoryItem[]>([]);
const messageIdCounterRef = useRef(0);
// Generates a unique message ID based on a timestamp and a counter.
const getNextMessageId = useCallback((baseTimestamp: number): number => {
// Increment *before* adding to ensure uniqueness against the base timestamp
messageIdCounterRef.current += 1;
return baseTimestamp + messageIdCounterRef.current;
}, []);
// Adds a new item to the history state with a unique ID and returns the ID.
const addItemToHistory = useCallback(
// Adds a new item to the history state with a unique ID.
const addItem = useCallback(
(itemData: Omit<HistoryItem, 'id'>, baseTimestamp: number): number => {
const id = getNextMessageId(baseTimestamp);
const newItem: HistoryItem = { ...itemData, id } as HistoryItem;
@ -49,22 +50,36 @@ export function useHistoryManager(): UseHistoryManagerReturn {
);
// Updates an existing history item identified by its ID.
const updateHistoryItem = useCallback(
(id: number, updates: Partial<Omit<HistoryItem, 'id'>>) => {
const updateItem = useCallback(
(
id: number,
updates: Partial<Omit<HistoryItem, 'id'>> | HistoryItemUpdater,
) => {
setHistory((prevHistory) =>
prevHistory.map((item) =>
item.id === id ? ({ ...item, ...updates } as HistoryItem) : item,
),
prevHistory.map((item) => {
if (item.id === id) {
// Apply updates based on whether it's an object or a function
const newUpdates =
typeof updates === 'function' ? updates(item) : updates;
return { ...item, ...newUpdates } as HistoryItem;
}
return item;
}),
);
},
[],
);
// Clears the entire history state.
const clearHistory = useCallback(() => {
// Clears the entire history state and resets the ID counter.
const clearItems = useCallback(() => {
setHistory([]);
messageIdCounterRef.current = 0; // Reset counter when history is cleared
messageIdCounterRef.current = 0;
}, []);
return { history, addItemToHistory, updateHistoryItem, clearHistory };
return {
history,
addItem,
updateItem,
clearItems,
};
}