ai-gateway-go/internal/audit/recorder.go

package audit

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"strings"
	"sync/atomic"
	"time"

	platformid "aigateway.local/core/internal/platform/id"
	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgtype"
	"github.com/jackc/pgx/v5/pgxpool"
)

type Event struct {
	TenantID         *string
	RequestID        string
	APIKeyID         *string
	ProviderCode     string
	Model            string
	Protocol         string
	StatusCode       int
	PromptTokens     int64
	CompletionTokens int64
	CostMicrounits   int64
	PriceID          string
	Currency         string
	LatencyMS        int
	Labels           map[string]any
	RecordedAt       time.Time
}

type Stats struct {
	Accepted uint64
	Dropped  uint64
	Written  uint64
	Failures uint64
}

type Recorder struct {
	pool          *pgxpool.Pool
	logger        *slog.Logger
	queue         chan Event
	batchSize     int
	flushInterval time.Duration
	accepted      atomic.Uint64
	dropped       atomic.Uint64
	written       atomic.Uint64
	failures      atomic.Uint64
}

func NewRecorder(pool *pgxpool.Pool, logger *slog.Logger, queueSize, batchSize int, flushInterval time.Duration) *Recorder {
	if queueSize < 1 {
		queueSize = 4096
	}
	if batchSize < 1 {
		batchSize = 200
	}
	if flushInterval <= 0 {
		flushInterval = time.Second
	}
	return &Recorder{pool: pool, logger: logger, queue: make(chan Event, queueSize), batchSize: batchSize, flushInterval: flushInterval}
}

func (r *Recorder) Record(event Event) bool {
	if r == nil || r.pool == nil {
		return false
	}
	if event.RecordedAt.IsZero() {
		event.RecordedAt = time.Now().UTC()
	}
	select {
	case r.queue <- event:
		r.accepted.Add(1)
		return true
	default:
		r.dropped.Add(1)
		return false
	}
}

func (r *Recorder) Stats() Stats {
	if r == nil {
		return Stats{}
	}
	return Stats{Accepted: r.accepted.Load(), Dropped: r.dropped.Load(), Written: r.written.Load(), Failures: r.failures.Load()}
}

func (r *Recorder) Prometheus() string {
	stats := r.Stats()
	return fmt.Sprintf("gateway_audit_events_accepted_total %d\ngateway_audit_events_dropped_total %d\ngateway_audit_events_written_total %d\ngateway_audit_flush_failures_total %d\n", stats.Accepted, stats.Dropped, stats.Written, stats.Failures)
}

func (r *Recorder) Run(ctx context.Context) {
	ticker := time.NewTicker(r.flushInterval)
	defer ticker.Stop()
	pending := make([]Event, 0, r.batchSize)
	flushAllowed := true
	for {
		var events <-chan Event = r.queue
		if len(pending) >= cap(r.queue)+r.batchSize {
			events = nil
		}
		select {
		case event := <-events:
			pending = append(pending, event)
			if len(pending) >= r.batchSize && flushAllowed {
				if r.flush(ctx, pending) == nil {
					pending = pending[:0]
				} else {
					flushAllowed = false
				}
			}
		case <-ticker.C:
			flushAllowed = true
			if len(pending) > 0 {
				if r.flush(ctx, pending) == nil {
					pending = pending[:0]
				} else {
					flushAllowed = false
				}
			}
		case <-ctx.Done():
			for {
				select {
				case event := <-r.queue:
					pending = append(pending, event)
				default:
					// 关闭前的最后一次落盘:数据库短暂不可用时重试有限次数,
					// 而不是只试一次就把整批审计事件静默丢弃。
					flushCtx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
					if len(pending) > 0 {
						lastErr := r.flush(flushCtx, pending)
						for attempt := 0; lastErr != nil && attempt < 3; attempt++ {
							select {
							case <-time.After(2 * time.Second):
							case <-flushCtx.Done():
								lastErr = flushCtx.Err()
							}
							if flushCtx.Err() != nil {
								break
							}
							lastErr = r.flush(flushCtx, pending)
						}
						if lastErr != nil && r.logger != nil {
							r.logger.Error("audit drain failed; events lost", "events", len(pending), "error", lastErr)
						}
					}
					cancel()
					return
				}
			}
		}
	}
}

type dailyKey struct {
	date, apiKeyID, provider, model, currency string
}

type dailyValue struct {
	requests, failed, prompt, completion, cost int64
}

type policyAlert struct {
	eventID, requestID string
	tenantID           *string
	payload            []byte
}

