summaryrefslogtreecommitdiff
path: root/packages/cli/src/ui/hooks
diff options
context:
space:
mode:
Diffstat (limited to 'packages/cli/src/ui/hooks')
-rw-r--r--packages/cli/src/ui/hooks/useGeminiStream.test.tsx16
-rw-r--r--packages/cli/src/ui/hooks/useGeminiStream.ts16
2 files changed, 32 insertions, 0 deletions
diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx
index fc6f93c5..62ade50f 100644
--- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx
+++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx
@@ -301,6 +301,8 @@ describe('useGeminiStream', () => {
getUsageStatisticsEnabled: () => true,
getDebugMode: () => false,
addHistory: vi.fn(),
+ setQuotaErrorOccurred: vi.fn(),
+ getQuotaErrorOccurred: vi.fn(() => false),
} as unknown as Config;
mockOnDebugMessage = vi.fn();
mockHandleSlashCommand = vi.fn().mockResolvedValue(false);
@@ -386,6 +388,8 @@ describe('useGeminiStream', () => {
() => 'vscode' as EditorType,
() => {},
() => Promise.resolve(),
+ false,
+ () => {},
);
},
{
@@ -518,6 +522,8 @@ describe('useGeminiStream', () => {
() => 'vscode' as EditorType,
() => {},
() => Promise.resolve(),
+ false,
+ () => {},
),
);
@@ -582,6 +588,8 @@ describe('useGeminiStream', () => {
() => 'vscode' as EditorType,
() => {},
() => Promise.resolve(),
+ false,
+ () => {},
),
);
@@ -675,6 +683,8 @@ describe('useGeminiStream', () => {
() => 'vscode' as EditorType,
() => {},
() => Promise.resolve(),
+ false,
+ () => {},
),
);
@@ -775,6 +785,8 @@ describe('useGeminiStream', () => {
() => 'vscode' as EditorType,
() => {},
() => Promise.resolve(),
+ false,
+ () => {},
),
);
@@ -1063,6 +1075,8 @@ describe('useGeminiStream', () => {
() => 'vscode' as EditorType,
() => {},
mockPerformMemoryRefresh,
+ false,
+ () => {},
),
);
@@ -1113,6 +1127,8 @@ describe('useGeminiStream', () => {
() => 'vscode' as EditorType,
() => {},
() => Promise.resolve(),
+ false,
+ () => {},
),
);
diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts
index 550cab86..d32c9ffa 100644
--- a/packages/cli/src/ui/hooks/useGeminiStream.ts
+++ b/packages/cli/src/ui/hooks/useGeminiStream.ts
@@ -90,6 +90,8 @@ export const useGeminiStream = (
getPreferredEditor: () => EditorType | undefined,
onAuthError: () => void,
performMemoryRefresh: () => Promise<void>,
+ modelSwitchedFromQuotaError: boolean,
+ setModelSwitchedFromQuotaError: React.Dispatch<React.SetStateAction<boolean>>,
) => {
const [initError, setInitError] = useState<string | null>(null);
const abortControllerRef = useRef<AbortController | null>(null);
@@ -494,6 +496,12 @@ export const useGeminiStream = (
const userMessageTimestamp = Date.now();
setShowHelp(false);
+ // 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;
@@ -552,6 +560,7 @@ export const useGeminiStream = (
[
streamingState,
setShowHelp,
+ setModelSwitchedFromQuotaError,
prepareQueryForGemini,
processGeminiStreamEvents,
pendingHistoryItemRef,
@@ -668,6 +677,12 @@ export const useGeminiStream = (
);
markToolsAsSubmitted(callIdsToMarkAsSubmitted);
+
+ // Don't continue if model was switched due to quota error
+ if (modelSwitchedFromQuotaError) {
+ return;
+ }
+
submitQuery(mergePartListUnions(responsesToSend), {
isContinuation: true,
});
@@ -678,6 +693,7 @@ export const useGeminiStream = (
markToolsAsSubmitted,
geminiClient,
performMemoryRefresh,
+ modelSwitchedFromQuotaError,
],
);