208 lines · 5.8 KB · Go
RawDownload
package 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) {	return parseCodexWebsocketErrorWithCooling(payload, false)} func parseCodexWebsocketErrorWithCooling(payload []byte, modelLevelCooling bool) (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)	isUsageLimit := isCodexUsageLimitError(out)	statusError := statusErr{		code:             status,		msg:              string(out),		credentialScoped: isUsageLimit && !modelLevelCooling,	}	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)	}}