Merge pull request #247 from huangye123/worktree-fix-ws-relay-stability-pr2

fix: 修复 WebSocket relay 稳定性问题

Author: james-6-23

proxy/wsrelay/error_passthrough_test.go

@@ -26,6 +26,59 @@ func TestBuildErrorEvent_UpstreamErrorBecomesFailedEvent(t *testing.T) {
 	}
 }
 
+// TestBuildErrorEvent_UpstreamStatusBecomesResponseStatusCode verifies an upstream
+// top-level status is copied to response.status_code even when the preserved
+// upstream error object does not contain status_code itself.
+func TestBuildErrorEvent_UpstreamStatusBecomesResponseStatusCode(t *testing.T) {
+	r := &WsResponse{}
+	upstream := []byte(`{"type":"error","status":429,"error":{"type":"rate_limit_error","message":"too many requests","param":"input"}}`)
+
+	event, isErr := r.buildErrorEvent(upstream)
+	if !isErr {
+		t.Fatal("expected error frame to be detected")
+	}
+	if statusCode := gjson.GetBytes(event, "response.status_code").Int(); statusCode != 429 {
+		t.Fatalf("response.status_code = %d, want 429; event=%s", statusCode, event)
+	}
+	if gjson.GetBytes(event, "response.error.status_code").Exists() {
+		t.Fatalf("response.error should preserve upstream error object without injecting status_code: %s", event)
+	}
+	if msg := gjson.GetBytes(event, "response.error.message").String(); msg != "too many requests" {
+		t.Fatalf("response.error.message = %q, want preserved upstream message", msg)
+	}
+	if typ := gjson.GetBytes(event, "response.error.type").String(); typ != "rate_limit_error" {
+		t.Fatalf("response.error.type = %q, want preserved upstream type", typ)
+	}
+	if param := gjson.GetBytes(event, "response.error.param").String(); param != "input" {
+		t.Fatalf("response.error.param = %q, want preserved upstream param", param)
+	}
+}
+
+// TestBuildErrorEvent_UpstreamStatusCodeBecomesResponseStatusCode verifies an
+// upstream top-level status_code is copied to response.status_code while the
+// original response.error payload remains unchanged.
+func TestBuildErrorEvent_UpstreamStatusCodeBecomesResponseStatusCode(t *testing.T) {
+	r := &WsResponse{}
+	upstream := []byte(`{"type":"error","status_code":503,"error":{"message":"service unavailable","code":"server_error"}}`)
+
+	event, isErr := r.buildErrorEvent(upstream)
+	if !isErr {
+		t.Fatal("expected error frame to be detected")
+	}
+	if statusCode := gjson.GetBytes(event, "response.status_code").Int(); statusCode != 503 {
+		t.Fatalf("response.status_code = %d, want 503; event=%s", statusCode, event)
+	}
+	if gjson.GetBytes(event, "response.error.status_code").Exists() {
+		t.Fatalf("response.error should preserve upstream error object without injecting status_code: %s", event)
+	}
+	if msg := gjson.GetBytes(event, "response.error.message").String(); msg != "service unavailable" {
+		t.Fatalf("response.error.message = %q, want preserved upstream message", msg)
+	}
+	if code := gjson.GetBytes(event, "response.error.code").String(); code != "server_error" {
+		t.Fatalf("response.error.code = %q, want preserved upstream code", code)
+	}
+}
+
 // TestBuildErrorEvent_NonErrorPassthrough 验证非错误帧不被识别为错误。
 func TestBuildErrorEvent_NonErrorPassthrough(t *testing.T) {
 	r := &WsResponse{}

proxy/wsrelay/executor.go

@@ -121,6 +121,7 @@ func (e *Executor) ExecuteRequestViaWebsocket(
 	// 获取或创建连接。无显式会话的请求（stateless 连接 ID）在确定性 cache key
 	// 的槽位池内复用连接，避免持续高 RPM 下逐请求握手触发上游限流。
 	poolSessionID := sessionID
+	effectiveProxy := effectiveProxyURL(account, proxyOverride)
 	var wc *WsConnection
 	var pr *PendingRequest
 	var err2 error
@@ -137,7 +138,7 @@ func (e *Executor) ExecuteRequestViaWebsocket(
 	sendErr := e.sendRequest(wc, wsBody, pr.RequestID)
 	for retries := 0; sendErr != nil && retries < 2; retries++ {
 		wc.session.RemovePendingRequest(pr.RequestID)
-		e.manager.RemoveConnection(account.ID(), wsURL, poolSessionID, proxyOverride)
+		e.manager.RemoveConnection(account.ID(), wsURL, poolSessionID, effectiveProxy)
 
 		// 短暂退避，避免瞬间重连风暴
 		select {
@@ -154,7 +155,7 @@ func (e *Executor) ExecuteRequestViaWebsocket(
 	}
 	if sendErr != nil {
 		wc.session.RemovePendingRequest(pr.RequestID)
-		e.manager.ReleaseConnection(wc)
+		e.manager.RemoveConnection(account.ID(), wsURL, poolSessionID, effectiveProxy)
 		return nil, fmt.Errorf("发送 WebSocket 请求失败: %w", sendErr)
 	}
 
@@ -393,6 +394,9 @@ func (r *WsResponse) buildErrorEvent(payload []byte) ([]byte, bool) {
 		errObj = fmt.Sprintf(`{"message":%q,"code":%d}`, errMsg, status)
 	}
 	event := fmt.Sprintf(`{"type":"response.failed","response":{"status":"failed","error":%s}}`, errObj)
+	if status > 0 {
+		event = fmt.Sprintf(`{"type":"response.failed","response":{"status":"failed","status_code":%d,"error":%s}}`, status, errObj)
+	}
 	return []byte(event), true
 }
 

proxy/wsrelay/executor_test.go

@@ -1,14 +1,21 @@
 package wsrelay
 
 import (
+	"bufio"
 	"context"
+	"crypto/sha1"
+	"encoding/base64"
+	"fmt"
 	"io"
+	"net"
 	"net/http"
 	"net/http/httptest"
+	"net/url"
 	"strings"
 	"testing"
 	"time"
 
+	"github.com/codex2api/auth"
 	"github.com/codex2api/proxy"
 	"github.com/gorilla/websocket"
 	"github.com/tidwall/gjson"
@@ -226,6 +233,81 @@ func TestWebsocketResponseToHTTPClosesBodyOnContextCancel(t *testing.T) {
 	}
 }
 
+func newClosedTestWebsocketConn(t *testing.T) *websocket.Conn {
+	t.Helper()
+	clientConn, serverConn := net.Pipe()
+	handshakeDone := make(chan struct{})
+	go func() {
+		defer close(handshakeDone)
+		defer serverConn.Close()
+		req, err := http.ReadRequest(bufio.NewReader(serverConn))
+		if err != nil {
+			return
+		}
+		acceptHash := sha1.Sum([]byte(req.Header.Get("Sec-Websocket-Key") + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"))
+		_, _ = fmt.Fprintf(serverConn, "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: %s\r\n\r\n", base64.StdEncoding.EncodeToString(acceptHash[:]))
+	}()
+
+	wsURL, err := url.Parse("ws://example.test/responses")
+	if err != nil {
+		t.Fatalf("parse websocket URL: %v", err)
+	}
+	conn, _, err := websocket.NewClient(clientConn, wsURL, nil, 1024, 1024)
+	if err != nil {
+		t.Fatalf("create test websocket client: %v", err)
+	}
+	<-handshakeDone
+	return conn
+}
+
+func TestExecuteRequestViaWebsocketSendFailureRemovesEffectiveProxyConnection(t *testing.T) {
+	manager := NewManager()
+	t.Cleanup(manager.Stop)
+
+	account := &auth.Account{
+		DBID:        42,
+		AccessToken: "token-123",
+		ProxyURL:    "http://account-proxy.test:8080",
+	}
+	sessionID := "session-1"
+	wsURL, err := buildWebsocketURL(proxy.CodexBaseURL + CodexWsEndpoint)
+	if err != nil {
+		t.Fatalf("buildWebsocketURL: %v", err)
+	}
+	effectiveProxy := effectiveProxyURL(account, "")
+	key := manager.poolKey(account.ID(), wsURL, sessionID, effectiveProxy)
+	session := NewSession(account.ID(), manager)
+	session.SetConnected(true)
+	conn := &WsConnection{
+		conn:    newClosedTestWebsocketConn(t),
+		session: session,
+		URL:     wsURL,
+		PoolKey: key,
+	}
+	conn.SetState(StateConnected)
+	conn.Touch()
+	manager.connections.Store(key, conn)
+	manager.sessions.Store(key, session)
+	manager.probeFunc = func(wc *WsConnection) bool { return true }
+
+	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
+	defer cancel()
+	exec := NewExecutorWithManager(manager)
+	_, err = exec.ExecuteRequestViaWebsocket(ctx, account, []byte(`{"model":"gpt-5.4","input":"hi"}`), sessionID, "", "", nil, http.Header{})
+	if err == nil {
+		t.Fatal("expected final send failure")
+	}
+	if _, ok := manager.connections.Load(key); ok {
+		t.Fatal("expected failed connection keyed by effective account proxy to be removed")
+	}
+	if _, ok := manager.sessions.Load(key); ok {
+		t.Fatal("expected failed session keyed by effective account proxy to be removed")
+	}
+	if conn.IsConnected() {
+		t.Fatal("expected failed connection to be closed")
+	}
+}
+
 func TestSendRequestWritesResponseCreatePayloadDirectly(t *testing.T) {
 	received := make(chan []byte, 1)
 	upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}

proxy/wsrelay/manager.go

@@ -169,6 +169,7 @@ type Manager struct {
 	// 清理定时器
 	cleanupTicker *time.Ticker
 	stopCleanup   chan struct{}
+	stopOnce      sync.Once
 
 	// 连接回调
 	onConnected    func(accountID int64, session *Session)
@@ -252,8 +253,10 @@ func (m *Manager) evictExpired() {
 
 // Stop 停止管理器
 func (m *Manager) Stop() {
-	close(m.stopCleanup)
-	m.closeAll()
+	m.stopOnce.Do(func() {
+		close(m.stopCleanup)
+		m.closeAll()
+	})
 }
 
 // closeAll 关闭所有连接

proxy/wsrelay/manager_test.go

@@ -3,12 +3,100 @@ package wsrelay
 import (
 	"context"
 	"net/http"
+	"sync"
 	"testing"
 	"time"
 
 	"github.com/codex2api/auth"
 )
 
+func TestManagerStopIdempotent(t *testing.T) {
+	manager := NewManager()
+
+	for i := 0; i < 3; i++ {
+		func() {
+			defer func() {
+				if r := recover(); r != nil {
+					t.Fatalf("Stop call %d panicked: %v", i+1, r)
+				}
+			}()
+			manager.Stop()
+		}()
+	}
+}
+
+func TestManagerStopConcurrent(t *testing.T) {
+	manager := NewManager()
+
+	const callers = 32
+	start := make(chan struct{})
+	panicCh := make(chan any, callers)
+	done := make(chan struct{})
+	var wg sync.WaitGroup
+	wg.Add(callers)
+
+	for i := 0; i < callers; i++ {
+		go func() {
+			defer wg.Done()
+			defer func() {
+				if r := recover(); r != nil {
+					panicCh <- r
+				}
+			}()
+			<-start
+			manager.Stop()
+		}()
+	}
+
+	close(start)
+	go func() {
+		wg.Wait()
+		close(done)
+	}()
+
+	select {
+	case <-done:
+	case <-time.After(2 * time.Second):
+		t.Fatal("concurrent Stop calls timed out")
+	}
+	close(panicCh)
+
+	for r := range panicCh {
+		t.Fatalf("concurrent Stop panicked: %v", r)
+	}
+}
+
+func TestRemoveConnectionUsesEffectiveProxyKey(t *testing.T) {
+	manager := NewManager()
+	t.Cleanup(manager.Stop)
+
+	account := &auth.Account{DBID: 42, ProxyURL: " http://proxy-a.example:8080 "}
+	wsURL := "wss://example.test/responses"
+	sessionKey := "session-1"
+	proxyURL := effectiveProxyURL(account, "")
+	key := manager.poolKey(account.ID(), wsURL, sessionKey, proxyURL)
+
+	session := NewSession(account.ID(), manager)
+	session.SetConnected(true)
+	conn := &WsConnection{session: session, URL: wsURL, PoolKey: key}
+	conn.SetState(StateConnected)
+	conn.Touch()
+	manager.connections.Store(key, conn)
+	manager.sessions.Store(key, session)
+
+	manager.RemoveConnection(account.ID(), wsURL, sessionKey, effectiveProxyURL(account, ""))
+
+	if _, ok := manager.connections.Load(key); ok {
+		t.Fatal("expected connection stored under effective proxy key to be removed")
+	}
+	if _, ok := manager.sessions.Load(key); ok {
+		t.Fatal("expected session stored under effective proxy key to be removed")
+	}
+	if conn.IsConnected() {
+		t.Fatal("expected removed connection to be closed")
+	}
+}
+
 func TestAcquireConnectionReusesIdleConnectedConnection(t *testing.T) {
 	manager := NewManager()
 	t.Cleanup(manager.Stop)

proxy/wsrelay/session.go

@@ -2,6 +2,7 @@ package wsrelay
 
 import (
 	"context"
+	"log"
 	"sync"
 	"time"
 
@@ -34,7 +35,7 @@ const (
 	// 复用同一 session 的连接时，等待其空闲的轮询退避参数。
 	AcquireInitialBackoff = 10 * time.Millisecond  // 初始退避
 	AcquireMaxBackoff     = 200 * time.Millisecond // 退避封顶
-	AcquireMaxWait        = 30 * time.Second        // 最大累计等待，超时返回错误
+	AcquireMaxWait        = 30 * time.Second       // 最大累计等待，超时返回错误
 )
 
 // ==================== Pending 请求管理 ====================
@@ -244,7 +245,8 @@ func (s *Session) DeliverStreamChunk(msg *Message) bool {
 			pr.closeMu.Unlock()
 			return true
 		default:
-			// 通道已满，丢弃旧数据
+			// 通道已满，丢弃当前流式数据块并返回 false
+			log.Printf("DeliverStreamChunk stream channel full: account=%d session=%s requestID=%s capacity=%d length=%d; dropping current chunk", s.AccountID, s.ID, msg.RequestID, cap(pr.StreamChan), len(pr.StreamChan))
 			pr.closeMu.Unlock()
 			return false
 		}