diff --git a/backend/internal/service/openai_gateway_grok.go b/backend/internal/service/openai_gateway_grok.go index c093029332..198055f20e 100644 --- a/backend/internal/service/openai_gateway_grok.go +++ b/backend/internal/service/openai_gateway_grok.go @@ -75,10 +75,6 @@ func (s *OpenAIGatewayService) forwardGrokResponses( upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx) defer releaseUpstreamCtx() - upstreamReq, err := buildGrokResponsesRequest(upstreamCtx, c, account, patchedBody, token, cacheIdentity, s.cfg) - if err != nil { - return nil, err - } proxyURL := "" if account.ProxyID != nil && account.Proxy != nil { @@ -86,10 +82,45 @@ func (s *OpenAIGatewayService) forwardGrokResponses( } upstreamStart := time.Now() - resp, err := s.httpUpstream.Do(upstreamReq, proxyURL, account.ID, account.Concurrency) - SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds()) - if err != nil { - return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, false) + var resp *http.Response + for attempt := 0; ; attempt++ { + upstreamReq, buildErr := buildGrokResponsesRequest(upstreamCtx, c, account, patchedBody, token, cacheIdentity, s.cfg) + if buildErr != nil { + return nil, buildErr + } + + resp, err = s.httpUpstream.Do(upstreamReq, proxyURL, account.ID, account.Concurrency) + SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds()) + if err != nil { + return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, false) + } + + // xAI can reject encrypted reasoning copied from a response produced under + // another account or cache identity. Retry once with the same routing and + // credential after removing only the rejected encrypted reasoning payload. + if attempt > 0 || resp.StatusCode != http.StatusBadRequest { + break + } + respBody := s.readUpstreamErrorBody(resp) + if resp.Body != nil { + _ = resp.Body.Close() + } + if !isGrokInvalidEncryptedContentResponse(resp.StatusCode, respBody) { + resp.Body = io.NopCloser(bytes.NewReader(respBody)) + break + } + + retryBody, changed, trimErr := trimGrokInvalidEncryptedContentRetryBody(patchedBody) + if trimErr != nil { + return nil, fmt.Errorf("prepare Grok invalid encrypted_content retry: %w", trimErr) + } + if !changed { + resp.Body = io.NopCloser(bytes.NewReader(respBody)) + break + } + + patchedBody = retryBody + slog.Info("grok_invalid_encrypted_content_retry", "account_id", account.ID, "cache_identity_present", cacheIdentity != "") } defer func() { _ = resp.Body.Close() }() @@ -162,6 +193,57 @@ func (s *OpenAIGatewayService) forwardGrokResponses( }, nil } +func isGrokInvalidEncryptedContentResponse(statusCode int, body []byte) bool { + if statusCode != http.StatusBadRequest { + return false + } + + code := gjson.GetBytes(body, "code") + message := gjson.GetBytes(body, "error") + if code.Type != gjson.String || message.Type != gjson.String || + !strings.EqualFold(strings.TrimSpace(code.String()), "invalid-argument") { + return false + } + + normalizedMessage := strings.ToLower(message.String()) + return strings.Contains(normalizedMessage, "decrypt") && strings.Contains(normalizedMessage, "encrypted_content") +} + +func trimGrokInvalidEncryptedContentRetryBody(body []byte) ([]byte, bool, error) { + input := gjson.GetBytes(body, "input") + items := input.Array() + if input.IsObject() { + items = []gjson.Result{input} + } + + hasEncryptedReasoning := false + for _, item := range items { + if strings.TrimSpace(item.Get("type").String()) == "reasoning" && item.Get("encrypted_content").Exists() { + hasEncryptedReasoning = true + break + } + } + if !hasEncryptedReasoning { + return body, false, nil + } + + var requestBody map[string]any + decoder := json.NewDecoder(bytes.NewReader(body)) + decoder.UseNumber() + if err := decoder.Decode(&requestBody); err != nil { + return nil, false, err + } + if !trimOpenAIEncryptedReasoningItems(requestBody) { + return body, false, nil + } + + retryBody, err := marshalOpenAIUpstreamJSON(requestBody) + if err != nil { + return nil, false, err + } + return retryBody, true, nil +} + func patchGrokResponsesBody(body []byte, upstreamModel string) ([]byte, error) { if !json.Valid(body) { return nil, fmt.Errorf("invalid json request body") diff --git a/backend/internal/service/openai_gateway_grok_test.go b/backend/internal/service/openai_gateway_grok_test.go index 8995c3b1de..806b9d5f19 100644 --- a/backend/internal/service/openai_gateway_grok_test.go +++ b/backend/internal/service/openai_gateway_grok_test.go @@ -1186,6 +1186,201 @@ func TestForwardGrokResponsesAPIKeyUsesXAIResponses(t *testing.T) { require.Equal(t, 1, result.Usage.OutputTokens) } +func TestForwardGrokResponsesRetriesInvalidEncryptedContentOnce(t *testing.T) { + gin.SetMode(gin.TestMode) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + body := []byte(`{ + "model":"grok", + "input":[ + {"type":"reasoning","summary":[{"type":"summary_text","text":"keep this summary"}],"encrypted_content":"encrypted-reasoning"}, + {"type":"message","role":"user","content":[{"type":"input_text","text":"hi"}]} + ], + "metadata":{"large_id":9007199254740993}, + "stream":false + }`) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body)) + c.Set("api_key", &APIKey{ID: 4535}) + + account := &Account{ + ID: 4535, + Name: "grok-api-key", + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 2, + Credentials: map[string]any{ + "api_key": "same-token", + "base_url": "https://api.x.ai/v1", + }, + } + upstream := &httpUpstreamRecorder{responses: []*http.Response{ + { + StatusCode: http.StatusBadRequest, + Header: http.Header{ + "Content-Type": []string{"application/json"}, + "Xai-Request-Id": []string{"recoverable-first"}, + }, + Body: io.NopCloser(strings.NewReader(`{"code":"invalid-argument","error":"Could not decrypt the provided encrypted_content. Ensure the value is unmodified."}`)), + }, + { + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": []string{"application/json"}, + "Xai-Request-Id": []string{"recovered-second"}, + }, + Body: io.NopCloser(strings.NewReader(`{"id":"resp_recovered","object":"response","model":"grok-4.5","status":"completed","output":[],"usage":{"input_tokens":2,"output_tokens":1}}`)), + }, + }} + svc := &OpenAIGatewayService{httpUpstream: upstream} + + result, err := svc.forwardGrokResponses(context.Background(), c, account, body, "grok", false, time.Now()) + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, "resp_recovered", result.ResponseID) + require.Equal(t, "recovered-second", result.RequestID) + require.Len(t, upstream.requests, 2) + require.Len(t, upstream.bodies, 2) + + require.Equal(t, "reasoning", gjson.GetBytes(upstream.bodies[0], "input.0.type").String()) + require.Equal(t, "encrypted-reasoning", gjson.GetBytes(upstream.bodies[0], "input.0.encrypted_content").String()) + require.Equal(t, "reasoning", gjson.GetBytes(upstream.bodies[1], "input.0.type").String()) + require.False(t, gjson.GetBytes(upstream.bodies[1], "input.0.encrypted_content").Exists()) + require.Equal(t, "keep this summary", gjson.GetBytes(upstream.bodies[1], "input.0.summary.0.text").String()) + require.Equal(t, "message", gjson.GetBytes(upstream.bodies[1], "input.1.type").String()) + require.Equal(t, "9007199254740993", gjson.GetBytes(upstream.bodies[0], "metadata.large_id").Raw) + require.Equal(t, "9007199254740993", gjson.GetBytes(upstream.bodies[1], "metadata.large_id").Raw) + + firstIdentity := gjson.GetBytes(upstream.bodies[0], "prompt_cache_key").String() + secondIdentity := gjson.GetBytes(upstream.bodies[1], "prompt_cache_key").String() + require.NotEmpty(t, firstIdentity) + require.Equal(t, firstIdentity, secondIdentity) + for _, req := range upstream.requests { + require.Equal(t, "Bearer same-token", req.Header.Get("Authorization")) + require.Equal(t, firstIdentity, req.Header.Get(grokConversationIDHeader)) + } + require.Equal(t, StatusActive, account.Status) + _, hasUpstreamErrors := c.Get(OpsUpstreamErrorsKey) + require.False(t, hasUpstreamErrors) + _, hasTerminalStatus := c.Get(OpsUpstreamStatusCodeKey) + require.False(t, hasTerminalStatus) +} + +func TestForwardGrokResponsesInvalidEncryptedContentRecoveryDoesNotOvermatch(t *testing.T) { + gin.SetMode(gin.TestMode) + + matchingError := `{"code":"invalid-argument","error":"Could not decrypt the provided encrypted_content."}` + tests := []struct { + name string + requestBody string + responseBody string + }{ + { + name: "different top-level code", + requestBody: `{"model":"grok","input":[{"type":"reasoning","encrypted_content":"cipher"}],"stream":false}`, + responseBody: `{"code":"bad-request","error":"Could not decrypt the provided encrypted_content."}`, + }, + { + name: "message does not mention decryption", + requestBody: `{"model":"grok","input":[{"type":"reasoning","encrypted_content":"cipher"}],"stream":false}`, + responseBody: `{"code":"invalid-argument","error":"The provided encrypted_content is invalid."}`, + }, + { + name: "nested OpenAI error shape", + requestBody: `{"model":"grok","input":[{"type":"reasoning","encrypted_content":"cipher"}],"stream":false}`, + responseBody: `{"code":"invalid-argument","error":{"message":"Could not decrypt the provided encrypted_content."}}`, + }, + { + name: "request has no encrypted reasoning", + requestBody: `{"model":"grok","input":[{"type":"message","role":"user","content":"hi"}],"stream":false}`, + responseBody: matchingError, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + body := []byte(tt.requestBody) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body)) + + account := &Account{ + ID: 4536, + Name: "grok-api-key", + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{"api_key": "token", "base_url": "https://api.x.ai/v1"}, + } + upstream := &httpUpstreamRecorder{responses: []*http.Response{{ + StatusCode: http.StatusBadRequest, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(tt.responseBody)), + }}} + svc := &OpenAIGatewayService{httpUpstream: upstream} + + result, err := svc.forwardGrokResponses(context.Background(), c, account, body, "grok", false, time.Now()) + require.Nil(t, result) + require.Error(t, err) + require.Len(t, upstream.requests, 1) + require.Len(t, upstream.bodies, 1) + }) + } +} + +func TestForwardGrokResponsesInvalidEncryptedContentRetryFailureIsTerminal(t *testing.T) { + gin.SetMode(gin.TestMode) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + body := []byte(`{"model":"grok","input":[{"type":"reasoning","encrypted_content":"cipher"},{"type":"message","role":"user","content":"hi"}],"stream":false}`) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body)) + + account := &Account{ + ID: 4537, + Name: "grok-api-key", + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{"api_key": "same-token", "base_url": "https://api.x.ai/v1"}, + } + newInvalidEncryptedResponse := func(requestID string) *http.Response { + return &http.Response{ + StatusCode: http.StatusBadRequest, + Header: http.Header{ + "Content-Type": []string{"application/json"}, + "Xai-Request-Id": []string{requestID}, + }, + Body: io.NopCloser(strings.NewReader(`{"code":"invalid-argument","error":"Could not decrypt the provided encrypted_content."}`)), + } + } + upstream := &httpUpstreamRecorder{responses: []*http.Response{ + newInvalidEncryptedResponse("recoverable-first"), + newInvalidEncryptedResponse("terminal-second"), + }} + svc := &OpenAIGatewayService{httpUpstream: upstream} + + result, err := svc.forwardGrokResponses(context.Background(), c, account, body, "grok", false, time.Now()) + require.Nil(t, result) + require.Error(t, err) + require.Len(t, upstream.requests, 2) + require.Len(t, upstream.bodies, 2) + require.True(t, gjson.GetBytes(upstream.bodies[0], "input.0.encrypted_content").Exists()) + require.False(t, gjson.GetBytes(upstream.bodies[1], `input.#(type=="reasoning")`).Exists()) + + rawEvents, ok := c.Get(OpsUpstreamErrorsKey) + require.True(t, ok) + events, ok := rawEvents.([]*OpsUpstreamErrorEvent) + require.True(t, ok) + require.NotEmpty(t, events) + for _, event := range events { + require.NotEqual(t, "recoverable-first", event.UpstreamRequestID) + } + require.Equal(t, http.StatusBadRequest, c.GetInt(OpsUpstreamStatusCodeKey)) +} + func TestForwardAsChatCompletionsForGrokAPIKeyUsesConfiguredRawEndpointWithoutOAuthIdentity(t *testing.T) { gin.SetMode(gin.TestMode) recorder := httptest.NewRecorder()