From 9b88808fc75ac012afb84955947aee0a0a62d4ad Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Thu, 27 Aug 2026 01:08:35 +0800 Subject: [PATCH] fix(xai): fold namespace tools and restore dispatcher tool calls - Introduce `xaiNamespaceRestorer` to track and restore folded dispatcher tool calls across SSE and WebSocket response events. - Support unwrapping dispatcher tool calls and arguments in `output_item.added` and `function_call_arguments.done` events. - Normalize historical input namespace tool calls to dispatcher format when namespace folding is active. Closes: #5214 --- internal/runtime/executor/xai_executor.go | 1 + .../runtime/executor/xai_executor_execute.go | 3 +- .../runtime/executor/xai_executor_request.go | 277 +++++++++- .../runtime/executor/xai_executor_response.go | 197 ++++++- .../runtime/executor/xai_executor_stream.go | 3 +- .../runtime/executor/xai_executor_test.go | 479 ++++++++++++++++++ .../executor/xai_websockets_executor.go | 3 +- 7 files changed, 938 insertions(+), 25 deletions(-) diff --git a/internal/runtime/executor/xai_executor.go b/internal/runtime/executor/xai_executor.go index 6e3a6b41..6772b080 100644 --- a/internal/runtime/executor/xai_executor.go +++ b/internal/runtime/executor/xai_executor.go @@ -27,6 +27,7 @@ const ( xaiToolSearchType = "tool_search" xaiWebSearchToolType = "web_search" xaiXSearchToolType = "x_search" + xaiMaxTools = 200 // Codex Desktop injects codex_app.automation_update with a large oneOf+$ref // schema. xAI's free/build Responses path accepts the HTTP request but never // emits SSE when that schema is present, so Desktop hangs on "thinking". diff --git a/internal/runtime/executor/xai_executor_execute.go b/internal/runtime/executor/xai_executor_execute.go index f62bddcb..90583370 100644 --- a/internal/runtime/executor/xai_executor_execute.go +++ b/internal/runtime/executor/xai_executor_execute.go @@ -84,12 +84,13 @@ func (e *XAIExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req outputItemsByIndex := make(map[int64][]byte) var outputItemsFallback [][]byte responseFilter := newXAIInternalXSearchResponseFilter(prepared.filterInternalXSearch, prepared.clientDeclaredTools) + namespaceRestorer := newXAINamespaceRestorer(prepared.namespaceTools) for _, line := range bytes.Split(data, []byte("\n")) { if !bytes.HasPrefix(line, xaiDataTag) { continue } eventData := xaiNormalizeReasoningSummaryData(bytes.TrimSpace(line[len(xaiDataTag):])) - eventData = restoreXAINamespaceToolCalls(eventData, prepared.namespaceTools) + eventData = namespaceRestorer.restore(eventData) eventData = responseFilter.apply(eventData) if len(eventData) == 0 { continue diff --git a/internal/runtime/executor/xai_executor_request.go b/internal/runtime/executor/xai_executor_request.go index 36cad26f..c92f17f7 100644 --- a/internal/runtime/executor/xai_executor_request.go +++ b/internal/runtime/executor/xai_executor_request.go @@ -37,8 +37,9 @@ type xaiPreparedRequest struct { } type xaiNamespaceToolRef struct { - namespace string - name string + namespace string + name string + isDispatcher bool } // xaiClientToolKey identifies a client-declared callable tool using the @@ -88,15 +89,17 @@ func (e *XAIExecutor) prepareResponsesRequestTo(ctx context.Context, req cliprox body, _ = sjson.DeleteBytes(body, "safety_identifier") body, _ = sjson.DeleteBytes(body, "stream_options") body = helps.RewriteCodexMultiAgentV2Input(ctx, opts.Headers, body, e.cfg) - namespaceTools := collectXAINamespaceToolRefs(body) + willInjectXSearch := e.cfg != nil && e.cfg.XAI.InjectXSearch + shouldFold := xaiShouldFoldNamespaceTools(body, willInjectXSearch) + namespaceTools := collectXAINamespaceToolRefsWithFold(body, shouldFold) // Collect before normalizeXAITools flattens namespace wrappers so keys match // the post-restore (namespace, short-name) shape used by the response filter. clientDeclaredTools := collectXAIClientDeclaredToolKeys(body) - body = normalizeXAITools(body) + body = normalizeXAIToolsWithFold(body, shouldFold) body = promoteXAIAdditionalTools(body) // Drop choices that point at tools removed by normalizeXAITools before any // configured x_search injection, so no surviving choice references a deleted tool. - body = normalizeXAINamespaceToolChoice(body) + body = normalizeXAINamespaceToolChoiceWithFold(body, shouldFold) body = normalizeXAIForcedWebSearchToolChoice(body) // Prune before rewriting image_generation choices so older models that still // strip the tool do not keep a leftover "required" selection. @@ -109,13 +112,14 @@ func (e *XAIExecutor) prepareResponsesRequestTo(ctx context.Context, req cliprox if e.cfg != nil && e.cfg.XAI.InjectXSearch && !xaiToolChoiceRequiresImageGenerationOnly(body) { body = ensureXAINativeXSearchTool(body) } + body = clampXAIToolsLimit(body, xaiMaxTools, namespaceTools) var replayScope xaiReasoningReplayScope body, replayScope, err = applyXAIReasoningReplayCacheRequired(ctx, from, req, opts, body) if err != nil { return nil, err } body = normalizeXAIInputCustomToolCalls(body) - body = normalizeXAIInputNamespaceToolCalls(body) + body = normalizeXAIInputNamespaceToolCallsWithFold(body, shouldFold) body = normalizeXAIInputReasoningItems(body) body = sanitizeXAIInputEncryptedContent(body) body = normalizeCodexInstructions(body) @@ -925,7 +929,159 @@ func xaiToolChoiceMatchesAvailable(choice gjson.Result, available map[xaiToolCho return ok } +func xaiCountFlattenedTools(tools gjson.Result) int { + if !tools.Exists() || !tools.IsArray() { + return 0 + } + count := 0 + for _, tool := range tools.Array() { + switch tool.Get("type").String() { + case xaiNamespaceToolType: + if nestedTools := tool.Get("tools"); nestedTools.IsArray() { + count += len(nestedTools.Array()) + } else { + count++ + } + case xaiToolSearchType: + // Tool search is stripped by normalizeXAITool + default: + count++ + } + } + return count +} + +func xaiTotalFlattenedToolsCount(body []byte, willInjectXSearch bool) int { + count := xaiCountFlattenedTools(gjson.GetBytes(body, "tools")) + input := gjson.GetBytes(body, "input") + if input.Exists() && input.IsArray() { + for _, item := range input.Array() { + if item.Get("type").String() == "additional_tools" { + count += xaiCountFlattenedTools(item.Get("tools")) + } + } + } + if willInjectXSearch && !xaiRequestHasNativeXSearch(body) && !xaiToolChoiceRequiresImageGenerationOnly(body) { + count++ + } + return count +} + +func xaiShouldFoldNamespaceTools(body []byte, willInjectXSearch bool) bool { + return xaiTotalFlattenedToolsCount(body, willInjectXSearch) > xaiMaxTools +} + +func buildXAINamespaceDispatcherTool(tool gjson.Result) []byte { + namespaceName := strings.TrimSpace(tool.Get("name").String()) + if namespaceName == "" { + return nil + } + description := strings.TrimSpace(tool.Get("description").String()) + + var toolNames []string + var toolDescriptions []string + if nestedTools := tool.Get("tools"); nestedTools.IsArray() { + for _, child := range nestedTools.Array() { + childName := strings.TrimSpace(child.Get("name").String()) + if childName == "" { + continue + } + toolNames = append(toolNames, childName) + childDesc := strings.TrimSpace(child.Get("description").String()) + + params := child.Get("parameters") + if !params.Exists() { + params = child.Get("input_schema") + } + + var paramStr string + if params.Exists() && params.Raw != "" { + rawParams := strings.TrimSpace(params.Raw) + if rawParams != "" && rawParams != "{}" && rawParams != `{"type":"object","properties":{}}` { + inlined := util.InlineLocalRefs(rawParams) + if gjson.Valid(inlined) { + cleaned := []byte(inlined) + if gjson.GetBytes(cleaned, "$defs").Exists() { + cleaned, _ = sjson.DeleteBytes(cleaned, "$defs") + } + if gjson.GetBytes(cleaned, "definitions").Exists() { + cleaned, _ = sjson.DeleteBytes(cleaned, "definitions") + } + paramStr = string(cleaned) + } else { + paramStr = inlined + } + } + } + + var entry string + if childDesc != "" { + if paramStr != "" { + entry = fmt.Sprintf("- %s: %s\n Parameters: %s", childName, childDesc, paramStr) + } else { + entry = fmt.Sprintf("- %s: %s", childName, childDesc) + } + } else { + if paramStr != "" { + entry = fmt.Sprintf("- %s\n Parameters: %s", childName, paramStr) + } else { + entry = fmt.Sprintf("- %s", childName) + } + } + toolDescriptions = append(toolDescriptions, entry) + } + } + + fullDescription := description + if len(toolDescriptions) > 0 { + catalog := "Available tools in this namespace:\n" + strings.Join(toolDescriptions, "\n") + if fullDescription != "" { + fullDescription += "\n\n" + catalog + } else { + fullDescription = fmt.Sprintf("Tools in namespace %s.\n\n%s", namespaceName, catalog) + } + } else if fullDescription == "" { + fullDescription = fmt.Sprintf("Tools in namespace %s.", namespaceName) + } + + nameProp := map[string]any{ + "type": "string", + "description": fmt.Sprintf("Child tool name to execute in namespace %s", namespaceName), + } + if len(toolNames) > 0 { + nameProp["enum"] = toolNames + } + + dispatcher := map[string]any{ + "type": xaiFunctionToolType, + "name": namespaceName, + "description": fullDescription, + "parameters": map[string]any{ + "type": "object", + "properties": map[string]any{ + "name": nameProp, + "arguments": map[string]any{ + "type": "object", + "description": "Arguments object matching the parameter schema of the selected child tool", + "additionalProperties": true, + }, + }, + "required": []string{"name"}, + }, + } + + raw, errMarshal := json.Marshal(dispatcher) + if errMarshal != nil { + return nil + } + return raw +} + func normalizeXAITools(body []byte) []byte { + return normalizeXAIToolsWithFold(body, xaiShouldFoldNamespaceTools(body, false)) +} + +func normalizeXAIToolsWithFold(body []byte, shouldFold bool) []byte { if !gjson.ValidBytes(body) { return body } @@ -936,7 +1092,7 @@ func normalizeXAITools(body []byte) []byte { if !tools.Exists() || !tools.IsArray() { return true } - filtered, changed, ok := normalizeXAIToolArray(tools, keepImageGeneration) + filtered, changed, ok := normalizeXAIToolArray(tools, keepImageGeneration, shouldFold) if !ok { return false } @@ -968,6 +1124,75 @@ func normalizeXAITools(body []byte) []byte { return body } +func xaiHasFunctionToolNamed(body []byte, name string) bool { + if name == "" { + return false + } + tools := gjson.GetBytes(body, "tools") + if tools.IsArray() { + for _, tool := range tools.Array() { + if tool.Get("type").String() == xaiFunctionToolType && tool.Get("name").String() == name { + return true + } + } + } + input := gjson.GetBytes(body, "input") + if input.IsArray() { + for _, item := range input.Array() { + if item.Get("type").String() == "additional_tools" { + for _, tool := range item.Get("tools").Array() { + if tool.Get("type").String() == xaiFunctionToolType && tool.Get("name").String() == name { + return true + } + } + } + } + } + return false +} + +func clampXAIToolsLimit(body []byte, maxTools int, refs map[string]xaiNamespaceToolRef) []byte { + tools := gjson.GetBytes(body, "tools") + if !tools.IsArray() || len(tools.Array()) <= maxTools { + return body + } + allTools := tools.Array() + var dispatcherTools []json.RawMessage + var regularTools []json.RawMessage + for _, tool := range allTools { + name := strings.TrimSpace(tool.Get("name").String()) + if ref, ok := refs[name]; ok && ref.isDispatcher { + dispatcherTools = append(dispatcherTools, json.RawMessage(tool.Raw)) + } else { + regularTools = append(regularTools, json.RawMessage(tool.Raw)) + } + } + + capped := make([]json.RawMessage, 0, maxTools) + if len(dispatcherTools) >= maxTools { + capped = append(capped, dispatcherTools[:maxTools]...) + } else { + capped = append(capped, dispatcherTools...) + remainingSlots := maxTools - len(dispatcherTools) + if len(regularTools) > remainingSlots { + capped = append(capped, regularTools[:remainingSlots]...) + } else { + capped = append(capped, regularTools...) + } + } + + raw, errMarshal := json.Marshal(capped) + if errMarshal != nil { + return body + } + updated, errSet := sjson.SetRawBytes(body, "tools", raw) + if errSet != nil { + return body + } + updated = pruneXAIOrphanedToolChoice(updated) + return normalizeXAIToolChoiceForTools(updated) +} + // 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 { @@ -1026,7 +1251,7 @@ func promoteXAIAdditionalTools(body []byte) []byte { return updated } -func normalizeXAIToolArray(tools gjson.Result, keepImageGeneration bool) ([]byte, bool, bool) { +func normalizeXAIToolArray(tools gjson.Result, keepImageGeneration, shouldFold bool) ([]byte, bool, bool) { toolItems := tools.Array() filtered := make([][]byte, 0, len(toolItems)) changed := false @@ -1034,6 +1259,12 @@ func normalizeXAIToolArray(tools gjson.Result, keepImageGeneration bool) ([]byte toolType := tool.Get("type").String() if toolType == xaiNamespaceToolType { changed = true + if shouldFold { + if dispatcher := buildXAINamespaceDispatcherTool(tool); len(dispatcher) > 0 { + filtered = append(filtered, dispatcher) + } + continue + } namespaceName := tool.Get("name").String() if namespaceTools := tool.Get("tools"); namespaceTools.IsArray() { for _, nestedTool := range namespaceTools.Array() { @@ -1102,6 +1333,10 @@ func normalizeXAIToolChoiceForTools(body []byte) []byte { // the same names sent in the flattened tools list. xAI does not accept the // Responses namespace field on tool choices. func normalizeXAINamespaceToolChoice(body []byte) []byte { + return normalizeXAINamespaceToolChoiceWithFold(body, xaiShouldFoldNamespaceTools(body, false)) +} + +func normalizeXAINamespaceToolChoiceWithFold(body []byte, shouldFold bool) []byte { if !gjson.ValidBytes(body) { return body } @@ -1113,11 +1348,24 @@ func normalizeXAINamespaceToolChoice(body []byte) []byte { } namespaceName := strings.TrimSpace(toolChoice.Get("namespace").String()) toolName := strings.TrimSpace(toolChoice.Get("name").String()) + if namespaceName == "" { + return true + } qualifiedName := qualifyXAINamespaceToolName(namespaceName, toolName) - if namespaceName == "" || qualifiedName == "" { + var targetName string + if xaiHasFunctionToolNamed(body, namespaceName) { + targetName = namespaceName + } else if xaiHasFunctionToolNamed(body, qualifiedName) { + targetName = qualifiedName + } else if shouldFold { + targetName = namespaceName + } else { + targetName = qualifiedName + } + if targetName == "" { return true } - updated, errSet := sjson.SetBytes(body, path+".name", qualifiedName) + updated, errSet := sjson.SetBytes(body, path+".name", targetName) if errSet != nil { return false } @@ -1261,6 +1509,10 @@ func qualifyXAINamespaceToolName(namespaceName, toolName string) string { } func collectXAINamespaceToolRefs(body []byte) map[string]xaiNamespaceToolRef { + return collectXAINamespaceToolRefsWithFold(body, xaiShouldFoldNamespaceTools(body, false)) +} + +func collectXAINamespaceToolRefsWithFold(body []byte, shouldFold bool) map[string]xaiNamespaceToolRef { refs := make(map[string]xaiNamespaceToolRef) collect := func(tools gjson.Result) { if !tools.Exists() || !tools.IsArray() { @@ -1274,13 +1526,16 @@ func collectXAINamespaceToolRefs(body []byte) map[string]xaiNamespaceToolRef { if namespaceName == "" { continue } + if shouldFold { + refs[namespaceName] = xaiNamespaceToolRef{namespace: namespaceName, name: "", isDispatcher: true} + } for _, nestedTool := range tool.Get("tools").Array() { toolName := strings.TrimSpace(nestedTool.Get("name").String()) qualifiedName := qualifyXAINamespaceToolName(namespaceName, toolName) if qualifiedName == "" { continue } - refs[qualifiedName] = xaiNamespaceToolRef{namespace: namespaceName, name: toolName} + refs[qualifiedName] = xaiNamespaceToolRef{namespace: namespaceName, name: toolName, isDispatcher: false} } } } diff --git a/internal/runtime/executor/xai_executor_response.go b/internal/runtime/executor/xai_executor_response.go index 9b19db29..96cfdfee 100644 --- a/internal/runtime/executor/xai_executor_response.go +++ b/internal/runtime/executor/xai_executor_response.go @@ -281,6 +281,10 @@ func (f *xaiInternalXSearchResponseFilter) filterCompletedOutput(eventData []byt } func normalizeXAIInputNamespaceToolCalls(body []byte) []byte { + return normalizeXAIInputNamespaceToolCallsWithFold(body, xaiShouldFoldNamespaceTools(body, false)) +} + +func normalizeXAIInputNamespaceToolCallsWithFold(body []byte, shouldFold bool) []byte { if !gjson.ValidBytes(body) { return body } @@ -294,8 +298,55 @@ func normalizeXAIInputNamespaceToolCalls(body []byte) []byte { } namespaceName := strings.TrimSpace(item.Get("namespace").String()) toolName := strings.TrimSpace(item.Get("name").String()) + if namespaceName == "" { + continue + } qualifiedName := qualifyXAINamespaceToolName(namespaceName, toolName) - if namespaceName == "" || qualifiedName == "" { + var isFolded bool + if xaiHasFunctionToolNamed(body, namespaceName) { + isFolded = true + } else if xaiHasFunctionToolNamed(body, qualifiedName) { + isFolded = false + } else { + isFolded = shouldFold + } + if isFolded { + namePath := fmt.Sprintf("input.%d.name", index) + namespacePath := fmt.Sprintf("input.%d.namespace", index) + argsPath := fmt.Sprintf("input.%d.arguments", index) + + dispatcherArgs := map[string]any{ + "name": toolName, + } + if rawArgs := item.Get("arguments").String(); rawArgs != "" { + if gjson.Valid(rawArgs) { + dispatcherArgs["arguments"] = json.RawMessage(rawArgs) + } else { + dispatcherArgs["arguments"] = rawArgs + } + } + encodedArgs, errMarshal := json.Marshal(dispatcherArgs) + if errMarshal != nil { + continue + } + + updated, errSet := sjson.SetBytes(body, namePath, namespaceName) + if errSet != nil { + continue + } + updated, errSet = sjson.SetBytes(updated, argsPath, string(encodedArgs)) + if errSet != nil { + continue + } + updated, errDelete := sjson.DeleteBytes(updated, namespacePath) + if errDelete != nil { + continue + } + body = updated + continue + } + + if qualifiedName == "" { continue } namePath := fmt.Sprintf("input.%d.name", index) @@ -313,29 +364,95 @@ func normalizeXAIInputNamespaceToolCalls(body []byte) []byte { return body } -func restoreXAINamespaceToolCalls(data []byte, refs map[string]xaiNamespaceToolRef) []byte { - if len(refs) == 0 || len(data) == 0 || !gjson.ValidBytes(data) { +type xaiNamespaceRestorer struct { + refs map[string]xaiNamespaceToolRef + dispatcherItemIDs map[string]string +} + +func newXAINamespaceRestorer(refs map[string]xaiNamespaceToolRef) *xaiNamespaceRestorer { + return &xaiNamespaceRestorer{ + refs: refs, + dispatcherItemIDs: make(map[string]string), + } +} + +func (r *xaiNamespaceRestorer) restore(data []byte) []byte { + if r == nil || len(r.refs) == 0 || len(data) == 0 || !gjson.ValidBytes(data) { return data } - data = restoreXAINamespaceToolCallAtPath(data, "item", refs) - output := gjson.GetBytes(data, "response.output") - if output.Exists() && output.IsArray() { - for index := range output.Array() { - data = restoreXAINamespaceToolCallAtPath(data, fmt.Sprintf("response.output.%d", index), refs) + eventType := gjson.GetBytes(data, "type").String() + switch eventType { + case "response.output_item.added": + item := gjson.GetBytes(data, "item") + if item.Get("type").String() == "function_call" { + name := strings.TrimSpace(item.Get("name").String()) + itemID := strings.TrimSpace(item.Get("id").String()) + if ref, ok := r.refs[name]; ok && ref.isDispatcher { + if itemID != "" { + r.dispatcherItemIDs[itemID] = ref.namespace + } + data, _ = sjson.SetBytes(data, "item.namespace", ref.namespace) + } + } + return data + + case "response.function_call_arguments.done": + itemID := strings.TrimSpace(gjson.GetBytes(data, "item_id").String()) + if namespaceName, isDisp := r.dispatcherItemIDs[itemID]; isDisp { + rawArgs := gjson.GetBytes(data, "arguments").String() + if _, childArgs, ok := unwrapXAIDispatcherArguments(rawArgs, namespaceName, r.refs); ok { + updated, errSet := sjson.SetBytes(data, "arguments", string(childArgs)) + if errSet == nil { + data = updated + } + } } + return data + + default: + data = r.restoreAtPath(data, "item") + output := gjson.GetBytes(data, "response.output") + if output.Exists() && output.IsArray() { + for index := range output.Array() { + data = r.restoreAtPath(data, fmt.Sprintf("response.output.%d", index)) + } + } + return data } - return data } -func restoreXAINamespaceToolCallAtPath(data []byte, path string, refs map[string]xaiNamespaceToolRef) []byte { +func (r *xaiNamespaceRestorer) restoreAtPath(data []byte, path string) []byte { if gjson.GetBytes(data, path+".type").String() != "function_call" { return data } qualifiedName := strings.TrimSpace(gjson.GetBytes(data, path+".name").String()) - ref, ok := refs[qualifiedName] + ref, ok := r.refs[qualifiedName] if !ok { return data } + if ref.isDispatcher { + rawArgs := gjson.GetBytes(data, path+".arguments").String() + childName, childArgs, unwrapped := unwrapXAIDispatcherArguments(rawArgs, ref.namespace, r.refs) + if !unwrapped && childName == "" { + childName = ref.name + } + updated, errSet := sjson.SetBytes(data, path+".namespace", ref.namespace) + if errSet != nil { + return data + } + if childName != "" { + if updatedName, errSetName := sjson.SetBytes(updated, path+".name", childName); errSetName == nil { + updated = updatedName + } + } + if len(childArgs) > 0 { + if updatedArgs, errSetArgs := sjson.SetBytes(updated, path+".arguments", string(childArgs)); errSetArgs == nil { + updated = updatedArgs + } + } + return updated + } + updated, errSet := sjson.SetBytes(data, path+".name", ref.name) if errSet != nil { return data @@ -347,6 +464,64 @@ func restoreXAINamespaceToolCallAtPath(data []byte, path string, refs map[string return updated } +func unwrapXAIDispatcherArguments(rawArgs string, namespaceName string, refs map[string]xaiNamespaceToolRef) (string, []byte, bool) { + if !gjson.Valid(rawArgs) { + return "", nil, false + } + argsParsed := gjson.Parse(rawArgs) + nameField := argsParsed.Get("name") + if !nameField.Exists() || nameField.Type != gjson.String { + return "", nil, false + } + childName := strings.TrimSpace(nameField.String()) + if childName == "" { + return "", nil, false + } + + if namespaceName != "" { + qualified := qualifyXAINamespaceToolName(namespaceName, childName) + if ref, exists := refs[qualified]; exists && ref.isDispatcher { + return "", nil, false + } + } else { + isChildOfDispatcher := false + for _, ref := range refs { + if ref.isDispatcher && (ref.name == childName || ref.namespace == childName) { + isChildOfDispatcher = true + break + } + } + if !isChildOfDispatcher && !argsParsed.Get("arguments").Exists() { + return "", nil, false + } + } + + var childArgs []byte + if argsField := argsParsed.Get("arguments"); argsField.Exists() { + if argsField.Type == gjson.String { + childArgs = []byte(argsField.String()) + } else { + childArgs = []byte(argsField.Raw) + } + } else { + cleaned, errDel := sjson.DeleteBytes([]byte(rawArgs), "name") + if errDel == nil && len(cleaned) > 0 && string(cleaned) != "{}" { + childArgs = cleaned + } else { + childArgs = []byte("{}") + } + } + if len(childArgs) == 0 { + childArgs = []byte("{}") + } + return childName, childArgs, true +} + +func restoreXAINamespaceToolCalls(data []byte, refs map[string]xaiNamespaceToolRef) []byte { + restorer := newXAINamespaceRestorer(refs) + return restorer.restore(data) +} + // normalizeXAIObjectRootUnionBranchTypes makes untyped root union branches // explicitly object-only when the parameter root already permits only objects. // This preserves the original schema semantics while satisfying xAI validation. diff --git a/internal/runtime/executor/xai_executor_stream.go b/internal/runtime/executor/xai_executor_stream.go index 2b5f3e86..f06c1c43 100644 --- a/internal/runtime/executor/xai_executor_stream.go +++ b/internal/runtime/executor/xai_executor_stream.go @@ -81,6 +81,7 @@ func (e *XAIExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth outputItemsByIndex := make(map[int64][]byte) var outputItemsFallback [][]byte responseFilter := newXAIInternalXSearchResponseFilter(prepared.filterInternalXSearch, prepared.clientDeclaredTools) + namespaceRestorer := newXAINamespaceRestorer(prepared.namespaceTools) var pendingEventLine []byte emitTranslatedLine := func(translatedLine []byte) bool { chunks := helps.TranslateStreamWithClaudeInputTokens(ctx, prepared.to, prepared.responseFormat, req.Model, prepared.originalPayload, prepared.body, translatedLine, ¶m, claudeInputTokens) @@ -109,7 +110,7 @@ func (e *XAIExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth eventDataList := xaiNormalizeReasoningSummaryDataEvents(bytes.TrimSpace(line[len(xaiDataTag):])) hasPendingEventLine := pendingEventLine != nil for i, eventData := range eventDataList { - eventData = restoreXAINamespaceToolCalls(eventData, prepared.namespaceTools) + eventData = namespaceRestorer.restore(eventData) eventData = responseFilter.apply(eventData) if len(eventData) == 0 { if hasPendingEventLine && i == 0 { diff --git a/internal/runtime/executor/xai_executor_test.go b/internal/runtime/executor/xai_executor_test.go index 6407fff6..693cd002 100644 --- a/internal/runtime/executor/xai_executor_test.go +++ b/internal/runtime/executor/xai_executor_test.go @@ -379,6 +379,185 @@ func TestXAIExecutorPrepareResponsesRequestRewritesCodexAgentMessage(t *testing. } } +func TestXAIExecutorExecuteFoldsNamespacesWhenToolsExceed200(t *testing.T) { + var gotBody []byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var errRead error + gotBody, errRead = io.ReadAll(r.Body) + if errRead != nil { + t.Fatalf("read body: %v", errRead) + } + w.Header().Set("Content-Type", "text/event-stream") + _, _ = w.Write([]byte("data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"type\":\"function_call\",\"name\":\"mcp__app_0\",\"call_id\":\"call_1\",\"arguments\":\"{\\\"name\\\":\\\"tool_2\\\",\\\"arguments\\\":{\\\"q\\\":\\\"test\\\"}}\"}}\n\n")) + _, _ = w.Write([]byte("data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_1\",\"object\":\"response\",\"status\":\"completed\",\"model\":\"grok-4.6\",\"output\":[{\"type\":\"function_call\",\"name\":\"mcp__app_0\",\"call_id\":\"call_1\",\"arguments\":\"{\\\"name\\\":\\\"tool_2\\\",\\\"arguments\\\":{\\\"q\\\":\\\"test\\\"}}\"}],\"usage\":{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2}}}\n\n")) + })) + defer server.Close() + + var nsList []string + for i := 0; i < 47; i++ { + var childTools []string + for j := 0; j < 10; j++ { + childTools = append(childTools, fmt.Sprintf(`{"type":"function","name":"tool_%d","description":"child tool %d","parameters":{"type":"object","properties":{"q":{"type":"string"}}}}`, j, j)) + } + nsList = append(nsList, fmt.Sprintf(`{"type":"namespace","name":"mcp__app_%d","description":"App %d tools","tools":[%s]}`, i, i, strings.Join(childTools, ","))) + } + + exec := NewXAIExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{ + Provider: "xai", + Attributes: map[string]string{"base_url": server.URL}, + Metadata: map[string]any{"access_token": "xai-token"}, + } + + turn1Payload := fmt.Sprintf(`{ + "model":"grok-4.6", + "tools":[%s], + "input":[{"role":"user","content":"call tool_2"}] + }`, strings.Join(nsList, ",")) + + resp, err := exec.Execute(context.Background(), auth, cliproxyexecutor.Request{ + Model: "grok-4.6", + Payload: []byte(turn1Payload), + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + Stream: false, + }) + if err != nil { + t.Fatalf("Execute() error = %v", err) + } + + tools := gjson.GetBytes(gotBody, "tools").Array() + if len(tools) != 47 { + t.Fatalf("upstream tools count = %d, want 47 dispatcher tools; body=%s", len(tools), string(gotBody)) + } + if got := tools[0].Get("name").String(); got != "mcp__app_0" { + t.Fatalf("upstream tools[0].name = %q, want mcp__app_0", got) + } + if got := tools[0].Get("type").String(); got != "function" { + t.Fatalf("upstream tools[0].type = %q, want function", got) + } + + output := gjson.GetBytes(resp.Payload, "output.0") + if got := output.Get("name").String(); got != "tool_2" { + t.Fatalf("response output name = %q, want tool_2; payload=%s", got, resp.Payload) + } + if got := output.Get("namespace").String(); got != "mcp__app_0" { + t.Fatalf("response output namespace = %q, want mcp__app_0; payload=%s", got, resp.Payload) + } + if got := output.Get("arguments").String(); got != `{"q":"test"}` { + t.Fatalf("response output arguments = %q, want {\"q\":\"test\"}; payload=%s", got, resp.Payload) + } + + // Turn 2: client passes back the restored function_call and function_call_output + turn2Payload := fmt.Sprintf(`{ + "model":"grok-4.6", + "tools":[%s], + "input":[ + {"role":"user","content":"call tool_2"}, + {"type":"function_call","name":"tool_2","namespace":"mcp__app_0","call_id":"call_1","arguments":"{\"q\":\"test\"}"}, + {"type":"function_call_output","call_id":"call_1","output":"ok"} + ] + }`, strings.Join(nsList, ",")) + + _, err2 := exec.Execute(context.Background(), auth, cliproxyexecutor.Request{ + Model: "grok-4.6", + Payload: []byte(turn2Payload), + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + Stream: false, + }) + if err2 != nil { + t.Fatalf("Execute() turn 2 error = %v", err2) + } + + input := gjson.GetBytes(gotBody, "input").Array() + if len(input) != 3 { + t.Fatalf("upstream input count = %d, want 3; body=%s", len(input), string(gotBody)) + } + fcItem := input[1] + if got := fcItem.Get("name").String(); got != "mcp__app_0" { + t.Fatalf("turn 2 upstream fc name = %q, want mcp__app_0; body=%s", got, string(gotBody)) + } + if fcItem.Get("namespace").Exists() { + t.Fatalf("turn 2 upstream fc namespace should be removed; body=%s", string(gotBody)) + } + if gotArgs := fcItem.Get("arguments").String(); gotArgs != `{"arguments":{"q":"test"},"name":"tool_2"}` && gotArgs != `{"name":"tool_2","arguments":{"q":"test"}}` { + t.Fatalf("turn 2 upstream fc arguments = %q, want wrapped dispatcher args; body=%s", gotArgs, string(gotBody)) + } +} + +func TestXAIExecutorExecuteStreamFoldsNamespacesWhenToolsExceed200(t *testing.T) { + var gotBody []byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var errRead error + gotBody, errRead = io.ReadAll(r.Body) + if errRead != nil { + t.Fatalf("read body: %v", errRead) + } + w.Header().Set("Content-Type", "text/event-stream") + _, _ = w.Write([]byte("data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"type\":\"function_call\",\"name\":\"mcp__app_0\",\"call_id\":\"call_1\",\"arguments\":\"{\\\"name\\\":\\\"tool_2\\\",\\\"arguments\\\":{\\\"q\\\":\\\"stream_test\\\"}}\"}}\n\n")) + _, _ = w.Write([]byte("data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_1\",\"object\":\"response\",\"status\":\"completed\",\"model\":\"grok-4.6\",\"output\":[{\"type\":\"function_call\",\"name\":\"mcp__app_0\",\"call_id\":\"call_1\",\"arguments\":\"{\\\"name\\\":\\\"tool_2\\\",\\\"arguments\\\":{\\\"q\\\":\\\"stream_test\\\"}}\"}],\"usage\":{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2}}}\n\n")) + })) + defer server.Close() + + var nsList []string + for i := 0; i < 47; i++ { + var childTools []string + for j := 0; j < 10; j++ { + childTools = append(childTools, fmt.Sprintf(`{"type":"function","name":"tool_%d","parameters":{"type":"object","properties":{"q":{"type":"string"}}}}`, j)) + } + nsList = append(nsList, fmt.Sprintf(`{"type":"namespace","name":"mcp__app_%d","tools":[%s]}`, i, strings.Join(childTools, ","))) + } + + exec := NewXAIExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{ + Provider: "xai", + Attributes: map[string]string{"base_url": server.URL}, + Metadata: map[string]any{"access_token": "xai-token"}, + } + + turn1Payload := fmt.Sprintf(`{ + "model":"grok-4.6", + "tools":[%s], + "input":[{"role":"user","content":"call tool_2"}] + }`, strings.Join(nsList, ",")) + + result, err := exec.ExecuteStream(context.Background(), auth, cliproxyexecutor.Request{ + Model: "grok-4.6", + Payload: []byte(turn1Payload), + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + Stream: true, + }) + if err != nil { + t.Fatalf("ExecuteStream() error = %v", err) + } + + var chunks [][]byte + for chunk := range result.Chunks { + if chunk.Err != nil { + t.Fatalf("Stream chunk error = %v", chunk.Err) + } + chunks = append(chunks, chunk.Payload) + } + + tools := gjson.GetBytes(gotBody, "tools").Array() + if len(tools) != 47 { + t.Fatalf("upstream tools count = %d, want 47 dispatcher tools; body=%s", len(tools), string(gotBody)) + } + + allStreamText := string(bytes.Join(chunks, nil)) + if !strings.Contains(allStreamText, `"namespace":"mcp__app_0"`) { + t.Fatalf("stream chunks missing restored namespace: %s", allStreamText) + } + if !strings.Contains(allStreamText, `"name":"tool_2"`) { + t.Fatalf("stream chunks missing restored tool name: %s", allStreamText) + } +} + func TestXAIExecutorExecuteRestoresAdditionalToolsNamespaceCalls(t *testing.T) { var gotBody []byte server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -3998,6 +4177,306 @@ func TestNormalizeXAITools_QualifiesSameNamedNamespaceTools(t *testing.T) { } } +func TestNormalizeXAITools_WhenFlattenedCountExceedsLimit_FoldsNamespaces(t *testing.T) { + // 47 namespaces with 10 tools each = 470 tools total (> 200 limit) + var nsList []string + for i := 0; i < 47; i++ { + var childTools []string + for j := 0; j < 10; j++ { + childTools = append(childTools, fmt.Sprintf(`{"type":"function","name":"tool_%d","description":"child tool %d","parameters":{"type":"object","properties":{"p":{"type":"string"}}}}`, j, j)) + } + nsList = append(nsList, fmt.Sprintf(`{"type":"namespace","name":"mcp__app_%d","tools":[%s]}`, i, strings.Join(childTools, ","))) + } + body := []byte(fmt.Sprintf(`{"tools":[%s]}`, strings.Join(nsList, ","))) + + refs := collectXAINamespaceToolRefs(body) + out := normalizeXAITools(body) + + tools := gjson.GetBytes(out, "tools").Array() + if len(tools) > 200 { + t.Fatalf("tools length = %d, must be <= 200; body length=%d", len(tools), len(out)) + } + if len(tools) != 47 { + t.Fatalf("tools length = %d, want 47 dispatcher tools", len(tools)) + } + if got := tools[0].Get("name").String(); got != "mcp__app_0" { + t.Fatalf("tools.0.name = %q, want mcp__app_0", got) + } + desc := tools[0].Get("description").String() + if !strings.Contains(desc, "Parameters:") || !strings.Contains(desc, `"properties":{"p":{"type":"string"}}`) { + t.Fatalf("dispatcher description missing child parameter schemas: %s", desc) + } + + // Verify restore of a dispatcher tool call + event := []byte(`{"type":"response.output_item.done","item":{"type":"function_call","name":"mcp__app_0","call_id":"call_1","arguments":"{\"name\":\"tool_3\",\"arguments\":{\"p\":\"val\"}}"}}`) + restored := restoreXAINamespaceToolCalls(event, refs) + if got := gjson.GetBytes(restored, "item.name").String(); got != "tool_3" { + t.Fatalf("restored item.name = %q, want tool_3; event=%s", got, string(restored)) + } + if got := gjson.GetBytes(restored, "item.namespace").String(); got != "mcp__app_0" { + t.Fatalf("restored item.namespace = %q, want mcp__app_0; event=%s", got, string(restored)) + } + if got := gjson.GetBytes(restored, "item.arguments").String(); got != `{"p":"val"}` { + t.Fatalf("restored item.arguments = %q, want {\"p\":\"val\"}; event=%s", got, string(restored)) + } + + // Verify normalization of historical tool call in next turn + turn2Body := []byte(fmt.Sprintf(`{"tools":[%s],"input":[{"type":"function_call","name":"tool_3","namespace":"mcp__app_0","call_id":"call_1","arguments":"{\"p\":\"val\"}"}]}`, strings.Join(nsList, ","))) + turn2Normalized := normalizeXAIInputNamespaceToolCalls(turn2Body) + if got := gjson.GetBytes(turn2Normalized, "input.0.name").String(); got != "mcp__app_0" { + t.Fatalf("turn2 normalized name = %q, want mcp__app_0; body=%s", got, string(turn2Normalized)) + } + if gjson.GetBytes(turn2Normalized, "input.0.namespace").Exists() { + t.Fatalf("turn2 namespace should be removed; body=%s", string(turn2Normalized)) + } + if got := gjson.GetBytes(turn2Normalized, "input.0.arguments").String(); got != `{"arguments":{"p":"val"},"name":"tool_3"}` && got != `{"name":"tool_3","arguments":{"p":"val"}}` { + t.Fatalf("turn2 normalized arguments = %q; body=%s", got, string(turn2Normalized)) + } +} + +func TestNormalizeXAINamespaceToolChoice_WhenFolding(t *testing.T) { + // Create payload with > 200 tools + var nsList []string + for i := 0; i < 25; i++ { + var childTools []string + for j := 0; j < 10; j++ { + childTools = append(childTools, fmt.Sprintf(`{"type":"function","name":"tool_%d","parameters":{"type":"object"}}`, j)) + } + nsList = append(nsList, fmt.Sprintf(`{"type":"namespace","name":"mcp__app_%d","tools":[%s]}`, i, strings.Join(childTools, ","))) + } + body := []byte(fmt.Sprintf(`{"tools":[%s],"tool_choice":{"type":"function","name":"tool_3","namespace":"mcp__app_0"}}`, strings.Join(nsList, ","))) + + out := normalizeXAITools(body) + out = normalizeXAINamespaceToolChoice(out) + + if got := gjson.GetBytes(out, "tool_choice.name").String(); got != "mcp__app_0" { + t.Fatalf("tool_choice.name = %q, want mcp__app_0; body=%s", got, string(out)) + } + if gjson.GetBytes(out, "tool_choice.namespace").Exists() { + t.Fatalf("tool_choice.namespace should be removed; body=%s", string(out)) + } +} + +func TestRestoreXAINamespaceToolCalls_DispatcherVariants(t *testing.T) { + refs := map[string]xaiNamespaceToolRef{ + "mcp__app_0": {namespace: "mcp__app_0", name: "", isDispatcher: true}, + } + + // Variant 1: arguments flattened in root + event1 := []byte(`{"type":"response.output_item.done","item":{"type":"function_call","name":"mcp__app_0","call_id":"c1","arguments":"{\"name\":\"tool_a\",\"arg1\":\"val1\"}"}}`) + restored1 := restoreXAINamespaceToolCalls(event1, refs) + if got := gjson.GetBytes(restored1, "item.name").String(); got != "tool_a" { + t.Fatalf("variant 1 name = %q, want tool_a", got) + } + if got := gjson.GetBytes(restored1, "item.namespace").String(); got != "mcp__app_0" { + t.Fatalf("variant 1 namespace = %q, want mcp__app_0", got) + } + if got := gjson.GetBytes(restored1, "item.arguments").String(); got != `{"arg1":"val1"}` { + t.Fatalf("variant 1 arguments = %q, want {\"arg1\":\"val1\"}", got) + } + + // Variant 2: arguments in nested string + event2 := []byte(`{"type":"response.output_item.done","item":{"type":"function_call","name":"mcp__app_0","call_id":"c2","arguments":"{\"name\":\"tool_b\",\"arguments\":\"{\\\"k\\\":\\\"v\\\"}\"}"}}`) + restored2 := restoreXAINamespaceToolCalls(event2, refs) + if got := gjson.GetBytes(restored2, "item.name").String(); got != "tool_b" { + t.Fatalf("variant 2 name = %q, want tool_b", got) + } + if got := gjson.GetBytes(restored2, "item.arguments").String(); got != `{"k":"v"}` { + t.Fatalf("variant 2 arguments = %q, want {\"k\":\"v\"}", got) + } + + // Variant 3: no arguments property + event3 := []byte(`{"type":"response.output_item.done","item":{"type":"function_call","name":"mcp__app_0","call_id":"c3","arguments":"{\"name\":\"tool_c\"}"}}`) + restored3 := restoreXAINamespaceToolCalls(event3, refs) + if got := gjson.GetBytes(restored3, "item.name").String(); got != "tool_c" { + t.Fatalf("variant 3 name = %q, want tool_c", got) + } + if got := gjson.GetBytes(restored3, "item.arguments").String(); got != `{}` { + t.Fatalf("variant 3 arguments = %q, want {}", got) + } +} + +func TestRestoreXAINamespaceToolCalls_FunctionCallArgumentsDone(t *testing.T) { + refs := map[string]xaiNamespaceToolRef{ + "mcp__app_0": {namespace: "mcp__app_0", name: "", isDispatcher: true}, + } + restorer := newXAINamespaceRestorer(refs) + + addedEvent := []byte(`{"type":"response.output_item.added","output_index":0,"item":{"id":"item_1","type":"function_call","name":"mcp__app_0"}}`) + restoredAdded := restorer.restore(addedEvent) + if got := gjson.GetBytes(restoredAdded, "item.namespace").String(); got != "mcp__app_0" { + t.Fatalf("restored item.namespace = %q, want mcp__app_0", got) + } + + doneArgsEvent := []byte(`{"type":"response.function_call_arguments.done","item_id":"item_1","output_index":0,"arguments":"{\"name\":\"tool_x\",\"arguments\":{\"count\":42}}"}`) + restoredArgs := restorer.restore(doneArgsEvent) + if got := gjson.GetBytes(restoredArgs, "arguments").String(); got != `{"count":42}` { + t.Fatalf("restored arguments.done = %q, want {\"count\":42}", got) + } + + doneItemEvent := []byte(`{"type":"response.output_item.done","output_index":0,"item":{"id":"item_1","type":"function_call","name":"mcp__app_0","arguments":"{\"name\":\"tool_x\",\"arguments\":{\"count\":42}}"}}`) + restoredItem := restorer.restore(doneItemEvent) + if got := gjson.GetBytes(restoredItem, "item.name").String(); got != "tool_x" { + t.Fatalf("restored item.name = %q, want tool_x", got) + } + if got := gjson.GetBytes(restoredItem, "item.namespace").String(); got != "mcp__app_0" { + t.Fatalf("restored item.namespace = %q, want mcp__app_0", got) + } +} + +func TestNormalizeXAIInputNamespaceToolCalls_PreservesLargeInts(t *testing.T) { + var nsList []string + for i := 0; i < 25; i++ { + var childTools []string + for j := 0; j < 10; j++ { + childTools = append(childTools, fmt.Sprintf(`{"type":"function","name":"tool_%d","parameters":{"type":"object"}}`, j)) + } + nsList = append(nsList, fmt.Sprintf(`{"type":"namespace","name":"mcp__app_%d","tools":[%s]}`, i, strings.Join(childTools, ","))) + } + body := []byte(fmt.Sprintf(`{"tools":[%s],"input":[{"type":"function_call","name":"tool_1","namespace":"mcp__app_0","call_id":"c1","arguments":"{\"large_id\":9223372036854775807}"}]}`, strings.Join(nsList, ","))) + + normalized := normalizeXAIInputNamespaceToolCalls(body) + args := gjson.GetBytes(normalized, "input.0.arguments").String() + if !strings.Contains(args, "9223372036854775807") { + t.Fatalf("large int lost precision in normalized arguments: %s", args) + } +} + +func TestRestoreXAINamespaceToolCalls_FoldModePreservesNonDispatcherArgumentsDone(t *testing.T) { + refs := map[string]xaiNamespaceToolRef{ + "mcp__app_0": {namespace: "mcp__app_0", name: "", isDispatcher: true}, + } + restorer := newXAINamespaceRestorer(refs) + + // Step 1: added event for a non-dispatcher tool (e.g. web_search or regular function) + addedNonDisp := []byte(`{"type":"response.output_item.added","output_index":0,"item":{"id":"non_disp_1","type":"function_call","name":"web_search"}}`) + restoredAdded := restorer.restore(addedNonDisp) + if got := gjson.GetBytes(restoredAdded, "item.name").String(); got != "web_search" { + t.Fatalf("restored name = %q, want web_search", got) + } + + // Step 2: arguments.done for non_disp_1 containing a "name" property + doneNonDisp := []byte(`{"type":"response.function_call_arguments.done","item_id":"non_disp_1","output_index":0,"arguments":"{\"name\":\"golang\",\"query\":\"test\"}"}`) + restoredDone := restorer.restore(doneNonDisp) + if got := gjson.GetBytes(restoredDone, "arguments").String(); got != `{"name":"golang","query":"test"}` { + t.Fatalf("non-dispatcher arguments.done was incorrectly mutated: %s", got) + } +} + +func TestPrepareResponsesRequest_CapsAt200WithInjectXSearch(t *testing.T) { + // 20 namespaces with 10 tools each = 200 tools. + // With InjectXSearch = true, total count = 201 > 200, so it must fold into 20 dispatchers + 1 x_search = 21 tools! + var nsList []string + for i := 0; i < 20; i++ { + var childTools []string + for j := 0; j < 10; j++ { + childTools = append(childTools, fmt.Sprintf(`{"type":"function","name":"tool_%d","parameters":{"type":"object"}}`, j)) + } + nsList = append(nsList, fmt.Sprintf(`{"type":"namespace","name":"mcp__app_%d","tools":[%s]}`, i, strings.Join(childTools, ","))) + } + payload := []byte(fmt.Sprintf(`{"model":"grok-4.6","tools":[%s],"input":[{"role":"user","content":"hi"}]}`, strings.Join(nsList, ","))) + + exec := NewXAIExecutor(&config.Config{ + XAI: config.XAIConfig{InjectXSearch: true}, + }) + prepared, err := exec.prepareResponsesRequestTo(context.Background(), cliproxyexecutor.Request{ + Model: "grok-4.6", + Payload: payload, + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + }, false, sdktranslator.FormatOpenAIResponse) + if err != nil { + t.Fatalf("prepareResponsesRequestTo error = %v", err) + } + + tools := gjson.GetBytes(prepared.body, "tools").Array() + if len(tools) > 200 { + t.Fatalf("prepared tools length = %d, want <= 200", len(tools)) + } + if len(tools) != 21 { + t.Fatalf("prepared tools length = %d, want 21 (20 dispatchers + 1 x_search); body=%s", len(tools), string(prepared.body)) + } + hasXSearch := false + for _, tool := range tools { + if tool.Get("type").String() == "x_search" { + hasXSearch = true + break + } + } + if !hasXSearch { + t.Fatalf("x_search tool was not preserved; body=%s", string(prepared.body)) + } +} + +func TestRestoreXAINamespaceToolCalls_FlattenModePreservesNameInArgumentsDone(t *testing.T) { + // In flatten mode (<= 200), refs isDispatcher = false + refs := map[string]xaiNamespaceToolRef{ + "mcp__github__create_repo": {namespace: "mcp__github", name: "create_repo", isDispatcher: false}, + } + + // Normal tool call arguments that happen to contain a "name" property + event := []byte(`{"type":"response.function_call_arguments.done","item_id":"item_1","output_index":0,"arguments":"{\"name\":\"my-awesome-repo\",\"private\":true}"}`) + restored := restoreXAINamespaceToolCalls(event, refs) + + if got := gjson.GetBytes(restored, "arguments").String(); got != `{"name":"my-awesome-repo","private":true}` { + t.Fatalf("arguments.done was mutated in flatten mode: %s", got) + } +} + +func TestRestoreXAINamespaceToolCalls_OutputItemAddedInDispatcherMode(t *testing.T) { + refs := map[string]xaiNamespaceToolRef{ + "mcp__app_0": {namespace: "mcp__app_0", name: "", isDispatcher: true}, + } + + // At output_item.added time, arguments is empty or in progress + event := []byte(`{"type":"response.output_item.added","output_index":0,"item":{"id":"fc_1","type":"function_call","name":"mcp__app_0","status":"in_progress"}}`) + restored := restoreXAINamespaceToolCalls(event, refs) + + if got := gjson.GetBytes(restored, "item.namespace").String(); got != "mcp__app_0" { + t.Fatalf("item.namespace = %q, want mcp__app_0", got) + } + if got := gjson.GetBytes(restored, "item.name").String(); got != "mcp__app_0" { + t.Fatalf("item.name = %q, want mcp__app_0 at added phase", got) + } +} + +func TestClampXAIToolsLimit_PreservesDispatchersOverRegularTools(t *testing.T) { + var toolList []string + refs := make(map[string]xaiNamespaceToolRef) + // 47 dispatchers + for i := 0; i < 47; i++ { + name := fmt.Sprintf("mcp__app_%d", i) + refs[name] = xaiNamespaceToolRef{namespace: name, name: "", isDispatcher: true} + toolList = append(toolList, fmt.Sprintf(`{"type":"function","name":"%s","parameters":{"type":"object"}}`, name)) + } + // 180 regular tools + for i := 0; i < 180; i++ { + toolList = append(toolList, fmt.Sprintf(`{"type":"function","name":"plain_fn_%d","parameters":{"type":"object"}}`, i)) + } + body := []byte(fmt.Sprintf(`{"tools":[%s]}`, strings.Join(toolList, ","))) + out := clampXAIToolsLimit(body, xaiMaxTools, refs) + + tools := gjson.GetBytes(out, "tools").Array() + if len(tools) != 200 { + t.Fatalf("clamped tools count = %d, want 200", len(tools)) + } + + // Verify all 47 dispatchers are preserved + for i := 0; i < 47; i++ { + name := fmt.Sprintf("mcp__app_%d", i) + found := false + for _, tool := range tools { + if tool.Get("name").String() == name { + found = true + break + } + } + if !found { + t.Fatalf("dispatcher %s was dropped by clamp", name) + } + } +} + func TestPromoteXAIAdditionalTools(t *testing.T) { body := []byte(`{ "tools":[{"type":"function","name":"lookup","parameters":{"type":"object"}}], diff --git a/internal/runtime/executor/xai_websockets_executor.go b/internal/runtime/executor/xai_websockets_executor.go index a8bc15b9..fbe779d1 100644 --- a/internal/runtime/executor/xai_websockets_executor.go +++ b/internal/runtime/executor/xai_websockets_executor.go @@ -699,6 +699,7 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox outputItemsByIndex := make(map[int64][]byte) var outputItemsFallback [][]byte responseFilter := newXAIInternalXSearchResponseFilter(prepared.filterInternalXSearch, prepared.clientDeclaredTools) + namespaceRestorer := newXAINamespaceRestorer(prepared.namespaceTools) recordedTranscript := false for { if ctx != nil && ctx.Err() != nil { @@ -759,7 +760,7 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox } for _, payload := range xaiNormalizeReasoningSummaryDataEvents(payload) { - payload = restoreXAINamespaceToolCalls(payload, prepared.namespaceTools) + payload = namespaceRestorer.restore(payload) payload = responseFilter.apply(payload) if len(payload) == 0 { continue -- 2.51.2