diff --git a/internal/execution/wsnative/session.go b/internal/execution/wsnative/session.go index 77369afe..afa095ee 100644 --- a/internal/execution/wsnative/session.go +++ b/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 = "" } diff --git a/internal/gateway/websocket_retry_test.go b/internal/gateway/websocket_retry_test.go new file mode 100644 index 00000000..5cad9ada --- /dev/null +++ b/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) + } +} diff --git a/internal/gateway/websocket_test.go b/internal/gateway/websocket_test.go index be4468d9..90b63710 100644 --- a/internal/gateway/websocket_test.go +++ b/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{} diff --git a/internal/gateway/websocket_turn.go b/internal/gateway/websocket_turn.go index 0e114af6..c1dd5eab 100644 --- a/internal/gateway/websocket_turn.go +++ b/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{}