fix(xai): prevent Desktop hang on complex tool schemas and set free-usage Retry-After

Codex Desktop injects codex_app.automation_update with a large oneOf+$ref
schema. On xAI free/build Responses, that request accepts HTTP but never
emits SSE, so clients hang until cancel.

- Simplify automation_update (and similar large oneOf+$ref schemas) to a
  permissive object schema before upstream xAI
- Map subscription free-usage-exhausted 429s to a 24h RetryAfter for
  account rotation/cooldown
- Add unit tests for both behaviors
This commit is contained in:
Wraient
2026-07-09 21:49:01 +05:30
parent bc279c6130
commit 0df267adbb
3 changed files with 161 additions and 5 deletions

View File

@@ -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
}

View File

@@ -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)

View File

@@ -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())
}
}