3e43233e0e
Add persistent storage for agent allowance quotas using go-store (SQLite KV). SQLiteStore implements all 11 AllowanceStore methods with JSON serialisation, mutex-guarded read-modify-write for atomic increments, and proper time.Duration handling via nanosecond int64 encoding. - allowance_sqlite.go: full AllowanceStore implementation with 4 KV groups - allowance_sqlite_test.go: 26 tests covering CRUD, persistence across close/reopen, concurrent access (10 goroutines), and AllowanceService integration - config.go: AllowanceConfig struct + NewAllowanceStoreFromConfig factory - go.mod: add forge.lthn.ai/core/go-store dependency Co-Authored-By: Virgil <virgil@lethean.io> Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
465 lines
12 KiB
Go
465 lines
12 KiB
Go
package agentic
|
|
|
|
import (
|
|
"path/filepath"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
// newTestSQLiteStore creates a SQLiteStore in a temp directory.
|
|
func newTestSQLiteStore(t *testing.T) *SQLiteStore {
|
|
t.Helper()
|
|
dbPath := filepath.Join(t.TempDir(), "test.db")
|
|
s, err := NewSQLiteStore(dbPath)
|
|
require.NoError(t, err)
|
|
t.Cleanup(func() { _ = s.Close() })
|
|
return s
|
|
}
|
|
|
|
// --- SetAllowance / GetAllowance ---
|
|
|
|
func TestSQLiteStore_SetGetAllowance_Good(t *testing.T) {
|
|
s := newTestSQLiteStore(t)
|
|
|
|
a := &AgentAllowance{
|
|
AgentID: "agent-1",
|
|
DailyTokenLimit: 100000,
|
|
DailyJobLimit: 10,
|
|
ConcurrentJobs: 2,
|
|
MaxJobDuration: 30 * time.Minute,
|
|
ModelAllowlist: []string{"claude-sonnet-4-5-20250929"},
|
|
}
|
|
|
|
err := s.SetAllowance(a)
|
|
require.NoError(t, err)
|
|
|
|
got, err := s.GetAllowance("agent-1")
|
|
require.NoError(t, err)
|
|
assert.Equal(t, a.AgentID, got.AgentID)
|
|
assert.Equal(t, a.DailyTokenLimit, got.DailyTokenLimit)
|
|
assert.Equal(t, a.DailyJobLimit, got.DailyJobLimit)
|
|
assert.Equal(t, a.ConcurrentJobs, got.ConcurrentJobs)
|
|
assert.Equal(t, a.MaxJobDuration, got.MaxJobDuration)
|
|
assert.Equal(t, a.ModelAllowlist, got.ModelAllowlist)
|
|
}
|
|
|
|
func TestSQLiteStore_GetAllowance_Bad_NotFound(t *testing.T) {
|
|
s := newTestSQLiteStore(t)
|
|
_, err := s.GetAllowance("nonexistent")
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), "allowance not found")
|
|
}
|
|
|
|
func TestSQLiteStore_SetAllowance_Good_Overwrite(t *testing.T) {
|
|
s := newTestSQLiteStore(t)
|
|
|
|
_ = s.SetAllowance(&AgentAllowance{AgentID: "agent-1", DailyTokenLimit: 100})
|
|
_ = s.SetAllowance(&AgentAllowance{AgentID: "agent-1", DailyTokenLimit: 200})
|
|
|
|
got, err := s.GetAllowance("agent-1")
|
|
require.NoError(t, err)
|
|
assert.Equal(t, int64(200), got.DailyTokenLimit)
|
|
}
|
|
|
|
// --- GetUsage / IncrementUsage ---
|
|
|
|
func TestSQLiteStore_GetUsage_Good_Default(t *testing.T) {
|
|
s := newTestSQLiteStore(t)
|
|
|
|
u, err := s.GetUsage("agent-1")
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "agent-1", u.AgentID)
|
|
assert.Equal(t, int64(0), u.TokensUsed)
|
|
assert.Equal(t, 0, u.JobsStarted)
|
|
assert.Equal(t, 0, u.ActiveJobs)
|
|
}
|
|
|
|
func TestSQLiteStore_IncrementUsage_Good(t *testing.T) {
|
|
s := newTestSQLiteStore(t)
|
|
|
|
err := s.IncrementUsage("agent-1", 5000, 1)
|
|
require.NoError(t, err)
|
|
|
|
u, err := s.GetUsage("agent-1")
|
|
require.NoError(t, err)
|
|
assert.Equal(t, int64(5000), u.TokensUsed)
|
|
assert.Equal(t, 1, u.JobsStarted)
|
|
assert.Equal(t, 1, u.ActiveJobs)
|
|
}
|
|
|
|
func TestSQLiteStore_IncrementUsage_Good_Accumulates(t *testing.T) {
|
|
s := newTestSQLiteStore(t)
|
|
|
|
_ = s.IncrementUsage("agent-1", 1000, 1)
|
|
_ = s.IncrementUsage("agent-1", 2000, 1)
|
|
_ = s.IncrementUsage("agent-1", 3000, 0)
|
|
|
|
u, err := s.GetUsage("agent-1")
|
|
require.NoError(t, err)
|
|
assert.Equal(t, int64(6000), u.TokensUsed)
|
|
assert.Equal(t, 2, u.JobsStarted)
|
|
assert.Equal(t, 2, u.ActiveJobs)
|
|
}
|
|
|
|
// --- DecrementActiveJobs ---
|
|
|
|
func TestSQLiteStore_DecrementActiveJobs_Good(t *testing.T) {
|
|
s := newTestSQLiteStore(t)
|
|
|
|
_ = s.IncrementUsage("agent-1", 0, 2)
|
|
_ = s.DecrementActiveJobs("agent-1")
|
|
|
|
u, _ := s.GetUsage("agent-1")
|
|
assert.Equal(t, 1, u.ActiveJobs)
|
|
}
|
|
|
|
func TestSQLiteStore_DecrementActiveJobs_Good_FloorAtZero(t *testing.T) {
|
|
s := newTestSQLiteStore(t)
|
|
|
|
_ = s.DecrementActiveJobs("agent-1") // no usage record yet
|
|
_ = s.IncrementUsage("agent-1", 0, 0)
|
|
_ = s.DecrementActiveJobs("agent-1") // should stay at 0
|
|
|
|
u, _ := s.GetUsage("agent-1")
|
|
assert.Equal(t, 0, u.ActiveJobs)
|
|
}
|
|
|
|
// --- ReturnTokens ---
|
|
|
|
func TestSQLiteStore_ReturnTokens_Good(t *testing.T) {
|
|
s := newTestSQLiteStore(t)
|
|
|
|
_ = s.IncrementUsage("agent-1", 10000, 0)
|
|
err := s.ReturnTokens("agent-1", 5000)
|
|
require.NoError(t, err)
|
|
|
|
u, _ := s.GetUsage("agent-1")
|
|
assert.Equal(t, int64(5000), u.TokensUsed)
|
|
}
|
|
|
|
func TestSQLiteStore_ReturnTokens_Good_FloorAtZero(t *testing.T) {
|
|
s := newTestSQLiteStore(t)
|
|
|
|
_ = s.IncrementUsage("agent-1", 1000, 0)
|
|
_ = s.ReturnTokens("agent-1", 5000) // more than used
|
|
|
|
u, _ := s.GetUsage("agent-1")
|
|
assert.Equal(t, int64(0), u.TokensUsed)
|
|
}
|
|
|
|
func TestSQLiteStore_ReturnTokens_Good_NoRecord(t *testing.T) {
|
|
s := newTestSQLiteStore(t)
|
|
|
|
// Return tokens for agent with no usage record -- should create one at 0
|
|
err := s.ReturnTokens("agent-1", 500)
|
|
require.NoError(t, err)
|
|
|
|
u, _ := s.GetUsage("agent-1")
|
|
assert.Equal(t, int64(0), u.TokensUsed)
|
|
}
|
|
|
|
// --- ResetUsage ---
|
|
|
|
func TestSQLiteStore_ResetUsage_Good(t *testing.T) {
|
|
s := newTestSQLiteStore(t)
|
|
|
|
_ = s.IncrementUsage("agent-1", 50000, 5)
|
|
|
|
err := s.ResetUsage("agent-1")
|
|
require.NoError(t, err)
|
|
|
|
u, _ := s.GetUsage("agent-1")
|
|
assert.Equal(t, int64(0), u.TokensUsed)
|
|
assert.Equal(t, 0, u.JobsStarted)
|
|
assert.Equal(t, 0, u.ActiveJobs)
|
|
}
|
|
|
|
// --- ModelQuota ---
|
|
|
|
func TestSQLiteStore_GetModelQuota_Bad_NotFound(t *testing.T) {
|
|
s := newTestSQLiteStore(t)
|
|
_, err := s.GetModelQuota("nonexistent")
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), "model quota not found")
|
|
}
|
|
|
|
func TestSQLiteStore_SetGetModelQuota_Good(t *testing.T) {
|
|
s := newTestSQLiteStore(t)
|
|
|
|
q := &ModelQuota{
|
|
Model: "claude-opus-4-6",
|
|
DailyTokenBudget: 500000,
|
|
HourlyRateLimit: 100,
|
|
CostCeiling: 10000,
|
|
}
|
|
err := s.SetModelQuota(q)
|
|
require.NoError(t, err)
|
|
|
|
got, err := s.GetModelQuota("claude-opus-4-6")
|
|
require.NoError(t, err)
|
|
assert.Equal(t, q.Model, got.Model)
|
|
assert.Equal(t, q.DailyTokenBudget, got.DailyTokenBudget)
|
|
assert.Equal(t, q.HourlyRateLimit, got.HourlyRateLimit)
|
|
assert.Equal(t, q.CostCeiling, got.CostCeiling)
|
|
}
|
|
|
|
// --- ModelUsage ---
|
|
|
|
func TestSQLiteStore_ModelUsage_Good(t *testing.T) {
|
|
s := newTestSQLiteStore(t)
|
|
|
|
_ = s.IncrementModelUsage("claude-sonnet", 10000)
|
|
_ = s.IncrementModelUsage("claude-sonnet", 5000)
|
|
|
|
usage, err := s.GetModelUsage("claude-sonnet")
|
|
require.NoError(t, err)
|
|
assert.Equal(t, int64(15000), usage)
|
|
}
|
|
|
|
func TestSQLiteStore_GetModelUsage_Good_Default(t *testing.T) {
|
|
s := newTestSQLiteStore(t)
|
|
|
|
usage, err := s.GetModelUsage("unknown-model")
|
|
require.NoError(t, err)
|
|
assert.Equal(t, int64(0), usage)
|
|
}
|
|
|
|
// --- Persistence: close and reopen ---
|
|
|
|
func TestSQLiteStore_Persistence_Good(t *testing.T) {
|
|
dbPath := filepath.Join(t.TempDir(), "persist.db")
|
|
|
|
// Phase 1: write data
|
|
s1, err := NewSQLiteStore(dbPath)
|
|
require.NoError(t, err)
|
|
|
|
_ = s1.SetAllowance(&AgentAllowance{
|
|
AgentID: "agent-1",
|
|
DailyTokenLimit: 100000,
|
|
MaxJobDuration: 15 * time.Minute,
|
|
})
|
|
_ = s1.IncrementUsage("agent-1", 25000, 3)
|
|
_ = s1.SetModelQuota(&ModelQuota{Model: "claude-opus-4-6", DailyTokenBudget: 500000})
|
|
_ = s1.IncrementModelUsage("claude-opus-4-6", 42000)
|
|
require.NoError(t, s1.Close())
|
|
|
|
// Phase 2: reopen and verify
|
|
s2, err := NewSQLiteStore(dbPath)
|
|
require.NoError(t, err)
|
|
defer func() { _ = s2.Close() }()
|
|
|
|
a, err := s2.GetAllowance("agent-1")
|
|
require.NoError(t, err)
|
|
assert.Equal(t, int64(100000), a.DailyTokenLimit)
|
|
assert.Equal(t, 15*time.Minute, a.MaxJobDuration)
|
|
|
|
u, err := s2.GetUsage("agent-1")
|
|
require.NoError(t, err)
|
|
assert.Equal(t, int64(25000), u.TokensUsed)
|
|
assert.Equal(t, 3, u.JobsStarted)
|
|
assert.Equal(t, 3, u.ActiveJobs)
|
|
|
|
q, err := s2.GetModelQuota("claude-opus-4-6")
|
|
require.NoError(t, err)
|
|
assert.Equal(t, int64(500000), q.DailyTokenBudget)
|
|
|
|
mu, err := s2.GetModelUsage("claude-opus-4-6")
|
|
require.NoError(t, err)
|
|
assert.Equal(t, int64(42000), mu)
|
|
}
|
|
|
|
// --- Concurrent access ---
|
|
|
|
func TestSQLiteStore_ConcurrentIncrementUsage_Good(t *testing.T) {
|
|
s := newTestSQLiteStore(t)
|
|
|
|
const goroutines = 10
|
|
const tokensEach = 1000
|
|
|
|
var wg sync.WaitGroup
|
|
wg.Add(goroutines)
|
|
for i := 0; i < goroutines; i++ {
|
|
go func() {
|
|
defer wg.Done()
|
|
err := s.IncrementUsage("agent-1", tokensEach, 1)
|
|
assert.NoError(t, err)
|
|
}()
|
|
}
|
|
wg.Wait()
|
|
|
|
u, err := s.GetUsage("agent-1")
|
|
require.NoError(t, err)
|
|
assert.Equal(t, int64(goroutines*tokensEach), u.TokensUsed)
|
|
assert.Equal(t, goroutines, u.JobsStarted)
|
|
assert.Equal(t, goroutines, u.ActiveJobs)
|
|
}
|
|
|
|
func TestSQLiteStore_ConcurrentModelUsage_Good(t *testing.T) {
|
|
s := newTestSQLiteStore(t)
|
|
|
|
const goroutines = 10
|
|
const tokensEach int64 = 500
|
|
|
|
var wg sync.WaitGroup
|
|
wg.Add(goroutines)
|
|
for i := 0; i < goroutines; i++ {
|
|
go func() {
|
|
defer wg.Done()
|
|
err := s.IncrementModelUsage("claude-opus-4-6", tokensEach)
|
|
assert.NoError(t, err)
|
|
}()
|
|
}
|
|
wg.Wait()
|
|
|
|
usage, err := s.GetModelUsage("claude-opus-4-6")
|
|
require.NoError(t, err)
|
|
assert.Equal(t, goroutines*tokensEach, usage)
|
|
}
|
|
|
|
func TestSQLiteStore_ConcurrentMixed_Good(t *testing.T) {
|
|
s := newTestSQLiteStore(t)
|
|
|
|
_ = s.SetAllowance(&AgentAllowance{
|
|
AgentID: "agent-1",
|
|
DailyTokenLimit: 1000000,
|
|
DailyJobLimit: 100,
|
|
ConcurrentJobs: 50,
|
|
})
|
|
|
|
const goroutines = 10
|
|
var wg sync.WaitGroup
|
|
wg.Add(goroutines * 3)
|
|
|
|
// Increment usage
|
|
for i := 0; i < goroutines; i++ {
|
|
go func() {
|
|
defer wg.Done()
|
|
_ = s.IncrementUsage("agent-1", 100, 1)
|
|
}()
|
|
}
|
|
|
|
// Decrement active jobs
|
|
for i := 0; i < goroutines; i++ {
|
|
go func() {
|
|
defer wg.Done()
|
|
_ = s.DecrementActiveJobs("agent-1")
|
|
}()
|
|
}
|
|
|
|
// Return tokens
|
|
for i := 0; i < goroutines; i++ {
|
|
go func() {
|
|
defer wg.Done()
|
|
_ = s.ReturnTokens("agent-1", 10)
|
|
}()
|
|
}
|
|
|
|
wg.Wait()
|
|
|
|
// Just verify no panics and data is consistent
|
|
u, err := s.GetUsage("agent-1")
|
|
require.NoError(t, err)
|
|
assert.GreaterOrEqual(t, u.TokensUsed, int64(0))
|
|
assert.GreaterOrEqual(t, u.ActiveJobs, 0)
|
|
}
|
|
|
|
// --- AllowanceService integration via SQLiteStore ---
|
|
|
|
func TestSQLiteStore_AllowanceServiceCheck_Good(t *testing.T) {
|
|
s := newTestSQLiteStore(t)
|
|
svc := NewAllowanceService(s)
|
|
|
|
_ = s.SetAllowance(&AgentAllowance{
|
|
AgentID: "agent-1",
|
|
DailyTokenLimit: 100000,
|
|
DailyJobLimit: 10,
|
|
ConcurrentJobs: 2,
|
|
})
|
|
|
|
result, err := svc.Check("agent-1", "")
|
|
require.NoError(t, err)
|
|
assert.True(t, result.Allowed)
|
|
assert.Equal(t, AllowanceOK, result.Status)
|
|
}
|
|
|
|
func TestSQLiteStore_AllowanceServiceRecordUsage_Good(t *testing.T) {
|
|
s := newTestSQLiteStore(t)
|
|
svc := NewAllowanceService(s)
|
|
|
|
_ = s.SetAllowance(&AgentAllowance{
|
|
AgentID: "agent-1",
|
|
DailyTokenLimit: 100000,
|
|
})
|
|
|
|
// Start job
|
|
err := svc.RecordUsage(UsageReport{
|
|
AgentID: "agent-1",
|
|
JobID: "job-1",
|
|
Event: QuotaEventJobStarted,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Complete job
|
|
err = svc.RecordUsage(UsageReport{
|
|
AgentID: "agent-1",
|
|
JobID: "job-1",
|
|
Model: "claude-sonnet",
|
|
TokensIn: 1000,
|
|
TokensOut: 500,
|
|
Event: QuotaEventJobCompleted,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
u, _ := s.GetUsage("agent-1")
|
|
assert.Equal(t, int64(1500), u.TokensUsed)
|
|
assert.Equal(t, 0, u.ActiveJobs)
|
|
}
|
|
|
|
// --- Config-based factory ---
|
|
|
|
func TestNewAllowanceStoreFromConfig_Good_Memory(t *testing.T) {
|
|
cfg := AllowanceConfig{StoreBackend: "memory"}
|
|
s, err := NewAllowanceStoreFromConfig(cfg)
|
|
require.NoError(t, err)
|
|
_, ok := s.(*MemoryStore)
|
|
assert.True(t, ok, "expected MemoryStore")
|
|
}
|
|
|
|
func TestNewAllowanceStoreFromConfig_Good_Default(t *testing.T) {
|
|
cfg := AllowanceConfig{} // empty defaults to memory
|
|
s, err := NewAllowanceStoreFromConfig(cfg)
|
|
require.NoError(t, err)
|
|
_, ok := s.(*MemoryStore)
|
|
assert.True(t, ok, "expected MemoryStore for empty config")
|
|
}
|
|
|
|
func TestNewAllowanceStoreFromConfig_Good_SQLite(t *testing.T) {
|
|
dbPath := filepath.Join(t.TempDir(), "factory.db")
|
|
cfg := AllowanceConfig{
|
|
StoreBackend: "sqlite",
|
|
StorePath: dbPath,
|
|
}
|
|
s, err := NewAllowanceStoreFromConfig(cfg)
|
|
require.NoError(t, err)
|
|
ss, ok := s.(*SQLiteStore)
|
|
assert.True(t, ok, "expected SQLiteStore")
|
|
_ = ss.Close()
|
|
}
|
|
|
|
func TestNewAllowanceStoreFromConfig_Bad_UnknownBackend(t *testing.T) {
|
|
cfg := AllowanceConfig{StoreBackend: "redis"}
|
|
_, err := NewAllowanceStoreFromConfig(cfg)
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), "unsupported store backend")
|
|
}
|
|
|
|
// --- NewSQLiteStore error case ---
|
|
|
|
func TestNewSQLiteStore_Bad_InvalidPath(t *testing.T) {
|
|
_, err := NewSQLiteStore("/nonexistent/deeply/nested/dir/test.db")
|
|
require.Error(t, err)
|
|
}
|