Merge pull request #4157 from Wraient/fix/xai-desktop-tool-schema-and-free-usage-retry

internal/runtime/executor/xai_executor.go

@@ -44,6 +44,12 @@ const (
 	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".
+	xaiAutomationUpdateToolName = "automation_update"
+	// Permissive placeholder schema: keeps the tool callable without the hang.
+	xaiSafeFunctionParameters   = `{"type":"object","properties":{},"additionalProperties":true}`
 	xaiImagesGenerationsPath    = "/images/generations"
 	xaiImagesEditsPath          = "/images/edits"
 	xaiDefaultImageEndpointPath = xaiImagesGenerationsPath
@@ -157,7 +163,7 @@ func (e *XAIExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req
 		}
 		helps.AppendAPIResponseChunk(ctx, e.cfg, data)
 		helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data))
-		return resp, statusErr{code: httpResp.StatusCode, msg: string(data)}
+		return resp, xaiStatusErr(httpResp.StatusCode, data)
 	}
 
 	data, err := io.ReadAll(httpResp.Body)
@@ -253,7 +259,7 @@ func (e *XAIExecutor) executeCompactRequest(ctx context.Context, auth *cliproxya
 
 	if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
 		helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data))
-		err = statusErr{code: httpResp.StatusCode, msg: string(data)}
+		err = xaiStatusErr(httpResp.StatusCode, data)
 		return nil, nil, nil, err
 	}
 
@@ -493,7 +499,7 @@ func (e *XAIExecutor) executeImages(ctx context.Context, auth *cliproxyauth.Auth
 
 	if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
 		helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data))
-		return resp, statusErr{code: httpResp.StatusCode, msg: string(data)}
+		return resp, xaiStatusErr(httpResp.StatusCode, data)
 	}
 
 	return cliproxyexecutor.Response{Payload: data, Headers: httpResp.Header.Clone()}, nil
@@ -558,7 +564,7 @@ func (e *XAIExecutor) executeVideos(ctx context.Context, auth *cliproxyauth.Auth
 
 	if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
 		helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data))
-		return resp, statusErr{code: httpResp.StatusCode, msg: string(data)}
+		return resp, xaiStatusErr(httpResp.StatusCode, data)
 	}
 
 	return cliproxyexecutor.Response{Payload: data, Headers: httpResp.Header.Clone()}, nil
@@ -613,7 +619,7 @@ func (e *XAIExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth
 		}
 		helps.AppendAPIResponseChunk(ctx, e.cfg, data)
 		helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data))
-		return nil, statusErr{code: httpResp.StatusCode, msg: string(data)}
+		return nil, xaiStatusErr(httpResp.StatusCode, data)
 	}
 
 	out := make(chan cliproxyexecutor.StreamChunk)
@@ -1143,9 +1149,43 @@ 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) {
+		updatedTool, errSet := sjson.SetRawBytes(raw, "parameters", []byte(xaiSafeFunctionParameters))
+		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())
+	}
 	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
+}
+
 func sanitizeXAIInputEncryptedContent(body []byte) []byte {
 	input := gjson.GetBytes(body, "input")
 	if !input.Exists() || !input.IsArray() {
@@ -1585,3 +1625,30 @@ func xaiPatchCompletedOutput(eventData []byte, outputItemsByIndex map[int64][]by
 	patched, _ := sjson.SetRawBytes(eventData, "response.output", outputArray)
 	return patched
 }
+
+// xaiFreeUsageExhaustedCooldown is the free-tier rolling window advertised by
+// cli-chat-proxy ("Usage resets over a rolling 24-hour window").
+const xaiFreeUsageExhaustedCooldown = 24 * time.Hour
+
+// xaiStatusErr wraps upstream error bodies so free-tier exhaustion
+// (subscription:free-usage-exhausted) carries a 24h RetryAfter hint for
+// auth cooldown / account rotation. Generic 429s stay without an explicit
+// retry hint so conductor backoff still applies.
+func xaiStatusErr(code int, body []byte) statusErr {
+	err := statusErr{code: code, msg: string(body)}
+	if code != http.StatusTooManyRequests || len(body) == 0 {
+		return err
+	}
+	codeStr := strings.ToLower(gjson.GetBytes(body, "code").String())
+	msg := strings.ToLower(gjson.GetBytes(body, "error").String())
+	if msg == "" {
+		msg = strings.ToLower(string(body))
+	}
+	if strings.Contains(codeStr, "free-usage-exhausted") ||
+		strings.Contains(msg, "free-usage-exhausted") ||
+		strings.Contains(msg, "included free usage") {
+		d := xaiFreeUsageExhaustedCooldown
+		err.retryAfter = &d
+	}
+	return err
+}

internal/runtime/executor/xai_executor_test.go

@@ -1060,6 +1060,58 @@ func TestXAIExecutorExecuteVideosUsesNativeEndpointFromRequestPath(t *testing.T)
 	}
 }
 
