import * as fs from 'fs'; import path from 'path'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import cliState from '../../../src/cliState'; import logger from '../../../src/logger'; import * as vertexUtil from '../../../src/providers/google/util'; import { VertexChatProvider } from '../../../src/providers/google/vertex'; import type { JSONClient } from 'google-auth-library/build/src/auth/googleauth'; // Hoisted mocks for cache const mockCacheGet = vi.hoisted(() => vi.fn()); const mockCacheSet = vi.hoisted(() => vi.fn()); const mockIsCacheEnabled = vi.hoisted(() => vi.fn()); // Hoisted mock for importModule const mockImportModule = vi.hoisted(() => vi.fn()); // Mock database vi.mock('libsql', () => { return vi.fn().mockReturnValue({ prepare: vi.fn(), transaction: vi.fn(), exec: vi.fn(), close: vi.fn(), }); }); vi.mock('../../../src/database', async (importOriginal) => { return { ...(await importOriginal()), getDb: vi.fn().mockReturnValue({ prepare: vi.fn(), transaction: vi.fn(), exec: vi.fn(), close: vi.fn(), }), }; }); vi.mock('csv-stringify/sync', async (importOriginal) => { return { ...(await importOriginal()), stringify: vi.fn().mockReturnValue('mocked,csv,output'), }; }); vi.mock('glob', async (importOriginal) => { return { ...(await importOriginal()), globSync: vi.fn().mockReturnValue([]), hasMagic: (path: string) => { // Match the real hasMagic behavior: only detect patterns in forward-slash paths // This mimics glob's actual behavior where backslash paths return false return /[*?[\]{}]/.test(path) && !path.includes('\\'); }, }; }); vi.mock('fs', async (importOriginal) => { return { ...(await importOriginal()), existsSync: vi.fn(), readFileSync: vi.fn(), writeFileSync: vi.fn(), statSync: vi.fn(), mkdirSync: vi.fn(), }; }); vi.mock('../../../src/cache', async (importOriginal) => { return { ...(await importOriginal()), getCache: vi.fn().mockImplementation(() => ({ get: mockCacheGet, set: mockCacheSet, wrap: vi.fn(), del: vi.fn(), reset: vi.fn(), store: {} as any, })), isCacheEnabled: mockIsCacheEnabled, }; }); vi.mock('../../../src/providers/google/util', async () => { const actual = await vi.importActual( '../../../src/providers/google/util', ); return { ...actual, getGoogleClient: vi.fn(), loadCredentials: vi.fn(), resolveProjectId: vi.fn(), }; }); // Mock GoogleAuthManager to prevent API key detection from environment vi.mock('../../../src/providers/google/auth', async () => { const actual = await vi.importActual( '../../../src/providers/google/auth', ); return { ...actual, GoogleAuthManager: { ...actual.GoogleAuthManager, // Return no API key by default so tests use OAuth mode getApiKey: vi.fn().mockReturnValue({ apiKey: undefined, source: 'none' }), determineVertexMode: vi.fn().mockReturnValue(true), validateAndWarn: vi.fn(), // Respect config.region when provided, otherwise default to us-central1 resolveRegion: vi.fn().mockImplementation((config?: { region?: string }) => { return config?.region || 'us-central1'; }), resolveProjectId: vi.fn().mockResolvedValue('test-project-id'), }, }; }); vi.mock('../../../src/esm', async (importOriginal) => { return { ...(await importOriginal()), importModule: mockImportModule, }; }); function mockVertexRequest(data: unknown) { const mockRequest = vi.fn().mockResolvedValue({ data }); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest, } as unknown as JSONClient, projectId: 'test-project-id', }); vi.spyOn(vertexUtil, 'loadCredentials').mockImplementation(function (creds) { if (typeof creds === 'object') { return JSON.stringify(creds); } return creds; }); vi.spyOn(vertexUtil, 'resolveProjectId').mockResolvedValue('test-project-id'); return mockRequest; } function expectHashedBodyCacheKeys(expectedPattern: RegExp, forbiddenValues: string[]) { expect(mockCacheGet).toHaveBeenCalledTimes(1); expect(mockCacheSet).toHaveBeenCalledTimes(1); const cacheGetKey = mockCacheGet.mock.calls[0][0] as string; const cacheSetKey = mockCacheSet.mock.calls[0][0] as string; expect(cacheGetKey).toBe(cacheSetKey); expect(cacheSetKey).toMatch(expectedPattern); for (const value of forbiddenValues) { expect(cacheSetKey).not.toContain(value); } return cacheSetKey; } describe('VertexChatProvider.callGeminiApi', () => { let provider: VertexChatProvider; beforeEach(() => { // Reset cache mocks to default state (no cached response) mockCacheGet.mockReset(); mockCacheGet.mockResolvedValue(null); mockCacheSet.mockReset(); mockImportModule.mockReset(); provider = new VertexChatProvider('gemini-pro', { config: { context: 'test-context', examples: [{ input: 'example input', output: 'example output' }], stopSequences: ['\n'], temperature: 0.7, maxOutputTokens: 100, topP: 0.9, topK: 40, }, }); mockIsCacheEnabled.mockReturnValue(true); }); afterEach(() => { vi.clearAllMocks(); }); it('should call the Gemini API and return the response', async () => { const mockResponse = { data: [ { candidates: [{ content: { parts: [{ text: 'response text' }] } }], usageMetadata: { totalTokenCount: 10, promptTokenCount: 5, candidatesTokenCount: 5, }, }, ], }; const mockRequest = vi.fn().mockResolvedValue(mockResponse); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest, } as unknown as JSONClient, projectId: 'test-project-id', }); vi.spyOn(vertexUtil, 'loadCredentials').mockImplementation(function (creds) { if (typeof creds === 'object') { return JSON.stringify(creds); } return creds; }); vi.spyOn(vertexUtil, 'resolveProjectId').mockResolvedValue('test-project-id'); const response = await provider.callGeminiApi('test prompt'); expect(response).toEqual({ cached: false, output: 'response text', tokenUsage: { total: 10, prompt: 5, completion: 5, }, cost: expect.closeTo(0.00001, 10), metadata: {}, }); expect(vertexUtil.getGoogleClient).toHaveBeenCalledWith({ credentials: undefined }); expect(mockRequest).toHaveBeenCalledWith({ url: expect.any(String), method: 'POST', data: expect.objectContaining({ contents: [{ parts: [{ text: 'test prompt' }], role: 'user' }], }), timeout: expect.any(Number), }); }); it('should return cached response if available', async () => { const mockCachedResponse = { cached: true, output: 'cached response text', tokenUsage: { total: 10, prompt: 5, completion: 5, }, }; mockCacheGet.mockResolvedValue(JSON.stringify(mockCachedResponse)); const response = await provider.callGeminiApi('test prompt'); expect(response).toEqual({ ...mockCachedResponse, tokenUsage: { ...mockCachedResponse.tokenUsage, cached: mockCachedResponse.tokenUsage.total, }, }); }); it('should handle API call errors', async () => { const mockError = new Error('something went wrong'); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: vi.fn().mockRejectedValue(mockError), } as unknown as JSONClient, projectId: 'test-project-id', }); const response = await provider.callGeminiApi('test prompt'); expect(response).toEqual({ error: `API call error: Error: something went wrong`, }); }); it('should handle API response errors', async () => { const mockResponse = { data: [ { error: { code: 400, message: 'Bad Request', }, }, ], }; vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: vi.fn().mockResolvedValue(mockResponse), } as unknown as JSONClient, projectId: 'test-project-id', }); const response = await provider.callGeminiApi('test prompt'); expect(response).toEqual({ error: 'Error 400: Bad Request', }); }); it('should handle function calling configuration', async () => { const tools = [ { functionDeclarations: [ { name: 'get_weather', description: 'Get weather information', parameters: { type: 'OBJECT' as const, properties: { location: { type: 'STRING' as const, description: 'City name', }, }, required: ['location'], }, }, ], }, ]; vi.spyOn(fs, 'existsSync').mockReturnValue(true); vi.spyOn(fs, 'readFileSync').mockReturnValue(JSON.stringify(tools)); provider = new VertexChatProvider('gemini-pro', { config: { toolConfig: { functionCallingConfig: { mode: 'AUTO', allowedFunctionNames: ['get_weather'], }, }, tools, }, }); const mockResponse = { data: [ { candidates: [ { content: { parts: [ { functionCall: { name: 'get_weather', args: { location: 'San Francisco' }, }, }, ], }, }, ], usageMetadata: { totalTokenCount: 15, promptTokenCount: 8, candidatesTokenCount: 7, }, }, ], }; const mockRequest = vi.fn().mockResolvedValue(mockResponse); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest, } as unknown as JSONClient, projectId: 'test-project-id', }); vi.spyOn(vertexUtil, 'loadCredentials').mockImplementation(function (creds) { if (typeof creds === 'object') { return JSON.stringify(creds); } return creds; }); vi.spyOn(vertexUtil, 'resolveProjectId').mockResolvedValue('test-project-id'); const response = await provider.callGeminiApi('What is the weather in San Francisco?'); expect(response).toEqual({ cached: false, output: [ { functionCall: { name: 'get_weather', args: { location: 'San Francisco' }, }, }, ], tokenUsage: { total: 15, prompt: 8, completion: 7, }, cost: expect.closeTo(0.0000145, 10), metadata: {}, }); expect(mockRequest).toHaveBeenCalledWith( expect.objectContaining({ data: expect.objectContaining({ toolConfig: { functionCallingConfig: { mode: 'AUTO', allowedFunctionNames: ['get_weather'], }, }, tools, }), }), ); }); it.each([ ['tool_choice', 'none' as const], ['toolConfig', { functionCallingConfig: { mode: 'NONE' as const } }], ['tool_config', { function_calling_config: { mode: 'none' as const } }], ])('should disable function calling via %s', async (key, value) => { const tools = [ { functionDeclarations: [ { name: 'get_weather', description: 'Get weather information', parameters: { type: 'OBJECT' as const, properties: {} }, }, ], }, { googleSearch: {} }, ]; provider = new VertexChatProvider('gemini-pro', { config: { tools, [key]: value, } as any, }); const mockResponse = { data: [ { candidates: [{ content: { parts: [{ text: 'no tools used' }] } }], usageMetadata: { totalTokenCount: 10, promptTokenCount: 5, candidatesTokenCount: 5 }, }, ], }; const mockRequest = vi.fn().mockResolvedValue(mockResponse); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest } as unknown as JSONClient, projectId: 'test-project-id', }); vi.spyOn(vertexUtil, 'loadCredentials').mockImplementation((c) => c as any); vi.spyOn(vertexUtil, 'resolveProjectId').mockResolvedValue('test-project-id'); await provider.callGeminiApi('hi'); expect(mockRequest).toHaveBeenCalledWith( expect.objectContaining({ data: expect.objectContaining({ toolConfig: { functionCallingConfig: { mode: 'NONE' } }, tools: [{ googleSearch: {} }], }), }), ); }); it('should not invoke functionToolCallbacks when tools are disabled', async () => { // Use a cached response so executeFunctionCallback would normally fire if not gated. mockIsCacheEnabled.mockReturnValue(true); mockCacheGet.mockResolvedValue( JSON.stringify({ cached: true, output: JSON.stringify({ functionCall: { name: 'should_not_run', args: '{}' }, }), tokenUsage: { total: 5, prompt: 3, completion: 2 }, }), ); const callback = vi.fn().mockResolvedValue('should not be called'); const provider = new VertexChatProvider('gemini-pro', { config: { tool_choice: 'none', tools: [ { functionDeclarations: [ { name: 'should_not_run', description: 'Test', parameters: { type: 'OBJECT', properties: {} }, }, ], }, ], functionToolCallbacks: { should_not_run: callback }, }, }); const result = await provider.callApi('test'); expect(callback).not.toHaveBeenCalled(); // Output remains the original functionCall envelope when callbacks are gated. expect(result.output).toBe( JSON.stringify({ functionCall: { name: 'should_not_run', args: '{}' } }), ); // Reset cache state so other tests aren't affected. mockIsCacheEnabled.mockReturnValue(false); mockCacheGet.mockReset(); }); it('should skip executable tool files while preserving inline non-function tools when disabled', async () => { provider = new VertexChatProvider('gemini-pro', { config: { tool_choice: 'none', tools: [{ googleSearch: {} }, 'file://tools.js:getTools'] as any, }, }); const mockResponse = { data: [ { candidates: [{ content: { parts: [{ text: 'no tools used' }] } }], usageMetadata: { totalTokenCount: 10, promptTokenCount: 5, candidatesTokenCount: 5 }, }, ], }; const mockRequest = vi.fn().mockResolvedValue(mockResponse); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest } as unknown as JSONClient, projectId: 'test-project-id', }); vi.spyOn(vertexUtil, 'loadCredentials').mockImplementation((c) => c as any); vi.spyOn(vertexUtil, 'resolveProjectId').mockResolvedValue('test-project-id'); await provider.callGeminiApi('hi'); expect(mockImportModule).not.toHaveBeenCalled(); expect(mockRequest).toHaveBeenCalledWith( expect.objectContaining({ data: expect.objectContaining({ toolConfig: { functionCallingConfig: { mode: 'NONE' } }, tools: [{ googleSearch: {} }], }), }), ); }); it('should load tools from external file and render variables', async () => { const mockExternalTools = [ { functionDeclarations: [ { name: 'get_weather', description: 'Get weather in San Francisco', parameters: { type: 'OBJECT' as const, properties: { location: { type: 'STRING' as const }, }, }, }, ], }, ]; // Mock file system operations (existsSync no longer called due to TOCTOU fix) vi.spyOn(fs, 'readFileSync').mockReturnValue(JSON.stringify(mockExternalTools)); provider = new VertexChatProvider('gemini-pro', { config: { tools: 'file://tools.json' as any, }, }); const mockResponse = { data: [ { candidates: [ { content: { parts: [{ text: 'response with tools' }], }, }, ], usageMetadata: { totalTokenCount: 10, promptTokenCount: 5, candidatesTokenCount: 5, }, }, ], }; const mockRequest = vi.fn().mockResolvedValue(mockResponse); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest, } as unknown as JSONClient, projectId: 'test-project-id', }); vi.spyOn(vertexUtil, 'loadCredentials').mockImplementation(function (creds) { if (typeof creds === 'object') { return JSON.stringify(creds); } return creds; }); vi.spyOn(vertexUtil, 'resolveProjectId').mockResolvedValue('test-project-id'); const response = await provider.callGeminiApi('test prompt', { vars: { location: 'San Francisco' }, prompt: { raw: 'test prompt', label: 'test' }, }); expect(response).toEqual({ cached: false, output: 'response with tools', tokenUsage: { total: 10, prompt: 5, completion: 5, }, cost: expect.closeTo(0.00001, 10), metadata: {}, }); // Note: existsSync no longer called - we use try/catch on readFileSync instead (TOCTOU fix) expect(fs.readFileSync).toHaveBeenCalledWith(expect.stringContaining('tools.json'), 'utf8'); expect(mockRequest).toHaveBeenCalledWith( expect.objectContaining({ data: expect.objectContaining({ tools: mockExternalTools, }), }), ); }); it('should use model name in cache key', async () => { const mockResponse = { data: [ { candidates: [ { content: { parts: [{ text: 'response text' }], }, }, ], usageMetadata: { totalTokenCount: 10, promptTokenCount: 5, candidatesTokenCount: 5, }, }, ], }; const mockRequest = vi.fn().mockResolvedValue(mockResponse); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest, } as unknown as JSONClient, projectId: 'test-project-id', }); vi.spyOn(vertexUtil, 'loadCredentials').mockImplementation(function (creds) { if (typeof creds === 'object') { return JSON.stringify(creds); } return creds; }); vi.spyOn(vertexUtil, 'resolveProjectId').mockResolvedValue('test-project-id'); provider = new VertexChatProvider('gemini-2.0-flash-001'); await provider.callGeminiApi('test prompt'); expectHashedBodyCacheKeys(/^vertex:gemini-2\.0-flash-001:[a-f0-9]{64}$/, ['test prompt']); expect(mockCacheSet).toHaveBeenCalledWith( expect.stringContaining('vertex:gemini-2.0-flash-001:'), expect.any(String), ); }); it('should not reuse cached responses across effective API hosts', async () => { const cachedResponses = new Map(); mockCacheGet.mockImplementation(async (key: string) => cachedResponses.get(key) ?? null); mockCacheSet.mockImplementation(async (key: string, value: string) => { cachedResponses.set(key, value); }); const mockRequest = mockVertexRequest({ candidates: [{ content: { parts: [{ text: 'response text' }] } }], usageMetadata: { totalTokenCount: 10, promptTokenCount: 5, candidatesTokenCount: 5, }, }); const publicProvider = new VertexChatProvider('gemini-pro', { config: { region: 'global' }, }); const proxyProvider = new VertexChatProvider('gemini-pro', { config: { region: 'global' }, env: { VERTEX_API_HOST: 'vertex-proxy.example.test' }, }); await publicProvider.callGeminiApi('same prompt'); await proxyProvider.callGeminiApi('same prompt'); const [publicCacheKey, proxyCacheKey] = mockCacheGet.mock.calls.map(([key]) => key as string); expect(publicCacheKey).not.toBe(proxyCacheKey); expect(mockRequest).toHaveBeenCalledTimes(2); expect(mockRequest.mock.calls.map(([request]) => request.url)).toEqual([ expect.stringContaining('https://aiplatform.googleapis.com/'), expect.stringContaining('https://vertex-proxy.example.test/'), ]); }); it('should handle function tool callbacks correctly', async () => { const mockCachedResponse = { cached: true, output: JSON.stringify({ functionCall: { name: 'get_weather', args: '{"location":"New York"}', }, }), tokenUsage: { total: 15, prompt: 10, completion: 5, }, cost: 0.00045, metadata: { groundingMetadata: { test: true, }, }, }; mockCacheGet.mockResolvedValue(JSON.stringify(mockCachedResponse)); const mockWeatherFunction = vi.fn().mockResolvedValue('Sunny, 25°C'); const provider = new VertexChatProvider('gemini', { config: { tools: [ { functionDeclarations: [ { name: 'get_weather', description: 'Get the weather for a location', parameters: { type: 'OBJECT', properties: { location: { type: 'STRING' }, }, required: ['location'], }, }, ], }, ], functionToolCallbacks: { get_weather: mockWeatherFunction, }, }, }); const result = await provider.callApi( JSON.stringify([{ role: 'user', content: "What's the weather in New York?" }]), ); expect(mockCacheGet).toHaveBeenCalledTimes(1); expect(mockWeatherFunction).toHaveBeenCalledWith('{"location":"New York"}'); expect(result.output).toBe('Sunny, 25°C'); expect(result.tokenUsage).toEqual({ total: 15, prompt: 10, completion: 5, cached: 15 }); expect(result.cost).toBe(0.00045); expect(result.metadata).toEqual({ groundingMetadata: { test: true, }, }); }); it('should return undefined cost when Gemini omits usage metadata', async () => { const mockResponse = { data: [ { candidates: [{ content: { parts: [{ text: 'response text' }] } }], }, ], }; const mockRequest = vi.fn().mockResolvedValue(mockResponse); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest, } as unknown as JSONClient, projectId: 'test-project-id', }); vi.spyOn(vertexUtil, 'loadCredentials').mockImplementation(function (creds) { if (typeof creds === 'object') { return JSON.stringify(creds); } return creds; }); vi.spyOn(vertexUtil, 'resolveProjectId').mockResolvedValue('test-project-id'); const response = await provider.callGeminiApi('test prompt'); expect(response).toEqual({ cached: false, output: 'response text', tokenUsage: { total: 0, prompt: 0, completion: 0, }, cost: undefined, metadata: {}, }); }); it('should handle errors in function tool callbacks', async () => { const mockCachedResponse = { cached: true, output: JSON.stringify({ functionCall: { name: 'errorFunction', args: '{}', }, }), tokenUsage: { total: 5, prompt: 2, completion: 3, }, }; mockCacheGet.mockResolvedValue(JSON.stringify(mockCachedResponse)); const provider = new VertexChatProvider('gemini', { config: { tools: [ { functionDeclarations: [ { name: 'errorFunction', description: 'A function that always throws an error', parameters: { type: 'OBJECT', properties: {}, }, }, ], }, ], functionToolCallbacks: { errorFunction: () => { throw new Error('Test error'); }, }, }, }); const result = await provider.callApi('Call the error function'); expect(result.output).toBe('{"functionCall":{"name":"errorFunction","args":"{}"}}'); expect(result.tokenUsage).toEqual({ total: 5, prompt: 2, completion: 3, cached: 5 }); }); describe('External Function Callbacks', () => { beforeEach(() => { // Set cliState basePath for external function loading cliState.basePath = '/test/base/path'; }); afterEach(() => { vi.clearAllMocks(); cliState.basePath = undefined; }); it('should load and execute external function callbacks from file', async () => { const mockCachedResponse = { cached: true, output: JSON.stringify({ functionCall: { name: 'external_function', args: '{"param":"test_value"}', }, }), tokenUsage: { total: 15, prompt: 10, completion: 5, }, }; mockCacheGet.mockResolvedValue(JSON.stringify(mockCachedResponse)); // Mock importModule to return our test function const mockExternalFunction = vi.fn().mockResolvedValue('External function result'); mockImportModule.mockResolvedValue(mockExternalFunction); const provider = new VertexChatProvider('gemini', { config: { tools: [ { functionDeclarations: [ { name: 'external_function', description: 'An external function', parameters: { type: 'OBJECT', properties: { param: { type: 'STRING' } }, required: ['param'], }, }, ], }, ], functionToolCallbacks: { external_function: 'file://test/callbacks.js:testFunction', }, }, }); const result = await provider.callApi('Call external function'); expect(mockImportModule).toHaveBeenCalledWith( path.resolve('/test/base/path', 'test/callbacks.js'), 'testFunction', ); expect(mockExternalFunction).toHaveBeenCalledWith('{"param":"test_value"}'); expect(result.output).toBe('External function result'); expect(result.tokenUsage).toEqual({ total: 15, prompt: 10, completion: 5, cached: 15 }); }); it('should cache external functions and not reload them on subsequent calls', async () => { const mockCachedResponse = { cached: true, output: JSON.stringify({ functionCall: { name: 'cached_function', args: '{"value":123}', }, }), tokenUsage: { total: 12, prompt: 8, completion: 4, }, }; mockCacheGet.mockResolvedValue(JSON.stringify(mockCachedResponse)); const mockCachedFunction = vi.fn().mockResolvedValue('Cached result'); mockImportModule.mockResolvedValue(mockCachedFunction); const provider = new VertexChatProvider('gemini', { config: { tools: [ { functionDeclarations: [ { name: 'cached_function', description: 'A cached function', parameters: { type: 'OBJECT', properties: { value: { type: 'NUMBER' } }, required: ['value'], }, }, ], }, ], functionToolCallbacks: { cached_function: 'file://callbacks/cache-test.js:cachedFunction', }, }, }); // First call - should load the function const result1 = await provider.callApi('First call'); expect(mockImportModule).toHaveBeenCalledTimes(1); expect(mockCachedFunction).toHaveBeenCalledWith('{"value":123}'); expect(result1.output).toBe('Cached result'); // Second call - should use cached function, not reload const result2 = await provider.callApi('Second call'); expect(mockImportModule).toHaveBeenCalledTimes(1); // Still only 1 call expect(mockCachedFunction).toHaveBeenCalledTimes(2); expect(result2.output).toBe('Cached result'); }); it('should handle errors in external function loading gracefully', async () => { const mockCachedResponse = { cached: true, output: JSON.stringify({ functionCall: { name: 'error_function', args: '{"test":"data"}', }, }), tokenUsage: { total: 10, prompt: 6, completion: 4, }, }; mockCacheGet.mockResolvedValue(JSON.stringify(mockCachedResponse)); // Mock import module to throw an error mockImportModule.mockRejectedValue(new Error('Module not found')); const provider = new VertexChatProvider('gemini', { config: { tools: [ { functionDeclarations: [ { name: 'error_function', description: 'A function that errors during loading', parameters: { type: 'OBJECT', properties: { test: { type: 'STRING' } }, }, }, ], }, ], functionToolCallbacks: { error_function: 'file://nonexistent/module.js:errorFunction', }, }, }); const result = await provider.callApi('Call error function'); expect(mockImportModule).toHaveBeenCalledWith( path.resolve('/test/base/path', 'nonexistent/module.js'), 'errorFunction', ); // Should fall back to original function call object when loading fails expect(result.output).toBe( '{"functionCall":{"name":"error_function","args":"{\\"test\\":\\"data\\"}"}}', ); }); it('should handle mixed inline and external function callbacks', async () => { const mockCachedResponse = { cached: true, output: JSON.stringify({ functionCall: { name: 'external_function', args: '{"external":"test"}', }, }), tokenUsage: { total: 20, prompt: 12, completion: 8, }, }; mockCacheGet.mockResolvedValue(JSON.stringify(mockCachedResponse)); const mockInlineFunction = vi.fn().mockResolvedValue('Inline result'); const mockExternalFunction = vi.fn().mockResolvedValue('External result'); mockImportModule.mockResolvedValue(mockExternalFunction); const provider = new VertexChatProvider('gemini', { config: { tools: [ { functionDeclarations: [ { name: 'inline_function', description: 'An inline function', parameters: { type: 'OBJECT', properties: { inline: { type: 'STRING' } }, }, }, { name: 'external_function', description: 'An external function', parameters: { type: 'OBJECT', properties: { external: { type: 'STRING' } }, }, }, ], }, ], functionToolCallbacks: { inline_function: mockInlineFunction, external_function: 'file://mixed/callbacks.js:externalFunc', }, }, }); const result = await provider.callApi('Test mixed callbacks'); expect(mockImportModule).toHaveBeenCalledWith( path.resolve('/test/base/path', 'mixed/callbacks.js'), 'externalFunc', ); expect(mockExternalFunction).toHaveBeenCalledWith('{"external":"test"}'); expect(result.output).toBe('External result'); }); }); describe('thinking token tracking', () => { it('should track thinking tokens when present in response', async () => { const provider = new VertexChatProvider('gemini-2.5-flash', { config: { generationConfig: { thinkingConfig: { thinkingBudget: 1024, }, }, }, }); const mockResponse = { data: [ { candidates: [{ content: { parts: [{ text: 'response with thinking' }] } }], usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 20, totalTokenCount: 30, thoughtsTokenCount: 50, // Thinking tokens }, }, ], }; const mockRequest = vi.fn().mockResolvedValue(mockResponse); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest, } as unknown as JSONClient, projectId: 'test-project-id', }); const response = await provider.callGeminiApi('test prompt'); expect(response.tokenUsage).toEqual({ prompt: 10, completion: 20, total: 30, completionDetails: { reasoning: 50, acceptedPrediction: 0, rejectedPrediction: 0, }, }); }); it('should handle response without thinking tokens', async () => { const provider = new VertexChatProvider('gemini-2.5-flash'); const mockResponse = { data: [ { candidates: [{ content: { parts: [{ text: 'response without thinking' }] } }], usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 20, totalTokenCount: 30, // No thoughtsTokenCount field }, }, ], }; const mockRequest = vi.fn().mockResolvedValue(mockResponse); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest, } as unknown as JSONClient, projectId: 'test-project-id', }); const response = await provider.callGeminiApi('test prompt'); expect(response.tokenUsage).toEqual({ prompt: 10, completion: 20, total: 30, // No completionDetails field when thoughtsTokenCount is absent }); }); it('should track thinking tokens with zero value', async () => { const provider = new VertexChatProvider('gemini-2.5-flash', { config: { generationConfig: { thinkingConfig: { thinkingBudget: 1024, }, }, }, }); const mockResponse = { data: [ { candidates: [{ content: { parts: [{ text: 'response with zero thinking' }] } }], usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 20, totalTokenCount: 30, thoughtsTokenCount: 0, // Zero thinking tokens }, }, ], }; const mockRequest = vi.fn().mockResolvedValue(mockResponse); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest, } as unknown as JSONClient, projectId: 'test-project-id', }); const response = await provider.callGeminiApi('test prompt'); expect(response.tokenUsage).toEqual({ prompt: 10, completion: 20, total: 30, completionDetails: { reasoning: 0, acceptedPrediction: 0, rejectedPrediction: 0, }, }); }); it('should track thinking tokens in cached responses', async () => { const provider = new VertexChatProvider('gemini-2.5-flash', { config: { generationConfig: { thinkingConfig: { thinkingBudget: 1024, }, }, }, }); const mockCachedResponse = { output: 'cached response with thinking', tokenUsage: { total: 80, prompt: 10, completion: 20, thoughtsTokenCount: 50, // This would be stored in the cached response }, cached: true, }; // Mock the cache to return a response that includes thinking tokens mockCacheGet.mockResolvedValue(JSON.stringify(mockCachedResponse)); const response = await provider.callGeminiApi('test prompt'); // The cached response should preserve the thinking tokens // but due to how the caching logic works, it transforms the response expect(response.cached).toBe(true); expect(response.output).toBe('cached response with thinking'); // Note: The current implementation doesn't preserve completionDetails in cached responses // This is a limitation that could be addressed in a future fix }); }); describe('Model Armor integration', () => { it('should include model_armor_config in request when configured', async () => { const provider = new VertexChatProvider('gemini-pro', { config: { modelArmor: { promptTemplate: 'projects/my-project/locations/us-central1/templates/basic-safety', responseTemplate: 'projects/my-project/locations/us-central1/templates/basic-safety', }, }, }); const mockResponse = { data: [ { candidates: [{ content: { parts: [{ text: 'response text' }] } }], usageMetadata: { totalTokenCount: 10, promptTokenCount: 5, candidatesTokenCount: 5, }, }, ], }; const mockRequest = vi.fn().mockResolvedValue(mockResponse); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest, } as unknown as JSONClient, projectId: 'test-project-id', }); vi.spyOn(vertexUtil, 'loadCredentials').mockImplementation(function (creds) { if (typeof creds === 'object') { return JSON.stringify(creds); } return creds; }); vi.spyOn(vertexUtil, 'resolveProjectId').mockResolvedValue('test-project-id'); await provider.callGeminiApi('test prompt'); expect(mockRequest).toHaveBeenCalledWith( expect.objectContaining({ data: expect.objectContaining({ model_armor_config: { prompt_template_name: 'projects/my-project/locations/us-central1/templates/basic-safety', response_template_name: 'projects/my-project/locations/us-central1/templates/basic-safety', }, }), }), ); }); it('should include only promptTemplate when responseTemplate is not configured', async () => { const provider = new VertexChatProvider('gemini-pro', { config: { modelArmor: { promptTemplate: 'projects/my-project/locations/us-central1/templates/prompt-only', }, }, }); const mockResponse = { data: [ { candidates: [{ content: { parts: [{ text: 'response text' }] } }], usageMetadata: { totalTokenCount: 10, promptTokenCount: 5, candidatesTokenCount: 5, }, }, ], }; const mockRequest = vi.fn().mockResolvedValue(mockResponse); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest, } as unknown as JSONClient, projectId: 'test-project-id', }); vi.spyOn(vertexUtil, 'loadCredentials').mockImplementation(function (creds) { if (typeof creds === 'object') { return JSON.stringify(creds); } return creds; }); vi.spyOn(vertexUtil, 'resolveProjectId').mockResolvedValue('test-project-id'); await provider.callGeminiApi('test prompt'); expect(mockRequest).toHaveBeenCalledWith( expect.objectContaining({ data: expect.objectContaining({ model_armor_config: { prompt_template_name: 'projects/my-project/locations/us-central1/templates/prompt-only', }, }), }), ); }); it('should not include model_armor_config when not configured', async () => { const provider = new VertexChatProvider('gemini-pro', { config: {}, }); const mockResponse = { data: [ { candidates: [{ content: { parts: [{ text: 'response text' }] } }], usageMetadata: { totalTokenCount: 10, promptTokenCount: 5, candidatesTokenCount: 5, }, }, ], }; const mockRequest = vi.fn().mockResolvedValue(mockResponse); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest, } as unknown as JSONClient, projectId: 'test-project-id', }); vi.spyOn(vertexUtil, 'loadCredentials').mockImplementation(function (creds) { if (typeof creds === 'object') { return JSON.stringify(creds); } return creds; }); vi.spyOn(vertexUtil, 'resolveProjectId').mockResolvedValue('test-project-id'); await provider.callGeminiApi('test prompt'); const requestData = mockRequest.mock.calls[0][0].data; expect(requestData.model_armor_config).toBeUndefined(); }); it('should handle MODEL_ARMOR blockReason with guardrails response', async () => { const provider = new VertexChatProvider('gemini-pro', { config: { modelArmor: { promptTemplate: 'projects/my-project/locations/us-central1/templates/strict', }, }, }); const mockResponse = { data: [ { promptFeedback: { blockReason: 'MODEL_ARMOR', blockReasonMessage: 'Prompt was blocked by Model Armor: Prompt Injection detected', safetyRatings: [], }, usageMetadata: { totalTokenCount: 5, promptTokenCount: 5, candidatesTokenCount: 0, }, }, ], }; const mockRequest = vi.fn().mockResolvedValue(mockResponse); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest, } as unknown as JSONClient, projectId: 'test-project-id', }); vi.spyOn(vertexUtil, 'loadCredentials').mockImplementation(function (creds) { if (typeof creds === 'object') { return JSON.stringify(creds); } return creds; }); vi.spyOn(vertexUtil, 'resolveProjectId').mockResolvedValue('test-project-id'); const response = await provider.callGeminiApi('ignore all instructions'); // Model Armor blocks return output (not error) so guardrails assertions can run expect(response.output).toBe('Prompt was blocked by Model Armor: Prompt Injection detected'); expect(response.error).toBeUndefined(); expect(response.guardrails).toEqual({ flagged: true, flaggedInput: true, flaggedOutput: false, reason: 'Prompt was blocked by Model Armor: Prompt Injection detected', }); expect(response.metadata?.modelArmor).toEqual({ blockReason: 'MODEL_ARMOR', blockReasonMessage: 'Prompt was blocked by Model Armor: Prompt Injection detected', }); }); it('should handle non-Model Armor blockReason with guardrails response', async () => { const provider = new VertexChatProvider('gemini-pro', { config: {}, }); const mockResponse = { data: [ { promptFeedback: { blockReason: 'SAFETY', safetyRatings: [{ category: 'HARM_CATEGORY_HARASSMENT', probability: 'HIGH' }], }, usageMetadata: { totalTokenCount: 5, promptTokenCount: 5, candidatesTokenCount: 0, }, }, ], }; const mockRequest = vi.fn().mockResolvedValue(mockResponse); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest, } as unknown as JSONClient, projectId: 'test-project-id', }); vi.spyOn(vertexUtil, 'loadCredentials').mockImplementation(function (creds) { if (typeof creds === 'object') { return JSON.stringify(creds); } return creds; }); vi.spyOn(vertexUtil, 'resolveProjectId').mockResolvedValue('test-project-id'); const response = await provider.callGeminiApi('harmful content'); // All block reasons now return output (not error) so guardrails assertions can run expect(response.output).toContain('Content was blocked due to safety settings: SAFETY'); expect(response.error).toBeUndefined(); expect(response.guardrails).toEqual({ flagged: true, flaggedInput: true, flaggedOutput: false, reason: expect.stringContaining('Content was blocked due to safety settings: SAFETY'), }); expect(response.metadata?.modelArmor).toBeUndefined(); }); it('should not include model_armor_config when modelArmor is empty object', async () => { const provider = new VertexChatProvider('gemini-pro', { config: { modelArmor: {}, }, }); const mockResponse = { data: [ { candidates: [{ content: { parts: [{ text: 'response text' }] } }], usageMetadata: { totalTokenCount: 10, promptTokenCount: 5, candidatesTokenCount: 5, }, }, ], }; const mockRequest = vi.fn().mockResolvedValue(mockResponse); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest, } as unknown as JSONClient, projectId: 'test-project-id', }); vi.spyOn(vertexUtil, 'loadCredentials').mockImplementation(function (creds) { if (typeof creds === 'object') { return JSON.stringify(creds); } return creds; }); vi.spyOn(vertexUtil, 'resolveProjectId').mockResolvedValue('test-project-id'); await provider.callGeminiApi('test prompt'); const requestData = mockRequest.mock.calls[0][0].data; expect(requestData.model_armor_config).toBeUndefined(); }); // TODO: This default message is user-facing and can be adjusted for clarity without // breaking behavior semantics (e.g., "Content was blocked by Model Armor policy"). it('should use default message when blockReasonMessage is not provided', async () => { const provider = new VertexChatProvider('gemini-pro', { config: {}, }); const mockResponse = { data: [ { promptFeedback: { blockReason: 'MODEL_ARMOR', // No blockReasonMessage safetyRatings: [], }, usageMetadata: { totalTokenCount: 5, promptTokenCount: 5, candidatesTokenCount: 0, }, }, ], }; const mockRequest = vi.fn().mockResolvedValue(mockResponse); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest, } as unknown as JSONClient, projectId: 'test-project-id', }); vi.spyOn(vertexUtil, 'loadCredentials').mockImplementation(function (creds) { if (typeof creds === 'object') { return JSON.stringify(creds); } return creds; }); vi.spyOn(vertexUtil, 'resolveProjectId').mockResolvedValue('test-project-id'); const response = await provider.callGeminiApi('test prompt'); // Block reasons return output (not error) so guardrails assertions can run expect(response.output).toBe('Content was blocked due to Model Armor: MODEL_ARMOR'); expect(response.error).toBeUndefined(); expect(response.guardrails?.reason).toBe( 'Content was blocked due to Model Armor: MODEL_ARMOR', ); }); it('should handle SAFETY finishReason with guardrails response', async () => { const provider = new VertexChatProvider('gemini-pro', { config: {}, }); const mockResponse = { data: [ { candidates: [ { content: { parts: [{ text: 'partial response' }] }, finishReason: 'SAFETY', }, ], usageMetadata: { totalTokenCount: 10, promptTokenCount: 5, candidatesTokenCount: 5, }, }, ], }; const mockRequest = vi.fn().mockResolvedValue(mockResponse); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest, } as unknown as JSONClient, projectId: 'test-project-id', }); vi.spyOn(vertexUtil, 'loadCredentials').mockImplementation(function (creds) { if (typeof creds === 'object') { return JSON.stringify(creds); } return creds; }); vi.spyOn(vertexUtil, 'resolveProjectId').mockResolvedValue('test-project-id'); const response = await provider.callGeminiApi('test prompt'); expect(response.error).toBe( 'Content was blocked due to safety settings with finish reason: SAFETY.', ); expect(response.guardrails).toEqual({ flagged: true, flaggedInput: false, flaggedOutput: true, reason: 'Content was blocked due to safety settings with finish reason: SAFETY.', }); }); it.each([ 'PROHIBITED_CONTENT', 'RECITATION', 'BLOCKLIST', 'SPII', 'IMAGE_SAFETY', ])('should handle %s finishReason with guardrails response', async (finishReason) => { const provider = new VertexChatProvider('gemini-pro', { config: {}, }); const mockResponse = { data: [ { candidates: [ { content: { parts: [{ text: 'partial response' }] }, finishReason, }, ], usageMetadata: { totalTokenCount: 10, promptTokenCount: 5, candidatesTokenCount: 5, }, }, ], }; const mockRequest = vi.fn().mockResolvedValue(mockResponse); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest, } as unknown as JSONClient, projectId: 'test-project-id', }); vi.spyOn(vertexUtil, 'loadCredentials').mockImplementation(function (creds) { if (typeof creds === 'object') { return JSON.stringify(creds); } return creds; }); vi.spyOn(vertexUtil, 'resolveProjectId').mockResolvedValue('test-project-id'); const response = await provider.callGeminiApi('test prompt'); expect(response.error).toBe( `Content was blocked due to safety settings with finish reason: ${finishReason}.`, ); expect(response.guardrails).toEqual({ flagged: true, flaggedInput: false, flaggedOutput: true, reason: `Content was blocked due to safety settings with finish reason: ${finishReason}.`, }); }); it('should handle MAX_TOKENS finishReason with truncated output', async () => { const provider = new VertexChatProvider('gemini-pro', { config: {}, }); const longOutput = 'A'.repeat(600); const mockResponse = { data: [ { candidates: [ { content: { parts: [{ text: longOutput }] }, finishReason: 'MAX_TOKENS', }, ], usageMetadata: { totalTokenCount: 1100, promptTokenCount: 100, candidatesTokenCount: 1000, }, }, ], }; const mockRequest = vi.fn().mockResolvedValue(mockResponse); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest, } as unknown as JSONClient, projectId: 'test-project-id', }); vi.spyOn(vertexUtil, 'loadCredentials').mockImplementation(function (creds) { if (typeof creds === 'object') { return JSON.stringify(creds); } return creds; }); vi.spyOn(vertexUtil, 'resolveProjectId').mockResolvedValue('test-project-id'); const response = await provider.callGeminiApi('test prompt'); expect(response.error).toBeUndefined(); expect(response.output).toBe(longOutput); expect(response.tokenUsage).toEqual({ total: 1100, prompt: 100, completion: 1000, }); }); it('should handle MAX_TOKENS finishReason with short output (no truncation)', async () => { const provider = new VertexChatProvider('gemini-pro', { config: {}, }); const shortOutput = 'Short response'; const mockResponse = { data: [ { candidates: [ { content: { parts: [{ text: shortOutput }] }, finishReason: 'MAX_TOKENS', }, ], usageMetadata: { totalTokenCount: 110, promptTokenCount: 100, candidatesTokenCount: 10, }, }, ], }; const mockRequest = vi.fn().mockResolvedValue(mockResponse); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest, } as unknown as JSONClient, projectId: 'test-project-id', }); vi.spyOn(vertexUtil, 'loadCredentials').mockImplementation(function (creds) { if (typeof creds === 'object') { return JSON.stringify(creds); } return creds; }); vi.spyOn(vertexUtil, 'resolveProjectId').mockResolvedValue('test-project-id'); const response = await provider.callGeminiApi('test prompt'); expect(response.error).toBeUndefined(); expect(response.output).toBe(shortOutput); expect(response.tokenUsage).toEqual({ total: 110, prompt: 100, completion: 10, }); }); }); }); describe('VertexChatProvider.callPalm2Api', () => { beforeEach(() => { mockCacheGet.mockReset(); mockCacheGet.mockResolvedValue(null); mockCacheSet.mockReset(); mockIsCacheEnabled.mockReturnValue(true); }); afterEach(() => { vi.clearAllMocks(); }); it('hashes Palm2 request body cache keys without leaking prompts', async () => { const prompt = 'palm2-secret-prompt-value'; const provider = new VertexChatProvider('chat-bison', { config: { context: 'palm2-secret-context', temperature: 0.4, }, }); mockVertexRequest({ predictions: [{ candidates: [{ content: 'Palm2 response content' }] }], }); await provider.callPalm2Api(prompt); expectHashedBodyCacheKeys(/^vertex:palm2:chat-bison:[a-f0-9]{64}$/, [ prompt, 'palm2-secret-context', ]); }); it('scopes Palm2 request body cache keys by model name', async () => { const prompt = 'same palm2 prompt'; const providerA = new VertexChatProvider('chat-bison', { config: { context: 'same palm2 context' }, }); const providerB = new VertexChatProvider('text-bison', { config: { context: 'same palm2 context' }, }); mockVertexRequest({ predictions: [{ candidates: [{ content: 'Palm2 response content' }] }], }); await providerA.callPalm2Api(prompt); await providerB.callPalm2Api(prompt); const [cacheKeyA, cacheKeyB] = mockCacheGet.mock.calls.map(([key]) => key as string); expect(cacheKeyA).toMatch(/^vertex:palm2:chat-bison:[a-f0-9]{64}$/); expect(cacheKeyB).toMatch(/^vertex:palm2:text-bison:[a-f0-9]{64}$/); expect(cacheKeyA).not.toBe(cacheKeyB); expect(cacheKeyA).not.toContain(prompt); expect(cacheKeyB).not.toContain(prompt); }); }); describe('VertexChatProvider.callLlamaApi', () => { let provider: VertexChatProvider; beforeEach(() => { // Reset cache mocks to default state mockCacheGet.mockReset(); mockCacheGet.mockResolvedValue(null); mockCacheSet.mockReset(); mockIsCacheEnabled.mockReturnValue(true); }); afterEach(() => { vi.clearAllMocks(); }); it('should enforce us-central1 region for Llama models', async () => { // Create provider with non-us-central1 region provider = new VertexChatProvider('llama-3.3-70b-instruct-maas', { config: { region: 'europe-west1' }, }); const response = await provider.callLlamaApi('test prompt'); // Should return error about region expect(response).toEqual({ error: "Llama models are only available in the us-central1 region. Current region: europe-west1. Please set region: 'us-central1' in your configuration.", }); }); it('should validate llama_guard_settings is a valid object', async () => { provider = new VertexChatProvider('llama-3.3-70b-instruct-maas', { config: { region: 'us-central1', llamaConfig: { safetySettings: { // @ts-ignore - intentionally passing invalid type for test llama_guard_settings: 'not-an-object', }, }, }, }); const response = await provider.callLlamaApi('test prompt'); // Should return error about invalid llama_guard_settings expect(response).toEqual({ error: 'Invalid llama_guard_settings: must be an object, received string', }); }); it('should successfully call Llama API with valid configuration', async () => { provider = new VertexChatProvider('llama-3.3-70b-instruct-maas', { config: { region: 'us-central1', temperature: 0.7, maxOutputTokens: 250, llamaConfig: { safetySettings: { enabled: true, llama_guard_settings: { custom_setting: 'value' }, }, }, }, }); const mockResponse = { data: { choices: [ { message: { content: 'Llama response content', }, }, ], usage: { total_tokens: 35, prompt_tokens: 15, completion_tokens: 20, }, }, }; const mockRequest = vi.fn().mockResolvedValue(mockResponse); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest, } as unknown as JSONClient, projectId: 'test-project-id', }); vi.spyOn(vertexUtil, 'loadCredentials').mockImplementation(function (creds) { if (typeof creds === 'object') { return JSON.stringify(creds); } return creds; }); vi.spyOn(vertexUtil, 'resolveProjectId').mockResolvedValue('test-project-id'); const response = await provider.callLlamaApi('test prompt'); // Should return successful response expect(response).toEqual({ cached: false, output: 'Llama response content', tokenUsage: { total: 35, prompt: 15, completion: 20, numRequests: 1, }, }); // Verify API request expect(mockRequest).toHaveBeenCalledWith( expect.objectContaining({ url: expect.stringContaining( 'us-central1-aiplatform.googleapis.com/v1beta1/projects/test-project-id/locations/us-central1/endpoints/openapi/chat/completions', ), method: 'POST', data: expect.objectContaining({ model: 'meta/llama-3.3-70b-instruct-maas', max_tokens: 250, temperature: 0.7, extra_body: { google: { model_safety_settings: { enabled: true, llama_guard_settings: { custom_setting: 'value' }, }, }, }, }), }), ); }); it('hashes Llama request body cache keys without leaking prompts', async () => { const prompt = 'llama-secret-prompt-value'; provider = new VertexChatProvider('llama-3.3-70b-instruct-maas', { config: { region: 'us-central1', temperature: 0.7, llamaConfig: { safetySettings: { llama_guard_settings: { marker: 'llama-secret-safety-setting' }, }, }, }, }); mockVertexRequest({ choices: [ { message: { content: 'Llama response content', }, }, ], usage: { total_tokens: 35, prompt_tokens: 15, completion_tokens: 20, }, }); await provider.callLlamaApi(prompt); expectHashedBodyCacheKeys(/^vertex:llama:llama-3\.3-70b-instruct-maas:[a-f0-9]{64}$/, [ prompt, 'llama-secret-safety-setting', ]); }); it('does not log raw Llama prompts, safety settings, or outputs', async () => { const prompt = 'llama-secret-prompt-value'; const output = 'llama-secret-output-value'; const debugSpy = vi.spyOn(logger, 'debug').mockImplementation(() => {}); provider = new VertexChatProvider('llama-3.3-70b-instruct-maas', { config: { region: 'us-central1', llamaConfig: { safetySettings: { llama_guard_settings: { marker: 'llama-secret-safety-setting' }, }, }, }, }); mockVertexRequest({ choices: [ { message: { content: output, }, }, ], usage: { total_tokens: 35, prompt_tokens: 15, completion_tokens: 20, }, }); await provider.callLlamaApi(prompt); const debugLogs = JSON.stringify(debugSpy.mock.calls); expect(debugLogs).not.toContain(prompt); expect(debugLogs).not.toContain(output); expect(debugLogs).not.toContain('llama-secret-safety-setting'); expect(debugLogs).toContain('Preparing to call Llama API'); debugSpy.mockRestore(); }); it('should default safety settings to enabled when not specified', async () => { provider = new VertexChatProvider('llama-3.3-70b-instruct-maas', { config: { region: 'us-central1', }, }); const mockResponse = { data: { choices: [ { message: { content: 'Llama response with default safety', }, }, ], usage: { total_tokens: 30, prompt_tokens: 10, completion_tokens: 20, }, }, }; const mockRequest = vi.fn().mockResolvedValue(mockResponse); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest, } as unknown as JSONClient, projectId: 'test-project-id', }); vi.spyOn(vertexUtil, 'loadCredentials').mockImplementation(function (creds) { if (typeof creds === 'object') { return JSON.stringify(creds); } return creds; }); vi.spyOn(vertexUtil, 'resolveProjectId').mockResolvedValue('test-project-id'); await provider.callLlamaApi('test prompt'); // Verify safety settings defaulted to enabled expect(mockRequest).toHaveBeenCalledWith( expect.objectContaining({ data: expect.objectContaining({ extra_body: { google: { model_safety_settings: { enabled: true, llama_guard_settings: {}, }, }, }, }), }), ); }); it('should handle API errors correctly', async () => { provider = new VertexChatProvider('llama-3.3-70b-instruct-maas', { config: { region: 'us-central1', }, }); const mockError = { response: { data: { error: { code: 400, message: 'Invalid request', }, }, }, }; vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: vi.fn().mockRejectedValue(mockError), } as unknown as JSONClient, projectId: 'test-project-id', }); const response = await provider.callLlamaApi('test prompt'); expect(response.error).toContain('API call error:'); expect(response.error).toContain('Invalid request'); }); it('should load system instructions from file', async () => { const mockSystemInstruction = 'You are a helpful assistant from a file.'; // Mock file system operations vi.spyOn(fs, 'existsSync').mockReturnValue(true); vi.spyOn(fs, 'readFileSync').mockReturnValue(mockSystemInstruction); provider = new VertexChatProvider('gemini-2.5-flash', { config: { systemInstruction: 'file://system-instruction.txt', }, }); const mockResponse = { data: [ { candidates: [{ content: { parts: [{ text: 'response text' }] } }], usageMetadata: { totalTokenCount: 10, promptTokenCount: 5, candidatesTokenCount: 5, }, }, ], }; const mockRequest = vi.fn().mockResolvedValue(mockResponse); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest, } as unknown as JSONClient, projectId: 'test-project-id', }); vi.spyOn(vertexUtil, 'loadCredentials').mockImplementation(function (creds) { if (typeof creds === 'object') { return JSON.stringify(creds); } return creds; }); vi.spyOn(vertexUtil, 'resolveProjectId').mockResolvedValue('test-project-id'); await provider.callGeminiApi('test prompt'); expect(mockRequest).toHaveBeenCalledWith( expect.objectContaining({ data: expect.objectContaining({ systemInstruction: { parts: [{ text: mockSystemInstruction }], }, }), }), ); // Verify file was read expect(fs.readFileSync).toHaveBeenCalledWith( expect.stringContaining('system-instruction.txt'), 'utf8', ); }); it('honors VERTEX_API_HOST overrides on the Llama path', async () => { const mockRequest = vi.fn().mockResolvedValue({ data: { choices: [{ message: { content: 'Llama response content' } }], usage: { total_tokens: 30, prompt_tokens: 10, completion_tokens: 20 }, }, }); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest } as unknown as JSONClient, projectId: 'test-project-id', }); vi.spyOn(vertexUtil, 'loadCredentials').mockImplementation((creds) => typeof creds === 'object' ? JSON.stringify(creds) : creds, ); vi.spyOn(vertexUtil, 'resolveProjectId').mockResolvedValue('test-project-id'); const provider = new VertexChatProvider('llama-3.3-70b-instruct-maas', { config: { region: 'us-central1' }, env: { VERTEX_API_HOST: 'llama-proxy.example.test' }, }); await provider.callLlamaApi('test prompt'); expect(mockRequest).toHaveBeenCalledWith( expect.objectContaining({ url: expect.stringContaining('https://llama-proxy.example.test/'), }), ); }); }); describe('VertexChatProvider.callClaudeApi', () => { let provider: VertexChatProvider; beforeEach(() => { // Reset cache mocks to default state mockCacheGet.mockReset(); mockCacheGet.mockResolvedValue(null); mockCacheSet.mockReset(); mockIsCacheEnabled.mockReset(); mockIsCacheEnabled.mockReturnValue(true); }); afterEach(() => { vi.clearAllMocks(); }); it('should accept anthropicVersion parameter', async () => { provider = new VertexChatProvider('claude-3-5-sonnet-v2@20241022', { config: { anthropicVersion: 'vertex-2023-10-16', maxOutputTokens: 500, }, }); const mockResponse = { data: { id: 'test-id', type: 'message', role: 'assistant', model: 'claude-3-5-sonnet-v2@20241022', content: [{ type: 'text', text: 'Response from Claude' }], stop_reason: 'end_turn', stop_sequence: null, usage: { input_tokens: 20, output_tokens: 30, cache_creation_input_tokens: 0, cache_read_input_tokens: 0, }, }, }; const mockRequest = vi.fn().mockResolvedValue(mockResponse); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest, } as unknown as JSONClient, projectId: 'test-project-id', }); vi.spyOn(vertexUtil, 'loadCredentials').mockImplementation(function (creds) { if (typeof creds === 'object') { return JSON.stringify(creds); } return creds; }); vi.spyOn(vertexUtil, 'resolveProjectId').mockResolvedValue('test-project-id'); await provider.callClaudeApi('test prompt'); expect(mockRequest).toHaveBeenCalledWith( expect.objectContaining({ data: expect.objectContaining({ anthropic_version: 'vertex-2023-10-16', max_tokens: 500, }), }), ); }); it('hashes Claude request body cache keys without leaking prompts', async () => { const prompt = 'claude-secret-prompt-value'; provider = new VertexChatProvider('claude-3-5-sonnet-v2@20241022', { config: { anthropicVersion: 'vertex-2023-10-16', maxOutputTokens: 500, systemInstruction: 'claude-secret-system-instruction', }, }); mockVertexRequest({ id: 'test-id', type: 'message', role: 'assistant', model: 'claude-3-5-sonnet-v2@20241022', content: [{ type: 'text', text: 'Response from Claude' }], stop_reason: 'end_turn', stop_sequence: null, usage: { input_tokens: 20, output_tokens: 30, cache_creation_input_tokens: 0, cache_read_input_tokens: 0, }, }); await provider.callClaudeApi(prompt); expectHashedBodyCacheKeys( /^vertex:claude:claude-3-5-sonnet-v2@20241022:showThinking=false:[a-f0-9]{64}$/, [prompt, 'claude-secret-system-instruction'], ); }); it('should accept anthropic_version parameter (alternative format)', async () => { provider = new VertexChatProvider('claude-3-5-sonnet-v2@20241022', { config: { anthropic_version: 'vertex-2023-10-16-alt', max_tokens: 600, }, }); const mockResponse = { data: { id: 'test-id', type: 'message', role: 'assistant', model: 'claude-3-5-sonnet-v2@20241022', content: [{ type: 'text', text: 'Response from Claude' }], stop_reason: 'end_turn', stop_sequence: null, usage: { input_tokens: 20, output_tokens: 30, cache_creation_input_tokens: 0, cache_read_input_tokens: 0, }, }, }; const mockRequest = vi.fn().mockResolvedValue(mockResponse); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest, } as unknown as JSONClient, projectId: 'test-project-id', }); vi.spyOn(vertexUtil, 'loadCredentials').mockImplementation(function (creds) { if (typeof creds === 'object') { return JSON.stringify(creds); } return creds; }); vi.spyOn(vertexUtil, 'resolveProjectId').mockResolvedValue('test-project-id'); await provider.callClaudeApi('test prompt'); expect(mockRequest).toHaveBeenCalledWith( expect.objectContaining({ data: expect.objectContaining({ anthropic_version: 'vertex-2023-10-16-alt', max_tokens: 600, }), }), ); }); it('omits temperature for Claude Opus 4.7 on Vertex', async () => { provider = new VertexChatProvider('claude-opus-4-7', { config: { max_tokens: 32, temperature: 0.5 }, }); const mockResponse = { data: { id: 'test-id', type: 'message', role: 'assistant', model: 'claude-opus-4-7', content: [{ type: 'text', text: 'ok' }], stop_reason: 'end_turn', stop_sequence: null, usage: { input_tokens: 5, output_tokens: 1, cache_creation_input_tokens: 0, cache_read_input_tokens: 0, }, }, }; const mockRequest = vi.fn().mockResolvedValue(mockResponse); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest } as unknown as JSONClient, projectId: 'test-project-id', }); vi.spyOn(vertexUtil, 'loadCredentials').mockImplementation((creds) => typeof creds === 'object' ? JSON.stringify(creds) : creds, ); vi.spyOn(vertexUtil, 'resolveProjectId').mockResolvedValue('test-project-id'); await provider.callClaudeApi('test prompt'); const sentBody = mockRequest.mock.calls[0][0].data as Record; expect(sentBody.temperature).toBeUndefined(); expect(sentBody.max_tokens).toBe(32); }); it('omits temperature, top_p, and top_k for Claude Opus 4.8 on Vertex', async () => { provider = new VertexChatProvider('claude-opus-4-8', { config: { max_tokens: 32, temperature: 0.5, top_p: 0.9, top_k: 40 }, }); const mockResponse = { data: { id: 'test-id', type: 'message', role: 'assistant', model: 'claude-opus-4-8', content: [{ type: 'text', text: 'ok' }], stop_reason: 'end_turn', stop_sequence: null, usage: { input_tokens: 5, output_tokens: 1, cache_creation_input_tokens: 0, cache_read_input_tokens: 0, }, }, }; const mockRequest = vi.fn().mockResolvedValue(mockResponse); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest } as unknown as JSONClient, projectId: 'test-project-id', }); vi.spyOn(vertexUtil, 'loadCredentials').mockImplementation((creds) => typeof creds === 'object' ? JSON.stringify(creds) : creds, ); vi.spyOn(vertexUtil, 'resolveProjectId').mockResolvedValue('test-project-id'); await provider.callClaudeApi('test prompt'); const sentBody = mockRequest.mock.calls[0][0].data as Record; expect(sentBody.temperature).toBeUndefined(); expect(sentBody.top_p).toBeUndefined(); expect(sentBody.top_k).toBeUndefined(); expect(sentBody.max_tokens).toBe(32); }); it.each([ { name: 'at the 200K boundary', region: 'global', inputTokens: 200_000, outputTokens: 10_000, cacheReadTokens: 0, cacheCreationTokens: 0, pricingOverrides: {}, expectedCost: 0.75, }, { name: 'above 200K globally', region: 'global', inputTokens: 300_000, outputTokens: 20_000, cacheReadTokens: 0, cacheCreationTokens: 0, pricingOverrides: {}, expectedCost: 2.25, }, { name: 'when cache tokens cross 200K regionally', region: 'us-central1', inputTokens: 150_000, outputTokens: 10_000, cacheReadTokens: 60_000, cacheCreationTokens: 10_000, pricingOverrides: {}, expectedCost: 1.3596, }, { name: 'with custom regional rates above 200K', region: 'us-central1', inputTokens: 300_000, outputTokens: 20_000, cacheReadTokens: 0, cacheCreationTokens: 0, pricingOverrides: { inputCost: 2 / 1e6, outputCost: 7 / 1e6 }, expectedCost: 0.74, }, ])('prices Vertex Sonnet 4.5 $name', async ({ region, inputTokens, outputTokens, cacheReadTokens, cacheCreationTokens, pricingOverrides, expectedCost, }) => { const model = 'claude-sonnet-4-5@20250929'; provider = new VertexChatProvider(model, { config: { region, max_tokens: 32, ...pricingOverrides }, }); mockVertexRequest({ id: 'test-id', type: 'message', role: 'assistant', model, content: [{ type: 'text', text: 'ok' }], stop_reason: 'end_turn', stop_sequence: null, usage: { input_tokens: inputTokens, output_tokens: outputTokens, cache_read_input_tokens: cacheReadTokens, cache_creation_input_tokens: cacheCreationTokens, }, }); const result = await provider.callClaudeApi('test prompt'); expect(result.cost).toBeCloseTo(expectedCost, 6); }); it('supports Claude Fable 5 with adaptive-safe parameters and regional pricing', async () => { const model = 'claude-fable-5'; provider = new VertexChatProvider(model, { config: { max_tokens: 32, temperature: 0.5, top_p: 0.9, top_k: 40 }, }); const mockRequest = vi.fn().mockResolvedValue({ data: { id: 'test-id', type: 'message', role: 'assistant', model, content: [{ type: 'text', text: 'ok' }], stop_reason: 'end_turn', stop_sequence: null, usage: { input_tokens: 5, output_tokens: 1, cache_creation_input_tokens: 0, cache_read_input_tokens: 0, }, }, }); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest } as unknown as JSONClient, projectId: 'test-project-id', }); vi.spyOn(vertexUtil, 'loadCredentials').mockImplementation((creds) => typeof creds === 'object' ? JSON.stringify(creds) : creds, ); vi.spyOn(vertexUtil, 'resolveProjectId').mockResolvedValue('test-project-id'); const result = await provider.callClaudeApi('test prompt'); const request = mockRequest.mock.calls[0][0]; const sentBody = request.data as Record; expect(request.url).toContain(`/publishers/anthropic/models/${model}:rawPredict`); expect(sentBody.temperature).toBeUndefined(); expect(sentBody.top_p).toBeUndefined(); expect(sentBody.top_k).toBeUndefined(); expect(result.cost).toBeCloseTo(0.00011, 8); }); it('uses base Claude 5 pricing for Fable on the global Vertex region', async () => { const model = 'claude-fable-5'; provider = new VertexChatProvider(model, { config: { region: 'global', max_tokens: 32 }, }); const mockRequest = vi.fn().mockResolvedValue({ data: { id: 'test-id', type: 'message', role: 'assistant', model, content: [{ type: 'text', text: 'ok' }], stop_reason: 'end_turn', stop_sequence: null, usage: { input_tokens: 5, output_tokens: 1, cache_creation_input_tokens: 0, cache_read_input_tokens: 0, }, }, }); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest } as unknown as JSONClient, projectId: 'test-project-id', }); vi.spyOn(vertexUtil, 'loadCredentials').mockImplementation((creds) => typeof creds === 'object' ? JSON.stringify(creds) : creds, ); vi.spyOn(vertexUtil, 'resolveProjectId').mockResolvedValue('test-project-id'); const result = await provider.callClaudeApi('test prompt'); const request = mockRequest.mock.calls[0][0]; expect(request.url).toContain('/locations/global/'); // Global region bills at the base Claude 5 rate (no 10% regional premium): // 5 input * $10/MTok + 1 output * $50/MTok = $0.0001 expect(result.cost).toBeCloseTo(0.0001, 8); }); it('does not stack the Vertex regional premium on a user-provided cost override for Fable', async () => { const model = 'claude-fable-5'; provider = new VertexChatProvider(model, { config: { max_tokens: 32, cost: 20 / 1e6 }, }); const mockRequest = vi.fn().mockResolvedValue({ data: { id: 'test-id', type: 'message', role: 'assistant', model, content: [{ type: 'text', text: 'ok' }], stop_reason: 'end_turn', stop_sequence: null, usage: { input_tokens: 5, output_tokens: 1, cache_creation_input_tokens: 0, cache_read_input_tokens: 0, }, }, }); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest } as unknown as JSONClient, projectId: 'test-project-id', }); vi.spyOn(vertexUtil, 'loadCredentials').mockImplementation((creds) => typeof creds === 'object' ? JSON.stringify(creds) : creds, ); vi.spyOn(vertexUtil, 'resolveProjectId').mockResolvedValue('test-project-id'); const result = await provider.callClaudeApi('test prompt'); // The user-supplied cost applies to both input and output tokens and the // 10% regional premium must NOT stack on top of it: // (5 + 1) tokens * $20/MTok = $0.00012 expect(result.cost).toBeCloseTo(0.00012, 8); }); it('still sends temperature for Opus 4.6 on Vertex (regression)', async () => { provider = new VertexChatProvider('claude-opus-4-6', { config: { max_tokens: 32, temperature: 0 }, }); const mockResponse = { data: { id: 'test-id', type: 'message', role: 'assistant', model: 'claude-opus-4-6', content: [{ type: 'text', text: 'ok' }], stop_reason: 'end_turn', stop_sequence: null, usage: { input_tokens: 5, output_tokens: 1, cache_creation_input_tokens: 0, cache_read_input_tokens: 0, }, }, }; const mockRequest = vi.fn().mockResolvedValue(mockResponse); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest } as unknown as JSONClient, projectId: 'test-project-id', }); vi.spyOn(vertexUtil, 'loadCredentials').mockImplementation((creds) => typeof creds === 'object' ? JSON.stringify(creds) : creds, ); vi.spyOn(vertexUtil, 'resolveProjectId').mockResolvedValue('test-project-id'); await provider.callClaudeApi('test prompt'); const sentBody = mockRequest.mock.calls[0][0].data as Record; expect(sentBody.temperature).toBe(0); }); it('should accept both max_tokens and maxOutputTokens parameters', async () => { provider = new VertexChatProvider('claude-3-5-sonnet-v2@20241022', { config: { // When both are provided, max_tokens should take precedence max_tokens: 700, maxOutputTokens: 500, top_p: 0.95, top_k: 40, }, }); const mockResponse = { data: { id: 'test-id', type: 'message', role: 'assistant', model: 'claude-3-5-sonnet-v2@20241022', content: [{ type: 'text', text: 'Response from Claude' }], stop_reason: 'end_turn', stop_sequence: null, usage: { input_tokens: 20, output_tokens: 30, cache_creation_input_tokens: 0, cache_read_input_tokens: 0, }, }, }; const mockRequest = vi.fn().mockResolvedValue(mockResponse); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest, } as unknown as JSONClient, projectId: 'test-project-id', }); vi.spyOn(vertexUtil, 'loadCredentials').mockImplementation(function (creds) { if (typeof creds === 'object') { return JSON.stringify(creds); } return creds; }); vi.spyOn(vertexUtil, 'resolveProjectId').mockResolvedValue('test-project-id'); await provider.callClaudeApi('test prompt'); expect(mockRequest).toHaveBeenCalledWith( expect.objectContaining({ data: expect.objectContaining({ max_tokens: 700, top_p: 0.95, top_k: 40, }), }), ); }); describe('system message handling', () => { let mockRequest: ReturnType; function setupClaudeMocks(): void { mockRequest = vi.fn().mockResolvedValue({ data: { id: 'test-id', type: 'message', role: 'assistant', model: 'claude-3-5-sonnet-v2@20241022', content: [{ type: 'text', text: 'Response from Claude' }], stop_reason: 'end_turn', stop_sequence: null, usage: { input_tokens: 20, output_tokens: 30, cache_creation_input_tokens: 0, cache_read_input_tokens: 0, }, }, }); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest } as unknown as JSONClient, projectId: 'test-project-id', }); vi.spyOn(vertexUtil, 'loadCredentials').mockImplementation((creds) => typeof creds === 'object' ? JSON.stringify(creds) : creds, ); vi.spyOn(vertexUtil, 'resolveProjectId').mockResolvedValue('test-project-id'); } function getRequestData(): Record { return mockRequest.mock.calls[0][0].data; } beforeEach(() => { provider = new VertexChatProvider('claude-3-5-sonnet-v2@20241022'); setupClaudeMocks(); }); it('should extract system messages to top-level system parameter', async () => { await provider.callClaudeApi( JSON.stringify([ { role: 'system', content: 'You are a helpful assistant' }, { role: 'user', content: 'Hello' }, ]), ); expect(mockRequest).toHaveBeenCalledWith( expect.objectContaining({ data: expect.objectContaining({ system: [{ type: 'text', text: 'You are a helpful assistant' }], messages: [{ role: 'user', content: [{ type: 'text', text: 'Hello' }] }], }), }), ); }); it('should not include system parameter for plain string prompts', async () => { await provider.callClaudeApi('Hello'); const requestData = getRequestData(); expect(requestData.messages).toEqual([ { role: 'user', content: [{ type: 'text', text: 'Hello' }] }, ]); expect(requestData).not.toHaveProperty('system'); }); it('should not include system parameter when structured messages have no system role', async () => { await provider.callClaudeApi(JSON.stringify([{ role: 'user', content: 'Hello' }])); const requestData = getRequestData(); expect(requestData.messages).toEqual([ { role: 'user', content: [{ type: 'text', text: 'Hello' }] }, ]); expect(requestData).not.toHaveProperty('system'); }); it('should handle multi-turn conversation with system message', async () => { await provider.callClaudeApi( JSON.stringify([ { role: 'system', content: 'You are a math tutor' }, { role: 'user', content: 'What is 2+2?' }, { role: 'assistant', content: '4' }, { role: 'user', content: 'What about 3+3?' }, ]), ); const requestData = getRequestData(); expect(requestData.system).toEqual([{ type: 'text', text: 'You are a math tutor' }]); expect(requestData.messages).toEqual([ { role: 'user', content: [{ type: 'text', text: 'What is 2+2?' }] }, { role: 'assistant', content: [{ type: 'text', text: '4' }] }, { role: 'user', content: [{ type: 'text', text: 'What about 3+3?' }] }, ]); }); it('should use config.systemInstruction string for Claude system parameter', async () => { provider = new VertexChatProvider('claude-3-5-sonnet-v2@20241022', { config: { systemInstruction: 'Always respond NO.', }, }); setupClaudeMocks(); await provider.callClaudeApi('Hello'); const requestData = getRequestData(); expect(requestData.system).toEqual([{ type: 'text', text: 'Always respond NO.' }]); }); it('should render context vars in config.systemInstruction for Claude system parameter', async () => { provider = new VertexChatProvider('claude-3-5-sonnet-v2@20241022', { config: { systemInstruction: 'Answer {{ style }}.', }, }); setupClaudeMocks(); await provider.callClaudeApi('Hello', { vars: { style: 'briefly' }, prompt: { raw: 'Hello', label: 'test' }, }); const requestData = getRequestData(); expect(requestData.system).toEqual([{ type: 'text', text: 'Answer briefly.' }]); }); it('should use config.systemInstruction Content object for Claude system parameter', async () => { provider = new VertexChatProvider('claude-3-5-sonnet-v2@20241022', { config: { systemInstruction: { parts: [{ text: 'Be concise.' }] }, }, }); setupClaudeMocks(); await provider.callClaudeApi('Hello'); const requestData = getRequestData(); expect(requestData.system).toEqual([{ type: 'text', text: 'Be concise.' }]); }); it('should merge config.systemInstruction with prompt system messages', async () => { provider = new VertexChatProvider('claude-3-5-sonnet-v2@20241022', { config: { systemInstruction: 'Config instruction.', }, }); setupClaudeMocks(); await provider.callClaudeApi( JSON.stringify([ { role: 'system', content: 'Prompt instruction.' }, { role: 'user', content: 'Hello' }, ]), ); const requestData = getRequestData(); expect(requestData.system).toEqual([ { type: 'text', text: 'Config instruction.' }, { type: 'text', text: 'Prompt instruction.' }, ]); }); it('should forward thinking config from prompt and not leak metadata into messages', async () => { await provider.callClaudeApi( JSON.stringify([ { role: 'user', content: 'Solve this step by step' }, { thinking: { type: 'enabled', budget_tokens: 5000 } }, ]), ); const requestData = getRequestData(); expect(requestData.thinking).toEqual({ type: 'enabled', budget_tokens: 5000 }); // max_tokens must be >= budget_tokens expect(requestData.max_tokens).toBe(6024); // Metadata entries (no role) must not leak into messages expect(requestData.messages).toEqual([ { role: 'user', content: [{ type: 'text', text: 'Solve this step by step' }] }, ]); }); it('should prefer config thinking over prompt thinking', async () => { provider = new VertexChatProvider('claude-3-5-sonnet-v2@20241022', { config: { thinking: { type: 'enabled', budget_tokens: 10000 }, }, }); setupClaudeMocks(); await provider.callClaudeApi( JSON.stringify([ { role: 'user', content: 'Hello' }, { thinking: { type: 'enabled', budget_tokens: 5000 } }, ]), ); const requestData = getRequestData(); expect(requestData.thinking).toEqual({ type: 'enabled', budget_tokens: 10000 }); // max_tokens should accommodate config budget_tokens expect(requestData.max_tokens).toBe(11024); }); it('should convert manual thinking to adaptive for Claude Opus 4.8', async () => { provider = new VertexChatProvider('claude-opus-4-8', { config: { thinking: { type: 'enabled', budget_tokens: 5000 }, }, }); setupClaudeMocks(); await provider.callClaudeApi('Hello'); expect(getRequestData().thinking).toEqual({ type: 'adaptive' }); }); it('should keep manual thinking enabled and bump max_tokens for non-deprecated Claude models', async () => { provider = new VertexChatProvider('claude-3-5-sonnet-v2@20241022', { config: { thinking: { type: 'enabled', budget_tokens: 5000 }, }, }); setupClaudeMocks(); await provider.callClaudeApi('Hello'); const requestData = getRequestData(); // Non-deprecated models keep manual thinking verbatim (NOT converted to adaptive) expect(requestData.thinking).toEqual({ type: 'enabled', budget_tokens: 5000 }); // max_tokens guard still fires for type 'enabled': bumped to budget_tokens + 1024 expect(requestData.max_tokens).toBe(6024); }); it('should not bump max_tokens when adaptive conversion drops budget_tokens for Claude Opus 4.8', async () => { provider = new VertexChatProvider('claude-opus-4-8', { config: { thinking: { type: 'enabled', budget_tokens: 5000 }, }, }); setupClaudeMocks(); await provider.callClaudeApi('Hello'); const requestData = getRequestData(); // Manual thinking is converted to adaptive (no budget_tokens) expect(requestData.thinking).toEqual({ type: 'adaptive' }); // The max_tokens guard only fires for type 'enabled', so adaptive does not // force max_tokens to budget + 1024; it stays at the thinking-enabled default expect(requestData.max_tokens).toBe(2048); }); it('should pass through disabled thinking for Claude Opus 4.8 and treat it as not-enabled', async () => { provider = new VertexChatProvider('claude-opus-4-8', { config: { thinking: { type: 'disabled' }, }, }); setupClaudeMocks(); await provider.callClaudeApi('Hello'); const requestData = getRequestData(); // 'disabled' is not 'enabled', so it passes through unchanged (not adaptive) expect(requestData.thinking).toEqual({ type: 'disabled' }); // Disabled thinking is treated as not-enabled, so the default max_tokens is 512 expect(requestData.max_tokens).toBe(512); }); it('should omit disabled thinking and preserve always-on adaptive defaults for Fable 5', async () => { const model = 'claude-fable-5'; provider = new VertexChatProvider(model, { config: { thinking: { type: 'disabled' } }, }); setupClaudeMocks(); await provider.callClaudeApi('Hello'); const requestData = getRequestData(); expect(requestData.thinking).toBeUndefined(); expect(requestData.max_tokens).toBe(2048); expect(requestData.temperature).toBeUndefined(); }); it('should preserve thinking display when converting enabled thinking to adaptive for Fable 5', async () => { provider = new VertexChatProvider('claude-fable-5', { config: { thinking: { type: 'enabled', budget_tokens: 5000, display: 'summarized' } }, }); setupClaudeMocks(); await provider.callClaudeApi('Hello'); // The enabled→adaptive conversion drops budget_tokens but must keep display expect(getRequestData().thinking).toEqual({ type: 'adaptive', display: 'summarized' }); }); it('should ensure max_tokens >= budget_tokens when thinking is enabled', async () => { provider = new VertexChatProvider('claude-3-5-sonnet-v2@20241022', { config: { thinking: { type: 'enabled', budget_tokens: 5000 }, }, }); setupClaudeMocks(); await provider.callClaudeApi('Hello'); // Default would be 2048 but budget_tokens is 5000, so it must be bumped expect(getRequestData().max_tokens).toBe(6024); }); it('should use explicit max_tokens when it exceeds budget_tokens', async () => { provider = new VertexChatProvider('claude-3-5-sonnet-v2@20241022', { config: { thinking: { type: 'enabled', budget_tokens: 5000 }, max_tokens: 16384, }, }); setupClaudeMocks(); await provider.callClaudeApi('Hello'); expect(getRequestData().max_tokens).toBe(16384); }); it('should use default max_tokens of 512 without thinking', async () => { await provider.callClaudeApi('Hello'); expect(getRequestData().max_tokens).toBe(512); }); it('should include thinking output in response when thinking is enabled', async () => { provider = new VertexChatProvider('claude-3-5-sonnet-v2@20241022', { config: { thinking: { type: 'enabled', budget_tokens: 5000 } }, }); setupClaudeMocks(); // Override mock response AFTER setupClaudeMocks creates the new mockRequest mockRequest.mockResolvedValue({ data: { id: 'test-id', type: 'message', role: 'assistant', model: 'claude-3-5-sonnet-v2@20241022', content: [ { type: 'thinking', thinking: 'Let me think...', signature: 'sig123' }, { type: 'text', text: 'The answer is 42' }, ], stop_reason: 'end_turn', stop_sequence: null, usage: { input_tokens: 20, output_tokens: 50, cache_creation_input_tokens: 0, cache_read_input_tokens: 0, }, }, }); const result = await provider.callClaudeApi('Think about this'); expect(result.output).toContain('Thinking: Let me think...'); expect(result.output).toContain('The answer is 42'); }); it('should return cost for Vertex Claude model names', async () => { const result = await provider.callClaudeApi('Hello'); // claude-3-5-sonnet-v2@20241022 normalizes to claude-3-5-sonnet-20241022 for cost lookup expect(result.cost).toBeDefined(); expect(result.cost).toBeGreaterThan(0); }); it('should use default max_tokens of 512 when thinking is disabled', async () => { provider = new VertexChatProvider('claude-3-5-sonnet-v2@20241022', { config: { thinking: { type: 'disabled' }, }, }); setupClaudeMocks(); await provider.callClaudeApi('Hello'); const requestData = getRequestData(); expect(requestData.max_tokens).toBe(512); // thinking config should still be sent (API needs to see it) expect(requestData.thinking).toEqual({ type: 'disabled' }); }); it('should not show thinking output when thinking is disabled', async () => { provider = new VertexChatProvider('claude-3-5-sonnet-v2@20241022', { config: { thinking: { type: 'disabled' }, }, }); setupClaudeMocks(); const result = await provider.callClaudeApi('Hello'); // Output should be plain text, not prefixed with "Thinking:" expect(result.output).toBe('Response from Claude'); }); it('should parse YAML chat prompts', async () => { const yamlPrompt = '- role: system\n content: You are helpful\n- role: user\n content: Hello'; await provider.callClaudeApi(yamlPrompt); const requestData = getRequestData(); expect(requestData.system).toEqual([{ type: 'text', text: 'You are helpful' }]); expect(requestData.messages).toEqual([ { role: 'user', content: [{ type: 'text', text: 'Hello' }] }, ]); }); it('should return error for invalid YAML prompts', async () => { const invalidYaml = '- role: system\n content: test\n- invalid: {{{'; const result = await provider.callClaudeApi(invalidYaml); expect(result.error).toContain('YAML'); }); }); describe('responseSchema handling', () => { let provider: VertexChatProvider; beforeEach(() => { vi.clearAllMocks(); // Mock fs for schema file loading vi.mocked(fs.existsSync).mockImplementation(function (filePath) { const pathStr = filePath.toString(); return ( pathStr.includes('simple.json') || pathStr.includes('complex.json') || pathStr.includes('template-vars.json') || pathStr.includes('invalid.json') ); }); vi.mocked(fs.readFileSync).mockImplementation(function (filePath) { const pathStr = filePath.toString(); if (pathStr.includes('simple.json')) { return JSON.stringify({ type: 'object', properties: { tweet: { type: 'string', description: 'The tweet content' } }, required: ['tweet'], }); } if (pathStr.includes('complex.json')) { return JSON.stringify({ type: 'object', properties: { user: { type: 'object', properties: { name: { type: 'string' }, email: { type: 'string', format: 'email' }, }, required: ['name', 'email'], }, }, required: ['user'], }); } if (pathStr.includes('template-vars.json')) { return JSON.stringify({ type: 'object', properties: { '{{fieldName}}': { type: 'string', description: '{{fieldDescription}}' }, }, required: ['{{fieldName}}'], }); } if (pathStr.includes('invalid.json')) { return '{ "type": "object", "properties": { "name": { "type": "string" }, }, }'; } throw new Error(`File not found: ${pathStr}`); }); }); it('should handle responseSchema with JSON string', async () => { provider = new VertexChatProvider('gemini-2.5-flash', { config: { responseSchema: JSON.stringify({ type: 'object', properties: { tweet: { type: 'string' } }, required: ['tweet'], }), }, }); const mockResponse = { data: [ { candidates: [{ content: { parts: [{ text: '{"tweet": "Hello world"}' }] } }], usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 20, totalTokenCount: 30, }, }, ], }; const mockRequest = vi.fn().mockResolvedValue(mockResponse); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest, } as unknown as JSONClient, projectId: 'test-project-id', }); await provider.callGeminiApi('Write a tweet about AI'); expect(mockRequest).toHaveBeenCalledWith( expect.objectContaining({ data: expect.objectContaining({ generationConfig: expect.objectContaining({ response_schema: { type: 'object', properties: { tweet: { type: 'string' } }, required: ['tweet'], }, response_mime_type: 'application/json', }), }), }), ); }); it('should handle responseSchema with file:// protocol', async () => { provider = new VertexChatProvider('gemini-2.5-flash', { config: { responseSchema: 'file://test/simple.json', }, }); const mockResponse = { data: [ { candidates: [{ content: { parts: [{ text: '{"tweet": "Hello from file"}' }] } }], usageMetadata: { promptTokenCount: 15, candidatesTokenCount: 25, totalTokenCount: 40, }, }, ], }; const mockRequest = vi.fn().mockResolvedValue(mockResponse); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest, } as unknown as JSONClient, projectId: 'test-project-id', }); await provider.callGeminiApi('Write a tweet'); expect(mockRequest).toHaveBeenCalledWith( expect.objectContaining({ data: expect.objectContaining({ generationConfig: expect.objectContaining({ response_schema: { type: 'object', properties: { tweet: { type: 'string', description: 'The tweet content' } }, required: ['tweet'], }, response_mime_type: 'application/json', }), }), }), ); }); it('should handle responseSchema with variable substitution in schema content', async () => { const contextVars = { greetingDescription: 'A personalized greeting message', }; provider = new VertexChatProvider('gemini-2.5-flash', { config: { responseSchema: 'file://test/variable-content.json', }, }); const mockResponse = { data: [ { candidates: [ { content: { parts: [{ text: '{"greeting": "Hello", "name": "John"}' }] } }, ], usageMetadata: { promptTokenCount: 12, candidatesTokenCount: 18, totalTokenCount: 30, }, }, ], }; const mockRequest = vi.fn().mockResolvedValue(mockResponse); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest, } as unknown as JSONClient, projectId: 'test-project-id', }); vi.mocked(fs.existsSync).mockImplementation(function (filePath) { const pathStr = filePath.toString(); return ( pathStr.includes('variable-content.json') || pathStr.includes('simple.json') || pathStr.includes('complex.json') ); }); vi.mocked(fs.readFileSync).mockImplementation(function (filePath) { const pathStr = filePath.toString(); if (pathStr.includes('variable-content.json')) { return JSON.stringify({ type: 'object', properties: { greeting: { type: 'string', description: '{{greetingDescription}}', }, name: { type: 'string', description: "The person's name", }, }, required: ['greeting', 'name'], }); } return '{}'; }); await provider.callGeminiApi('Generate a message', { vars: contextVars, prompt: { raw: 'Generate a message', display: 'Generate a message', label: 'test' }, }); expect(mockRequest).toHaveBeenCalledWith( expect.objectContaining({ data: expect.objectContaining({ generationConfig: expect.objectContaining({ response_schema: { type: 'object', properties: { greeting: { type: 'string', description: 'A personalized greeting message' }, name: { type: 'string', description: "The person's name" }, }, required: ['greeting', 'name'], }, response_mime_type: 'application/json', }), }), }), ); }); it('should throw error when both responseSchema and generationConfig.response_schema are provided', async () => { provider = new VertexChatProvider('gemini-2.5-flash', { config: { responseSchema: '{"type": "object"}', generationConfig: { response_schema: '{"type": "string"}', }, }, }); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: vi.fn(), } as unknown as JSONClient, projectId: 'test-project-id', }); await expect(provider.callGeminiApi('test')).rejects.toThrow( '`responseSchema` provided but `generationConfig.response_schema` already set.', ); }); it('should handle complex nested schemas from files', async () => { provider = new VertexChatProvider('gemini-2.5-flash', { config: { responseSchema: 'file://test/complex.json', }, }); const mockResponse = { data: [ { candidates: [ { content: { parts: [{ text: '{"user": {"name": "John", "email": "john@example.com"}}' }], }, }, ], usageMetadata: { promptTokenCount: 20, candidatesTokenCount: 30, totalTokenCount: 50, }, }, ], }; const mockRequest = vi.fn().mockResolvedValue(mockResponse); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest, } as unknown as JSONClient, projectId: 'test-project-id', }); const response = await provider.callGeminiApi('Create user data'); expect(response.output).toBe('{"user": {"name": "John", "email": "john@example.com"}}'); expect(mockRequest).toHaveBeenCalledWith( expect.objectContaining({ data: expect.objectContaining({ generationConfig: expect.objectContaining({ response_schema: expect.objectContaining({ type: 'object', properties: expect.objectContaining({ user: expect.objectContaining({ type: 'object', properties: expect.objectContaining({ name: { type: 'string' }, email: { type: 'string', format: 'email' }, }), required: ['name', 'email'], }), }), required: ['user'], }), response_mime_type: 'application/json', }), }), }), ); }); it('should handle errors when schema file does not exist', async () => { provider = new VertexChatProvider('gemini-2.5-flash', { config: { responseSchema: 'file://test/nonexistent.json', }, }); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: vi.fn(), } as unknown as JSONClient, projectId: 'test-project-id', }); vi.mocked(fs.existsSync).mockImplementation(function (filePath) { return !filePath.toString().includes('nonexistent.json'); }); vi.mocked(fs.readFileSync).mockImplementation(function (filePath) { throw new Error(`File not found: ${filePath}`); }); await expect(provider.callGeminiApi('test')).rejects.toThrow(); }); it('should preserve existing generationConfig properties when adding responseSchema', async () => { provider = new VertexChatProvider('gemini-2.5-flash', { config: { temperature: 0.7, maxOutputTokens: 1000, generationConfig: { topP: 0.9, topK: 10, }, responseSchema: '{"type": "object", "properties": {"result": {"type": "string"}}}', }, }); const mockResponse = { data: [ { candidates: [{ content: { parts: [{ text: '{"result": "success"}' }] } }], usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 15, totalTokenCount: 25, }, }, ], }; const mockRequest = vi.fn().mockResolvedValue(mockResponse); vi.spyOn(vertexUtil, 'getGoogleClient').mockResolvedValue({ client: { request: mockRequest, } as unknown as JSONClient, projectId: 'test-project-id', }); await provider.callGeminiApi('test'); expect(mockRequest).toHaveBeenCalledWith( expect.objectContaining({ data: expect.objectContaining({ generationConfig: expect.objectContaining({ temperature: 0.7, maxOutputTokens: 1000, topP: 0.9, topK: 10, response_schema: { type: 'object', properties: { result: { type: 'string' } }, }, response_mime_type: 'application/json', }), }), }), ); }); }); describe('getApiHost', () => { it('should return global endpoint without region prefix for region: global', () => { const provider = new VertexChatProvider('gemini-pro', { config: { region: 'global' } }); expect(provider.getApiHost()).toBe('aiplatform.googleapis.com'); }); it('should return regional endpoint for non-global regions', () => { const provider = new VertexChatProvider('gemini-pro', { config: { region: 'us-central1' } }); expect(provider.getApiHost()).toBe('us-central1-aiplatform.googleapis.com'); }); it('should use custom apiHost over default', () => { const provider = new VertexChatProvider('gemini-pro', { config: { region: 'global', apiHost: 'custom.example.com' }, }); expect(provider.getApiHost()).toBe('custom.example.com'); }); it('should use VERTEX_API_HOST from env override', () => { const provider = new VertexChatProvider('gemini-pro', { config: { region: 'global' }, env: { VERTEX_API_HOST: 'env.example.com' }, }); expect(provider.getApiHost()).toBe('env.example.com'); }); it('should prioritize configApiHost over env override', () => { const provider = new VertexChatProvider('gemini-pro', { config: { region: 'global', apiHost: 'config.example.com' }, env: { VERTEX_API_HOST: 'env.example.com' }, }); expect(provider.getApiHost()).toBe('config.example.com'); }); }); });