Something went wrong. Try again.
sloproxy
Something went wrong. Try again.
22 kB · 818 lines
Go
at commit 291cfb87
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819package 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 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}
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)}
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 { 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 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 (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() if errClose := closer.Close(); errClose != nil { log.Errorf("codex websockets executor: close websocket error: %v", errClose) } return previous, previousCloser, nil, nil } sess.conn = conn sess.connCloser = closer sess.multiAgentV2OptimizedConn = nil sess.wsURL = wsURL sess.authID = authID sess.readerConn = conn sess.connMu.Unlock()
sess.configureConn(conn) go e.readUpstreamLoop(sess, conn) logCodexWebsocketConnected(sess.sessionID, authID, wsURL) return conn, closer, resp, nil}
func (e *CodexWebsocketsExecutor) readUpstreamLoop(sess *codexWebsocketSession, conn *websocket.Conn) { if e == nil || sess == nil || conn == nil { return } for { _ = conn.SetReadDeadline(time.Now().Add(codexResponsesWebsocketIdleTimeout)) msgType, payload, errRead := conn.ReadMessage() if errRead != nil { invalidate := func() { e.invalidateUpstreamConn(sess, conn, "upstream_disconnected", errRead) } invalidated := false ch, done := sess.activeForConn(conn) if ch != nil { invalidated = sendTerminalWebsocketRead(ch, done, codexWebsocketRead{conn: conn, err: errRead}, invalidate) if sess.clearActive(conn, ch) { close(ch) } } if !invalidated { invalidate() } return }
if msgType != websocket.TextMessage { if msgType == websocket.BinaryMessage { errBinary := fmt.Errorf("codex websockets executor: unexpected binary message") invalidate := func() { e.invalidateUpstreamConn(sess, conn, "unexpected_binary", errBinary) } invalidated := false ch, done := sess.activeForConn(conn) if ch != nil { invalidated = sendTerminalWebsocketRead(ch, done, codexWebsocketRead{conn: conn, err: errBinary}, invalidate) if sess.clearActive(conn, ch) { close(ch) } } if !invalidated { invalidate() } return } continue }
ch, done := sess.activeForConn(conn) if ch == nil { continue } select { case ch <- codexWebsocketRead{conn: conn, msgType: msgType, payload: payload}: case <-done: } }}
func (e *CodexWebsocketsExecutor) invalidateUpstreamConn(sess *codexWebsocketSession, conn *websocket.Conn, reason string, err error) { e.invalidateUpstreamConnWithNotify(sess, conn, reason, err, true)}
func (e *CodexWebsocketsExecutor) invalidateUpstreamConnWithoutDisconnectNotify(sess *codexWebsocketSession, conn *websocket.Conn, reason string, err error) { e.invalidateUpstreamConnWithNotify(sess, conn, reason, err, false)}
func (e *CodexWebsocketsExecutor) invalidateUpstreamConnWithNotify(sess *codexWebsocketSession, conn *websocket.Conn, reason string, err error, notify bool) { if sess == nil || conn == nil { return }
sess.connMu.Lock() current := sess.conn authID := sess.authID wsURL := sess.wsURL sessionID := sess.sessionID if current == nil || current != conn { sess.connMu.Unlock() return } 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 } sess.connMu.Unlock()
logCodexWebsocketDisconnected(sessionID, authID, wsURL, reason, err) if notify { sess.notifyUpstreamDisconnect(err) } if closer != nil { if errClose := closer.Close(); errClose != nil { log.Errorf("codex websockets executor: close websocket error: %v", errClose) } } if lifecycle != nil { lifecycle.End(reason) }}
func (e *CodexWebsocketsExecutor) CloseExecutionSession(sessionID string) { sessionID = strings.TrimSpace(sessionID) if e == nil { return } if sessionID == "" { return } if sessionID == cliproxyauth.CloseAllExecutionSessionsID { e.closeAllExecutionSessions("executor_shutdown") return }
store := e.store if store == nil { store = globalCodexWebsocketSessionStore } store.mu.Lock() sess := store.sessions[sessionID] delete(store.sessions, sessionID) store.mu.Unlock()
e.closeExecutionSession(sess, "session_closed")}
func (e *CodexWebsocketsExecutor) closeAllExecutionSessions(reason string) { if e == nil { return }
store := e.store if store == nil { store = globalCodexWebsocketSessionStore } store.mu.Lock() sessions := make([]*codexWebsocketSession, 0, len(store.sessions)) for sessionID, sess := range store.sessions { delete(store.sessions, sessionID) if sess != nil { sessions = append(sessions, sess) } } store.mu.Unlock()
for i := range sessions { e.closeExecutionSession(sessions[i], reason) }}
func (e *CodexWebsocketsExecutor) closeExecutionSession(sess *codexWebsocketSession, reason string) { closeCodexWebsocketSession(sess, reason)}
func closeCodexWebsocketSession(sess *codexWebsocketSession, reason string) { if sess == nil { return } reason = strings.TrimSpace(reason) if reason == "" { reason = "session_closed" }
sess.connMu.Lock() conn := sess.conn authID := sess.authID wsURL := 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 } sessionID := sess.sessionID sess.connMu.Unlock()
if conn != nil { logCodexWebsocketDisconnected(sessionID, authID, wsURL, reason, nil) if closer != nil { if errClose := closer.Close(); errClose != nil { log.Errorf("codex websockets executor: close websocket error: %v", errClose) } } } if lifecycle != nil { lifecycle.End(reason) }}
func logCodexWebsocketConnected(sessionID string, authID string, wsURL string) { log.Infof("codex websockets: upstream connected session=%s auth=%s url=%s", strings.TrimSpace(sessionID), strings.TrimSpace(authID), strings.TrimSpace(wsURL))}
func logCodexWebsocketDisconnected(sessionID string, authID string, wsURL string, reason string, err error) { if err != nil { log.Infof("codex websockets: upstream disconnected session=%s auth=%s url=%s reason=%s err=%v", strings.TrimSpace(sessionID), strings.TrimSpace(authID), strings.TrimSpace(wsURL), strings.TrimSpace(reason), err) return } log.Infof("codex websockets: upstream disconnected session=%s auth=%s url=%s reason=%s", strings.TrimSpace(sessionID), strings.TrimSpace(authID), strings.TrimSpace(wsURL), strings.TrimSpace(reason))}
// CloseCodexWebsocketSessionsForAuthID closes all active Codex upstream websocket sessions// associated with the supplied auth ID.func CloseCodexWebsocketSessionsForAuthID(authID string, reason string) { authID = strings.TrimSpace(authID) if authID == "" { return } reason = strings.TrimSpace(reason) if reason == "" { reason = "auth_removed" }
store := globalCodexWebsocketSessionStore if store == nil { return }
type sessionItem struct { sessionID string sess *codexWebsocketSession }
store.mu.Lock() items := make([]sessionItem, 0, len(store.sessions)) for sessionID, sess := range store.sessions { items = append(items, sessionItem{sessionID: sessionID, sess: sess}) } store.mu.Unlock()
matches := make([]sessionItem, 0) for i := range items { sess := items[i].sess if sess == nil { continue } sess.connMu.Lock() sessAuthID := strings.TrimSpace(sess.authID) sess.connMu.Unlock() if sessAuthID == authID { matches = append(matches, items[i]) } } if len(matches) == 0 { return }
toClose := make([]*codexWebsocketSession, 0, len(matches)) store.mu.Lock() for i := range matches { current, ok := store.sessions[matches[i].sessionID] if !ok || current == nil || current != matches[i].sess { continue } delete(store.sessions, matches[i].sessionID) toClose = append(toClose, current) } store.mu.Unlock()
for i := range toClose { closeCodexWebsocketSession(toClose[i], reason) }}