785 lines · 22.9 KB · Go
RawDownload
package helps import (	"bytes"	"context"	"fmt"	"html"	"net/http"	"net/url"	"sort"	"strings"	"time" 	"github.com/gin-gonic/gin"	"github.com/router-for-me/CLIProxyAPI/v7/internal/config"	"github.com/router-for-me/CLIProxyAPI/v7/internal/logging"	"github.com/router-for-me/CLIProxyAPI/v7/internal/util"	log "github.com/sirupsen/logrus"	"github.com/tidwall/gjson") const (	apiAttemptsKey                 = "API_UPSTREAM_ATTEMPTS"	apiRequestKey                  = "API_REQUEST"	apiResponseKey                 = "API_RESPONSE"	apiWebsocketTimelineKey        = "API_WEBSOCKET_TIMELINE"	deferredAPIRequestBytesKey     = "DEFERRED_API_REQUEST_BYTES"	creditsUsedKey                 = "__antigravity_credits_used__"	maxDeferredAPIRequestBodyBytes = 32 << 20 // 32 MiB) // UpstreamRequestLog captures the outbound upstream request details for logging.type UpstreamRequestLog struct {	URL       string	Method    string	Headers   http.Header	Body      []byte	Provider  string	AuthID    string	AuthLabel string	AuthType  string	AuthValue string} type upstreamAttempt struct {	index                int	request              string	response             *strings.Builder	responseSource       *logging.FileBodySource	responseIntroWritten bool	statusWritten        bool	headersWritten       bool	bodyStarted          bool	bodyHasContent       bool	prevWasSSEEvent      bool	errorWritten         bool	trailingNewlines     int} func requestLogCaptureEnabled(cfg *config.Config) bool {	return cfg != nil && cfg.RequestLog && !cfg.CommercialMode} // RecordAPIRequest stores the upstream request metadata in Gin context for request logging.func RecordAPIRequest(ctx context.Context, cfg *config.Config, info UpstreamRequestLog) {	if cfg == nil || cfg.CommercialMode {		return	}	ginCtx := ginContextFrom(ctx)	if ginCtx == nil {		return	}	if !cfg.RequestLog {		deferAPIRequest(ginCtx, info)		return	} 	attempts := getAttempts(ginCtx)	index := len(attempts) + 1	builder := newAPIRequestLogBuilder(index, info, time.Now()) 	requestText := ""	if source, ok := apiRequestSource(ginCtx); ok {		if errWrite := source.AppendBytes([]byte(builder.String())); errWrite == nil {			if len(info.Body) > 0 {				if errBody := source.AppendBytes(info.Body); errBody != nil {					log.WithError(errBody).Warn("failed to append api request body log part")				}			} else if errEmpty := source.AppendBytes([]byte("<empty>")); errEmpty != nil {				log.WithError(errEmpty).Warn("failed to append empty api request log part")			}			if errEnd := source.AppendBytes([]byte("\n\n")); errEnd != nil {				log.WithError(errEnd).Warn("failed to append api request log terminator")			}		} else {			log.WithError(errWrite).Warn("failed to append api request log part")			if len(info.Body) > 0 {				builder.WriteString(string(info.Body))			} else {				builder.WriteString("<empty>")			}			builder.WriteString("\n\n")			requestText = builder.String()		}	} else {		if len(info.Body) > 0 {			builder.WriteString(string(info.Body))		} else {			builder.WriteString("<empty>")		}		builder.WriteString("\n\n")		requestText = builder.String()	} 	attempt := &upstreamAttempt{		index:          index,		request:        requestText,		response:       &strings.Builder{},		responseSource: apiResponseSourceOrNil(ginCtx),	}	attempts = append(attempts, attempt)	ginCtx.Set(apiAttemptsKey, attempts)	if requestText != "" {		updateAggregatedRequest(ginCtx, attempts)	}} func deferAPIRequest(ginCtx *gin.Context, info UpstreamRequestLog) {	if ginCtx == nil {		return	}	var requests []logging.DeferredAPIRequest	if value, exists := ginCtx.Get(logging.DeferredAPIRequestContextKey); exists {		requests, _ = value.([]logging.DeferredAPIRequest)	}	index := len(requests) + 1	capturedInfo := info	capturedAt := time.Now()	capturedBytes, _ := ginCtx.Get(deferredAPIRequestBytesKey)	bytesUsed, _ := capturedBytes.(int)	remaining := maxDeferredAPIRequestBodyBytes - bytesUsed	if remaining < 0 {		remaining = 0	}	captureLength := len(info.Body)	if captureLength > remaining {		captureLength = remaining	}	capturedInfo.Body = bytes.Clone(info.Body[:captureLength])	bodyEmpty := len(info.Body) == 0	bodyTruncated := captureLength < len(info.Body)	ginCtx.Set(deferredAPIRequestBytesKey, bytesUsed+captureLength)	requests = append(requests, func() []byte {		builder := newAPIRequestLogBuilder(index, capturedInfo, capturedAt)		if bodyEmpty {			builder.WriteString("<empty>")		} else {			builder.Write(capturedInfo.Body)			if bodyTruncated {				builder.WriteString(fmt.Sprintf("\n[API REQUEST BODY TRUNCATED: captured first %d bytes]", captureLength))			}		}		builder.WriteString("\n\n")		return []byte(builder.String())	})	ginCtx.Set(logging.DeferredAPIRequestContextKey, requests)} func newAPIRequestLogBuilder(index int, info UpstreamRequestLog, timestamp time.Time) *strings.Builder {	builder := &strings.Builder{}	builder.WriteString(fmt.Sprintf("=== API REQUEST %d ===\n", index))	builder.WriteString(fmt.Sprintf("Timestamp: %s\n", timestamp.Format(time.RFC3339Nano)))	if info.URL != "" {		builder.WriteString(fmt.Sprintf("Upstream URL: %s\n", info.URL))	} else {		builder.WriteString("Upstream URL: <unknown>\n")	}	if info.Method != "" {		builder.WriteString(fmt.Sprintf("HTTP Method: %s\n", info.Method))	}	if auth := formatAuthInfo(info); auth != "" {		builder.WriteString(fmt.Sprintf("Auth: %s\n", auth))	}	builder.WriteString("\nHeaders:\n")	writeHeaders(builder, info.Headers)	builder.WriteString("\nBody:\n")	return builder} // RecordAPIResponseMetadata captures upstream response status/header information for the latest attempt.func RecordAPIResponseMetadata(ctx context.Context, cfg *config.Config, status int, headers http.Header) {	logging.SetResponseHeaders(ctx, headers)	if !requestLogCaptureEnabled(cfg) {		return	}	ginCtx := ginContextFrom(ctx)	if ginCtx == nil {		return	}	attempts, attempt := ensureAttempt(ginCtx)	ensureResponseIntro(ginCtx, attempt) 	if status > 0 && !attempt.statusWritten {		writeAttemptResponse(ginCtx, attempt, []byte(fmt.Sprintf("Status: %d\n", status)))		attempt.statusWritten = true	}	if !attempt.headersWritten {		builder := &strings.Builder{}		builder.WriteString("Headers:\n")		writeHeaders(builder, headers)		writeAttemptResponse(ginCtx, attempt, []byte(builder.String()))		attempt.headersWritten = true		writeAttemptResponse(ginCtx, attempt, []byte("\n"))	} 	updateAggregatedResponseIfMemoryBacked(ginCtx, attempts)} // RecordAPIResponseError adds an error entry for the latest attempt when no HTTP response is available.func RecordAPIResponseError(ctx context.Context, cfg *config.Config, err error) {	if !requestLogCaptureEnabled(cfg) || err == nil {		return	}	ginCtx := ginContextFrom(ctx)	if ginCtx == nil {		return	}	attempts, attempt := ensureAttempt(ginCtx)	ensureResponseIntro(ginCtx, attempt) 	if attempt.bodyStarted && !attempt.bodyHasContent {		// Ensure body does not stay empty marker if error arrives first.		attempt.bodyStarted = false	}	if attempt.errorWritten {		writeAttemptResponse(ginCtx, attempt, []byte("\n"))	}	writeAttemptResponse(ginCtx, attempt, []byte(fmt.Sprintf("Error: %s\n", err.Error())))	attempt.errorWritten = true 	updateAggregatedResponseIfMemoryBacked(ginCtx, attempts)} // AppendAPIResponseChunk appends an upstream response chunk to Gin context for request logging.func AppendAPIResponseChunk(ctx context.Context, cfg *config.Config, chunk []byte) {	if !requestLogCaptureEnabled(cfg) {		return	}	data := bytes.TrimSpace(chunk)	if len(data) == 0 {		return	}	ginCtx := ginContextFrom(ctx)	if ginCtx == nil {		return	}	attempts, attempt := ensureAttempt(ginCtx)	ensureResponseIntro(ginCtx, attempt) 	if !attempt.headersWritten {		builder := &strings.Builder{}		builder.WriteString("Headers:\n")		writeHeaders(builder, nil)		writeAttemptResponse(ginCtx, attempt, []byte(builder.String()))		attempt.headersWritten = true		writeAttemptResponse(ginCtx, attempt, []byte("\n"))	}	if !attempt.bodyStarted {		writeAttemptResponse(ginCtx, attempt, []byte("Body:\n"))		attempt.bodyStarted = true	}	currentChunkIsSSEEvent := bytes.HasPrefix(data, []byte("event:"))	currentChunkIsSSEData := bytes.HasPrefix(data, []byte("data:"))	if attempt.bodyHasContent {		separator := "\n\n"		if attempt.prevWasSSEEvent && currentChunkIsSSEData {			separator = "\n"		}		writeAttemptResponse(ginCtx, attempt, []byte(separator))	}	writeAttemptResponse(ginCtx, attempt, data)	attempt.bodyHasContent = true	attempt.prevWasSSEEvent = currentChunkIsSSEEvent 	updateAggregatedResponseIfMemoryBacked(ginCtx, attempts)} // RecordAPIWebsocketRequest stores an upstream websocket request event in Gin context.func RecordAPIWebsocketRequest(ctx context.Context, cfg *config.Config, info UpstreamRequestLog) {	if !requestLogCaptureEnabled(cfg) {		return	}	ginCtx := ginContextFrom(ctx)	if ginCtx == nil {		return	} 	builder := &strings.Builder{}	builder.WriteString(fmt.Sprintf("Timestamp: %s\n", time.Now().Format(time.RFC3339Nano)))	builder.WriteString("Event: api.websocket.request\n")	if info.URL != "" {		builder.WriteString(fmt.Sprintf("Upstream URL: %s\n", info.URL))	}	if auth := formatAuthInfo(info); auth != "" {		builder.WriteString(fmt.Sprintf("Auth: %s\n", auth))	}	builder.WriteString("Headers:\n")	writeHeaders(builder, info.Headers)	builder.WriteString("\nBody:\n")	if len(info.Body) > 0 {		builder.Write(info.Body)	} else {		builder.WriteString("<empty>")	}	builder.WriteString("\n") 	appendAPIWebsocketTimeline(ginCtx, []byte(builder.String()))} // RecordAPIWebsocketHandshake stores the upstream websocket handshake response metadata.func RecordAPIWebsocketHandshake(ctx context.Context, cfg *config.Config, status int, headers http.Header) {	logging.SetResponseHeaders(ctx, headers)	if !requestLogCaptureEnabled(cfg) {		return	}	ginCtx := ginContextFrom(ctx)	if ginCtx == nil {		return	} 	builder := &strings.Builder{}	builder.WriteString(fmt.Sprintf("Timestamp: %s\n", time.Now().Format(time.RFC3339Nano)))	builder.WriteString("Event: api.websocket.handshake\n")	if status > 0 {		builder.WriteString(fmt.Sprintf("Status: %d\n", status))	}	builder.WriteString("Headers:\n")	writeHeaders(builder, headers)	builder.WriteString("\n") 	appendAPIWebsocketTimeline(ginCtx, []byte(builder.String()))} // RecordAPIWebsocketUpgradeRejection stores a rejected websocket upgrade as an HTTP attempt.func RecordAPIWebsocketUpgradeRejection(ctx context.Context, cfg *config.Config, info UpstreamRequestLog, status int, headers http.Header, body []byte) {	logging.SetResponseHeaders(ctx, headers)	if !requestLogCaptureEnabled(cfg) {		return	}	ginCtx := ginContextFrom(ctx)	if ginCtx == nil {		return	} 	RecordAPIRequest(ctx, cfg, info)	RecordAPIResponseMetadata(ctx, cfg, status, headers)	AppendAPIResponseChunk(ctx, cfg, body)} // WebsocketUpgradeRequestURL converts a websocket URL back to its HTTP handshake URL for logging.func WebsocketUpgradeRequestURL(rawURL string) string {	trimmedURL := strings.TrimSpace(rawURL)	if trimmedURL == "" {		return ""	}	parsed, err := url.Parse(trimmedURL)	if err != nil {		return trimmedURL	}	switch strings.ToLower(parsed.Scheme) {	case "ws":		parsed.Scheme = "http"	case "wss":		parsed.Scheme = "https"	}	return parsed.String()} // AppendAPIWebsocketResponse stores an upstream websocket response frame in Gin context.func AppendAPIWebsocketResponse(ctx context.Context, cfg *config.Config, payload []byte) {	if !requestLogCaptureEnabled(cfg) {		return	}	data := bytes.TrimSpace(payload)	if len(data) == 0 {		return	}	ginCtx := ginContextFrom(ctx)	if ginCtx == nil {		return	}	markAPIResponseTimestamp(ginCtx) 	builder := &strings.Builder{}	builder.WriteString(fmt.Sprintf("Timestamp: %s\n", time.Now().Format(time.RFC3339Nano)))	builder.WriteString("Event: api.websocket.response\n")	builder.Write(data)	builder.WriteString("\n") 	appendAPIWebsocketTimeline(ginCtx, []byte(builder.String()))} // AppendCodexAPIWebsocketResponse stores a codex upstream websocket response frame and merges any// quota event headers carried by the frame into the request log.func AppendCodexAPIWebsocketResponse(ctx context.Context, cfg *config.Config, payload []byte) {	logging.MergeResponseHeaders(ctx, ParseCodexQuotaEventHeaders(payload))	AppendAPIWebsocketResponse(ctx, cfg, payload)} // RecordAPIWebsocketError stores an upstream websocket error event in Gin context.func RecordAPIWebsocketError(ctx context.Context, cfg *config.Config, stage string, err error) {	if !requestLogCaptureEnabled(cfg) || err == nil {		return	}	ginCtx := ginContextFrom(ctx)	if ginCtx == nil {		return	}	markAPIResponseTimestamp(ginCtx) 	builder := &strings.Builder{}	builder.WriteString(fmt.Sprintf("Timestamp: %s\n", time.Now().Format(time.RFC3339Nano)))	builder.WriteString("Event: api.websocket.error\n")	if trimmed := strings.TrimSpace(stage); trimmed != "" {		builder.WriteString(fmt.Sprintf("Stage: %s\n", trimmed))	}	builder.WriteString(fmt.Sprintf("Error: %s\n", err.Error())) 	appendAPIWebsocketTimeline(ginCtx, []byte(builder.String()))} func ginContextFrom(ctx context.Context) *gin.Context {	ginCtx, _ := ctx.Value("gin").(*gin.Context)	return ginCtx} func getAttempts(ginCtx *gin.Context) []*upstreamAttempt {	if ginCtx == nil {		return nil	}	if value, exists := ginCtx.Get(apiAttemptsKey); exists {		if attempts, ok := value.([]*upstreamAttempt); ok {			return attempts		}	}	return nil} func ensureAttempt(ginCtx *gin.Context) ([]*upstreamAttempt, *upstreamAttempt) {	attempts := getAttempts(ginCtx)	if len(attempts) == 0 {		attempt := &upstreamAttempt{			index:          1,			response:       &strings.Builder{},			responseSource: apiResponseSourceOrNil(ginCtx),		}		if source, ok := apiRequestSource(ginCtx); ok {			if errWrite := source.AppendBytes([]byte("=== API REQUEST 1 ===\n<missing>\n\n")); errWrite != nil {				log.WithError(errWrite).Warn("failed to append missing api request log part")				attempt.request = "=== API REQUEST 1 ===\n<missing>\n\n"			}		} else {			attempt.request = "=== API REQUEST 1 ===\n<missing>\n\n"		}		attempts = []*upstreamAttempt{attempt}		ginCtx.Set(apiAttemptsKey, attempts)		if attempt.request != "" {			updateAggregatedRequest(ginCtx, attempts)		}	}	return attempts, attempts[len(attempts)-1]} func ensureResponseIntro(ginCtx *gin.Context, attempt *upstreamAttempt) {	if attempt == nil || attempt.response == nil || attempt.responseIntroWritten {		return	}	attempts := getAttempts(ginCtx)	for i := len(attempts) - 1; i >= 0; i-- {		previousAttempt := attempts[i]		if previousAttempt == nil || previousAttempt == attempt || !previousAttempt.responseIntroWritten {			continue		}		if missingNewlines := 2 - previousAttempt.trailingNewlines; missingNewlines > 0 {			writeAttemptResponse(ginCtx, attempt, []byte(strings.Repeat("\n", missingNewlines)))		}		break	}	writeAttemptResponse(ginCtx, attempt, []byte(fmt.Sprintf("=== API RESPONSE %d ===\n", attempt.index)))	writeAttemptResponse(ginCtx, attempt, []byte(fmt.Sprintf("Timestamp: %s\n", time.Now().Format(time.RFC3339Nano))))	writeAttemptResponse(ginCtx, attempt, []byte("\n"))	attempt.responseIntroWritten = true} func writeAttemptResponse(ginCtx *gin.Context, attempt *upstreamAttempt, payload []byte) {	if attempt == nil || len(payload) == 0 {		return	}	trailingNewlines := 0	for i := len(payload) - 1; i >= 0 && payload[i] == '\n'; i-- {

Lines 1–500 of 785.