feat(wsrelay): reuse idle codex websocket sessions

Author: DeliciousBuding

proxy/wsrelay/executor.go

@@ -86,7 +86,7 @@ func (e *Executor) ExecuteRequestViaWebsocket(
 	}
 
 	// 准备请求头
-	headers := e.prepareWebsocketHeaders(accessToken, accountIDStr, apiKey, deviceCfg, ginHeaders)
+	headers := e.prepareWebsocketHeaders(accessToken, accountIDStr, sessionID, apiKey, deviceCfg, ginHeaders)
 
 	// 获取或创建连接
 	wc, err := e.manager.AcquireConnection(ctx, account, wsURL, headers, proxyOverride)
@@ -162,7 +162,7 @@ func (e *Executor) prepareWebsocketBody(body []byte, sessionID string) []byte {
 }
 
 // prepareWebsocketHeaders 准备 WebSocket 请求头
-func (e *Executor) prepareWebsocketHeaders(accessToken, accountID, apiKey string, deviceCfg *proxy.DeviceProfileConfig, ginHeaders http.Header) http.Header {
+func (e *Executor) prepareWebsocketHeaders(accessToken, accountID, sessionID, apiKey string, deviceCfg *proxy.DeviceProfileConfig, ginHeaders http.Header) http.Header {
 	headers := http.Header{}
 
 	// 认证头
@@ -207,6 +207,9 @@ func (e *Executor) prepareWebsocketHeaders(accessToken, accountID, apiKey string
 	if accountID != "" {
 		headers.Set("Chatgpt-Account-Id", accountID)
 	}
+	if sessionID = strings.TrimSpace(sessionID); sessionID != "" {
+		headers.Set("Conversation_id", sessionID)
+	}
 
 	return headers
 }

proxy/wsrelay/executor_test.go

@@ -22,7 +22,7 @@ func TestPrepareWebsocketHeadersUsesConfiguredDefaultsAndBetaFeatures(t *testing
 		"Originator": []string{"custom-originator"},
 	}
 
-	headers := exec.prepareWebsocketHeaders("token-123", "42", "api-key-1", cfg, ginHeaders)
+	headers := exec.prepareWebsocketHeaders("token-123", "42", "session-123", "api-key-1", cfg, ginHeaders)
 
 	if got := headers.Get("Authorization"); got != "Bearer token-123" {
 		t.Fatalf("Authorization = %q", got)
@@ -45,4 +45,7 @@ func TestPrepareWebsocketHeadersUsesConfiguredDefaultsAndBetaFeatures(t *testing
 	if got := headers.Get("Chatgpt-Account-Id"); got != "42" {
 		t.Fatalf("Chatgpt-Account-Id = %q", got)
 	}
+	if got := headers.Get("Conversation_id"); got != "session-123" {
+		t.Fatalf("Conversation_id = %q", got)
+	}
 }

proxy/wsrelay/manager.go

@@ -90,7 +90,10 @@ func (wc *WsConnection) Close() error {
 		if wc.onDisconnected != nil && wc.session != nil {
 			wc.onDisconnected(wc.session.AccountID)
 		}
-		return wc.conn.Close()
+		if wc.conn != nil {
+			return wc.conn.Close()
+		}
+		return nil
 	}
 	return nil
 }
@@ -274,9 +277,13 @@ func (m *Manager) AcquireConnection(
 ) (*WsConnection, error) {
 	key := m.poolKey(account.ID(), wsURL)
 
-	// 清理可能存在的旧连接（避免泄漏）
-	if v, ok := m.connections.LoadAndDelete(key); ok {
+	if v, ok := m.connections.Load(key); ok {
 		wc := v.(*WsConnection)
+		if wc.IsConnected() && !wc.IsExpired() && wc.session != nil && wc.session.PendingCount() == 0 {
+			wc.Touch()
+			return wc, nil
+		}
+		m.connections.Delete(key)
 		wc.Close()
 	}
 
@@ -473,4 +480,4 @@ func ShutdownManager() {
 	if globalManager != nil {
 		globalManager.Stop()
 	}
-}
+}

proxy/wsrelay/manager_test.go

@@ -0,0 +1,38 @@
+package wsrelay
+
+import (
+	"context"
+	"net/http"
+	"testing"
+
+	"github.com/codex2api/auth"
+)
+
+func TestAcquireConnectionReusesIdleConnectedConnection(t *testing.T) {
+	manager := NewManager()
+	t.Cleanup(manager.Stop)
+
+	account := &auth.Account{DBID: 42}
+	wsURL := "wss://example.test/responses"
+	key := manager.poolKey(account.ID(), wsURL)
+
+	session := NewSession(account.ID(), manager)
+	session.SetConnected(true)
+	conn := &WsConnection{
+		session:  session,
+		URL:      wsURL,
+		httpResp: &http.Response{StatusCode: http.StatusSwitchingProtocols},
+	}
+	conn.SetState(StateConnected)
+	conn.Touch()
+	manager.connections.Store(key, conn)
+	manager.sessions.Store(key, session)
+
+	got, err := manager.AcquireConnection(context.Background(), account, wsURL, http.Header{}, "")
+	if err != nil {
+		t.Fatalf("AcquireConnection() error = %v", err)
+	}
+	if got != conn {
+		t.Fatal("expected existing connection to be reused")
+	}
+}

proxy/wsrelay/session.go

@@ -56,7 +56,7 @@ type PendingRequest struct {
 	Cancel context.CancelFunc
 
 	// 关闭标志，防止重复关闭
-	closed bool
+	closed  bool
 	closeMu sync.Mutex
 }
 
@@ -190,6 +190,19 @@ func (s *Session) RemovePendingRequest(requestID string) {
 	}
 }
 
+// PendingCount returns the number of in-flight requests bound to this session.
+func (s *Session) PendingCount() int {
+	if s == nil {
+		return 0
+	}
+	count := 0
+	s.pending.Range(func(_, _ any) bool {
+		count++
+		return true
+	})
+	return count
+}
+
 // DeliverResponse 投递响应到等待请求
 func (s *Session) DeliverResponse(msg *Message) bool {
 	if pr, ok := s.GetPendingRequest(msg.RequestID); ok {
@@ -323,4 +336,4 @@ func (s *Session) ClearPendingRequests() {
 		s.pending.Delete(key)
 		return true
 	})
-}
+}