From 8364af1e482741df7864ff0055c196244b12cc37 Mon Sep 17 00:00:00 2001 From: Alex Newman Date: Thu, 25 Dec 2025 19:47:41 -0500 Subject: [PATCH] feat(gemini): update Gemini model to 2.5 versions and add billing toggle in settings --- docs/public/configuration.mdx | 2 +- docs/public/usage/gemini-provider.mdx | 1 - .../worker/http/routes/SettingsRoutes.ts | 4 +- .../components/ContextSettingsModal.tsx | 3 +- src/ui/viewer/types.ts | 2 +- tests/gemini_agent.test.ts | 298 ++++++++++++++++++ 6 files changed, 304 insertions(+), 6 deletions(-) create mode 100644 tests/gemini_agent.test.ts diff --git a/docs/public/configuration.mdx b/docs/public/configuration.mdx index cd8fe8d5..be0fc8eb 100644 --- a/docs/public/configuration.mdx +++ b/docs/public/configuration.mdx @@ -25,7 +25,7 @@ Settings are managed in `~/.claude-mem/settings.json`. The file is auto-created | Setting | Default | Description | |-------------------------------|---------------------------------|---------------------------------------| | `CLAUDE_MEM_GEMINI_API_KEY` | — | Gemini API key ([get free key](https://aistudio.google.com/app/apikey)) | -| `CLAUDE_MEM_GEMINI_MODEL` | `gemini-2.0-flash-exp` | Gemini model: `gemini-2.0-flash-exp`, `gemini-1.5-flash`, `gemini-1.5-pro` | +| `CLAUDE_MEM_GEMINI_MODEL` | `gemini-2.5-flash-lite` | Gemini model: `gemini-2.5-flash-lite`, `gemini-2.5-flash`, `gemini-3-flash` | See [Gemini Provider](usage/gemini-provider) for detailed configuration and free tier information. diff --git a/docs/public/usage/gemini-provider.mdx b/docs/public/usage/gemini-provider.mdx index f47c6e3f..40dd29c2 100644 --- a/docs/public/usage/gemini-provider.mdx +++ b/docs/public/usage/gemini-provider.mdx @@ -162,7 +162,6 @@ If you hit rate limits: ### Observation Quality If observations seem lower quality with Gemini: -- Try `gemini-1.5-pro` for more capable extraction - Note that Claude typically produces slightly higher quality observations - Consider using Gemini for cost savings and Claude for important projects diff --git a/src/services/worker/http/routes/SettingsRoutes.ts b/src/services/worker/http/routes/SettingsRoutes.ts index fbcd894a..ef6e6a52 100644 --- a/src/services/worker/http/routes/SettingsRoutes.ts +++ b/src/services/worker/http/routes/SettingsRoutes.ts @@ -224,9 +224,9 @@ export class SettingsRoutes extends BaseRouteHandler { // Validate CLAUDE_MEM_GEMINI_MODEL if (settings.CLAUDE_MEM_GEMINI_MODEL) { - const validGeminiModels = ['gemini-2.0-flash-exp', 'gemini-1.5-flash', 'gemini-1.5-pro']; + const validGeminiModels = ['gemini-2.5-flash-lite', 'gemini-2.5-flash', 'gemini-3-flash']; if (!validGeminiModels.includes(settings.CLAUDE_MEM_GEMINI_MODEL)) { - return { valid: false, error: 'CLAUDE_MEM_GEMINI_MODEL must be one of: gemini-2.0-flash-exp, gemini-1.5-flash, gemini-1.5-pro' }; + return { valid: false, error: 'CLAUDE_MEM_GEMINI_MODEL must be one of: gemini-2.5-flash-lite, gemini-2.5-flash, gemini-3-flash' }; } } diff --git a/src/ui/viewer/components/ContextSettingsModal.tsx b/src/ui/viewer/components/ContextSettingsModal.tsx index 3ba8050a..26575272 100644 --- a/src/ui/viewer/components/ContextSettingsModal.tsx +++ b/src/ui/viewer/components/ContextSettingsModal.tsx @@ -481,8 +481,9 @@ export function ContextSettingsModal({
updateSetting('CLAUDE_MEM_GEMINI_BILLING_ENABLED', checked ? 'true' : 'false')} /> diff --git a/src/ui/viewer/types.ts b/src/ui/viewer/types.ts index aa703f56..f60e5cb2 100644 --- a/src/ui/viewer/types.ts +++ b/src/ui/viewer/types.ts @@ -63,7 +63,7 @@ export interface Settings { // AI Provider Configuration CLAUDE_MEM_PROVIDER?: string; // 'claude' | 'gemini' CLAUDE_MEM_GEMINI_API_KEY?: string; - CLAUDE_MEM_GEMINI_MODEL?: string; // 'gemini-2.0-flash-exp' | 'gemini-1.5-flash' | 'gemini-1.5-pro' + CLAUDE_MEM_GEMINI_MODEL?: string; // 'gemini-2.5-flash-lite' | 'gemini-2.5-flash' | 'gemini-3-flash' // Token Economics Display CLAUDE_MEM_CONTEXT_SHOW_READ_TOKENS?: string; diff --git a/tests/gemini_agent.test.ts b/tests/gemini_agent.test.ts new file mode 100644 index 00000000..8793d8a7 --- /dev/null +++ b/tests/gemini_agent.test.ts @@ -0,0 +1,298 @@ +import { describe, it, expect, beforeEach, afterEach, spyOn, mock } from 'bun:test'; +import { GeminiAgent } from '../src/services/worker/GeminiAgent'; +import { DatabaseManager } from '../src/services/worker/DatabaseManager'; +import { SessionManager } from '../src/services/worker/SessionManager'; +import { ModeManager } from '../src/services/worker/domain/ModeManager'; +import { SettingsDefaultsManager } from '../src/shared/SettingsDefaultsManager'; + +let billingEnabled = 'true'; + +// Mock SettingsDefaultsManager +mock.module('../src/shared/SettingsDefaultsManager', () => ({ + SettingsDefaultsManager: { + loadFromFile: () => ({ + CLAUDE_MEM_GEMINI_API_KEY: 'test-api-key', + CLAUDE_MEM_GEMINI_MODEL: 'gemini-2.5-flash-lite', + CLAUDE_MEM_GEMINI_BILLING_ENABLED: billingEnabled + }), + get: (key: string) => { + if (key === 'CLAUDE_MEM_LOG_LEVEL') return 'INFO'; + return ''; + } + } +})); + +// Mock ModeManager +const mockMode = { + name: 'code', + prompts: { + init: 'init prompt', + observation: 'obs prompt', + summary: 'summary prompt' + }, + observation_types: [{ id: 'discovery' }, { id: 'bugfix' }], + observation_concepts: [] +}; + +mock.module('../src/services/domain/ModeManager', () => ({ + ModeManager: { + getInstance: () => ({ + getActiveMode: () => mockMode + }) + } +})); + +describe('GeminiAgent', () => { + let agent: GeminiAgent; + let originalFetch: typeof global.fetch; + + // Mocks + let mockStoreObservation: any; + let mockStoreSummary: any; + let mockMarkSessionCompleted: any; + let mockSyncObservation: any; + let mockSyncSummary: any; + let mockMarkProcessed: any; + let mockCleanupProcessed: any; + let mockResetStuckMessages: any; + let mockDbManager: DatabaseManager; + let mockSessionManager: SessionManager; + + beforeEach(() => { + // Reset billing for each test default + billingEnabled = 'true'; + + // Initialize mocks + mockStoreObservation = mock(() => ({ id: 1, createdAtEpoch: Date.now() })); + mockStoreSummary = mock(() => ({ id: 1, createdAtEpoch: Date.now() })); + mockMarkSessionCompleted = mock(() => {}); + mockSyncObservation = mock(() => Promise.resolve()); + mockSyncSummary = mock(() => Promise.resolve()); + mockMarkProcessed = mock(() => {}); + mockCleanupProcessed = mock(() => 0); + mockResetStuckMessages = mock(() => 0); + + const mockSessionStore = { + storeObservation: mockStoreObservation, + storeSummary: mockStoreSummary, + markSessionCompleted: mockMarkSessionCompleted + }; + + const mockChromaSync = { + syncObservation: mockSyncObservation, + syncSummary: mockSyncSummary + }; + + mockDbManager = { + getSessionStore: () => mockSessionStore, + getChromaSync: () => mockChromaSync + } as unknown as DatabaseManager; + + const mockPendingMessageStore = { + markProcessed: mockMarkProcessed, + cleanupProcessed: mockCleanupProcessed, + resetStuckMessages: mockResetStuckMessages + }; + + mockSessionManager = { + getMessageIterator: async function* () { yield* []; }, + getPendingMessageStore: () => mockPendingMessageStore + } as unknown as SessionManager; + + agent = new GeminiAgent(mockDbManager, mockSessionManager); + originalFetch = global.fetch; + }); + + afterEach(() => { + global.fetch = originalFetch; + mock.restore(); + }); + + it('should initialize with correct config', async () => { + const session = { + sessionDbId: 1, + claudeSessionId: 'test-session', + sdkSessionId: 'test-sdk', + project: 'test-project', + userPrompt: 'test prompt', + conversationHistory: [], + lastPromptNumber: 1, + cumulativeInputTokens: 0, + cumulativeOutputTokens: 0, + pendingProcessingIds: new Set(), + startTime: Date.now() + } as any; + + global.fetch = mock(() => Promise.resolve(new Response(JSON.stringify({ + candidates: [{ + content: { + parts: [{ text: 'discoveryTest' }] + } + }], + usageMetadata: { totalTokenCount: 100 } + })))); + + await agent.startSession(session); + + expect(global.fetch).toHaveBeenCalledTimes(1); + const url = (global.fetch as any).mock.calls[0][0]; + expect(url).toContain('https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent'); + expect(url).toContain('key=test-api-key'); + }); + + it('should handle multi-turn conversation', async () => { + const session = { + sessionDbId: 1, + claudeSessionId: 'test-session', + sdkSessionId: 'test-sdk', + project: 'test-project', + userPrompt: 'test prompt', + conversationHistory: [{ role: 'user', content: 'prev context' }, { role: 'assistant', content: 'prev response' }], + lastPromptNumber: 2, + cumulativeInputTokens: 0, + cumulativeOutputTokens: 0, + pendingProcessingIds: new Set(), + startTime: Date.now() + } as any; + + global.fetch = mock(() => Promise.resolve(new Response(JSON.stringify({ + candidates: [{ content: { parts: [{ text: 'response' }] } }] + })))); + + await agent.startSession(session); + + const body = JSON.parse((global.fetch as any).mock.calls[0][1].body); + expect(body.contents).toHaveLength(3); + expect(body.contents[0].role).toBe('user'); + expect(body.contents[1].role).toBe('model'); + expect(body.contents[2].role).toBe('user'); + }); + + it('should process observations and store them', async () => { + const session = { + sessionDbId: 1, + claudeSessionId: 'test-session', + sdkSessionId: 'test-sdk', + project: 'test-project', + userPrompt: 'test prompt', + conversationHistory: [], + lastPromptNumber: 1, + cumulativeInputTokens: 0, + cumulativeOutputTokens: 0, + pendingProcessingIds: new Set(), + startTime: Date.now() + } as any; + + const observationXml = ` + + discovery + Found bug + Null pointer + Found a null pointer in the code + Null check missing + bug + src/main.ts + + + `; + + global.fetch = mock(() => Promise.resolve(new Response(JSON.stringify({ + candidates: [{ content: { parts: [{ text: observationXml }] } }], + usageMetadata: { totalTokenCount: 50 } + })))); + + await agent.startSession(session); + + expect(mockStoreObservation).toHaveBeenCalled(); + expect(mockSyncObservation).toHaveBeenCalled(); + expect(session.cumulativeInputTokens).toBeGreaterThan(0); + }); + + it('should fallback to Claude on rate limit error', async () => { + const session = { + sessionDbId: 1, + claudeSessionId: 'test-session', + sdkSessionId: 'test-sdk', + project: 'test-project', + userPrompt: 'test prompt', + conversationHistory: [], + lastPromptNumber: 1, + cumulativeInputTokens: 0, + cumulativeOutputTokens: 0, + pendingProcessingIds: new Set(), + startTime: Date.now() + } as any; + + global.fetch = mock(() => Promise.resolve(new Response('Resource has been exhausted (e.g. check quota).', { status: 429 }))); + + const fallbackAgent = { + startSession: mock(() => Promise.resolve()) + }; + agent.setFallbackAgent(fallbackAgent); + + await agent.startSession(session); + + expect(fallbackAgent.startSession).toHaveBeenCalledWith(session, undefined); + expect(mockResetStuckMessages).toHaveBeenCalled(); + }); + + it('should NOT fallback on other errors', async () => { + const session = { + sessionDbId: 1, + claudeSessionId: 'test-session', + sdkSessionId: 'test-sdk', + project: 'test-project', + userPrompt: 'test prompt', + conversationHistory: [], + lastPromptNumber: 1, + cumulativeInputTokens: 0, + cumulativeOutputTokens: 0, + pendingProcessingIds: new Set(), + startTime: Date.now() + } as any; + + global.fetch = mock(() => Promise.resolve(new Response('Invalid argument', { status: 400 }))); + + const fallbackAgent = { + startSession: mock(() => Promise.resolve()) + }; + agent.setFallbackAgent(fallbackAgent); + + expect(agent.startSession(session)).rejects.toThrow('Gemini API error: 400 - Invalid argument'); + expect(fallbackAgent.startSession).not.toHaveBeenCalled(); + }); + + it('should respect rate limits when billing disabled', async () => { + billingEnabled = 'false'; + const originalSetTimeout = global.setTimeout; + const mockSetTimeout = mock((cb: any) => cb()); + global.setTimeout = mockSetTimeout as any; + + try { + const session = { + sessionDbId: 1, + claudeSessionId: 'test-session', + sdkSessionId: 'test-sdk', + project: 'test-project', + userPrompt: 'test prompt', + conversationHistory: [], + lastPromptNumber: 1, + cumulativeInputTokens: 0, + cumulativeOutputTokens: 0, + pendingProcessingIds: new Set(), + startTime: Date.now() + } as any; + + global.fetch = mock(() => Promise.resolve(new Response(JSON.stringify({ + candidates: [{ content: { parts: [{ text: 'ok' }] } }] + })))); + + await agent.startSession(session); + await agent.startSession(session); + + expect(mockSetTimeout).toHaveBeenCalled(); + } finally { + global.setTimeout = originalSetTimeout; + } + }); +}); \ No newline at end of file