diff --git a/TODO.md b/TODO.md index eabf1f0..d594763 100644 --- a/TODO.md +++ b/TODO.md @@ -10,10 +10,10 @@ ## Phase 2: Allowance Persistence -- [ ] MemoryStore is in-memory only -- state lost on restart +- [x] MemoryStore is in-memory only -- state lost on restart - [ ] Add Redis backend for `AllowanceStore` interface (multi-process safe) -- [ ] Add SQLite backend for `AllowanceStore` interface (single-node fallback) -- [ ] Config already supports YAML -- wire backend selection into config loader +- [x] Add SQLite backend for `AllowanceStore` interface (single-node fallback) +- [x] Config already supports YAML -- wire backend selection into config loader ## Phase 3: Multi-Agent Coordination diff --git a/allowance_sqlite.go b/allowance_sqlite.go new file mode 100644 index 0000000..a6eee3d --- /dev/null +++ b/allowance_sqlite.go @@ -0,0 +1,296 @@ +package agentic + +import ( + "encoding/json" + "errors" + "sync" + "time" + + "forge.lthn.ai/core/go-store" +) + +// SQLite group names for namespacing data in the KV store. +const ( + groupAllowances = "allowances" + groupUsage = "usage" + groupModelQuota = "model_quotas" + groupModelUsage = "model_usage" +) + +// SQLiteStore implements AllowanceStore using go-store (SQLite KV). +// It provides persistent storage that survives process restarts. +type SQLiteStore struct { + db *store.Store + mu sync.Mutex // serialises read-modify-write operations +} + +// NewSQLiteStore creates a new SQLite-backed allowance store at the given path. +// Use ":memory:" for tests that do not need persistence. +func NewSQLiteStore(dbPath string) (*SQLiteStore, error) { + db, err := store.New(dbPath) + if err != nil { + return nil, &APIError{Code: 500, Message: "failed to open SQLite store: " + err.Error()} + } + return &SQLiteStore{db: db}, nil +} + +// Close releases the underlying SQLite database. +func (s *SQLiteStore) Close() error { + return s.db.Close() +} + +// GetAllowance returns the quota limits for an agent. +func (s *SQLiteStore) GetAllowance(agentID string) (*AgentAllowance, error) { + val, err := s.db.Get(groupAllowances, agentID) + if err != nil { + if errors.Is(err, store.ErrNotFound) { + return nil, &APIError{Code: 404, Message: "allowance not found for agent: " + agentID} + } + return nil, &APIError{Code: 500, Message: "failed to get allowance: " + err.Error()} + } + var a allowanceJSON + if err := json.Unmarshal([]byte(val), &a); err != nil { + return nil, &APIError{Code: 500, Message: "failed to unmarshal allowance: " + err.Error()} + } + return a.toAgentAllowance(), nil +} + +// SetAllowance persists quota limits for an agent. +func (s *SQLiteStore) SetAllowance(a *AgentAllowance) error { + aj := newAllowanceJSON(a) + data, err := json.Marshal(aj) + if err != nil { + return &APIError{Code: 500, Message: "failed to marshal allowance: " + err.Error()} + } + if err := s.db.Set(groupAllowances, a.AgentID, string(data)); err != nil { + return &APIError{Code: 500, Message: "failed to set allowance: " + err.Error()} + } + return nil +} + +// GetUsage returns the current usage record for an agent. +func (s *SQLiteStore) GetUsage(agentID string) (*UsageRecord, error) { + val, err := s.db.Get(groupUsage, agentID) + if err != nil { + if errors.Is(err, store.ErrNotFound) { + return &UsageRecord{ + AgentID: agentID, + PeriodStart: startOfDay(time.Now().UTC()), + }, nil + } + return nil, &APIError{Code: 500, Message: "failed to get usage: " + err.Error()} + } + var u UsageRecord + if err := json.Unmarshal([]byte(val), &u); err != nil { + return nil, &APIError{Code: 500, Message: "failed to unmarshal usage: " + err.Error()} + } + return &u, nil +} + +// IncrementUsage atomically adds to an agent's usage counters. +func (s *SQLiteStore) IncrementUsage(agentID string, tokens int64, jobs int) error { + s.mu.Lock() + defer s.mu.Unlock() + + u, err := s.getUsageLocked(agentID) + if err != nil { + return err + } + u.TokensUsed += tokens + u.JobsStarted += jobs + if jobs > 0 { + u.ActiveJobs += jobs + } + return s.putUsageLocked(u) +} + +// DecrementActiveJobs reduces the active job count by 1. +func (s *SQLiteStore) DecrementActiveJobs(agentID string) error { + s.mu.Lock() + defer s.mu.Unlock() + + u, err := s.getUsageLocked(agentID) + if err != nil { + return err + } + if u.ActiveJobs > 0 { + u.ActiveJobs-- + } + return s.putUsageLocked(u) +} + +// ReturnTokens adds tokens back to the agent's remaining quota. +func (s *SQLiteStore) ReturnTokens(agentID string, tokens int64) error { + s.mu.Lock() + defer s.mu.Unlock() + + u, err := s.getUsageLocked(agentID) + if err != nil { + return err + } + u.TokensUsed -= tokens + if u.TokensUsed < 0 { + u.TokensUsed = 0 + } + return s.putUsageLocked(u) +} + +// ResetUsage clears usage counters for an agent. +func (s *SQLiteStore) ResetUsage(agentID string) error { + s.mu.Lock() + defer s.mu.Unlock() + + u := &UsageRecord{ + AgentID: agentID, + PeriodStart: startOfDay(time.Now().UTC()), + } + return s.putUsageLocked(u) +} + +// GetModelQuota returns global limits for a model. +func (s *SQLiteStore) GetModelQuota(model string) (*ModelQuota, error) { + val, err := s.db.Get(groupModelQuota, model) + if err != nil { + if errors.Is(err, store.ErrNotFound) { + return nil, &APIError{Code: 404, Message: "model quota not found: " + model} + } + return nil, &APIError{Code: 500, Message: "failed to get model quota: " + err.Error()} + } + var q ModelQuota + if err := json.Unmarshal([]byte(val), &q); err != nil { + return nil, &APIError{Code: 500, Message: "failed to unmarshal model quota: " + err.Error()} + } + return &q, nil +} + +// GetModelUsage returns current token usage for a model. +func (s *SQLiteStore) GetModelUsage(model string) (int64, error) { + val, err := s.db.Get(groupModelUsage, model) + if err != nil { + if errors.Is(err, store.ErrNotFound) { + return 0, nil + } + return 0, &APIError{Code: 500, Message: "failed to get model usage: " + err.Error()} + } + var tokens int64 + if err := json.Unmarshal([]byte(val), &tokens); err != nil { + return 0, &APIError{Code: 500, Message: "failed to unmarshal model usage: " + err.Error()} + } + return tokens, nil +} + +// IncrementModelUsage atomically adds to a model's usage counter. +func (s *SQLiteStore) IncrementModelUsage(model string, tokens int64) error { + s.mu.Lock() + defer s.mu.Unlock() + + current, err := s.getModelUsageLocked(model) + if err != nil { + return err + } + current += tokens + data, err := json.Marshal(current) + if err != nil { + return &APIError{Code: 500, Message: "failed to marshal model usage: " + err.Error()} + } + if err := s.db.Set(groupModelUsage, model, string(data)); err != nil { + return &APIError{Code: 500, Message: "failed to set model usage: " + err.Error()} + } + return nil +} + +// SetModelQuota persists global limits for a model. +func (s *SQLiteStore) SetModelQuota(q *ModelQuota) error { + data, err := json.Marshal(q) + if err != nil { + return &APIError{Code: 500, Message: "failed to marshal model quota: " + err.Error()} + } + if err := s.db.Set(groupModelQuota, q.Model, string(data)); err != nil { + return &APIError{Code: 500, Message: "failed to set model quota: " + err.Error()} + } + return nil +} + +// --- internal helpers (must be called with mu held) --- + +// getUsageLocked reads a UsageRecord from the store. Caller must hold s.mu. +func (s *SQLiteStore) getUsageLocked(agentID string) (*UsageRecord, error) { + val, err := s.db.Get(groupUsage, agentID) + if err != nil { + if errors.Is(err, store.ErrNotFound) { + return &UsageRecord{ + AgentID: agentID, + PeriodStart: startOfDay(time.Now().UTC()), + }, nil + } + return nil, &APIError{Code: 500, Message: "failed to get usage: " + err.Error()} + } + var u UsageRecord + if err := json.Unmarshal([]byte(val), &u); err != nil { + return nil, &APIError{Code: 500, Message: "failed to unmarshal usage: " + err.Error()} + } + return &u, nil +} + +// putUsageLocked writes a UsageRecord to the store. Caller must hold s.mu. +func (s *SQLiteStore) putUsageLocked(u *UsageRecord) error { + data, err := json.Marshal(u) + if err != nil { + return &APIError{Code: 500, Message: "failed to marshal usage: " + err.Error()} + } + if err := s.db.Set(groupUsage, u.AgentID, string(data)); err != nil { + return &APIError{Code: 500, Message: "failed to set usage: " + err.Error()} + } + return nil +} + +// getModelUsageLocked reads model usage from the store. Caller must hold s.mu. +func (s *SQLiteStore) getModelUsageLocked(model string) (int64, error) { + val, err := s.db.Get(groupModelUsage, model) + if err != nil { + if errors.Is(err, store.ErrNotFound) { + return 0, nil + } + return 0, &APIError{Code: 500, Message: "failed to get model usage: " + err.Error()} + } + var tokens int64 + if err := json.Unmarshal([]byte(val), &tokens); err != nil { + return 0, &APIError{Code: 500, Message: "failed to unmarshal model usage: " + err.Error()} + } + return tokens, nil +} + +// --- JSON serialisation helper for AgentAllowance --- +// time.Duration does not have a stable JSON representation. We serialise it +// as an int64 (nanoseconds) to avoid locale-dependent string parsing. + +type allowanceJSON struct { + AgentID string `json:"agent_id"` + DailyTokenLimit int64 `json:"daily_token_limit"` + DailyJobLimit int `json:"daily_job_limit"` + ConcurrentJobs int `json:"concurrent_jobs"` + MaxJobDurationNs int64 `json:"max_job_duration_ns"` + ModelAllowlist []string `json:"model_allowlist,omitempty"` +} + +func newAllowanceJSON(a *AgentAllowance) *allowanceJSON { + return &allowanceJSON{ + AgentID: a.AgentID, + DailyTokenLimit: a.DailyTokenLimit, + DailyJobLimit: a.DailyJobLimit, + ConcurrentJobs: a.ConcurrentJobs, + MaxJobDurationNs: int64(a.MaxJobDuration), + ModelAllowlist: a.ModelAllowlist, + } +} + +func (aj *allowanceJSON) toAgentAllowance() *AgentAllowance { + return &AgentAllowance{ + AgentID: aj.AgentID, + DailyTokenLimit: aj.DailyTokenLimit, + DailyJobLimit: aj.DailyJobLimit, + ConcurrentJobs: aj.ConcurrentJobs, + MaxJobDuration: time.Duration(aj.MaxJobDurationNs), + ModelAllowlist: aj.ModelAllowlist, + } +} diff --git a/allowance_sqlite_test.go b/allowance_sqlite_test.go new file mode 100644 index 0000000..dfdb16c --- /dev/null +++ b/allowance_sqlite_test.go @@ -0,0 +1,465 @@ +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) +} diff --git a/config.go b/config.go index 1907534..7c785dc 100644 --- a/config.go +++ b/config.go @@ -195,3 +195,45 @@ func ConfigPath() (string, error) { } return filepath.Join(homeDir, ".core", configFileName), nil } + +// AllowanceConfig controls allowance store backend selection. +type AllowanceConfig struct { + // StoreBackend is the storage backend: "memory" or "sqlite". Default: "memory". + StoreBackend string `yaml:"store_backend" json:"store_backend"` + // StorePath is the file path for the SQLite database. + // Default: ~/.config/agentic/allowance.db (only used when StoreBackend is "sqlite"). + StorePath string `yaml:"store_path" json:"store_path"` +} + +// DefaultAllowanceStorePath returns the default SQLite path for allowance data. +func DefaultAllowanceStorePath() (string, error) { + homeDir, err := os.UserHomeDir() + if err != nil { + return "", errors.E("agentic.DefaultAllowanceStorePath", "failed to get home directory", err) + } + return filepath.Join(homeDir, ".config", "agentic", "allowance.db"), nil +} + +// NewAllowanceStoreFromConfig creates an AllowanceStore based on the given config. +// It returns a MemoryStore for "memory" (or empty) backend and a SQLiteStore for "sqlite". +func NewAllowanceStoreFromConfig(cfg AllowanceConfig) (AllowanceStore, error) { + switch cfg.StoreBackend { + case "", "memory": + return NewMemoryStore(), nil + case "sqlite": + dbPath := cfg.StorePath + if dbPath == "" { + var err error + dbPath, err = DefaultAllowanceStorePath() + if err != nil { + return nil, err + } + } + return NewSQLiteStore(dbPath) + default: + return nil, &APIError{ + Code: 400, + Message: "unsupported store backend: " + cfg.StoreBackend, + } + } +} diff --git a/go.mod b/go.mod index c4a4f82..f1b0385 100644 --- a/go.mod +++ b/go.mod @@ -4,13 +4,28 @@ go 1.25.5 require ( forge.lthn.ai/core/go v0.0.0 + forge.lthn.ai/core/go-store v0.0.0-00010101000000-000000000000 github.com/stretchr/testify v1.11.1 gopkg.in/yaml.v3 v3.0.1 ) require ( github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect + github.com/dustin/go-humanize v1.0.1 // indirect + github.com/google/uuid v1.6.0 // indirect + github.com/kr/pretty v0.3.1 // indirect + github.com/mattn/go-isatty v0.0.20 // indirect + github.com/ncruces/go-strftime v1.0.0 // indirect github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect + github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect + golang.org/x/exp v0.0.0-20260212183809-81e46e3db34a // indirect + golang.org/x/sys v0.41.0 // indirect + modernc.org/libc v1.67.7 // indirect + modernc.org/mathutil v1.7.1 // indirect + modernc.org/memory v1.11.0 // indirect + modernc.org/sqlite v1.46.1 // indirect ) replace forge.lthn.ai/core/go => ../go + +replace forge.lthn.ai/core/go-store => ../go-store diff --git a/go.sum b/go.sum index e575324..c422646 100644 --- a/go.sum +++ b/go.sum @@ -1,10 +1,73 @@ +github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= +github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= +github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e h1:ijClszYn+mADRFY17kjQEVQ1XRhq2/JR1M3sGqeJoxs= +github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e/go.mod h1:boTsfXsheKC2y+lKOCMpSfarhxDeIzfZG1jqGcPl3cA= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k= +github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM= +github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= +github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= +github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= +github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= +github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= +github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= +github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w= +github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls= +github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA= github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U= github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE= +github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo= +github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs= +github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ= +github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc= github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= -gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= +golang.org/x/exp v0.0.0-20260212183809-81e46e3db34a h1:ovFr6Z0MNmU7nH8VaX5xqw+05ST2uO1exVfZPVqRC5o= +golang.org/x/exp v0.0.0-20260212183809-81e46e3db34a/go.mod h1:K79w1Vqn7PoiZn+TkNpx3BUWUQksGO3JcVX6qIjytmA= +golang.org/x/mod v0.33.0 h1:tHFzIWbBifEmbwtGz65eaWyGiGZatSrT9prnU8DbVL8= +golang.org/x/mod v0.33.0/go.mod h1:swjeQEj+6r7fODbD2cqrnje9PnziFuw4bmLbBZFrQ5w= +golang.org/x/sync v0.19.0 h1:vV+1eWNmZ5geRlYjzm2adRgW2/mcpevXNg50YZtPCE4= +golang.org/x/sync v0.19.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI= +golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.41.0 h1:Ivj+2Cp/ylzLiEU89QhWblYnOE9zerudt9Ftecq2C6k= +golang.org/x/sys v0.41.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks= +golang.org/x/tools v0.42.0 h1:uNgphsn75Tdz5Ji2q36v/nsFSfR/9BRFvqhGBaJGd5k= +golang.org/x/tools v0.42.0/go.mod h1:Ma6lCIwGZvHK6XtgbswSoWroEkhugApmsXyrUmBhfr0= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +modernc.org/cc/v4 v4.27.1 h1:9W30zRlYrefrDV2JE2O8VDtJ1yPGownxciz5rrbQZis= +modernc.org/cc/v4 v4.27.1/go.mod h1:uVtb5OGqUKpoLWhqwNQo/8LwvoiEBLvZXIQ/SmO6mL0= +modernc.org/ccgo/v4 v4.30.1 h1:4r4U1J6Fhj98NKfSjnPUN7Ze2c6MnAdL0hWw6+LrJpc= +modernc.org/ccgo/v4 v4.30.1/go.mod h1:bIOeI1JL54Utlxn+LwrFyjCx2n2RDiYEaJVSrgdrRfM= +modernc.org/fileutil v1.3.40 h1:ZGMswMNc9JOCrcrakF1HrvmergNLAmxOPjizirpfqBA= +modernc.org/fileutil v1.3.40/go.mod h1:HxmghZSZVAz/LXcMNwZPA/DRrQZEVP9VX0V4LQGQFOc= +modernc.org/gc/v2 v2.6.5 h1:nyqdV8q46KvTpZlsw66kWqwXRHdjIlJOhG6kxiV/9xI= +modernc.org/gc/v2 v2.6.5/go.mod h1:YgIahr1ypgfe7chRuJi2gD7DBQiKSLMPgBQe9oIiito= +modernc.org/gc/v3 v3.1.1 h1:k8T3gkXWY9sEiytKhcgyiZ2L0DTyCQ/nvX+LoCljoRE= +modernc.org/gc/v3 v3.1.1/go.mod h1:HFK/6AGESC7Ex+EZJhJ2Gni6cTaYpSMmU/cT9RmlfYY= +modernc.org/goabi0 v0.2.0 h1:HvEowk7LxcPd0eq6mVOAEMai46V+i7Jrj13t4AzuNks= +modernc.org/goabi0 v0.2.0/go.mod h1:CEFRnnJhKvWT1c1JTI3Avm+tgOWbkOu5oPA8eH8LnMI= +modernc.org/libc v1.67.7 h1:H+gYQw2PyidyxwxQsGTwQw6+6H+xUk+plvOKW7+d3TI= +modernc.org/libc v1.67.7/go.mod h1:UjCSJFl2sYbJbReVQeVpq/MgzlbmDM4cRHIYFelnaDk= +modernc.org/mathutil v1.7.1 h1:GCZVGXdaN8gTqB1Mf/usp1Y/hSqgI2vAGGP4jZMCxOU= +modernc.org/mathutil v1.7.1/go.mod h1:4p5IwJITfppl0G4sUEDtCr4DthTaT47/N3aT6MhfgJg= +modernc.org/memory v1.11.0 h1:o4QC8aMQzmcwCK3t3Ux/ZHmwFPzE6hf2Y5LbkRs+hbI= +modernc.org/memory v1.11.0/go.mod h1:/JP4VbVC+K5sU2wZi9bHoq2MAkCnrt2r98UGeSK7Mjw= +modernc.org/opt v0.1.4 h1:2kNGMRiUjrp4LcaPuLY2PzUfqM/w9N23quVwhKt5Qm8= +modernc.org/opt v0.1.4/go.mod h1:03fq9lsNfvkYSfxrfUhZCWPk1lm4cq4N+Bh//bEtgns= +modernc.org/sortutil v1.2.1 h1:+xyoGf15mM3NMlPDnFqrteY07klSFxLElE2PVuWIJ7w= +modernc.org/sortutil v1.2.1/go.mod h1:7ZI3a3REbai7gzCLcotuw9AC4VZVpYMjDzETGsSMqJE= +modernc.org/sqlite v1.46.1 h1:eFJ2ShBLIEnUWlLy12raN0Z1plqmFX9Qe3rjQTKt6sU= +modernc.org/sqlite v1.46.1/go.mod h1:CzbrU2lSB1DKUusvwGz7rqEKIq+NUd8GWuBBZDs9/nA= +modernc.org/strutil v1.2.1 h1:UneZBkQA+DX2Rp35KcM69cSsNES9ly8mQWD71HKlOA0= +modernc.org/strutil v1.2.1/go.mod h1:EHkiggD70koQxjVdSBM3JKM7k6L0FbGE5eymy9i3B9A= +modernc.org/token v1.1.0 h1:Xl7Ap9dKaEs5kLoOQeQmPWevfnk/DM5qcLcYlA8ys6Y= +modernc.org/token v1.1.0/go.mod h1:UGzOrNV1mAFSEB63lOFHIpNRUVMvYTc6yu1SMY/XTDM=