Merge pull request #124 from Wibias/codex/fix-xai-nested-tool-schema

fix(xai): normalize nested tool schemas

Author: lidge-jun

src/adapters/openai-chat.ts

@@ -125,7 +125,49 @@ function safeToolName(name: string | undefined): string {
   return sanitized;
 }
 
-function toolsToChatFormat(parsed: OcxParsedRequest): unknown[] | undefined {
+const XAI_SCHEMA_BASE_URLS = new Set(["api.x.ai", "cli-chat-proxy.grok.com"]);
+
+function isXaiSchemaTarget(provider: OcxProviderConfig): boolean {
+  try {
+    return XAI_SCHEMA_BASE_URLS.has(new URL(provider.baseUrl).hostname);
+  } catch {
+    return false;
+  }
+}
+
+function expandXaiRootObjectSchemas(schema: unknown): Record<string, unknown>[] | undefined {
+  if (!schema || typeof schema !== "object" || Array.isArray(schema)) return undefined;
+  const obj = schema as Record<string, unknown>;
+  const compositionKey = ["oneOf", "anyOf"].find(key => Array.isArray(obj[key]));
+  if (!compositionKey) {
+    if (obj.type !== undefined && obj.type !== "object") return undefined;
+    return [{ ...obj, type: "object" }];
+  }
+
+  const siblings = Object.fromEntries(Object.entries(obj).filter(([key]) => key !== compositionKey));
+  const branches = obj[compositionKey];
+  if (!Array.isArray(branches)) return undefined;
+  const expanded: Record<string, unknown>[] = [];
+  for (const branch of branches) {
+    const variants = expandXaiRootObjectSchemas(branch);
+    if (!variants) return undefined;
+    for (const variant of variants) expanded.push({ ...siblings, ...variant });
+  }
+  return expanded.length > 0 ? expanded : undefined;
+}
+
+function normalizeXaiToolParameters(parameters: unknown): Record<string, unknown> | undefined {
+  const variants = expandXaiRootObjectSchemas(parameters);
+  if (!variants) return undefined;
+  if (variants.length === 1) return variants[0];
+  const root = parameters && typeof parameters === "object" && !Array.isArray(parameters)
+    ? parameters as Record<string, unknown>
+    : {};
+  const metadata = Object.fromEntries(Object.entries(root).filter(([key]) => key !== "oneOf" && key !== "anyOf" && key !== "type"));
+  return { ...metadata, oneOf: variants };
+}
+
+function toolsToChatFormat(parsed: OcxParsedRequest, provider: OcxProviderConfig): unknown[] | undefined {
   if (!parsed.context.tools || parsed.context.tools.length === 0) return undefined;
   const allowed = isAllowedToolChoice(parsed.options.toolChoice)
     ? new Set(parsed.options.toolChoice.allowedTools)
@@ -134,15 +176,21 @@ function toolsToChatFormat(parsed: OcxParsedRequest): unknown[] | undefined {
     ? parsed.context.tools.filter(t => toolAllowedByChoice(t, allowed))
     : parsed.context.tools;
   if (tools.length === 0) return undefined;
-  return tools.map(t => ({
+  const xaiTarget = isXaiSchemaTarget(provider);
+  const formatted = tools.flatMap(t => {
+    const parameters = xaiTarget ? normalizeXaiToolParameters(t.parameters) : t.parameters;
+    if (parameters === undefined) return [];
+    return [{
     type: "function",
     function: {
       name: namespacedToolName(t.namespace, t.name),
       description: t.description,
-      parameters: t.parameters,
+      parameters,
       ...(t.strict !== undefined ? { strict: t.strict } : {}),
     },
-  }));
+    }];
+  });
+  return formatted.length > 0 ? formatted : undefined;
 }
 
 function toolChoiceToChatFormat(tc: OcxParsedRequest["options"]["toolChoice"], tools: OcxParsedRequest["context"]["tools"]): unknown {
@@ -190,7 +238,7 @@ export function createOpenAIChatAdapter(provider: OcxProviderConfig): ProviderAd
       }
 
       const messages = messagesToChatFormat(parsed, provider);
-      const tools = toolsToChatFormat(parsed);
+      const tools = toolsToChatFormat(parsed, provider);
       const toolChoice = toolChoiceToChatFormat(parsed.options.toolChoice, parsed.context.tools);
 
       const body: Record<string, unknown> = {

tests/xai-transport.test.ts

@@ -1,5 +1,6 @@
 import { describe, expect, test } from "bun:test";
 import { createOpenAIChatAdapter } from "../src/adapters/openai-chat";
+import { parseRequest } from "../src/responses/parser";
 import { buildModelsRequest } from "../src/oauth";
 import {
   resolveProviderTransport,
@@ -85,6 +86,67 @@ describe("xAI auth-mode transport selection", () => {
       "x-xai-token-auth": "xai-grok-cli",
     });
   });
+
+  test("flattens nested root tool unions for xAI without changing other providers", () => {
+    const schema = {
+      oneOf: [
+        { type: "object", properties: { mode: { type: "string", enum: ["view"] } } },
+        { oneOf: [{ type: "object", properties: { path: { type: "string" } } }, { type: "object", properties: {} }] },
+      ],
+      $defs: { shared: { type: "string" } },
+    };
+    const request = createOpenAIChatAdapter(provider("key")).buildRequest({
+      ...parsed(),
+      context: { messages: [], tools: [{ name: "automation_update", description: "Update", parameters: schema }] },
+    });
+    const xaiParameters = (JSON.parse(request.body) as { tools: Array<{ function: { parameters: Record<string, unknown> } }> }).tools[0].function.parameters;
+
+    expect(xaiParameters.type).toBeUndefined();
+    expect(xaiParameters.oneOf).toHaveLength(3);
+    expect((xaiParameters.oneOf as Record<string, unknown>[]).every(branch => branch.type === "object")).toBe(true);
+    expect(xaiParameters.$defs).toEqual(schema.$defs);
+
+    const otherRequest = createOpenAIChatAdapter({ ...provider("key"), baseUrl: "https://example.test/v1" }).buildRequest({
+      ...parsed(),
+      context: { messages: [], tools: [{ name: "automation_update", description: "Update", parameters: schema }] },
+    });
+    expect((JSON.parse(otherRequest.body) as { tools: Array<{ function: { parameters: unknown } }> }).tools[0].function.parameters).toEqual(schema);
+  });
+
+  test("omits an xAI tool whose root schema cannot be normalized safely", () => {
+    const request = createOpenAIChatAdapter(provider("key")).buildRequest({
+      ...parsed(),
+      context: { messages: [], tools: [{ name: "unsafe", description: "Unsafe", parameters: { oneOf: [{ type: "string" }] } }] },
+    });
+    expect(JSON.parse(request.body).tools).toBeUndefined();
+  });
+
+  test("normalizes a tool loaded from tool_search history on later turns", () => {
+    const parsedRequest = parseRequest({
+      model: "xai/grok-4.5",
+      input: [
+        { type: "tool_search_call", call_id: "search-1", arguments: { query: "automation" } },
+        {
+          type: "tool_search_output",
+          call_id: "search-1",
+          status: "completed",
+          tools: [{
+            type: "function",
+            name: "automation_update",
+            description: "Update an automation",
+            parameters: { oneOf: [{ type: "object", properties: {} }, { oneOf: [{ type: "object", properties: {} }] }] },
+          }],
+        },
+        { type: "message", role: "user", content: [{ type: "input_text", text: "continue" }] },
+      ],
+    });
+    const request = createOpenAIChatAdapter(provider("key")).buildRequest(parsedRequest);
+    const body = JSON.parse(request.body) as { tools: Array<{ function: { name: string; parameters: Record<string, unknown> } }> };
+    const tool = body.tools.find(entry => entry.function.name === "automation_update");
+
+    expect(tool?.function.parameters.oneOf).toHaveLength(2);
+    expect((tool?.function.parameters.oneOf as Record<string, unknown>[]).every(branch => branch.type === "object")).toBe(true);
+  });
 });
 
 describe("xAI prompt-cache conv-id affinity", () => {