fix(gateway): 修复 WebSocket 首错重试与故障转移 (#646)

internal/execution/wsnative/session.go

@@ -337,6 +337,11 @@ func (s *session) ExecuteTurn(ctx context.Context, payload []byte, emit func(con
 						code = "upstream_response_failed"
 					}
 					result.Error = failure(execution.ErrorKindProvider, code, event.Status)
+					if event.Type == "error" || event.Type == "response.failed" ||
+						(event.Type == "response.done" && event.Response.Status == "failed") {
+						// 明确错误事件交由共享规则判断；断流和不完整响应仍保持结果未知。
+						result.Error.ReplaySafety = ""
+					}
 					// 原生请求错误的健康作用域由既有分类器结合错误码确定。
 					result.Error.ScopeHint = ""
 				}

internal/gateway/websocket_retry_test.go

@@ -0,0 +1,230 @@
+package gateway
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/gorilla/websocket"
+
+	"gpt-load/internal/channel"
+	"gpt-load/internal/execution"
+	"gpt-load/internal/platform/config"
+	"gpt-load/internal/state"
+)
+
+func TestWebsocketFirstRejectionRetry(t *testing.T) {
+	for _, tc := range []struct {
+		name                                          string
+		retries, credentials, wantAttempts            int
+		partial, unknown, exhausted, bound, multiplex bool
+		host, invalid                                 bool
+	}{
+		{name: "quota switches credential", retries: 2, credentials: 2, wantAttempts: 2},
+		{name: "host error skips group", retries: 2, credentials: 2, wantAttempts: 2, host: true},
+		{name: "explicit request rejection stops", retries: 2, credentials: 2, wantAttempts: 1, invalid: true},
+		{name: "zero budget preserves error", credentials: 2, wantAttempts: 1},
+		{name: "no candidate preserves error", retries: 2, credentials: 1, wantAttempts: 1},
+		{name: "exhausted budget preserves last error", retries: 1, credentials: 2, wantAttempts: 2, exhausted: true},
+		{name: "output already committed", retries: 2, credentials: 2, wantAttempts: 1, partial: true},
+		{name: "unknown replay safety", retries: 2, credentials: 2, wantAttempts: 1, unknown: true},
+		{name: "continuation stays bound", retries: 2, credentials: 2, wantAttempts: 1, bound: true},
+		{name: "multiplex stays bound", retries: 2, credentials: 2, wantAttempts: 1, multiplex: true},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			id := channel.CLIProxyAPI
+			if tc.multiplex {
+				id = channel.OpenAI
+			}
+			h, engine, input := websocketTestHandler(t, "http://127.0.0.1:1", id)
+			input.SystemSettings = config.Settings{state.SettingRetryCount: tc.retries}
+			registry := h.registry.(*state.CredentialRegistry)
+			entries := []state.CredentialEntry{testCredentialEntry(t, h.encryption, 1, 1, "first-key")}
+			if tc.credentials == 2 {
+				input.Credentials = append(input.Credentials, testCredentialConfig(2, 1))
+				entries = append(entries, testCredentialEntry(t, h.encryption, 2, 1, "second-key"))
+			}
+			if tc.host {
+				backup := input.Groups[0]
+				backup.ID, backup.Name = 2, "backup"
+				input.Groups = append(input.Groups, backup)
+				input.Credentials = append(input.Credentials, testCredentialConfig(3, 2))
+				entries = append(entries, testCredentialEntry(t, h.encryption, 3, 2, "third-key"))
+			}
+			if _, err := h.manager.Publish(input); err != nil {
+				t.Fatal(err)
+			}
+			if err := registry.ReplaceCredentials(entries); err != nil {
+				t.Fatal(err)
+			}
+			sink := &recordingRequestLogSink{}
+			h.requestLogSink = sink
+			var attempted []uint
+			h.forwarder = websocketScriptForwarder{AttemptForwarder: h.forwarder, open: func(_ context.Context, in ForwardInput) (execution.WebsocketSession, execution.WebsocketResult) {
+				session := &websocketScriptSession{done: make(chan struct{})}
+				turns := 0
+				session.turn = func(ctx context.Context, _ []byte, emit func(context.Context, []byte) error) execution.WebsocketResult {
+					turns++
+					if tc.bound && turns == 1 {
+						_ = emit(ctx, websocketCompleted("resp_prior", ""))
+						return execution.WebsocketResult{DispatchState: execution.DispatchMaybeSent}
+					}
+					attempted = append(attempted, in.Credential.ID)
+					if in.Credential.ID != 1 && !tc.exhausted {
+						_ = emit(ctx, websocketCompleted("resp_recovered", ""))
+						return execution.WebsocketResult{DispatchState: execution.DispatchMaybeSent}
+					}
+					if tc.partial {
+						_ = emit(ctx, []byte(`{"type":"response.created","response":{"id":"resp_partial","object":"response","status":"in_progress","model":"upstream"}}`))
+					}
+					status, errorType := 429, "usage_limit_reached"
+					if tc.host {
+						status, errorType = 503, "server_error"
+					}
+					if tc.invalid {
+						status, errorType = 400, "invalid_parameter"
+					}
+					payload := []byte(fmt.Sprintf(`{"type":"error","status":%d,"error":{"type":%q,"message":"quota-%d %s","resets_in_seconds":60}}`, status, errorType, in.Credential.ID, in.APIKey))
+					if err := emit(ctx, payload); err != nil {
+						return execution.WebsocketResult{DispatchState: execution.DispatchMaybeSent, Error: &execution.ErrorEvidence{Kind: execution.ErrorKindCanceled}}
+					}
+					safety := execution.ReplaySafety("")
+					if tc.unknown {
+						safety = execution.ReplaySafetyUnknown
+					}
+					return execution.WebsocketResult{DispatchState: execution.DispatchMaybeSent, Error: &execution.ErrorEvidence{Kind: execution.ErrorKindHTTP, StatusCode: status, Code: "upstream_error", OriginHint: execution.ErrorOriginUpstream, ReplaySafety: safety, Summary: "upstream rejected request"}}
+				}
+				return session, execution.WebsocketResult{DispatchState: execution.DispatchNotSent}
+			}}
+			server := httptest.NewServer(engine)
+			defer server.Close()
+			conn := dialGatewayWebsocket(t, server.URL)
+			defer conn.Close()
+			request := map[string]any{"type": "response.create", "model": "public", "input": "hello", "store": false}
+			if tc.bound {
+				if err := conn.WriteJSON(request); err != nil {
+					t.Fatal(err)
+				}
+				if _, _, err := conn.ReadMessage(); err != nil {
+					t.Fatal(err)
+				}
+				waitWebsocketLogs(t, sink, 1)
+				request["previous_response_id"] = "resp_prior"
+			}
+			if err := conn.WriteJSON(request); err != nil {
+				t.Fatal(err)
+			}
+			_, body, err := conn.ReadMessage()
+			if err != nil {
+				t.Fatal(err)
+			}
+			if tc.partial {
+				if !strings.Contains(string(body), "response.created") {
+					t.Fatalf("missing prior output: %s", body)
+				}
+				_, body, err = conn.ReadMessage()
+				if err != nil {
+					t.Fatal(err)
+				}
+			}
+			count := 1
+			if tc.bound {
+				count++
+			}
+			events := waitWebsocketLogs(t, sink, count)
+			event := events[count-1]
+			if len(attempted) != tc.wantAttempts || len(event.Attempts) != tc.wantAttempts {
+				t.Fatalf("attempted=%v logs=%+v response=%s", attempted, event.Attempts, body)
+			}
+			if registry.ModelCooldowns(1, time.Now())["upstream"].After(time.Now()) != (!tc.host && !tc.invalid) {
+				t.Fatal("unexpected model cooldown")
+			}
+			if tc.host && attempted[1] != 3 {
+				t.Fatalf("failed group was not skipped: %v", attempted)
+			}
+			if until, _ := registry.CredentialCooldownUntil(1); !until.IsZero() {
+				t.Fatal("entire credential cooled")
+			}
+			var response map[string]any
+			if err := json.Unmarshal(body, &response); err != nil {
+				t.Fatal(err)
+			}
+			recovered := tc.wantAttempts == 2 && !tc.exhausted
+			if recovered {
+				if response["type"] != "response.completed" || !strings.Contains(string(body), "resp_recovered") {
+					t.Fatalf("error leaked instead of recovery: %s", body)
+				}
+				if !event.Attempts[0].WillRetry || event.Attempts[0].Committed || attempted[0] == attempted[1] {
+					t.Fatalf("retry accounting=%+v", event.Attempts)
+				}
+			} else {
+				if response["type"] != "error" || !strings.Contains(string(body), fmt.Sprintf("quota-%d", attempted[len(attempted)-1])) || strings.Contains(string(body), "first-key") {
+					t.Fatalf("original sanitized error lost: %s", body)
+				}
+			}
+		})
+	}
+}
+
+func TestWebsocketNativeFirstQuotaErrorRetries(t *testing.T) {
+	var mu sync.Mutex
+	var attempts []string
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil)
+		if err != nil {
+			return
+		}
+		defer conn.Close()
+		if _, _, err = conn.ReadMessage(); err != nil {
+			return
+		}
+		key := r.Header.Get("Authorization")
+		mu.Lock()
+		attempts = append(attempts, key)
+		mu.Unlock()
+		body := websocketCompleted("resp_recovered", "")
+		if key == "Bearer first-key" {
+			body = []byte(`{"type":"error","status":429,"error":{"type":"usage_limit_reached","message":"quota exhausted"}}`)
+		}
+		_ = conn.WriteMessage(websocket.TextMessage, body)
+		_, _, _ = conn.ReadMessage()
+	}))
+	defer upstream.Close()
+	h, engine, input := websocketTestHandler(t, upstream.URL+"/v1", channel.CLIProxyAPI)
+	input.Credentials = append(input.Credentials, testCredentialConfig(2, 1))
+	if _, err := h.manager.Publish(input); err != nil {
+		t.Fatal(err)
+	}
+	registry := h.registry.(*state.CredentialRegistry)
+	if err := registry.ReplaceCredentials([]state.CredentialEntry{testCredentialEntry(t, h.encryption, 1, 1, "first-key"), testCredentialEntry(t, h.encryption, 2, 1, "second-key")}); err != nil {
+		t.Fatal(err)
+	}
+	sink := &recordingRequestLogSink{}
+	h.requestLogSink = sink
+	server := httptest.NewServer(engine)
+	defer server.Close()
+	conn := dialGatewayWebsocket(t, server.URL)
+	defer conn.Close()
+	if err := conn.WriteJSON(map[string]any{"type": "response.create", "model": "public", "input": "hello", "store": false}); err != nil {
+		t.Fatal(err)
+	}
+	_, body, err := conn.ReadMessage()
+	if err != nil {
+		t.Fatal(err)
+	}
+	logs := waitWebsocketLogs(t, sink, 1)
+	mu.Lock()
+	defer mu.Unlock()
+	if strings.Join(attempts, ",") != "Bearer first-key,Bearer second-key" || !strings.Contains(string(body), "resp_recovered") {
+		t.Fatalf("attempts=%v response=%s", attempts, body)
+	}
+	if len(logs[0].Attempts) != 2 || logs[0].Attempts[0].Committed || !logs[0].Attempts[0].WillRetry {
+		t.Fatalf("logs=%+v", logs)
+	}
+}

