mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-21 06:56:58 +00:00
feat(ollama): async rate-limit reset for Ollama Cloud usage windows
Real Ollama Cloud 429s are driven by ollama.com usage windows, not by response headers that OpenAI/Anthropic/CN rate-limit parsing understands. Handle them separately: apply a never-shortening temporary cooldown first, then schedule an async probe that learns the true usage-window reset point with a controlled, coalesced per-api_key refresh. - ratelimit_service: route real Ollama Cloud usage accounts to a dedicated 429 path before the CN-provider branch. - ollama_cloud_usage: add a background probe coordinator (started in Start, waited by Stop) that fans out model-429 probes to one controlled refresh per api_key group, reusing the recent group snapshot for coalesced members. - repository/account_repo: add SetRateLimitedIfUnchanged, a write-back CAS counterpart to ClearRateLimitIfObserved, so an async reset never overwrites a newer 429, an admin clear, a re-armed generation, or a key/state change observed by another writer between read and write. - wire: inject OllamaCloudUsageService into RateLimitService. Tests: unit tests for service and repository pass; targeted race tests pass; build and vet clean; git diff --check clean. The CAS integration test compiles but is not run (no Docker available).
This commit is contained in:
alfadb
@@ -117,10 +117,12 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
openAI403CounterCache := repository.NewOpenAI403CounterCache(redisClient)
|
||||
geminiTokenCache := repository.NewGeminiTokenCache(redisClient)
|
||||
compositeTokenCacheInvalidator := service.NewCompositeTokenCacheInvalidator(geminiTokenCache)
|
||||
rateLimitService := service.ProvideRateLimitService(accountRepository, usageLogRepository, configConfig, geminiQuotaService, tempUnschedCache, timeoutCounterCache, openAI403CounterCache, settingService, compositeTokenCacheInvalidator)
|
||||
httpUpstream := repository.NewHTTPUpstream(configConfig)
|
||||
leaderLockCache := repository.NewLeaderLockCache(redisClient)
|
||||
ollamaCloudUsageService := service.ProvideOllamaCloudUsageService(accountRepository, httpUpstream, settingService, secretEncryptor, configConfig, leaderLockCache, db)
|
||||
rateLimitService := service.ProvideRateLimitService(accountRepository, usageLogRepository, configConfig, geminiQuotaService, tempUnschedCache, timeoutCounterCache, openAI403CounterCache, settingService, compositeTokenCacheInvalidator, ollamaCloudUsageService)
|
||||
identityCache := repository.NewIdentityCache(redisClient)
|
||||
identityService := service.NewIdentityService(identityCache)
|
||||
httpUpstream := repository.NewHTTPUpstream(configConfig)
|
||||
timingWheelService, err := service.ProvideTimingWheelService()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -183,7 +185,6 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
dashboardAggregationRepository := repository.NewDashboardAggregationRepository(db)
|
||||
dashboardStatsCache := repository.NewDashboardCache(redisClient, configConfig)
|
||||
dashboardService := service.NewDashboardService(usageLogRepository, dashboardAggregationRepository, dashboardStatsCache, configConfig)
|
||||
leaderLockCache := repository.NewLeaderLockCache(redisClient)
|
||||
dashboardAggregationService := service.ProvideDashboardAggregationService(dashboardAggregationRepository, timingWheelService, leaderLockCache, db, configConfig)
|
||||
dashboardHandler := admin.NewDashboardHandler(dashboardService, dashboardAggregationService)
|
||||
adminGroupRepository := repository.NewAdminGroupRepository(client, db)
|
||||
@@ -284,7 +285,6 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
auditLogService := service.ProvideAuditLogService(auditLogRepository, settingService)
|
||||
auditLogHandler := admin.NewAuditLogHandler(auditLogService, totpService)
|
||||
upstreamBillingProbeService := service.ProvideUpstreamBillingProbeService(accountRepository, accountTestService, settingService, leaderLockCache, db)
|
||||
ollamaCloudUsageService := service.ProvideOllamaCloudUsageService(accountRepository, httpUpstream, settingService, secretEncryptor, configConfig, leaderLockCache, db)
|
||||
adminHandlers := handler.ProvideAdminHandlers(dashboardHandler, adminUserHandler, groupHandler, accountHandler, adminAnnouncementHandler, dataManagementHandler, backupHandler, oAuthHandler, openAIOAuthHandler, geminiOAuthHandler, antigravityOAuthHandler, grokOAuthHandler, cnProviderHandler, proxyHandler, adminRedeemHandler, promoHandler, settingHandler, opsHandler, systemHandler, adminSubscriptionHandler, adminUsageHandler, userAttributeHandler, errorPassthroughHandler, tlsFingerprintProfileHandler, pluginHandler, adminAPIKeyHandler, scheduledTestHandler, channelHandler, channelMonitorHandler, channelMonitorRequestTemplateHandler, contentModerationHandler, promptAdminHandler, paymentHandler, affiliateHandler, complianceHandler, auditLogHandler, upstreamBillingProbeService, ollamaCloudUsageService)
|
||||
usageRecordWorkerPool := service.NewUsageRecordWorkerPool(configConfig)
|
||||
userMsgQueueCache := repository.NewUserMsgQueueCache(redisClient)
|
||||
|
||||
@@ -2257,6 +2257,61 @@ func (r *accountRepository) ClearRateLimitIfObserved(ctx context.Context, id int
|
||||
return true, nil
|
||||
}
|
||||
|
||||
// SetRateLimitedIfUnchanged atomically applies a rate-limit reset only while the
|
||||
// account still carries exactly the generation the caller observed: its
|
||||
// UpdatedAt row version, its RateLimitedAt and its RateLimitResetAt (nil means
|
||||
// that field is currently unset). It is the write-back CAS counterpart to
|
||||
// ClearRateLimitIfObserved: an async rate-limit reset (e.g. an Ollama Cloud
|
||||
// usage probe) must not overwrite a newer 429, an admin clear, a re-armed
|
||||
// generation, or a key/state change observed by another writer between the
|
||||
// caller's read and this write. The whole update is a single statement, so the
|
||||
// write itself is race-free. updated reports whether the write happened, and the
|
||||
// caller must ONLY send its scheduling notification when updated == true (this
|
||||
// method already performed the DB update; no further SetRateLimited call is
|
||||
// allowed, as a second unconditional write would reintroduce the race). No new
|
||||
// migration is required.
|
||||
func (r *accountRepository) SetRateLimitedIfUnchanged(
|
||||
ctx context.Context,
|
||||
id int64,
|
||||
expectedUpdatedAt time.Time,
|
||||
expectedLimitedAt, expectedResetAt *time.Time,
|
||||
newResetAt time.Time,
|
||||
) (bool, error) {
|
||||
preds := []dbpredicate.Account{dbaccount.IDEQ(id), dbaccount.UpdatedAtEQ(expectedUpdatedAt)}
|
||||
if expectedLimitedAt == nil {
|
||||
preds = append(preds, dbaccount.RateLimitedAtIsNil())
|
||||
} else {
|
||||
preds = append(preds, dbaccount.RateLimitedAtEQ(*expectedLimitedAt))
|
||||
}
|
||||
if expectedResetAt == nil {
|
||||
preds = append(preds, dbaccount.RateLimitResetAtIsNil())
|
||||
} else {
|
||||
preds = append(preds, dbaccount.RateLimitResetAtEQ(*expectedResetAt))
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
updated, err := r.client.Account.Update().
|
||||
Where(preds...).
|
||||
SetRateLimitedAt(now).
|
||||
SetRateLimitResetAt(newResetAt).
|
||||
Save(ctx)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
if updated == 0 {
|
||||
// The generation changed concurrently (cleared, re-armed, or the account
|
||||
// was otherwise updated elsewhere): do not announce anything, just
|
||||
// refresh the local scheduler snapshot.
|
||||
r.syncSchedulerAccountSnapshot(ctx, id)
|
||||
return false, nil
|
||||
}
|
||||
if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
|
||||
logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue rate limit failed: account=%d err=%v", id, err)
|
||||
}
|
||||
r.syncSchedulerAccountSnapshot(ctx, id)
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func (r *accountRepository) SetModelRateLimit(ctx context.Context, id int64, scope string, resetAt time.Time, reason ...string) error {
|
||||
if scope == "" {
|
||||
return nil
|
||||
|
||||
@@ -0,0 +1,96 @@
|
||||
//go:build integration
|
||||
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// TestAccountRepositorySetRateLimitedIfUnchanged exercises the atomic generation
|
||||
// CAS backing the Ollama Cloud usage probe write-back (SetRateLimitedIfUnchanged).
|
||||
// It must only write when the row still carries the exact UpdatedAt /
|
||||
// RateLimitedAt / RateLimitResetAt the caller observed, so a stale async reset can
|
||||
// never overwrite a newer 429, an admin clear, or a re-armed generation.
|
||||
func TestAccountRepositorySetRateLimitedIfUnchanged(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testEntTx(t)
|
||||
repo := newAccountRepositoryWithSQL(tx.Client(), tx, nil)
|
||||
|
||||
makeAcct := func(name string) *service.Account {
|
||||
return mustCreateAccount(t, tx.Client(), &service.Account{
|
||||
Name: name, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey,
|
||||
Credentials: map[string]any{"base_url": "https://www.ollama.com", "api_key": name},
|
||||
})
|
||||
}
|
||||
|
||||
// --- Matching generation applies the new reset. ---
|
||||
acct := makeAcct("cas-match")
|
||||
base := time.Now().Add(5 * time.Second)
|
||||
require.NoError(t, repo.SetRateLimited(ctx, acct.ID, base))
|
||||
|
||||
cur, err := repo.GetByID(ctx, acct.ID)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, cur)
|
||||
require.NotNil(t, cur.RateLimitedAt)
|
||||
require.NotNil(t, cur.RateLimitResetAt)
|
||||
expUpdated, expLimited, expReset := cur.UpdatedAt, cur.RateLimitedAt, cur.RateLimitResetAt
|
||||
|
||||
newReset := time.Now().Add(2 * time.Hour)
|
||||
updated, err := repo.SetRateLimitedIfUnchanged(ctx, acct.ID, expUpdated, expLimited, expReset, newReset)
|
||||
require.NoError(t, err)
|
||||
require.True(t, updated, "a write matching the observed generation must apply")
|
||||
|
||||
after, err := repo.GetByID(ctx, acct.ID)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, after.RateLimitResetAt)
|
||||
require.WithinDuration(t, newReset, *after.RateLimitResetAt, 2*time.Second)
|
||||
|
||||
// --- Stale UpdatedAt (row version moved by the write above) must NOT apply. ---
|
||||
stale := time.Now().Add(3 * time.Hour)
|
||||
updated, err = repo.SetRateLimitedIfUnchanged(ctx, acct.ID, expUpdated, expLimited, expReset, stale)
|
||||
require.NoError(t, err)
|
||||
require.False(t, updated, "a write whose row version moved on must not apply")
|
||||
|
||||
// --- Current UpdatedAt but OLD limited/reset must NOT apply: the rate-limit
|
||||
// generation was re-armed by the successful write, so the CAS must reject
|
||||
// purely on the limited/reset predicates (isolated from UpdatedAt). ---
|
||||
updated, err = repo.SetRateLimitedIfUnchanged(ctx, acct.ID, after.UpdatedAt, expLimited, expReset, stale)
|
||||
require.NoError(t, err)
|
||||
require.False(t, updated, "re-armed limited/reset generation must not accept the old one even with the current row version")
|
||||
|
||||
// --- Current UpdatedAt but nil limited/reset must NOT apply against a set
|
||||
// generation (admin-clear expectation only matches an unset state). ---
|
||||
clearedReset := time.Now().Add(30 * time.Minute)
|
||||
updated, err = repo.SetRateLimitedIfUnchanged(ctx, acct.ID, after.UpdatedAt, nil, nil, clearedReset)
|
||||
require.NoError(t, err)
|
||||
require.False(t, updated, "nil-limited expectation must not match a set generation")
|
||||
|
||||
// --- Stale UpdatedAt alone (with nil/nil) must also fail on the version. ---
|
||||
updated, err = repo.SetRateLimitedIfUnchanged(ctx, acct.ID, expUpdated, nil, nil, clearedReset)
|
||||
require.NoError(t, err)
|
||||
require.False(t, updated, "stale row version must not apply even with nil limited/reset")
|
||||
|
||||
// --- Never-limited account: nil generation CAS applies. ---
|
||||
fresh := makeAcct("cas-fresh")
|
||||
freshCur, err := repo.GetByID(ctx, fresh.ID)
|
||||
require.NoError(t, err)
|
||||
require.Nil(t, freshCur.RateLimitedAt)
|
||||
require.Nil(t, freshCur.RateLimitResetAt)
|
||||
firstReset := time.Now().Add(time.Hour)
|
||||
updated, err = repo.SetRateLimitedIfUnchanged(ctx, fresh.ID, freshCur.UpdatedAt, nil, nil, firstReset)
|
||||
require.NoError(t, err)
|
||||
require.True(t, updated, "nil generation CAS must apply to a never-limited account")
|
||||
|
||||
// Admin-clear style: clear, then the old (set) generation must no longer match.
|
||||
require.NoError(t, repo.ClearRateLimit(ctx, fresh.ID))
|
||||
afterClear, err := repo.GetByID(ctx, fresh.ID)
|
||||
require.NoError(t, err)
|
||||
updated, err = repo.SetRateLimitedIfUnchanged(ctx, fresh.ID, afterClear.UpdatedAt, nil, nil, time.Now().Add(time.Hour))
|
||||
require.NoError(t, err)
|
||||
require.True(t, updated, "cleared account with nil generation accepts a fresh write")
|
||||
}
|
||||
@@ -393,6 +393,16 @@ type OllamaCloudUsageService struct {
|
||||
lockCache LeaderLockCache
|
||||
db *sql.DB
|
||||
instanceID string
|
||||
|
||||
// Event-triggered rate-limit probe coordination. A single background
|
||||
// coordinator (started in Start, waited by Stop) fans out model-429 probe
|
||||
// requests to one controlled refresh per api_key group, reusing the recent
|
||||
// group snapshot for coalesced members so each account still gets its own
|
||||
// result without a second upstream fetch. See ollama_cloud_usage_rate_limit_probe.go.
|
||||
probeMu sync.Mutex
|
||||
probeQueue []ollamaCloudUsageProbeRequest
|
||||
probeWake chan struct{}
|
||||
probeGroups map[string]ollamaCloudUsageProbeGroupEntry
|
||||
}
|
||||
|
||||
func NewOllamaCloudUsageService(
|
||||
@@ -414,6 +424,8 @@ func NewOllamaCloudUsageService(
|
||||
refreshSlots: make(chan struct{}, ollamaCloudUsageConcurrency),
|
||||
now: time.Now,
|
||||
instanceID: uuid.NewString(),
|
||||
probeWake: make(chan struct{}, 1),
|
||||
probeGroups: make(map[string]ollamaCloudUsageProbeGroupEntry),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -444,9 +456,10 @@ func (s *OllamaCloudUsageService) Start() {
|
||||
return
|
||||
}
|
||||
s.started = true
|
||||
s.wg.Add(1)
|
||||
s.wg.Add(2)
|
||||
s.mu.Unlock()
|
||||
go s.runLoop()
|
||||
go s.probeLoop()
|
||||
}
|
||||
|
||||
func (s *OllamaCloudUsageService) Stop() {
|
||||
|
||||
@@ -0,0 +1,385 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
)
|
||||
|
||||
// This file adds the event-triggered rate-limit probe to OllamaCloudUsageService:
|
||||
// a model-429 handler calls ScheduleOllamaCloudUsageRateLimitProbe to prompt a
|
||||
// controlled usage refresh now (instead of waiting for the regular cycle) and to
|
||||
// learn the recovery horizon when the account's usage is actually exhausted.
|
||||
//
|
||||
// The probe reuses refreshAccount wholesale (singleflight, bounded concurrency
|
||||
// slot, group resolution, settings-HTML fetch); no HTTP/parsing is duplicated.
|
||||
// Work runs on a single background coordinator owned by Start/Stop with a
|
||||
// detached, timeout-bounded context, so a model request never blocks on a probe
|
||||
// and no unbounded goroutines are introduced.
|
||||
|
||||
const (
|
||||
// ollamaCloudUsageProbeTimeout bounds one detached probe task. It must exceed
|
||||
// the underlying settings-HTML request timeout.
|
||||
ollamaCloudUsageProbeTimeout = 45 * time.Second
|
||||
|
||||
// ollamaCloudUsageProbeMaxQueue bounds queued (not yet serviced) probe
|
||||
// requests. A duplicate request for an already-queued account still merges
|
||||
// into the latest callback even when the queue is full.
|
||||
ollamaCloudUsageProbeMaxQueue = 256
|
||||
|
||||
// ollamaCloudUsageProbeGroupRetention bounds how long finished probe results
|
||||
// are retained for coalescing before being pruned.
|
||||
ollamaCloudUsageProbeGroupRetention = 2 * ollamaCloudUsageManualRefreshInterval
|
||||
)
|
||||
|
||||
// OllamaCloudUsageRateLimitProbeCallback is invoked (on the coordinator goroutine)
|
||||
// when an account's Ollama Cloud usage is confirmed exhausted from a fresh
|
||||
// successful snapshot; resetAt is the earliest wall-clock rollover of the
|
||||
// exhausted window(s). It is never invoked for accounts that are not configured
|
||||
// Ollama usage accounts, lack a web session, are in failure backoff, or whose most
|
||||
// recent snapshot is not exhausted.
|
||||
//
|
||||
// The callback runs on the single coordinator, so it must not block indefinitely:
|
||||
// the next serialized probe (and Stop) would be held back. It MAY do bounded
|
||||
// persistence — e.g. a limited DB write-back under its own timeout context — and
|
||||
// must not spawn unbounded goroutines. Stop waits the coordinator, so it may wait
|
||||
// out an in-flight callback's own bounded write-back timeout (shorter than the
|
||||
// probe task timeout) rather than a strict probe timeout; keep any persistence
|
||||
// bounded so Stop returns promptly. The 429-wiring author is responsible for how
|
||||
// the callback maps to a rate-limit write-back.
|
||||
type OllamaCloudUsageRateLimitProbeCallback func(accountID int64, resetAt time.Time)
|
||||
|
||||
// ollamaCloudUsageProbeRequest is one queued, per-account probe request.
|
||||
type ollamaCloudUsageProbeRequest struct {
|
||||
accountID int64
|
||||
onExhausted OllamaCloudUsageRateLimitProbeCallback
|
||||
}
|
||||
|
||||
// ollamaCloudUsageProbeGroupEntry is the most recent probe fetch outcome for an
|
||||
// api_key group, kept briefly so coalesced members reuse the same upstream result
|
||||
// instead of each issuing its own fetch.
|
||||
type ollamaCloudUsageProbeGroupEntry struct {
|
||||
attemptAt time.Time
|
||||
snapshot *OllamaCloudUsageSnapshot
|
||||
}
|
||||
|
||||
// ScheduleOllamaCloudUsageRateLimitProbe schedules an asynchronous probe of the
|
||||
// account's Ollama Cloud usage and reports whether the request was accepted. It
|
||||
// never blocks on network or DB work: it only records the request for the
|
||||
// background coordinator and returns, so it is safe to call from a model-429 path.
|
||||
// The probe runs under a detached, timeout-bounded context derived from the
|
||||
// service lifecycle.
|
||||
//
|
||||
// Accepted == true means the request was queued, not that a result will follow:
|
||||
// an ineligible account quietly yields no callback. Accepted == false means the
|
||||
// service is not running (not started / stopped) or the queue is full for a brand
|
||||
// new account; a duplicate for an already-queued account is always merged.
|
||||
func (s *OllamaCloudUsageService) ScheduleOllamaCloudUsageRateLimitProbe(
|
||||
accountID int64,
|
||||
onExhausted OllamaCloudUsageRateLimitProbeCallback,
|
||||
) bool {
|
||||
if s == nil || onExhausted == nil || accountID <= 0 {
|
||||
return false
|
||||
}
|
||||
s.mu.Lock()
|
||||
running := s.started && !s.stopped
|
||||
s.mu.Unlock()
|
||||
if !running {
|
||||
return false
|
||||
}
|
||||
|
||||
s.probeMu.Lock()
|
||||
// Coalesce a duplicate for the same account first (keep the latest callback)
|
||||
// so a full queue still accepts a merge.
|
||||
for i := range s.probeQueue {
|
||||
if s.probeQueue[i].accountID == accountID {
|
||||
s.probeQueue[i].onExhausted = onExhausted
|
||||
s.probeMu.Unlock()
|
||||
s.wakeProbeLoop()
|
||||
return true
|
||||
}
|
||||
}
|
||||
if len(s.probeQueue) >= ollamaCloudUsageProbeMaxQueue {
|
||||
s.probeMu.Unlock()
|
||||
return false
|
||||
}
|
||||
s.probeQueue = append(s.probeQueue, ollamaCloudUsageProbeRequest{
|
||||
accountID: accountID,
|
||||
onExhausted: onExhausted,
|
||||
})
|
||||
s.probeMu.Unlock()
|
||||
|
||||
s.wakeProbeLoop()
|
||||
return true
|
||||
}
|
||||
|
||||
// wakeProbeLoop nudges the coordinator; a dropped wake is safe because the
|
||||
// coordinator re-checks the queue after every wake.
|
||||
func (s *OllamaCloudUsageService) wakeProbeLoop() {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
select {
|
||||
case s.probeWake <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
// probeLoop is the single background worker. It is added to the service WaitGroup
|
||||
// in Start and returns when parentCtx is cancelled in Stop, so Stop waits for any
|
||||
// in-flight probe (including a bounded callback) to finish under its own timeout.
|
||||
func (s *OllamaCloudUsageService) probeLoop() {
|
||||
defer s.wg.Done()
|
||||
for {
|
||||
s.probeMu.Lock()
|
||||
if len(s.probeQueue) == 0 {
|
||||
s.probeMu.Unlock()
|
||||
select {
|
||||
case <-s.probeWake:
|
||||
case <-s.parentCtx.Done():
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
batch := s.probeQueue
|
||||
s.probeQueue = nil
|
||||
s.probeMu.Unlock()
|
||||
|
||||
for _, req := range batch {
|
||||
if s.parentCtx.Err() != nil {
|
||||
return
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(s.parentCtx, ollamaCloudUsageProbeTimeout)
|
||||
s.runOllamaCloudUsageProbe(ctx, req.accountID, req.onExhausted)
|
||||
cancel()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *OllamaCloudUsageService) probeGroupResult(key string) (ollamaCloudUsageProbeGroupEntry, bool) {
|
||||
if s == nil {
|
||||
return ollamaCloudUsageProbeGroupEntry{}, false
|
||||
}
|
||||
s.probeMu.Lock()
|
||||
defer s.probeMu.Unlock()
|
||||
entry, ok := s.probeGroups[key]
|
||||
return entry, ok
|
||||
}
|
||||
|
||||
func (s *OllamaCloudUsageService) storeProbeGroupResult(key string, entry ollamaCloudUsageProbeGroupEntry) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.probeMu.Lock()
|
||||
defer s.probeMu.Unlock()
|
||||
if s.probeGroups == nil {
|
||||
s.probeGroups = make(map[string]ollamaCloudUsageProbeGroupEntry)
|
||||
}
|
||||
pruneAt := entry.attemptAt.Add(-ollamaCloudUsageProbeGroupRetention)
|
||||
for groupKey, existing := range s.probeGroups {
|
||||
if existing.attemptAt.Before(pruneAt) {
|
||||
delete(s.probeGroups, groupKey)
|
||||
}
|
||||
}
|
||||
s.probeGroups[key] = entry
|
||||
}
|
||||
|
||||
// runOllamaCloudUsageProbe resolves the account and, when appropriate, performs a
|
||||
// controlled group refresh and reports exhaustion. All decisions are best-effort;
|
||||
// ineligible/missing-session accounts, transient failures, backoff and
|
||||
// non-exhaustion are skipped silently.
|
||||
//
|
||||
// This is a 429-event recovery query: it deliberately runs regardless of the
|
||||
// periodic auto_refresh switch (a model was just refused, so a one-off read is
|
||||
// warranted even when auto-refresh is off). It still respects the presence of a
|
||||
// configured session/cookie and any scrape-429 failure backoff.
|
||||
func (s *OllamaCloudUsageService) runOllamaCloudUsageProbe(
|
||||
ctx context.Context,
|
||||
accountID int64,
|
||||
onExhausted OllamaCloudUsageRateLimitProbeCallback,
|
||||
) {
|
||||
if s == nil || s.accountRepo == nil || onExhausted == nil || ctx.Err() != nil {
|
||||
return
|
||||
}
|
||||
account, err := s.accountRepo.GetByID(ctx, accountID)
|
||||
if err != nil || account == nil {
|
||||
return
|
||||
}
|
||||
if !IsOllamaCloudUsageAccount(account) {
|
||||
return
|
||||
}
|
||||
// Resolve the api_key group so the account sees the group's managed session
|
||||
// and most recent snapshot even when it is not the row that last refreshed.
|
||||
if err := s.ResolveAccounts(ctx, []*Account{account}); err != nil {
|
||||
return
|
||||
}
|
||||
if !ollamaCloudUsageConfigured(account) {
|
||||
return
|
||||
}
|
||||
key, valid := ollamaCloudUsageGroupFingerprint(account)
|
||||
if !valid {
|
||||
return
|
||||
}
|
||||
|
||||
window := ollamaCloudUsageManualRefreshInterval
|
||||
accountSnapshot := decodeOllamaCloudUsageSnapshot(account.Extra)
|
||||
cached, hasCached := s.probeGroupResult(key)
|
||||
|
||||
// The most recent observation wins, so an older success never masks a newer
|
||||
// one (success or failure) and never drives a stale exhaustion report.
|
||||
now := s.currentTime()
|
||||
if newest := ollamaCloudUsageProbeNewestSuccess(accountSnapshot, cached, hasCached, now, window); newest != nil {
|
||||
maybeOllamaCloudUsageProbeExhaustion(ctx, accountID, newest, s.currentTime(), window, onExhausted)
|
||||
return
|
||||
}
|
||||
|
||||
// Failure backoff: a model 429 must not bypass a scrape 429's NextRefreshAt.
|
||||
if horizon := ollamaCloudUsageProbeBackoffHorizon(accountSnapshot, cached, hasCached); !horizon.IsZero() && now.Before(horizon) {
|
||||
return
|
||||
}
|
||||
// 30s group cooldown after the most recent group attempt.
|
||||
if hasCached && now.Before(cached.attemptAt.Add(window)) {
|
||||
return
|
||||
}
|
||||
|
||||
settings, settingsErr := s.GetSettings(ctx)
|
||||
if settingsErr != nil {
|
||||
return
|
||||
}
|
||||
fetched, refreshErr := s.refreshAccount(ctx, accountID, settings, false)
|
||||
// A nil success means refreshAccount piggybacked on a singleflight run owned
|
||||
// by another request (a timed-cycle refresh or a sibling probe). That owner
|
||||
// already stored the group result, so this no-op must not be recorded as a
|
||||
// fresh attempt (it would wrongly start a 30s group cooldown or mask state).
|
||||
if refreshErr == nil && fetched == nil {
|
||||
return
|
||||
}
|
||||
// Evaluate against when the fetch actually completed: a reset that rolled over
|
||||
// during a slow fetch must not be reported as still-future, and the group
|
||||
// cooldown must not start before the attempt finished.
|
||||
doneNow := s.currentTime()
|
||||
s.storeProbeGroupResult(key, ollamaCloudUsageProbeGroupEntry{attemptAt: doneNow, snapshot: fetched})
|
||||
if refreshErr != nil {
|
||||
return
|
||||
}
|
||||
maybeOllamaCloudUsageProbeExhaustion(ctx, accountID, fetched, doneNow, window, onExhausted)
|
||||
}
|
||||
|
||||
// ollamaCloudUsageProbeObservedAt returns the wall-clock time at which a snapshot
|
||||
// was actually observed: FetchedAt for a success (a failure carries a stale
|
||||
// FetchedAt copied from a prior success), otherwise LastAttemptAt.
|
||||
func ollamaCloudUsageProbeObservedAt(snapshot *OllamaCloudUsageSnapshot) (time.Time, bool) {
|
||||
if snapshot == nil {
|
||||
return time.Time{}, false
|
||||
}
|
||||
if snapshot.Status == OllamaCloudUsageStatusOK && snapshot.FetchedAt != nil && !snapshot.FetchedAt.IsZero() {
|
||||
return snapshot.FetchedAt.UTC(), true
|
||||
}
|
||||
if !snapshot.LastAttemptAt.IsZero() {
|
||||
return snapshot.LastAttemptAt.UTC(), true
|
||||
}
|
||||
return time.Time{}, false
|
||||
}
|
||||
|
||||
// ollamaCloudUsageProbeNewestSuccess returns the most recent observation when it
|
||||
// is a fresh, successful snapshot that may feed the exhaustion decision, or nil
|
||||
// otherwise. The candidate with the latest actual observation time is picked, so
|
||||
// neither the account row nor the coordinator cache is unconditionally preferred,
|
||||
// and a newer failure (or an unknown recent attempt) suppresses an older success
|
||||
// rather than letting it drive a stale exhaustion report.
|
||||
func ollamaCloudUsageProbeNewestSuccess(
|
||||
accountSnapshot *OllamaCloudUsageSnapshot,
|
||||
cached ollamaCloudUsageProbeGroupEntry,
|
||||
hasCached bool,
|
||||
now time.Time,
|
||||
window time.Duration,
|
||||
) *OllamaCloudUsageSnapshot {
|
||||
var newest *OllamaCloudUsageSnapshot
|
||||
var newestAt time.Time
|
||||
consider := func(snapshot *OllamaCloudUsageSnapshot) {
|
||||
if snapshot == nil {
|
||||
return
|
||||
}
|
||||
at, ok := ollamaCloudUsageProbeObservedAt(snapshot)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if newest == nil || at.After(newestAt) {
|
||||
newest = snapshot
|
||||
newestAt = at
|
||||
}
|
||||
}
|
||||
consider(accountSnapshot)
|
||||
if hasCached {
|
||||
if cached.snapshot != nil {
|
||||
consider(cached.snapshot)
|
||||
} else if !cached.attemptAt.IsZero() && (newest == nil || cached.attemptAt.After(newestAt)) {
|
||||
// The most recent group attempt produced no snapshot (an errored
|
||||
// fetch); it supersedes any older success.
|
||||
newest = nil
|
||||
newestAt = cached.attemptAt
|
||||
}
|
||||
}
|
||||
if newest == nil || newest.Status != OllamaCloudUsageStatusOK {
|
||||
return nil
|
||||
}
|
||||
if newest.FetchedAt == nil || newest.FetchedAt.IsZero() || newest.FetchedAt.Before(now.Add(-window)) {
|
||||
return nil
|
||||
}
|
||||
return newest
|
||||
}
|
||||
|
||||
// ollamaCloudUsageProbeBackoffHorizon returns the furthest NextRefreshAt among
|
||||
// failed/unauthorized snapshots (the account's persisted one and the coordinator's
|
||||
// most recent group outcome), or the zero time when no failure backoff is in force.
|
||||
func ollamaCloudUsageProbeBackoffHorizon(
|
||||
accountSnapshot *OllamaCloudUsageSnapshot,
|
||||
cached ollamaCloudUsageProbeGroupEntry,
|
||||
hasCached bool,
|
||||
) time.Time {
|
||||
var horizon time.Time
|
||||
consider := func(snapshot *OllamaCloudUsageSnapshot) {
|
||||
if snapshot == nil {
|
||||
return
|
||||
}
|
||||
if snapshot.Status != OllamaCloudUsageStatusFailed && snapshot.Status != OllamaCloudUsageStatusUnauthorized {
|
||||
return
|
||||
}
|
||||
if snapshot.NextRefreshAt.IsZero() {
|
||||
return
|
||||
}
|
||||
if horizon.IsZero() || snapshot.NextRefreshAt.After(horizon) {
|
||||
horizon = snapshot.NextRefreshAt.UTC()
|
||||
}
|
||||
}
|
||||
consider(accountSnapshot)
|
||||
if hasCached {
|
||||
consider(cached.snapshot)
|
||||
}
|
||||
return horizon
|
||||
}
|
||||
|
||||
// maybeOllamaCloudUsageProbeExhaustion invokes onExhausted only when the snapshot
|
||||
// is a fresh, successful, genuinely exhausted one (per
|
||||
// ollamaCloudUsageExhaustionResetAt). Non-exhausted or unusable snapshots yield
|
||||
// no callback.
|
||||
func maybeOllamaCloudUsageProbeExhaustion(
|
||||
ctx context.Context,
|
||||
accountID int64,
|
||||
snapshot *OllamaCloudUsageSnapshot,
|
||||
now time.Time,
|
||||
window time.Duration,
|
||||
onExhausted OllamaCloudUsageRateLimitProbeCallback,
|
||||
) {
|
||||
if snapshot == nil || onExhausted == nil {
|
||||
return
|
||||
}
|
||||
resetAt, exhausted := ollamaCloudUsageExhaustionResetAt(snapshot, now, now.Add(-window))
|
||||
if !exhausted {
|
||||
return
|
||||
}
|
||||
if ctx.Err() != nil {
|
||||
// Probe was cancelled/stopped: do not deliver a stale callback.
|
||||
return
|
||||
}
|
||||
onExhausted(accountID, resetAt)
|
||||
}
|
||||
@@ -0,0 +1,495 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// ollamaCloudProbeUsageBody renders a minimal settings page exposing only the
|
||||
// five-hour usage window at the given percent and an absolute reset time.
|
||||
func ollamaCloudProbeUsageBody(percent int, resetRFC3339 string) []byte {
|
||||
return []byte(fmt.Sprintf(`
|
||||
<section>
|
||||
<div><span>5 hour usage</span><span>%d%% used</span></div>
|
||||
<time datetime="%s"></time>
|
||||
</section>`, percent, resetRFC3339))
|
||||
}
|
||||
|
||||
// TestOllamaCloudProbeUsageBodiesParse pins the crafted HTML fixtures used by the
|
||||
// probe tests: 100% usage is exhausted and must carry the future reset.
|
||||
func TestOllamaCloudProbeUsageBodiesParse(t *testing.T) {
|
||||
reset := time.Date(2026, 8, 1, 0, 0, 0, 0, time.UTC)
|
||||
exhausted, err := parseOllamaCloudUsageHTML(ollamaCloudProbeUsageBody(100, reset.Format(time.RFC3339)))
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, exhausted.FiveHour)
|
||||
require.Equal(t, float64(100), exhausted.FiveHour.UsedPercent)
|
||||
require.NotNil(t, exhausted.FiveHour.ResetAt)
|
||||
require.True(t, exhausted.FiveHour.ResetAt.Equal(reset), "reset = %v", exhausted.FiveHour.ResetAt)
|
||||
|
||||
mild, err := parseOllamaCloudUsageHTML(ollamaCloudProbeUsageBody(5, reset.Format(time.RFC3339)))
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, mild.FiveHour)
|
||||
require.Equal(t, float64(5), mild.FiveHour.UsedPercent)
|
||||
}
|
||||
|
||||
// ollamaCloudProbeFixture builds a started service with a fixed clock so exhaustion
|
||||
// and reset decisions are deterministic.
|
||||
func ollamaCloudProbeFixture(t *testing.T, repo AccountRepository, upstream HTTPUpstream, fixedNow time.Time) *OllamaCloudUsageService {
|
||||
t.Helper()
|
||||
svc := NewOllamaCloudUsageService(repo, upstream, NewSettingService(&upstreamBillingProbeSettingRepo{}, nil), ollamaUsageTestEncryptor{}, true)
|
||||
svc.now = func() time.Time { return fixedNow }
|
||||
svc.Start()
|
||||
t.Cleanup(svc.Stop)
|
||||
return svc
|
||||
}
|
||||
|
||||
func ollamaCloudProbeReset(t *testing.T, fixedNow time.Time) time.Time {
|
||||
t.Helper()
|
||||
return fixedNow.Add(2 * time.Hour).UTC()
|
||||
}
|
||||
|
||||
func TestOllamaCloudUsageScheduleRateLimitProbeReturnsWhileProbeRuns(t *testing.T) {
|
||||
// A gated GetByID proves ScheduleRateLimitProbe does not synchronously block
|
||||
// on the probe's database/network work.
|
||||
fixedNow := time.Date(2026, 7, 20, 12, 0, 0, 0, time.UTC)
|
||||
account := ollamaUsageAccount(1)
|
||||
account.Extra[OllamaCloudUsageSessionExtraKey] = "cipher:wos-session=secret"
|
||||
base := &ollamaUsageTestRepo{upstreamBillingProbeAccountRepo: &upstreamBillingProbeAccountRepo{
|
||||
accounts: map[int64]*Account{account.ID: account},
|
||||
}}
|
||||
repo := &ollamaCloudProbeGateRepo{ollamaUsageTestRepo: base, entered: make(chan struct{}), release: make(chan struct{})}
|
||||
repo.gating.Store(true)
|
||||
upstream := &ollamaUsageHTTPStub{body: ollamaCloudProbeUsageBody(100, ollamaCloudProbeReset(t, fixedNow).Format(time.RFC3339))}
|
||||
svc := ollamaCloudProbeFixture(t, repo, upstream, fixedNow)
|
||||
|
||||
called := make(chan time.Time, 1)
|
||||
accepted := svc.ScheduleOllamaCloudUsageRateLimitProbe(account.ID, func(id int64, reset time.Time) {
|
||||
called <- reset
|
||||
})
|
||||
require.True(t, accepted, "running service must accept the probe")
|
||||
|
||||
// The fetch is still blocked on the gate, yet Schedule already returned:
|
||||
// this is only possible if the probe runs asynchronously on the coordinator.
|
||||
require.Eventually(t, func() bool {
|
||||
select {
|
||||
case <-repo.entered:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}, 5*time.Second, time.Millisecond, "coordinator must reach the gated GetByID")
|
||||
require.Zero(t, upstream.calls.Load(), "no upstream fetch while the probe is gated")
|
||||
select {
|
||||
case <-called:
|
||||
t.Fatal("callback must not fire before the probe completes")
|
||||
default:
|
||||
}
|
||||
|
||||
repo.gating.Store(false)
|
||||
close(repo.release)
|
||||
select {
|
||||
case reset := <-called:
|
||||
require.True(t, reset.Equal(ollamaCloudProbeReset(t, fixedNow)))
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("expected an exhaustion callback after the probe completed")
|
||||
}
|
||||
require.Equal(t, int64(1), upstream.calls.Load())
|
||||
}
|
||||
|
||||
func TestOllamaCloudUsageRateLimitProbeExhaustedReportsReset(t *testing.T) {
|
||||
fixedNow := time.Date(2026, 7, 20, 12, 0, 0, 0, time.UTC)
|
||||
account := ollamaUsageAccount(11)
|
||||
account.Extra[OllamaCloudUsageSessionExtraKey] = "cipher:wos-session=secret"
|
||||
repo := &ollamaUsageTestRepo{upstreamBillingProbeAccountRepo: &upstreamBillingProbeAccountRepo{
|
||||
accounts: map[int64]*Account{account.ID: account},
|
||||
}}
|
||||
wantReset := ollamaCloudProbeReset(t, fixedNow)
|
||||
upstream := &ollamaUsageHTTPStub{body: ollamaCloudProbeUsageBody(100, wantReset.Format(time.RFC3339))}
|
||||
svc := ollamaCloudProbeFixture(t, repo, upstream, fixedNow)
|
||||
|
||||
called := make(chan time.Time, 1)
|
||||
require.True(t, svc.ScheduleOllamaCloudUsageRateLimitProbe(account.ID, func(id int64, reset time.Time) {
|
||||
require.Equal(t, account.ID, id)
|
||||
called <- reset
|
||||
}))
|
||||
select {
|
||||
case reset := <-called:
|
||||
require.True(t, reset.Equal(wantReset))
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("expected a callback for an exhausted account")
|
||||
}
|
||||
require.Equal(t, int64(1), upstream.calls.Load())
|
||||
}
|
||||
|
||||
func TestOllamaCloudUsageRateLimitProbeNotExhaustedNoCallback(t *testing.T) {
|
||||
fixedNow := time.Date(2026, 7, 20, 12, 0, 0, 0, time.UTC)
|
||||
account := ollamaUsageAccount(12)
|
||||
account.Extra[OllamaCloudUsageSessionExtraKey] = "cipher:wos-session=secret"
|
||||
repo := &ollamaUsageTestRepo{upstreamBillingProbeAccountRepo: &upstreamBillingProbeAccountRepo{
|
||||
accounts: map[int64]*Account{account.ID: account},
|
||||
}}
|
||||
upstream := &ollamaUsageHTTPStub{body: ollamaCloudProbeUsageBody(5, ollamaCloudProbeReset(t, fixedNow).Format(time.RFC3339))}
|
||||
svc := ollamaCloudProbeFixture(t, repo, upstream, fixedNow)
|
||||
|
||||
called := make(chan struct{}, 1)
|
||||
require.True(t, svc.ScheduleOllamaCloudUsageRateLimitProbe(account.ID, func(int64, time.Time) {
|
||||
called <- struct{}{}
|
||||
}))
|
||||
select {
|
||||
case <-called:
|
||||
t.Fatal("not-exhausted usage must not invoke the callback")
|
||||
case <-time.After(300 * time.Millisecond):
|
||||
}
|
||||
require.Equal(t, int64(1), upstream.calls.Load(), "a fresh probe must still refresh once")
|
||||
}
|
||||
|
||||
func TestOllamaCloudUsageRateLimitProbeSameGroupCoalescesIntoOneFetchPerAccount(t *testing.T) {
|
||||
fixedNow := time.Date(2026, 7, 20, 12, 0, 0, 0, time.UTC)
|
||||
first := ollamaUsageAccount(21)
|
||||
first.Credentials["api_key"] = "shared-key"
|
||||
first.Extra[OllamaCloudUsageSessionExtraKey] = "cipher:wos-session=shared"
|
||||
second := ollamaUsageAccount(22)
|
||||
second.Platform = PlatformAnthropic
|
||||
second.Credentials = map[string]any{"base_url": "https://www.ollama.com/v1", "api_key": "shared-key"}
|
||||
second.Extra[OllamaCloudUsageSessionExtraKey] = "cipher:wos-session=shared"
|
||||
repo := &ollamaUsageTestRepo{upstreamBillingProbeAccountRepo: &upstreamBillingProbeAccountRepo{
|
||||
accounts: map[int64]*Account{first.ID: first, second.ID: second},
|
||||
}}
|
||||
wantReset := ollamaCloudProbeReset(t, fixedNow)
|
||||
upstream := &ollamaUsageHTTPStub{body: ollamaCloudProbeUsageBody(100, wantReset.Format(time.RFC3339))}
|
||||
svc := ollamaCloudProbeFixture(t, repo, upstream, fixedNow)
|
||||
|
||||
results := make(map[int64]time.Time)
|
||||
var mu sync.Mutex
|
||||
got := make(chan int64, 2)
|
||||
cb := func(id int64, reset time.Time) {
|
||||
mu.Lock()
|
||||
results[id] = reset
|
||||
mu.Unlock()
|
||||
got <- id
|
||||
}
|
||||
require.True(t, svc.ScheduleOllamaCloudUsageRateLimitProbe(first.ID, cb))
|
||||
require.True(t, svc.ScheduleOllamaCloudUsageRateLimitProbe(second.ID, cb))
|
||||
|
||||
for range 2 {
|
||||
select {
|
||||
case <-got:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("both same-group accounts must each receive their own callback")
|
||||
}
|
||||
}
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
require.Contains(t, results, first.ID)
|
||||
require.Contains(t, results, second.ID)
|
||||
for _, reset := range results {
|
||||
require.True(t, reset.Equal(wantReset))
|
||||
}
|
||||
require.Equal(t, int64(1), upstream.calls.Load(), "one shared group must fetch at most once for the coalesced probes")
|
||||
}
|
||||
|
||||
func TestOllamaCloudUsageRateLimitProbeSkipsMissingCookieAndNonOllamaAccount(t *testing.T) {
|
||||
fixedNow := time.Date(2026, 7, 20, 12, 0, 0, 0, time.UTC)
|
||||
noCookie := ollamaUsageAccount(31)
|
||||
noCookie.Extra = map[string]any{} // no session
|
||||
nonOllama := ollamaUsageAccount(32)
|
||||
nonOllama.Credentials["base_url"] = "https://api.openai.com"
|
||||
repo := &ollamaUsageTestRepo{upstreamBillingProbeAccountRepo: &upstreamBillingProbeAccountRepo{
|
||||
accounts: map[int64]*Account{noCookie.ID: noCookie, nonOllama.ID: nonOllama},
|
||||
}}
|
||||
upstream := &ollamaUsageHTTPStub{body: ollamaCloudProbeUsageBody(100, ollamaCloudProbeReset(t, fixedNow).Format(time.RFC3339))}
|
||||
svc := ollamaCloudProbeFixture(t, repo, upstream, fixedNow)
|
||||
|
||||
called := make(chan struct{}, 1)
|
||||
cb := func(int64, time.Time) { called <- struct{}{} }
|
||||
require.True(t, svc.ScheduleOllamaCloudUsageRateLimitProbe(noCookie.ID, cb))
|
||||
require.True(t, svc.ScheduleOllamaCloudUsageRateLimitProbe(nonOllama.ID, cb))
|
||||
|
||||
select {
|
||||
case <-called:
|
||||
t.Fatal("ineligible accounts must be skipped quietly")
|
||||
case <-time.After(300 * time.Millisecond):
|
||||
}
|
||||
require.Zero(t, upstream.calls.Load(), "no upstream fetch for ineligible accounts")
|
||||
}
|
||||
|
||||
func TestOllamaCloudUsageRateLimitProbeRespectsFailureBackoff(t *testing.T) {
|
||||
fixedNow := time.Date(2026, 7, 20, 12, 0, 0, 0, time.UTC)
|
||||
account := ollamaUsageAccount(41)
|
||||
account.Extra[OllamaCloudUsageSessionExtraKey] = "cipher:wos-session=secret"
|
||||
backoffUntil := fixedNow.Add(30 * time.Minute)
|
||||
account.Extra[OllamaCloudUsageSnapshotExtraKey] = &OllamaCloudUsageSnapshot{
|
||||
Status: OllamaCloudUsageStatusFailed,
|
||||
LastAttemptAt: fixedNow.Add(-time.Minute),
|
||||
NextRefreshAt: backoffUntil,
|
||||
FailureCount: 2,
|
||||
}
|
||||
repo := &ollamaUsageTestRepo{upstreamBillingProbeAccountRepo: &upstreamBillingProbeAccountRepo{
|
||||
accounts: map[int64]*Account{account.ID: account},
|
||||
}}
|
||||
upstream := &ollamaUsageHTTPStub{body: ollamaCloudProbeUsageBody(100, ollamaCloudProbeReset(t, fixedNow).Format(time.RFC3339))}
|
||||
svc := ollamaCloudProbeFixture(t, repo, upstream, fixedNow)
|
||||
|
||||
called := make(chan struct{}, 1)
|
||||
require.True(t, svc.ScheduleOllamaCloudUsageRateLimitProbe(account.ID, func(int64, time.Time) {
|
||||
called <- struct{}{}
|
||||
}))
|
||||
select {
|
||||
case <-called:
|
||||
t.Fatal("a model 429 must not bypass scrape-429 backoff")
|
||||
case <-time.After(300 * time.Millisecond):
|
||||
}
|
||||
require.Zero(t, upstream.calls.Load(), "no fetch while the snapshot is in failure backoff")
|
||||
}
|
||||
|
||||
func TestOllamaCloudUsageRateLimitProbeStopCancelsInFlightAndRejectsNew(t *testing.T) {
|
||||
fixedNow := time.Date(2026, 7, 20, 12, 0, 0, 0, time.UTC)
|
||||
account := ollamaUsageAccount(51)
|
||||
account.Extra[OllamaCloudUsageSessionExtraKey] = "cipher:wos-session=secret"
|
||||
base := &ollamaUsageTestRepo{upstreamBillingProbeAccountRepo: &upstreamBillingProbeAccountRepo{
|
||||
accounts: map[int64]*Account{account.ID: account},
|
||||
}}
|
||||
repo := &ollamaCloudProbeGateRepo{ollamaUsageTestRepo: base, entered: make(chan struct{}), release: make(chan struct{})}
|
||||
repo.gating.Store(true)
|
||||
svc := ollamaCloudProbeFixture(t, repo, &ollamaUsageHTTPStub{body: ollamaCloudProbeUsageBody(100, ollamaCloudProbeReset(t, fixedNow).Format(time.RFC3339))}, fixedNow)
|
||||
|
||||
called := make(chan struct{}, 1)
|
||||
require.True(t, svc.ScheduleOllamaCloudUsageRateLimitProbe(account.ID, func(int64, time.Time) {
|
||||
called <- struct{}{}
|
||||
}))
|
||||
require.Eventually(t, func() bool {
|
||||
select {
|
||||
case <-repo.entered:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}, 5*time.Second, time.Millisecond, "coordinator must be blocked in the probe")
|
||||
|
||||
// Stop cancels parentCtx; the in-flight probe's detached context is derived
|
||||
// from it, so the gated GetByID is aborted and Stop returns promptly.
|
||||
done := make(chan struct{})
|
||||
go func() { svc.Stop(); close(done) }()
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("Stop must not hang while a probe is in flight")
|
||||
}
|
||||
select {
|
||||
case <-called:
|
||||
t.Fatal("no callback may be delivered for a probe cancelled by Stop")
|
||||
default:
|
||||
}
|
||||
|
||||
// A stopped service rejects further probes and performs no further fetch.
|
||||
require.False(t, svc.ScheduleOllamaCloudUsageRateLimitProbe(account.ID, func(int64, time.Time) {
|
||||
called <- struct{}{}
|
||||
}), "a stopped service must reject new probes")
|
||||
}
|
||||
|
||||
// ollamaCloudProbeGateRepo blocks GetByID until released (or its context is
|
||||
// cancelled), letting tests prove probes run asynchronously and that Stop aborts
|
||||
// in-flight work.
|
||||
type ollamaCloudProbeGateRepo struct {
|
||||
*ollamaUsageTestRepo
|
||||
entered chan struct{}
|
||||
release chan struct{}
|
||||
gating atomic.Bool
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func (r *ollamaCloudProbeGateRepo) GetByID(ctx context.Context, id int64) (*Account, error) {
|
||||
if r.gating.Load() {
|
||||
r.once.Do(func() { close(r.entered) })
|
||||
select {
|
||||
case <-r.release:
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
}
|
||||
return r.ollamaUsageTestRepo.GetByID(ctx, id)
|
||||
}
|
||||
|
||||
// TestOllamaCloudUsageRateLimitProbeSlowFetchDoesNotReportExpiredReset checks that
|
||||
// an exhaustion decision is evaluated at fetch completion time: when the reset
|
||||
// rolls over during a slow fetch, the probe must not report a horizon that is
|
||||
// already in the past (which would otherwise trigger a needless long-lived ban).
|
||||
func TestOllamaCloudUsageRateLimitProbeSlowFetchDoesNotReportExpiredReset(t *testing.T) {
|
||||
base := time.Date(2026, 7, 20, 12, 0, 0, 0, time.UTC)
|
||||
account := ollamaUsageAccount(61)
|
||||
account.Extra[OllamaCloudUsageSessionExtraKey] = "cipher:wos-session=secret"
|
||||
repo := &ollamaUsageTestRepo{upstreamBillingProbeAccountRepo: &upstreamBillingProbeAccountRepo{
|
||||
accounts: map[int64]*Account{account.ID: account},
|
||||
}}
|
||||
|
||||
var mu sync.Mutex
|
||||
cur := base
|
||||
reset := base.Add(time.Second) // in the future at start, expired by completion
|
||||
advanced := false
|
||||
upstream := &ollamaUsageHTTPStub{body: ollamaCloudProbeUsageBody(100, reset.Format(time.RFC3339)), beforeResponse: func(*http.Request) {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
if !advanced {
|
||||
cur = base.Add(2 * time.Second) // fetch completes after the reset
|
||||
advanced = true
|
||||
}
|
||||
}}
|
||||
svc := NewOllamaCloudUsageService(repo, upstream, NewSettingService(&upstreamBillingProbeSettingRepo{}, nil), ollamaUsageTestEncryptor{}, true)
|
||||
svc.now = func() time.Time { mu.Lock(); defer mu.Unlock(); return cur }
|
||||
svc.Start()
|
||||
t.Cleanup(svc.Stop)
|
||||
|
||||
called := make(chan struct{}, 1)
|
||||
require.True(t, svc.ScheduleOllamaCloudUsageRateLimitProbe(account.ID, func(int64, time.Time) {
|
||||
called <- struct{}{}
|
||||
}))
|
||||
select {
|
||||
case <-called:
|
||||
t.Fatal("a reset that already rolled over during the fetch must not be reported")
|
||||
case <-time.After(400 * time.Millisecond):
|
||||
}
|
||||
require.Equal(t, int64(1), upstream.calls.Load(), "the slow probe must still fetch once")
|
||||
|
||||
// The snapshot really was exhausted; only the expired reset suppressed the
|
||||
// callback, which is what this test pins. Read it through the repo (which
|
||||
// clones under its lock) to avoid racing the coordinator's snapshot write.
|
||||
require.Eventually(t, func() bool {
|
||||
loaded, err := repo.GetByID(context.Background(), account.ID)
|
||||
if err != nil || loaded == nil {
|
||||
return false
|
||||
}
|
||||
snapshot := decodeOllamaCloudUsageSnapshot(loaded.Extra)
|
||||
return snapshot != nil && snapshot.Status == OllamaCloudUsageStatusOK &&
|
||||
snapshot.Data != nil && snapshot.Data.FiveHour != nil && snapshot.Data.FiveHour.UsedPercent == 100
|
||||
}, 5*time.Second, 5*time.Millisecond, "coordinator must persist the exhausted snapshot")
|
||||
}
|
||||
|
||||
// TestOllamaCloudUsageRateLimitProbeQueueFullStillMergesExistingAccount verifies
|
||||
// that a full queue still accepts (and merges into the latest callback) a probe
|
||||
// for an account that is already queued, while a brand-new account is rejected.
|
||||
func TestOllamaCloudUsageRateLimitProbeQueueFullStillMergesExistingAccount(t *testing.T) {
|
||||
base := time.Date(2026, 7, 20, 12, 0, 0, 0, time.UTC)
|
||||
repo := &ollamaCloudProbeGateRepo{ollamaUsageTestRepo: &ollamaUsageTestRepo{
|
||||
upstreamBillingProbeAccountRepo: &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{}},
|
||||
}, entered: make(chan struct{}), release: make(chan struct{})}
|
||||
repo.gating.Store(true)
|
||||
svc := NewOllamaCloudUsageService(repo, &ollamaUsageHTTPStub{body: ollamaCloudProbeUsageBody(100, base.Add(time.Hour).Format(time.RFC3339))}, NewSettingService(&upstreamBillingProbeSettingRepo{}, nil), ollamaUsageTestEncryptor{}, true)
|
||||
svc.now = func() time.Time { return base }
|
||||
svc.Start()
|
||||
t.Cleanup(svc.Stop)
|
||||
|
||||
noop := func(int64, time.Time) {}
|
||||
// The first probe blocks the single coordinator on its gated GetByID.
|
||||
require.True(t, svc.ScheduleOllamaCloudUsageRateLimitProbe(9001, noop))
|
||||
require.Eventually(t, func() bool {
|
||||
select {
|
||||
case <-repo.entered:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}, 5*time.Second, time.Millisecond, "coordinator must be blocked")
|
||||
|
||||
// Fill the pending queue to capacity (item 9001 is in flight, not pending).
|
||||
for id := int64(9002); id <= 9001+ollamaCloudUsageProbeMaxQueue; id++ {
|
||||
require.True(t, svc.ScheduleOllamaCloudUsageRateLimitProbe(id, noop), "filling the queue must be accepted")
|
||||
}
|
||||
// Queue is now full: a brand-new account is rejected, but an already-queued
|
||||
// account is still merged into its latest callback.
|
||||
require.False(t, svc.ScheduleOllamaCloudUsageRateLimitProbe(9999, noop), "a new account must be rejected when the queue is full")
|
||||
require.True(t, svc.ScheduleOllamaCloudUsageRateLimitProbe(9002, noop), "an already-queued account must still merge when the queue is full")
|
||||
}
|
||||
|
||||
func probeSnap(status string, fetchedAt, lastAttempt time.Time) *OllamaCloudUsageSnapshot {
|
||||
s := &OllamaCloudUsageSnapshot{Status: status, LastAttemptAt: lastAttempt}
|
||||
if !fetchedAt.IsZero() {
|
||||
s.FetchedAt = &fetchedAt
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// TestOllamaCloudUsageProbeNewestSuccess pins the "most recent observation wins"
|
||||
// rule: neither the account row nor the coordinator cache is preferred outright,
|
||||
// an older success never masks a newer success, and a newer failure suppresses an
|
||||
// older success so it cannot drive a stale exhaustion report.
|
||||
func TestOllamaCloudUsageProbeNewestSuccess(t *testing.T) {
|
||||
now := time.Date(2026, 7, 20, 12, 0, 0, 0, time.UTC)
|
||||
window := ollamaCloudUsageManualRefreshInterval
|
||||
fresh := now.Add(-5 * time.Second)
|
||||
newer := now.Add(-time.Second)
|
||||
stale := now.Add(-time.Hour)
|
||||
failed := now.Add(-2 * time.Second) // LastAttemptAt for a failure
|
||||
entry := func(snapshot *OllamaCloudUsageSnapshot, at time.Time) ollamaCloudUsageProbeGroupEntry {
|
||||
return ollamaCloudUsageProbeGroupEntry{attemptAt: at, snapshot: snapshot}
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
accountSnapshot *OllamaCloudUsageSnapshot
|
||||
cached ollamaCloudUsageProbeGroupEntry
|
||||
hasCached bool
|
||||
wantFetchedAt time.Time // zero means nil result
|
||||
}{
|
||||
{
|
||||
name: "no snapshots",
|
||||
wantFetchedAt: time.Time{},
|
||||
},
|
||||
{
|
||||
name: "cached newer success wins over account success",
|
||||
accountSnapshot: probeSnap(OllamaCloudUsageStatusOK, fresh, fresh),
|
||||
cached: entry(probeSnap(OllamaCloudUsageStatusOK, newer, newer), newer),
|
||||
hasCached: true,
|
||||
wantFetchedAt: newer,
|
||||
},
|
||||
{
|
||||
name: "account newer success wins over cached success",
|
||||
accountSnapshot: probeSnap(OllamaCloudUsageStatusOK, newer, newer),
|
||||
cached: entry(probeSnap(OllamaCloudUsageStatusOK, fresh, fresh), fresh),
|
||||
hasCached: true,
|
||||
wantFetchedAt: newer,
|
||||
},
|
||||
{
|
||||
name: "newer account failure suppresses older cached success",
|
||||
accountSnapshot: probeSnap(OllamaCloudUsageStatusFailed, time.Time{}, failed),
|
||||
cached: entry(probeSnap(OllamaCloudUsageStatusOK, fresh, fresh), fresh),
|
||||
hasCached: true,
|
||||
wantFetchedAt: time.Time{},
|
||||
},
|
||||
{
|
||||
name: "newer cached failure suppresses older account success",
|
||||
accountSnapshot: probeSnap(OllamaCloudUsageStatusOK, fresh, fresh),
|
||||
cached: entry(probeSnap(OllamaCloudUsageStatusFailed, time.Time{}, failed), failed),
|
||||
hasCached: true,
|
||||
wantFetchedAt: time.Time{},
|
||||
},
|
||||
{
|
||||
name: "stale success alone is not usable",
|
||||
accountSnapshot: probeSnap(OllamaCloudUsageStatusOK, stale, stale),
|
||||
wantFetchedAt: time.Time{},
|
||||
},
|
||||
{
|
||||
name: "only a fresh account success is usable",
|
||||
accountSnapshot: probeSnap(OllamaCloudUsageStatusOK, fresh, fresh),
|
||||
wantFetchedAt: fresh,
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := ollamaCloudUsageProbeNewestSuccess(tt.accountSnapshot, tt.cached, tt.hasCached, now, window)
|
||||
if tt.wantFetchedAt.IsZero() {
|
||||
require.Nil(t, got)
|
||||
return
|
||||
}
|
||||
require.NotNil(t, got)
|
||||
require.NotNil(t, got.FetchedAt)
|
||||
require.True(t, got.FetchedAt.Equal(tt.wantFetchedAt), "FetchedAt = %v, want %v", got.FetchedAt, tt.wantFetchedAt)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,65 @@
|
||||
package service
|
||||
|
||||
import "time"
|
||||
|
||||
// ollamaCloudUsageExhaustionResetAt returns the wall-clock time at which an
|
||||
// exhausted Ollama Cloud usage window is expected to roll over, i.e. the
|
||||
// earliest moment a quota-driven upstream block (a 429 or equivalent) can be
|
||||
// expected to clear on its own. It is a pure helper and performs no I/O and no
|
||||
// routing; it does not itself trigger anything.
|
||||
//
|
||||
// Snapshot usability (conservative): the observation is trusted only when it is
|
||||
// successful (Status == OK) and fresh. Freshness is expressed as a caller-chosen
|
||||
// validity bound minFetchedAt: the snapshot must carry a FetchedAt that is
|
||||
// present and not older than minFetchedAt. Passing now.Add(-freshness) keeps
|
||||
// old snapshots out of the decision without reading any configuration here.
|
||||
//
|
||||
// Window handling:
|
||||
// - Only windows with UsedPercent >= 100 count as exhausted.
|
||||
// - A window that is not exhausted (or absent from Data) never affects the
|
||||
// result and is ignored even when its ResetAt is missing or in the past.
|
||||
// - The reset of every exhausted window must be known (non-nil, non-zero) and
|
||||
// strictly in the future relative to now. If any exhausted window cannot
|
||||
// produce such a reset, no complete recovery time can be claimed and ok is
|
||||
// false — the caller must not assume a known recovery horizon.
|
||||
// - When more than one window is exhausted, the later of their resets is
|
||||
// returned: the account cannot recover until the last exhausted window rolls.
|
||||
//
|
||||
// ok is false (zero time) when there is nothing exhausted or the recovery time
|
||||
// cannot be determined completely.
|
||||
func ollamaCloudUsageExhaustionResetAt(
|
||||
snapshot *OllamaCloudUsageSnapshot,
|
||||
now time.Time,
|
||||
minFetchedAt time.Time,
|
||||
) (reset time.Time, ok bool) {
|
||||
if snapshot == nil || snapshot.Status != OllamaCloudUsageStatusOK || snapshot.Data == nil {
|
||||
return time.Time{}, false
|
||||
}
|
||||
if snapshot.FetchedAt == nil || snapshot.FetchedAt.IsZero() || snapshot.FetchedAt.Before(minFetchedAt) {
|
||||
return time.Time{}, false
|
||||
}
|
||||
|
||||
var latest time.Time
|
||||
for _, window := range []*OllamaCloudUsageWindow{
|
||||
snapshot.Data.FiveHour,
|
||||
snapshot.Data.SevenDay,
|
||||
} {
|
||||
if window == nil || window.UsedPercent < 100 {
|
||||
// Not exhausted: must not influence the recovery decision.
|
||||
continue
|
||||
}
|
||||
if window.ResetAt == nil || window.ResetAt.IsZero() || !window.ResetAt.After(now) {
|
||||
// Exhausted but without a known future reset: cannot claim a complete
|
||||
// recovery time, so give up on the whole snapshot.
|
||||
return time.Time{}, false
|
||||
}
|
||||
candidate := window.ResetAt.UTC()
|
||||
if latest.IsZero() || candidate.After(latest) {
|
||||
latest = candidate
|
||||
}
|
||||
}
|
||||
if latest.IsZero() {
|
||||
return time.Time{}, false
|
||||
}
|
||||
return latest, true
|
||||
}
|
||||
@@ -0,0 +1,243 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func mustT(t time.Time) *time.Time { return &t }
|
||||
|
||||
func TestOllamaCloudUsageExhaustionResetAt(t *testing.T) {
|
||||
now := time.Date(2026, 7, 20, 12, 0, 0, 0, time.UTC)
|
||||
minFetchedAt := now.Add(-time.Hour)
|
||||
|
||||
at := func(h, m int) *time.Time {
|
||||
return mustT(now.Add(time.Duration(h)*time.Hour + time.Duration(m)*time.Minute))
|
||||
}
|
||||
window := func(percent float64, reset *time.Time) *OllamaCloudUsageWindow {
|
||||
return &OllamaCloudUsageWindow{UsedPercent: percent, ResetAt: reset}
|
||||
}
|
||||
okSnapshot := func(data *OllamaCloudUsageData) *OllamaCloudUsageSnapshot {
|
||||
return &OllamaCloudUsageSnapshot{
|
||||
Status: OllamaCloudUsageStatusOK,
|
||||
Data: data,
|
||||
FetchedAt: mustT(now.Add(-time.Minute)),
|
||||
}
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
snapshot *OllamaCloudUsageSnapshot
|
||||
now time.Time
|
||||
minFresh time.Time
|
||||
want time.Time
|
||||
wantOK bool
|
||||
}{
|
||||
{
|
||||
name: "nil snapshot",
|
||||
wantOK: false,
|
||||
},
|
||||
{
|
||||
name: "failed status unusable",
|
||||
snapshot: &OllamaCloudUsageSnapshot{
|
||||
Status: OllamaCloudUsageStatusFailed,
|
||||
Data: &OllamaCloudUsageData{FiveHour: window(100, at(2, 0))},
|
||||
},
|
||||
wantOK: false,
|
||||
},
|
||||
{
|
||||
name: "nil data unusable",
|
||||
snapshot: &OllamaCloudUsageSnapshot{
|
||||
Status: OllamaCloudUsageStatusOK,
|
||||
FetchedAt: mustT(now.Add(-time.Minute)),
|
||||
},
|
||||
wantOK: false,
|
||||
},
|
||||
{
|
||||
name: "stale snapshot rejected",
|
||||
snapshot: func() *OllamaCloudUsageSnapshot {
|
||||
s := okSnapshot(&OllamaCloudUsageData{FiveHour: window(100, at(2, 0))})
|
||||
s.FetchedAt = mustT(now.Add(-3 * time.Hour))
|
||||
return s
|
||||
}(),
|
||||
minFresh: minFetchedAt,
|
||||
wantOK: false,
|
||||
},
|
||||
{
|
||||
name: "no window exhausted yields no recovery",
|
||||
snapshot: okSnapshot(&OllamaCloudUsageData{
|
||||
FiveHour: window(50, at(2, 0)),
|
||||
SevenDay: window(30, at(48, 0)),
|
||||
}),
|
||||
minFresh: minFetchedAt,
|
||||
wantOK: false,
|
||||
},
|
||||
{
|
||||
name: "five hour exhausted below boundary not counted",
|
||||
snapshot: okSnapshot(&OllamaCloudUsageData{
|
||||
FiveHour: window(99.999, at(2, 0)),
|
||||
}),
|
||||
minFresh: minFetchedAt,
|
||||
wantOK: false,
|
||||
},
|
||||
{
|
||||
name: "five hour exactly 100 future reset",
|
||||
snapshot: okSnapshot(&OllamaCloudUsageData{
|
||||
FiveHour: window(100, at(2, 0)),
|
||||
}),
|
||||
minFresh: minFetchedAt,
|
||||
want: now.Add(2 * time.Hour),
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "five hour exhausted past reset yields no recovery",
|
||||
snapshot: okSnapshot(&OllamaCloudUsageData{
|
||||
FiveHour: window(100, at(-1, 0)),
|
||||
}),
|
||||
minFresh: minFetchedAt,
|
||||
wantOK: false,
|
||||
},
|
||||
{
|
||||
name: "five hour exhausted missing reset yields no recovery",
|
||||
snapshot: okSnapshot(&OllamaCloudUsageData{
|
||||
FiveHour: window(100, nil),
|
||||
}),
|
||||
minFresh: minFetchedAt,
|
||||
wantOK: false,
|
||||
},
|
||||
{
|
||||
name: "seven day exhausted alone",
|
||||
snapshot: okSnapshot(&OllamaCloudUsageData{
|
||||
SevenDay: window(120, at(30, 0)),
|
||||
}),
|
||||
minFresh: minFetchedAt,
|
||||
want: now.Add(30 * time.Hour),
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "seven day exhausted missing reset blocks complete recovery",
|
||||
snapshot: okSnapshot(&OllamaCloudUsageData{
|
||||
FiveHour: window(100, at(2, 0)),
|
||||
SevenDay: window(100, nil),
|
||||
}),
|
||||
minFresh: minFetchedAt,
|
||||
wantOK: false,
|
||||
},
|
||||
{
|
||||
name: "seven day exhausted past reset blocks complete recovery",
|
||||
snapshot: okSnapshot(&OllamaCloudUsageData{
|
||||
FiveHour: window(100, at(2, 0)),
|
||||
SevenDay: window(100, at(-1, 0)),
|
||||
}),
|
||||
minFresh: minFetchedAt,
|
||||
wantOK: false,
|
||||
},
|
||||
{
|
||||
name: "both exhausted returns later reset",
|
||||
snapshot: okSnapshot(&OllamaCloudUsageData{
|
||||
FiveHour: window(100, at(2, 0)),
|
||||
SevenDay: window(100, at(48, 0)),
|
||||
}),
|
||||
minFresh: minFetchedAt,
|
||||
want: now.Add(48 * time.Hour),
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "both exhausted returns later even when seven day is sooner",
|
||||
snapshot: okSnapshot(&OllamaCloudUsageData{
|
||||
FiveHour: window(100, at(48, 0)),
|
||||
SevenDay: window(100, at(2, 0)),
|
||||
}),
|
||||
minFresh: minFetchedAt,
|
||||
want: now.Add(48 * time.Hour),
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "non exhausted window reset must not affect",
|
||||
snapshot: okSnapshot(&OllamaCloudUsageData{
|
||||
FiveHour: window(100, at(2, 0)),
|
||||
SevenDay: window(40, at(1, 0)), // past + not exhausted -> ignored
|
||||
}),
|
||||
minFresh: minFetchedAt,
|
||||
want: now.Add(2 * time.Hour),
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "non exhausted window with missing reset must not affect",
|
||||
snapshot: okSnapshot(&OllamaCloudUsageData{
|
||||
FiveHour: window(100, at(2, 0)),
|
||||
SevenDay: window(40, nil),
|
||||
}),
|
||||
minFresh: minFetchedAt,
|
||||
want: now.Add(2 * time.Hour),
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "absent seven day does not block",
|
||||
snapshot: okSnapshot(&OllamaCloudUsageData{
|
||||
FiveHour: window(100, at(2, 0)),
|
||||
}),
|
||||
minFresh: minFetchedAt,
|
||||
want: now.Add(2 * time.Hour),
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "future reset at now boundary treated as not future",
|
||||
snapshot: okSnapshot(&OllamaCloudUsageData{
|
||||
FiveHour: window(100, at(0, 0)),
|
||||
}),
|
||||
minFresh: minFetchedAt,
|
||||
wantOK: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
testNow := tt.now
|
||||
if testNow.IsZero() {
|
||||
testNow = now
|
||||
}
|
||||
got, ok := ollamaCloudUsageExhaustionResetAt(tt.snapshot, testNow, tt.minFresh)
|
||||
require.Equal(t, tt.wantOK, ok)
|
||||
if tt.wantOK {
|
||||
require.True(t, got.Equal(tt.want), "reset = %v, want %v", got, tt.want)
|
||||
} else {
|
||||
require.True(t, got.IsZero())
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestOllamaCloudUsageExhaustionResetAtMissingFetchedAt checks the fetched-at
|
||||
// gate separately because the okSnapshot helper always sets FetchedAt.
|
||||
func TestOllamaCloudUsageExhaustionResetAtMissingFetchedAt(t *testing.T) {
|
||||
now := time.Date(2026, 7, 20, 12, 0, 0, 0, time.UTC)
|
||||
snapshot := &OllamaCloudUsageSnapshot{
|
||||
Status: OllamaCloudUsageStatusOK,
|
||||
Data: &OllamaCloudUsageData{
|
||||
FiveHour: &OllamaCloudUsageWindow{UsedPercent: 100, ResetAt: mustT(now.Add(2 * time.Hour))},
|
||||
},
|
||||
}
|
||||
got, ok := ollamaCloudUsageExhaustionResetAt(snapshot, now, now.Add(-time.Hour))
|
||||
require.False(t, ok)
|
||||
require.True(t, got.IsZero())
|
||||
}
|
||||
|
||||
// TestOllamaCloudUsageExhaustionResetAtFetchedBoundary pins the inclusive
|
||||
// minFetchedAt boundary: a snapshot fetched exactly at the bound is usable.
|
||||
func TestOllamaCloudUsageExhaustionResetAtFetchedBoundary(t *testing.T) {
|
||||
now := time.Date(2026, 7, 20, 12, 0, 0, 0, time.UTC)
|
||||
bound := now.Add(-time.Hour)
|
||||
snapshot := &OllamaCloudUsageSnapshot{
|
||||
Status: OllamaCloudUsageStatusOK,
|
||||
FetchedAt: &bound, // exactly at the validity bound -> accepted
|
||||
Data: &OllamaCloudUsageData{
|
||||
FiveHour: &OllamaCloudUsageWindow{UsedPercent: 100, ResetAt: mustT(now.Add(2 * time.Hour))},
|
||||
},
|
||||
}
|
||||
got, ok := ollamaCloudUsageExhaustionResetAt(snapshot, now, bound)
|
||||
require.True(t, ok)
|
||||
require.True(t, got.Equal(now.Add(2*time.Hour)))
|
||||
}
|
||||
@@ -31,6 +31,11 @@ type RateLimitService struct {
|
||||
settingService *SettingService
|
||||
tokenCacheInvalidator TokenCacheInvalidator
|
||||
runtimeBlocker AccountRuntimeBlocker
|
||||
// ollamaCloudUsageProbe is the optional Ollama Cloud usage probe scheduler
|
||||
// injected via SetOllamaCloudUsageProbeScheduler. See
|
||||
// ratelimit_service_ollama_429.go for how real-Ollama 429s schedule an async
|
||||
// probe to learn the true usage-window reset.
|
||||
ollamaCloudUsageProbe ollamaCloudUsageProbeScheduler
|
||||
usageCacheMu sync.RWMutex
|
||||
usageCache map[int64]*geminiUsageCacheEntry
|
||||
|
||||
@@ -1152,6 +1157,14 @@ func (s *RateLimitService) handle429(ctx context.Context, account *Account, head
|
||||
}
|
||||
return
|
||||
}
|
||||
// 真实 Ollama Cloud 用量账号（credentials base_url 指向 ollama.com）的 429 由
|
||||
// ollama.com 的用量窗口驱动。其响应头不得被当作 OpenAI codex / Anthropic /
|
||||
// CN 限流来解析，故在国产供应商分支之前单独处理：先设置永不缩短的临时冷却，
|
||||
// 再调度异步 probe 学习真实重置点（详见 ratelimit_service_ollama_429.go）。
|
||||
if account != nil && IsOllamaCloudUsageAccount(account) {
|
||||
s.handleOllamaCloudUsage429(ctx, account, headers)
|
||||
return
|
||||
}
|
||||
// 国产供应商（kimi/zhipu/deepseek）的 429 走专用可恢复路径：余额不足 → 临时停调，
|
||||
// Coding Plan 窗口耗尽 → 冷却到快照重置点。未命中则继续默认 429 逻辑。
|
||||
if account.IsCNProvider() || account.IsOpenCodeGo() {
|
||||
|
||||
@@ -0,0 +1,218 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Real Ollama Cloud usage accounts (API-key accounts whose base_url points at
|
||||
// ollama.com) have quota-driven 429s; their headers must not be re-read as an
|
||||
// OpenAI codex / Anthropic / CN limit. handleOllamaCloudUsage429 applies a
|
||||
// never-shrinking immediate cooldown (valid Retry-After, else the seconds
|
||||
// fallback), announces runtime scheduling from the authoritative row, then
|
||||
// schedules an async usage probe (ollama_cloud_usage_rate_limit_probe.go) to
|
||||
// learn the true window reset. The probe callback only writes back while the
|
||||
// account still matches the trigger's fingerprint and rate-limit generation,
|
||||
// enforced atomically in the repository CAS, so a stale async result cannot
|
||||
// overwrite a newer 429, an admin clear/re-arm, a disabled account, or a key
|
||||
// swap. No in-memory per-account guard is kept; each pending probe carries its
|
||||
// own version in its callback closure.
|
||||
|
||||
const ollamaCloudUsageProbeWritebackTimeout = 10 * time.Second
|
||||
|
||||
// ollamaCloudUsageProbeScheduler is the single-method surface RateLimitService
|
||||
// needs from the Ollama Cloud usage service. It is optional.
|
||||
type ollamaCloudUsageProbeScheduler interface {
|
||||
ScheduleOllamaCloudUsageRateLimitProbe(accountID int64, onExhausted OllamaCloudUsageRateLimitProbeCallback) bool
|
||||
}
|
||||
|
||||
// ollamaCloudUsageRateLimitExtender extends (never shrinks) an account-level
|
||||
// rate limit via the concrete repository (SetRateLimitedIfLater).
|
||||
type ollamaCloudUsageRateLimitExtender interface {
|
||||
SetRateLimitedIfLater(ctx context.Context, id int64, resetAt time.Time) error
|
||||
}
|
||||
|
||||
// ollamaCloudUsageRateLimitSetterIfGeneration atomically writes a new reset only
|
||||
// while the account still carries the exact generation the caller observed
|
||||
// (UpdatedAt == expectedUpdatedAt AND RateLimitedAt == expectedLimitedAt AND
|
||||
// RateLimitResetAt == expectedResetAt, where nil means unset), returning whether
|
||||
// it updated. Provided by the concrete repository (SetRateLimitedIfUnchanged).
|
||||
type ollamaCloudUsageRateLimitSetterIfGeneration interface {
|
||||
SetRateLimitedIfUnchanged(ctx context.Context, id int64, expectedUpdatedAt time.Time, expectedLimitedAt, expectedResetAt *time.Time, newResetAt time.Time) (bool, error)
|
||||
}
|
||||
|
||||
// SetOllamaCloudUsageProbeScheduler injects the optional probe scheduler. When
|
||||
// nil, handle429 still cools the account and announces runtime scheduling; only
|
||||
// the async probe is skipped. Wire calls this once at startup.
|
||||
func (s *RateLimitService) SetOllamaCloudUsageProbeScheduler(scheduler ollamaCloudUsageProbeScheduler) {
|
||||
s.ollamaCloudUsageProbe = scheduler
|
||||
}
|
||||
|
||||
func (s *RateLimitService) handleOllamaCloudUsage429(ctx context.Context, account *Account, headers http.Header) {
|
||||
if s == nil || account == nil || account.ID <= 0 || s.accountRepo == nil {
|
||||
return
|
||||
}
|
||||
|
||||
var shortReset time.Time
|
||||
now := time.Now()
|
||||
if d := retryAfter(headers, now); d > 0 {
|
||||
shortReset = now.Add(d)
|
||||
} else if cooldown, enabled := s.get429FallbackCooldown(ctx, account); enabled {
|
||||
shortReset = now.Add(cooldown)
|
||||
} else {
|
||||
slog.Info("rate_limit_ollama_429_fallback_ignored", "account_id", account.ID, "platform", account.Platform)
|
||||
}
|
||||
|
||||
if !shortReset.IsZero() {
|
||||
s.applyOllamaCloudUsageImmediateCooldown(ctx, account, shortReset)
|
||||
}
|
||||
|
||||
// Read the authoritative row (IfLater may have kept a longer pre-existing
|
||||
// reset) and announce runtime scheduling from it. This is independent of the
|
||||
// optional probe scheduler.
|
||||
authoritative, err := s.accountRepo.GetByID(ctx, account.ID)
|
||||
if err != nil || authoritative == nil {
|
||||
slog.Warn("ollama_cloud_usage_authoritative_load_failed", "account_id", account.ID, "error", err)
|
||||
return
|
||||
}
|
||||
if authoritative.RateLimitResetAt != nil && authoritative.RateLimitResetAt.After(now) {
|
||||
s.notifyAccountSchedulingBlocked(authoritative, *authoritative.RateLimitResetAt, "ollama_429")
|
||||
}
|
||||
|
||||
if s.ollamaCloudUsageProbe == nil {
|
||||
return
|
||||
}
|
||||
s.scheduleOllamaCloudUsageProbe(authoritative)
|
||||
}
|
||||
|
||||
// applyOllamaCloudUsageImmediateCooldown persists the short cooldown with the
|
||||
// atomic never-shrink IfLater write when available. Runtime announcement is done
|
||||
// later from the authoritative row, never here, so an IfLater no-op never sends a
|
||||
// shorter runtime-block notification.
|
||||
func (s *RateLimitService) applyOllamaCloudUsageImmediateCooldown(ctx context.Context, account *Account, shortReset time.Time) {
|
||||
if extender, ok := s.accountRepo.(ollamaCloudUsageRateLimitExtender); ok {
|
||||
if err := extender.SetRateLimitedIfLater(ctx, account.ID, shortReset); err != nil {
|
||||
slog.Warn("rate_limit_ollama_429_iflater_failed", "account_id", account.ID, "error", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
reset := shortReset
|
||||
if account.RateLimitResetAt != nil && account.RateLimitResetAt.After(reset) {
|
||||
reset = *account.RateLimitResetAt
|
||||
}
|
||||
if err := s.accountRepo.SetRateLimited(ctx, account.ID, reset); err != nil {
|
||||
slog.Warn("rate_limit_set_failed", "account_id", account.ID, "error", err)
|
||||
}
|
||||
}
|
||||
|
||||
// scheduleOllamaCloudUsageProbe captures the trigger fingerprint and rate-limit
|
||||
// generation from the authoritative row and schedules the async probe bound to
|
||||
// them.
|
||||
func (s *RateLimitService) scheduleOllamaCloudUsageProbe(account *Account) {
|
||||
if s == nil || account == nil || s.ollamaCloudUsageProbe == nil {
|
||||
return
|
||||
}
|
||||
if !IsOllamaCloudUsageAccount(account) {
|
||||
return
|
||||
}
|
||||
fingerprint, valid := ollamaCloudUsageGroupFingerprint(account)
|
||||
if !valid {
|
||||
return
|
||||
}
|
||||
expectedLimitedAt := cloneTimePtr(account.RateLimitedAt)
|
||||
expectedResetAt := cloneTimePtr(account.RateLimitResetAt)
|
||||
accepted := s.ollamaCloudUsageProbe.ScheduleOllamaCloudUsageRateLimitProbe(
|
||||
account.ID,
|
||||
func(accountID int64, resetAt time.Time) {
|
||||
s.applyOllamaCloudUsageProbeReset(accountID, fingerprint, expectedLimitedAt, expectedResetAt, resetAt)
|
||||
},
|
||||
)
|
||||
if !accepted {
|
||||
// Service stopped or the probe queue is full; the account is already
|
||||
// cooled by the immediate cooldown, so no retry or other action is needed.
|
||||
slog.Debug("ollama_cloud_usage_probe_schedule_rejected", "account_id", account.ID)
|
||||
}
|
||||
}
|
||||
|
||||
// applyOllamaCloudUsageProbeReset writes the reset learned by an async usage
|
||||
// probe, but only while the account is still the same active Ollama Cloud usage
|
||||
// account with the same rate-limit generation observed at scheduling time. It
|
||||
// runs on the probe coordinator with a detached, timeout-bounded context and does
|
||||
// at most one bounded re-read plus one bounded atomic CAS.
|
||||
//
|
||||
// The result is dropped when any of the following holds:
|
||||
// - the reset is not in the future, or the account is gone;
|
||||
// - the account is no longer an active/schedulable Ollama Cloud usage account
|
||||
// (disabled/errored since the trigger) or its group fingerprint changed;
|
||||
// - the reset does not extend the trigger's cooldown floor (an explicit
|
||||
// Retry-After or an already-confirmed cooldown must never be shortened);
|
||||
// - the CAS reports no update (the rate-limit generation moved on: a newer 429,
|
||||
// an admin clear, a re-arm, or a concurrent update between the re-read and
|
||||
// this write, guarded by the row's UpdatedAt).
|
||||
//
|
||||
// On success the CAS already updated the DB; only the runtime scheduling
|
||||
// notification is sent here.
|
||||
func (s *RateLimitService) applyOllamaCloudUsageProbeReset(
|
||||
accountID int64,
|
||||
expectedFingerprint string,
|
||||
expectedLimitedAt, expectedResetAt *time.Time,
|
||||
resetAt time.Time,
|
||||
) {
|
||||
now := time.Now()
|
||||
if s == nil || accountID <= 0 || s.accountRepo == nil || !resetAt.After(now) {
|
||||
return
|
||||
}
|
||||
// Never write an earlier reset than the trigger floor captured at schedule.
|
||||
if expectedResetAt != nil && !resetAt.After(*expectedResetAt) {
|
||||
return
|
||||
}
|
||||
|
||||
bgCtx, cancel := context.WithTimeout(context.WithoutCancel(context.Background()), ollamaCloudUsageProbeWritebackTimeout)
|
||||
defer cancel()
|
||||
|
||||
account, err := s.accountRepo.GetByID(bgCtx, accountID)
|
||||
if err != nil || account == nil {
|
||||
slog.Warn("ollama_cloud_usage_probe_reset_load_failed", "account_id", accountID, "error", err)
|
||||
return
|
||||
}
|
||||
if !IsOllamaCloudUsageAccount(account) {
|
||||
return
|
||||
}
|
||||
if !account.IsActive() || !account.Schedulable {
|
||||
return
|
||||
}
|
||||
currentFingerprint, valid := ollamaCloudUsageGroupFingerprint(account)
|
||||
if !valid || currentFingerprint != expectedFingerprint {
|
||||
return
|
||||
}
|
||||
|
||||
setter, ok := s.accountRepo.(ollamaCloudUsageRateLimitSetterIfGeneration)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
// Bind the CAS to the version just re-read (not the scheduling-time version,
|
||||
// which the probe's own snapshot persist legitimately bumps via UpdatedAt).
|
||||
// Any change between this re-read and the CAS therefore fails the write.
|
||||
updated, err := setter.SetRateLimitedIfUnchanged(bgCtx, accountID, account.UpdatedAt, expectedLimitedAt, expectedResetAt, resetAt)
|
||||
if err != nil {
|
||||
slog.Warn("rate_limit_set_failed", "account_id", accountID, "error", err)
|
||||
return
|
||||
}
|
||||
if !updated {
|
||||
// The generation moved on (a newer 429, admin clear/re-arm, or an
|
||||
// unrelated write bumped UpdatedAt): conservatively skip, never override.
|
||||
// No credentials are logged. A later 429 for the same account schedules a
|
||||
// fresh probe with a fresh generation, so no retry is needed here.
|
||||
slog.Debug("ollama_cloud_usage_probe_reset_skipped_stale", "account_id", accountID)
|
||||
return
|
||||
}
|
||||
|
||||
s.notifyAccountSchedulingBlocked(account, resetAt, "ollama_cloud_usage_429_probe")
|
||||
slog.Info("ollama_cloud_account_rate_limited_probe",
|
||||
"account_id", accountID,
|
||||
"reset_at", resetAt.UTC(),
|
||||
"reset_in", time.Until(resetAt).Truncate(time.Second),
|
||||
)
|
||||
}
|
||||
@@ -0,0 +1,721 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// This file unit-tests the Ollama Cloud usage 429 branch added to
|
||||
// RateLimitService.handle429 (see ratelimit_service_ollama_429.go): early
|
||||
// platform diversion, immediate never-shrink cooldown (Retry-After vs seconds
|
||||
// fallback), scheduling rejection / scheduler-absent runtime notification, and,
|
||||
// for the async probe callback, that a stale result can never overwrite a newer
|
||||
// 429, admin clear, key swap, disabled account, or a shorter cooldown floor,
|
||||
// guarded atomically via the repository CAS.
|
||||
|
||||
func ollama429Account(id int64, platform string) *Account {
|
||||
return &Account{
|
||||
ID: id,
|
||||
Platform: platform,
|
||||
Type: AccountTypeAPIKey,
|
||||
Status: StatusActive,
|
||||
Schedulable: true,
|
||||
Credentials: map[string]any{
|
||||
"base_url": "https://www.ollama.com",
|
||||
"api_key": "ollama429-key",
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// ollama429SchedulerStub records Schedule calls so tests can drive the recorded
|
||||
// probe callback (as the background coordinator would) deterministically.
|
||||
type ollama429SchedulerStub struct {
|
||||
mu sync.Mutex
|
||||
accept bool
|
||||
scheduled []int64
|
||||
callbacks map[int64]OllamaCloudUsageRateLimitProbeCallback
|
||||
}
|
||||
|
||||
func newOllama429SchedulerStub(accept bool) *ollama429SchedulerStub {
|
||||
return &ollama429SchedulerStub{accept: accept, callbacks: map[int64]OllamaCloudUsageRateLimitProbeCallback{}}
|
||||
}
|
||||
|
||||
func (p *ollama429SchedulerStub) ScheduleOllamaCloudUsageRateLimitProbe(accountID int64, cb OllamaCloudUsageRateLimitProbeCallback) bool {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
if !p.accept {
|
||||
return false
|
||||
}
|
||||
p.scheduled = append(p.scheduled, accountID)
|
||||
p.callbacks[accountID] = cb // latest wins, mirroring the coordinator coalescing
|
||||
return true
|
||||
}
|
||||
|
||||
func (p *ollama429SchedulerStub) fire(accountID int64, resetAt time.Time) {
|
||||
p.mu.Lock()
|
||||
cb := p.callbacks[accountID]
|
||||
p.mu.Unlock()
|
||||
if cb != nil {
|
||||
cb(accountID, resetAt)
|
||||
}
|
||||
}
|
||||
|
||||
func (p *ollama429SchedulerStub) count() int {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
return len(p.scheduled)
|
||||
}
|
||||
|
||||
// ollama429BlockRec records an in-memory runtime scheduling block.
|
||||
type ollama429BlockRec struct {
|
||||
accountID int64
|
||||
until time.Time
|
||||
reason string
|
||||
}
|
||||
|
||||
type ollama429BlockerStub struct {
|
||||
mu sync.Mutex
|
||||
blocks []ollama429BlockRec
|
||||
}
|
||||
|
||||
func (b *ollama429BlockerStub) BlockAccountScheduling(account *Account, until time.Time, reason string) {
|
||||
b.mu.Lock()
|
||||
defer b.mu.Unlock()
|
||||
if account != nil {
|
||||
b.blocks = append(b.blocks, ollama429BlockRec{accountID: account.ID, until: until, reason: reason})
|
||||
}
|
||||
}
|
||||
|
||||
func (b *ollama429BlockerStub) ClearAccountSchedulingBlock(int64) {}
|
||||
|
||||
func (b *ollama429BlockerStub) last() (ollama429BlockRec, bool) {
|
||||
b.mu.Lock()
|
||||
defer b.mu.Unlock()
|
||||
if len(b.blocks) == 0 {
|
||||
return ollama429BlockRec{}, false
|
||||
}
|
||||
return b.blocks[len(b.blocks)-1], true
|
||||
}
|
||||
|
||||
// ollama429Repo is a controllable AccountRepository fake implementing the narrow
|
||||
// IfLater + generation-CAS contracts the production *accountRepository provides.
|
||||
// updatedAt is simulated by bumping Account.UpdatedAt forward on every write.
|
||||
type ollama429Repo struct {
|
||||
mockAccountRepoForGemini
|
||||
mu sync.Mutex
|
||||
accounts map[int64]*Account
|
||||
casUpdated int
|
||||
casSkipped int
|
||||
ifLaterWrites []time.Time
|
||||
uncondWrites []time.Time
|
||||
versionCounter int64
|
||||
}
|
||||
|
||||
func newOllama429Repo(acct *Account) *ollama429Repo {
|
||||
r := &ollama429Repo{accounts: map[int64]*Account{}}
|
||||
r.accounts[acct.ID] = acct
|
||||
return r
|
||||
}
|
||||
|
||||
func (r *ollama429Repo) bump(acct *Account) {
|
||||
r.versionCounter++
|
||||
acct.UpdatedAt = time.Unix(0, r.versionCounter).UTC()
|
||||
}
|
||||
|
||||
func (r *ollama429Repo) GetByID(_ context.Context, id int64) (*Account, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
return r.accounts[id], nil
|
||||
}
|
||||
|
||||
func (r *ollama429Repo) currentReset(id int64) *time.Time {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
a := r.accounts[id]
|
||||
if a == nil || a.RateLimitResetAt == nil {
|
||||
return nil
|
||||
}
|
||||
v := *a.RateLimitResetAt
|
||||
return &v
|
||||
}
|
||||
|
||||
// mutate lets a test simulate an external change (admin clear / re-arm / key swap).
|
||||
func (r *ollama429Repo) mutate(id int64, fn func(*Account)) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
if a, ok := r.accounts[id]; ok {
|
||||
fn(a)
|
||||
r.bump(a)
|
||||
}
|
||||
}
|
||||
|
||||
// bumpVersion simulates the real probe persisting its usage snapshot: it advances
|
||||
// the row's UpdatedAt (TimeMixin bump) without touching rate-limit fields, which
|
||||
// is exactly what updateSnapshot does between scheduling and callback delivery.
|
||||
func (r *ollama429Repo) bumpVersion(id int64) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
if a, ok := r.accounts[id]; ok {
|
||||
r.bump(a)
|
||||
}
|
||||
}
|
||||
|
||||
// SetRateLimitedIfLater implements the never-shrink narrow extender.
|
||||
func (r *ollama429Repo) SetRateLimitedIfLater(_ context.Context, id int64, resetAt time.Time) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
a := r.accounts[id]
|
||||
if a == nil {
|
||||
return nil
|
||||
}
|
||||
cur := a.RateLimitResetAt
|
||||
if cur == nil || resetAt.After(*cur) {
|
||||
a.RateLimitResetAt = ollama429TimePtr(resetAt)
|
||||
now := time.Now()
|
||||
a.RateLimitedAt = &now
|
||||
r.bump(a)
|
||||
r.ifLaterWrites = append(r.ifLaterWrites, resetAt)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *ollama429Repo) SetRateLimited(_ context.Context, id int64, resetAt time.Time) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
a := r.accounts[id]
|
||||
if a == nil {
|
||||
return nil
|
||||
}
|
||||
a.RateLimitResetAt = ollama429TimePtr(resetAt)
|
||||
now := time.Now()
|
||||
a.RateLimitedAt = &now
|
||||
r.bump(a)
|
||||
r.uncondWrites = append(r.uncondWrites, resetAt)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *ollama429Repo) SetRateLimitedIfUnchanged(
|
||||
_ context.Context, id int64,
|
||||
expectedUpdatedAt time.Time,
|
||||
expectedLimitedAt, expectedResetAt *time.Time,
|
||||
newResetAt time.Time,
|
||||
) (bool, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
a := r.accounts[id]
|
||||
if a == nil {
|
||||
r.casSkipped++
|
||||
return false, nil
|
||||
}
|
||||
limitedMatch := timePtrEqual(expectedLimitedAt, a.RateLimitedAt)
|
||||
resetMatch := timePtrEqual(expectedResetAt, a.RateLimitResetAt)
|
||||
if !a.UpdatedAt.Equal(expectedUpdatedAt) || !limitedMatch || !resetMatch {
|
||||
r.casSkipped++
|
||||
return false, nil
|
||||
}
|
||||
a.RateLimitedAt = ollama429TimePtr(time.Now())
|
||||
a.RateLimitResetAt = ollama429TimePtr(newResetAt)
|
||||
r.bump(a)
|
||||
r.casUpdated++
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func ollama429TimePtr(t time.Time) *time.Time { return &t }
|
||||
|
||||
func timePtrEqual(a, b *time.Time) bool {
|
||||
if a == nil || b == nil {
|
||||
return a == nil && b == nil
|
||||
}
|
||||
return a.Equal(*b)
|
||||
}
|
||||
|
||||
// ollama429Fixture builds a started-ish RateLimitService over the given repo and
|
||||
// scheduler. runtime blocker is injected so scheduling notifications are visible.
|
||||
func ollama429Fixture(t *testing.T, repo *ollama429Repo, scheduler *ollama429SchedulerStub) (*RateLimitService, *ollama429BlockerStub) {
|
||||
t.Helper()
|
||||
svc := NewRateLimitService(repo, nil, nil, nil, nil)
|
||||
blocker := &ollama429BlockerStub{}
|
||||
svc.SetAccountRuntimeBlocker(blocker)
|
||||
svc.SetOllamaCloudUsageProbeScheduler(scheduler)
|
||||
return svc, blocker
|
||||
}
|
||||
|
||||
func TestHandle429_OllamaEarlyBranchAcrossPlatforms(t *testing.T) {
|
||||
platforms := []string{PlatformOpenAI, PlatformAnthropic, PlatformKimi, PlatformZhipu, PlatformDeepseek}
|
||||
for _, platform := range platforms {
|
||||
t.Run(platform, func(t *testing.T) {
|
||||
acct := ollama429Account(101, platform)
|
||||
repo := newOllama429Repo(acct)
|
||||
scheduler := newOllama429SchedulerStub(true)
|
||||
svc, _ := ollama429Fixture(t, repo, scheduler)
|
||||
|
||||
svc.handle429(context.Background(), acct, http.Header{}, nil)
|
||||
|
||||
require.Equal(t, 1, scheduler.count(), "real-Ollama %s account must schedule a probe", platform)
|
||||
require.Greater(t, len(repo.ifLaterWrites), 0, "immediate never-shrink cooldown must be applied")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandle429_NonOllamaUnchanged(t *testing.T) {
|
||||
acct := &Account{ID: 202, Platform: PlatformAnthropic, Type: AccountTypeAPIKey,
|
||||
Credentials: map[string]any{"base_url": "https://api.anthropic.com", "api_key": "k"}}
|
||||
repo := newOllama429Repo(acct)
|
||||
scheduler := newOllama429SchedulerStub(true)
|
||||
svc, _ := ollama429Fixture(t, repo, scheduler)
|
||||
|
||||
svc.handle429(context.Background(), acct, http.Header{}, nil)
|
||||
|
||||
require.Zero(t, scheduler.count(), "non-Ollama account must not schedule an Ollama probe")
|
||||
require.Zero(t, repo.casUpdated)
|
||||
}
|
||||
|
||||
func TestHandle429_OllamaUsesValidRetryAfter(t *testing.T) {
|
||||
acct := ollama429Account(301, PlatformOpenAI)
|
||||
repo := newOllama429Repo(acct)
|
||||
scheduler := newOllama429SchedulerStub(true)
|
||||
svc, _ := ollama429Fixture(t, repo, scheduler)
|
||||
|
||||
before := time.Now()
|
||||
svc.handle429(context.Background(), acct, http.Header{"Retry-After": []string{"120"}}, nil)
|
||||
after := time.Now()
|
||||
|
||||
require.Equal(t, 1, scheduler.count())
|
||||
reset := repo.currentReset(acct.ID)
|
||||
require.NotNil(t, reset)
|
||||
require.False(t, reset.Before(before.Add(120*time.Second)), "reset %v < now+120s", reset)
|
||||
require.False(t, reset.After(after.Add(120*time.Second)), "reset %v > now+120s", reset)
|
||||
}
|
||||
|
||||
func TestHandle429_OllamaFallsBackToSecondsCooldown(t *testing.T) {
|
||||
acct := ollama429Account(302, PlatformAnthropic)
|
||||
repo := newOllama429Repo(acct)
|
||||
scheduler := newOllama429SchedulerStub(true)
|
||||
svc, _ := ollama429Fixture(t, repo, scheduler)
|
||||
|
||||
before := time.Now()
|
||||
svc.handle429(context.Background(), acct, http.Header{}, nil) // no Retry-After
|
||||
after := time.Now()
|
||||
|
||||
require.Equal(t, 1, scheduler.count())
|
||||
reset := repo.currentReset(acct.ID)
|
||||
require.NotNil(t, reset)
|
||||
def := defaultRateLimit429CooldownSeconds
|
||||
require.False(t, reset.Before(before.Add(time.Duration(def)*time.Second)))
|
||||
require.False(t, reset.After(after.Add(time.Duration(def)*time.Second)))
|
||||
}
|
||||
|
||||
func TestHandle429_OllamaScheduleRejectedStillCooled(t *testing.T) {
|
||||
acct := ollama429Account(303, PlatformOpenAI)
|
||||
repo := newOllama429Repo(acct)
|
||||
scheduler := newOllama429SchedulerStub(false) // queue full / service stopped
|
||||
svc, _ := ollama429Fixture(t, repo, scheduler)
|
||||
|
||||
svc.handle429(context.Background(), acct, http.Header{}, nil)
|
||||
|
||||
require.Zero(t, scheduler.count())
|
||||
require.Greater(t, len(repo.ifLaterWrites), 0, "immediate cooldown must still be applied when scheduling is rejected")
|
||||
}
|
||||
|
||||
func TestHandle429_OllamaDoesNotShrinkConfirmedLongCooldown(t *testing.T) {
|
||||
acct := ollama429Account(304, PlatformOpenAI)
|
||||
repo := newOllama429Repo(acct)
|
||||
long := time.Now().Add(time.Hour)
|
||||
repo.mutate(acct.ID, func(a *Account) { a.RateLimitResetAt = &long })
|
||||
scheduler := newOllama429SchedulerStub(true)
|
||||
svc, _ := ollama429Fixture(t, repo, scheduler)
|
||||
|
||||
// A 429 with no useful Retry-After arrives while a +1h reset is already set;
|
||||
// the 5s fallback must not shrink it.
|
||||
svc.handle429(context.Background(), acct, http.Header{}, nil)
|
||||
|
||||
reset := repo.currentReset(acct.ID)
|
||||
require.NotNil(t, reset)
|
||||
require.GreaterOrEqual(t, reset.Unix(), long.Unix(), "already-confirmed long cooldown must not be shrunk to fallback seconds")
|
||||
require.Equal(t, 1, scheduler.count())
|
||||
}
|
||||
|
||||
func TestOllamaProbeCallback_ValidUpdatesDBAndMemory(t *testing.T) {
|
||||
acct := ollama429Account(401, PlatformOpenAI)
|
||||
repo := newOllama429Repo(acct)
|
||||
scheduler := newOllama429SchedulerStub(true)
|
||||
svc, blocker := ollama429Fixture(t, repo, scheduler)
|
||||
|
||||
svc.handle429(context.Background(), acct, http.Header{}, nil) // schedules; immediate 5s
|
||||
require.Equal(t, 1, scheduler.count())
|
||||
|
||||
probeReset := time.Now().Add(2 * time.Hour)
|
||||
scheduler.fire(acct.ID, probeReset)
|
||||
|
||||
require.Equal(t, 1, repo.casUpdated, "valid callback must update DB once")
|
||||
require.Zero(t, repo.casSkipped)
|
||||
reset := repo.currentReset(acct.ID)
|
||||
require.NotNil(t, reset)
|
||||
require.True(t, reset.Equal(probeReset), "DB reset = %v, want %v", reset, probeReset)
|
||||
rec, ok := blocker.last()
|
||||
require.True(t, ok)
|
||||
require.Equal(t, "ollama_cloud_usage_429_probe", rec.reason)
|
||||
require.True(t, rec.until.Equal(probeReset))
|
||||
}
|
||||
|
||||
func TestOllamaProbeCallback_DuplicateNoDoubleUpdate(t *testing.T) {
|
||||
acct := ollama429Account(402, PlatformOpenAI)
|
||||
repo := newOllama429Repo(acct)
|
||||
scheduler := newOllama429SchedulerStub(true)
|
||||
svc, blocker := ollama429Fixture(t, repo, scheduler)
|
||||
|
||||
svc.handle429(context.Background(), acct, http.Header{}, nil)
|
||||
probeReset := time.Now().Add(2 * time.Hour)
|
||||
scheduler.fire(acct.ID, probeReset)
|
||||
require.Equal(t, 1, repo.casUpdated)
|
||||
|
||||
// Same generation fired again (a duplicate/out-of-order delivery): the CAS
|
||||
// generation has moved on, so no second update and no second notification.
|
||||
scheduler.fire(acct.ID, probeReset.Add(time.Hour))
|
||||
require.Equal(t, 1, repo.casUpdated)
|
||||
rec, _ := blocker.last()
|
||||
require.Equal(t, "ollama_cloud_usage_429_probe", rec.reason)
|
||||
require.True(t, rec.until.Equal(probeReset))
|
||||
}
|
||||
|
||||
func TestOllamaProbeCallback_StaleLongDoesNotOverrideNewShort(t *testing.T) {
|
||||
acct := ollama429Account(403, PlatformOpenAI)
|
||||
repo := newOllama429Repo(acct)
|
||||
scheduler := newOllama429SchedulerStub(true)
|
||||
svc, blocker := ollama429Fixture(t, repo, scheduler)
|
||||
|
||||
svc.handle429(context.Background(), acct, http.Header{}, nil) // schedules gen @5s
|
||||
require.Equal(t, 1, scheduler.count())
|
||||
|
||||
// An admin / newer policy re-arms the account to a fresh SHORT cooldown.
|
||||
newShort := time.Now().Add(5 * time.Second)
|
||||
repo.mutate(acct.ID, func(a *Account) { a.RateLimitResetAt = ollama429TimePtr(newShort) })
|
||||
|
||||
// The old async result reports a long 7d reset; it must not override.
|
||||
oldLong := time.Now().Add(7 * 24 * time.Hour)
|
||||
scheduler.fire(acct.ID, oldLong)
|
||||
|
||||
require.Zero(t, repo.casUpdated, "stale long callback must not pass the CAS")
|
||||
require.Greater(t, repo.casSkipped, 0)
|
||||
reset := repo.currentReset(acct.ID)
|
||||
require.NotNil(t, reset)
|
||||
require.True(t, reset.Equal(newShort), "current reset must remain the new short %v, got %v", newShort, reset)
|
||||
rec, ok := blocker.last()
|
||||
// The only scheduling notification should be the immediate cooldown one, not
|
||||
// an ollama_cloud_usage_429_probe block for the stale 7d reset.
|
||||
if ok && rec.reason == "ollama_cloud_usage_429_probe" {
|
||||
t.Fatalf("stale long callback must not announce scheduling; got block to %v", rec.until)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOllamaProbeCallback_AdminClearSkips(t *testing.T) {
|
||||
acct := ollama429Account(404, PlatformOpenAI)
|
||||
repo := newOllama429Repo(acct)
|
||||
scheduler := newOllama429SchedulerStub(true)
|
||||
svc, blocker := ollama429Fixture(t, repo, scheduler)
|
||||
|
||||
svc.handle429(context.Background(), acct, http.Header{}, nil)
|
||||
// Admin clears the rate limit entirely.
|
||||
repo.mutate(acct.ID, func(a *Account) {
|
||||
a.RateLimitedAt = nil
|
||||
a.RateLimitResetAt = nil
|
||||
})
|
||||
|
||||
scheduler.fire(acct.ID, time.Now().Add(2*time.Hour))
|
||||
require.Zero(t, repo.casUpdated, "cleared account must not be re-limited by a stale callback")
|
||||
require.Nil(t, repo.currentReset(acct.ID))
|
||||
rec, ok := blocker.last()
|
||||
if ok && rec.reason == "ollama_cloud_usage_429_probe" {
|
||||
t.Fatalf("cleared account must not be re-blocked by a stale callback")
|
||||
}
|
||||
}
|
||||
|
||||
func TestOllamaProbeCallback_KeySwapSkips(t *testing.T) {
|
||||
acct := ollama429Account(405, PlatformOpenAI)
|
||||
repo := newOllama429Repo(acct)
|
||||
scheduler := newOllama429SchedulerStub(true)
|
||||
svc, _ := ollama429Fixture(t, repo, scheduler)
|
||||
|
||||
svc.handle429(context.Background(), acct, http.Header{}, nil) // captures fingerprint of k1
|
||||
// The api_key is swapped to another key (still on ollama.com).
|
||||
repo.mutate(acct.ID, func(a *Account) { a.Credentials["api_key"] = "ollama429-key-2" })
|
||||
|
||||
scheduler.fire(acct.ID, time.Now().Add(2*time.Hour))
|
||||
require.Zero(t, repo.casUpdated, "key-swapped account must not be limited by the old key's probe")
|
||||
// The fingerprint guard returns before any CAS is attempted, so the current
|
||||
// reset stays at the immediate cooldown, not the probe's far-future reset.
|
||||
reset := repo.currentReset(acct.ID)
|
||||
require.NotNil(t, reset)
|
||||
require.False(t, reset.After(time.Now().Add(time.Hour)), "key-swapped account must keep only its immediate cooldown, got %v", reset)
|
||||
}
|
||||
|
||||
func TestOllamaProbeCallback_ResetAlreadyPastSkips(t *testing.T) {
|
||||
acct := ollama429Account(406, PlatformOpenAI)
|
||||
repo := newOllama429Repo(acct)
|
||||
scheduler := newOllama429SchedulerStub(true)
|
||||
svc, _ := ollama429Fixture(t, repo, scheduler)
|
||||
|
||||
svc.handle429(context.Background(), acct, http.Header{}, nil)
|
||||
repo.mutate(acct.ID, func(a *Account) {
|
||||
// keep same reset/state, just ensure callback reset is in the past
|
||||
})
|
||||
scheduler.fire(acct.ID, time.Now().Add(-time.Minute))
|
||||
|
||||
require.Zero(t, repo.casUpdated)
|
||||
require.Zero(t, repo.casSkipped, "an already-passed reset is dropped before any CAS is attempted")
|
||||
}
|
||||
|
||||
func TestOllamaProbeCallback_SnapshotPersistBumpStillApplies(t *testing.T) {
|
||||
// The real probe persists its usage snapshot between scheduling and callback
|
||||
// delivery, which advances the row's UpdatedAt (TimeMixin) without touching
|
||||
// the rate-limit generation. The write-back must still succeed: it binds the
|
||||
// CAS to the version re-read in the callback, not the scheduling-time one.
|
||||
acct := ollama429Account(501, PlatformOpenAI)
|
||||
repo := newOllama429Repo(acct)
|
||||
scheduler := newOllama429SchedulerStub(true)
|
||||
svc, blocker := ollama429Fixture(t, repo, scheduler)
|
||||
|
||||
svc.handle429(context.Background(), acct, http.Header{}, nil)
|
||||
require.Equal(t, 1, scheduler.count())
|
||||
|
||||
repo.bumpVersion(acct.ID) // simulate updateSnapshot persisting the snapshot
|
||||
|
||||
probeReset := time.Now().Add(2 * time.Hour)
|
||||
scheduler.fire(acct.ID, probeReset)
|
||||
|
||||
require.Equal(t, 1, repo.casUpdated, "snapshot persist bumping UpdatedAt must not fail the exhaustion CAS")
|
||||
require.Zero(t, repo.casSkipped)
|
||||
reset := repo.currentReset(acct.ID)
|
||||
require.NotNil(t, reset)
|
||||
require.True(t, reset.Equal(probeReset), "DB reset = %v, want %v", reset, probeReset)
|
||||
rec, ok := blocker.last()
|
||||
require.True(t, ok)
|
||||
require.Equal(t, "ollama_cloud_usage_429_probe", rec.reason)
|
||||
}
|
||||
|
||||
func TestOllamaProbeCallback_NeverShortensCooldownFloor(t *testing.T) {
|
||||
// An explicit Retry-After sets a 120s floor. A later async probe that reports
|
||||
// a shorter (still-future) reset must not shrink that floor back down.
|
||||
acct := ollama429Account(502, PlatformOpenAI)
|
||||
repo := newOllama429Repo(acct)
|
||||
scheduler := newOllama429SchedulerStub(true)
|
||||
svc, _ := ollama429Fixture(t, repo, scheduler)
|
||||
|
||||
before := time.Now()
|
||||
svc.handle429(context.Background(), acct, http.Header{"Retry-After": []string{"120"}}, nil)
|
||||
after := time.Now()
|
||||
floor := repo.currentReset(acct.ID)
|
||||
require.NotNil(t, floor)
|
||||
require.False(t, floor.Before(before.Add(120*time.Second)))
|
||||
require.False(t, floor.After(after.Add(120*time.Second)))
|
||||
|
||||
// Probe reports a reset still in the future but sooner than the 120s floor.
|
||||
scheduler.fire(acct.ID, time.Now().Add(30*time.Second))
|
||||
|
||||
require.Zero(t, repo.casUpdated, "a probe reset sooner than the explicit floor must be dropped")
|
||||
still := repo.currentReset(acct.ID)
|
||||
require.True(t, still.Equal(*floor), "the explicit Retry-After floor must be preserved")
|
||||
}
|
||||
|
||||
func TestOllamaProbeCallback_DisabledSkips(t *testing.T) {
|
||||
acct := ollama429Account(503, PlatformOpenAI)
|
||||
repo := newOllama429Repo(acct)
|
||||
scheduler := newOllama429SchedulerStub(true)
|
||||
svc, _ := ollama429Fixture(t, repo, scheduler)
|
||||
|
||||
svc.handle429(context.Background(), acct, http.Header{}, nil)
|
||||
repo.mutate(acct.ID, func(a *Account) { a.Status = StatusDisabled })
|
||||
|
||||
scheduler.fire(acct.ID, time.Now().Add(2*time.Hour))
|
||||
require.Zero(t, repo.casUpdated, "a disabled account must not be re-limited by a stale callback")
|
||||
}
|
||||
|
||||
func TestHandle429_OllamaSchedulerAbsentStillNotifiesRuntime(t *testing.T) {
|
||||
acct := ollama429Account(504, PlatformOpenAI)
|
||||
repo := newOllama429Repo(acct)
|
||||
svc := NewRateLimitService(repo, nil, nil, nil, nil)
|
||||
blocker := &ollama429BlockerStub{}
|
||||
svc.SetAccountRuntimeBlocker(blocker)
|
||||
// No probe scheduler injected.
|
||||
|
||||
before := time.Now()
|
||||
svc.handle429(context.Background(), acct, http.Header{}, nil)
|
||||
after := time.Now()
|
||||
|
||||
// Immediate cooldown applied and runtime scheduling announced even without a
|
||||
// scheduler; only the async probe is skipped.
|
||||
require.Greater(t, len(repo.ifLaterWrites), 0)
|
||||
rec, ok := blocker.last()
|
||||
require.True(t, ok)
|
||||
require.Equal(t, "ollama_429", rec.reason)
|
||||
require.False(t, rec.until.Before(before.Add(defaultRateLimit429CooldownSeconds*time.Second)))
|
||||
require.False(t, rec.until.After(after.Add(defaultRateLimit429CooldownSeconds*time.Second)))
|
||||
}
|
||||
|
||||
// ollama429LinkRepo drives a real probe->callback linkage through the actual
|
||||
// OllamaCloudUsageService. It wraps the shared ollama feature test repo, adds the
|
||||
// narrow IfLater/CAS surface RateLimitService needs, and makes the real snapshot
|
||||
// persist (UpdateOllamaCloudUsageSnapshot) advance the row's UpdatedAt exactly as
|
||||
// the production TimeMixin-backed repository does.
|
||||
type ollama429LinkRepo struct {
|
||||
*ollamaUsageTestRepo
|
||||
mu sync.Mutex
|
||||
versionCounter int64
|
||||
casUpdated int
|
||||
linkIfLater []time.Time
|
||||
}
|
||||
|
||||
func (r *ollama429LinkRepo) bump(a *Account) {
|
||||
r.versionCounter++
|
||||
a.UpdatedAt = time.Unix(0, r.versionCounter).UTC()
|
||||
}
|
||||
|
||||
func (r *ollama429LinkRepo) UpdateOllamaCloudUsageSnapshot(ctx context.Context, expected *Account, snapshot *OllamaCloudUsageSnapshot) error {
|
||||
if err := r.ollamaUsageTestRepo.UpdateOllamaCloudUsageSnapshot(ctx, expected, snapshot); err != nil {
|
||||
return err
|
||||
}
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
if a, ok := r.upstreamBillingProbeAccountRepo.accounts[expected.ID]; ok {
|
||||
r.bump(a)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *ollama429LinkRepo) SetRateLimitedIfLater(_ context.Context, id int64, resetAt time.Time) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
a := r.upstreamBillingProbeAccountRepo.accounts[id]
|
||||
if a == nil {
|
||||
return nil
|
||||
}
|
||||
cur := a.RateLimitResetAt
|
||||
if cur == nil || resetAt.After(*cur) {
|
||||
a.RateLimitResetAt = ollama429TimePtr(resetAt)
|
||||
now := time.Now()
|
||||
a.RateLimitedAt = &now
|
||||
r.bump(a)
|
||||
r.linkIfLater = append(r.linkIfLater, resetAt)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *ollama429LinkRepo) SetRateLimitedIfUnchanged(
|
||||
_ context.Context, id int64,
|
||||
expectedUpdatedAt time.Time,
|
||||
expectedLimitedAt, expectedResetAt *time.Time,
|
||||
newResetAt time.Time,
|
||||
) (bool, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
a := r.upstreamBillingProbeAccountRepo.accounts[id]
|
||||
if a == nil {
|
||||
return false, nil
|
||||
}
|
||||
if !a.UpdatedAt.Equal(expectedUpdatedAt) ||
|
||||
!timePtrEqual(expectedLimitedAt, a.RateLimitedAt) ||
|
||||
!timePtrEqual(expectedResetAt, a.RateLimitResetAt) {
|
||||
return false, nil
|
||||
}
|
||||
a.RateLimitedAt = ollama429TimePtr(time.Now())
|
||||
a.RateLimitResetAt = ollama429TimePtr(newResetAt)
|
||||
r.bump(a)
|
||||
r.casUpdated++
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func (r *ollama429LinkRepo) currentReset(id int64) *time.Time {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
a := r.upstreamBillingProbeAccountRepo.accounts[id]
|
||||
if a == nil || a.RateLimitResetAt == nil {
|
||||
return nil
|
||||
}
|
||||
v := *a.RateLimitResetAt
|
||||
return &v
|
||||
}
|
||||
|
||||
func (r *ollama429LinkRepo) casUpdatedCount() int {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
return r.casUpdated
|
||||
}
|
||||
|
||||
// TestOllama429RealProbeLinkage_SnapshotPersistThenWriteBack drives the full real
|
||||
// chain: handle429 schedules the async probe on a real started
|
||||
// OllamaCloudUsageService; the probe fetches an exhausted settings page, persists
|
||||
// its snapshot (which advances UpdatedAt through the link repo, mirroring the
|
||||
// production TimeMixin bump), and the exhaustion callback then writes the reset
|
||||
// back through the generation CAS. The CAS must still succeed even though the
|
||||
// snapshot persist moved UpdatedAt, because it is bound to the version re-read at
|
||||
// callback time, not the scheduling-time version.
|
||||
func TestOllama429RealProbeLinkage_SnapshotPersistThenWriteBack(t *testing.T) {
|
||||
account := ollamaUsageAccount(701)
|
||||
account.Extra[OllamaCloudUsageSessionExtraKey] = "cipher:wos-session=secret"
|
||||
// The probe is a 429-event recovery query and must run even when the periodic
|
||||
// auto_refresh switch is off (it still needs the configured cookie/session).
|
||||
account.Extra[OllamaCloudUsageAutoRefreshExtraKey] = false
|
||||
|
||||
reset := time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339)
|
||||
body := []byte(`
|
||||
<section>
|
||||
<div><span>5 hour usage</span><span>100% used</span></div>
|
||||
<time datetime="` + reset + `"></time>
|
||||
</section>`)
|
||||
|
||||
repo := &ollama429LinkRepo{ollamaUsageTestRepo: &ollamaUsageTestRepo{
|
||||
upstreamBillingProbeAccountRepo: &upstreamBillingProbeAccountRepo{
|
||||
accounts: map[int64]*Account{account.ID: account},
|
||||
},
|
||||
}}
|
||||
upstream := &ollamaUsageHTTPStub{body: body}
|
||||
|
||||
usageSvc := NewOllamaCloudUsageService(repo, upstream, NewSettingService(&upstreamBillingProbeSettingRepo{}, nil), ollamaUsageTestEncryptor{}, true)
|
||||
|
||||
// With auto_refresh off, a requireEnabled refresh is a (nil,nil) no-op. When
|
||||
// the 429-event probe piggybacks such a timed-cycle no-op through the shared
|
||||
// singleflight, runOllamaCloudUsageProbe must not record it as a real attempt.
|
||||
noop, noopErr := usageSvc.refreshAccount(context.Background(), account.ID, defaultOllamaCloudUsageSettings(), true)
|
||||
require.NoError(t, noopErr)
|
||||
require.Nil(t, noop, "auto_refresh-disabled refresh is a nil,nil no-op")
|
||||
|
||||
usageSvc.Start()
|
||||
t.Cleanup(usageSvc.Stop)
|
||||
|
||||
blocker := &ollama429BlockerStub{}
|
||||
rlSvc := NewRateLimitService(repo, nil, nil, nil, nil)
|
||||
rlSvc.SetAccountRuntimeBlocker(blocker)
|
||||
rlSvc.SetOllamaCloudUsageProbeScheduler(usageSvc)
|
||||
|
||||
rlSvc.handle429(context.Background(), account, http.Header{}, nil)
|
||||
require.Greater(t, len(repo.linkIfLater), 0, "immediate cooldown must be applied")
|
||||
|
||||
// The real probe coordinator (async) fetches, persists (bumping UpdatedAt),
|
||||
// reports exhaustion, and the write-back must then succeed. Reads go through
|
||||
// lock-protected accessors because the coordinator mutates the shared repo.
|
||||
require.Eventually(t, func() bool {
|
||||
if repo.casUpdatedCount() != 1 {
|
||||
return false
|
||||
}
|
||||
rec, ok := blocker.last()
|
||||
return ok && rec.reason == "ollama_cloud_usage_429_probe"
|
||||
}, 10*time.Second, 5*time.Millisecond, "exhaustion write-back through the real probe must update once and notify")
|
||||
|
||||
// The 429-event probe actually performed a usage fetch despite auto_refresh=off.
|
||||
require.Greater(t, upstream.calls.Load(), int64(0), "probe must query usage even with auto_refresh disabled")
|
||||
|
||||
rec, ok := blocker.last()
|
||||
require.True(t, ok)
|
||||
require.Equal(t, "ollama_cloud_usage_429_probe", rec.reason)
|
||||
require.True(t, rec.until.After(time.Now()), "probe reset must be in the future")
|
||||
resetAfter := repo.currentReset(account.ID)
|
||||
require.NotNil(t, resetAfter)
|
||||
require.True(t, resetAfter.After(time.Now()), "DB reset must be extended by the real probe write-back")
|
||||
}
|
||||
@@ -486,6 +486,7 @@ func ProvideRateLimitService(
|
||||
openAI403CounterCache OpenAI403CounterCache,
|
||||
settingService *SettingService,
|
||||
tokenCacheInvalidator TokenCacheInvalidator,
|
||||
ollamaCloudUsage *OllamaCloudUsageService,
|
||||
) *RateLimitService {
|
||||
svc := NewRateLimitService(accountRepo, usageRepo, cfg, geminiQuotaService, tempUnschedCache)
|
||||
if healthCache, ok := tempUnschedCache.(OpenAIAPIKeyHealthCache); ok {
|
||||
@@ -495,6 +496,7 @@ func ProvideRateLimitService(
|
||||
svc.SetOpenAI403CounterCache(openAI403CounterCache)
|
||||
svc.SetSettingService(settingService)
|
||||
svc.SetTokenCacheInvalidator(tokenCacheInvalidator)
|
||||
svc.SetOllamaCloudUsageProbeScheduler(ollamaCloudUsage)
|
||||
return svc
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user