Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 0 additions & 20 deletions packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -166,11 +166,6 @@
const intent = await promise!;
expect(intent).toBe('retry_always');

// Verify activateFallbackMode was called
expect(mockConfig.activateFallbackMode).toHaveBeenCalledWith(
'gemini-flash',
);

// The pending request should be cleared from the state
expect(result.current.proQuotaRequest).toBeNull();
expect(mockHistoryManager.addItem).toHaveBeenCalledTimes(1);
Expand Down Expand Up @@ -282,11 +277,6 @@
const intent = await promise!;
expect(intent).toBe('retry_always');

// Verify activateFallbackMode was called
expect(mockConfig.activateFallbackMode).toHaveBeenCalledWith(
'model-B',
);

// The pending request should be cleared from the state
expect(result.current.proQuotaRequest).toBeNull();
expect(mockConfig.setQuotaErrorOccurred).toHaveBeenCalledWith(true);
Expand Down Expand Up @@ -317,21 +307,21 @@
const error = new ModelNotFoundError('model not found', 404);

act(() => {
promise = handler('gemini-3-pro-preview', 'gemini-2.5-pro', error);

Check warning on line 310 in packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts

View workflow job for this annotation

GitHub Actions / Lint

Found sensitive keyword "gemini-3". Please make sure this change is appropriate to submit.
});

// The hook should now have a pending request for the UI to handle
const request = result.current.proQuotaRequest;
expect(request).not.toBeNull();
expect(request?.failedModel).toBe('gemini-3-pro-preview');

Check warning on line 316 in packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts

View workflow job for this annotation

GitHub Actions / Lint

Found sensitive keyword "gemini-3". Please make sure this change is appropriate to submit.
expect(request?.isTerminalQuotaError).toBe(false);
expect(request?.isModelNotFoundError).toBe(true);

const message = request!.message;
expect(message).toBe(
`It seems like you don't have access to gemini-3-pro-preview.

Check warning on line 322 in packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts

View workflow job for this annotation

GitHub Actions / Lint

Found sensitive keyword "gemini-3". Please make sure this change is appropriate to submit.
Learn more at https://goo.gle/enable-preview-features
To disable gemini-3-pro-preview, disable "Preview features" in /settings.`,

Check warning on line 324 in packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts

View workflow job for this annotation

GitHub Actions / Lint

Found sensitive keyword "gemini-3". Please make sure this change is appropriate to submit.
);

// Simulate the user choosing to switch
Expand All @@ -342,11 +332,6 @@
const intent = await promise!;
expect(intent).toBe('retry_always');

// Verify activateFallbackMode was called
expect(mockConfig.activateFallbackMode).toHaveBeenCalledWith(
'gemini-2.5-pro',
);

expect(result.current.proQuotaRequest).toBeNull();
});
});
Expand Down Expand Up @@ -430,11 +415,6 @@
expect(intent).toBe('retry_always');
expect(result.current.proQuotaRequest).toBeNull();

// Verify activateFallbackMode was called
expect(mockConfig.activateFallbackMode).toHaveBeenCalledWith(
'gemini-flash',
);

