pi-mono/packages/ai/test/xai-responses.test.ts

import { afterEach, describe, expect, it, vi } from "vitest";
import { stream as streamOpenAICompletions } from "../src/api/openai-completions.ts";
import type { OpenAIResponsesOptions } from "../src/api/openai-responses.ts";
import { stream as streamOpenAIResponses } from "../src/api/openai-responses.ts";
import { getSupportedThinkingLevels } from "../src/models.ts";
import { XAI_MODELS } from "../src/providers/xai.models.ts";
import { xaiProvider } from "../src/providers/xai.ts";
import type { Context, Model } from "../src/types.ts";
import { formatUserAgent } from "../src/utils/identity.ts";
import { normalizeContext } from "../src/utils/transcript.ts";

const PI_USER_AGENT = formatUserAgent();

type CapturedRequest = {
	url: string;
	headers: Headers;
	body: Record<string, unknown>;
};

function completedResponse(): Response {
	const event = {
		type: "response.completed",
		sequence_number: 0,
		response: {
			id: "resp_xai_test",
			status: "completed",
			output: [],
			usage: {
				input_tokens: 1,
				output_tokens: 1,
				total_tokens: 2,
				input_tokens_details: { cached_tokens: 0 },
			},
		},
	};
	return new Response(`data: ${JSON.stringify(event)}\n\ndata: [DONE]\n\n`, {
		status: 200,
		headers: { "content-type": "text/event-stream" },
	});
}

const customCompletionsModel: Model<"openai-completions"> = {
	id: "grok-custom",
	name: "Grok Custom",
	api: "openai-completions",
	provider: "xai",
	baseUrl: "https://api.x.ai/v1",
	reasoning: false,
	input: ["text"],
	cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
	contextWindow: 128000,
	maxTokens: 16384,
};

async function captureCompletionsUserAgent(headers?: Record<string, string>): Promise<string | null> {
	let userAgent: string | null = null;
	vi.spyOn(globalThis, "fetch").mockImplementation(async (input, init) => {
		userAgent = new Request(input, init).headers.get("user-agent");
		const chunks = [
			{ id: "chatcmpl-ua", choices: [{ delta: { content: "ok" }, finish_reason: null, index: 0 }] },
			{
				id: "chatcmpl-ua",
				choices: [{ delta: {}, finish_reason: "stop", index: 0 }],
				usage: {
					prompt_tokens: 1,
					completion_tokens: 1,
					prompt_tokens_details: { cached_tokens: 0 },
					completion_tokens_details: { reasoning_tokens: 0 },
				},
			},
		];
		const body = `${chunks.map((chunk) => `data: ${JSON.stringify(chunk)}`).join("\n\n")}\n\ndata: [DONE]\n\n`;
		return new Response(body, {
			status: 200,
			headers: { "content-type": "text/event-stream" },
		});
	});

	const result = await streamOpenAICompletions(
		customCompletionsModel,
		normalizeContext({ messages: [{ role: "user", content: "hello", timestamp: 1 }] }),
		{ apiKey: "xai-test-token", headers },
	).result();

	expect(result.stopReason, result.errorMessage).toBe("stop");
	return userAgent;
}

async function captureRequest(
	model: Model<"openai-responses">,
	context: Context,
	options: OpenAIResponsesOptions,
): Promise<CapturedRequest> {
	let captured: CapturedRequest | undefined;
	vi.spyOn(globalThis, "fetch").mockImplementation(async (input, init) => {
		const request = new Request(input, init);
		captured = {
			url: request.url,
			headers: request.headers,
			body: JSON.parse(await request.clone().text()) as Record<string, unknown>,
		};
		return completedResponse();
	});

	const result = await xaiProvider().stream(model, normalizeContext(context), options).result();
	expect(result.stopReason, result.errorMessage).toBe("stop");
	expect(captured).toBeDefined();
	return captured!;
}

