package executor import (	"bytes"	"context"	"fmt"	"net/http"	"strings"	"sync"	"time" 	"github.com/gorilla/websocket"	cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"	cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"	log "github.com/sirupsen/logrus"	"github.com/tidwall/gjson") type codexWebsocketSessionStore struct {	mu       sync.Mutex	sessions map[string]*codexWebsocketSession} var globalCodexWebsocketSessionStore = &codexWebsocketSessionStore{	sessions: make(map[string]*codexWebsocketSession),} type websocketConnectionCloser struct {	conn *websocket.Conn	once sync.Once	err  error} func newWebsocketConnectionCloser(conn *websocket.Conn) *websocketConnectionCloser {	if conn == nil {		return nil	}	return &websocketConnectionCloser{conn: conn}} func (c *websocketConnectionCloser) Close() error {	if c == nil || c.conn == nil {		return nil	}	c.once.Do(func() {		c.err = c.conn.Close()	})	return c.err} type codexWebsocketSession struct {	sessionID string 	reqMu sync.Mutex 	connMu                    sync.Mutex	conn                      *websocket.Conn	connCloser                *websocketConnectionCloser	wsURL                     string	authID                    string	multiAgentV2OptimizedConn *websocket.Conn	lifecycleBindMu           sync.Mutex	lifecycle                 cliproxyexecutor.ExecutionLifecycle	lifecycleModel            string 	writeMu sync.Mutex 	activeMu     sync.Mutex	activeConn   *websocket.Conn	activeCh     chan codexWebsocketRead	activeDone   <-chan struct{}	activeCancel context.CancelFunc 	readerConn *websocket.Conn 	upstreamDisconnectOnce    sync.Once	upstreamDisconnectCh      chan error	upstreamDisconnectErrMu   sync.RWMutex	upstreamDisconnectErrConn *websocket.Conn	upstreamDisconnectErr     error 	lastEventMu   sync.Mutex	lastEventConn *websocket.Conn	lastEventType string} type codexWebsocketRead struct {	conn    *websocket.Conn	msgType int	payload []byte	err     error} func (s *codexWebsocketSession) setActive(conn *websocket.Conn, ch chan codexWebsocketRead) {	if s == nil {		return	}	s.activeMu.Lock()	if s.activeCancel != nil {		s.activeCancel()		s.activeCancel = nil		s.activeDone = nil	}	s.activeConn = conn	s.activeCh = ch	if conn != nil && ch != nil {		activeCtx, activeCancel := context.WithCancel(context.Background())		s.activeDone = activeCtx.Done()		s.activeCancel = activeCancel	}	s.activeMu.Unlock()} func (s *codexWebsocketSession) activate(conn *websocket.Conn) chan codexWebsocketRead {	if s == nil || conn == nil {		return nil	}	ch := make(chan codexWebsocketRead, 4096)	s.setActive(conn, ch)	return ch} func (s *codexWebsocketSession) activeForConn(conn *websocket.Conn) (chan codexWebsocketRead, <-chan struct{}) {	if s == nil || conn == nil {		return nil, nil	}	s.activeMu.Lock()	defer s.activeMu.Unlock()	if s.activeConn != conn {		return nil, nil	}	return s.activeCh, s.activeDone} func clearRetryActiveState(sess *codexWebsocketSession, conn *websocket.Conn, ch chan codexWebsocketRead) bool {	if sess == nil {		return false	}	return sess.clearActive(conn, ch)} func (s *codexWebsocketSession) clearActive(conn *websocket.Conn, ch chan codexWebsocketRead) bool {	if s == nil {		return false	}	s.activeMu.Lock()	defer s.activeMu.Unlock()	if s.activeConn != conn || s.activeCh != ch {		return false	}	s.activeConn = nil	s.activeCh = nil	if s.activeCancel != nil {		s.activeCancel()	}	s.activeCancel = nil	s.activeDone = nil	return true} const codexWebsocketWriteChunkSize = 32 * 1024 var (	testWebsocketWritePayloadHook func(conn *websocket.Conn)	testWebsocketWriteChunkHook   func(chunkIndex int, totalChunks int)) func (s *codexWebsocketSession) writeMessage(conn *websocket.Conn, msgType int, payload []byte) error {	if s == nil {		return fmt.Errorf("codex websockets executor: session is nil")	}	if conn == nil {		return fmt.Errorf("codex websockets executor: websocket conn is nil")	}	s.writeMu.Lock()	defer s.writeMu.Unlock()	if testWebsocketWritePayloadHook != nil {		testWebsocketWritePayloadHook(conn)	}	if len(payload) <= codexWebsocketWriteChunkSize {		return conn.WriteMessage(msgType, payload)	}	w, errNext := conn.NextWriter(msgType)	if errNext != nil {		return errNext	}	totalChunks := (len(payload) + codexWebsocketWriteChunkSize - 1) / codexWebsocketWriteChunkSize	for i := 0; i < len(payload); i += codexWebsocketWriteChunkSize {		end := i + codexWebsocketWriteChunkSize		if end > len(payload) {			end = len(payload)		}		chunkIdx := i / codexWebsocketWriteChunkSize		if testWebsocketWriteChunkHook != nil {			testWebsocketWriteChunkHook(chunkIdx, totalChunks)		}		if _, errWrite := w.Write(payload[i:end]); errWrite != nil {			_ = w.Close()			return errWrite		}	}	return w.Close()} func (s *codexWebsocketSession) setMultiAgentV2Optimized(conn *websocket.Conn, optimized bool) {	if s == nil || conn == nil {		return	}	s.connMu.Lock()	if s.conn == conn {		if optimized {			s.multiAgentV2OptimizedConn = conn		} else {			s.multiAgentV2OptimizedConn = nil		}	}	s.connMu.Unlock()} func (s *codexWebsocketSession) isMultiAgentV2Optimized(conn *websocket.Conn) bool {	if s == nil || conn == nil {		return false	}	s.connMu.Lock()	defer s.connMu.Unlock()	return s.conn == conn && s.multiAgentV2OptimizedConn == conn} // sendTerminalWebsocketRead reports whether it invalidated a full channel's connection before waiting.func sendTerminalWebsocketRead(ch chan<- codexWebsocketRead, done <-chan struct{}, event codexWebsocketRead, invalidate func()) bool {	select {	case ch <- event:		return false	case <-done:		return false	default:	} 	invalidated := invalidate != nil	if invalidated {		invalidate()	}	select {	case ch <- event:	case <-done:	}	return invalidated} func (s *codexWebsocketSession) configureConn(conn *websocket.Conn) {	if s == nil || conn == nil {		return	}	s.resetUpstreamDisconnectError(conn)	conn.SetPingHandler(func(appData string) error {		sessionID := ""		if s != nil {			sessionID = s.sessionID		}		sessionKind := sessionObjectKind(s)		log.Debugf("codex websockets: upstream ping received session=%s session_object=%s ping_bytes=%d", sessionID, sessionKind, len(appData))		log.Debugf("codex websockets: upstream pong write started session=%s session_object=%s", sessionID, sessionKind)		start := time.Now()		// Gorilla websocket allows concurrent WriteControl with WriteMessage.		// Avoid writeMu here so keepalive pongs are not starved by long payload writes.		errPong := conn.WriteControl(websocket.PongMessage, []byte(appData), time.Now().Add(10*time.Second))		if errPong != nil {			log.Warnf("codex websockets: upstream pong write failed session=%s session_object=%s duration=%v err=%v", sessionID, sessionKind, time.Since(start), errPong)		} else {			log.Debugf("codex websockets: upstream pong replied session=%s session_object=%s duration=%v", sessionID, sessionKind, time.Since(start))		}		return errPong	})	defaultCloseHandler := conn.CloseHandler()	conn.SetCloseHandler(func(code int, text string) error {		s.setUpstreamDisconnectError(conn, &websocket.CloseError{Code: code, Text: text})		return defaultCloseHandler(code, text)	})} func (s *codexWebsocketSession) bindExecutionLifecycle(opts cliproxyexecutor.Options, conn *websocket.Conn, closer *websocketConnectionCloser, model string) error {	if closer == nil {		return fmt.Errorf("codex websockets executor: websocket connection closer is nil")	}	if s == nil {		return cliproxyexecutor.BindExecutionResource(opts, closer)	}	lifecycle := opts.ExecutionLifecycle	if lifecycle == nil || conn == nil {		return nil	} 	s.lifecycleBindMu.Lock()	defer s.lifecycleBindMu.Unlock() 	s.connMu.Lock()	if s.conn == conn && s.connCloser == nil {		s.connCloser = closer	}	alreadyBound := s.conn == conn && s.connCloser == closer && s.lifecycle == lifecycle	s.connMu.Unlock()	if alreadyBound {		return nil	} 	if errBind := lifecycle.Bind(func() error {		return s.closeBoundConnection(conn, closer, lifecycle)	}); errBind != nil {		return errBind	}	if retained, ok := lifecycle.(interface{ Retain() }); ok {		retained.Retain()	} 	s.connMu.Lock()	if s.conn != conn || s.connCloser != closer {		s.connMu.Unlock()		return fmt.Errorf("codex websockets executor: websocket connection closed during lifecycle bind")	}	previous := s.lifecycle	s.lifecycle = lifecycle	s.lifecycleModel = strings.TrimSpace(model)	s.connMu.Unlock()	if previous != nil && previous != lifecycle {		previous.End("target_replaced")	}	return nil} func (s *codexWebsocketSession) closeBoundConnection(conn *websocket.Conn, closer *websocketConnectionCloser, lifecycle cliproxyexecutor.ExecutionLifecycle) error {	if s == nil || conn == nil {		return nil	}	s.detachConnection(conn, lifecycle)	errClose := closer.Close()	go lifecycle.End("connection_closed")	return errClose} func (s *codexWebsocketSession) detachConnection(conn *websocket.Conn, lifecycle cliproxyexecutor.ExecutionLifecycle) *websocketConnectionCloser {	if s == nil || conn == nil {		return nil	}	s.connMu.Lock()	var closer *websocketConnectionCloser	matched := s.conn == conn	if matched {		closer = s.connCloser		s.conn = nil		s.connCloser = nil		s.multiAgentV2OptimizedConn = nil		if s.readerConn == conn {			s.readerConn = nil		}	}	if (lifecycle == nil && matched) || (lifecycle != nil && s.lifecycle == lifecycle) {		s.lifecycle = nil		s.lifecycleModel = ""	}	s.connMu.Unlock()	return closer} func closeWebsocketAfterBindFailure(sess *codexWebsocketSession, conn *websocket.Conn, closer *websocketConnectionCloser) {	if conn == nil || closer == nil {		return	}	if sess != nil {		sess.detachConnection(conn, nil)	}	if errClose := closer.Close(); errClose != nil {		log.Errorf("websockets executor: close lifecycle bind failure connection error: %v", errClose)	}} func websocketSessionTargetChanged(sess *codexWebsocketSession, authID string, wsURL string) bool {	if sess == nil {		return false	} 	sess.connMu.Lock()	defer sess.connMu.Unlock()	if strings.TrimSpace(sess.authID) == "" && strings.TrimSpace(sess.wsURL) == "" {		return false	}	return strings.TrimSpace(sess.authID) != strings.TrimSpace(authID) || strings.TrimSpace(sess.wsURL) != strings.TrimSpace(wsURL)} func existingWebsocketSessionConn(sess *codexWebsocketSession, authID string, wsURL string) (*websocket.Conn, *websocketConnectionCloser) {	if sess == nil {		return nil, nil	}	sess.connMu.Lock()	conn := sess.conn	closer := sess.connCloser	matches := conn != nil && closer != nil &&		strings.TrimSpace(sess.authID) == strings.TrimSpace(authID) &&		strings.TrimSpace(sess.wsURL) == strings.TrimSpace(wsURL)	sess.connMu.Unlock()	if !matches || sess.upstreamDisconnectError(conn) != nil {		return nil, nil	}	return conn, closer} func detachMismatchedWebsocketSessionConn(sess *codexWebsocketSession, authID string, wsURL string) (*websocket.Conn, *websocketConnectionCloser, string, string, cliproxyexecutor.ExecutionLifecycle) {	if sess == nil {		return nil, nil, "", "", nil	} 	sess.connMu.Lock()	defer sess.connMu.Unlock()	conn := sess.conn	if conn == nil || (strings.TrimSpace(sess.authID) == strings.TrimSpace(authID) && strings.TrimSpace(sess.wsURL) == strings.TrimSpace(wsURL)) {		return nil, nil, "", "", nil	} 	previousAuthID := sess.authID	previousWSURL := sess.wsURL	lifecycle := sess.lifecycle	closer := sess.connCloser	sess.lifecycle = nil	sess.lifecycleModel = ""	sess.conn = nil	sess.connCloser = nil	sess.multiAgentV2OptimizedConn = nil	if sess.readerConn == conn {		sess.readerConn = nil	}	return conn, closer, previousAuthID, previousWSURL, lifecycle} func configureRawCodexWebsocketConn(conn *websocket.Conn, authID string, wsURL string) {	if conn == nil {		return	}	conn.SetPingHandler(func(appData string) error {		log.Debugf("codex websockets: upstream ping received session= session_object=none ping_bytes=%d", len(appData))		log.Debugf("codex websockets: upstream pong write started session= session_object=none")		start := time.Now()		errPong := conn.WriteControl(websocket.PongMessage, []byte(appData), time.Now().Add(10*time.Second))		if errPong != nil {			log.Warnf("codex websockets: upstream pong write failed session= session_object=none duration=%v err=%v", time.Since(start), errPong)		} else {			log.Debugf("codex websockets: upstream pong replied session= session_object=none duration=%v", time.Since(start))		}		return errPong	})} func (s *codexWebsocketSession) setLastEventType(conn *websocket.Conn, eventType string) {	if s == nil || conn == nil || eventType == "" {		return	}	s.lastEventMu.Lock()	if s.lastEventConn == conn {		s.lastEventType = eventType	}	s.lastEventMu.Unlock()} func (s *codexWebsocketSession) getLastEventType(conn *websocket.Conn) string {	if s == nil || conn == nil {		return ""	}	s.lastEventMu.Lock()	defer s.lastEventMu.Unlock()	if s.lastEventConn != conn {		return ""	}	return s.lastEventType} func newEphemeralCodexWebsocketSession() *codexWebsocketSession {	return &codexWebsocketSession{		sessionID:            "",		upstreamDisconnectCh: make(chan error, 1),	}} func (s *codexWebsocketSession) resetUpstreamDisconnectError(conn *websocket.Conn) {	if s == nil || conn == nil {		return	}	s.upstreamDisconnectErrMu.Lock()	s.upstreamDisconnectErrConn = conn	s.upstreamDisconnectErr = nil	s.upstreamDisconnectErrMu.Unlock() 	s.lastEventMu.Lock()	s.lastEventConn = conn	s.lastEventType = ""	s.lastEventMu.Unlock()} func (s *codexWebsocketSession) setUpstreamDisconnectError(conn *websocket.Conn, err error) {	if s == nil || conn == nil || err == nil {		return	}	s.upstreamDisconnectErrMu.Lock()

Lines 1–500 of 983. Load more