diff --git a/backend/cmd/server/wire_gen.go b/backend/cmd/server/wire_gen.go
index b68a06ac54..e2adb3a17a 100644
--- a/backend/cmd/server/wire_gen.go
+++ b/backend/cmd/server/wire_gen.go
@@ -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)
diff --git a/backend/internal/repository/account_repo.go b/backend/internal/repository/account_repo.go
index 2e3961423a..db43cc7422 100644
--- a/backend/internal/repository/account_repo.go
+++ b/backend/internal/repository/account_repo.go
@@ -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
diff --git a/backend/internal/repository/account_repo_ratelimit_cas_integration_test.go b/backend/internal/repository/account_repo_ratelimit_cas_integration_test.go
new file mode 100644
index 0000000000..250a964034
--- /dev/null
+++ b/backend/internal/repository/account_repo_ratelimit_cas_integration_test.go
@@ -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")
+}
diff --git a/backend/internal/service/ollama_cloud_usage.go b/backend/internal/service/ollama_cloud_usage.go
index 08f1cbbcc6..a797246499 100644
--- a/backend/internal/service/ollama_cloud_usage.go
+++ b/backend/internal/service/ollama_cloud_usage.go
@@ -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() {
diff --git a/backend/internal/service/ollama_cloud_usage_rate_limit_probe.go b/backend/internal/service/ollama_cloud_usage_rate_limit_probe.go
new file mode 100644
index 0000000000..9d20e344ed
--- /dev/null
+++ b/backend/internal/service/ollama_cloud_usage_rate_limit_probe.go
@@ -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)
+}
diff --git a/backend/internal/service/ollama_cloud_usage_rate_limit_probe_test.go b/backend/internal/service/ollama_cloud_usage_rate_limit_probe_test.go
new file mode 100644
index 0000000000..d4fb57ba12
--- /dev/null
+++ b/backend/internal/service/ollama_cloud_usage_rate_limit_probe_test.go
@@ -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(`
+
+
5 hour usage%d%% used
+
+ `, 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)
+ })
+ }
+}
diff --git a/backend/internal/service/ollama_cloud_usage_reset.go b/backend/internal/service/ollama_cloud_usage_reset.go
new file mode 100644
index 0000000000..6de4880799
--- /dev/null
+++ b/backend/internal/service/ollama_cloud_usage_reset.go
@@ -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
+}
diff --git a/backend/internal/service/ollama_cloud_usage_reset_test.go b/backend/internal/service/ollama_cloud_usage_reset_test.go
new file mode 100644
index 0000000000..c717438cbd
--- /dev/null
+++ b/backend/internal/service/ollama_cloud_usage_reset_test.go
@@ -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)))
+}
diff --git a/backend/internal/service/ratelimit_service.go b/backend/internal/service/ratelimit_service.go
index 4158d0559c..370546af99 100644
--- a/backend/internal/service/ratelimit_service.go
+++ b/backend/internal/service/ratelimit_service.go
@@ -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() {
diff --git a/backend/internal/service/ratelimit_service_ollama_429.go b/backend/internal/service/ratelimit_service_ollama_429.go
new file mode 100644
index 0000000000..39116737bf
--- /dev/null
+++ b/backend/internal/service/ratelimit_service_ollama_429.go
@@ -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),
+ )
+}
diff --git a/backend/internal/service/ratelimit_service_ollama_429_test.go b/backend/internal/service/ratelimit_service_ollama_429_test.go
new file mode 100644
index 0000000000..dca1c28608
--- /dev/null
+++ b/backend/internal/service/ratelimit_service_ollama_429_test.go
@@ -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(`
+
+
5 hour usage100% used
+
+ `)
+
+ 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")
+}
diff --git a/backend/internal/service/wire.go b/backend/internal/service/wire.go
index e7ff2dc4d3..3944ca5f07 100644
--- a/backend/internal/service/wire.go
+++ b/backend/internal/service/wire.go
@@ -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
}