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/v6/internal/config" "github.com/router-for-me/CLIProxyAPI/v6/internal/logging" "github.com/router-for-me/CLIProxyAPI/v6/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" // maxErrorLogResponseBodySize limits cached response body when request-log is disabled. // Prevents unbounded memory growth for large/streaming responses in error-only mode. maxErrorLogResponseBodySize = 32 * 1024 // 32KB // maxErrorLogRequestBodySize limits materialized request body in error-only mode. // Prevents OOM from large payloads (e.g. base64 images) when full request logging is off. maxErrorLogRequestBodySize = 32 * 1024 // 32KB ) // 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 deferredBody []byte // lazy body reference; only materialized on error response *strings.Builder responseIntroWritten bool statusWritten bool headersWritten bool bodyStarted bool bodyHasContent bool prevWasSSEEvent bool errorWritten bool bodyBytesWritten int bodyTruncated bool } // RecordAPIRequest stores the upstream request metadata in Gin context for request logging. func RecordAPIRequest(ctx context.Context, cfg *config.Config, info UpstreamRequestLog) { ginCtx := ginContextFrom(ctx) if ginCtx == nil { return } attempts := getAttempts(ginCtx) index := len(attempts) + 1 requestLogEnabled := cfg != nil && cfg.RequestLog builder := &strings.Builder{} builder.WriteString(fmt.Sprintf("=== API REQUEST %d ===\n", index)) builder.WriteString(fmt.Sprintf("Timestamp: %s\n", time.Now().Format(time.RFC3339Nano))) if info.URL != "" { builder.WriteString(fmt.Sprintf("Upstream URL: %s\n", info.URL)) } else { builder.WriteString("Upstream URL: \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") if requestLogEnabled { // Full request logging: format body inline if len(info.Body) > 0 { builder.WriteString(string(info.Body)) } else { builder.WriteString("") } } else { // Error-only mode: defer body to avoid allocating copies for the 99% success path if len(info.Body) > 0 { builder.WriteString(fmt.Sprintf("", len(info.Body))) } else { builder.WriteString("") } } builder.WriteString("\n\n") attempt := &upstreamAttempt{ index: index, request: builder.String(), response: &strings.Builder{}, } if !requestLogEnabled && len(info.Body) > 0 { attempt.deferredBody = info.Body } attempts = append(attempts, attempt) ginCtx.Set(apiAttemptsKey, attempts) updateAggregatedRequest(ginCtx, attempts) } // RecordAPIResponseMetadata captures upstream response status/header information for the latest attempt. func RecordAPIResponseMetadata(ctx context.Context, cfg *config.Config, status int, headers http.Header) { ginCtx := ginContextFrom(ctx) if ginCtx == nil { return } attempts, attempt := ensureAttempt(ginCtx) // Materialize deferred request body when upstream returns an error. // Success responses (2xx) skip this — their deferred body is dropped with gin context. if status >= http.StatusBadRequest { materializeDeferredBodies(ginCtx, attempts) } else { for _, a := range attempts { a.deferredBody = nil } } ensureResponseIntro(attempt) if status > 0 && !attempt.statusWritten { attempt.response.WriteString(fmt.Sprintf("Status: %d\n", status)) attempt.statusWritten = true } if !attempt.headersWritten { attempt.response.WriteString("Headers:\n") writeHeaders(attempt.response, headers) attempt.headersWritten = true attempt.response.WriteString("\n") } updateAggregatedResponse(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 err == nil { return } ginCtx := ginContextFrom(ctx) if ginCtx == nil { return } attempts, attempt := ensureAttempt(ginCtx) // Materialize deferred request body on error — this is the only path that // actually needs the body. Success path (99%) never pays for body copies. materializeDeferredBodies(ginCtx, attempts) ensureResponseIntro(attempt) if attempt.bodyStarted && !attempt.bodyHasContent { // Ensure body does not stay empty marker if error arrives first. attempt.bodyStarted = false } if attempt.errorWritten { attempt.response.WriteString("\n") } attempt.response.WriteString(fmt.Sprintf("Error: %s\n", err.Error())) attempt.errorWritten = true updateAggregatedResponse(ginCtx, attempts) } // AppendAPIResponseChunk appends an upstream response chunk to Gin context for request logging. func AppendAPIResponseChunk(ctx context.Context, cfg *config.Config, chunk []byte) { data := bytes.TrimSpace(chunk) if len(data) == 0 { return } ginCtx := ginContextFrom(ctx) if ginCtx == nil { return } attempts, attempt := ensureAttempt(ginCtx) ensureResponseIntro(attempt) requestLogEnabled := cfg != nil && cfg.RequestLog if !requestLogEnabled && attempt.bodyTruncated { return } if !attempt.headersWritten { attempt.response.WriteString("Headers:\n") writeHeaders(attempt.response, nil) attempt.headersWritten = true attempt.response.WriteString("\n") } if !attempt.bodyStarted { attempt.response.WriteString("Body:\n") attempt.bodyStarted = true } // Cap response body size when full request-log is disabled to prevent memory growth if !requestLogEnabled { remaining := maxErrorLogResponseBodySize - attempt.bodyBytesWritten if remaining <= 0 { attempt.bodyTruncated = true attempt.response.WriteString("\n") updateAggregatedResponse(ginCtx, attempts) return } if len(data) > remaining { data = data[:remaining] attempt.bodyTruncated = 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" } attempt.response.WriteString(separator) } attempt.response.WriteString(string(data)) attempt.bodyBytesWritten += len(data) attempt.bodyHasContent = true attempt.prevWasSSEEvent = currentChunkIsSSEEvent if attempt.bodyTruncated { attempt.response.WriteString("\n") } updateAggregatedResponse(ginCtx, attempts) } // RecordAPIWebsocketRequest stores an upstream websocket request event in Gin context. func RecordAPIWebsocketRequest(ctx context.Context, cfg *config.Config, info UpstreamRequestLog) { if cfg == nil || !cfg.RequestLog { 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("") } 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) { if cfg == nil || !cfg.RequestLog { 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) { if cfg == nil || !cfg.RequestLog { 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 cfg == nil || !cfg.RequestLog { 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())) } // RecordAPIWebsocketError stores an upstream websocket error event in Gin context. func RecordAPIWebsocketError(ctx context.Context, cfg *config.Config, stage string, err error) { if cfg == nil || !cfg.RequestLog || 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 } const creditsUsedKey = "__antigravity_credits_used__" // MarkCreditsUsed flags the request as having used AI credits for billing. func MarkCreditsUsed(ctx context.Context) { if ginCtx := ginContextFrom(ctx); ginCtx != nil { ginCtx.Set(creditsUsedKey, true) } } // CreditsUsed returns true if the request used AI credits. func CreditsUsed(ctx context.Context) bool { if ginCtx := ginContextFrom(ctx); ginCtx != nil { if val, exists := ginCtx.Get(creditsUsedKey); exists { if b, ok := val.(bool); ok { return b } } } return false } 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 } // materializeDeferredBodies replaces deferred body placeholders with actual // (truncated) body content. Called only on the error path so the 99% success // path pays zero allocation cost for request body logging. func materializeDeferredBodies(ginCtx *gin.Context, attempts []*upstreamAttempt) { changed := false for _, attempt := range attempts { if attempt.deferredBody == nil { continue } body := attempt.deferredBody attempt.deferredBody = nil // release reference to allow GC of full payload placeholder := fmt.Sprintf("", len(body)) var replacement string if len(body) > maxErrorLogRequestBodySize { replacement = string(body[:maxErrorLogRequestBodySize]) + fmt.Sprintf("\n", len(body), maxErrorLogRequestBodySize) } else { replacement = string(body) } attempt.request = strings.Replace(attempt.request, placeholder, replacement, 1) changed = true } if changed { updateAggregatedRequest(ginCtx, attempts) } } func ensureAttempt(ginCtx *gin.Context) ([]*upstreamAttempt, *upstreamAttempt) { attempts := getAttempts(ginCtx) if len(attempts) == 0 { attempt := &upstreamAttempt{ index: 1, request: "=== API REQUEST 1 ===\n\n\n", response: &strings.Builder{}, } attempts = []*upstreamAttempt{attempt} ginCtx.Set(apiAttemptsKey, attempts) updateAggregatedRequest(ginCtx, attempts) } return attempts, attempts[len(attempts)-1] } func ensureResponseIntro(attempt *upstreamAttempt) { if attempt == nil || attempt.response == nil || attempt.responseIntroWritten { return } attempt.response.WriteString(fmt.Sprintf("=== API RESPONSE %d ===\n", attempt.index)) attempt.response.WriteString(fmt.Sprintf("Timestamp: %s\n", time.Now().Format(time.RFC3339Nano))) attempt.response.WriteString("\n") attempt.responseIntroWritten = true } func updateAggregatedRequest(ginCtx *gin.Context, attempts []*upstreamAttempt) { if ginCtx == nil { return } var builder strings.Builder for _, attempt := range attempts { builder.WriteString(attempt.request) } ginCtx.Set(apiRequestKey, []byte(builder.String())) } func updateAggregatedResponse(ginCtx *gin.Context, attempts []*upstreamAttempt) { if ginCtx == nil { return } var builder strings.Builder for idx, attempt := range attempts { if attempt == nil || attempt.response == nil { continue } responseText := attempt.response.String() if responseText == "" { continue } builder.WriteString(responseText) if !strings.HasSuffix(responseText, "\n") { builder.WriteString("\n") } if idx < len(attempts)-1 { builder.WriteString("\n") } } ginCtx.Set(apiResponseKey, []byte(builder.String())) } func appendAPIWebsocketTimeline(ginCtx *gin.Context, chunk []byte) { if ginCtx == nil { return } data := bytes.TrimSpace(chunk) if len(data) == 0 { return } if existing, exists := ginCtx.Get(apiWebsocketTimelineKey); exists { if existingBytes, ok := existing.([]byte); ok && len(existingBytes) > 0 { combined := make([]byte, 0, len(existingBytes)+len(data)+2) combined = append(combined, existingBytes...) if !bytes.HasSuffix(existingBytes, []byte("\n")) { combined = append(combined, '\n') } combined = append(combined, '\n') combined = append(combined, data...) ginCtx.Set(apiWebsocketTimelineKey, combined) return } } ginCtx.Set(apiWebsocketTimelineKey, bytes.Clone(data)) } func markAPIResponseTimestamp(ginCtx *gin.Context) { if ginCtx == nil { return } if _, exists := ginCtx.Get("API_RESPONSE_TIMESTAMP"); exists { return } ginCtx.Set("API_RESPONSE_TIMESTAMP", time.Now()) } func writeHeaders(builder *strings.Builder, headers http.Header) { if builder == nil { return } if len(headers) == 0 { builder.WriteString("\n") return } keys := make([]string, 0, len(headers)) for key := range headers { keys = append(keys, key) } sort.Strings(keys) for _, key := range keys { values := headers[key] if len(values) == 0 { builder.WriteString(fmt.Sprintf("%s:\n", key)) continue } for _, value := range values { masked := util.MaskSensitiveHeaderValue(key, value) builder.WriteString(fmt.Sprintf("%s: %s\n", key, masked)) } } } func formatAuthInfo(info UpstreamRequestLog) string { var parts []string if trimmed := strings.TrimSpace(info.Provider); trimmed != "" { parts = append(parts, fmt.Sprintf("provider=%s", trimmed)) } if trimmed := strings.TrimSpace(info.AuthID); trimmed != "" { parts = append(parts, fmt.Sprintf("auth_id=%s", trimmed)) } if trimmed := strings.TrimSpace(info.AuthLabel); trimmed != "" { parts = append(parts, fmt.Sprintf("label=%s", trimmed)) } authType := strings.ToLower(strings.TrimSpace(info.AuthType)) authValue := strings.TrimSpace(info.AuthValue) switch authType { case "api_key": if authValue != "" { parts = append(parts, fmt.Sprintf("type=api_key value=%s", util.HideAPIKey(authValue))) } else { parts = append(parts, "type=api_key") } case "oauth": parts = append(parts, "type=oauth") default: if authType != "" { if authValue != "" { parts = append(parts, fmt.Sprintf("type=%s value=%s", authType, authValue)) } else { parts = append(parts, fmt.Sprintf("type=%s", authType)) } } } return strings.Join(parts, ", ") } func SummarizeErrorBody(contentType string, body []byte) string { isHTML := strings.Contains(strings.ToLower(contentType), "text/html") if !isHTML { trimmed := bytes.TrimSpace(bytes.ToLower(body)) if bytes.HasPrefix(trimmed, []byte("') if gt == -1 { return "" } start += gt + 1 end := bytes.Index(lower[start:], []byte("")) if end == -1 { return "" } title := string(body[start : start+end]) title = html.UnescapeString(title) title = strings.TrimSpace(title) if title == "" { return "" } return strings.Join(strings.Fields(title), " ") } // extractJSONErrorMessage attempts to extract error.message from JSON error responses func extractJSONErrorMessage(body []byte) string { result := gjson.GetBytes(body, "error.message") if result.Exists() && result.String() != "" { return result.String() } return "" } // logWithRequestID returns a logrus Entry with request_id field populated from context. // If no request ID is found in context, it returns the standard logger. func LogWithRequestID(ctx context.Context) *log.Entry { if ctx == nil { return log.NewEntry(log.StandardLogger()) } requestID := logging.GetRequestID(ctx) if requestID == "" { return log.NewEntry(log.StandardLogger()) } return log.WithField("request_id", requestID) }