From ca67caf0872eef65f8eaedf2e127cf214a49621d Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Fri, 10 Jul 2026 13:11:55 +0800 Subject: [PATCH] **fix(xai): improve namespace-specific tool parameter handling and simplify Codex automation schema** - Refined tool normalization logic to pass namespace context (`namespaceName`) to `normalizeXAITool`. - Limited schema simplification to `codex_app.automation_update` to avoid unintended modifications to unrelated tools or namespaces. - Updated `xaiFunctionParametersNeedSimplification` to match tools by exact namespace and name. - Enhanced tests to validate correct schema simplification for `codex_app.automation_update` while preserving unrelated tool definitions. --- internal/runtime/executor/xai_executor.go | 67 ++++++------- .../runtime/executor/xai_executor_test.go | 73 ++++++++++++-- .../executor/xai_websockets_executor.go | 15 ++- .../executor/xai_websockets_executor_test.go | 96 +++++++++++++++++++ 4 files changed, 207 insertions(+), 44 deletions(-) diff --git a/internal/runtime/executor/xai_executor.go b/internal/runtime/executor/xai_executor.go index a21d852b4..2118aae1c 100644 --- a/internal/runtime/executor/xai_executor.go +++ b/internal/runtime/executor/xai_executor.go @@ -36,17 +36,18 @@ var ( ) const ( - xaiImageHandlerType = "openai-image" - xaiVideoHandlerType = "openai-video" - xaiCustomToolType = "custom" - xaiFunctionToolType = "function" - xaiImageGenerationToolType = "image_generation" - xaiNamespaceToolType = "namespace" - xaiToolSearchType = "tool_search" - xaiWebSearchToolType = "web_search" + xaiImageHandlerType = "openai-image" + xaiVideoHandlerType = "openai-video" + xaiCustomToolType = "custom" + xaiFunctionToolType = "function" + xaiImageGenerationToolType = "image_generation" + xaiNamespaceToolType = "namespace" + xaiToolSearchType = "tool_search" + xaiWebSearchToolType = "web_search" // 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". + xaiCodexAppNamespaceName = "codex_app" xaiAutomationUpdateToolName = "automation_update" // Permissive placeholder schema: keeps the tool callable without the hang. xaiSafeFunctionParameters = `{"type":"object","properties":{},"additionalProperties":true}` @@ -1049,9 +1050,10 @@ func normalizeXAITools(body []byte) []byte { toolType := tool.Get("type").String() if toolType == xaiNamespaceToolType { changed = true + namespaceName := tool.Get("name").String() if namespaceTools := tool.Get("tools"); namespaceTools.IsArray() { for _, nestedTool := range namespaceTools.Array() { - nestedRaw, nestedChanged, ok := normalizeXAITool(nestedTool) + nestedRaw, nestedChanged, ok := normalizeXAITool(nestedTool, namespaceName) if !ok { return body } @@ -1068,7 +1070,7 @@ func normalizeXAITools(body []byte) []byte { } continue } - raw, toolChanged, ok := normalizeXAITool(tool) + raw, toolChanged, ok := normalizeXAITool(tool, "") if !ok { return body } @@ -1114,7 +1116,7 @@ func normalizeXAIToolChoiceForTools(body []byte) []byte { return body } -func normalizeXAITool(tool gjson.Result) ([]byte, bool, bool) { +func normalizeXAITool(tool gjson.Result, namespaceName string) ([]byte, bool, bool) { toolType := tool.Get("type").String() changed := false if toolType == xaiToolSearchType || toolType == xaiImageGenerationToolType { @@ -1149,41 +1151,34 @@ func normalizeXAITool(tool gjson.Result) ([]byte, bool, bool) { raw = updatedTool changed = true } - // Codex Desktop's automation_update schema (and similar large oneOf+$ref - // function schemas) hang xAI free/build streaming. Simplify parameters so - // the request still carries the tool name but does not stall the stream. - if toolType == xaiFunctionToolType && xaiFunctionParametersNeedSimplification(tool) { + // Codex Desktop's codex_app.automation_update schema hangs xAI free/build + // streaming. Limit the workaround to that exact namespaced tool so unrelated + // tools keep their parameter contracts. + if toolType == xaiFunctionToolType && xaiFunctionParametersNeedSimplification(tool, namespaceName) { updatedTool, errSet := sjson.SetRawBytes(raw, "parameters", []byte(xaiSafeFunctionParameters)) if errSet != nil { return nil, false, false } raw = updatedTool + if strict := tool.Get("strict"); strict.Exists() && strict.Bool() { + updatedTool, errSet = sjson.SetBytes(raw, "strict", false) + if errSet != nil { + return nil, false, false + } + raw = updatedTool + } changed = true - log.Debugf("xai: simplified parameters for tool %s to avoid upstream hang", tool.Get("name").String()) + log.Debugf("xai: simplified parameters for tool %s.%s to avoid upstream hang", namespaceName, tool.Get("name").String()) } return raw, changed, true } -// xaiFunctionParametersNeedSimplification reports whether a function tool's -// JSON Schema is known/likely to hang xAI Responses streaming. -func xaiFunctionParametersNeedSimplification(tool gjson.Result) bool { - name := strings.TrimSpace(tool.Get("name").String()) - if strings.EqualFold(name, xaiAutomationUpdateToolName) { - return true - } - params := tool.Get("parameters") - if !params.Exists() { - return false - } - raw := params.Raw - // Heuristic: large schemas combining oneOf with $ref/$defs hang the free - // Grok Responses path (no SSE events until client cancel). - if len(raw) < 1500 { - return false - } - hasOneOf := strings.Contains(raw, `"oneOf"`) - hasRef := strings.Contains(raw, `"$ref"`) || strings.Contains(raw, `"$defs"`) - return hasOneOf && hasRef +// 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 { + return strings.EqualFold(strings.TrimSpace(tool.Get("type").String()), xaiFunctionToolType) && + strings.EqualFold(strings.TrimSpace(namespaceName), xaiCodexAppNamespaceName) && + strings.EqualFold(strings.TrimSpace(tool.Get("name").String()), xaiAutomationUpdateToolName) } func sanitizeXAIInputEncryptedContent(body []byte) []byte { diff --git a/internal/runtime/executor/xai_executor_test.go b/internal/runtime/executor/xai_executor_test.go index 9ef2b18c4..77d3d6506 100644 --- a/internal/runtime/executor/xai_executor_test.go +++ b/internal/runtime/executor/xai_executor_test.go @@ -1060,10 +1060,10 @@ func TestXAIExecutorExecuteVideosUsesNativeEndpointFromRequestPath(t *testing.T) } } -func TestNormalizeXAITools_SimplifiesAutomationUpdateSchema(t *testing.T) { +func TestNormalizeXAITools_SimplifiesCodexAppAutomationUpdateSchema(t *testing.T) { // Large oneOf+$ref schema mimicking Codex Desktop codex_app.automation_update. params := `{"oneOf":[{"type":"object","properties":{"mode":{"type":"string"}}}],"$defs":{"a":{"type":"string"}},"x":"` + strings.Repeat("y", 1600) + `"}` - body := []byte(`{"model":"grok-4.5","tools":[{"type":"namespace","name":"codex_app","tools":[{"type":"function","name":"automation_update","description":"sched","strict":false,"parameters":` + params + `}]},{"type":"function","name":"exec_command","parameters":{"type":"object","properties":{"cmd":{"type":"string"}}}}]}`) + body := []byte(`{"model":"grok-4.5","tools":[{"type":"namespace","name":"codex_app","tools":[{"type":"function","name":"automation_update","description":"sched","strict":true,"parameters":` + params + `}]},{"type":"function","name":"exec_command","parameters":{"type":"object","properties":{"cmd":{"type":"string"}}}}]}`) out := normalizeXAITools(body) tools := gjson.GetBytes(out, "tools") @@ -1086,6 +1086,9 @@ func TestNormalizeXAITools_SimplifiesAutomationUpdateSchema(t *testing.T) { if tool.Get("parameters.additionalProperties").Type != gjson.True { t.Fatalf("automation_update parameters should allow additionalProperties: %s", paramsRaw) } + if tool.Get("strict").Type != gjson.False { + t.Fatalf("automation_update strict = %s, want false", tool.Get("strict").Raw) + } case "exec_command": foundExec = true if got := tool.Get("parameters.properties.cmd.type").String(); got != "string" { @@ -1101,14 +1104,72 @@ func TestNormalizeXAITools_SimplifiesAutomationUpdateSchema(t *testing.T) { } } +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 { + name string + body []byte + }{ + { + name: "top-level automation_update", + body: []byte(`{"tools":[{"type":"function","name":"automation_update","strict":true,"parameters":{"type":"object","properties":{"cron":{"type":"string"}},"required":["cron"],"additionalProperties":false}}]}`), + }, + { + name: "automation_update in another namespace", + body: []byte(`{"tools":[{"type":"namespace","name":"calendar","tools":[{"type":"function","name":"automation_update","strict":true,"parameters":{"type":"object","properties":{"cron":{"type":"string"}},"required":["cron"],"additionalProperties":false}}]}]}`), + }, + { + name: "custom automation_update in codex_app", + body: []byte(`{"tools":[{"type":"namespace","name":"codex_app","tools":[{"type":"custom","name":"automation_update","strict":true,"parameters":{"type":"object","properties":{"cron":{"type":"string"}},"required":["cron"],"additionalProperties":false}}]}]}`), + }, + { + name: "large schema on another codex_app function", + body: []byte(`{"tools":[{"type":"namespace","name":"codex_app","tools":[{"type":"function","name":"exec_command","strict":true,"parameters":` + largeParams + `}]}]}`), + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + out := normalizeXAITools(tt.body) + tool := gjson.GetBytes(out, "tools.0") + if tool.Get("strict").Type != gjson.True { + t.Fatalf("strict changed for unrelated tool: %s", string(out)) + } + params := tool.Get("parameters") + if tt.name == "large schema on another codex_app function" { + if !params.Get("oneOf").Exists() || !params.Get("$defs").Exists() { + t.Fatalf("large schema was simplified: %s", string(out)) + } + return + } + if got := params.Get("properties.cron.type").String(); got != "string" { + t.Fatalf("schema was simplified, cron type = %q: %s", got, string(out)) + } + if params.Get("additionalProperties").Type != gjson.False { + t.Fatalf("additionalProperties changed: %s", string(out)) + } + }) + } +} + func TestXAIFunctionParametersNeedSimplification(t *testing.T) { auto := gjson.Parse(`{"type":"function","name":"automation_update","parameters":{"type":"object"}}`) - if !xaiFunctionParametersNeedSimplification(auto) { - t.Fatal("automation_update should always need simplification") + if !xaiFunctionParametersNeedSimplification(auto, "codex_app") { + t.Fatal("codex_app.automation_update should need simplification") + } + if xaiFunctionParametersNeedSimplification(auto, "calendar") { + t.Fatal("automation_update outside codex_app should not need simplification") + } + if xaiFunctionParametersNeedSimplification(auto, "") { + t.Fatal("top-level automation_update should not need simplification") + } + custom := gjson.Parse(`{"type":"custom","name":"automation_update","parameters":{"type":"object"}}`) + if xaiFunctionParametersNeedSimplification(custom, "codex_app") { + t.Fatal("custom codex_app.automation_update should not need simplification") } safe := gjson.Parse(`{"type":"function","name":"exec_command","parameters":{"type":"object","properties":{"cmd":{"type":"string"}}}}`) - if xaiFunctionParametersNeedSimplification(safe) { - t.Fatal("simple schema should not need simplification") + if xaiFunctionParametersNeedSimplification(safe, "codex_app") { + t.Fatal("unrelated codex_app function should not need simplification") } } diff --git a/internal/runtime/executor/xai_websockets_executor.go b/internal/runtime/executor/xai_websockets_executor.go index d7bed502f..057845446 100644 --- a/internal/runtime/executor/xai_websockets_executor.go +++ b/internal/runtime/executor/xai_websockets_executor.go @@ -471,7 +471,7 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox if sess != nil { sess.reqMu.Unlock() } - return nil, statusErr{code: respHS.StatusCode, msg: string(bodyErr)} + return nil, xaiStatusErr(respHS.StatusCode, bodyErr) } helps.RecordAPIWebsocketError(ctx, e.cfg, "dial", errDial) if sess != nil { @@ -498,10 +498,14 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox e.invalidateUpstreamConn(sess, conn, "send_error", errSend) connRetry, respHSRetry, errDialRetry := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders) if errDialRetry != nil || connRetry == nil { + bodyErrRetry := websocketHandshakeBody(respHSRetry) closeHTTPResponseBody(respHSRetry, "xai websockets executor: close handshake response body error") helps.RecordAPIWebsocketError(ctx, e.cfg, "dial_retry", errDialRetry) sess.clearActive(readCh) sess.reqMu.Unlock() + if respHSRetry != nil && respHSRetry.StatusCode > 0 { + return nil, xaiStatusErr(respHSRetry.StatusCode, bodyErrRetry) + } return nil, errDialRetry } wsReqBodyRetry := buildXAIWebsocketRequestBody(prepared.body) @@ -815,6 +819,13 @@ func buildXAIWebsocketWarmupCompletedPayload(createdPayload []byte) []byte { func parseXAIWebsocketError(payload []byte) (error, bool) { if wsErr, ok := parseCodexWebsocketError(payload); ok { + if statusError, okStatus := wsErr.(statusErrWithHeaders); okStatus { + xaiError := xaiStatusErr(statusError.code, payload) + if xaiError.retryAfter != nil { + statusError.retryAfter = xaiError.retryAfter + } + return statusError, true + } return wsErr, true } if len(payload) == 0 || !gjson.GetBytes(payload, "error").Exists() { @@ -833,7 +844,7 @@ func parseXAIWebsocketError(payload []byte) (error, bool) { if errNode := gjson.GetBytes(payload, "error"); errNode.Exists() { out, _ = sjson.SetRawBytes(out, "error", []byte(errNode.Raw)) } - return statusErr{code: status, msg: string(out)}, true + return xaiStatusErr(status, out), true } func xaiBareWebsocketErrorStatus(payload []byte) int { diff --git a/internal/runtime/executor/xai_websockets_executor_test.go b/internal/runtime/executor/xai_websockets_executor_test.go index 237e8ab5c..d9de647ea 100644 --- a/internal/runtime/executor/xai_websockets_executor_test.go +++ b/internal/runtime/executor/xai_websockets_executor_test.go @@ -712,6 +712,102 @@ func TestXAIWebsocketsExecuteStreamCompletesGenerateFalseWarmup(t *testing.T) { } } +func TestXAIWebsocketsExecuteStreamHandshakeFreeUsageExhaustedSetsRetryAfter(t *testing.T) { + body := []byte(`{"code":"subscription:free-usage-exhausted","error":"You've used all the included free usage for now."}`) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusTooManyRequests) + if _, errWrite := w.Write(body); errWrite != nil { + t.Errorf("write handshake rejection: %v", errWrite) + } + })) + defer server.Close() + + exec := NewXAIWebsocketsExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{ + ID: "xai-auth-free-usage", + 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":"hello"}`), + } + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + } + + _, err := exec.ExecuteStream(context.Background(), auth, req, opts) + if err == nil { + t.Fatal("ExecuteStream() error = nil, want handshake rejection") + } + status, ok := err.(interface{ StatusCode() int }) + if !ok || status.StatusCode() != http.StatusTooManyRequests { + t.Fatalf("status = %#v, want 429", err) + } + retryable, ok := err.(interface{ RetryAfter() *time.Duration }) + if !ok || retryable.RetryAfter() == nil { + t.Fatalf("expected RetryAfter for free-usage-exhausted handshake error: %#v", err) + } + if got := *retryable.RetryAfter(); got != 24*time.Hour { + t.Fatalf("RetryAfter = %v, want 24h", got) + } + if got := err.Error(); got != string(body) { + t.Fatalf("error payload = %q, want %q", got, body) + } +} + +func TestParseXAIWebsocketErrorFreeUsageExhaustedSetsRetryAfter(t *testing.T) { + payload := []byte(`{"type":"error","status":429,"error":{"code":"subscription:free-usage-exhausted","message":"You've used all the included free usage for now."}}`) + err, ok := parseXAIWebsocketError(payload) + if !ok { + t.Fatal("expected xAI websocket error") + } + + retryable, ok := err.(interface{ RetryAfter() *time.Duration }) + if !ok || retryable.RetryAfter() == nil { + t.Fatalf("expected RetryAfter for free-usage-exhausted websocket event: %#v", err) + } + if got := *retryable.RetryAfter(); got != 24*time.Hour { + t.Fatalf("RetryAfter = %v, want 24h", got) + } + parsed := gjson.Parse(err.Error()) + if got := parsed.Get("status").Int(); got != http.StatusTooManyRequests { + t.Fatalf("error status = %d, want 429; payload=%s", got, err) + } + if got := parsed.Get("error.code").String(); got != "subscription:free-usage-exhausted" { + t.Fatalf("error code = %q, want free-usage-exhausted; payload=%s", got, err) + } +} + +func TestParseXAIWebsocketBareErrorFreeUsageExhaustedSetsRetryAfter(t *testing.T) { + payload := []byte(`{"status":429,"error":{"code":"subscription:free-usage-exhausted","message":"You've used all the included free usage for now."}}`) + err, ok := parseXAIWebsocketError(payload) + if !ok { + t.Fatal("expected bare xAI websocket error") + } + + retryable, ok := err.(interface{ RetryAfter() *time.Duration }) + if !ok || retryable.RetryAfter() == nil { + t.Fatalf("expected RetryAfter for bare free-usage-exhausted websocket event: %#v", err) + } + if got := *retryable.RetryAfter(); got != 24*time.Hour { + t.Fatalf("RetryAfter = %v, want 24h", got) + } + parsed := gjson.Parse(err.Error()) + if got := parsed.Get("type").String(); got != "error" { + t.Fatalf("error type = %q, want error; payload=%s", got, err) + } + if got := parsed.Get("error.code").String(); got != "subscription:free-usage-exhausted" { + t.Fatalf("error code = %q, want free-usage-exhausted; payload=%s", got, err) + } +} + func TestXAIWebsocketsExecuteStreamStopsOnBareErrorPayload(t *testing.T) { upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} releaseServer := make(chan struct{})