func (r *Recorder) flush(ctx context.Context, events []Event) error {
	if len(events) == 0 {
		return nil
	}
	tx, err := r.pool.Begin(ctx)
	if err != nil {
		return r.flushError(err)
	}
	defer func() { _ = tx.Rollback(ctx) }()
	rows := make([][]any, 0, len(events))
	daily := make(map[dailyKey]dailyValue)
	alerts := make([]policyAlert, 0)
	for _, event := range events {
		id, err := platformid.NewUUID()
		if err != nil {
			return r.flushError(err)
		}
		labels, _ := json.Marshal(event.Labels)
		if policies, matched := event.Labels["content_policies"]; matched {
			alertID, idErr := platformid.NewUUID()
			if idErr != nil {
				return r.flushError(idErr)
			}
			payload, _ := json.Marshal(map[string]any{
				"request_id": event.RequestID, "protocol": event.Protocol, "status_code": event.StatusCode,
				"provider_code": event.ProviderCode, "model": event.Model, "content_policies": policies,
				"content_redacted": event.Labels["content_redacted"], "recorded_at": event.RecordedAt,
			})
			alerts = append(alerts, policyAlert{eventID: alertID, requestID: event.RequestID, tenantID: event.TenantID, payload: payload})
		}
		rows = append(rows, []any{
			uuidValue(id), uuidPointer(event.TenantID), event.RequestID, nil, uuidPointer(event.APIKeyID),
			nullString(event.ProviderCode), nullString(event.Model), event.Protocol, event.StatusCode,
			nullToken(event.PromptTokens), nullToken(event.CompletionTokens), nullCost(event.CostMicrounits), event.LatencyMS,
			nil, nil, labels, event.RecordedAt,
		})
		if event.APIKeyID != nil && *event.APIKeyID != "" {
			// 仅聚合合法 UUID 的 API Key:usage_daily.api_key_id 是
			// REFERENCES api_keys 的 uuid 列,bootstrap 等非 UUID 身份写入
			// 会让整批事务失败、审计管线永久卡死。审计事件本身仍落 audit_events。
			if _, uuidErr := uuid.Parse(strings.TrimSpace(*event.APIKeyID)); uuidErr != nil {
				continue
			}
			// 成本按币种独立聚合:不同货币的价格不能相加成单一数字。
			currency := strings.ToUpper(strings.TrimSpace(event.Currency))
			if currency == "" {
				currency = "USD"
			}
			key := dailyKey{date: event.RecordedAt.UTC().Format("2006-01-02"), apiKeyID: *event.APIKeyID, provider: event.ProviderCode, model: event.Model, currency: currency}
			value := daily[key]
			value.requests++
			if event.StatusCode >= 400 {
				value.failed++
			}
			value.prompt += max(event.PromptTokens, 0)
			value.completion += max(event.CompletionTokens, 0)
			value.cost += max(event.CostMicrounits, 0)
			daily[key] = value
		}
	}
	_, err = tx.CopyFrom(ctx, pgx.Identifier{"gateway", "audit_events"}, []string{
		"id", "tenant_id", "request_id", "actor_id", "api_key_id", "provider_code", "model", "protocol",
		"status_code", "prompt_tokens", "completion_tokens", "cost_microunits", "latency_ms",
		"request_preview", "response_preview", "labels", "recorded_at",
	}, pgx.CopyFromRows(rows))
	if err != nil {
		return r.flushError(err)
	}
	batch := &pgx.Batch{}
	for key, value := range daily {
		batch.Queue(`INSERT INTO gateway.usage_daily(usage_date,api_key_id,provider_code,model,currency,requests,failed_requests,prompt_tokens,completion_tokens,cost_microunits)
			VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10)
			ON CONFLICT(usage_date,api_key_id,provider_code,model,currency) DO UPDATE SET
			requests=gateway.usage_daily.requests+EXCLUDED.requests,
			failed_requests=gateway.usage_daily.failed_requests+EXCLUDED.failed_requests,
			prompt_tokens=gateway.usage_daily.prompt_tokens+EXCLUDED.prompt_tokens,
			completion_tokens=gateway.usage_daily.completion_tokens+EXCLUDED.completion_tokens,
			cost_microunits=gateway.usage_daily.cost_microunits+EXCLUDED.cost_microunits,
			updated_at=clock_timestamp()`, key.date, key.apiKeyID, key.provider, key.model, key.currency, value.requests, value.failed, value.prompt, value.completion, value.cost)
	}
	for _, alert := range alerts {
		batch.Queue(`INSERT INTO gateway.outbox_events(event_id,event_type,event_version,tenant_id,aggregate_type,aggregate_id,payload)
			VALUES($1,'content_policy.matched',1,$2,'gateway_request',$3,$4)`, alert.eventID, uuidPointer(alert.tenantID), alert.requestID, alert.payload)
	}
	results := tx.SendBatch(ctx, batch)
	if err := results.Close(); err != nil {
		return r.flushError(err)
	}
	if err := tx.Commit(ctx); err != nil {
		return r.flushError(err)
	}
	r.written.Add(uint64(len(events)))
	return nil
}

func (r *Recorder) flushError(err error) error {
	r.failures.Add(1)
	if r.logger != nil {
		r.logger.Error("audit batch flush failed", "error", err)
	}
	return err
}

func uuidValue(value string) pgtype.UUID {
	var result pgtype.UUID
	_ = result.Scan(value)
	return result
}

func uuidPointer(value *string) any {
	if value == nil || strings.TrimSpace(*value) == "" {
		return nil
	}
	// 非 UUID 身份(bootstrap key 等)返回 nil:audit_events.api_key_id 允许
	// NULL,写零值/非法值只会掩盖问题。
	if _, err := uuid.Parse(strings.TrimSpace(*value)); err != nil {
		return nil
	}
	return uuidValue(*value)
}

func nullString(value string) any {
	if value == "" {
		return nil
	}
	return value
}

func nullToken(value int64) any {
	if value <= 0 {
		return nil
	}
	return value
}

func nullCost(value int64) any {
	if value <= 0 {
		return nil
	}
	return value
}