Merge pull request #4251 from huahai0202/fix/xai-namespace-tool-routing

internal/runtime/executor/xai_executor.go

@@ -188,6 +188,7 @@ func (e *XAIExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req
 			continue
 		}
 		eventData := xaiNormalizeReasoningSummaryData(bytes.TrimSpace(line[len(xaiDataTag):]))
+		eventData = restoreXAINamespaceToolCalls(eventData, prepared.namespaceTools)
 		switch gjson.GetBytes(eventData, "type").String() {
 		case "response.output_item.done":
 			xaiCollectOutputItemDone(eventData, outputItemsByIndex, &outputItemsFallback)
@@ -672,6 +673,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)
 					normalizedEventName := gjson.GetBytes(eventData, "type").String()
 					switch normalizedEventName {
 					case "response.output_item.done":
@@ -817,10 +819,16 @@ type xaiPreparedRequest struct {
 	to              sdktranslator.Format
 	originalPayload []byte
 	body            []byte
+	namespaceTools  map[string]xaiNamespaceToolRef
 	sessionID       string
 	replayScope     xaiReasoningReplayScope
 }
 
+type xaiNamespaceToolRef struct {
+	namespace string
+	name      string
+}
+
 func (e *XAIExecutor) prepareResponsesRequest(ctx context.Context, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, stream bool) (*xaiPreparedRequest, error) {
 	return e.prepareResponsesRequestTo(ctx, req, opts, stream, sdktranslator.FormatCodex)
 }
@@ -852,13 +860,16 @@ func (e *XAIExecutor) prepareResponsesRequestTo(ctx context.Context, req cliprox
 	body, _ = sjson.DeleteBytes(body, "prompt_cache_retention")
 	body, _ = sjson.DeleteBytes(body, "safety_identifier")
 	body, _ = sjson.DeleteBytes(body, "stream_options")
+	namespaceTools := collectXAINamespaceToolRefs(body)
 	body = normalizeXAITools(body)
+	body = normalizeXAINamespaceToolChoice(body)
 	body = normalizeXAIToolChoiceForTools(body)
 	var replayScope xaiReasoningReplayScope
 	body, replayScope, err = applyXAIReasoningReplayCacheRequired(ctx, from, req, opts, body)
 	if err != nil {
 		return nil, err
 	}
+	body = normalizeXAIInputNamespaceToolCalls(body)
 	body = normalizeXAIInputReasoningItems(body)
 	body = sanitizeXAIInputEncryptedContent(body)
 	body = normalizeCodexInstructions(body)
@@ -879,6 +890,7 @@ func (e *XAIExecutor) prepareResponsesRequestTo(ctx context.Context, req cliprox
 		to:              to,
 		originalPayload: originalPayload,
 		body:            body,
+		namespaceTools:  namespaceTools,
 		sessionID:       sessionID,
 		replayScope:     replayScope,
 	}, nil
@@ -1158,11 +1170,48 @@ func sanitizeXAIResponsesBody(body []byte, model string) []byte {
 }
 
 func normalizeXAITools(body []byte) []byte {
-	tools := gjson.GetBytes(body, "tools")
-	if !tools.Exists() || !tools.IsArray() {
+	if !gjson.ValidBytes(body) {
 		return body
 	}
+	original := body
+	normalizeAtPath := func(path string) bool {
+		tools := gjson.GetBytes(body, path)
+		if !tools.Exists() || !tools.IsArray() {
+			return true
+		}
+		filtered, changed, ok := normalizeXAIToolArray(tools)
+		if !ok {
+			return false
+		}
+		if !changed {
+			return true
+		}
+		updated, errSet := sjson.SetRawBytes(body, path, filtered)
+		if errSet != nil {
+			return false
+		}
+		body = updated
+		return true
+	}
 
+	if !normalizeAtPath("tools") {
+		return original
+	}
+	input := gjson.GetBytes(body, "input")
+	if input.Exists() && input.IsArray() {
+		for index, item := range input.Array() {
+			if item.Get("type").String() != "additional_tools" {
+				continue
+			}
+			if !normalizeAtPath(fmt.Sprintf("input.%d.tools", index)) {
+				return original
+			}
+		}
+	}
+	return body
+}
+
+func normalizeXAIToolArray(tools gjson.Result) ([]byte, bool, bool) {
 	changed := false
 	filtered := []byte(`[]`)
 	for _, tool := range tools.Array() {
@@ -1174,7 +1223,7 @@ func normalizeXAITools(body []byte) []byte {
 				for _, nestedTool := range namespaceTools.Array() {
 					nestedRaw, nestedChanged, ok := normalizeXAITool(nestedTool, namespaceName)
 					if !ok {
-						return body
+						return nil, false, false
 					}
 					changed = changed || nestedChanged
 					if len(nestedRaw) == 0 {
@@ -1182,7 +1231,7 @@ func normalizeXAITools(body []byte) []byte {
 					}
 					updated, errSet := sjson.SetRawBytes(filtered, "-1", nestedRaw)
 					if errSet != nil {
-						return body
+						return nil, false, false
 					}
 					filtered = updated
 				}
@@ -1191,7 +1240,7 @@ func normalizeXAITools(body []byte) []byte {
 		}
 		raw, toolChanged, ok := normalizeXAITool(tool, "")
 		if !ok {
-			return body
+			return nil, false, false
 		}
 		changed = changed || toolChanged
 		if len(raw) == 0 {
@@ -1199,18 +1248,11 @@ func normalizeXAITools(body []byte) []byte {
 		}
 		updated, errSet := sjson.SetRawBytes(filtered, "-1", raw)
 		if errSet != nil {
-			return body
+			return nil, false, false
 		}
 		filtered = updated
 	}
-	if !changed {
-		return body
-	}
-	updated, errSet := sjson.SetRawBytes(body, "tools", filtered)
-	if errSet != nil {
-		return body
-	}
-	return updated
+	return filtered, changed, true
 }
 
 // normalizeXAIToolChoiceForTools drops tool_choice and parallel_tool_calls
@@ -1220,6 +1262,18 @@ func normalizeXAITools(body []byte) []byte {
 func normalizeXAIToolChoiceForTools(body []byte) []byte {
 	tools := gjson.GetBytes(body, "tools")
 	hasTools := tools.Exists() && tools.IsArray() && len(tools.Array()) > 0
+	if !hasTools {
+		input := gjson.GetBytes(body, "input")
+		if input.Exists() && input.IsArray() {
+			for _, item := range input.Array() {
+				additionalTools := item.Get("tools")
+				if item.Get("type").String() == "additional_tools" && additionalTools.IsArray() && len(additionalTools.Array()) > 0 {
+					hasTools = true
+					break
+				}
+			}
+		}
+	}
 	if hasTools {
 		return body
 	}
@@ -1235,6 +1289,34 @@ func normalizeXAIToolChoiceForTools(body []byte) []byte {
 	return body
 }
 
+// normalizeXAINamespaceToolChoice qualifies a forced namespaced function choice
+// using the same name sent in the flattened tools list. xAI does not accept the
+// Responses namespace field on tool_choice.
+func normalizeXAINamespaceToolChoice(body []byte) []byte {
+	if !gjson.ValidBytes(body) {
+		return body
+	}
+	toolChoice := gjson.GetBytes(body, "tool_choice")
+	if !toolChoice.IsObject() || toolChoice.Get("type").String() != xaiFunctionToolType {
+		return body
+	}
+	namespaceName := strings.TrimSpace(toolChoice.Get("namespace").String())
+	toolName := strings.TrimSpace(toolChoice.Get("name").String())
+	qualifiedName := qualifyXAINamespaceToolName(namespaceName, toolName)
+	if namespaceName == "" || qualifiedName == "" {
+		return body
+	}
+	updated, errSet := sjson.SetBytes(body, "tool_choice.name", qualifiedName)
+	if errSet != nil {
+		return body
+	}
+	updated, errDelete := sjson.DeleteBytes(updated, "tool_choice.namespace")
+	if errDelete != nil {
+		return body
+	}
+	return updated
+}
+
 func normalizeXAITool(tool gjson.Result, namespaceName string) ([]byte, bool, bool) {
 	toolType := tool.Get("type").String()
 	changed := false
@@ -1289,9 +1371,140 @@ func normalizeXAITool(tool gjson.Result, namespaceName string) ([]byte, bool, bo
 		changed = true
 		log.Debugf("xai: simplified parameters for tool %s.%s to avoid upstream hang", namespaceName, tool.Get("name").String())
 	}
+	if toolType == xaiFunctionToolType && strings.TrimSpace(namespaceName) != "" {
+		qualifiedName := qualifyXAINamespaceToolName(namespaceName, tool.Get("name").String())
+		if qualifiedName == "" {
+			return nil, false, false
+		}
+		updatedTool, errSet := sjson.SetBytes(raw, "name", qualifiedName)
+		if errSet != nil {
+			return nil, false, false
+		}
+		raw = updatedTool
+		changed = true
+	}
 	return raw, changed, true
 }
 
+func qualifyXAINamespaceToolName(namespaceName, toolName string) string {
+	namespaceName = strings.TrimSpace(namespaceName)
+	toolName = strings.TrimSpace(toolName)
+	if namespaceName == "" || toolName == "" || strings.HasPrefix(toolName, "mcp__") {
+		return toolName
+	}
+	prefix := namespaceName
+	if !strings.HasSuffix(prefix, "__") {
+		prefix += "__"
+	}
+	if strings.HasPrefix(toolName, prefix) {
+		return toolName
+	}
+	return prefix + toolName
+}
+
+func collectXAINamespaceToolRefs(body []byte) map[string]xaiNamespaceToolRef {
+	refs := make(map[string]xaiNamespaceToolRef)
+	collect := func(tools gjson.Result) {
+		if !tools.Exists() || !tools.IsArray() {
+			return
+		}
+		for _, tool := range tools.Array() {
+			if tool.Get("type").String() != xaiNamespaceToolType {
+				continue
+			}
+			namespaceName := strings.TrimSpace(tool.Get("name").String())
+			if namespaceName == "" {
+				continue
+			}
+			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}
+			}
+		}
+	}
+	collect(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" {
+				collect(item.Get("tools"))
+			}
+		}
+	}
+	return refs
+}
+
+func normalizeXAIInputNamespaceToolCalls(body []byte) []byte {
+	if !gjson.ValidBytes(body) {
+		return body
+	}
+	input := gjson.GetBytes(body, "input")
+	if !input.Exists() || !input.IsArray() {
+		return body
+	}
+	for index, item := range input.Array() {
+		if item.Get("type").String() != "function_call" {
+			continue
+		}
+		namespaceName := strings.TrimSpace(item.Get("namespace").String())
+		toolName := strings.TrimSpace(item.Get("name").String())
+		qualifiedName := qualifyXAINamespaceToolName(namespaceName, toolName)
+		if namespaceName == "" || qualifiedName == "" {
+			continue
+		}
+		namePath := fmt.Sprintf("input.%d.name", index)
+		namespacePath := fmt.Sprintf("input.%d.namespace", index)
+		updated, errSet := sjson.SetBytes(body, namePath, qualifiedName)
+		if errSet != nil {
+			continue
+		}
+		updated, errDelete := sjson.DeleteBytes(updated, namespacePath)
+		if errDelete != nil {
+			continue
+		}
+		body = updated
+	}
+	return body
+}
+
+func restoreXAINamespaceToolCalls(data []byte, refs map[string]xaiNamespaceToolRef) []byte {
+	if len(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)
+		}
+	}
+	return data
+}
+
+func restoreXAINamespaceToolCallAtPath(data []byte, path string, refs map[string]xaiNamespaceToolRef) []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]
+	if !ok {
+		return data
+	}
+	updated, errSet := sjson.SetBytes(data, path+".name", ref.name)
+	if errSet != nil {
+		return data
+	}
+	updated, errSet = sjson.SetBytes(updated, path+".namespace", ref.namespace)
+	if errSet != nil {
+		return data
+	}
+	return updated
+}
+
 // xaiFunctionParametersNeedSimplification reports whether a function tool is
 // the Codex Desktop automation tool known to hang xAI Responses streaming.
 func xaiFunctionParametersNeedSimplification(tool gjson.Result, namespaceName string) bool {

internal/runtime/executor/xai_executor_test.go

@@ -148,9 +148,9 @@ func TestXAIExecutorExecuteShapesResponsesRequest(t *testing.T) {
 			t.Fatalf("tools.%d.name = apply_patch, want removed; body=%s", i, string(gotBody))
 		}
 		switch tool.Get("name").String() {
-		case "automation_update":
+		case "codex_app__automation_update":
 			foundAutomationUpdate = true
-		case "namespace_custom":
+		case "codex_app__namespace_custom":
 			foundNamespaceCustom = true
 		}
 		if toolType == "web_search" {
@@ -180,6 +180,63 @@ func TestXAIExecutorExecuteShapesResponsesRequest(t *testing.T) {
 	}
 }
 
+func TestXAIExecutorExecuteRestoresAdditionalToolsNamespaceCalls(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__exa__web_search_exa\",\"call_id\":\"call_1\",\"arguments\":\"{}\"}}\n\n"))
+		_, _ = w.Write([]byte("data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_1\",\"object\":\"response\",\"status\":\"completed\",\"model\":\"grok-4.3\",\"output\":[],\"usage\":{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2}}}\n\n"))
+	}))
+	defer server.Close()
+
+	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"},
+	}
+	resp, err := exec.Execute(context.Background(), auth, cliproxyexecutor.Request{
+		Model: "grok-4.3",
+		Payload: []byte(`{
+			"model":"grok-4.3",
+			"input":[
+				{"type":"additional_tools","role":"developer","tools":[{"type":"namespace","name":"mcp__exa","tools":[{"type":"function","name":"web_search_exa","parameters":{"type":"object"}}]}]},
+				{"role":"user","content":"use Exa"}
+			]
+		}`),
+	}, cliproxyexecutor.Options{
+		SourceFormat:   sdktranslator.FormatOpenAIResponse,
+		ResponseFormat: sdktranslator.FormatOpenAIResponse,
+		Stream:         false,
+	})
+	if err != nil {
+		t.Fatalf("Execute() error = %v", err)
+	}
+
+	tool := gjson.GetBytes(gotBody, "input.0.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)
+	}
+	if got := tool.Get("type").String(); got != "function" {
+		t.Fatalf("upstream additional 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)
+	}
+	output := gjson.GetBytes(resp.Payload, "output.0")
+	if got := output.Get("name").String(); got != "web_search_exa" {
+		t.Fatalf("response output name = %q, want child name; payload=%s", got, resp.Payload)
+	}
+	if got := output.Get("namespace").String(); got != "mcp__exa" {
+		t.Fatalf("response output namespace = %q, want mcp__exa; payload=%s", got, resp.Payload)
+	}
+}
+
 func TestXAIExecutorComposerSessionIsolation(t *testing.T) {
 	exec := NewXAIExecutor(&config.Config{})
 	auth := &cliproxyauth.Auth{
@@ -786,9 +843,9 @@ func TestXAIExecutorExecuteStreamFiltersToolSearchTool(t *testing.T) {
 			t.Fatalf("tools.%d.name = apply_patch, want removed; body=%s", i, string(gotBody))
 		}
 		switch tool.Get("name").String() {
-		case "automation_update":
+		case "codex_app__automation_update":
 			foundAutomationUpdate = true
-		case "namespace_custom":
+		case "codex_app__namespace_custom":
 			foundNamespaceCustom = true
 		}
 		if toolType == "web_search" {
@@ -1203,7 +1260,7 @@ func TestNormalizeXAITools_SimplifiesCodexAppAutomationUpdateSchema(t *testing.T
 	foundExec := false
 	for _, tool := range tools.Array() {
 		switch tool.Get("name").String() {
-		case "automation_update":
+		case "codex_app__automation_update":
 			foundAuto = true
 			paramsRaw := tool.Get("parameters").Raw
 			if strings.Contains(paramsRaw, `"oneOf"`) || strings.Contains(paramsRaw, `"$defs"`) {
@@ -1233,6 +1290,108 @@ func TestNormalizeXAITools_SimplifiesCodexAppAutomationUpdateSchema(t *testing.T
 	}
 }
 
+func TestNormalizeXAITools_QualifiesSameNamedNamespaceTools(t *testing.T) {
+	body := []byte(`{
+		"tools":[
+			{"type":"namespace","name":"mcp__exa","tools":[{"type":"function","name":"search","parameters":{"type":"object"}}]},
+			{"type":"namespace","name":"mcp__docs","tools":[{"type":"function","name":"search","parameters":{"type":"object"}}]}
+		]
+	}`)
+	out := normalizeXAITools(body)
+
+	tools := gjson.GetBytes(out, "tools").Array()
+	if len(tools) != 2 {
+		t.Fatalf("tools length = %d, want 2; body=%s", len(tools), string(out))
+	}
+	if got := tools[0].Get("name").String(); got != "mcp__exa__search" {
+		t.Fatalf("tools.0.name = %q, want mcp__exa__search; body=%s", got, string(out))
+	}
+	if got := tools[1].Get("name").String(); got != "mcp__docs__search" {
+		t.Fatalf("tools.1.name = %q, want mcp__docs__search; body=%s", got, string(out))
+	}
+}
+
+func TestNormalizeXAITools_AdditionalToolsNamespace(t *testing.T) {
+	body := []byte(`{
+		"input":[
+			{"type":"additional_tools","role":"developer","tools":[{"type":"namespace","name":"mcp__exa","tools":[{"type":"function","name":"search","parameters":{"type":"object"}}]}]},
+			{"role":"user","content":"hello"}
+		]
+	}`)
+	out := 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))
+	}
+	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))
+	}
+	if got := tools[0].Get("type").String(); got != "function" {
+		t.Fatalf("additional tool type = %q, want function; body=%s", got, string(out))
+	}
+}
+
+func TestNormalizeXAINamespaceToolChoice(t *testing.T) {
+	body := []byte(`{
+		"tools":[{"type":"namespace","name":"mcp__exa","tools":[{"type":"function","name":"search","parameters":{"type":"object"}}]}],
+		"tool_choice":{"type":"function","name":"search","namespace":"mcp__exa"}
+	}`)
+	out := normalizeXAITools(body)
+	out = normalizeXAINamespaceToolChoice(out)
+
+	if got := gjson.GetBytes(out, "tools.0.name").String(); got != "mcp__exa__search" {
+		t.Fatalf("tools.0.name = %q, want mcp__exa__search; body=%s", got, string(out))
+	}
+	if got := gjson.GetBytes(out, "tool_choice.name").String(); got != "mcp__exa__search" {
+		t.Fatalf("tool_choice.name = %q, want mcp__exa__search; body=%s", got, string(out))
+	}
+	if gjson.GetBytes(out, "tool_choice.namespace").Exists() {
+		t.Fatalf("tool_choice.namespace should be removed for xAI upstream: %s", string(out))
+	}
+}
+
+func TestNormalizeXAINamespaceToolChoice_PreservesOtherChoices(t *testing.T) {
+	tests := []struct {
+		name string
+		body []byte
+	}{
+		{name: "automatic choice", body: []byte(`{"tool_choice":"auto"}`)},
+		{name: "top-level function", body: []byte(`{"tool_choice":{"type":"function","name":"search"}}`)},
+		{name: "non-function choice", body: []byte(`{"tool_choice":{"type":"web_search","name":"search","namespace":"mcp__exa"}}`)},
+		{name: "malformed payload", body: []byte(`{"tool_choice":{"type":"function"`)},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := normalizeXAINamespaceToolChoice(tt.body); !bytes.Equal(got, tt.body) {
+				t.Fatalf("payload changed: got=%q want=%q", got, tt.body)
+			}
+		})
+	}
+}
+
+func TestQualifyXAINamespaceToolNamePreservesQualifiedNames(t *testing.T) {
+	tests := []struct {
+		name      string
+		namespace string
+		tool      string
+		want      string
+	}{
+		{name: "plain child", namespace: "mcp__exa", tool: "search", want: "mcp__exa__search"},
+		{name: "prequalified MCP child", namespace: "mcp__exa", tool: "mcp__exa__search", want: "mcp__exa__search"},
+		{name: "prequalified generic child", namespace: "collaboration", tool: "collaboration__send_message", want: "collaboration__send_message"},
+		{name: "namespace with separator", namespace: "collaboration__", tool: "send_message", want: "collaboration__send_message"},
+		{name: "partial prefix is not qualified", namespace: "exa", tool: "example_tool", want: "exa__example_tool"},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := qualifyXAINamespaceToolName(tt.namespace, tt.tool); got != tt.want {
+				t.Fatalf("qualifyXAINamespaceToolName(%q, %q) = %q, want %q", tt.namespace, tt.tool, got, tt.want)
+			}
+		})
+	}
+}
+
 func TestNormalizeXAITools_PreservesUnrelatedSchemas(t *testing.T) {
 	largeParams := `{"oneOf":[{"type":"object","properties":{"mode":{"type":"string"}}}],"$defs":{"a":{"type":"string"}},"x":"` + strings.Repeat("y", 1600) + `"}`
 	tests := []struct {
@@ -1302,6 +1461,55 @@ func TestXAIFunctionParametersNeedSimplification(t *testing.T) {
 	}
 }
 
+func TestNormalizeXAIInputNamespaceToolCalls(t *testing.T) {
+	body := []byte(`{"input":[{"type":"function_call","name":"web_search_exa","namespace":"mcp__exa","call_id":"call_1","arguments":"{}"},{"type":"function_call","name":"plain_tool","call_id":"call_2","arguments":"{}"}]}`)
+	out := normalizeXAIInputNamespaceToolCalls(body)
+
+	if got := gjson.GetBytes(out, "input.0.name").String(); got != "mcp__exa__web_search_exa" {
+		t.Fatalf("input.0.name = %q, want qualified namespace name; body=%s", got, string(out))
+	}
+	if gjson.GetBytes(out, "input.0.namespace").Exists() {
+		t.Fatalf("input.0.namespace should be removed for xAI upstream: %s", string(out))
+	}
+	if got := gjson.GetBytes(out, "input.1.name").String(); got != "plain_tool" {
+		t.Fatalf("plain function call name changed to %q", got)
+	}
+}
+
+func TestRestoreXAINamespaceToolCalls(t *testing.T) {
+	request := []byte(`{"tools":[{"type":"namespace","name":"mcp__exa","tools":[{"type":"function","name":"web_search_exa","parameters":{"type":"object"}}]}]}`)
+	refs := collectXAINamespaceToolRefs(request)
+
+	event := []byte(`{"type":"response.output_item.done","item":{"type":"function_call","name":"mcp__exa__web_search_exa","call_id":"call_1","arguments":"{}"}}`)
+	restoredEvent := restoreXAINamespaceToolCalls(event, refs)
+	if got := gjson.GetBytes(restoredEvent, "item.name").String(); got != "web_search_exa" {
+		t.Fatalf("item.name = %q, want child name; event=%s", got, string(restoredEvent))
+	}
+	if got := gjson.GetBytes(restoredEvent, "item.namespace").String(); got != "mcp__exa" {
+		t.Fatalf("item.namespace = %q, want mcp__exa; event=%s", got, string(restoredEvent))
+	}
+
+	completed := []byte(`{"type":"response.completed","response":{"output":[{"type":"function_call","name":"mcp__exa__web_search_exa","call_id":"call_1","arguments":"{}"}]}}`)
+	restoredCompleted := restoreXAINamespaceToolCalls(completed, refs)
+	if got := gjson.GetBytes(restoredCompleted, "response.output.0.name").String(); got != "web_search_exa" {
+		t.Fatalf("response.output.0.name = %q, want child name; event=%s", got, string(restoredCompleted))
+	}
+	if got := gjson.GetBytes(restoredCompleted, "response.output.0.namespace").String(); got != "mcp__exa" {
+		t.Fatalf("response.output.0.namespace = %q, want mcp__exa; event=%s", got, string(restoredCompleted))
+	}
+}
+
+func TestRestoreXAINamespaceToolCallsPreservesMalformedPayload(t *testing.T) {
+	data := []byte(`{"item":{"type":"function_call","name":"mcp__exa__web_search_exa"`)
+	refs := map[string]xaiNamespaceToolRef{
+		"mcp__exa__web_search_exa": {namespace: "mcp__exa", name: "web_search_exa"},
+	}
+
+	if got := restoreXAINamespaceToolCalls(data, refs); !bytes.Equal(got, data) {
+		t.Fatalf("malformed payload changed: got=%q want=%q", got, data)
+	}
+}
+
 func TestNormalizeXAIToolChoiceForTools_DropsWhenToolsEmpty(t *testing.T) {
 	body := []byte(`{"model":"grok-4","tools":[],"tool_choice":"auto","parallel_tool_calls":true,"input":"hi"}`)
 	out := normalizeXAIToolChoiceForTools(body)
@@ -1347,6 +1555,18 @@ func TestNormalizeXAIToolChoiceForTools_KeepsWhenToolsPresent(t *testing.T) {
 	}
 }
 
+func TestNormalizeXAIToolChoiceForTools_KeepsWhenAdditionalToolsPresent(t *testing.T) {
+	body := []byte(`{"model":"grok-4","input":[{"type":"additional_tools","tools":[{"type":"function","name":"Bash"}]}],"tool_choice":"auto","parallel_tool_calls":true}`)
+	out := normalizeXAIToolChoiceForTools(body)
+
+	if got := gjson.GetBytes(out, "tool_choice").String(); got != "auto" {
+		t.Fatalf("tool_choice = %q, want auto: %s", got, string(out))
+	}
+	if !gjson.GetBytes(out, "parallel_tool_calls").Bool() {
+		t.Fatalf("parallel_tool_calls should be kept: %s", string(out))
+	}
+}
+
 func TestNormalizeXAIToolChoiceForTools_NoOpWhenBothAbsent(t *testing.T) {
 	body := []byte(`{"model":"grok-4","input":"hi"}`)
 	out := normalizeXAIToolChoiceForTools(body)

internal/runtime/executor/xai_websockets_executor.go

@@ -637,6 +637,7 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox
 			}
 
 			for _, payload := range xaiNormalizeReasoningSummaryDataEvents(payload) {
+				payload = restoreXAINamespaceToolCalls(payload, prepared.namespaceTools)
 				eventType := gjson.GetBytes(payload, "type").String()
 				isTerminalEvent := eventType == "response.completed" || eventType == "response.done" || eventType == "error"
 				warmupCompletedPayload := []byte(nil)

internal/runtime/executor/xai_websockets_executor_test.go

@@ -121,6 +121,112 @@ func TestXAIWebsocketsExecuteStreamSendsResponseCreateWithPreviousResponseID(t *
 	}
 }
 
+func TestXAIWebsocketsExecuteStreamRestoresNamespaceToolCalls(t *testing.T) {
+	upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }}
+	capturedPayload := make(chan []byte, 1)
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		conn, err := upgrader.Upgrade(w, r, nil)
+		if err != nil {
+			t.Errorf("upgrade websocket: %v", err)
+			return
+		}
+		defer func() { _ = conn.Close() }()
+
+		_, payload, errRead := conn.ReadMessage()
+		if errRead != nil {
+			t.Errorf("read upstream websocket message: %v", errRead)
+			return
+		}
+		capturedPayload <- bytes.Clone(payload)
+
+		events := [][]byte{
+			[]byte(`{"type":"response.output_item.done","output_index":0,"item":{"type":"function_call","name":"mcp__exa__web_search_exa","call_id":"call_1","arguments":"{}"}}`),
+			[]byte(`{"type":"response.completed","response":{"id":"resp_1","output":[],"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}}`),
+		}
+		for _, event := range events {
+			if errWrite := conn.WriteMessage(websocket.TextMessage, event); errWrite != nil {
+				t.Errorf("write websocket event: %v", errWrite)
+				return
+			}
+		}
+	}))
+	defer server.Close()
+
+	exec := NewXAIWebsocketsExecutor(&config.Config{})
+	auth := &cliproxyauth.Auth{
+		Provider: "xai",
+		Attributes: map[string]string{
+			"base_url":   server.URL,
+			"websockets": "true",
+		},
+		Metadata: map[string]any{"access_token": "xai-token"},
+	}
+	req := cliproxyexecutor.Request{
+		Model: "grok-4.3",
+		Payload: []byte(`{
+			"model":"grok-4.3",
+			"input":[
+				{"type":"additional_tools","role":"developer","tools":[{
+					"type":"namespace",
+					"name":"mcp__exa",
+					"tools":[{"type":"function","name":"web_search_exa","parameters":{"type":"object"}}]
+				}]},
+				{"role":"user","content":"use Exa"}
+			]
+		}`),
+	}
+	opts := cliproxyexecutor.Options{
+		SourceFormat:   sdktranslator.FormatOpenAIResponse,
+		ResponseFormat: sdktranslator.FormatOpenAIResponse,
+		Stream:         true,
+	}
+	ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background())
+
+	result, err := exec.ExecuteStream(ctx, auth, req, opts)
+	if err != nil {
+		t.Fatalf("ExecuteStream() error = %v", err)
+	}
+
+	select {
+	case payload := <-capturedPayload:
+		tool := gjson.GetBytes(payload, "input.0.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)
+		}
+		if tool.Get("tools").Exists() {
+			t.Fatalf("upstream tool should not contain namespace children: %s", payload)
+		}
+	case <-time.After(5 * time.Second):
+		t.Fatal("timed out waiting for upstream websocket payload")
+	}
+
+	var outputItemDone, completed gjson.Result
+	for chunk := range result.Chunks {
+		if chunk.Err != nil {
+			t.Fatalf("stream chunk error = %v", chunk.Err)
+		}
+		payload := gjson.ParseBytes(bytes.TrimSpace(chunk.Payload))
+		switch payload.Get("type").String() {
+		case "response.output_item.done":
+			outputItemDone = payload
+		case "response.completed":
+			completed = payload
+		}
+	}
+
+	for label, item := range map[string]gjson.Result{
+		"output_item.done": outputItemDone.Get("item"),
+		"completed":        completed.Get("response.output.0"),
+	} {
+		if got := item.Get("name").String(); got != "web_search_exa" {
+			t.Fatalf("%s name = %q, want child name; item=%s", label, got, item.Raw)
+		}
+		if got := item.Get("namespace").String(); got != "mcp__exa" {
+			t.Fatalf("%s namespace = %q, want mcp__exa; item=%s", label, got, item.Raw)
+		}
+	}
+}
+
 func TestXAIWebsocketsExecuteStreamNormalizesReasoningTextEvents(t *testing.T) {
 	upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }}
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {