ollama_cloud_usage.go 43 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235
  1. package service
  2. import (
  3. "context"
  4. "crypto/sha256"
  5. "database/sql"
  6. "encoding/hex"
  7. "encoding/json"
  8. "errors"
  9. "fmt"
  10. "io"
  11. "maps"
  12. "math/rand/v2"
  13. "net/http"
  14. "net/url"
  15. "strconv"
  16. "strings"
  17. "sync"
  18. "time"
  19. "github.com/Wei-Shaw/sub2api/internal/config"
  20. infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
  21. "github.com/Wei-Shaw/sub2api/internal/pkg/logger"
  22. "github.com/google/uuid"
  23. "golang.org/x/net/http/httpguts"
  24. "golang.org/x/sync/errgroup"
  25. "golang.org/x/sync/singleflight"
  26. )
  27. const (
  28. OllamaCloudUsageSessionExtraKey = "ollama_cloud_usage_session"
  29. OllamaCloudUsageAutoRefreshExtraKey = "ollama_cloud_usage_auto_refresh"
  30. OllamaCloudUsageSnapshotExtraKey = "ollama_cloud_usage_snapshot"
  31. // OllamaCloudUsageMinFetchInterval is the hard floor between two successful
  32. // fetches of the same group, mirroring the floor nextOllamaCloudUsageDelay
  33. // applies to next_refresh_at. Activity may bring a refresh forward to this
  34. // bound but never past it. Exported so the repository can apply the same
  35. // floor inside the SQL due filter.
  36. OllamaCloudUsageMinFetchInterval = ollamaCloudUsageMinIntervalMinutes * time.Minute
  37. ollamaCloudUsageSettingsURL = "https://ollama.com/settings"
  38. ollamaCloudUsageDefaultIntervalMinutes = 60
  39. ollamaCloudUsageMinIntervalMinutes = 15
  40. ollamaCloudUsageMaxIntervalMinutes = 24 * 60
  41. ollamaCloudUsageDefaultDebounceMinutes = 1
  42. ollamaCloudUsageMinDebounceMinutes = 1
  43. ollamaCloudUsageMaxDebounceMinutes = 60
  44. ollamaCloudUsageCycleInterval = time.Minute
  45. ollamaCloudUsageManualRefreshInterval = 30 * time.Second
  46. ollamaCloudUsageRequestTimeout = 15 * time.Second
  47. ollamaCloudUsageMaxBodyBytes = 512 * 1024
  48. ollamaCloudUsageMaxSessionBytes = 16 * 1024
  49. ollamaCloudUsageMaxPerCycle = 20
  50. ollamaCloudUsageConcurrency = 4
  51. ollamaCloudUsageMaxDelay = 24 * time.Hour
  52. ollamaCloudUsageLeaderLockKey = "ollama:cloud:usage:leader"
  53. ollamaCloudUsageLeaderLockTTL = 2 * time.Minute
  54. )
  55. var (
  56. ErrOllamaCloudUsageUnavailable = infraerrors.ServiceUnavailable(
  57. "OLLAMA_CLOUD_USAGE_UNAVAILABLE", "Ollama Cloud usage is unavailable",
  58. )
  59. ErrOllamaCloudUsageAccountInvalid = infraerrors.BadRequest(
  60. "OLLAMA_CLOUD_USAGE_ACCOUNT_INVALID", "account must be an OpenAI or Anthropic API key account using https://ollama.com",
  61. )
  62. ErrOllamaCloudUsageSessionRequired = infraerrors.BadRequest(
  63. "OLLAMA_CLOUD_USAGE_SESSION_REQUIRED", "an Ollama web session must be configured first",
  64. )
  65. ErrOllamaCloudUsageEncryptionKey = infraerrors.BadRequest(
  66. "OLLAMA_CLOUD_USAGE_ENCRYPTION_KEY_NOT_CONFIGURED", "cannot store an Ollama web session without a fixed TOTP_ENCRYPTION_KEY",
  67. )
  68. ErrOllamaCloudUsageIdentityChanged = infraerrors.Conflict(
  69. "OLLAMA_CLOUD_USAGE_IDENTITY_CHANGED", "account identity or Ollama web session changed during refresh; retry",
  70. )
  71. ErrOllamaCloudUsageRefreshRateLimited = infraerrors.TooManyRequests(
  72. "OLLAMA_CLOUD_USAGE_REFRESH_RATE_LIMITED", "Ollama Cloud usage can be refreshed manually once every 30 seconds",
  73. )
  74. errOllamaCloudUsageUnauthorizedHTML = errors.New("settings HTML is a sign-in page")
  75. )
  76. const (
  77. OllamaCloudUsageStatusOK = "ok"
  78. OllamaCloudUsageStatusUnauthorized = "unauthorized"
  79. OllamaCloudUsageStatusFailed = "failed"
  80. )
  81. // OllamaCloudUsageSettings controls the opt-in request-driven refresh runner.
  82. //
  83. // IntervalMinutes is the max-wait bound: when model requests keep arriving and
  84. // the trailing debounce keeps sliding, a refresh is forced after this long.
  85. // DebounceMinutes is the quiet period after the latest request in a group.
  86. type OllamaCloudUsageSettings struct {
  87. Enabled bool `json:"enabled"`
  88. IntervalMinutes int `json:"interval_minutes"` // max wait while requests continue
  89. DebounceMinutes int `json:"debounce_minutes"` // trailing quiet period after last request
  90. }
  91. // OllamaCloudUsageWindow is a narrow, sanitized view of one official usage window.
  92. type OllamaCloudUsageWindow struct {
  93. UsedPercent float64 `json:"used_percent"`
  94. ResetAt *time.Time `json:"reset_at,omitempty"`
  95. ResetText string `json:"reset_text,omitempty"`
  96. }
  97. // OllamaCloudUsageModelWindow identifies the official window for a model count.
  98. type OllamaCloudUsageModelWindow string
  99. const (
  100. OllamaCloudUsageModelWindowFiveHour OllamaCloudUsageModelWindow = "five_hour"
  101. OllamaCloudUsageModelWindowSevenDay OllamaCloudUsageModelWindow = "seven_day"
  102. )
  103. // OllamaCloudUsageModel is the window-scoped model/request pair exposed by Ollama's usage DOM.
  104. type OllamaCloudUsageModel struct {
  105. Model string `json:"model"`
  106. Window OllamaCloudUsageModelWindow `json:"window"`
  107. Requests int64 `json:"requests"`
  108. }
  109. // OllamaCloudUsageData intentionally excludes raw HTML and browser-session data.
  110. type OllamaCloudUsageData struct {
  111. Plan string `json:"plan,omitempty"`
  112. FiveHour *OllamaCloudUsageWindow `json:"five_hour,omitempty"`
  113. SevenDay *OllamaCloudUsageWindow `json:"seven_day,omitempty"`
  114. Balance string `json:"balance,omitempty"`
  115. Models []OllamaCloudUsageModel `json:"models,omitempty"`
  116. }
  117. // OllamaCloudUsageSnapshot is the only usage observation persisted in account extra.
  118. //
  119. // NextRefreshAt remains a persisted compatibility field. For status=ok it is a
  120. // max-wait horizon marker only; automatic success refreshes are driven by model
  121. // request activity (group last_used_at + debounce/max-wait), not by this field
  122. // alone. For failed/unauthorized snapshots it is the failure not-before time
  123. // (Retry-After / exponential backoff) and is enforced as max(activityDue, NextRefreshAt).
  124. type OllamaCloudUsageSnapshot struct {
  125. Status string `json:"status"`
  126. Data *OllamaCloudUsageData `json:"data,omitempty"`
  127. FetchedAt *time.Time `json:"fetched_at,omitempty"`
  128. LastAttemptAt time.Time `json:"last_attempt_at"`
  129. NextRefreshAt time.Time `json:"next_refresh_at"`
  130. FailureCount int `json:"failure_count,omitempty"`
  131. HTTPStatus int `json:"http_status,omitempty"`
  132. LastError string `json:"last_error,omitempty"`
  133. }
  134. // OllamaCloudUsageState is the dedicated DTO exposed to administrators.
  135. type OllamaCloudUsageState struct {
  136. AccountID int64 `json:"account_id"`
  137. Eligible bool `json:"eligible"`
  138. Configured bool `json:"configured"`
  139. AutoRefreshEnabled bool `json:"auto_refresh_enabled"`
  140. EncryptionKeyConfigured bool `json:"encryption_key_configured"`
  141. Snapshot *OllamaCloudUsageSnapshot `json:"snapshot,omitempty"`
  142. }
  143. type ollamaCloudUsageRepository interface {
  144. ListOllamaCloudUsageGroupAccounts(context.Context, []*Account) ([]Account, error)
  145. SaveOllamaCloudUsageSession(context.Context, *Account, string, bool) error
  146. DeleteOllamaCloudUsageSession(context.Context, *Account) error
  147. SetOllamaCloudUsageAutoRefresh(context.Context, *Account, bool) error
  148. UpdateOllamaCloudUsageSnapshot(context.Context, *Account, *OllamaCloudUsageSnapshot) error
  149. DisableOllamaCloudUsageAutoRefresh(context.Context, *Account) error
  150. ListDueOllamaCloudUsageAccounts(context.Context, time.Time, time.Duration, time.Duration, int) ([]Account, error)
  151. }
  152. // GetOllamaCloudUsageSettings returns fail-safe defaults when the setting is absent.
  153. func (s *SettingService) GetOllamaCloudUsageSettings(ctx context.Context) (*OllamaCloudUsageSettings, error) {
  154. defaults := defaultOllamaCloudUsageSettings()
  155. if s == nil || s.settingRepo == nil {
  156. return defaults, nil
  157. }
  158. raw, err := s.settingRepo.GetValue(ctx, SettingKeyOllamaCloudUsageSettings)
  159. if err != nil {
  160. if errors.Is(err, ErrSettingNotFound) {
  161. return defaults, nil
  162. }
  163. return nil, fmt.Errorf("get Ollama Cloud usage settings: %w", err)
  164. }
  165. if strings.TrimSpace(raw) == "" {
  166. return defaults, nil
  167. }
  168. settings := *defaults
  169. if err := json.Unmarshal([]byte(raw), &settings); err != nil {
  170. return nil, fmt.Errorf("parse Ollama Cloud usage settings: %w", err)
  171. }
  172. if settings.IntervalMinutes == 0 {
  173. settings.IntervalMinutes = defaults.IntervalMinutes
  174. }
  175. if settings.DebounceMinutes == 0 {
  176. settings.DebounceMinutes = defaults.DebounceMinutes
  177. }
  178. normalizeOllamaCloudUsageSettings(&settings)
  179. return &settings, nil
  180. }
  181. func (s *SettingService) SetOllamaCloudUsageSettings(ctx context.Context, settings *OllamaCloudUsageSettings) error {
  182. if s == nil || s.settingRepo == nil {
  183. return ErrOllamaCloudUsageUnavailable
  184. }
  185. if settings == nil {
  186. return infraerrors.BadRequest("INVALID_OLLAMA_CLOUD_USAGE_SETTINGS", "settings cannot be nil")
  187. }
  188. if settings.DebounceMinutes == 0 {
  189. // Legacy clients that omit debounce_minutes keep the fail-safe default.
  190. settings.DebounceMinutes = ollamaCloudUsageDefaultDebounceMinutes
  191. }
  192. if settings.IntervalMinutes < ollamaCloudUsageMinIntervalMinutes || settings.IntervalMinutes > ollamaCloudUsageMaxIntervalMinutes {
  193. return infraerrors.BadRequest(
  194. "INVALID_OLLAMA_CLOUD_USAGE_INTERVAL",
  195. fmt.Sprintf("interval_minutes must be between %d and %d", ollamaCloudUsageMinIntervalMinutes, ollamaCloudUsageMaxIntervalMinutes),
  196. )
  197. }
  198. if settings.DebounceMinutes < ollamaCloudUsageMinDebounceMinutes || settings.DebounceMinutes > ollamaCloudUsageMaxDebounceMinutes {
  199. return infraerrors.BadRequest(
  200. "INVALID_OLLAMA_CLOUD_USAGE_DEBOUNCE",
  201. fmt.Sprintf("debounce_minutes must be between %d and %d", ollamaCloudUsageMinDebounceMinutes, ollamaCloudUsageMaxDebounceMinutes),
  202. )
  203. }
  204. // The due time is min(lastUsed+debounce, fetchedAt+maxWait). Once the debounce
  205. // reaches the max wait the debounce term can never win, so the knob would be
  206. // silently inert instead of doing what the operator asked for.
  207. if settings.DebounceMinutes >= settings.IntervalMinutes {
  208. return infraerrors.BadRequest(
  209. "INVALID_OLLAMA_CLOUD_USAGE_DEBOUNCE",
  210. fmt.Sprintf("debounce_minutes (%d) must be less than interval_minutes (%d)", settings.DebounceMinutes, settings.IntervalMinutes),
  211. )
  212. }
  213. normalizeOllamaCloudUsageSettings(settings)
  214. data, err := json.Marshal(settings)
  215. if err != nil {
  216. return fmt.Errorf("marshal Ollama Cloud usage settings: %w", err)
  217. }
  218. return s.settingRepo.Set(ctx, SettingKeyOllamaCloudUsageSettings, string(data))
  219. }
  220. func defaultOllamaCloudUsageSettings() *OllamaCloudUsageSettings {
  221. return &OllamaCloudUsageSettings{
  222. Enabled: false,
  223. IntervalMinutes: ollamaCloudUsageDefaultIntervalMinutes,
  224. DebounceMinutes: ollamaCloudUsageDefaultDebounceMinutes,
  225. }
  226. }
  227. func normalizeOllamaCloudUsageSettings(settings *OllamaCloudUsageSettings) {
  228. if settings.IntervalMinutes < ollamaCloudUsageMinIntervalMinutes {
  229. settings.IntervalMinutes = ollamaCloudUsageMinIntervalMinutes
  230. }
  231. if settings.IntervalMinutes > ollamaCloudUsageMaxIntervalMinutes {
  232. settings.IntervalMinutes = ollamaCloudUsageMaxIntervalMinutes
  233. }
  234. if settings.DebounceMinutes <= 0 {
  235. settings.DebounceMinutes = ollamaCloudUsageDefaultDebounceMinutes
  236. }
  237. if settings.DebounceMinutes < ollamaCloudUsageMinDebounceMinutes {
  238. settings.DebounceMinutes = ollamaCloudUsageMinDebounceMinutes
  239. }
  240. if settings.DebounceMinutes > ollamaCloudUsageMaxDebounceMinutes {
  241. settings.DebounceMinutes = ollamaCloudUsageMaxDebounceMinutes
  242. }
  243. }
  244. func ollamaCloudUsageDurations(settings *OllamaCloudUsageSettings) (debounce, maxWait time.Duration) {
  245. normalized := defaultOllamaCloudUsageSettings()
  246. if settings != nil {
  247. *normalized = *settings
  248. }
  249. normalizeOllamaCloudUsageSettings(normalized)
  250. return time.Duration(normalized.DebounceMinutes) * time.Minute,
  251. time.Duration(normalized.IntervalMinutes) * time.Minute
  252. }
  253. // ollamaCloudUsageIsAutoRefreshDue decides whether a configured auto-refresh
  254. // group should fetch now. groupLastUsedAt must be MAX(last_used_at) across the
  255. // exact api_key group so shared multi-platform accounts do not miss activity.
  256. //
  257. // Success: a request must be newer than fetched_at; dueAt = min(lastUsed+debounce, fetchedAt+maxWait).
  258. // Failure: a request must be newer than last_attempt_at; activity due uses the same min formula,
  259. // then dueAt = max(activityDue, next_refresh_at) so Retry-After / exponential backoff win.
  260. // Missing or invalid snapshots fail open to a first fetch.
  261. func ollamaCloudUsageIsAutoRefreshDue(
  262. snapshot *OllamaCloudUsageSnapshot,
  263. groupLastUsedAt *time.Time,
  264. now time.Time,
  265. debounce, maxWait time.Duration,
  266. ) bool {
  267. dueAt, ok := ollamaCloudUsageAutoRefreshDueAt(snapshot, groupLastUsedAt, debounce, maxWait)
  268. if !ok {
  269. return false
  270. }
  271. return !now.Before(dueAt)
  272. }
  273. func ollamaCloudUsageAutoRefreshDueAt(
  274. snapshot *OllamaCloudUsageSnapshot,
  275. groupLastUsedAt *time.Time,
  276. debounce, maxWait time.Duration,
  277. ) (time.Time, bool) {
  278. if debounce <= 0 {
  279. debounce = time.Duration(ollamaCloudUsageDefaultDebounceMinutes) * time.Minute
  280. }
  281. if maxWait <= 0 {
  282. maxWait = time.Duration(ollamaCloudUsageDefaultIntervalMinutes) * time.Minute
  283. }
  284. if snapshot == nil {
  285. return time.Time{}, true
  286. }
  287. switch snapshot.Status {
  288. case OllamaCloudUsageStatusOK:
  289. if snapshot.FetchedAt == nil || snapshot.FetchedAt.IsZero() {
  290. return time.Time{}, true
  291. }
  292. fetchedAt := snapshot.FetchedAt.UTC()
  293. if groupLastUsedAt == nil || !groupLastUsedAt.After(fetchedAt) {
  294. return time.Time{}, false
  295. }
  296. lastUsed := groupLastUsedAt.UTC()
  297. dueAt := minTime(lastUsed.Add(debounce), fetchedAt.Add(maxWait))
  298. // Keep the pre-existing hard floor between successful fetches. The success
  299. // path no longer consults next_refresh_at, which is where
  300. // nextOllamaCloudUsageDelay used to apply ollamaCloudUsageMinIntervalMinutes;
  301. // without this, request traffic spaced slightly wider than the debounce
  302. // drives the group's outbound rate far above the previous minimum.
  303. if floor := fetchedAt.Add(OllamaCloudUsageMinFetchInterval); dueAt.Before(floor) {
  304. return floor, true
  305. }
  306. return dueAt, true
  307. case OllamaCloudUsageStatusFailed, OllamaCloudUsageStatusUnauthorized:
  308. if snapshot.LastAttemptAt.IsZero() {
  309. return time.Time{}, true
  310. }
  311. lastAttempt := snapshot.LastAttemptAt.UTC()
  312. if groupLastUsedAt == nil || !groupLastUsedAt.After(lastAttempt) {
  313. return time.Time{}, false
  314. }
  315. lastUsed := groupLastUsedAt.UTC()
  316. activityDue := minTime(lastUsed.Add(debounce), lastAttempt.Add(maxWait))
  317. if !snapshot.NextRefreshAt.IsZero() && snapshot.NextRefreshAt.UTC().After(activityDue) {
  318. return snapshot.NextRefreshAt.UTC(), true
  319. }
  320. return activityDue, true
  321. default:
  322. return time.Time{}, true
  323. }
  324. }
  325. // maxOllamaCloudUsageGroupLastUsed returns the newest last_used_at among group members.
  326. func maxOllamaCloudUsageGroupLastUsed(accounts []Account) *time.Time {
  327. var latest *time.Time
  328. for i := range accounts {
  329. candidate := accounts[i].LastUsedAt
  330. if candidate == nil || candidate.IsZero() {
  331. continue
  332. }
  333. if latest == nil || candidate.After(*latest) {
  334. ts := candidate.UTC()
  335. latest = &ts
  336. }
  337. }
  338. return latest
  339. }
  340. // scheduleOllamaCloudUsageActivity records that an Ollama Cloud API-key account
  341. // actually attempted an upstream model request (including 429/5xx/transport errors).
  342. // Local auth/validation failures must not call this. DeferredService dedupes writes.
  343. func scheduleOllamaCloudUsageActivity(deferred *DeferredService, account *Account) {
  344. if deferred == nil || account == nil || !IsOllamaCloudUsageAccount(account) {
  345. return
  346. }
  347. deferred.ScheduleLastUsedUpdate(account.ID)
  348. }
  349. // OllamaCloudUsageService refreshes the official settings HTML without affecting routing state.
  350. type OllamaCloudUsageService struct {
  351. accountRepo AccountRepository
  352. httpUpstream HTTPUpstream
  353. settingService *SettingService
  354. encryptor SecretEncryptor
  355. encryptionKeyConfigured bool
  356. parentCtx context.Context
  357. parentCancel context.CancelFunc
  358. wg sync.WaitGroup
  359. mu sync.Mutex
  360. started bool
  361. stopped bool
  362. cycleMu sync.Mutex
  363. refreshGroup singleflight.Group
  364. refreshSlots chan struct{}
  365. now func() time.Time
  366. lockCache LeaderLockCache
  367. db *sql.DB
  368. instanceID string
  369. }
  370. func NewOllamaCloudUsageService(
  371. accountRepo AccountRepository,
  372. httpUpstream HTTPUpstream,
  373. settingService *SettingService,
  374. encryptor SecretEncryptor,
  375. encryptionKeyConfigured bool,
  376. ) *OllamaCloudUsageService {
  377. ctx, cancel := context.WithCancel(context.Background())
  378. return &OllamaCloudUsageService{
  379. accountRepo: accountRepo,
  380. httpUpstream: httpUpstream,
  381. settingService: settingService,
  382. encryptor: encryptor,
  383. encryptionKeyConfigured: encryptionKeyConfigured,
  384. parentCtx: ctx,
  385. parentCancel: cancel,
  386. refreshSlots: make(chan struct{}, ollamaCloudUsageConcurrency),
  387. now: time.Now,
  388. instanceID: uuid.NewString(),
  389. }
  390. }
  391. func ProvideOllamaCloudUsageService(
  392. accountRepo AccountRepository,
  393. httpUpstream HTTPUpstream,
  394. settingService *SettingService,
  395. encryptor SecretEncryptor,
  396. cfg *config.Config,
  397. lockCache LeaderLockCache,
  398. db *sql.DB,
  399. ) *OllamaCloudUsageService {
  400. keyConfigured := cfg != nil && cfg.Totp.EncryptionKeyConfigured
  401. svc := NewOllamaCloudUsageService(accountRepo, httpUpstream, settingService, encryptor, keyConfigured)
  402. svc.lockCache = lockCache
  403. svc.db = db
  404. svc.Start()
  405. return svc
  406. }
  407. func (s *OllamaCloudUsageService) Start() {
  408. if s == nil {
  409. return
  410. }
  411. s.mu.Lock()
  412. if s.started || s.stopped {
  413. s.mu.Unlock()
  414. return
  415. }
  416. s.started = true
  417. s.wg.Add(1)
  418. s.mu.Unlock()
  419. go s.runLoop()
  420. }
  421. func (s *OllamaCloudUsageService) Stop() {
  422. if s == nil {
  423. return
  424. }
  425. s.mu.Lock()
  426. if s.stopped {
  427. s.mu.Unlock()
  428. return
  429. }
  430. s.stopped = true
  431. s.parentCancel()
  432. s.mu.Unlock()
  433. s.wg.Wait()
  434. }
  435. func (s *OllamaCloudUsageService) runLoop() {
  436. defer s.wg.Done()
  437. _ = s.RunDue(s.parentCtx)
  438. ticker := time.NewTicker(ollamaCloudUsageCycleInterval)
  439. defer ticker.Stop()
  440. for {
  441. select {
  442. case <-s.parentCtx.Done():
  443. return
  444. case <-ticker.C:
  445. if err := s.RunDue(s.parentCtx); err != nil {
  446. logger.LegacyPrintf("service.ollama_cloud_usage", "run_due_failed: err=%v", err)
  447. }
  448. }
  449. }
  450. }
  451. func (s *OllamaCloudUsageService) GetSettings(ctx context.Context) (*OllamaCloudUsageSettings, error) {
  452. if s == nil || s.settingService == nil {
  453. return defaultOllamaCloudUsageSettings(), nil
  454. }
  455. return s.settingService.GetOllamaCloudUsageSettings(ctx)
  456. }
  457. func (s *OllamaCloudUsageService) UpdateSettings(ctx context.Context, settings *OllamaCloudUsageSettings) error {
  458. if s == nil || s.settingService == nil {
  459. return ErrOllamaCloudUsageUnavailable
  460. }
  461. return s.settingService.SetOllamaCloudUsageSettings(ctx, settings)
  462. }
  463. func (s *OllamaCloudUsageService) GetState(ctx context.Context, accountID int64) (*OllamaCloudUsageState, error) {
  464. if s == nil || s.accountRepo == nil {
  465. return nil, ErrOllamaCloudUsageUnavailable
  466. }
  467. account, err := s.accountRepo.GetByID(ctx, accountID)
  468. if err != nil {
  469. return nil, err
  470. }
  471. if err := s.ResolveAccounts(ctx, []*Account{account}); err != nil {
  472. return nil, err
  473. }
  474. state := OllamaCloudUsageStateFromAccount(account)
  475. s.EnrichState(state)
  476. return state, nil
  477. }
  478. // ResolveAccounts overlays group-owned managed state onto the supplied account
  479. // objects. The repository resolves all matching siblings in one bounded query,
  480. // so account-list responses do not issue one query per row.
  481. func (s *OllamaCloudUsageService) ResolveAccounts(ctx context.Context, accounts []*Account) error {
  482. if s == nil || s.accountRepo == nil || len(accounts) == 0 {
  483. return nil
  484. }
  485. writer, ok := s.accountRepo.(ollamaCloudUsageRepository)
  486. if !ok {
  487. return nil
  488. }
  489. eligible := make([]*Account, 0, len(accounts))
  490. for _, account := range accounts {
  491. if _, ok := ollamaCloudUsageGroupFingerprint(account); ok {
  492. eligible = append(eligible, account)
  493. }
  494. }
  495. if len(eligible) == 0 {
  496. return nil
  497. }
  498. siblings, err := writer.ListOllamaCloudUsageGroupAccounts(ctx, eligible)
  499. if err != nil {
  500. return fmt.Errorf("resolve Ollama Cloud usage groups: %w", err)
  501. }
  502. sources := make(map[string]*Account)
  503. for index := range siblings {
  504. candidate := &siblings[index]
  505. fingerprint, valid := ollamaCloudUsageGroupFingerprint(candidate)
  506. if !valid || !ollamaCloudUsageConfigured(candidate) {
  507. continue
  508. }
  509. current := sources[fingerprint]
  510. if current == nil || candidate.UpdatedAt.After(current.UpdatedAt) ||
  511. (candidate.UpdatedAt.Equal(current.UpdatedAt) && candidate.ID < current.ID) {
  512. sources[fingerprint] = candidate
  513. }
  514. }
  515. resolvedSources := make(map[string]*Account, len(sources))
  516. for fingerprint, source := range sources {
  517. clone := *source
  518. clone.Extra = make(map[string]any, len(source.Extra))
  519. maps.Copy(clone.Extra, source.Extra)
  520. resolvedSources[fingerprint] = &clone
  521. }
  522. for index := range siblings {
  523. candidate := &siblings[index]
  524. fingerprint, valid := ollamaCloudUsageGroupFingerprint(candidate)
  525. source := resolvedSources[fingerprint]
  526. if !valid || source == nil || !sameOllamaCloudUsageSession(source, candidate) {
  527. continue
  528. }
  529. candidateSnapshot := decodeOllamaCloudUsageSnapshot(candidate.Extra)
  530. currentSnapshot := decodeOllamaCloudUsageSnapshot(source.Extra)
  531. if candidateSnapshot != nil && (currentSnapshot == nil || candidateSnapshot.LastAttemptAt.After(currentSnapshot.LastAttemptAt)) {
  532. source.Extra[OllamaCloudUsageSnapshotExtraKey] = candidate.Extra[OllamaCloudUsageSnapshotExtraKey]
  533. }
  534. }
  535. for _, account := range eligible {
  536. fingerprint, _ := ollamaCloudUsageGroupFingerprint(account)
  537. applyOllamaCloudUsageManagedExtra(account, resolvedSources[fingerprint])
  538. }
  539. return nil
  540. }
  541. func sameOllamaCloudUsageSession(left, right *Account) bool {
  542. if left == nil || right == nil || left.Extra == nil || right.Extra == nil {
  543. return false
  544. }
  545. leftSession, leftOK := left.Extra[OllamaCloudUsageSessionExtraKey].(string)
  546. rightSession, rightOK := right.Extra[OllamaCloudUsageSessionExtraKey].(string)
  547. return leftOK && rightOK && leftSession != "" && leftSession == rightSession
  548. }
  549. func applyOllamaCloudUsageManagedExtra(target, source *Account) {
  550. if target == nil {
  551. return
  552. }
  553. if target.Extra == nil {
  554. target.Extra = make(map[string]any)
  555. }
  556. for _, key := range []string{
  557. OllamaCloudUsageSessionExtraKey,
  558. OllamaCloudUsageAutoRefreshExtraKey,
  559. OllamaCloudUsageSnapshotExtraKey,
  560. } {
  561. delete(target.Extra, key)
  562. if source != nil && source.Extra != nil {
  563. if value, ok := source.Extra[key]; ok {
  564. target.Extra[key] = value
  565. }
  566. }
  567. }
  568. }
  569. func (s *OllamaCloudUsageService) SaveSession(ctx context.Context, accountID int64, session string) (*OllamaCloudUsageState, error) {
  570. if s == nil || s.accountRepo == nil || s.encryptor == nil {
  571. return nil, ErrOllamaCloudUsageUnavailable
  572. }
  573. if !s.encryptionKeyConfigured {
  574. return nil, ErrOllamaCloudUsageEncryptionKey
  575. }
  576. normalized, err := normalizeOllamaCloudUsageCookie(session)
  577. if err != nil {
  578. return nil, infraerrors.BadRequest("INVALID_OLLAMA_CLOUD_USAGE_SESSION", err.Error())
  579. }
  580. account, err := s.accountRepo.GetByID(ctx, accountID)
  581. if err != nil {
  582. return nil, err
  583. }
  584. if !IsOllamaCloudUsageAccount(account) {
  585. return nil, ErrOllamaCloudUsageAccountInvalid
  586. }
  587. if err := s.ResolveAccounts(ctx, []*Account{account}); err != nil {
  588. return nil, err
  589. }
  590. ciphertext, err := s.encryptor.Encrypt(normalized)
  591. if err != nil {
  592. return nil, fmt.Errorf("encrypt Ollama web session: %w", err)
  593. }
  594. writer, ok := s.accountRepo.(ollamaCloudUsageRepository)
  595. if !ok {
  596. return nil, ErrOllamaCloudUsageUnavailable
  597. }
  598. preserveAutoRefresh := ollamaCloudUsageConfigured(account) && ollamaCloudUsageAutoRefreshEnabled(account)
  599. if err := writer.SaveOllamaCloudUsageSession(ctx, account, ciphertext, preserveAutoRefresh); err != nil {
  600. return nil, err
  601. }
  602. return s.GetState(ctx, accountID)
  603. }
  604. func (s *OllamaCloudUsageService) DeleteSession(ctx context.Context, accountID int64) (*OllamaCloudUsageState, error) {
  605. if s == nil || s.accountRepo == nil {
  606. return nil, ErrOllamaCloudUsageUnavailable
  607. }
  608. account, err := s.accountRepo.GetByID(ctx, accountID)
  609. if err != nil {
  610. return nil, err
  611. }
  612. if !IsOllamaCloudUsageAccount(account) {
  613. return nil, ErrOllamaCloudUsageAccountInvalid
  614. }
  615. if err := s.ResolveAccounts(ctx, []*Account{account}); err != nil {
  616. return nil, err
  617. }
  618. writer, ok := s.accountRepo.(ollamaCloudUsageRepository)
  619. if !ok {
  620. return nil, ErrOllamaCloudUsageUnavailable
  621. }
  622. if err := writer.DeleteOllamaCloudUsageSession(ctx, account); err != nil {
  623. return nil, err
  624. }
  625. return s.GetState(ctx, accountID)
  626. }
  627. func (s *OllamaCloudUsageService) SetAutoRefresh(ctx context.Context, accountID int64, enabled bool) (*OllamaCloudUsageState, error) {
  628. if s == nil || s.accountRepo == nil {
  629. return nil, ErrOllamaCloudUsageUnavailable
  630. }
  631. account, err := s.accountRepo.GetByID(ctx, accountID)
  632. if err != nil {
  633. return nil, err
  634. }
  635. if !IsOllamaCloudUsageAccount(account) {
  636. return nil, ErrOllamaCloudUsageAccountInvalid
  637. }
  638. if err := s.ResolveAccounts(ctx, []*Account{account}); err != nil {
  639. return nil, err
  640. }
  641. if enabled && !ollamaCloudUsageConfigured(account) {
  642. return nil, ErrOllamaCloudUsageSessionRequired
  643. }
  644. writer, ok := s.accountRepo.(ollamaCloudUsageRepository)
  645. if !ok {
  646. return nil, ErrOllamaCloudUsageUnavailable
  647. }
  648. if err := writer.SetOllamaCloudUsageAutoRefresh(ctx, account, enabled); err != nil {
  649. return nil, err
  650. }
  651. return s.GetState(ctx, accountID)
  652. }
  653. func (s *OllamaCloudUsageService) Refresh(ctx context.Context, accountID int64) (*OllamaCloudUsageState, error) {
  654. settings, err := s.GetSettings(ctx)
  655. if err != nil {
  656. return nil, err
  657. }
  658. if _, err := s.refreshAccount(ctx, accountID, settings, false); err != nil {
  659. return nil, err
  660. }
  661. return s.GetState(ctx, accountID)
  662. }
  663. func (s *OllamaCloudUsageService) RunDue(ctx context.Context) error {
  664. if s == nil || s.accountRepo == nil {
  665. return nil
  666. }
  667. s.cycleMu.Lock()
  668. defer s.cycleMu.Unlock()
  669. settings, err := s.GetSettings(ctx)
  670. if err != nil {
  671. return err
  672. }
  673. if !settings.Enabled {
  674. return nil
  675. }
  676. release, acquired := tryAcquireSingletonLeaderLock(ctx, s.lockCache, s.db, ollamaCloudUsageLeaderLockKey, s.instanceID, ollamaCloudUsageLeaderLockTTL)
  677. if !acquired {
  678. return nil
  679. }
  680. defer release()
  681. writer, ok := s.accountRepo.(ollamaCloudUsageRepository)
  682. if !ok {
  683. return ErrOllamaCloudUsageUnavailable
  684. }
  685. now := s.currentTime()
  686. debounce, maxWait := ollamaCloudUsageDurations(settings)
  687. accounts, err := writer.ListDueOllamaCloudUsageAccounts(ctx, now, debounce, maxWait, ollamaCloudUsageMaxPerCycle)
  688. if err != nil {
  689. return fmt.Errorf("list due Ollama Cloud usage accounts: %w", err)
  690. }
  691. var group errgroup.Group
  692. seenGroups := make(map[string]struct{}, len(accounts))
  693. for index := range accounts {
  694. account := accounts[index]
  695. fingerprint, valid := ollamaCloudUsageGroupFingerprint(&account)
  696. if !valid || !account.IsActive() || !ollamaCloudUsageConfigured(&account) || !ollamaCloudUsageAutoRefreshEnabled(&account) {
  697. continue
  698. }
  699. if _, duplicate := seenGroups[fingerprint]; duplicate {
  700. continue
  701. }
  702. seenGroups[fingerprint] = struct{}{}
  703. snapshot := decodeOllamaCloudUsageSnapshot(account.Extra)
  704. // ListDue stamps Account.LastUsedAt with the api_key group MAX(last_used_at).
  705. if !ollamaCloudUsageIsAutoRefreshDue(snapshot, account.LastUsedAt, now, debounce, maxWait) {
  706. continue
  707. }
  708. accountID := account.ID
  709. expected := account
  710. group.Go(func() error {
  711. if _, refreshErr := s.refreshAccount(ctx, accountID, settings, true); refreshErr != nil {
  712. if errors.Is(refreshErr, ErrOllamaCloudUsageIdentityChanged) {
  713. if disableErr := writer.DisableOllamaCloudUsageAutoRefresh(ctx, &expected); disableErr != nil {
  714. logger.LegacyPrintf("service.ollama_cloud_usage", "disable_auto_refresh_failed: account_id=%d err=%v", accountID, disableErr)
  715. }
  716. return nil
  717. }
  718. logger.LegacyPrintf("service.ollama_cloud_usage", "refresh_due_failed: account_id=%d err=%v", accountID, refreshErr)
  719. }
  720. return nil
  721. })
  722. }
  723. return group.Wait()
  724. }
  725. func (s *OllamaCloudUsageService) refreshAccount(ctx context.Context, accountID int64, settings *OllamaCloudUsageSettings, requireEnabled bool) (*OllamaCloudUsageSnapshot, error) {
  726. if s == nil || s.accountRepo == nil {
  727. return nil, ErrOllamaCloudUsageUnavailable
  728. }
  729. if settings == nil {
  730. settings = defaultOllamaCloudUsageSettings()
  731. }
  732. intervalMinutes := settings.IntervalMinutes
  733. debounce, maxWait := ollamaCloudUsageDurations(settings)
  734. anchor, err := s.accountRepo.GetByID(ctx, accountID)
  735. if err != nil {
  736. return nil, err
  737. }
  738. key, valid := ollamaCloudUsageGroupFingerprint(anchor)
  739. if !valid {
  740. return nil, ErrOllamaCloudUsageAccountInvalid
  741. }
  742. value, err, _ := s.refreshGroup.Do(key, func() (any, error) {
  743. select {
  744. case s.refreshSlots <- struct{}{}:
  745. defer func() { <-s.refreshSlots }()
  746. case <-ctx.Done():
  747. return nil, ctx.Err()
  748. }
  749. account, loadErr := s.accountRepo.GetByID(ctx, accountID)
  750. if loadErr != nil {
  751. return nil, loadErr
  752. }
  753. currentKey, currentValid := ollamaCloudUsageGroupFingerprint(account)
  754. if !currentValid {
  755. return nil, ErrOllamaCloudUsageAccountInvalid
  756. }
  757. if currentKey != key {
  758. return nil, ErrOllamaCloudUsageIdentityChanged
  759. }
  760. if err := s.ResolveAccounts(ctx, []*Account{account}); err != nil {
  761. return nil, err
  762. }
  763. if !ollamaCloudUsageConfigured(account) {
  764. return nil, ErrOllamaCloudUsageSessionRequired
  765. }
  766. if !requireEnabled {
  767. if snapshot := decodeOllamaCloudUsageSnapshot(account.Extra); snapshot != nil && !snapshot.LastAttemptAt.IsZero() {
  768. retryAt := snapshot.LastAttemptAt.Add(ollamaCloudUsageManualRefreshInterval)
  769. if now := s.currentTime(); now.Before(retryAt) {
  770. remaining := retryAt.Sub(now)
  771. seconds := int((remaining + time.Second - 1) / time.Second)
  772. return nil, ErrOllamaCloudUsageRefreshRateLimited.WithMetadata(map[string]string{
  773. "retry_after_seconds": strconv.Itoa(seconds),
  774. })
  775. }
  776. }
  777. }
  778. if requireEnabled {
  779. if !account.IsActive() || !ollamaCloudUsageAutoRefreshEnabled(account) {
  780. return nil, nil
  781. }
  782. groupLastUsed := account.LastUsedAt
  783. if writer, ok := s.accountRepo.(ollamaCloudUsageRepository); ok {
  784. siblings, listErr := writer.ListOllamaCloudUsageGroupAccounts(ctx, []*Account{account})
  785. if listErr != nil {
  786. // Fall back to this account's own last_used_at. That is a narrower
  787. // activity signal than the group maximum, so the due check may skip a
  788. // refresh it would otherwise have run; surface it rather than
  789. // silently changing the due semantics.
  790. logger.LegacyPrintf("service.ollama_cloud_usage",
  791. "group_last_used_lookup_failed: account_id=%d err=%v", account.ID, listErr)
  792. } else {
  793. groupLastUsed = maxOllamaCloudUsageGroupLastUsed(siblings)
  794. }
  795. }
  796. if !ollamaCloudUsageIsAutoRefreshDue(decodeOllamaCloudUsageSnapshot(account.Extra), groupLastUsed, s.currentTime(), debounce, maxWait) {
  797. return nil, nil
  798. }
  799. }
  800. return s.refreshLoadedAccount(ctx, account, intervalMinutes)
  801. })
  802. if err != nil || value == nil {
  803. return nil, err
  804. }
  805. snapshot, ok := value.(*OllamaCloudUsageSnapshot)
  806. if !ok {
  807. return nil, fmt.Errorf("invalid Ollama Cloud usage refresh result")
  808. }
  809. return snapshot, nil
  810. }
  811. func (s *OllamaCloudUsageService) refreshLoadedAccount(ctx context.Context, account *Account, intervalMinutes int) (*OllamaCloudUsageSnapshot, error) {
  812. now := s.currentTime().UTC()
  813. ciphertext, _ := account.Extra[OllamaCloudUsageSessionExtraKey].(string)
  814. if ciphertext == "" {
  815. return nil, ErrOllamaCloudUsageSessionRequired
  816. }
  817. if !s.encryptionKeyConfigured || s.encryptor == nil {
  818. return nil, ErrOllamaCloudUsageEncryptionKey
  819. }
  820. cookie, err := s.encryptor.Decrypt(ciphertext)
  821. if err != nil {
  822. return nil, infraerrors.ServiceUnavailable("OLLAMA_CLOUD_USAGE_SESSION_DECRYPT_FAILED", "stored Ollama web session cannot be decrypted")
  823. }
  824. cookie, err = normalizeOllamaCloudUsageCookie(cookie)
  825. if err != nil {
  826. return nil, infraerrors.ServiceUnavailable("OLLAMA_CLOUD_USAGE_SESSION_INVALID", "stored Ollama web session is invalid")
  827. }
  828. if s.httpUpstream == nil {
  829. return nil, ErrOllamaCloudUsageUnavailable
  830. }
  831. proxyURL := ""
  832. if account.ProxyID != nil {
  833. if account.Proxy == nil || account.Proxy.ID != *account.ProxyID {
  834. return nil, ErrOllamaCloudUsageIdentityChanged
  835. }
  836. proxyURL = account.Proxy.URL()
  837. }
  838. requestCtx, cancel := context.WithTimeout(WithHTTPUpstreamRedirectsDisabled(ctx), ollamaCloudUsageRequestTimeout)
  839. defer cancel()
  840. req, err := http.NewRequestWithContext(requestCtx, http.MethodGet, ollamaCloudUsageSettingsURL, nil)
  841. if err != nil || !isExactOllamaCloudSettingsURL(req.URL) {
  842. return nil, ErrOllamaCloudUsageUnavailable
  843. }
  844. req.Header.Set("Accept", "text/html,application/xhtml+xml")
  845. req.Header.Set("Cookie", cookie)
  846. req.Header.Set("User-Agent", "sub2api-ollama-usage/1")
  847. resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, account.Concurrency)
  848. if err != nil {
  849. return s.persistFailure(ctx, account, intervalMinutes, now, 0, "request_failed", 0, false)
  850. }
  851. if resp == nil || resp.Body == nil {
  852. return s.persistFailure(ctx, account, intervalMinutes, now, 0, "empty_response", 0, false)
  853. }
  854. defer func() { _ = resp.Body.Close() }()
  855. if resp.Request != nil && !isExactOllamaCloudSettingsURL(resp.Request.URL) {
  856. return s.persistFailure(ctx, account, intervalMinutes, now, resp.StatusCode, "response_host_mismatch", 0, false)
  857. }
  858. if resp.StatusCode >= 300 && resp.StatusCode < 400 {
  859. return s.persistFailure(ctx, account, intervalMinutes, now, resp.StatusCode, "redirect_blocked", retryAfter(resp.Header, now), false)
  860. }
  861. if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden {
  862. return s.persistFailure(ctx, account, intervalMinutes, now, resp.StatusCode, "unauthorized", retryAfter(resp.Header, now), true)
  863. }
  864. if resp.StatusCode < 200 || resp.StatusCode >= 300 {
  865. return s.persistFailure(ctx, account, intervalMinutes, now, resp.StatusCode, "http_error", retryAfter(resp.Header, now), false)
  866. }
  867. body, readErr := io.ReadAll(io.LimitReader(resp.Body, ollamaCloudUsageMaxBodyBytes+1))
  868. if readErr != nil {
  869. return s.persistFailure(ctx, account, intervalMinutes, now, resp.StatusCode, "response_read_failed", 0, false)
  870. }
  871. if len(body) > ollamaCloudUsageMaxBodyBytes {
  872. return s.persistFailure(ctx, account, intervalMinutes, now, resp.StatusCode, "response_too_large", 0, false)
  873. }
  874. data, parseErr := parseOllamaCloudUsageHTML(body)
  875. if errors.Is(parseErr, errOllamaCloudUsageUnauthorizedHTML) {
  876. return s.persistFailure(ctx, account, intervalMinutes, now, resp.StatusCode, "unauthorized", 0, true)
  877. }
  878. if parseErr != nil {
  879. return s.persistFailure(ctx, account, intervalMinutes, now, resp.StatusCode, "invalid_html", 0, false)
  880. }
  881. snapshot := &OllamaCloudUsageSnapshot{
  882. Status: OllamaCloudUsageStatusOK,
  883. Data: data,
  884. FetchedAt: &now,
  885. LastAttemptAt: now,
  886. NextRefreshAt: now.Add(nextOllamaCloudUsageDelay(intervalMinutes, 0, 0)),
  887. HTTPStatus: resp.StatusCode,
  888. }
  889. if err := s.updateSnapshot(ctx, account, snapshot); err != nil {
  890. return nil, err
  891. }
  892. return snapshot, nil
  893. }
  894. func (s *OllamaCloudUsageService) persistFailure(
  895. ctx context.Context,
  896. account *Account,
  897. intervalMinutes int,
  898. now time.Time,
  899. httpStatus int,
  900. reason string,
  901. retryAfterDuration time.Duration,
  902. unauthorized bool,
  903. ) (*OllamaCloudUsageSnapshot, error) {
  904. previous := decodeOllamaCloudUsageSnapshot(account.Extra)
  905. failureCount := 1
  906. if previous != nil {
  907. failureCount = previous.FailureCount + 1
  908. }
  909. status := OllamaCloudUsageStatusFailed
  910. if unauthorized {
  911. status = OllamaCloudUsageStatusUnauthorized
  912. }
  913. snapshot := &OllamaCloudUsageSnapshot{
  914. Status: status,
  915. LastAttemptAt: now,
  916. NextRefreshAt: now.Add(nextOllamaCloudUsageDelay(intervalMinutes, failureCount, retryAfterDuration)),
  917. FailureCount: failureCount,
  918. HTTPStatus: httpStatus,
  919. LastError: reason,
  920. }
  921. if previous != nil {
  922. snapshot.Data = previous.Data
  923. snapshot.FetchedAt = previous.FetchedAt
  924. }
  925. if err := s.updateSnapshot(ctx, account, snapshot); err != nil {
  926. return nil, err
  927. }
  928. return snapshot, nil
  929. }
  930. func (s *OllamaCloudUsageService) updateSnapshot(ctx context.Context, account *Account, snapshot *OllamaCloudUsageSnapshot) error {
  931. writer, ok := s.accountRepo.(ollamaCloudUsageRepository)
  932. if !ok {
  933. return ErrOllamaCloudUsageUnavailable
  934. }
  935. return writer.UpdateOllamaCloudUsageSnapshot(ctx, account, snapshot)
  936. }
  937. // EnrichState adds service-owned runtime configuration to an account-derived state.
  938. func (s *OllamaCloudUsageService) EnrichState(state *OllamaCloudUsageState) {
  939. if state == nil {
  940. return
  941. }
  942. state.EncryptionKeyConfigured = s != nil && s.encryptionKeyConfigured
  943. }
  944. func OllamaCloudUsageStateFromAccount(account *Account) *OllamaCloudUsageState {
  945. state := &OllamaCloudUsageState{}
  946. if account == nil {
  947. return state
  948. }
  949. state.AccountID = account.ID
  950. state.Eligible = IsOllamaCloudUsageAccount(account)
  951. if !state.Eligible {
  952. return state
  953. }
  954. state.Configured = ollamaCloudUsageConfigured(account)
  955. state.AutoRefreshEnabled = state.Configured && ollamaCloudUsageAutoRefreshEnabled(account)
  956. state.Snapshot = decodeOllamaCloudUsageSnapshot(account.Extra)
  957. return state
  958. }
  959. // isOllamaCloudUsagePlatform 收敛 Ollama Cloud 用量窗口的平台白名单。官方
  960. // ollama.com base_url 除官方两平台外,也允许挂在经 OpenAI 网关转发的国产
  961. // OpenAI 兼容平台下(用户把 Ollama Cloud key 挂在 kimi/zhipu/deepseek 分组
  962. // 里跑托管的 glm/kimi/deepseek 模型)。repository 侧 SQL 白名单
  963. // (ollamaCloudUsagePlatformsSQL)是本列表的镜像,两侧必须同步修改。
  964. func isOllamaCloudUsagePlatform(platform string) bool {
  965. switch platform {
  966. case PlatformOpenAI, PlatformAnthropic, PlatformKimi, PlatformZhipu, PlatformDeepseek, PlatformMiniMax:
  967. return true
  968. default:
  969. return false
  970. }
  971. }
  972. func IsOllamaCloudUsageAccount(account *Account) bool {
  973. if account == nil || account.Type != AccountTypeAPIKey || !isOllamaCloudUsagePlatform(account.Platform) {
  974. return false
  975. }
  976. baseURL, _ := account.Credentials["base_url"].(string)
  977. return isOllamaCloudBaseURL(baseURL)
  978. }
  979. func isOllamaCloudBaseURL(raw string) bool {
  980. raw = strings.TrimSpace(raw)
  981. if raw == "" || strings.ContainsAny(raw, "?#") {
  982. return false
  983. }
  984. parsed, err := url.Parse(raw)
  985. if err != nil || parsed.Opaque != "" || !strings.EqualFold(parsed.Scheme, "https") || parsed.User != nil || parsed.ForceQuery || parsed.RawQuery != "" || parsed.Fragment != "" || parsed.RawFragment != "" {
  986. return false
  987. }
  988. hostname := strings.ToLower(parsed.Hostname())
  989. if hostname != "ollama.com" && hostname != "www.ollama.com" {
  990. return false
  991. }
  992. authority := strings.ToLower(parsed.Host)
  993. if authority != hostname && authority != hostname+":443" {
  994. return false
  995. }
  996. if parsed.RawPath != "" {
  997. return false
  998. }
  999. return parsed.Path == "" || parsed.Path == "/v1"
  1000. }
  1001. func ollamaCloudUsageIdentity(account *Account) map[string]any {
  1002. if !IsOllamaCloudUsageAccount(account) {
  1003. return nil
  1004. }
  1005. apiKey, ok := account.Credentials["api_key"].(string)
  1006. if !ok || apiKey == "" {
  1007. return nil
  1008. }
  1009. return map[string]any{"host": "ollama.com", "api_key": apiKey}
  1010. }
  1011. func ollamaCloudUsageGroupFingerprint(account *Account) (string, bool) {
  1012. identity := ollamaCloudUsageIdentity(account)
  1013. if identity == nil {
  1014. return "", false
  1015. }
  1016. apiKey, _ := identity["api_key"].(string)
  1017. sum := sha256.Sum256([]byte("ollama.com\x00" + apiKey))
  1018. return hex.EncodeToString(sum[:]), true
  1019. }
  1020. func isExactOllamaCloudSettingsURL(parsed *url.URL) bool {
  1021. return parsed != nil && parsed.Scheme == "https" && parsed.Host == "ollama.com" && parsed.Path == "/settings" &&
  1022. parsed.User == nil && parsed.RawQuery == "" && parsed.Fragment == "" && parsed.RawPath == ""
  1023. }
  1024. func normalizeOllamaCloudUsageCookie(raw string) (string, error) {
  1025. if len(raw) > ollamaCloudUsageMaxSessionBytes {
  1026. return "", errors.New("session is too large")
  1027. }
  1028. raw = strings.TrimSpace(raw)
  1029. if strings.ContainsAny(raw, "\r\n") {
  1030. return "", errors.New("session contains invalid header characters")
  1031. }
  1032. if raw == "" {
  1033. return "", errors.New("session cannot be empty")
  1034. }
  1035. if !httpguts.ValidHeaderFieldValue(raw) {
  1036. return "", errors.New("session contains invalid header characters")
  1037. }
  1038. blockedAttributes := map[string]struct{}{
  1039. "domain": {}, "path": {}, "expires": {}, "max-age": {}, "samesite": {}, "secure": {}, "httponly": {}, "partitioned": {},
  1040. }
  1041. parts := strings.Split(raw, ";")
  1042. normalized := make([]string, 0, len(parts))
  1043. seen := make(map[string]struct{}, len(parts))
  1044. for _, part := range parts {
  1045. part = strings.TrimSpace(part)
  1046. name, value, ok := strings.Cut(part, "=")
  1047. name = strings.TrimSpace(name)
  1048. value = strings.TrimSpace(value)
  1049. if !ok || name == "" || value == "" || !httpguts.ValidHeaderFieldName(name) || strings.HasPrefix(name, "$") {
  1050. return "", errors.New("session must be a Cookie header containing name=value pairs")
  1051. }
  1052. lowerName := strings.ToLower(name)
  1053. if _, blocked := blockedAttributes[lowerName]; blocked {
  1054. return "", errors.New("paste a Cookie header, not a Set-Cookie value with attributes")
  1055. }
  1056. if _, duplicate := seen[lowerName]; duplicate {
  1057. return "", errors.New("session contains duplicate cookie names")
  1058. }
  1059. if strings.ContainsAny(value, ";\r\n") {
  1060. return "", errors.New("session contains an invalid cookie value")
  1061. }
  1062. seen[lowerName] = struct{}{}
  1063. if isAllowedOllamaCloudSessionCookie(name) {
  1064. normalized = append(normalized, name+"="+value)
  1065. }
  1066. }
  1067. if len(normalized) == 0 {
  1068. return "", errors.New("session does not contain an allowed Ollama session cookie")
  1069. }
  1070. return strings.Join(normalized, "; "), nil
  1071. }
  1072. func isAllowedOllamaCloudSessionCookie(name string) bool {
  1073. switch name {
  1074. case "wos-session", "__Secure-session", "session", "ollama_session", "__Host-ollama_session":
  1075. return true
  1076. }
  1077. for _, base := range []string{
  1078. "next-auth.session-token",
  1079. "__Secure-next-auth.session-token",
  1080. "authjs.session-token",
  1081. "__Secure-authjs.session-token",
  1082. } {
  1083. if name == base {
  1084. return true
  1085. }
  1086. if suffix, ok := strings.CutPrefix(name, base+"."); ok && suffix != "" {
  1087. validShard := true
  1088. for _, char := range suffix {
  1089. if char < '0' || char > '9' {
  1090. validShard = false
  1091. break
  1092. }
  1093. }
  1094. if validShard {
  1095. return true
  1096. }
  1097. }
  1098. }
  1099. return false
  1100. }
  1101. func ollamaCloudUsageConfigured(account *Account) bool {
  1102. if account == nil || account.Extra == nil {
  1103. return false
  1104. }
  1105. value, ok := account.Extra[OllamaCloudUsageSessionExtraKey].(string)
  1106. return ok && strings.TrimSpace(value) != ""
  1107. }
  1108. func ollamaCloudUsageAutoRefreshEnabled(account *Account) bool {
  1109. if account == nil || account.Extra == nil {
  1110. return false
  1111. }
  1112. enabled, ok := account.Extra[OllamaCloudUsageAutoRefreshExtraKey].(bool)
  1113. return ok && enabled
  1114. }
  1115. func decodeOllamaCloudUsageSnapshot(extra map[string]any) *OllamaCloudUsageSnapshot {
  1116. if extra == nil {
  1117. return nil
  1118. }
  1119. value, ok := extra[OllamaCloudUsageSnapshotExtraKey]
  1120. if !ok || value == nil {
  1121. return nil
  1122. }
  1123. raw, err := json.Marshal(value)
  1124. if err != nil {
  1125. return nil
  1126. }
  1127. var snapshot OllamaCloudUsageSnapshot
  1128. if err := json.Unmarshal(raw, &snapshot); err != nil {
  1129. return nil
  1130. }
  1131. if snapshot.Status != OllamaCloudUsageStatusOK && snapshot.Status != OllamaCloudUsageStatusUnauthorized && snapshot.Status != OllamaCloudUsageStatusFailed {
  1132. return nil
  1133. }
  1134. return &snapshot
  1135. }
  1136. func nextOllamaCloudUsageDelay(intervalMinutes, failureCount int, retryAfterDuration time.Duration) time.Duration {
  1137. minimumDelay := retryAfterDuration
  1138. base := time.Duration(intervalMinutes) * time.Minute
  1139. if base < ollamaCloudUsageMinIntervalMinutes*time.Minute {
  1140. base = ollamaCloudUsageMinIntervalMinutes * time.Minute
  1141. }
  1142. if failureCount > 0 {
  1143. shift := min(failureCount-1, 6)
  1144. base *= time.Duration(1 << shift)
  1145. }
  1146. if base > ollamaCloudUsageMaxDelay {
  1147. base = ollamaCloudUsageMaxDelay
  1148. }
  1149. if retryAfterDuration > base {
  1150. base = retryAfterDuration
  1151. }
  1152. jitterRange := base / 10
  1153. if jitterRange > 5*time.Minute {
  1154. jitterRange = 5 * time.Minute
  1155. }
  1156. if jitterRange > 0 {
  1157. base += time.Duration(rand.Int64N(int64(jitterRange)*2+1)) - jitterRange
  1158. }
  1159. if base < minimumDelay {
  1160. return minimumDelay
  1161. }
  1162. if base < time.Minute {
  1163. return time.Minute
  1164. }
  1165. return base
  1166. }
  1167. func (s *OllamaCloudUsageService) currentTime() time.Time {
  1168. if s != nil && s.now != nil {
  1169. return s.now()
  1170. }
  1171. return time.Now()
  1172. }