go-agentic/allowance_edge_test.go
Claude 23aa635c91
test: achieve 85.6% coverage with 7 new test files
Phase 1 complete: coverage from 70.1% to 85.6% (130+ tests, all passing).

New test files:
- lifecycle_test.go: full claim -> process -> complete integration flows
- allowance_edge_test.go: boundary conditions for token/job/concurrent limits
- allowance_error_test.go: mock errorStore covering all RecordUsage error paths
- embed_test.go: Prompt() hit/miss and content trimming
- service_test.go: DefaultServiceOptions, TaskPrompt, TaskCommit type coverage
- completion_git_test.go: real git repos for AutoCommit, CreateBranch, CommitAndSync
- context_git_test.go: findRelatedCode with keyword search, file limits, truncation

Updated config_test.go with YAML fallback, env override, and empty-dir paths.

Co-Authored-By: Charon <developers@lethean.io>
2026-02-20 00:59:58 +00:00

662 lines
19 KiB
Go

package agentic
import (
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// --- Allowance exhaustion edge cases ---
func TestAllowanceExhaustion_ExactlyAtTokenLimit(t *testing.T) {
store := NewMemoryStore()
svc := NewAllowanceService(store)
_ = store.SetAllowance(&AgentAllowance{
AgentID: "edge-agent",
DailyTokenLimit: 10000,
})
// Use exactly the limit.
_ = store.IncrementUsage("edge-agent", 10000, 0)
result, err := svc.Check("edge-agent", "")
require.NoError(t, err)
assert.False(t, result.Allowed, "should be denied at exactly the limit")
assert.Equal(t, AllowanceExceeded, result.Status)
assert.Equal(t, int64(0), result.RemainingTokens)
assert.Contains(t, result.Reason, "daily token limit exceeded")
}
func TestAllowanceExhaustion_OneOverTokenLimit(t *testing.T) {
store := NewMemoryStore()
svc := NewAllowanceService(store)
_ = store.SetAllowance(&AgentAllowance{
AgentID: "edge-agent",
DailyTokenLimit: 10000,
})
_ = store.IncrementUsage("edge-agent", 10001, 0)
result, err := svc.Check("edge-agent", "")
require.NoError(t, err)
assert.False(t, result.Allowed)
assert.Equal(t, AllowanceExceeded, result.Status)
assert.True(t, result.RemainingTokens < 0, "remaining should be negative")
}
func TestAllowanceExhaustion_OneUnderTokenLimit(t *testing.T) {
store := NewMemoryStore()
svc := NewAllowanceService(store)
_ = store.SetAllowance(&AgentAllowance{
AgentID: "edge-agent",
DailyTokenLimit: 10000,
})
_ = store.IncrementUsage("edge-agent", 9999, 0)
result, err := svc.Check("edge-agent", "")
require.NoError(t, err)
assert.True(t, result.Allowed, "should be allowed with 1 token remaining")
assert.Equal(t, AllowanceWarning, result.Status, "99.99% usage should be warning")
assert.Equal(t, int64(1), result.RemainingTokens)
}
func TestAllowanceExhaustion_ZeroAllowance(t *testing.T) {
store := NewMemoryStore()
svc := NewAllowanceService(store)
// DailyTokenLimit=0 means unlimited.
_ = store.SetAllowance(&AgentAllowance{
AgentID: "unlimited-agent",
DailyTokenLimit: 0,
DailyJobLimit: 0,
ConcurrentJobs: 0,
})
_ = store.IncrementUsage("unlimited-agent", 999999999, 999)
result, err := svc.Check("unlimited-agent", "")
require.NoError(t, err)
assert.True(t, result.Allowed, "unlimited agent should always be allowed")
assert.Equal(t, AllowanceOK, result.Status)
assert.Equal(t, int64(-1), result.RemainingTokens, "unlimited should show -1")
assert.Equal(t, -1, result.RemainingJobs, "unlimited should show -1")
}
func TestAllowanceExhaustion_ExactlyAtJobLimit(t *testing.T) {
store := NewMemoryStore()
svc := NewAllowanceService(store)
_ = store.SetAllowance(&AgentAllowance{
AgentID: "edge-agent",
DailyJobLimit: 5,
})
_ = store.IncrementUsage("edge-agent", 0, 5)
result, err := svc.Check("edge-agent", "")
require.NoError(t, err)
assert.False(t, result.Allowed, "should be denied at exactly the job limit")
assert.Equal(t, AllowanceExceeded, result.Status)
assert.Equal(t, 0, result.RemainingJobs)
assert.Contains(t, result.Reason, "daily job limit exceeded")
}
func TestAllowanceExhaustion_OneUnderJobLimit(t *testing.T) {
store := NewMemoryStore()
svc := NewAllowanceService(store)
_ = store.SetAllowance(&AgentAllowance{
AgentID: "edge-agent",
DailyJobLimit: 5,
})
_ = store.IncrementUsage("edge-agent", 0, 4)
result, err := svc.Check("edge-agent", "")
require.NoError(t, err)
assert.True(t, result.Allowed, "should be allowed with 1 job remaining")
assert.Equal(t, 1, result.RemainingJobs)
}
func TestAllowanceExhaustion_ConcurrentJobsExactlyAtLimit(t *testing.T) {
store := NewMemoryStore()
svc := NewAllowanceService(store)
_ = store.SetAllowance(&AgentAllowance{
AgentID: "edge-agent",
ConcurrentJobs: 2,
})
// Start 2 concurrent jobs.
_ = store.IncrementUsage("edge-agent", 0, 2)
result, err := svc.Check("edge-agent", "")
require.NoError(t, err)
assert.False(t, result.Allowed, "should be denied at concurrent limit")
assert.Contains(t, result.Reason, "concurrent job limit reached")
}
func TestAllowanceExhaustion_ConcurrentJobsOneUnderLimit(t *testing.T) {
store := NewMemoryStore()
svc := NewAllowanceService(store)
_ = store.SetAllowance(&AgentAllowance{
AgentID: "edge-agent",
ConcurrentJobs: 3,
})
_ = store.IncrementUsage("edge-agent", 0, 2)
result, err := svc.Check("edge-agent", "")
require.NoError(t, err)
assert.True(t, result.Allowed, "should be allowed with 1 concurrent slot remaining")
}
func TestAllowanceExhaustion_ConcurrentJobsFreedByCompletion(t *testing.T) {
store := NewMemoryStore()
svc := NewAllowanceService(store)
_ = store.SetAllowance(&AgentAllowance{
AgentID: "edge-agent",
ConcurrentJobs: 1,
})
// Start a job - fills the slot.
_ = svc.RecordUsage(UsageReport{
AgentID: "edge-agent",
JobID: "job-1",
Event: QuotaEventJobStarted,
})
result, err := svc.Check("edge-agent", "")
require.NoError(t, err)
assert.False(t, result.Allowed, "should be denied, 1 active job")
// Complete the job - frees the slot.
_ = svc.RecordUsage(UsageReport{
AgentID: "edge-agent",
JobID: "job-1",
TokensIn: 100,
TokensOut: 50,
Event: QuotaEventJobCompleted,
})
result, err = svc.Check("edge-agent", "")
require.NoError(t, err)
assert.True(t, result.Allowed, "should be allowed after job completes")
}
func TestAllowanceExhaustion_TokenWarningThreshold(t *testing.T) {
tests := []struct {
name string
limit int64
used int64
expectedStatus AllowanceStatus
expectedAllow bool
}{
{
name: "79% usage is OK",
limit: 10000,
used: 7900,
expectedStatus: AllowanceOK,
expectedAllow: true,
},
{
name: "80% usage is warning",
limit: 10000,
used: 8000,
expectedStatus: AllowanceWarning,
expectedAllow: true,
},
{
name: "90% usage is warning",
limit: 10000,
used: 9000,
expectedStatus: AllowanceWarning,
expectedAllow: true,
},
{
name: "99% usage is warning",
limit: 10000,
used: 9999,
expectedStatus: AllowanceWarning,
expectedAllow: true,
},
{
name: "100% usage is exceeded",
limit: 10000,
used: 10000,
expectedStatus: AllowanceExceeded,
expectedAllow: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
store := NewMemoryStore()
svc := NewAllowanceService(store)
_ = store.SetAllowance(&AgentAllowance{
AgentID: "threshold-agent",
DailyTokenLimit: tt.limit,
})
_ = store.IncrementUsage("threshold-agent", tt.used, 0)
result, err := svc.Check("threshold-agent", "")
require.NoError(t, err)
assert.Equal(t, tt.expectedAllow, result.Allowed)
assert.Equal(t, tt.expectedStatus, result.Status)
})
}
}
func TestAllowanceExhaustion_ResetRestoresCapacity(t *testing.T) {
store := NewMemoryStore()
svc := NewAllowanceService(store)
_ = store.SetAllowance(&AgentAllowance{
AgentID: "reset-agent",
DailyTokenLimit: 10000,
DailyJobLimit: 5,
})
// Exhaust all limits.
_ = store.IncrementUsage("reset-agent", 10000, 5)
result, err := svc.Check("reset-agent", "")
require.NoError(t, err)
assert.False(t, result.Allowed, "should be denied when exhausted")
// Reset the agent (simulates midnight reset).
err = svc.ResetAgent("reset-agent")
require.NoError(t, err)
result, err = svc.Check("reset-agent", "")
require.NoError(t, err)
assert.True(t, result.Allowed, "should be allowed after reset")
assert.Equal(t, int64(10000), result.RemainingTokens)
assert.Equal(t, 5, result.RemainingJobs)
}
func TestAllowanceExhaustion_GlobalModelBudgetExactlyAtLimit(t *testing.T) {
store := NewMemoryStore()
svc := NewAllowanceService(store)
_ = store.SetAllowance(&AgentAllowance{
AgentID: "model-edge-agent",
})
store.SetModelQuota(&ModelQuota{
Model: "claude-opus-4-6",
DailyTokenBudget: 50000,
})
_ = store.IncrementModelUsage("claude-opus-4-6", 50000)
result, err := svc.Check("model-edge-agent", "claude-opus-4-6")
require.NoError(t, err)
assert.False(t, result.Allowed, "should be denied at exact model budget")
assert.Contains(t, result.Reason, "global model token budget exceeded")
}
func TestAllowanceExhaustion_GlobalModelBudgetOneUnder(t *testing.T) {
store := NewMemoryStore()
svc := NewAllowanceService(store)
_ = store.SetAllowance(&AgentAllowance{
AgentID: "model-edge-agent",
})
store.SetModelQuota(&ModelQuota{
Model: "claude-opus-4-6",
DailyTokenBudget: 50000,
})
_ = store.IncrementModelUsage("claude-opus-4-6", 49999)
result, err := svc.Check("model-edge-agent", "claude-opus-4-6")
require.NoError(t, err)
assert.True(t, result.Allowed, "should be allowed with 1 token remaining in model budget")
}
func TestAllowanceExhaustion_FailedJobWithZeroTokens(t *testing.T) {
store := NewMemoryStore()
svc := NewAllowanceService(store)
_ = svc.RecordUsage(UsageReport{
AgentID: "zero-token-agent",
JobID: "job-1",
Event: QuotaEventJobStarted,
})
// Job fails but consumed zero tokens.
err := svc.RecordUsage(UsageReport{
AgentID: "zero-token-agent",
JobID: "job-1",
Model: "claude-sonnet",
TokensIn: 0,
TokensOut: 0,
Event: QuotaEventJobFailed,
})
require.NoError(t, err)
usage, _ := store.GetUsage("zero-token-agent")
assert.Equal(t, int64(0), usage.TokensUsed, "no tokens should be charged")
assert.Equal(t, 0, usage.ActiveJobs)
// Model usage should be zero too (50% of 0 = 0).
modelUsage, _ := store.GetModelUsage("claude-sonnet")
assert.Equal(t, int64(0), modelUsage)
}
func TestAllowanceExhaustion_CancelledJobWithZeroTokens(t *testing.T) {
store := NewMemoryStore()
svc := NewAllowanceService(store)
_ = svc.RecordUsage(UsageReport{
AgentID: "zero-token-agent",
JobID: "job-2",
Event: QuotaEventJobStarted,
})
// Job cancelled with zero tokens.
err := svc.RecordUsage(UsageReport{
AgentID: "zero-token-agent",
JobID: "job-2",
TokensIn: 0,
TokensOut: 0,
Event: QuotaEventJobCancelled,
})
require.NoError(t, err)
usage, _ := store.GetUsage("zero-token-agent")
assert.Equal(t, int64(0), usage.TokensUsed)
assert.Equal(t, 0, usage.ActiveJobs)
}
func TestAllowanceExhaustion_CompletedJobWithNoModel(t *testing.T) {
store := NewMemoryStore()
svc := NewAllowanceService(store)
_ = svc.RecordUsage(UsageReport{
AgentID: "no-model-agent",
JobID: "job-1",
Event: QuotaEventJobStarted,
})
// Complete with empty model -- should skip model-level usage recording.
err := svc.RecordUsage(UsageReport{
AgentID: "no-model-agent",
JobID: "job-1",
Model: "",
TokensIn: 500,
TokensOut: 200,
Event: QuotaEventJobCompleted,
})
require.NoError(t, err)
usage, _ := store.GetUsage("no-model-agent")
assert.Equal(t, int64(700), usage.TokensUsed)
assert.Equal(t, 0, usage.ActiveJobs)
}
func TestAllowanceExhaustion_FailedJobWithNoModel(t *testing.T) {
store := NewMemoryStore()
svc := NewAllowanceService(store)
_ = svc.RecordUsage(UsageReport{
AgentID: "no-model-fail-agent",
JobID: "job-1",
Event: QuotaEventJobStarted,
})
// Fail with empty model.
err := svc.RecordUsage(UsageReport{
AgentID: "no-model-fail-agent",
JobID: "job-1",
Model: "",
TokensIn: 600,
TokensOut: 400,
Event: QuotaEventJobFailed,
})
require.NoError(t, err)
usage, _ := store.GetUsage("no-model-fail-agent")
// 1000 tokens used, 500 returned = 500 net.
assert.Equal(t, int64(500), usage.TokensUsed)
assert.Equal(t, 0, usage.ActiveJobs)
}
func TestAllowanceExhaustion_MultipleChecksWithIncrementalUsage(t *testing.T) {
store := NewMemoryStore()
svc := NewAllowanceService(store)
_ = store.SetAllowance(&AgentAllowance{
AgentID: "incremental-agent",
DailyTokenLimit: 1000,
})
// First check: fresh agent.
result, err := svc.Check("incremental-agent", "")
require.NoError(t, err)
assert.True(t, result.Allowed)
assert.Equal(t, AllowanceOK, result.Status)
assert.Equal(t, int64(1000), result.RemainingTokens)
// Use 500 tokens.
_ = store.IncrementUsage("incremental-agent", 500, 0)
result, err = svc.Check("incremental-agent", "")
require.NoError(t, err)
assert.True(t, result.Allowed)
assert.Equal(t, AllowanceOK, result.Status)
assert.Equal(t, int64(500), result.RemainingTokens)
// Use another 300 tokens (total 800, at 80% threshold).
_ = store.IncrementUsage("incremental-agent", 300, 0)
result, err = svc.Check("incremental-agent", "")
require.NoError(t, err)
assert.True(t, result.Allowed)
assert.Equal(t, AllowanceWarning, result.Status)
assert.Equal(t, int64(200), result.RemainingTokens)
// Use remaining 200 tokens (total 1000, at 100%).
_ = store.IncrementUsage("incremental-agent", 200, 0)
result, err = svc.Check("incremental-agent", "")
require.NoError(t, err)
assert.False(t, result.Allowed)
assert.Equal(t, AllowanceExceeded, result.Status)
assert.Equal(t, int64(0), result.RemainingTokens)
}
// --- MemoryStore additional edge cases ---
func TestMemoryStore_GetUsage_NewAgentReturnsDefaults(t *testing.T) {
store := NewMemoryStore()
usage, err := store.GetUsage("brand-new-agent")
require.NoError(t, err)
assert.Equal(t, "brand-new-agent", usage.AgentID)
assert.Equal(t, int64(0), usage.TokensUsed)
assert.Equal(t, 0, usage.JobsStarted)
assert.Equal(t, 0, usage.ActiveJobs)
assert.Equal(t, startOfDay(time.Now().UTC()), usage.PeriodStart)
}
func TestMemoryStore_ReturnTokens_NonexistentAgent(t *testing.T) {
store := NewMemoryStore()
// ReturnTokens on a nonexistent agent should be a no-op.
err := store.ReturnTokens("ghost-agent", 5000)
require.NoError(t, err)
}
func TestMemoryStore_DecrementActiveJobs_NonexistentAgent(t *testing.T) {
store := NewMemoryStore()
// DecrementActiveJobs on a nonexistent agent should be a no-op.
err := store.DecrementActiveJobs("ghost-agent")
require.NoError(t, err)
}
func TestMemoryStore_GetModelQuota_NotFound(t *testing.T) {
store := NewMemoryStore()
_, err := store.GetModelQuota("nonexistent-model")
require.Error(t, err)
assert.Contains(t, err.Error(), "model quota not found")
}
func TestMemoryStore_GetModelUsage_NewModelReturnsZero(t *testing.T) {
store := NewMemoryStore()
usage, err := store.GetModelUsage("brand-new-model")
require.NoError(t, err)
assert.Equal(t, int64(0), usage)
}
func TestMemoryStore_SetAllowance_Overwrite(t *testing.T) {
store := NewMemoryStore()
_ = store.SetAllowance(&AgentAllowance{
AgentID: "overwrite-agent",
DailyTokenLimit: 5000,
})
_ = store.SetAllowance(&AgentAllowance{
AgentID: "overwrite-agent",
DailyTokenLimit: 9000,
})
a, err := store.GetAllowance("overwrite-agent")
require.NoError(t, err)
assert.Equal(t, int64(9000), a.DailyTokenLimit, "should have overwritten the old allowance")
}
func TestMemoryStore_SetAllowance_IsolatesOriginal(t *testing.T) {
store := NewMemoryStore()
original := &AgentAllowance{
AgentID: "isolated-agent",
DailyTokenLimit: 5000,
}
_ = store.SetAllowance(original)
// Mutate the original.
original.DailyTokenLimit = 99999
got, err := store.GetAllowance("isolated-agent")
require.NoError(t, err)
assert.Equal(t, int64(5000), got.DailyTokenLimit, "store should hold a copy, not the original")
}
func TestMemoryStore_GetAllowance_IsolatesReturn(t *testing.T) {
store := NewMemoryStore()
_ = store.SetAllowance(&AgentAllowance{
AgentID: "isolated-agent",
DailyTokenLimit: 5000,
})
got1, _ := store.GetAllowance("isolated-agent")
got1.DailyTokenLimit = 99999
got2, _ := store.GetAllowance("isolated-agent")
assert.Equal(t, int64(5000), got2.DailyTokenLimit, "returned value should be a copy")
}
func TestMemoryStore_IncrementUsage_MultipleIncrements(t *testing.T) {
store := NewMemoryStore()
_ = store.IncrementUsage("multi-agent", 100, 1)
_ = store.IncrementUsage("multi-agent", 200, 1)
_ = store.IncrementUsage("multi-agent", 300, 0)
usage, err := store.GetUsage("multi-agent")
require.NoError(t, err)
assert.Equal(t, int64(600), usage.TokensUsed)
assert.Equal(t, 2, usage.JobsStarted)
assert.Equal(t, 2, usage.ActiveJobs)
}
func TestMemoryStore_IncrementUsage_ZeroJobsDoesNotIncrementActive(t *testing.T) {
store := NewMemoryStore()
_ = store.IncrementUsage("token-only-agent", 5000, 0)
usage, err := store.GetUsage("token-only-agent")
require.NoError(t, err)
assert.Equal(t, int64(5000), usage.TokensUsed)
assert.Equal(t, 0, usage.JobsStarted)
assert.Equal(t, 0, usage.ActiveJobs, "zero jobs should not increment active count")
}
// --- AllowanceService Check priority ordering ---
func TestAllowanceServiceCheck_ModelAllowlistCheckedFirst(t *testing.T) {
store := NewMemoryStore()
svc := NewAllowanceService(store)
// Agent is over token limit AND using a disallowed model.
_ = store.SetAllowance(&AgentAllowance{
AgentID: "order-agent",
DailyTokenLimit: 1000,
ModelAllowlist: []string{"claude-haiku"},
})
_ = store.IncrementUsage("order-agent", 2000, 0)
result, err := svc.Check("order-agent", "claude-opus-4-6")
require.NoError(t, err)
assert.False(t, result.Allowed)
// Model allowlist is checked first in the code, so it should be the reason.
assert.Contains(t, result.Reason, "model not in allowlist")
}
func TestAllowanceServiceCheck_EmptyModelAllowlistPermitsAll(t *testing.T) {
store := NewMemoryStore()
svc := NewAllowanceService(store)
_ = store.SetAllowance(&AgentAllowance{
AgentID: "any-model-agent",
ModelAllowlist: nil,
})
result, err := svc.Check("any-model-agent", "any-model-at-all")
require.NoError(t, err)
assert.True(t, result.Allowed)
}
// --- QuotaEvent constants ---
func TestQuotaEvent_Values(t *testing.T) {
assert.Equal(t, QuotaEvent("job_started"), QuotaEventJobStarted)
assert.Equal(t, QuotaEvent("job_completed"), QuotaEventJobCompleted)
assert.Equal(t, QuotaEvent("job_failed"), QuotaEventJobFailed)
assert.Equal(t, QuotaEvent("job_cancelled"), QuotaEventJobCancelled)
}
func TestAllowanceExhaustion_FailedJobWithOddTokenCount(t *testing.T) {
store := NewMemoryStore()
svc := NewAllowanceService(store)
_ = svc.RecordUsage(UsageReport{
AgentID: "odd-agent",
JobID: "job-1",
Event: QuotaEventJobStarted,
})
// Odd total: 7 tokens. 50% return = 3 (integer division).
err := svc.RecordUsage(UsageReport{
AgentID: "odd-agent",
JobID: "job-1",
Model: "claude-sonnet",
TokensIn: 4,
TokensOut: 3,
Event: QuotaEventJobFailed,
})
require.NoError(t, err)
usage, _ := store.GetUsage("odd-agent")
// 7 charged - 3 returned = 4 net.
assert.Equal(t, int64(4), usage.TokensUsed)
// Model gets 7 - 3 = 4.
modelUsage, _ := store.GetModelUsage("claude-sonnet")
assert.Equal(t, int64(4), modelUsage)
}