promptfoo--promptfoo/test/providers/xai/responses.test.ts

import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import { XAIResponsesProvider } from '../../../src/providers/xai/responses';

const mockMaybeLoadToolsFromExternalFile = vi.hoisted(() => vi.fn());

const mockFetchWithCache = vi.hoisted(() => vi.fn());
const mockFetchWithProxy = vi.hoisted(() => vi.fn());

const DEFAULT_TEST_MODEL = 'grok-4.3';
const createMockResponseData = (model: string = DEFAULT_TEST_MODEL) => ({
  id: 'resp_123',
  model,
  output: [
    {
      type: 'message',
      role: 'assistant',
      content: [{ type: 'output_text', text: 'hello' }],
    },
  ],
  usage: {
    input_tokens: 10,
    output_tokens: 5,
    total_tokens: 15,
  },
});

class TestableXAIResponsesProvider extends XAIResponsesProvider {
  public getResolvedApiKey(): string | undefined {
    return this.getApiKey();
  }

  public getResolvedApiUrl(): string {
    return this.getApiUrl();
  }
}

vi.mock('../../../src/cache', async (importOriginal) => ({
  ...(await importOriginal()),
  fetchWithCache: (...args: any[]) => mockFetchWithCache(...args),
}));

vi.mock('../../../src/util/fetch/index', async (importOriginal) => ({
  ...(await importOriginal()),
  fetchWithProxy: (...args: any[]) => mockFetchWithProxy(...args),
}));

vi.mock('../../../src/util/index', async (importOriginal) => ({
  ...(await importOriginal()),
  maybeLoadToolsFromExternalFile: mockMaybeLoadToolsFromExternalFile,
}));

vi.mock('../../../src/logger');

function createSSEStream(events: string) {
  const encoder = new TextEncoder();
  return new ReadableStream({
    start(controller) {
      controller.enqueue(encoder.encode(events));
      controller.close();
    },
  });
}