+func TestNormalizeXAITools_SimplifiesAutomationUpdateSchema(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"}}}}]}`)
+	out := normalizeXAITools(body)
+
+	tools := gjson.GetBytes(out, "tools")
+	if !tools.IsArray() {
+		t.Fatalf("tools missing: %s", string(out))
+	}
+	foundAuto := false
+	foundExec := false
+	for _, tool := range tools.Array() {
+		switch tool.Get("name").String() {
+		case "automation_update":
+			foundAuto = true
+			paramsRaw := tool.Get("parameters").Raw
+			if strings.Contains(paramsRaw, `"oneOf"`) || strings.Contains(paramsRaw, `"$defs"`) {
+				t.Fatalf("automation_update parameters were not simplified: %s", paramsRaw)
+			}
+			if tool.Get("parameters.type").String() != "object" {
+				t.Fatalf("automation_update parameters.type = %q, want object", tool.Get("parameters.type").String())
+			}
+			if tool.Get("parameters.additionalProperties").Type != gjson.True {
+				t.Fatalf("automation_update parameters should allow additionalProperties: %s", paramsRaw)
+			}
+		case "exec_command":
+			foundExec = true
+			if got := tool.Get("parameters.properties.cmd.type").String(); got != "string" {
+				t.Fatalf("exec_command schema should be preserved, got %q in %s", got, tool.Raw)
+			}
+		}
+	}
+	if !foundAuto {
+		t.Fatalf("automation_update tool missing after normalize: %s", string(out))
+	}
+	if !foundExec {
+		t.Fatalf("exec_command tool missing after normalize: %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")
+	}
+	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")
+	}
+}
+
 func TestNormalizeXAIToolChoiceForTools_DropsWhenToolsEmpty(t *testing.T) {
 	body := []byte(`{"model":"grok-4","tools":[],"tool_choice":"auto","parallel_tool_calls":true,"input":"hi"}`)
 	out := normalizeXAIToolChoiceForTools(body)

internal/runtime/executor/xai_status_err_test.go

@@ -0,0 +1,37 @@
+package executor
+
+import (
+	"net/http"
+	"testing"
+	"time"
+)
+
+func TestXAIStatusErr_FreeUsageExhaustedSets24hRetryAfter(t *testing.T) {
+	body := []byte(`{"code":"subscription:free-usage-exhausted","error":"You've used all the included free usage for model grok-4.5-build-free for now. Usage resets over a rolling 24-hour window — tokens (actual/limit): 1065387/1000000."}`)
+	err := xaiStatusErr(http.StatusTooManyRequests, body)
+	if err.StatusCode() != http.StatusTooManyRequests {
+		t.Fatalf("status = %d, want 429", err.StatusCode())
+	}
+	if err.RetryAfter() == nil {
+		t.Fatal("expected RetryAfter for free-usage-exhausted")
+	}
+	if *err.RetryAfter() != 24*time.Hour {
+		t.Fatalf("RetryAfter = %v, want 24h", *err.RetryAfter())
+	}
+}
+
+func TestXAIStatusErr_Generic429HasNoRetryAfter(t *testing.T) {
+	body := []byte(`{"code":"rate_limit","error":"too many requests"}`)
+	err := xaiStatusErr(http.StatusTooManyRequests, body)
+	if err.RetryAfter() != nil {
+		t.Fatalf("expected nil RetryAfter for generic 429, got %v", *err.RetryAfter())
+	}
+}
+
+func TestXAIStatusErr_Non429Unchanged(t *testing.T) {
+	body := []byte(`{"error":"nope"}`)
+	err := xaiStatusErr(http.StatusBadRequest, body)
+	if err.RetryAfter() != nil {
+		t.Fatalf("expected nil RetryAfter for 400, got %v", *err.RetryAfter())
+	}
+}