diff --git a/internal/provider/claude_usage_cli.go b/internal/provider/claude_usage_cli.go index 596b6bf6..c72bae0c 100644 --- a/internal/provider/claude_usage_cli.go +++ b/internal/provider/claude_usage_cli.go @@ -11,6 +11,7 @@ import ( "slices" "strconv" "strings" + "sync" "time" // the zone /usage names its resets in: Windows has no zone database, // and without one a reset was read in the machine's own zone @@ -19,19 +20,37 @@ import ( // claudeCLIUsage runs Claude Code's /usage for the account it is signed in // to and is what it printed. The gateway, which runs Claude Code, sets it -// (UsageClaudeVia). -var claudeCLIUsage func(ctx context.Context) (string, error) +// (UsageClaudeVia). It is held for reading while it runs, as +// loginUsageFor is: Allowances reads Claude's usage in the background, and +// tests swap it. +var claudeCLIUsage struct { + sync.RWMutex + f func(ctx context.Context) (string, error) +} -// UsageClaudeVia sets how Claude Code's /usage runs. -func UsageClaudeVia(f func(ctx context.Context) (string, error)) { claudeCLIUsage = f } +// UsageClaudeVia sets how Claude Code's /usage runs. It returns once no +// run is still going through the one it replaces. +func UsageClaudeVia(f func(ctx context.Context) (string, error)) { + claudeCLIUsage.Lock() + claudeCLIUsage.f = f + claudeCLIUsage.Unlock() +} + +// runClaudeUsage is what Claude Code's /usage printed, through +// claudeCLIUsage; errClaudeCannotRun when nothing set it. +func runClaudeUsage(ctx context.Context) (string, error) { + claudeCLIUsage.RLock() + defer claudeCLIUsage.RUnlock() + if claudeCLIUsage.f == nil { + return "", errClaudeCannotRun + } + return claudeCLIUsage.f(ctx) +} // readClaudeUsage is the allowance of the account Claude Code is signed in // to, from its /usage. func readClaudeUsage(ctx context.Context) ([]QuotaWindow, error) { - if claudeCLIUsage == nil { - return []QuotaWindow{}, errClaudeCannotRun - } - text, err := claudeCLIUsage(ctx) + text, err := runClaudeUsage(ctx) if err != nil { return []QuotaWindow{}, err } diff --git a/internal/provider/claude_usage_test.go b/internal/provider/claude_usage_test.go index 6d36b114..156f5892 100644 --- a/internal/provider/claude_usage_test.go +++ b/internal/provider/claude_usage_test.go @@ -18,7 +18,10 @@ import ( func fakeClaudeUsage(t *testing.T, out *atomic.Value, fail *atomic.Bool) *atomic.Int32 { t.Helper() var runs atomic.Int32 - old, oldAsked := claudeCLIUsage, claudeAsked.Load() + claudeCLIUsage.RLock() + old := claudeCLIUsage.f + claudeCLIUsage.RUnlock() + oldAsked := claudeAsked.Load() UsageClaudeVia(func(context.Context) (string, error) { runs.Add(1) if fail != nil && fail.Load() { @@ -38,7 +41,8 @@ func fakeClaudeUsage(t *testing.T, out *atomic.Value, fail *atomic.Bool) *atomic claudeUsage.m = nil claudeUsage.Unlock() t.Cleanup(func() { - claudeCLIUsage, claudeBase, claudeUsageWait, claudeUsedSince = old, oldBase, oldWait, oldUsed + UsageClaudeVia(old) + claudeBase, claudeUsageWait, claudeUsedSince = oldBase, oldWait, oldUsed claudeAsked.Store(oldAsked) srv.Close() claudeUsage.Lock() diff --git a/internal/provider/login_usage.go b/internal/provider/login_usage.go index 31576002..5af5d2e0 100644 --- a/internal/provider/login_usage.go +++ b/internal/provider/login_usage.go @@ -35,12 +35,31 @@ type loginRead struct { } // loginUsageFor, when set, stands in for LoginUsage's readings (tests), -// as UsageClaudeVia does for Claude's own. -var loginUsageFor func(ctx context.Context, agent string) map[string]SubscriptionQuota +// as UsageClaudeVia does for Claude's own. It is held for reading while it +// runs: Allowances reads in the background and returns without waiting, +// so a reading it began can outlast the test that set the stand-in. +var loginUsageFor struct { + sync.RWMutex + f func(ctx context.Context, agent string) map[string]SubscriptionQuota +} -// LoginUsageVia has tests stand in for LoginUsage's readings. +// LoginUsageVia has tests stand in for LoginUsage's readings. It returns +// once no reading is still running through the one it replaces, and none +// starts through it after. func LoginUsageVia(f func(ctx context.Context, agent string) map[string]SubscriptionQuota) { - loginUsageFor = f + loginUsageFor.Lock() + loginUsageFor.f = f + loginUsageFor.Unlock() +} + +// loginUsageStandIn is the stand-in's reading, and false when there is none. +func loginUsageStandIn(ctx context.Context, agent string) (map[string]SubscriptionQuota, bool) { + loginUsageFor.RLock() + defer loginUsageFor.RUnlock() + if loginUsageFor.f == nil { + return nil, false + } + return loginUsageFor.f(ctx, agent), true } // LoginUsage is the allowance used by each of an agent's accounts, by @@ -57,8 +76,8 @@ func LoginUsage(ctx context.Context, agent string) map[string]SubscriptionQuota // counts its own minute from then, not from when it asked (#1295). func loginUsageAt(ctx context.Context, agent string) (map[string]SubscriptionQuota, time.Time) { asked := time.Now() - if loginUsageFor != nil { - return loginUsageFor(ctx, agent), asked + if out, ok := loginUsageStandIn(ctx, agent); ok { + return out, asked } out := map[string]SubscriptionQuota{} if agent == "grok" { diff --git a/internal/provider/login_usage_via_test.go b/internal/provider/login_usage_via_test.go new file mode 100644 index 00000000..b443f6b3 --- /dev/null +++ b/internal/provider/login_usage_via_test.go @@ -0,0 +1,69 @@ +package provider + +import ( + "context" + "testing" + "time" +) + +// swapWaits says whether swap, run while a call through the stand-in it +// replaces is held, waits for that call: Allowances reads in the +// background and returns without waiting, so a reading it began outlasts +// the test that set the stand-in, and the test's swap back raced it (the +// ubuntu race shard of run 37942751214, TestClaudeAccountBackOnceItsWindowRenews). +func swapWaits(t *testing.T, set func(block func()), call func(), swap func()) { + t.Helper() + running, release := make(chan struct{}), make(chan struct{}) + set(func() { + close(running) + <-release + }) + called := make(chan struct{}) + go func() { defer close(called); call() }() + select { + case <-running: + case <-time.After(5 * time.Second): + close(release) + t.Fatal("the stand-in wasn't called") + } + swapped := make(chan struct{}) + go func() { defer close(swapped); swap() }() + select { + case <-swapped: + close(release) + <-called + t.Fatal("swapped while a call was still running through the stand-in it replaced") + case <-time.After(200 * time.Millisecond): + } + close(release) + for _, c := range []chan struct{}{called, swapped} { + select { + case <-c: + case <-time.After(5 * time.Second): + t.Fatal("the swap didn't go on once the call was done") + } + } +} + +func TestLoginUsageViaWaitsForReadingsThroughTheOld(t *testing.T) { + t.Cleanup(func() { LoginUsageVia(nil) }) + swapWaits(t, func(block func()) { + LoginUsageVia(func(context.Context, string) map[string]SubscriptionQuota { + block() + return map[string]SubscriptionQuota{} + }) + }, func() { LoginUsage(context.Background(), "swap-waits") }, func() { LoginUsageVia(nil) }) +} + +func TestUsageClaudeViaWaitsForRunsThroughTheOld(t *testing.T) { + claudeCLIUsage.RLock() + old := claudeCLIUsage.f + claudeCLIUsage.RUnlock() + t.Cleanup(func() { UsageClaudeVia(old) }) + swapWaits(t, func(block func()) { + UsageClaudeVia(func(context.Context) (string, error) { + block() + return "", nil + }) + }, func() { _, _ = readClaudeUsage(context.Background()) }, func() { UsageClaudeVia(nil) }) +}