describe("xAI Responses provider", () => {
	afterEach(() => {
		vi.restoreAllMocks();
	});

	it("excludes retired and redundant models from the built-in catalog", () => {
		for (const modelId of [
			"grok-3",
			"grok-3-fast",
			"grok-4.20-0309-non-reasoning",
			"grok-4.20-0309-reasoning",
			"grok-build-0.1",
			"grok-code-fast-1",
		]) {
			expect(Object.keys(XAI_MODELS)).not.toContain(modelId);
		}
	});

	it("routes every built-in xAI model through Responses", () => {
		for (const model of Object.values(XAI_MODELS)) {
			expect(model.api, model.id).toBe("openai-responses");
		}
		expect(getSupportedThinkingLevels(XAI_MODELS["grok-4.5"])).toEqual(["low", "medium", "high", "xhigh"]);
		expect(getSupportedThinkingLevels(XAI_MODELS["grok-4.6"])).toEqual(["low", "medium", "high", "xhigh"]);
		expect(getSupportedThinkingLevels(XAI_MODELS["grok-4.7"])).toEqual(["low", "medium", "high", "xhigh"]);
		expect(getSupportedThinkingLevels(XAI_MODELS["grok-4.3"])).toEqual(["off", "low", "medium", "high", "xhigh"]);
	});

	it("includes Grok 4.7 capabilities and long-context pricing", () => {
		expect(XAI_MODELS["grok-4.7"]).toMatchObject({
			api: "openai-responses",
			reasoning: true,
			input: ["text", "image"],
			contextWindow: 500000,
			maxTokens: 500000,
			cost: {
				input: 2,
				output: 6,
				cacheRead: 0.5,
				cacheWrite: 0,
				tiers: [
					{
						inputTokensAbove: 200000,
						input: 4,
						output: 12,
						cacheRead: 1,
						cacheWrite: 0,
					},
				],
			},
		});
	});

	it("uses /responses with bearer auth and xAI-compatible request fields", async () => {
		const captured = await captureRequest(
			XAI_MODELS["grok-4.5"],
			{
				systemPrompt: "You are a careful coding assistant.",
				messages: [{ role: "user", content: "hello", timestamp: 1 }],
			},
			{
				apiKey: "xai-test-token",
				sessionId: "pi-session-123",
				cacheRetention: "long",
				reasoningEffort: "medium",
			},
		);

		expect(captured.url).toBe("https://api.x.ai/v1/responses");
		expect(captured.headers.get("authorization")).toBe("Bearer xai-test-token");
		expect(captured.headers.get("user-agent")).toBe(PI_USER_AGENT);
		expect(captured.headers.get("session-id")).toBe("pi-session-123");
		expect(captured.headers.get("session_id")).toBeNull();
		expect(captured.body).toMatchObject({
			model: "grok-4.5",
			store: false,
			stream: true,
			prompt_cache_key: "pi-session-123",
			reasoning: { effort: "medium" },
			include: ["reasoning.encrypted_content"],
		});
		expect(captured.body).not.toHaveProperty("prompt_cache_retention");
		expect(captured.body.input).toEqual(
			expect.arrayContaining([
				expect.objectContaining({
					role: "developer",
					content: "You are a careful coding assistant.",
				}),
			]),
		);
	});

	it("requests encrypted reasoning without an effort override", async () => {
		const captured = await captureRequest(
			XAI_MODELS["grok-4.5"],
			{ messages: [{ role: "user", content: "hello", timestamp: 1 }] },
			{ apiKey: "xai-test-token" },
		);

		expect(captured.body).toMatchObject({
			model: "grok-4.5",
			store: false,
			include: ["reasoning.encrypted_content"],
		});
		expect(captured.body).not.toHaveProperty("reasoning");
	});

	it("uses /responses for Grok 4.7 with xhigh effort and encrypted reasoning", async () => {
		const captured = await captureRequest(
			XAI_MODELS["grok-4.7"],
			{
				systemPrompt: "You are a careful coding assistant.",
				messages: [{ role: "user", content: "hello", timestamp: 1 }],
			},
			{
				apiKey: "xai-test-token",
				reasoningEffort: "xhigh",
			},
		);

		expect(captured.url).toBe("https://api.x.ai/v1/responses");
		expect(captured.body).toMatchObject({
			model: "grok-4.7",
			store: false,
			stream: true,
			reasoning: { effort: "xhigh" },
			include: ["reasoning.encrypted_content"],
		});
	});

	it("uses /responses for Grok 4.3", async () => {
		const captured = await captureRequest(
			XAI_MODELS["grok-4.3"],
			{
				messages: [{ role: "user", content: "hello", timestamp: 1 }],
			},
			{
				apiKey: "xai-test-token",
				reasoningEffort: "low",
			},
		);

		expect(captured.url).toBe("https://api.x.ai/v1/responses");
		expect(captured.body).toMatchObject({
			model: "grok-4.3",
			store: false,
			include: ["reasoning.encrypted_content"],
			reasoning: { effort: "low" },
		});
	});

	it("uses pi's User-Agent by default for Responses requests", async () => {
		let userAgent: string | null = null;
		vi.spyOn(globalThis, "fetch").mockImplementation(async (input, init) => {
			userAgent = new Request(input, init).headers.get("user-agent");
			return completedResponse();
		});

		const openaiModel: Model<"openai-responses"> = {
			...XAI_MODELS["grok-4.5"],
			provider: "openai",
			baseUrl: "https://api.openai.com/v1",
		};
		const result = await streamOpenAIResponses(
			openaiModel,
			normalizeContext({ messages: [{ role: "user", content: "hello", timestamp: 1 }] }),
			{ apiKey: "test-token" },
		).result();

		expect(result.stopReason, result.errorMessage).toBe("stop");
		expect(userAgent).toBe(PI_USER_AGENT);
	});

	it("lets explicit headers override the default Responses User-Agent", async () => {
		const captured = await captureRequest(
			XAI_MODELS["grok-4.5"],
			{ messages: [{ role: "user", content: "hello", timestamp: 1 }] },
			{ apiKey: "xai-test-token", headers: { "User-Agent": "custom-agent" } },
		);

		expect(captured.headers.get("user-agent")).toBe("custom-agent");
	});

	it("uses pi's User-Agent by default for Completions requests", async () => {
		expect(await captureCompletionsUserAgent()).toBe(PI_USER_AGENT);
	});

	it("lets explicit headers override the default Completions User-Agent", async () => {
		expect(await captureCompletionsUserAgent({ "User-Agent": "custom-agent" })).toBe("custom-agent");
	});
});