From de349187d91f5d458199d332aa90966d732bd55d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=99=BD=E5=AE=A6=E6=88=90?= Date: Fri, 7 Aug 2026 13:11:08 +0800 Subject: [PATCH] fix(openai): harden priority routing hints --- .../openai_compact_service_tier_test.go | 100 ++++++ .../service/openai_gateway_forward.go | 1 + .../service/openai_gateway_passthrough.go | 36 +++ .../service/openai_gateway_request_body.go | 1 + .../service/openai_gateway_service_test.go | 4 +- .../internal/service/openai_routing_hint.go | 82 ++++- .../service/openai_routing_hint_test.go | 190 +++++++++++- .../service/openai_ws_forwarder_payload.go | 9 + backend/internal/service/openai_ws_pool.go | 181 +++++++++-- .../internal/service/openai_ws_pool_test.go | 285 ++++++++++++++++++ 10 files changed, 850 insertions(+), 39 deletions(-) create mode 100644 backend/internal/service/openai_compact_service_tier_test.go diff --git a/backend/internal/service/openai_compact_service_tier_test.go b/backend/internal/service/openai_compact_service_tier_test.go new file mode 100644 index 0000000000..7e92d8d2c2 --- /dev/null +++ b/backend/internal/service/openai_compact_service_tier_test.go @@ -0,0 +1,100 @@ +package service + +import ( + "bytes" + "context" + "net/http" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" +) + +func TestNormalizeOpenAICompactRequestBodyPreservesServiceTier(t *testing.T) { + body := []byte(`{ + "model":"gpt-5.6-sol", + "input":[{"type":"message","role":"user","content":"hello"}], + "service_tier":"priority", + "prompt_cache_key":"compact-cache-key", + "store":false, + "stream":true + }`) + + normalized, changed, err := normalizeOpenAICompactRequestBody(body) + require.NoError(t, err) + require.True(t, changed) + require.Equal(t, "gpt-5.6-sol", gjson.GetBytes(normalized, "model").String()) + require.Equal(t, "priority", gjson.GetBytes(normalized, "service_tier").String()) + require.False(t, gjson.GetBytes(normalized, "prompt_cache_key").Exists()) + require.False(t, gjson.GetBytes(normalized, "store").Exists()) + require.False(t, gjson.GetBytes(normalized, "stream").Exists()) +} + +func TestOpenAIOAuthCompactHTTPBuildersUsePreservedServiceTierInRoutingHint(t *testing.T) { + gin.SetMode(gin.TestMode) + body := []byte(`{ + "model":"gpt-5.6-sol", + "input":[{"type":"message","role":"user","content":"hello"}], + "service_tier":"priority", + "stream":true + }`) + normalized, changed, err := normalizeOpenAICompactRequestBody(body) + require.NoError(t, err) + require.True(t, changed) + require.Equal(t, "priority", gjson.GetBytes(normalized, "service_tier").String()) + + account := &Account{ + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Credentials: map[string]any{ + "chatgpt_account_id": "test-account", + }, + } + svc := &OpenAIGatewayService{} + + tests := []struct { + name string + build func(*gin.Context) (*http.Request, error) + }{ + { + name: "ordinary", + build: func(c *gin.Context) (*http.Request, error) { + return svc.buildUpstreamRequest( + context.Background(), c, account, normalized, "test-token", + false, "", true, + ) + }, + }, + { + name: "passthrough", + build: func(c *gin.Context) (*http.Request, error) { + return svc.buildUpstreamRequestOpenAIPassthrough( + context.Background(), c, account, normalized, "test-token", + ) + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest( + http.MethodPost, + "/v1/responses/compact", + bytes.NewReader(normalized), + ) + + req, buildErr := tt.build(c) + require.NoError(t, buildErr) + require.Equal( + t, + "model=gpt-5.6-sol;tier=priority", + req.Header.Get(openAICodexRoutingHintHeader), + ) + require.Equal(t, "priority", gjson.GetBytes(normalized, "service_tier").String()) + }) + } +} diff --git a/backend/internal/service/openai_gateway_forward.go b/backend/internal/service/openai_gateway_forward.go index 41dd5d4956..b700eb39f7 100644 --- a/backend/internal/service/openai_gateway_forward.go +++ b/backend/internal/service/openai_gateway_forward.go @@ -1111,6 +1111,7 @@ func (s *OpenAIGatewayService) buildUpstreamRequest(ctx context.Context, c *gin. // 账号级请求头覆写（仅 openai api_key 账号启用时生效；OAuth 路径 no-op） account.ApplyHeaderOverrides(req.Header) setOpenAICodexRoutingHintFromBody(req.Header, account, body) + logOpenAIRoutingDiagnosticsFromBody(ctx, account, "http", req.Header, body, "not_applicable") return req, nil } diff --git a/backend/internal/service/openai_gateway_passthrough.go b/backend/internal/service/openai_gateway_passthrough.go index 1592fd41da..fa1b46a3d3 100644 --- a/backend/internal/service/openai_gateway_passthrough.go +++ b/backend/internal/service/openai_gateway_passthrough.go @@ -390,6 +390,10 @@ func (s *OpenAIGatewayService) buildUpstreamRequestOpenAIPassthrough( // OAuth 透传到 ChatGPT internal API 时补齐必要头。 if account.Type == AccountTypeOAuth { + // Current Codex OAuth HTTP no longer negotiates the legacy Responses + // experiment. Passthrough may receive it from an older client, so remove + // only that token while preserving any independent beta negotiation. + stripOpenAILegacyResponsesBeta(req.Header) promptCacheKey := strings.TrimSpace(gjson.GetBytes(body, "prompt_cache_key").String()) req.Host = "chatgpt.com" if err := resolveAndSetOpenAIChatGPTAccountHeaders(ctx, s.accountRepo, req.Header, account); err != nil { @@ -454,10 +458,42 @@ func (s *OpenAIGatewayService) buildUpstreamRequestOpenAIPassthrough( // 账号级请求头覆写（仅 openai api_key 账号启用时生效；OAuth 路径 no-op） account.ApplyHeaderOverrides(req.Header) setOpenAICodexRoutingHintFromBody(req.Header, account, body) + logOpenAIRoutingDiagnosticsFromBody(ctx, account, "http_passthrough", req.Header, body, "not_applicable") return req, nil } +func stripOpenAILegacyResponsesBeta(headers http.Header) { + if headers == nil { + return + } + + preserved := make([]string, 0) + for key, values := range headers { + if !strings.EqualFold(strings.TrimSpace(key), "OpenAI-Beta") { + continue + } + delete(headers, key) + for _, value := range values { + parts := strings.Split(value, ",") + kept := parts[:0] + for _, part := range parts { + part = strings.TrimSpace(part) + if part == "" || strings.EqualFold(part, "responses=experimental") { + continue + } + kept = append(kept, part) + } + if len(kept) > 0 { + preserved = append(preserved, strings.Join(kept, ", ")) + } + } + } + for _, value := range preserved { + headers.Add("OpenAI-Beta", value) + } +} + func shouldFailoverOpenAIPassthroughResponse(account *Account, statusCode int, responseBody []byte) bool { if isOpenAIContextWindowError("", responseBody) { return false diff --git a/backend/internal/service/openai_gateway_request_body.go b/backend/internal/service/openai_gateway_request_body.go index 2bb586ce52..04b3e7e050 100644 --- a/backend/internal/service/openai_gateway_request_body.go +++ b/backend/internal/service/openai_gateway_request_body.go @@ -308,6 +308,7 @@ func normalizeOpenAICompactRequestBody(body []byte) ([]byte, bool, error) { "tools", "parallel_tool_calls", "reasoning", + "service_tier", "text", "previous_response_id", } { diff --git a/backend/internal/service/openai_gateway_service_test.go b/backend/internal/service/openai_gateway_service_test.go index c821a48639..48d1808245 100644 --- a/backend/internal/service/openai_gateway_service_test.go +++ b/backend/internal/service/openai_gateway_service_test.go @@ -2888,7 +2888,7 @@ func TestOpenAIBuildUpstreamRequestOpenAIPassthroughPreservesCompactPath(t *test require.Equal(t, chatgptCodexURL+"/compact", req.URL.String()) require.Equal(t, "application/json", req.Header.Get("Accept")) require.Equal(t, codexCLIVersion, req.Header.Get("Version")) - require.Empty(t, req.Header.Get("OpenAI-Beta"), "Codex 0.147.0 OAuth HTTP must not synthesize the legacy responses beta header") + require.Empty(t, req.Header.Get("OpenAI-Beta"), "Codex OAuth HTTP must not synthesize the legacy responses beta header") require.NotEmpty(t, req.Header.Get("Session_Id")) require.Equal(t, HTTPUpstreamProfileOpenAI, HTTPUpstreamProfileFromContext(req.Context())) } @@ -2929,7 +2929,7 @@ func TestOpenAIBuildUpstreamRequestCompactForcesJSONAcceptForOAuth(t *testing.T) require.Equal(t, chatgptCodexURL+"/compact", req.URL.String()) require.Equal(t, "application/json", req.Header.Get("Accept")) require.Equal(t, codexCLIVersion, req.Header.Get("Version")) - require.Empty(t, req.Header.Get("OpenAI-Beta"), "Codex 0.147.0 OAuth HTTP must not synthesize the legacy responses beta header") + require.Empty(t, req.Header.Get("OpenAI-Beta"), "Codex OAuth HTTP must not synthesize the legacy responses beta header") require.NotEmpty(t, req.Header.Get("Session_Id")) require.Equal(t, HTTPUpstreamProfileOpenAI, HTTPUpstreamProfileFromContext(req.Context())) } diff --git a/backend/internal/service/openai_routing_hint.go b/backend/internal/service/openai_routing_hint.go index e8c35ed412..0a1cc928cd 100644 --- a/backend/internal/service/openai_routing_hint.go +++ b/backend/internal/service/openai_routing_hint.go @@ -1,10 +1,13 @@ package service import ( + "context" "net/http" "strings" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" "github.com/tidwall/gjson" + "go.uber.org/zap" "golang.org/x/net/http/httpguts" ) @@ -14,15 +17,21 @@ const openAICodexRoutingHintHeader = "x-codex-routing-hint" // OpenAI OAuth requests. The request model must already be the final upstream // slug and serviceTier must already reflect any local policy rewrite/filter. func setOpenAICodexRoutingHint(headers http.Header, account *Account, model string, serviceTier string) { - if headers == nil || account == nil || !account.IsOpenAIOAuth() { + if headers == nil { + return + } + + // The routing hint is gateway-owned. Strip every spelling before deciding + // whether to synthesize it so API-key/provider credential paths cannot pass + // through a caller- or account-override-supplied hint. http.Header.Del only + // removes the canonical map key; inbound maps can contain raw lowercase keys. + deleteOpenAIHeaderEqualFold(headers, openAICodexRoutingHintHeader) + if account == nil || !account.IsOpenAIOAuth() { return } - // OAuth headers are built afresh for every request/dial, but delete first so - // callers updating a cached header set cannot retain a stale routing hint. - headers.Del(openAICodexRoutingHintHeader) model = strings.TrimSpace(model) - if model == "" { + if model == "" || strings.ContainsAny(model, ";=") { return } @@ -50,7 +59,70 @@ func setOpenAICodexRoutingHint(headers http.Header, account *Account, model stri headers.Set(openAICodexRoutingHintHeader, hint) } +func deleteOpenAIHeaderEqualFold(headers http.Header, name string) { + if headers == nil { + return + } + name = strings.TrimSpace(name) + for key := range headers { + if strings.EqualFold(strings.TrimSpace(key), name) { + delete(headers, key) + } + } +} + func setOpenAICodexRoutingHintFromBody(headers http.Header, account *Account, body []byte) { fields := gjson.GetManyBytes(body, "model", "service_tier") setOpenAICodexRoutingHint(headers, account, fields[0].String(), fields[1].String()) } + +// logOpenAIRoutingDiagnostics records only gateway-derived routing state. In +// particular, it deliberately does not include any header values, tokens, or +// credentials because these diagnostics run on authentication-bearing paths. +func logOpenAIRoutingDiagnostics( + ctx context.Context, + account *Account, + transport string, + model string, + serviceTier string, + hintGenerated bool, + wsAffinityDecision string, +) { + if ctx == nil { + ctx = context.Background() + } + accountID := int64(0) + if account != nil { + accountID = account.ID + } + + logger.FromContext(ctx).Debug("openai routing decision", + zap.String("component", "service.openai_routing"), + zap.String("transport", strings.TrimSpace(transport)), + zap.Int64("account_id", accountID), + zap.String("final_model", strings.TrimSpace(model)), + zap.String("final_service_tier", normalizedOpenAIServiceTierValue(serviceTier)), + zap.Bool("routing_hint_generated", hintGenerated), + zap.String("ws_affinity_decision", strings.TrimSpace(wsAffinityDecision)), + ) +} + +func logOpenAIRoutingDiagnosticsFromBody( + ctx context.Context, + account *Account, + transport string, + headers http.Header, + body []byte, + wsAffinityDecision string, +) { + fields := gjson.GetManyBytes(body, "model", "service_tier") + logOpenAIRoutingDiagnostics( + ctx, + account, + transport, + fields[0].String(), + fields[1].String(), + strings.TrimSpace(headers.Get(openAICodexRoutingHintHeader)) != "", + wsAffinityDecision, + ) +} diff --git a/backend/internal/service/openai_routing_hint_test.go b/backend/internal/service/openai_routing_hint_test.go index a7b56bcb7d..5801b15fa4 100644 --- a/backend/internal/service/openai_routing_hint_test.go +++ b/backend/internal/service/openai_routing_hint_test.go @@ -5,6 +5,7 @@ import ( "context" "net/http" "net/http/httptest" + "strings" "testing" "github.com/Wei-Shaw/sub2api/internal/config" @@ -44,11 +45,30 @@ func TestSetOpenAICodexRoutingHintCanonicalizesOfficialServiceTiers(t *testing.T require.Empty(t, headers.Get(openAICodexRoutingHintHeader)) }) - t.Run("api key is untouched", func(t *testing.T) { + for _, model := range []string{"gpt-5.6;evil", "gpt=5.6"} { + t.Run("delimiter in model is omitted: "+model, func(t *testing.T) { + headers := make(http.Header) + setOpenAICodexRoutingHint(headers, oauthAccount, model, "priority") + require.Empty(t, headers.Get(openAICodexRoutingHintHeader)) + }) + } + + t.Run("api key strips gateway-owned hint in every key casing", func(t *testing.T) { headers := make(http.Header) - headers.Set(openAICodexRoutingHintHeader, "caller-owned") + headers[openAICodexRoutingHintHeader] = []string{"lowercase-spoof"} + headers["X-Codex-Routing-Hint"] = []string{"canonical-spoof"} setOpenAICodexRoutingHint(headers, &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}, "gpt-5.6", "priority") - require.Equal(t, "caller-owned", headers.Get(openAICodexRoutingHintHeader)) + for key := range headers { + require.False(t, strings.EqualFold(key, openAICodexRoutingHintHeader)) + } + }) + + t.Run("oauth replaces spoofed lowercase hint", func(t *testing.T) { + headers := make(http.Header) + headers[openAICodexRoutingHintHeader] = []string{"model=spoof;tier=flex"} + setOpenAICodexRoutingHint(headers, oauthAccount, "gpt-5.6", "priority") + require.Equal(t, "model=gpt-5.6;tier=priority", headers.Get(openAICodexRoutingHintHeader)) + require.Len(t, headers, 1) }) } @@ -101,6 +121,67 @@ func TestOpenAIOAuthHTTPBuildersSendRoutingHintFromFinalBody(t *testing.T) { } } +func TestOpenAIHTTPPassthroughStripsOnlyOAuthLegacyResponsesBeta(t *testing.T) { + gin.SetMode(gin.TestMode) + svc := &OpenAIGatewayService{cfg: &config.Config{ + Security: config.SecurityConfig{ + URLAllowlist: config.URLAllowlistConfig{Enabled: false}, + }, + }} + body := []byte(`{"model":"gpt-5.6-codex","service_tier":"priority"}`) + + build := func(t *testing.T, account *Account, betaValues []string, rawLowercaseKey bool) http.Header { + t.Helper() + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body)) + if rawLowercaseKey { + c.Request.Header["openai-beta"] = append([]string(nil), betaValues...) + } else { + for _, value := range betaValues { + c.Request.Header.Add("OpenAI-Beta", value) + } + } + + req, err := svc.buildUpstreamRequestOpenAIPassthrough(context.Background(), c, account, body, "test-token") + require.NoError(t, err) + return req.Header + } + + oauth := &Account{ + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Credentials: map[string]any{ + "chatgpt_account_id": "test-account", + }, + } + apiKey := &Account{ + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Credentials: map[string]any{ + "api_key": "test-api-key", + }, + } + + t.Run("oauth legacy only is removed including raw lowercase key", func(t *testing.T) { + headers := build(t, oauth, []string{"responses=experimental"}, true) + require.Empty(t, headers.Values("OpenAI-Beta")) + }) + + t.Run("oauth mixed beta preserves independent tokens", func(t *testing.T) { + headers := build(t, oauth, []string{ + "responses=experimental, future_feature=v1", + "another_feature=v2, RESPONSES=EXPERIMENTAL", + }, false) + require.Equal(t, []string{"future_feature=v1", "another_feature=v2"}, headers.Values("OpenAI-Beta")) + }) + + t.Run("api key explicit beta remains caller controlled", func(t *testing.T) { + headers := build(t, apiKey, []string{"responses=experimental, future_feature=v1"}, false) + require.Equal(t, []string{"responses=experimental, future_feature=v1"}, headers.Values("OpenAI-Beta")) + }) +} + func TestBuildOpenAIWSHeadersSendsOAuthRoutingHintOnly(t *testing.T) { gin.SetMode(gin.TestMode) recorder := httptest.NewRecorder() @@ -139,7 +220,98 @@ func TestBuildOpenAIWSHeadersSendsOAuthRoutingHintOnly(t *testing.T) { require.Empty(t, build(t, &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}, "priority").Get(openAICodexRoutingHintHeader)) } -func TestOpenAIWSConnPoolDoesNotReuseDifferentRoutingHints(t *testing.T) { +func TestOpenAIRoutingDiagnosticsUseFinalDerivedValuesOnly(t *testing.T) { + gin.SetMode(gin.TestMode) + logSink, restore := captureStructuredLog(t) + defer restore() + + account := &Account{ + ID: 917, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Credentials: map[string]any{ + "chatgpt_account_id": "chatgpt-account", + }, + } + body := []byte(`{"model":"gpt-5.6-codex","service_tier":"fast"}`) + svc := &OpenAIGatewayService{} + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body)) + c.Request.Header.Set("Authorization", "Bearer caller-secret") + c.Request.Header.Set(openAICodexRoutingHintHeader, "model=caller-secret") + _, err := svc.buildUpstreamRequest(context.Background(), c, account, body, "oauth-secret", false, "", true) + require.NoError(t, err) + + decision := OpenAIWSProtocolDecision{Transport: OpenAIUpstreamTransportResponsesWebsocketV2} + _, _, err = svc.buildOpenAIWSHeaders( + context.Background(), c, account, "oauth-secret", decision, true, + "", "", "", "gpt-5.6-codex", "fast", + ) + require.NoError(t, err) + + require.True(t, logSink.ContainsMessageAtLevel("openai routing decision", "debug")) + require.True(t, logSink.ContainsFieldValue("account_id", "917")) + require.True(t, logSink.ContainsFieldValue("final_model", "gpt-5.6-codex")) + require.True(t, logSink.ContainsFieldValue("final_service_tier", "priority")) + require.True(t, logSink.ContainsFieldValue("routing_hint_generated", "true")) + require.True(t, logSink.ContainsFieldValue("transport", "http")) + require.True(t, logSink.ContainsFieldValue("transport", string(OpenAIUpstreamTransportResponsesWebsocketV2))) + require.True(t, logSink.ContainsFieldValue("ws_affinity_decision", "not_applicable")) + require.True(t, logSink.ContainsFieldValue("ws_affinity_decision", "soft_routing_hint")) + require.False(t, logSink.ContainsFieldValue("authorization", "caller-secret")) + require.False(t, logSink.ContainsFieldValue("credentials", "oauth-secret")) + require.False(t, logSink.ContainsFieldValue("routing_hint", "caller-secret")) +} + +func TestOpenAIWSConnPoolPreferredContinuationIgnoresRoutingHintChanges(t *testing.T) { + cfg := &config.Config{} + cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2 + cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0 + cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 2 + + pool := newOpenAIWSConnPool(cfg) + dialer := &openAIWSCountingDialer{} + pool.setClientDialerForTest(dialer) + account := &Account{ID: 913, Platform: PlatformOpenAI, Type: AccountTypeOAuth} + + acquire := func(t *testing.T, hint, preferred string, forcePreferred bool) *openAIWSConnLease { + t.Helper() + headers := make(http.Header) + if hint != "" { + headers.Set(openAICodexRoutingHintHeader, hint) + } + lease, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{ + Account: account, + WSURL: "wss://example.com/v1/responses", + Headers: headers, + PreferredConnID: preferred, + ForcePreferredConn: forcePreferred, + }) + require.NoError(t, err) + require.NotNil(t, lease) + return lease + } + + standard := acquire(t, "model=gpt-5.6-codex", "", false) + connID := standard.ConnID() + standard.Release() + + priority := acquire(t, "model=gpt-5.6-codex;tier=priority", connID, true) + require.True(t, priority.Reused()) + require.Equal(t, connID, priority.ConnID()) + priority.Release() + + standardAgain := acquire(t, "model=gpt-5.6-codex", connID, true) + require.True(t, standardAgain.Reused()) + require.Equal(t, connID, standardAgain.ConnID()) + standardAgain.Release() + + require.Equal(t, 1, dialer.DialCount(), "routing hint is dial-time advisory, not continuation compatibility") +} + +func TestOpenAIWSConnPoolUsesRoutingHintAsSoftDialAffinity(t *testing.T) { cfg := &config.Config{} cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 4 cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0 @@ -166,16 +338,6 @@ func TestOpenAIWSConnPoolDoesNotReuseDifferentRoutingHints(t *testing.T) { priority := acquire(t, "model=gpt-5.6-codex;tier=priority") priorityConnID := priority.ConnID() priority.Release() - flexHeaders := make(http.Header) - flexHeaders.Set(openAICodexRoutingHintHeader, "model=gpt-5.6-codex;tier=flex") - _, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{ - Account: account, - WSURL: "wss://example.com/v1/responses", - Headers: flexHeaders, - PreferredConnID: priorityConnID, - ForcePreferredConn: true, - }) - require.ErrorIs(t, err, errOpenAIWSPreferredConnUnavailable) priorityAgain := acquire(t, "model=gpt-5.6-codex;tier=priority") require.True(t, priorityAgain.Reused()) diff --git a/backend/internal/service/openai_ws_forwarder_payload.go b/backend/internal/service/openai_ws_forwarder_payload.go index bdd8776203..c53d9e8241 100644 --- a/backend/internal/service/openai_ws_forwarder_payload.go +++ b/backend/internal/service/openai_ws_forwarder_payload.go @@ -160,6 +160,15 @@ func (s *OpenAIGatewayService) buildOpenAIWSHeaders( // 覆盖所有 WS 模式（ctx_pool/dedicated/passthrough）的握手头。 account.ApplyHeaderOverrides(headers) setOpenAICodexRoutingHint(headers, account, routingModel, routingServiceTier) + logOpenAIRoutingDiagnostics( + ctx, + account, + string(decision.Transport), + routingModel, + routingServiceTier, + strings.TrimSpace(headers.Get(openAICodexRoutingHintHeader)) != "", + "soft_routing_hint", + ) return headers, sessionResolution, nil } diff --git a/backend/internal/service/openai_ws_pool.go b/backend/internal/service/openai_ws_pool.go index 187c5284ea..048031d640 100644 --- a/backend/internal/service/openai_ws_pool.go +++ b/backend/internal/service/openai_ws_pool.go @@ -78,7 +78,6 @@ type openAIWSAcquireRequest struct { type openAIWSHandshakeCompatibilityKey struct { betaFeatures string - routingHint string } type openAIWSConnLease struct { @@ -247,6 +246,7 @@ type openAIWSConn struct { handshakeHeaders http.Header handshakeCompatibility openAIWSHandshakeCompatibilityKey + routingAffinity string leaseCh chan struct{} closedCh chan struct{} @@ -310,6 +310,14 @@ func (c *openAIWSConn) acquire(ctx context.Context) error { case <-c.closedCh: return errOpenAIWSConnClosed case <-c.leaseCh: + // A cancellation and a lease delivery can become ready together. Once + // the semaphore token has been consumed, check the context again and + // return it before reporting cancellation so a canceled waiter cannot + // strand a pooled connection. + if err := ctx.Err(); err != nil { + c.release() + return err + } select { case <-c.closedCh: c.release() @@ -534,6 +542,10 @@ func (c *openAIWSConn) matchesHandshakeCompatibility(compatibility openAIWSHands return c != nil && c.handshakeCompatibility == compatibility } +func (c *openAIWSConn) matchesRoutingAffinity(routingAffinity string) bool { + return c != nil && c.routingAffinity == routingAffinity +} + func (c *openAIWSConn) isPrewarmed() bool { if c == nil { return false @@ -844,6 +856,7 @@ func (p *openAIWSConnPool) acquire(ctx context.Context, req openAIWSAcquireReque retryAcquire: accountID := req.Account.ID compatibility := normalizeOpenAIWSHandshakeCompatibility(req.Headers) + routingAffinity := normalizeOpenAIWSRoutingAffinity(req.Headers) effectiveMaxConns := p.effectiveMaxConnsByAccount(req.Account) if effectiveMaxConns <= 0 { return nil, errOpenAIWSConnQueueFull @@ -851,7 +864,7 @@ retryAcquire: var evicted []*openAIWSConn ap := p.getOrCreateAccountPool(accountID) ap.mu.Lock() - ap.lastAcquire = cloneOpenAIWSAcquireRequestPtr(&req) + acquireGeneration := ap.generation now := time.Now() if ap.lastCleanupAt.IsZero() || now.Sub(ap.lastCleanupAt) >= openAIWSAcquireCleanupInterval { evicted = p.cleanupAccountLocked(ap, now, effectiveMaxConns) @@ -900,6 +913,7 @@ retryAcquire: reused: true, } p.metrics.acquireReuseTotal.Add(1) + p.recordLastSuccessfulAcquire(accountID, acquireGeneration, req) p.ensureTargetIdleAsync(accountID) return lease, nil } @@ -947,6 +961,7 @@ retryAcquire: reused: true, } p.metrics.acquireReuseTotal.Add(1) + p.recordLastSuccessfulAcquire(accountID, acquireGeneration, req) p.ensureTargetIdleAsync(accountID) return lease, nil } @@ -969,12 +984,16 @@ retryAcquire: } lease := &openAIWSConnLease{pool: p, accountID: accountID, conn: conn, connPick: connPick, reused: true} p.metrics.acquireReuseTotal.Add(1) + p.recordLastSuccessfulAcquire(accountID, acquireGeneration, req) p.ensureTargetIdleAsync(accountID) return lease, nil } } - best := p.pickLeastBusyConnLocked(ap, "", compatibility) + // A routing hint is advisory at WebSocket dial time. Prefer a pooled + // connection whose handshake used the same hint, but do not make that + // preference a continuation compatibility requirement. + best := p.pickLeastBusyConnWithRoutingAffinityLocked(ap, compatibility, routingAffinity) if best != nil && best.tryAcquire() { connPick := time.Since(pickStartedAt) p.recordConnPickDuration(connPick) @@ -992,11 +1011,12 @@ retryAcquire: } lease := &openAIWSConnLease{pool: p, accountID: accountID, conn: best, connPick: connPick, reused: true} p.metrics.acquireReuseTotal.Add(1) + p.recordLastSuccessfulAcquire(accountID, acquireGeneration, req) p.ensureTargetIdleAsync(accountID) return lease, nil } for _, conn := range ap.conns { - if conn == nil || conn == best || !conn.matchesHandshakeCompatibility(compatibility) { + if conn == nil || conn == best || !conn.matchesHandshakeCompatibility(compatibility) || !conn.matchesRoutingAffinity(routingAffinity) { continue } if conn.tryAcquire() { @@ -1016,6 +1036,7 @@ retryAcquire: } lease := &openAIWSConnLease{pool: p, accountID: accountID, conn: conn, connPick: connPick, reused: true} p.metrics.acquireReuseTotal.Add(1) + p.recordLastSuccessfulAcquire(accountID, acquireGeneration, req) p.ensureTargetIdleAsync(accountID) return lease, nil } @@ -1023,12 +1044,18 @@ retryAcquire: } if !req.ForceNewConn && len(ap.conns)+ap.creating >= effectiveMaxConns { - compatible := p.pickLeastBusyConnLocked(ap, "", compatibility) - if idle := p.pickOldestIdleConnWithDifferentHandshakeCompatibilityLocked(ap, compatibility); idle != nil { + affine := p.pickLeastBusyConnWithRoutingAffinityLocked(ap, compatibility, routingAffinity) + if idle := p.pickOldestIdleConnWithoutHandshakeCompatibilityOrRoutingAffinityLocked(ap, compatibility, routingAffinity); idle != nil { delete(ap.conns, idle.id) evicted = append(evicted, idle) p.metrics.scaleDownTotal.Add(1) - } else if compatible == nil { + } else if affine == nil { + compatible := p.pickLeastBusyConnLocked(ap, "", compatibility) + if compatible != nil { + // Capacity is full and every compatible connection is busy. The + // hint remains soft here: queue on a compatible connection below. + goto acquireAtCapacity + } hasConnection := false for _, conn := range ap.conns { if conn != nil { @@ -1073,6 +1100,17 @@ retryAcquire: ap = p.getOrCreateAccountPool(accountID) ap.mu.Lock() ap.creating-- + if ap.generation != acquireGeneration { + ap.signalChangedLocked() + ap.mu.Unlock() + if conn != nil { + conn.close() + } + if retry < 1 { + return p.acquire(ctx, req, retry+1) + } + return nil, errOpenAIWSConnClosed + } if dialErr != nil { ap.prewarmFails++ ap.prewarmFailAt = time.Now() @@ -1080,20 +1118,26 @@ retryAcquire: ap.mu.Unlock() return nil, dialErr } + // Claim the freshly dialed connection before publishing it. Otherwise a + // topology waiter awakened below can take the free semaphore first and + // make the caller that paid for the dial queue behind it. + if !conn.tryAcquire() { + ap.signalChangedLocked() + ap.mu.Unlock() + conn.close() + return nil, errOpenAIWSConnClosed + } ap.conns[conn.id] = conn ap.prewarmFails = 0 ap.prewarmFailAt = time.Time{} + // Wake acquires that observed creating>0 with no compatible connection. + // Without this signal they can remain asleep until the new lease is + // released, even though the pool topology already changed. + ap.signalChangedLocked() ap.mu.Unlock() p.metrics.acquireCreateTotal.Add(1) - - if !conn.tryAcquire() { - if err := conn.acquire(ctx); err != nil { - conn.close() - p.evictConn(accountID, conn.id) - return nil, err - } - } lease := &openAIWSConnLease{pool: p, accountID: accountID, conn: conn, connPick: connPick} + p.recordLastSuccessfulAcquire(accountID, acquireGeneration, req) p.ensureTargetIdleAsync(accountID) return lease, nil } @@ -1105,6 +1149,7 @@ retryAcquire: return nil, errOpenAIWSConnQueueFull } +acquireAtCapacity: target := p.pickLeastBusyConnLocked(ap, req.PreferredConnID, compatibility) connPick := time.Since(pickStartedAt) p.recordConnPickDuration(connPick) @@ -1147,6 +1192,7 @@ retryAcquire: p.metrics.acquireQueueWaitMs.Add(queueWait.Milliseconds()) lease := &openAIWSConnLease{pool: p, accountID: accountID, conn: target, queueWait: queueWait, connPick: connPick, reused: true} p.metrics.acquireReuseTotal.Add(1) + p.recordLastSuccessfulAcquire(accountID, acquireGeneration, req) p.ensureTargetIdleAsync(accountID) return lease, nil } @@ -1162,6 +1208,23 @@ func (p *openAIWSConnPool) recordConnPickDuration(duration time.Duration) { p.metrics.connPickMs.Add(duration.Milliseconds()) } +func (p *openAIWSConnPool) recordLastSuccessfulAcquire(accountID int64, generation uint64, req openAIWSAcquireRequest) { + if p == nil || accountID <= 0 { + return + } + ap, ok := p.getAccountPool(accountID) + if !ok || ap == nil { + return + } + ap.mu.Lock() + if ap.generation != generation { + ap.mu.Unlock() + return + } + ap.lastAcquire = cloneOpenAIWSAcquireRequestPtr(&req) + ap.mu.Unlock() +} + func (p *openAIWSConnPool) pickOldestIdleConnLocked(ap *openAIWSAccountPool) *openAIWSConn { if ap == nil || len(ap.conns) == 0 { return nil @@ -1178,16 +1241,19 @@ func (p *openAIWSConnPool) pickOldestIdleConnLocked(ap *openAIWSAccountPool) *op return oldest } -func (p *openAIWSConnPool) pickOldestIdleConnWithDifferentHandshakeCompatibilityLocked( +func (p *openAIWSConnPool) pickOldestIdleConnWithoutHandshakeCompatibilityOrRoutingAffinityLocked( ap *openAIWSAccountPool, compatibility openAIWSHandshakeCompatibilityKey, + routingAffinity string, ) *openAIWSConn { if ap == nil || len(ap.conns) == 0 { return nil } var oldest *openAIWSConn for _, conn := range ap.conns { - if conn == nil || conn.matchesHandshakeCompatibility(compatibility) || conn.isLeased() || conn.waiters.Load() > 0 || p.isConnPinnedLocked(ap, conn.id) { + if conn == nil || + (conn.matchesHandshakeCompatibility(compatibility) && conn.matchesRoutingAffinity(routingAffinity)) || + conn.isLeased() || conn.waiters.Load() > 0 || p.isConnPinnedLocked(ap, conn.id) { continue } if oldest == nil || conn.lastUsedAt().Before(oldest.lastUsedAt()) { @@ -1372,6 +1438,36 @@ func (p *openAIWSConnPool) pickLeastBusyConnLocked( return best } +func (p *openAIWSConnPool) pickLeastBusyConnWithRoutingAffinityLocked( + ap *openAIWSAccountPool, + compatibility openAIWSHandshakeCompatibilityKey, + routingAffinity string, +) *openAIWSConn { + if ap == nil || len(ap.conns) == 0 { + return nil + } + var best *openAIWSConn + var bestWaiters int32 + var bestLastUsed time.Time + for _, conn := range ap.conns { + if conn == nil || + !conn.matchesHandshakeCompatibility(compatibility) || + !conn.matchesRoutingAffinity(routingAffinity) { + continue + } + waiters := conn.waiters.Load() + lastUsed := conn.lastUsedAt() + if best == nil || + waiters < bestWaiters || + (waiters == bestWaiters && lastUsed.Before(bestLastUsed)) { + best = conn + bestWaiters = waiters + bestLastUsed = lastUsed + } + } + return best +} + func accountPoolLoadLocked(ap *openAIWSAccountPool) (inflight int, waiters int) { if ap == nil { return 0, 0 @@ -1500,12 +1596,20 @@ func (p *openAIWSConnPool) prewarmConns(accountID int64, req openAIWSAcquireRequ if len(generations) > 0 { generation = generations[0] } + staleTarget := false defer func() { if ap, ok := p.getAccountPool(accountID); ok && ap != nil { ap.mu.Lock() ap.prewarmActive = false + ap.signalChangedLocked() ap.mu.Unlock() } + if staleTarget { + // A newer acquire arrived while the old dial was in flight. Re-run + // target selection only after clearing prewarmActive so the latest + // beta/hint target can fill the idle budget. + p.ensureTargetIdleAsync(accountID) + } }() for i := 0; i < total; i++ { @@ -1536,6 +1640,13 @@ func (p *openAIWSConnPool) prewarmConns(accountID int64, req openAIWSAcquireRequ conn.close() continue } + if !sameOpenAIWSPrewarmTarget(req, *ap.lastAcquire) { + staleTarget = true + ap.signalChangedLocked() + ap.mu.Unlock() + conn.close() + continue + } if len(ap.conns) >= p.effectiveMaxConnsByAccount(req.Account) { ap.signalChangedLocked() ap.mu.Unlock() @@ -1575,6 +1686,7 @@ func (p *openAIWSConnPool) ClearAccount(accountID int64) { ap.prewarmUntil = time.Time{} ap.prewarmFails = 0 ap.prewarmFailAt = time.Time{} + ap.signalChangedLocked() ap.mu.Unlock() closeOpenAIWSConns(conns) } @@ -1689,6 +1801,7 @@ func (p *openAIWSConnPool) dialConn(ctx context.Context, req openAIWSAcquireRequ id := p.nextConnID(req.Account.ID) pooledConn := newOpenAIWSConn(id, req.Account.ID, conn, handshakeHeaders) pooledConn.handshakeCompatibility = normalizeOpenAIWSHandshakeCompatibility(req.Headers) + pooledConn.routingAffinity = normalizeOpenAIWSRoutingAffinity(req.Headers) return pooledConn, nil } @@ -1867,6 +1980,12 @@ func cloneOpenAIWSAcquireRequestPtr(req *openAIWSAcquireRequest) *openAIWSAcquir return &copied } +func sameOpenAIWSPrewarmTarget(a, b openAIWSAcquireRequest) bool { + return stringsTrim(a.WSURL) == stringsTrim(b.WSURL) && + stringsTrim(a.ProxyURL) == stringsTrim(b.ProxyURL) && + normalizeOpenAIWSHandshakeCompatibility(a.Headers) == normalizeOpenAIWSHandshakeCompatibility(b.Headers) +} + func normalizeOpenAIWSBetaFeatures(headers http.Header) string { features := make(map[string]struct{}) for name, values := range headers { @@ -1895,10 +2014,36 @@ func normalizeOpenAIWSBetaFeatures(headers http.Header) string { func normalizeOpenAIWSHandshakeCompatibility(headers http.Header) openAIWSHandshakeCompatibilityKey { return openAIWSHandshakeCompatibilityKey{ betaFeatures: normalizeOpenAIWSBetaFeatures(headers), - routingHint: strings.TrimSpace(headers.Get(openAICodexRoutingHintHeader)), } } +func normalizeOpenAIWSRoutingAffinity(headers http.Header) string { + canonicalName := http.CanonicalHeaderKey(openAICodexRoutingHintHeader) + if values, ok := headers[canonicalName]; ok { + for _, value := range values { + if value = strings.TrimSpace(value); value != "" { + return value + } + } + } + + variantNames := make([]string, 0) + for name := range headers { + if name != canonicalName && strings.EqualFold(strings.TrimSpace(name), openAICodexRoutingHintHeader) { + variantNames = append(variantNames, name) + } + } + sort.Strings(variantNames) + for _, name := range variantNames { + for _, value := range headers[name] { + if value = strings.TrimSpace(value); value != "" { + return value + } + } + } + return "" +} + func cloneHeader(src http.Header) http.Header { if src == nil { return nil diff --git a/backend/internal/service/openai_ws_pool_test.go b/backend/internal/service/openai_ws_pool_test.go index 8d339359ee..f04b81d40d 100644 --- a/backend/internal/service/openai_ws_pool_test.go +++ b/backend/internal/service/openai_ws_pool_test.go @@ -61,6 +61,18 @@ func TestOpenAIWSConnPool_AcquireCleanupInterval(t *testing.T) { require.Less(t, openAIWSAcquireCleanupInterval, openAIWSBackgroundSweepTicker) } +func TestNormalizeOpenAIWSRoutingAffinityPrefersCanonicalAndSortsVariants(t *testing.T) { + headers := http.Header{ + "X-CODEX-ROUTING-HINT": []string{" variant-uppercase "}, + "X-Codex-Routing-Hint": []string{" ", " canonical "}, + } + + require.Equal(t, "canonical", normalizeOpenAIWSRoutingAffinity(headers)) + + delete(headers, "X-Codex-Routing-Hint") + require.Equal(t, "variant-uppercase", normalizeOpenAIWSRoutingAffinity(headers)) +} + func TestOpenAIWSConnLease_WriteJSONAndGuards(t *testing.T) { conn := newOpenAIWSConn("lease_write", 1, &openAIWSFakeConn{}, nil) lease := &openAIWSConnLease{conn: conn} @@ -310,6 +322,219 @@ func TestOpenAIWSConnPool_AcquireQueueWaitMetrics(t *testing.T) { require.GreaterOrEqual(t, metrics.ConnPickTotal, int64(1)) } +func TestOpenAIWSConnPool_DialSuccessWakesTopologyWaiterAndCanceledWaiterDoesNotLoseLease(t *testing.T) { + cfg := &config.Config{} + cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1 + cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0 + cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1 + cfg.Gateway.OpenAIWS.QueueLimitPerConn = 4 + + pool := newOpenAIWSConnPool(cfg) + dialer := newOpenAIWSFirstDialBlockingCaptureDialer() + pool.setClientDialerForTest(dialer) + account := &Account{ID: 991, Platform: PlatformOpenAI, Type: AccountTypeOAuth} + req := openAIWSAcquireRequest{Account: account, WSURL: "wss://example.com/v1/responses"} + + type result struct { + lease *openAIWSConnLease + err error + } + firstCh := make(chan result, 1) + go func() { + lease, err := pool.Acquire(context.Background(), req) + firstCh <- result{lease: lease, err: err} + }() + <-dialer.firstStarted + + waitCtx, cancelWait := context.WithCancel(context.Background()) + secondCh := make(chan result, 1) + waitReq := req + waitReq.Headers = http.Header{openAICodexRoutingHintHeader: {"model=gpt-5.6-codex;tier=priority"}} + go func() { + lease, err := pool.Acquire(waitCtx, waitReq) + secondCh <- result{lease: lease, err: err} + }() + + close(dialer.releaseFirst) + first := <-firstCh + require.NoError(t, first.err) + require.NotNil(t, first.lease) + + // The second acquire initially waits on the account topology channel while + // the first dial is in flight. Dial success must wake it immediately so it + // can queue on the newly-created (still leased) connection. + require.Eventually(t, func() bool { + ap, ok := pool.getAccountPool(account.ID) + if !ok || ap == nil { + return false + } + ap.mu.Lock() + defer ap.mu.Unlock() + for _, conn := range ap.conns { + if conn != nil && conn.waiters.Load() == 1 { + return true + } + } + return false + }, time.Second, 5*time.Millisecond) + + cancelWait() + second := <-secondCh + require.ErrorIs(t, second.err, context.Canceled) + require.Nil(t, second.lease) + ap, ok := pool.getAccountPool(account.ID) + require.True(t, ok) + ap.mu.Lock() + require.NotNil(t, ap.lastAcquire) + require.Empty(t, normalizeOpenAIWSRoutingAffinity(ap.lastAcquire.Headers), "a canceled acquire must not replace the successful prewarm target") + ap.mu.Unlock() + first.lease.Release() + + third, err := pool.Acquire(context.Background(), req) + require.NoError(t, err) + require.True(t, third.Reused(), "a canceled waiter must not consume the released semaphore token") + require.Equal(t, first.lease.ConnID(), third.ConnID()) + third.Release() +} + +func TestOpenAIWSConnPool_PrewarmHintChangeDoesNotInvalidateHealthyDial(t *testing.T) { + cfg := &config.Config{} + cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1 + cfg.Gateway.OpenAIWS.MinIdlePerAccount = 1 + cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1 + cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 2 + + pool := newOpenAIWSConnPool(cfg) + dialer := newOpenAIWSFirstDialBlockingCaptureDialer() + pool.setClientDialerForTest(dialer) + account := &Account{ID: 992, Platform: PlatformOpenAI, Type: AccountTypeOAuth} + oldHeaders := make(http.Header) + oldHeaders.Set(openAICodexRoutingHintHeader, "model=gpt-5.6-codex") + newHeaders := make(http.Header) + newHeaders.Set(openAICodexRoutingHintHeader, "model=gpt-5.6-codex;tier=priority") + ap := pool.getOrCreateAccountPool(account.ID) + ap.mu.Lock() + ap.lastAcquire = &openAIWSAcquireRequest{ + Account: account, + WSURL: "wss://example.com/v1/responses", + Headers: oldHeaders, + } + ap.mu.Unlock() + + pool.ensureTargetIdleAsync(account.ID) + <-dialer.firstStarted + + // Simulate a newer priority target arriving while the old model-only + // prewarm dial is in flight. Routing hints are advisory, so this alone must + // not discard an otherwise compatible connection. + ap.mu.Lock() + ap.lastAcquire = &openAIWSAcquireRequest{ + Account: account, + WSURL: "wss://example.com/v1/responses", + Headers: newHeaders, + } + ap.mu.Unlock() + close(dialer.releaseFirst) + + require.Eventually(t, func() bool { + ap.mu.Lock() + defer ap.mu.Unlock() + if ap.prewarmActive || len(ap.conns) != 1 { + return false + } + for _, conn := range ap.conns { + return conn != nil && conn.routingAffinity == "model=gpt-5.6-codex" + } + return false + }, 2*time.Second, 10*time.Millisecond) + require.Equal(t, 1, dialer.DialCount(), "routing-hint-only changes must not turn advisory metadata into hard reconnects") +} + +func TestOpenAIWSConnPool_ClearAccountWakesIncompatibleTopologyWaiter(t *testing.T) { + cfg := &config.Config{} + cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1 + cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0 + cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1 + + pool := newOpenAIWSConnPool(cfg) + dialer := &openAIWSCountingDialer{} + pool.setClientDialerForTest(dialer) + account := &Account{ID: 993, Platform: PlatformOpenAI, Type: AccountTypeOAuth} + baseReq := openAIWSAcquireRequest{Account: account, WSURL: "wss://example.com/v1/responses"} + betaAReq := baseReq + betaAReq.Headers = http.Header{"X-Codex-Beta-Features": {"feature_a"}} + betaBReq := baseReq + betaBReq.Headers = http.Header{"X-Codex-Beta-Features": {"feature_b"}} + + busy, err := pool.Acquire(context.Background(), betaAReq) + require.NoError(t, err) + require.Equal(t, 1, dialer.DialCount()) + + type result struct { + lease *openAIWSConnLease + err error + } + resultCh := make(chan result, 1) + waitCtx, cancelWait := context.WithTimeout(context.Background(), time.Second) + defer cancelWait() + go func() { + lease, acquireErr := pool.Acquire(waitCtx, betaBReq) + resultCh <- result{lease: lease, err: acquireErr} + }() + + require.Never(t, func() bool { return dialer.DialCount() > 1 }, 50*time.Millisecond, 5*time.Millisecond) + pool.ClearAccount(account.ID) + + resultB := <-resultCh + require.NoError(t, resultB.err) + require.NotNil(t, resultB.lease) + require.False(t, resultB.lease.Reused()) + require.Equal(t, 2, dialer.DialCount(), "ClearAccount must wake the waiter to redial immediately") + resultB.lease.Release() + busy.Release() +} + +func TestOpenAIWSConnPool_ClearAccountDoesNotReviveInFlightDialGeneration(t *testing.T) { + cfg := &config.Config{} + cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1 + cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0 + cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1 + + pool := newOpenAIWSConnPool(cfg) + dialer := newOpenAIWSFirstDialBlockingCaptureDialer() + pool.setClientDialerForTest(dialer) + account := &Account{ID: 994, Platform: PlatformOpenAI, Type: AccountTypeOAuth} + req := openAIWSAcquireRequest{Account: account, WSURL: "wss://example.com/v1/responses"} + + type result struct { + lease *openAIWSConnLease + err error + } + resultCh := make(chan result, 1) + go func() { + lease, err := pool.Acquire(context.Background(), req) + resultCh <- result{lease: lease, err: err} + }() + <-dialer.firstStarted + + pool.ClearAccount(account.ID) + close(dialer.releaseFirst) + got := <-resultCh + require.NoError(t, got.err) + require.NotNil(t, got.lease) + require.Equal(t, 2, dialer.DialCount(), "the pre-clear dial must be discarded and retried in the new generation") + require.True(t, strings.HasSuffix(got.lease.ConnID(), "_2")) + + ap, ok := pool.getAccountPool(account.ID) + require.True(t, ok) + ap.mu.Lock() + require.Equal(t, uint64(1), ap.generation) + require.Len(t, ap.conns, 1) + require.NotNil(t, ap.lastAcquire, "only the post-clear successful acquire may restore the prewarm target") + ap.mu.Unlock() + got.lease.Release() +} + func TestOpenAIWSConnPool_ForceNewConnSkipsReuse(t *testing.T) { cfg := &config.Config{} cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2 @@ -1468,6 +1693,21 @@ func TestOpenAIWSConn_AdditionalGuardBranches(t *testing.T) { closeOpenAIWSConns([]*openAIWSConn{nil, connOK}) } +func TestOpenAIWSConnPool_CanceledWaiterReturnsDeliveredLease(t *testing.T) { + conn := newOpenAIWSConn("cancelled_delivery", 1, &openAIWSFakeConn{}, nil) + + // Both branches of acquire's select are ready. Before the post-delivery + // cancellation check this intermittently returned nil after consuming the + // only lease token, which made the next pool acquire block forever. + for range 64 { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + require.ErrorIs(t, conn.acquire(ctx), context.Canceled) + require.True(t, conn.tryAcquire(), "a canceled waiter must return a delivered lease token") + conn.release() + } +} + func TestOpenAIWSConnLease_MarkBrokenEvictsConn(t *testing.T) { pool := newOpenAIWSConnPool(&config.Config{}) accountID := int64(5001) @@ -1602,6 +1842,21 @@ type openAIWSCountingDialer struct { dialCount int } +type openAIWSFirstDialBlockingCaptureDialer struct { + mu sync.Mutex + dialCount int + headers []http.Header + firstStarted chan struct{} + releaseFirst chan struct{} +} + +func newOpenAIWSFirstDialBlockingCaptureDialer() *openAIWSFirstDialBlockingCaptureDialer { + return &openAIWSFirstDialBlockingCaptureDialer{ + firstStarted: make(chan struct{}), + releaseFirst: make(chan struct{}), + } +} + type openAIWSAlwaysFailDialer struct { mu sync.Mutex dialCount int @@ -1671,6 +1926,36 @@ func (d *openAIWSCountingDialer) Dial( return &openAIWSFakeConn{}, 0, nil, nil } +func (d *openAIWSFirstDialBlockingCaptureDialer) Dial( + ctx context.Context, + wsURL string, + headers http.Header, + proxyURL string, +) (openAIWSClientConn, int, http.Header, error) { + _ = wsURL + _ = proxyURL + d.mu.Lock() + d.dialCount++ + dialNumber := d.dialCount + d.headers = append(d.headers, cloneHeader(headers)) + d.mu.Unlock() + if dialNumber == 1 { + close(d.firstStarted) + select { + case <-ctx.Done(): + return nil, 0, nil, ctx.Err() + case <-d.releaseFirst: + } + } + return &openAIWSFakeConn{}, 0, nil, nil +} + +func (d *openAIWSFirstDialBlockingCaptureDialer) DialCount() int { + d.mu.Lock() + defer d.mu.Unlock() + return d.dialCount +} + func (d *openAIWSCountingDialer) DialCount() int { d.mu.Lock() defer d.mu.Unlock()