describe('XAIResponsesProvider', () => {
  beforeEach(() => {
    vi.clearAllMocks();
    mockFetchWithCache.mockReset();
    mockFetchWithProxy.mockReset();
    // Default passthrough: return whatever tools array is passed in
    mockMaybeLoadToolsFromExternalFile.mockImplementation((tools: any) => Promise.resolve(tools));
    mockFetchWithCache.mockResolvedValue({
      data: createMockResponseData(),
      cached: false,
      status: 200,
      statusText: 'OK',
    });
  });

  afterEach(() => {
    vi.resetAllMocks();
  });

  it('calls maybeLoadToolsFromExternalFile when tools are configured', async () => {
    const tools = [
      { type: 'code_execution' },
      {
        type: 'file_search',
        vector_store_ids: ['collection_123'],
        max_num_results: 3,
      },
    ];
    const provider = new XAIResponsesProvider('grok-4.3', {
      config: {
        apiKey: 'test-key',
        tools,
      },
    });

    mockMaybeLoadToolsFromExternalFile.mockResolvedValue(tools);

    await provider.callApi('hello');

    expect(mockMaybeLoadToolsFromExternalFile).toHaveBeenCalledTimes(1);
    expect(mockMaybeLoadToolsFromExternalFile).toHaveBeenCalledWith(tools, undefined);
  });

  it('builds current xAI Responses tool payloads', async () => {
    const provider = new XAIResponsesProvider('grok-4.3', {
      config: {
        apiKey: 'test-key',
        tools: [
          { type: 'code_execution' },
          {
            type: 'file_search',
            vector_store_ids: ['collection_123'],
            max_num_results: 3,
          },
          {
            type: 'collections_search',
            collection_ids: ['collection_456'],
            max_num_results: 2,
          },
          {
            type: 'mcp',
            server_url: 'https://mcp.example.com',
            server_label: 'example',
            server_description: 'Example server',
            authorization: 'Bearer token',
          },
        ],
        include: ['reasoning.encrypted_content'],
        reasoning: { effort: 'high' },
        max_tool_calls: 4,
      },
    });

    await provider.callApi('hello');

    const [, request] = mockFetchWithCache.mock.calls[0];
    const body = JSON.parse(request.body);

    expect(body.tools).toEqual([
      { type: 'code_execution' },
      {
        type: 'file_search',
        vector_store_ids: ['collection_123'],
        max_num_results: 3,
      },
      {
        type: 'collections_search',
        collection_ids: ['collection_456'],
        max_num_results: 2,
      },
      {
        type: 'mcp',
        server_url: 'https://mcp.example.com',
        server_label: 'example',
        server_description: 'Example server',
        authorization: 'Bearer token',
      },
    ]);
    expect(body.include).toEqual(['reasoning.encrypted_content']);
    expect(body.reasoning).toEqual({ effort: 'high' });
    expect(body.max_tool_calls).toBe(4);
  });

  it('uses regional endpoints when configured', () => {
    const provider = new TestableXAIResponsesProvider('grok-4.3', {
      config: {
        region: 'eu-west-1',
      },
    });

    expect(provider.getResolvedApiUrl()).toBe('https://eu-west-1.api.x.ai/v1');
  });

  it('honors env overrides for authentication and base URL', () => {
    const provider = new TestableXAIResponsesProvider('grok-4.3', {
      env: {
        XAI_API_KEY: 'env-key',
        XAI_API_BASE_URL: 'https://env.api.x.ai/v1',
      },
    });

    expect(provider.getResolvedApiKey()).toBe('env-key');
    expect(provider.getResolvedApiUrl()).toBe('https://env.api.x.ai/v1');
  });

  it('uses the exact billed xAI cost when the API returns cost ticks', async () => {
    mockFetchWithCache.mockResolvedValueOnce({
      data: {
        id: 'resp_123',
        model: 'grok-4.3',
        output: [
          {
            type: 'message',
            role: 'assistant',
            content: [{ type: 'output_text', text: 'hello' }],
          },
        ],
        usage: {
          input_tokens: 10,
          output_tokens: 5,
          total_tokens: 15,
          cost_in_usd_ticks: 12_500_000_000,
        },
      },
      cached: false,
      status: 200,
      statusText: 'OK',
    });

    const provider = new XAIResponsesProvider('grok-4.3', {
      config: { apiKey: 'test-key' },
    });

    const result = await provider.callApi('hello');

    expect(result.cost).toBe(1.25);
  });

  it('uses fallback pricing for Grok 4.20 responses', async () => {
    mockFetchWithCache.mockResolvedValueOnce({
      data: createMockResponseData('grok-4.20'),
      cached: false,
      status: 200,
      statusText: 'OK',
    });

    const provider = new XAIResponsesProvider('grok-4.20', {
      config: { apiKey: 'test-key' },
    });

    const result = await provider.callApi('hello');

    expect(result.cost).toBeCloseTo(0.000025, 8);
  });

  it('prefers billed cost ticks over calculated cost when reasoning tokens are present', async () => {
    mockFetchWithCache.mockResolvedValueOnce({
      data: {
        ...createMockResponseData(),
        id: 'resp_124',
        output: [
          {
            type: 'message',
            role: 'assistant',
            content: [{ type: 'output_text', text: 'hello with reasoning' }],
          },
        ],
        usage: {
          input_tokens: 10,
          output_tokens: 5,
          total_tokens: 15,
          output_tokens_details: {
            reasoning_tokens: 100,
          },
          cost_in_usd_ticks: 12_500_000_000,
        },
      },
      cached: false,
      status: 200,
      statusText: 'OK',
    });

    const provider = new XAIResponsesProvider(DEFAULT_TEST_MODEL, {
      config: { apiKey: 'test-key' },
    });

    const result = await provider.callApi('hello');

    expect(result.cost).toBe(1.25);
  });

  it('falls back to calculated xAI pricing when cost ticks are absent', async () => {
    mockFetchWithCache.mockResolvedValueOnce({
      data: {
        id: 'resp_124',
        model: 'grok-4.3',
        output: [
          {
            type: 'message',
            role: 'assistant',
            content: [{ type: 'output_text', text: 'hello' }],
          },
        ],
        usage: {
          input_tokens: 10,
          output_tokens: 5,
          total_tokens: 18,
          output_tokens_details: {
            reasoning_tokens: 3,
          },
        },
      },
      cached: false,
      status: 200,
      statusText: 'OK',
    });

    const provider = new XAIResponsesProvider('grok-4.3', {
      config: { apiKey: 'test-key' },
    });

    const result = await provider.callApi('hello');

    // output_tokens (5) already includes the 3 reasoning tokens, so they are
    // not billed twice: 10 input @ $1.25/M + 5 output @ $2.50/M.
    expect(result.cost).toBeCloseTo(0.000025, 8);
    expect(result.tokenUsage?.completionDetails?.reasoning).toBe(3);
  });

  it('applies cache-read pricing when cost ticks are absent', async () => {
    mockFetchWithCache.mockResolvedValueOnce({
      data: {
        id: 'resp_cached',
        model: 'grok-4.3',
        output: [
          {
            type: 'message',
            role: 'assistant',
            content: [{ type: 'output_text', text: 'hello' }],
          },
        ],
        usage: {
          input_tokens: 10,
          output_tokens: 5,
          total_tokens: 15,
          input_tokens_details: {
            cached_tokens: 8,
          },
        },
      },
      cached: false,
      status: 200,
      statusText: 'OK',
    });

    const provider = new XAIResponsesProvider('grok-4.3', {
      config: { apiKey: 'test-key' },
    });

    const result = await provider.callApi('hello');

    // 2 uncached input @ $1.25/M + 8 cached input @ $0.20/M + 5 output @ $2.50/M.
    expect(result.cost).toBeCloseTo(0.0000166, 10);
    expect(result.tokenUsage?.completionDetails?.cacheReadInputTokens).toBe(8);
  });

  it('honors explicit cache-read pricing overrides', async () => {
    mockFetchWithCache.mockResolvedValueOnce({
      data: {
        id: 'resp_cached_override',
        model: 'grok-4.3',
        output: [
          {
            type: 'message',
            role: 'assistant',
            content: [{ type: 'output_text', text: 'hello' }],
          },
        ],
        usage: {
          input_tokens: 10,
          output_tokens: 5,
          total_tokens: 15,
          input_tokens_details: { cached_tokens: 8 },
        },
      },
      cached: false,
      status: 200,
      statusText: 'OK',
    });

    const provider = new XAIResponsesProvider('grok-4.3', {
      config: {
        apiKey: 'test-key',
        inputCost: 3e-6,
        outputCost: 15e-6,
        cacheReadCost: 0.75e-6,
      },
    });

    const result = await provider.callApi('hello');

    // 2 uncached input @ $3/M + 8 cached input @ $0.75/M + 5 output @ $15/M.
    expect(result.cost).toBeCloseTo(0.000087, 10);
  });

  it('keeps fallback xAI pricing when input tokens are zero', async () => {
    mockFetchWithCache.mockResolvedValueOnce({
      data: {
        id: 'resp_zero_input',
        model: 'grok-4.3',
        output: [
          {
            type: 'message',
            role: 'assistant',
            content: [{ type: 'output_text', text: 'hello' }],
          },
        ],
        usage: {
          input_tokens: 0,
          output_tokens: 5,
          total_tokens: 8,
          output_tokens_details: {
            reasoning_tokens: 3,
          },
        },
      },
      cached: false,
      status: 200,
      statusText: 'OK',
    });

    const provider = new XAIResponsesProvider('grok-4.3', {
      config: { apiKey: 'test-key' },
    });

    const result = await provider.callApi('hello');

    // 0 input @ $1.25/M + 5 output @ $2.50/M (output_tokens already includes
    // the 3 reasoning tokens, so reasoning is not added on top).
    expect(result.cost).toBeCloseTo(0.0000125, 10);
  });

  it('parses streamed Responses API events', async () => {
    mockFetchWithProxy.mockResolvedValueOnce({
      status: 200,
      statusText: 'OK',
      body: createSSEStream(
        [
          'data: {"type":"response.output_text.delta","delta":"hel"}',
          '',
          'data: {"type":"response.output_text.delta","delta":"lo"}',
          '',
          'data: {"type":"response.completed","response":{"id":"resp_stream","model":"grok-4.3","output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"hello"}]}],"usage":{"input_tokens":10,"output_tokens":5,"total_tokens":15}}}',
          '',
          'data: [DONE]',
          '',
        ].join('\n'),
      ),
    });

    const provider = new XAIResponsesProvider('grok-4.3', {
      config: { apiKey: 'test-key', stream: true },
    });

    const result = await provider.callApi('hello');

    expect(mockFetchWithCache).not.toHaveBeenCalled();
    expect(mockFetchWithProxy).toHaveBeenCalledWith(
      'https://api.x.ai/v1/responses',
      expect.objectContaining({
        body: expect.stringContaining('"stream":true'),
      }),
    );
    expect(result.output).toBe('hello');
    expect(result.tokenUsage).toEqual({
      total: 15,
      prompt: 10,
      completion: 5,
      numRequests: 1,
    });
  });

  it('ignores malformed SSE events when later streamed output is valid', async () => {
    mockFetchWithProxy.mockResolvedValueOnce({
      status: 200,
      statusText: 'OK',
      body: createSSEStream(
        [
          'data: {"type":"response.output_text.delta","delta":"hel"}',
          '',
          'data: not-json',
          '',
          'data: {"type":"response.output_text.delta","delta":"lo"}',
          '',
          'data: {"type":"response.completed","response":{"id":"resp_stream","model":"grok-4.3","output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"final from completed"}]}],"usage":{"input_tokens":10,"output_tokens":5,"total_tokens":15}}}',
          '',
          'data: [DONE]',
          '',
        ].join('\n'),
      ),
    });

    const provider = new XAIResponsesProvider('grok-4.3', {
      config: { apiKey: 'test-key', stream: true },
    });

    const result = await provider.callApi('hello');

    expect(result.output).toBe('final from completed');
  });

  it('accumulates nested output_text deltas from streamed Responses API events', async () => {
    mockFetchWithProxy.mockResolvedValueOnce({
      status: 200,
      statusText: 'OK',
      body: createSSEStream(
        [
          'data: {"type":"response.output_text.delta","output_text":{"delta":"hel"}}',
          '',
          'data: {"type":"response.output_text.delta","output_text":{"delta":"lo"}}',
          '',
          'data: [DONE]',
          '',
        ].join('\n'),
      ),
    });

    const provider = new XAIResponsesProvider('grok-4.3', {
      config: { apiKey: 'test-key', stream: true },
    });

    const result = await provider.callApi('hello');

    expect(result.output).toBe('hello');
  });

  it('ignores streamed reasoning summary deltas when rebuilding output text', async () => {
    mockFetchWithProxy.mockResolvedValueOnce({
      status: 200,
      statusText: 'OK',
      body: createSSEStream(
        [
          'data: {"type":"response.reasoning_summary_text.delta","delta":"internal reasoning"}',
          '',
          'data: {"type":"response.output_text.delta","delta":"final answer"}',
          '',
          'data: [DONE]',
          '',
        ].join('\n'),
      ),
    });

    const provider = new XAIResponsesProvider('grok-4.3', {
      config: { apiKey: 'test-key', stream: true },
    });

    const result = await provider.callApi('hello');

    expect(result.output).toBe('final answer');
  });

  it('prefers streamed final-answer deltas over completed reasoning summaries', async () => {
    mockFetchWithProxy.mockResolvedValueOnce({
      status: 200,
      statusText: 'OK',
      body: createSSEStream(
        [
          'data: {"type":"response.reasoning_summary_text.delta","delta":"internal reasoning"}',
          '',
          'data: {"type":"response.output_text.delta","delta":"final answer"}',
          '',
          'data: {"type":"response.completed","response":{"id":"resp_stream","model":"grok-4.3","output":[{"type":"reasoning","summary":[{"text":"internal reasoning"}]},{"type":"message","role":"assistant","content":[{"type":"output_text","text":"final answer"}]}],"usage":{"input_tokens":10,"output_tokens":5,"total_tokens":15}}}',
          '',
          'data: [DONE]',
          '',
        ].join('\n'),
      ),
    });

    const provider = new XAIResponsesProvider('grok-4.3', {
      config: { apiKey: 'test-key', stream: true },
    });

    const result = await provider.callApi('hello');

    expect(result.output).toBe('final answer');
    expect(result.output).not.toContain('Reasoning:');
    expect(result.tokenUsage).toEqual({
      total: 15,
      prompt: 10,
      completion: 5,
      numRequests: 1,
    });
  });

  it('preserves encrypted reasoning output in raw streamed responses without exposing the summary', async () => {
    mockFetchWithProxy.mockResolvedValueOnce({
      status: 200,
      statusText: 'OK',
      body: createSSEStream(
        [
          'data: {"type":"response.output_text.delta","delta":"final answer"}',
          '',
          'data: {"type":"response.completed","response":{"id":"resp_stream","model":"grok-4.3","output":[{"type":"reasoning","summary":[{"text":"internal reasoning"}],"encrypted_content":"opaque-payload"},{"type":"message","role":"assistant","content":[{"type":"output_text","text":"final answer"}]}],"usage":{"input_tokens":10,"output_tokens":5,"total_tokens":15}}}',
          '',
          'data: [DONE]',
          '',
        ].join('\n'),
      ),
    });

    const provider = new XAIResponsesProvider('grok-4.3', {
      config: {
        apiKey: 'test-key',
        stream: true,
        include: ['reasoning.encrypted_content'],
      },
    });

    const result = await provider.callApi('hello');

    expect(result.output).toBe('final answer');
    expect(result.output).not.toContain('Reasoning:');
    expect(result.output).not.toContain('internal reasoning');
    expect(result.raw?.output).toEqual(
      expect.arrayContaining([
        expect.objectContaining({
          type: 'reasoning',
          encrypted_content: 'opaque-payload',
        }),
      ]),
    );
  });

  it('uses the completed streamed response when malformed deltas would truncate fallback text', async () => {
    mockFetchWithProxy.mockResolvedValueOnce({
      status: 200,
      statusText: 'OK',
      body: createSSEStream(
        [
          'data: {"type":"response.output_text.delta","delta":"hel"}',
          '',
          'data: not-json',
          '',
          'data: {"type":"response.completed","response":{"id":"resp_stream","model":"grok-4.3","output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"hello"}]}],"usage":{"input_tokens":10,"output_tokens":5,"total_tokens":15}}}',
          '',
          'data: [DONE]',
          '',
        ].join('\n'),
      ),
    });

    const provider = new XAIResponsesProvider('grok-4.3', {
      config: { apiKey: 'test-key', stream: true },
    });

    const result = await provider.callApi('hello');

    expect(result.output).toBe('hello');
  });

  it('preserves completed-response annotations for streamed output text', async () => {
    mockFetchWithProxy.mockResolvedValueOnce({
      status: 200,
      statusText: 'OK',
      body: createSSEStream(
        [
          'data: {"type":"response.output_text.delta","delta":"final answer"}',
          '',
          'data: {"type":"response.completed","response":{"id":"resp_stream","model":"grok-4.3","output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"final answer","annotations":[{"url":"https://example.com","title":"Example"}]}]}],"usage":{"input_tokens":10,"output_tokens":5,"total_tokens":15}}}',
          '',
          'data: [DONE]',
          '',
        ].join('\n'),
      ),
    });

    const provider = new XAIResponsesProvider('grok-4.3', {
      config: { apiKey: 'test-key', stream: true },
    });

    const result = await provider.callApi('hello');

    expect(result.output).toBe('final answer');
    expect(result.metadata?.annotations).toEqual([
      { url: 'https://example.com', title: 'Example' },
    ]);
    expect(result.raw?.annotations).toEqual([{ url: 'https://example.com', title: 'Example' }]);
  });

  it('preserves non-message output items (web_search_call) in streamed completed responses', async () => {
    mockFetchWithProxy.mockResolvedValueOnce({
      status: 200,
      statusText: 'OK',
      body: createSSEStream(
        [
          'data: {"type":"response.output_text.delta","delta":"per docs, the answer is 42"}',
          '',
          'data: {"type":"response.completed","response":{"id":"resp_stream","model":"grok-4.3","output":[{"type":"web_search_call","id":"ws_1","status":"completed"},{"type":"message","role":"assistant","content":[{"type":"output_text","text":"per docs, the answer is 42","annotations":[{"url":"https://example.com","title":"Example"}]}]}],"usage":{"input_tokens":10,"output_tokens":7,"total_tokens":17}}}',
          '',
          'data: [DONE]',
          '',
        ].join('\n'),
      ),
    });

    const provider = new XAIResponsesProvider('grok-4.3', {
      config: {
        apiKey: 'test-key',
        stream: true,
        tools: [{ type: 'web_search' }],
      },
    });

    const result = await provider.callApi('hello');

    // The processor renders web_search_call items into the output text; the
    // key signal is that the web_search_call survives streamed processing
    // rather than being silently dropped alongside the assistant message.
    expect(result.output).toContain('Web Search Call');
    expect(result.output).toContain('per docs, the answer is 42');
    expect(result.metadata?.annotations).toEqual([
      { url: 'https://example.com', title: 'Example' },
    ]);
  });

  it('returns an error when a streamed response has no output content', async () => {
    mockFetchWithProxy.mockResolvedValueOnce({
      status: 200,
      statusText: 'OK',
      body: createSSEStream(['data: not-json', '', 'data: [DONE]', ''].join('\n')),
    });

    const provider = new XAIResponsesProvider('grok-4.3', {
      config: { apiKey: 'test-key', stream: true },
    });

    const result = await provider.callApi('hello');

    expect(result.error).toBe(
      'xAI API error: xAI streaming response did not include output content\n\nIf this persists, verify your API key at https://x.ai/',
    );
  });
});