diff --git a/backend/internal/service/account_test_service.go b/backend/internal/service/account_test_service.go index f83a99a7e2..5a0bc40f40 100644 --- a/backend/internal/service/account_test_service.go +++ b/backend/internal/service/account_test_service.go @@ -655,6 +655,7 @@ func (s *AccountTestService) testOpenAIAccountConnection(c *gin.Context, account if resp.StatusCode != http.StatusOK { body, _ := io.ReadAll(resp.Body) + body = redactAgentIdentitySensitiveBodyForAccount(ctx, s.accountRepo, credentialAccount, body) if resp.StatusCode == http.StatusTooManyRequests { s.reconcileOpenAI429State(ctx, account, resp.Header, body) } @@ -924,6 +925,7 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account defer func() { _ = resp.Body.Close() }() body, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + body = redactAgentIdentitySensitiveBodyForAccount(ctx, s.accountRepo, credentialAccount, body) if s.accountRepo != nil { updates := buildOpenAICompactProbeExtraUpdates(resp, body, nil, time.Now()) @@ -1800,6 +1802,7 @@ func (s *AccountTestService) testOpenAIImageOAuth(c *gin.Context, ctx context.Co }() if resp.StatusCode >= 400 { body, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) + body = redactAgentIdentitySensitiveBodyForAccount(ctx, s.accountRepo, credentialAccount, body) message := strings.TrimSpace(extractUpstreamErrorMessage(body)) if message == "" { message = fmt.Sprintf("Responses API returned %d", resp.StatusCode) @@ -1811,6 +1814,7 @@ func (s *AccountTestService) testOpenAIImageOAuth(c *gin.Context, ctx context.Co if err != nil { return s.sendErrorAndEnd(c, fmt.Sprintf("Failed to read image response: %s", err.Error())) } + body = redactAgentIdentitySensitiveBodyForAccount(ctx, s.accountRepo, credentialAccount, body) results, _, _, _, _, err := collectOpenAIImagesFromResponsesBody(body) if err != nil { diff --git a/backend/internal/service/openai_agent_identity.go b/backend/internal/service/openai_agent_identity.go index c8818cd6f3..8383a10e9c 100644 --- a/backend/internal/service/openai_agent_identity.go +++ b/backend/internal/service/openai_agent_identity.go @@ -29,6 +29,8 @@ const ( var openAIAgentIdentityAuthAPIBaseURL = agentIdentityAuthAPIBaseURL +var agentIdentityTaskLocks sync.Map // map[int64]*sync.Mutex + type agentIdentityKey struct { runtimeID string privateKey ed25519.PrivateKey @@ -251,8 +253,14 @@ func ensureAgentIdentityTaskForAccount(ctx context.Context, repo AccountReposito if taskMu == nil { return errors.New("agent identity task lock is unavailable") } - taskMu.Lock() - defer taskMu.Unlock() + sharedTaskMu := taskMu + if credAccount.ID > 0 { + candidate := &sync.Mutex{} + actual, _ := agentIdentityTaskLocks.LoadOrStore(credAccount.ID, candidate) + sharedTaskMu = actual.(*sync.Mutex) + } + sharedTaskMu.Lock() + defer sharedTaskMu.Unlock() currentTaskID = strings.TrimSpace(credAccount.GetCredential("task_id")) if currentTaskID != "" && (expectedTaskID == "" || currentTaskID != expectedTaskID) { return nil @@ -302,6 +310,10 @@ func isAgentIdentityTaskInvalidHTTPResponse(statusCode int, body []byte) bool { return false } +func isAgentIdentityTaskInvalidWSDialError(err *openAIWSDialError) bool { + return err != nil && isAgentIdentityTaskInvalidHTTPResponse(err.StatusCode, err.ResponseBody) +} + func (s *OpenAIGatewayService) buildOpenAIAuthenticationHeaders(ctx context.Context, account *Account, token string) (http.Header, error) { if account == nil { return nil, errors.New("account is nil") @@ -401,16 +413,19 @@ func (s *OpenAIGatewayService) isAgentIdentityAccount(ctx context.Context, accou // upstream error can reach logs, ops events, or returned error text. Agent // Identity responses should not echo these values, but keeping this boundary // defensive prevents accidental disclosure if an upstream error does. -func (s *OpenAIGatewayService) redactAgentIdentitySensitiveBody(ctx context.Context, account *Account, body []byte) []byte { - if !s.isAgentIdentityAccount(ctx, account) || len(body) == 0 { +func redactAgentIdentitySensitiveBodyForAccount(ctx context.Context, repo AccountRepository, account *Account, body []byte) []byte { + if account == nil || len(body) == 0 { return body } credAccount := account if account != nil && account.IsShadow() { - if resolved, err := resolveCredentialAccount(ctx, s.accountRepo, account); err == nil && resolved != nil { + if resolved, err := resolveCredentialAccount(ctx, repo, account); err == nil && resolved != nil { credAccount = resolved } } + if credAccount == nil || !credAccount.IsOpenAIAgentIdentity() { + return body + } redacted := string(body) for _, key := range []string{ "agent_private_key", @@ -427,5 +442,23 @@ func (s *OpenAIGatewayService) redactAgentIdentitySensitiveBody(ctx context.Cont redacted = strings.ReplaceAll(redacted, value, "[redacted]") } } + for { + start := strings.Index(redacted, "AgentAssertion ") + if start < 0 { + break + } + end := start + len("AgentAssertion ") + for end < len(redacted) && !strings.ContainsRune(" \t\r\n\"',}", rune(redacted[end])) { + end++ + } + redacted = redacted[:start] + "AgentAssertion [redacted]" + redacted[end:] + } return []byte(redacted) } + +func (s *OpenAIGatewayService) redactAgentIdentitySensitiveBody(ctx context.Context, account *Account, body []byte) []byte { + if !s.isAgentIdentityAccount(ctx, account) { + return body + } + return redactAgentIdentitySensitiveBodyForAccount(ctx, s.accountRepo, account, body) +} diff --git a/backend/internal/service/openai_agent_identity_compat_test.go b/backend/internal/service/openai_agent_identity_compat_test.go index 5c0c996c42..c034af2bfa 100644 --- a/backend/internal/service/openai_agent_identity_compat_test.go +++ b/backend/internal/service/openai_agent_identity_compat_test.go @@ -104,13 +104,45 @@ func TestOpenAIAgentIdentityErrorRedactionDoesNotLeakCredentialValues(t *testing } svc := &OpenAIGatewayService{} oauthValue := account.GetCredential("access_token") - redacted := svc.redactAgentIdentitySensitiveBody(context.Background(), account, []byte(`{"message":"runtime-test task-test `+oauthValue+`"}`)) + redacted := svc.redactAgentIdentitySensitiveBody(context.Background(), account, []byte(`{"message":"runtime-test task-test `+oauthValue+` AgentAssertion abc123"}`)) require.NotContains(t, string(redacted), key.runtimeID) require.NotContains(t, string(redacted), key.taskID) require.NotContains(t, string(redacted), oauthValue) + require.NotContains(t, string(redacted), "AgentAssertion abc123") require.Contains(t, string(redacted), "[redacted]") } +func TestOpenAIAuthenticationHeadersPreserveOAuthPATAndAPIKeyBearerModes(t *testing.T) { + svc := &OpenAIGatewayService{} + tests := []struct { + name string + account *Account + token string + }{ + {name: "oauth", account: &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth}, token: "oauth-runtime-token"}, + {name: "personal access token", account: &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth, Credentials: map[string]any{"auth_mode": OpenAIAuthModePersonalAccessToken}}, token: "pat-runtime-token"}, + {name: "api key", account: &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}, token: "api-key-runtime-token"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + headers, err := svc.buildOpenAIAuthenticationHeaders(context.Background(), tt.account, tt.token) + require.NoError(t, err) + require.Equal(t, "Bearer "+tt.token, headers.Get("Authorization")) + }) + } +} + +func TestOpenAIWSAgentIdentityRecoveryRequiresTaskInvalidBody(t *testing.T) { + require.False(t, isAgentIdentityTaskInvalidWSDialError(&openAIWSDialError{ + StatusCode: http.StatusUnauthorized, + ResponseBody: []byte(`{"error":{"code":"invalid_signature"}}`), + })) + require.True(t, isAgentIdentityTaskInvalidWSDialError(&openAIWSDialError{ + StatusCode: http.StatusUnauthorized, + ResponseBody: []byte(`{"error":{"code":"invalid_task_id"}}`), + })) +} + func TestOpenAIWSConnPoolHeadersFactoryRunsAtDialAndStalePrewarmIsDiscarded(t *testing.T) { cfg := &config.Config{} cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1 @@ -218,6 +250,21 @@ func TestOpenAIAgentIdentityTaskInvalidRetriesExactlyOnce(t *testing.T) { require.Error(t, err) require.Equal(t, 2, registerCalls) require.Len(t, upstream.requests, 4) + + // Passthrough uses the same one-shot task recovery contract. + account.Extra = map[string]any{"openai_passthrough": true} + account.Credentials["task_id"] = "task-old-passthrough" + upstream.responses = []*http.Response{ + {StatusCode: http.StatusUnauthorized, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: io.NopCloser(strings.NewReader(`{"error":{"code":"invalid_task_id"}}`))}, + {StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: io.NopCloser(strings.NewReader(successBody))}, + } + rec3 := httptest.NewRecorder() + c3, _ := gin.CreateTestContext(rec3) + c3.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"gpt-5.4","instructions":"Reply OK","input":[],"stream":false}`)) + _, err = svc.Forward(context.Background(), c3, account, []byte(`{"model":"gpt-5.4","instructions":"Reply OK","input":[],"stream":false}`)) + require.NoError(t, err) + require.Equal(t, 3, registerCalls) + require.Len(t, upstream.requests, 6) } func decodeAgentAssertionTask(t *testing.T, header string) string { diff --git a/backend/internal/service/openai_agent_identity_test.go b/backend/internal/service/openai_agent_identity_test.go index 73ac0ddf08..6e94f53abe 100644 --- a/backend/internal/service/openai_agent_identity_test.go +++ b/backend/internal/service/openai_agent_identity_test.go @@ -156,6 +156,44 @@ func TestEnsureAgentIdentityTaskPersistsAndRedactsCredentials(t *testing.T) { require.NotContains(t, string(mustJSON(t, redacted)), privateKey) } +func TestEnsureAgentIdentityTaskSharesLockAcrossServicesForSameAccount(t *testing.T) { + key, privateKey := newTestAgentIdentityKey(t) + account := &Account{ID: 9001, Type: AccountTypeOAuth, Platform: PlatformOpenAI, Credentials: map[string]any{ + "auth_mode": OpenAIAuthModeAgentIdentity, + "agent_runtime_id": key.runtimeID, + "agent_private_key": privateKey, + }} + repo := &agentIdentityCredentialsRepo{} + registerCalls := 0 + var registerMu sync.Mutex + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + registerMu.Lock() + registerCalls++ + registerMu.Unlock() + _, _ = w.Write([]byte(`{"task_id":"task-shared"}`)) + })) + defer server.Close() + oldBase := openAIAgentIdentityAuthAPIBaseURL + openAIAgentIdentityAuthAPIBaseURL = server.URL + t.Cleanup(func() { openAIAgentIdentityAuthAPIBaseURL = oldBase }) + + start := make(chan struct{}) + errors := make(chan error, 2) + for range 2 { + go func() { + <-start + errors <- ensureAgentIdentityTaskForAccount(context.Background(), repo, nil, &sync.Mutex{}, account, "") + }() + } + close(start) + require.NoError(t, <-errors) + require.NoError(t, <-errors) + registerMu.Lock() + defer registerMu.Unlock() + require.Equal(t, 1, registerCalls) + require.Equal(t, "task-shared", account.GetCredential("task_id")) +} + type agentIdentityCredentialsRepo struct { AccountRepository credentials map[string]any diff --git a/backend/internal/service/openai_gateway_passthrough.go b/backend/internal/service/openai_gateway_passthrough.go index 2ea914b537..f280f37841 100644 --- a/backend/internal/service/openai_gateway_passthrough.go +++ b/backend/internal/service/openai_gateway_passthrough.go @@ -11,6 +11,7 @@ import ( "encoding/json" "errors" "fmt" + "io" "net/http" "sort" "strings" @@ -161,13 +162,6 @@ func (s *OpenAIGatewayService) forwardOpenAIPassthrough( return nil, err } - upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx) - upstreamReq, err := s.buildUpstreamRequestOpenAIPassthrough(upstreamCtx, c, account, body, token) - releaseUpstreamCtx() - if err != nil { - return nil, err - } - proxyURL := "" if account.ProxyID != nil && account.Proxy != nil { proxyURL = account.Proxy.URL() @@ -177,18 +171,42 @@ func (s *OpenAIGatewayService) forwardOpenAIPassthrough( c.Set("openai_passthrough", true) } - upstreamStart := time.Now() - resp, err := s.httpUpstream.Do(upstreamReq, proxyURL, account.ID, account.Concurrency) - SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds()) - if err != nil { - // Transport-level failure (proxy/DNS/TCP/TLS — no HTTP response). Convert to - // a failover so the handler switches to a healthy account, and temporarily - // unschedule the account on durable faults (e.g. rejected proxy credentials). - return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, true) - } - defer func() { _ = resp.Body.Close() }() + agentTaskRecoveryTried := false + var resp *http.Response + for { + upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx) + upstreamReq, buildErr := s.buildUpstreamRequestOpenAIPassthrough(upstreamCtx, c, account, body, token) + releaseUpstreamCtx() + if buildErr != nil { + return nil, buildErr + } + + upstreamStart := time.Now() + resp, err = s.httpUpstream.Do(upstreamReq, proxyURL, account.ID, account.Concurrency) + SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds()) + if err != nil { + // Transport-level failure (proxy/DNS/TCP/TLS — no HTTP response). Convert to + // a failover so the handler switches to a healthy account. + return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, true) + } + if resp.StatusCode < 400 { + break + } + + // Peek only to identify an invalid task. Restore the body so the existing + // passthrough error handling sees the same response after recovery fails. + probeBody := s.readUpstreamErrorBody(resp) + _ = resp.Body.Close() + resp.Body = io.NopCloser(bytes.NewReader(probeBody)) + if !agentTaskRecoveryTried && s.isAgentIdentityAccount(ctx, account) && isAgentIdentityTaskInvalidHTTPResponse(resp.StatusCode, probeBody) { + agentTaskRecoveryTried = true + expectedTaskID := account.GetCredential("task_id") + if recoveryErr := s.recoverAgentIdentityTask(ctx, account, expectedTaskID); recoveryErr != nil { + return nil, fmt.Errorf("agent identity task recovery failed: %w", recoveryErr) + } + continue + } - if resp.StatusCode >= 400 { // 透传模式默认保持原样代理；但 429/529 属于网关必须兜底的 // 上游容量类错误，应先触发多账号 failover 以维持基础 SLA。 if shouldFailoverOpenAIPassthroughResponse(resp.StatusCode) { @@ -196,6 +214,7 @@ func (s *OpenAIGatewayService) forwardOpenAIPassthrough( } return nil, s.handleErrorResponsePassthrough(ctx, resp, c, account, body) } + defer func() { _ = resp.Body.Close() }() serviceTier := extractOpenAIServiceTierFromBody(body) diff --git a/backend/internal/service/openai_quota_service.go b/backend/internal/service/openai_quota_service.go index b8b766fcff..cebab0a1d7 100644 --- a/backend/internal/service/openai_quota_service.go +++ b/backend/internal/service/openai_quota_service.go @@ -171,7 +171,7 @@ func (s *OpenAIQuotaService) QueryUsage(ctx context.Context, accountID int64) (* } if !resp.IsSuccessState() { status := resp.StatusCode - body := truncate(resp.String(), 240) + body := truncate(s.redactQuotaErrorBody(ctx, accountID, resp.String()), 240) slog.Warn("openai_quota_query_failed", "account_id", accountID, "status", status, "body", body) return nil, infraerrors.Newf(mapUpstreamStatus(status), "OPENAI_QUOTA_UPSTREAM_ERROR", "upstream returned %d: %s", status, body) } @@ -268,7 +268,7 @@ func (s *OpenAIQuotaService) ResetCredit(ctx context.Context, accountID int64) ( } if !resp.IsSuccessState() { status := resp.StatusCode - body := truncate(resp.String(), 240) + body := truncate(s.redactQuotaErrorBody(callCtx, accountID, resp.String()), 240) slog.Warn("openai_quota_reset_failed", "account_id", accountID, "status", status, "body", body) return nil, infraerrors.Newf(mapUpstreamStatus(status), "OPENAI_QUOTA_RESET_UPSTREAM_ERROR", "upstream returned %d: %s", status, body) } @@ -393,6 +393,17 @@ func (s *OpenAIQuotaService) buildCodexQuotaHeaders(ctx context.Context, account return headers, nil } +func (s *OpenAIQuotaService) redactQuotaErrorBody(ctx context.Context, accountID int64, body string) string { + if s == nil || s.accountRepo == nil { + return body + } + account, err := s.accountRepo.GetByID(ctx, accountID) + if err != nil || account == nil { + return body + } + return string(redactAgentIdentitySensitiveBodyForAccount(ctx, s.accountRepo, account, []byte(body))) +} + // buildCodexCommonHeaders sets the request headers expected by the chatgpt.com // backend so calls succeed past Cloudflare/WASM checks. func buildCodexCommonHeaders(accessToken, chatGPTAccountID string, fedRAMP bool) map[string]string { diff --git a/backend/internal/service/openai_ws_client.go b/backend/internal/service/openai_ws_client.go index 80b7553083..d30c4a1cb3 100644 --- a/backend/internal/service/openai_ws_client.go +++ b/backend/internal/service/openai_ws_client.go @@ -4,6 +4,7 @@ import ( "context" "errors" "fmt" + "io" "net/http" "net/url" "strings" @@ -61,6 +62,28 @@ type coderOpenAIWSClientDialer struct { proxyMisses atomic.Int64 } +// openAIWSHandshakeError keeps a bounded, non-logged HTTP error body so the +// Agent Identity recovery path can distinguish an invalid task from other +// 401 handshake failures. +type openAIWSHandshakeError struct { + Body []byte + Err error +} + +func (e *openAIWSHandshakeError) Error() string { + if e == nil || e.Err == nil { + return "openai ws handshake failed" + } + return e.Err.Error() +} + +func (e *openAIWSHandshakeError) Unwrap() error { + if e == nil { + return nil + } + return e.Err +} + type openAIWSProxyClientEntry struct { client *http.Client lastUsedUnixNano int64 @@ -97,7 +120,12 @@ func (d *coderOpenAIWSClientDialer) Dial( status = resp.StatusCode respHeaders = cloneHeader(resp.Header) } - return nil, status, respHeaders, err + var body []byte + if resp != nil && resp.Body != nil { + body, _ = io.ReadAll(io.LimitReader(resp.Body, 8<<10)) + _ = resp.Body.Close() + } + return nil, status, respHeaders, &openAIWSHandshakeError{Body: body, Err: err} } // coder/websocket 默认单消息读取上限为 32KB，Codex WS 事件（如 rate_limits/大 delta） // 可能超过该阈值，需显式提高上限，避免本地 read_fail(message too big)。 diff --git a/backend/internal/service/openai_ws_forwarder_ingress.go b/backend/internal/service/openai_ws_forwarder_ingress.go index 9a6453892a..169919b8b9 100644 --- a/backend/internal/service/openai_ws_forwarder_ingress.go +++ b/backend/internal/service/openai_ws_forwarder_ingress.go @@ -617,7 +617,7 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( lease, acquireErr := pool.Acquire(acquireCtx, req) acquireCancel() var dialErr *openAIWSDialError - if acquireErr != nil && s.isAgentIdentityAccount(ctx, account) && errors.As(acquireErr, &dialErr) && dialErr != nil && dialErr.StatusCode == http.StatusUnauthorized && !agentTaskRecoveryTried { + if acquireErr != nil && s.isAgentIdentityAccount(ctx, account) && errors.As(acquireErr, &dialErr) && isAgentIdentityTaskInvalidWSDialError(dialErr) && !agentTaskRecoveryTried { agentTaskRecoveryTried = true if recoveryErr := s.recoverAgentIdentityTask(ctx, account, account.GetCredential("task_id")); recoveryErr != nil { return nil, fmt.Errorf("agent identity task recovery failed: %w", recoveryErr) diff --git a/backend/internal/service/openai_ws_forwarder_v2.go b/backend/internal/service/openai_ws_forwarder_v2.go index 65a6add7a5..90151be93a 100644 --- a/backend/internal/service/openai_ws_forwarder_v2.go +++ b/backend/internal/service/openai_ws_forwarder_v2.go @@ -192,7 +192,7 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2( }) if err != nil { var agentDialErr *openAIWSDialError - if s.isAgentIdentityAccount(ctx, account) && errors.As(err, &agentDialErr) && agentDialErr != nil && agentDialErr.StatusCode == http.StatusUnauthorized && agentTaskRecoveryTried != nil && !*agentTaskRecoveryTried { + if s.isAgentIdentityAccount(ctx, account) && errors.As(err, &agentDialErr) && isAgentIdentityTaskInvalidWSDialError(agentDialErr) && agentTaskRecoveryTried != nil && !*agentTaskRecoveryTried { *agentTaskRecoveryTried = true if recoveryErr := s.recoverAgentIdentityTask(ctx, account, account.GetCredential("task_id")); recoveryErr != nil { return nil, fmt.Errorf("agent identity task recovery failed: %w", recoveryErr) diff --git a/backend/internal/service/openai_ws_pool.go b/backend/internal/service/openai_ws_pool.go index be81d611f9..3a02da6fe5 100644 --- a/backend/internal/service/openai_ws_pool.go +++ b/backend/internal/service/openai_ws_pool.go @@ -39,6 +39,7 @@ var ( type openAIWSDialError struct { StatusCode int ResponseHeaders http.Header + ResponseBody []byte Err error } @@ -1540,9 +1541,15 @@ func (p *openAIWSConnPool) dialConn(ctx context.Context, req openAIWSAcquireRequ } conn, status, handshakeHeaders, err := p.clientDialer.Dial(ctx, req.WSURL, headers, req.ProxyURL) if err != nil { + var handshakeErr *openAIWSHandshakeError + var responseBody []byte + if errors.As(err, &handshakeErr) && handshakeErr != nil { + responseBody = append([]byte(nil), handshakeErr.Body...) + } return nil, &openAIWSDialError{ StatusCode: status, ResponseHeaders: cloneHeader(handshakeHeaders), + ResponseBody: responseBody, Err: err, } } diff --git a/backend/internal/service/openai_ws_v2_passthrough_adapter.go b/backend/internal/service/openai_ws_v2_passthrough_adapter.go index 01b66c376b..bb097e540b 100644 --- a/backend/internal/service/openai_ws_v2_passthrough_adapter.go +++ b/backend/internal/service/openai_ws_v2_passthrough_adapter.go @@ -374,7 +374,8 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough( if err == nil { break } - if s.isAgentIdentityAccount(ctx, account) && statusCode == http.StatusUnauthorized && !agentTaskRecoveryTried { + var dialErr *openAIWSDialError + if s.isAgentIdentityAccount(ctx, account) && errors.As(err, &dialErr) && isAgentIdentityTaskInvalidWSDialError(dialErr) && !agentTaskRecoveryTried { agentTaskRecoveryTried = true if recoveryErr := s.recoverAgentIdentityTask(ctx, account, account.GetCredential("task_id")); recoveryErr != nil { return fmt.Errorf("agent identity task recovery failed: %w", recoveryErr) @@ -696,9 +697,15 @@ func (s *OpenAIGatewayService) mapOpenAIWSPassthroughDialError( wrappedErr := err var dialErr *openAIWSDialError if !errors.As(err, &dialErr) { + var handshakeErr *openAIWSHandshakeError + var responseBody []byte + if errors.As(err, &handshakeErr) && handshakeErr != nil { + responseBody = append([]byte(nil), handshakeErr.Body...) + } wrappedErr = &openAIWSDialError{ StatusCode: statusCode, ResponseHeaders: cloneHeader(handshakeHeaders), + ResponseBody: responseBody, Err: err, } }