mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-07 16:37:52 +08:00
All four pass on a quiet box and fail on a loaded one, each for its own reason. None of them is testing the clock, so none of them should be failing on it. - ollama_cloud_usage_test.go: the second caller was released with `close(release)` right after its goroutine was started, not after it had reached the singleflight group. When the first refresh won that race, the second became a new singleflight execution, re-read the account, saw the LastAttemptAt the first one had just written, and came back with the 30-second manual-refresh 429 at line 685. It now counts account loads and waits for the second caller's own load, which happens right before it joins the group. Adds a counting GetByID to the test repo. - gateway_hotpath_optimization_test.go: a 20ms sleep was meant to let all 12 callers reach the cache before the loader was released. A caller that arrived after the load had finished got a hit, not a miss, so the miss count came out 11 of 12. It now waits on the miss counter itself, which is the value the test asserts on. - token_refresh_pool_health_test.go: the floor was `configuredSpacing` minus 10ms, i.e. 40ms out of 50ms. Each start timestamp is taken after the rate gate releases the goroutine, so scheduler delay can compress one observed gap with the gate behaving correctly — seen at 37ms and again at 13ms. The floor is now a tenth of the configured spacing. Measured with providerQPS=20, 8 attempts, concurrency 2: gate at 50ms gives a minimum gap of 49.97ms, gate at 0 gives 22µs. So an unpaced gate sits three orders of magnitude under the 5ms floor and is still caught, while jitter has room to move. A comment warns against replacing this with an assertion on the total span of the starts: the span is set by how long each attempt takes under the concurrency limit, not by the gate — 471ms paced against 241ms unpaced — so a span check passes with the gate disabled. - prompt_guard_test.go: the bound only has to show the failover shared the first endpoint's 70ms deadline and did not take the second endpoint's own 500ms one. An unshared deadline lands near 535ms, so 350ms still fails loudly (seen: 224ms against a 180ms bound). Tests only; no product code is touched. Each fix was checked in both directions: it passes with the behaviour intact, and it still fails when the behaviour is broken on purpose (for the QPS one, by swapping the shared rate gate for a zero-interval one).
1109 lines
49 KiB
Go
1109 lines
49 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"os"
|
|
"reflect"
|
|
"strings"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
type ollamaUsageTestEncryptor struct{}
|
|
|
|
func (ollamaUsageTestEncryptor) Encrypt(value string) (string, error) { return "cipher:" + value, nil }
|
|
func (ollamaUsageTestEncryptor) Decrypt(value string) (string, error) {
|
|
if !strings.HasPrefix(value, "cipher:") {
|
|
return "", errors.New("authentication failed")
|
|
}
|
|
return strings.TrimPrefix(value, "cipher:"), nil
|
|
}
|
|
|
|
type ollamaUsageTestRepo struct {
|
|
*upstreamBillingProbeAccountRepo
|
|
due []Account
|
|
beforeSnapshot func()
|
|
disableAutoAttempts atomic.Int64
|
|
disableAutoCalls atomic.Int64
|
|
groupResolveCalls atomic.Int64
|
|
getByIDCalls atomic.Int64
|
|
}
|
|
|
|
// GetByID counts loads so a test can wait for a caller to reach the point just
|
|
// before the singleflight group, instead of guessing with a sleep.
|
|
func (r *ollamaUsageTestRepo) GetByID(ctx context.Context, id int64) (*Account, error) {
|
|
r.getByIDCalls.Add(1)
|
|
return r.upstreamBillingProbeAccountRepo.GetByID(ctx, id)
|
|
}
|
|
|
|
func (r *ollamaUsageTestRepo) ListOllamaCloudUsageGroupAccounts(_ context.Context, anchors []*Account) ([]Account, error) {
|
|
r.groupResolveCalls.Add(1)
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
wanted := make(map[string]struct{}, len(anchors))
|
|
for _, anchor := range anchors {
|
|
if fingerprint, ok := ollamaCloudUsageGroupFingerprint(anchor); ok {
|
|
wanted[fingerprint] = struct{}{}
|
|
}
|
|
}
|
|
result := make([]Account, 0, len(r.accounts))
|
|
for _, account := range r.accounts {
|
|
fingerprint, ok := ollamaCloudUsageGroupFingerprint(account)
|
|
if _, match := wanted[fingerprint]; !ok || !match {
|
|
continue
|
|
}
|
|
result = append(result, cloneOllamaUsageTestAccount(*account))
|
|
}
|
|
return result, nil
|
|
}
|
|
|
|
// cloneOllamaUsageTestAccount 深拷贝共享 map,模拟真实仓储每次查询返回全新行:
|
|
// 组写在 r.mu 下改成员 map,浅拷贝会让 RunDue 过滤循环无锁读到同一 map 而竞争。
|
|
func cloneOllamaUsageTestAccount(account Account) Account {
|
|
account.Credentials = mergeMap(nil, account.Credentials)
|
|
account.Extra = mergeMap(nil, account.Extra)
|
|
return account
|
|
}
|
|
|
|
func (r *ollamaUsageTestRepo) SaveOllamaCloudUsageSession(_ context.Context, expected *Account, ciphertext string, autoRefresh bool) error {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
members, err := r.ollamaGroupMembersLocked(expected)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
for _, account := range members {
|
|
account.Extra[OllamaCloudUsageSessionExtraKey] = ciphertext
|
|
account.Extra[OllamaCloudUsageAutoRefreshExtraKey] = autoRefresh
|
|
delete(account.Extra, OllamaCloudUsageSnapshotExtraKey)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (r *ollamaUsageTestRepo) DeleteOllamaCloudUsageSession(_ context.Context, expected *Account) error {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
members, err := r.ollamaGroupMembersLocked(expected)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
for _, account := range members {
|
|
delete(account.Extra, OllamaCloudUsageSessionExtraKey)
|
|
delete(account.Extra, OllamaCloudUsageAutoRefreshExtraKey)
|
|
delete(account.Extra, OllamaCloudUsageSnapshotExtraKey)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (r *ollamaUsageTestRepo) SetOllamaCloudUsageAutoRefresh(_ context.Context, expected *Account, enabled bool) error {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
members, err := r.ollamaGroupMembersLocked(expected)
|
|
if err != nil || !r.ollamaExpectedSessionExistsLocked(members, expected) {
|
|
return ErrOllamaCloudUsageIdentityChanged
|
|
}
|
|
for _, account := range members {
|
|
applyOllamaUsageTestManagedExtra(account, expected)
|
|
account.Extra[OllamaCloudUsageAutoRefreshExtraKey] = enabled
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (r *ollamaUsageTestRepo) UpdateOllamaCloudUsageSnapshot(_ context.Context, expected *Account, snapshot *OllamaCloudUsageSnapshot) error {
|
|
if r.beforeSnapshot != nil {
|
|
r.beforeSnapshot()
|
|
}
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
members, err := r.ollamaGroupMembersLocked(expected)
|
|
if err != nil || !r.ollamaExpectedSessionExistsLocked(members, expected) {
|
|
return ErrOllamaCloudUsageIdentityChanged
|
|
}
|
|
for _, account := range members {
|
|
applyOllamaUsageTestManagedExtra(account, expected)
|
|
account.Extra[OllamaCloudUsageSnapshotExtraKey] = snapshot
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (r *ollamaUsageTestRepo) DisableOllamaCloudUsageAutoRefresh(_ context.Context, expected *Account) error {
|
|
r.disableAutoAttempts.Add(1)
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
members, err := r.ollamaGroupMembersLocked(expected)
|
|
if err != nil || !r.ollamaExpectedSessionExistsLocked(members, expected) {
|
|
return ErrOllamaCloudUsageIdentityChanged
|
|
}
|
|
for _, account := range members {
|
|
applyOllamaUsageTestManagedExtra(account, expected)
|
|
account.Extra[OllamaCloudUsageAutoRefreshExtraKey] = false
|
|
delete(account.Extra, OllamaCloudUsageSnapshotExtraKey)
|
|
}
|
|
r.disableAutoCalls.Add(1)
|
|
return nil
|
|
}
|
|
|
|
func (r *ollamaUsageTestRepo) ollamaGroupMembersLocked(expected *Account) ([]*Account, error) {
|
|
anchor := r.accounts[expected.ID]
|
|
if !sameOllamaUsageTestIdentity(anchor, expected) {
|
|
return nil, ErrOllamaCloudUsageIdentityChanged
|
|
}
|
|
fingerprint, ok := ollamaCloudUsageGroupFingerprint(expected)
|
|
if !ok {
|
|
return nil, ErrOllamaCloudUsageAccountInvalid
|
|
}
|
|
members := make([]*Account, 0, len(r.accounts))
|
|
for _, account := range r.accounts {
|
|
candidate, valid := ollamaCloudUsageGroupFingerprint(account)
|
|
if valid && candidate == fingerprint {
|
|
if account.Extra == nil {
|
|
account.Extra = make(map[string]any)
|
|
}
|
|
members = append(members, account)
|
|
}
|
|
}
|
|
return members, nil
|
|
}
|
|
|
|
func (r *ollamaUsageTestRepo) ollamaExpectedSessionExistsLocked(members []*Account, expected *Account) bool {
|
|
for _, member := range members {
|
|
if member.Extra[OllamaCloudUsageSessionExtraKey] == expected.Extra[OllamaCloudUsageSessionExtraKey] {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func applyOllamaUsageTestManagedExtra(account, source *Account) {
|
|
for _, key := range []string{OllamaCloudUsageSessionExtraKey, OllamaCloudUsageAutoRefreshExtraKey, OllamaCloudUsageSnapshotExtraKey} {
|
|
delete(account.Extra, key)
|
|
if value, ok := source.Extra[key]; ok {
|
|
account.Extra[key] = value
|
|
}
|
|
}
|
|
}
|
|
|
|
func (r *ollamaUsageTestRepo) ListDueOllamaCloudUsageAccounts(_ context.Context, _ time.Time, _, _ time.Duration, limit int) ([]Account, error) {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
if len(r.due) > 0 {
|
|
out := make([]Account, 0, min(limit, len(r.due)))
|
|
for _, account := range r.due[:min(limit, len(r.due))] {
|
|
out = append(out, cloneOllamaUsageTestAccount(account))
|
|
}
|
|
return out, nil
|
|
}
|
|
out := make([]Account, 0, len(r.accounts))
|
|
for _, account := range r.accounts {
|
|
out = append(out, cloneOllamaUsageTestAccount(*account))
|
|
if len(out) == limit {
|
|
break
|
|
}
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
type ollamaRefreshPreflightIdentityChangeRepo struct {
|
|
*ollamaUsageTestRepo
|
|
getCalls atomic.Int64
|
|
}
|
|
|
|
func (r *ollamaRefreshPreflightIdentityChangeRepo) GetByID(ctx context.Context, id int64) (*Account, error) {
|
|
if r.getCalls.Add(1) == 2 {
|
|
r.mu.Lock()
|
|
r.accounts[id].Credentials["api_key"] = "rotated-before-refresh"
|
|
r.mu.Unlock()
|
|
}
|
|
return r.upstreamBillingProbeAccountRepo.GetByID(ctx, id)
|
|
}
|
|
|
|
type ollamaManagedExtraUpdateRepo struct {
|
|
AccountRepository
|
|
account *Account
|
|
updated *Account
|
|
}
|
|
|
|
func (r *ollamaManagedExtraUpdateRepo) GetByID(_ context.Context, _ int64) (*Account, error) {
|
|
return r.account, nil
|
|
}
|
|
|
|
func (r *ollamaManagedExtraUpdateRepo) Update(_ context.Context, account *Account) error {
|
|
r.updated = account
|
|
return nil
|
|
}
|
|
|
|
func sameOllamaUsageTestIdentity(left, right *Account) bool {
|
|
return left != nil && right != nil && left.Platform == right.Platform && left.Type == right.Type &&
|
|
reflect.DeepEqual(left.Credentials, right.Credentials) && reflect.DeepEqual(left.ProxyID, right.ProxyID)
|
|
}
|
|
|
|
type ollamaUsageHTTPStub struct {
|
|
status int
|
|
body []byte
|
|
header http.Header
|
|
calls atomic.Int64
|
|
active atomic.Int64
|
|
maxActive atomic.Int64
|
|
beforeResponse func(*http.Request)
|
|
lastRequest *http.Request
|
|
lastProxyURL string
|
|
mu sync.Mutex
|
|
}
|
|
|
|
func (s *ollamaUsageHTTPStub) Do(req *http.Request, proxyURL string, _ int64, _ int) (*http.Response, error) {
|
|
s.calls.Add(1)
|
|
active := s.active.Add(1)
|
|
defer s.active.Add(-1)
|
|
for {
|
|
peak := s.maxActive.Load()
|
|
if active <= peak || s.maxActive.CompareAndSwap(peak, active) {
|
|
break
|
|
}
|
|
}
|
|
s.mu.Lock()
|
|
s.lastRequest = req
|
|
s.lastProxyURL = proxyURL
|
|
s.mu.Unlock()
|
|
if s.beforeResponse != nil {
|
|
s.beforeResponse(req)
|
|
}
|
|
status := s.status
|
|
if status == 0 {
|
|
status = http.StatusOK
|
|
}
|
|
header := s.header
|
|
if header == nil {
|
|
header = http.Header{"Content-Type": []string{"text/html; charset=utf-8"}}
|
|
}
|
|
return &http.Response{StatusCode: status, Header: header, Body: io.NopCloser(strings.NewReader(string(s.body))), Request: req}, nil
|
|
}
|
|
|
|
func (s *ollamaUsageHTTPStub) DoWithTLS(req *http.Request, proxyURL string, accountID int64, concurrency int, _ *tlsfingerprint.Profile) (*http.Response, error) {
|
|
return s.Do(req, proxyURL, accountID, concurrency)
|
|
}
|
|
|
|
func ollamaUsageAccount(id int64) *Account {
|
|
return &Account{
|
|
ID: id, Name: fmt.Sprintf("ollama-%d", id), Platform: PlatformOpenAI, Type: AccountTypeAPIKey,
|
|
Credentials: map[string]any{"base_url": "https://ollama.com", "api_key": fmt.Sprintf("key-%d", id)},
|
|
Extra: map[string]any{}, Status: StatusActive, Schedulable: true, Concurrency: 1,
|
|
}
|
|
}
|
|
|
|
func newOllamaUsageTestService(t *testing.T, repo *ollamaUsageTestRepo, upstream HTTPUpstream, settingsRepo SettingRepository, fixedKey bool) *OllamaCloudUsageService {
|
|
t.Helper()
|
|
svc := NewOllamaCloudUsageService(repo, upstream, NewSettingService(settingsRepo, nil), ollamaUsageTestEncryptor{}, fixedKey)
|
|
t.Cleanup(svc.Stop)
|
|
return svc
|
|
}
|
|
|
|
func ollamaUsageFixture(t *testing.T) []byte {
|
|
t.Helper()
|
|
body, err := os.ReadFile("testdata/ollama_settings_usage.html")
|
|
require.NoError(t, err)
|
|
return body
|
|
}
|
|
|
|
func TestOllamaCloudUsageSettingsDefaultOffAndValidation(t *testing.T) {
|
|
repo := &upstreamBillingProbeSettingRepo{}
|
|
settingsService := NewSettingService(repo, nil)
|
|
settings, err := settingsService.GetOllamaCloudUsageSettings(context.Background())
|
|
require.NoError(t, err)
|
|
require.False(t, settings.Enabled)
|
|
require.Equal(t, 60, settings.IntervalMinutes)
|
|
require.Equal(t, 1, settings.DebounceMinutes)
|
|
|
|
err = settingsService.SetOllamaCloudUsageSettings(context.Background(), &OllamaCloudUsageSettings{Enabled: true, IntervalMinutes: 14, DebounceMinutes: 1})
|
|
require.Error(t, err)
|
|
err = settingsService.SetOllamaCloudUsageSettings(context.Background(), &OllamaCloudUsageSettings{Enabled: true, IntervalMinutes: 90, DebounceMinutes: 61})
|
|
require.Error(t, err)
|
|
// DebounceMinutes=0 (legacy omit) defaults to 1 on write.
|
|
err = settingsService.SetOllamaCloudUsageSettings(context.Background(), &OllamaCloudUsageSettings{Enabled: true, IntervalMinutes: 90, DebounceMinutes: 0})
|
|
require.NoError(t, err)
|
|
settings, err = settingsService.GetOllamaCloudUsageSettings(context.Background())
|
|
require.NoError(t, err)
|
|
require.Equal(t, 1, settings.DebounceMinutes)
|
|
err = settingsService.SetOllamaCloudUsageSettings(context.Background(), &OllamaCloudUsageSettings{Enabled: true, IntervalMinutes: 90, DebounceMinutes: 2})
|
|
require.NoError(t, err)
|
|
settings, err = settingsService.GetOllamaCloudUsageSettings(context.Background())
|
|
require.NoError(t, err)
|
|
require.True(t, settings.Enabled)
|
|
require.Equal(t, 90, settings.IntervalMinutes)
|
|
require.Equal(t, 2, settings.DebounceMinutes)
|
|
|
|
// debounce >= interval would make the debounce term unreachable in
|
|
// min(lastUsed+debounce, fetchedAt+maxWait), silently ignoring the operator's
|
|
// setting, so it is rejected rather than accepted and dropped.
|
|
err = settingsService.SetOllamaCloudUsageSettings(context.Background(), &OllamaCloudUsageSettings{Enabled: true, IntervalMinutes: 15, DebounceMinutes: 15})
|
|
require.Error(t, err, "debounce equal to interval must be rejected")
|
|
err = settingsService.SetOllamaCloudUsageSettings(context.Background(), &OllamaCloudUsageSettings{Enabled: true, IntervalMinutes: 15, DebounceMinutes: 60})
|
|
require.Error(t, err, "debounce greater than interval must be rejected")
|
|
err = settingsService.SetOllamaCloudUsageSettings(context.Background(), &OllamaCloudUsageSettings{Enabled: true, IntervalMinutes: 16, DebounceMinutes: 15})
|
|
require.NoError(t, err, "debounce below interval stays valid")
|
|
|
|
// Legacy JSON without debounce_minutes defaults to 1.
|
|
repo.values[SettingKeyOllamaCloudUsageSettings] = `{"enabled":true,"interval_minutes":45}`
|
|
settings, err = settingsService.GetOllamaCloudUsageSettings(context.Background())
|
|
require.NoError(t, err)
|
|
require.Equal(t, 45, settings.IntervalMinutes)
|
|
require.Equal(t, 1, settings.DebounceMinutes)
|
|
}
|
|
|
|
func TestOllamaCloudUsageIsAutoRefreshDue(t *testing.T) {
|
|
debounce := time.Minute
|
|
maxWait := time.Hour
|
|
now := time.Date(2026, time.July, 25, 12, 0, 0, 0, time.UTC)
|
|
fetched := now.Add(-30 * time.Minute)
|
|
ptr := func(ts time.Time) *time.Time { return &ts }
|
|
|
|
require.True(t, ollamaCloudUsageIsAutoRefreshDue(nil, nil, now, debounce, maxWait), "missing snapshot first due")
|
|
require.True(t, ollamaCloudUsageIsAutoRefreshDue(&OllamaCloudUsageSnapshot{Status: "bogus"}, nil, now, debounce, maxWait), "invalid status first due")
|
|
|
|
okSnap := &OllamaCloudUsageSnapshot{
|
|
Status: OllamaCloudUsageStatusOK, FetchedAt: ptr(fetched),
|
|
LastAttemptAt: fetched, NextRefreshAt: fetched.Add(maxWait),
|
|
}
|
|
require.False(t, ollamaCloudUsageIsAutoRefreshDue(okSnap, nil, now, debounce, maxWait), "no request after success")
|
|
require.False(t, ollamaCloudUsageIsAutoRefreshDue(okSnap, ptr(fetched), now, debounce, maxWait), "request not after fetched_at")
|
|
require.False(t, ollamaCloudUsageIsAutoRefreshDue(okSnap, ptr(now.Add(-30*time.Second)), now, debounce, maxWait), "debounce not elapsed")
|
|
require.True(t, ollamaCloudUsageIsAutoRefreshDue(okSnap, ptr(now.Add(-time.Minute)), now, debounce, maxWait), "single request quiet for debounce")
|
|
|
|
// Continuous requests: last used is now, but max-wait from old fetch forces due.
|
|
oldFetched := now.Add(-2 * time.Hour)
|
|
oldSnap := &OllamaCloudUsageSnapshot{
|
|
Status: OllamaCloudUsageStatusOK, FetchedAt: ptr(oldFetched),
|
|
LastAttemptAt: oldFetched, NextRefreshAt: oldFetched.Add(maxWait),
|
|
}
|
|
require.True(t, ollamaCloudUsageIsAutoRefreshDue(oldSnap, ptr(now), now, debounce, maxWait), "max-wait forces due while requests continue")
|
|
// First request after a very old snapshot is immediately due because fetched+maxWait is past.
|
|
require.True(t, ollamaCloudUsageIsAutoRefreshDue(oldSnap, ptr(now.Add(-time.Second)), now, debounce, maxWait), "stale snapshot first request immediate")
|
|
|
|
failSnap := &OllamaCloudUsageSnapshot{
|
|
Status: OllamaCloudUsageStatusFailed, FetchedAt: ptr(fetched),
|
|
LastAttemptAt: now.Add(-10 * time.Minute), NextRefreshAt: now.Add(20 * time.Minute),
|
|
}
|
|
require.False(t, ollamaCloudUsageIsAutoRefreshDue(failSnap, nil, now, debounce, maxWait), "failure without new request")
|
|
require.False(t, ollamaCloudUsageIsAutoRefreshDue(failSnap, ptr(now.Add(-time.Minute)), now, debounce, maxWait), "failure blocked by backoff")
|
|
failSnap.NextRefreshAt = now.Add(-time.Second)
|
|
require.True(t, ollamaCloudUsageIsAutoRefreshDue(failSnap, ptr(now.Add(-time.Minute)), now, debounce, maxWait), "failure after backoff with new request")
|
|
|
|
require.True(t, ollamaCloudUsageIsAutoRefreshDue(&OllamaCloudUsageSnapshot{
|
|
Status: OllamaCloudUsageStatusOK, LastAttemptAt: now,
|
|
}, nil, now, debounce, maxWait), "ok without fetched_at fails open")
|
|
}
|
|
|
|
// The success path stopped consulting next_refresh_at, which is where
|
|
// nextOllamaCloudUsageDelay used to apply the minimum interval. Activity may pull
|
|
// a refresh forward only as far as that floor, otherwise request traffic spaced
|
|
// just wider than the debounce drives the group's outbound rate far above the
|
|
// pre-existing minimum.
|
|
func TestOllamaCloudUsageAutoRefreshDueAtHonoursMinFetchInterval(t *testing.T) {
|
|
debounce := time.Minute
|
|
maxWait := time.Hour
|
|
now := time.Date(2026, time.July, 25, 12, 0, 0, 0, time.UTC)
|
|
ptr := func(ts time.Time) *time.Time { return &ts }
|
|
|
|
// Debounce elapsed, but the last successful fetch is inside the floor.
|
|
recent := now.Add(-5 * time.Minute)
|
|
recentSnap := &OllamaCloudUsageSnapshot{
|
|
Status: OllamaCloudUsageStatusOK, FetchedAt: ptr(recent), LastAttemptAt: recent,
|
|
}
|
|
dueAt, ok := ollamaCloudUsageAutoRefreshDueAt(recentSnap, ptr(now.Add(-2*time.Minute)), debounce, maxWait)
|
|
require.True(t, ok)
|
|
require.Equal(t, recent.Add(OllamaCloudUsageMinFetchInterval), dueAt,
|
|
"due time must be clamped to fetched_at + min fetch interval")
|
|
require.False(t, ollamaCloudUsageIsAutoRefreshDue(recentSnap, ptr(now.Add(-2*time.Minute)), now, debounce, maxWait),
|
|
"debounce alone must not refresh within the min fetch interval")
|
|
|
|
// Once the floor has passed the debounce governs again.
|
|
atFloor := now.Add(-OllamaCloudUsageMinFetchInterval)
|
|
floorSnap := &OllamaCloudUsageSnapshot{
|
|
Status: OllamaCloudUsageStatusOK, FetchedAt: ptr(atFloor), LastAttemptAt: atFloor,
|
|
}
|
|
require.True(t, ollamaCloudUsageIsAutoRefreshDue(floorSnap, ptr(now.Add(-2*time.Minute)), now, debounce, maxWait),
|
|
"past the floor a quiet debounce window is due")
|
|
|
|
// The floor never delays a refresh that max-wait has already forced.
|
|
stale := now.Add(-2 * time.Hour)
|
|
staleSnap := &OllamaCloudUsageSnapshot{
|
|
Status: OllamaCloudUsageStatusOK, FetchedAt: ptr(stale), LastAttemptAt: stale,
|
|
}
|
|
require.True(t, ollamaCloudUsageIsAutoRefreshDue(staleSnap, ptr(now), now, debounce, maxWait),
|
|
"max-wait still forces due on a stale snapshot")
|
|
}
|
|
|
|
func TestScheduleOllamaCloudUsageActivityOnlyForOllama(t *testing.T) {
|
|
deferred := NewDeferredService(nil, nil, time.Second)
|
|
ollama := ollamaUsageAccount(1)
|
|
other := ollamaUsageAccount(2)
|
|
other.Credentials["base_url"] = "https://api.openai.com"
|
|
|
|
scheduleOllamaCloudUsageActivity(deferred, ollama)
|
|
scheduleOllamaCloudUsageActivity(deferred, other)
|
|
scheduleOllamaCloudUsageActivity(nil, ollama)
|
|
|
|
_, ok := deferred.lastUsedUpdates.Load(int64(1))
|
|
require.True(t, ok)
|
|
_, ok = deferred.lastUsedUpdates.Load(int64(2))
|
|
require.False(t, ok)
|
|
}
|
|
|
|
func TestIsOllamaCloudUsageAccountStrictOfficialHost(t *testing.T) {
|
|
tests := []struct {
|
|
baseURL string
|
|
platform string
|
|
want bool
|
|
}{
|
|
{"https://ollama.com", PlatformOpenAI, true},
|
|
{"HTTPS://OLLAMA.COM", PlatformAnthropic, true},
|
|
{"https://www.OLLAMA.com:443/v1", PlatformOpenAI, true},
|
|
{"https://ollama.com:443", PlatformOpenAI, true},
|
|
{"https://ollama.com/", PlatformAnthropic, false},
|
|
{"https://ollama.com/v1/", PlatformOpenAI, false},
|
|
{"http://ollama.com", PlatformOpenAI, false},
|
|
{"https://ollama.com.evil.test", PlatformOpenAI, false},
|
|
{"https://ollama.com:444", PlatformOpenAI, false},
|
|
{"https://user@ollama.com", PlatformOpenAI, false},
|
|
{"https://ollama.com/v2", PlatformOpenAI, false},
|
|
{"https://ollama.com?next=https://evil.test", PlatformOpenAI, false},
|
|
{"https://ollama.com#usage", PlatformOpenAI, false},
|
|
}
|
|
for _, test := range tests {
|
|
t.Run(test.baseURL+test.platform, func(t *testing.T) {
|
|
account := ollamaUsageAccount(1)
|
|
account.Platform = test.platform
|
|
account.Credentials["base_url"] = test.baseURL
|
|
require.Equal(t, test.want, IsOllamaCloudUsageAccount(account))
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestNormalizeOllamaCloudUsageCookieAllowlist(t *testing.T) {
|
|
normalized, err := normalizeOllamaCloudUsageCookie(" tracking=discard ; wos-session=secret ; __Secure-authjs.session-token.0=part-a ; device=discard ")
|
|
require.NoError(t, err)
|
|
require.Equal(t, "wos-session=secret; __Secure-authjs.session-token.0=part-a", normalized)
|
|
|
|
normalized, err = normalizeOllamaCloudUsageCookie(" \t\r\nwos-session=secret; tracking=discard\r\n\t ")
|
|
require.NoError(t, err)
|
|
require.Equal(t, "wos-session=secret", normalized)
|
|
|
|
_, err = normalizeOllamaCloudUsageCookie("wos-session=secret\r\nHost: evil.test")
|
|
require.ErrorContains(t, err, "invalid header")
|
|
|
|
for _, allowed := range []string{
|
|
"wos-session", "__Secure-session", "session", "ollama_session", "__Host-ollama_session",
|
|
"next-auth.session-token", "next-auth.session-token.0", "__Secure-next-auth.session-token.12",
|
|
"authjs.session-token", "__Secure-authjs.session-token.1",
|
|
} {
|
|
normalized, err := normalizeOllamaCloudUsageCookie(allowed + "=value")
|
|
require.NoError(t, err, allowed)
|
|
require.Equal(t, allowed+"=value", normalized)
|
|
}
|
|
|
|
for _, invalid := range []string{
|
|
"", "Domain=ollama.com; wos-session=x", "wos-session=x; Path=/",
|
|
"wos-session=x; wos-session=y", "Secure", "tracking=only", "__session=arbitrary",
|
|
"authjs.session-token.bad=not-a-shard", "Authjs.session-token=wrong-case",
|
|
} {
|
|
_, err := normalizeOllamaCloudUsageCookie(invalid)
|
|
require.Error(t, err, invalid)
|
|
}
|
|
_, err = normalizeOllamaCloudUsageCookie("wos-session=" + strings.Repeat("x", ollamaCloudUsageMaxSessionBytes))
|
|
require.ErrorContains(t, err, "too large")
|
|
}
|
|
|
|
func TestParseOllamaCloudUsageHTMLFixture(t *testing.T) {
|
|
data, err := parseOllamaCloudUsageHTML(ollamaUsageFixture(t))
|
|
require.NoError(t, err)
|
|
require.Equal(t, "max", data.Plan)
|
|
require.NotNil(t, data.FiveHour)
|
|
require.Equal(t, 5.6, data.FiveHour.UsedPercent)
|
|
require.NotNil(t, data.FiveHour.ResetAt)
|
|
require.Equal(t, time.Date(2026, time.July, 23, 3, 0, 0, 0, time.UTC), *data.FiveHour.ResetAt)
|
|
require.NotNil(t, data.SevenDay)
|
|
require.Equal(t, 14.2, data.SevenDay.UsedPercent)
|
|
require.NotNil(t, data.SevenDay.ResetAt)
|
|
require.Equal(t, time.Date(2026, time.July, 29, 0, 0, 0, 0, time.UTC), *data.SevenDay.ResetAt)
|
|
require.Equal(t, "$0", data.Balance)
|
|
require.Equal(t, []OllamaCloudUsageModel{
|
|
{Model: "gpt-oss:120b-cloud", Window: OllamaCloudUsageModelWindowFiveHour, Requests: 2},
|
|
{Model: "qwen3-coder:480b-cloud", Window: OllamaCloudUsageModelWindowFiveHour, Requests: 3},
|
|
{Model: "gpt-oss:120b-cloud", Window: OllamaCloudUsageModelWindowSevenDay, Requests: 12},
|
|
{Model: "qwen3-coder:480b-cloud", Window: OllamaCloudUsageModelWindowSevenDay, Requests: 13},
|
|
}, data.Models)
|
|
|
|
_, err = parseOllamaCloudUsageHTML([]byte(`<html><body><main>Sign in to Ollama</main></body></html>`))
|
|
require.ErrorIs(t, err, errOllamaCloudUsageUnauthorizedHTML)
|
|
_, err = parseOllamaCloudUsageHTML([]byte(`<html><body><p>5 hour usage 42% used</p><form>Sign in to Ollama</form></body></html>`))
|
|
require.ErrorIs(t, err, errOllamaCloudUsageUnauthorizedHTML)
|
|
_, err = parseOllamaCloudUsageHTML([]byte(`<html><body><main>unrelated settings</main></body></html>`))
|
|
require.Error(t, err)
|
|
}
|
|
|
|
func TestParseOllamaCloudUsageHTMLMissingOptionalFieldsAndCSSWidthFallback(t *testing.T) {
|
|
data, err := parseOllamaCloudUsageHTML([]byte(`
|
|
<section>
|
|
<p>5 hour usage</p>
|
|
<div data-usage-track>
|
|
<div data-usage-segment style="width: 23.5%"><span data-model="model-a" data-requests="1,234"></span></div>
|
|
<div data-usage-segment data-model="model-a" data-requests="9,999" style="width: 0%"></div>
|
|
</div>
|
|
</section>`))
|
|
require.NoError(t, err)
|
|
require.Equal(t, 23.5, data.FiveHour.UsedPercent)
|
|
require.Nil(t, data.FiveHour.ResetAt)
|
|
require.Empty(t, data.Plan)
|
|
require.Nil(t, data.SevenDay)
|
|
require.Empty(t, data.Balance)
|
|
require.Equal(t, []OllamaCloudUsageModel{{
|
|
Model: "model-a", Window: OllamaCloudUsageModelWindowFiveHour, Requests: 1234,
|
|
}}, data.Models)
|
|
}
|
|
|
|
func TestParseOllamaCloudUsageHTMLResetElementVariants(t *testing.T) {
|
|
const want = "2026-07-23T03:00:00Z"
|
|
for name, element := range map[string]string{
|
|
"time datetime": `<time datetime="` + want + `">2 hours.</time>`,
|
|
"custom element": `<local-time data-time="` + want + `">2 hours.</local-time>`,
|
|
"class token": `<span class="text-xs local-time tabular-nums" data-time="` + want + `">2 hours.</span>`,
|
|
} {
|
|
t.Run(name, func(t *testing.T) {
|
|
data, err := parseOllamaCloudUsageHTML([]byte(
|
|
`<div><div><span>Session usage</span><span>1% used</span></div><div>Resets in ` + element + `</div></div>`,
|
|
))
|
|
require.NoError(t, err)
|
|
require.NotNil(t, data.FiveHour)
|
|
require.NotNil(t, data.FiveHour.ResetAt)
|
|
require.Equal(t, want, data.FiveHour.ResetAt.Format(time.RFC3339))
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestParseOllamaCloudUsageHTMLPlanAndBalanceFallbacks(t *testing.T) {
|
|
data, err := parseOllamaCloudUsageHTML([]byte(`
|
|
<section>
|
|
<h2><span>Cloud usage</span><span>max</span></h2>
|
|
<div><span>Plan</span><span>Pro</span></div>
|
|
<p>Credits currently available: USD $9.50</p>
|
|
</section>`))
|
|
require.NoError(t, err)
|
|
require.Equal(t, "max", data.Plan)
|
|
require.Equal(t, "USD$9.50", data.Balance)
|
|
|
|
data, err = parseOllamaCloudUsageHTML([]byte(`<div><span>Subscription</span><span>Pro</span></div>`))
|
|
require.NoError(t, err)
|
|
require.Equal(t, "Pro", data.Plan)
|
|
}
|
|
|
|
func TestOllamaCloudUsageManagedExtraCannotBeImported(t *testing.T) {
|
|
remoteExtra := map[string]any{
|
|
OllamaCloudUsageSessionExtraKey: "remote-ciphertext",
|
|
OllamaCloudUsageAutoRefreshExtraKey: true,
|
|
OllamaCloudUsageSnapshotExtraKey: map[string]any{"status": "forged"},
|
|
}
|
|
created, err := buildAccountForCreate(&CreateAccountInput{
|
|
Name: "ollama", Platform: PlatformOpenAI, Type: AccountTypeAPIKey,
|
|
Credentials: map[string]any{"base_url": "https://ollama.com", "api_key": "key"},
|
|
Concurrency: 1,
|
|
}, mergeMap(nil, remoteExtra))
|
|
require.NoError(t, err)
|
|
require.NotContains(t, created.Extra, OllamaCloudUsageSessionExtraKey)
|
|
require.NotContains(t, created.Extra, OllamaCloudUsageAutoRefreshExtraKey)
|
|
require.NotContains(t, created.Extra, OllamaCloudUsageSnapshotExtraKey)
|
|
|
|
existing := ollamaUsageAccount(6)
|
|
existing.Extra = map[string]any{
|
|
OllamaCloudUsageSessionExtraKey: "local-ciphertext",
|
|
OllamaCloudUsageAutoRefreshExtraKey: false,
|
|
OllamaCloudUsageSnapshotExtraKey: map[string]any{"status": OllamaCloudUsageStatusOK},
|
|
}
|
|
targetExtra := mergeMap(existing.Extra, remoteExtra)
|
|
reconcileCRSUpstreamBillingProbeExtra(existing, existing.Platform, existing.Type, mergeMap(existing.Credentials, nil), targetExtra)
|
|
require.Equal(t, "local-ciphertext", targetExtra[OllamaCloudUsageSessionExtraKey])
|
|
require.Equal(t, false, targetExtra[OllamaCloudUsageAutoRefreshExtraKey])
|
|
require.Equal(t, map[string]any{"status": OllamaCloudUsageStatusOK}, targetExtra[OllamaCloudUsageSnapshotExtraKey])
|
|
|
|
changedCredentials := mergeMap(existing.Credentials, map[string]any{"api_key": "rotated"})
|
|
targetExtra = mergeMap(existing.Extra, remoteExtra)
|
|
reconcileCRSUpstreamBillingProbeExtra(existing, existing.Platform, existing.Type, changedCredentials, targetExtra)
|
|
require.NotContains(t, targetExtra, OllamaCloudUsageSessionExtraKey)
|
|
require.NotContains(t, targetExtra, OllamaCloudUsageAutoRefreshExtraKey)
|
|
require.NotContains(t, targetExtra, OllamaCloudUsageSnapshotExtraKey)
|
|
}
|
|
|
|
func TestAccountServiceUpdateStripsOllamaManagedExtra(t *testing.T) {
|
|
account := ollamaUsageAccount(61)
|
|
account.Extra = map[string]any{
|
|
OllamaCloudUsageSessionExtraKey: "local-ciphertext",
|
|
OllamaCloudUsageAutoRefreshExtraKey: true,
|
|
OllamaCloudUsageSnapshotExtraKey: map[string]any{"status": OllamaCloudUsageStatusOK},
|
|
}
|
|
repo := &ollamaManagedExtraUpdateRepo{account: account}
|
|
svc := NewAccountService(repo, nil)
|
|
requestedExtra := map[string]any{
|
|
"note": "preserved",
|
|
OllamaCloudUsageSessionExtraKey: "forged-ciphertext",
|
|
OllamaCloudUsageAutoRefreshExtraKey: nil,
|
|
OllamaCloudUsageSnapshotExtraKey: nil,
|
|
}
|
|
|
|
_, err := svc.Update(context.Background(), account.ID, UpdateAccountRequest{Extra: &requestedExtra})
|
|
require.NoError(t, err)
|
|
require.Equal(t, "preserved", repo.updated.Extra["note"])
|
|
require.NotContains(t, repo.updated.Extra, OllamaCloudUsageSessionExtraKey)
|
|
require.NotContains(t, repo.updated.Extra, OllamaCloudUsageAutoRefreshExtraKey)
|
|
require.NotContains(t, repo.updated.Extra, OllamaCloudUsageSnapshotExtraKey)
|
|
// The request map is not mutated while managed fields are stripped.
|
|
require.Contains(t, requestedExtra, OllamaCloudUsageSessionExtraKey)
|
|
}
|
|
|
|
func TestOllamaCloudUsageSessionEncryptionFailClosedAndWriteOnlyState(t *testing.T) {
|
|
account := ollamaUsageAccount(7)
|
|
repo := &ollamaUsageTestRepo{upstreamBillingProbeAccountRepo: &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{7: account}}}
|
|
settings := &upstreamBillingProbeSettingRepo{}
|
|
|
|
ephemeral := newOllamaUsageTestService(t, repo, &ollamaUsageHTTPStub{}, settings, false)
|
|
_, err := ephemeral.SaveSession(context.Background(), 7, "wos-session=plaintext-secret")
|
|
require.ErrorIs(t, err, ErrOllamaCloudUsageEncryptionKey)
|
|
require.NotContains(t, account.Extra, OllamaCloudUsageSessionExtraKey)
|
|
|
|
svc := newOllamaUsageTestService(t, repo, &ollamaUsageHTTPStub{}, settings, true)
|
|
_, err = svc.SaveSession(context.Background(), 7, "tracking=arbitrary-only")
|
|
require.Error(t, err)
|
|
require.NotContains(t, account.Extra, OllamaCloudUsageSessionExtraKey)
|
|
|
|
state, err := svc.SaveSession(context.Background(), 7, "tracking=must-not-persist; wos-session=plaintext-secret")
|
|
require.NoError(t, err)
|
|
require.True(t, state.Configured)
|
|
stored, ok := account.Extra[OllamaCloudUsageSessionExtraKey].(string)
|
|
require.True(t, ok)
|
|
require.Equal(t, "cipher:wos-session=plaintext-secret", stored)
|
|
require.NotContains(t, stored, "tracking")
|
|
raw, err := json.Marshal(state)
|
|
require.NoError(t, err)
|
|
require.NotContains(t, string(raw), "plaintext-secret")
|
|
require.NotContains(t, string(raw), "cipher:")
|
|
|
|
account.Extra[OllamaCloudUsageSessionExtraKey] = "plaintext-secret"
|
|
_, err = svc.Refresh(context.Background(), 7)
|
|
require.ErrorContains(t, err, "cannot be decrypted")
|
|
}
|
|
|
|
func TestOllamaCloudUsageGroupSharesAcrossPlatformsURLVariantsAndDynamicSiblings(t *testing.T) {
|
|
source := ollamaUsageAccount(71)
|
|
source.Credentials["api_key"] = "shared-key"
|
|
source.Extra[OllamaCloudUsageSessionExtraKey] = "cipher:wos-session=shared"
|
|
source.Extra[OllamaCloudUsageAutoRefreshExtraKey] = true
|
|
source.Extra[OllamaCloudUsageSnapshotExtraKey] = &OllamaCloudUsageSnapshot{
|
|
Status: OllamaCloudUsageStatusOK,
|
|
Data: &OllamaCloudUsageData{Plan: "pro"},
|
|
}
|
|
source.UpdatedAt = time.Now().Add(-time.Minute)
|
|
sibling := ollamaUsageAccount(72)
|
|
sibling.Platform = PlatformAnthropic
|
|
sibling.Credentials = map[string]any{"base_url": "HTTPS://WWW.OLLAMA.COM:443/v1", "api_key": "shared-key"}
|
|
sibling.Extra[OllamaCloudUsageSessionExtraKey] = "cipher:wos-session=shared"
|
|
sibling.Extra[OllamaCloudUsageAutoRefreshExtraKey] = true
|
|
sibling.UpdatedAt = time.Now()
|
|
different := ollamaUsageAccount(73)
|
|
different.Credentials["api_key"] = "different-key"
|
|
repo := &ollamaUsageTestRepo{upstreamBillingProbeAccountRepo: &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{
|
|
source.ID: source, sibling.ID: sibling, different.ID: different,
|
|
}}}
|
|
svc := newOllamaUsageTestService(t, repo, &ollamaUsageHTTPStub{}, &upstreamBillingProbeSettingRepo{}, true)
|
|
|
|
state, err := svc.GetState(context.Background(), sibling.ID)
|
|
require.NoError(t, err)
|
|
require.True(t, state.Configured)
|
|
require.True(t, state.AutoRefreshEnabled)
|
|
require.Equal(t, "pro", state.Snapshot.Data.Plan)
|
|
|
|
differentState, err := svc.GetState(context.Background(), different.ID)
|
|
require.NoError(t, err)
|
|
require.False(t, differentState.Configured)
|
|
|
|
newSibling := ollamaUsageAccount(74)
|
|
newSibling.Platform = PlatformAnthropic
|
|
newSibling.Credentials = map[string]any{"base_url": "https://ollama.com:443", "api_key": "shared-key"}
|
|
repo.mu.Lock()
|
|
repo.accounts[newSibling.ID] = newSibling
|
|
repo.mu.Unlock()
|
|
newState, err := svc.GetState(context.Background(), newSibling.ID)
|
|
require.NoError(t, err)
|
|
require.True(t, newState.Configured)
|
|
require.Equal(t, state.Snapshot, newState.Snapshot)
|
|
|
|
before := repo.groupResolveCalls.Load()
|
|
require.NoError(t, svc.ResolveAccounts(context.Background(), []*Account{source, sibling, different, newSibling}))
|
|
require.Equal(t, before+1, repo.groupResolveCalls.Load(), "one list batch must issue one group lookup")
|
|
}
|
|
|
|
func TestOllamaCloudUsageSaveAutoRefreshAndDeleteAreGroupScoped(t *testing.T) {
|
|
first := ollamaUsageAccount(81)
|
|
first.Credentials["api_key"] = "shared-key"
|
|
second := ollamaUsageAccount(82)
|
|
second.Platform = PlatformAnthropic
|
|
second.Credentials = map[string]any{"base_url": "https://www.ollama.com/v1", "api_key": "shared-key"}
|
|
different := ollamaUsageAccount(83)
|
|
different.Credentials["api_key"] = "different-key"
|
|
repo := &ollamaUsageTestRepo{upstreamBillingProbeAccountRepo: &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{
|
|
first.ID: first, second.ID: second, different.ID: different,
|
|
}}}
|
|
svc := newOllamaUsageTestService(t, repo, &ollamaUsageHTTPStub{}, &upstreamBillingProbeSettingRepo{}, true)
|
|
|
|
state, err := svc.SaveSession(context.Background(), second.ID, "wos-session=shared-browser")
|
|
require.NoError(t, err)
|
|
require.True(t, state.Configured)
|
|
require.Equal(t, "cipher:wos-session=shared-browser", first.Extra[OllamaCloudUsageSessionExtraKey])
|
|
require.Equal(t, first.Extra[OllamaCloudUsageSessionExtraKey], second.Extra[OllamaCloudUsageSessionExtraKey])
|
|
require.NotContains(t, different.Extra, OllamaCloudUsageSessionExtraKey)
|
|
|
|
state, err = svc.SetAutoRefresh(context.Background(), first.ID, true)
|
|
require.NoError(t, err)
|
|
require.True(t, state.AutoRefreshEnabled)
|
|
require.Equal(t, true, first.Extra[OllamaCloudUsageAutoRefreshExtraKey])
|
|
require.Equal(t, true, second.Extra[OllamaCloudUsageAutoRefreshExtraKey])
|
|
|
|
state, err = svc.DeleteSession(context.Background(), second.ID)
|
|
require.NoError(t, err)
|
|
require.False(t, state.Configured)
|
|
for _, member := range []*Account{first, second} {
|
|
require.NotContains(t, member.Extra, OllamaCloudUsageSessionExtraKey)
|
|
require.NotContains(t, member.Extra, OllamaCloudUsageAutoRefreshExtraKey)
|
|
require.NotContains(t, member.Extra, OllamaCloudUsageSnapshotExtraKey)
|
|
}
|
|
}
|
|
|
|
func TestOllamaCloudUsageRefreshSingleflightAndRunnerDeduplicateSharedGroup(t *testing.T) {
|
|
first := ollamaUsageAccount(91)
|
|
first.Credentials["api_key"] = "shared-key"
|
|
first.Extra[OllamaCloudUsageSessionExtraKey] = "cipher:wos-session=shared"
|
|
first.Extra[OllamaCloudUsageAutoRefreshExtraKey] = true
|
|
second := ollamaUsageAccount(92)
|
|
second.Platform = PlatformAnthropic
|
|
second.Credentials = map[string]any{"base_url": "https://www.ollama.com:443/v1", "api_key": "shared-key"}
|
|
second.Extra[OllamaCloudUsageSessionExtraKey] = "cipher:wos-session=shared"
|
|
second.Extra[OllamaCloudUsageAutoRefreshExtraKey] = true
|
|
repo := &ollamaUsageTestRepo{
|
|
upstreamBillingProbeAccountRepo: &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{first.ID: first, second.ID: second}},
|
|
due: []Account{*first, *second},
|
|
}
|
|
settingsRepo := &upstreamBillingProbeSettingRepo{values: map[string]string{
|
|
SettingKeyOllamaCloudUsageSettings: `{"enabled":true,"interval_minutes":60}`,
|
|
}}
|
|
started := make(chan struct{})
|
|
release := make(chan struct{})
|
|
var once sync.Once
|
|
upstream := &ollamaUsageHTTPStub{body: ollamaUsageFixture(t), beforeResponse: func(*http.Request) {
|
|
once.Do(func() { close(started) })
|
|
<-release
|
|
}}
|
|
svc := newOllamaUsageTestService(t, repo, upstream, settingsRepo, true)
|
|
|
|
errs := make(chan error, 2)
|
|
go func() { _, err := svc.Refresh(context.Background(), first.ID); errs <- err }()
|
|
<-started
|
|
// The first caller is now parked in the stub, having loaded the account twice
|
|
// (once to build the group key, once inside the singleflight function).
|
|
loadsBeforeSecond := repo.getByIDCalls.Load()
|
|
go func() { _, err := svc.Refresh(context.Background(), second.ID); errs <- err }()
|
|
// Only release the first caller once the second one has loaded its own
|
|
// account, which happens immediately before it joins the singleflight group.
|
|
// Releasing right after starting the goroutine raced: if the first refresh
|
|
// finished first, the second became a fresh singleflight execution, re-read
|
|
// the account, saw the LastAttemptAt just written, and failed with the 30s
|
|
// manual-refresh 429 instead of sharing the in-flight result.
|
|
require.Eventually(t, func() bool {
|
|
return repo.getByIDCalls.Load() > loadsBeforeSecond
|
|
}, 5*time.Second, time.Millisecond, "the second caller must reach the singleflight group before the first is released")
|
|
close(release)
|
|
require.NoError(t, <-errs)
|
|
require.NoError(t, <-errs)
|
|
require.Equal(t, int64(1), upstream.calls.Load())
|
|
require.NotNil(t, decodeOllamaCloudUsageSnapshot(first.Extra))
|
|
require.Equal(t, decodeOllamaCloudUsageSnapshot(first.Extra), decodeOllamaCloudUsageSnapshot(second.Extra))
|
|
|
|
delete(first.Extra, OllamaCloudUsageSnapshotExtraKey)
|
|
delete(second.Extra, OllamaCloudUsageSnapshotExtraKey)
|
|
upstream.beforeResponse = nil
|
|
require.NoError(t, svc.RunDue(context.Background()))
|
|
require.Equal(t, int64(2), upstream.calls.Load(), "RunDue must issue one request for the shared group")
|
|
}
|
|
|
|
func TestOllamaCloudUsageRefreshRejectsGroupChangeBeforeUpstreamRequest(t *testing.T) {
|
|
account := ollamaUsageAccount(94)
|
|
account.Extra[OllamaCloudUsageSessionExtraKey] = "cipher:wos-session=secret"
|
|
base := &ollamaUsageTestRepo{upstreamBillingProbeAccountRepo: &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{account.ID: account}}}
|
|
repo := &ollamaRefreshPreflightIdentityChangeRepo{ollamaUsageTestRepo: base}
|
|
upstream := &ollamaUsageHTTPStub{body: ollamaUsageFixture(t)}
|
|
svc := NewOllamaCloudUsageService(repo, upstream, NewSettingService(&upstreamBillingProbeSettingRepo{}, nil), ollamaUsageTestEncryptor{}, true)
|
|
t.Cleanup(svc.Stop)
|
|
|
|
_, err := svc.Refresh(context.Background(), account.ID)
|
|
|
|
require.ErrorIs(t, err, ErrOllamaCloudUsageIdentityChanged)
|
|
require.Zero(t, upstream.calls.Load())
|
|
require.NotContains(t, account.Extra, OllamaCloudUsageSnapshotExtraKey)
|
|
}
|
|
|
|
func TestOllamaCloudUsageRefreshUsesFixedURLCookieAndNoRedirects(t *testing.T) {
|
|
account := ollamaUsageAccount(8)
|
|
account.Extra[OllamaCloudUsageSessionExtraKey] = "cipher:wos-session=browser-secret; tracking=must-not-send"
|
|
repo := &ollamaUsageTestRepo{upstreamBillingProbeAccountRepo: &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{8: account}}}
|
|
upstream := &ollamaUsageHTTPStub{body: ollamaUsageFixture(t)}
|
|
svc := newOllamaUsageTestService(t, repo, upstream, &upstreamBillingProbeSettingRepo{}, true)
|
|
fixedNow := time.Date(2026, time.July, 22, 15, 0, 0, 0, time.UTC)
|
|
svc.now = func() time.Time { return fixedNow }
|
|
|
|
state, err := svc.Refresh(context.Background(), 8)
|
|
require.NoError(t, err)
|
|
require.Equal(t, OllamaCloudUsageStatusOK, state.Snapshot.Status)
|
|
require.Equal(t, "https://ollama.com/settings", upstream.lastRequest.URL.String())
|
|
require.Equal(t, "ollama.com", upstream.lastRequest.Host)
|
|
require.Equal(t, "wos-session=browser-secret", upstream.lastRequest.Header.Get("Cookie"))
|
|
require.NotContains(t, upstream.lastRequest.Header.Get("Cookie"), "tracking")
|
|
require.Empty(t, upstream.lastRequest.Header.Get("Authorization"))
|
|
require.True(t, HTTPUpstreamRedirectsDisabled(upstream.lastRequest.Context()))
|
|
}
|
|
|
|
func TestOllamaCloudUsageManualRefreshUsesShortIndependentInterval(t *testing.T) {
|
|
account := ollamaUsageAccount(12)
|
|
account.Extra[OllamaCloudUsageSessionExtraKey] = "cipher:wos-session=initial"
|
|
repo := &ollamaUsageTestRepo{upstreamBillingProbeAccountRepo: &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{12: account}}}
|
|
upstream := &ollamaUsageHTTPStub{body: ollamaUsageFixture(t)}
|
|
svc := newOllamaUsageTestService(t, repo, upstream, &upstreamBillingProbeSettingRepo{}, true)
|
|
fixedNow := time.Date(2026, time.July, 22, 15, 0, 0, 0, time.UTC)
|
|
svc.now = func() time.Time { return fixedNow }
|
|
|
|
_, err := svc.Refresh(context.Background(), 12)
|
|
require.NoError(t, err)
|
|
_, err = svc.Refresh(context.Background(), 12)
|
|
require.ErrorIs(t, err, ErrOllamaCloudUsageRefreshRateLimited)
|
|
require.Equal(t, int64(1), upstream.calls.Load())
|
|
|
|
// Saving a repaired session clears the prior snapshot, so the global 60-minute
|
|
// next_refresh_at does not block immediate administrator verification.
|
|
_, err = svc.SaveSession(context.Background(), 12, "wos-session=repaired")
|
|
require.NoError(t, err)
|
|
_, err = svc.Refresh(context.Background(), 12)
|
|
require.NoError(t, err)
|
|
require.Equal(t, int64(2), upstream.calls.Load())
|
|
}
|
|
|
|
func TestOllamaCloudUsageRefreshUsesHydratedProxyIdentity(t *testing.T) {
|
|
account := ollamaUsageAccount(13)
|
|
account.Extra[OllamaCloudUsageSessionExtraKey] = "cipher:wos-session=secret"
|
|
proxyID := int64(4)
|
|
account.ProxyID = &proxyID
|
|
account.Proxy = &Proxy{
|
|
ID: proxyID, Protocol: "http", Host: "127.0.0.1", Port: 3128,
|
|
Username: "proxy-user", Password: "proxy-pass", Status: StatusActive,
|
|
}
|
|
repo := &ollamaUsageTestRepo{upstreamBillingProbeAccountRepo: &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{13: account}}}
|
|
upstream := &ollamaUsageHTTPStub{body: ollamaUsageFixture(t)}
|
|
svc := newOllamaUsageTestService(t, repo, upstream, &upstreamBillingProbeSettingRepo{}, true)
|
|
|
|
_, err := svc.Refresh(context.Background(), 13)
|
|
require.NoError(t, err)
|
|
require.Equal(t, account.Proxy.URL(), upstream.lastProxyURL)
|
|
}
|
|
|
|
func TestOllamaCloudUsageRedirectAndBodyLimitArePersistedSafely(t *testing.T) {
|
|
for _, test := range []struct {
|
|
name string
|
|
status int
|
|
body []byte
|
|
reason string
|
|
}{
|
|
{"redirect", http.StatusFound, nil, "redirect_blocked"},
|
|
{"body limit", http.StatusOK, make([]byte, ollamaCloudUsageMaxBodyBytes+1), "response_too_large"},
|
|
} {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
account := ollamaUsageAccount(9)
|
|
account.Extra[OllamaCloudUsageSessionExtraKey] = "cipher:wos-session=secret"
|
|
repo := &ollamaUsageTestRepo{upstreamBillingProbeAccountRepo: &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{9: account}}}
|
|
svc := newOllamaUsageTestService(t, repo, &ollamaUsageHTTPStub{status: test.status, body: test.body}, &upstreamBillingProbeSettingRepo{}, true)
|
|
state, err := svc.Refresh(context.Background(), 9)
|
|
require.NoError(t, err)
|
|
require.Equal(t, OllamaCloudUsageStatusFailed, state.Snapshot.Status)
|
|
require.Equal(t, test.reason, state.Snapshot.LastError)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestOllamaCloudUsageRefreshRejectsIdentityChange(t *testing.T) {
|
|
account := ollamaUsageAccount(10)
|
|
account.Extra[OllamaCloudUsageSessionExtraKey] = "cipher:wos-session=secret"
|
|
repo := &ollamaUsageTestRepo{upstreamBillingProbeAccountRepo: &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{10: account}}}
|
|
repo.beforeSnapshot = func() { account.Credentials["api_key"] = "rotated" }
|
|
svc := newOllamaUsageTestService(t, repo, &ollamaUsageHTTPStub{body: ollamaUsageFixture(t)}, &upstreamBillingProbeSettingRepo{}, true)
|
|
_, err := svc.Refresh(context.Background(), 10)
|
|
require.ErrorIs(t, err, ErrOllamaCloudUsageIdentityChanged)
|
|
require.NotContains(t, account.Extra, OllamaCloudUsageSnapshotExtraKey)
|
|
}
|
|
|
|
func TestOllamaCloudUsageRunnerHonorsLeaderLockAndBackoff(t *testing.T) {
|
|
account := ollamaUsageAccount(11)
|
|
account.Extra[OllamaCloudUsageSessionExtraKey] = "cipher:wos-session=secret"
|
|
account.Extra[OllamaCloudUsageAutoRefreshExtraKey] = true
|
|
repo := &ollamaUsageTestRepo{upstreamBillingProbeAccountRepo: &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{11: account}}}
|
|
upstream := &ollamaUsageHTTPStub{body: ollamaUsageFixture(t)}
|
|
settingsRepo := &upstreamBillingProbeSettingRepo{values: map[string]string{
|
|
SettingKeyOllamaCloudUsageSettings: `{"enabled":true,"interval_minutes":60}`,
|
|
}}
|
|
cache := &fakeLeaderLockCache{}
|
|
_, acquired := tryAcquireSingletonLeaderLock(context.Background(), cache, nil, ollamaCloudUsageLeaderLockKey, "peer", time.Minute)
|
|
require.True(t, acquired)
|
|
svc := newOllamaUsageTestService(t, repo, upstream, settingsRepo, true)
|
|
svc.lockCache = cache
|
|
require.NoError(t, svc.RunDue(context.Background()))
|
|
require.Zero(t, upstream.calls.Load())
|
|
require.NoError(t, cache.ReleaseLeaderLock(context.Background(), ollamaCloudUsageLeaderLockKey, "peer"))
|
|
require.NoError(t, svc.RunDue(context.Background()))
|
|
require.Equal(t, int64(1), upstream.calls.Load())
|
|
|
|
firstFailure := nextOllamaCloudUsageDelay(60, 1, 0)
|
|
thirdFailure := nextOllamaCloudUsageDelay(60, 3, 0)
|
|
require.Greater(t, thirdFailure, firstFailure)
|
|
require.GreaterOrEqual(t, nextOllamaCloudUsageDelay(60, 1, 3*time.Hour), 3*time.Hour)
|
|
require.LessOrEqual(t, nextOllamaCloudUsageDelay(60, 20, 0), ollamaCloudUsageMaxDelay+5*time.Minute)
|
|
}
|
|
|
|
func TestOllamaCloudUsageRunnerDisablesAutoRefreshAfterUnpersistableIdentityError(t *testing.T) {
|
|
account := ollamaUsageAccount(14)
|
|
account.Extra[OllamaCloudUsageSessionExtraKey] = "cipher:wos-session=secret"
|
|
account.Extra[OllamaCloudUsageAutoRefreshExtraKey] = true
|
|
missingProxyID := int64(99)
|
|
account.ProxyID = &missingProxyID
|
|
account.Proxy = nil
|
|
repo := &ollamaUsageTestRepo{upstreamBillingProbeAccountRepo: &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{14: account}}}
|
|
settingsRepo := &upstreamBillingProbeSettingRepo{values: map[string]string{
|
|
SettingKeyOllamaCloudUsageSettings: `{"enabled":true,"interval_minutes":60}`,
|
|
}}
|
|
upstream := &ollamaUsageHTTPStub{body: ollamaUsageFixture(t)}
|
|
svc := newOllamaUsageTestService(t, repo, upstream, settingsRepo, true)
|
|
|
|
require.NoError(t, svc.RunDue(context.Background()))
|
|
require.Equal(t, int64(1), repo.disableAutoCalls.Load())
|
|
require.Equal(t, false, account.Extra[OllamaCloudUsageAutoRefreshExtraKey])
|
|
require.Zero(t, upstream.calls.Load())
|
|
|
|
require.NoError(t, svc.RunDue(context.Background()))
|
|
require.Equal(t, int64(1), repo.disableAutoCalls.Load())
|
|
require.Zero(t, upstream.calls.Load())
|
|
}
|
|
|
|
func TestOllamaCloudUsageRunnerIdentityChangePreservesOldGroupAndDoesNotLoop(t *testing.T) {
|
|
anchor := ollamaUsageAccount(15)
|
|
anchor.Credentials["api_key"] = "shared-before-rotation"
|
|
anchor.Extra[OllamaCloudUsageSessionExtraKey] = "cipher:wos-session=secret"
|
|
anchor.Extra[OllamaCloudUsageAutoRefreshExtraKey] = true
|
|
sibling := ollamaUsageAccount(16)
|
|
sibling.Platform = PlatformAnthropic
|
|
sibling.Credentials = map[string]any{"api_key": "shared-before-rotation", "base_url": "https://www.ollama.com:443/v1"}
|
|
sibling.Extra[OllamaCloudUsageSessionExtraKey] = "cipher:wos-session=secret"
|
|
sibling.Extra[OllamaCloudUsageAutoRefreshExtraKey] = true
|
|
dueAnchor := *anchor
|
|
dueAnchor.Credentials = mergeMap(nil, anchor.Credentials)
|
|
dueAnchor.Extra = mergeMap(nil, anchor.Extra)
|
|
repo := &ollamaUsageTestRepo{
|
|
upstreamBillingProbeAccountRepo: &upstreamBillingProbeAccountRepo{accounts: map[int64]*Account{
|
|
anchor.ID: anchor, sibling.ID: sibling,
|
|
}},
|
|
due: []Account{dueAnchor},
|
|
}
|
|
var rotateOnce sync.Once
|
|
repo.beforeSnapshot = func() {
|
|
rotateOnce.Do(func() {
|
|
repo.mu.Lock()
|
|
defer repo.mu.Unlock()
|
|
anchor.Credentials["api_key"] = "rotated-account-key"
|
|
delete(anchor.Extra, OllamaCloudUsageSessionExtraKey)
|
|
delete(anchor.Extra, OllamaCloudUsageAutoRefreshExtraKey)
|
|
delete(anchor.Extra, OllamaCloudUsageSnapshotExtraKey)
|
|
})
|
|
}
|
|
settingsRepo := &upstreamBillingProbeSettingRepo{values: map[string]string{
|
|
SettingKeyOllamaCloudUsageSettings: `{"enabled":true,"interval_minutes":60}`,
|
|
}}
|
|
upstream := &ollamaUsageHTTPStub{body: ollamaUsageFixture(t)}
|
|
svc := newOllamaUsageTestService(t, repo, upstream, settingsRepo, true)
|
|
|
|
require.NoError(t, svc.RunDue(context.Background()))
|
|
require.Equal(t, int64(1), repo.disableAutoAttempts.Load())
|
|
require.Zero(t, repo.disableAutoCalls.Load(), "the stale anchor CAS must not disable the old sibling group")
|
|
require.Equal(t, true, sibling.Extra[OllamaCloudUsageAutoRefreshExtraKey])
|
|
require.NotContains(t, anchor.Extra, OllamaCloudUsageAutoRefreshExtraKey)
|
|
|
|
repo.due = []Account{*anchor, *sibling}
|
|
require.NoError(t, svc.RunDue(context.Background()))
|
|
require.Equal(t, int64(1), repo.disableAutoAttempts.Load(), "the changed account must not be retried")
|
|
require.Equal(t, true, sibling.Extra[OllamaCloudUsageAutoRefreshExtraKey])
|
|
require.NotNil(t, decodeOllamaCloudUsageSnapshot(sibling.Extra), "the still-valid sibling must refresh normally")
|
|
require.Equal(t, int64(2), upstream.calls.Load())
|
|
}
|
|
|
|
func TestOllamaCloudUsageSingleflightConcurrencyAndRunnerSwitches(t *testing.T) {
|
|
accounts := make(map[int64]*Account)
|
|
for id := int64(1); id <= 7; id++ {
|
|
account := ollamaUsageAccount(id)
|
|
account.Extra[OllamaCloudUsageSessionExtraKey] = "cipher:wos-session=secret"
|
|
account.Extra[OllamaCloudUsageAutoRefreshExtraKey] = true
|
|
accounts[id] = account
|
|
}
|
|
repo := &ollamaUsageTestRepo{upstreamBillingProbeAccountRepo: &upstreamBillingProbeAccountRepo{accounts: accounts}}
|
|
unblock := make(chan struct{})
|
|
entered := make(chan struct{}, 10)
|
|
upstream := &ollamaUsageHTTPStub{body: ollamaUsageFixture(t), beforeResponse: func(*http.Request) {
|
|
entered <- struct{}{}
|
|
<-unblock
|
|
}}
|
|
settingsRepo := &upstreamBillingProbeSettingRepo{values: map[string]string{}}
|
|
svc := newOllamaUsageTestService(t, repo, upstream, settingsRepo, true)
|
|
|
|
// Global automatic refresh is fail-safe off by default.
|
|
require.NoError(t, svc.RunDue(context.Background()))
|
|
require.Zero(t, upstream.calls.Load())
|
|
|
|
settingsRepo.values[SettingKeyOllamaCloudUsageSettings] = `{"enabled":true,"interval_minutes":60}`
|
|
var singleflight sync.WaitGroup
|
|
singleflight.Add(2)
|
|
for range 2 {
|
|
go func() {
|
|
defer singleflight.Done()
|
|
_, _ = svc.Refresh(context.Background(), 1)
|
|
}()
|
|
}
|
|
<-entered
|
|
close(unblock)
|
|
singleflight.Wait()
|
|
require.Equal(t, int64(1), upstream.calls.Load())
|
|
|
|
// Clear snapshots so all accounts are due, then verify the shared four-slot bound.
|
|
for _, account := range accounts {
|
|
delete(account.Extra, OllamaCloudUsageSnapshotExtraKey)
|
|
}
|
|
unblock2 := make(chan struct{})
|
|
upstream.beforeResponse = func(*http.Request) { <-unblock2 }
|
|
done := make(chan struct{})
|
|
go func() {
|
|
_ = svc.RunDue(context.Background())
|
|
close(done)
|
|
}()
|
|
require.Eventually(t, func() bool { return upstream.active.Load() == ollamaCloudUsageConcurrency }, time.Second, 10*time.Millisecond)
|
|
close(unblock2)
|
|
<-done
|
|
require.LessOrEqual(t, upstream.maxActive.Load(), int64(ollamaCloudUsageConcurrency))
|
|
require.Equal(t, int64(8), upstream.calls.Load())
|
|
}
|