From f175c084d327c34e70fbcff4d54bbd327139eb71 Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Sun, 19 Jul 2026 16:32:47 +0800 Subject: [PATCH] feat(executor): promote `additional_tools` to top-level `tools` in XAI requests - Added `promoteXAIAdditionalTools` to normalize and migrate `additional_tools` to the top-level `tools` array in XAI request bodies. - Updated test cases to validate the promotion logic and ensure unsupported `additional_tools` items are excluded from input payloads. Closes: #4434 --- internal/runtime/executor/xai_executor.go | 59 +++++++++++++++++++ .../runtime/executor/xai_executor_test.go | 51 +++++++++++----- .../executor/xai_websockets_executor_test.go | 10 +++- 3 files changed, 105 insertions(+), 15 deletions(-) diff --git a/internal/runtime/executor/xai_executor.go b/internal/runtime/executor/xai_executor.go index c5fd3533..5c2561a0 100644 --- a/internal/runtime/executor/xai_executor.go +++ b/internal/runtime/executor/xai_executor.go @@ -915,6 +915,7 @@ func (e *XAIExecutor) prepareResponsesRequestTo(ctx context.Context, req cliprox // the post-restore (namespace, short-name) shape used by the response filter. clientDeclaredTools := collectXAIClientDeclaredToolKeys(body) body = normalizeXAITools(body) + body = promoteXAIAdditionalTools(body) // Drop choices that point at tools removed by normalizeXAITools before we // inject native x_search, so a surviving allowed_tools / forced choice is not // left pointing at a deleted tool once only x_search remains. @@ -1537,6 +1538,64 @@ func normalizeXAITools(body []byte) []byte { return body } +// promoteXAIAdditionalTools moves Responses Lite tool declarations to the +// top-level tools array because xAI does not accept additional_tools input items. +func promoteXAIAdditionalTools(body []byte) []byte { + if !gjson.ValidBytes(body) { + return body + } + input := gjson.GetBytes(body, "input") + if !input.IsArray() { + return body + } + + inputItems := input.Array() + remainingInput := make([]json.RawMessage, 0, len(inputItems)) + promotedTools := make([]json.RawMessage, 0) + for _, item := range inputItems { + if item.Get("type").String() != "additional_tools" { + remainingInput = append(remainingInput, json.RawMessage(item.Raw)) + continue + } + for _, tool := range item.Get("tools").Array() { + promotedTools = append(promotedTools, json.RawMessage(tool.Raw)) + } + } + if len(remainingInput) == len(inputItems) { + return body + } + + rawInput, errMarshalInput := json.Marshal(remainingInput) + if errMarshalInput != nil { + return body + } + updated, errSetInput := sjson.SetRawBytes(body, "input", rawInput) + if errSetInput != nil { + return body + } + if len(promotedTools) == 0 { + return updated + } + + topLevelTools := gjson.GetBytes(updated, "tools") + tools := make([]json.RawMessage, 0, len(topLevelTools.Array())+len(promotedTools)) + if topLevelTools.IsArray() { + for _, tool := range topLevelTools.Array() { + tools = append(tools, json.RawMessage(tool.Raw)) + } + } + tools = append(tools, promotedTools...) + rawTools, errMarshalTools := json.Marshal(tools) + if errMarshalTools != nil { + return body + } + updated, errSetTools := sjson.SetRawBytes(updated, "tools", rawTools) + if errSetTools != nil { + return body + } + return updated +} + func normalizeXAIToolArray(tools gjson.Result) ([]byte, bool, bool) { changed := false filtered := []byte(`[]`) diff --git a/internal/runtime/executor/xai_executor_test.go b/internal/runtime/executor/xai_executor_test.go index 9dba8da1..cd752c18 100644 --- a/internal/runtime/executor/xai_executor_test.go +++ b/internal/runtime/executor/xai_executor_test.go @@ -253,15 +253,23 @@ func TestXAIExecutorExecuteRestoresAdditionalToolsNamespaceCalls(t *testing.T) { t.Fatalf("Execute() error = %v", err) } - tool := gjson.GetBytes(gotBody, "input.0.tools.0") + for _, item := range gjson.GetBytes(gotBody, "input").Array() { + if got := item.Get("type").String(); got == "additional_tools" { + t.Fatalf("upstream input contains unsupported additional_tools item: %s", gotBody) + } + } + if got := gjson.GetBytes(gotBody, "input.0.role").String(); got != "user" { + t.Fatalf("input.0.role = %q, want user; body=%s", got, gotBody) + } + tool := gjson.GetBytes(gotBody, "tools.0") if got := tool.Get("name").String(); got != "mcp__exa__web_search_exa" { - t.Fatalf("upstream additional tool name = %q, want qualified name; body=%s", got, gotBody) + t.Fatalf("upstream tool name = %q, want qualified name; body=%s", got, gotBody) } if got := tool.Get("type").String(); got != "function" { - t.Fatalf("upstream additional tool type = %q, want function; body=%s", got, gotBody) + t.Fatalf("upstream tool type = %q, want function; body=%s", got, gotBody) } if tool.Get("tools").Exists() { - t.Fatalf("upstream additional tool should not contain namespace children: %s", gotBody) + t.Fatalf("upstream tool should not contain namespace children: %s", gotBody) } output := gjson.GetBytes(resp.Payload, "output.0") if got := output.Get("name").String(); got != "web_search_exa" { @@ -2907,24 +2915,39 @@ func TestNormalizeXAITools_QualifiesSameNamedNamespaceTools(t *testing.T) { } } -func TestNormalizeXAITools_AdditionalToolsNamespace(t *testing.T) { +func TestPromoteXAIAdditionalTools(t *testing.T) { body := []byte(`{ + "tools":[{"type":"function","name":"lookup","parameters":{"type":"object"}}], "input":[ {"type":"additional_tools","role":"developer","tools":[{"type":"namespace","name":"mcp__exa","tools":[{"type":"function","name":"search","parameters":{"type":"object"}}]}]}, - {"role":"user","content":"hello"} + {"role":"user","content":"hello"}, + {"type":"additional_tools","role":"developer","tools":[{"type":"custom","name":"custom_lookup"}]} ] }`) - out := normalizeXAITools(body) + out := promoteXAIAdditionalTools(normalizeXAITools(body)) - tools := gjson.GetBytes(out, "input.0.tools").Array() - if len(tools) != 1 { - t.Fatalf("additional tools length = %d, want 1; body=%s", len(tools), string(out)) + input := gjson.GetBytes(out, "input").Array() + if len(input) != 1 || input[0].Get("role").String() != "user" { + t.Fatalf("input should contain only the user message: %s", string(out)) } - if got := tools[0].Get("name").String(); got != "mcp__exa__search" { - t.Fatalf("additional tool name = %q, want mcp__exa__search; body=%s", got, string(out)) + tools := gjson.GetBytes(out, "tools").Array() + if len(tools) != 3 { + t.Fatalf("tools length = %d, want 3; body=%s", len(tools), string(out)) } - if got := tools[0].Get("type").String(); got != "function" { - t.Fatalf("additional tool type = %q, want function; body=%s", got, string(out)) + if got := tools[0].Get("name").String(); got != "lookup" { + t.Fatalf("tools.0.name = %q, want lookup; body=%s", got, string(out)) + } + if got := tools[1].Get("name").String(); got != "mcp__exa__search" { + t.Fatalf("tools.1.name = %q, want mcp__exa__search; body=%s", got, string(out)) + } + if got := tools[2].Get("name").String(); got != "custom_lookup" { + t.Fatalf("tools.2.name = %q, want custom_lookup; body=%s", got, string(out)) + } + if got := tools[2].Get("type").String(); got != "function" { + t.Fatalf("tools.2.type = %q, want function; body=%s", got, string(out)) + } + if !tools[2].Get("parameters").Exists() { + t.Fatalf("tools.2.parameters missing: %s", string(out)) } } diff --git a/internal/runtime/executor/xai_websockets_executor_test.go b/internal/runtime/executor/xai_websockets_executor_test.go index edc694d8..60fe4b14 100644 --- a/internal/runtime/executor/xai_websockets_executor_test.go +++ b/internal/runtime/executor/xai_websockets_executor_test.go @@ -202,7 +202,15 @@ func TestXAIWebsocketsExecuteStreamRestoresNamespaceToolCalls(t *testing.T) { select { case payload := <-capturedPayload: - tool := gjson.GetBytes(payload, "input.0.tools.0") + for _, item := range gjson.GetBytes(payload, "input").Array() { + if got := item.Get("type").String(); got == "additional_tools" { + t.Fatalf("upstream input contains unsupported additional_tools item: %s", payload) + } + } + if got := gjson.GetBytes(payload, "input.0.role").String(); got != "user" { + t.Fatalf("input.0.role = %q, want user; payload=%s", got, payload) + } + tool := gjson.GetBytes(payload, "tools.0") if got := tool.Get("name").String(); got != "mcp__exa__web_search_exa" { t.Fatalf("upstream tool name = %q, want qualified name; payload=%s", got, payload) }