refactor: break submitQuery into smaller functions (#350)
This commit is contained in:
Brandon Keiji
GitHub
parent
1245fe4885
commit
521708e294
|
|
@ -9,6 +9,12 @@ import { useInput } from 'ink';
|
|||
import {
|
||||
GeminiClient,
|
||||
GeminiEventType as ServerGeminiEventType,
|
||||
ServerGeminiStreamEvent as GeminiEvent,
|
||||
ServerGeminiContentEvent as ContentEvent,
|
||||
ServerGeminiToolCallRequestEvent as ToolCallRequestEvent,
|
||||
ServerGeminiToolCallResponseEvent as ToolCallResponseEvent,
|
||||
ServerGeminiToolCallConfirmationEvent as ToolCallConfirmationEvent,
|
||||
ServerGeminiErrorEvent as ErrorEvent,
|
||||
getErrorMessage,
|
||||
isNodeError,
|
||||
Config,
|
||||
|
|
@ -26,6 +32,7 @@ import {
|
|||
IndividualToolCallDisplay,
|
||||
ToolCallStatus,
|
||||
HistoryItemWithoutId,
|
||||
HistoryItemToolGroup,
|
||||
MessageType,
|
||||
} from '../types.js';
|
||||
import { isAtCommand } from '../utils/commandUtils.js';
|
||||
|
|
@ -35,6 +42,17 @@ import { findLastSafeSplitPoint } from '../utils/markdownUtilities.js';
|
|||
import { useStateAndRef } from './useStateAndRef.js';
|
||||
import { UseHistoryManagerReturn } from './useHistoryManager.js';
|
||||
|
||||
enum StreamProcessingStatus {
|
||||
Completed,
|
||||
PausedForConfirmation,
|
||||
UserCancelled,
|
||||
Error,
|
||||
}
|
||||
|
||||
/**
|
||||
* Hook to manage the Gemini stream, handle user input, process commands,
|
||||
* and interact with the Gemini API and history manager.
|
||||
*/
|
||||
export const useGeminiStream = (
|
||||
addItem: UseHistoryManagerReturn['addItem'],
|
||||
_clearItems: UseHistoryManagerReturn['clearItems'],
|
||||
|
|
@ -82,240 +100,490 @@ export const useGeminiStream = (
|
|||
}
|
||||
});
|
||||
|
||||
const prepareQueryForGemini = async (
|
||||
query: PartListUnion,
|
||||
userMessageTimestamp: number,
|
||||
signal: AbortSignal,
|
||||
): Promise<{ queryToSend: PartListUnion | null; shouldProceed: boolean }> => {
|
||||
if (typeof query === 'string' && query.trim().length === 0) {
|
||||
return { queryToSend: null, shouldProceed: false };
|
||||
}
|
||||
|
||||
let localQueryToSendToGemini: PartListUnion | null = null;
|
||||
|
||||
if (typeof query === 'string') {
|
||||
const trimmedQuery = query.trim();
|
||||
onDebugMessage(`User query: '${trimmedQuery}'`);
|
||||
|
||||
// Handle UI-only commands first
|
||||
if (handleSlashCommand(trimmedQuery)) {
|
||||
return { queryToSend: null, shouldProceed: false };
|
||||
}
|
||||
if (handleShellCommand(trimmedQuery)) {
|
||||
return { queryToSend: null, shouldProceed: false };
|
||||
}
|
||||
|
||||
// Handle @-commands (which might involve tool calls)
|
||||
if (isAtCommand(trimmedQuery)) {
|
||||
const atCommandResult = await handleAtCommand({
|
||||
query: trimmedQuery,
|
||||
config,
|
||||
addItem,
|
||||
onDebugMessage,
|
||||
messageId: userMessageTimestamp,
|
||||
signal,
|
||||
});
|
||||
if (!atCommandResult.shouldProceed) {
|
||||
return { queryToSend: null, shouldProceed: false };
|
||||
}
|
||||
localQueryToSendToGemini = atCommandResult.processedQuery;
|
||||
} else {
|
||||
// Normal query for Gemini
|
||||
addItem(
|
||||
{ type: MessageType.USER, text: trimmedQuery },
|
||||
userMessageTimestamp,
|
||||
);
|
||||
localQueryToSendToGemini = trimmedQuery;
|
||||
}
|
||||
} else {
|
||||
// It's a function response (PartListUnion that isn't a string)
|
||||
localQueryToSendToGemini = query;
|
||||
}
|
||||
|
||||
if (localQueryToSendToGemini === null) {
|
||||
onDebugMessage(
|
||||
'Query processing resulted in null, not sending to Gemini.',
|
||||
);
|
||||
return { queryToSend: null, shouldProceed: false };
|
||||
}
|
||||
return { queryToSend: localQueryToSendToGemini, shouldProceed: true };
|
||||
};
|
||||
|
||||
const ensureChatSession = async (): Promise<{
|
||||
client: GeminiClient | null;
|
||||
chat: Chat | null;
|
||||
}> => {
|
||||
const currentClient = geminiClientRef.current;
|
||||
if (!currentClient) {
|
||||
const errorMsg = 'Gemini client is not available.';
|
||||
setInitError(errorMsg);
|
||||
addItem({ type: MessageType.ERROR, text: errorMsg }, Date.now());
|
||||
return { client: null, chat: null };
|
||||
}
|
||||
|
||||
if (!chatSessionRef.current) {
|
||||
try {
|
||||
chatSessionRef.current = await currentClient.startChat();
|
||||
} catch (err: unknown) {
|
||||
const errorMsg = `Failed to start chat: ${getErrorMessage(err)}`;
|
||||
setInitError(errorMsg);
|
||||
addItem({ type: MessageType.ERROR, text: errorMsg }, Date.now());
|
||||
setStreamingState(StreamingState.Idle);
|
||||
return { client: currentClient, chat: null };
|
||||
}
|
||||
}
|
||||
return { client: currentClient, chat: chatSessionRef.current };
|
||||
};
|
||||
|
||||
// --- UI Helper Functions (used by event handlers) ---
|
||||
const updateFunctionResponseUI = (
|
||||
toolResponse: ToolCallResponseInfo,
|
||||
status: ToolCallStatus,
|
||||
) => {
|
||||
setPendingHistoryItem((item) =>
|
||||
item?.type === 'tool_group'
|
||||
? {
|
||||
...item,
|
||||
tools: item.tools.map((tool) => {
|
||||
if (tool.callId === toolResponse.callId) {
|
||||
return {
|
||||
...tool,
|
||||
status,
|
||||
resultDisplay: toolResponse.resultDisplay,
|
||||
};
|
||||
} else {
|
||||
return tool;
|
||||
}
|
||||
}),
|
||||
}
|
||||
: null,
|
||||
);
|
||||
};
|
||||
|
||||
const updateConfirmingFunctionStatusUI = (
|
||||
callId: string,
|
||||
confirmationDetails: ToolCallConfirmationDetails | undefined,
|
||||
) => {
|
||||
if (pendingHistoryItemRef.current?.type !== 'tool_group') return;
|
||||
setPendingHistoryItem((item) =>
|
||||
item?.type === 'tool_group'
|
||||
? {
|
||||
...item,
|
||||
tools: item.tools.map((tool) =>
|
||||
tool.callId === callId
|
||||
? {
|
||||
...tool,
|
||||
status: ToolCallStatus.Confirming,
|
||||
confirmationDetails,
|
||||
}
|
||||
: tool,
|
||||
),
|
||||
}
|
||||
: null,
|
||||
);
|
||||
};
|
||||
|
||||
// This function will be fully refactored in a later step
|
||||
const wireConfirmationSubmission = (
|
||||
confirmationDetails: ServerToolCallConfirmationDetails,
|
||||
): ToolCallConfirmationDetails => {
|
||||
const originalConfirmationDetails = confirmationDetails.details;
|
||||
const request = confirmationDetails.request;
|
||||
const resubmittingConfirm = async (outcome: ToolConfirmationOutcome) => {
|
||||
originalConfirmationDetails.onConfirm(outcome);
|
||||
if (pendingHistoryItemRef?.current?.type === 'tool_group') {
|
||||
setPendingHistoryItem((item) =>
|
||||
item?.type === 'tool_group'
|
||||
? {
|
||||
...item,
|
||||
tools: item.tools.map((tool) =>
|
||||
tool.callId === request.callId
|
||||
? {
|
||||
...tool,
|
||||
confirmationDetails: undefined,
|
||||
status: ToolCallStatus.Executing,
|
||||
}
|
||||
: tool,
|
||||
),
|
||||
}
|
||||
: item,
|
||||
);
|
||||
refreshStatic();
|
||||
}
|
||||
|
||||
if (outcome === ToolConfirmationOutcome.Cancel) {
|
||||
declineToolExecution(
|
||||
'User rejected function call.',
|
||||
ToolCallStatus.Error,
|
||||
request,
|
||||
originalConfirmationDetails,
|
||||
);
|
||||
} else {
|
||||
const tool = toolRegistry.getTool(request.name);
|
||||
if (!tool) {
|
||||
throw new Error(
|
||||
`Tool "${request.name}" not found or is not registered.`,
|
||||
);
|
||||
}
|
||||
try {
|
||||
abortControllerRef.current = new AbortController();
|
||||
const result = await tool.execute(
|
||||
request.args,
|
||||
abortControllerRef.current.signal,
|
||||
);
|
||||
if (abortControllerRef.current.signal.aborted) {
|
||||
declineToolExecution(
|
||||
result.llmContent,
|
||||
ToolCallStatus.Canceled,
|
||||
request,
|
||||
originalConfirmationDetails,
|
||||
);
|
||||
return;
|
||||
}
|
||||
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);
|
||||
if (pendingHistoryItemRef.current) {
|
||||
addItem(pendingHistoryItemRef.current, Date.now());
|
||||
setPendingHistoryItem(null);
|
||||
}
|
||||
setStreamingState(StreamingState.Idle);
|
||||
await submitQuery(functionResponse); // Recursive call
|
||||
} finally {
|
||||
if (streamingState !== StreamingState.WaitingForConfirmation) {
|
||||
abortControllerRef.current = null;
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// Extracted declineToolExecution to be part of wireConfirmationSubmission's closure
|
||||
// or could be a standalone helper if more params are passed.
|
||||
function declineToolExecution(
|
||||
declineMessage: string,
|
||||
status: ToolCallStatus,
|
||||
request: ServerToolCallConfirmationDetails['request'],
|
||||
originalDetails: ServerToolCallConfirmationDetails['details'],
|
||||
) {
|
||||
let resultDisplay: ToolResultDisplay | undefined;
|
||||
if ('fileDiff' in originalDetails) {
|
||||
resultDisplay = {
|
||||
fileDiff: (originalDetails as ToolEditConfirmationDetails).fileDiff,
|
||||
};
|
||||
} else {
|
||||
resultDisplay = `~~${(originalDetails as ToolExecuteConfirmationDetails).command}~~`;
|
||||
}
|
||||
const functionResponse: Part = {
|
||||
functionResponse: {
|
||||
id: request.callId,
|
||||
name: request.name,
|
||||
response: { error: declineMessage },
|
||||
},
|
||||
};
|
||||
const responseInfo: ToolCallResponseInfo = {
|
||||
callId: request.callId,
|
||||
responsePart: functionResponse,
|
||||
resultDisplay,
|
||||
error: new Error(declineMessage),
|
||||
};
|
||||
const history = chatSessionRef.current?.getHistory();
|
||||
if (history) {
|
||||
history.push({ role: 'model', parts: [functionResponse] });
|
||||
}
|
||||
updateFunctionResponseUI(responseInfo, status);
|
||||
if (pendingHistoryItemRef.current) {
|
||||
addItem(pendingHistoryItemRef.current, Date.now());
|
||||
setPendingHistoryItem(null);
|
||||
}
|
||||
setStreamingState(StreamingState.Idle);
|
||||
}
|
||||
|
||||
return { ...originalConfirmationDetails, onConfirm: resubmittingConfirm };
|
||||
};
|
||||
|
||||
// --- Stream Event Handlers ---
|
||||
const handleContentEvent = (
|
||||
eventValue: ContentEvent['value'],
|
||||
currentGeminiMessageBuffer: string,
|
||||
userMessageTimestamp: number,
|
||||
): string => {
|
||||
let newGeminiMessageBuffer = currentGeminiMessageBuffer + eventValue;
|
||||
if (
|
||||
pendingHistoryItemRef.current?.type !== 'gemini' &&
|
||||
pendingHistoryItemRef.current?.type !== 'gemini_content'
|
||||
) {
|
||||
if (pendingHistoryItemRef.current) {
|
||||
addItem(pendingHistoryItemRef.current, userMessageTimestamp);
|
||||
}
|
||||
setPendingHistoryItem({ type: 'gemini', text: '' });
|
||||
newGeminiMessageBuffer = eventValue;
|
||||
}
|
||||
// Split large messages for better rendering performance. Ideally,
|
||||
// we should maximize the amount of output sent to <Static />.
|
||||
const splitPoint = findLastSafeSplitPoint(newGeminiMessageBuffer);
|
||||
if (splitPoint === newGeminiMessageBuffer.length) {
|
||||
// Update the existing message with accumulated content
|
||||
setPendingHistoryItem((item) => ({
|
||||
type: item?.type as 'gemini' | 'gemini_content',
|
||||
text: newGeminiMessageBuffer,
|
||||
}));
|
||||
} else {
|
||||
// This indicates that we need to split up this Gemini Message.
|
||||
// Splitting a message is primarily a performance consideration. There is a
|
||||
// <Static> component at the root of App.tsx which takes care of rendering
|
||||
// content statically or dynamically. Everything but the last message is
|
||||
// treated as static in order to prevent re-rendering an entire message history
|
||||
// multiple times per-second (as streaming occurs). Prior to this change you'd
|
||||
// see heavy flickering of the terminal. This ensures that larger messages get
|
||||
// broken up so that there are more "statically" rendered.
|
||||
const beforeText = newGeminiMessageBuffer.substring(0, splitPoint);
|
||||
const afterText = newGeminiMessageBuffer.substring(splitPoint);
|
||||
addItem(
|
||||
{
|
||||
type: pendingHistoryItemRef.current?.type as
|
||||
| 'gemini'
|
||||
| 'gemini_content',
|
||||
text: beforeText,
|
||||
},
|
||||
userMessageTimestamp,
|
||||
);
|
||||
setPendingHistoryItem({ type: 'gemini_content', text: afterText });
|
||||
newGeminiMessageBuffer = afterText;
|
||||
}
|
||||
return newGeminiMessageBuffer;
|
||||
};
|
||||
|
||||
const handleToolCallRequestEvent = (
|
||||
eventValue: ToolCallRequestEvent['value'],
|
||||
userMessageTimestamp: number,
|
||||
) => {
|
||||
const { callId, name, args } = eventValue;
|
||||
const cliTool = toolRegistry.getTool(name);
|
||||
if (!cliTool) {
|
||||
console.error(`CLI Tool "${name}" not found!`);
|
||||
return; // Skip this event if tool is not found
|
||||
}
|
||||
if (pendingHistoryItemRef.current?.type !== 'tool_group') {
|
||||
if (pendingHistoryItemRef.current) {
|
||||
addItem(pendingHistoryItemRef.current, userMessageTimestamp);
|
||||
}
|
||||
setPendingHistoryItem({ type: 'tool_group', tools: [] });
|
||||
}
|
||||
let description: string;
|
||||
try {
|
||||
description = cliTool.getDescription(args);
|
||||
} catch (e) {
|
||||
description = `Error: Unable to get description: ${getErrorMessage(e)}`;
|
||||
}
|
||||
const toolCallDisplay: IndividualToolCallDisplay = {
|
||||
callId,
|
||||
name: cliTool.displayName,
|
||||
description,
|
||||
status: ToolCallStatus.Pending,
|
||||
resultDisplay: undefined,
|
||||
confirmationDetails: undefined,
|
||||
};
|
||||
setPendingHistoryItem((pending) =>
|
||||
pending?.type === 'tool_group'
|
||||
? { ...pending, tools: [...pending.tools, toolCallDisplay] }
|
||||
: null,
|
||||
);
|
||||
};
|
||||
|
||||
const handleToolCallResponseEvent = (
|
||||
eventValue: ToolCallResponseEvent['value'],
|
||||
) => {
|
||||
const status = eventValue.error
|
||||
? ToolCallStatus.Error
|
||||
: ToolCallStatus.Success;
|
||||
updateFunctionResponseUI(eventValue, status);
|
||||
};
|
||||
|
||||
const handleToolCallConfirmationEvent = (
|
||||
eventValue: ToolCallConfirmationEvent['value'],
|
||||
) => {
|
||||
const confirmationDetails = wireConfirmationSubmission(eventValue);
|
||||
updateConfirmingFunctionStatusUI(
|
||||
eventValue.request.callId,
|
||||
confirmationDetails,
|
||||
);
|
||||
setStreamingState(StreamingState.WaitingForConfirmation);
|
||||
};
|
||||
|
||||
const handleUserCancelledEvent = (userMessageTimestamp: number) => {
|
||||
if (pendingHistoryItemRef.current) {
|
||||
if (pendingHistoryItemRef.current.type === 'tool_group') {
|
||||
const updatedTools = pendingHistoryItemRef.current.tools.map((tool) =>
|
||||
tool.status === ToolCallStatus.Pending ||
|
||||
tool.status === ToolCallStatus.Confirming ||
|
||||
tool.status === ToolCallStatus.Executing
|
||||
? { ...tool, status: ToolCallStatus.Canceled }
|
||||
: tool,
|
||||
);
|
||||
const pendingItem: HistoryItemToolGroup = {
|
||||
...pendingHistoryItemRef.current,
|
||||
tools: updatedTools,
|
||||
};
|
||||
addItem(pendingItem, userMessageTimestamp);
|
||||
} else {
|
||||
addItem(pendingHistoryItemRef.current, userMessageTimestamp);
|
||||
}
|
||||
setPendingHistoryItem(null);
|
||||
}
|
||||
addItem(
|
||||
{ type: MessageType.INFO, text: 'User cancelled the request.' },
|
||||
userMessageTimestamp,
|
||||
);
|
||||
setStreamingState(StreamingState.Idle);
|
||||
};
|
||||
|
||||
const handleErrorEvent = (
|
||||
eventValue: ErrorEvent['value'],
|
||||
userMessageTimestamp: number,
|
||||
) => {
|
||||
if (pendingHistoryItemRef.current) {
|
||||
addItem(pendingHistoryItemRef.current, userMessageTimestamp);
|
||||
setPendingHistoryItem(null);
|
||||
}
|
||||
addItem(
|
||||
{ type: MessageType.ERROR, text: `[API Error: ${eventValue.message}]` },
|
||||
userMessageTimestamp,
|
||||
);
|
||||
};
|
||||
|
||||
const processGeminiStreamEvents = async (
|
||||
stream: AsyncIterable<GeminiEvent>,
|
||||
userMessageTimestamp: number,
|
||||
): Promise<StreamProcessingStatus> => {
|
||||
let geminiMessageBuffer = '';
|
||||
|
||||
for await (const event of stream) {
|
||||
if (event.type === ServerGeminiEventType.Content) {
|
||||
geminiMessageBuffer = handleContentEvent(
|
||||
event.value,
|
||||
geminiMessageBuffer,
|
||||
userMessageTimestamp,
|
||||
);
|
||||
} else if (event.type === ServerGeminiEventType.ToolCallRequest) {
|
||||
handleToolCallRequestEvent(event.value, userMessageTimestamp);
|
||||
} else if (event.type === ServerGeminiEventType.ToolCallResponse) {
|
||||
handleToolCallResponseEvent(event.value);
|
||||
} else if (event.type === ServerGeminiEventType.ToolCallConfirmation) {
|
||||
handleToolCallConfirmationEvent(event.value);
|
||||
return StreamProcessingStatus.PausedForConfirmation; // Explicit return as this pauses the stream
|
||||
} else if (event.type === ServerGeminiEventType.UserCancelled) {
|
||||
handleUserCancelledEvent(userMessageTimestamp);
|
||||
return StreamProcessingStatus.UserCancelled;
|
||||
} else if (event.type === ServerGeminiEventType.Error) {
|
||||
handleErrorEvent(event.value, userMessageTimestamp);
|
||||
return StreamProcessingStatus.Error;
|
||||
}
|
||||
}
|
||||
return StreamProcessingStatus.Completed;
|
||||
};
|
||||
|
||||
const submitQuery = useCallback(
|
||||
async (query: PartListUnion) => {
|
||||
if (streamingState === StreamingState.Responding) return;
|
||||
if (typeof query === 'string' && query.trim().length === 0) return;
|
||||
|
||||
const userMessageTimestamp = Date.now();
|
||||
let queryToSendToGemini: PartListUnion | null = null;
|
||||
|
||||
setShowHelp(false);
|
||||
|
||||
abortControllerRef.current ??= new AbortController();
|
||||
const signal = abortControllerRef.current.signal;
|
||||
|
||||
if (typeof query === 'string') {
|
||||
const trimmedQuery = query.trim();
|
||||
onDebugMessage(`User query: '${trimmedQuery}'`);
|
||||
const { queryToSend, shouldProceed } = await prepareQueryForGemini(
|
||||
query,
|
||||
userMessageTimestamp,
|
||||
signal,
|
||||
);
|
||||
|
||||
if (handleSlashCommand(trimmedQuery)) return;
|
||||
if (handleShellCommand(trimmedQuery)) return;
|
||||
|
||||
if (isAtCommand(trimmedQuery)) {
|
||||
const atCommandResult = await handleAtCommand({
|
||||
query: trimmedQuery,
|
||||
config,
|
||||
addItem,
|
||||
onDebugMessage,
|
||||
messageId: userMessageTimestamp,
|
||||
signal,
|
||||
});
|
||||
if (!atCommandResult.shouldProceed) return;
|
||||
queryToSendToGemini = atCommandResult.processedQuery;
|
||||
} else {
|
||||
addItem(
|
||||
{ type: MessageType.USER, text: trimmedQuery },
|
||||
userMessageTimestamp,
|
||||
);
|
||||
queryToSendToGemini = trimmedQuery;
|
||||
}
|
||||
} else {
|
||||
queryToSendToGemini = query;
|
||||
}
|
||||
|
||||
if (queryToSendToGemini === null) {
|
||||
onDebugMessage(
|
||||
'Query processing resulted in null, not sending to Gemini.',
|
||||
);
|
||||
if (!shouldProceed || queryToSend === null) {
|
||||
return;
|
||||
}
|
||||
|
||||
const client = geminiClientRef.current;
|
||||
if (!client) {
|
||||
const errorMsg = 'Gemini client is not available.';
|
||||
setInitError(errorMsg);
|
||||
addItem({ type: MessageType.ERROR, text: errorMsg }, Date.now());
|
||||
return;
|
||||
}
|
||||
const { client, chat } = await ensureChatSession();
|
||||
|
||||
if (!chatSessionRef.current) {
|
||||
try {
|
||||
chatSessionRef.current = await client.startChat();
|
||||
} catch (err: unknown) {
|
||||
const errorMsg = `Failed to start chat: ${getErrorMessage(err)}`;
|
||||
setInitError(errorMsg);
|
||||
addItem({ type: MessageType.ERROR, text: errorMsg }, Date.now());
|
||||
setStreamingState(StreamingState.Idle);
|
||||
return;
|
||||
}
|
||||
if (!client || !chat) {
|
||||
return;
|
||||
}
|
||||
|
||||
setStreamingState(StreamingState.Responding);
|
||||
setInitError(null);
|
||||
const chat = chatSessionRef.current;
|
||||
|
||||
try {
|
||||
const stream = client.sendMessageStream(
|
||||
chat,
|
||||
queryToSendToGemini,
|
||||
signal,
|
||||
const stream = client.sendMessageStream(chat, queryToSend, signal);
|
||||
const processingStatus = await processGeminiStreamEvents(
|
||||
stream,
|
||||
userMessageTimestamp,
|
||||
);
|
||||
|
||||
let geminiMessageBuffer = '';
|
||||
|
||||
for await (const event of stream) {
|
||||
if (event.type === ServerGeminiEventType.Content) {
|
||||
if (
|
||||
pendingHistoryItemRef.current?.type !== 'gemini' &&
|
||||
pendingHistoryItemRef.current?.type !== 'gemini_content'
|
||||
) {
|
||||
if (pendingHistoryItemRef.current) {
|
||||
addItem(pendingHistoryItemRef.current, userMessageTimestamp);
|
||||
}
|
||||
setPendingHistoryItem({
|
||||
type: 'gemini',
|
||||
text: '',
|
||||
});
|
||||
geminiMessageBuffer = '';
|
||||
}
|
||||
|
||||
geminiMessageBuffer += event.value;
|
||||
|
||||
// Split large messages for better rendering performance. Ideally,
|
||||
// we should maximize the amount of output sent to <Static />.
|
||||
const splitPoint = findLastSafeSplitPoint(geminiMessageBuffer);
|
||||
if (splitPoint === geminiMessageBuffer.length) {
|
||||
// Update the existing message with accumulated content
|
||||
setPendingHistoryItem((item) => ({
|
||||
type: item?.type as 'gemini' | 'gemini_content',
|
||||
text: geminiMessageBuffer,
|
||||
}));
|
||||
} else {
|
||||
// This indicates that we need to split up this Gemini Message.
|
||||
// Splitting a message is primarily a performance consideration. There is a
|
||||
// <Static> component at the root of App.tsx which takes care of rendering
|
||||
// content statically or dynamically. Everything but the last message is
|
||||
// treated as static in order to prevent re-rendering an entire message history
|
||||
// multiple times per-second (as streaming occurs). Prior to this change you'd
|
||||
// see heavy flickering of the terminal. This ensures that larger messages get
|
||||
// broken up so that there are more "statically" rendered.
|
||||
const beforeText = geminiMessageBuffer.substring(0, splitPoint);
|
||||
const afterText = geminiMessageBuffer.substring(splitPoint);
|
||||
geminiMessageBuffer = afterText;
|
||||
addItem(
|
||||
{
|
||||
type: pendingHistoryItemRef.current?.type as
|
||||
| 'gemini'
|
||||
| 'gemini_content',
|
||||
text: beforeText,
|
||||
},
|
||||
userMessageTimestamp,
|
||||
);
|
||||
setPendingHistoryItem({
|
||||
type: 'gemini_content',
|
||||
text: afterText,
|
||||
});
|
||||
}
|
||||
} else if (event.type === ServerGeminiEventType.ToolCallRequest) {
|
||||
const { callId, name, args } = event.value;
|
||||
const cliTool = toolRegistry.getTool(name);
|
||||
if (!cliTool) {
|
||||
console.error(`CLI Tool "${name}" not found!`);
|
||||
continue;
|
||||
}
|
||||
|
||||
if (pendingHistoryItemRef.current?.type !== 'tool_group') {
|
||||
if (pendingHistoryItemRef.current) {
|
||||
addItem(pendingHistoryItemRef.current, userMessageTimestamp);
|
||||
}
|
||||
setPendingHistoryItem({
|
||||
type: 'tool_group',
|
||||
tools: [],
|
||||
});
|
||||
}
|
||||
|
||||
let description: string;
|
||||
try {
|
||||
description = cliTool.getDescription(args);
|
||||
} catch (e) {
|
||||
description = `Error: Unable to get description: ${getErrorMessage(e)}`;
|
||||
}
|
||||
|
||||
const toolCallDisplay: IndividualToolCallDisplay = {
|
||||
callId,
|
||||
name: cliTool.displayName,
|
||||
description,
|
||||
status: ToolCallStatus.Pending,
|
||||
resultDisplay: undefined,
|
||||
confirmationDetails: undefined,
|
||||
};
|
||||
|
||||
setPendingHistoryItem((pending) =>
|
||||
pending?.type === 'tool_group'
|
||||
? {
|
||||
...pending,
|
||||
tools: [...pending.tools, toolCallDisplay],
|
||||
}
|
||||
: null,
|
||||
);
|
||||
} else if (event.type === ServerGeminiEventType.ToolCallResponse) {
|
||||
const status = event.value.error
|
||||
? ToolCallStatus.Error
|
||||
: ToolCallStatus.Success;
|
||||
updateFunctionResponseUI(event.value, status);
|
||||
} else if (
|
||||
event.type === ServerGeminiEventType.ToolCallConfirmation
|
||||
) {
|
||||
const confirmationDetails = wireConfirmationSubmission(event.value);
|
||||
updateConfirmingFunctionStatusUI(
|
||||
event.value.request.callId,
|
||||
confirmationDetails,
|
||||
);
|
||||
setStreamingState(StreamingState.WaitingForConfirmation);
|
||||
return;
|
||||
} else if (event.type === ServerGeminiEventType.UserCancelled) {
|
||||
if (pendingHistoryItemRef.current) {
|
||||
if (pendingHistoryItemRef.current.type === 'tool_group') {
|
||||
const updatedTools = pendingHistoryItemRef.current.tools.map(
|
||||
(tool) => {
|
||||
if (
|
||||
tool.status === ToolCallStatus.Pending ||
|
||||
tool.status === ToolCallStatus.Confirming ||
|
||||
tool.status === ToolCallStatus.Executing
|
||||
) {
|
||||
return { ...tool, status: ToolCallStatus.Canceled };
|
||||
}
|
||||
return tool;
|
||||
},
|
||||
);
|
||||
const pendingHistoryItem = pendingHistoryItemRef.current;
|
||||
pendingHistoryItem.tools = updatedTools;
|
||||
addItem(pendingHistoryItem, userMessageTimestamp);
|
||||
} else {
|
||||
addItem(pendingHistoryItemRef.current, userMessageTimestamp);
|
||||
}
|
||||
setPendingHistoryItem(null);
|
||||
}
|
||||
addItem(
|
||||
{ type: MessageType.INFO, text: 'User cancelled the request.' },
|
||||
userMessageTimestamp,
|
||||
);
|
||||
setStreamingState(StreamingState.Idle);
|
||||
return;
|
||||
} else if (event.type === ServerGeminiEventType.Error) {
|
||||
if (pendingHistoryItemRef.current) {
|
||||
addItem(pendingHistoryItemRef.current, userMessageTimestamp);
|
||||
setPendingHistoryItem(null);
|
||||
}
|
||||
addItem(
|
||||
{
|
||||
type: MessageType.ERROR,
|
||||
text: `[API Error: ${event.value.message}]`,
|
||||
},
|
||||
userMessageTimestamp,
|
||||
);
|
||||
}
|
||||
if (
|
||||
processingStatus === StreamProcessingStatus.PausedForConfirmation ||
|
||||
processingStatus === StreamProcessingStatus.UserCancelled
|
||||
) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (pendingHistoryItemRef.current) {
|
||||
|
|
@ -323,7 +591,12 @@ export const useGeminiStream = (
|
|||
setPendingHistoryItem(null);
|
||||
}
|
||||
|
||||
setStreamingState(StreamingState.Idle);
|
||||
if (
|
||||
processingStatus === StreamProcessingStatus.Completed ||
|
||||
processingStatus === StreamProcessingStatus.Error
|
||||
) {
|
||||
setStreamingState(StreamingState.Idle);
|
||||
}
|
||||
} catch (error: unknown) {
|
||||
if (!isNodeError(error) || error.name !== 'AbortError') {
|
||||
addItem(
|
||||
|
|
@ -336,191 +609,12 @@ export const useGeminiStream = (
|
|||
}
|
||||
setStreamingState(StreamingState.Idle);
|
||||
} finally {
|
||||
abortControllerRef.current = null;
|
||||
}
|
||||
|
||||
function updateConfirmingFunctionStatusUI(
|
||||
callId: string,
|
||||
confirmationDetails: ToolCallConfirmationDetails | undefined,
|
||||
) {
|
||||
if (pendingHistoryItemRef.current?.type !== 'tool_group') return;
|
||||
setPendingHistoryItem((item) =>
|
||||
item?.type === 'tool_group'
|
||||
? {
|
||||
...item,
|
||||
tools: item.tools.map((tool) =>
|
||||
tool.callId === callId
|
||||
? {
|
||||
...tool,
|
||||
status: ToolCallStatus.Confirming,
|
||||
confirmationDetails,
|
||||
}
|
||||
: tool,
|
||||
),
|
||||
}
|
||||
: null,
|
||||
);
|
||||
}
|
||||
|
||||
function updateFunctionResponseUI(
|
||||
toolResponse: ToolCallResponseInfo,
|
||||
status: ToolCallStatus,
|
||||
) {
|
||||
setPendingHistoryItem((item) =>
|
||||
item?.type === 'tool_group'
|
||||
? {
|
||||
...item,
|
||||
tools: item.tools.map((tool) => {
|
||||
if (tool.callId === toolResponse.callId) {
|
||||
return {
|
||||
...tool,
|
||||
status,
|
||||
resultDisplay: toolResponse.resultDisplay,
|
||||
};
|
||||
} else {
|
||||
return tool;
|
||||
}
|
||||
}),
|
||||
}
|
||||
: null,
|
||||
);
|
||||
}
|
||||
|
||||
function wireConfirmationSubmission(
|
||||
confirmationDetails: ServerToolCallConfirmationDetails,
|
||||
): ToolCallConfirmationDetails {
|
||||
const originalConfirmationDetails = confirmationDetails.details;
|
||||
const request = confirmationDetails.request;
|
||||
const resubmittingConfirm = async (
|
||||
outcome: ToolConfirmationOutcome,
|
||||
) => {
|
||||
originalConfirmationDetails.onConfirm(outcome);
|
||||
|
||||
if (pendingHistoryItemRef?.current?.type === 'tool_group') {
|
||||
setPendingHistoryItem((item) =>
|
||||
item?.type === 'tool_group'
|
||||
? {
|
||||
...item,
|
||||
tools: item.tools.map((tool) =>
|
||||
tool.callId === request.callId
|
||||
? {
|
||||
...tool,
|
||||
confirmationDetails: undefined,
|
||||
status: ToolCallStatus.Executing,
|
||||
}
|
||||
: tool,
|
||||
),
|
||||
}
|
||||
: item,
|
||||
);
|
||||
refreshStatic();
|
||||
}
|
||||
|
||||
if (outcome === ToolConfirmationOutcome.Cancel) {
|
||||
declineToolExecution(
|
||||
'User rejected function call.',
|
||||
ToolCallStatus.Error,
|
||||
);
|
||||
} else {
|
||||
const tool = toolRegistry.getTool(request.name);
|
||||
if (!tool) {
|
||||
throw new Error(
|
||||
`Tool "${request.name}" not found or is not registered.`,
|
||||
);
|
||||
}
|
||||
|
||||
try {
|
||||
abortControllerRef.current = new AbortController();
|
||||
const result = await tool.execute(
|
||||
request.args,
|
||||
abortControllerRef.current.signal,
|
||||
);
|
||||
|
||||
if (abortControllerRef.current.signal.aborted) {
|
||||
declineToolExecution(
|
||||
result.llmContent,
|
||||
ToolCallStatus.Canceled,
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
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);
|
||||
if (pendingHistoryItemRef.current) {
|
||||
addItem(pendingHistoryItemRef.current, Date.now());
|
||||
setPendingHistoryItem(null);
|
||||
}
|
||||
setStreamingState(StreamingState.Idle);
|
||||
await submitQuery(functionResponse);
|
||||
} finally {
|
||||
abortControllerRef.current = null;
|
||||
}
|
||||
}
|
||||
|
||||
function declineToolExecution(
|
||||
declineMessage: string,
|
||||
status: ToolCallStatus,
|
||||
) {
|
||||
let resultDisplay: ToolResultDisplay | undefined;
|
||||
if ('fileDiff' in originalConfirmationDetails) {
|
||||
resultDisplay = {
|
||||
fileDiff: (
|
||||
originalConfirmationDetails as ToolEditConfirmationDetails
|
||||
).fileDiff,
|
||||
};
|
||||
} else {
|
||||
resultDisplay = `~~${(originalConfirmationDetails as ToolExecuteConfirmationDetails).command}~~`;
|
||||
}
|
||||
const functionResponse: Part = {
|
||||
functionResponse: {
|
||||
id: request.callId,
|
||||
name: request.name,
|
||||
response: { error: declineMessage },
|
||||
},
|
||||
};
|
||||
const responseInfo: ToolCallResponseInfo = {
|
||||
callId: request.callId,
|
||||
responsePart: functionResponse,
|
||||
resultDisplay,
|
||||
error: new Error(declineMessage),
|
||||
};
|
||||
|
||||
const history = chatSessionRef.current?.getHistory();
|
||||
if (history) {
|
||||
history.push({
|
||||
role: 'model',
|
||||
parts: [functionResponse],
|
||||
});
|
||||
}
|
||||
|
||||
updateFunctionResponseUI(responseInfo, status);
|
||||
if (pendingHistoryItemRef.current) {
|
||||
addItem(pendingHistoryItemRef.current, Date.now());
|
||||
setPendingHistoryItem(null);
|
||||
}
|
||||
setStreamingState(StreamingState.Idle);
|
||||
}
|
||||
};
|
||||
|
||||
return {
|
||||
...originalConfirmationDetails,
|
||||
onConfirm: resubmittingConfirm,
|
||||
};
|
||||
if (streamingState !== StreamingState.WaitingForConfirmation) {
|
||||
abortControllerRef.current = null;
|
||||
}
|
||||
}
|
||||
},
|
||||
// eslint-disable-next-line react-hooks/exhaustive-deps
|
||||
[
|
||||
streamingState,
|
||||
setShowHelp,
|
||||
|
|
@ -528,11 +622,10 @@ export const useGeminiStream = (
|
|||
handleShellCommand,
|
||||
config,
|
||||
addItem,
|
||||
pendingHistoryItemRef,
|
||||
setPendingHistoryItem,
|
||||
toolRegistry,
|
||||
refreshStatic,
|
||||
onDebugMessage,
|
||||
refreshStatic,
|
||||
setInitError,
|
||||
setStreamingState,
|
||||
],
|
||||
);
|
||||
|
||||
|
|
|
|||
|
|
@ -55,18 +55,47 @@ export interface HistoryItemBase {
|
|||
text?: string; // Text content for user/gemini/info/error messages
|
||||
}
|
||||
|
||||
export type HistoryItemUser = HistoryItemBase & {
|
||||
type: 'user';
|
||||
text: string;
|
||||
};
|
||||
|
||||
export type HistoryItemGemini = HistoryItemBase & {
|
||||
type: 'gemini';
|
||||
text: string;
|
||||
};
|
||||
|
||||
export type HistoryItemGeminiContent = HistoryItemBase & {
|
||||
type: 'gemini_content';
|
||||
text: string;
|
||||
};
|
||||
|
||||
export type HistoryItemInfo = HistoryItemBase & {
|
||||
type: 'info';
|
||||
text: string;
|
||||
};
|
||||
|
||||
export type HistoryItemError = HistoryItemBase & {
|
||||
type: 'error';
|
||||
text: string;
|
||||
};
|
||||
|
||||
export type HistoryItemToolGroup = HistoryItemBase & {
|
||||
type: 'tool_group';
|
||||
tools: IndividualToolCallDisplay[];
|
||||
};
|
||||
|
||||
// Using Omit<HistoryItem, 'id'> seems to have some issues with typescript's
|
||||
// type inference e.g. historyItem.type === 'tool_group' isn't auto-inferring that
|
||||
// 'tools' in historyItem.
|
||||
export type HistoryItemWithoutId = HistoryItemBase &
|
||||
(
|
||||
| { type: 'user'; text: string }
|
||||
| { type: 'gemini'; text: string }
|
||||
| { type: 'gemini_content'; text: string }
|
||||
| { type: 'info'; text: string }
|
||||
| { type: 'error'; text: string }
|
||||
| { type: 'tool_group'; tools: IndividualToolCallDisplay[] }
|
||||
);
|
||||
// Individually exported types extending HistoryItemBase
|
||||
export type HistoryItemWithoutId =
|
||||
| HistoryItemUser
|
||||
| HistoryItemGemini
|
||||
| HistoryItemGeminiContent
|
||||
| HistoryItemInfo
|
||||
| HistoryItemError
|
||||
| HistoryItemToolGroup;
|
||||
|
||||
export type HistoryItem = HistoryItemWithoutId & { id: number };
|
||||
|
||||
|
|
|
|||
|
|
@ -78,16 +78,43 @@ export interface ServerToolCallConfirmationDetails {
|
|||
details: ToolCallConfirmationDetails;
|
||||
}
|
||||
|
||||
export type ServerGeminiContentEvent = {
|
||||
type: GeminiEventType.Content;
|
||||
value: string;
|
||||
};
|
||||
|
||||
export type ServerGeminiToolCallRequestEvent = {
|
||||
type: GeminiEventType.ToolCallRequest;
|
||||
value: ToolCallRequestInfo;
|
||||
};
|
||||
|
||||
export type ServerGeminiToolCallResponseEvent = {
|
||||
type: GeminiEventType.ToolCallResponse;
|
||||
value: ToolCallResponseInfo;
|
||||
};
|
||||
|
||||
export type ServerGeminiToolCallConfirmationEvent = {
|
||||
type: GeminiEventType.ToolCallConfirmation;
|
||||
value: ServerToolCallConfirmationDetails;
|
||||
};
|
||||
|
||||
export type ServerGeminiUserCancelledEvent = {
|
||||
type: GeminiEventType.UserCancelled;
|
||||
};
|
||||
|
||||
export type ServerGeminiErrorEvent = {
|
||||
type: GeminiEventType.Error;
|
||||
value: GeminiErrorEventValue;
|
||||
};
|
||||
|
||||
// The original union type, now composed of the individual types
|
||||
export type ServerGeminiStreamEvent =
|
||||
| { type: GeminiEventType.Content; value: string }
|
||||
| { type: GeminiEventType.ToolCallRequest; value: ToolCallRequestInfo }
|
||||
| { type: GeminiEventType.ToolCallResponse; value: ToolCallResponseInfo }
|
||||
| {
|
||||
type: GeminiEventType.ToolCallConfirmation;
|
||||
value: ServerToolCallConfirmationDetails;
|
||||
}
|
||||
| { type: GeminiEventType.UserCancelled }
|
||||
| { type: GeminiEventType.Error; value: GeminiErrorEventValue };
|
||||
| ServerGeminiContentEvent
|
||||
| ServerGeminiToolCallRequestEvent
|
||||
| ServerGeminiToolCallResponseEvent
|
||||
| ServerGeminiToolCallConfirmationEvent
|
||||
| ServerGeminiUserCancelledEvent
|
||||
| ServerGeminiErrorEvent;
|
||||
|
||||
// A turn manages the agentic loop turn within the server context.
|
||||
export class Turn {
|
||||
|
|
|
|||
Loading…
Reference in New Issue