diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 0000000..f71254f --- /dev/null +++ b/.gitmodules @@ -0,0 +1,4 @@ +[submodule "external/go"] + path = external/go + url = https://github.com/dappcore/go.git + branch = dev diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..11fc604 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,11 @@ + + +# go-ratelimit Agent Notes + +This repository follows the core/go v0.9.0 consumer layout. + +- The Go module lives under `go/`. +- Use core/go primitives from `dappco.re/go` instead of direct banned stdlib imports where wrappers exist. +- Keep public-symbol tests as `Test__{Good,Bad,Ugly}` in the matching `_test.go`. +- Keep examples in the matching `_example_test.go`. +- Do not edit `.core/`, `external/`, or `/Users/snider/Code/core/go`. diff --git a/ax7_public_test.go b/ax7_public_test.go deleted file mode 100644 index 5cfbb09..0000000 --- a/ax7_public_test.go +++ /dev/null @@ -1,704 +0,0 @@ -// SPDX-License-Identifier: EUPL-1.2 - -package ratelimit_test - -import ( - "context" - "io" - "net/http" - "os" - "time" - - . "dappco.re/go" - ratelimit "dappco.re/go/ratelimit" -) - -type ax7RoundTrip func(*http.Request) (*http.Response, error) - -func (f ax7RoundTrip) RoundTrip(req *http.Request) (*http.Response, error) { - return f(req) -} - -func ax7YAMLPath(t *T) string { - t.Helper() - return Path(t.TempDir(), "ratelimits.yaml") -} - -func ax7Limiter(t *T) *ratelimit.RateLimiter { - t.Helper() - rl, err := ratelimit.NewWithConfig(ratelimit.Config{ - FilePath: ax7YAMLPath(t), - Providers: []ratelimit.Provider{ratelimit.ProviderLocal}, - }) - RequireNoError(t, err) - return rl -} - -func TestAX7_New_Good(t *T) { - rl, err := ratelimit.New() - - RequireNoError(t, err) - AssertTrue(t, rl.CanSend("gemini-3-pro-preview", 1)) -} - -func TestAX7_New_Bad(t *T) { - t.Setenv("CORE_HOME", "") - t.Setenv("HOME", "") - t.Setenv("home", "") - t.Setenv("USERPROFILE", "") - - rl, err := ratelimit.New() - AssertError(t, err) - AssertNil(t, rl) -} - -func TestAX7_New_Ugly(t *T) { - rl1, err := ratelimit.New() - RequireNoError(t, err) - rl2, err := ratelimit.New() - RequireNoError(t, err) - - rl1.RecordUsage("gemini-3-pro-preview", 1, 1) - AssertEqual(t, 1, rl1.Stats("gemini-3-pro-preview").RPD) - AssertEqual(t, 0, rl2.Stats("gemini-3-pro-preview").RPD) -} - -func TestAX7_NewWithConfig_Good(t *T) { - rl, err := ratelimit.NewWithConfig(ratelimit.Config{ - FilePath: ax7YAMLPath(t), - Providers: []ratelimit.Provider{ratelimit.ProviderOpenAI}, - Quotas: map[string]ratelimit.ModelQuota{ - "custom": {MaxRPM: 7, MaxTPM: 70, MaxRPD: 700}, - }, - }) - - RequireNoError(t, err) - AssertEqual(t, 500, rl.Stats("gpt-4o").MaxRPM) - AssertEqual(t, 7, rl.Stats("custom").MaxRPM) -} - -func TestAX7_NewWithConfig_Bad(t *T) { - rl, err := ratelimit.NewWithConfig(ratelimit.Config{ - FilePath: ax7YAMLPath(t), - Backend: "bogus", - }) - - AssertError(t, err) - AssertNil(t, rl) -} - -func TestAX7_NewWithConfig_Ugly(t *T) { - rl, err := ratelimit.NewWithConfig(ratelimit.Config{FilePath: ax7YAMLPath(t)}) - - RequireNoError(t, err) - AssertEqual(t, 150, rl.Stats("gemini-3-pro-preview").MaxRPM) - AssertEqual(t, 0, rl.Stats("missing-model").MaxRPM) -} - -func TestAX7_RateLimiter_SetQuota_Good(t *T) { - rl := ax7Limiter(t) - rl.SetQuota("model-a", ratelimit.ModelQuota{MaxRPM: 3, MaxTPM: 30, MaxRPD: 300}) - - stats := rl.Stats("model-a") - AssertEqual(t, 3, stats.MaxRPM) - AssertEqual(t, 30, stats.MaxTPM) - AssertEqual(t, 300, stats.MaxRPD) -} - -func TestAX7_RateLimiter_SetQuota_Bad(t *T) { - rl := ax7Limiter(t) - rl.SetQuota("model-a", ratelimit.ModelQuota{MaxRPM: 1}) - rl.SetQuota("model-a", ratelimit.ModelQuota{}) - - decision := rl.Decide("model-a", 999999) - AssertTrue(t, decision.Allowed) - AssertEqual(t, ratelimit.DecisionUnlimited, decision.Code) -} - -func TestAX7_RateLimiter_SetQuota_Ugly(t *T) { - rl := ax7Limiter(t) - rl.SetQuota("", ratelimit.ModelQuota{MaxRPM: 1}) - - AssertEqual(t, 1, rl.Stats("").MaxRPM) - AssertFalse(t, rl.CanSend("", -1)) -} - -func TestAX7_RateLimiter_AddProvider_Good(t *T) { - rl := ax7Limiter(t) - rl.AddProvider(ratelimit.ProviderAnthropic) - - AssertEqual(t, 50, rl.Stats("claude-opus-4").MaxRPM) - AssertEqual(t, 40000, rl.Stats("claude-opus-4").MaxTPM) -} - -func TestAX7_RateLimiter_AddProvider_Bad(t *T) { - rl := ax7Limiter(t) - rl.AddProvider(ratelimit.Provider("unknown")) - - models := make([]string, 0) - for model := range rl.Models() { - models = append(models, model) - } - AssertEmpty(t, models) -} - -func TestAX7_RateLimiter_AddProvider_Ugly(t *T) { - rl := ax7Limiter(t) - rl.SetQuota("local-model", ratelimit.ModelQuota{MaxRPM: 2}) - rl.AddProvider(ratelimit.ProviderLocal) - - AssertEqual(t, 2, rl.Stats("local-model").MaxRPM) - AssertTrue(t, rl.CanSend("unknown-local", 1)) -} - -func TestAX7_RateLimiter_Load_Good(t *T) { - path := ax7YAMLPath(t) - rl, err := ratelimit.NewWithConfig(ratelimit.Config{ - FilePath: path, - Quotas: map[string]ratelimit.ModelQuota{"model-a": {MaxRPM: 5}}, - }) - RequireNoError(t, err) - rl.RecordUsage("model-a", 1, 2) - RequireNoError(t, rl.Persist()) - - loaded, err := ratelimit.NewWithConfig(ratelimit.Config{FilePath: path, Providers: []ratelimit.Provider{ratelimit.ProviderLocal}}) - RequireNoError(t, err) - AssertNoError(t, loaded.Load()) - AssertEqual(t, 1, loaded.Stats("model-a").RPD) -} - -func TestAX7_RateLimiter_Load_Bad(t *T) { - path := ax7YAMLPath(t) - RequireNoError(t, os.WriteFile(path, []byte("{{{not yaml"), 0o600)) - rl, err := ratelimit.NewWithConfig(ratelimit.Config{FilePath: path, Providers: []ratelimit.Provider{ratelimit.ProviderLocal}}) - RequireNoError(t, err) - - err = rl.Load() - AssertError(t, err) - AssertContains(t, err.Error(), "yaml") -} - -func TestAX7_RateLimiter_Load_Ugly(t *T) { - rl, err := ratelimit.NewWithConfig(ratelimit.Config{ - FilePath: Path(t.TempDir(), "missing.yaml"), - Providers: []ratelimit.Provider{ratelimit.ProviderLocal}, - }) - RequireNoError(t, err) - - AssertNoError(t, rl.Load()) - AssertEmpty(t, rl.AllStats()) -} - -func TestAX7_RateLimiter_Persist_Good(t *T) { - path := ax7YAMLPath(t) - rl, err := ratelimit.NewWithConfig(ratelimit.Config{FilePath: path, Quotas: map[string]ratelimit.ModelQuota{"model-a": {MaxRPM: 2}}}) - RequireNoError(t, err) - rl.RecordUsage("model-a", 1, 1) - - AssertNoError(t, rl.Persist()) - _, err = os.Stat(path) - AssertNoError(t, err) -} - -func TestAX7_RateLimiter_Persist_Bad(t *T) { - rl, err := ratelimit.NewWithConfig(ratelimit.Config{ - FilePath: t.TempDir(), - Providers: []ratelimit.Provider{ratelimit.ProviderLocal}, - }) - RequireNoError(t, err) - rl.RecordUsage("model-a", 1, 1) - - err = rl.Persist() - AssertError(t, err) -} - -func TestAX7_RateLimiter_Persist_Ugly(t *T) { - path := ax7YAMLPath(t) - rl, err := ratelimit.NewWithConfig(ratelimit.Config{FilePath: path, Providers: []ratelimit.Provider{ratelimit.ProviderLocal}}) - RequireNoError(t, err) - - AssertNoError(t, rl.Persist()) - loaded, err := ratelimit.NewWithConfig(ratelimit.Config{FilePath: path, Providers: []ratelimit.Provider{ratelimit.ProviderLocal}}) - RequireNoError(t, err) - AssertNoError(t, loaded.Load()) -} - -func TestAX7_RateLimiter_BackgroundPrune_Good(t *T) { - rl := ax7Limiter(t) - stop := rl.BackgroundPrune(10 * time.Millisecond) - time.Sleep(20 * time.Millisecond) - - AssertNotPanics(t, stop) - AssertEmpty(t, rl.AllStats()) -} - -func TestAX7_RateLimiter_BackgroundPrune_Bad(t *T) { - rl := ax7Limiter(t) - stop := rl.BackgroundPrune(0) - rl.RecordUsage("model-a", 1, 1) - - AssertNotPanics(t, stop) - AssertEqual(t, 1, rl.Stats("model-a").RPD) -} - -func TestAX7_RateLimiter_BackgroundPrune_Ugly(t *T) { - rl := ax7Limiter(t) - stop := rl.BackgroundPrune(1 * time.Millisecond) - - AssertNotPanics(t, stop) - AssertNotPanics(t, stop) -} - -func TestAX7_RateLimiter_CanSend_Good(t *T) { - rl := ax7Limiter(t) - rl.SetQuota("model-a", ratelimit.ModelQuota{MaxRPM: 2, MaxTPM: 100, MaxRPD: 5}) - - AssertTrue(t, rl.CanSend("model-a", 10)) - AssertEqual(t, ratelimit.DecisionAllowed, rl.Decide("model-a", 10).Code) -} - -func TestAX7_RateLimiter_CanSend_Bad(t *T) { - rl := ax7Limiter(t) - rl.SetQuota("model-a", ratelimit.ModelQuota{MaxRPM: 1, MaxTPM: 100, MaxRPD: 5}) - rl.RecordUsage("model-a", 1, 1) - - AssertFalse(t, rl.CanSend("model-a", 1)) - AssertEqual(t, ratelimit.DecisionRPMLimit, rl.Decide("model-a", 1).Code) -} - -func TestAX7_RateLimiter_CanSend_Ugly(t *T) { - rl := ax7Limiter(t) - rl.SetQuota("model-a", ratelimit.ModelQuota{MaxRPM: 1, MaxTPM: 100, MaxRPD: 5}) - - AssertFalse(t, rl.CanSend("model-a", -1)) - AssertTrue(t, rl.CanSend("unknown-model", 999999)) -} - -func TestAX7_RateLimiter_RecordUsage_Good(t *T) { - rl := ax7Limiter(t) - rl.RecordUsage("model-a", 10, 15) - - stats := rl.Stats("model-a") - AssertEqual(t, 1, stats.RPD) - AssertEqual(t, 25, stats.TPM) -} - -func TestAX7_RateLimiter_RecordUsage_Bad(t *T) { - rl := ax7Limiter(t) - rl.RecordUsage("model-a", -10, 15) - - stats := rl.Stats("model-a") - AssertEqual(t, 1, stats.RPD) - AssertEqual(t, 15, stats.TPM) -} - -func TestAX7_RateLimiter_RecordUsage_Ugly(t *T) { - rl := ax7Limiter(t) - rl.RecordUsage("", 0, 0) - - stats := rl.Stats("") - AssertEqual(t, 1, stats.RPD) - AssertEqual(t, 0, stats.TPM) -} - -func TestAX7_RateLimiter_WaitForCapacity_Good(t *T) { - rl := ax7Limiter(t) - rl.SetQuota("model-a", ratelimit.ModelQuota{MaxRPM: 1, MaxTPM: 100, MaxRPD: 5}) - - err := rl.WaitForCapacity(context.Background(), "model-a", 1) - AssertNoError(t, err) -} - -func TestAX7_RateLimiter_WaitForCapacity_Bad(t *T) { - rl := ax7Limiter(t) - err := rl.WaitForCapacity(context.Background(), "model-a", -1) - - AssertError(t, err) - AssertContains(t, err.Error(), "negative tokens") -} - -func TestAX7_RateLimiter_WaitForCapacity_Ugly(t *T) { - rl := ax7Limiter(t) - rl.SetQuota("model-a", ratelimit.ModelQuota{MaxRPM: 1, MaxTPM: 100, MaxRPD: 5}) - rl.RecordUsage("model-a", 1, 1) - ctx, cancel := context.WithCancel(context.Background()) - cancel() - - err := rl.WaitForCapacity(ctx, "model-a", 1) - AssertErrorIs(t, err, context.Canceled) -} - -func TestAX7_RateLimiter_Reset_Good(t *T) { - rl := ax7Limiter(t) - rl.RecordUsage("model-a", 1, 1) - rl.Reset("model-a") - - stats := rl.Stats("model-a") - AssertEqual(t, 0, stats.RPD) - AssertEqual(t, 0, stats.TPM) -} - -func TestAX7_RateLimiter_Reset_Bad(t *T) { - rl := ax7Limiter(t) - rl.RecordUsage("model-a", 1, 1) - rl.Reset("missing-model") - - AssertEqual(t, 1, rl.Stats("model-a").RPD) - AssertEqual(t, 0, rl.Stats("missing-model").RPD) -} - -func TestAX7_RateLimiter_Reset_Ugly(t *T) { - rl := ax7Limiter(t) - rl.RecordUsage("model-a", 1, 1) - rl.RecordUsage("model-b", 1, 1) - rl.Reset("") - - AssertEqual(t, 0, rl.Stats("model-a").RPD) - AssertEqual(t, 0, rl.Stats("model-b").RPD) -} - -func TestAX7_RateLimiter_Models_Good(t *T) { - rl, err := ratelimit.NewWithConfig(ratelimit.Config{ - FilePath: ax7YAMLPath(t), - Quotas: map[string]ratelimit.ModelQuota{ - "model-b": {MaxRPM: 2}, - "model-a": {MaxRPM: 1}, - }, - }) - RequireNoError(t, err) - - var models []string - for model := range rl.Models() { - models = append(models, model) - } - AssertEqual(t, []string{"model-a", "model-b"}, models) -} - -func TestAX7_RateLimiter_Models_Bad(t *T) { - rl := ax7Limiter(t) - var models []string - for model := range rl.Models() { - models = append(models, model) - } - - AssertEmpty(t, models) - AssertEqual(t, 0, rl.Stats("missing").MaxRPM) -} - -func TestAX7_RateLimiter_Models_Ugly(t *T) { - rl, err := ratelimit.NewWithConfig(ratelimit.Config{ - FilePath: ax7YAMLPath(t), - Quotas: map[string]ratelimit.ModelQuota{ - "model-a": {MaxRPM: 1}, - "model-b": {MaxRPM: 2}, - "model-c": {MaxRPM: 3}, - }, - }) - RequireNoError(t, err) - - var first string - for model := range rl.Models() { - first = model - break - } - AssertEqual(t, "model-a", first) -} - -func TestAX7_RateLimiter_Iter_Good(t *T) { - rl, err := ratelimit.NewWithConfig(ratelimit.Config{ - FilePath: ax7YAMLPath(t), - Quotas: map[string]ratelimit.ModelQuota{ - "model-b": {MaxRPM: 2}, - "model-a": {MaxRPM: 1}, - }, - }) - RequireNoError(t, err) - rl.RecordUsage("model-a", 1, 1) - - var models []string - for model, stats := range rl.Iter() { - models = append(models, model) - if model == "model-a" { - AssertEqual(t, 1, stats.RPD) - } - } - AssertEqual(t, []string{"model-a", "model-b"}, models) -} - -func TestAX7_RateLimiter_Iter_Bad(t *T) { - rl := ax7Limiter(t) - var models []string - for model := range rl.Iter() { - models = append(models, model) - } - - AssertEmpty(t, models) - AssertEmpty(t, rl.AllStats()) -} - -func TestAX7_RateLimiter_Iter_Ugly(t *T) { - rl, err := ratelimit.NewWithConfig(ratelimit.Config{ - FilePath: ax7YAMLPath(t), - Quotas: map[string]ratelimit.ModelQuota{ - "model-a": {MaxRPM: 1}, - "model-b": {MaxRPM: 2}, - "model-c": {MaxRPM: 3}, - }, - }) - RequireNoError(t, err) - - var seen []string - for model := range rl.Iter() { - seen = append(seen, model) - if len(seen) == 2 { - break - } - } - AssertEqual(t, []string{"model-a", "model-b"}, seen) -} - -func TestAX7_RateLimiter_Stats_Good(t *T) { - rl := ax7Limiter(t) - rl.SetQuota("model-a", ratelimit.ModelQuota{MaxRPM: 10, MaxTPM: 100, MaxRPD: 5}) - rl.RecordUsage("model-a", 2, 3) - - stats := rl.Stats("model-a") - AssertEqual(t, 1, stats.RPM) - AssertEqual(t, 5, stats.TPM) - AssertEqual(t, 10, stats.MaxRPM) -} - -func TestAX7_RateLimiter_Stats_Bad(t *T) { - rl := ax7Limiter(t) - stats := rl.Stats("missing-model") - - AssertEqual(t, 0, stats.RPM) - AssertEqual(t, 0, stats.MaxRPM) - AssertTrue(t, stats.DayStart.IsZero()) -} - -func TestAX7_RateLimiter_Stats_Ugly(t *T) { - rl := ax7Limiter(t) - rl.RecordUsage("model-a", 1, 1) - rl.Reset("model-a") - - stats := rl.Stats("model-a") - AssertEqual(t, 0, stats.RPD) - AssertEqual(t, 0, stats.TPM) -} - -func TestAX7_RateLimiter_AllStats_Good(t *T) { - rl := ax7Limiter(t) - rl.SetQuota("model-a", ratelimit.ModelQuota{MaxRPM: 10}) - rl.RecordUsage("model-b", 1, 1) - - all := rl.AllStats() - AssertContains(t, all, "model-a") - AssertContains(t, all, "model-b") - AssertEqual(t, 2, all["model-b"].TPM) -} - -func TestAX7_RateLimiter_AllStats_Bad(t *T) { - rl := ax7Limiter(t) - all := rl.AllStats() - - AssertEmpty(t, all) - AssertEqual(t, 0, len(all)) -} - -func TestAX7_RateLimiter_AllStats_Ugly(t *T) { - rl := ax7Limiter(t) - rl.SetQuota("model-a", ratelimit.ModelQuota{MaxRPM: 10}) - all := rl.AllStats() - - AssertContains(t, all, "model-a") - AssertEqual(t, 0, all["model-a"].RPM) - AssertEqual(t, 10, all["model-a"].MaxRPM) -} - -func TestAX7_RateLimiter_Decide_Good(t *T) { - rl := ax7Limiter(t) - rl.SetQuota("model-a", ratelimit.ModelQuota{MaxRPM: 2, MaxTPM: 100, MaxRPD: 5}) - - decision := rl.Decide("model-a", 10) - AssertTrue(t, decision.Allowed) - AssertEqual(t, ratelimit.DecisionAllowed, decision.Code) -} - -func TestAX7_RateLimiter_Decide_Bad(t *T) { - rl := ax7Limiter(t) - decision := rl.Decide("model-a", -1) - - AssertFalse(t, decision.Allowed) - AssertEqual(t, ratelimit.DecisionInvalidTokens, decision.Code) - AssertContains(t, decision.Reason, "non-negative") -} - -func TestAX7_RateLimiter_Decide_Ugly(t *T) { - rl := ax7Limiter(t) - rl.SetQuota("model-a", ratelimit.ModelQuota{MaxRPM: 1, MaxTPM: 100, MaxRPD: 5}) - rl.RecordUsage("model-a", 1, 1) - - decision := rl.Decide("model-a", 1) - AssertFalse(t, decision.Allowed) - AssertEqual(t, ratelimit.DecisionRPMLimit, decision.Code) -} - -func TestAX7_NewWithSQLite_Good(t *T) { - rl, err := ratelimit.NewWithSQLite(Path(t.TempDir(), "limits.db")) - - RequireNoError(t, err) - defer rl.Close() - rl.RecordUsage("gemini-3-pro-preview", 1, 1) - AssertNoError(t, rl.Persist()) -} - -func TestAX7_NewWithSQLite_Bad(t *T) { - rl, err := ratelimit.NewWithSQLite(Path(t.TempDir(), "missing", "limits.db")) - - AssertError(t, err) - AssertNil(t, rl) -} - -func TestAX7_NewWithSQLite_Ugly(t *T) { - path := Path(t.TempDir(), "limits with spaces.db") - rl, err := ratelimit.NewWithSQLite(path) - RequireNoError(t, err) - rl.RecordUsage("gemini-3-pro-preview", 1, 1) - - AssertNoError(t, rl.Persist()) - AssertNoError(t, rl.Close()) -} - -func TestAX7_NewWithSQLiteConfig_Good(t *T) { - rl, err := ratelimit.NewWithSQLiteConfig(Path(t.TempDir(), "limits.db"), ratelimit.Config{ - Providers: []ratelimit.Provider{ratelimit.ProviderOpenAI}, - Quotas: map[string]ratelimit.ModelQuota{"custom": {MaxRPM: 9}}, - }) - - RequireNoError(t, err) - defer rl.Close() - AssertEqual(t, 500, rl.Stats("gpt-4o").MaxRPM) - AssertEqual(t, 9, rl.Stats("custom").MaxRPM) -} - -func TestAX7_NewWithSQLiteConfig_Bad(t *T) { - rl, err := ratelimit.NewWithSQLiteConfig(Path(t.TempDir(), "missing", "limits.db"), ratelimit.Config{}) - - AssertError(t, err) - AssertNil(t, rl) -} - -func TestAX7_NewWithSQLiteConfig_Ugly(t *T) { - rl, err := ratelimit.NewWithSQLiteConfig(Path(t.TempDir(), "limits.db"), ratelimit.Config{ - Backend: "yaml", - Providers: []ratelimit.Provider{ratelimit.ProviderLocal}, - }) - - RequireNoError(t, err) - defer rl.Close() - rl.RecordUsage("local", 1, 1) - AssertNoError(t, rl.Persist()) -} - -func TestAX7_RateLimiter_Close_Good(t *T) { - rl := ax7Limiter(t) - err := rl.Close() - - AssertNoError(t, err) - AssertNoError(t, rl.Close()) -} - -func TestAX7_RateLimiter_Close_Bad(t *T) { - rl, err := ratelimit.NewWithSQLite(Path(t.TempDir(), "limits.db")) - RequireNoError(t, err) - RequireNoError(t, rl.Close()) - - err = rl.Persist() - AssertError(t, err) -} - -func TestAX7_RateLimiter_Close_Ugly(t *T) { - rl, err := ratelimit.NewWithSQLite(Path(t.TempDir(), "limits.db")) - RequireNoError(t, err) - - AssertNoError(t, rl.Close()) - AssertNoError(t, rl.Close()) -} - -func TestAX7_MigrateYAMLToSQLite_Good(t *T) { - yamlPath := ax7YAMLPath(t) - sqlitePath := Path(t.TempDir(), "limits.db") - rl, err := ratelimit.NewWithConfig(ratelimit.Config{ - FilePath: yamlPath, - Quotas: map[string]ratelimit.ModelQuota{"model-a": {MaxRPM: 4}}, - }) - RequireNoError(t, err) - rl.RecordUsage("model-a", 2, 3) - RequireNoError(t, rl.Persist()) - - AssertNoError(t, ratelimit.MigrateYAMLToSQLite(yamlPath, sqlitePath)) - loaded, err := ratelimit.NewWithSQLite(sqlitePath) - RequireNoError(t, err) - defer loaded.Close() - AssertNoError(t, loaded.Load()) - AssertEqual(t, 1, loaded.Stats("model-a").RPD) -} - -func TestAX7_MigrateYAMLToSQLite_Bad(t *T) { - err := ratelimit.MigrateYAMLToSQLite(Path(t.TempDir(), "missing.yaml"), Path(t.TempDir(), "limits.db")) - - AssertError(t, err) - AssertContains(t, err.Error(), "read") -} - -func TestAX7_MigrateYAMLToSQLite_Ugly(t *T) { - yamlPath := ax7YAMLPath(t) - sqlitePath := Path(t.TempDir(), "limits.db") - RequireNoError(t, os.WriteFile(yamlPath, []byte("{{{not yaml"), 0o600)) - - err := ratelimit.MigrateYAMLToSQLite(yamlPath, sqlitePath) - AssertError(t, err) - AssertContains(t, err.Error(), "unmarshal") -} - -func TestAX7_CountTokens_Good(t *T) { - oldTransport := http.DefaultTransport - http.DefaultTransport = ax7RoundTrip(func(req *http.Request) (*http.Response, error) { - AssertContains(t, req.URL.Path, "/v1beta/models/gemini-3-pro-preview:countTokens") - return &http.Response{ - StatusCode: http.StatusOK, - Body: io.NopCloser(NewReader(`{"totalTokens":7}`)), - Header: make(http.Header), - }, nil - }) - defer func() { http.DefaultTransport = oldTransport }() - - tokens, err := ratelimit.CountTokens(context.Background(), "key", "gemini-3-pro-preview", "hello") - AssertNoError(t, err) - AssertEqual(t, 7, tokens) -} - -func TestAX7_CountTokens_Bad(t *T) { - tokens, err := ratelimit.CountTokens(context.Background(), "key", "", "hello") - - AssertError(t, err) - AssertEqual(t, 0, tokens) - AssertContains(t, err.Error(), "empty model") -} - -func TestAX7_CountTokens_Ugly(t *T) { - oldTransport := http.DefaultTransport - http.DefaultTransport = ax7RoundTrip(func(req *http.Request) (*http.Response, error) { - return &http.Response{ - StatusCode: http.StatusOK, - Body: io.NopCloser(NewReader(`{"totalTokens":"bad"}`)), - Header: make(http.Header), - }, nil - }) - defer func() { http.DefaultTransport = oldTransport }() - - tokens, err := ratelimit.CountTokens(context.Background(), "key", "gemini-3-pro-preview", "hello") - AssertError(t, err) - AssertEqual(t, 0, tokens) -} diff --git a/external/go b/external/go new file mode 160000 index 0000000..d661b70 --- /dev/null +++ b/external/go @@ -0,0 +1 @@ +Subproject commit d661b703e16183b3cbab101de189f688888a1174 diff --git a/go.work b/go.work new file mode 100644 index 0000000..2682675 --- /dev/null +++ b/go.work @@ -0,0 +1,12 @@ +// SPDX-License-Identifier: EUPL-1.2 + +go 1.26.0 + +// Workspace mode for development: pulls local sources from external/ submodules. +// +// CI: GOWORK=off uses go/go.mod tags for reproducible resolution. + +use ( + ./go + ./external/go +) diff --git a/assertions_test.go b/go/assertions_test.go similarity index 78% rename from assertions_test.go rename to go/assertions_test.go index 6bfabac..48cd878 100644 --- a/assertions_test.go +++ b/go/assertions_test.go @@ -3,11 +3,9 @@ package ratelimit import ( - "errors" - "fmt" + core "dappco.re/go" "math" "reflect" - "strings" "time" ) @@ -16,9 +14,9 @@ func testFailureMessage(defaultMsg string, msgAndArgs ...any) string { return defaultMsg } - msg := fmt.Sprint(msgAndArgs...) + msg := core.Sprint(msgAndArgs...) if format, ok := msgAndArgs[0].(string); ok && len(msgAndArgs) > 1 { - msg = fmt.Sprintf(format, msgAndArgs[1:]...) + msg = core.Sprintf(format, msgAndArgs[1:]...) } if defaultMsg == "" { return msg @@ -27,7 +25,7 @@ func testFailureMessage(defaultMsg string, msgAndArgs ...any) string { } func testUnexpectedErrorMessage(err error, msgAndArgs ...any) string { - return testFailureMessage(fmt.Sprintf("unexpected error: %v", err), msgAndArgs...) + return testFailureMessage(core.Sprintf("unexpected error: %v", err), msgAndArgs...) } func testExpectedErrorMessage(msgAndArgs ...any) string { @@ -35,7 +33,7 @@ func testExpectedErrorMessage(msgAndArgs ...any) string { } func testWantGotMessage(want, got any, msgAndArgs ...any) string { - return testFailureMessage(fmt.Sprintf("want %v, got %v", want, got), msgAndArgs...) + return testFailureMessage(core.Sprintf("want %v, got %v", want, got), msgAndArgs...) } func testExpectedTrueMessage(msgAndArgs ...any) string { @@ -47,7 +45,7 @@ func testExpectedFalseMessage(msgAndArgs ...any) string { } func testExpectedNilMessage(v any, msgAndArgs ...any) string { - return testFailureMessage(fmt.Sprintf("expected nil, got %v", v), msgAndArgs...) + return testFailureMessage(core.Sprintf("expected nil, got %v", v), msgAndArgs...) } func testExpectedNonNilMessage(msgAndArgs ...any) string { @@ -55,46 +53,46 @@ func testExpectedNonNilMessage(msgAndArgs ...any) string { } func testContainsMessage(container, elem any, msgAndArgs ...any) string { - return testFailureMessage(fmt.Sprintf("expected %v to contain %v", container, elem), msgAndArgs...) + return testFailureMessage(core.Sprintf("expected %v to contain %v", container, elem), msgAndArgs...) } func testNotContainsMessage(container, elem any, msgAndArgs ...any) string { - return testFailureMessage(fmt.Sprintf("expected %v not to contain %v", container, elem), msgAndArgs...) + return testFailureMessage(core.Sprintf("expected %v not to contain %v", container, elem), msgAndArgs...) } func testEmptyMessage(v any, msgAndArgs ...any) string { - return testFailureMessage(fmt.Sprintf("expected empty, got %v", v), msgAndArgs...) + return testFailureMessage(core.Sprintf("expected empty, got %v", v), msgAndArgs...) } func testLenMessage(v any, want int, msgAndArgs ...any) string { if got, ok := testLenOf(v); ok { - return testFailureMessage(fmt.Sprintf("expected length %d, got %d", want, got), msgAndArgs...) + return testFailureMessage(core.Sprintf("expected length %d, got %d", want, got), msgAndArgs...) } - return testFailureMessage(fmt.Sprintf("expected length %d, got non-len value %v", want, v), msgAndArgs...) + return testFailureMessage(core.Sprintf("expected length %d, got non-len value %v", want, v), msgAndArgs...) } func testErrorIsMessage(err, target error, msgAndArgs ...any) string { - return testFailureMessage(fmt.Sprintf("expected error %v to match %v", err, target), msgAndArgs...) + return testFailureMessage(core.Sprintf("expected error %v to match %v", err, target), msgAndArgs...) } func testInDeltaMessage(want, got, delta any, msgAndArgs ...any) string { - return testFailureMessage(fmt.Sprintf("expected %v and %v to be within %v", want, got, delta), msgAndArgs...) + return testFailureMessage(core.Sprintf("expected %v and %v to be within %v", want, got, delta), msgAndArgs...) } func testZeroMessage(v any, msgAndArgs ...any) string { - return testFailureMessage(fmt.Sprintf("expected zero value, got %v", v), msgAndArgs...) + return testFailureMessage(core.Sprintf("expected zero value, got %v", v), msgAndArgs...) } func testUnexpectedPanicMessage(recovered any, msgAndArgs ...any) string { - return testFailureMessage(fmt.Sprintf("unexpected panic: %v", recovered), msgAndArgs...) + return testFailureMessage(core.Sprintf("unexpected panic: %v", recovered), msgAndArgs...) } func testGreaterOrEqualMessage(a, b any, msgAndArgs ...any) string { - return testFailureMessage(fmt.Sprintf("expected %v >= %v", a, b), msgAndArgs...) + return testFailureMessage(core.Sprintf("expected %v >= %v", a, b), msgAndArgs...) } func testGreaterMessage(a, b any, msgAndArgs ...any) string { - return testFailureMessage(fmt.Sprintf("expected %v > %v", a, b), msgAndArgs...) + return testFailureMessage(core.Sprintf("expected %v > %v", a, b), msgAndArgs...) } func testEventuallyMessage(msgAndArgs ...any) string { @@ -126,9 +124,9 @@ func testContains(container, elem any) bool { if s, ok := container.(string); ok { needle, ok := elem.(string) if !ok { - needle = fmt.Sprint(elem) + needle = core.Sprint(elem) } - return strings.Contains(s, needle) + return core.Contains(s, needle) } cv := reflect.ValueOf(container) @@ -186,7 +184,7 @@ func testHasLen(v any, want int) bool { } func testErrorIs(err, target error) bool { - return errors.Is(err, target) + return core.Is(err, target) } func testInDelta(want, got, delta any) bool { diff --git a/error_test.go b/go/error_test.go similarity index 95% rename from error_test.go rename to go/error_test.go index 148f12f..0857abe 100644 --- a/error_test.go +++ b/go/error_test.go @@ -8,7 +8,7 @@ import ( "time" ) -func TestError_SQLiteErrorPaths_Bad(t *testing.T) { +func TestError_SQLiteErrorPaths_Case(t *testing.T) { dbPath := testPath(t.TempDir(), "error.db") rl, err := NewWithSQLite(dbPath) if err != nil { @@ -47,7 +47,7 @@ func TestError_SQLiteErrorPaths_Bad(t *testing.T) { }) } -func TestError_SQLiteInitErrors_Bad(t *testing.T) { +func TestError_SQLiteInitErrors_Case(t *testing.T) { dbPath := testPath(t.TempDir(), "closed-schema.db") store, err := newSQLiteStore(dbPath) if err != nil { @@ -63,7 +63,7 @@ func TestError_SQLiteInitErrors_Bad(t *testing.T) { } } -func TestError_PersistYAML_Good(t *testing.T) { +func TestError_PersistYAML_Case(t *testing.T) { t.Run("successful YAML persist and load", func(t *testing.T) { tmpDir := t.TempDir() path := testPath(tmpDir, "ratelimits.yaml") @@ -89,7 +89,7 @@ func TestError_PersistYAML_Good(t *testing.T) { }) } -func TestError_SQLiteLoadViaLimiter_Bad(t *testing.T) { +func TestError_SQLiteLoadViaLimiter_Case(t *testing.T) { t.Run("Load returns error when SQLite DB is closed", func(t *testing.T) { dbPath := testPath(t.TempDir(), "load-err.db") rl, err := NewWithSQLite(dbPath) @@ -133,7 +133,7 @@ func TestError_SQLiteLoadViaLimiter_Bad(t *testing.T) { }) } -func TestError_SQLitePersistViaLimiter_Bad(t *testing.T) { +func TestError_SQLitePersistViaLimiter_Case(t *testing.T) { t.Run("Persist returns error when SQLite saveQuotas fails", func(t *testing.T) { dbPath := testPath(t.TempDir(), "persist-err.db") rl, err := NewWithSQLite(dbPath) @@ -201,7 +201,7 @@ func TestError_NewWithSQLite_Bad(t *testing.T) { }) } -func TestError_SQLiteSaveState_Bad(t *testing.T) { +func TestError_SQLiteSaveState_Case(t *testing.T) { t.Run("saveState fails when tokens table is dropped", func(t *testing.T) { dbPath := testPath(t.TempDir(), "tokens-err.db") store, err := newSQLiteStore(dbPath) @@ -334,7 +334,7 @@ func TestError_SQLiteSaveState_Bad(t *testing.T) { }) } -func TestError_SQLiteLoadState_Bad(t *testing.T) { +func TestError_SQLiteLoadState_Case(t *testing.T) { t.Run("loadState fails when requests table is dropped", func(t *testing.T) { dbPath := testPath(t.TempDir(), "req-err.db") store, err := newSQLiteStore(dbPath) @@ -425,7 +425,7 @@ func TestError_SQLiteLoadState_Bad(t *testing.T) { }) } -func TestError_SQLiteSaveQuotasExec_Bad(t *testing.T) { +func TestError_SQLiteSaveQuotasExec_Case(t *testing.T) { t.Run("saveQuotas fails with renamed column at prepare", func(t *testing.T) { dbPath := testPath(t.TempDir(), "quota-exec-err.db") store, err := newSQLiteStore(dbPath) @@ -475,7 +475,7 @@ func TestError_SQLiteSaveQuotasExec_Bad(t *testing.T) { }) } -func TestError_SQLiteSaveStateExec_Bad(t *testing.T) { +func TestError_SQLiteSaveStateExec_Case(t *testing.T) { t.Run("request insert exec fails via trigger", func(t *testing.T) { dbPath := testPath(t.TempDir(), "trigger-req.db") store, err := newSQLiteStore(dbPath) @@ -569,7 +569,7 @@ func TestError_SQLiteSaveStateExec_Bad(t *testing.T) { }) } -func TestError_SQLiteLoadQuotasScan_Bad(t *testing.T) { +func TestError_SQLiteLoadQuotasScan_Case(t *testing.T) { t.Run("loadQuotas fails with renamed column", func(t *testing.T) { dbPath := testPath(t.TempDir(), "quota-scan-err.db") store, err := newSQLiteStore(dbPath) @@ -599,7 +599,7 @@ func TestError_SQLiteLoadQuotasScan_Bad(t *testing.T) { }) } -func TestError_NewSQLiteStoreInReadOnlyDir_Bad(t *testing.T) { +func TestError_NewSQLiteStoreInReadOnlyDir_Case(t *testing.T) { if isRootUser() { t.Skip("chmod restrictions do not apply to root") } @@ -621,7 +621,7 @@ func TestError_NewSQLiteStoreInReadOnlyDir_Bad(t *testing.T) { }) } -func TestError_SQLiteCreateSchema_Bad(t *testing.T) { +func TestError_SQLiteCreateSchema_Case(t *testing.T) { t.Run("createSchema fails on closed DB", func(t *testing.T) { dbPath := testPath(t.TempDir(), "schema-err.db") store, err := newSQLiteStore(dbPath) @@ -642,7 +642,7 @@ func TestError_SQLiteCreateSchema_Bad(t *testing.T) { }) } -func TestError_SQLiteLoadStateScan_Bad(t *testing.T) { +func TestError_SQLiteLoadStateScan_Case(t *testing.T) { t.Run("scan daily fails with NULL values", func(t *testing.T) { dbPath := testPath(t.TempDir(), "scan-daily.db") store, err := newSQLiteStore(dbPath) @@ -746,7 +746,7 @@ func TestError_SQLiteLoadStateScan_Bad(t *testing.T) { }) } -func TestError_SQLiteLoadQuotasScanWithBadSchema_Bad(t *testing.T) { +func TestError_SQLiteLoadQuotasScanWithBadSchema_Case(t *testing.T) { t.Run("scan fails with NULL quota values", func(t *testing.T) { dbPath := testPath(t.TempDir(), "scan-quota.db") store, err := newSQLiteStore(dbPath) @@ -775,7 +775,7 @@ func TestError_SQLiteLoadQuotasScanWithBadSchema_Bad(t *testing.T) { }) } -func TestError_MigrateYAMLToSQLiteWithSaveErrors_Bad(t *testing.T) { +func TestError_MigrateYAMLToSQLiteWithSaveErrors_Case(t *testing.T) { t.Run("saveQuotas failure during migration via trigger", func(t *testing.T) { tmpDir := t.TempDir() yamlPath := testPath(tmpDir, "with-quotas.yaml") @@ -843,7 +843,7 @@ func TestError_MigrateYAMLToSQLiteWithSaveErrors_Bad(t *testing.T) { }) } -func TestError_MigrateYAMLToSQLiteNilQuotasAndState_Good(t *testing.T) { +func TestError_MigrateYAMLToSQLiteNilQuotasAndState_Case(t *testing.T) { t.Run("YAML with empty quotas and state migrates cleanly", func(t *testing.T) { tmpDir := t.TempDir() yamlPath := testPath(tmpDir, "empty.yaml") @@ -878,7 +878,7 @@ func TestError_MigrateYAMLToSQLiteNilQuotasAndState_Good(t *testing.T) { }) } -func TestError_NewWithConfigHomeUnavailable_Bad(t *testing.T) { +func TestError_NewWithConfigHomeUnavailable_Case(t *testing.T) { // Clear all supported home env vars so defaultStatePath cannot resolve a home directory. t.Setenv("CORE_HOME", "") t.Setenv("HOME", "") @@ -891,7 +891,7 @@ func TestError_NewWithConfigHomeUnavailable_Bad(t *testing.T) { } } -func TestError_PersistMarshal_Good(t *testing.T) { +func TestError_PersistMarshal_Case(t *testing.T) { // yaml.Marshal on a struct with map[string]ModelQuota and map[string]*UsageStats // should not fail in practice. We test the error path by using a type that // yaml.Marshal cannot handle: a channel. @@ -905,7 +905,7 @@ func TestError_PersistMarshal_Good(t *testing.T) { } } -func TestError_MigrateErrorsExtended_Bad(t *testing.T) { +func TestError_MigrateErrorsExtended_Case(t *testing.T) { t.Run("unmarshal failure", func(t *testing.T) { tmpDir := t.TempDir() path := testPath(tmpDir, "bad.yaml") diff --git a/go.mod b/go/go.mod similarity index 100% rename from go.mod rename to go/go.mod diff --git a/go.sum b/go/go.sum similarity index 100% rename from go.sum rename to go/go.sum diff --git a/iter_test.go b/go/iter_test.go similarity index 95% rename from iter_test.go rename to go/iter_test.go index 3593bec..b1c1abe 100644 --- a/iter_test.go +++ b/go/iter_test.go @@ -9,7 +9,7 @@ import ( "testing" ) -func TestIter_Iterators_Good(t *testing.T) { +func TestIter_Iterators_Case(t *testing.T) { rl, err := NewWithConfig(Config{ Quotas: map[string]ModelQuota{ "model-c": {MaxRPM: 10}, @@ -97,7 +97,7 @@ func TestIter_Iterators_Good(t *testing.T) { }) } -func TestIter_Iterators_Bad(t *testing.T) { +func TestIter_Iterators_Case_2(t *testing.T) { rl, err := NewWithConfig(Config{ FilePath: testPath(t.TempDir(), "iter-empty.yaml"), Providers: []Provider{ProviderLocal}, @@ -123,7 +123,7 @@ func TestIter_Iterators_Bad(t *testing.T) { } } -func TestIter_Iterators_Ugly(t *testing.T) { +func TestIter_Iterators_Case_3(t *testing.T) { rl, err := NewWithConfig(Config{ FilePath: testPath(t.TempDir(), "iter-break.yaml"), Quotas: map[string]ModelQuota{ @@ -160,7 +160,7 @@ func TestIter_Iterators_Ugly(t *testing.T) { } } -func TestIter_IterEarlyBreak_Good(t *testing.T) { +func TestIter_IterEarlyBreak_Case(t *testing.T) { rl, err := NewWithConfig(Config{ Quotas: map[string]ModelQuota{ "model-a": {MaxRPM: 10}, @@ -199,7 +199,7 @@ func TestIter_IterEarlyBreak_Good(t *testing.T) { }) } -func TestIter_CountTokensFull_Ugly(t *testing.T) { +func TestIter_CountTokensFull_Case(t *testing.T) { t.Run("empty model is rejected", func(t *testing.T) { _, err := CountTokens(context.Background(), "key", "", "text") if err == nil { diff --git a/ratelimit.go b/go/ratelimit.go similarity index 95% rename from ratelimit.go rename to go/ratelimit.go index a1450c1..6871755 100644 --- a/ratelimit.go +++ b/go/ratelimit.go @@ -170,7 +170,7 @@ func DefaultProfiles() map[Provider]ProviderProfile { // This preserves backward compatibility -- existing callers are unaffected. // // rl, err := New() -func New() (*RateLimiter, error) { +func New() (*RateLimiter, error) /* core result boundary */ { return NewWithConfig(Config{ Providers: []Provider{ProviderGemini}, }) @@ -180,7 +180,7 @@ func New() (*RateLimiter, error) { // If no providers or quotas are specified, Gemini defaults are used. // // rl, err := NewWithConfig(Config{Providers: []Provider{ProviderAnthropic}}) -func NewWithConfig(cfg Config) (*RateLimiter, error) { +func NewWithConfig(cfg Config) (*RateLimiter, error) /* core result boundary */ { backend, err := normaliseBackend(cfg.Backend) if err != nil { return nil, err @@ -234,7 +234,7 @@ func (rl *RateLimiter) AddProvider(provider Provider) { // Load reads the state from disk (YAML) or database (SQLite). // // if err := rl.Load(); err != nil { /* handle error */ } -func (rl *RateLimiter) Load() error { +func (rl *RateLimiter) Load() error /* core result boundary */ { rl.mu.Lock() defer rl.mu.Unlock() @@ -260,7 +260,7 @@ func (rl *RateLimiter) Load() error { // loadSQLite reads quotas and state from the SQLite backend. // Caller must hold the lock. -func (rl *RateLimiter) loadSQLite() error { +func (rl *RateLimiter) loadSQLite() error /* core result boundary */ { quotas, err := rl.sqlite.loadQuotas() if err != nil { return err @@ -283,7 +283,7 @@ func (rl *RateLimiter) loadSQLite() error { // It clones the state under a lock and performs I/O without blocking other callers. // // if err := rl.Persist(); err != nil { /* handle error */ } -func (rl *RateLimiter) Persist() error { +func (rl *RateLimiter) Persist() error /* core result boundary */ { rl.mu.Lock() quotas := maps.Clone(rl.Quotas) state := make(map[string]*UsageStats, len(rl.State)) @@ -437,7 +437,7 @@ func (rl *RateLimiter) RecordUsage(model string, promptTokens, outputTokens int) // WaitForCapacity blocks until capacity is available or context is cancelled. // // err := rl.WaitForCapacity(ctx, "gemini-3-pro-preview", 1200) -func (rl *RateLimiter) WaitForCapacity(ctx context.Context, model string, tokens int) error { +func (rl *RateLimiter) WaitForCapacity(ctx context.Context, model string, tokens int) error /* core result boundary */ { if tokens < 0 { return core.E("ratelimit.WaitForCapacity", "negative tokens", nil) } @@ -703,7 +703,7 @@ func (rl *RateLimiter) snapshotLocked(model string) ModelStats { // release the database connection when finished. // // rl, err := NewWithSQLite("/tmp/ratelimits.db") -func NewWithSQLite(dbPath string) (*RateLimiter, error) { +func NewWithSQLite(dbPath string) (*RateLimiter, error) /* core result boundary */ { return NewWithSQLiteConfig(dbPath, Config{ Providers: []Provider{ProviderGemini}, }) @@ -714,7 +714,7 @@ func NewWithSQLite(dbPath string) (*RateLimiter, error) { // release the database connection when finished. // // rl, err := NewWithSQLiteConfig("/tmp/ratelimits.db", Config{Providers: []Provider{ProviderOpenAI}}) -func NewWithSQLiteConfig(dbPath string, cfg Config) (*RateLimiter, error) { +func NewWithSQLiteConfig(dbPath string, cfg Config) (*RateLimiter, error) /* core result boundary */ { store, err := newSQLiteStore(dbPath) if err != nil { return nil, err @@ -730,7 +730,7 @@ func NewWithSQLiteConfig(dbPath string, cfg Config) (*RateLimiter, error) { // database connection. // // defer rl.Close() -func (rl *RateLimiter) Close() error { +func (rl *RateLimiter) Close() error /* core result boundary */ { if rl.sqlite != nil { return rl.sqlite.close() } @@ -742,7 +742,7 @@ func (rl *RateLimiter) Close() error { // database is created if it does not exist. // // err := MigrateYAMLToSQLite("ratelimits.yaml", "ratelimits.db") -func MigrateYAMLToSQLite(yamlPath, sqlitePath string) error { +func MigrateYAMLToSQLite(yamlPath, sqlitePath string) error /* core result boundary */ { // Load from YAML. content, err := readLocalFile(yamlPath) if err != nil { @@ -770,11 +770,11 @@ func MigrateYAMLToSQLite(yamlPath, sqlitePath string) error { // CountTokens calls the Google API to count tokens for a prompt. // // tokens, err := CountTokens(ctx, apiKey, "gemini-3-pro-preview", prompt) -func CountTokens(ctx context.Context, apiKey, model, text string) (int, error) { +func CountTokens(ctx context.Context, apiKey, model, text string) (int, error) /* core result boundary */ { return countTokensWithClient(ctx, http.DefaultClient, "https://generativelanguage.googleapis.com", apiKey, model, text) } -func countTokensWithClient(ctx context.Context, client *http.Client, baseURL, apiKey, model, text string) (int, error) { +func countTokensWithClient(ctx context.Context, client *http.Client, baseURL, apiKey, model, text string) (int, error) /* core result boundary */ { requestURL, err := countTokensURL(baseURL, model) if err != nil { return 0, core.E("ratelimit.CountTokens", "build url", err) @@ -879,7 +879,7 @@ func applyConfig(rl *RateLimiter, cfg Config) { maps.Copy(rl.Quotas, cfg.Quotas) } -func normaliseBackend(backend string) (string, error) { +func normaliseBackend(backend string) (string, error) /* core result boundary */ { switch core.Lower(core.Trim(backend)) { case "", backendYAML: return backendYAML, nil @@ -890,7 +890,7 @@ func normaliseBackend(backend string) (string, error) { } } -func defaultStatePath(backend string) (string, error) { +func defaultStatePath(backend string) (string, error) /* core result boundary */ { home := currentHomeDir() if home == "" { return "", core.E("ratelimit.defaultStatePath", "home dir unavailable", nil) @@ -997,7 +997,7 @@ func nonNegativeDuration(value time.Duration) time.Duration { return value } -func countTokensURL(baseURL, model string) (string, error) { +func countTokensURL(baseURL, model string) (string, error) /* core result boundary */ { if core.Trim(model) == "" { return "", core.E("ratelimit.countTokensURL", "empty model", nil) } @@ -1027,7 +1027,7 @@ func hasURLAuthority(rawURL string) bool { return authority != "" } -func readLimitedBody(r io.Reader, limit int64) (string, error) { +func readLimitedBody(r io.Reader, limit int64) (string, error) /* core result boundary */ { body, err := io.ReadAll(io.LimitReader(r, limit+1)) if err != nil { return "", err @@ -1045,7 +1045,7 @@ func readLimitedBody(r io.Reader, limit int64) (string, error) { return result, nil } -func readLocalFile(path string) (string, error) { +func readLocalFile(path string) (string, error) /* core result boundary */ { var fs core.Fs result := fs.Read(path) if !result.OK { @@ -1059,17 +1059,17 @@ func readLocalFile(path string) (string, error) { return content, nil } -func writeLocalFile(path, content string) error { +func writeLocalFile(path, content string) error /* core result boundary */ { var fs core.Fs return resultError(fs.Write(path, content)) } -func ensureDir(path string) error { +func ensureDir(path string) error /* core result boundary */ { var fs core.Fs return resultError(fs.EnsureDir(path)) } -func resultError(result core.Result) error { +func resultError(result core.Result) error /* core result boundary */ { if result.OK { return nil } diff --git a/go/ratelimit_example_test.go b/go/ratelimit_example_test.go new file mode 100644 index 0000000..048b0b9 --- /dev/null +++ b/go/ratelimit_example_test.go @@ -0,0 +1,182 @@ +// SPDX-License-Identifier: EUPL-1.2 + +package ratelimit + +import ( + "context" + "time" +) + +func ExampleDefaultProfiles() { + profiles := DefaultProfiles() + _ = profiles[ProviderGemini] +} + +func ExampleNew() { + rl, err := New() + if err != nil { + return + } + defer rl.Close() + _ = rl.CanSend("gemini-3-pro-preview", 1) +} + +func ExampleNewWithConfig() { + rl, err := NewWithConfig(Config{Providers: []Provider{ProviderLocal}}) + if err != nil { + return + } + defer rl.Close() + rl.SetQuota("local-model", ModelQuota{MaxRPM: 1}) +} + +func ExampleRateLimiter_SetQuota() { + rl, err := NewWithConfig(Config{Providers: []Provider{ProviderLocal}}) + if err != nil { + return + } + rl.SetQuota("local-model", ModelQuota{MaxRPM: 10}) +} + +func ExampleRateLimiter_AddProvider() { + rl, err := NewWithConfig(Config{Providers: []Provider{ProviderLocal}}) + if err != nil { + return + } + rl.AddProvider(ProviderOpenAI) +} + +func ExampleRateLimiter_Load() { + rl, err := NewWithConfig(Config{Providers: []Provider{ProviderLocal}}) + if err != nil { + return + } + _ = rl.Load() +} + +func ExampleRateLimiter_Persist() { + rl, err := NewWithConfig(Config{Providers: []Provider{ProviderLocal}}) + if err != nil { + return + } + _ = rl.Persist() +} + +func ExampleRateLimiter_BackgroundPrune() { + rl, err := NewWithConfig(Config{Providers: []Provider{ProviderLocal}}) + if err != nil { + return + } + stop := rl.BackgroundPrune(time.Minute) + defer stop() +} + +func ExampleRateLimiter_CanSend() { + rl, err := NewWithConfig(Config{Providers: []Provider{ProviderLocal}}) + if err != nil { + return + } + _ = rl.CanSend("local-model", 1) +} + +func ExampleRateLimiter_RecordUsage() { + rl, err := NewWithConfig(Config{Providers: []Provider{ProviderLocal}}) + if err != nil { + return + } + rl.RecordUsage("local-model", 1, 1) +} + +func ExampleRateLimiter_WaitForCapacity() { + rl, err := NewWithConfig(Config{Providers: []Provider{ProviderLocal}}) + if err != nil { + return + } + _ = rl.WaitForCapacity(context.Background(), "local-model", 1) +} + +func ExampleRateLimiter_Reset() { + rl, err := NewWithConfig(Config{Providers: []Provider{ProviderLocal}}) + if err != nil { + return + } + rl.Reset("") +} + +func ExampleRateLimiter_Models() { + rl, err := NewWithConfig(Config{Providers: []Provider{ProviderGemini}}) + if err != nil { + return + } + for model := range rl.Models() { + _ = model + break + } +} + +func ExampleRateLimiter_Iter() { + rl, err := NewWithConfig(Config{Providers: []Provider{ProviderGemini}}) + if err != nil { + return + } + for model, stats := range rl.Iter() { + _, _ = model, stats + break + } +} + +func ExampleRateLimiter_Stats() { + rl, err := NewWithConfig(Config{Providers: []Provider{ProviderLocal}}) + if err != nil { + return + } + _ = rl.Stats("local-model") +} + +func ExampleRateLimiter_AllStats() { + rl, err := NewWithConfig(Config{Providers: []Provider{ProviderLocal}}) + if err != nil { + return + } + _ = rl.AllStats() +} + +func ExampleRateLimiter_Decide() { + rl, err := NewWithConfig(Config{Providers: []Provider{ProviderLocal}}) + if err != nil { + return + } + _ = rl.Decide("local-model", 1) +} + +func ExampleNewWithSQLite() { + rl, err := NewWithSQLite(":memory:") + if err != nil { + return + } + defer rl.Close() +} + +func ExampleNewWithSQLiteConfig() { + rl, err := NewWithSQLiteConfig(":memory:", Config{Providers: []Provider{ProviderLocal}}) + if err != nil { + return + } + defer rl.Close() +} + +func ExampleRateLimiter_Close() { + rl, err := NewWithSQLite(":memory:") + if err != nil { + return + } + _ = rl.Close() +} + +func ExampleMigrateYAMLToSQLite() { + _ = MigrateYAMLToSQLite("ratelimits.yaml", "ratelimits.db") +} + +func ExampleCountTokens() { + _, _ = CountTokens(context.Background(), "api-key", "gemini-3-pro-preview", "hello") +} diff --git a/ratelimit_test.go b/go/ratelimit_test.go similarity index 80% rename from ratelimit_test.go rename to go/ratelimit_test.go index 60c008c..6548c7b 100644 --- a/ratelimit_test.go +++ b/go/ratelimit_test.go @@ -4,7 +4,6 @@ package ratelimit import ( "context" - "fmt" "io" "net/http" "net/http/httptest" @@ -495,7 +494,7 @@ func TestRatelimit_Decide_Good(t *testing.T) { // --- Phase 0: Sliding window / prune tests --- -func TestRatelimit_Prune_Good(t *testing.T) { +func TestRatelimit_Prune_Case(t *testing.T) { t.Run("removes old entries", func(t *testing.T) { rl := newTestLimiter(t) model := "test-prune" @@ -849,7 +848,7 @@ func TestRatelimit_WaitForCapacity_Good(t *testing.T) { }) } -func TestRatelimit_NilUsageStats_Ugly(t *testing.T) { +func TestRatelimit_NilUsageStats_Case(t *testing.T) { t.Run("CanSend replaces nil state without panicking", func(t *testing.T) { rl := newTestLimiter(t) model := "nil-cansend" @@ -1065,7 +1064,7 @@ func TestRatelimit_AllStats_Good(t *testing.T) { // --- Phase 0: Persist and Load --- -func TestRatelimit_PersistAndLoad_Ugly(t *testing.T) { +func TestRatelimit_PersistAndLoad_Case(t *testing.T) { t.Run("round-trip preserves state", func(t *testing.T) { tmpDir := t.TempDir() path := testPath(tmpDir, "ratelimits.yaml") @@ -1186,7 +1185,7 @@ func TestRatelimit_PersistAndLoad_Ugly(t *testing.T) { // --- Phase 0: Default quotas --- -func TestRatelimit_DefaultQuotas_Good(t *testing.T) { +func TestRatelimit_DefaultQuotas_Case(t *testing.T) { rl := newTestLimiter(t) tests := []struct { @@ -1206,7 +1205,7 @@ func TestRatelimit_DefaultQuotas_Good(t *testing.T) { t.Run(tt.model, func(t *testing.T) { q, ok := rl.Quotas[tt.model] if !ok { - t.Fatal(testExpectedTrueMessage(fmt.Sprintf("quota should exist for %s", tt.model))) + t.Fatal(testExpectedTrueMessage(core.Sprintf("quota should exist for %s", tt.model))) } if !testEqual(tt.maxRPM, q.MaxRPM) { t.Fatal(testWantGotMessage(tt.maxRPM, q.MaxRPM)) @@ -1223,7 +1222,7 @@ func TestRatelimit_DefaultQuotas_Good(t *testing.T) { // --- Phase 0: Concurrent access (race test) --- -func TestRatelimit_ConcurrentAccess_Good(t *testing.T) { +func TestRatelimit_ConcurrentAccess_Case(t *testing.T) { rl := newTestLimiter(t) model := "concurrent-test" rl.Quotas[model] = ModelQuota{MaxRPM: 1000, MaxTPM: 10000000, MaxRPD: 10000} @@ -1251,7 +1250,7 @@ func TestRatelimit_ConcurrentAccess_Good(t *testing.T) { } } -func TestRatelimit_ConcurrentResetAndRecord_Ugly(t *testing.T) { +func TestRatelimit_ConcurrentResetAndRecord_Case(t *testing.T) { rl := newTestLimiter(t) model := "concurrent-reset" rl.Quotas[model] = ModelQuota{MaxRPM: 10000, MaxTPM: 100000000, MaxRPD: 100000} @@ -1334,7 +1333,7 @@ func TestRatelimit_BackgroundPrune_Good(t *testing.T) { // --- Phase 0: CountTokens (with mock HTTP server) --- -func TestRatelimit_CountTokens_Ugly(t *testing.T) { +func TestRatelimit_CountTokens_ExistingCase(t *testing.T) { t.Run("successful token count", func(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if !testEqual(http.MethodPost, r.Method) { @@ -1522,7 +1521,7 @@ func TestRatelimit_CountTokens_Ugly(t *testing.T) { }) } -func TestRatelimit_PersistSkipsNilState_Good(t *testing.T) { +func TestRatelimit_PersistSkipsNilState_Case1524(t *testing.T) { path := testPath(t.TempDir(), "nil-state.yaml") rl, err := New() @@ -1548,7 +1547,7 @@ func TestRatelimit_PersistSkipsNilState_Good(t *testing.T) { } } -func TestRatelimit_TokenTotals_Good(t *testing.T) { +func TestRatelimit_TokenTotals_Case1550(t *testing.T) { maxInt := int(^uint(0) >> 1) if !testEqual(25, safeTokenSum(-100, 25)) { t.Fatal(testWantGotMessage(25, safeTokenSum(-100, 25))) @@ -1745,7 +1744,7 @@ func TestRatelimit_DefaultProfiles_Good(t *testing.T) { }) } -func TestRatelimit_NewWithConfig_Ugly(t *testing.T) { +func TestRatelimit_NewWithConfig_ExistingCase(t *testing.T) { t.Run("empty config defaults to Gemini", func(t *testing.T) { rl, err := NewWithConfig(Config{ FilePath: testPath(t.TempDir(), "test.yaml"), @@ -1930,7 +1929,7 @@ func TestRatelimit_NewWithConfig_Ugly(t *testing.T) { }) } -func TestRatelimit_NewBackwardCompatibility_Good(t *testing.T) { +func TestRatelimit_NewBackwardCompatibility_Case1932(t *testing.T) { // New() should produce the exact same result as before Phase 1 rl, err := New() if err != nil { @@ -2082,7 +2081,7 @@ func TestRatelimit_AddProvider_Good(t *testing.T) { }) } -func TestRatelimit_ProviderConstants_Good(t *testing.T) { +func TestRatelimit_ProviderConstants_Case2084(t *testing.T) { // Verify the string values are stable (they may be used in YAML configs) if !testEqual(Provider("gemini"), ProviderGemini) { t.Fatal(testWantGotMessage(Provider("gemini"), ProviderGemini)) @@ -2100,7 +2099,7 @@ func TestRatelimit_ProviderConstants_Good(t *testing.T) { // --- Phase 0 addendum: Additional concurrent and multi-model race tests --- -func TestRatelimit_ConcurrentMultipleModels_Good(t *testing.T) { +func TestRatelimit_ConcurrentMultipleModels_Case2102(t *testing.T) { rl := newTestLimiter(t) models := []string{"model-a", "model-b", "model-c", "model-d", "model-e"} for _, m := range models { @@ -2132,7 +2131,7 @@ func TestRatelimit_ConcurrentMultipleModels_Good(t *testing.T) { } } -func TestRatelimit_ConcurrentPersistAndLoad_Ugly(t *testing.T) { +func TestRatelimit_ConcurrentPersistAndLoad_Case2134(t *testing.T) { tmpDir := t.TempDir() path := testPath(tmpDir, "concurrent.yaml") @@ -2166,7 +2165,7 @@ func TestRatelimit_ConcurrentPersistAndLoad_Ugly(t *testing.T) { // No panics or data races = pass } -func TestRatelimit_ConcurrentAllStatsAndRecordUsage_Good(t *testing.T) { +func TestRatelimit_ConcurrentAllStatsAndRecordUsage_Case2168(t *testing.T) { rl := newTestLimiter(t) models := []string{"stats-a", "stats-b", "stats-c"} for _, m := range models { @@ -2197,7 +2196,7 @@ func TestRatelimit_ConcurrentAllStatsAndRecordUsage_Good(t *testing.T) { wg.Wait() } -func TestRatelimit_ConcurrentWaitForCapacityAndRecordUsage_Good(t *testing.T) { +func TestRatelimit_ConcurrentWaitForCapacityAndRecordUsage_Case2199(t *testing.T) { rl := newTestLimiter(t) model := "race-wait" rl.Quotas[model] = ModelQuota{MaxRPM: 100, MaxTPM: 10000000, MaxRPD: 10000} @@ -2315,7 +2314,7 @@ func BenchmarkPersist(b *testing.B) { } } -func TestRatelimit_EndToEndMultiProvider_Good(t *testing.T) { +func TestRatelimit_EndToEndMultiProvider_Case2317(t *testing.T) { // Simulate a real-world scenario: limiter for both Gemini and Anthropic rl, err := NewWithConfig(Config{ FilePath: testPath(t.TempDir(), "multi.yaml"), @@ -2477,7 +2476,7 @@ func TestRatelimit_Decide_Ugly(t *testing.T) { } } -func TestRatelimit_Prune_Bad(t *testing.T) { +func TestRatelimit_Prune_Case2479(t *testing.T) { rl := newTestLimiter(t) model := "prune-recent-over-quota" now := time.Unix(1_700_000_000, 0) @@ -2498,7 +2497,7 @@ func TestRatelimit_Prune_Bad(t *testing.T) { } } -func TestRatelimit_Prune_Ugly(t *testing.T) { +func TestRatelimit_Prune_Case2500(t *testing.T) { rl := newTestLimiter(t) model := "prune-future-clock-skew" now := time.Unix(1_700_000_000, 0) @@ -2797,7 +2796,7 @@ func TestRatelimit_AllStats_Ugly(t *testing.T) { wg.Go(func() { for i := range 50 { rl.RecordUsage(model, i, 1) - rl.SetQuota(fmt.Sprintf("allstats-dynamic-%d", i), ModelQuota{MaxRPM: i + 1}) + rl.SetQuota(core.Sprintf("allstats-dynamic-%d", i), ModelQuota{MaxRPM: i + 1}) } }) wg.Go(func() { @@ -2825,7 +2824,7 @@ func TestRatelimit_AllStats_Ugly(t *testing.T) { } } -func TestRatelimit_DefaultQuotas_Bad(t *testing.T) { +func TestRatelimit_DefaultQuotas_Case2827(t *testing.T) { rl, err := NewWithConfig(Config{ FilePath: testPath(t.TempDir(), "missing-provider.yaml"), Providers: []Provider{Provider("missing-provider")}, @@ -2839,7 +2838,7 @@ func TestRatelimit_DefaultQuotas_Bad(t *testing.T) { } } -func TestRatelimit_DefaultQuotas_Ugly(t *testing.T) { +func TestRatelimit_DefaultQuotas_Case2841(t *testing.T) { rl := newTestLimiter(t) model := "gemini-3-pro-preview" @@ -2866,7 +2865,7 @@ func TestRatelimit_DefaultQuotas_Ugly(t *testing.T) { } } -func TestRatelimit_ConcurrentAccess_Bad(t *testing.T) { +func TestRatelimit_ConcurrentAccess_Case2868(t *testing.T) { rl := newTestLimiter(t) model := "concurrent-denied" rl.Quotas[model] = ModelQuota{MaxRPM: 1, MaxTPM: 1000, MaxRPD: 100} @@ -2889,7 +2888,7 @@ func TestRatelimit_ConcurrentAccess_Bad(t *testing.T) { } } -func TestRatelimit_ConcurrentAccess_Ugly(t *testing.T) { +func TestRatelimit_ConcurrentAccess_Case2891(t *testing.T) { rl := newTestLimiter(t) model := "concurrent-nil-state" rl.Quotas[model] = ModelQuota{MaxRPM: 1000, MaxTPM: 1000000, MaxRPD: 1000} @@ -2937,7 +2936,7 @@ func TestRatelimit_BackgroundPrune_Ugly(t *testing.T) { stop() } -func TestRatelimit_PersistSkipsNilState_Bad(t *testing.T) { +func TestRatelimit_PersistSkipsNilState_Case2939(t *testing.T) { path := testPath(t.TempDir(), "nil-state.yaml") rl := newTestLimiter(t) rl.filePath = path @@ -2955,7 +2954,7 @@ func TestRatelimit_PersistSkipsNilState_Bad(t *testing.T) { } } -func TestRatelimit_PersistSkipsNilState_Ugly(t *testing.T) { +func TestRatelimit_PersistSkipsNilState_Case2957(t *testing.T) { path := testPath(t.TempDir(), "mixed-state.yaml") rl := newTestLimiter(t) rl.filePath = path @@ -2979,7 +2978,7 @@ func TestRatelimit_PersistSkipsNilState_Ugly(t *testing.T) { } } -func TestRatelimit_TokenTotals_Bad(t *testing.T) { +func TestRatelimit_TokenTotals_Case2981(t *testing.T) { if !testEqual(maxInt(), safeTokenSum(maxInt()-1, 1)) { t.Fatal(testWantGotMessage(maxInt(), safeTokenSum(maxInt()-1, 1), "exact max boundary should be representable")) } @@ -2992,7 +2991,7 @@ func TestRatelimit_TokenTotals_Bad(t *testing.T) { } } -func TestRatelimit_TokenTotals_Ugly(t *testing.T) { +func TestRatelimit_TokenTotals_Case2994(t *testing.T) { tokens := []TokenEntry{ {}, {Count: 0}, @@ -3033,7 +3032,7 @@ func TestRatelimit_DefaultProfiles_Ugly(t *testing.T) { } } -func TestRatelimit_NewBackwardCompatibility_Bad(t *testing.T) { +func TestRatelimit_NewBackwardCompatibility_Case3035(t *testing.T) { t.Setenv("CORE_HOME", "") t.Setenv("HOME", "") t.Setenv("home", "") @@ -3045,7 +3044,7 @@ func TestRatelimit_NewBackwardCompatibility_Bad(t *testing.T) { } } -func TestRatelimit_NewBackwardCompatibility_Ugly(t *testing.T) { +func TestRatelimit_NewBackwardCompatibility_Case3047(t *testing.T) { coreHome := t.TempDir() home := t.TempDir() t.Setenv("CORE_HOME", coreHome) @@ -3112,14 +3111,14 @@ func TestRatelimit_AddProvider_Ugly(t *testing.T) { } } -func TestRatelimit_ProviderConstants_Bad(t *testing.T) { +func TestRatelimit_ProviderConstants_Case3114(t *testing.T) { unknown := Provider("gemini ") if testEqual(ProviderGemini, unknown) { t.Fatal(testExpectedFalseMessage("provider constants should not silently normalize unknown strings")) } } -func TestRatelimit_ProviderConstants_Ugly(t *testing.T) { +func TestRatelimit_ProviderConstants_Case3121(t *testing.T) { seen := map[Provider]bool{ ProviderGemini: true, ProviderOpenAI: true, @@ -3134,7 +3133,7 @@ func TestRatelimit_ProviderConstants_Ugly(t *testing.T) { } } -func TestRatelimit_ConcurrentMultipleModels_Bad(t *testing.T) { +func TestRatelimit_ConcurrentMultipleModels_Case3136(t *testing.T) { rl := newTestLimiter(t) rl.Quotas["blocked"] = ModelQuota{MaxRPM: 1, MaxTPM: 1000, MaxRPD: 10} rl.Quotas["open"] = ModelQuota{MaxRPM: 100, MaxTPM: 1000, MaxRPD: 10} @@ -3167,7 +3166,7 @@ func TestRatelimit_ConcurrentMultipleModels_Bad(t *testing.T) { } } -func TestRatelimit_ConcurrentMultipleModels_Ugly(t *testing.T) { +func TestRatelimit_ConcurrentMultipleModels_Case3169(t *testing.T) { rl := newTestLimiter(t) models := []string{"", "model/with/slash", "model with spaces"} for _, model := range models { @@ -3193,7 +3192,7 @@ func TestRatelimit_ConcurrentMultipleModels_Ugly(t *testing.T) { } } -func TestRatelimit_ConcurrentAllStatsAndRecordUsage_Bad(t *testing.T) { +func TestRatelimit_ConcurrentAllStatsAndRecordUsage_Case3195(t *testing.T) { rl := newTestLimiter(t) model := "allstats-record-nil" rl.Quotas[model] = ModelQuota{MaxRPM: 1000, MaxTPM: 100000, MaxRPD: 1000} @@ -3219,7 +3218,7 @@ func TestRatelimit_ConcurrentAllStatsAndRecordUsage_Bad(t *testing.T) { } } -func TestRatelimit_ConcurrentAllStatsAndRecordUsage_Ugly(t *testing.T) { +func TestRatelimit_ConcurrentAllStatsAndRecordUsage_Case3221(t *testing.T) { rl := newTestLimiter(t) model := "allstats-record-reset" rl.Quotas[model] = ModelQuota{MaxRPM: 1000, MaxTPM: 100000, MaxRPD: 1000} @@ -3245,7 +3244,7 @@ func TestRatelimit_ConcurrentAllStatsAndRecordUsage_Ugly(t *testing.T) { wg.Wait() } -func TestRatelimit_ConcurrentWaitForCapacityAndRecordUsage_Bad(t *testing.T) { +func TestRatelimit_ConcurrentWaitForCapacityAndRecordUsage_Case3247(t *testing.T) { rl := newTestLimiter(t) model := "wait-record-denied" rl.Quotas[model] = ModelQuota{MaxRPM: 1, MaxTPM: 1000, MaxRPD: 100} @@ -3275,7 +3274,7 @@ func TestRatelimit_ConcurrentWaitForCapacityAndRecordUsage_Bad(t *testing.T) { } } -func TestRatelimit_ConcurrentWaitForCapacityAndRecordUsage_Ugly(t *testing.T) { +func TestRatelimit_ConcurrentWaitForCapacityAndRecordUsage_Case3277(t *testing.T) { rl := newTestLimiter(t) model := "wait-record-zero" rl.Quotas[model] = ModelQuota{MaxRPM: 1000, MaxTPM: 100000, MaxRPD: 1000} @@ -3306,7 +3305,7 @@ func TestRatelimit_ConcurrentWaitForCapacityAndRecordUsage_Ugly(t *testing.T) { } } -func TestRatelimit_EndToEndMultiProvider_Bad(t *testing.T) { +func TestRatelimit_EndToEndMultiProvider_Case3308(t *testing.T) { rl, err := NewWithConfig(Config{ FilePath: testPath(t.TempDir(), "multi-bad.yaml"), Providers: []Provider{ProviderGemini, ProviderAnthropic}, @@ -3331,7 +3330,7 @@ func TestRatelimit_EndToEndMultiProvider_Bad(t *testing.T) { } } -func TestRatelimit_EndToEndMultiProvider_Ugly(t *testing.T) { +func TestRatelimit_EndToEndMultiProvider_Case3333(t *testing.T) { path := testPath(t.TempDir(), "multi-ugly.yaml") model := "local/custom model" quota := ModelQuota{MaxRPM: 5, MaxTPM: 500, MaxRPD: 0} @@ -3404,3 +3403,694 @@ func (testRequireShim) False(tb testing.TB, value bool, msgAndArgs ...any) { tb.Fatal(testExpectedFalseMessage(msgAndArgs...)) } } + +// Public-symbol v0.9.0 triplets migrated from the old AX7 aggregate file. +type publicRoundTrip func(*http.Request) (*http.Response, error) + +func (f publicRoundTrip) RoundTrip(req *http.Request) (*http.Response, error) { + return f(req) +} + +func publicYAMLPath(t *core.T) string { + t.Helper() + return core.Path(t.TempDir(), "ratelimits.yaml") +} + +func publicLimiter(t *core.T) *RateLimiter { + t.Helper() + rl, err := NewWithConfig(Config{ + FilePath: publicYAMLPath(t), + Providers: []Provider{ProviderLocal}, + }) + core.RequireNoError(t, err) + return rl +} + +func TestRatelimit_New_Good(t *core.T) { + rl, err := New() + + core.RequireNoError(t, err) + core.AssertTrue(t, rl.CanSend("gemini-3-pro-preview", 1)) +} + +func TestRatelimit_New_Bad(t *core.T) { + t.Setenv("CORE_HOME", "") + t.Setenv("HOME", "") + t.Setenv("home", "") + t.Setenv("USERPROFILE", "") + + rl, err := New() + core.AssertError(t, err) + core.AssertNil(t, rl) +} + +func TestRatelimit_New_Ugly(t *core.T) { + rl1, err := New() + core.RequireNoError(t, err) + rl2, err := New() + core.RequireNoError(t, err) + + rl1.RecordUsage("gemini-3-pro-preview", 1, 1) + core.AssertEqual(t, 1, rl1.Stats("gemini-3-pro-preview").RPD) + core.AssertEqual(t, 0, rl2.Stats("gemini-3-pro-preview").RPD) +} + +func TestRatelimit_NewWithConfig_Good(t *core.T) { + rl, err := NewWithConfig(Config{ + FilePath: publicYAMLPath(t), + Providers: []Provider{ProviderOpenAI}, + Quotas: map[string]ModelQuota{ + "custom": {MaxRPM: 7, MaxTPM: 70, MaxRPD: 700}, + }, + }) + + core.RequireNoError(t, err) + core.AssertEqual(t, 500, rl.Stats("gpt-4o").MaxRPM) + core.AssertEqual(t, 7, rl.Stats("custom").MaxRPM) +} + +func TestRatelimit_NewWithConfig_Bad(t *core.T) { + rl, err := NewWithConfig(Config{ + FilePath: publicYAMLPath(t), + Backend: "bogus", + }) + + core.AssertError(t, err) + core.AssertNil(t, rl) +} + +func TestRatelimit_NewWithConfig_Ugly(t *core.T) { + rl, err := NewWithConfig(Config{FilePath: publicYAMLPath(t)}) + + core.RequireNoError(t, err) + core.AssertEqual(t, 150, rl.Stats("gemini-3-pro-preview").MaxRPM) + core.AssertEqual(t, 0, rl.Stats("missing-model").MaxRPM) +} + +func TestRatelimit_RateLimiter_SetQuota_Good(t *core.T) { + rl := publicLimiter(t) + rl.SetQuota("model-a", ModelQuota{MaxRPM: 3, MaxTPM: 30, MaxRPD: 300}) + + stats := rl.Stats("model-a") + core.AssertEqual(t, 3, stats.MaxRPM) + core.AssertEqual(t, 30, stats.MaxTPM) + core.AssertEqual(t, 300, stats.MaxRPD) +} + +func TestRatelimit_RateLimiter_SetQuota_Bad(t *core.T) { + rl := publicLimiter(t) + rl.SetQuota("model-a", ModelQuota{MaxRPM: 1}) + rl.SetQuota("model-a", ModelQuota{}) + + decision := rl.Decide("model-a", 999999) + core.AssertTrue(t, decision.Allowed) + core.AssertEqual(t, DecisionUnlimited, decision.Code) +} + +func TestRatelimit_RateLimiter_SetQuota_Ugly(t *core.T) { + rl := publicLimiter(t) + rl.SetQuota("", ModelQuota{MaxRPM: 1}) + + core.AssertEqual(t, 1, rl.Stats("").MaxRPM) + core.AssertFalse(t, rl.CanSend("", -1)) +} + +func TestRatelimit_RateLimiter_AddProvider_Good(t *core.T) { + rl := publicLimiter(t) + rl.AddProvider(ProviderAnthropic) + + core.AssertEqual(t, 50, rl.Stats("claude-opus-4").MaxRPM) + core.AssertEqual(t, 40000, rl.Stats("claude-opus-4").MaxTPM) +} + +func TestRatelimit_RateLimiter_AddProvider_Bad(t *core.T) { + rl := publicLimiter(t) + rl.AddProvider(Provider("unknown")) + + models := make([]string, 0) + for model := range rl.Models() { + models = append(models, model) + } + core.AssertEmpty(t, models) +} + +func TestRatelimit_RateLimiter_AddProvider_Ugly(t *core.T) { + rl := publicLimiter(t) + rl.SetQuota("local-model", ModelQuota{MaxRPM: 2}) + rl.AddProvider(ProviderLocal) + + core.AssertEqual(t, 2, rl.Stats("local-model").MaxRPM) + core.AssertTrue(t, rl.CanSend("unknown-local", 1)) +} + +func TestRatelimit_RateLimiter_Load_Good(t *core.T) { + path := publicYAMLPath(t) + rl, err := NewWithConfig(Config{ + FilePath: path, + Quotas: map[string]ModelQuota{"model-a": {MaxRPM: 5}}, + }) + core.RequireNoError(t, err) + rl.RecordUsage("model-a", 1, 2) + core.RequireNoError(t, rl.Persist()) + + loaded, err := NewWithConfig(Config{FilePath: path, Providers: []Provider{ProviderLocal}}) + core.RequireNoError(t, err) + core.AssertNoError(t, loaded.Load()) + core.AssertEqual(t, 1, loaded.Stats("model-a").RPD) +} + +func TestRatelimit_RateLimiter_Load_Bad(t *core.T) { + path := publicYAMLPath(t) + core.RequireTrue(t, core.WriteFile(path, []byte("{{{not yaml"), 0o600).OK) + rl, err := NewWithConfig(Config{FilePath: path, Providers: []Provider{ProviderLocal}}) + core.RequireNoError(t, err) + + err = rl.Load() + core.AssertError(t, err) + core.AssertContains(t, err.Error(), "yaml") +} + +func TestRatelimit_RateLimiter_Load_Ugly(t *core.T) { + rl, err := NewWithConfig(Config{ + FilePath: core.Path(t.TempDir(), "missing.yaml"), + Providers: []Provider{ProviderLocal}, + }) + core.RequireNoError(t, err) + + core.AssertNoError(t, rl.Load()) + core.AssertEmpty(t, rl.AllStats()) +} + +func TestRatelimit_RateLimiter_Persist_Good(t *core.T) { + path := publicYAMLPath(t) + rl, err := NewWithConfig(Config{FilePath: path, Quotas: map[string]ModelQuota{"model-a": {MaxRPM: 2}}}) + core.RequireNoError(t, err) + rl.RecordUsage("model-a", 1, 1) + + core.AssertNoError(t, rl.Persist()) + stat := core.Stat(path) + core.AssertTrue(t, stat.OK) +} + +func TestRatelimit_RateLimiter_Persist_Bad(t *core.T) { + rl, err := NewWithConfig(Config{ + FilePath: t.TempDir(), + Providers: []Provider{ProviderLocal}, + }) + core.RequireNoError(t, err) + rl.RecordUsage("model-a", 1, 1) + + err = rl.Persist() + core.AssertError(t, err) +} + +func TestRatelimit_RateLimiter_Persist_Ugly(t *core.T) { + path := publicYAMLPath(t) + rl, err := NewWithConfig(Config{FilePath: path, Providers: []Provider{ProviderLocal}}) + core.RequireNoError(t, err) + + core.AssertNoError(t, rl.Persist()) + loaded, err := NewWithConfig(Config{FilePath: path, Providers: []Provider{ProviderLocal}}) + core.RequireNoError(t, err) + core.AssertNoError(t, loaded.Load()) +} + +func TestRatelimit_RateLimiter_BackgroundPrune_Good(t *core.T) { + rl := publicLimiter(t) + stop := rl.BackgroundPrune(10 * time.Millisecond) + time.Sleep(20 * time.Millisecond) + + core.AssertNotPanics(t, stop) + core.AssertEmpty(t, rl.AllStats()) +} + +func TestRatelimit_RateLimiter_BackgroundPrune_Bad(t *core.T) { + rl := publicLimiter(t) + stop := rl.BackgroundPrune(0) + rl.RecordUsage("model-a", 1, 1) + + core.AssertNotPanics(t, stop) + core.AssertEqual(t, 1, rl.Stats("model-a").RPD) +} + +func TestRatelimit_RateLimiter_BackgroundPrune_Ugly(t *core.T) { + rl := publicLimiter(t) + stop := rl.BackgroundPrune(1 * time.Millisecond) + + core.AssertNotPanics(t, stop) + core.AssertNotPanics(t, stop) +} + +func TestRatelimit_RateLimiter_CanSend_Good(t *core.T) { + rl := publicLimiter(t) + rl.SetQuota("model-a", ModelQuota{MaxRPM: 2, MaxTPM: 100, MaxRPD: 5}) + + core.AssertTrue(t, rl.CanSend("model-a", 10)) + core.AssertEqual(t, DecisionAllowed, rl.Decide("model-a", 10).Code) +} + +func TestRatelimit_RateLimiter_CanSend_Bad(t *core.T) { + rl := publicLimiter(t) + rl.SetQuota("model-a", ModelQuota{MaxRPM: 1, MaxTPM: 100, MaxRPD: 5}) + rl.RecordUsage("model-a", 1, 1) + + core.AssertFalse(t, rl.CanSend("model-a", 1)) + core.AssertEqual(t, DecisionRPMLimit, rl.Decide("model-a", 1).Code) +} + +func TestRatelimit_RateLimiter_CanSend_Ugly(t *core.T) { + rl := publicLimiter(t) + rl.SetQuota("model-a", ModelQuota{MaxRPM: 1, MaxTPM: 100, MaxRPD: 5}) + + core.AssertFalse(t, rl.CanSend("model-a", -1)) + core.AssertTrue(t, rl.CanSend("unknown-model", 999999)) +} + +func TestRatelimit_RateLimiter_RecordUsage_Good(t *core.T) { + rl := publicLimiter(t) + rl.RecordUsage("model-a", 10, 15) + + stats := rl.Stats("model-a") + core.AssertEqual(t, 1, stats.RPD) + core.AssertEqual(t, 25, stats.TPM) +} + +func TestRatelimit_RateLimiter_RecordUsage_Bad(t *core.T) { + rl := publicLimiter(t) + rl.RecordUsage("model-a", -10, 15) + + stats := rl.Stats("model-a") + core.AssertEqual(t, 1, stats.RPD) + core.AssertEqual(t, 15, stats.TPM) +} + +func TestRatelimit_RateLimiter_RecordUsage_Ugly(t *core.T) { + rl := publicLimiter(t) + rl.RecordUsage("", 0, 0) + + stats := rl.Stats("") + core.AssertEqual(t, 1, stats.RPD) + core.AssertEqual(t, 0, stats.TPM) +} + +func TestRatelimit_RateLimiter_WaitForCapacity_Good(t *core.T) { + rl := publicLimiter(t) + rl.SetQuota("model-a", ModelQuota{MaxRPM: 1, MaxTPM: 100, MaxRPD: 5}) + + err := rl.WaitForCapacity(context.Background(), "model-a", 1) + core.AssertNoError(t, err) +} + +func TestRatelimit_RateLimiter_WaitForCapacity_Bad(t *core.T) { + rl := publicLimiter(t) + err := rl.WaitForCapacity(context.Background(), "model-a", -1) + + core.AssertError(t, err) + core.AssertContains(t, err.Error(), "negative tokens") +} + +func TestRatelimit_RateLimiter_WaitForCapacity_Ugly(t *core.T) { + rl := publicLimiter(t) + rl.SetQuota("model-a", ModelQuota{MaxRPM: 1, MaxTPM: 100, MaxRPD: 5}) + rl.RecordUsage("model-a", 1, 1) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + err := rl.WaitForCapacity(ctx, "model-a", 1) + core.AssertErrorIs(t, err, context.Canceled) +} + +func TestRatelimit_RateLimiter_Reset_Good(t *core.T) { + rl := publicLimiter(t) + rl.RecordUsage("model-a", 1, 1) + rl.Reset("model-a") + + stats := rl.Stats("model-a") + core.AssertEqual(t, 0, stats.RPD) + core.AssertEqual(t, 0, stats.TPM) +} + +func TestRatelimit_RateLimiter_Reset_Bad(t *core.T) { + rl := publicLimiter(t) + rl.RecordUsage("model-a", 1, 1) + rl.Reset("missing-model") + + core.AssertEqual(t, 1, rl.Stats("model-a").RPD) + core.AssertEqual(t, 0, rl.Stats("missing-model").RPD) +} + +func TestRatelimit_RateLimiter_Reset_Ugly(t *core.T) { + rl := publicLimiter(t) + rl.RecordUsage("model-a", 1, 1) + rl.RecordUsage("model-b", 1, 1) + rl.Reset("") + + core.AssertEqual(t, 0, rl.Stats("model-a").RPD) + core.AssertEqual(t, 0, rl.Stats("model-b").RPD) +} + +func TestRatelimit_RateLimiter_Models_Good(t *core.T) { + rl, err := NewWithConfig(Config{ + FilePath: publicYAMLPath(t), + Quotas: map[string]ModelQuota{ + "model-b": {MaxRPM: 2}, + "model-a": {MaxRPM: 1}, + }, + }) + core.RequireNoError(t, err) + + var models []string + for model := range rl.Models() { + models = append(models, model) + } + core.AssertEqual(t, []string{"model-a", "model-b"}, models) +} + +func TestRatelimit_RateLimiter_Models_Bad(t *core.T) { + rl := publicLimiter(t) + var models []string + for model := range rl.Models() { + models = append(models, model) + } + + core.AssertEmpty(t, models) + core.AssertEqual(t, 0, rl.Stats("missing").MaxRPM) +} + +func TestRatelimit_RateLimiter_Models_Ugly(t *core.T) { + rl, err := NewWithConfig(Config{ + FilePath: publicYAMLPath(t), + Quotas: map[string]ModelQuota{ + "model-a": {MaxRPM: 1}, + "model-b": {MaxRPM: 2}, + "model-c": {MaxRPM: 3}, + }, + }) + core.RequireNoError(t, err) + + var first string + for model := range rl.Models() { + first = model + break + } + core.AssertEqual(t, "model-a", first) +} + +func TestRatelimit_RateLimiter_Iter_Good(t *core.T) { + rl, err := NewWithConfig(Config{ + FilePath: publicYAMLPath(t), + Quotas: map[string]ModelQuota{ + "model-b": {MaxRPM: 2}, + "model-a": {MaxRPM: 1}, + }, + }) + core.RequireNoError(t, err) + rl.RecordUsage("model-a", 1, 1) + + var models []string + for model, stats := range rl.Iter() { + models = append(models, model) + if model == "model-a" { + core.AssertEqual(t, 1, stats.RPD) + } + } + core.AssertEqual(t, []string{"model-a", "model-b"}, models) +} + +func TestRatelimit_RateLimiter_Iter_Bad(t *core.T) { + rl := publicLimiter(t) + var models []string + for model := range rl.Iter() { + models = append(models, model) + } + + core.AssertEmpty(t, models) + core.AssertEmpty(t, rl.AllStats()) +} + +func TestRatelimit_RateLimiter_Iter_Ugly(t *core.T) { + rl, err := NewWithConfig(Config{ + FilePath: publicYAMLPath(t), + Quotas: map[string]ModelQuota{ + "model-a": {MaxRPM: 1}, + "model-b": {MaxRPM: 2}, + "model-c": {MaxRPM: 3}, + }, + }) + core.RequireNoError(t, err) + + var seen []string + for model := range rl.Iter() { + seen = append(seen, model) + if len(seen) == 2 { + break + } + } + core.AssertEqual(t, []string{"model-a", "model-b"}, seen) +} + +func TestRatelimit_RateLimiter_Stats_Good(t *core.T) { + rl := publicLimiter(t) + rl.SetQuota("model-a", ModelQuota{MaxRPM: 10, MaxTPM: 100, MaxRPD: 5}) + rl.RecordUsage("model-a", 2, 3) + + stats := rl.Stats("model-a") + core.AssertEqual(t, 1, stats.RPM) + core.AssertEqual(t, 5, stats.TPM) + core.AssertEqual(t, 10, stats.MaxRPM) +} + +func TestRatelimit_RateLimiter_Stats_Bad(t *core.T) { + rl := publicLimiter(t) + stats := rl.Stats("missing-model") + + core.AssertEqual(t, 0, stats.RPM) + core.AssertEqual(t, 0, stats.MaxRPM) + core.AssertTrue(t, stats.DayStart.IsZero()) +} + +func TestRatelimit_RateLimiter_Stats_Ugly(t *core.T) { + rl := publicLimiter(t) + rl.RecordUsage("model-a", 1, 1) + rl.Reset("model-a") + + stats := rl.Stats("model-a") + core.AssertEqual(t, 0, stats.RPD) + core.AssertEqual(t, 0, stats.TPM) +} + +func TestRatelimit_RateLimiter_AllStats_Good(t *core.T) { + rl := publicLimiter(t) + rl.SetQuota("model-a", ModelQuota{MaxRPM: 10}) + rl.RecordUsage("model-b", 1, 1) + + all := rl.AllStats() + core.AssertContains(t, all, "model-a") + core.AssertContains(t, all, "model-b") + core.AssertEqual(t, 2, all["model-b"].TPM) +} + +func TestRatelimit_RateLimiter_AllStats_Bad(t *core.T) { + rl := publicLimiter(t) + all := rl.AllStats() + + core.AssertEmpty(t, all) + core.AssertEqual(t, 0, len(all)) +} + +func TestRatelimit_RateLimiter_AllStats_Ugly(t *core.T) { + rl := publicLimiter(t) + rl.SetQuota("model-a", ModelQuota{MaxRPM: 10}) + all := rl.AllStats() + + core.AssertContains(t, all, "model-a") + core.AssertEqual(t, 0, all["model-a"].RPM) + core.AssertEqual(t, 10, all["model-a"].MaxRPM) +} + +func TestRatelimit_RateLimiter_Decide_Good(t *core.T) { + rl := publicLimiter(t) + rl.SetQuota("model-a", ModelQuota{MaxRPM: 2, MaxTPM: 100, MaxRPD: 5}) + + decision := rl.Decide("model-a", 10) + core.AssertTrue(t, decision.Allowed) + core.AssertEqual(t, DecisionAllowed, decision.Code) +} + +func TestRatelimit_RateLimiter_Decide_Bad(t *core.T) { + rl := publicLimiter(t) + decision := rl.Decide("model-a", -1) + + core.AssertFalse(t, decision.Allowed) + core.AssertEqual(t, DecisionInvalidTokens, decision.Code) + core.AssertContains(t, decision.Reason, "non-negative") +} + +func TestRatelimit_RateLimiter_Decide_Ugly(t *core.T) { + rl := publicLimiter(t) + rl.SetQuota("model-a", ModelQuota{MaxRPM: 1, MaxTPM: 100, MaxRPD: 5}) + rl.RecordUsage("model-a", 1, 1) + + decision := rl.Decide("model-a", 1) + core.AssertFalse(t, decision.Allowed) + core.AssertEqual(t, DecisionRPMLimit, decision.Code) +} + +func TestRatelimit_NewWithSQLite_Good(t *core.T) { + rl, err := NewWithSQLite(core.Path(t.TempDir(), "limits.db")) + + core.RequireNoError(t, err) + defer rl.Close() + rl.RecordUsage("gemini-3-pro-preview", 1, 1) + core.AssertNoError(t, rl.Persist()) +} + +func TestRatelimit_NewWithSQLite_Bad(t *core.T) { + rl, err := NewWithSQLite(core.Path(t.TempDir(), "missing", "limits.db")) + + core.AssertError(t, err) + core.AssertNil(t, rl) +} + +func TestRatelimit_NewWithSQLite_Ugly(t *core.T) { + path := core.Path(t.TempDir(), "limits with spaces.db") + rl, err := NewWithSQLite(path) + core.RequireNoError(t, err) + rl.RecordUsage("gemini-3-pro-preview", 1, 1) + + core.AssertNoError(t, rl.Persist()) + core.AssertNoError(t, rl.Close()) +} + +func TestRatelimit_NewWithSQLiteConfig_Good(t *core.T) { + rl, err := NewWithSQLiteConfig(core.Path(t.TempDir(), "limits.db"), Config{ + Providers: []Provider{ProviderOpenAI}, + Quotas: map[string]ModelQuota{"custom": {MaxRPM: 9}}, + }) + + core.RequireNoError(t, err) + defer rl.Close() + core.AssertEqual(t, 500, rl.Stats("gpt-4o").MaxRPM) + core.AssertEqual(t, 9, rl.Stats("custom").MaxRPM) +} + +func TestRatelimit_NewWithSQLiteConfig_Bad(t *core.T) { + rl, err := NewWithSQLiteConfig(core.Path(t.TempDir(), "missing", "limits.db"), Config{}) + + core.AssertError(t, err) + core.AssertNil(t, rl) +} + +func TestRatelimit_NewWithSQLiteConfig_Ugly(t *core.T) { + rl, err := NewWithSQLiteConfig(core.Path(t.TempDir(), "limits.db"), Config{ + Backend: "yaml", + Providers: []Provider{ProviderLocal}, + }) + + core.RequireNoError(t, err) + defer rl.Close() + rl.RecordUsage("local", 1, 1) + core.AssertNoError(t, rl.Persist()) +} + +func TestRatelimit_RateLimiter_Close_Good(t *core.T) { + rl := publicLimiter(t) + err := rl.Close() + + core.AssertNoError(t, err) + core.AssertNoError(t, rl.Close()) +} + +func TestRatelimit_RateLimiter_Close_Bad(t *core.T) { + rl, err := NewWithSQLite(core.Path(t.TempDir(), "limits.db")) + core.RequireNoError(t, err) + core.RequireNoError(t, rl.Close()) + + err = rl.Persist() + core.AssertError(t, err) +} + +func TestRatelimit_RateLimiter_Close_Ugly(t *core.T) { + rl, err := NewWithSQLite(core.Path(t.TempDir(), "limits.db")) + core.RequireNoError(t, err) + + core.AssertNoError(t, rl.Close()) + core.AssertNoError(t, rl.Close()) +} + +func TestRatelimit_MigrateYAMLToSQLite_Good(t *core.T) { + yamlPath := publicYAMLPath(t) + sqlitePath := core.Path(t.TempDir(), "limits.db") + rl, err := NewWithConfig(Config{ + FilePath: yamlPath, + Quotas: map[string]ModelQuota{"model-a": {MaxRPM: 4}}, + }) + core.RequireNoError(t, err) + rl.RecordUsage("model-a", 2, 3) + core.RequireNoError(t, rl.Persist()) + + core.AssertNoError(t, MigrateYAMLToSQLite(yamlPath, sqlitePath)) + loaded, err := NewWithSQLite(sqlitePath) + core.RequireNoError(t, err) + defer loaded.Close() + core.AssertNoError(t, loaded.Load()) + core.AssertEqual(t, 1, loaded.Stats("model-a").RPD) +} + +func TestRatelimit_MigrateYAMLToSQLite_Bad(t *core.T) { + err := MigrateYAMLToSQLite(core.Path(t.TempDir(), "missing.yaml"), core.Path(t.TempDir(), "limits.db")) + + core.AssertError(t, err) + core.AssertContains(t, err.Error(), "read") +} + +func TestRatelimit_MigrateYAMLToSQLite_Ugly(t *core.T) { + yamlPath := publicYAMLPath(t) + sqlitePath := core.Path(t.TempDir(), "limits.db") + core.RequireTrue(t, core.WriteFile(yamlPath, []byte("{{{not yaml"), 0o600).OK) + + err := MigrateYAMLToSQLite(yamlPath, sqlitePath) + core.AssertError(t, err) + core.AssertContains(t, err.Error(), "unmarshal") +} + +func TestRatelimit_CountTokens_Good(t *core.T) { + oldTransport := http.DefaultTransport + http.DefaultTransport = publicRoundTrip(func(req *http.Request) (*http.Response, error) { + core.AssertContains(t, req.URL.Path, "/v1beta/models/gemini-3-pro-preview:countTokens") + return &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(core.NewReader(`{"totalTokens":7}`)), + Header: make(http.Header), + }, nil + }) + defer func() { http.DefaultTransport = oldTransport }() + + tokens, err := CountTokens(context.Background(), "key", "gemini-3-pro-preview", "hello") + core.AssertNoError(t, err) + core.AssertEqual(t, 7, tokens) +} + +func TestRatelimit_CountTokens_Bad(t *core.T) { + tokens, err := CountTokens(context.Background(), "key", "", "hello") + + core.AssertError(t, err) + core.AssertEqual(t, 0, tokens) + core.AssertContains(t, err.Error(), "empty model") +} + +func TestRatelimit_CountTokens_Ugly(t *core.T) { + oldTransport := http.DefaultTransport + http.DefaultTransport = publicRoundTrip(func(req *http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(core.NewReader(`{"totalTokens":"bad"}`)), + Header: make(http.Header), + }, nil + }) + defer func() { http.DefaultTransport = oldTransport }() + + tokens, err := CountTokens(context.Background(), "key", "gemini-3-pro-preview", "hello") + core.AssertError(t, err) + core.AssertEqual(t, 0, tokens) +} diff --git a/sqlite.go b/go/sqlite.go similarity index 93% rename from sqlite.go rename to go/sqlite.go index e2ad994..5036d59 100644 --- a/sqlite.go +++ b/go/sqlite.go @@ -18,7 +18,7 @@ type sqliteStore struct { // newSQLiteStore opens (or creates) a SQLite database at dbPath and initialises // the schema. It follows the go-store pattern: single connection, WAL journal // mode, and a 5-second busy timeout for contention handling. -func newSQLiteStore(dbPath string) (*sqliteStore, error) { +func newSQLiteStore(dbPath string) (*sqliteStore, error) /* core result boundary */ { db, err := sql.Open("sqlite", dbPath) if err != nil { return nil, core.E("ratelimit.newSQLiteStore", "open", err) @@ -45,7 +45,7 @@ func newSQLiteStore(dbPath string) (*sqliteStore, error) { } // createSchema creates the tables and indices if they do not already exist. -func createSchema(db *sql.DB) error { +func createSchema(db *sql.DB) error /* core result boundary */ { stmts := []string{ `CREATE TABLE IF NOT EXISTS quotas ( model TEXT PRIMARY KEY, @@ -80,7 +80,7 @@ func createSchema(db *sql.DB) error { } // saveQuotas writes a complete quota snapshot to the quotas table. -func (s *sqliteStore) saveQuotas(quotas map[string]ModelQuota) error { +func (s *sqliteStore) saveQuotas(quotas map[string]ModelQuota) error /* core result boundary */ { tx, err := s.db.Begin() if err != nil { return core.E("ratelimit.saveQuotas", "begin", err) @@ -99,7 +99,7 @@ func (s *sqliteStore) saveQuotas(quotas map[string]ModelQuota) error { } // loadQuotas reads all rows from the quotas table. -func (s *sqliteStore) loadQuotas() (map[string]ModelQuota, error) { +func (s *sqliteStore) loadQuotas() (map[string]ModelQuota, error) /* core result boundary */ { rows, err := s.db.Query("SELECT model, max_rpm, max_tpm, max_rpd FROM quotas") if err != nil { return nil, core.E("ratelimit.loadQuotas", "query", err) @@ -122,7 +122,7 @@ func (s *sqliteStore) loadQuotas() (map[string]ModelQuota, error) { } // saveSnapshot writes quotas and state as a single atomic snapshot. -func (s *sqliteStore) saveSnapshot(quotas map[string]ModelQuota, state map[string]*UsageStats) error { +func (s *sqliteStore) saveSnapshot(quotas map[string]ModelQuota, state map[string]*UsageStats) error /* core result boundary */ { tx, err := s.db.Begin() if err != nil { return core.E("ratelimit.saveSnapshot", "begin", err) @@ -146,7 +146,7 @@ func (s *sqliteStore) saveSnapshot(quotas map[string]ModelQuota, state map[strin // saveState writes all usage state to SQLite in a single transaction. // It uses a truncate-and-insert approach for simplicity in this version, // but ensures atomicity via a single transaction. -func (s *sqliteStore) saveState(state map[string]*UsageStats) error { +func (s *sqliteStore) saveState(state map[string]*UsageStats) error /* core result boundary */ { tx, err := s.db.Begin() if err != nil { return core.E("ratelimit.saveState", "begin", err) @@ -164,7 +164,7 @@ func (s *sqliteStore) saveState(state map[string]*UsageStats) error { return commitTx(tx, "ratelimit.saveState") } -func clearSnapshotTables(tx *sql.Tx, includeQuotas bool) error { +func clearSnapshotTables(tx *sql.Tx, includeQuotas bool) error /* core result boundary */ { if includeQuotas { if _, err := tx.Exec("DELETE FROM quotas"); err != nil { return core.E("ratelimit.saveSnapshot", "clear quotas", err) @@ -182,7 +182,7 @@ func clearSnapshotTables(tx *sql.Tx, includeQuotas bool) error { return nil } -func insertQuotas(tx *sql.Tx, quotas map[string]ModelQuota) error { +func insertQuotas(tx *sql.Tx, quotas map[string]ModelQuota) error /* core result boundary */ { stmt, err := tx.Prepare("INSERT INTO quotas (model, max_rpm, max_tpm, max_rpd) VALUES (?, ?, ?, ?)") if err != nil { return core.E("ratelimit.saveQuotas", "prepare", err) @@ -197,7 +197,7 @@ func insertQuotas(tx *sql.Tx, quotas map[string]ModelQuota) error { return nil } -func insertState(tx *sql.Tx, state map[string]*UsageStats) error { +func insertState(tx *sql.Tx, state map[string]*UsageStats) error /* core result boundary */ { reqStmt, err := tx.Prepare("INSERT INTO requests (model, ts) VALUES (?, ?)") if err != nil { return core.E("ratelimit.saveState", "prepare requests", err) @@ -237,7 +237,7 @@ func insertState(tx *sql.Tx, state map[string]*UsageStats) error { return nil } -func commitTx(tx *sql.Tx, scope string) error { +func commitTx(tx *sql.Tx, scope string) error /* core result boundary */ { if err := tx.Commit(); err != nil { return core.E(scope, "commit", err) } @@ -245,7 +245,7 @@ func commitTx(tx *sql.Tx, scope string) error { } // loadState reconstructs the UsageStats map from SQLite tables. -func (s *sqliteStore) loadState() (map[string]*UsageStats, error) { +func (s *sqliteStore) loadState() (map[string]*UsageStats, error) /* core result boundary */ { result := make(map[string]*UsageStats) // Load daily counters first (these define which models have state). @@ -323,6 +323,6 @@ func (s *sqliteStore) loadState() (map[string]*UsageStats, error) { } // close closes the underlying database connection. -func (s *sqliteStore) close() error { +func (s *sqliteStore) close() error /* core result boundary */ { return s.db.Close() } diff --git a/sqlite_test.go b/go/sqlite_test.go similarity index 95% rename from sqlite_test.go rename to go/sqlite_test.go index b4c6ca2..39901cf 100644 --- a/sqlite_test.go +++ b/go/sqlite_test.go @@ -3,7 +3,7 @@ package ratelimit import ( - "fmt" + core "dappco.re/go" "sync" "testing" "time" @@ -13,7 +13,7 @@ import ( // --- Phase 2: SQLite basic tests --- -func TestSQLite_NewSQLiteStore_Good(t *testing.T) { +func TestSQLite_NewSQLiteStore_Case(t *testing.T) { dbPath := testPath(t.TempDir(), "test.db") store, err := newSQLiteStore(dbPath) if err != nil { @@ -27,7 +27,7 @@ func TestSQLite_NewSQLiteStore_Good(t *testing.T) { } } -func TestSQLite_NewSQLiteStore_Bad(t *testing.T) { +func TestSQLite_NewSQLiteStore_Case_2(t *testing.T) { t.Run("invalid path returns error", func(t *testing.T) { // Path inside a non-existent directory with no parent. _, err := newSQLiteStore("/nonexistent/deep/nested/dir/test.db") @@ -37,7 +37,7 @@ func TestSQLite_NewSQLiteStore_Bad(t *testing.T) { }) } -func TestSQLite_QuotasRoundTrip_Good(t *testing.T) { +func TestSQLite_QuotasRoundTrip_Case(t *testing.T) { dbPath := testPath(t.TempDir(), "quotas.db") store, err := newSQLiteStore(dbPath) if err != nil { @@ -64,7 +64,7 @@ func TestSQLite_QuotasRoundTrip_Good(t *testing.T) { for model, expected := range quotas { actual, ok := loaded[model] if !ok { - t.Fatal(testExpectedTrueMessage(fmt.Sprintf("loaded quotas should contain %s", model))) + t.Fatal(testExpectedTrueMessage(core.Sprintf("loaded quotas should contain %s", model))) } if !testEqual(expected.MaxRPM, actual.MaxRPM) { t.Fatal(testWantGotMessage(expected.MaxRPM, actual.MaxRPM)) @@ -78,7 +78,7 @@ func TestSQLite_QuotasRoundTrip_Good(t *testing.T) { } } -func TestSQLite_QuotasOverwrite_Good(t *testing.T) { +func TestSQLite_QuotasOverwrite_Case(t *testing.T) { dbPath := testPath(t.TempDir(), "overwrite.db") store, err := newSQLiteStore(dbPath) if err != nil { @@ -119,7 +119,7 @@ func TestSQLite_QuotasOverwrite_Good(t *testing.T) { } } -func TestSQLite_StateRoundTrip_Good(t *testing.T) { +func TestSQLite_StateRoundTrip_Case(t *testing.T) { dbPath := testPath(t.TempDir(), "state.db") store, err := newSQLiteStore(dbPath) if err != nil { @@ -166,40 +166,40 @@ func TestSQLite_StateRoundTrip_Good(t *testing.T) { for model, expected := range state { actual, ok := loaded[model] if !ok { - t.Fatal(testExpectedTrueMessage(fmt.Sprintf("loaded state should contain %s", model))) + t.Fatal(testExpectedTrueMessage(core.Sprintf("loaded state should contain %s", model))) } if !testHasLen(actual.Requests, len(expected.Requests)) { - t.Fatal(testLenMessage(actual.Requests, len(expected.Requests), fmt.Sprintf("request count for %s", model))) + t.Fatal(testLenMessage(actual.Requests, len(expected.Requests), core.Sprintf("request count for %s", model))) } if !testHasLen(actual.Tokens, len(expected.Tokens)) { - t.Fatal(testLenMessage(actual.Tokens, len(expected.Tokens), fmt.Sprintf("token count for %s", model))) + t.Fatal(testLenMessage(actual.Tokens, len(expected.Tokens), core.Sprintf("token count for %s", model))) } if !testEqual(expected.DayCount, actual.DayCount) { - t.Fatal(testWantGotMessage(expected.DayCount, actual.DayCount, fmt.Sprintf("day count for %s", model))) + t.Fatal(testWantGotMessage(expected.DayCount, actual.DayCount, core.Sprintf("day count for %s", model))) } // Time comparison with nanosecond precision (UnixNano round-trip). if !testEqual(expected.DayStart.UnixNano(), actual.DayStart.UnixNano()) { - t.Fatal(testWantGotMessage(expected.DayStart.UnixNano(), actual.DayStart.UnixNano(), fmt.Sprintf("day start for %s", model))) + t.Fatal(testWantGotMessage(expected.DayStart.UnixNano(), actual.DayStart.UnixNano(), core.Sprintf("day start for %s", model))) } for i, req := range expected.Requests { if !testEqual(req.UnixNano(), actual.Requests[i].UnixNano()) { - t.Fatal(testWantGotMessage(req.UnixNano(), actual.Requests[i].UnixNano(), fmt.Sprintf("request %d for %s", i, model))) + t.Fatal(testWantGotMessage(req.UnixNano(), actual.Requests[i].UnixNano(), core.Sprintf("request %d for %s", i, model))) } } for i, tok := range expected.Tokens { if !testEqual(tok.Time.UnixNano(), actual.Tokens[i].Time.UnixNano()) { - t.Fatal(testWantGotMessage(tok.Time.UnixNano(), actual.Tokens[i].Time.UnixNano(), fmt.Sprintf("token time %d for %s", i, model))) + t.Fatal(testWantGotMessage(tok.Time.UnixNano(), actual.Tokens[i].Time.UnixNano(), core.Sprintf("token time %d for %s", i, model))) } if !testEqual(tok.Count, actual.Tokens[i].Count) { - t.Fatal(testWantGotMessage(tok.Count, actual.Tokens[i].Count, fmt.Sprintf("token count %d for %s", i, model))) + t.Fatal(testWantGotMessage(tok.Count, actual.Tokens[i].Count, core.Sprintf("token count %d for %s", i, model))) } } } } -func TestSQLite_StateOverwrite_Good(t *testing.T) { +func TestSQLite_StateOverwrite_Case(t *testing.T) { dbPath := testPath(t.TempDir(), "overwrite.db") store, err := newSQLiteStore(dbPath) if err != nil { @@ -255,7 +255,7 @@ func TestSQLite_StateOverwrite_Good(t *testing.T) { } } -func TestSQLite_EmptyState_Good(t *testing.T) { +func TestSQLite_EmptyState_Case(t *testing.T) { dbPath := testPath(t.TempDir(), "empty.db") store, err := newSQLiteStore(dbPath) if err != nil { @@ -281,7 +281,7 @@ func TestSQLite_EmptyState_Good(t *testing.T) { } } -func TestSQLite_Close_Good(t *testing.T) { +func TestSQLite_Close_Case(t *testing.T) { dbPath := testPath(t.TempDir(), "close.db") store, err := newSQLiteStore(dbPath) if err != nil { @@ -344,7 +344,7 @@ func TestSQLite_NewWithSQLiteConfig_Good(t *testing.T) { } } -func TestSQLite_PersistAndLoad_Good(t *testing.T) { +func TestSQLite_PersistAndLoad_Case(t *testing.T) { dbPath := testPath(t.TempDir(), "persist.db") rl, err := NewWithSQLite(dbPath) if err != nil { @@ -393,7 +393,7 @@ func TestSQLite_PersistAndLoad_Good(t *testing.T) { } } -func TestSQLite_PersistMultipleModels_Good(t *testing.T) { +func TestSQLite_PersistMultipleModels_Case(t *testing.T) { dbPath := testPath(t.TempDir(), "multi.db") rl, err := NewWithSQLiteConfig(dbPath, Config{ Providers: []Provider{ProviderGemini, ProviderAnthropic}, @@ -439,7 +439,7 @@ func TestSQLite_PersistMultipleModels_Good(t *testing.T) { } } -func TestSQLite_RecordUsageThenPersistReload_Good(t *testing.T) { +func TestSQLite_RecordUsageThenPersistReload_Case(t *testing.T) { dbPath := testPath(t.TempDir(), "record.db") rl, err := NewWithSQLite(dbPath) if err != nil { @@ -496,7 +496,7 @@ func TestSQLite_RecordUsageThenPersistReload_Good(t *testing.T) { } } -func TestSQLite_CloseNoOp_Good(t *testing.T) { +func TestSQLite_CloseNoOp_Case(t *testing.T) { // Close on YAML-backed limiter is a no-op. rl := newTestLimiter(t) if err := rl.Close(); err != nil { @@ -571,7 +571,7 @@ func TestSQLite_Concurrent_Good(t *testing.T) { // --- Phase 2: YAML backward compatibility --- -func TestSQLite_YAMLBackwardCompat_Good(t *testing.T) { +func TestSQLite_YAMLBackwardCompat_Case(t *testing.T) { // Verify that the default YAML backend still works after SQLite additions. tmpDir := t.TempDir() path := testPath(tmpDir, "compat.yaml") @@ -611,7 +611,7 @@ func TestSQLite_YAMLBackwardCompat_Good(t *testing.T) { } } -func TestSQLite_ConfigBackendDefault_Good(t *testing.T) { +func TestSQLite_ConfigBackendDefault_Case(t *testing.T) { // Empty Backend string should default to YAML behaviour. rl, err := NewWithConfig(Config{ FilePath: testPath(t.TempDir(), "default.yaml"), @@ -624,7 +624,7 @@ func TestSQLite_ConfigBackendDefault_Good(t *testing.T) { } } -func TestSQLite_ConfigBackendSQLite_Good(t *testing.T) { +func TestSQLite_ConfigBackendSQLite_Case(t *testing.T) { dbPath := testPath(t.TempDir(), "config-backend.db") rl, err := NewWithConfig(Config{ Backend: backendSQLite, @@ -649,7 +649,7 @@ func TestSQLite_ConfigBackendSQLite_Good(t *testing.T) { } } -func TestSQLite_ConfigBackendSQLiteDefaultPath_Good(t *testing.T) { +func TestSQLite_ConfigBackendSQLiteDefaultPath_Case(t *testing.T) { home := t.TempDir() t.Setenv("HOME", home) t.Setenv("USERPROFILE", "") @@ -882,10 +882,10 @@ func TestSQLite_MigrateYAMLToSQLitePreservesAllGeminiModels_Good(t *testing.T) { for model := range rl.Quotas { q, ok := rl2.Quotas[model] if !ok { - t.Fatal(testExpectedTrueMessage(fmt.Sprintf("migrated quota should exist for %s", model))) + t.Fatal(testExpectedTrueMessage(core.Sprintf("migrated quota should exist for %s", model))) } if !testEqual(rl.Quotas[model], q) { - t.Fatal(testWantGotMessage(rl.Quotas[model], q, fmt.Sprintf("quota values should match for %s", model))) + t.Fatal(testWantGotMessage(rl.Quotas[model], q, core.Sprintf("quota values should match for %s", model))) } } } diff --git a/tests/cli/ratelimit/Taskfile.yaml b/go/tests/cli/ratelimit/Taskfile.yaml similarity index 100% rename from tests/cli/ratelimit/Taskfile.yaml rename to go/tests/cli/ratelimit/Taskfile.yaml diff --git a/go/tests/cli/ratelimit/main.go b/go/tests/cli/ratelimit/main.go new file mode 100644 index 0000000..ab34018 --- /dev/null +++ b/go/tests/cli/ratelimit/main.go @@ -0,0 +1,68 @@ +// AX-10 CLI driver for go-ratelimit. Exercises the public RateLimiter API +// without depending on the package's own test files. +// +// task -d tests/cli/ratelimit test +// go run ./tests/cli/ratelimit +package main + +import ( + "time" + + core "dappco.re/go" + "dappco.re/go/ratelimit" +) + +func main() { + if r := run(); !r.OK { + core.Print(core.Stderr(), "%s\n", r.Error()) + core.Exit(1) + } +} + +func run() core.Result { + dirResult := core.MkdirTemp("", "go-ratelimit-ax10-") + if !dirResult.OK { + return core.Fail(core.E("ratelimit.cli", "mkdtemp", resultFailure(dirResult))) + } + dir := dirResult.Value.(string) + defer core.RemoveAll(dir) + + cfg := ratelimit.Config{ + FilePath: core.Path(dir, "ratelimits.yaml"), + Providers: []ratelimit.Provider{ratelimit.ProviderGemini}, + } + + rl, err := ratelimit.NewWithConfig(cfg) + if err != nil { + return core.Fail(core.E("ratelimit.cli", "new ratelimiter", err)) + } + + allowed := rl.CanSend("gemini-2.5-flash", 100) + decision := rl.Decide("gemini-2.5-flash", 100) + out := map[string]any{ + "allowed": allowed, + "code": string(decision.Code), + "reason": decision.Reason, + "retryAfterMs": decision.RetryAfter / time.Millisecond, + } + encoded := core.JSONMarshalIndent(out, "", " ") + if !encoded.OK { + return core.Fail(core.E("ratelimit.cli", "encode decision", resultFailure(encoded))) + } + if write := core.WriteString(core.Stdout(), string(encoded.Value.([]byte))+"\n"); !write.OK { + return write + } + + if err := rl.Persist(); err != nil { + return core.Fail(core.E("ratelimit.cli", "persist", err)) + } + + return core.Ok(nil) +} + +func resultFailure(r core.Result) error /* core result boundary */ { + if err, ok := r.Value.(error); ok { + return err + } + return core.E("ratelimit.cli", r.Error(), nil) +} diff --git a/tests/cli/ratelimit/main.go b/tests/cli/ratelimit/main.go deleted file mode 100644 index 4323f53..0000000 --- a/tests/cli/ratelimit/main.go +++ /dev/null @@ -1,61 +0,0 @@ -// AX-10 CLI driver for go-ratelimit. Exercises the public RateLimiter API -// without depending on the package's own test files. -// -// task -d tests/cli/ratelimit test -// go run ./tests/cli/ratelimit -package main - -import ( - "encoding/json" - "fmt" - "os" - "path/filepath" - "time" - - "dappco.re/go/ratelimit" -) - -func main() { - if err := run(); err != nil { - fmt.Fprintln(os.Stderr, err) - os.Exit(1) - } -} - -func run() error { - dir, err := os.MkdirTemp("", "go-ratelimit-ax10-") - if err != nil { - return fmt.Errorf("mkdtemp: %w", err) - } - defer os.RemoveAll(dir) - - cfg := ratelimit.Config{ - FilePath: filepath.Join(dir, "ratelimits.yaml"), - Providers: []ratelimit.Provider{ratelimit.ProviderGemini}, - } - - rl, err := ratelimit.NewWithConfig(cfg) - if err != nil { - return fmt.Errorf("new ratelimiter: %w", err) - } - - allowed := rl.CanSend("gemini-2.5-flash", 100) - decision := rl.Decide("gemini-2.5-flash", 100) - out := map[string]any{ - "allowed": allowed, - "code": string(decision.Code), - "reason": decision.Reason, - "retryAfterMs": decision.RetryAfter / time.Millisecond, - } - enc := json.NewEncoder(os.Stdout) - enc.SetIndent("", " ") - if err := enc.Encode(out); err != nil { - return fmt.Errorf("encode decision: %w", err) - } - - if err := rl.Persist(); err != nil { - return fmt.Errorf("persist: %w", err) - } - - return nil -}