Add projects

backend/internal/runtime/executor/codex_websockets_errors.go

@@ -0,0 +1,199 @@
+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) {
+	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)
+	}
+}