mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-03-25 05:21:03 -07:00
feat: added basic dev otel trace instrumentation (#11690)
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
Pavel Jbanov
@@ -35,6 +35,7 @@ import {
|
||||
promptIdContext,
|
||||
WRITE_FILE_TOOL_NAME,
|
||||
tokenLimit,
|
||||
runInDevTraceSpan,
|
||||
} from '@google/gemini-cli-core';
|
||||
import { type Part, type PartListUnion, FinishReason } from '@google/genai';
|
||||
import type {
|
||||
@@ -805,152 +806,160 @@ export const useGeminiStream = (
|
||||
query: PartListUnion,
|
||||
options?: { isContinuation: boolean },
|
||||
prompt_id?: string,
|
||||
) => {
|
||||
const queryId = `${Date.now()}-${Math.random()}`;
|
||||
activeQueryIdRef.current = queryId;
|
||||
if (
|
||||
(streamingState === StreamingState.Responding ||
|
||||
streamingState === StreamingState.WaitingForConfirmation) &&
|
||||
!options?.isContinuation
|
||||
)
|
||||
return;
|
||||
|
||||
const userMessageTimestamp = Date.now();
|
||||
|
||||
// Reset quota error flag when starting a new query (not a continuation)
|
||||
if (!options?.isContinuation) {
|
||||
setModelSwitchedFromQuotaError(false);
|
||||
config.setQuotaErrorOccurred(false);
|
||||
}
|
||||
|
||||
abortControllerRef.current = new AbortController();
|
||||
const abortSignal = abortControllerRef.current.signal;
|
||||
turnCancelledRef.current = false;
|
||||
|
||||
if (!prompt_id) {
|
||||
prompt_id = config.getSessionId() + '########' + getPromptCount();
|
||||
}
|
||||
return promptIdContext.run(prompt_id, async () => {
|
||||
const { queryToSend, shouldProceed } = await prepareQueryForGemini(
|
||||
query,
|
||||
userMessageTimestamp,
|
||||
abortSignal,
|
||||
prompt_id,
|
||||
);
|
||||
|
||||
if (!shouldProceed || queryToSend === null) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (!options?.isContinuation) {
|
||||
if (typeof queryToSend === 'string') {
|
||||
// logging the text prompts only for now
|
||||
const promptText = queryToSend;
|
||||
logUserPrompt(
|
||||
config,
|
||||
new UserPromptEvent(
|
||||
promptText.length,
|
||||
prompt_id,
|
||||
config.getContentGeneratorConfig()?.authType,
|
||||
promptText,
|
||||
),
|
||||
);
|
||||
}
|
||||
startNewPrompt();
|
||||
setThought(null); // Reset thought when starting a new prompt
|
||||
}
|
||||
|
||||
setIsResponding(true);
|
||||
setInitError(null);
|
||||
|
||||
// Store query and prompt_id for potential retry on loop detection
|
||||
lastQueryRef.current = queryToSend;
|
||||
lastPromptIdRef.current = prompt_id;
|
||||
|
||||
try {
|
||||
const stream = geminiClient.sendMessageStream(
|
||||
queryToSend,
|
||||
abortSignal,
|
||||
prompt_id,
|
||||
);
|
||||
const processingStatus = await processGeminiStreamEvents(
|
||||
stream,
|
||||
userMessageTimestamp,
|
||||
abortSignal,
|
||||
);
|
||||
|
||||
if (processingStatus === StreamProcessingStatus.UserCancelled) {
|
||||
) =>
|
||||
runInDevTraceSpan(
|
||||
{ name: 'submitQuery' },
|
||||
async ({ metadata: spanMetadata }) => {
|
||||
spanMetadata.input = query;
|
||||
const queryId = `${Date.now()}-${Math.random()}`;
|
||||
activeQueryIdRef.current = queryId;
|
||||
if (
|
||||
(streamingState === StreamingState.Responding ||
|
||||
streamingState === StreamingState.WaitingForConfirmation) &&
|
||||
!options?.isContinuation
|
||||
)
|
||||
return;
|
||||
|
||||
const userMessageTimestamp = Date.now();
|
||||
|
||||
// Reset quota error flag when starting a new query (not a continuation)
|
||||
if (!options?.isContinuation) {
|
||||
setModelSwitchedFromQuotaError(false);
|
||||
config.setQuotaErrorOccurred(false);
|
||||
}
|
||||
|
||||
if (pendingHistoryItemRef.current) {
|
||||
addItem(pendingHistoryItemRef.current, userMessageTimestamp);
|
||||
setPendingHistoryItem(null);
|
||||
}
|
||||
if (loopDetectedRef.current) {
|
||||
loopDetectedRef.current = false;
|
||||
// Show the confirmation dialog to choose whether to disable loop detection
|
||||
setLoopDetectionConfirmationRequest({
|
||||
onComplete: (result: { userSelection: 'disable' | 'keep' }) => {
|
||||
setLoopDetectionConfirmationRequest(null);
|
||||
abortControllerRef.current = new AbortController();
|
||||
const abortSignal = abortControllerRef.current.signal;
|
||||
turnCancelledRef.current = false;
|
||||
|
||||
if (result.userSelection === 'disable') {
|
||||
config
|
||||
.getGeminiClient()
|
||||
.getLoopDetectionService()
|
||||
.disableForSession();
|
||||
addItem(
|
||||
{
|
||||
type: 'info',
|
||||
text: `Loop detection has been disabled for this session. Retrying request...`,
|
||||
},
|
||||
Date.now(),
|
||||
);
|
||||
|
||||
if (lastQueryRef.current && lastPromptIdRef.current) {
|
||||
submitQuery(
|
||||
lastQueryRef.current,
|
||||
{ isContinuation: true },
|
||||
lastPromptIdRef.current,
|
||||
);
|
||||
}
|
||||
} else {
|
||||
addItem(
|
||||
{
|
||||
type: 'info',
|
||||
text: `A potential loop was detected. This can happen due to repetitive tool calls or other model behavior. The request has been halted.`,
|
||||
},
|
||||
Date.now(),
|
||||
);
|
||||
}
|
||||
},
|
||||
});
|
||||
if (!prompt_id) {
|
||||
prompt_id = config.getSessionId() + '########' + getPromptCount();
|
||||
}
|
||||
} catch (error: unknown) {
|
||||
if (error instanceof UnauthorizedError) {
|
||||
onAuthError('Session expired or is unauthorized.');
|
||||
} else if (!isNodeError(error) || error.name !== 'AbortError') {
|
||||
addItem(
|
||||
{
|
||||
type: MessageType.ERROR,
|
||||
text: parseAndFormatApiError(
|
||||
getErrorMessage(error) || 'Unknown error',
|
||||
config.getContentGeneratorConfig()?.authType,
|
||||
undefined,
|
||||
config.getModel(),
|
||||
DEFAULT_GEMINI_FLASH_MODEL,
|
||||
),
|
||||
},
|
||||
return promptIdContext.run(prompt_id, async () => {
|
||||
const { queryToSend, shouldProceed } = await prepareQueryForGemini(
|
||||
query,
|
||||
userMessageTimestamp,
|
||||
abortSignal,
|
||||
prompt_id!,
|
||||
);
|
||||
}
|
||||
} finally {
|
||||
if (activeQueryIdRef.current === queryId) {
|
||||
setIsResponding(false);
|
||||
}
|
||||
}
|
||||
});
|
||||
},
|
||||
|
||||
if (!shouldProceed || queryToSend === null) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (!options?.isContinuation) {
|
||||
if (typeof queryToSend === 'string') {
|
||||
// logging the text prompts only for now
|
||||
const promptText = queryToSend;
|
||||
logUserPrompt(
|
||||
config,
|
||||
new UserPromptEvent(
|
||||
promptText.length,
|
||||
prompt_id!,
|
||||
config.getContentGeneratorConfig()?.authType,
|
||||
promptText,
|
||||
),
|
||||
);
|
||||
}
|
||||
startNewPrompt();
|
||||
setThought(null); // Reset thought when starting a new prompt
|
||||
}
|
||||
|
||||
setIsResponding(true);
|
||||
setInitError(null);
|
||||
|
||||
// Store query and prompt_id for potential retry on loop detection
|
||||
lastQueryRef.current = queryToSend;
|
||||
lastPromptIdRef.current = prompt_id!;
|
||||
|
||||
try {
|
||||
const stream = geminiClient.sendMessageStream(
|
||||
queryToSend,
|
||||
abortSignal,
|
||||
prompt_id!,
|
||||
);
|
||||
const processingStatus = await processGeminiStreamEvents(
|
||||
stream,
|
||||
userMessageTimestamp,
|
||||
abortSignal,
|
||||
);
|
||||
|
||||
if (processingStatus === StreamProcessingStatus.UserCancelled) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (pendingHistoryItemRef.current) {
|
||||
addItem(pendingHistoryItemRef.current, userMessageTimestamp);
|
||||
setPendingHistoryItem(null);
|
||||
}
|
||||
if (loopDetectedRef.current) {
|
||||
loopDetectedRef.current = false;
|
||||
// Show the confirmation dialog to choose whether to disable loop detection
|
||||
setLoopDetectionConfirmationRequest({
|
||||
onComplete: (result: {
|
||||
userSelection: 'disable' | 'keep';
|
||||
}) => {
|
||||
setLoopDetectionConfirmationRequest(null);
|
||||
|
||||
if (result.userSelection === 'disable') {
|
||||
config
|
||||
.getGeminiClient()
|
||||
.getLoopDetectionService()
|
||||
.disableForSession();
|
||||
addItem(
|
||||
{
|
||||
type: 'info',
|
||||
text: `Loop detection has been disabled for this session. Retrying request...`,
|
||||
},
|
||||
Date.now(),
|
||||
);
|
||||
|
||||
if (lastQueryRef.current && lastPromptIdRef.current) {
|
||||
submitQuery(
|
||||
lastQueryRef.current,
|
||||
{ isContinuation: true },
|
||||
lastPromptIdRef.current,
|
||||
);
|
||||
}
|
||||
} else {
|
||||
addItem(
|
||||
{
|
||||
type: 'info',
|
||||
text: `A potential loop was detected. This can happen due to repetitive tool calls or other model behavior. The request has been halted.`,
|
||||
},
|
||||
Date.now(),
|
||||
);
|
||||
}
|
||||
},
|
||||
});
|
||||
}
|
||||
} catch (error: unknown) {
|
||||
spanMetadata.error = error;
|
||||
if (error instanceof UnauthorizedError) {
|
||||
onAuthError('Session expired or is unauthorized.');
|
||||
} else if (!isNodeError(error) || error.name !== 'AbortError') {
|
||||
addItem(
|
||||
{
|
||||
type: MessageType.ERROR,
|
||||
text: parseAndFormatApiError(
|
||||
getErrorMessage(error) || 'Unknown error',
|
||||
config.getContentGeneratorConfig()?.authType,
|
||||
undefined,
|
||||
config.getModel(),
|
||||
DEFAULT_GEMINI_FLASH_MODEL,
|
||||
),
|
||||
},
|
||||
userMessageTimestamp,
|
||||
);
|
||||
}
|
||||
} finally {
|
||||
if (activeQueryIdRef.current === queryId) {
|
||||
setIsResponding(false);
|
||||
}
|
||||
}
|
||||
});
|
||||
},
|
||||
),
|
||||
[
|
||||
streamingState,
|
||||
setModelSwitchedFromQuotaError,
|
||||
|
||||
Reference in New Issue
Block a user