chore: initial commit

2026-10-23 18:35:44 +08:00
commit f6921b7b23
69 changed files with 12148 additions and 0 deletions
@@ -0,0 +1,158 @@
+package server
+
+import (
+	"context"
+	"errors"
+	"log"
+	"strings"
+	"time"
+	"unicode/utf8"
+)
+
+const (
+	maxLogFieldBytes = 512
+	maxLogErrorBytes = 4096
+	maxLogUsageBytes = 16384
+	logCleanupBatch  = 1000
+)
+
+func (s *Server) SetLogRetention(days int) {
+	if days < 0 {
+		days = 0
+	}
+	s.logRetentionDays.Store(int64(days))
+}
+
+func truncateLogField(value string, limit int) string {
+	if len(value) <= limit {
+		return value
+	}
+	for limit > 0 && !utf8.RuneStart(value[limit]) {
+		limit--
+	}
+	return strings.Clone(value[:limit])
+}
+
+func boundedLog(l requestLog) requestLog {
+	for _, field := range []*string{&l.CreatedAt, &l.KeyID, &l.KeyName, &l.Channel, &l.RequestModel, &l.ActualModel, &l.Downstream, &l.Upstream} {
+		*field = truncateLogField(*field, maxLogFieldBytes)
+	}
+	l.Error = truncateLogField(l.Error, maxLogErrorBytes)
+	if len(l.Usage.Raw) > maxLogUsageBytes {
+		// Keep valid JSON without retaining an attacker-controlled oversized object.
+		l.Usage.Raw = []byte(`{"truncated":true}`)
+	}
+	l.RawUsage = nil
+	return l
+}
+
+func (s *Server) purgeExpiredLogs(ctx context.Context, days int) error {
+	if days <= 0 {
+		return nil
+	}
+	cutoff := time.Now().UTC().AddDate(0, 0, -days).Truncate(time.Second)
+	_, err := s.db.ExecContext(ctx, "DELETE FROM logs WHERE id IN (SELECT id FROM logs WHERE created_at<? AND julianday(created_at)<julianday(?) LIMIT ?)", cutoff.Format("2006-01-02T15:04:05"), cutoff.Format(time.RFC3339), logCleanupBatch)
+	return err
+}
+
+type logEntry struct {
+	value   requestLog
+	barrier chan error
+}
+
+func (s *Server) record(l requestLog) {
+	s.logMu.RLock()
+	defer s.logMu.RUnlock()
+	if s.logsClosed {
+		return
+	}
+	l = boundedLog(l)
+	l.Usage.Raw = append([]byte(nil), l.Usage.Raw...)
+	s.logQueue <- logEntry{value: l}
+}
+
+func (s *Server) flushLogs() error {
+	s.logMu.RLock()
+	if s.logsClosed {
+		s.logMu.RUnlock()
+		<-s.logDone
+		return s.logErr
+	}
+	done := make(chan error, 1)
+	s.logQueue <- logEntry{barrier: done}
+	s.logMu.RUnlock()
+	return <-done
+}
+
+func (s *Server) writeLogs() {
+	defer close(s.logDone)
+	ticker := time.NewTicker(100 * time.Millisecond)
+	defer ticker.Stop()
+	cleanup := time.NewTicker(time.Minute)
+	defer cleanup.Stop()
+	batch := make([]requestLog, 0, 100)
+	flush := func() {
+		if len(batch) == 0 {
+			return
+		}
+		for attempt := 0; ; attempt++ {
+			if s.persistLogs(batch) == nil {
+				break
+			}
+			if attempt == 3 {
+				s.logErr = errors.New("request log persistence failed after 3 retries")
+				log.Printf("ai-router: request log persistence failed; batch_size=%d retries=3", len(batch))
+				break
+			}
+			log.Printf("ai-router: retrying request log persistence; batch_size=%d retry=%d", len(batch), attempt+1)
+			time.Sleep(time.Duration(attempt+1) * 10 * time.Millisecond)
+		}
+		clear(batch)
+		batch = batch[:0]
+	}
+	for {
+		select {
+		case entry, ok := <-s.logQueue:
+			if !ok {
+				flush()
+				return
+			}
+			if entry.barrier != nil {
+				flush()
+				entry.barrier <- s.logErr
+				continue
+			}
+			batch = append(batch, entry.value)
+			if len(batch) == cap(batch) {
+				flush()
+			}
+		case <-ticker.C:
+			flush()
+		case <-cleanup.C:
+			if days := s.logRetentionDays.Load(); days > 0 {
+				ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+				err := s.purgeExpiredLogs(ctx, int(days))
+				cancel()
+				if err != nil {
+					log.Print("ai-router: request log retention cleanup failed")
+				}
+			}
+		}
+	}
+}
+
+func (s *Server) persistLogs(batch []requestLog) error {
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	tx, err := s.db.BeginTx(ctx, nil)
+	if err != nil {
+		return err
+	}
+	defer tx.Rollback()
+	for _, l := range batch {
+		if err := insertLog(ctx, tx, l); err != nil {
+			return err
+		}
+	}
+	return tx.Commit()
+}