promptfoo--promptfoo/test/providers/google/provider.test.ts

import * as fs from 'fs';

import { beforeEach, describe, expect, it, vi } from 'vitest';
import * as cache from '../../../src/cache';
import cliState from '../../../src/cliState';
import { GoogleProvider } from '../../../src/providers/google/provider';
import * as util from '../../../src/providers/google/util';
import * as fetchUtil from '../../../src/util/fetch/index';
import { getNunjucksEngineForFilePath } from '../../../src/util/file';
import * as templates from '../../../src/util/templates';
import { mockProcessEnv } from '../../util/utils';

vi.mock('../../../src/cache', async (importOriginal) => {
  return {
    ...(await importOriginal()),
    fetchWithCache: vi.fn(),
  };
});

vi.mock('../../../src/util/fetch/index', async (importOriginal) => {
  return {
    ...(await importOriginal()),
    fetchWithProxy: vi.fn(),
  };
});

vi.mock('../../../src/providers/google/util', async () => ({
  ...(await vi.importActual('../../../src/providers/google/util')),
  maybeCoerceToGeminiFormat: vi.fn(),
  getGoogleClient: vi.fn().mockResolvedValue({
    client: {
      request: vi.fn().mockResolvedValue({
        data: {
          candidates: [{ content: { parts: [{ text: 'test response' }] } }],
          usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 20, totalTokenCount: 30 },
        },
      }),
    },
    credentials: {},
  }),
  loadCredentials: vi.fn().mockReturnValue({}),
  createAuthCacheDiscriminator: vi.fn().mockReturnValue(''),
}));

vi.mock('../../../src/util/templates', async (importOriginal) => {
  return {
    ...(await importOriginal()),
    getNunjucksEngine: vi.fn(() => ({
      renderString: vi.fn((str) => str),
    })),
  };
});

// Hoisted mocks for file loading functions
const mockMaybeLoadToolsFromExternalFile = vi.hoisted(() => vi.fn((input) => input));
const mockMaybeLoadFromExternalFile = vi.hoisted(() => vi.fn((input) => input));

vi.mock('../../../src/util/file', async (importOriginal) => {
  return {
    ...(await importOriginal()),
    getNunjucksEngineForFilePath: vi.fn(),
    maybeLoadToolsFromExternalFile: mockMaybeLoadToolsFromExternalFile,
    maybeLoadFromExternalFile: mockMaybeLoadFromExternalFile,
  };
});

// Also mock the barrel file since the provider imports from util/index
vi.mock('../../../src/util/index', async (importOriginal) => {
  return {
    ...(await importOriginal()),
    maybeLoadToolsFromExternalFile: mockMaybeLoadToolsFromExternalFile,
  };
});

vi.mock('glob', async (importOriginal) => {
  return {
    ...(await importOriginal()),
    globSync: vi.fn().mockReturnValue([]),
  };
});

// Mock envars to control API key availability in tests
vi.mock('../../../src/envars', async (importOriginal) => {
  return {
    ...(await importOriginal()),
    getEnvString: vi.fn().mockReturnValue(undefined),
  };
});

vi.mock('fs', async (importOriginal) => {
  return {
    ...(await importOriginal()),
    existsSync: vi.fn(),
    readFileSync: vi.fn(),
    writeFileSync: vi.fn(),
    statSync: vi.fn(),
  };
});

