Initial open-source release (MIT): image2api AI gateway
Full Go backend + Vue 3 frontend, OpenAI-compatible API, multi-provider account pools, billing/admin, Docker one-command deploy with auto HTTPS. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,282 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"backend/internal/model"
|
||||
"backend/internal/repo"
|
||||
"backend/internal/storage"
|
||||
)
|
||||
|
||||
// MaintenanceService runs the periodic self-healing sweep that the Python
|
||||
// original did via a 60s daemon thread plus read-time lazy cleanup. Without it
|
||||
// the Go token pool only ever loses capacity: tokens never re-activate after a
|
||||
// quota reset, cookies never auto-renew, stale pending events permanently block
|
||||
// a user's generation gate, and old media/logs accumulate unbounded.
|
||||
type MaintenanceService struct {
|
||||
tokens *repo.TokenRepository
|
||||
tokenSvc *TokenService
|
||||
events *repo.EventRepository
|
||||
users *repo.UserRepository
|
||||
refresh *RefreshProfileService
|
||||
settings *repo.SiteSettingRepository
|
||||
store *storage.Client
|
||||
inflight *InflightRegistry
|
||||
showcase *repo.ShowcaseRepository
|
||||
interval time.Duration
|
||||
stalePending time.Duration
|
||||
mediaPruneEvery time.Duration
|
||||
lastMediaPrune time.Time
|
||||
}
|
||||
|
||||
func NewMaintenanceService(tokens *repo.TokenRepository, tokenSvc *TokenService, events *repo.EventRepository, users *repo.UserRepository, refresh *RefreshProfileService, settings *repo.SiteSettingRepository, store *storage.Client, inflight *InflightRegistry, showcase *repo.ShowcaseRepository) *MaintenanceService {
|
||||
return &MaintenanceService{
|
||||
tokens: tokens,
|
||||
tokenSvc: tokenSvc,
|
||||
events: events,
|
||||
users: users,
|
||||
refresh: refresh,
|
||||
settings: settings,
|
||||
store: store,
|
||||
inflight: inflight,
|
||||
showcase: showcase,
|
||||
interval: 60 * time.Second,
|
||||
stalePending: 600 * time.Second,
|
||||
mediaPruneEvery: 60 * time.Second,
|
||||
}
|
||||
}
|
||||
|
||||
// Run drives the sweep every interval until ctx is cancelled. It runs one sweep
|
||||
// immediately on startup so a freshly restarted process heals stuck state right
|
||||
// away rather than after the first tick.
|
||||
func (m *MaintenanceService) Run(ctx context.Context) {
|
||||
ticker := time.NewTicker(m.interval)
|
||||
defer ticker.Stop()
|
||||
m.tick(ctx)
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
m.tick(ctx)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// syncRecoveredQuota re-probes each just-recovered account so its displayed
|
||||
// balance reflects the post-reset value (these providers only sync quota when
|
||||
// accessed). krea additionally needs /app (Activate) to actually grant the daily
|
||||
// free balance before billing-data reports it. Bounded concurrency avoids a
|
||||
// thundering herd at the daily reset.
|
||||
func (m *MaintenanceService) syncRecoveredQuota(accs []model.TokenAccount) {
|
||||
sem := make(chan struct{}, 4)
|
||||
var wg sync.WaitGroup
|
||||
for _, acc := range accs {
|
||||
switch acc.Pool {
|
||||
case "chatgpt", "leonardo", "krea", "imagine":
|
||||
default:
|
||||
continue
|
||||
}
|
||||
wg.Add(1)
|
||||
sem <- struct{}{}
|
||||
go func(a model.TokenAccount) {
|
||||
defer wg.Done()
|
||||
defer func() { <-sem }()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
|
||||
defer cancel()
|
||||
if a.Pool == "krea" && m.tokenSvc.krea != nil {
|
||||
m.tokenSvc.krea.Activate(ctx, a.Value)
|
||||
}
|
||||
_, _ = m.tokenSvc.Quota(ctx, a.Pool, a.ID)
|
||||
}(acc)
|
||||
}
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
func (m *MaintenanceService) tick(ctx context.Context) {
|
||||
// 1. Re-activate quota-exhausted tokens whose reset time has passed, then
|
||||
// auto-sync their real balance — these providers only refresh quota when
|
||||
// accessed, so recovery alone would leave a stale 0/—. For krea the sync
|
||||
// must first load /app (Activate) to grant the daily free balance.
|
||||
if recovered, err := m.tokens.RecoverQuota(ctx); err != nil {
|
||||
log.Printf("maintenance: recover_quota: %v", err)
|
||||
} else if len(recovered) > 0 {
|
||||
log.Printf("maintenance: recovered %d quota token(s)", len(recovered))
|
||||
if m.tokenSvc != nil {
|
||||
go m.syncRecoveredQuota(recovered)
|
||||
}
|
||||
}
|
||||
|
||||
// 1a. Roll the 恢复时间 marker of ACTIVE daily-reset accounts forward to the next
|
||||
// future reset (same time-of-day, +1 day) so the column never shows a stale
|
||||
// past time. Limited accounts are intentionally skipped (RecoverQuota owns
|
||||
// their marker). adobe/leonardo/krea/imagine all renew daily.
|
||||
if _, err := m.tokens.RollResetMarkers(ctx, []string{"adobe", "leonardo", "krea", "imagine"}); err != nil {
|
||||
log.Printf("maintenance: roll_reset: %v", err)
|
||||
}
|
||||
|
||||
// 1b. Runway tokens have no refresh — once the JWT expiry (its reset marker)
|
||||
// passes, mark them dead directly instead of letting them 401 on next use.
|
||||
if n, err := m.tokens.ExpireByReset(ctx, "runway"); err != nil {
|
||||
log.Printf("maintenance: expire_runway: %v", err)
|
||||
} else if n > 0 {
|
||||
log.Printf("maintenance: expired %d runway token(s)", n)
|
||||
}
|
||||
|
||||
// 1c. Proactively renew krea/imagine sessions ~10min before expiry so a
|
||||
// dormant account's rotating refresh_token never lapses (a dead token
|
||||
// can't be recovered and, for krea, blocks the daily free-credit meter
|
||||
// from being re-created). Only near-expiry accounts hit the network.
|
||||
if m.tokenSvc != nil {
|
||||
m.tokenSvc.RefreshExpiringTokens(ctx)
|
||||
// 1d. Once-per-day krea /app activation for accounts not yet synced since the
|
||||
// daily reset — krea only grants the free balance after /app loads, so an
|
||||
// always-active account (never went 限额) would otherwise read 0 / 402
|
||||
// after each reset. Self-guarded + background; no-op once all are done.
|
||||
m.tokenSvc.ActivateKreaDue(ctx)
|
||||
}
|
||||
|
||||
// 2. Auto-renew Adobe cookies whose refresh interval has elapsed.
|
||||
if m.refresh != nil {
|
||||
if n, err := m.refresh.RefreshDue(ctx); err != nil {
|
||||
log.Printf("maintenance: refresh_due: %v", err)
|
||||
} else if n > 0 {
|
||||
log.Printf("maintenance: refreshed %d cookie profile(s)", n)
|
||||
}
|
||||
}
|
||||
|
||||
// 3. Fail long-pending events so they stop blocking the per-user gate, and
|
||||
// refund the credits debited up-front for each abandoned generation (the
|
||||
// normal failure-refund path never ran for a process-restart orphan).
|
||||
if purged, err := m.events.PurgeStale(ctx, m.stalePending); err != nil {
|
||||
log.Printf("maintenance: purge_stale: %v", err)
|
||||
} else if len(purged) > 0 {
|
||||
refunded := 0
|
||||
cancelled := 0
|
||||
for _, e := range purged {
|
||||
// Stop the generation goroutine if it's still running, so it doesn't
|
||||
// keep grinding for minutes and surface a late "success" on this
|
||||
// just-abandoned event.
|
||||
if m.inflight != nil && m.inflight.Cancel(e.ID) {
|
||||
cancelled++
|
||||
}
|
||||
// Attribute the abandoned failure back to the account it was using
|
||||
// (the normal markTokenFailure path never ran for an orphaned job).
|
||||
if e.AccountID != "" {
|
||||
if err := m.tokens.IncrementFail(ctx, e.AccountID); err != nil {
|
||||
log.Printf("maintenance: fail-count abandoned event %s (account %s): %v", e.ID, e.AccountID, err)
|
||||
}
|
||||
}
|
||||
if e.UserID == "" || e.Cost <= 0 {
|
||||
continue
|
||||
}
|
||||
// Exactly-once: only refund if we win the claim (the in-flight request
|
||||
// may have already refunded itself on its own failure path).
|
||||
claimed, err := m.events.MarkRefunded(ctx, e.ID)
|
||||
if err != nil {
|
||||
log.Printf("maintenance: claim refund %s: %v", e.ID, err)
|
||||
continue
|
||||
}
|
||||
if !claimed {
|
||||
continue
|
||||
}
|
||||
if _, err := m.users.AdjustCredits(ctx, e.UserID, e.Cost); err != nil {
|
||||
log.Printf("maintenance: refund abandoned event %s (user %s, %.0f): %v", e.ID, e.UserID, e.Cost, err)
|
||||
} else {
|
||||
refunded++
|
||||
}
|
||||
}
|
||||
log.Printf("maintenance: marked %d stale pending event(s) failed, refunded %d, cancelled %d in-flight", len(purged), refunded, cancelled)
|
||||
}
|
||||
|
||||
// 4. Enforce the admin-configured log retention window.
|
||||
m.pruneLogs(ctx)
|
||||
|
||||
// 5. Enforce the media retention window. Runs every 60s like the log prune;
|
||||
// mediaPruneEvery still gates it in case the interval is ever shortened.
|
||||
if time.Since(m.lastMediaPrune) >= m.mediaPruneEvery {
|
||||
m.pruneMedia(ctx)
|
||||
m.lastMediaPrune = time.Now()
|
||||
}
|
||||
}
|
||||
|
||||
func (m *MaintenanceService) pruneLogs(ctx context.Context) {
|
||||
days := m.retentionDays(ctx, "logs.retention_days")
|
||||
if days <= 0 {
|
||||
return
|
||||
}
|
||||
if _, err := m.events.PurgeOlderThan(ctx, time.Duration(days)*24*time.Hour); err != nil {
|
||||
log.Printf("maintenance: purge_older_than: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func (m *MaintenanceService) pruneMedia(ctx context.Context) {
|
||||
if m.store == nil || !m.store.Configured() {
|
||||
return
|
||||
}
|
||||
days := m.retentionDays(ctx, "media.retention_days")
|
||||
if days <= 0 {
|
||||
return
|
||||
}
|
||||
cutoff := time.Now().Add(-time.Duration(days) * 24 * time.Hour)
|
||||
objs, err := m.store.List(ctx, "")
|
||||
if err != nil {
|
||||
log.Printf("maintenance: list media: %v", err)
|
||||
return
|
||||
}
|
||||
// Files referenced by the homepage showcase are kept forever, no matter how
|
||||
// old — deleting them would break the public landing page.
|
||||
var pinned map[string]struct{}
|
||||
if m.showcase != nil {
|
||||
if pinned, err = m.showcase.PublicFileSet(ctx); err != nil {
|
||||
log.Printf("maintenance: showcase file set: %v", err)
|
||||
pinned = nil
|
||||
}
|
||||
}
|
||||
removed, skipped := 0, 0
|
||||
var clearedKeys []string
|
||||
for _, o := range objs {
|
||||
if !o.LastModified.Before(cutoff) {
|
||||
continue
|
||||
}
|
||||
if _, ok := pinned[strings.TrimLeft(o.Key, "/")]; ok {
|
||||
skipped++
|
||||
continue
|
||||
}
|
||||
if err := m.store.Delete(ctx, o.Key); err != nil {
|
||||
log.Printf("maintenance: delete %s: %v", o.Key, err)
|
||||
continue
|
||||
}
|
||||
removed++
|
||||
// event_log.file stores the same key — blank those rows so the log views
|
||||
// don't dangle a 404 preview.
|
||||
clearedKeys = append(clearedKeys, o.Key)
|
||||
}
|
||||
if removed > 0 || skipped > 0 {
|
||||
log.Printf("maintenance: pruned %d expired media object(s), kept %d showcase-pinned", removed, skipped)
|
||||
}
|
||||
if len(clearedKeys) > 0 {
|
||||
if n, err := m.events.ClearFiles(ctx, clearedKeys); err != nil {
|
||||
log.Printf("maintenance: clear_files: %v", err)
|
||||
} else if n > 0 {
|
||||
log.Printf("maintenance: cleared file ref on %d log row(s)", n)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (m *MaintenanceService) retentionDays(ctx context.Context, key string) int {
|
||||
raw, err := m.settings.GetValue(ctx, key)
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
days, err := strconv.Atoi(strings.TrimSpace(raw))
|
||||
if err != nil || days <= 0 {
|
||||
return 0
|
||||
}
|
||||
return days
|
||||
}
|
||||
Reference in New Issue
Block a user