test(wsrelay): lock websocket reuse rules

Author: DeliciousBuding

proxy/wsrelay/manager.go

@@ -284,7 +284,7 @@ func (m *Manager) AcquireConnection(
 
 	if v, ok := m.connections.Load(key); ok {
 		wc := v.(*WsConnection)
-		if wc.IsConnected() && !wc.IsExpired() && wc.session != nil && wc.session.PendingCount() == 0 {
+		if canReuseConnection(wc) {
 			wc.Touch()
 			return wc, nil
 		}
@@ -309,6 +309,19 @@ func (m *Manager) AcquireConnection(
 	return wc, nil
 }
 
+func canReuseConnection(wc *WsConnection) bool {
+	if wc == nil {
+		return false
+	}
+	if !wc.IsConnected() || wc.IsExpired() {
+		return false
+	}
+	if wc.session == nil {
+		return false
+	}
+	return wc.session.PendingCount() == 0
+}
+
 // createConnection 创建新 WebSocket 连接
 func (m *Manager) createConnection(
 	ctx context.Context,

proxy/wsrelay/manager_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"net/http"
 	"testing"
+	"time"
 
 	"github.com/codex2api/auth"
 )
@@ -59,3 +60,47 @@ func TestPoolKeyKeepsSameSessionStable(t *testing.T) {
 		t.Fatal("expected identical session keys to produce the same pool key")
 	}
 }
+
+func TestCanReuseConnection(t *testing.T) {
+	manager := NewManager()
+	t.Cleanup(manager.Stop)
+
+	t.Run("idle connected session can be reused", func(t *testing.T) {
+		session := NewSession(42, manager)
+		session.SetConnected(true)
+		conn := &WsConnection{session: session}
+		conn.SetState(StateConnected)
+		conn.Touch()
+
+		if !canReuseConnection(conn) {
+			t.Fatal("expected connection to be reusable")
+		}
+	})
+
+	t.Run("pending request blocks reuse", func(t *testing.T) {
+		session := NewSession(42, manager)
+		session.SetConnected(true)
+		pending := session.AddPendingRequest("session-a")
+		t.Cleanup(func() { session.RemovePendingRequest(pending.RequestID) })
+
+		conn := &WsConnection{session: session}
+		conn.SetState(StateConnected)
+		conn.Touch()
+
+		if canReuseConnection(conn) {
+			t.Fatal("expected connection with pending request to be non-reusable")
+		}
+	})
+
+	t.Run("expired connection cannot be reused", func(t *testing.T) {
+		session := NewSession(42, manager)
+		session.SetConnected(true)
+		conn := &WsConnection{session: session}
+		conn.SetState(StateConnected)
+		conn.lastUsed.Store(time.Now().Add(-IdleTimeout - time.Second).UnixNano())
+
+		if canReuseConnection(conn) {
+			t.Fatal("expected expired connection to be non-reusable")
+		}
+	})
+}