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,248 @@
+package server
+
+import (
+	"ai-router/internal/protocol"
+	"bytes"
+	"context"
+	"database/sql"
+	"log"
+	"net/http"
+	"net/http/httptest"
+	"path/filepath"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+)
+
+func waitLogCondition(t *testing.T, condition func() bool) {
+	t.Helper()
+	deadline := time.Now().Add(3 * time.Second)
+	for !condition() {
+		if time.Now().After(deadline) {
+			t.Fatal("condition timed out")
+		}
+		time.Sleep(time.Millisecond)
+	}
+}
+
+func logCount(t *testing.T, db *sql.DB) int {
+	t.Helper()
+	var n int
+	if err := db.QueryRow("SELECT COUNT(*) FROM logs").Scan(&n); err != nil {
+		t.Fatal(err)
+	}
+	return n
+}
+
+func TestLogAsyncAndRawCopy(t *testing.T) {
+	s := testServer(t)
+	conn, err := s.db.Conn(context.Background())
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer conn.Close()
+	raw := []byte(`{"input_tokens":12}`)
+	done := make(chan struct{})
+	go func() {
+		s.record(requestLog{Usage: protocol.Usage{Raw: raw}})
+		close(done)
+	}()
+	select {
+	case <-done:
+	case <-time.After(time.Second):
+		t.Fatal("record waited for database")
+	}
+	raw[0] = 'x'
+	conn.Close()
+	if err := s.flushLogs(); err != nil {
+		t.Fatal(err)
+	}
+	var got string
+	if err := s.db.QueryRow("SELECT raw_usage FROM logs").Scan(&got); err != nil {
+		t.Fatal(err)
+	}
+	if got != `{"input_tokens":12}` {
+		t.Fatal(got)
+	}
+}
+
+func TestLogTimerAndBatchDrain(t *testing.T) {
+	s := testServer(t)
+	s.record(requestLog{})
+	waitLogCondition(t, func() bool { return logCount(t, s.db) == 1 })
+	for i := 0; i < 257; i++ {
+		s.record(requestLog{})
+	}
+	if err := s.flushLogs(); err != nil {
+		t.Fatal(err)
+	}
+	if n := logCount(t, s.db); n != 258 {
+		t.Fatal(n)
+	}
+}
+
+func TestLogBackpressureAndConcurrentClose(t *testing.T) {
+	path := filepath.Join(t.TempDir(), "logs.db")
+	s, err := New(path, "password", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { s.Close() })
+	conn, err := s.db.Conn(context.Background())
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer conn.Close()
+	s.record(requestLog{})
+	waitLogCondition(t, func() bool { return s.db.Stats().WaitCount > 0 })
+	if cap(s.logQueue) != 1024 {
+		t.Fatal(cap(s.logQueue))
+	}
+	for i := 0; i < 1024; i++ {
+		s.record(requestLog{})
+	}
+	enqueued := make(chan struct{})
+	go func() { s.record(requestLog{}); close(enqueued) }()
+	select {
+	case <-enqueued:
+		t.Fatal("full queue did not apply backpressure")
+	case <-time.After(20 * time.Millisecond):
+	}
+	var wg sync.WaitGroup
+	for i := 0; i < 8; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			if err := s.Close(); err != nil {
+				t.Error(err)
+			}
+		}()
+	}
+	conn.Close()
+	done := make(chan struct{})
+	go func() { wg.Wait(); close(done) }()
+	select {
+	case <-done:
+	case <-time.After(5 * time.Second):
+		t.Fatal("Close deadlocked")
+	}
+	<-enqueued
+	s.record(requestLog{})
+	if err := s.flushLogs(); err != nil {
+		t.Fatal(err)
+	}
+	db, err := sql.Open("sqlite", path)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer db.Close()
+	if n := logCount(t, db); n != 1026 {
+		t.Fatal(n)
+	}
+}
+
+func TestLogPersistenceErrorIsStickyAndRedacted(t *testing.T) {
+	s := testServer(t)
+	var output bytes.Buffer
+	previous := log.Writer()
+	log.SetOutput(&output)
+	defer log.SetOutput(previous)
+	if _, err := s.db.Exec(`CREATE TRIGGER reject_log BEFORE INSERT ON logs WHEN NEW.key_name='secret' BEGIN SELECT RAISE(FAIL,'private database secret'); END`); err != nil {
+		t.Fatal(err)
+	}
+	s.record(requestLog{KeyName: "ok"})
+	s.record(requestLog{KeyName: "secret"})
+	start := time.Now()
+	err := s.flushLogs()
+	if err == nil {
+		t.Fatal("missing persistence error")
+	}
+	if time.Since(start) > time.Second {
+		t.Fatal("unbounded retries")
+	}
+	if n := logCount(t, s.db); n != 0 {
+		t.Fatalf("partial transaction: %d", n)
+	}
+	if strings.Contains(output.String(), "secret") || strings.Contains(err.Error(), "secret") {
+		t.Fatal("error leaked metadata")
+	}
+	if n := strings.Count(output.String(), "retrying"); n != 3 {
+		t.Fatalf("retries: %d", n)
+	}
+	if _, err := s.db.Exec("DROP TRIGGER reject_log"); err != nil {
+		t.Fatal(err)
+	}
+	s.record(requestLog{})
+	if err := s.flushLogs(); err == nil {
+		t.Fatal("error was not sticky")
+	}
+	if n := logCount(t, s.db); n != 1 {
+		t.Fatal(n)
+	}
+	if err := s.Close(); err == nil {
+		t.Fatal("Close lost persistence error")
+	}
+	if err := s.Close(); err == nil {
+		t.Fatal("second Close lost error")
+	}
+}
+
+func TestCloseWaitsForActiveRequest(t *testing.T) {
+	path := filepath.Join(t.TempDir(), "logs.db")
+	s, err := New(path, "test-password", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { s.Close() })
+	started, release := make(chan struct{}), make(chan struct{})
+	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		close(started)
+		<-release
+		w.Header().Set("Content-Type", "application/json")
+		w.Write([]byte(integrationResponse("openai-chat")))
+	}))
+	defer up.Close()
+	defer func() {
+		select {
+		case <-release:
+		default:
+			close(release)
+		}
+	}()
+	cookie := loginCookie(t, s)
+	ch := integrationChannel(t, s, cookie, "abc", "openai-chat", up.URL)
+	key := integrationCreateKey(t, s, cookie, "active", ch)
+	requestDone := make(chan struct{})
+	go func() { integrationProxy(s, key, "openai-chat", "abc/gpt-5", false); close(requestDone) }()
+	<-started
+	closed := make(chan error, 1)
+	go func() { closed <- s.Close() }()
+	waitLogCondition(t, func() bool { s.lifecycleMu.Lock(); defer s.lifecycleMu.Unlock(); return s.closing })
+	if w := call(s, "GET", "/v1/models", "", nil); w.Code != 503 {
+		t.Fatal(w.Code)
+	}
+	select {
+	case <-closed:
+		t.Fatal("Close skipped active request")
+	default:
+	}
+	close(release)
+	<-requestDone
+	select {
+	case err := <-closed:
+		if err != nil {
+			t.Fatal(err)
+		}
+	case <-time.After(3 * time.Second):
+		t.Fatal("Close did not drain")
+	}
+	db, err := sql.Open("sqlite", path)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer db.Close()
+	if n := logCount(t, db); n != 1 {
+		t.Fatal(n)
+	}
+}