From 059b0db6c32eb87ca72c06158f4afe0b80dbf99d Mon Sep 17 00:00:00 2001 From: Roo Code Date: Fri, 20 Mar 2026 00:02:42 +0000 Subject: [PATCH] fix: address review issues in xAI Responses API migration - Use base provider convertToolSchemaForOpenAI() for tool schema hardening - Handle MCP tools with isMcpTool() (strict: false for MCP, true otherwise) - Respect metadata.tool_choice instead of always using "auto" - Add parallel_tool_calls pass-through from metadata - Fix assistant string content to use output_text format (not input_text) - Remove unused modelInfo param from createUsageNormalizer - Add cacheWriteTokens extraction to usage normalizer - Fix test() -> it() inconsistency in xai.spec.ts - Update tests to match all changes --- src/api/providers/__tests__/xai.spec.ts | 4 ++- src/api/providers/xai.ts | 29 ++++++++++++------ .../__tests__/responses-api-input.spec.ts | 10 +++++-- .../__tests__/responses-api-stream.spec.ts | 30 ++++++++++++------- src/api/transform/responses-api-input.ts | 17 ++++++++--- src/api/transform/responses-api-stream.ts | 6 ++-- 6 files changed, 66 insertions(+), 30 deletions(-) diff --git a/src/api/providers/__tests__/xai.spec.ts b/src/api/providers/__tests__/xai.spec.ts index ff4b14cf14..d760c0fd47 100644 --- a/src/api/providers/__tests__/xai.spec.ts +++ b/src/api/providers/__tests__/xai.spec.ts @@ -82,7 +82,7 @@ describe("XAIHandler", () => { expect(model.info).toEqual(xaiModels[xaiDefaultModelId]) }) - test("should return specified model when valid model is provided", () => { + it("should return specified model when valid model is provided", () => { const testModelId = "grok-3" const handlerWithModel = new XAIHandler({ apiModelId: testModelId }) const model = handlerWithModel.getModel() @@ -227,9 +227,11 @@ describe("XAIHandler", () => { type: "function", name: "test_tool", description: "A test tool", + strict: true, }), ], tool_choice: "auto", + parallel_tool_calls: true, }), ) }) diff --git a/src/api/providers/xai.ts b/src/api/providers/xai.ts index a750dfbeef..daaad88d6c 100644 --- a/src/api/providers/xai.ts +++ b/src/api/providers/xai.ts @@ -14,6 +14,7 @@ import { DEFAULT_HEADERS } from "./constants" import { BaseProvider } from "./base-provider" import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" import { handleOpenAIError } from "./utils/openai-error-handler" +import { isMcpTool } from "../../utils/mcp-name" const XAI_DEFAULT_TEMPERATURE = 0 @@ -48,6 +49,9 @@ export class XAIHandler extends BaseProvider implements SingleCompletionHandler * Convert tools from OpenAI Chat Completions format to Responses API format. * Chat Completions: { type: "function", function: { name, description, parameters } } * Responses API: { type: "function", name, description, parameters } + * + * Uses base provider's convertToolSchemaForOpenAI() for schema hardening + * (additionalProperties: false, ensureAllRequired) and handles MCP tools. */ private mapResponseTools(tools?: any[]): any[] | undefined { if (!tools?.length) { @@ -55,13 +59,18 @@ export class XAIHandler extends BaseProvider implements SingleCompletionHandler } return tools .filter((tool) => tool?.type === "function") - .map((tool) => ({ - type: "function", - name: tool.function.name, - description: tool.function.description, - parameters: tool.function.parameters ?? null, - strict: false, - })) + .map((tool) => { + const isMcp = isMcpTool(tool.function.name) + return { + type: "function", + name: tool.function.name, + description: tool.function.description, + parameters: isMcp + ? tool.function.parameters + : this.convertToolSchemaForOpenAI(tool.function.parameters), + strict: !isMcp, + } + }) } override async *createMessage( @@ -86,7 +95,9 @@ export class XAIHandler extends BaseProvider implements SingleCompletionHandler stream: true, store: false, // Don't store responses server-side for privacy tools: responseTools, - tool_choice: responseTools ? "auto" : undefined, + // Cast tool_choice since metadata uses Chat Completions types but Responses API has its own type + tool_choice: (metadata?.tool_choice ?? (responseTools ? "auto" : undefined)) as any, + parallel_tool_calls: metadata?.parallelToolCalls ?? true, include: ["reasoning.encrypted_content"], }) } catch (error) { @@ -96,7 +107,7 @@ export class XAIHandler extends BaseProvider implements SingleCompletionHandler throw handleOpenAIError(error, this.providerName) } - const normalizeUsage = createUsageNormalizer(model.info) + const normalizeUsage = createUsageNormalizer() yield* processResponsesApiStream(stream, normalizeUsage) } diff --git a/src/api/transform/__tests__/responses-api-input.spec.ts b/src/api/transform/__tests__/responses-api-input.spec.ts index c78345c789..c57b61d895 100644 --- a/src/api/transform/__tests__/responses-api-input.spec.ts +++ b/src/api/transform/__tests__/responses-api-input.spec.ts @@ -15,12 +15,18 @@ describe("convertToResponsesApiInput", () => { expect(result).toEqual([{ role: "user", content: [{ type: "input_text", text: "Hello" }] }]) }) - it("should convert assistant string content", () => { + it("should convert assistant string content to output_text message format", () => { const messages: Anthropic.Messages.MessageParam[] = [{ role: "assistant", content: "Hi there" }] const result = convertToResponsesApiInput(messages) - expect(result).toEqual([{ role: "assistant", content: [{ type: "input_text", text: "Hi there" }] }]) + expect(result).toEqual([ + { + type: "message", + role: "assistant", + content: [{ type: "output_text", text: "Hi there" }], + }, + ]) }) }) diff --git a/src/api/transform/__tests__/responses-api-stream.spec.ts b/src/api/transform/__tests__/responses-api-stream.spec.ts index 4d90ebaacb..6abdbddb65 100644 --- a/src/api/transform/__tests__/responses-api-stream.spec.ts +++ b/src/api/transform/__tests__/responses-api-stream.spec.ts @@ -291,16 +291,14 @@ describe("processResponsesApiStream", () => { }) describe("createUsageNormalizer", () => { - const mockModelInfo = { contextWindow: 128000, supportsPromptCache: false } as any - it("should return undefined for null/undefined usage", () => { - const normalize = createUsageNormalizer(mockModelInfo) + const normalize = createUsageNormalizer() expect(normalize(null)).toBeUndefined() expect(normalize(undefined)).toBeUndefined() }) it("should extract input and output tokens", () => { - const normalize = createUsageNormalizer(mockModelInfo) + const normalize = createUsageNormalizer() const result = normalize({ input_tokens: 100, output_tokens: 50 }) @@ -314,7 +312,7 @@ describe("createUsageNormalizer", () => { }) it("should extract cached tokens from input_tokens_details", () => { - const normalize = createUsageNormalizer(mockModelInfo) + const normalize = createUsageNormalizer() const result = normalize({ input_tokens: 100, @@ -325,8 +323,20 @@ describe("createUsageNormalizer", () => { expect(result?.cacheReadTokens).toBe(30) }) + it("should extract cache write tokens", () => { + const normalize = createUsageNormalizer() + + const result = normalize({ + input_tokens: 100, + output_tokens: 50, + cache_creation_input_tokens: 15, + }) + + expect(result?.cacheWriteTokens).toBe(15) + }) + it("should extract reasoning tokens from output_tokens_details", () => { - const normalize = createUsageNormalizer(mockModelInfo) + const normalize = createUsageNormalizer() const result = normalize({ input_tokens: 100, @@ -338,7 +348,7 @@ describe("createUsageNormalizer", () => { }) it("should not include reasoningTokens when not present", () => { - const normalize = createUsageNormalizer(mockModelInfo) + const normalize = createUsageNormalizer() const result = normalize({ input_tokens: 100, output_tokens: 50 }) @@ -347,7 +357,7 @@ describe("createUsageNormalizer", () => { it("should compute totalCost when calculateCost is provided", () => { const calculateCost = (input: number, output: number, cached: number) => 0.42 - const normalize = createUsageNormalizer(mockModelInfo, calculateCost) + const normalize = createUsageNormalizer(calculateCost) const result = normalize({ input_tokens: 100, output_tokens: 50 }) @@ -355,7 +365,7 @@ describe("createUsageNormalizer", () => { }) it("should not include totalCost when calculateCost is not provided", () => { - const normalize = createUsageNormalizer(mockModelInfo) + const normalize = createUsageNormalizer() const result = normalize({ input_tokens: 100, output_tokens: 50 }) @@ -363,7 +373,7 @@ describe("createUsageNormalizer", () => { }) it("should handle Chat Completions style field names as fallback", () => { - const normalize = createUsageNormalizer(mockModelInfo) + const normalize = createUsageNormalizer() const result = normalize({ prompt_tokens: 100, diff --git a/src/api/transform/responses-api-input.ts b/src/api/transform/responses-api-input.ts index bfcc52b542..a766dfef6e 100644 --- a/src/api/transform/responses-api-input.ts +++ b/src/api/transform/responses-api-input.ts @@ -18,10 +18,19 @@ export function convertToResponsesApiInput(messages: Anthropic.Messages.MessageP for (const message of messages) { if (typeof message.content === "string") { - input.push({ - role: message.role, - content: [{ type: "input_text", text: message.content }], - }) + if (message.role === "assistant") { + // Assistant messages use output_text in the Responses API format + input.push({ + type: "message", + role: "assistant", + content: [{ type: "output_text", text: message.content }], + }) + } else { + input.push({ + role: message.role, + content: [{ type: "input_text", text: message.content }], + }) + } continue } diff --git a/src/api/transform/responses-api-stream.ts b/src/api/transform/responses-api-stream.ts index ab0cdfbf94..1ea9744180 100644 --- a/src/api/transform/responses-api-stream.ts +++ b/src/api/transform/responses-api-stream.ts @@ -1,5 +1,3 @@ -import type { ModelInfo } from "@roo-code/types" - import type { ApiStream, ApiStreamUsageChunk } from "./stream" /** @@ -106,11 +104,9 @@ export async function* processResponsesApiStream( * Creates a standard usage normalizer for providers with per-token pricing. * Extracts input/output tokens, cache tokens, reasoning tokens, and computes cost. * - * @param modelInfo - Model info with pricing details * @param calculateCost - Optional function to compute total cost from token counts */ export function createUsageNormalizer( - modelInfo: ModelInfo, calculateCost?: (inputTokens: number, outputTokens: number, cacheReadTokens: number) => number, ): (usage: any) => ApiStreamUsageChunk | undefined { return (usage: any): ApiStreamUsageChunk | undefined => { @@ -122,6 +118,7 @@ export function createUsageNormalizer( const inputTokens = usage.input_tokens ?? usage.prompt_tokens ?? 0 const outputTokens = usage.output_tokens ?? usage.completion_tokens ?? 0 const cacheReadTokens = usage.cache_read_input_tokens ?? cachedTokens ?? 0 + const cacheWriteTokens = usage.cache_creation_input_tokens ?? usage.cache_write_tokens ?? 0 const reasoningTokens = typeof usage.output_tokens_details?.reasoning_tokens === "number" @@ -135,6 +132,7 @@ export function createUsageNormalizer( inputTokens, outputTokens, cacheReadTokens, + cacheWriteTokens, ...(typeof reasoningTokens === "number" ? { reasoningTokens } : {}), ...(typeof totalCost === "number" ? { totalCost } : {}), }