WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes

2026-09-22 16:19:06 +08:00
commit f7fc992f57
3929 changed files with 1006607 additions and 0 deletions
@@ -0,0 +1,367 @@
+package workbench
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"errors"
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/gin-gonic/gin"
+	"github.com/google/uuid"
+	"gorm.io/driver/postgres"
+	"gorm.io/gorm"
+	"gorm.io/gorm/logger"
+)
+
+func seedQuota(t *testing.T, s *Service, quota UsageQuota) UsageQuota {
+	t.Helper()
+	quota.ID = uuid.NewString()
+	if quota.Period == "" {
+		quota.Period = "daily"
+	}
+	if quota.OveragePolicy == "" {
+		quota.OveragePolicy = "hard_limit"
+	}
+	if err := s.db.Create(&quota).Error; err != nil {
+		t.Fatal(err)
+	}
+	return quota
+}
+
+func TestQuotaReservationPoliciesAndRollback(t *testing.T) {
+	for _, mode := range []string{"token", "points", "concurrency", "soft", "warn", "disabled", "unlimited"} {
+		t.Run(mode, func(t *testing.T) {
+			s := sqliteWorkbench(t)
+			ctx := workbenchContext(1, "alice", types.TenantRoleContributor)
+			seedBalance(t, s, 1, "alice", 100000)
+			quota := UsageQuota{TenantID: 1, UserID: "alice", TokenLimit: 100, Enabled: true, WarningRatio: .8}
+			switch mode {
+			case "points":
+				quota.TokenLimit, quota.PointsLimit = 0, 1000
+			case "concurrency":
+				quota.TokenLimit, quota.ConcurrentLimit = 0, 1
+			case "soft":
+				quota.OveragePolicy = "soft_limit"
+			case "warn":
+				quota.OveragePolicy = "warn"
+			case "disabled":
+				quota.Enabled = false
+			case "unlimited":
+				quota.TokenLimit, quota.WarningRatio = 0, 0
+			}
+			seedQuota(t, s, quota)
+			first, err := s.reserveBalance(ctx, 900)
+			if err != nil {
+				t.Fatal(err)
+			}
+			second, err := s.reserveBalance(ctx, 200)
+			blocked := mode == "token" || mode == "points" || mode == "concurrency" || mode == "soft"
+			if blocked != errors.Is(err, ErrInvalid) {
+				t.Fatalf("blocked=%v, error=%v", blocked, err)
+			}
+			wantReserved := int64(900)
+			if !blocked {
+				wantReserved += 200
+			}
+			if got := accountBalance(t, s, 1, "alice"); got.Reserved != wantReserved || got.Balance != 100000 {
+				t.Fatalf("failed reservation leaked: %+v", got)
+			}
+			settled := Usage{ID: uuid.NewString(), CallID: first.ID, TenantID: 1, OwnerID: "alice", InputTokens: 30, Points: 300, CreatedAt: first.CreatedAt, Metadata: JSON{}}
+			if err := s.settleBalance(ctx, first, settled); err != nil {
+				t.Fatal(err)
+			}
+			if blocked {
+				if _, err := s.reserveBalance(ctx, 200); err != nil {
+					t.Fatalf("released capacity unavailable: %v", err)
+				}
+			} else if second == nil {
+				t.Fatal("missing accepted hold")
+			}
+			var warnings int64
+			s.db.Model(&Event{}).Count(&warnings)
+			if mode == "unlimited" || mode == "disabled" || mode == "concurrency" {
+				if warnings != 0 {
+					t.Fatal("unlimited or disabled quota warned")
+				}
+			} else if warnings != 1 {
+				t.Fatalf("expected one warning per period, got %d", warnings)
+			}
+		})
+	}
+}
+
+func TestQuotaPeriodsIsolationAndExpiredCalls(t *testing.T) {
+	s := sqliteWorkbench(t)
+	now := time.Date(2026, 9, 30, 16, 0, 0, 0, time.UTC) // October 1 midnight in Shanghai.
+	for _, period := range []string{"daily", "monthly"} {
+		start, end := quotaPeriod(period, now)
+		if !start.Equal(now) || end.In(quotaLocation).Hour() != 0 {
+			t.Fatalf("wrong period: %v %v", start, end)
+		}
+		if period == "monthly" && end.In(quotaLocation).Month() != time.November {
+			t.Fatal("month boundary wrong")
+		}
+	}
+	month, monthEnd := quotaPeriod("monthly", time.Date(2028, 2, 29, 6, 0, 0, 0, time.UTC))
+	if month.Day() != 1 || monthEnd.Month() != time.March || monthEnd.Sub(month) != 29*24*time.Hour {
+		t.Fatal("leap year boundary wrong")
+	}
+	quota := UsageQuota{TenantID: 1, UserID: "alice", Period: "daily", TokenLimit: 100, PointsLimit: 1000, ConcurrentLimit: 2, WarningRatio: .8, Enabled: true, OveragePolicy: "hard_limit"}
+	_, end := quotaPeriod("daily", now)
+	for _, row := range []Usage{
+		{TenantID: 1, OwnerID: "alice", InputTokens: 1000, CreatedAt: now.Add(-time.Nanosecond)},
+		{TenantID: 1, OwnerID: "alice", InputTokens: 20, Points: 200, CreatedAt: now},
+		{TenantID: 1, OwnerID: "alice", InputTokens: 1000, CreatedAt: end},
+		{TenantID: 2, OwnerID: "alice", InputTokens: 1000, CreatedAt: now},
+		{TenantID: 1, OwnerID: "bob", InputTokens: 1000, CreatedAt: now},
+		{TenantID: 1, OwnerID: "alice", InputTokens: 1000, CreatedAt: now, Status: "void"},
+	} {
+		row.ID, row.CallID, row.Metadata = uuid.NewString(), uuid.NewString(), JSON{}
+		if err := s.db.Create(&row).Error; err != nil {
+			t.Fatal(err)
+		}
+	}
+	for _, row := range []BalanceHold{
+		{TenantID: 1, UserID: "alice", Amount: 500, Status: "reserved", CreatedAt: now, ExpiresAt: now.Add(time.Minute)},
+		{TenantID: 1, UserID: "alice", Amount: 900, Status: "reserved", CreatedAt: now.Add(-time.Minute), ExpiresAt: now.Add(time.Minute)},
+		{TenantID: 1, UserID: "alice", Amount: 900, Status: "reserved", CreatedAt: now, ExpiresAt: now},
+		{TenantID: 2, UserID: "alice", Amount: 900, Status: "reserved", CreatedAt: now, ExpiresAt: now.Add(time.Minute)},
+	} {
+		row.ID = uuid.NewString()
+		if err := s.db.Create(&row).Error; err != nil {
+			t.Fatal(err)
+		}
+	}
+	state, err := loadQuotaState(s.db, quota, now)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if state.Used.Tokens != 20 || state.Used.Points != 200 || state.Used.ReservedTokens != 50 || state.Used.ReservedPoints != 500 || state.Used.ActiveCalls != 2 || state.Allowed || state.Warning {
+		t.Fatalf("bad state: %+v", state)
+	}
+	quota.WarningRatio = 0
+	state, err = loadQuotaState(s.db, quota, now)
+	if err != nil || !state.Warning {
+		t.Fatalf("zero warning threshold ignored: %+v %v", state, err)
+	}
+}
+
+func TestQuotaSettlementNormalizesCallerTimezone(t *testing.T) {
+	s := sqliteWorkbench(t)
+	ctx := workbenchContext(1, "alice", types.TenantRoleContributor)
+	now := time.Date(2026, 10, 1, 0, 0, 0, 0, quotaLocation)
+	for i, at := range []time.Time{now.Add(-time.Nanosecond), now, now.AddDate(0, 0, 1)} {
+		usage := Usage{ID: uuid.NewString(), CallID: uuid.NewString(), TenantID: 1, OwnerID: "alice", InputTokens: int64(i + 1), Points: int64((i + 1) * 10), CreatedAt: at, Metadata: JSON{}}
+		if err := s.settleBalance(ctx, nil, usage); err != nil {
+			t.Fatal(err)
+		}
+	}
+	quota := UsageQuota{TenantID: 1, UserID: "alice", Period: "daily"}
+	for _, clock := range []time.Time{now, now.UTC(), now.In(time.FixedZone("west", -7*60*60))} {
+		state, err := loadQuotaState(s.db, quota, clock)
+		if err != nil || state.Used.Tokens != 2 || state.Used.Points != 20 {
+			t.Fatalf("timezone changed period usage: %+v, %v", state, err)
+		}
+	}
+}
+
+func quotaRequest(t *testing.T, s *Service, ctx context.Context, method, path string, body any) *httptest.ResponseRecorder {
+	t.Helper()
+	gin.SetMode(gin.TestMode)
+	router := gin.New()
+	h := NewHandler(s)
+	router.GET("/quotas", h.Quotas)
+	router.PUT("/quotas/:userID", h.SaveQuota)
+	router.DELETE("/quotas/:id", h.DeleteQuota)
+	encoded, _ := json.Marshal(body)
+	request := httptest.NewRequest(method, path, bytes.NewReader(encoded)).WithContext(ctx)
+	request.Header.Set("Content-Type", "application/json")
+	response := httptest.NewRecorder()
+	router.ServeHTTP(response, request)
+	return response
+}
+
+func TestQuotaAdminValidationUpsertAndTenantIsolation(t *testing.T) {
+	s := sqliteWorkbench(t)
+	if err := s.db.AutoMigrate(&types.TenantMember{}, &types.User{}); err != nil {
+		t.Fatal(err)
+	}
+	if err := s.db.Create(&types.User{ID: "alice", Username: "Alice", Email: "alice@example.test"}).Error; err != nil {
+		t.Fatal(err)
+	}
+	if err := s.db.Create(&types.TenantMember{TenantID: 1, UserID: "alice", Role: types.TenantRoleContributor}).Error; err != nil {
+		t.Fatal(err)
+	}
+	admin := workbenchContext(1, "admin", types.TenantRoleAdmin)
+	body := UsageQuota{Period: "daily", Enabled: true, TokenLimit: 100, WarningRatio: .8, OveragePolicy: "hard_limit"}
+	for _, role := range []types.TenantRole{types.TenantRoleViewer, types.TenantRoleContributor} {
+		ctx := workbenchContext(1, "alice", role)
+		for _, method := range []string{http.MethodGet, http.MethodPut, http.MethodDelete} {
+			path := "/quotas"
+			if method != http.MethodGet {
+				path += "/alice"
+			}
+			if got := quotaRequest(t, s, ctx, method, path, body).Code; got != 403 {
+				t.Fatalf("%s role %s status %d", method, role, got)
+			}
+		}
+	}
+	if response := quotaRequest(t, s, admin, "PUT", "/quotas/alice", body); response.Code != 200 {
+		t.Fatal(response.Body.String())
+	}
+	body.TokenLimit, body.Enabled, body.WarningRatio = 0, false, 0
+	if response := quotaRequest(t, s, admin, "PUT", "/quotas/alice", body); response.Code != 200 {
+		t.Fatal(response.Body.String())
+	}
+	var rows []UsageQuota
+	s.db.Find(&rows)
+	if len(rows) != 1 || rows[0].TokenLimit != 0 || rows[0].Enabled || rows[0].WarningRatio != 0 {
+		t.Fatalf("upsert duplicated or zero fields lost: %+v", rows)
+	}
+	if response := quotaRequest(t, s, admin, "GET", "/quotas", nil); response.Code != 200 || !strings.Contains(response.Body.String(), "Alice") {
+		t.Fatal(response.Body.String())
+	}
+	other := workbenchContext(2, "admin", types.TenantRoleAdmin)
+	if response := quotaRequest(t, s, other, "GET", "/quotas", nil); response.Code != 200 || strings.Contains(response.Body.String(), "Alice") {
+		t.Fatal("cross tenant list")
+	}
+	if got := quotaRequest(t, s, other, "DELETE", "/quotas/"+rows[0].ID, nil).Code; got != 404 {
+		t.Fatalf("cross tenant delete: %d", got)
+	}
+	if got := quotaRequest(t, s, other, "PUT", "/quotas/alice", body).Code; got != 404 {
+		t.Fatalf("cross tenant upsert: %d", got)
+	}
+	for _, invalid := range []UsageQuota{
+		{Period: "weekly", OveragePolicy: "warn"}, {Period: "daily", OveragePolicy: "free"}, {Period: "daily", OveragePolicy: "warn", PointsLimit: -1}, {Period: "daily", OveragePolicy: "warn", WarningRatio: 1.1}, {Period: "daily", OveragePolicy: "warn", ConcurrentLimit: 10001},
+	} {
+		if got := quotaRequest(t, s, admin, "PUT", "/quotas/alice", invalid).Code; got != 400 {
+			t.Fatalf("invalid quota accepted: %+v", invalid)
+		}
+	}
+	if got := quotaRequest(t, s, admin, "DELETE", "/quotas/"+rows[0].ID, nil).Code; got != 200 {
+		t.Fatalf("delete status %d", got)
+	}
+}
+
+func TestPostgresQuotaReservationRaces(t *testing.T) {
+	dsn := os.Getenv("WORKBENCH_TEST_POSTGRES_DSN")
+	if dsn == "" {
+		t.Skip("WORKBENCH_TEST_POSTGRES_DSN is not set")
+	}
+	db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)})
+	if err != nil {
+		t.Fatal("PostgreSQL unavailable")
+	}
+	connection, _ := db.DB()
+	defer connection.Close()
+	schema := "wb_quota_test_" + strings.ReplaceAll(uuid.NewString(), "-", "")
+	if err := db.Exec("CREATE SCHEMA " + schema).Error; err != nil {
+		t.Fatal(err)
+	}
+	defer db.Exec("DROP SCHEMA " + schema + " CASCADE")
+	testDSN := dsn + " search_path=" + schema
+	if strings.Contains(dsn, "://") {
+		separator := "?"
+		if strings.Contains(dsn, "?") {
+			separator = "&"
+		}
+		testDSN = dsn + separator + "search_path=" + schema
+	}
+	testDB, err := gorm.Open(postgres.Open(testDSN), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)})
+	if err != nil {
+		t.Fatal("isolated PostgreSQL connection unavailable")
+	}
+	pool, _ := testDB.DB()
+	defer pool.Close()
+	if err := testDB.AutoMigrate(&Balance{}, &BalanceHold{}, &Usage{}, &UsageQuota{}, &Transaction{}, &Event{}); err != nil {
+		t.Fatal(err)
+	}
+	s := NewService(testDB)
+	for _, mode := range []string{"token", "points", "concurrency", "unlimited"} {
+		t.Run(mode, func(t *testing.T) {
+			ctx := workbenchContext(1, mode, types.TenantRoleContributor)
+			seedBalance(t, s, 1, mode, 1000000)
+			q := UsageQuota{TenantID: 1, UserID: mode, Enabled: true, WarningRatio: .8}
+			want := 2
+			switch mode {
+			case "token":
+				q.TokenLimit = 200
+			case "points":
+				q.PointsLimit = 2000
+			case "concurrency":
+				q.ConcurrentLimit = 2
+			case "unlimited":
+				want = 12
+			}
+			seedQuota(t, s, q)
+			var workers sync.WaitGroup
+			results := make(chan *BalanceHold, 12)
+			failures := make(chan error, 12)
+			start := make(chan struct{})
+			for i := 0; i < 12; i++ {
+				workers.Add(1)
+				go func() {
+					defer workers.Done()
+					<-start
+					hold, err := s.reserveBalance(ctx, 1000)
+					if err != nil {
+						failures <- err
+					} else {
+						results <- hold
+					}
+				}()
+			}
+			close(start)
+			workers.Wait()
+			close(results)
+			close(failures)
+			if len(results) != want {
+				t.Fatalf("accepted %d calls; want %d", len(results), want)
+			}
+			for err := range failures {
+				if !errors.Is(err, ErrInvalid) {
+					t.Fatalf("unexpected reservation error: %v", err)
+				}
+			}
+			if balance := accountBalance(t, s, 1, mode); balance.Reserved != int64(want*1000) {
+				t.Fatalf("hold/account mismatch: %+v", balance)
+			}
+			var holds []*BalanceHold
+			for hold := range results {
+				holds = append(holds, hold)
+			}
+			// Concurrent settlements release all slots and remain idempotent.
+			settleErrors := make(chan error, len(holds)*2)
+			for _, hold := range holds {
+				for repeat := 0; repeat < 2; repeat++ {
+					workers.Add(1)
+					go func(hold *BalanceHold) {
+						defer workers.Done()
+						settleErrors <- s.settleBalance(ctx, hold, Usage{ID: hold.ID, CallID: hold.ID, TenantID: 1, OwnerID: mode, InputTokens: 10, Points: 100, CreatedAt: hold.CreatedAt, Metadata: JSON{}})
+					}(hold)
+				}
+			}
+			workers.Wait()
+			close(settleErrors)
+			for err := range settleErrors {
+				if err != nil {
+					t.Fatal(err)
+				}
+			}
+			if balance := accountBalance(t, s, 1, mode); balance.Reserved != 0 || balance.Balance != 1000000-int64(want*100) {
+				t.Fatalf("settlement mismatch: %+v", balance)
+			}
+			if _, err := s.reserveBalance(ctx, 1000); err != nil {
+				t.Fatalf("settled capacity not reusable: %v", err)
+			}
+		})
+	}
+}