internal/gateway/websocket_test.go

@@ -861,12 +861,15 @@ func TestWebsocketRetriesOnlyUnsentUnboundTurns(t *testing.T) {
 	for _, dispatch := range []execution.DispatchState{execution.DispatchNotSent, execution.DispatchMaybeSent} {
 		t.Run(string(dispatch), func(t *testing.T) {
 			h, engine, input := websocketTestHandler(t, "http://127.0.0.1:1", channel.CLIProxyAPI)
-			input.Credentials = append(input.Credentials, testCredentialConfig(2, 1))
+			backup := input.Groups[0]
+			backup.ID, backup.Name = 2, "backup"
+			input.Groups = append(input.Groups, backup)
+			input.Credentials = append(input.Credentials, testCredentialConfig(2, 2))
 			if _, err := h.manager.Publish(input); err != nil {
 				t.Fatal(err)
 			}
 			registry := h.registry.(*state.CredentialRegistry)
-			if err := registry.ReplaceCredentials([]state.CredentialEntry{testCredentialEntry(t, h.encryption, 1, 1, "first-key"), testCredentialEntry(t, h.encryption, 2, 1, "second-key")}); err != nil {
+			if err := registry.ReplaceCredentials([]state.CredentialEntry{testCredentialEntry(t, h.encryption, 1, 1, "first-key"), testCredentialEntry(t, h.encryption, 2, 2, "second-key")}); err != nil {
 				t.Fatal(err)
 			}
 			sink := &recordingRequestLogSink{}

internal/gateway/websocket_turn.go

@@ -245,6 +245,7 @@ func (s *websocketConnection) executeTurn(turn websocketTurn) {
 	var refreshSelection *scheduler.Selection
 	var refreshRef state.CredentialRef
 	authRefreshUsed := false
+	var finishRejectedAttempt func()
 	for sequence := 1; sequence <= limit; sequence++ {
 		if s.ctx.Err() != nil {
 			recorder.completeCanceled(s.ctx, 0, -1)
@@ -272,6 +273,10 @@ func (s *websocketConnection) executeTurn(turn websocketTurn) {
 		} else {
 			selection, err = iterator.Next()
 			if err != nil {
+				if finishRejectedAttempt != nil {
+					finishRejectedAttempt()
+					return
+				}
 				reject(reasonNoCandidate)
 				return
 			}
@@ -418,7 +423,8 @@ func (s *websocketConnection) executeTurn(turn websocketTurn) {
 			result.Err = s.ctx.Err()
 			result.ExecutionError = &execution.ErrorEvidence{Kind: execution.ErrorKindCanceled, OriginHint: execution.ErrorOriginDownstream, Code: "websocket_canceled"}
 		} else if binding != nil {
-			result = s.runWebsocketAttempt(ctx, cancel, binding, turn.lane, selection, ref, input, recorder, unlock, firstByteDeadline)
+			bufferFirstError := newBinding && requiredRef == nil && !binding.capabilities.Multiplex
+			result = s.runWebsocketAttempt(ctx, cancel, binding, turn.lane, selection, ref, input, recorder, unlock, firstByteDeadline, bufferFirstError)
 		} else {
 			result.Err = executionFailureError(ctx, wsResult.Error)
 			result.StatusCode = wsResult.Error.StatusCode
@@ -434,15 +440,30 @@ func (s *websocketConnection) executeTurn(turn websocketTurn) {
 		decision := judgeUpstreamResult(result, h.now(), health.DecisionContext{DefaultRateLimitCooldown: fixedCooldown, CredentialRefreshable: selection.Group.ConnectionType == "subscription", Method: http.MethodPost, Operation: execution.OperationResponsesCreate})
 		index := recorder.recordStreamAttempt(selection, credential.secrets, result, decision, started, recorder.now())
 		h.applyGroupDecisionEffect(selection.Group, ref, 0, decision, result.StatusCode, h.now(), input.UpstreamModelID)
+		if decision.Effect == health.EffectSkipGroup {
+			iterator.SkipGroup(selection.GroupID)
+		}
+		if len(result.Body) > 0 {
+			finishRejectedAttempt = func() {
+				unlock()
+				if err := s.emit(s.ctx, result.Body); err != nil {
+					recorder.completeCanceled(s.ctx, result.StatusCode, index)
+					return
+				}
+				recorder.completeStream(result, input.UpstreamModelID, index)
+			}
+		}
 		if result.Stream.EndReason == StreamEndCleanEOF {
 			h.recordCredentialSuccess(ref, h.now())
 			if requiredRef == nil {
 				h.recordAffinitySuccess(affinity, selection, ref)
 			}
 		}
-		// 延迟拨号的串行 Session 若首轮确实未发送，可释放失败句柄后继续既有候选迭代。
-		// 多流一旦共享连接，不能以某一轮未发送推断其他轮也未发送。
-		if result.DispatchState == execution.DispatchNotSent && !result.Committed && newBinding &&
+		// 未绑定的串行首轮允许重试未发送失败或尚未交付的上游错误事件。
+		// 是否可重放仍由错误规则决定；续接和共享连接不能随某一轮切换身份。
+		retryableTurn := result.DispatchState == execution.DispatchNotSent ||
+			(len(result.Body) > 0 && result.Stream.EndReason == StreamEndSSEError)
+		if retryableTurn && !result.Committed && newBinding &&
 			(binding == nil || !binding.capabilities.Multiplex) && decision.Retry != health.RetryNone &&
 			sequence < limit && requiredRef == nil && !authRefreshUsed && s.ctx.Err() == nil {
 			if decision.Retry == health.RetryRefreshCredential {
@@ -462,6 +483,8 @@ func (s *websocketConnection) executeTurn(turn websocketTurn) {
 			recorder.completeStream(result, input.UpstreamModelID, index)
 		} else if s.ctx.Err() != nil {
 			recorder.completeCanceled(s.ctx, result.StatusCode, index)
+		} else if len(result.Body) > 0 {
+			finishRejectedAttempt()
 		} else {
 			value := reasonUpstreamConnect
 			if result.ExecutionError != nil {
@@ -536,7 +559,7 @@ type websocketCancelCloser struct {
 
 func (c websocketCancelCloser) Close() error { c.timedOut.Store(true); c.cancel(); return nil }
 
-func (s *websocketConnection) runWebsocketAttempt(ctx context.Context, cancel context.CancelFunc, binding *websocketBinding, lane string, selection scheduler.Selection, ref state.CredentialRef, input ForwardInput, recorder *requestRecorder, unlock func(), firstByteDeadline time.Time) UpstreamResult {
+func (s *websocketConnection) runWebsocketAttempt(ctx context.Context, cancel context.CancelFunc, binding *websocketBinding, lane string, selection scheduler.Selection, ref state.CredentialRef, input ForwardInput, recorder *requestRecorder, unlock func(), firstByteDeadline time.Time, bufferFirstError bool) UpstreamResult {
 	observer := newStreamEventObserver(input.Dialect, newUsageCaptureBoundary().newStreamForRequest(input.Dialect, input.ObserveUsage))
 	result := UpstreamResult{UpstreamProtocol: protocol.OpenAIResponses}
 	var responseID string
@@ -644,6 +667,23 @@ func (s *websocketConnection) runWebsocketAttempt(ctx context.Context, cancel co
 				}
 			}
 		}
+		if responsealias.Needs(input.ExternalModel, input.UpstreamModelID) {
+			body, err = responsealias.RewriteJSON(protocol.OpenAIResponses, body, input.ExternalModel)
+			if err != nil {
+				return err
+			}
+		}
+		if response.Model != "" {
+			result.UpstreamReportedModel = response.Model
+			result.ResponseModelObserved = true
+			result.ResponseModelMismatch = response.Model != input.UpstreamModelID
+		}
+		if providerError && !result.Committed && bufferFirstError {
+			// 先保留已脱敏的错误，待分类及候选选择结束后再决定是否交付。
+			// 失败响应不能登记为后续请求的续接目标。
+			result.Body = append([]byte(nil), body...)
+			return nil
+		}
 		if len(event.Response) > 0 {
 			if response.ID != "" {
 				if len(response.ID) > 4096 {
@@ -665,11 +705,6 @@ func (s *websocketConnection) runWebsocketAttempt(ctx context.Context, cancel co
 				}
 				s.mu.Unlock()
 			}
-			if response.Model != "" {
-				result.UpstreamReportedModel = response.Model
-				result.ResponseModelObserved = true
-				result.ResponseModelMismatch = response.Model != input.UpstreamModelID
-			}
 			if !providerError && binding.capabilities.StoredResponses && onResponse != nil {
 				if err := onResponse(event.Response); err != nil {
 					return err
@@ -679,12 +714,6 @@ func (s *websocketConnection) runWebsocketAttempt(ctx context.Context, cancel co
 		if !result.Committed {
 			recorder.recordFirstResponse()
 		}
-		if responsealias.Needs(input.ExternalModel, input.UpstreamModelID) {
-			body, err = responsealias.RewriteJSON(protocol.OpenAIResponses, body, input.ExternalModel)
-			if err != nil {
-				return err
-			}
-		}
 		unlock()
 		if err = s.emit(ctx, body); err != nil {
 			return err
@@ -729,6 +758,13 @@ func (s *websocketConnection) runWebsocketAttempt(ctx context.Context, cancel co
 		}
 	}
 	knownUpstreamError := wsResult.Error != nil && (wsResult.Error.Kind == execution.ErrorKindHTTP || wsResult.Error.Kind == execution.ErrorKindProvider)
+	if len(result.Body) > 0 && knownUpstreamError {
+		replaySafety := result.ExecutionError.ReplaySafety
+		result.ExecutionError = firstStreamErrorEvidence(result.Body, result.ExecutionError,
+			result.StatusCode, "", result.ErrorSummary, recorder.redactor, input.CredentialSecrets)
+		// 仅复用错误字段解析，不引入 HTTP 启动缓冲特有的容量拒绝证明。
+		result.ExecutionError.ReplaySafety = replaySafety
+	}
 	if !observer.terminalForwarded && !knownUpstreamError && (timedOut.Load() || errors.Is(ctx.Err(), context.DeadlineExceeded)) {
 		result.Stream = streamTerminalObservation(StreamEndIdleTimeout)
 		result.Err = upstreamExecutionTimeoutError{}