diff options
26 files changed, 289 insertions, 72 deletions
diff --git a/packages/cli/src/gemini.tsx b/packages/cli/src/gemini.tsx index d9f864c6..84a3da62 100644 --- a/packages/cli/src/gemini.tsx +++ b/packages/cli/src/gemini.tsx @@ -196,10 +196,12 @@ export async function main() { process.exit(1); } + const prompt_id = Math.random().toString(16).slice(2); logUserPrompt(config, { 'event.name': 'user_prompt', 'event.timestamp': new Date().toISOString(), prompt: input, + prompt_id, prompt_length: input.length, }); @@ -210,7 +212,7 @@ export async function main() { settings, ); - await runNonInteractive(nonInteractiveConfig, input); + await runNonInteractive(nonInteractiveConfig, input, prompt_id); process.exit(0); } diff --git a/packages/cli/src/nonInteractiveCli.test.ts b/packages/cli/src/nonInteractiveCli.test.ts index 2a32cacb..14352f53 100644 --- a/packages/cli/src/nonInteractiveCli.test.ts +++ b/packages/cli/src/nonInteractiveCli.test.ts @@ -81,15 +81,18 @@ describe('runNonInteractive', () => { })(); mockChat.sendMessageStream.mockResolvedValue(inputStream); - await runNonInteractive(mockConfig, 'Test input'); + await runNonInteractive(mockConfig, 'Test input', 'prompt-id-1'); - expect(mockChat.sendMessageStream).toHaveBeenCalledWith({ - message: [{ text: 'Test input' }], - config: { - abortSignal: expect.any(AbortSignal), - tools: [{ functionDeclarations: [] }], + expect(mockChat.sendMessageStream).toHaveBeenCalledWith( + { + message: [{ text: 'Test input' }], + config: { + abortSignal: expect.any(AbortSignal), + tools: [{ functionDeclarations: [] }], + }, }, - }); + expect.any(String), + ); expect(mockProcessStdoutWrite).toHaveBeenCalledWith('Hello'); expect(mockProcessStdoutWrite).toHaveBeenCalledWith(' World'); expect(mockProcessStdoutWrite).toHaveBeenCalledWith('\n'); @@ -131,7 +134,7 @@ describe('runNonInteractive', () => { .mockResolvedValueOnce(stream1) .mockResolvedValueOnce(stream2); - await runNonInteractive(mockConfig, 'Use a tool'); + await runNonInteractive(mockConfig, 'Use a tool', 'prompt-id-2'); expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(2); expect(mockCoreExecuteToolCall).toHaveBeenCalledWith( @@ -144,6 +147,7 @@ describe('runNonInteractive', () => { expect.objectContaining({ message: [toolResponsePart], }), + expect.any(String), ); expect(mockProcessStdoutWrite).toHaveBeenCalledWith('Final answer'); }); @@ -190,7 +194,7 @@ describe('runNonInteractive', () => { .spyOn(console, 'error') .mockImplementation(() => {}); - await runNonInteractive(mockConfig, 'Trigger tool error'); + await runNonInteractive(mockConfig, 'Trigger tool error', 'prompt-id-3'); expect(mockCoreExecuteToolCall).toHaveBeenCalled(); expect(consoleErrorSpy).toHaveBeenCalledWith( @@ -200,6 +204,7 @@ describe('runNonInteractive', () => { expect.objectContaining({ message: [errorResponsePart], }), + expect.any(String), ); expect(mockProcessStdoutWrite).toHaveBeenCalledWith( 'Could not complete request.', @@ -213,7 +218,7 @@ describe('runNonInteractive', () => { .spyOn(console, 'error') .mockImplementation(() => {}); - await runNonInteractive(mockConfig, 'Initial fail'); + await runNonInteractive(mockConfig, 'Initial fail', 'prompt-id-4'); expect(consoleErrorSpy).toHaveBeenCalledWith( '[API Error: API connection failed]', @@ -265,7 +270,11 @@ describe('runNonInteractive', () => { .spyOn(console, 'error') .mockImplementation(() => {}); - await runNonInteractive(mockConfig, 'Trigger tool not found'); + await runNonInteractive( + mockConfig, + 'Trigger tool not found', + 'prompt-id-5', + ); expect(consoleErrorSpy).toHaveBeenCalledWith( 'Error executing tool nonExistentTool: Tool "nonExistentTool" not found in registry.', @@ -278,6 +287,7 @@ describe('runNonInteractive', () => { expect.objectContaining({ message: [errorResponsePart], }), + expect.any(String), ); expect(mockProcessStdoutWrite).toHaveBeenCalledWith( diff --git a/packages/cli/src/nonInteractiveCli.ts b/packages/cli/src/nonInteractiveCli.ts index 92fce058..b8b8ac3f 100644 --- a/packages/cli/src/nonInteractiveCli.ts +++ b/packages/cli/src/nonInteractiveCli.ts @@ -46,6 +46,7 @@ function getResponseText(response: GenerateContentResponse): string | null { export async function runNonInteractive( config: Config, input: string, + prompt_id: string, ): Promise<void> { await config.initialize(); // Handle EPIPE errors when the output is piped to a command that closes early. @@ -67,15 +68,18 @@ export async function runNonInteractive( while (true) { const functionCalls: FunctionCall[] = []; - const responseStream = await chat.sendMessageStream({ - message: currentMessages[0]?.parts || [], // Ensure parts are always provided - config: { - abortSignal: abortController.signal, - tools: [ - { functionDeclarations: toolRegistry.getFunctionDeclarations() }, - ], + const responseStream = await chat.sendMessageStream( + { + message: currentMessages[0]?.parts || [], // Ensure parts are always provided + config: { + abortSignal: abortController.signal, + tools: [ + { functionDeclarations: toolRegistry.getFunctionDeclarations() }, + ], + }, }, - }); + prompt_id, + ); for await (const resp of responseStream) { if (abortController.signal.aborted) { @@ -101,6 +105,7 @@ export async function runNonInteractive( name: fc.name as string, args: (fc.args ?? {}) as Record<string, unknown>, isClientInitiated: false, + prompt_id, }; const toolResponse = await executeToolCall( diff --git a/packages/cli/src/ui/components/ModelStatsDisplay.test.tsx b/packages/cli/src/ui/components/ModelStatsDisplay.test.tsx index 6c41b775..57382d91 100644 --- a/packages/cli/src/ui/components/ModelStatsDisplay.test.tsx +++ b/packages/cli/src/ui/components/ModelStatsDisplay.test.tsx @@ -27,7 +27,11 @@ const renderWithMockedStats = (metrics: SessionMetrics) => { sessionStartTime: new Date(), metrics, lastPromptTokenCount: 0, + promptCount: 5, }, + + getPromptCount: () => 5, + startNewPrompt: vi.fn(), }); return render(<ModelStatsDisplay />); diff --git a/packages/cli/src/ui/components/SessionSummaryDisplay.test.tsx b/packages/cli/src/ui/components/SessionSummaryDisplay.test.tsx index f3c0764e..38400caf 100644 --- a/packages/cli/src/ui/components/SessionSummaryDisplay.test.tsx +++ b/packages/cli/src/ui/components/SessionSummaryDisplay.test.tsx @@ -26,7 +26,11 @@ const renderWithMockedStats = (metrics: SessionMetrics) => { sessionStartTime: new Date(), metrics, lastPromptTokenCount: 0, + promptCount: 5, }, + + getPromptCount: () => 5, + startNewPrompt: vi.fn(), }); return render(<SessionSummaryDisplay duration="1h 23m 45s" />); diff --git a/packages/cli/src/ui/components/StatsDisplay.test.tsx b/packages/cli/src/ui/components/StatsDisplay.test.tsx index a62815d9..a0ed3858 100644 --- a/packages/cli/src/ui/components/StatsDisplay.test.tsx +++ b/packages/cli/src/ui/components/StatsDisplay.test.tsx @@ -27,7 +27,11 @@ const renderWithMockedStats = (metrics: SessionMetrics) => { sessionStartTime: new Date(), metrics, lastPromptTokenCount: 0, + promptCount: 5, }, + + getPromptCount: () => 5, + startNewPrompt: vi.fn(), }); return render(<StatsDisplay duration="1s" />); @@ -288,7 +292,11 @@ describe('<StatsDisplay />', () => { sessionStartTime: new Date(), metrics: zeroMetrics, lastPromptTokenCount: 0, + promptCount: 5, }, + + getPromptCount: () => 5, + startNewPrompt: vi.fn(), }); const { lastFrame } = render( diff --git a/packages/cli/src/ui/components/ToolStatsDisplay.test.tsx b/packages/cli/src/ui/components/ToolStatsDisplay.test.tsx index 54902788..e48fcc83 100644 --- a/packages/cli/src/ui/components/ToolStatsDisplay.test.tsx +++ b/packages/cli/src/ui/components/ToolStatsDisplay.test.tsx @@ -27,7 +27,11 @@ const renderWithMockedStats = (metrics: SessionMetrics) => { sessionStartTime: new Date(), metrics, lastPromptTokenCount: 0, + promptCount: 5, }, + + getPromptCount: () => 5, + startNewPrompt: vi.fn(), }); return render(<ToolStatsDisplay />); diff --git a/packages/cli/src/ui/contexts/SessionContext.tsx b/packages/cli/src/ui/contexts/SessionContext.tsx index 320df324..942af8b5 100644 --- a/packages/cli/src/ui/contexts/SessionContext.tsx +++ b/packages/cli/src/ui/contexts/SessionContext.tsx @@ -6,6 +6,7 @@ import React, { createContext, + useCallback, useContext, useState, useMemo, @@ -26,6 +27,7 @@ export interface SessionStatsState { sessionStartTime: Date; metrics: SessionMetrics; lastPromptTokenCount: number; + promptCount: number; } export interface ComputedSessionStats { @@ -46,6 +48,8 @@ export interface ComputedSessionStats { // and the functions to update it. interface SessionStatsContextValue { stats: SessionStatsState; + startNewPrompt: () => void; + getPromptCount: () => number; } // --- Context Definition --- @@ -63,6 +67,7 @@ export const SessionStatsProvider: React.FC<{ children: React.ReactNode }> = ({ sessionStartTime: new Date(), metrics: uiTelemetryService.getMetrics(), lastPromptTokenCount: 0, + promptCount: 0, }); useEffect(() => { @@ -92,11 +97,25 @@ export const SessionStatsProvider: React.FC<{ children: React.ReactNode }> = ({ }; }, []); + const startNewPrompt = useCallback(() => { + setStats((prevState) => ({ + ...prevState, + promptCount: prevState.promptCount + 1, + })); + }, []); + + const getPromptCount = useCallback( + () => stats.promptCount, + [stats.promptCount], + ); + const value = useMemo( () => ({ stats, + startNewPrompt, + getPromptCount, }), - [stats], + [stats, startNewPrompt, getPromptCount], ); return ( diff --git a/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts b/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts index 137098df..45f52074 100644 --- a/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts +++ b/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts @@ -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({ diff --git a/packages/cli/src/ui/hooks/slashCommandProcessor.ts b/packages/cli/src/ui/hooks/slashCommandProcessor.ts index 66cf4e39..f53bdc12 100644 --- a/packages/cli/src/ui/hooks/slashCommandProcessor.ts +++ b/packages/cli/src/ui/hooks/slashCommandProcessor.ts @@ -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, diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx index 62ade50f..e0e21f55 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx +++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx @@ -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, diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index d32c9ffa..b82b0cb2 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -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, diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index cd77a3f7..44828a74 100644 --- a/packages/core/src/core/client.test.ts +++ b/packages/core/src/core/client.test.ts @@ -450,7 +450,7 @@ describe('Gemini Client (client.ts)', () => { }); const initialChat = client.getChat(); - const result = await client.tryCompressChat(); + const result = await client.tryCompressChat('prompt-id-2'); const newChat = client.getChat(); expect(tokenLimit).toHaveBeenCalled(); @@ -476,7 +476,7 @@ describe('Gemini Client (client.ts)', () => { }); const initialChat = client.getChat(); - const result = await client.tryCompressChat(); + const result = await client.tryCompressChat('prompt-id-3'); const newChat = client.getChat(); expect(tokenLimit).toHaveBeenCalled(); @@ -507,7 +507,7 @@ describe('Gemini Client (client.ts)', () => { }); const initialChat = client.getChat(); - const result = await client.tryCompressChat(true); // force = true + const result = await client.tryCompressChat('prompt-id-1', true); // force = true const newChat = client.getChat(); expect(mockSendMessage).toHaveBeenCalled(); @@ -545,6 +545,7 @@ describe('Gemini Client (client.ts)', () => { const stream = client.sendMessageStream( [{ text: 'Hi' }], new AbortController().signal, + 'prompt-id-1', ); // Consume the stream manually to get the final return value. @@ -597,6 +598,7 @@ describe('Gemini Client (client.ts)', () => { const stream = client.sendMessageStream( [{ text: 'Start conversation' }], signal, + 'prompt-id-2', ); // Count how many stream events we get @@ -697,6 +699,7 @@ describe('Gemini Client (client.ts)', () => { const stream = client.sendMessageStream( [{ text: 'Start conversation' }], signal, + 'prompt-id-3', Number.MAX_SAFE_INTEGER, // Bypass the MAX_TURNS protection ); @@ -806,7 +809,7 @@ describe('Gemini Client (client.ts)', () => { client['contentGenerator'] = mockGenerator as ContentGenerator; client['startChat'] = vi.fn().mockResolvedValue(mockChat); - const result = await client.tryCompressChat(true); + const result = await client.tryCompressChat('prompt-id-4', true); expect(mockCountTokens).toHaveBeenCalledTimes(2); expect(mockCountTokens).toHaveBeenNthCalledWith(1, { diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index 51aab961..eee52cb4 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -261,23 +261,25 @@ export class GeminiClient { async *sendMessageStream( request: PartListUnion, signal: AbortSignal, + prompt_id: string, turns: number = this.MAX_TURNS, originalModel?: string, ): AsyncGenerator<ServerGeminiStreamEvent, Turn> { // Ensure turns never exceeds MAX_TURNS to prevent infinite loops const boundedTurns = Math.min(turns, this.MAX_TURNS); if (!boundedTurns) { - return new Turn(this.getChat()); + return new Turn(this.getChat(), prompt_id); } // Track the original model from the first call to detect model switching const initialModel = originalModel || this.config.getModel(); - const compressed = await this.tryCompressChat(); + const compressed = await this.tryCompressChat(prompt_id); + if (compressed) { yield { type: GeminiEventType.ChatCompressed, value: compressed }; } - const turn = new Turn(this.getChat()); + const turn = new Turn(this.getChat(), prompt_id); const resultStream = turn.run(request, signal); for await (const event of resultStream) { yield event; @@ -303,6 +305,7 @@ export class GeminiClient { yield* this.sendMessageStream( nextRequest, signal, + prompt_id, boundedTurns - 1, initialModel, ); @@ -492,6 +495,7 @@ export class GeminiClient { } async tryCompressChat( + prompt_id: string, force: boolean = false, ): Promise<ChatCompressionInfo | null> { const curatedHistory = this.getChat().getHistory(true); @@ -538,14 +542,17 @@ export class GeminiClient { this.getChat().setHistory(historyToCompress); - const { text: summary } = await this.getChat().sendMessage({ - message: { - text: 'First, reason in your scratchpad. Then, generate the <state_snapshot>.', - }, - config: { - systemInstruction: { text: getCompressionPrompt() }, + const { text: summary } = await this.getChat().sendMessage( + { + message: { + text: 'First, reason in your scratchpad. Then, generate the <state_snapshot>.', + }, + config: { + systemInstruction: { text: getCompressionPrompt() }, + }, }, - }); + prompt_id, + ); this.chat = await this.startChat([ { role: 'user', diff --git a/packages/core/src/core/coreToolScheduler.test.ts b/packages/core/src/core/coreToolScheduler.test.ts index d6030d6f..0b2c5124 100644 --- a/packages/core/src/core/coreToolScheduler.test.ts +++ b/packages/core/src/core/coreToolScheduler.test.ts @@ -139,6 +139,7 @@ describe('CoreToolScheduler', () => { name: 'mockTool', args: {}, isClientInitiated: false, + prompt_id: 'prompt-id-1', }; abortController.abort(); @@ -206,6 +207,7 @@ describe('CoreToolScheduler with payload', () => { name: 'mockModifiableTool', args: {}, isClientInitiated: false, + prompt_id: 'prompt-id-2', }; await scheduler.schedule([request], abortController.signal); diff --git a/packages/core/src/core/geminiChat.test.ts b/packages/core/src/core/geminiChat.test.ts index 35e6bf6c..39dd883e 100644 --- a/packages/core/src/core/geminiChat.test.ts +++ b/packages/core/src/core/geminiChat.test.ts @@ -77,7 +77,7 @@ describe('GeminiChat', () => { } as unknown as GenerateContentResponse; vi.mocked(mockModelsModule.generateContent).mockResolvedValue(response); - await chat.sendMessage({ message: 'hello' }); + await chat.sendMessage({ message: 'hello' }, 'prompt-id-1'); expect(mockModelsModule.generateContent).toHaveBeenCalledWith({ model: 'gemini-pro', @@ -109,7 +109,7 @@ describe('GeminiChat', () => { response, ); - await chat.sendMessageStream({ message: 'hello' }); + await chat.sendMessageStream({ message: 'hello' }, 'prompt-id-1'); expect(mockModelsModule.generateContentStream).toHaveBeenCalledWith({ model: 'gemini-pro', diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index 2c149e93..f57425a3 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -151,13 +151,18 @@ export class GeminiChat { private async _logApiRequest( contents: Content[], model: string, + prompt_id: string, ): Promise<void> { const requestText = this._getRequestTextFromContents(contents); - logApiRequest(this.config, new ApiRequestEvent(model, requestText)); + logApiRequest( + this.config, + new ApiRequestEvent(model, prompt_id, requestText), + ); } private async _logApiResponse( durationMs: number, + prompt_id: string, usageMetadata?: GenerateContentResponseUsageMetadata, responseText?: string, ): Promise<void> { @@ -166,13 +171,18 @@ export class GeminiChat { new ApiResponseEvent( this.config.getModel(), durationMs, + prompt_id, usageMetadata, responseText, ), ); } - private _logApiError(durationMs: number, error: unknown): void { + private _logApiError( + durationMs: number, + error: unknown, + prompt_id: string, + ): void { const errorMessage = error instanceof Error ? error.message : String(error); const errorType = error instanceof Error ? error.name : 'unknown'; @@ -182,6 +192,7 @@ export class GeminiChat { this.config.getModel(), errorMessage, durationMs, + prompt_id, errorType, ), ); @@ -255,12 +266,13 @@ export class GeminiChat { */ async sendMessage( params: SendMessageParameters, + prompt_id: string, ): Promise<GenerateContentResponse> { await this.sendPromise; const userContent = createUserContent(params.message); const requestContents = this.getHistory(true).concat(userContent); - this._logApiRequest(requestContents, this.config.getModel()); + this._logApiRequest(requestContents, this.config.getModel(), prompt_id); const startTime = Date.now(); let response: GenerateContentResponse; @@ -301,6 +313,7 @@ export class GeminiChat { const durationMs = Date.now() - startTime; await this._logApiResponse( durationMs, + prompt_id, response.usageMetadata, getStructuredResponse(response), ); @@ -332,7 +345,7 @@ export class GeminiChat { return response; } catch (error) { const durationMs = Date.now() - startTime; - this._logApiError(durationMs, error); + this._logApiError(durationMs, error, prompt_id); this.sendPromise = Promise.resolve(); throw error; } @@ -362,11 +375,12 @@ export class GeminiChat { */ async sendMessageStream( params: SendMessageParameters, + prompt_id: string, ): Promise<AsyncGenerator<GenerateContentResponse>> { await this.sendPromise; const userContent = createUserContent(params.message); const requestContents = this.getHistory(true).concat(userContent); - this._logApiRequest(requestContents, this.config.getModel()); + this._logApiRequest(requestContents, this.config.getModel(), prompt_id); const startTime = Date.now(); @@ -420,11 +434,12 @@ export class GeminiChat { streamResponse, userContent, startTime, + prompt_id, ); return result; } catch (error) { const durationMs = Date.now() - startTime; - this._logApiError(durationMs, error); + this._logApiError(durationMs, error, prompt_id); this.sendPromise = Promise.resolve(); throw error; } @@ -496,6 +511,7 @@ export class GeminiChat { streamResponse: AsyncGenerator<GenerateContentResponse>, inputContent: Content, startTime: number, + prompt_id: string, ) { const outputContent: Content[] = []; const chunks: GenerateContentResponse[] = []; @@ -519,7 +535,7 @@ export class GeminiChat { } catch (error) { errorOccurred = true; const durationMs = Date.now() - startTime; - this._logApiError(durationMs, error); + this._logApiError(durationMs, error, prompt_id); throw error; } @@ -534,6 +550,7 @@ export class GeminiChat { const fullText = getStructuredResponseFromParts(allParts); await this._logApiResponse( durationMs, + prompt_id, this.getFinalUsageMetadata(chunks), fullText, ); diff --git a/packages/core/src/core/nonInteractiveToolExecutor.test.ts b/packages/core/src/core/nonInteractiveToolExecutor.test.ts index 80a8bdaf..14b048b4 100644 --- a/packages/core/src/core/nonInteractiveToolExecutor.test.ts +++ b/packages/core/src/core/nonInteractiveToolExecutor.test.ts @@ -67,6 +67,7 @@ describe('executeToolCall', () => { name: 'testTool', args: { param1: 'value1' }, isClientInitiated: false, + prompt_id: 'prompt-id-1', }; const toolResult: ToolResult = { llmContent: 'Tool executed successfully', @@ -105,6 +106,7 @@ describe('executeToolCall', () => { name: 'nonExistentTool', args: {}, isClientInitiated: false, + prompt_id: 'prompt-id-2', }; vi.mocked(mockToolRegistry.getTool).mockReturnValue(undefined); @@ -140,6 +142,7 @@ describe('executeToolCall', () => { name: 'testTool', args: { param1: 'value1' }, isClientInitiated: false, + prompt_id: 'prompt-id-3', }; const executionError = new Error('Tool execution failed'); vi.mocked(mockToolRegistry.getTool).mockReturnValue(mockTool); @@ -172,6 +175,7 @@ describe('executeToolCall', () => { name: 'testTool', args: { param1: 'value1' }, isClientInitiated: false, + prompt_id: 'prompt-id-4', }; const cancellationError = new Error('Operation cancelled'); vi.mocked(mockToolRegistry.getTool).mockReturnValue(mockTool); @@ -215,6 +219,7 @@ describe('executeToolCall', () => { name: 'testTool', args: {}, isClientInitiated: false, + prompt_id: 'prompt-id-5', }; const imageDataPart: Part = { inlineData: { mimeType: 'image/png', data: 'base64data' }, diff --git a/packages/core/src/core/nonInteractiveToolExecutor.ts b/packages/core/src/core/nonInteractiveToolExecutor.ts index 8efb58e0..38c55c10 100644 --- a/packages/core/src/core/nonInteractiveToolExecutor.ts +++ b/packages/core/src/core/nonInteractiveToolExecutor.ts @@ -40,6 +40,7 @@ export async function executeToolCall( duration_ms: durationMs, success: false, error: error.message, + prompt_id: toolCallRequest.prompt_id, }); // Ensure the response structure matches what the API expects for an error return { @@ -75,6 +76,7 @@ export async function executeToolCall( function_args: toolCallRequest.args, duration_ms: durationMs, success: true, + prompt_id: toolCallRequest.prompt_id, }); const response = convertToFunctionResponse( @@ -100,6 +102,7 @@ export async function executeToolCall( duration_ms: durationMs, success: false, error: error.message, + prompt_id: toolCallRequest.prompt_id, }); return { callId: toolCallRequest.callId, diff --git a/packages/core/src/core/turn.test.ts b/packages/core/src/core/turn.test.ts index bfbd6e17..b0c27f7e 100644 --- a/packages/core/src/core/turn.test.ts +++ b/packages/core/src/core/turn.test.ts @@ -55,7 +55,7 @@ describe('Turn', () => { sendMessageStream: mockSendMessageStream, getHistory: mockGetHistory, }; - turn = new Turn(mockChatInstance as unknown as GeminiChat); + turn = new Turn(mockChatInstance as unknown as GeminiChat, 'prompt-id-1'); mockGetHistory.mockReturnValue([]); mockSendMessageStream.mockResolvedValue((async function* () {})()); }); @@ -92,10 +92,13 @@ describe('Turn', () => { events.push(event); } - expect(mockSendMessageStream).toHaveBeenCalledWith({ - message: reqParts, - config: { abortSignal: expect.any(AbortSignal) }, - }); + expect(mockSendMessageStream).toHaveBeenCalledWith( + { + message: reqParts, + config: { abortSignal: expect.any(AbortSignal) }, + }, + 'prompt-id-1', + ); expect(events).toEqual([ { type: GeminiEventType.Content, value: 'Hello' }, diff --git a/packages/core/src/core/turn.ts b/packages/core/src/core/turn.ts index 4f93247b..aeeaa889 100644 --- a/packages/core/src/core/turn.ts +++ b/packages/core/src/core/turn.ts @@ -64,6 +64,7 @@ export interface ToolCallRequestInfo { name: string; args: Record<string, unknown>; isClientInitiated: boolean; + prompt_id: string; } export interface ToolCallResponseInfo { @@ -143,7 +144,10 @@ export class Turn { readonly pendingToolCalls: ToolCallRequestInfo[]; private debugResponses: GenerateContentResponse[]; - constructor(private readonly chat: GeminiChat) { + constructor( + private readonly chat: GeminiChat, + private readonly prompt_id: string, + ) { this.pendingToolCalls = []; this.debugResponses = []; } @@ -153,12 +157,15 @@ export class Turn { signal: AbortSignal, ): AsyncGenerator<ServerGeminiStreamEvent> { try { - const responseStream = await this.chat.sendMessageStream({ - message: req, - config: { - abortSignal: signal, + const responseStream = await this.chat.sendMessageStream( + { + message: req, + config: { + abortSignal: signal, + }, }, - }); + this.prompt_id, + ); for await (const resp of responseStream) { if (signal?.aborted) { @@ -252,6 +259,7 @@ export class Turn { name, args, isClientInitiated: false, + prompt_id: this.prompt_id, }; this.pendingToolCalls.push(toolCallRequest); diff --git a/packages/core/src/telemetry/clearcut-logger/clearcut-logger.ts b/packages/core/src/telemetry/clearcut-logger/clearcut-logger.ts index 73c82f23..a64a9795 100644 --- a/packages/core/src/telemetry/clearcut-logger/clearcut-logger.ts +++ b/packages/core/src/telemetry/clearcut-logger/clearcut-logger.ts @@ -265,6 +265,10 @@ export class ClearcutLogger { gemini_cli_key: EventMetadataKey.GEMINI_CLI_USER_PROMPT_LENGTH, value: JSON.stringify(event.prompt_length), }, + { + gemini_cli_key: EventMetadataKey.GEMINI_CLI_PROMPT_ID, + value: JSON.stringify(event.prompt_id), + }, ]; this.enqueueLogEvent(this.createLogEvent(new_prompt_event_name, data)); @@ -280,6 +284,10 @@ export class ClearcutLogger { value: JSON.stringify(event.function_name), }, { + gemini_cli_key: EventMetadataKey.GEMINI_CLI_PROMPT_ID, + value: JSON.stringify(event.prompt_id), + }, + { gemini_cli_key: EventMetadataKey.GEMINI_CLI_TOOL_CALL_DECISION, value: JSON.stringify(event.decision), }, @@ -313,6 +321,10 @@ export class ClearcutLogger { gemini_cli_key: EventMetadataKey.GEMINI_CLI_API_REQUEST_MODEL, value: JSON.stringify(event.model), }, + { + gemini_cli_key: EventMetadataKey.GEMINI_CLI_PROMPT_ID, + value: JSON.stringify(event.prompt_id), + }, ]; this.enqueueLogEvent(this.createLogEvent(api_request_event_name, data)); @@ -328,6 +340,10 @@ export class ClearcutLogger { value: JSON.stringify(event.model), }, { + gemini_cli_key: EventMetadataKey.GEMINI_CLI_PROMPT_ID, + value: JSON.stringify(event.prompt_id), + }, + { gemini_cli_key: EventMetadataKey.GEMINI_CLI_API_RESPONSE_STATUS_CODE, value: JSON.stringify(event.status_code), }, @@ -379,6 +395,10 @@ export class ClearcutLogger { value: JSON.stringify(event.model), }, { + gemini_cli_key: EventMetadataKey.GEMINI_CLI_PROMPT_ID, + value: JSON.stringify(event.prompt_id), + }, + { gemini_cli_key: EventMetadataKey.GEMINI_CLI_API_ERROR_TYPE, value: JSON.stringify(event.error_type), }, diff --git a/packages/core/src/telemetry/clearcut-logger/event-metadata-key.ts b/packages/core/src/telemetry/clearcut-logger/event-metadata-key.ts index 146dcdeb..e8a74936 100644 --- a/packages/core/src/telemetry/clearcut-logger/event-metadata-key.ts +++ b/packages/core/src/telemetry/clearcut-logger/event-metadata-key.ts @@ -137,6 +137,13 @@ export enum EventMetadataKey { // Logs the end of a session. GEMINI_CLI_END_SESSION_ID = 34, + + // ========================================================================== + // Shared Keys + // =========================================================================== + + // Logs the Prompt Id + GEMINI_CLI_PROMPT_ID = 35, } export function getEventMetadataKey( diff --git a/packages/core/src/telemetry/loggers.test.ts b/packages/core/src/telemetry/loggers.test.ts index 5b922333..e7dd721f 100644 --- a/packages/core/src/telemetry/loggers.test.ts +++ b/packages/core/src/telemetry/loggers.test.ts @@ -127,7 +127,7 @@ describe('loggers', () => { } as unknown as Config; it('should log a user prompt', () => { - const event = new UserPromptEvent(11, 'test-prompt'); + const event = new UserPromptEvent(11, 'prompt-id-8', 'test-prompt'); logUserPrompt(mockConfig, event); @@ -201,6 +201,7 @@ describe('loggers', () => { const event = new ApiResponseEvent( 'test-model', 100, + 'prompt-id-1', usageData, 'test-response', ); @@ -224,6 +225,7 @@ describe('loggers', () => { tool_token_count: 2, total_token_count: 0, response_text: 'test-response', + prompt_id: 'prompt-id-1', }, }); @@ -260,6 +262,7 @@ describe('loggers', () => { const event = new ApiResponseEvent( 'test-model', 100, + 'prompt-id-1', usageData, 'test-response', 'test-error', @@ -296,7 +299,11 @@ describe('loggers', () => { } as Config; it('should log an API request with request_text', () => { - const event = new ApiRequestEvent('test-model', 'This is a test request'); + const event = new ApiRequestEvent( + 'test-model', + 'prompt-id-7', + 'This is a test request', + ); logApiRequest(mockConfig, event); @@ -308,12 +315,13 @@ describe('loggers', () => { 'event.timestamp': '2025-01-01T00:00:00.000Z', model: 'test-model', request_text: 'This is a test request', + prompt_id: 'prompt-id-7', }, }); }); it('should log an API request without request_text', () => { - const event = new ApiRequestEvent('test-model'); + const event = new ApiRequestEvent('test-model', 'prompt-id-6'); logApiRequest(mockConfig, event); @@ -324,6 +332,7 @@ describe('loggers', () => { 'event.name': EVENT_API_REQUEST, 'event.timestamp': '2025-01-01T00:00:00.000Z', model: 'test-model', + prompt_id: 'prompt-id-6', }, }); }); @@ -394,6 +403,7 @@ describe('loggers', () => { }, callId: 'test-call-id', isClientInitiated: true, + prompt_id: 'prompt-id-1', }, response: { callId: 'test-call-id', @@ -427,6 +437,7 @@ describe('loggers', () => { duration_ms: 100, success: true, decision: ToolCallDecision.ACCEPT, + prompt_id: 'prompt-id-1', }, }); @@ -455,6 +466,7 @@ describe('loggers', () => { }, callId: 'test-call-id', isClientInitiated: true, + prompt_id: 'prompt-id-2', }, response: { callId: 'test-call-id', @@ -487,6 +499,7 @@ describe('loggers', () => { duration_ms: 100, success: false, decision: ToolCallDecision.REJECT, + prompt_id: 'prompt-id-2', }, }); @@ -516,6 +529,7 @@ describe('loggers', () => { }, callId: 'test-call-id', isClientInitiated: true, + prompt_id: 'prompt-id-3', }, response: { callId: 'test-call-id', @@ -549,6 +563,7 @@ describe('loggers', () => { duration_ms: 100, success: true, decision: ToolCallDecision.MODIFY, + prompt_id: 'prompt-id-3', }, }); @@ -578,6 +593,7 @@ describe('loggers', () => { }, callId: 'test-call-id', isClientInitiated: true, + prompt_id: 'prompt-id-4', }, response: { callId: 'test-call-id', @@ -609,6 +625,7 @@ describe('loggers', () => { ), duration_ms: 100, success: true, + prompt_id: 'prompt-id-4', }, }); @@ -638,6 +655,7 @@ describe('loggers', () => { }, callId: 'test-call-id', isClientInitiated: true, + prompt_id: 'prompt-id-5', }, response: { callId: 'test-call-id', @@ -675,6 +693,7 @@ describe('loggers', () => { 'error.message': 'test-error', error_type: 'test-error-type', 'error.type': 'test-error-type', + prompt_id: 'prompt-id-5', }, }); diff --git a/packages/core/src/telemetry/types.ts b/packages/core/src/telemetry/types.ts index 9883111a..46f86b89 100644 --- a/packages/core/src/telemetry/types.ts +++ b/packages/core/src/telemetry/types.ts @@ -95,12 +95,14 @@ export class UserPromptEvent { 'event.name': 'user_prompt'; 'event.timestamp': string; // ISO 8601 prompt_length: number; + prompt_id: string; prompt?: string; - constructor(prompt_length: number, prompt?: string) { + constructor(prompt_length: number, prompt_Id: string, prompt?: string) { this['event.name'] = 'user_prompt'; this['event.timestamp'] = new Date().toISOString(); this.prompt_length = prompt_length; + this.prompt_id = prompt_Id; this.prompt = prompt; } } @@ -115,6 +117,7 @@ export class ToolCallEvent { decision?: ToolCallDecision; error?: string; error_type?: string; + prompt_id: string; constructor(call: CompletedToolCall) { this['event.name'] = 'tool_call'; @@ -128,6 +131,7 @@ export class ToolCallEvent { : undefined; this.error = call.response.error?.message; this.error_type = call.response.error?.name; + this.prompt_id = call.request.prompt_id; } } @@ -135,12 +139,14 @@ export class ApiRequestEvent { 'event.name': 'api_request'; 'event.timestamp': string; // ISO 8601 model: string; + prompt_id: string; request_text?: string; - constructor(model: string, request_text?: string) { + constructor(model: string, prompt_id: string, request_text?: string) { this['event.name'] = 'api_request'; this['event.timestamp'] = new Date().toISOString(); this.model = model; + this.prompt_id = prompt_id; this.request_text = request_text; } } @@ -153,11 +159,13 @@ export class ApiErrorEvent { error_type?: string; status_code?: number | string; duration_ms: number; + prompt_id: string; constructor( model: string, error: string, duration_ms: number, + prompt_id: string, error_type?: string, status_code?: number | string, ) { @@ -168,6 +176,7 @@ export class ApiErrorEvent { this.error_type = error_type; this.status_code = status_code; this.duration_ms = duration_ms; + this.prompt_id = prompt_id; } } @@ -185,10 +194,12 @@ export class ApiResponseEvent { tool_token_count: number; total_token_count: number; response_text?: string; + prompt_id: string; constructor( model: string, duration_ms: number, + prompt_id: string, usage_data?: GenerateContentResponseUsageMetadata, response_text?: string, error?: string, @@ -206,6 +217,7 @@ export class ApiResponseEvent { this.total_token_count = usage_data?.totalTokenCount ?? 0; this.response_text = response_text; this.error = error; + this.prompt_id = prompt_id; } } diff --git a/packages/core/src/telemetry/uiTelemetry.test.ts b/packages/core/src/telemetry/uiTelemetry.test.ts index 9643ed97..34a2fe22 100644 --- a/packages/core/src/telemetry/uiTelemetry.test.ts +++ b/packages/core/src/telemetry/uiTelemetry.test.ts @@ -36,6 +36,7 @@ const createFakeCompletedToolCall = ( name, args: { foo: 'bar' }, isClientInitiated: false, + prompt_id: 'prompt-id-1', }; if (success) { |