package executor import ( "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") 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 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} 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} 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() return conn.WriteMessage(msgType, payload)} // 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 { s.writeMu.Lock() defer s.writeMu.Unlock() // Reply pongs from the same write lock to avoid concurrent writes. return conn.WriteControl(websocket.PongMessage, []byte(appData), time.Now().Add(10*time.Second)) }) 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 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 if sess.readerConn == conn { sess.readerConn = nil } return conn, closer, previousAuthID, previousWSURL, lifecycle} 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()} func (s *codexWebsocketSession) setUpstreamDisconnectError(conn *websocket.Conn, err error) { if s == nil || conn == nil || err == nil { return } s.upstreamDisconnectErrMu.Lock() if s.upstreamDisconnectErrConn == conn && s.upstreamDisconnectErr == nil { s.upstreamDisconnectErr = err } s.upstreamDisconnectErrMu.Unlock()} func (s *codexWebsocketSession) upstreamDisconnectError(conn *websocket.Conn) error { if s == nil || conn == nil { return nil } s.upstreamDisconnectErrMu.RLock() defer s.upstreamDisconnectErrMu.RUnlock() if s.upstreamDisconnectErrConn != conn { return nil } return s.upstreamDisconnectErr} func (s *codexWebsocketSession) notifyUpstreamDisconnect(err error) { if s == nil { return } s.upstreamDisconnectOnce.Do(func() { if s.upstreamDisconnectCh == nil { return } select { case s.upstreamDisconnectCh <- err: default: } close(s.upstreamDisconnectCh) })} func executionSessionIDFromOptions(opts cliproxyexecutor.Options) string { if len(opts.Metadata) == 0 { return "" } raw, ok := opts.Metadata[cliproxyexecutor.ExecutionSessionMetadataKey] if !ok || raw == nil { return "" } switch v := raw.(type) { case string: return strings.TrimSpace(v) case []byte: return strings.TrimSpace(string(v)) default: return "" }} func (e *CodexWebsocketsExecutor) getOrCreateSession(sessionID string) *codexWebsocketSession { sessionID = strings.TrimSpace(sessionID) if sessionID == "" { return nil } if e == nil { return nil } store := e.store if store == nil { store = globalCodexWebsocketSessionStore } store.mu.Lock() defer store.mu.Unlock() if store.sessions == nil { store.sessions = make(map[string]*codexWebsocketSession) } if sess, ok := store.sessions[sessionID]; ok && sess != nil { return sess } sess := &codexWebsocketSession{ sessionID: sessionID, upstreamDisconnectCh: make(chan error, 1), } store.sessions[sessionID] = sess return sess} func (e *CodexWebsocketsExecutor) UpstreamDisconnectChan(sessionID string) <-chan error { sess := e.getOrCreateSession(sessionID) if sess == nil { return nil } return sess.upstreamDisconnectCh} func (e *CodexWebsocketsExecutor) ensureUpstreamConn(ctx context.Context, auth *cliproxyauth.Auth, sess *codexWebsocketSession, authID string, wsURL string, headers http.Header) (*websocket.Conn, *websocketConnectionCloser, *http.Response, error) { if sess == nil { return e.dialCodexWebsocket(ctx, auth, wsURL, headers) } if staleConn, staleCloser, staleAuthID, staleWSURL, staleLifecycle := detachMismatchedWebsocketSessionConn(sess, authID, wsURL); staleConn != nil { logCodexWebsocketDisconnected(sess.sessionID, staleAuthID, staleWSURL, "target_changed", nil) if staleCloser != nil { if errClose := staleCloser.Close(); errClose != nil { log.Errorf("codex websockets executor: close stale websocket error: %v", errClose) } } if staleLifecycle != nil { staleLifecycle.End("target_changed") } } sess.connMu.Lock() conn := sess.conn closer := sess.connCloser readerConn := sess.readerConn sess.connMu.Unlock() if conn != nil { if readerConn != conn { sess.connMu.Lock() sess.readerConn = conn sess.connMu.Unlock() sess.configureConn(conn) go e.readUpstreamLoop(sess, conn) } return conn, closer, nil, nil } conn, closer, resp, errDial := e.dialCodexWebsocket(ctx, auth, wsURL, headers) if errDial != nil { return nil, closer, resp, errDial } sess.connMu.Lock() if sess.conn != nil { previous := sess.conn previousCloser := sess.connCloser sess.connMu.Unlock()Lines 1–500 of 788. Load more
package executor import ( "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" ) 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 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 } 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 } 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() return conn.WriteMessage(msgType, payload) } // 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 { s.writeMu.Lock() defer s.writeMu.Unlock() // Reply pongs from the same write lock to avoid concurrent writes. return conn.WriteControl(websocket.PongMessage, []byte(appData), time.Now().Add(10*time.Second)) }) 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 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 if sess.readerConn == conn { sess.readerConn = nil } return conn, closer, previousAuthID, previousWSURL, lifecycle } 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() } func (s *codexWebsocketSession) setUpstreamDisconnectError(conn *websocket.Conn, err error) { if s == nil || conn == nil || err == nil { return } s.upstreamDisconnectErrMu.Lock() if s.upstreamDisconnectErrConn == conn && s.upstreamDisconnectErr == nil { s.upstreamDisconnectErr = err } s.upstreamDisconnectErrMu.Unlock() } func (s *codexWebsocketSession) upstreamDisconnectError(conn *websocket.Conn) error { if s == nil || conn == nil { return nil } s.upstreamDisconnectErrMu.RLock() defer s.upstreamDisconnectErrMu.RUnlock() if s.upstreamDisconnectErrConn != conn { return nil } return s.upstreamDisconnectErr } func (s *codexWebsocketSession) notifyUpstreamDisconnect(err error) { if s == nil { return } s.upstreamDisconnectOnce.Do(func() { if s.upstreamDisconnectCh == nil { return } select { case s.upstreamDisconnectCh <- err: default: } close(s.upstreamDisconnectCh) }) } func executionSessionIDFromOptions(opts cliproxyexecutor.Options) string { if len(opts.Metadata) == 0 { return "" } raw, ok := opts.Metadata[cliproxyexecutor.ExecutionSessionMetadataKey] if !ok || raw == nil { return "" } switch v := raw.(type) { case string: return strings.TrimSpace(v) case []byte: return strings.TrimSpace(string(v)) default: return "" } } func (e *CodexWebsocketsExecutor) getOrCreateSession(sessionID string) *codexWebsocketSession { sessionID = strings.TrimSpace(sessionID) if sessionID == "" { return nil } if e == nil { return nil } store := e.store if store == nil { store = globalCodexWebsocketSessionStore } store.mu.Lock() defer store.mu.Unlock() if store.sessions == nil { store.sessions = make(map[string]*codexWebsocketSession) } if sess, ok := store.sessions[sessionID]; ok && sess != nil { return sess } sess := &codexWebsocketSession{ sessionID: sessionID, upstreamDisconnectCh: make(chan error, 1), } store.sessions[sessionID] = sess return sess } func (e *CodexWebsocketsExecutor) UpstreamDisconnectChan(sessionID string) <-chan error { sess := e.getOrCreateSession(sessionID) if sess == nil { return nil } return sess.upstreamDisconnectCh } func (e *CodexWebsocketsExecutor) ensureUpstreamConn(ctx context.Context, auth *cliproxyauth.Auth, sess *codexWebsocketSession, authID string, wsURL string, headers http.Header) (*websocket.Conn, *websocketConnectionCloser, *http.Response, error) { if sess == nil { return e.dialCodexWebsocket(ctx, auth, wsURL, headers) } if staleConn, staleCloser, staleAuthID, staleWSURL, staleLifecycle := detachMismatchedWebsocketSessionConn(sess, authID, wsURL); staleConn != nil { logCodexWebsocketDisconnected(sess.sessionID, staleAuthID, staleWSURL, "target_changed", nil) if staleCloser != nil { if errClose := staleCloser.Close(); errClose != nil { log.Errorf("codex websockets executor: close stale websocket error: %v", errClose) } } if staleLifecycle != nil { staleLifecycle.End("target_changed") } } sess.connMu.Lock() conn := sess.conn closer := sess.connCloser readerConn := sess.readerConn sess.connMu.Unlock() if conn != nil { if readerConn != conn { sess.connMu.Lock() sess.readerConn = conn sess.connMu.Unlock() sess.configureConn(conn) go e.readUpstreamLoop(sess, conn) } return conn, closer, nil, nil } conn, closer, resp, errDial := e.dialCodexWebsocket(ctx, auth, wsURL, headers) if errDial != nil { return nil, closer, resp, errDial } sess.connMu.Lock() if sess.conn != nil { previous := sess.conn previousCloser := sess.connCloser sess.connMu.Unlock()