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

package gateway

import (
	"bufio"
	"encoding/json"
	"io"
	"net"
	"net/http"
	"sync"
	"time"

	"aigateway.local/core/internal/apikey"
	auditpkg "aigateway.local/core/internal/audit"
	"aigateway.local/core/internal/contentpolicy"
	"aigateway.local/core/internal/pricing"
)

const auditRequestCaptureBytes = 64 << 10

type AuditRecorder interface {
	Record(auditpkg.Event) bool
}

type auditSpan struct {
	recorder AuditRecorder
	started  time.Time
	event    auditpkg.Event
	mu       sync.Mutex
	capture  *captureReadCloser
	target   string
	pricing  *pricing.Service
	policies []contentpolicy.Match
	redacted bool
}

func newAuditSpan(recorder AuditRecorder, principal apikey.Principal, request *http.Request, started time.Time) *auditSpan {
	if recorder == nil {
		return nil
	}
	event := auditpkg.Event{TenantID: principal.TenantID, RequestID: RequestID(request.Context()), Protocol: request.URL.Path, RecordedAt: time.Now().UTC()}
	if principal.APIKeyID != "" {
		id := principal.APIKeyID
		event.APIKeyID = &id
	}
	return &auditSpan{recorder: recorder, started: started, event: event}
}

func (s *auditSpan) setModel(model string) {
	if s != nil && model != "" {
		s.mu.Lock()
		s.event.Model = model
		s.mu.Unlock()
	}
}

func (s *auditSpan) setRoute(providerCode, targetModel string) {
	if s == nil {
		return
	}
	s.mu.Lock()
	s.event.ProviderCode = providerCode
	s.target = targetModel
	s.mu.Unlock()
}

func (s *auditSpan) setUsage(usage TokenUsage) {
	if s == nil {
		return
	}
	s.mu.Lock()
	s.event.PromptTokens = usage.Input
	s.event.CompletionTokens = usage.Output
	s.calculateCostLocked()
	s.mu.Unlock()
}

func (s *auditSpan) setContentPolicy(result contentpolicy.Result) {
	if s == nil {
		return
	}
	s.mu.Lock()
	s.policies = append(s.policies, result.Matches...)
	s.redacted = result.Redacted
	s.mu.Unlock()
}

func (s *auditSpan) captureBody(body io.ReadCloser) io.ReadCloser {
	if s == nil || body == nil || body == http.NoBody {
		return body
	}
	capture := &captureReadCloser{ReadCloser: body, limit: auditRequestCaptureBytes}
	s.capture = capture
	return capture
}

func (s *auditSpan) finish(status int) {
	if s == nil {
		return
	}
	s.mu.Lock()
	if s.event.Model == "" && s.capture != nil {
		var payload map[string]json.RawMessage
		if json.Unmarshal(s.capture.buffer, &payload) == nil {
			_ = json.Unmarshal(payload["model"], &s.event.Model)
		}
	}
	s.calculateCostLocked()
	s.event.StatusCode = status
	s.event.LatencyMS = int(time.Since(s.started).Milliseconds())
	s.event.RecordedAt = time.Now().UTC()
	labels := map[string]any{}
	if s.target != "" && s.target != s.event.Model {
		labels["target_model"] = s.target
		labels["routed"] = true
	}
	if len(s.policies) > 0 {
		labels["content_policies"] = s.policies
	}
	if s.redacted {
		labels["content_redacted"] = true
	}
	if s.event.PriceID != "" {
		labels["price_id"] = s.event.PriceID
		labels["currency"] = s.event.Currency
	}
	s.event.Labels = labels
	event := s.event
	s.mu.Unlock()
	s.recorder.Record(event)
}

func (s *auditSpan) calculateCostLocked() {
	if s.pricing == nil || s.event.ProviderCode == "" || s.event.Model == "" {
		return
	}
	model := s.event.Model
	if s.target != "" {
		model = s.target
	}
	cost := s.pricing.Calculate(s.event.ProviderCode, model, s.event.PromptTokens, s.event.CompletionTokens, s.started.UTC())
	s.event.CostMicrounits = cost.Microunits
	s.event.PriceID = cost.PriceID
	s.event.Currency = cost.Currency
}

type captureReadCloser struct {
	io.ReadCloser
	buffer []byte
	limit  int
}

func (r *captureReadCloser) Read(buffer []byte) (int, error) {
	n, err := r.ReadCloser.Read(buffer)
	if n > 0 && len(r.buffer) < r.limit {
		remaining := r.limit - len(r.buffer)
		r.buffer = append(r.buffer, buffer[:min(n, remaining)]...)
	}
	return n, err
}

type statusResponseWriter struct {
	http.ResponseWriter
	status int
}

func (w *statusResponseWriter) WriteHeader(status int) {
	if w.status != 0 {
		return
	}
	w.status = status
	w.ResponseWriter.WriteHeader(status)
}

func (w *statusResponseWriter) Write(buffer []byte) (int, error) {
	if w.status == 0 {
		w.WriteHeader(http.StatusOK)
	}
	return w.ResponseWriter.Write(buffer)
}

func (w *statusResponseWriter) Status() int {
	if w.status == 0 {
		return http.StatusOK
	}
	return w.status
}

func (w *statusResponseWriter) Flush() {
	if w.status == 0 {
		w.WriteHeader(http.StatusOK)
	}
	if flusher, ok := w.ResponseWriter.(http.Flusher); ok {
		flusher.Flush()
	}
}

func (w *statusResponseWriter) Unwrap() http.ResponseWriter { return w.ResponseWriter }

func (w *statusResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	if hijacker, ok := w.ResponseWriter.(http.Hijacker); ok {
		return hijacker.Hijack()
	}
	return nil, nil, http.ErrNotSupported
}

func (w *statusResponseWriter) Push(target string, options *http.PushOptions) error {
	if pusher, ok := w.ResponseWriter.(http.Pusher); ok {
		return pusher.Push(target, options)
	}
	return http.ErrNotSupported
}

func (w *statusResponseWriter) ReadFrom(reader io.Reader) (int64, error) {
	if w.status == 0 {
		w.WriteHeader(http.StatusOK)
	}
	if readerFrom, ok := w.ResponseWriter.(io.ReaderFrom); ok {
		return readerFrom.ReadFrom(reader)
	}
	return io.Copy(struct{ io.Writer }{w.ResponseWriter}, reader)
}