diff --git a/internal/gateway/plugin_test.go b/internal/gateway/plugin_test.go index 2180d368..44208a05 100644 --- a/internal/gateway/plugin_test.go +++ b/internal/gateway/plugin_test.go @@ -8,6 +8,7 @@ import ( "net/http/httptest" "os/exec" "path/filepath" + "slices" "strings" "sync" "testing" @@ -165,6 +166,14 @@ func TestPluginProvider(t *testing.T) { mu.Lock() limited = "Bearer k1" mu.Unlock() + var staled []string + staleAllowance = func(agent, user string) { + mu.Lock() + staled = append(staled, agent) + mu.Unlock() + provider.StaleAllowance(agent, user) + } + defer func() { staleAllowance = provider.StaleAllowance }() code, body := postAs(t, s, "", `{"model":"fakeco/fake-1","messages":[{"role":"user","content":"hello"}]}`) mu.Lock() last := got[len(got)-1] @@ -172,6 +181,23 @@ func TestPluginProvider(t *testing.T) { if code != 200 || !strings.Contains(body, "hi from chat") || last.auth != "Bearer k2" { t.Fatalf("with the first account limited: %d %s, upstream saw %+v", code, body, last) } + // the limited account's allowance is asked again as its usage is kept, + // and it rests as that account + mu.Lock() + asked := slices.Clone(staled) + mu.Unlock() + if !slices.Contains(asked, "plugin:fakeco") { + t.Fatalf("the limited account's allowance went stale as %q, not plugin:fakeco", asked) + } + restingUntil.Lock() + var rests []string + for _, r := range restingUntil.note { + rests = append(rests, r.agent) + } + restingUntil.Unlock() + if !slices.Contains(rests, "plugin:fakeco") { + t.Fatalf("the limited account rests as %q, not plugin:fakeco", rests) + } if err := plugin.SignOut(ctx, "fakeco", ""); err != nil { t.Fatal(err) diff --git a/internal/gateway/routing.go b/internal/gateway/routing.go index 48bcf311..83026862 100644 --- a/internal/gateway/routing.go +++ b/internal/gateway/routing.go @@ -207,6 +207,9 @@ func renewed(agent, user string) { func init() { provider.OnRenewed(renewed) } +// staleAllowance is provider.StaleAllowance, swapped in tests. +var staleAllowance = provider.StaleAllowance + // Unrest lifts the rest of what rests by key — an account just verified // with its vendor, say — so the next request asks it again. False when it // wasn't resting. @@ -320,12 +323,14 @@ func (s *Server) restAfterMarked(c candidate, status int, header http.Header, bo d, r.By = t.Sub(now), "window" } } + // an account is kept by the agent its usage is asked of: a plugin's + // by the plugin's provider ("plugin:grok"), not by "plugin" if a := c.p.Account; a != nil && why != failOther && why != failVerify { - provider.StaleAllowance(a.Agent, a.User) // ask again what it has left + staleAllowance(a.UsageAgent(), a.User) // ask again what it has left } r.Until = now.Add(d) if a := c.p.Account; a != nil { - r.agent, r.user = a.Agent, a.User + r.agent, r.user = a.UsageAgent(), a.User } id := c.restKey() // OpenRouter identifies a provider's shared pool separately from its diff --git a/internal/provider/routing.go b/internal/provider/routing.go index 95b61ab7..fbd907b7 100644 --- a/internal/provider/routing.go +++ b/internal/provider/routing.go @@ -187,6 +187,8 @@ func StaleAllowance(agent, user string) { loginUsageCache.Lock() delete(loginUsageCache.m, agent+"/"+strings.ToLower(user)) loginUsageCache.Unlock() + // the built-in keeps Grok's usage by home; a Grok moved to its plugin + // keeps it as "plugin:grok"'s, the line above if agent == "grok" { gs := grokLogins() grokHomeUsage.Lock()