Adding TurnId to Tool call and API responses and error logs. (#3039)

Co-authored-by: Scott Densmore <scottdensmore@mac.com>
This commit is contained in:
uttamkanodia14
2025-07-10 00:19:30 +05:30
committed by GitHub
parent 6c12f9e0d9
commit 063481faa4
26 changed files with 289 additions and 72 deletions

View File

@@ -159,7 +159,7 @@ describe('useSlashCommandProcessor', () => {
stats: {
sessionStartTime: new Date('2025-01-01T00:00:00.000Z'),
cumulative: {
turnCount: 0,
promptCount: 0,
promptTokenCount: 0,
candidatesTokenCount: 0,
totalTokenCount: 0,
@@ -1311,7 +1311,10 @@ describe('useSlashCommandProcessor', () => {
hook.rerender();
});
expect(hook.result.current.pendingHistoryItems).toEqual([]);
expect(mockGeminiClient.tryCompressChat).toHaveBeenCalledWith(true);
expect(mockGeminiClient.tryCompressChat).toHaveBeenCalledWith(
'Prompt Id not set',
true,
);
expect(mockAddItem).toHaveBeenNthCalledWith(
2,
expect.objectContaining({

View File

@@ -880,7 +880,8 @@ export const useSlashCommandProcessor = (
try {
const compressed = await config!
.getGeminiClient()!
.tryCompressChat(true);
// TODO: Set Prompt id for CompressChat from SlashCommandProcessor.
.tryCompressChat('Prompt Id not set', true);
if (compressed) {
addMessage({
type: MessageType.COMPRESSION,

View File

@@ -109,12 +109,13 @@ vi.mock('./useLogger.js', () => ({
}),
}));
const mockStartNewTurn = vi.fn();
const mockStartNewPrompt = vi.fn();
const mockAddUsage = vi.fn();
vi.mock('../contexts/SessionContext.js', () => ({
useSessionStats: vi.fn(() => ({
startNewTurn: mockStartNewTurn,
startNewPrompt: mockStartNewPrompt,
addUsage: mockAddUsage,
getPromptCount: vi.fn(() => 5),
})),
}));
@@ -301,6 +302,9 @@ describe('useGeminiStream', () => {
getUsageStatisticsEnabled: () => true,
getDebugMode: () => false,
addHistory: vi.fn(),
getSessionId() {
return 'test-session-id';
},
setQuotaErrorOccurred: vi.fn(),
getQuotaErrorOccurred: vi.fn(() => false),
} as unknown as Config;
@@ -426,6 +430,7 @@ describe('useGeminiStream', () => {
name: 'tool1',
args: {},
isClientInitiated: false,
prompt_id: 'prompt-id-1',
},
status: 'success',
responseSubmittedToGemini: false,
@@ -444,7 +449,12 @@ describe('useGeminiStream', () => {
endTime: Date.now(),
} as TrackedCompletedToolCall,
{
request: { callId: 'call2', name: 'tool2', args: {} },
request: {
callId: 'call2',
name: 'tool2',
args: {},
prompt_id: 'prompt-id-1',
},
status: 'executing',
responseSubmittedToGemini: false,
tool: {
@@ -481,6 +491,7 @@ describe('useGeminiStream', () => {
name: 'tool1',
args: {},
isClientInitiated: false,
prompt_id: 'prompt-id-2',
},
status: 'success',
responseSubmittedToGemini: false,
@@ -492,6 +503,7 @@ describe('useGeminiStream', () => {
name: 'tool2',
args: {},
isClientInitiated: false,
prompt_id: 'prompt-id-2',
},
status: 'error',
responseSubmittedToGemini: false,
@@ -546,6 +558,7 @@ describe('useGeminiStream', () => {
expect(mockSendMessageStream).toHaveBeenCalledWith(
expectedMergedResponse,
expect.any(AbortSignal),
'prompt-id-2',
);
});
@@ -557,6 +570,7 @@ describe('useGeminiStream', () => {
name: 'testTool',
args: {},
isClientInitiated: false,
prompt_id: 'prompt-id-3',
},
status: 'cancelled',
response: { callId: '1', responseParts: [{ text: 'cancelled' }] },
@@ -618,6 +632,7 @@ describe('useGeminiStream', () => {
name: 'toolA',
args: {},
isClientInitiated: false,
prompt_id: 'prompt-id-7',
},
tool: {
name: 'toolA',
@@ -641,6 +656,7 @@ describe('useGeminiStream', () => {
name: 'toolB',
args: {},
isClientInitiated: false,
prompt_id: 'prompt-id-8',
},
tool: {
name: 'toolB',
@@ -731,6 +747,7 @@ describe('useGeminiStream', () => {
name: 'tool1',
args: {},
isClientInitiated: false,
prompt_id: 'prompt-id-4',
},
status: 'executing',
responseSubmittedToGemini: false,
@@ -824,6 +841,7 @@ describe('useGeminiStream', () => {
expect(mockSendMessageStream).toHaveBeenCalledWith(
toolCallResponseParts,
expect.any(AbortSignal),
'prompt-id-4',
);
});
@@ -1036,6 +1054,7 @@ describe('useGeminiStream', () => {
name: 'save_memory',
args: { fact: 'test' },
isClientInitiated: true,
prompt_id: 'prompt-id-6',
},
status: 'success',
responseSubmittedToGemini: false,

View File

@@ -53,6 +53,7 @@ import {
TrackedCompletedToolCall,
TrackedCancelledToolCall,
} from './useReactToolScheduler.js';
import { useSessionStats } from '../contexts/SessionContext.js';
export function mergePartListUnions(list: PartListUnion[]): PartListUnion {
const resultParts: PartListUnion = [];
@@ -101,6 +102,7 @@ export const useGeminiStream = (
const [pendingHistoryItemRef, setPendingHistoryItem] =
useStateAndRef<HistoryItemWithoutId | null>(null);
const processedMemoryToolsRef = useRef<Set<string>>(new Set());
const { startNewPrompt, getPromptCount } = useSessionStats();
const logger = useLogger();
const gitService = useMemo(() => {
if (!config.getProjectRoot()) {
@@ -203,6 +205,7 @@ export const useGeminiStream = (
query: PartListUnion,
userMessageTimestamp: number,
abortSignal: AbortSignal,
prompt_id: string,
): Promise<{
queryToSend: PartListUnion | null;
shouldProceed: boolean;
@@ -220,7 +223,7 @@ export const useGeminiStream = (
const trimmedQuery = query.trim();
logUserPrompt(
config,
new UserPromptEvent(trimmedQuery.length, trimmedQuery),
new UserPromptEvent(trimmedQuery.length, prompt_id, trimmedQuery),
);
onDebugMessage(`User query: '${trimmedQuery}'`);
await logger?.logMessage(MessageSenderType.USER, trimmedQuery);
@@ -236,6 +239,7 @@ export const useGeminiStream = (
name: toolName,
args: toolArgs,
isClientInitiated: true,
prompt_id,
};
scheduleToolCalls([toolCallRequest], abortSignal);
}
@@ -485,7 +489,11 @@ export const useGeminiStream = (
);
const submitQuery = useCallback(
async (query: PartListUnion, options?: { isContinuation: boolean }) => {
async (
query: PartListUnion,
options?: { isContinuation: boolean },
prompt_id?: string,
) => {
if (
(streamingState === StreamingState.Responding ||
streamingState === StreamingState.WaitingForConfirmation) &&
@@ -506,21 +514,34 @@ export const useGeminiStream = (
const abortSignal = abortControllerRef.current.signal;
turnCancelledRef.current = false;
if (!prompt_id) {
prompt_id = config.getSessionId() + '########' + getPromptCount();
}
const { queryToSend, shouldProceed } = await prepareQueryForGemini(
query,
userMessageTimestamp,
abortSignal,
prompt_id!,
);
if (!shouldProceed || queryToSend === null) {
return;
}
if (!options?.isContinuation) {
startNewPrompt();
}
setIsResponding(true);
setInitError(null);
try {
const stream = geminiClient.sendMessageStream(queryToSend, abortSignal);
const stream = geminiClient.sendMessageStream(
queryToSend,
abortSignal,
prompt_id!,
);
const processingStatus = await processGeminiStreamEvents(
stream,
userMessageTimestamp,
@@ -570,6 +591,8 @@ export const useGeminiStream = (
geminiClient,
onAuthError,
config,
startNewPrompt,
getPromptCount,
],
);
@@ -676,6 +699,10 @@ export const useGeminiStream = (
(toolCall) => toolCall.request.callId,
);
const prompt_ids = geminiTools.map(
(toolCall) => toolCall.request.prompt_id,
);
markToolsAsSubmitted(callIdsToMarkAsSubmitted);
// Don't continue if model was switched due to quota error
@@ -683,9 +710,13 @@ export const useGeminiStream = (
return;
}
submitQuery(mergePartListUnions(responsesToSend), {
isContinuation: true,
});
submitQuery(
mergePartListUnions(responsesToSend),
{
isContinuation: true,
},
prompt_ids[0],
);
},
[
isResponding,