lywsvip/CLIProxyAPI
1
0
Fork 0
mirror of https://github.com/router-for-me/CLIProxyAPI.git synced 2026-09-09 01:20:13 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #4522 from sususu98/fix/responses-ws-continuity

Browse Source
This commit is contained in:
Luis Pater
2026-07-23 14:31:41 +08:00
parent 285322cd97 a661172b1f
commit e5729872fa
13 changed files with 1987 additions and 238 deletions
Download Patch File Download Diff File Expand all files Collapse all files

85
internal/runtime/executor/codex_websockets_executor.go
View File

@@ -352,6 +352,23 @@ func websocketSessionTargetChanged(sess *codexWebsocketSession, authID string, w
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
@@ -532,14 +549,25 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut
}
helps.RecordAPIWebsocketRequest(ctx, e.cfg, wsReqLog)
conn, closer, respHS, errDial := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders)
var conn *websocket.Conn
var closer *websocketConnectionCloser
var respHS *http.Response
var errDial error
if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
conn, closer = existingWebsocketSessionConn(sess, authID, wsURL)
if conn == nil {
return resp, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
}
} else {
conn, closer, respHS, errDial = e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders)
}
if errDial != nil {
bodyErr := websocketHandshakeBody(respHS)
if respHS != nil {
helps.RecordAPIWebsocketUpgradeRejection(ctx, e.cfg, websocketUpgradeRequestLog(wsReqLog), respHS.StatusCode, respHS.Header.Clone(), bodyErr)
}
if respHS != nil && respHS.StatusCode == http.StatusUpgradeRequired {
if opts.ExecutionLifecycle != nil {
if opts.ExecutionLifecycle != nil || cliproxyexecutor.DownstreamWebsocket(ctx) {
return resp, statusErr{code: respHS.StatusCode, msg: string(bodyErr)}
}
return e.CodexExecutor.Execute(ctx, auth, req, opts)
@@ -582,6 +610,14 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut
if errSend := writeCodexWebsocketMessage(sess, conn, wsReqBody); errSend != nil {
errSend = mapCodexWebsocketWriteError(sess, conn, errSend)
if sess != nil {
if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
e.invalidateUpstreamConnWithoutDisconnectNotify(sess, conn, "send_error", errSend)
if !shouldRetryCodexWebsocketSend(errSend) {
helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend)
return resp, errSend
}
return resp, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
}
e.invalidateUpstreamConn(sess, conn, "send_error", errSend)
if !shouldRetryCodexWebsocketSend(errSend) {
helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend)
@@ -810,7 +846,21 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr
}
helps.RecordAPIWebsocketRequest(ctx, e.cfg, wsReqLog)
conn, closer, respHS, errDial := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders)
var conn *websocket.Conn
var closer *websocketConnectionCloser
var respHS *http.Response
var errDial error
if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
conn, closer = existingWebsocketSessionConn(sess, authID, wsURL)
if conn == nil {
if sess != nil {
sess.reqMu.Unlock()
}
return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
}
} else {
conn, closer, respHS, errDial = e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders)
}
var upstreamHeaders http.Header
if respHS != nil {
upstreamHeaders = respHS.Header.Clone()
@@ -824,7 +874,7 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr
if sess != nil {
sess.reqMu.Unlock()
}
if opts.ExecutionLifecycle != nil {
if opts.ExecutionLifecycle != nil || cliproxyexecutor.DownstreamWebsocket(ctx) {
return nil, statusErr{code: respHS.StatusCode, msg: string(bodyErr)}
}
return e.CodexExecutor.ExecuteStream(ctx, auth, req, opts)
@@ -864,6 +914,15 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr
errSend = mapCodexWebsocketWriteError(sess, conn, errSend)
helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend)
if sess != nil {
if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
e.invalidateUpstreamConnWithoutDisconnectNotify(sess, conn, "send_error", errSend)
sess.clearActive(conn, readCh)
sess.reqMu.Unlock()
if !shouldRetryCodexWebsocketSend(errSend) {
return nil, errSend
}
return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
}
e.invalidateUpstreamConn(sess, conn, "send_error", errSend)
if !shouldRetryCodexWebsocketSend(errSend) {
sess.clearActive(conn, readCh)
@@ -1964,6 +2023,14 @@ func (e *CodexWebsocketsExecutor) readUpstreamLoop(sess *codexWebsocketSession,
}
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
}
@@ -1989,7 +2056,9 @@ func (e *CodexWebsocketsExecutor) invalidateUpstreamConn(sess *codexWebsocketSes
sess.connMu.Unlock()
logCodexWebsocketDisconnected(sessionID, authID, wsURL, reason, err)
sess.notifyUpstreamDisconnect(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)
@@ -2206,6 +2275,9 @@ func (e *CodexAutoExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth
if cliproxyexecutor.DownstreamWebsocket(ctx) && codexWebsocketsEnabled(auth) {
return e.wsExec.Execute(ctx, auth, req, opts)
}
if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
return cliproxyexecutor.Response{}, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
}
return e.httpExec.Execute(ctx, auth, req, opts)
}
@@ -2216,6 +2288,9 @@ func (e *CodexAutoExecutor) ExecuteStream(ctx context.Context, auth *cliproxyaut
if cliproxyexecutor.DownstreamWebsocket(ctx) && codexWebsocketsEnabled(auth) {
return e.wsExec.ExecuteStream(ctx, auth, req, opts)
}
if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
}
return e.httpExec.ExecuteStream(ctx, auth, req, opts)
}