summaryrefslogtreecommitdiff
path: root/packages/cli/src
diff options
context:
space:
mode:
Diffstat (limited to 'packages/cli/src')
-rw-r--r--packages/cli/src/gemini.tsx4
-rw-r--r--packages/cli/src/nonInteractiveCli.test.ts32
-rw-r--r--packages/cli/src/nonInteractiveCli.ts21
-rw-r--r--packages/cli/src/ui/components/ModelStatsDisplay.test.tsx4
-rw-r--r--packages/cli/src/ui/components/SessionSummaryDisplay.test.tsx4
-rw-r--r--packages/cli/src/ui/components/StatsDisplay.test.tsx8
-rw-r--r--packages/cli/src/ui/components/ToolStatsDisplay.test.tsx4
-rw-r--r--packages/cli/src/ui/contexts/SessionContext.tsx21
-rw-r--r--packages/cli/src/ui/hooks/slashCommandProcessor.test.ts7
-rw-r--r--packages/cli/src/ui/hooks/slashCommandProcessor.ts3
-rw-r--r--packages/cli/src/ui/hooks/useGeminiStream.test.tsx25
-rw-r--r--packages/cli/src/ui/hooks/useGeminiStream.ts43
12 files changed, 143 insertions, 33 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,