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)?.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)?.Authorization).toBeUndefined(); } finally { vi.unstubAllEnvs(); } }); }); });