Something went wrong. Try again.
sloproxy
Something went wrong. Try again.
5.6 kB · 199 lines
Go
at commit 42eef103
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200package executor
import ( "context" "io" "net/http" "strings" "time"
"github.com/router-for-me/CLIProxyAPI/v7/internal/config" "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" log "github.com/sirupsen/logrus" "github.com/tidwall/gjson" "github.com/tidwall/sjson")
type statusErrWithHeaders struct { statusErr headers http.Header}
func (e statusErrWithHeaders) Headers() http.Header { if e.headers == nil { return nil } return e.headers.Clone()}
func parseCodexWebsocketError(payload []byte) (error, bool) { if len(payload) == 0 { return nil, false } if strings.TrimSpace(gjson.GetBytes(payload, "type").String()) != "error" { return nil, false } status := int(gjson.GetBytes(payload, "status").Int()) if status == 0 { status = int(gjson.GetBytes(payload, "status_code").Int()) } if status <= 0 { return nil, false }
out := buildCodexWebsocketErrorPayload(payload, status) headers := parseCodexWebsocketErrorHeaders(payload) statusError := statusErr{code: status, msg: string(out)} if retryAfter := parseCodexRetryAfter(status, out, time.Now()); retryAfter != nil { statusError.retryAfter = retryAfter } else if isCodexWebsocketConnectionLimitError(payload) { retryAfter := time.Duration(0) statusError.retryAfter = &retryAfter } return statusErrWithHeaders{ statusErr: statusError, headers: headers, }, true}
func clearCodexReasoningReplayOnWebsocketError(ctx context.Context, scope codexReasoningReplayScope, payload []byte) error { status := int(gjson.GetBytes(payload, "status").Int()) if status == 0 { status = int(gjson.GetBytes(payload, "status_code").Int()) } if status <= 0 { return nil } return clearCodexReasoningReplayOnInvalidSignature(ctx, scope, status, buildCodexWebsocketErrorPayload(payload, status))}
func buildCodexWebsocketErrorPayload(payload []byte, status int) []byte { out := []byte(`{}`) out, _ = sjson.SetBytes(out, "status", status)
if bodyNode := gjson.GetBytes(payload, "body"); bodyNode.Exists() { out, _ = sjson.SetRawBytes(out, "body", []byte(bodyNode.Raw)) if bodyErrorNode := bodyNode.Get("error"); bodyErrorNode.Exists() { out, _ = sjson.SetRawBytes(out, "error", []byte(bodyErrorNode.Raw)) return out } }
if errNode := gjson.GetBytes(payload, "error"); errNode.Exists() { out, _ = sjson.SetRawBytes(out, "error", []byte(errNode.Raw)) return out }
out, _ = sjson.SetBytes(out, "error.type", "server_error") out, _ = sjson.SetBytes(out, "error.message", http.StatusText(status)) return out}
func isCodexWebsocketConnectionLimitError(payload []byte) bool { if len(payload) == 0 { return false } for _, path := range []string{"error.code", "error.type", "body.error.code", "body.error.type", "code", "error"} { if strings.TrimSpace(gjson.GetBytes(payload, path).String()) == "websocket_connection_limit_reached" { return true } } return false}
func parseCodexWebsocketErrorHeaders(payload []byte) http.Header { headersNode := gjson.GetBytes(payload, "headers") if !headersNode.Exists() || !headersNode.IsObject() { return nil } mapped := make(http.Header) headersNode.ForEach(func(key, value gjson.Result) bool { name := strings.TrimSpace(key.String()) if name == "" { return true } switch value.Type { case gjson.String: if v := strings.TrimSpace(value.String()); v != "" { mapped.Set(name, v) } case gjson.Number, gjson.True, gjson.False: if v := strings.TrimSpace(value.Raw); v != "" { mapped.Set(name, v) } default: } return true }) if len(mapped) == 0 { return nil } return mapped}
func normalizeCodexWebsocketCompletion(payload []byte) []byte { if strings.TrimSpace(gjson.GetBytes(payload, "type").String()) == "response.done" { updated, err := sjson.SetBytes(payload, "type", "response.completed") if err == nil && len(updated) > 0 { return updated } } return payload}
func encodeCodexWebsocketAsSSE(payload []byte) []byte { if len(payload) == 0 { return nil } line := make([]byte, 0, len("data: ")+len(payload)) line = append(line, []byte("data: ")...) line = append(line, payload...) return line}
func websocketUpgradeRequestLog(info helps.UpstreamRequestLog) helps.UpstreamRequestLog { upgradeInfo := info upgradeInfo.URL = helps.WebsocketUpgradeRequestURL(info.URL) upgradeInfo.Method = http.MethodGet upgradeInfo.Body = nil upgradeInfo.Headers = info.Headers.Clone() if upgradeInfo.Headers == nil { upgradeInfo.Headers = make(http.Header) } if strings.TrimSpace(upgradeInfo.Headers.Get("Connection")) == "" { upgradeInfo.Headers.Set("Connection", "Upgrade") } if strings.TrimSpace(upgradeInfo.Headers.Get("Upgrade")) == "" { upgradeInfo.Headers.Set("Upgrade", "websocket") } return upgradeInfo}
func recordAPIWebsocketHandshake(ctx context.Context, cfg *config.Config, resp *http.Response) { if resp == nil { return } helps.RecordAPIWebsocketHandshake(ctx, cfg, resp.StatusCode, resp.Header.Clone()) closeHTTPResponseBody(resp, "codex websockets executor: close handshake response body error")}
func websocketHandshakeBody(resp *http.Response) []byte { if resp == nil || resp.Body == nil { return nil } body, _ := io.ReadAll(resp.Body) closeHTTPResponseBody(resp, "codex websockets executor: close handshake response body error") if len(body) == 0 { return nil } return body}
func closeHTTPResponseBody(resp *http.Response, logPrefix string) { if resp == nil || resp.Body == nil { return } if errClose := resp.Body.Close(); errClose != nil { log.Errorf("%s: %v", logPrefix, errClose) }}