feat: add Grok (grok.com) video provider; runway nano-banana image; unify runway/grok pool policy

- grok provider: imagine text/image-to-video (media.post.create → conversations/new),
  GetGrokCreditsConfig credit query (remaining = 100 - used) + weekly reset, spoofed
  x-statsig-id (no Cloudflare clearance needed), /api/auth/session email lookup,
  6 reference images, 10 concurrent jobs/account, no token refresh (dead = dead)
- runway: nano-banana-2 image flow (Nano Banana 2); drop pre-deduct + post-success
  reconcile; out-of-credits/403 → dead (no revive); 10-ratio support
- imagine: drop post-success credit reconcile (consistent with krea)
- account gate: per-account N-concurrency (grok=10, others=1)
- admin: provider health lists all 7 providers; frontend import auto-detects Grok SSO
- docs: README (CN/EN) updated to 7 providers + Grok

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
2026-06-28 04:01:06 +08:00
co-authored by Claude Opus 4.8
parent 8633addf64
commit 17cd289dfd
16 changed files with 1442 additions and 166 deletions
+4 -2
View File
@@ -12,6 +12,7 @@ import (
"backend/internal/model"
"backend/internal/provider/adobe"
"backend/internal/provider/chatgpt"
"backend/internal/provider/grok"
"backend/internal/provider/imagine"
"backend/internal/provider/krea"
"backend/internal/provider/leonardo"
@@ -108,14 +109,15 @@ func NewApp(ctx context.Context) (*App, error) {
leonardoClient := leonardo.NewClient("")
kreaClient := krea.NewClient("")
imagineClient := imagine.NewClient("")
v1Svc := service.NewV1Service(cfg, modelRepo, userRepo, eventRepo, tokenRepo, siteRepo, adobeClient, chatGPTClient, runwayClient, leonardoClient, kreaClient, imagineClient, rustfsClient)
grokClient := grok.NewClient("")
v1Svc := service.NewV1Service(cfg, modelRepo, userRepo, eventRepo, tokenRepo, siteRepo, adobeClient, chatGPTClient, runwayClient, leonardoClient, kreaClient, imagineClient, grokClient, rustfsClient)
siteSvc := service.NewSiteService(siteRepo, cfg.AppTitle)
showcaseSvc := service.NewShowcaseService(showcaseRepo)
adminReadSvc := service.NewAdminReadService(cfg, userRepo, modelRepo, eventRepo, siteRepo, tokenRepo, cdkRepo, rustfsClient)
adminWriteSvc := service.NewAdminWriteService(userRepo, showcaseRepo, modelRepo, eventRepo, apiKeyRepo)
cdkSvc := service.NewCDKService(cdkRepo, userRepo)
apiKeySvc := service.NewAPIKeyService(apiKeyRepo)
tokenSvc := service.NewTokenService(tokenRepo, refreshRepo, eventRepo, siteRepo, adobeClient, chatGPTClient, runwayClient, leonardoClient, kreaClient, imagineClient)
tokenSvc := service.NewTokenService(tokenRepo, refreshRepo, eventRepo, siteRepo, adobeClient, chatGPTClient, runwayClient, leonardoClient, kreaClient, imagineClient, grokClient)
refreshSvc := service.NewRefreshProfileService(refreshRepo, tokenRepo, adobeClient)
// Enable refresh-then-retry on a mid-request Adobe 401 (re-mint access token
// from the cookie). Wired post-construction to avoid a ctor init cycle.
@@ -102,6 +102,36 @@ func (h *ProviderAdminHandler) ImportRunwayToken(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"ok": true, "id": item.ID, "status": item.Status, "pending": item.Status == "pending"})
}
func (h *ProviderAdminHandler) ImportGrokToken(c *gin.Context) {
var body struct {
AccessToken string `json:"access_token"`
Value string `json:"value"`
SSO string `json:"sso"`
Name string `json:"name"`
ID string `json:"id"`
}
if err := c.ShouldBindJSON(&body); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"})
return
}
token := body.AccessToken
for _, v := range []string{body.Value, body.SSO} {
if token == "" {
token = v
}
}
name := body.Name
if name == "" {
name = body.ID
}
item, err := h.tokens.ImportGrokToken(c.Request.Context(), token, name)
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
return
}
c.JSON(http.StatusOK, gin.H{"ok": true, "id": item.ID, "status": item.Status, "pending": item.Status == "pending"})
}
func (h *ProviderAdminHandler) ImportKreaCookie(c *gin.Context) {
var body struct {
Cookie string `json:"cookie"`
+1
View File
@@ -108,6 +108,7 @@ func New(cfg *config.Config, auth *service.AuthService, handlers Handlers) *gin.
authed.POST("/tokens/import-leonardo-cookie", handlers.ProviderAdmin.ImportLeonardoCookie)
authed.POST("/tokens/import-krea-cookie", handlers.ProviderAdmin.ImportKreaCookie)
authed.POST("/tokens/import-imagine-token", handlers.ProviderAdmin.ImportImagineToken)
authed.POST("/tokens/import-grok-token", handlers.ProviderAdmin.ImportGrokToken)
authed.POST("/tokens/delete-bulk", handlers.ProviderAdmin.TokenDeleteBulk)
authed.PATCH("/tokens/:pool/:id", handlers.ProviderAdmin.TokenUpdate)
authed.DELETE("/tokens/:pool/:id", handlers.ProviderAdmin.TokenDelete)
+438
View File
@@ -0,0 +1,438 @@
// Package grok implements the Grok (grok.com / xAI) provider client. Auth is the
// website "sso" session cookie (a JWT whose only claim is a session_id — no exp,
// no refresh: when the session dies upstream the account is simply dead, never
// renewed). grok.com gates requests with an x-statsig-id header; the web app's
// value is just a base64-encoded fake JS TypeError string, which the upstream
// accepts — so we spoof it the same way (no Cloudflare clearance needed). Uses
// tls-client so the JA3/JA4 fingerprint matches Chrome.
package grok
import (
"context"
"encoding/base64"
"encoding/json"
"errors"
"fmt"
"io"
"math"
"math/rand/v2"
"strconv"
"strings"
http "github.com/bogdanfinn/fhttp"
tlsclient "github.com/bogdanfinn/tls-client"
"github.com/bogdanfinn/tls-client/profiles"
"github.com/google/uuid"
)
const (
apiBase = "https://grok.com"
origin = "https://grok.com"
userAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/133.0.0.0 Safari/537.36"
// fullCredits is the weekly grant — UI shows "100 满额".
fullCredits = 100
)
var (
ErrAuth = errors.New("grok auth failed")
ErrQuotaExhausted = errors.New("grok quota exhausted")
ErrTemporaryUpstream = errors.New("grok upstream temporary error")
)
type Client struct {
proxy string
}
func NewClient(proxy string) *Client {
return &Client{proxy: strings.TrimSpace(proxy)}
}
func (c *Client) SetProxy(proxy string) {
c.proxy = strings.TrimSpace(proxy)
}
// IsGrokToken reports whether a JWT looks like a Grok website "sso" cookie: a
// payload whose ONLY claim is "session_id". That disambiguates it from a runway
// token (id + sso claims) or a chatgpt token (openai.com claims).
func IsGrokToken(token string) bool {
claims := decodeJWTPayload(token)
if len(claims) == 0 {
return false
}
if _, ok := claims["session_id"]; !ok {
return false
}
// Reject tokens that ALSO carry other-provider markers.
for k := range claims {
if k == "session_id" {
continue
}
if k == "sso" || k == "id" || strings.HasPrefix(k, "https://api.openai.com/") {
return false
}
}
return true
}
// SessionIDFromToken returns the sso session id (for dedup / display).
func SessionIDFromToken(token string) string {
return strings.TrimSpace(stringValue(decodeJWTPayload(token)["session_id"]))
}
// ExtractAccountInfo returns the free (no-network) account view. grok sso has no
// email/exp claim, so identity falls back to the session id.
func ExtractAccountInfo(token string) map[string]any {
sid := SessionIDFromToken(token)
return map[string]any{
"email": emptyStringNil(sid),
"session_id": emptyStringNil(sid),
"expires_at": nil,
}
}
// FetchCreditsBalance reads the account's live credit balance via the billing
// gRPC-web endpoint GetGrokCreditsConfig (empty request). The response carries
// the remaining credits (field 1, a float32) and the weekly reset timestamp
// (field 5). A 401/403 maps to ErrAuth (the session is dead). Returns the
// normalized map the TokenService quota plumbing expects.
func (c *Client) FetchCreditsBalance(ctx context.Context, token string) (map[string]any, error) {
token = strings.TrimSpace(strings.TrimPrefix(token, "Bearer "))
if token == "" {
return unknownBalance("empty token"), nil
}
client, err := c.newTLSClient()
if err != nil {
return nil, err
}
// gRPC-web empty message frame: 1-byte flag + 4-byte length (both zero).
body := []byte{0, 0, 0, 0, 0}
req, err := http.NewRequest(http.MethodPost, apiBase+"/grok_api_v2.GrokBuildBilling/GetGrokCreditsConfig", strings.NewReader(string(body)))
if err != nil {
return nil, err
}
req = req.WithContext(ctx)
c.applyHeaders(req, token, map[string]string{
"content-type": "application/grpc-web+proto",
"x-grpc-web": "1",
"accept": "application/grpc-web+proto",
})
resp, err := client.Do(req)
if err != nil {
return unknownBalance("network: " + err.Error()), nil
}
defer resp.Body.Close()
raw, err := io.ReadAll(resp.Body)
if err != nil {
return nil, err
}
if resp.StatusCode == 401 || resp.StatusCode == 403 {
return nil, ErrAuth
}
if resp.StatusCode != 200 {
return unknownBalance(fmt.Sprintf("http %d: %s", resp.StatusCode, clip(raw, 160))), nil
}
// GetGrokCreditsConfig field #1 is the credits USED this period (not remaining):
// an exhausted account reads 100, a fresh one reads ~0. Remaining = 100 - used.
used, reset, ok := parseCreditsConfig(raw)
if !ok {
return unknownBalance("unparsable credits config"), nil
}
if used < 0 {
used = 0
}
if used > fullCredits {
used = fullCredits
}
remaining := fullCredits - used
return map[string]any{
"remaining": remaining,
"used": used,
"total": fullCredits,
"reset_after": emptyStringNil(reset),
"unknown": false,
"error": nil,
}, nil
}
// FetchSession reads the account profile via GET /api/auth/session and returns
// (email, userID). A 401/403 means the sso session is dead → ErrAuth.
func (c *Client) FetchSession(ctx context.Context, token string) (email, userID string, err error) {
token = strings.TrimSpace(strings.TrimPrefix(token, "Bearer "))
if token == "" {
return "", "", ErrAuth
}
client, err := c.newTLSClient()
if err != nil {
return "", "", err
}
req, err := http.NewRequest(http.MethodGet, apiBase+"/api/auth/session", nil)
if err != nil {
return "", "", err
}
req = req.WithContext(ctx)
c.applyHeaders(req, token, nil)
resp, err := client.Do(req)
if err != nil {
return "", "", fmt.Errorf("%w: %v", ErrTemporaryUpstream, err)
}
defer resp.Body.Close()
raw, _ := io.ReadAll(resp.Body)
if resp.StatusCode == 401 || resp.StatusCode == 403 {
return "", "", ErrAuth
}
if resp.StatusCode != 200 {
return "", "", fmt.Errorf("%w: session http %d", ErrTemporaryUpstream, resp.StatusCode)
}
var body struct {
Session struct {
Email string `json:"email"`
UserID string `json:"userId"`
} `json:"session"`
}
if err := json.Unmarshal(raw, &body); err != nil {
return "", "", fmt.Errorf("%w: session non-json", ErrTemporaryUpstream)
}
return strings.TrimSpace(body.Session.Email), strings.TrimSpace(body.Session.UserID), nil
}
// statsigID mirrors grok2api's _statsig_id: base64 of a fake JS TypeError string.
// The upstream's anti-bot check accepts this spoofed value.
func statsigID() string {
const charset = "abcdefghijklmnopqrstuvwxyz0123456789"
b := make([]byte, 5)
for i := range b {
b[i] = charset[rand.IntN(len(charset))]
}
msg := fmt.Sprintf("x1:TypeError: Cannot read properties of null (reading 'children['%s']')", string(b))
return base64.StdEncoding.EncodeToString([]byte(msg))
}
// applyHeaders sets the browser-like header set + sso cookie + spoofed statsig id.
// extra overrides/adds per-request headers (e.g. content-type).
func (c *Client) applyHeaders(req *http.Request, token string, extra map[string]string) {
h := http.Header{
"accept": {"*/*"},
"accept-language": {"en-US,en;q=0.9"},
"content-type": {"application/json"},
"origin": {origin},
"referer": {origin + "/"},
"user-agent": {userAgent},
"x-statsig-id": {statsigID()},
"x-xai-request-id": {uuid.NewString()},
"sec-ch-ua": {`"Chromium";v="133", "Not(A:Brand";v="99"`},
"sec-ch-ua-mobile": {"?0"},
"sec-ch-ua-platform": {`"Windows"`},
"sec-fetch-dest": {"empty"},
"sec-fetch-mode": {"cors"},
"sec-fetch-site": {"same-origin"},
"cookie": {"sso=" + token + "; sso-rw=" + token},
}
for k, v := range extra {
h[k] = []string{v}
}
h[http.HeaderOrderKey] = []string{
"accept", "accept-language", "content-type", "origin", "referer",
"user-agent", "x-statsig-id", "x-xai-request-id", "x-grpc-web",
"sec-ch-ua", "sec-ch-ua-mobile", "sec-ch-ua-platform",
"sec-fetch-dest", "sec-fetch-mode", "sec-fetch-site", "cookie",
}
req.Header = h
}
func (c *Client) newTLSClient() (tlsclient.HttpClient, error) {
options := []tlsclient.HttpClientOption{
tlsclient.WithTimeoutSeconds(120),
tlsclient.WithClientProfile(profiles.Chrome_133),
tlsclient.WithRandomTLSExtensionOrder(),
}
if c.proxy != "" {
options = append(options, tlsclient.WithProxyUrl(c.proxy))
}
return tlsclient.NewHttpClient(tlsclient.NewNoopLogger(), options...)
}
// --- gRPC-web / protobuf decoding for GetGrokCreditsConfig ---
// parseCreditsConfig extracts (remaining credits, reset RFC3339-ish unix string)
// from the gRPC-web framed protobuf. Layout (reverse-engineered):
//
// frame: 1-byte flag + 4-byte big-endian length + payload
// payload: field 1 (message) {
// field 1: float32 -> remaining credits
// field 5: message { field 1: varint -> reset unix seconds }
// }
func parseCreditsConfig(buf []byte) (remaining int, resetUnix string, ok bool) {
for len(buf) >= 5 {
flag := buf[0]
ln := int(buf[1])<<24 | int(buf[2])<<16 | int(buf[3])<<8 | int(buf[4])
buf = buf[5:]
if ln > len(buf) {
break
}
payload := buf[:ln]
buf = buf[ln:]
if flag&0x80 != 0 { // trailer frame (grpc-status), skip
continue
}
// payload: expect field 1 (wire type 2) wrapping the config message.
fn, wt, val, rest, good := readField(payload)
if !good || fn != 1 || wt != 2 {
continue
}
rem, reset, found := scanConfigMessage(val)
_ = rest
if found {
return rem, reset, true
}
}
return 0, "", false
}
func scanConfigMessage(msg []byte) (remaining int, resetUnix string, ok bool) {
var remF float32
haveRem := false
for len(msg) > 0 {
fn, wt, val, rest, good := readField(msg)
if !good {
break
}
msg = rest
switch {
case fn == 1 && wt == 5: // float32 remaining credits
remF = float32FromLE(val)
haveRem = true
case fn == 5 && wt == 2: // reset timestamp message { #1 varint=seconds }
if sec, sok := firstVarint(val); sok {
resetUnix = strconv.FormatInt(sec, 10)
}
}
}
if haveRem {
return int(remF), resetUnix, true
}
return 0, resetUnix, false
}
// readField reads one protobuf field: returns (fieldNum, wireType, value, rest, ok).
// For wt 2 value is the length-delimited bytes; wt 5 the 4 LE bytes; wt 0 the
// raw varint bytes; wt 1 the 8 bytes.
func readField(b []byte) (fn int, wt int, val []byte, rest []byte, ok bool) {
tag, n := readVarint(b)
if n == 0 {
return 0, 0, nil, b, false
}
b = b[n:]
fn = int(tag >> 3)
wt = int(tag & 7)
switch wt {
case 0:
_, m := readVarint(b)
if m == 0 {
return 0, 0, nil, b, false
}
return fn, wt, b[:m], b[m:], true
case 1:
if len(b) < 8 {
return 0, 0, nil, b, false
}
return fn, wt, b[:8], b[8:], true
case 2:
ln, m := readVarint(b)
if m == 0 || int(ln) > len(b)-m {
return 0, 0, nil, b, false
}
return fn, wt, b[m : m+int(ln)], b[m+int(ln):], true
case 5:
if len(b) < 4 {
return 0, 0, nil, b, false
}
return fn, wt, b[:4], b[4:], true
default:
return 0, 0, nil, b, false
}
}
func firstVarint(b []byte) (int64, bool) {
fn, wt, val, _, ok := readField(b)
if !ok || fn != 1 || wt != 0 {
return 0, false
}
v, _ := readVarint(val)
return int64(v), true
}
func readVarint(b []byte) (uint64, int) {
var v uint64
var s uint
for i := 0; i < len(b); i++ {
v |= uint64(b[i]&0x7f) << s
if b[i]&0x80 == 0 {
return v, i + 1
}
s += 7
}
return 0, 0
}
func float32FromLE(b []byte) float32 {
if len(b) < 4 {
return 0
}
bits := uint32(b[0]) | uint32(b[1])<<8 | uint32(b[2])<<16 | uint32(b[3])<<24
return math.Float32frombits(bits)
}
// --- small helpers (mirror the other provider clients) ---
func decodeJWTPayload(token string) map[string]any {
parts := strings.Split(strings.TrimSpace(strings.TrimPrefix(token, "Bearer ")), ".")
if len(parts) < 2 {
return map[string]any{}
}
raw, err := base64.RawURLEncoding.DecodeString(parts[1])
if err != nil {
return map[string]any{}
}
var out map[string]any
if err := json.Unmarshal(raw, &out); err != nil {
return map[string]any{}
}
return out
}
func stringValue(v any) string {
switch x := v.(type) {
case string:
return x
case nil:
return ""
default:
b, _ := json.Marshal(x)
return strings.TrimSpace(string(b))
}
}
func emptyStringNil(v string) any {
v = strings.TrimSpace(v)
if v == "" {
return nil
}
return v
}
func unknownBalance(reason string) map[string]any {
return map[string]any{
"remaining": nil, "used": nil, "total": nil,
"unknown": true, "error": reason,
}
}
func clip(b []byte, n int) string {
s := strings.TrimSpace(string(b))
if len(s) > n {
return s[:n]
}
return s
}
+278
View File
@@ -0,0 +1,278 @@
package grok
import (
"context"
"encoding/base64"
"encoding/json"
"fmt"
"io"
"regexp"
"strings"
http "github.com/bogdanfinn/fhttp"
tlsclient "github.com/bogdanfinn/tls-client"
)
// assetBase is where generated media artifacts live (the stream returns a
// path like "users/<uid>/generated/<id>/generated_video.mp4").
const assetBase = "https://assets.grok.com/"
var videoURLRe = regexp.MustCompile(`"videoUrl":"([^"]+)"`)
// GenerateVideo runs grok's imagine video pipeline:
// 1. POST /rest/media/post/create -> a media post id (parentPostId)
// 2. POST /rest/app-chat/conversations/new with modelName "imagine-video-gen"
// and a videoGenModelConfig referencing that post id; the streaming response
// reports progress and, at completion, the artifact videoUrl.
//
// frames (optional, up to the model's max) enable image-to-video: each image is
// uploaded to grok and referenced as an imageReference. aspectRatio is passed
// through ("9:16" etc.); resolution is the tier ("720p"); seconds is the clip
// length (6 or 10). When downloadResult is false, returns nil bytes and the
// artifact URL in meta["video_url"]; otherwise downloads the mp4.
func (c *Client) GenerateVideo(ctx context.Context, token, prompt, aspectRatio, resolution string, seconds int, frames [][]byte, downloadResult bool) ([]byte, map[string]any, error) {
token = strings.TrimSpace(strings.TrimPrefix(token, "Bearer "))
if token == "" {
return nil, nil, ErrAuth
}
if strings.TrimSpace(prompt) == "" {
return nil, nil, fmt.Errorf("grok: prompt required")
}
if strings.TrimSpace(aspectRatio) == "" {
aspectRatio = "16:9"
}
if strings.TrimSpace(resolution) == "" {
resolution = "720p"
}
if seconds != 6 && seconds != 10 {
seconds = 10
}
client, err := c.newTLSClient()
if err != nil {
return nil, nil, err
}
// Image-to-video: upload each reference frame and collect its asset URL.
var imageRefs []string
for _, f := range frames {
if len(f) == 0 {
continue
}
url, upErr := c.uploadImage(ctx, client, token, f)
if upErr != nil {
return nil, nil, upErr
}
imageRefs = append(imageRefs, url)
}
postID, err := c.createPost(ctx, client, token, prompt)
if err != nil {
return nil, nil, err
}
videoCfg := map[string]any{
"parentPostId": postID,
"aspectRatio": aspectRatio,
"videoLength": seconds,
"resolutionName": resolution,
"isReferenceToVideo": len(imageRefs) > 0,
}
if len(imageRefs) > 0 {
videoCfg["imageReferences"] = imageRefs
}
payload := map[string]any{
"temporary": true,
"modelName": "imagine-video-gen",
"message": prompt + " --mode=custom",
"enableSideBySide": true,
"responseMetadata": map[string]any{
"modelConfigOverride": map[string]any{
"modelMap": map[string]any{"videoGenModelConfig": videoCfg},
},
},
}
body, err := c.postStream(ctx, client, token, "/rest/app-chat/conversations/new", payload)
if err != nil {
return nil, nil, err
}
// Out-of-credits surfaces as a stream error (HTTP is still 200).
if strings.Contains(body, "usagePoolExhausted") || strings.Contains(body, "media generation credits") {
return nil, nil, fmt.Errorf("%w: media generation credits exhausted", ErrQuotaExhausted)
}
// The artifact path appears as "videoUrl":"users/.../generated_video.mp4".
var artifact string
for _, m := range videoURLRe.FindAllStringSubmatch(body, -1) {
if v := strings.TrimSpace(m[1]); v != "" {
artifact = v // keep the last (progress=100) one
}
}
if artifact == "" {
return nil, nil, fmt.Errorf("%w: no video artifact in response: %s", ErrTemporaryUpstream, clip([]byte(body), 200))
}
fullURL := artifact
if !strings.HasPrefix(fullURL, "http") {
fullURL = assetBase + strings.TrimPrefix(artifact, "/")
}
meta := map[string]any{
"provider": "grok",
"post_id": postID,
"video_url": fullURL,
}
if !downloadResult {
return nil, meta, nil
}
data, err := c.download(ctx, client, token, fullURL)
if err != nil {
return nil, nil, err
}
return data, meta, nil
}
// uploadImage uploads one reference frame via /rest/app-chat/upload-file (JSON
// with base64 content) and returns its asset content URL for imageReferences.
func (c *Client) uploadImage(ctx context.Context, client tlsclient.HttpClient, token string, img []byte) (string, error) {
res, err := c.postJSON(ctx, client, token, "/rest/app-chat/upload-file", map[string]any{
"fileName": "ref.png",
"fileMimeType": "image/png",
"content": base64.StdEncoding.EncodeToString(img),
})
if err != nil {
return "", err
}
fileURI := strings.TrimSpace(stringValue(res["fileUri"]))
if fileURI == "" {
return "", fmt.Errorf("%w: upload missing fileUri", ErrTemporaryUpstream)
}
if strings.HasPrefix(fileURI, "http") {
return fileURI, nil
}
return assetBase + strings.TrimPrefix(fileURI, "/"), nil
}
// createPost registers a video media post and returns its id (parentPostId).
func (c *Client) createPost(ctx context.Context, client tlsclient.HttpClient, token, prompt string) (string, error) {
res, err := c.postJSON(ctx, client, token, "/rest/media/post/create", map[string]any{
"mediaType": "MEDIA_POST_TYPE_VIDEO",
"prompt": prompt,
})
if err != nil {
return "", err
}
post, _ := res["post"].(map[string]any)
id := strings.TrimSpace(stringValue(post["id"]))
if id == "" {
return "", fmt.Errorf("%w: media post missing id", ErrTemporaryUpstream)
}
return id, nil
}
// postJSON does an authed JSON POST and parses a single JSON object response.
func (c *Client) postJSON(ctx context.Context, client tlsclient.HttpClient, token, path string, body any) (map[string]any, error) {
raw, status, err := c.doPost(ctx, client, token, path, body)
if err != nil {
return nil, err
}
if e := mapStatus(path, status, raw); e != nil {
return nil, e
}
var out map[string]any
if len(raw) == 0 {
return map[string]any{}, nil
}
if err := json.Unmarshal(raw, &out); err != nil {
return nil, fmt.Errorf("%w: %s non-json: %s", ErrTemporaryUpstream, path, clip(raw, 120))
}
return out, nil
}
// postStream does an authed JSON POST and returns the full (streamed) text body.
func (c *Client) postStream(ctx context.Context, client tlsclient.HttpClient, token, path string, body any) (string, error) {
raw, status, err := c.doPost(ctx, client, token, path, body)
if err != nil {
return "", err
}
if e := mapStatus(path, status, raw); e != nil {
return "", e
}
return string(raw), nil
}
func (c *Client) doPost(ctx context.Context, client tlsclient.HttpClient, token, path string, body any) ([]byte, int, error) {
var reader io.Reader
if body != nil {
b, _ := json.Marshal(body)
reader = strings.NewReader(string(b))
}
req, err := http.NewRequest(http.MethodPost, apiBase+path, reader)
if err != nil {
return nil, 0, err
}
req = req.WithContext(ctx)
c.applyHeaders(req, token, nil)
resp, err := client.Do(req)
if err != nil {
return nil, 0, fmt.Errorf("%w: %v", ErrTemporaryUpstream, err)
}
defer resp.Body.Close()
raw, err := io.ReadAll(resp.Body)
if err != nil {
return nil, resp.StatusCode, err
}
return raw, resp.StatusCode, nil
}
func (c *Client) download(ctx context.Context, client tlsclient.HttpClient, token, url string) ([]byte, error) {
req, err := http.NewRequest(http.MethodGet, url, nil)
if err != nil {
return nil, err
}
req = req.WithContext(ctx)
req.Header = http.Header{
"user-agent": {userAgent},
"referer": {origin + "/"},
"cookie": {"sso=" + token + "; sso-rw=" + token},
}
resp, err := client.Do(req)
if err != nil {
return nil, fmt.Errorf("%w: %v", ErrTemporaryUpstream, err)
}
defer resp.Body.Close()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
return nil, fmt.Errorf("%w: download %d", ErrTemporaryUpstream, resp.StatusCode)
}
data, err := io.ReadAll(resp.Body)
if err != nil {
return nil, err
}
if len(data) == 0 {
return nil, fmt.Errorf("%w: empty video download", ErrTemporaryUpstream)
}
return data, nil
}
// mapStatus maps an HTTP status to the shared provider error sentinels.
func mapStatus(path string, status int, raw []byte) error {
switch {
case status == 200:
return nil
case status == 401 || status == 403:
return fmt.Errorf("%w: %s %d %s", ErrAuth, path, status, clip(raw, 160))
case status == 429:
return fmt.Errorf("%w: %s 429 %s", ErrQuotaExhausted, path, clip(raw, 160))
case status >= 500:
return fmt.Errorf("%w: %s %d %s", ErrTemporaryUpstream, path, status, clip(raw, 160))
default:
if isCreditError(string(raw)) {
return fmt.Errorf("%w: %s", ErrQuotaExhausted, clip(raw, 160))
}
return fmt.Errorf("grok: %s %d %s", path, status, clip(raw, 160))
}
}
func isCreditError(s string) bool {
s = strings.ToLower(s)
return strings.Contains(s, "usagepoolexhausted") || strings.Contains(s, "credit") || strings.Contains(s, "insufficient") || strings.Contains(s, "quota")
}
@@ -139,6 +139,8 @@ func (c *Client) FetchCreditsBalance(ctx context.Context, token string) (map[str
if err != nil {
return nil, err
}
// Per ops decision a rate-limit (403) is treated as a dead account too, same as
// a 401 — a throttled Runway token is considered done.
if resp.StatusCode == 401 || resp.StatusCode == 403 {
return nil, ErrAuth
}
+139
View File
@@ -0,0 +1,139 @@
package runway
import (
"bytes"
"context"
"errors"
"fmt"
"image"
"strings"
"time"
http "github.com/bogdanfinn/fhttp"
tlsclient "github.com/bogdanfinn/tls-client"
"github.com/google/uuid"
)
// GenerateImage runs the Runway "Nano Banana 2" (gemini_3_1_flash_image)
// text/image-to-image pipeline: upload each reference image (DATASET +
// DATASET_PREVIEW → dataset) to obtain its {assetId, url}, create a gemini image
// task and poll it to completion, then download the rendered PNG. teamID is the
// workspace id; if empty it's derived from the token. aspectRatio is passed
// through as-is (e.g. "16:9"); imageSize is the "1K"/"2K"/"4K" tier. refs may be
// empty (pure text-to-image).
func (c *Client) GenerateImage(ctx context.Context, token, teamID, prompt, aspectRatio, imageSize string, refs [][]byte) ([]byte, map[string]any, error) {
token = strings.TrimSpace(strings.TrimPrefix(token, "Bearer "))
if token == "" {
return nil, nil, ErrAuth
}
if teamID == "" {
teamID = TeamIDFromToken(token)
}
if teamID == "" {
return nil, nil, errors.New("runway: no team id")
}
if strings.TrimSpace(aspectRatio) == "" {
aspectRatio = "16:9"
}
if strings.TrimSpace(imageSize) == "" {
imageSize = "1K"
}
client, err := c.newTLSClient()
if err != nil {
return nil, nil, err
}
var refImages []map[string]any
for i, raw := range refs {
if len(raw) == 0 {
continue
}
filename := fmt.Sprintf("ref_%s_%d.png", time.Now().UTC().Format("20060102_150405"), i+1)
assetID, url, upErr := c.uploadReference(ctx, client, token, teamID, filename, raw)
if upErr != nil {
return nil, nil, upErr
}
refImages = append(refImages, map[string]any{
"tag": fmt.Sprintf("IMG_%d", i+1),
"assetId": assetID,
"url": url,
})
}
taskID, err := c.createImageTask(ctx, client, token, teamID, prompt, aspectRatio, imageSize, refImages)
if err != nil {
return nil, nil, err
}
artifactURL, err := c.pollTask(ctx, client, token, teamID, taskID)
if err != nil {
return nil, nil, err
}
data, err := c.download(ctx, client, artifactURL)
if err != nil {
return nil, nil, err
}
meta := map[string]any{
"provider": "runway",
"task_id": taskID,
"team_id": teamID,
"image_url": artifactURL,
}
return data, meta, nil
}
// uploadReference uploads one reference image through the dataset pipeline
// (DATASET_PREVIEW + DATASET uploads → /v1/datasets) and returns its asset id
// (= dataset id) and the cloudfront URL the task references.
func (c *Client) uploadReference(ctx context.Context, client tlsclient.HttpClient, token, teamID, filename string, data []byte) (string, string, error) {
cfg, _, err := image.DecodeConfig(bytes.NewReader(data))
if err != nil {
return "", "", errors.New("runway: failed to decode reference image")
}
previewUploadID, _, err := c.uploadFile(ctx, client, token, teamID, filename, "DATASET_PREVIEW", data)
if err != nil {
return "", "", err
}
// The DATASET upload's completed URL is exactly what the task references.
datasetUploadID, refURL, err := c.uploadFile(ctx, client, token, teamID, filename, "DATASET", data)
if err != nil {
return "", "", err
}
assetID, _, err := c.createDataset(ctx, client, token, teamID, filename, datasetUploadID, previewUploadID, cfg.Width, cfg.Height)
if err != nil {
return "", "", err
}
return assetID, refURL, nil
}
// createImageTask creates a gemini_3_1_flash_image task and returns its id.
func (c *Client) createImageTask(ctx context.Context, client tlsclient.HttpClient, token, teamID, prompt, aspectRatio, imageSize string, refImages []map[string]any) (string, error) {
opts := map[string]any{
"name": "Nano Banana 2 - " + prompt,
"text_prompt": prompt,
"aspect_ratio": aspectRatio,
"num_images": 1,
"image_size": imageSize,
"model": "gemini-3.1-flash-image-preview",
"exploreMode": false,
"creationSource": "tool-mode",
}
if len(refImages) > 0 {
opts["reference_images"] = refImages
}
res, err := c.apiJSON(ctx, client, token, teamID, http.MethodPost, "/v1/tasks", map[string]any{
"taskType": "gemini_3_1_flash_image",
"options": opts,
"asTeamId": jsonNumberOrString(teamID),
"sessionId": uuid.NewString(),
})
if err != nil {
return "", err
}
task, _ := res["task"].(map[string]any)
id := strings.TrimSpace(stringValue(task["id"]))
if id == "" {
return "", fmt.Errorf("%w: image task missing id", ErrTemporaryUpstream)
}
return id, nil
}
@@ -300,6 +300,7 @@ func (c *Client) apiJSON(ctx context.Context, client tlsclient.HttpClient, token
}
switch {
case resp.StatusCode == 401 || resp.StatusCode == 403:
// Rate-limit (403) is treated as a dead account too, same as a 401.
return nil, fmt.Errorf("%w: %s %d %s", ErrAuth, path, resp.StatusCode, clip(raw, 200))
case resp.StatusCode == 429:
return nil, fmt.Errorf("%w: %s 429 %s", ErrQuotaExhausted, path, clip(raw, 200))
+5
View File
@@ -342,6 +342,11 @@ func (s *AdminReadService) Providers(ctx context.Context) ([]map[string]any, err
}{
{Name: "chatgpt", Pool: "chatgpt", Type: "openai"},
{Name: "adobe", Pool: "adobe", Type: "adobe"},
{Name: "runway", Pool: "runway", Type: "runway"},
{Name: "leonardo", Pool: "leonardo", Type: "leonardo"},
{Name: "krea", Pool: "krea", Type: "krea"},
{Name: "imagine", Pool: "imagine", Type: "imagine"},
{Name: "grok", Pool: "grok", Type: "grok"},
}
out := make([]map[string]any, 0, len(providers))
for _, item := range providers {
+160 -10
View File
@@ -18,6 +18,7 @@ import (
"backend/internal/provider/imagine"
"backend/internal/provider/krea"
"backend/internal/provider/leonardo"
"backend/internal/provider/grok"
"backend/internal/provider/runway"
"backend/internal/repo"
@@ -32,6 +33,7 @@ var validTokenPools = map[string]string{
"leonardo": "leonardo",
"krea": "krea",
"imagine": "imagine",
"grok": "grok",
}
type TokenService struct {
@@ -45,6 +47,7 @@ type TokenService struct {
leonardo *leonardo.Client
krea *krea.Client
imagine *imagine.Client
grok *grok.Client
// sem caps concurrent background pending-probe goroutines (mirrors Python's
// 10-worker _quota_check_pool) so a big paste doesn't fire hundreds of
// simultaneous upstream requests.
@@ -54,7 +57,7 @@ type TokenService struct {
kreaActivating atomic.Bool
}
func NewTokenService(tokens *repo.TokenRepository, refresh *repo.RefreshProfileRepository, events *repo.EventRepository, settings *repo.SiteSettingRepository, adobeClient *adobe.Client, chatGPTClient *chatgpt.Client, runwayClient *runway.Client, leonardoClient *leonardo.Client, kreaClient *krea.Client, imagineClient *imagine.Client) *TokenService {
func NewTokenService(tokens *repo.TokenRepository, refresh *repo.RefreshProfileRepository, events *repo.EventRepository, settings *repo.SiteSettingRepository, adobeClient *adobe.Client, chatGPTClient *chatgpt.Client, runwayClient *runway.Client, leonardoClient *leonardo.Client, kreaClient *krea.Client, imagineClient *imagine.Client, grokClient *grok.Client) *TokenService {
return &TokenService{
tokens: tokens,
refresh: refresh,
@@ -66,6 +69,7 @@ func NewTokenService(tokens *repo.TokenRepository, refresh *repo.RefreshProfileR
leonardo: leonardoClient,
krea: kreaClient,
imagine: imagineClient,
grok: grokClient,
sem: make(chan struct{}, 10),
}
}
@@ -91,6 +95,9 @@ func (s *TokenService) applyProxy(ctx context.Context) {
if s.imagine != nil {
s.imagine.SetProxy(proxy)
}
if s.grok != nil {
s.grok.SetProxy(proxy)
}
}
// RefreshExpiringTokens proactively renews krea/imagine sessions ~10min before
@@ -840,6 +847,106 @@ func (s *TokenService) checkPendingRunway(tokenID, accessToken string) {
s.finishPending(ctx, "runway", tokenID, "active", false, quotaMeta)
}
// ImportGrokToken lands a Grok website "sso" cookie (a JWT carrying only a
// session_id) as a pending account and probes its credit balance off-thread.
// Identity is the session id (grok sso has no email/exp claim). No refresh: a
// dead session just dies (失效就失效).
func (s *TokenService) ImportGrokToken(ctx context.Context, ssoToken, tokenID string) (*model.TokenAccount, error) {
ssoToken = strings.TrimSpace(strings.TrimPrefix(ssoToken, "Bearer "))
ssoToken = strings.TrimPrefix(ssoToken, "sso=")
if ssoToken == "" {
return nil, errors.New("sso token required")
}
if !grok.IsGrokToken(ssoToken) {
return nil, errors.New("not a grok sso token")
}
sid := grok.SessionIDFromToken(ssoToken)
// Resolve the real account email up front (GET /api/auth/session) for dedup +
// display — the sso session_id rotates per login, so email is the stable id.
email := ""
if s.grok != nil {
s.applyProxy(ctx)
if e, _, ferr := s.grok.FetchSession(ctx, ssoToken); ferr == nil {
email = e
}
}
idKey := email
if idKey == "" {
idKey = sid
}
// Identity is (pool, email): reuse the row for this account, else mint.
if existing, _ := s.tokens.GetByPoolEmail(ctx, "grok", idKey); existing != nil {
tokenID = existing.ID
} else if idKey != "" || tokenID == "" {
tokenID = newTokenID("grok")
}
meta := datatypes.JSONMap{"pending_check": true}
if sid != "" {
meta["session_id"] = sid
}
item, err := s.createToken(ctx, "grok", tokenID, ssoToken, "pending", meta)
if err != nil {
if errors.Is(err, gorm.ErrDuplicatedKey) {
if item, err = s.tokens.Update(ctx, "grok", tokenID, map[string]any{
"value": ssoToken, "status": "pending", "meta": meta,
}); err != nil {
return nil, err
}
} else {
return nil, err
}
}
if idKey != "" {
if updated, uerr := s.tokens.Update(ctx, "grok", tokenID, map[string]any{"account_email": idKey}); uerr == nil {
item = updated
}
}
go s.checkPendingGrok(tokenID, ssoToken)
return item, nil
}
func (s *TokenService) checkPendingGrok(tokenID, ssoToken string) {
defer func() {
if r := recover(); r != nil {
log.Printf("token import: grok pending check panicked for %s: %v", tokenID, r)
}
}()
s.sem <- struct{}{}
defer func() { <-s.sem }()
ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
defer cancel()
if s.grok == nil {
s.finishPending(ctx, "grok", tokenID, "active", false, nil)
return
}
s.applyProxy(ctx)
data, err := s.grok.FetchCreditsBalance(ctx, ssoToken)
if err != nil {
if errors.Is(err, grok.ErrAuth) {
s.finishPending(ctx, "grok", tokenID, "disabled", true, nil)
return
}
s.finishPending(ctx, "grok", tokenID, "active", false, nil)
return
}
quotaMeta := map[string]any{}
if rem, ok := data["remaining"].(int); ok {
quotaMeta["cached_quota_remaining"] = rem
quotaMeta["cached_quota_at"] = int(time.Now().Unix())
}
if used, ok := data["used"].(int); ok {
quotaMeta["cached_quota_used"] = used
}
if total, ok := data["total"].(int); ok {
quotaMeta["cached_quota_total"] = total
}
if reset := strings.TrimSpace(stringValue(data["reset_after"])); reset != "" {
_, _ = s.tokens.Update(ctx, "grok", tokenID, map[string]any{"cached_quota_reset_after": reset})
}
s.finishPending(ctx, "grok", tokenID, "active", false, quotaMeta)
}
// finishPending writes the terminal status/dead flag and clears the pending_check
// marker (merging any cached quota) for a background import probe.
func (s *TokenService) finishPending(ctx context.Context, pool, id, status string, dead bool, quotaMeta map[string]any) {
@@ -1224,13 +1331,10 @@ func (s *TokenService) Quota(ctx context.Context, pool, id string) (map[string]a
meta := cloneJSONMap(item.Meta)
meta["cached_quota_at"] = int(time.Now().Unix())
if remaining, ok := data["remaining"].(int); ok {
// Refresh only updates the displayed balance number — it never flips
// status. Out-of-credits is judged at generation time (dead/401), so a
// refresh can't sink a runway account into a revivable "quota" state.
meta["cached_quota_remaining"] = remaining
// Refreshing in account management: an account below the credit floor
// is sunk into "限额" (quota) so it stops being scheduled. Mark-down
// only — recovery is a separate (unwritten) path.
if remaining < runwayMinCredits && item.Status == "active" {
patch["status"] = "quota"
}
}
if used, ok := data["used"].(int); ok {
meta["cached_quota_used"] = used
@@ -1256,12 +1360,58 @@ func (s *TokenService) Quota(ctx context.Context, pool, id string) (map[string]a
"error": data["error"],
}, nil
}
if poolToType(item.Pool) == "grok" && s.grok != nil {
data, err := s.grok.FetchCreditsBalance(ctx, item.Value)
if err != nil {
if errors.Is(err, grok.ErrAuth) {
_, _ = s.tokens.Update(ctx, item.Pool, item.ID, map[string]any{
"status": "disabled",
"dead": true,
"fails": gorm.Expr("fails + 1"),
})
}
return nil, err
}
patch := map[string]any{}
meta := cloneJSONMap(item.Meta)
meta["cached_quota_at"] = int(time.Now().Unix())
if remaining, ok := data["remaining"].(int); ok {
// Refresh only updates the displayed credit number; never flips status.
// Out-of-credits is judged at generation time (dead/401, no renewal).
meta["cached_quota_remaining"] = remaining
}
if used, ok := data["used"].(int); ok {
meta["cached_quota_used"] = used
}
if total, ok := data["total"].(int); ok {
meta["cached_quota_total"] = total
}
patch["meta"] = meta
if reset := strings.TrimSpace(stringValue(data["reset_after"])); reset != "" {
patch["cached_quota_reset_after"] = reset
item.CachedQuotaResetAfter = reset
}
if updated, updateErr := s.tokens.Update(ctx, item.Pool, item.ID, patch); updateErr == nil {
item = updated
}
return map[string]any{
"supported": true,
"remaining": data["remaining"],
"used": data["used"],
"total": data["total"],
"reset_after": emptyToNil(item.CachedQuotaResetAfter),
"quota_cached_at": meta["cached_quota_at"],
"unchanged": false,
"unknown": boolValueWithDefault(data["unknown"], false),
"error": data["error"],
}, nil
}
remaining, hasRemaining := jsonMapInt(item.Meta, "cached_quota_remaining")
quotaAt, _ := jsonMapInt(item.Meta, "cached_quota_at")
typeLabel := poolToType(item.Pool)
return map[string]any{
"supported": typeLabel == "openai" || typeLabel == "adobe" || typeLabel == "runway",
"remaining": valueOrNil((typeLabel == "openai" || typeLabel == "runway") && hasRemaining, remaining),
"supported": typeLabel == "openai" || typeLabel == "adobe" || typeLabel == "runway" || typeLabel == "grok",
"remaining": valueOrNil((typeLabel == "openai" || typeLabel == "runway" || typeLabel == "grok") && hasRemaining, remaining),
"total": nil,
"reset_after": emptyToNil(item.CachedQuotaResetAfter),
"quota_cached_at": valueOrNil(quotaAt != 0, quotaAt),
@@ -1372,7 +1522,7 @@ func accountRow(item model.TokenAccount, inFlight int64) map[string]any {
if item.Meta != nil {
teamID = strings.TrimSpace(stringValue(item.Meta["team_id"]))
}
hasQuota := typeLabel == "openai" || typeLabel == "adobe" || typeLabel == "runway" || typeLabel == "leonardo" || typeLabel == "krea" || typeLabel == "imagine"
hasQuota := typeLabel == "openai" || typeLabel == "adobe" || typeLabel == "runway" || typeLabel == "leonardo" || typeLabel == "krea" || typeLabel == "imagine" || typeLabel == "grok"
return map[string]any{
"id": item.ID,
"pool": item.Pool,
+331 -130
View File
@@ -19,6 +19,7 @@ import (
"backend/internal/model"
"backend/internal/provider/adobe"
"backend/internal/provider/chatgpt"
"backend/internal/provider/grok"
"backend/internal/provider/imagine"
"backend/internal/provider/krea"
"backend/internal/provider/leonardo"
@@ -68,6 +69,7 @@ type V1Service struct {
leonardo *leonardo.Client
krea *krea.Client
imagine *imagine.Client
grok *grok.Client
store *storage.Client
// refresh re-mints an Adobe access token from its cookie when a request hits a
// 401 mid-flight (set via SetRefresh — wired after construction to avoid an
@@ -97,17 +99,37 @@ type V1Service struct {
// account isn't already running a generation; release frees it when done.
type accountGate struct{ m sync.Map } // accountID -> struct{} held while busy
func (g *accountGate) tryAcquire(id string) bool {
// tryAcquireN wins if the account has fewer than max in-flight jobs, atomically
// bumping its counter. max=1 is the default 1-job-per-account policy; some
// providers (grok) allow more.
func (g *accountGate) tryAcquireN(id string, max int) bool {
if id == "" {
return true
}
_, loaded := g.m.LoadOrStore(id, struct{}{})
return !loaded
if max < 1 {
max = 1
}
v, _ := g.m.LoadOrStore(id, new(int64))
cnt := v.(*int64)
for {
cur := atomic.LoadInt64(cnt)
if cur >= int64(max) {
return false
}
if atomic.CompareAndSwapInt64(cnt, cur, cur+1) {
return true
}
}
}
func (g *accountGate) tryAcquire(id string) bool { return g.tryAcquireN(id, 1) }
func (g *accountGate) release(id string) {
if id != "" {
g.m.Delete(id)
if id == "" {
return
}
if v, ok := g.m.Load(id); ok {
atomic.AddInt64(v.(*int64), -1)
}
}
@@ -173,7 +195,7 @@ type V1VideoRequest struct {
BaseURL string
}
func NewV1Service(cfg *config.Config, models *repo.ModelRepository, users *repo.UserRepository, events *repo.EventRepository, tokens *repo.TokenRepository, settings *repo.SiteSettingRepository, adobeClient *adobe.Client, chatGPTClient *chatgpt.Client, runwayClient *runway.Client, leonardoClient *leonardo.Client, kreaClient *krea.Client, imagineClient *imagine.Client, store *storage.Client) *V1Service {
func NewV1Service(cfg *config.Config, models *repo.ModelRepository, users *repo.UserRepository, events *repo.EventRepository, tokens *repo.TokenRepository, settings *repo.SiteSettingRepository, adobeClient *adobe.Client, chatGPTClient *chatgpt.Client, runwayClient *runway.Client, leonardoClient *leonardo.Client, kreaClient *krea.Client, imagineClient *imagine.Client, grokClient *grok.Client, store *storage.Client) *V1Service {
return &V1Service{
cfg: cfg,
models: models,
@@ -187,6 +209,7 @@ func NewV1Service(cfg *config.Config, models *repo.ModelRepository, users *repo.
leonardo: leonardoClient,
krea: kreaClient,
imagine: imagineClient,
grok: grokClient,
store: store,
inflight: &InflightRegistry{},
}
@@ -415,6 +438,23 @@ func (s *V1Service) prepareImageExecution(ctx context.Context, principal *APIPri
}
}
imageBytes = b
case "runway":
b, execErr := s.generateRunwayImage(genCtx, eventID, modelItem, in, aspectRatio, resolution)
if execErr != nil {
_ = s.refundIfNeeded(ctx, principal, eventID, price)
_ = s.events.UpdateStatus(ctx, eventID, "failed", execErr.Error(), 0)
switch {
case errors.Is(execErr, runway.ErrAuth):
return nil, ErrProviderAuth
case errors.Is(execErr, runway.ErrQuotaExhausted):
return nil, ErrProviderQuota
case errors.Is(execErr, runway.ErrTemporaryUpstream):
return nil, ErrProviderTemporary
default:
return nil, fmt.Errorf("%w: %v", ErrProviderExecution, execErr)
}
}
imageBytes = b
default:
_ = s.refundIfNeeded(ctx, principal, eventID, price)
_ = s.events.UpdateStatus(ctx, eventID, "failed", "provider not implemented", 0)
@@ -517,6 +557,8 @@ func (s *V1Service) prepareVideoExecution(ctx context.Context, principal *APIPri
videoBytes, _, execErr = s.generateAdobeVideo(genCtx, eventID, modelItem, in, aspectRatio, resolution, parseDurationSeconds(duration), true)
case "runway":
videoBytes, _, execErr = s.generateRunwayVideo(genCtx, eventID, modelItem, in, aspectRatio, parseDurationSeconds(duration), true)
case "grok":
videoBytes, _, execErr = s.generateGrokVideo(genCtx, eventID, modelItem, in, aspectRatio, resolution, parseDurationSeconds(duration), true)
default:
_ = s.refundIfNeeded(ctx, principal, eventID, price)
_ = s.events.UpdateStatus(ctx, eventID, "failed", "provider not implemented", 0)
@@ -528,11 +570,11 @@ func (s *V1Service) prepareVideoExecution(ctx context.Context, principal *APIPri
switch {
case errors.Is(execErr, ErrNoProviderAccount):
return nil, ErrNoProviderAccount
case errors.Is(execErr, adobe.ErrAuth), errors.Is(execErr, runway.ErrAuth):
case errors.Is(execErr, adobe.ErrAuth), errors.Is(execErr, runway.ErrAuth), errors.Is(execErr, grok.ErrAuth):
return nil, ErrProviderAuth
case errors.Is(execErr, adobe.ErrQuotaExhausted), errors.Is(execErr, runway.ErrQuotaExhausted):
case errors.Is(execErr, adobe.ErrQuotaExhausted), errors.Is(execErr, runway.ErrQuotaExhausted), errors.Is(execErr, grok.ErrQuotaExhausted):
return nil, ErrProviderQuota
case errors.Is(execErr, adobe.ErrTemporaryUpstream), errors.Is(execErr, runway.ErrTemporaryUpstream):
case errors.Is(execErr, adobe.ErrTemporaryUpstream), errors.Is(execErr, runway.ErrTemporaryUpstream), errors.Is(execErr, grok.ErrTemporaryUpstream):
return nil, ErrProviderTemporary
default:
return nil, fmt.Errorf("%w: %v", ErrProviderExecution, execErr)
@@ -621,6 +663,8 @@ func (s *V1Service) runVideoJob(ctx context.Context, principal *APIPrincipal, in
_, videoURL, execErr = s.generateAdobeVideo(genCtx, eventID, modelItem, in, aspectRatio, resolution, parseDurationSeconds(duration), false)
case "runway":
_, videoURL, execErr = s.generateRunwayVideo(genCtx, eventID, modelItem, in, aspectRatio, parseDurationSeconds(duration), false)
case "grok":
_, videoURL, execErr = s.generateGrokVideo(genCtx, eventID, modelItem, in, aspectRatio, resolution, parseDurationSeconds(duration), false)
default:
_ = s.refundIfNeeded(ctx, principal, eventID, price)
_ = s.events.UpdateStatus(ctx, eventID, "failed", "provider not implemented", 0)
@@ -1106,6 +1150,16 @@ func (s *V1Service) finishUnimplementedEvent(ctx context.Context, eventID string
// same account; account-level errors (auth/quota) skip straight to the next.
const maxSameAccountAttempts = 3
// grokConcurrencyPerAccount is how many simultaneous generations one grok account
// may run (grok tolerates 10, unlike the 1-per-account default elsewhere).
const grokConcurrencyPerAccount = 10
// maxTempDeadAccounts caps how many accounts the "temporary error = dead account"
// policy (tempAsDead, used by adobe) is allowed to mark dead + fail over before
// giving up, so an upstream-wide blip ("system under load") can't nuke the whole
// pool. After this many accounts fail this way, the request fails.
const maxTempDeadAccounts = 3
// runPoolWithFailover drives a generation across a round-robin-ordered account
// list with per-error-class behavior, so a bad request never burns the whole
// pool while genuinely limited accounts still fail over:
@@ -1115,9 +1169,13 @@ const maxSameAccountAttempts = 3
// - 认证失效 auth → refresh the token from its cookie and retry ONCE with the
// fresh token; if it still auth-fails (or there's nothing to refresh, e.g.
// chatgpt's JWT IS the credential), mark the account and fail over.
// - 上游临时 temporary → retry the SAME account up to maxSameAccountAttempts
// times (not counted); if still failing, STOP (no fan-out — an upstream-wide
// blip fails identically everywhere).
// - 上游临时 temporary → behavior depends on tempAsDead:
// • tempAsDead=false (default): retry the SAME account up to
// maxSameAccountAttempts times (not counted); if still failing, STOP
// (no fan-out — an upstream-wide blip fails identically everywhere).
// • tempAsDead=true (adobe): treat the temporary error as a DEAD account —
// mark it like a 401 and fail over to the next account, capped at
// maxTempDeadAccounts accounts so a pool-wide blip can't kill everything.
// - 参数错 / request-level (anything else) → return immediately, no retry, no
// account penalty (the account isn't at fault).
//
@@ -1130,9 +1188,11 @@ func (s *V1Service) runPoolWithFailover(ctx context.Context, eventID, pool strin
attempt func(token model.TokenAccount) ([]byte, error),
classify func(error) (isAuth, isQuota, isTemporary bool),
refreshOnAuth func(tokenID string) (model.TokenAccount, bool),
tempAsDead bool,
) ([]byte, error) {
var lastErr error
busy := 0
tempDeadCount := 0
for _, token := range active {
// 1 concurrent job per account: skip any account already generating.
if !s.gate.tryAcquire(token.ID) {
@@ -1140,14 +1200,23 @@ func (s *V1Service) runPoolWithFailover(ctx context.Context, eventID, pool strin
continue
}
// release via defer so a panic in tryAccount can't leak the 1-job slot.
data, err, failover := func() ([]byte, error, bool) {
data, err, failover, tempDead := func() ([]byte, error, bool, bool) {
defer s.gate.release(token.ID)
return s.tryAccount(ctx, eventID, pool, token, kind, attempt, classify, refreshOnAuth)
return s.tryAccount(ctx, eventID, pool, token, kind, attempt, classify, refreshOnAuth, tempAsDead)
}()
if err == nil {
return data, nil
}
lastErr = err
if tempDead {
// temp-as-dead policy: this account was marked dead for a temporary
// upstream error. Cap how many accounts that can burn before we stop,
// so an upstream-wide blip doesn't wipe the whole pool.
tempDeadCount++
if tempDeadCount >= maxTempDeadAccounts {
return nil, lastErr
}
}
if failover {
continue
}
@@ -1173,7 +1242,8 @@ func (s *V1Service) tryAccount(ctx context.Context, eventID, pool string, token
attempt func(token model.TokenAccount) ([]byte, error),
classify func(error) (isAuth, isQuota, isTemporary bool),
refreshOnAuth func(tokenID string) (model.TokenAccount, bool),
) ([]byte, error, bool) {
tempAsDead bool,
) ([]byte, error, bool, bool) {
_ = s.events.SetAccount(ctx, eventID, token.ID)
_ = s.tokens.TouchLastUsed(ctx, token.ID)
authRefreshed := false
@@ -1186,12 +1256,12 @@ func (s *V1Service) tryAccount(ctx context.Context, eventID, pool string, token
"success_total": gorm.Expr("success_total + 1"),
"fails": 0,
})
return data, nil, false
return data, nil, false, false
}
isAuth, isQuota, isTemp := classify(err)
if isQuota {
s.markTokenFailure(ctx, pool, token, kind, false, true)
return nil, err, true
return nil, err, true, false
}
if isAuth {
// Refresh from cookie and retry ONCE; otherwise the credential is dead.
@@ -1203,24 +1273,32 @@ func (s *V1Service) tryAccount(ctx context.Context, eventID, pool string, token
}
}
s.markTokenFailure(ctx, pool, token, kind, true, false)
return nil, err, true
return nil, err, true, false
}
if isTemp {
if tempAsDead {
// Ops policy (adobe): a temporary upstream error ("system under
// load" etc.) means this account is effectively dead — mark it
// like a 401 and fail over to the next account. The pool driver
// caps how many accounts this is allowed to burn.
s.markTokenFailure(ctx, pool, token, kind, true, false)
return nil, err, true, true
}
tempAttempts++
if tempAttempts < maxSameAccountAttempts {
// Short linear backoff (1s, 2s) so an overloaded/rate-limited upstream
// (e.g. adobe "system under load") gets a moment to recover before the
// same-account retry, instead of hammering it instantly.
// gets a moment to recover before the same-account retry, instead of
// hammering it instantly.
select {
case <-time.After(time.Duration(tempAttempts) * time.Second):
case <-ctx.Done():
return nil, err, false
return nil, err, false, false
}
continue
}
return nil, err, false // exhausted; no fan-out
return nil, err, false, false // exhausted; no fan-out
}
return nil, err, false // 参数错 / request-level
return nil, err, false, false // 参数错 / request-level
}
}
@@ -1262,8 +1340,10 @@ func (s *V1Service) generateAdobeImage(ctx context.Context, eventID string, mode
return nil, err
}
// Round-robin order; same-account retry on transient errors, fail over to the
// next account on auth/quota (see runPoolWithFailover).
// Round-robin order. Adobe uses tempAsDead=true: a temporary upstream error
// ("system under load") marks the account dead (like a 401) and fails over to
// the next account, capped at maxTempDeadAccounts; auth/quota also fail over
// (see runPoolWithFailover).
return s.runPoolWithFailover(ctx, eventID, "adobe", active, "image", func(token model.TokenAccount) ([]byte, error) {
var blobIDs []string
for _, ref := range refs {
@@ -1277,7 +1357,7 @@ func (s *V1Service) generateAdobeImage(ctx context.Context, eventID string, mode
return data, genErr
}, adobeErrClass, func(id string) (model.TokenAccount, bool) {
return s.refreshAdobeToken(ctx, id)
})
}, true)
}
func (s *V1Service) generateAdobeVideo(ctx context.Context, eventID string, modelItem *model.ModelConfig, in V1VideoRequest, aspectRatio, resolution string, durationSeconds int, downloadResult bool) ([]byte, string, error) {
@@ -1323,8 +1403,9 @@ func (s *V1Service) generateAdobeVideo(ctx context.Context, eventID string, mode
referenceMode := defaultString(strings.TrimSpace(modelItem.ReferenceMode), "frame")
// Round-robin order; same-account retry on transient errors, fail over to the
// next account on auth/quota (see runPoolWithFailover). videoURL is captured
// from the successful attempt's meta (the upstream presigned URL).
// next account on auth/quota; temporary upstream errors mark the account dead
// and fail over too (tempAsDead, capped at maxTempDeadAccounts). videoURL is
// captured from the successful attempt's meta (the upstream presigned URL).
var videoURL string
data, err := s.runPoolWithFailover(ctx, eventID, "adobe", active, "video", func(token model.TokenAccount) ([]byte, error) {
var blobIDs []string
@@ -1342,21 +1423,10 @@ func (s *V1Service) generateAdobeVideo(ctx context.Context, eventID string, mode
return bytes, genErr
}, adobeErrClass, func(id string) (model.TokenAccount, bool) {
return s.refreshAdobeToken(ctx, id)
})
}, true)
return data, videoURL, err
}
// runwayMinCredits gates account scheduling: a Runway account with fewer than
// this many credits remaining is treated as quota-limited and skipped, so we
// never dial upstream with an account that's about to run dry. Flat threshold
// (not per-duration) by request.
const runwayMinCredits = 50
// runwayCreditsPerSecond is Gen-4 Turbo's price (5 credits/sec → 5s=25, 10s=50),
// used to pre-reserve the exact render cost so concurrent picks of one account
// can't over-commit it.
const runwayCreditsPerSecond = 5
// leonardoMinCredits is the per-generation token cost (one Leonardo image = 30
// tokens). An account with fewer is treated as 限额 and skipped — it can't afford
// a generation. Daily renewal (tokenRenewalDate) drives auto-recovery.
@@ -1391,10 +1461,10 @@ func (s *V1Service) generateRunwayVideo(ctx context.Context, eventID string, mod
if item.Status != "active" || item.Dead || strings.TrimSpace(item.Value) == "" {
continue
}
// Skip accounts under the credit floor (treated as quota-limited). Only
// skip when we KNOW the balance is too low — an unknown balance gets the
// benefit of the doubt (upstream will reject if it's truly empty).
if rem, ok := jsonMapInt(item.Meta, "cached_quota_remaining"); ok && rem < runwayMinCredits {
// No pre-deduct (same policy as the image flow): skip only accounts we KNOW
// are out of credits (cached remaining <= 0) — those are treated as dead.
// Unknown balance gets the benefit of the doubt.
if rem, ok := jsonMapInt(item.Meta, "cached_quota_remaining"); ok && rem <= 0 {
continue
}
active = append(active, item)
@@ -1404,9 +1474,6 @@ func (s *V1Service) generateRunwayVideo(ctx context.Context, eventID string, mod
}
s.rotateRoundRobin("runway", active)
// Pre-reserve the exact render cost (5 credits/sec) so two concurrent renders
// can't over-commit the same account.
cost := durationSeconds * runwayCreditsPerSecond
var lastErr error
var videoURL string
busy := 0
@@ -1417,22 +1484,10 @@ func (s *V1Service) generateRunwayVideo(ctx context.Context, eventID string, mod
continue
}
var data []byte
ok := func() bool {
done, failover := func() (bool, bool) {
defer s.gate.release(token.ID)
_ = s.events.SetAccount(ctx, eventID, token.ID)
_ = s.tokens.TouchLastUsed(ctx, token.ID)
// Atomic pre-deduction (row-locked). Known-insufficient → sink to 限额
// and fail over to the next account.
allowed, deducted, rerr := s.tokens.ReserveQuota(ctx, "runway", token.ID, cost)
if rerr != nil {
lastErr = fmt.Errorf("%w: reserve: %v", runway.ErrTemporaryUpstream, rerr)
return false
}
if !allowed {
s.markTokenFailure(ctx, "runway", token, "video", false, true)
lastErr = runway.ErrQuotaExhausted
return false
}
teamID := ""
if token.Meta != nil {
teamID = strings.TrimSpace(stringValue(token.Meta["team_id"]))
@@ -1444,26 +1499,31 @@ func (s *V1Service) generateRunwayVideo(ctx context.Context, eventID string, mod
"success_total": gorm.Expr("success_total + 1"),
"fails": 0,
})
// Re-fetch the REAL balance after the render and sink to 限额 if below
// the floor. Best-effort — never fail an already-successful render.
s.reconcileRunwayCredits(ctx, token.ID, token.Value)
data = d
videoURL = strings.TrimSpace(stringValue(meta["video_url"]))
return true
}
// Release the hold so a failed render doesn't burn credits.
if deducted {
_ = s.tokens.RefundQuota(ctx, "runway", token.ID, cost)
return true, false
}
lastErr = genErr
s.markTokenFailure(ctx, "runway", token, "video",
errors.Is(genErr, runway.ErrAuth),
errors.Is(genErr, runway.ErrQuotaExhausted))
return false
switch {
case errors.Is(genErr, runway.ErrAuth), errors.Is(genErr, runway.ErrQuotaExhausted):
// 额度没了 / token 失效 → 当 401 判死(status=disabled, dead),换号。
s.markTokenFailure(ctx, "runway", token, "video", true, false)
return false, true
case errors.Is(genErr, runway.ErrTemporaryUpstream):
// 上游临时错误 → 直接换下一个号。
return false, true
default:
// 参数级错误(如 prompt 未过审)→ 直接失败,不换号。
return false, false
}
}()
if ok {
if done {
return data, videoURL, nil
}
if failover {
continue
}
return nil, "", lastErr
}
if lastErr == nil {
if busy > 0 {
@@ -1474,41 +1534,209 @@ func (s *V1Service) generateRunwayVideo(ctx context.Context, eventID string, mod
return nil, "", lastErr
}
// reconcileRunwayCredits re-fetches an account's authoritative credit balance
// (after a render) and writes it back, flipping the account to "限额" (quota)
// when it's below the floor. Concurrency-safe: every write stores a freshly
// observed real balance — no local arithmetic that could lose updates under
// concurrent renders. Best-effort; marks down only (recovery is unwritten).
func (s *V1Service) reconcileRunwayCredits(ctx context.Context, tokenID, tokenValue string) {
// generateGrokVideo runs grok's imagine video pipeline across the grok pool.
// Mirrors the runway policy: no pre-deduct, skip accounts known out of credits
// (cached remaining <= 0), and treat an out-of-credits / auth failure as a dead
// account (the grok sso can't be renewed — 失效就失效). Text-to-video only for
// now (grok reference-image upload isn't wired yet).
func (s *V1Service) generateGrokVideo(ctx context.Context, eventID string, modelItem *model.ModelConfig, in V1VideoRequest, aspectRatio, resolution string, durationSeconds int, downloadResult bool) ([]byte, string, error) {
if s.grok == nil {
return nil, "", errors.New("grok client not configured")
}
if s.settings != nil {
if proxy, err := s.settings.GetValue(ctx, "proxy.url"); err == nil {
s.grok.SetProxy(proxy)
}
}
// Optional reference frames (image-to-video), up to the model's max.
frames, err := decodeReferenceImages(in.ReferenceImages, max(1, modelItem.MaxReferenceImages))
if err != nil {
return nil, "", err
}
items, err := s.tokens.ListByPool(ctx, "grok")
if err != nil {
return nil, "", err
}
var active []model.TokenAccount
for _, item := range items {
if item.Status != "active" || item.Dead || strings.TrimSpace(item.Value) == "" {
continue
}
if rem, ok := jsonMapInt(item.Meta, "cached_quota_remaining"); ok && rem <= 0 {
continue
}
active = append(active, item)
}
if len(active) == 0 {
return nil, "", ErrNoProviderAccount
}
s.rotateRoundRobin("grok", active)
res := strings.TrimSpace(resolution)
if res == "" {
res = "720p"
}
var lastErr error
var videoURL string
busy := 0
for _, token := range active {
// grok allows 10 concurrent jobs per account (unlike the 1-per-account
// default of the other pools).
if !s.gate.tryAcquireN(token.ID, grokConcurrencyPerAccount) {
busy++
continue
}
var data []byte
done, failover := func() (bool, bool) {
defer s.gate.release(token.ID)
_ = s.events.SetAccount(ctx, eventID, token.ID)
_ = s.tokens.TouchLastUsed(ctx, token.ID)
d, meta, genErr := s.grok.GenerateVideo(ctx, token.Value, in.Prompt, aspectRatio, res, durationSeconds, frames, downloadResult)
if genErr == nil {
_, _ = s.tokens.Update(ctx, "grok", token.ID, map[string]any{
"last_used_at": time.Now(),
"success_total": gorm.Expr("success_total + 1"),
"fails": 0,
})
data = d
videoURL = strings.TrimSpace(stringValue(meta["video_url"]))
return true, false
}
lastErr = genErr
switch {
case errors.Is(genErr, grok.ErrAuth), errors.Is(genErr, grok.ErrQuotaExhausted):
// 失效 / 额度没了 → 当 401 判死(不续期),换号。
s.markTokenFailure(ctx, "grok", token, "video", true, false)
return false, true
case errors.Is(genErr, grok.ErrTemporaryUpstream):
return false, true
default:
return false, false
}
}()
if done {
return data, videoURL, nil
}
if failover {
continue
}
return nil, "", lastErr
}
if lastErr == nil {
if busy > 0 {
return nil, "", ErrConcurrencyFull
}
lastErr = ErrProviderExecution
}
return nil, "", lastErr
}
// generateRunwayImage runs the Runway "Nano Banana 2" (gemini_3_1_flash_image)
// image pipeline across the runway pool. Unlike the video path it does NOT
// pre-deduct credits: it simply round-robins the pool and generates. Per ops
// decision an out-of-credits account is treated like a dead 401 — marked
// dead (status=disabled) and skipped — because Runway credits don't refill
// daily, so a "quota" mark (which the maintenance loop would revive) is wrong.
// Reference images (up to the model's max) are uploaded per attempt.
func (s *V1Service) generateRunwayImage(ctx context.Context, eventID string, modelItem *model.ModelConfig, in V1ImageRequest, aspectRatio, resolution string) ([]byte, error) {
if s.runway == nil {
return
return nil, errors.New("runway client not configured")
}
data, err := s.runway.FetchCreditsBalance(ctx, tokenValue)
if s.settings != nil {
if proxy, err := s.settings.GetValue(ctx, "proxy.url"); err == nil {
s.runway.SetProxy(proxy)
}
}
refs, err := decodeReferenceImages(in.ReferenceImages, max(1, modelItem.MaxReferenceImages))
if err != nil {
return
return nil, err
}
rem, ok := data["remaining"].(int)
if !ok {
return
}
item, err := s.tokens.Get(ctx, "runway", tokenID)
items, err := s.tokens.ListByPool(ctx, "runway")
if err != nil {
return
return nil, err
}
meta := cloneJSONMap(item.Meta)
meta["cached_quota_remaining"] = rem
meta["cached_quota_at"] = int(time.Now().Unix())
if used, ok := data["used"].(int); ok {
meta["cached_quota_used"] = used
var active []model.TokenAccount
for _, item := range items {
if item.Status != "active" || item.Dead || strings.TrimSpace(item.Value) == "" {
continue
}
// No pre-deduct: skip only accounts we KNOW are out of credits
// (cached remaining <= 0); they're treated as dead. Unknown balance gets
// the benefit of the doubt — upstream rejects if it's truly empty.
if rem, ok := jsonMapInt(item.Meta, "cached_quota_remaining"); ok && rem <= 0 {
continue
}
active = append(active, item)
}
if total, ok := data["total"].(int); ok {
meta["cached_quota_total"] = total
if len(active) == 0 {
return nil, ErrNoProviderAccount
}
patch := map[string]any{"meta": meta}
if rem < runwayMinCredits && item.Status == "active" {
patch["status"] = "quota"
s.rotateRoundRobin("runway", active)
imageSize := strings.TrimSpace(resolution)
if imageSize == "" {
imageSize = "1K"
}
_, _ = s.tokens.Update(ctx, "runway", tokenID, patch)
var lastErr error
busy := 0
for _, token := range active {
// 1 concurrent job per account: skip any account already generating.
if !s.gate.tryAcquire(token.ID) {
busy++
continue
}
var data []byte
done, failover := func() (bool, bool) {
defer s.gate.release(token.ID)
_ = s.events.SetAccount(ctx, eventID, token.ID)
_ = s.tokens.TouchLastUsed(ctx, token.ID)
teamID := ""
if token.Meta != nil {
teamID = strings.TrimSpace(stringValue(token.Meta["team_id"]))
}
d, _, genErr := s.runway.GenerateImage(ctx, token.Value, teamID, in.Prompt, aspectRatio, imageSize, refs)
if genErr == nil {
_, _ = s.tokens.Update(ctx, "runway", token.ID, map[string]any{
"last_used_at": time.Now(),
"success_total": gorm.Expr("success_total + 1"),
"fails": 0,
})
data = d
return true, false
}
lastErr = genErr
switch {
case errors.Is(genErr, runway.ErrAuth), errors.Is(genErr, runway.ErrQuotaExhausted):
// 额度没了 / token 失效 → 当 401 判死(status=disabled, dead),换号。
s.markTokenFailure(ctx, "runway", token, "image", true, false)
return false, true
case errors.Is(genErr, runway.ErrTemporaryUpstream):
// 上游临时错误 → 直接换下一个号。
return false, true
default:
// 参数级错误(如 prompt 未过审)→ 直接失败,不换号。
return false, false
}
}()
if done {
return data, nil
}
if failover {
continue
}
return nil, lastErr
}
if lastErr == nil {
if busy > 0 {
return nil, ErrConcurrencyFull
}
lastErr = ErrProviderExecution
}
return nil, lastErr
}
// reconcileChatGPTQuota re-reads OpenAI's image_gen remaining right after a
@@ -1590,7 +1818,7 @@ func (s *V1Service) generateChatGPTImage(ctx context.Context, eventID string, mo
return data, genErr
}, func(e error) (bool, bool, bool) {
return errors.Is(e, chatgpt.ErrAuth), errors.Is(e, chatgpt.ErrQuotaExhausted), errors.Is(e, chatgpt.ErrTemporaryUpstream)
}, nil) // chatgpt token IS the credential — no cookie to refresh
}, nil, false) // chatgpt token IS the credential — no cookie to refresh
}
// leonardoResetAfter returns when a Leonardo account's daily free tokens renew.
@@ -1717,7 +1945,7 @@ func (s *V1Service) generateLeonardoImage(ctx context.Context, eventID string, m
return data, nil
}, func(e error) (bool, bool, bool) {
return errors.Is(e, leonardo.ErrAuth), errors.Is(e, leonardo.ErrQuotaExhausted), errors.Is(e, leonardo.ErrTemporaryUpstream)
}, nil)
}, nil, false)
}
// reconcileLeonardoCredits re-fetches an account's real token balance after a
@@ -1847,7 +2075,7 @@ func (s *V1Service) generateKreaImage(ctx context.Context, eventID string, model
return data, genErr
}, func(e error) (bool, bool, bool) {
return errors.Is(e, krea.ErrAuth), errors.Is(e, krea.ErrQuotaExhausted), errors.Is(e, krea.ErrTemporaryUpstream)
}, nil)
}, nil, false)
}
// imagineRefreshAndPersist ensures the account's Imagine credential has a valid
@@ -1916,37 +2144,10 @@ func (s *V1Service) generateImagineImage(ctx context.Context, eventID string, mo
if genErr != nil {
return nil, genErr
}
// Success → re-sync the displayed balance (best-effort).
s.reconcileImagineCredits(ctx, token.ID, cred)
return data, nil
}, func(e error) (bool, bool, bool) {
return errors.Is(e, imagine.ErrAuth), errors.Is(e, imagine.ErrQuotaExhausted), errors.Is(e, imagine.ErrTemporaryUpstream)
}, nil)
}
// reconcileImagineCredits re-fetches the account's real balance after a render
// and writes it back (best-effort; never fails a done render). Imagine credits
// don't daily-reset, so there's no reset marker to advance.
func (s *V1Service) reconcileImagineCredits(ctx context.Context, tokenID, cred string) {
if s.imagine == nil {
return
}
data, err := s.imagine.FetchCreditsBalance(ctx, cred)
if err != nil {
return
}
rem, ok := data["remaining"].(int)
if !ok {
return
}
item, err := s.tokens.Get(ctx, "imagine", tokenID)
if err != nil {
return
}
meta := cloneJSONMap(item.Meta)
meta["cached_quota_remaining"] = rem
meta["cached_quota_at"] = int(time.Now().Unix())
_, _ = s.tokens.Update(ctx, "imagine", tokenID, map[string]any{"meta": meta})
}, nil, false)
}
func (s *V1Service) refundIfNeeded(ctx context.Context, principal *APIPrincipal, eventID string, price float64) error {