fix(openai): harden priority routing hints

This commit is contained in:
白宦成
2026-08-07 21:47:54 +08:00
parent 815035fcc9
commit de349187d9
10 changed files with 850 additions and 39 deletions
@@ -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())
})
}
}
@@ -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
}
@@ -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
@@ -308,6 +308,7 @@ func normalizeOpenAICompactRequestBody(body []byte) ([]byte, bool, error) {
"tools",
"parallel_tool_calls",
"reasoning",
"service_tier",
"text",
"previous_response_id",
} {
@@ -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()))
}
@@ -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,
)
}
@@ -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())
@@ -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
}
+163 -18
View File
@@ -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
@@ -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()