describe('GoogleProvider', () => {
  beforeEach(() => {
    cliState.config = undefined;
    vi.clearAllMocks();
    // Reset hoisted mocks to default pass-through behavior
    mockMaybeLoadToolsFromExternalFile.mockReset().mockImplementation((input) => input);
    mockMaybeLoadFromExternalFile.mockReset().mockImplementation((input) => input);
    vi.mocked(templates.getNunjucksEngine).mockImplementation(function () {
      return {
        renderString: vi.fn((str) => str),
      } as any;
    });
    vi.mocked(fs.existsSync).mockReset();
    vi.mocked(fs.readFileSync).mockReset();
    vi.mocked(fs.writeFileSync).mockReset();
    vi.mocked(fs.statSync).mockReset();
    vi.mocked(getNunjucksEngineForFilePath).mockImplementation(function () {
      return {
        renderString: vi.fn((str) => str),
      } as any;
    });
  });

  describe('constructor and mode determination', () => {
    it('should default to AI Studio mode (vertexai: false)', () => {
      const provider = new GoogleProvider('gemini-pro', {
        config: { apiKey: 'test-key' },
      });

      expect(provider.id()).toBe('google:gemini-pro');
      expect((provider as any).isVertexMode).toBe(false);
    });

    it('should use Vertex AI mode when vertexai: true', () => {
      const provider = new GoogleProvider('gemini-pro', {
        config: { vertexai: true, projectId: 'my-project' },
      });

      expect(provider.id()).toBe('vertex:gemini-pro');
      expect((provider as any).isVertexMode).toBe(true);
    });

    it('should detect Vertex mode from projectId presence', () => {
      const provider = new GoogleProvider('gemini-pro', {
        config: { projectId: 'my-project' },
      });

      expect(provider.id()).toBe('vertex:gemini-pro');
      expect((provider as any).isVertexMode).toBe(true);
    });

    it('should detect Vertex mode from credentials presence', () => {
      const provider = new GoogleProvider('gemini-pro', {
        config: {
          credentials: JSON.stringify({ client_email: 'test@test.iam.gserviceaccount.com' }),
        },
      });

      expect(provider.id()).toBe('vertex:gemini-pro');
      expect((provider as any).isVertexMode).toBe(true);
    });

    it('should respect explicit vertexai: false even with projectId', () => {
      const provider = new GoogleProvider('gemini-pro', {
        config: {
          vertexai: false,
          projectId: 'my-project',
          apiKey: 'test-key',
        },
      });

      expect(provider.id()).toBe('google:gemini-pro');
      expect((provider as any).isVertexMode).toBe(false);
    });
  });

  describe('AI Studio mode', () => {
    let provider: GoogleProvider;

    beforeEach(() => {
      provider = new GoogleProvider('gemini-pro', {
        config: {
          apiKey: 'test-key',
          temperature: 0.7,
          maxOutputTokens: 100,
        },
      });
    });

    it('should resolve API endpoint correctly', () => {
      const endpoint = provider.getApiEndpoint('generateContent');
      expect(endpoint).toContain('/v1beta/models/gemini-pro:generateContent');
      expect(endpoint).toContain('generativelanguage.googleapis.com');
    });

    it('should use v1alpha for thinking models', () => {
      const thinkingProvider = new GoogleProvider('gemini-2.0-flash-thinking-exp', {
        config: { apiKey: 'test-key' },
      });

      const endpoint = thinkingProvider.getApiEndpoint('generateContent');
      expect(endpoint).toContain('/v1alpha/');
    });

    it('should use v1beta for gemini-3 models', () => {
      // Regression: dash-named gemini-3-* models resolve to v1beta, the same
      // as dotted gemini-3.x IDs. v1beta is Google's primary Gemini 3 endpoint.
      const gemini3Provider = new GoogleProvider('gemini-3-pro', {
        config: { apiKey: 'test-key' },
      });

      const endpoint = gemini3Provider.getApiEndpoint('generateContent');
      expect(endpoint).toContain('/v1beta/');
    });

    it('should use v1beta for gemini-3.1 models', () => {
      const gemini31Provider = new GoogleProvider('gemini-3.1-pro-preview', {
        config: { apiKey: 'test-key' },
      });

      const endpoint = gemini31Provider.getApiEndpoint('generateContent');
      expect(endpoint).toContain('/v1beta/');
    });

    it('should allow explicit apiVersion override in AI Studio mode', () => {
      // Test that config.apiVersion takes precedence over auto-detection
      const overrideProvider = new GoogleProvider('gemini-pro', {
        config: {
          apiKey: 'test-key',
          apiVersion: 'v1', // Override default v1beta
        },
      });

      const endpoint = overrideProvider.getApiEndpoint('generateContent');
      expect(endpoint).toContain('/v1/models/gemini-pro:generateContent');
      expect(endpoint).not.toContain('v1beta');
    });

    it('should use custom apiHost in endpoint', () => {
      const customProvider = new GoogleProvider('gemini-pro', {
        config: {
          apiKey: 'test-key',
          apiHost: 'custom.host.com',
        },
      });

      const endpoint = customProvider.getApiEndpoint('generateContent');
      expect(endpoint).toContain('https://custom.host.com');
    });

    it('should pass API key in x-goog-api-key header', async () => {
      vi.mocked(cache.fetchWithCache).mockResolvedValueOnce({
        data: {
          candidates: [{ content: { parts: [{ text: 'test response' }] } }],
          usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 20, totalTokenCount: 30 },
        },
        cached: false,
        status: 200,
        statusText: 'OK',
      });

      await provider.callApi('test prompt');

      const calledUrl = vi.mocked(cache.fetchWithCache).mock.calls[0][0] as string;
      expect(calledUrl).not.toContain('?key=');
      expect(calledUrl).not.toContain('&key=');

      const calledOptions = vi.mocked(cache.fetchWithCache).mock.calls[0][1] as any;
      expect(calledOptions.headers['x-goog-api-key']).toBe('test-key');
    });

    it('should throw error when API key is missing in AI Studio mode', async () => {
      // Delete all possible API key env vars
      mockProcessEnv({ GEMINI_API_KEY: undefined });
      mockProcessEnv({ GOOGLE_API_KEY: undefined });
      mockProcessEnv({ PALM_API_KEY: undefined });
      mockProcessEnv({ VERTEX_API_KEY: undefined });
      // Also delete project-related env vars that would trigger Vertex mode detection
      mockProcessEnv({ GOOGLE_PROJECT_ID: undefined });
      mockProcessEnv({ VERTEX_PROJECT_ID: undefined });
      mockProcessEnv({ GOOGLE_CLOUD_PROJECT: undefined });
      mockProcessEnv({ GOOGLE_GENAI_USE_VERTEXAI: undefined });

      // Explicitly set vertexai: false to ensure AI Studio mode regardless of env vars
      const noKeyProvider = new GoogleProvider('gemini-pro', {
        config: { vertexai: false },
      });

      await expect(noKeyProvider.callApi('test prompt')).rejects.toThrow(
        'Google API key is not set. Set the GOOGLE_API_KEY or GEMINI_API_KEY environment variable or add `apiKey` to the provider config.',
      );
    });

    it('should call API and return response', async () => {
      vi.mocked(cache.fetchWithCache).mockResolvedValueOnce({
        data: {
          candidates: [{ content: { parts: [{ text: 'Hello, world!' }] } }],
          usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 5, totalTokenCount: 15 },
        },
        cached: false,
        status: 200,
        statusText: 'OK',
      });

      const result = await provider.callApi('test prompt');

      expect(result.error).toBeUndefined();
      expect(result.output).toBe('Hello, world!');
      expect(result.tokenUsage).toEqual({
        prompt: 10,
        completion: 5,
        total: 15,
        numRequests: 1,
      });
    });

    it('should handle safety blocked response', async () => {
      vi.mocked(cache.fetchWithCache).mockResolvedValueOnce({
        data: {
          promptFeedback: {
            blockReason: 'SAFETY',
            safetyRatings: [{ category: 'HARM_CATEGORY_HARASSMENT', probability: 'HIGH' }],
          },
          usageMetadata: { promptTokenCount: 10, totalTokenCount: 10 },
        },
        cached: false,
        status: 200,
        statusText: 'OK',
      });

      const result = await provider.callApi('test prompt');

      expect(result.guardrails?.flagged).toBe(true);
      expect(result.guardrails?.flaggedInput).toBe(true);
    });
  });

  describe('Vertex AI mode', () => {
    describe('OAuth mode', () => {
      let provider: GoogleProvider;

      beforeEach(() => {
        provider = new GoogleProvider('gemini-pro', {
          config: {
            vertexai: true,
            projectId: 'my-project',
            region: 'us-central1',
          },
        });
      });

      it('should resolve API endpoint correctly', () => {
        const endpoint = provider.getApiEndpoint('generateContent');
        expect(endpoint).toContain('us-central1-aiplatform.googleapis.com');
        expect(endpoint).toContain('/publishers/google/models/gemini-pro:generateContent');
      });

      it('should use global region endpoint when region is global', () => {
        const globalProvider = new GoogleProvider('gemini-pro', {
          config: {
            vertexai: true,
            projectId: 'my-project',
            region: 'global',
          },
        });

        expect(globalProvider.getApiHost()).toBe('aiplatform.googleapis.com');
      });

      it('should call API using Google client for OAuth mode', async () => {
        const getProjectIdSpy = vi
          .spyOn(provider as any, 'getProjectId')
          .mockResolvedValue('my-project');

        try {
          await provider.callApi('test prompt');
        } finally {
          getProjectIdSpy.mockRestore();
        }

        expect(vi.mocked(util.getGoogleClient)).toHaveBeenCalled();
      });
    });

    describe('Express mode', () => {
      let provider: GoogleProvider;

      beforeEach(() => {
        provider = new GoogleProvider('gemini-pro', {
          config: {
            vertexai: true,
            apiKey: 'vertex-api-key',
            // expressMode not needed - automatic when API key is present
          },
        });
      });

      it('should use express mode automatically when API key is present', () => {
        // Express mode is invisible to users - just provide an API key and it works
        expect((provider as any).isExpressMode()).toBe(true);
      });

      it('should not use express mode when no API key is available', () => {
        const noApiKeyProvider = new GoogleProvider('gemini-pro', {
          config: {
            vertexai: true,
            projectId: 'my-project',
            // No API key - will use OAuth/ADC
          },
        });

        expect((noApiKeyProvider as any).isExpressMode()).toBe(false);
      });

      it('should not use express mode when expressMode: false (opt-out)', () => {
        // Users can explicitly opt-out if they need OAuth features
        const noExpressProvider = new GoogleProvider('gemini-pro', {
          config: {
            vertexai: true,
            apiKey: 'vertex-api-key',
            expressMode: false,
          },
        });

        expect((noExpressProvider as any).isExpressMode()).toBe(false);
      });

      it('should pass API key in header for express mode', async () => {
        const mockResponse = {
          ok: true,
          json: vi.fn().mockResolvedValue({
            candidates: [{ content: { parts: [{ text: 'test response' }] } }],
            usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 20, totalTokenCount: 30 },
          }),
        };
        vi.mocked(fetchUtil.fetchWithProxy).mockResolvedValueOnce(mockResponse as any);

        await provider.callApi('test prompt');

        const calledUrl = vi.mocked(fetchUtil.fetchWithProxy).mock.calls[0][0] as string;
        expect(calledUrl).not.toContain('?key=');
        expect(calledUrl).not.toContain('&key=');

        const calledOptions = vi.mocked(fetchUtil.fetchWithProxy).mock.calls[0][1] as any;
        expect(calledOptions.headers['x-goog-api-key']).toBe('vertex-api-key');
      });

      it('should use aiplatform.googleapis.com endpoint for express mode', async () => {
        const mockResponse = {
          ok: true,
          json: vi.fn().mockResolvedValue({
            candidates: [{ content: { parts: [{ text: 'test response' }] } }],
            usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 20, totalTokenCount: 30 },
          }),
        };
        vi.mocked(fetchUtil.fetchWithProxy).mockResolvedValueOnce(mockResponse as any);

        await provider.callApi('test prompt');

        const calledUrl = vi.mocked(fetchUtil.fetchWithProxy).mock.calls[0][0] as string;
        expect(calledUrl).toContain('aiplatform.googleapis.com');
        expect(calledUrl).toContain('/publishers/google/models/gemini-pro:generateContent');
      });
    });
  });

  describe('response parsing', () => {
    let provider: GoogleProvider;

    beforeEach(() => {
      provider = new GoogleProvider('gemini-pro', {
        config: { apiKey: 'test-key' },
      });
    });

    it('should handle cached response', async () => {
      vi.mocked(cache.fetchWithCache).mockResolvedValueOnce({
        data: {
          candidates: [{ content: { parts: [{ text: 'cached response' }] } }],
          usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 5, totalTokenCount: 15 },
        },
        cached: true,
        status: 200,
        statusText: 'OK',
      });

      const result = await provider.callApi('test prompt');

      expect(result.cached).toBe(true);
      expect(result.tokenUsage).toEqual({
        cached: 15,
        total: 15,
        numRequests: 0,
      });
    });

    it('should extract grounding metadata from response', async () => {
      vi.mocked(cache.fetchWithCache).mockResolvedValueOnce({
        data: {
          candidates: [
            {
              content: { parts: [{ text: 'response with grounding' }] },
              groundingMetadata: { searchQueries: ['test query'] },
              webSearchQueries: ['test search'],
            },
          ],
          usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 5, totalTokenCount: 15 },
        },
        cached: false,
        status: 200,
        statusText: 'OK',
      });

      const result = await provider.callApi('test prompt');

      expect(result.metadata?.groundingMetadata).toEqual({ searchQueries: ['test query'] });
      expect(result.metadata?.webSearchQueries).toEqual(['test search']);
    });

    it('should handle Model Armor block reason', async () => {
      vi.mocked(cache.fetchWithCache).mockResolvedValueOnce({
        data: {
          promptFeedback: {
            blockReason: 'MODEL_ARMOR',
            blockReasonMessage: 'Content blocked by Model Armor',
          },
          usageMetadata: { promptTokenCount: 10, totalTokenCount: 10 },
        },
        cached: false,
        status: 200,
        statusText: 'OK',
      });

      const result = await provider.callApi('test prompt');

      expect(result.guardrails?.flagged).toBe(true);
      expect(result.metadata?.modelArmor?.blockReason).toBe('MODEL_ARMOR');
    });

    it('should handle MAX_TOKENS finish reason as success', async () => {
      vi.mocked(cache.fetchWithCache).mockResolvedValueOnce({
        data: {
          candidates: [
            {
              content: { parts: [{ text: 'truncated response' }] },
              finishReason: 'MAX_TOKENS',
            },
          ],
          usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 100, totalTokenCount: 110 },
        },
        cached: false,
        status: 200,
        statusText: 'OK',
      });

      const result = await provider.callApi('test prompt');

      expect(result.error).toBeUndefined();
      expect(result.output).toBe('truncated response');
    });

    it('should ignore metadata-only chunks in streaming responses', async () => {
      vi.mocked(cache.fetchWithCache).mockResolvedValueOnce({
        data: [
          {
            candidates: [{ content: { parts: [{ text: 'streamed response' }] } }],
          },
          {
            usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 5, totalTokenCount: 15 },
          },
        ],
        cached: false,
        status: 200,
        statusText: 'OK',
      });

      const result = await provider.callApi('test prompt');

      expect(result.error).toBeUndefined();
      expect(result.output).toBe('streamed response');
      expect(result.tokenUsage).toEqual({
        prompt: 10,
        completion: 5,
        total: 15,
        numRequests: 1,
      });
    });

    it('should preserve prompt safety ratings from separate streaming chunks', async () => {
      vi.mocked(cache.fetchWithCache).mockResolvedValueOnce({
        data: [
          {
            promptFeedback: {
              safetyRatings: [{ category: 'HARM_CATEGORY_HARASSMENT', probability: 'HIGH' }],
            },
          },
          {
            candidates: [
              {
                content: { parts: [{ text: 'streamed response' }] },
                safetyRatings: [
                  { category: 'HARM_CATEGORY_HARASSMENT', probability: 'NEGLIGIBLE' },
                ],
              },
            ],
          },
          {
            usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 5, totalTokenCount: 15 },
          },
        ],
        cached: false,
        status: 200,
        statusText: 'OK',
      });

      const result = await provider.callApi('test prompt');

      expect(result.error).toBeUndefined();
      expect(result.output).toBe('streamed response');
      expect(result.guardrails).toEqual({
        flaggedInput: true,
        flaggedOutput: false,
        flagged: true,
      });
    });

    it('should accumulate incremental chunks in streaming responses', async () => {
      vi.mocked(cache.fetchWithCache).mockResolvedValueOnce({
        data: [
          {
            candidates: [{ content: { parts: [{ text: 'Hello ' }] } }],
          },
          {
            candidates: [{ content: { parts: [{ text: 'world' }] } }],
          },
          {
            usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 2, totalTokenCount: 12 },
          },
        ],
        cached: false,
        status: 200,
        statusText: 'OK',
      });

      const result = await provider.callApi('test prompt');

      expect(result.error).toBeUndefined();
      expect(result.output).toBe('Hello world');
    });

    it('should not mutate raw multipart chunks when accumulating streaming responses', async () => {
      const firstPart = { functionCall: { name: 'look_up', args: { query: 'weather' } } };
      const firstChunk = {
        candidates: [{ content: { parts: [firstPart] } }],
      };
      const responseData = [
        firstChunk,
        {
          candidates: [{ content: { parts: [{ text: 'done' }] } }],
        },
        {
          usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 2, totalTokenCount: 12 },
        },
      ];
      vi.mocked(cache.fetchWithCache).mockResolvedValueOnce({
        data: responseData,
        cached: false,
        status: 200,
        statusText: 'OK',
      });

      const result = await provider.callApi('test prompt');

      expect(result.error).toBeUndefined();
      expect(result.output).toEqual([firstPart, { text: 'done' }]);
      expect(firstChunk.candidates[0].content.parts).toEqual([firstPart]);
      expect(result.raw).toBe(responseData);
    });

    it('should reject streaming responses that never provide output', async () => {
      vi.mocked(cache.fetchWithCache).mockResolvedValueOnce({
        data: [
          {
            usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 0, totalTokenCount: 10 },
          },
        ],
        cached: false,
        status: 200,
        statusText: 'OK',
      });

      const result = await provider.callApi('test prompt');

      expect(result.error).toContain('No output found in response');
    });

    it('should return an error for safety finish reasons outside scorable evaluations', async () => {
      const responseData = {
        candidates: [{ content: { parts: [{ text: '' }] }, finishReason: 'SAFETY' }],
        usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 0, totalTokenCount: 10 },
      };
      vi.mocked(cache.fetchWithCache).mockResolvedValueOnce({
        data: responseData,
        cached: true,
        status: 200,
        statusText: 'OK',
      });

      const result = await provider.callApi('test prompt');

      expect(result).toEqual(
        expect.objectContaining({
          error: 'Content was blocked due to safety settings with finish reason: SAFETY.',
          guardrails: expect.objectContaining({
            flagged: true,
            flaggedOutput: true,
          }),
          cached: true,
          raw: responseData,
        }),
      );
    });

    it('should return safety finish reasons as output during redteam evaluations', async () => {
      cliState.config = { redteam: {} } as any;
      vi.mocked(cache.fetchWithCache).mockResolvedValueOnce({
        data: {
          candidates: [{ content: { parts: [{ text: '' }] }, finishReason: 'SAFETY' }],
          usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 0, totalTokenCount: 10 },
        },
        cached: false,
        status: 200,
        statusText: 'OK',
      });

      const result = await provider.callApi('test prompt');

      expect(result.error).toBeUndefined();
      expect(result.output).toBe(
        'Content was blocked due to safety settings with finish reason: SAFETY.',
      );
      expect(result.guardrails?.flagged).toBe(true);
      expect(result.cached).toBe(false);
      expect(result.raw).toEqual({
        candidates: [{ content: { parts: [{ text: '' }] }, finishReason: 'SAFETY' }],
        usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 0, totalTokenCount: 10 },
      });
    });

    it('should expose safety finish reasons to guardrails assertions', async () => {
      vi.mocked(cache.fetchWithCache).mockResolvedValueOnce({
        data: {
          candidates: [{ content: { parts: [{ text: '' }] }, finishReason: 'SAFETY' }],
          usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 0, totalTokenCount: 10 },
        },
        cached: false,
        status: 200,
        statusText: 'OK',
      });

      const result = await provider.callApi('test prompt', {
        prompt: { raw: 'test prompt', label: 'test prompt' },
        vars: {},
        test: { assert: [{ type: 'not-guardrails' }] },
      } as any);

      expect(result.error).toBeUndefined();
      expect(result.output).toBe(
        'Content was blocked due to safety settings with finish reason: SAFETY.',
      );
      expect(result.guardrails?.flagged).toBe(true);
    });

    it('should expose safety finish reasons for exported redteam test metadata', async () => {
      vi.mocked(cache.fetchWithCache).mockResolvedValueOnce({
        data: {
          candidates: [{ content: { parts: [{ text: '' }] }, finishReason: 'SAFETY' }],
          usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 0, totalTokenCount: 10 },
        },
        cached: false,
        status: 200,
        statusText: 'OK',
      });

      const result = await provider.callApi('test prompt', {
        prompt: { raw: 'test prompt', label: 'test prompt' },
        vars: {},
        test: {
          metadata: { pluginId: 'ascii-smuggling', goal: 'exfiltrate data' },
        },
      } as any);

      expect(result.error).toBeUndefined();
      expect(result.output).toBe(
        'Content was blocked due to safety settings with finish reason: SAFETY.',
      );
      expect(result.guardrails?.flagged).toBe(true);
    });

    it('should expose safety finish reasons for nested redteam assertions', async () => {
      vi.mocked(cache.fetchWithCache).mockResolvedValueOnce({
        data: {
          candidates: [{ content: { parts: [{ text: '' }] }, finishReason: 'SAFETY' }],
          usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 0, totalTokenCount: 10 },
        },
        cached: false,
        status: 200,
        statusText: 'OK',
      });

      const result = await provider.callApi('test prompt', {
        prompt: { raw: 'test prompt', label: 'test prompt' },
        vars: {},
        test: {
          assert: [
            {
              type: 'assert-set',
              assert: [{ type: 'promptfoo:redteam:ascii-smuggling' }],
            },
          ],
        },
      } as any);

      expect(result.error).toBeUndefined();
      expect(result.output).toBe(
        'Content was blocked due to safety settings with finish reason: SAFETY.',
      );
      expect(result.guardrails?.flagged).toBe(true);
    });

    it('should handle thinking tokens in response', async () => {
      vi.mocked(cache.fetchWithCache).mockResolvedValueOnce({
        data: {
          candidates: [{ content: { parts: [{ text: 'thinking response' }] } }],
          usageMetadata: {
            promptTokenCount: 10,
            candidatesTokenCount: 20,
            totalTokenCount: 30,
            thoughtsTokenCount: 100,
          },
        },
        cached: false,
        status: 200,
        statusText: 'OK',
      });

      const result = await provider.callApi('test prompt');

      expect(result.tokenUsage?.completionDetails?.reasoning).toBe(100);
    });
  });

  describe('cost calculation', () => {
    it('should return cost for AI Studio mode with known model', async () => {
      const provider = new GoogleProvider('gemini-pro', {
        config: { apiKey: 'test-key' },
      });

      vi.mocked(cache.fetchWithCache).mockResolvedValueOnce({
        data: {
          candidates: [{ content: { parts: [{ text: 'response' }] } }],
          usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 5, totalTokenCount: 15 },
        },
        cached: false,
        status: 200,
        statusText: 'OK',
      });

      const result = await provider.callApi('test prompt');

      // gemini-pro: input=0.5/1e6, output=1.5/1e6
      // cost = 0.5e-6 * 10 + 1.5e-6 * 5 = 1.25e-5
      expect(result.cost).toBeCloseTo(1.25e-5, 10);
    });

    it('should return cost for Vertex AI mode with known model', async () => {
      const provider = new GoogleProvider('gemini-pro', {
        config: { vertexai: true, apiKey: 'test-vertex-key' },
      });

      vi.mocked(fetchUtil.fetchWithProxy).mockResolvedValueOnce({
        ok: true,
        json: vi.fn().mockResolvedValue({
          candidates: [{ content: { parts: [{ text: 'response' }] } }],
          usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 20, totalTokenCount: 30 },
        }),
      } as any);

      const result = await provider.callApi('test prompt');
      // gemini-pro: input 0.5/1e6, output 1.5/1e6
      // 10 prompt + 20 completion = 0.000035
      expect(result.cost).toBeCloseTo(0.000035, 10);
    });

    it('should use Vertex-specific pricing when it differs from AI Studio', async () => {
      const provider = new GoogleProvider('gemini-2.0-flash', {
        config: { vertexai: true, apiKey: 'test-vertex-key' },
      });

      vi.mocked(fetchUtil.fetchWithProxy).mockResolvedValueOnce({
        ok: true,
        json: vi.fn().mockResolvedValue({
          candidates: [{ content: { parts: [{ text: 'response' }] } }],
          usageMetadata: {
            promptTokenCount: 1000,
            candidatesTokenCount: 500,
            totalTokenCount: 1500,
          },
        }),
      } as any);

      const result = await provider.callApi('test prompt');
      // Vertex pricing for gemini-2.0-flash: input $0.15/1M, output $0.60/1M
      // 1000 * 0.15/1e6 + 500 * 0.60/1e6 = 0.00045
      expect(result.cost).toBeCloseTo(0.00045, 10);
    });

    it('should return undefined cost for cached responses', async () => {
      const provider = new GoogleProvider('gemini-pro', {
        config: { apiKey: 'test-key' },
      });

      vi.mocked(cache.fetchWithCache).mockResolvedValueOnce({
        data: {
          candidates: [{ content: { parts: [{ text: 'cached response' }] } }],
          usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 5, totalTokenCount: 15 },
        },
        cached: true,
        status: 200,
        statusText: 'OK',
      });

      const result = await provider.callApi('test prompt');

      expect(result.cost).toBeUndefined();
    });

    it('should use tiered pricing when prompt tokens exceed threshold', async () => {
      const provider = new GoogleProvider('gemini-2.5-pro', {
        config: { apiKey: 'test-key' },
      });

      vi.mocked(cache.fetchWithCache).mockResolvedValueOnce({
        data: {
          candidates: [{ content: { parts: [{ text: 'response' }] } }],
          usageMetadata: {
            promptTokenCount: 250_000,
            candidatesTokenCount: 1000,
            totalTokenCount: 251_000,
          },
        },
        cached: false,
        status: 200,
        statusText: 'OK',
      });

      const result = await provider.callApi('test prompt');

      // gemini-2.5-pro tiered: input=2.5/1e6, output=15.0/1e6 (above 200k threshold)
      // cost = 2.5e-6 * 250000 + 15.0e-6 * 1000 = 0.625 + 0.015 = 0.64
      expect(result.cost).toBeCloseTo(0.64, 5);
    });

    it('should use standard pricing when prompt tokens are below threshold', async () => {
      const provider = new GoogleProvider('gemini-2.5-pro', {
        config: { apiKey: 'test-key' },
      });

      vi.mocked(cache.fetchWithCache).mockResolvedValueOnce({
        data: {
          candidates: [{ content: { parts: [{ text: 'response' }] } }],
          usageMetadata: {
            promptTokenCount: 100_000,
            candidatesTokenCount: 1000,
            totalTokenCount: 101_000,
          },
        },
        cached: false,
        status: 200,
        statusText: 'OK',
      });

      const result = await provider.callApi('test prompt');

      // gemini-2.5-pro standard: input=1.25/1e6, output=10.0/1e6 (below 200k threshold)
      // cost = 1.25e-6 * 100000 + 10.0e-6 * 1000 = 0.125 + 0.01 = 0.135
      expect(result.cost).toBeCloseTo(0.135, 5);
    });

    it('should use config.cost override when provided', async () => {
      const provider = new GoogleProvider('gemini-pro', {
        config: { apiKey: 'test-key', cost: 0.001 },
      });

      vi.mocked(cache.fetchWithCache).mockResolvedValueOnce({
        data: {
          candidates: [{ content: { parts: [{ text: 'response' }] } }],
          usageMetadata: { promptTokenCount: 100, candidatesTokenCount: 50, totalTokenCount: 150 },
        },
        cached: false,
        status: 200,
        statusText: 'OK',
      });

      const result = await provider.callApi('test prompt');

      // config.cost=0.001 applied to both input and output
      // cost = 0.001 * 100 + 0.001 * 50 = 0.15
      expect(result.cost).toBeCloseTo(0.15, 5);
    });

    it('should include thinking tokens in cost calculation', async () => {
      const provider = new GoogleProvider('gemini-2.5-flash', {
        config: { apiKey: 'test-key' },
      });

      vi.mocked(cache.fetchWithCache).mockResolvedValueOnce({
        data: {
          candidates: [{ content: { parts: [{ text: 'response' }] } }],
          usageMetadata: {
            promptTokenCount: 10,
            candidatesTokenCount: 5,
            totalTokenCount: 315,
            thoughtsTokenCount: 300,
          },
        },
        cached: false,
        status: 200,
        statusText: 'OK',
      });

      const result = await provider.callApi('test prompt');

      // gemini-2.5-flash: input=0.3/1e6, output=2.5/1e6
      // completionForCost = candidatesTokenCount + thoughtsTokenCount = 5 + 300 = 305
      // cost = 0.3e-6 * 10 + 2.5e-6 * 305 = 0.000003 + 0.0007625 = 0.0007655
      expect(result.cost).toBeCloseTo(0.0007655, 10);
    });

    it('should not double-count when thoughtsTokenCount is zero', async () => {
      const provider = new GoogleProvider('gemini-2.5-flash', {
        config: { apiKey: 'test-key' },
      });

      vi.mocked(cache.fetchWithCache).mockResolvedValueOnce({
        data: {
          candidates: [{ content: { parts: [{ text: 'response' }] } }],
          usageMetadata: {
            promptTokenCount: 10,
            candidatesTokenCount: 5,
            totalTokenCount: 15,
            thoughtsTokenCount: 0,
          },
        },
        cached: false,
        status: 200,
        statusText: 'OK',
      });

      const result = await provider.callApi('test prompt');

      // gemini-2.5-flash: input=0.3/1e6, output=2.5/1e6
      // completionForCost = 5 + 0 = 5
      // cost = 0.3e-6 * 10 + 2.5e-6 * 5 = 0.000003 + 0.0000125 = 0.0000155
      expect(result.cost).toBeCloseTo(0.0000155, 10);
    });
  });

  describe('tool handling', () => {
    it('should include tools in request body', async () => {
      const provider = new GoogleProvider('gemini-pro', {
        config: {
          apiKey: 'test-key',
          tools: [
            {
              functionDeclarations: [
                {
                  name: 'test_function',
                  description: 'A test function',
                  parameters: { type: 'OBJECT', properties: {} },
                },
              ],
            },
          ],
        },
      });

      vi.mocked(cache.fetchWithCache).mockResolvedValueOnce({
        data: {
          candidates: [{ content: { parts: [{ text: 'response' }] } }],
          usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 5, totalTokenCount: 15 },
        },
        cached: false,
        status: 200,
        statusText: 'OK',
      });

      await provider.callApi('test prompt');

      const calledOptions = vi.mocked(cache.fetchWithCache).mock.calls[0][1] as any;
      const body = JSON.parse(calledOptions.body);
      expect(body.tools).toBeDefined();
      expect(body.tools[0].functionDeclarations[0].name).toBe('test_function');
    });

    it('should include toolConfig in request body', async () => {
      const provider = new GoogleProvider('gemini-pro', {
        config: {
          apiKey: 'test-key',
          toolConfig: {
            functionCallingConfig: {
              mode: 'ANY',
            },
          },
        },
      });

      vi.mocked(cache.fetchWithCache).mockResolvedValueOnce({
        data: {
          candidates: [{ content: { parts: [{ text: 'response' }] } }],
          usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 5, totalTokenCount: 15 },
        },
        cached: false,
        status: 200,
        statusText: 'OK',
      });

      await provider.callApi('test prompt');

      const calledOptions = vi.mocked(cache.fetchWithCache).mock.calls[0][1] as any;
      const body = JSON.parse(calledOptions.body);
      expect(body.toolConfig).toEqual({ functionCallingConfig: { mode: 'ANY' } });
    });

    it.each([
      [{ tool_choice: 'none' }],
      [{ toolConfig: { functionCallingConfig: { mode: 'none' } } }],
      [{ tool_config: { function_calling_config: { mode: 'none' } } }],
    ])('should enforce documented no-tools controls: %j', async (toolDisableConfig) => {
      const callback = vi.fn(async () => 'callback result');
      const provider = new GoogleProvider('gemini-pro', {
        config: {
          apiKey: 'test-key',
          tools: [
            {
              functionDeclarations: [
                {
                  name: 'test_function',
                  description: 'Test function',
                  parameters: { type: 'OBJECT', properties: {} },
                },
              ],
            },
          ],
          functionToolCallbacks: { test_function: callback },
          ...toolDisableConfig,
        } as any,
      });

      vi.mocked(cache.fetchWithCache).mockResolvedValueOnce({
        data: {
          candidates: [
            {
              content: {
                parts: [
                  {
                    text: JSON.stringify({
                      functionCall: { name: 'test_function', args: {} },
                    }),
                  },
                ],
              },
            },
          ],
          usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 5, totalTokenCount: 15 },
        },
        cached: false,
        status: 200,
        statusText: 'OK',
      });

      const result = await provider.callApi('test prompt');

      const calledOptions = vi.mocked(cache.fetchWithCache).mock.calls.at(-1)?.[1] as any;
      const body = JSON.parse(calledOptions.body);
      expect(body.toolConfig).toEqual({ functionCallingConfig: { mode: 'NONE' } });
      expect(body.tools).toBeUndefined();
      expect(callback).not.toHaveBeenCalled();
      expect(result.output).toBe(
        JSON.stringify({ functionCall: { name: 'test_function', args: {} } }),
      );
    });

    it('should fall back to tool_choice when explicit toolConfig is invalid', async () => {
      const provider = new GoogleProvider('gemini-pro', {
        config: {
          apiKey: 'test-key',
          toolConfig: { functionCallingConfig: { mode: 'invalid' as any } },
          tool_choice: 'required',
        },
      });

      await provider.callApi('test prompt');

      const calledOptions = vi.mocked(cache.fetchWithCache).mock.calls.at(-1)?.[1] as any;
      const body = JSON.parse(calledOptions.body);
      expect(body.toolConfig).toEqual({ functionCallingConfig: { mode: 'ANY' } });
    });

    it('should honor prompt-level snake_case no-tools overrides over provider toolConfig', async () => {
      const provider = new GoogleProvider('gemini-pro', {
        config: {
          apiKey: 'test-key',
          tools: [
            {
              functionDeclarations: [
                {
                  name: 'test_function',
                  description: 'Test function',
                  parameters: { type: 'OBJECT', properties: {} },
                },
              ],
            },
          ],
          toolConfig: { functionCallingConfig: { mode: 'AUTO' } },
        },
      });

      await provider.callApi('test prompt', {
        prompt: {
          config: {
            tool_config: { function_calling_config: { mode: 'none' } },
          },
        },
      } as any);

      const calledOptions = vi.mocked(cache.fetchWithCache).mock.calls.at(-1)?.[1] as any;
      const body = JSON.parse(calledOptions.body);
      expect(body.toolConfig).toEqual({ functionCallingConfig: { mode: 'NONE' } });
      expect(body.tools).toBeUndefined();
    });

    it('should preserve non-function Google tools when function calling is disabled', async () => {
      const provider = new GoogleProvider('gemini-pro', {
        config: {
          apiKey: 'test-key',
          tools: [{ googleSearch: {} }],
          toolConfig: { functionCallingConfig: { mode: 'NONE' } },
        },
      });

      await provider.callApi('test prompt');

      const calledOptions = vi.mocked(cache.fetchWithCache).mock.calls.at(-1)?.[1] as any;
      const body = JSON.parse(calledOptions.body);
      expect(body.toolConfig).toEqual({ functionCallingConfig: { mode: 'NONE' } });
      expect(body.tools).toEqual([{ googleSearch: {} }]);
    });

    it('should skip executable tool files while preserving inline non-function tools when disabled', async () => {
      const provider = new GoogleProvider('gemini-pro', {
        config: {
          apiKey: 'test-key',
          tool_choice: 'none',
          tools: [{ googleSearch: {} }, 'file://tools.js:getTools'] as any,
        },
      });

      await provider.callApi('test prompt');

      expect(mockMaybeLoadToolsFromExternalFile).toHaveBeenCalledWith(
        [{ googleSearch: {} }],
        undefined,
      );
      const calledOptions = vi.mocked(cache.fetchWithCache).mock.calls.at(-1)?.[1] as any;
      const body = JSON.parse(calledOptions.body);
      expect(body.toolConfig).toEqual({ functionCallingConfig: { mode: 'NONE' } });
      expect(body.tools).toEqual([{ googleSearch: {} }]);
    });

    it('should preserve non-function tools loaded from data files when disabled', async () => {
      mockMaybeLoadToolsFromExternalFile.mockResolvedValueOnce([
        {
          functionDeclarations: [
            {
              name: 'get_weather',
              description: 'Get weather information',
              parameters: { type: 'OBJECT' as const, properties: {} },
            },
          ],
        },
        { googleSearch: {} },
      ]);
      const provider = new GoogleProvider('gemini-pro', {
        config: {
          apiKey: 'test-key',
          tool_choice: 'none',
          tools: 'file://tools.json' as any,
        },
      });

      await provider.callApi('test prompt');

      expect(mockMaybeLoadToolsFromExternalFile).toHaveBeenCalledWith(
        'file://tools.json',
        undefined,
      );
      const calledOptions = vi.mocked(cache.fetchWithCache).mock.calls.at(-1)?.[1] as any;
      const body = JSON.parse(calledOptions.body);
      expect(body.toolConfig).toEqual({ functionCallingConfig: { mode: 'NONE' } });
      expect(body.tools).toEqual([{ googleSearch: {} }]);
    });

    it('should preserve supported explicit Google toolConfig fields', async () => {
      const provider = new GoogleProvider('gemini-pro', {
        config: {
          apiKey: 'test-key',
          toolConfig: {
            functionCallingConfig: {
              mode: 'VALIDATED',
              streamFunctionCallArguments: true,
            },
          },
        },
      });

      await provider.callApi('test prompt');

      const calledOptions = vi.mocked(cache.fetchWithCache).mock.calls.at(-1)?.[1] as any;
      const body = JSON.parse(calledOptions.body);
      expect(body.toolConfig).toEqual({
        functionCallingConfig: {
          mode: 'VALIDATED',
          streamFunctionCallArguments: true,
        },
      });
    });
  });

  describe('error handling', () => {
    let provider: GoogleProvider;

    beforeEach(() => {
      provider = new GoogleProvider('gemini-pro', {
        config: { apiKey: 'test-key' },
      });
    });

    it('should return error for API call failure', async () => {
      vi.mocked(cache.fetchWithCache).mockRejectedValueOnce(new Error('Network error'));

      const result = await provider.callApi('test prompt');

      expect(result.error).toContain('API call error');
      expect(result.error).toContain('Network error');
    });

    it('should return error for API error response', async () => {
      vi.mocked(cache.fetchWithCache).mockResolvedValueOnce({
        data: {
          error: {
            code: 400,
            message: 'Invalid request',
          },
        },
        cached: false,
        status: 400,
        statusText: 'Bad Request',
      });

      const result = await provider.callApi('test prompt');

      expect(result.error).toContain('Error 400: Invalid request');
    });

    it('should handle express mode API error response', async () => {
      const provider = new GoogleProvider('gemini-pro', {
        config: {
          vertexai: true,
          apiKey: 'vertex-api-key',
          expressMode: true, // Explicit for test clarity; auto-enabled when apiKey is present
        },
      });

      const mockResponse = {
        ok: false,
        status: 401,
        statusText: 'Unauthorized',
        json: vi.fn().mockResolvedValue({ error: { message: 'Invalid API key' } }),
      };
      vi.mocked(fetchUtil.fetchWithProxy).mockResolvedValueOnce(mockResponse as any);

      const result = await provider.callApi('test prompt');

      expect(result.error).toContain('API call error: 401 Unauthorized');
    });
  });

  describe('cleanup', () => {
    it('should have cleanup method', () => {
      const provider = new GoogleProvider('gemini-pro', {
        config: { apiKey: 'test-key' },
      });

      expect(typeof provider.cleanup).toBe('function');
    });
  });
});