refactor(logging): strip unrelated deferred body changes, keep credits-only logging

Remove deferred body optimization and maxErrorLog constants that were
unrelated to credits fallback. Keep only MarkCreditsUsed/CreditsUsed
helpers for flagging requests that consumed AI credits.
This commit is contained in:
sususu98
2026-04-23 17:41:54 +08:00
parent e75daa299b
commit 920b6efffa

View File

@@ -24,14 +24,7 @@ const (
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
creditsUsedKey = "__antigravity_credits_used__"
)
// UpstreamRequestLog captures the outbound upstream request details for logging.
@@ -50,7 +43,6 @@ type UpstreamRequestLog struct {
type upstreamAttempt struct {
index int
request string
deferredBody []byte // lazy body reference; only materialized on error
response *strings.Builder
responseIntroWritten bool
statusWritten bool
@@ -59,12 +51,13 @@ type upstreamAttempt struct {
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) {
if cfg == nil || !cfg.RequestLog {
return
}
ginCtx := ginContextFrom(ctx)
if ginCtx == nil {
return
@@ -73,8 +66,6 @@ func RecordAPIRequest(ctx context.Context, cfg *config.Config, info UpstreamRequ
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)))
@@ -92,20 +83,10 @@ func RecordAPIRequest(ctx context.Context, cfg *config.Config, info UpstreamRequ
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("<empty>")
}
if len(info.Body) > 0 {
builder.WriteString(string(info.Body))
} else {
// Error-only mode: defer body to avoid allocating copies for the 99% success path
if len(info.Body) > 0 {
builder.WriteString(fmt.Sprintf("<deferred: %d bytes>", len(info.Body)))
} else {
builder.WriteString("<empty>")
}
builder.WriteString("<empty>")
}
builder.WriteString("\n\n")
@@ -114,9 +95,6 @@ func RecordAPIRequest(ctx context.Context, cfg *config.Config, info UpstreamRequ
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)
@@ -124,22 +102,14 @@ func RecordAPIRequest(ctx context.Context, cfg *config.Config, info UpstreamRequ
// RecordAPIResponseMetadata captures upstream response status/header information for the latest attempt.
func RecordAPIResponseMetadata(ctx context.Context, cfg *config.Config, status int, headers http.Header) {
if cfg == nil || !cfg.RequestLog {
return
}
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 {
@@ -158,7 +128,7 @@ func RecordAPIResponseMetadata(ctx context.Context, cfg *config.Config, status i
// 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 {
if cfg == nil || !cfg.RequestLog || err == nil {
return
}
ginCtx := ginContextFrom(ctx)
@@ -166,11 +136,6 @@ func RecordAPIResponseError(ctx context.Context, cfg *config.Config, err error)
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 {
@@ -188,6 +153,9 @@ func RecordAPIResponseError(ctx context.Context, cfg *config.Config, err error)
// AppendAPIResponseChunk appends an upstream response chunk to Gin context for request logging.
func AppendAPIResponseChunk(ctx context.Context, cfg *config.Config, chunk []byte) {
if cfg == nil || !cfg.RequestLog {
return
}
data := bytes.TrimSpace(chunk)
if len(data) == 0 {
return
@@ -199,11 +167,6 @@ func AppendAPIResponseChunk(ctx context.Context, cfg *config.Config, chunk []byt
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)
@@ -214,22 +177,6 @@ func AppendAPIResponseChunk(ctx context.Context, cfg *config.Config, chunk []byt
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<truncated: response body exceeded 32KB limit for error log>")
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 {
@@ -240,14 +187,9 @@ func AppendAPIResponseChunk(ctx context.Context, cfg *config.Config, chunk []byt
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<truncated: response body exceeded 32KB limit for error log>")
}
updateAggregatedResponse(ginCtx, attempts)
}
@@ -391,27 +333,6 @@ func ginContextFrom(ctx context.Context) *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
@@ -424,34 +345,6 @@ func getAttempts(ginCtx *gin.Context) []*upstreamAttempt {
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("<deferred: %d bytes>", len(body))
var replacement string
if len(body) > maxErrorLogRequestBodySize {
replacement = string(body[:maxErrorLogRequestBodySize]) +
fmt.Sprintf("\n<truncated: request body %d bytes, showing first %d>", 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 {
@@ -676,3 +569,24 @@ func LogWithRequestID(ctx context.Context) *log.Entry {
}
return log.WithField("request_id", requestID)
}
// MarkCreditsUsed flags the request as having used AI credits for billing.
func MarkCreditsUsed(ctx context.Context) {
ginCtx := ginContextFrom(ctx)
if ginCtx != nil {
ginCtx.Set(creditsUsedKey, true)
}
}
// CreditsUsed returns true if the request used AI credits.
func CreditsUsed(ctx context.Context) bool {
ginCtx := ginContextFrom(ctx)
if ginCtx != nil {
if val, exists := ginCtx.Get(creditsUsedKey); exists {
if b, ok := val.(bool); ok {
return b
}
}
}
return false
}