diff --git a/backend/internal/service/openai_ws_client.go b/backend/internal/service/openai_ws_client.go index c9b92d22e8..7444e41d07 100644 --- a/backend/internal/service/openai_ws_client.go +++ b/backend/internal/service/openai_ws_client.go @@ -47,6 +47,22 @@ type openAIWSIdlePingCapable interface { SupportsIdlePingWithoutReader() bool } +// openAIWSReaderLoopCapable 声明该实现的控制帧只在阻塞读期间被消费， +// 连接池需为其常驻一个读循环，否则空闲连接无法应答上游 ping。 +type openAIWSReaderLoopCapable interface { + RequiresReaderLoop() bool +} + +// openAIWSUpstreamPingCounter 报告连接收到过多少个上游 ping 帧，用于核对读循环是否在应答保活。 +type openAIWSUpstreamPingCounter interface { + UpstreamPingCount() int64 +} + +// openAIWSForceCloser 不做关闭握手直接切断连接，用于对端已不响应的场景。 +type openAIWSForceCloser interface { + CloseNow() error +} + // openAIWSClientDialer 抽象 WS 建连器。 type openAIWSClientDialer interface { Dial(ctx context.Context, wsURL string, headers http.Header, proxyURL string) (openAIWSClientConn, int, http.Header, error) @@ -107,9 +123,14 @@ func (d *coderOpenAIWSClientDialer) Dial( return nil, 0, nil, errors.New("ws url is empty") } + wrapped := &coderOpenAIWSClientConn{} opts := &coderws.DialOptions{ HTTPHeader: cloneHeader(headers), CompressionMode: coderws.CompressionContextTakeover, + OnPingReceived: func(context.Context, []byte) bool { + wrapped.upstreamPings.Add(1) + return true + }, } if proxy := strings.TrimSpace(proxyURL); proxy != "" { proxyClient, err := d.proxyHTTPClient(proxy) @@ -141,7 +162,8 @@ func (d *coderOpenAIWSClientDialer) Dial( if resp != nil { respHeaders = cloneHeader(resp.Header) } - return &coderOpenAIWSClientConn{conn: conn}, 0, respHeaders, nil + wrapped.conn = conn + return wrapped, 0, respHeaders, nil } func (d *coderOpenAIWSClientDialer) proxyHTTPClient(proxy string) (*http.Client, error) { @@ -267,7 +289,15 @@ func (d *coderOpenAIWSClientDialer) SnapshotTransportMetrics() OpenAIWSTransport } type coderOpenAIWSClientConn struct { - conn *coderws.Conn + conn *coderws.Conn + upstreamPings atomic.Int64 +} + +func (c *coderOpenAIWSClientConn) UpstreamPingCount() int64 { + if c == nil { + return 0 + } + return c.upstreamPings.Load() } var _ openaiwsv2.FrameConn = (*coderOpenAIWSClientConn)(nil) @@ -338,12 +368,17 @@ func (c *coderOpenAIWSClientConn) Ping(ctx context.Context) error { // SupportsIdlePingWithoutReader reports the actual coder/websocket contract. // Conn.Ping waits for a pong, while control frames are only consumed by Read. -// The pool deliberately has no reader on an idle connection, so using Ping as -// a health probe would deterministically time out a healthy socket. +// Without a reader, using Ping as a health probe would deterministically time +// out a healthy socket; the pool compensates with a resident reader loop. func (*coderOpenAIWSClientConn) SupportsIdlePingWithoutReader() bool { return false } +// RequiresReaderLoop 让池为 coder/websocket 连接常驻读循环，空闲期也能应答 ping。 +func (*coderOpenAIWSClientConn) RequiresReaderLoop() bool { + return true +} + func (c *coderOpenAIWSClientConn) Close() error { if c == nil || c.conn == nil { return nil @@ -353,3 +388,11 @@ func (c *coderOpenAIWSClientConn) Close() error { _ = c.conn.CloseNow() return nil } + +func (c *coderOpenAIWSClientConn) CloseNow() error { + if c == nil || c.conn == nil { + return nil + } + _ = c.conn.CloseNow() + return nil +} diff --git a/backend/internal/service/openai_ws_forwarder_ingress.go b/backend/internal/service/openai_ws_forwarder_ingress.go index f24dc25697..889f921595 100644 --- a/backend/internal/service/openai_ws_forwarder_ingress.go +++ b/backend/internal/service/openai_ws_forwarder_ingress.go @@ -853,13 +853,14 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( } agentTaskRecoveryTried := false - var acquireTurnLease func(int, string, bool) (*openAIWSConnLease, error) - acquireTurnLease = func(turn int, preferred string, forcePreferredConn bool) (*openAIWSConnLease, error) { + var acquireTurnLease func(int, string, bool, bool) (*openAIWSConnLease, error) + acquireTurnLease = func(turn int, preferred string, forcePreferredConn bool, forceNewConn bool) (*openAIWSConnLease, error) { req := cloneOpenAIWSAcquireRequest(baseAcquireReq) req.PreferredConnID = strings.TrimSpace(preferred) req.ForcePreferredConn = forcePreferredConn - // dedicated 模式下每次获取均新建连接，避免跨会话复用残留上下文。 - req.ForceNewConn = dedicatedMode + // dedicated 模式下每次获取均新建连接，避免跨会话复用残留上下文； + // 上游读写失败后的重试同样新建，避免再拿到同批陈旧的空闲连接。 + req.ForceNewConn = dedicatedMode || forceNewConn acquireCtx, acquireCancel := context.WithTimeout(ctx, acquireTimeout) lease, acquireErr := pool.Acquire(acquireCtx, req) acquireCancel() @@ -869,7 +870,7 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( if recoveryErr := s.recoverAgentIdentityTask(ctx, account, account.GetCredential("task_id")); recoveryErr != nil { return nil, fmt.Errorf("agent identity task recovery failed: %w", recoveryErr) } - return acquireTurnLease(turn, preferred, forcePreferredConn) + return acquireTurnLease(turn, preferred, forcePreferredConn, forceNewConn) } if acquireErr != nil { if isOpenAIWSSessionPreempted(ctx) { @@ -933,11 +934,14 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( baseAcquireReq.Headers = updatedHeaders } logOpenAIWSModeInfo( - "ingress_ws_upstream_connected account_id=%d turn=%d conn_id=%s conn_reused=%v conn_pick_ms=%d queue_wait_ms=%d preferred_conn_id=%s", + "ingress_ws_upstream_connected account_id=%d turn=%d conn_id=%s conn_reused=%v conn_idle_ms=%d conn_age_ms=%d upstream_pings=%d conn_pick_ms=%d queue_wait_ms=%d preferred_conn_id=%s", account.ID, turn, truncateOpenAIWSLogValue(connID, openAIWSIDValueMaxLen), lease.Reused(), + lease.IdleBefore().Milliseconds(), + lease.AgeBefore().Milliseconds(), + lease.UpstreamPingCount(), lease.ConnPickDuration().Milliseconds(), lease.QueueWaitDuration().Milliseconds(), truncateOpenAIWSLogValue(preferred, openAIWSIDValueMaxLen), @@ -1628,7 +1632,7 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( } forcePreferredConn := isStrictAffinityTurn(currentPayload) if sessionLease == nil { - acquiredLease, acquireErr := acquireTurnLease(turn, preferredConnID, forcePreferredConn) + acquiredLease, acquireErr := acquireTurnLease(turn, preferredConnID, forcePreferredConn, turnRetry > 0) if acquireErr != nil { return fmt.Errorf("acquire upstream websocket: %w", acquireErr) } @@ -1647,7 +1651,7 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( } } if shouldPreflightPing { - if pingErr := sessionLease.PingWithTimeout(openAIWSConnHealthCheckTO); pingErr != nil { + if pingErr := sessionLease.PingWithTimeout(openAIWSProbePingTO); pingErr != nil { logOpenAIWSModeInfo( "ingress_ws_upstream_preflight_ping_fail account_id=%d turn=%d conn_id=%s cause=%s", account.ID, @@ -1736,7 +1740,7 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( } resetSessionLease(true) - acquiredLease, acquireErr := acquireTurnLease(turn, preferredConnID, forcePreferredConn) + acquiredLease, acquireErr := acquireTurnLease(turn, preferredConnID, forcePreferredConn, false) if acquireErr != nil { return fmt.Errorf("acquire upstream websocket after preflight ping fail: %w", acquireErr) } diff --git a/backend/internal/service/openai_ws_forwarder_ingress_retry_test.go b/backend/internal/service/openai_ws_forwarder_ingress_retry_test.go new file mode 100644 index 0000000000..c0c1c65548 --- /dev/null +++ b/backend/internal/service/openai_ws_forwarder_ingress_retry_test.go @@ -0,0 +1,320 @@ +package service + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + coderws "github.com/coder/websocket" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" +) + +type openAIWSSlowPongCaptureConn struct { + *openAIWSCaptureConn + delay time.Duration + pings atomic.Int32 +} + +func (c *openAIWSSlowPongCaptureConn) Ping(ctx context.Context) error { + c.pings.Add(1) + timer := time.NewTimer(c.delay) + defer timer.Stop() + select { + case <-timer.C: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +// 轮次间隔后的预检 ping 经代理可能 2 秒多才回 pong：慢 pong 不应被判为连接失效而换连。 +func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_PreflightPingToleratesSlowPong(t *testing.T) { + gin.SetMode(gin.TestMode) + prevPreflightPingIdle := openAIWSIngressPreflightPingIdle + openAIWSIngressPreflightPingIdle = 0 + defer func() { + openAIWSIngressPreflightPingIdle = prevPreflightPingIdle + }() + + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + cfg.Gateway.OpenAIWS.Enabled = true + cfg.Gateway.OpenAIWS.OAuthEnabled = true + cfg.Gateway.OpenAIWS.APIKeyEnabled = true + cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true + cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2 + cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0 + cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 2 + cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8 + cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3 + cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 5 + cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3 + + slowConn := &openAIWSSlowPongCaptureConn{ + openAIWSCaptureConn: &openAIWSCaptureConn{ + events: [][]byte{ + []byte(`{"type":"response.completed","response":{"id":"resp_slow_pong_1","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`), + []byte(`{"type":"response.completed","response":{"id":"resp_slow_pong_2","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`), + }, + }, + delay: 2500 * time.Millisecond, + } + fallback := &openAIWSCaptureConn{ + events: [][]byte{ + []byte(`{"type":"response.completed","response":{"id":"resp_slow_pong_fallback","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`), + }, + } + dialer := &openAIWSQueueDialer{conns: []openAIWSClientConn{slowConn, fallback}} + pool := newOpenAIWSConnPool(cfg) + pool.setClientDialerForTest(dialer) + + svc := &OpenAIGatewayService{ + cfg: cfg, + httpUpstream: &httpUpstreamRecorder{}, + cache: &stubGatewayCache{}, + openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), + toolCorrector: NewCodexToolCorrector(), + openaiWSPool: pool, + } + account := &Account{ + ID: 444, + Name: "openai-ingress-preflight-slow-pong", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Credentials: map[string]any{"api_key": "sk-test"}, + Extra: map[string]any{"responses_websockets_v2_enabled": true}, + } + + serverErrCh := make(chan error, 1) + wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{CompressionMode: coderws.CompressionContextTakeover}) + if err != nil { + serverErrCh <- err + return + } + defer func() { _ = conn.CloseNow() }() + rec := httptest.NewRecorder() + ginCtx, _ := gin.CreateTestContext(rec) + req := r.Clone(r.Context()) + req.Header = req.Header.Clone() + req.Header.Set("User-Agent", "unit-test-agent/1.0") + ginCtx.Request = req + readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second) + msgType, firstMessage, readErr := conn.Read(readCtx) + cancel() + if readErr != nil { + serverErrCh <- readErr + return + } + if msgType != coderws.MessageText && msgType != coderws.MessageBinary { + serverErrCh <- errors.New("unsupported websocket client message type") + return + } + serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil) + })) + defer wsServer.Close() + + dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second) + clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil) + cancelDial() + require.NoError(t, err) + defer func() { _ = clientConn.CloseNow() }() + + writeMessage := func(payload string) { + writeCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + require.NoError(t, clientConn.Write(writeCtx, coderws.MessageText, []byte(payload))) + } + readMessage := func() []byte { + readCtx, cancel := context.WithTimeout(context.Background(), 8*time.Second) + defer cancel() + msgType, message, readErr := clientConn.Read(readCtx) + require.NoError(t, readErr) + require.Equal(t, coderws.MessageText, msgType) + return message + } + + writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"input":[{"type":"input_text","text":"hello"}]}`) + require.Equal(t, "resp_slow_pong_1", gjson.GetBytes(readMessage(), "response.id").String()) + writeMessage(`{"type":"response.create","model":"gpt-5.1","stream":false,"input":[{"type":"input_text","text":"world"}]}`) + require.Equal(t, "resp_slow_pong_2", gjson.GetBytes(readMessage(), "response.id").String(), "慢 pong 的健康会话连接应继续使用") + + require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done")) + select { + case serverErr := <-serverErrCh: + require.NoError(t, serverErr) + case <-time.After(5 * time.Second): + t.Fatal("等待 ingress websocket 结束超时") + } + require.GreaterOrEqual(t, slowConn.pings.Load(), int32(1), "第二轮前应执行预检 ping") + require.Equal(t, 1, dialer.DialCount(), "预检 ping 慢于 2 秒不应换连") +} + +// 场景：池内两条空闲连接同批陈旧（首读即失败）。首轮拿到其一失败后重试， +// 重试必须新建连接，而不是再从池里拿另一条同样陈旧的连接。 +func TestOpenAIGatewayService_ProxyResponsesWebSocketFromClient_TurnRetryForcesFreshConn(t *testing.T) { + gin.SetMode(gin.TestMode) + + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + cfg.Gateway.OpenAIWS.Enabled = true + cfg.Gateway.OpenAIWS.OAuthEnabled = true + cfg.Gateway.OpenAIWS.APIKeyEnabled = true + cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true + cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2 + cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0 + cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 2 + cfg.Gateway.OpenAIWS.QueueLimitPerConn = 8 + cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3 + cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3 + cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3 + + // 首会话用完 events 后归还池内，再次借出时首读即 EOF，等价于上游已判死的空闲连接。 + staleA := &openAIWSCaptureConn{ + events: [][]byte{ + []byte(`{"type":"response.completed","response":{"id":"resp_retry_first","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`), + }, + } + staleB := &openAIWSCaptureConn{} + fresh := &openAIWSCaptureConn{ + events: [][]byte{ + []byte(`{"type":"response.completed","response":{"id":"resp_retry_fresh","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`), + }, + } + dialer := &openAIWSQueueDialer{ + conns: []openAIWSClientConn{staleA, fresh}, + } + pool := newOpenAIWSConnPool(cfg) + pool.setClientDialerForTest(dialer) + + svc := &OpenAIGatewayService{ + cfg: cfg, + httpUpstream: &httpUpstreamRecorder{}, + cache: &stubGatewayCache{}, + openaiWSResolver: NewOpenAIWSProtocolResolver(cfg), + toolCorrector: NewCodexToolCorrector(), + openaiWSPool: pool, + } + + account := &Account{ + ID: 443, + Name: "openai-ingress-turn-retry-fresh", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "sk-test", + }, + Extra: map[string]any{ + "responses_websockets_v2_enabled": true, + }, + } + + serverErrCh := make(chan error, 2) + wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{ + CompressionMode: coderws.CompressionContextTakeover, + }) + if err != nil { + serverErrCh <- err + return + } + defer func() { + _ = conn.CloseNow() + }() + + rec := httptest.NewRecorder() + ginCtx, _ := gin.CreateTestContext(rec) + req := r.Clone(r.Context()) + req.Header = req.Header.Clone() + req.Header.Set("User-Agent", "unit-test-agent/1.0") + ginCtx.Request = req + + readCtx, cancel := context.WithTimeout(r.Context(), 3*time.Second) + msgType, firstMessage, readErr := conn.Read(readCtx) + cancel() + if readErr != nil { + serverErrCh <- readErr + return + } + if msgType != coderws.MessageText && msgType != coderws.MessageBinary { + serverErrCh <- errors.New("unsupported websocket client message type") + return + } + + serverErrCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "sk-test", firstMessage, nil) + })) + defer wsServer.Close() + + runSingleTurnSession := func(expectedResponseID string) { + dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second) + clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http"), nil) + cancelDial() + require.NoError(t, err) + defer func() { + _ = clientConn.CloseNow() + }() + + writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second) + err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.1","stream":false}`)) + cancelWrite() + require.NoError(t, err) + + readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second) + msgType, event, readErr := clientConn.Read(readCtx) + cancelRead() + require.NoError(t, readErr) + require.Equal(t, coderws.MessageText, msgType) + require.Equal(t, expectedResponseID, gjson.GetBytes(event, "response.id").String()) + + require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done")) + + select { + case serverErr := <-serverErrCh: + require.NoError(t, serverErr) + case <-time.After(5 * time.Second): + t.Fatal("等待 ingress websocket 结束超时") + } + } + + runSingleTurnSession("resp_retry_first") + require.Equal(t, 1, dialer.DialCount()) + + // 用首会话的真实 acquire 参数构造第二条空闲连接，保证握手兼容键与后续请求一致。 + ap := pool.getOrCreateAccountPool(account.ID) + ap.mu.Lock() + lastAcquire := ap.lastAcquire + ap.mu.Unlock() + require.NotNil(t, lastAcquire) + staleConnB := newOpenAIWSConn(pool.nextConnID(account.ID), account.ID, staleB, nil) + staleConnB.handshakeCompatibility = normalizeOpenAIWSHandshakeCompatibility(lastAcquire.Account, lastAcquire.Headers) + staleConnB.routingAffinity = normalizeOpenAIWSRoutingAffinity(lastAcquire.Headers) + ap.mu.Lock() + ap.conns[staleConnB.id] = staleConnB + idleConns := len(ap.conns) + ap.mu.Unlock() + require.Equal(t, 2, idleConns, "第二会话开始前池内应有两条空闲连接") + + runSingleTurnSession("resp_retry_fresh") + + require.Equal(t, 2, dialer.DialCount(), "首读失败后的重试应新建连接") + fresh.mu.Lock() + freshWrites := len(fresh.writes) + fresh.mu.Unlock() + require.Equal(t, 1, freshWrites) +} diff --git a/backend/internal/service/openai_ws_forwarder_logutil.go b/backend/internal/service/openai_ws_forwarder_logutil.go index ddb37d8380..330ba33bb9 100644 --- a/backend/internal/service/openai_ws_forwarder_logutil.go +++ b/backend/internal/service/openai_ws_forwarder_logutil.go @@ -515,6 +515,10 @@ func logOpenAIWSModeInfo(format string, args ...any) { logger.LegacyPrintf("service.openai_gateway", "[OpenAI WS Mode][openai_ws_mode=true] "+format, args...) } +func logOpenAIWSModeWarn(format string, args ...any) { + logger.LegacyPrintf("service.openai_gateway", "[warn] [OpenAI WS Mode][openai_ws_mode=true] "+format, args...) +} + func isOpenAIWSModeDebugEnabled() bool { return logger.L().Core().Enabled(zap.DebugLevel) } diff --git a/backend/internal/service/openai_ws_forwarder_v2.go b/backend/internal/service/openai_ws_forwarder_v2.go index 9f04d42640..d37f41bbd4 100644 --- a/backend/internal/service/openai_ws_forwarder_v2.go +++ b/backend/internal/service/openai_ws_forwarder_v2.go @@ -261,12 +261,15 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2( }() connID := strings.TrimSpace(lease.ConnID()) logOpenAIWSModeDebug( - "connected account_id=%d account_type=%s transport=%s conn_id=%s conn_reused=%v conn_pick_ms=%d queue_wait_ms=%d has_previous_response_id=%v", + "connected account_id=%d account_type=%s transport=%s conn_id=%s conn_reused=%v conn_idle_ms=%d conn_age_ms=%d upstream_pings=%d conn_pick_ms=%d queue_wait_ms=%d has_previous_response_id=%v", account.ID, account.Type, normalizeOpenAIWSLogValue(string(decision.Transport)), connID, lease.Reused(), + lease.IdleBefore().Milliseconds(), + lease.AgeBefore().Milliseconds(), + lease.UpstreamPingCount(), lease.ConnPickDuration().Milliseconds(), lease.QueueWaitDuration().Milliseconds(), previousResponseID != "", diff --git a/backend/internal/service/openai_ws_pool.go b/backend/internal/service/openai_ws_pool.go index 5a04392399..ed86a36284 100644 --- a/backend/internal/service/openai_ws_pool.go +++ b/backend/internal/service/openai_ws_pool.go @@ -20,10 +20,13 @@ import ( const ( openAIWSConnMaxAge = 60 * time.Minute openAIWSConnHealthCheckIdle = 90 * time.Second - // coder/websocket cannot consume pong frames without a reader. Recycle - // unsupported idle sockets before the upstream keepalive window expires. - openAIWSConnIdleRecycleAfter = 90 * time.Second - openAIWSConnHealthCheckTO = 2 * time.Second + // 仅对没有常驻读循环的连接实现生效：这类连接空闲时无人应答上游 ping，须在 + // 上游保活窗口到期前回收。coder/websocket 连接由池常驻读循环应答 ping，不受此阈值约束。 + openAIWSConnIdleRecycleAfter = 90 * time.Second + openAIWSConnHealthCheckTO = 2 * time.Second + // 不在请求热路径上的探活（后台巡检、轮次间预检）给经代理链路的 pong 留足余量， + // 实测最大往返约 1.7s；误判的代价是换连甚至断会话，比多等几秒重得多。 + openAIWSProbePingTO = 10 * time.Second openAIWSConnPrewarmExtraDelay = 2 * time.Second openAIWSAcquireCleanupInterval = 3 * time.Second openAIWSBackgroundPingInterval = 30 * time.Second @@ -90,13 +93,15 @@ type openAIWSHandshakeCompatibilityKey struct { } type openAIWSConnLease struct { - pool *openAIWSConnPool - accountID int64 - conn *openAIWSConn - queueWait time.Duration - connPick time.Duration - reused bool - released atomic.Bool + pool *openAIWSConnPool + accountID int64 + conn *openAIWSConn + queueWait time.Duration + connPick time.Duration + idleBefore time.Duration + ageBefore time.Duration + reused bool + released atomic.Bool } func (l *openAIWSConnLease) activeConn() (*openAIWSConn, error) { @@ -137,6 +142,29 @@ func (l *openAIWSConnLease) Reused() bool { return l.reused } +// IdleBefore 返回借出时该连接已空闲的时长。 +func (l *openAIWSConnLease) IdleBefore() time.Duration { + if l == nil { + return 0 + } + return l.idleBefore +} + +// AgeBefore 返回借出时该连接自建立起的时长。 +func (l *openAIWSConnLease) AgeBefore() time.Duration { + if l == nil { + return 0 + } + return l.ageBefore +} + +func (l *openAIWSConnLease) UpstreamPingCount() int64 { + if l == nil || l.conn == nil { + return 0 + } + return l.conn.upstreamPingCount() +} + func (l *openAIWSConnLease) HandshakeHeader(name string) string { if l == nil || l.conn == nil { return "" @@ -264,6 +292,17 @@ type openAIWSConn struct { readMu sync.Mutex writeMu sync.Mutex + // readerLoopResults 非 nil 表示池为该连接常驻了读循环：coder/websocket 只在 + // 阻塞读期间应答上游 ping，空闲连接没有读循环会被上游按保活超时关闭。 + readerLoopResults chan []byte + readerLoopErrMu sync.Mutex + readerLoopErr error + readerLoopPeerClosed atomic.Bool + // onPeerClosed 由池在建连后设置：上游主动关闭时立刻把连接移出账号池，不等清理周期。 + onPeerClosed atomic.Pointer[func()] + // unusable 表示空闲期收到数据被判为脏连接：持有令牌不再借出，由池在锁外关闭。 + unusable atomic.Bool + waiters atomic.Int32 createdAtNano atomic.Int64 lastUsedNano atomic.Int64 @@ -282,9 +321,138 @@ func newOpenAIWSConn(id string, _ int64, ws openAIWSClientConn, handshakeHeaders conn.leaseCh <- struct{}{} conn.createdAtNano.Store(now.UnixNano()) conn.lastUsedNano.Store(now.UnixNano()) + if capable, ok := ws.(openAIWSReaderLoopCapable); ok && capable.RequiresReaderLoop() { + conn.readerLoopResults = make(chan []byte, 1) + go conn.runReaderLoop() + } return conn } +func (c *openAIWSConn) runReaderLoop() { + defer close(c.readerLoopResults) + for { + payload, err := c.ws.ReadMessage(context.Background()) + if err != nil { + c.readerLoopErrMu.Lock() + c.readerLoopErr = err + c.readerLoopErrMu.Unlock() + // 本地主动关闭时对端会回 close 帧，同样以读错误结束循环，不算上游事件。 + peerClosed := false + select { + case <-c.closedCh: + default: + peerClosed = true + c.readerLoopPeerClosed.Store(true) + now := time.Now() + // 空闲连接被上游断开是常态，池已当场出池，只记 info；借出中断开会影响请求，记 warn。 + logClosed := logOpenAIWSModeInfo + if c.isLeased() { + logClosed = logOpenAIWSModeWarn + } + logClosed( + "conn_reader_loop_closed conn_id=%s leased=%v idle_ms=%d age_ms=%d upstream_pings=%d cause=%s", + c.id, + c.isLeased(), + c.idleDuration(now).Milliseconds(), + c.age(now).Milliseconds(), + c.upstreamPingCount(), + truncateOpenAIWSLogValue(err.Error(), openAIWSLogValueMaxLen), + ) + } + c.close() + if evict := c.onPeerClosed.Load(); peerClosed && evict != nil { + (*evict)() + } + return + } + select { + case c.readerLoopResults <- payload: + case <-c.closedCh: + return + } + } +} + +func (c *openAIWSConn) hasReaderLoop() bool { + return c != nil && c.readerLoopResults != nil +} + +func (c *openAIWSConn) readerLoopClosedByPeer() bool { + return c != nil && c.readerLoopPeerClosed.Load() +} + +func (c *openAIWSConn) upstreamPingCount() int64 { + if c == nil || c.ws == nil { + return 0 + } + if counter, ok := c.ws.(openAIWSUpstreamPingCounter); ok { + return counter.UpstreamPingCount() + } + return 0 +} + +// readerLoopPending 报告空闲期是否已有数据消息被读循环缓存。len 不消费消息； +// 缓存满时读循环阻塞在投递上，因此最多只有这一条待接管消息。 +func (c *openAIWSConn) readerLoopPending() bool { + return c.hasReaderLoop() && len(c.readerLoopResults) > 0 +} + +func (c *openAIWSConn) readerLoopError() error { + c.readerLoopErrMu.Lock() + defer c.readerLoopErrMu.Unlock() + if c.readerLoopErr != nil { + return c.readerLoopErr + } + return errOpenAIWSConnClosed +} + +// leaseTokenUsable 在拿到租约令牌后确认连接仍可借出：已关闭的连接退回令牌； +// 空闲期收到过数据消息的连接状态已不可信，直接关闭而不交给借用者。 +func (c *openAIWSConn) leaseTokenUsable() bool { + select { + case <-c.closedCh: + c.release() + return false + default: + } + if c.readerLoopPending() { + // 只记事件类型，不记报文原文，避免模型输出进日志。 + eventType := "" + select { + case payload := <-c.readerLoopResults: + eventType = effectiveOpenAISSEEventType(payload, "") + default: + } + logOpenAIWSModeWarn( + "conn_idle_dirty_discard conn_id=%s idle_ms=%d upstream_pings=%d event=%s", + c.id, + c.idleDuration(time.Now()).Milliseconds(), + c.upstreamPingCount(), + normalizeOpenAIWSLogValue(eventType), + ) + // 关闭握手可能阻塞到一个 RTT，而 tryAcquire 在池锁内调用，这里只标记，出池后再关闭。 + c.unusable.Store(true) + return false + } + return true +} + +func (c *openAIWSConn) isClosed() bool { + if c == nil { + return true + } + select { + case <-c.closedCh: + return true + default: + return false + } +} + +func (c *openAIWSConn) isUnusable() bool { + return c != nil && c.unusable.Load() +} + func (c *openAIWSConn) tryAcquire() bool { if c == nil { return false @@ -296,13 +464,7 @@ func (c *openAIWSConn) tryAcquire() bool { } select { case <-c.leaseCh: - select { - case <-c.closedCh: - c.release() - return false - default: - } - return true + return c.leaseTokenUsable() default: return false } @@ -327,11 +489,8 @@ func (c *openAIWSConn) acquire(ctx context.Context) error { c.release() return err } - select { - case <-c.closedCh: - c.release() + if !c.leaseTokenUsable() { return errOpenAIWSConnClosed - default: } return nil } @@ -350,13 +509,27 @@ func (c *openAIWSConn) release() { } func (c *openAIWSConn) close() { + c.closeWith(false) +} + +// abort 不做关闭握手直接切断。读循环常驻持有读锁，礼貌关闭要等对端回 close 帧， +// 对端已不响应时库会等满 5 秒；读超时这类场景必须立即返回。 +func (c *openAIWSConn) abort() { + c.closeWith(true) +} + +func (c *openAIWSConn) closeWith(force bool) { if c == nil { return } c.closeOnce.Do(func() { close(c.closedCh) if c.ws != nil { - _ = c.ws.Close() + if forceCloser, ok := c.ws.(openAIWSForceCloser); ok && force { + _ = forceCloser.CloseNow() + } else { + _ = c.ws.Close() + } } select { case c.leaseCh <- struct{}{}: @@ -412,10 +585,13 @@ func (c *openAIWSConn) readMessageWithContextTimeout(parent context.Context, tim if c == nil { return nil, errOpenAIWSConnClosed } - select { - case <-c.closedCh: - return nil, errOpenAIWSConnClosed - default: + // 有读循环时连接关闭后缓冲里可能还有未取走的消息，交给 readMessage 先排空再报错。 + if c.readerLoopResults == nil { + select { + case <-c.closedCh: + return nil, errOpenAIWSConnClosed + default: + } } if parent == nil { @@ -438,12 +614,27 @@ func (c *openAIWSConn) readMessage(readCtx context.Context) ([]byte, error) { if readCtx == nil { readCtx = context.Background() } - payload, err := c.ws.ReadMessage(readCtx) - if err != nil { - return nil, err + if c.readerLoopResults == nil { + payload, err := c.ws.ReadMessage(readCtx) + if err != nil { + return nil, err + } + c.touch() + return payload, nil + } + select { + case payload, ok := <-c.readerLoopResults: + if !ok { + return nil, c.readerLoopError() + } + c.touch() + return payload, nil + case <-readCtx.Done(): + // 与库在 ctx 取消时切断连接的语义一致：读超时后消息边界已不可信，且对端多半 + // 已不响应，直接切断而不做关闭握手。 + c.abort() + return nil, readCtx.Err() } - c.touch() - return payload, nil } func (c *openAIWSConn) pingWithTimeout(timeout time.Duration) error { @@ -456,8 +647,8 @@ func (c *openAIWSConn) pingWithTimeout(timeout time.Duration) error { default: } - c.writeMu.Lock() - defer c.writeMu.Unlock() + // coder/websocket 除 Reader/Read 外的方法都可并发调用，控制帧由库内 writeFrameMu 串行化， + // 这里不持有 writeMu，避免等 pong 期间阻塞借用者写请求。 if c.ws == nil { return errOpenAIWSConnClosed } @@ -476,6 +667,9 @@ func (c *openAIWSConn) supportsIdlePingWithoutReader() bool { if c == nil || c.ws == nil { return false } + if c.readerLoopResults != nil { + return true + } capable, ok := c.ws.(openAIWSIdlePingCapable) // Test and alternate implementations keep the historical probe behavior // unless they explicitly declare it unsafe. @@ -761,9 +955,39 @@ func (p *openAIWSConnPool) runBackgroundPingSweep() { continue } g.Go(func() error { - if err := item.conn.pingWithTimeout(openAIWSConnHealthCheckTO); err != nil { + started := time.Now() + idleMs := item.conn.idleDuration(started).Milliseconds() + if err := item.conn.pingWithTimeout(openAIWSProbePingTO); err != nil { + // 只有拿到租约令牌才能剔除：判断与占有必须是同一个原子动作，否则 + // 等 pong 期间刚借出的连接会被从借用者手里关掉。已关闭或已判脏的连接 + // 没有借用者，照常剔除。 + if !item.conn.tryAcquire() && !item.conn.isClosed() && !item.conn.isUnusable() { + logOpenAIWSModeWarn( + "conn_background_ping_skip_leased conn_id=%s idle_ms=%d upstream_pings=%d cause=%s", + item.conn.id, + idleMs, + item.conn.upstreamPingCount(), + truncateOpenAIWSLogValue(err.Error(), openAIWSLogValueMaxLen), + ) + return nil + } + logOpenAIWSModeWarn( + "conn_background_ping_evict conn_id=%s idle_ms=%d upstream_pings=%d cause=%s", + item.conn.id, + idleMs, + item.conn.upstreamPingCount(), + truncateOpenAIWSLogValue(err.Error(), openAIWSLogValueMaxLen), + ) p.evictConn(item.accountID, item.conn.id) + return nil } + logOpenAIWSModeDebug( + "conn_background_ping_ok conn_id=%s idle_ms=%d rtt_ms=%d upstream_pings=%d", + item.conn.id, + idleMs, + time.Since(started).Milliseconds(), + item.conn.upstreamPingCount(), + ) return nil }) } @@ -851,7 +1075,13 @@ func (p *openAIWSConnPool) Acquire(ctx context.Context, req openAIWSAcquireReque if p != nil { p.metrics.acquireTotal.Add(1) } - return p.acquire(ctx, cloneOpenAIWSAcquireRequest(req), 0) + lease, err := p.acquire(ctx, cloneOpenAIWSAcquireRequest(req), 0) + if lease != nil && lease.conn != nil { + now := time.Now() + lease.idleBefore = lease.conn.idleDuration(now) + lease.ageBefore = lease.conn.age(now) + } + return lease, err } func (p *openAIWSConnPool) acquire(ctx context.Context, req openAIWSAcquireRequest, retry int) (*openAIWSConnLease, error) { @@ -927,6 +1157,12 @@ retryAcquire: return lease, nil } + if p.dropDeadConnLocked(ap, preferredConn, &evicted) { + p.recordConnPickDuration(time.Since(pickStartedAt)) + ap.mu.Unlock() + closeOpenAIWSConns(evicted) + return nil, errOpenAIWSPreferredConnUnavailable + } connPick := time.Since(pickStartedAt) p.recordConnPickDuration(connPick) if int(preferredConn.waiters.Load()) >= p.queueLimitPerConn() { @@ -942,8 +1178,11 @@ retryAcquire: p.metrics.acquireQueueWaitTotal.Add(1) if err := preferredConn.acquire(ctx); err != nil { - if errors.Is(err, errOpenAIWSConnClosed) && retry < 1 { - return p.acquire(ctx, req, retry+1) + if errors.Is(err, errOpenAIWSConnClosed) { + p.evictConn(accountID, preferredConn.id) + if retry < 1 { + return p.acquire(ctx, req, retry+1) + } } return nil, err } @@ -996,6 +1235,8 @@ retryAcquire: p.recordLastSuccessfulAcquire(accountID, acquireGeneration, req) p.ensureTargetIdleAsync(accountID) return lease, nil + } else if conn, ok := ap.conns[preferredConnID]; ok { + p.dropDeadConnLocked(ap, conn, &evicted) } } @@ -1023,6 +1264,8 @@ retryAcquire: p.recordLastSuccessfulAcquire(accountID, acquireGeneration, req) p.ensureTargetIdleAsync(accountID) return lease, nil + } else if best != nil { + p.dropDeadConnLocked(ap, best, &evicted) } if routingAffinity == "" || len(ap.conns)+ap.creating >= effectiveMaxConns { for _, conn := range ap.conns { @@ -1050,6 +1293,7 @@ retryAcquire: p.ensureTargetIdleAsync(accountID) return lease, nil } + p.dropDeadConnLocked(ap, conn, &evicted) } } } @@ -1182,8 +1426,11 @@ acquireAtCapacity: p.metrics.acquireQueueWaitTotal.Add(1) if err := target.acquire(ctx); err != nil { - if errors.Is(err, errOpenAIWSConnClosed) && retry < 1 { - return p.acquire(ctx, req, retry+1) + if errors.Is(err, errOpenAIWSConnClosed) { + p.evictConn(accountID, target.id) + if retry < 1 { + return p.acquire(ctx, req, retry+1) + } } return nil, err } @@ -1343,15 +1590,13 @@ func (p *openAIWSConnPool) cleanupAccountLocked(ap *openAIWSAccountPool, now tim } continue } - select { - case <-conn.closedCh: + if conn.isClosed() || conn.isUnusable() { delete(ap.conns, id) if len(ap.pinnedConns) > 0 { delete(ap.pinnedConns, id) } evicted = append(evicted, conn) continue - default: } if p.isConnPinnedLocked(ap, id) { continue @@ -1712,6 +1957,24 @@ func (p *openAIWSConnPool) ClearAccount(accountID int64) { closeOpenAIWSConns(conns) } +// dropDeadConnLocked 在池锁内把已关闭或已判脏的连接移出账号池并释放容量， +// 连接本身交给调用方在解锁后关闭。 +func (p *openAIWSConnPool) dropDeadConnLocked(ap *openAIWSAccountPool, conn *openAIWSConn, evicted *[]*openAIWSConn) bool { + if ap == nil || conn == nil || (!conn.isClosed() && !conn.isUnusable()) { + return false + } + if _, exists := ap.conns[conn.id]; !exists { + return false + } + delete(ap.conns, conn.id) + if len(ap.pinnedConns) > 0 { + delete(ap.pinnedConns, conn.id) + } + ap.signalChangedLocked() + *evicted = append(*evicted, conn) + return true +} + func (p *openAIWSConnPool) evictConn(accountID int64, connID string) { if p == nil || accountID <= 0 || stringsTrim(connID) == "" { return @@ -1821,6 +2084,9 @@ func (p *openAIWSConnPool) dialConn(ctx context.Context, req openAIWSAcquireRequ } id := p.nextConnID(req.Account.ID) pooledConn := newOpenAIWSConn(id, req.Account.ID, conn, handshakeHeaders) + accountID := req.Account.ID + evict := func() { p.evictConn(accountID, id) } + pooledConn.onPeerClosed.Store(&evict) pooledConn.handshakeCompatibility = normalizeOpenAIWSHandshakeCompatibility(req.Account, req.Headers) pooledConn.routingAffinity = normalizeOpenAIWSRoutingAffinity(req.Headers) return pooledConn, nil @@ -1840,6 +2106,10 @@ func (p *openAIWSConnPool) shouldHealthCheckConn(conn *openAIWSConn) bool { if conn == nil || !conn.supportsIdlePingWithoutReader() { return false } + // 有读循环的连接能即时感知上游关闭，半开探测已由后台巡检覆盖，借出前不再多付一个往返。 + if conn.hasReaderLoop() { + return false + } return conn.idleDuration(time.Now()) >= openAIWSConnHealthCheckIdle } diff --git a/backend/internal/service/openai_ws_pool_reader_loop_test.go b/backend/internal/service/openai_ws_pool_reader_loop_test.go new file mode 100644 index 0000000000..c36c53de4f --- /dev/null +++ b/backend/internal/service/openai_ws_pool_reader_loop_test.go @@ -0,0 +1,794 @@ +package service + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + coderws "github.com/coder/websocket" + "github.com/stretchr/testify/require" + + "github.com/Wei-Shaw/sub2api/internal/config" +) + +// openAIWSReaderLoopFakeConn 模拟 coder/websocket 的契约： +// 控制帧只在有人阻塞读时被消费，因此 Ping 只有在存在未返回的 ReadMessage 时才成功。 +type openAIWSReaderLoopFakeConn struct { + messages chan []byte + failCh chan struct{} + failOnce sync.Once + failErr error + closed chan struct{} + closeOnce sync.Once + readers atomic.Int32 + pings atomic.Int32 + pinging atomic.Int32 + pingDelay time.Duration + pingHold chan struct{} + pingErr error +} + +func newOpenAIWSReaderLoopFakeConn() *openAIWSReaderLoopFakeConn { + return &openAIWSReaderLoopFakeConn{ + messages: make(chan []byte, 8), + failCh: make(chan struct{}), + closed: make(chan struct{}), + } +} + +func (c *openAIWSReaderLoopFakeConn) WriteJSON(context.Context, any) error { + return nil +} + +func (c *openAIWSReaderLoopFakeConn) ReadMessage(ctx context.Context) ([]byte, error) { + c.readers.Add(1) + defer c.readers.Add(-1) + select { + case msg := <-c.messages: + return msg, nil + case <-c.failCh: + return nil, c.failErr + case <-c.closed: + return nil, errors.New("closed") + case <-ctx.Done(): + return nil, ctx.Err() + } +} + +func (c *openAIWSReaderLoopFakeConn) Ping(ctx context.Context) error { + c.pings.Add(1) + if c.readers.Load() == 0 { + return errors.New("ping without reader") + } + c.pinging.Add(1) + defer c.pinging.Add(-1) + if c.pingHold != nil { + select { + case <-c.pingHold: + case <-ctx.Done(): + return ctx.Err() + } + } + if c.pingDelay > 0 { + timer := time.NewTimer(c.pingDelay) + defer timer.Stop() + select { + case <-timer.C: + case <-ctx.Done(): + return ctx.Err() + } + } + return c.pingErr +} + +// openAIWSSlowPingConn 是没有读循环的旧式连接，用来确认请求路径上的健康检查语义不变。 +type openAIWSSlowPingConn struct { + openAIWSFakeConn + pingDelay time.Duration +} + +func (c *openAIWSSlowPingConn) Ping(ctx context.Context) error { + timer := time.NewTimer(c.pingDelay) + defer timer.Stop() + select { + case <-timer.C: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +func (c *openAIWSReaderLoopFakeConn) SupportsIdlePingWithoutReader() bool { + return false +} + +func (c *openAIWSReaderLoopFakeConn) RequiresReaderLoop() bool { + return true +} + +func (c *openAIWSReaderLoopFakeConn) Close() error { + c.closeOnce.Do(func() { close(c.closed) }) + return nil +} + +func (c *openAIWSReaderLoopFakeConn) failReads(err error) { + c.failOnce.Do(func() { + c.failErr = err + close(c.failCh) + }) +} + +func (c *openAIWSReaderLoopFakeConn) isClosed() bool { + select { + case <-c.closed: + return true + default: + return false + } +} + +type openAIWSReaderLoopFakeDialer struct { + mu sync.Mutex + conns []*openAIWSReaderLoopFakeConn +} + +func (d *openAIWSReaderLoopFakeDialer) Dial(context.Context, string, http.Header, string) (openAIWSClientConn, int, http.Header, error) { + conn := newOpenAIWSReaderLoopFakeConn() + d.mu.Lock() + d.conns = append(d.conns, conn) + d.mu.Unlock() + return conn, http.StatusSwitchingProtocols, nil, nil +} + +func (d *openAIWSReaderLoopFakeDialer) dialed() []*openAIWSReaderLoopFakeConn { + d.mu.Lock() + defer d.mu.Unlock() + return append([]*openAIWSReaderLoopFakeConn(nil), d.conns...) +} + +func requireConnClosed(t *testing.T, conn *openAIWSConn) { + t.Helper() + select { + case <-conn.closedCh: + case <-time.After(time.Second): + t.Fatal("连接应已关闭") + } +} + +func TestCoderOpenAIWSClientConn_RequiresReaderLoop(t *testing.T) { + require.True(t, (&coderOpenAIWSClientConn{}).RequiresReaderLoop()) +} + +func TestOpenAIWSConnReaderLoop_NotStartedForPlainConn(t *testing.T) { + conn := newOpenAIWSConn("rl_plain", 1, &openAIWSFakeConn{}, nil) + defer conn.close() + require.False(t, conn.hasReaderLoop()) +} + +func TestOpenAIWSConnReaderLoop_KeepsReaderWhileIdleSoPingSucceeds(t *testing.T) { + fake := newOpenAIWSReaderLoopFakeConn() + conn := newOpenAIWSConn("rl_idle_ping", 1, fake, nil) + + require.True(t, conn.hasReaderLoop()) + require.Eventually(t, func() bool { return fake.readers.Load() == 1 }, time.Second, 5*time.Millisecond) + require.True(t, conn.supportsIdlePingWithoutReader()) + require.NoError(t, conn.pingWithTimeout(time.Second)) + + conn.close() + require.Eventually(t, func() bool { return fake.readers.Load() == 0 }, time.Second, 5*time.Millisecond, "关闭后读循环应退出") +} + +func TestOpenAIWSConnReaderLoop_DeliversMessagesInOrder(t *testing.T) { + fake := newOpenAIWSReaderLoopFakeConn() + conn := newOpenAIWSConn("rl_order", 1, fake, nil) + defer conn.close() + require.True(t, conn.tryAcquire()) + + fake.messages <- []byte("m1") + fake.messages <- []byte("m2") + fake.messages <- []byte("m3") + for _, want := range []string{"m1", "m2", "m3"} { + got, err := conn.readMessageWithTimeout(time.Second) + require.NoError(t, err) + require.Equal(t, want, string(got)) + } +} + +func TestOpenAIWSConnReaderLoop_ReadTimeoutClosesConn(t *testing.T) { + fake := newOpenAIWSReaderLoopFakeConn() + conn := newOpenAIWSConn("rl_timeout", 1, fake, nil) + require.True(t, conn.tryAcquire()) + + _, err := conn.readMessageWithTimeout(30 * time.Millisecond) + require.ErrorIs(t, err, context.DeadlineExceeded) + requireConnClosed(t, conn) + require.True(t, fake.isClosed()) +} + +func TestOpenAIWSConnReaderLoop_UpstreamCloseWhileIdleClosesConn(t *testing.T) { + fake := newOpenAIWSReaderLoopFakeConn() + conn := newOpenAIWSConn("rl_idle_close", 1, fake, nil) + + fake.failReads(errors.New("received close frame: keepalive ping timeout")) + requireConnClosed(t, conn) + require.False(t, conn.tryAcquire(), "上游已关闭的连接不能再被借出") + require.True(t, conn.readerLoopClosedByPeer(), "上游主动关闭应被记为对端关闭") +} + +func TestOpenAIWSConnReaderLoop_LocalCloseIsNotPeerClose(t *testing.T) { + fake := newOpenAIWSReaderLoopFakeConn() + conn := newOpenAIWSConn("rl_local_close", 1, fake, nil) + require.Eventually(t, func() bool { return fake.readers.Load() == 1 }, time.Second, 5*time.Millisecond) + + conn.close() + require.Eventually(t, func() bool { return fake.readers.Load() == 0 }, time.Second, 5*time.Millisecond) + require.False(t, conn.readerLoopClosedByPeer(), "本地主动关闭引发的读错误不应记为对端关闭") +} + +func TestOpenAIWSConnReaderLoop_UpstreamCloseWhileReadingReturnsCause(t *testing.T) { + fake := newOpenAIWSReaderLoopFakeConn() + conn := newOpenAIWSConn("rl_read_close", 1, fake, nil) + require.True(t, conn.tryAcquire()) + + upstreamErr := errors.New("received close frame: keepalive ping timeout") + done := make(chan error, 1) + go func() { + _, err := conn.readMessageWithTimeout(time.Second) + done <- err + }() + require.Eventually(t, func() bool { + if conn.readMu.TryLock() { + conn.readMu.Unlock() + return false + } + return true + }, time.Second, 5*time.Millisecond, "借用者应已进入读等待") + fake.failReads(upstreamErr) + + select { + case err := <-done: + require.ErrorIs(t, err, upstreamErr) + case <-time.After(time.Second): + t.Fatal("借用者应收到上游关闭错误") + } + requireConnClosed(t, conn) +} + +func TestOpenAIWSConnReaderLoop_ReadAfterUpstreamCloseReturnsCause(t *testing.T) { + fake := newOpenAIWSReaderLoopFakeConn() + conn := newOpenAIWSConn("rl_late_read", 1, fake, nil) + require.True(t, conn.tryAcquire()) + + upstreamErr := errors.New("received close frame: keepalive ping timeout") + fake.failReads(upstreamErr) + requireConnClosed(t, conn) + + _, err := conn.readMessageWithTimeout(time.Second) + require.ErrorIs(t, err, upstreamErr, "持有租约期间上游关闭，之后的读应带回关闭原因") +} + +// 上游发完最后一条事件随即关闭连接时，借用者必须先拿到已缓存的事件，之后才收到关闭错误。 +func TestOpenAIWSConnReaderLoop_BufferedMessageDeliveredAfterPeerClose(t *testing.T) { + fake := newOpenAIWSReaderLoopFakeConn() + conn := newOpenAIWSConn("rl_buffered_close", 1, fake, nil) + require.True(t, conn.tryAcquire()) + + fake.messages <- []byte(`{"type":"response.completed"}`) + require.Eventually(t, conn.readerLoopPending, time.Second, 5*time.Millisecond) + upstreamErr := errors.New("received close frame: status = StatusNormalClosure") + fake.failReads(upstreamErr) + requireConnClosed(t, conn) + + payload, err := conn.readMessageWithTimeout(time.Second) + require.NoError(t, err, "连接关闭前已缓存的事件不能丢") + require.Equal(t, `{"type":"response.completed"}`, string(payload)) + + _, err = conn.readMessageWithTimeout(time.Second) + require.ErrorIs(t, err, upstreamErr) +} + +func TestOpenAIWSConnPool_PeerClosedIdleConnEvictedImmediately(t *testing.T) { + cfg := &config.Config{} + cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1 + cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1 + cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 1 + pool := newOpenAIWSConnPool(cfg) + defer pool.Close() + dialer := &openAIWSReaderLoopFakeDialer{} + pool.setClientDialerForTest(dialer) + + account := &Account{ID: 307, Platform: PlatformOpenAI, Type: AccountTypeAPIKey} + req := openAIWSAcquireRequest{Account: account, WSURL: "wss://example.com/v1/responses"} + + first, err := pool.Acquire(context.Background(), req) + require.NoError(t, err) + firstID := first.ConnID() + first.Release() + + dialer.dialed()[0].failReads(errors.New("received close frame: keepalive ping timeout")) + requireConnClosed(t, first.conn) + + ap, ok := pool.getAccountPool(account.ID) + require.True(t, ok) + require.Eventually(t, func() bool { + ap.mu.Lock() + defer ap.mu.Unlock() + _, exists := ap.conns[firstID] + return !exists + }, time.Second, 5*time.Millisecond, "上游关闭的空闲连接应立即移出账号池") + + acquireCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + second, err := pool.Acquire(acquireCtx, req) + require.NoError(t, err, "池满时死连接不应继续占用容量") + defer second.Release() + require.NotEqual(t, firstID, second.ConnID()) + require.Len(t, dialer.dialed(), 2) +} + +func newReaderLoopTestPool(t *testing.T, maxConns int) (*openAIWSConnPool, *openAIWSReaderLoopFakeDialer, openAIWSAcquireRequest) { + t.Helper() + cfg := &config.Config{} + cfg.Gateway.OpenAIWS.MaxConnsPerAccount = maxConns + cfg.Gateway.OpenAIWS.MaxIdlePerAccount = maxConns + cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 1 + pool := newOpenAIWSConnPool(cfg) + t.Cleanup(pool.Close) + dialer := &openAIWSReaderLoopFakeDialer{} + pool.setClientDialerForTest(dialer) + account := &Account{ID: 308, Platform: PlatformOpenAI, Type: AccountTypeAPIKey} + return pool, dialer, openAIWSAcquireRequest{Account: account, WSURL: "wss://example.com/v1/responses"} +} + +func requireConnGone(t *testing.T, pool *openAIWSConnPool, accountID int64, connID string) { + t.Helper() + ap, ok := pool.getAccountPool(accountID) + require.True(t, ok) + ap.mu.Lock() + defer ap.mu.Unlock() + _, exists := ap.conns[connID] + require.False(t, exists, "失效连接应已移出账号池") +} + +// 池上限为 1：唯一的空闲连接空闲期收到残留消息，本次获取必须直接新建，不能等清理周期。 +func TestOpenAIWSConnPool_DirtyIdleConnReplacedAtCapacity(t *testing.T) { + pool, dialer, req := newReaderLoopTestPool(t, 1) + + first, err := pool.Acquire(context.Background(), req) + require.NoError(t, err) + firstID := first.ConnID() + first.Release() + dialer.dialed()[0].messages <- []byte(`{"type":"response.output_text.delta"}`) + require.Eventually(t, first.conn.readerLoopPending, time.Second, 5*time.Millisecond) + + acquireCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + second, err := pool.Acquire(acquireCtx, req) + require.NoError(t, err) + defer second.Release() + require.NotEqual(t, firstID, second.ConnID()) + require.Len(t, dialer.dialed(), 2) + requireConnGone(t, pool, req.Account.ID, firstID) + requireConnClosed(t, first.conn) +} + +// 池已满且其余连接都在忙：脏的空闲连接要让位给新建连接，而不是把请求挂到忙连接上排队。 +func TestOpenAIWSConnPool_DirtyIdleConnReplacedWhileOthersBusy(t *testing.T) { + pool, dialer, req := newReaderLoopTestPool(t, 2) + + busy, err := pool.Acquire(context.Background(), req) + require.NoError(t, err) + defer busy.Release() + idle, err := pool.Acquire(context.Background(), req) + require.NoError(t, err) + idleID := idle.ConnID() + idle.Release() + dialer.dialed()[1].messages <- []byte(`{"type":"response.output_text.delta"}`) + require.Eventually(t, idle.conn.readerLoopPending, time.Second, 5*time.Millisecond) + + acquireCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + fresh, err := pool.Acquire(acquireCtx, req) + require.NoError(t, err) + defer fresh.Release() + require.NotEqual(t, idleID, fresh.ConnID()) + require.NotEqual(t, busy.ConnID(), fresh.ConnID()) + require.Len(t, dialer.dialed(), 3) + requireConnGone(t, pool, req.Account.ID, idleID) +} + +// 排队等待的获取在拿到令牌时发现连接已脏：应把它出池并新建，而不是失败。 +func TestOpenAIWSConnPool_QueuedAcquireReplacesDirtyConnOnHandoff(t *testing.T) { + pool, dialer, req := newReaderLoopTestPool(t, 1) + + first, err := pool.Acquire(context.Background(), req) + require.NoError(t, err) + firstID := first.ConnID() + + type acquireResult struct { + lease *openAIWSConnLease + err error + } + done := make(chan acquireResult, 1) + go func() { + acquireCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + lease, err := pool.Acquire(acquireCtx, req) + done <- acquireResult{lease: lease, err: err} + }() + require.Eventually(t, func() bool { return first.conn.waiters.Load() == 1 }, time.Second, 5*time.Millisecond) + + dialer.dialed()[0].messages <- []byte(`{"type":"response.output_text.delta"}`) + require.Eventually(t, first.conn.readerLoopPending, time.Second, 5*time.Millisecond) + first.Release() + + select { + case res := <-done: + require.NoError(t, res.err) + defer res.lease.Release() + require.NotEqual(t, firstID, res.lease.ConnID()) + case <-time.After(4 * time.Second): + t.Fatal("排队获取未返回") + } + require.Len(t, dialer.dialed(), 2) + requireConnGone(t, pool, req.Account.ID, firstID) +} + +func newSweepTestPool(conn *openAIWSConn, accountID int64) (*openAIWSConnPool, *openAIWSAccountPool) { + cfg := &config.Config{} + cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2 + cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 2 + pool := &openAIWSConnPool{cfg: cfg} + ap := &openAIWSAccountPool{conns: map[string]*openAIWSConn{conn.id: conn}} + pool.accounts.Store(accountID, ap) + return pool, ap +} + +func requireInPool(t *testing.T, ap *openAIWSAccountPool, connID string, want bool) { + t.Helper() + ap.mu.Lock() + defer ap.mu.Unlock() + _, exists := ap.conns[connID] + require.Equal(t, want, exists) +} + +// 后台巡检经代理往返可能超过 2 秒：慢 pong 不应被判死。 +func TestOpenAIWSConnPool_BackgroundPingSweepToleratesSlowPong(t *testing.T) { + fake := newOpenAIWSReaderLoopFakeConn() + fake.pingDelay = 2500 * time.Millisecond + conn := newOpenAIWSConn("rl_slow_pong", 309, fake, nil) + defer conn.close() + pool, ap := newSweepTestPool(conn, 309) + require.Eventually(t, func() bool { return fake.readers.Load() == 1 }, time.Second, 5*time.Millisecond) + + pool.runBackgroundPingSweep() + + requireInPool(t, ap, conn.id, true) + require.False(t, conn.isClosed(), "慢 pong 的健康连接不应被巡检剔除") + require.EqualValues(t, 1, fake.pings.Load()) +} + +// ping 期间连接被借出，随后 ping 失败：巡检不得剔除已借出的连接。 +func TestOpenAIWSConnPool_BackgroundPingSweepSkipsLeasedConnOnFailure(t *testing.T) { + fake := newOpenAIWSReaderLoopFakeConn() + fake.pingHold = make(chan struct{}) + fake.pingErr = errors.New("pong lost") + conn := newOpenAIWSConn("rl_leased_during_ping", 310, fake, nil) + defer conn.close() + pool, ap := newSweepTestPool(conn, 310) + require.Eventually(t, func() bool { return fake.readers.Load() == 1 }, time.Second, 5*time.Millisecond) + + done := make(chan struct{}) + go func() { + defer close(done) + pool.runBackgroundPingSweep() + }() + require.Eventually(t, func() bool { return fake.pinging.Load() == 1 }, time.Second, 5*time.Millisecond) + require.True(t, conn.tryAcquire(), "ping 进行中连接应仍可借出") + close(fake.pingHold) + select { + case <-done: + case <-time.After(3 * time.Second): + t.Fatal("巡检未结束") + } + + requireInPool(t, ap, conn.id, true) + require.False(t, conn.isClosed(), "已借出的连接不应被巡检剔除") +} + +// 巡检判定失败到真正剔除之间若被借出，绝不能关闭借用者手里的连接： +// 用 ap.mu 把 evictConn 卡在门口，在这个空隙里借出连接，复现"看一眼再动手"的竞态。 +func TestOpenAIWSConnPool_BackgroundPingSweepNeverClosesConnLeasedInEvictWindow(t *testing.T) { + fake := newOpenAIWSReaderLoopFakeConn() + fake.pingHold = make(chan struct{}) + fake.pingErr = errors.New("pong lost") + conn := newOpenAIWSConn("rl_evict_window", 312, fake, nil) + defer conn.close() + pool, ap := newSweepTestPool(conn, 312) + require.Eventually(t, func() bool { return fake.readers.Load() == 1 }, time.Second, 5*time.Millisecond) + + done := make(chan struct{}) + go func() { + defer close(done) + pool.runBackgroundPingSweep() + }() + require.Eventually(t, func() bool { return fake.pinging.Load() == 1 }, time.Second, 5*time.Millisecond) + + ap.mu.Lock() + close(fake.pingHold) + time.Sleep(50 * time.Millisecond) + leased := conn.tryAcquire() + ap.mu.Unlock() + select { + case <-done: + case <-time.After(3 * time.Second): + t.Fatal("巡检未结束") + } + + if leased { + require.False(t, conn.isClosed(), "借用者拿到租约后，巡检不得再关闭该连接") + requireInPool(t, ap, conn.id, true) + } else { + require.True(t, conn.isClosed(), "巡检拿到租约令牌后才允许剔除") + } +} + +// 读超时后必须立即切断连接：上游已不响应时礼貌关闭握手要等 5 秒，会拖慢重试。 +func TestOpenAIWSConnReaderLoop_ReadTimeoutAbortsWithoutCloseHandshake(t *testing.T) { + release := make(chan struct{}) + defer close(release) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + srvConn, err := coderws.Accept(w, r, nil) + if err != nil { + return + } + <-release + _ = srvConn.CloseNow() + })) + defer server.Close() + + ws, _, _, err := newDefaultOpenAIWSClientDialer().Dial(context.Background(), "ws"+strings.TrimPrefix(server.URL, "http"), nil, "") + require.NoError(t, err) + conn := newOpenAIWSConn("rl_read_timeout_abort", 1, ws, nil) + defer conn.close() + require.True(t, conn.tryAcquire()) + + started := time.Now() + _, err = conn.readMessageWithTimeout(20 * time.Millisecond) + elapsed := time.Since(started) + require.ErrorIs(t, err, context.DeadlineExceeded) + require.Less(t, elapsed, 500*time.Millisecond, "读超时后关闭不应等待关闭握手") + requireConnClosed(t, conn) +} + +// ping 等待 pong 期间不能阻塞写：借用者写请求不应等到 pong 回来。 +func TestOpenAIWSConnPingDoesNotBlockWrite(t *testing.T) { + fake := newOpenAIWSReaderLoopFakeConn() + fake.pingHold = make(chan struct{}) + conn := newOpenAIWSConn("rl_ping_write", 1, fake, nil) + defer conn.close() + defer close(fake.pingHold) + require.Eventually(t, func() bool { return fake.readers.Load() == 1 }, time.Second, 5*time.Millisecond) + + pingDone := make(chan error, 1) + go func() { pingDone <- conn.pingWithTimeout(5 * time.Second) }() + require.Eventually(t, func() bool { return fake.pinging.Load() == 1 }, time.Second, 5*time.Millisecond) + + started := time.Now() + require.NoError(t, conn.writeJSONWithTimeout(context.Background(), map[string]any{"type": "response.create"}, time.Second)) + require.Less(t, time.Since(started), 300*time.Millisecond, "写请求不应等待 ping 的 pong") + require.EqualValues(t, 1, fake.pinging.Load(), "写完成时 ping 仍应在等待 pong") +} + +// 有读循环的连接不再做借出前的空闲健康检查：上游关闭能被即时感知，巡检已覆盖半开探测。 +func TestOpenAIWSConnPool_AcquireSkipsIdleHealthCheckForReaderLoopConn(t *testing.T) { + pool, dialer, req := newReaderLoopTestPool(t, 1) + + first, err := pool.Acquire(context.Background(), req) + require.NoError(t, err) + firstID := first.ConnID() + first.Release() + first.conn.lastUsedNano.Store(time.Now().Add(-openAIWSConnHealthCheckIdle - time.Second).UnixNano()) + + second, err := pool.Acquire(context.Background(), req) + require.NoError(t, err) + defer second.Release() + require.Equal(t, firstID, second.ConnID()) + require.Zero(t, dialer.dialed()[0].pings.Load(), "有读循环的连接借出时不应再 ping") +} + +// 没有读循环的旧式连接保持原语义：空闲超过阈值借出前 ping，2 秒无 pong 即换连接。 +func TestOpenAIWSConnPool_AcquireIdleHealthCheckStillAppliesWithoutReaderLoop(t *testing.T) { + cfg := &config.Config{} + cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1 + cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1 + cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 1 + pool := newOpenAIWSConnPool(cfg) + defer pool.Close() + dialer := &openAIWSQueueDialer{conns: []openAIWSClientConn{ + &openAIWSSlowPingConn{pingDelay: 2200 * time.Millisecond}, + &openAIWSFakeConn{}, + }} + pool.setClientDialerForTest(dialer) + account := &Account{ID: 311, Platform: PlatformOpenAI, Type: AccountTypeAPIKey} + req := openAIWSAcquireRequest{Account: account, WSURL: "wss://example.com/v1/responses"} + + first, err := pool.Acquire(context.Background(), req) + require.NoError(t, err) + firstID := first.ConnID() + first.Release() + first.conn.lastUsedNano.Store(time.Now().Add(-openAIWSConnHealthCheckIdle - time.Second).UnixNano()) + + started := time.Now() + second, err := pool.Acquire(context.Background(), req) + require.NoError(t, err) + defer second.Release() + require.NotEqual(t, firstID, second.ConnID(), "2 秒内无 pong 的旧式连接应被换掉") + require.GreaterOrEqual(t, time.Since(started), openAIWSConnHealthCheckTO) + require.Equal(t, 2, dialer.DialCount()) +} + +func TestOpenAIWSConnReaderLoop_DataWhileIdleRejectsLease(t *testing.T) { + fake := newOpenAIWSReaderLoopFakeConn() + conn := newOpenAIWSConn("rl_dirty_try", 1, fake, nil) + + fake.messages <- []byte(`{"type":"response.output_text.delta"}`) + require.Eventually(t, conn.readerLoopPending, time.Second, 5*time.Millisecond) + + require.False(t, conn.tryAcquire(), "空闲期收到数据的连接不能借出") + require.True(t, conn.isUnusable(), "脏连接应标记为不可用，由池在锁外关闭") + require.False(t, conn.tryAcquire(), "脏连接的令牌不应归还，避免被其他获取者拿到") +} + +func TestOpenAIWSConnReaderLoop_DataWhileIdleRejectsBlockingAcquire(t *testing.T) { + fake := newOpenAIWSReaderLoopFakeConn() + conn := newOpenAIWSConn("rl_dirty_acquire", 1, fake, nil) + + fake.messages <- []byte(`{"type":"response.output_text.delta"}`) + require.Eventually(t, conn.readerLoopPending, time.Second, 5*time.Millisecond) + + require.ErrorIs(t, conn.acquire(context.Background()), errOpenAIWSConnClosed) + require.True(t, conn.isUnusable(), "脏连接应标记为不可用，由池在锁外关闭") + require.False(t, conn.tryAcquire(), "脏连接的令牌不应归还") +} + +func TestOpenAIWSConnPool_ReaderLoopConnNotIdleRecycled(t *testing.T) { + cfg := &config.Config{} + cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2 + cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 2 + pool := &openAIWSConnPool{cfg: cfg} + + accountID := int64(304) + ap := &openAIWSAccountPool{conns: make(map[string]*openAIWSConn)} + conn := newOpenAIWSConn("rl_no_recycle", accountID, newOpenAIWSReaderLoopFakeConn(), nil) + defer conn.close() + conn.lastUsedNano.Store(time.Now().Add(-openAIWSConnIdleRecycleAfter - time.Second).UnixNano()) + ap.conns[conn.id] = conn + pool.accounts.Store(accountID, ap) + + pool.runBackgroundCleanupSweep(time.Now()) + + ap.mu.Lock() + _, exists := ap.conns[conn.id] + ap.mu.Unlock() + require.True(t, exists, "有常驻读循环的连接不应按空闲阈值回收") +} + +// 用真实 coder/websocket 连接复现线上场景：空闲池化连接必须能应答服务端 ping， +// 服务端随后以 1011/keepalive ping timeout 关闭时，池要立即感知并保留关闭原因。 +func TestOpenAIWSConnReaderLoop_RealConnAnswersServerPingWhileIdle(t *testing.T) { + pingResult := make(chan error, 1) + closeNow := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + srvConn, err := coderws.Accept(w, r, nil) + if err != nil { + pingResult <- err + return + } + defer func() { _ = srvConn.CloseNow() }() + ctx, cancel := context.WithTimeout(r.Context(), 5*time.Second) + defer cancel() + readDone := make(chan struct{}) + go func() { + defer close(readDone) + _, _, _ = srvConn.Reader(ctx) + }() + pingResult <- srvConn.Ping(ctx) + <-closeNow + _ = srvConn.Write(ctx, coderws.MessageText, []byte(`{"type":"response.completed"}`)) + _ = srvConn.Close(coderws.StatusInternalError, "keepalive ping timeout") + <-readDone + })) + defer server.Close() + + ws, _, _, err := newDefaultOpenAIWSClientDialer().Dial(context.Background(), "ws"+strings.TrimPrefix(server.URL, "http"), nil, "") + require.NoError(t, err) + conn := newOpenAIWSConn("rl_real", 1, ws, nil) + defer conn.close() + require.True(t, conn.hasReaderLoop()) + + select { + case pingErr := <-pingResult: + require.NoError(t, pingErr, "空闲池化连接应由读循环应答服务端 ping") + case <-time.After(5 * time.Second): + t.Fatal("服务端 ping 未在期限内收到 pong") + } + require.Eventually(t, func() bool { return conn.upstreamPingCount() == 1 }, time.Second, 5*time.Millisecond) + + require.True(t, conn.tryAcquire()) + close(closeNow) + requireConnClosed(t, conn) + require.True(t, conn.readerLoopClosedByPeer()) + + payload, err := conn.readMessageWithTimeout(2 * time.Second) + require.NoError(t, err, "close 帧之前的数据帧必须先交给借用者") + require.Equal(t, `{"type":"response.completed"}`, string(payload)) + _, err = conn.readMessageWithTimeout(2 * time.Second) + require.Error(t, err) + require.Contains(t, err.Error(), "keepalive ping timeout") +} + +func TestOpenAIWSConnPool_AcquireSnapshotsIdleBefore(t *testing.T) { + cfg := &config.Config{} + cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2 + cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 2 + cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 1 + pool := newOpenAIWSConnPool(cfg) + defer pool.Close() + pool.setClientDialerForTest(&openAIWSReaderLoopFakeDialer{}) + + account := &Account{ID: 306, Platform: PlatformOpenAI, Type: AccountTypeAPIKey} + req := openAIWSAcquireRequest{Account: account, WSURL: "wss://example.com/v1/responses"} + + first, err := pool.Acquire(context.Background(), req) + require.NoError(t, err) + require.Less(t, first.IdleBefore(), 500*time.Millisecond) + first.Release() + first.conn.lastUsedNano.Store(time.Now().Add(-5 * time.Second).UnixNano()) + first.conn.createdAtNano.Store(time.Now().Add(-time.Minute).UnixNano()) + + second, err := pool.Acquire(context.Background(), req) + require.NoError(t, err) + defer second.Release() + require.True(t, second.Reused()) + require.GreaterOrEqual(t, second.IdleBefore(), 5*time.Second, "借出时应记录借出前的空闲时长") + require.GreaterOrEqual(t, second.AgeBefore(), time.Minute, "借出时应记录连接年龄") + require.Zero(t, second.UpstreamPingCount()) +} + +func TestOpenAIWSConnPool_AcquireSkipsDirtyIdleConn(t *testing.T) { + cfg := &config.Config{} + cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2 + cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 2 + cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 1 + pool := newOpenAIWSConnPool(cfg) + defer pool.Close() + dialer := &openAIWSReaderLoopFakeDialer{} + pool.setClientDialerForTest(dialer) + + account := &Account{ID: 305, Platform: PlatformOpenAI, Type: AccountTypeAPIKey} + req := openAIWSAcquireRequest{Account: account, WSURL: "wss://example.com/v1/responses"} + + first, err := pool.Acquire(context.Background(), req) + require.NoError(t, err) + firstID := first.ConnID() + first.Release() + + dialed := dialer.dialed() + require.Len(t, dialed, 1) + dialed[0].messages <- []byte(`{"type":"response.output_text.delta"}`) + require.Eventually(t, first.conn.readerLoopPending, time.Second, 5*time.Millisecond) + + second, err := pool.Acquire(context.Background(), req) + require.NoError(t, err) + defer second.Release() + require.NotEqual(t, firstID, second.ConnID(), "脏连接不应再被借出") + require.Len(t, dialer.dialed(), 2) + requireConnClosed(t, first.conn) +}