// Verify quota error flags are reset
expect(mockSetModelSwitchedFromQuotaError).toHaveBeenCalledWith(false);
expect(mockConfig.setQuotaErrorOccurred).toHaveBeenCalledWith(false);
Expand Down
4 changes: 0 additions & 4 deletions packages/cli/src/ui/hooks/useQuotaAndFallback.ts
Original file line number Diff line number Diff line change
Expand Up @@ -135,10 +135,6 @@ export function useQuotaAndFallback({
config.setQuotaErrorOccurred(false);

if (choice === 'retry_always') {
// Set the model to the fallback model for the current session.
// This ensures the Footer updates and future turns use this model.
// The change is not persisted, so the original model is restored on restart.
config.activateFallbackMode(proQuotaRequest.fallbackModel);
historyManager.addItem(
{
type: MessageType.INFO,
Expand Down
5 changes: 2 additions & 3 deletions packages/core/src/config/config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2029,9 +2029,8 @@ export class Config {
*/
async dispose(): Promise<void> {
coreEvents.off(CoreEvent.AgentsRefreshed, this.onAgentsRefreshed);
if (this.agentRegistry) {
this.agentRegistry.dispose();
}
this.agentRegistry?.dispose();
this.geminiClient?.dispose();
if (this.mcpClientManager) {
await this.mcpClientManager.stop();
}
Expand Down
51 changes: 51 additions & 0 deletions packages/core/src/core/client.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ import type {
import { ClearcutLogger } from '../telemetry/clearcut-logger/clearcut-logger.js';
import * as policyCatalog from '../availability/policyCatalog.js';
import { partToString } from '../utils/partUtils.js';
import { coreEvents } from '../utils/events.js';

// Mock fs module to prevent actual file system operations during tests
const mockFileSystem = new Map<string, string>();
Expand Down Expand Up @@ -281,6 +282,7 @@ describe('Gemini Client (client.ts)', () => {
});

afterEach(() => {
client.dispose();
vi.restoreAllMocks();
});

Expand Down Expand Up @@ -1757,6 +1759,55 @@ ${JSON.stringify(
expect.any(AbortSignal),
);
});

it('should re-route within the same prompt when the configured model changes', async () => {
mockTurnRunFn.mockClear();
mockTurnRunFn.mockImplementation(async function* () {
yield { type: 'content', value: 'Hello' };
});

mockRouterService.route.mockResolvedValueOnce({
model: 'original-model',
reason: 'test',
});

let stream = client.sendMessageStream(
[{ text: 'Hi' }],
new AbortController().signal,
'prompt-1',
);
await fromAsync(stream);

expect(mockRouterService.route).toHaveBeenCalledTimes(1);
expect(mockTurnRunFn).toHaveBeenNthCalledWith(
1,
{ model: 'original-model' },
[{ text: 'Hi' }],
expect.any(AbortSignal),
);

mockRouterService.route.mockResolvedValue({
model: 'fallback-model',
reason: 'test',
});
vi.mocked(mockConfig.getModel).mockReturnValue('gemini-2.5-flash');
coreEvents.emitModelChanged('gemini-2.5-flash');

stream = client.sendMessageStream(
[{ text: 'Continue' }],
new AbortController().signal,
'prompt-1',
);
await fromAsync(stream);

expect(mockRouterService.route).toHaveBeenCalledTimes(2);
expect(mockTurnRunFn).toHaveBeenNthCalledWith(
2,
{ model: 'fallback-model' },
[{ text: 'Continue' }],
expect.any(AbortSignal),
);
});
});

it('should use getGlobalMemory for system instruction when JIT is enabled', async () => {
Expand Down
11 changes: 11 additions & 0 deletions packages/core/src/core/client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@ import {
import { resolveModel } from '../config/models.js';
import type { RetryAvailabilityContext } from '../utils/retry.js';
import { partToString } from '../utils/partUtils.js';
import { coreEvents, CoreEvent } from '../utils/events.js';

const MAX_TURNS = 100;

Expand Down Expand Up @@ -94,8 +95,14 @@ export class GeminiClient {
this.loopDetector = new LoopDetectionService(config);
this.compressionService = new ChatCompressionService();
this.lastPromptId = this.config.getSessionId();

coreEvents.on(CoreEvent.ModelChanged, this.handleModelChanged);
}

private handleModelChanged = () => {
this.currentSequenceModel = null;
};

// Hook state to deduplicate BeforeAgent calls and track response for
// AfterAgent
private hookStateMap = new Map<
Expand Down Expand Up @@ -253,6 +260,10 @@ export class GeminiClient {
this.updateTelemetryTokenCount();
}

dispose() {
coreEvents.off(CoreEvent.ModelChanged, this.handleModelChanged);
}

async resumeChat(
history: Content[],
resumedSessionData?: ResumedSessionData,
Expand Down
20 changes: 12 additions & 8 deletions packages/core/src/fallback/handler.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,8 @@ const createMockConfig = (overrides: Partial<Config> = {}): Config =>
fallbackHandler: undefined,
getFallbackModelHandler: vi.fn(),
setActiveModel: vi.fn(),
setModel: vi.fn(),
activateFallbackMode: vi.fn(),
getModelAvailabilityService: vi.fn(() =>
createAvailabilityServiceMock({
selectedModel: FALLBACK_MODEL,
Expand Down Expand Up @@ -198,7 +200,7 @@ describe('handleFallback', () => {

expect(result).toBe(true);
expect(policyConfig.getFallbackModelHandler).not.toHaveBeenCalled();
expect(policyConfig.setActiveModel).toHaveBeenCalledWith(
expect(policyConfig.activateFallbackMode).toHaveBeenCalledWith(
DEFAULT_GEMINI_FLASH_MODEL,
);
} finally {
Expand Down Expand Up @@ -273,7 +275,7 @@ describe('handleFallback', () => {
expect(openBrowserSecurely).toHaveBeenCalledWith(
'https://goo.gle/set-up-gemini-code-assist',
);
expect(policyConfig.setActiveModel).not.toHaveBeenCalled();
expect(policyConfig.activateFallbackMode).not.toHaveBeenCalled();
});

it('should catch errors from the handler, log an error, and return null', async () => {
Expand Down Expand Up @@ -378,7 +380,7 @@ describe('handleFallback', () => {
);
});

it('calls setActiveModel and logs telemetry when handler returns "retry_always"', async () => {
it('calls activateFallbackMode when handler returns "retry_always"', async () => {
policyHandler.mockResolvedValue('retry_always');
vi.mocked(policyConfig.getModel).mockReturnValue(
DEFAULT_GEMINI_MODEL_AUTO,
Expand All @@ -391,11 +393,13 @@ describe('handleFallback', () => {
);

expect(result).toBe(true);
expect(policyConfig.setActiveModel).toHaveBeenCalledWith(FALLBACK_MODEL);
expect(policyConfig.activateFallbackMode).toHaveBeenCalledWith(
FALLBACK_MODEL,
);
// TODO: add logging expect statement
});

it('does NOT call setActiveModel when handler returns "stop"', async () => {
it('does NOT call activateFallbackMode when handler returns "stop"', async () => {
policyHandler.mockResolvedValue('stop');

const result = await handleFallback(
Expand All @@ -405,11 +409,11 @@ describe('handleFallback', () => {
);

expect(result).toBe(false);
expect(policyConfig.setActiveModel).not.toHaveBeenCalled();
expect(policyConfig.activateFallbackMode).not.toHaveBeenCalled();
// TODO: add logging expect statement
});

it('does NOT call setActiveModel when handler returns "retry_once"', async () => {
it('does NOT call activateFallbackMode when handler returns "retry_once"', async () => {
policyHandler.mockResolvedValue('retry_once');

const result = await handleFallback(
Expand All @@ -419,7 +423,7 @@ describe('handleFallback', () => {
);

expect(result).toBe(true);
expect(policyConfig.setActiveModel).not.toHaveBeenCalled();
expect(policyConfig.activateFallbackMode).not.toHaveBeenCalled();
});
});
});
2 changes: 1 addition & 1 deletion packages/core/src/fallback/handler.ts
Original file line number Diff line number Diff line change
Expand Up @@ -131,7 +131,7 @@ async function processIntent(
case 'retry_always':
// TODO(telemetry): Implement generic fallback event logging. Existing
// logFlashFallback is specific to a single Model.
config.setActiveModel(fallbackModel);
config.activateFallbackMode(fallbackModel);
return true;

case 'retry_once':
Expand Down
Loading