promptfoo--promptfoo/test/guardrails.test.ts

import { beforeEach, describe, expect, it, vi } from 'vitest';
import { fetchWithCache } from '../src/cache';
import { cloudConfig } from '../src/globalConfig/cloud';
import guardrails, { type AdaptiveRequest } from '../src/guardrails';

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

vi.mock('../src/globalConfig/cloud', () => ({
  cloudConfig: {
    isEnabled: vi.fn(),
    getApiHost: vi.fn(),
    getApiKey: vi.fn(),
  },
}));

describe('guardrails', () => {
  const mockFetchResponse = {
    data: {
      model: 'test-model',
      results: [
        {
          categories: {
            test_category: true,
          },
          category_scores: {
            test_category: 0.95,
          },
          flagged: true,
        },
      ],
    },
    cached: false,
    status: 200,
    statusText: 'OK',
  };

  beforeEach(() => {
    vi.clearAllMocks();
    vi.mocked(fetchWithCache).mockResolvedValue(mockFetchResponse);
    // Default to the non-cloud path so existing assertions stay deterministic
    // regardless of the local machine's cloud-login state.
    vi.mocked(cloudConfig.isEnabled).mockReturnValue(false);
    vi.mocked(cloudConfig.getApiHost).mockReturnValue('https://api.promptfoo.app');
    vi.mocked(cloudConfig.getApiKey).mockReturnValue(undefined);
  });

  describe('guard', () => {
    it('should make a request to the guard endpoint', async () => {
      const input = 'test input';
      await guardrails.guard(input);

      expect(fetchWithCache).toHaveBeenCalledWith(
        expect.stringContaining('/v1/guard'),
        {
          method: 'POST',
          headers: {
            'Content-Type': 'application/json',
          },
          body: JSON.stringify({ input }),
        },
        undefined,
        'json',
      );
    });

    it('should return the parsed guard result', async () => {
      const result = await guardrails.guard('test input');
      expect(result).toEqual(mockFetchResponse.data);
    });

    it('should handle API errors', async () => {
      const errorMessage = 'API Error';
      vi.mocked(fetchWithCache).mockRejectedValue(new Error(errorMessage));

      await expect(guardrails.guard('test input')).rejects.toThrow(errorMessage);
    });

    it('should handle empty API response', async () => {
      vi.mocked(fetchWithCache).mockResolvedValue({
        data: null,
        cached: false,
        status: 200,
        statusText: 'OK',
      });

      await expect(guardrails.guard('test input')).rejects.toThrow('No data returned from API');
    });
  });

  describe('pii', () => {
    const mockPiiResponse = {
      data: {
        model: 'test-model',
        results: [
          {
            categories: {
              pii: true,
            },
            category_scores: {
              pii: 1,
            },
            flagged: true,
            payload: {
              pii: [
                {
                  entity_type: 'EMAIL',
                  start: 0,
                  end: 17,
                  pii: 'test@example.com',
                },
              ],
            },
          },
        ],
      },
      cached: false,
      status: 200,
      statusText: 'OK',
    };

    beforeEach(() => {
      vi.mocked(fetchWithCache).mockResolvedValue(mockPiiResponse);
    });

    it('should make a request to the pii endpoint', async () => {
      const input = 'test@example.com';
      await guardrails.pii(input);

      expect(fetchWithCache).toHaveBeenCalledWith(
        expect.stringContaining('/v1/pii'),
        {
          method: 'POST',
          headers: {
            'Content-Type': 'application/json',
          },
          body: JSON.stringify({ input }),
        },
        undefined,
        'json',
      );
    });

    it('should return the parsed PII result', async () => {
      const result = await guardrails.pii('test@example.com');
      expect(result).toEqual(mockPiiResponse.data);
      expect(result.results[0].payload?.pii?.[0].entity_type).toBe('EMAIL');
    });

    it('should handle API errors', async () => {
      const errorMessage = 'API Error';
      vi.mocked(fetchWithCache).mockRejectedValue(new Error(errorMessage));

      await expect(guardrails.pii('test input')).rejects.toThrow(errorMessage);
    });
  });

  describe('harm', () => {
    const mockHarmResponse = {
      data: {
        model: 'test-model',
        results: [
          {
            categories: {
              hate: true,
              violent_crimes: false,
            },
            category_scores: {
              hate: 0.95,
              violent_crimes: 0.1,
            },
            flagged: true,
          },
        ],
      },
      cached: false,
      status: 200,
      statusText: 'OK',
    };

    beforeEach(() => {
      vi.mocked(fetchWithCache).mockResolvedValue(mockHarmResponse);
    });

    it('should make a request to the harm endpoint', async () => {
      const input = 'test input';
      await guardrails.harm(input);

      expect(fetchWithCache).toHaveBeenCalledWith(
        expect.stringContaining('/v1/harm'),
        {
          method: 'POST',
          headers: {
            'Content-Type': 'application/json',
          },
          body: JSON.stringify({ input }),
        },
        undefined,
        'json',
      );
    });

    it('should return the parsed harm result', async () => {
      const result = await guardrails.harm('test input');
      expect(result).toEqual(mockHarmResponse.data);
      expect(result.results[0].categories.hate).toBe(true);
      expect(result.results[0].category_scores.hate).toBe(0.95);
    });

    it('should handle API errors', async () => {
      const errorMessage = 'API Error';
      vi.mocked(fetchWithCache).mockRejectedValue(new Error(errorMessage));

      await expect(guardrails.harm('test input')).rejects.toThrow(errorMessage);
    });
  });

  describe('response structure', () => {
    it('should have correct guard result structure', async () => {
      const result = await guardrails.guard('test input');
      expect(result).toHaveProperty('model');
      expect(result).toHaveProperty('results');
      expect(Array.isArray(result.results)).toBe(true);
      expect(result.results[0]).toHaveProperty('categories');
      expect(result.results[0]).toHaveProperty('category_scores');
      expect(result.results[0]).toHaveProperty('flagged');
      expect(typeof result.results[0].flagged).toBe('boolean');
    });

    it('should have correct PII result structure with payload', async () => {
      vi.mocked(fetchWithCache).mockResolvedValue({
        data: {
          model: 'test-model',
          results: [
            {
              categories: {
                pii: true,
              },
              category_scores: {
                pii: 1,
              },
              flagged: true,
              payload: {
                pii: [
                  {
                    entity_type: 'EMAIL',
                    start: 0,
                    end: 17,
                    pii: 'test@example.com',
                  },
                ],
              },
            },
          ],
        },
        cached: false,
        status: 200,
        statusText: 'OK',
      });

      const result = await guardrails.pii('test input');
      expect(result).toHaveProperty('model');
      expect(result).toHaveProperty('results');
      expect(Array.isArray(result.results)).toBe(true);
      expect(result.results[0]).toHaveProperty('payload');
      expect(Array.isArray(result.results[0].payload?.pii)).toBe(true);

      const piiEntity = result.results[0].payload?.pii?.[0];
      expect(piiEntity).toBeDefined();
      expect(piiEntity).toEqual(
        expect.objectContaining({
          entity_type: expect.any(String),
          start: expect.any(Number),
          end: expect.any(Number),
          pii: expect.any(String),
        }),
      );
    });
  });

  describe('adaptive function', () => {
    it('should call fetchWithCache with correct parameters', async () => {
      const mockResponse = {
        data: {
          model: 'promptfoo-adaptive-prompt',
          adaptedPrompt: 'Adapted test input',
          modifications: [
            {
              type: 'substitution',
              reason: 'Policy compliance',
              original: 'test input',
              modified: 'Adapted test input',
            },
          ],
        },
        cached: false,
        status: 200,
        statusText: 'OK',
      };
      vi.mocked(fetchWithCache).mockResolvedValue(mockResponse);

      const request: AdaptiveRequest = {
        prompt: 'test input',
        policies: ['No harmful content'],
      };

      const result = await guardrails.adaptive(request);
      expect(fetchWithCache).toHaveBeenCalledWith(
        'https://api.promptfoo.app/v1/adaptive',
        {
          method: 'POST',
          headers: {
            'Content-Type': 'application/json',
          },
          body: JSON.stringify({
            prompt: 'test input',
            policies: ['No harmful content'],
          }),
        },
        undefined,
        'json',
      );
      expect(result).toEqual(mockResponse.data);
    });

    it('should handle missing policies parameter', async () => {
      const mockResponse = {
        data: {
          model: 'promptfoo-adaptive-prompt',
          adaptedPrompt: 'Adapted test input',
          modifications: [],
        },
        cached: false,
        status: 200,
        statusText: 'OK',
      };
      vi.mocked(fetchWithCache).mockResolvedValue(mockResponse);

      const request: AdaptiveRequest = {
        prompt: 'test input',
      };

      const result = await guardrails.adaptive(request);
      expect(fetchWithCache).toHaveBeenCalledWith(
        'https://api.promptfoo.app/v1/adaptive',
        {
          method: 'POST',
          headers: {
            'Content-Type': 'application/json',
          },
          body: JSON.stringify({
            prompt: 'test input',
            policies: [],
          }),
        },
        undefined,
        'json',
      );
      expect(result).toEqual(mockResponse.data);
    });
  });

  describe('on-prem / cloud-enabled host resolution', () => {
    const ONPREM_HOST = 'https://onprem.example.com';
    const ONPREM_KEY = 'test-onprem-key';

    beforeEach(() => {
      vi.mocked(cloudConfig.isEnabled).mockReturnValue(true);
      vi.mocked(cloudConfig.getApiHost).mockReturnValue(ONPREM_HOST);
      vi.mocked(cloudConfig.getApiKey).mockReturnValue(ONPREM_KEY);
    });

    it.each([
      ['guard', () => guardrails.guard('test input'), '/v1/guard'],
      ['pii', () => guardrails.pii('test input'), '/v1/pii'],
      ['harm', () => guardrails.harm('test input'), '/v1/harm'],
    ])('routes %s to the configured cloud host with a bearer token', async (_name, call, path) => {
      await call();

      expect(fetchWithCache).toHaveBeenCalledWith(
        `${ONPREM_HOST}${path}`,
        expect.objectContaining({
          method: 'POST',
          headers: expect.objectContaining({ Authorization: `Bearer ${ONPREM_KEY}` }),
        }),
        undefined,
        'json',
      );
      const [calledUrl] = vi.mocked(fetchWithCache).mock.calls[0];
      expect(calledUrl).not.toContain('api.promptfoo.app');
    });

    it('routes adaptive requests to the configured cloud host with a bearer token', async () => {
      vi.mocked(fetchWithCache).mockResolvedValue({
        data: { model: 'm', adaptedPrompt: 'a', modifications: [] },
        cached: false,
        status: 200,
        statusText: 'OK',
      });

      await guardrails.adaptive({ prompt: 'test input', policies: ['No harmful content'] });

      expect(fetchWithCache).toHaveBeenCalledWith(
        `${ONPREM_HOST}/v1/adaptive`,
        expect.objectContaining({
          method: 'POST',
          headers: expect.objectContaining({ Authorization: `Bearer ${ONPREM_KEY}` }),
        }),
        undefined,
        'json',
      );
      const [calledUrl] = vi.mocked(fetchWithCache).mock.calls[0];
      expect(calledUrl).not.toContain('api.promptfoo.app');
    });

    it('strips a trailing slash on the configured host (no //v1)', async () => {
      vi.mocked(cloudConfig.getApiHost).mockReturnValue('https://onprem.example.com/');

      await guardrails.guard('test input');

      const [calledUrl] = vi.mocked(fetchWithCache).mock.calls[0];
      expect(calledUrl).toBe('https://onprem.example.com/v1/guard');
    });

    it('falls back to the public share host (no auth) when cloud is not enabled', async () => {
      vi.mocked(cloudConfig.isEnabled).mockReturnValue(false);

      await guardrails.guard('test input');

      const [calledUrl, opts] = vi.mocked(fetchWithCache).mock.calls[0];
      expect(calledUrl).toBe('https://api.promptfoo.app/v1/guard');
      // getApiHost() must not be consulted on the non-cloud path, and no token is sent.
      expect(cloudConfig.getApiHost).not.toHaveBeenCalled();
      expect((opts?.headers as Record<string, string>)?.Authorization).toBeUndefined();
    });

    it('lets an explicit PROMPTFOO_REMOTE_API_BASE_URL override win, without leaking the token', async () => {
      // Regression guard: logged-in users who redirect guardrails to a private
      // endpoint must keep that override (codex P2 on the original PR). The cloud
      // bearer token must NOT be sent to the override host.
      vi.stubEnv('PROMPTFOO_REMOTE_API_BASE_URL', 'https://guardrails-override.example.com');
      vi.mocked(cloudConfig.isEnabled).mockReturnValue(true);
      vi.mocked(cloudConfig.getApiHost).mockReturnValue(ONPREM_HOST);

      try {
        await guardrails.guard('test input');

        const [calledUrl, opts] = vi.mocked(fetchWithCache).mock.calls[0];
        expect(calledUrl).toBe('https://guardrails-override.example.com/v1/guard');
        expect(calledUrl).not.toContain('onprem.example.com');
        expect((opts?.headers as Record<string, string>)?.Authorization).toBeUndefined();
      } finally {
        vi.unstubAllEnvs();
      }
    });
  });
});