diff --git a/backend/cmd/marklabel/main.go b/backend/cmd/marklabel/main.go deleted file mode 100644 index bb30b1c..0000000 --- a/backend/cmd/marklabel/main.go +++ /dev/null @@ -1,45 +0,0 @@ -package main - -import ( - "fmt" - "os" - - "gorm.io/driver/postgres" - "gorm.io/gorm" -) - -type row struct { - ID string - Pool string - Status string - AccountEmail string - ImageLimited bool - VideoLimited bool -} - -func main() { - dsn := os.Getenv("POSTGRES_DSN") - if dsn == "" { - fmt.Println("POSTGRES_DSN env is required, e.g. host=127.0.0.1 user=postgres password=... dbname=vivid_ai port=5432 sslmode=disable TimeZone=Asia/Shanghai") - os.Exit(1) - } - db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{}) - if err != nil { - fmt.Println("open err:", err) - os.Exit(1) - } - - var pick row - db.Raw(`SELECT id, pool, status, account_email, image_limited, video_limited FROM token_accounts WHERE pool='adobe' ORDER BY id LIMIT 1`).Scan(&pick) - fmt.Printf("picked: id=%s email=%s status=%s image_limited=%v video_limited=%v\n", pick.ID, pick.AccountEmail, pick.Status, pick.ImageLimited, pick.VideoLimited) - - if err := db.Exec(`UPDATE token_accounts SET video_limited=true, updated_at=now() WHERE id=?`, pick.ID).Error; err != nil { - fmt.Println("update err:", err) - os.Exit(1) - } - fmt.Println("-> set video_limited=true") - - var after row - db.Raw(`SELECT id, pool, status, account_email, image_limited, video_limited FROM token_accounts WHERE id=?`, pick.ID).Scan(&after) - fmt.Printf("after: id=%s status=%s image_limited=%v video_limited=%v\n", after.ID, after.Status, after.ImageLimited, after.VideoLimited) -} diff --git a/backend/internal/bootstrap/app.go b/backend/internal/bootstrap/app.go index b01d5e2..2b0cd77 100644 --- a/backend/internal/bootstrap/app.go +++ b/backend/internal/bootstrap/app.go @@ -5,25 +5,16 @@ import ( "fmt" "log" "os" - "strconv" "time" "backend/internal/config" "backend/internal/http/handler" "backend/internal/http/router" "backend/internal/model" - "backend/internal/provider/adobe" - "backend/internal/provider/chatgpt" - "backend/internal/provider/custom" - "backend/internal/provider/creativefabrica" - "backend/internal/provider/grok" - "backend/internal/provider/imagine" - "backend/internal/provider/krea" - "backend/internal/provider/leonardo" - "backend/internal/provider/runway" "backend/internal/repo" "backend/internal/service" "backend/internal/storage" + "backend/internal/upstream" "github.com/gin-gonic/gin" "github.com/redis/go-redis/v9" @@ -51,9 +42,9 @@ func NewApp(ctx context.Context) (*App, error) { return nil, fmt.Errorf("create generated root %s: %w", cfg.GeneratedRoot, err) } - // TranslateError: 把驱动层错误(如 Postgres 23505 唯一冲突)翻译成 gorm.ErrDuplicatedKey, - // 否则各 import-*（krea/adobe/leonardo/runway）里的 errors.Is(err, gorm.ErrDuplicatedKey) - // 兜底命不中,重复导入会直接抛原始错误 → 400,而不是按预期 Update 已有行。 + // TranslateError: 把驱动层错误(如 Postgres 23505 唯一冲突)翻译成 + // gorm.ErrDuplicatedKey,否则 errors.Is(err, gorm.ErrDuplicatedKey) 兜底命不中, + // 重复导入会直接抛原始错误 → 400,而不是按预期 Update 已有行。 db, err := gorm.Open(postgres.Open(cfg.PostgresDSN), &gorm.Config{TranslateError: true}) if err != nil { return nil, fmt.Errorf("open postgres: %w", err) @@ -97,8 +88,7 @@ func NewApp(ctx context.Context) (*App, error) { eventRepo := repo.NewEventRepository(db) cdkRepo := repo.NewCDKRepository(db) apiKeyRepo := repo.NewAPIKeyRepository(db) - tokenRepo := repo.NewTokenRepository(db) - refreshRepo := repo.NewRefreshProfileRepository(db) + channelRepo := repo.NewChannelRepository(db) cgroupRepo := repo.NewConcurrencyGroupRepository(db) // Seed the "默认并发" group (cap 10) and bind any ungrouped users to it. if err := cgroupRepo.EnsureDefault(ctx); err != nil { @@ -117,32 +107,21 @@ func NewApp(ctx context.Context) (*App, error) { authSvc := service.NewAuthService(userRepo, siteRepo, sessionSvc, emailCodeSvc, smtpSvc, cgroupRepo) appSettingsSvc := service.NewAppSettingsService(siteRepo, eventRepo, smtpSvc, rustfsClient) imageAccessSvc := service.NewImageAccessService(cfg.GeneratedRoot, showcaseRepo, authSvc) - adobeClient := adobe.NewClient("projectx_webapp", "") - chatGPTClient := chatgpt.NewClient("") - runwayClient := runway.NewClient("") - leonardoClient := leonardo.NewClient("") - kreaClient := krea.NewClient("") - imagineClient := imagine.NewClient("") - grokClient := grok.NewClient("") - // Keep grok's x-statsig-id anti-bot recipe current by reading grok's own - // headless-browser signer output. Event-driven: seed from the persisted - // recipe, capture once at startup, then re-capture only on an anti-bot 403 - // (a reship made the recipe stale). No polling. - startGrokStatsigRefresh(siteRepo) - customClient := custom.NewClient() - cfClient := creativefabrica.NewClient("") - v1Svc := service.NewV1Service(cfg, modelRepo, userRepo, eventRepo, tokenRepo, siteRepo, cgroupRepo, concSvc, adobeClient, chatGPTClient, runwayClient, leonardoClient, kreaClient, imagineClient, grokClient, customClient, cfClient, rustfsClient) + upClient := upstream.NewClient() + channelSecret := cfg.ChannelSecret + if channelSecret == "" { + // Stable fallback so a restarted process can still decrypt existing + // rows when CHANNEL_SECRET isn't set. + channelSecret = cfg.PostgresDSN + } + channelSvc := service.NewChannelService(channelRepo, modelRepo, upClient, channelSecret) + v1Svc := service.NewV1Service(cfg, modelRepo, userRepo, eventRepo, channelSvc, siteRepo, cgroupRepo, concSvc, rustfsClient) siteSvc := service.NewSiteService(siteRepo, cfg.AppTitle) showcaseSvc := service.NewShowcaseService(showcaseRepo) - adminReadSvc := service.NewAdminReadService(cfg, userRepo, modelRepo, eventRepo, siteRepo, tokenRepo, cdkRepo, rustfsClient, showcaseRepo) - adminWriteSvc := service.NewAdminWriteService(userRepo, showcaseRepo, modelRepo, eventRepo, apiKeyRepo, tokenRepo, orderRepo) + adminReadSvc := service.NewAdminReadService(cfg, userRepo, modelRepo, eventRepo, siteRepo, channelRepo, cdkRepo, rustfsClient, showcaseRepo) + adminWriteSvc := service.NewAdminWriteService(userRepo, showcaseRepo, modelRepo, eventRepo, apiKeyRepo, channelRepo, orderRepo) cdkSvc := service.NewCDKService(cdkRepo, userRepo, siteRepo, orderRepo) apiKeySvc := service.NewAPIKeyService(apiKeyRepo) - tokenSvc := service.NewTokenService(tokenRepo, refreshRepo, eventRepo, siteRepo, adobeClient, chatGPTClient, runwayClient, leonardoClient, kreaClient, imagineClient, grokClient, cfClient) - 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. - v1Svc.SetRefresh(refreshSvc) bannedWordRepo := repo.NewBannedWordRepository(db) v1Svc.SetBannedWords(bannedWordRepo) userGenSvc := service.NewUserGenerationService(v1Svc, eventRepo, userRepo, modelRepo) @@ -150,30 +129,30 @@ func NewApp(ctx context.Context) (*App, error) { canvasSvc := service.NewCanvasService(canvasRepo) engine := router.New(cfg, authSvc, router.Handlers{ - Health: handler.NewHealthHandler(), - Images: handler.NewImageHandler(cfg, imageAccessSvc, rustfsClient), - V1: handler.NewV1Handler(v1Svc), - Site: handler.NewSiteHandler(siteSvc), - Showcase: handler.NewShowcaseHandler(showcaseSvc), - Auth: handler.NewAuthHandler(cfg, authSvc, rateLimitSvc), - SiteSettings: handler.NewSiteSettingsHandler(siteSvc), - AppSettings: handler.NewAppSettingsHandler(appSettingsSvc), - AdminRead: handler.NewAdminReadHandler(adminReadSvc), - AdminWrite: handler.NewAdminWriteHandler(adminWriteSvc), - CDK: handler.NewCDKHandler(cdkSvc), - UserTools: handler.NewUserToolsHandler(apiKeySvc, cdkSvc), - UserGen: handler.NewUserGenerationHandler(userGenSvc, adminReadSvc), - ProviderAdmin: handler.NewProviderAdminHandler(tokenSvc, refreshSvc), - ConcGroups: handler.NewConcurrencyGroupHandler(cgroupSvc), - Announcement: handler.NewAnnouncementHandler(announcementSvc), - Payment: handler.NewPaymentHandler(paymentSvc), - BannedWords: handler.NewBannedWordsHandler(bannedWordRepo), - Canvas: handler.NewCanvasHandler(canvasSvc), + Health: handler.NewHealthHandler(), + Images: handler.NewImageHandler(cfg, imageAccessSvc, rustfsClient), + V1: handler.NewV1Handler(v1Svc), + Site: handler.NewSiteHandler(siteSvc), + Showcase: handler.NewShowcaseHandler(showcaseSvc), + Auth: handler.NewAuthHandler(cfg, authSvc, rateLimitSvc), + SiteSettings: handler.NewSiteSettingsHandler(siteSvc), + AppSettings: handler.NewAppSettingsHandler(appSettingsSvc), + AdminRead: handler.NewAdminReadHandler(adminReadSvc), + AdminWrite: handler.NewAdminWriteHandler(adminWriteSvc), + CDK: handler.NewCDKHandler(cdkSvc), + UserTools: handler.NewUserToolsHandler(apiKeySvc, cdkSvc), + UserGen: handler.NewUserGenerationHandler(userGenSvc, adminReadSvc), + ChannelAdmin: handler.NewChannelAdminHandler(channelSvc), + ConcGroups: handler.NewConcurrencyGroupHandler(cgroupSvc), + Announcement: handler.NewAnnouncementHandler(announcementSvc), + Payment: handler.NewPaymentHandler(paymentSvc), + BannedWords: handler.NewBannedWordsHandler(bannedWordRepo), + Canvas: handler.NewCanvasHandler(canvasSvc), }) - // Background self-healing sweep (quota recovery, cookie refresh, stale-pending - // cleanup, log retention) — the Go equivalent of the Python daemon thread. - maintenanceSvc := service.NewMaintenanceService(tokenRepo, tokenSvc, eventRepo, userRepo, refreshSvc, siteRepo, rustfsClient, v1Svc.Inflight(), showcaseRepo, orderRepo) + // Background self-healing sweep (order expiry, stale-pending cleanup+refund, + // log/media retention). + maintenanceSvc := service.NewMaintenanceService(eventRepo, userRepo, siteRepo, rustfsClient, v1Svc.Inflight(), showcaseRepo, orderRepo) loopCtx, loopCancel := context.WithCancel(context.Background()) go maintenanceSvc.Run(loopCtx) @@ -186,33 +165,6 @@ func NewApp(ctx context.Context) (*App, error) { }, nil } -// startGrokStatsigRefresh wires grok's headless x-statsig-id refresher to the -// site-settings store (persisted across restarts) with an app-lifetime context. -func startGrokStatsigRefresh(siteRepo *repo.SiteSettingRepository) { - const kHeader, kSuffix, kTrailer = "grok.statsig.header", "grok.statsig.suffix", "grok.statsig.trailer" - grok.StartStatsigAutoRefresh(context.Background(), 0, - func(ctx context.Context) (string, string, int, bool) { - h, _ := siteRepo.GetValue(ctx, kHeader) - s, _ := siteRepo.GetValue(ctx, kSuffix) - tv, _ := siteRepo.GetValue(ctx, kTrailer) - if h == "" || s == "" { - return "", "", 0, false - } - t, _ := strconv.Atoi(tv) - return h, s, t, true - }, - func(ctx context.Context, h, s string, t int) { - if err := siteRepo.UpsertValues(ctx, map[string]string{ - kHeader: h, - kSuffix: s, - kTrailer: strconv.Itoa(t), - }); err != nil { - log.Printf("grok statsig: persist recipe failed: %v", err) - } - }, - ) -} - func (a *App) Close() error { if a.maintenanceCancel != nil { a.maintenanceCancel() diff --git a/backend/internal/config/config.go b/backend/internal/config/config.go index 4ef2ef3..c408359 100644 --- a/backend/internal/config/config.go +++ b/backend/internal/config/config.go @@ -26,6 +26,9 @@ type Config struct { RustFSBucket string RustFSAccessKey string RustFSSecretKey string + // ChannelSecret encrypts channel upstream API keys at rest (AES-GCM). + // Optional — when empty a stable key is derived from the Postgres DSN. + ChannelSecret string } func Load() (*Config, error) { @@ -60,6 +63,7 @@ func Load() (*Config, error) { RustFSBucket: envString("RUSTFS_BUCKET", ""), RustFSAccessKey: envString("RUSTFS_ACCESS_KEY", ""), RustFSSecretKey: envString("RUSTFS_SECRET_KEY", ""), + ChannelSecret: envString("CHANNEL_SECRET", ""), } return cfg, nil diff --git a/backend/internal/http/handler/channel_admin.go b/backend/internal/http/handler/channel_admin.go new file mode 100644 index 0000000..5dc8614 --- /dev/null +++ b/backend/internal/http/handler/channel_admin.go @@ -0,0 +1,122 @@ +package handler + +import ( + "errors" + "net/http" + "strings" + + "backend/internal/service" + "backend/internal/upstream" + + "github.com/gin-gonic/gin" + "gorm.io/gorm" +) + +// ChannelAdminHandler exposes channel CRUD + upstream model import under +// /admin/api/channels*. Replaces the old token/account admin surface. +type ChannelAdminHandler struct { + channels *service.ChannelService +} + +func NewChannelAdminHandler(channels *service.ChannelService) *ChannelAdminHandler { + return &ChannelAdminHandler{channels: channels} +} + +func (h *ChannelAdminHandler) List(c *gin.Context) { + data, err := h.channels.List(c.Request.Context()) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load channels"}) + return + } + c.JSON(http.StatusOK, gin.H{"data": data}) +} + +func (h *ChannelAdminHandler) Save(c *gin.Context) { + var body service.SaveChannelDto + if err := c.ShouldBindJSON(&body); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"}) + return + } + body.BaseURL = strings.TrimSpace(body.BaseURL) + ch, err := h.channels.Save(c.Request.Context(), body) + if err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + c.JSON(http.StatusNotFound, gin.H{"detail": "channel not found"}) + return + } + c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"ok": true, "id": ch.ID}) +} + +func (h *ChannelAdminHandler) Delete(c *gin.Context) { + var body struct { + ID string `json:"id"` + } + if err := c.ShouldBindJSON(&body); err != nil || strings.TrimSpace(body.ID) == "" { + c.JSON(http.StatusBadRequest, gin.H{"detail": "channel id required"}) + return + } + if err := h.channels.Delete(c.Request.Context(), body.ID); err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + c.JSON(http.StatusNotFound, gin.H{"detail": "channel not found"}) + return + } + c.JSON(http.StatusInternalServerError, gin.H{"detail": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"ok": true}) +} + +// FetchModels proxies GET /v1/models on a channel (stored creds or ad-hoc +// base_url+api_key for a not-yet-saved channel form). +func (h *ChannelAdminHandler) FetchModels(c *gin.Context) { + var body service.FetchModelsDto + if err := c.ShouldBindJSON(&body); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"}) + return + } + items, err := h.channels.FetchModels(c.Request.Context(), body) + if err != nil { + writeUpstreamErr(c, err) + return + } + if items == nil { + items = []upstream.RemoteModel{} + } + c.JSON(http.StatusOK, gin.H{"data": items}) +} + +// ImportModels binds the ticked upstream models to the channel (creating local +// model rows on first import). +func (h *ChannelAdminHandler) ImportModels(c *gin.Context) { + var body struct { + ID string `json:"id"` + Models []service.ChannelModelSpec `json:"models"` + } + if err := c.ShouldBindJSON(&body); err != nil || strings.TrimSpace(body.ID) == "" { + c.JSON(http.StatusBadRequest, gin.H{"detail": "channel id required"}) + return + } + if err := h.channels.ImportModels(c.Request.Context(), body.ID, body.Models); err != nil { + writeUpstreamErr(c, err) + return + } + c.JSON(http.StatusOK, gin.H{"ok": true}) +} + +// writeUpstreamErr translates upstream-protocol errors into sensible admin +// responses (auth → 400 with the upstream's message, temporary → 502). +func writeUpstreamErr(c *gin.Context, err error) { + switch { + case errors.Is(err, upstream.ErrAuth): + c.JSON(http.StatusBadRequest, gin.H{"detail": "上游鉴权失败：" + err.Error()}) + case errors.Is(err, upstream.ErrInvalid): + c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()}) + case errors.Is(err, gorm.ErrRecordNotFound): + c.JSON(http.StatusNotFound, gin.H{"detail": "channel not found"}) + default: + c.JSON(http.StatusBadGateway, gin.H{"detail": "上游请求失败：" + err.Error()}) + } +} diff --git a/backend/internal/http/handler/provider_admin.go b/backend/internal/http/handler/provider_admin.go deleted file mode 100644 index 346836c..0000000 --- a/backend/internal/http/handler/provider_admin.go +++ /dev/null @@ -1,544 +0,0 @@ -package handler - -import ( - "errors" - "net/http" - "strings" - - "backend/internal/service" - - "github.com/gin-gonic/gin" - "gorm.io/gorm" -) - -type ProviderAdminHandler struct { - tokens *service.TokenService - refresh *service.RefreshProfileService -} - -func NewProviderAdminHandler(tokens *service.TokenService, refresh *service.RefreshProfileService) *ProviderAdminHandler { - return &ProviderAdminHandler{ - tokens: tokens, - refresh: refresh, - } -} - -func (h *ProviderAdminHandler) TokensList(c *gin.Context) { - data, err := h.tokens.List(c.Request.Context()) - if err != nil { - c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load tokens"}) - return - } - c.JSON(http.StatusOK, gin.H{"data": data}) -} - -func (h *ProviderAdminHandler) TokensCreate(c *gin.Context) { - var body struct { - Pool string `json:"pool"` - Value string `json:"value"` - ID string `json:"id"` - } - if err := c.ShouldBindJSON(&body); err != nil { - c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"}) - return - } - item, err := h.tokens.Add(c.Request.Context(), body.Pool, body.Value, body.ID) - if err != nil { - c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()}) - return - } - c.JSON(http.StatusOK, gin.H{"ok": true, "id": item.ID}) -} - -func (h *ProviderAdminHandler) ImportChatGPTToken(c *gin.Context) { - var body struct { - AccessToken string `json:"access_token"` - Value string `json:"value"` - 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 - if token == "" { - token = body.Value - } - name := body.Name - if name == "" { - name = body.ID - } - item, err := h.tokens.ImportChatGPTToken(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) ImportRunwayToken(c *gin.Context) { - var body struct { - AccessToken string `json:"access_token"` - Value string `json:"value"` - 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 - if token == "" { - token = body.Value - } - name := body.Name - if name == "" { - name = body.ID - } - item, err := h.tokens.ImportRunwayToken(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) 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) ImportCustomAccount(c *gin.Context) { - var body struct { - BaseURL string `json:"base_url"` - URL string `json:"url"` - Key string `json:"key"` - APIKey string `json:"api_key"` - Models string `json:"models"` - Name string `json:"name"` - Protocol string `json:"protocol"` // openai (default) | vivid - Weight int `json:"weight"` - Concurrency int `json:"concurrency"` - ID string `json:"id"` - } - if err := c.ShouldBindJSON(&body); err != nil { - c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"}) - return - } - baseURL := body.BaseURL - if baseURL == "" { - baseURL = body.URL - } - key := body.Key - if key == "" { - key = body.APIKey - } - item, err := h.tokens.ImportCustomAccount(c.Request.Context(), baseURL, key, body.Models, body.Name, body.Protocol, body.Weight, body.Concurrency, body.ID) - 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}) -} - -func (h *ProviderAdminHandler) ImportKreaCookie(c *gin.Context) { - var body struct { - Cookie string `json:"cookie"` - Value string `json:"value"` - 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 - } - cookie := body.Cookie - if cookie == "" { - cookie = body.Value - } - name := body.Name - if name == "" { - name = body.ID - } - item, err := h.tokens.ImportKreaCookie(c.Request.Context(), cookie, 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) ImportImagineToken(c *gin.Context) { - var body struct { - Cookie string `json:"cookie"` - Value string `json:"value"` - 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 - } - cred := body.Cookie - if cred == "" { - cred = body.Value - } - name := body.Name - if name == "" { - name = body.ID - } - item, err := h.tokens.ImportImagineToken(c.Request.Context(), cred, 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) ImportLeonardoCookie(c *gin.Context) { - var body struct { - Cookie string `json:"cookie"` - Value string `json:"value"` - 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 - } - cookie := body.Cookie - if cookie == "" { - cookie = body.Value - } - name := body.Name - if name == "" { - name = body.ID - } - item, err := h.tokens.ImportLeonardoCookie(c.Request.Context(), cookie, 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) ImportAdobeCookie(c *gin.Context) { - var body struct { - Cookie string `json:"cookie"` - Value string `json:"value"` - 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 - } - cookie := body.Cookie - if cookie == "" { - cookie = body.Value - } - name := body.Name - if name == "" { - name = body.ID - } - item, profile, err := h.tokens.ImportAdobeCookie(c.Request.Context(), cookie, name) - if err != nil { - c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()}) - return - } - c.JSON(http.StatusOK, gin.H{ - "ok": true, - "profile_id": profile.ID, - "id": item.ID, - "status": item.Status, - "pending": item.Status == "pending", - }) -} - -func (h *ProviderAdminHandler) ImportCreativeFabricaCookie(c *gin.Context) { - var body struct { - Cookie string `json:"cookie"` - Value string `json:"value"` - 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 - } - cookie := body.Cookie - if cookie == "" { - cookie = body.Value - } - name := body.Name - if name == "" { - name = body.ID - } - item, err := h.tokens.ImportCreativeFabricaCookie(c.Request.Context(), cookie, 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) TokenUpdate(c *gin.Context) { - var body map[string]any - if err := c.ShouldBindJSON(&body); err != nil { - c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"}) - return - } - item, err := h.tokens.Update(c.Request.Context(), c.Param("pool"), c.Param("id"), body) - if err != nil { - c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()}) - return - } - c.JSON(http.StatusOK, gin.H{"ok": true, "data": item}) -} - -func (h *ProviderAdminHandler) TokenDelete(c *gin.Context) { - if err := h.tokens.Delete(c.Request.Context(), c.Param("pool"), c.Param("id")); err != nil { - if errors.Is(err, service.ErrNotFound) { - c.JSON(http.StatusNotFound, gin.H{"detail": "token not found"}) - return - } - c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()}) - return - } - c.JSON(http.StatusOK, gin.H{"ok": true}) -} - -// TokenDeleteBulk removes multiple accounts in one call (account multi-select). -func (h *ProviderAdminHandler) TokenDeleteBulk(c *gin.Context) { - var body struct { - IDs []string `json:"ids"` - } - if err := c.ShouldBindJSON(&body); err != nil { - c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"}) - return - } - if len(body.IDs) == 0 { - c.JSON(http.StatusBadRequest, gin.H{"detail": "未选择任何账号"}) - return - } - n, err := h.tokens.DeleteBulk(c.Request.Context(), body.IDs) - if err != nil { - c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()}) - return - } - c.JSON(http.StatusOK, gin.H{"ok": true, "deleted": n}) -} - -func (h *ProviderAdminHandler) AccountsList(c *gin.Context) { - data, err := h.tokens.Accounts(c.Request.Context()) - if err != nil { - c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load accounts"}) - return - } - - // KPI stats are computed over the FULL set (每个类型的 成功/失败/限额), independent - // of the current filter/page — mirrors the old client-side `stats` computed. - stats := accountsStats(data) - - // Server-side filtering (类型 / 状态 / 搜索 邮箱·ID·类型 / dead) so pagination is - // correct across pages. ?dead=1 returns only 异常(已失效) accounts — used by - // 「删除异常账号」to collect every dead id regardless of the current page. - typeFilter := strings.TrimSpace(c.Query("type")) - statusFilter := strings.TrimSpace(c.Query("status")) - planFilter := strings.TrimSpace(c.Query("plan")) - deadOnly := c.Query("dead") == "1" - q := strings.ToLower(strings.TrimSpace(c.Query("q"))) - filtered := make([]map[string]any, 0, len(data)) - for _, row := range data { - if deadOnly && !rowBool(row, "dead") { - continue - } - if typeFilter != "" && rowStr(row, "type") != typeFilter { - continue - } - if statusFilter != "" && rowStr(row, "status") != statusFilter { - continue - } - // plan filter: vip=非free, free=free, sub=子号, master=vip且非子号 - if planFilter != "" { - plan := strings.ToLower(rowStr(row, "plan")) - isSub := rowBool(row, "sub_account") - isFree := plan == "" || plan == "free" - switch planFilter { - case "vip": - if isFree { - continue - } - case "free": - if !isFree { - continue - } - case "sub": - if !isSub { - continue - } - case "master": - if isFree || isSub { - continue - } - } - } - if q != "" { - email := strings.ToLower(rowStr(row, "email")) - id := strings.ToLower(rowStr(row, "id")) - typ := strings.ToLower(rowStr(row, "type")) - if !strings.Contains(email, q) && !strings.Contains(id, q) && !strings.Contains(typ, q) { - continue - } - } - filtered = append(filtered, row) - } - - total := len(filtered) - limit, offset := pageParams(c, 20) - page := pageSlice(filtered, limit, offset) - c.JSON(http.StatusOK, gin.H{"data": page, "total": total, "limit": limit, "offset": offset, "stats": stats}) -} - -// accountsStats reproduces the 账号 KPI strip: per-type 正常/失效/限额 counts plus a -// grand total and total dead count (drives 「删除异常账号 (N)」). -func accountsStats(rows []map[string]any) gin.H { - types := []string{"openai", "adobe", "runway", "leonardo", "krea", "imagine", "grok", "creativefabrica"} - by := map[string]*struct{ N, Ok, Dead, Quota int }{} - for _, t := range types { - by[t] = &struct{ N, Ok, Dead, Quota int }{} - } - deadTotal := 0 - for _, row := range rows { - dead := rowBool(row, "dead") - if dead { - deadTotal++ - } - g, ok := by[rowStr(row, "type")] - if !ok { - continue - } - g.N++ - status := rowStr(row, "status") - switch { - case status == "active": - g.Ok++ - case dead || status == "disabled": - g.Dead++ - case status == "quota": - g.Quota++ - } - } - out := gin.H{"total": len(rows), "dead_total": deadTotal} - for _, t := range types { - g := by[t] - out[t] = gin.H{"n": g.N, "ok": g.Ok, "dead": g.Dead, "quota": g.Quota} - } - return out -} - -func rowStr(m map[string]any, key string) string { - s, _ := m[key].(string) - return s -} - -func rowBool(m map[string]any, key string) bool { - b, _ := m[key].(bool) - return b -} - -func (h *ProviderAdminHandler) AccountQuota(c *gin.Context) { - data, err := h.tokens.Quota(c.Request.Context(), c.Param("pool"), c.Param("id")) - if err != nil { - if errors.Is(err, gorm.ErrRecordNotFound) { - c.JSON(http.StatusNotFound, gin.H{"detail": "account not found"}) - return - } - c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()}) - return - } - c.JSON(http.StatusOK, data) -} - -func (h *ProviderAdminHandler) AccountEmail(c *gin.Context) { - data, err := h.tokens.Email(c.Request.Context(), c.Param("pool"), c.Param("id")) - if err != nil { - if errors.Is(err, gorm.ErrRecordNotFound) { - c.JSON(http.StatusNotFound, gin.H{"detail": "account not found"}) - return - } - c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()}) - return - } - c.JSON(http.StatusOK, data) -} - -func (h *ProviderAdminHandler) RefreshProfiles(c *gin.Context) { - items, err := h.refresh.List(c.Request.Context()) - if err != nil { - c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load refresh profiles"}) - return - } - c.JSON(http.StatusOK, gin.H{"data": items}) -} - -func (h *ProviderAdminHandler) RefreshNow(c *gin.Context) { - if err := h.refresh.RefreshNow(c.Request.Context(), c.Param("profile_id")); err != nil { - c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()}) - return - } - c.JSON(http.StatusOK, gin.H{"ok": true}) -} - -func (h *ProviderAdminHandler) RefreshUpdate(c *gin.Context) { - var body map[string]any - if err := c.ShouldBindJSON(&body); err != nil { - c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"}) - return - } - item, err := h.refresh.Update(c.Request.Context(), c.Param("profile_id"), body) - if err != nil { - c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()}) - return - } - c.JSON(http.StatusOK, gin.H{"ok": true, "data": item}) -} - -func (h *ProviderAdminHandler) RefreshDelete(c *gin.Context) { - if err := h.refresh.Delete(c.Request.Context(), c.Param("profile_id")); err != nil { - c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()}) - return - } - c.JSON(http.StatusOK, gin.H{"ok": true}) -} diff --git a/backend/internal/http/handler/user_generation.go b/backend/internal/http/handler/user_generation.go index f7b746e..55a2cf5 100644 --- a/backend/internal/http/handler/user_generation.go +++ b/backend/internal/http/handler/user_generation.go @@ -3,6 +3,7 @@ package handler import ( "bytes" "errors" + "fmt" "image" _ "image/gif" _ "image/jpeg" @@ -13,6 +14,8 @@ import ( "sync" "time" + "backend/internal/model" + "backend/internal/repo" "backend/internal/service" "github.com/gin-gonic/gin" _ "golang.org/x/image/webp" @@ -248,7 +251,7 @@ func (h *UserGenerationHandler) Test(c *gin.Context) { Resolution string `json:"resolution"` Duration string `json:"duration"` ReferenceImages []string `json:"reference_images"` - AccountID string `json:"account_id"` + ChannelID string `json:"channel_id"` } if err := c.ShouldBindJSON(&body); err != nil { c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"}) @@ -262,7 +265,7 @@ func (h *UserGenerationHandler) Test(c *gin.Context) { Resolution: body.Resolution, Duration: body.Duration, ReferenceImages: body.ReferenceImages, - AccountID: body.AccountID, + ChannelID: body.ChannelID, }) if err != nil { switch { @@ -451,161 +454,19 @@ func (h *UserGenerationHandler) Logs(c *gin.Context) { } func (h *UserGenerationHandler) VideoPresets(c *gin.Context) { - c.JSON(http.StatusOK, gin.H{ - "data": []gin.H{ - { - "key": "gemini-veo31", - "label": "Veo31 Fast", - "type": "video", - "provider": "adobe", - "durations": []string{"4s", "6s", "8s"}, - "ratios": []string{"16:9", "9:16"}, - "resolutions": []string{"720p", "1080p"}, - "max_reference_images": 9, - "reference_mode": "frame", - }, - { - "key": "gemini-veo3.1-fast", - "label": "Veo 3.1 Fast", - "type": "video", - "provider": "adobe", - "durations": []string{"4s", "6s", "8s"}, - "ratios": []string{"16:9", "9:16"}, - "resolutions": []string{"720p", "1080p"}, - "max_reference_images": 9, - "reference_mode": "frame", - }, - { - "key": "gemini-veo3.1", - "label": "Veo 3.1", - "type": "video", - "provider": "adobe", - "durations": []string{"4s", "6s", "8s"}, - "ratios": []string{"16:9", "9:16"}, - "resolutions": []string{"720p", "1080p"}, - }, - { - "key": "adobe-seedance-2.0-fast", - "label": "Seedance 2.0 Fast", - "type": "video", - "provider": "adobe", - "durations": []string{"4s", "5s", "6s", "7s", "8s", "9s", "10s", "11s", "12s", "13s", "14s", "15s"}, - "ratios": []string{"16:9", "9:16"}, - "resolutions": []string{"720p"}, - "max_reference_images": 9, - "reference_mode": "style", - }, - { - "key": "adobe-seedance-2.0", - "label": "Seedance 2.0", - "type": "video", - "provider": "adobe", - "durations": []string{"4s", "5s", "6s", "7s", "8s", "9s", "10s", "11s", "12s", "13s", "14s", "15s"}, - "ratios": []string{"16:9", "9:16"}, - "resolutions": []string{"720p", "1080p"}, - "max_reference_images": 9, - "reference_mode": "style", - }, - { - "key": "seedance-2.0-fast", - "label": "Seedance 2.0 Fast (Creative Fabrica)", - "type": "video", - "provider": "creativefabrica", - "durations": []string{"14s"}, - "ratios": []string{"16:9", "9:16"}, - "resolutions": []string{"720p"}, - // Creative Fabrica 上游只有普通参考图(VIDEO_FRAME_TYPE_REFERENCE), - // 没有首尾帧,也不收视频/音频参考。 - "max_reference_images": 9, - "reference_mode": "asset", - "max_videos": 0, - "max_audios": 0, - }, - { - "key": "seedance-2.0", - "label": "Seedance 2.0 (Creative Fabrica)", - "type": "video", - "provider": "creativefabrica", - "durations": []string{"10s"}, - "ratios": []string{"16:9", "9:16"}, - "resolutions": []string{"720p"}, - // Creative Fabrica 上游只有普通参考图(VIDEO_FRAME_TYPE_REFERENCE), - // 没有首尾帧,也不收视频/音频参考。 - "max_reference_images": 9, - "reference_mode": "asset", - "max_videos": 0, - "max_audios": 0, - }, - { - "key": "seedance-2.0-不卡人脸", - "label": "Seedance 2.0 (Leonardo 私有)", - "type": "video", - "provider": "leonardo", - "durations": []string{"4s", "5s", "6s", "7s", "8s", "9s", "10s", "11s", "12s", "13s", "14s", "15s"}, - "ratios": []string{"16:9", "9:16"}, - "resolutions": []string{"720p"}, - "max_reference_images": 8, - "max_images": 4, - "max_videos": 3, - "max_audios": 1, - "reference_mode": "style", - }, - { - "key": "seedance-2.0-fast-不卡人脸", - "label": "Seedance 2.0 Fast (Leonardo 私有)", - "type": "video", - "provider": "leonardo", - "durations": []string{"4s", "5s", "6s", "7s", "8s", "9s", "10s", "11s", "12s", "13s", "14s", "15s"}, - "ratios": []string{"16:9", "9:16"}, - "resolutions": []string{"720p"}, - "max_reference_images": 8, - "max_images": 4, - "max_videos": 3, - "max_audios": 1, - "reference_mode": "style", - }, - { - "key": "minimax-h3", - "label": "MiniMax H3 (Leonardo 私有)", - "type": "video", - "provider": "leonardo", - "durations": []string{"5s", "6s", "7s", "8s", "9s", "10s", "11s", "12s", "13s", "14s", "15s"}, - "ratios": []string{"16:9", "9:16"}, - "resolutions": []string{"1440p"}, - "max_reference_images": 8, - "max_images": 5, - "max_videos": 0, - "max_audios": 3, - "max_audio_seconds": 15, - "reference_mode": "style", - }, - { - "key": "firefly-video", - "label": "Firefly Video", - "type": "video", - "provider": "adobe", - "durations": []string{"5s"}, - "ratios": []string{"16:9", "1:1", "9:16"}, - "resolutions": []string{"540p", "720p", "1080p"}, - "max_reference_images": 9, - "reference_mode": "frame", - }, - { - "key": "runway-gen4-turbo", - "label": "Runway Gen-4 Turbo", - "type": "video", - "provider": "runway", - "durations": []string{"5s", "10s"}, - "ratios": []string{"16:9", "9:16", "1:1", "4:3", "3:4", "21:9"}, - "resolutions": []string{"2K"}, - "max_reference_images": 1, - "reference_mode": "frame", - // Runway is strictly image-to-video — a first-frame image is required - // (no text2video), so the UI must block submit without one. - "requires_reference": true, - }, - }, - }) + models, err := h.admin.Models(c.Request.Context()) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load video presets"}) + return + } + out := make([]gin.H, 0, len(models)) + for _, m := range models { + if m.Type != "video" || !m.Enabled { + continue + } + out = append(out, modelPreset(m)) + } + c.JSON(http.StatusOK, gin.H{"data": out}) } func (h *UserGenerationHandler) Catalog(c *gin.Context) { @@ -620,7 +481,7 @@ func (h *UserGenerationHandler) Catalog(c *gin.Context) { } func (h *UserGenerationHandler) Models(c *gin.Context) { - items, err := h.publicModels() + items, err := h.publicModels(c) if err != nil { c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load models"}) return @@ -628,537 +489,85 @@ func (h *UserGenerationHandler) Models(c *gin.Context) { c.JSON(http.StatusOK, gin.H{"data": items}) } +// modelPreset projects one model_configs row into the preset shape the studio +// (画图台) / model pickers consume. Database-driven now: providers are gone and +// every model arrives via a channel import. +func modelPreset(m model.ModelConfig) gin.H { + durations := repo.JSONStrings(m.Durations) + if len(durations) == 0 && m.DurationMin > 0 && m.DurationMax >= m.DurationMin { + // Range spec → expand integer seconds (capped so absurd ranges can't + // flood a dropdown). + for d := m.DurationMin; d <= m.DurationMax && len(durations) < 61; d++ { + durations = append(durations, fmt.Sprintf("%ds", d)) + } + } + _, perSecond := m.DurationPrices["per_second"] + return gin.H{ + "key": m.EffectiveName(), + "id": m.ID, + "label": m.EffectiveName(), + "type": m.Type, + "provider": m.Provider, + "ratios": repo.JSONStrings(m.Ratios), + "resolutions": repo.JSONStrings(m.Resolutions), + "durations": durations, + "duration_min": m.DurationMin, + "duration_max": m.DurationMax, + "per_second": perSecond, + "max_reference_images": m.MaxReferenceImages, + "max_images": m.MaxReferenceImages, + "max_videos": m.MaxReferenceVideos, + "max_audios": m.MaxReferenceAudios, + "max_reference_videos": m.MaxReferenceVideos, + "max_reference_audios": m.MaxReferenceAudios, + "max_video_seconds": m.MaxRefVideoSeconds, + "max_audio_seconds": m.MaxRefAudioSeconds, + "reference_mode": m.ReferenceMode, + "image_to_image": m.ImageToImage, + "description": m.Name, + "added": true, + } +} + +// catalogEntries lists the managed models (previously a static picklist of +// offerable provider models). With channel imports the DB IS the catalog — +// rows are already "added", so the admin picklist degrades gracefully. func (h *UserGenerationHandler) catalogEntries(c *gin.Context) ([]gin.H, error) { - items := []gin.H{ - { - "id": "gpt-image-2", - "provider": "chatgpt", - "type": "image", - // ChatGPT web backend only reliably produces 1K and honors a limited - // ratio set; size params are advisory prompt hints. Mirrors the Python - // reference (providers/chatgpt/provider.py) — do not offer 2K/4K. - "ratios": []string{"1:1", "16:9", "9:16", "4:3", "3:4"}, - "resolutions": []string{"1K"}, - "image_to_image": true, - "max_reference_images": 6, - "description": "ChatGPT image generation", - }, - { - "id": "firefly-gpt-image-2", - "provider": "adobe", - "type": "image", - "ratios": []string{"1:1", "5:4", "9:16", "21:9", "16:9", "4:3", "3:2", "4:5", "3:4", "2:3"}, - "resolutions": []string{"1K", "2K", "4K"}, - "image_to_image": true, - "max_reference_images": 6, - "description": "Adobe Firefly GPT Image", - }, - { - "id": "firefly-image-5", - "provider": "adobe", - "type": "image", - "ratios": []string{"1:1", "16:9", "9:16", "4:3", "3:4"}, - "resolutions": []string{"1K", "2K"}, - "image_to_image": true, - "description": "Adobe Firefly Image 5", - }, - { - "id": "runway-nano-banana-2", - "provider": "runway", - "type": "image", - "ratios": []string{"1:1", "1:4", "1:8", "2:3", "3:2", "3:4", "4:1", "4:3", "4:5", "5:4", "8:1", "9:16", "16:9", "21:9"}, - "resolutions": []string{"1K", "2K", "4K"}, - "image_to_image": true, - "max_reference_images": 6, - "reference_mode": "asset", - "description": "Runway Nano Banana 2 (图/参考图)", - }, - { - "id": "runway-nano-banana-pro", - "provider": "runway", - "type": "image", - "ratios": []string{"1:1", "1:4", "1:8", "2:3", "3:2", "3:4", "4:1", "4:3", "4:5", "5:4", "8:1", "9:16", "16:9", "21:9"}, - "resolutions": []string{"1K", "2K", "4K"}, - "image_to_image": true, - "max_reference_images": 6, - "reference_mode": "asset", - "description": "Runway Nano Banana Pro (图/参考图)", - }, - { - "id": "nano-banana-2", - "provider": "adobe", - "type": "image", - "ratios": []string{"1:1", "1:4", "1:8", "2:3", "3:2", "3:4", "4:1", "4:3", "4:5", "5:4", "8:1", "9:16", "16:9", "21:9"}, - "resolutions": []string{"1K", "2K", "4K"}, - "image_to_image": true, - "max_reference_images": 6, - "reference_mode": "asset", - "description": "Nano Banana 2 (图/参考图)", - }, - { - "id": "nano-banana-pro", - "provider": "adobe", - "type": "image", - "ratios": []string{"1:1", "1:4", "1:8", "2:3", "3:2", "3:4", "4:1", "4:3", "4:5", "5:4", "8:1", "9:16", "16:9", "21:9"}, - "resolutions": []string{"1K", "2K", "4K"}, - "image_to_image": true, - "max_reference_images": 6, - "reference_mode": "asset", - "description": "Nano Banana Pro (图/参考图)", - }, - { - "id": "gemini-veo3.1-fast", - "provider": "adobe", - "type": "video", - "ratios": []string{"16:9", "9:16"}, - "resolutions": []string{"720p", "1080p"}, - "durations": []string{"4s", "6s", "8s"}, - "max_reference_images": 3, - "reference_mode": "frame", - "description": "Veo 3.1 Fast", - }, - { - "id": "gemini-veo3.1", - "provider": "adobe", - "type": "video", - "ratios": []string{"16:9", "9:16"}, - "resolutions": []string{"720p", "1080p"}, - "durations": []string{"4s", "6s", "8s"}, - "max_reference_images": 3, - "reference_mode": "style", - "description": "Veo 3.1", - }, - { - "id": "gemini-veo3.1", - "provider": "adobe", - "type": "video", - "ratios": []string{"16:9", "9:16"}, - "resolutions": []string{"720p", "1080p"}, - "durations": []string{"5s", "10s"}, - "max_reference_images": 2, - "reference_mode": "frame", - "description": "Luma Ray video", - }, - { - "id": "firefly-video", - "provider": "adobe", - "type": "video", - "ratios": []string{"16:9", "1:1", "9:16"}, - "resolutions": []string{"720p", "1080p"}, - "durations": []string{"5s"}, - "max_reference_images": 2, - "reference_mode": "frame", - "description": "Adobe Firefly Video", - }, - { - "id": "adobe-seedance-2.0-fast", - "provider": "adobe", - "type": "video", - "ratios": []string{"16:9", "9:16"}, - "resolutions": []string{"720p"}, - "durations": []string{"4s", "5s", "6s", "7s", "8s", "9s", "10s", "11s", "12s", "13s", "14s", "15s"}, - "per_second": true, - "max_reference_images": 9, - "reference_mode": "style", - "description": "Seedance 2.0 Fast", - }, - { - "id": "adobe-seedance-2.0", - "provider": "adobe", - "type": "video", - "ratios": []string{"16:9", "9:16"}, - "resolutions": []string{"720p", "1080p"}, - "durations": []string{"4s", "5s", "6s", "7s", "8s", "9s", "10s", "11s", "12s", "13s", "14s", "15s"}, - "per_second": true, - "max_reference_images": 9, - "reference_mode": "style", - "description": "Seedance 2.0", - }, - { - "id": "seedance-2.0-fast", - "provider": "creativefabrica", - "type": "video", - "ratios": []string{"16:9", "9:16"}, - "resolutions": []string{"720p"}, - // 一次性账号：积分刚好够一次生成，时长固定 14 秒。 - "durations": []string{"14s"}, - // Creative Fabrica 上游只有普通参考图,没有首尾帧,也不收视频/音频参考。 - "max_reference_images": 9, - "reference_mode": "asset", - "description": "Seedance 2.0 Fast (Creative Fabrica)", - }, - { - "id": "seedance-2.0", - "provider": "creativefabrica", - "type": "video", - "ratios": []string{"16:9", "9:16"}, - "resolutions": []string{"720p"}, - // 一次性账号：积分刚好够一次生成，时长固定 10 秒。 - "durations": []string{"10s"}, - // Creative Fabrica 上游只有普通参考图,没有首尾帧,也不收视频/音频参考。 - "max_reference_images": 9, - "reference_mode": "asset", - "description": "Seedance 2.0 (Creative Fabrica)", - }, - { - "id": "runway-gen4-turbo", - "provider": "runway", - "type": "video", - "ratios": []string{"16:9", "9:16", "1:1", "4:3", "3:4", "21:9"}, - "resolutions": []string{"720p"}, - "durations": []string{"5s", "10s"}, - "max_reference_images": 1, - "reference_mode": "frame", - "description": "Runway Gen-4 Turbo video (图生视频)", - }, - { - "id": "grok-video", - "provider": "grok", - "type": "video", - "ratios": []string{"2:3", "3:2", "1:1", "9:16", "16:9"}, - "resolutions": []string{"720p"}, - // Console 视频吃 1–15 秒的整数时长，所以按秒计价（同 seedance）； - // 参考图只有 1 张首帧（上游只有 image 字段，没有尾帧）。 - "durations": []string{"4s", "5s", "6s", "7s", "8s", "9s", "10s", "11s", "12s", "13s", "14s", "15s"}, - "per_second": true, - "max_reference_images": 1, - "reference_mode": "frame", - "description": "Grok Imagine video (文/图生视频)", - }, - { - "id": "grok-image", - "provider": "grok", - "type": "image", - "ratios": []string{"1:1", "2:3", "3:2", "3:4", "4:3", "9:16", "16:9"}, - "resolutions": []string{"1K", "2K"}, - "image_to_image": true, - "max_reference_images": 3, - "reference_mode": "asset", - "description": "Grok Imagine image (文生图 / 图生图)", - }, - { - "id": "seedream-4.5", - "provider": "leonardo", - "type": "image", - "ratios": []string{"2:3", "1:1", "16:9", "4:3", "4:5", "9:16", "2:1"}, - "resolutions": []string{"2K", "4K"}, - "image_to_image": true, - "max_reference_images": 6, - "description": "Leonardo Seedream 4.5 (生图 / 图生图)", - }, - { - "id": "seedance-2.0-不卡人脸", - "provider": "leonardo", - "type": "video", - "ratios": []string{"16:9", "9:16"}, - "resolutions": []string{"720p"}, - "durations": []string{"4s", "5s", "6s", "7s", "8s", "9s", "10s", "11s", "12s", "13s", "14s", "15s"}, - "per_second": true, - // 参考资产总上限 8 = 4 图 + 1 音频 + 3 视频（分类上限在服务端校验）。 - "max_reference_images": 8, - "reference_mode": "style", - "description": "Leonardo Seedance 2.0 (私有生成 / 图音视频参考)", - }, - { - "id": "seedance-2.0-fast-不卡人脸", - "provider": "leonardo", - "type": "video", - "ratios": []string{"16:9", "9:16"}, - "resolutions": []string{"720p"}, - "durations": []string{"4s", "5s", "6s", "7s", "8s", "9s", "10s", "11s", "12s", "13s", "14s", "15s"}, - "per_second": true, - "max_reference_images": 8, - "reference_mode": "style", - "description": "Leonardo Seedance 2.0 Fast (私有生成 / 图音视频参考)", - }, - { - "id": "minimax-h3", - "provider": "leonardo", - "type": "video", - "ratios": []string{"16:9", "9:16"}, - "resolutions": []string{"1440p"}, - // 上游吃 5–15 秒整数时长，所以按秒计价（后台只填一个 /s 价）。 - "durations": []string{"5s", "6s", "7s", "8s", "9s", "10s", "11s", "12s", "13s", "14s", "15s"}, - "per_second": true, - // 参考资产总上限 8 = 5 图 + 3 音频（分类上限在服务端校验）。 - "max_reference_images": 8, - "reference_mode": "style", - "description": "Leonardo MiniMax H3 / hailuo-03 (1440p / 图音参考)", - }, - { - "id": "flux-klein-2", - "provider": "krea", - "type": "image", - "ratios": []string{"1:1", "4:3", "3:4", "16:9", "9:16"}, - "resolutions": []string{"1K", "2K"}, - "image_to_image": true, - "max_reference_images": 4, - "description": "Krea Flux Klein (生图 / 图生图)", - }, - { - "id": "imagine-1.5", - "provider": "imagine", - "type": "image", - "ratios": []string{"1:3", "9:16", "2:3", "3:4", "1:1", "4:3", "3:2", "16:9", "3:1"}, - "resolutions": []string{"2K"}, - "max_reference_images": 0, - "description": "Imagine 1.5 (文生图)", - }, - { - "id": "imagine-1.5pro", - "provider": "imagine", - "type": "image", - "ratios": []string{"1:3", "9:16", "2:3", "3:4", "1:1", "4:3", "3:2", "16:9", "3:1"}, - "resolutions": []string{"4K"}, - "max_reference_images": 0, - "description": "Imagine 1.5 Pro (文生图)", - }, + models, err := h.admin.Models(c.Request.Context()) + if err != nil { + return nil, err } - existing := map[string]bool{} - if h.admin != nil { - models, err := h.admin.Models(c.Request.Context()) - if err != nil { - return nil, err - } - for _, item := range models { - existing[item.ID] = true - } - } - for i := range items { - items[i]["added"] = existing[items[i]["id"].(string)] + items := make([]gin.H, 0, len(models)) + for _, m := range models { + row := modelPreset(m) + row["id"] = m.EffectiveName() + items = append(items, row) } return items, nil } -func (h *UserGenerationHandler) publicModels() ([]gin.H, error) { - items := []gin.H{ - { - "id": "gpt-image-2", - "provider": "chatgpt", - "kind": "image", - // See catalogEntries — ChatGPT only reliably does 1K and a limited - // ratio set; matches the Python reference. Keep both lists in sync. - "ratios": []string{"1:1", "16:9", "9:16", "4:3", "3:4"}, - "resolutions": []string{"1K"}, - "description": "ChatGPT image generation", +// publicModels — GET /api/models for the public site (docs/landing pickers). +// Enabled models only, in the same shape as before minus the provider id. +func (h *UserGenerationHandler) publicModels(c *gin.Context) ([]gin.H, error) { + models, err := h.admin.Models(c.Request.Context()) + if err != nil { + return nil, err + } + items := make([]gin.H, 0, len(models)) + for _, m := range models { + if !m.Enabled { + continue + } + items = append(items, gin.H{ + "id": m.EffectiveName(), + "kind": m.Type, + "type": m.Type, + "ratios": repo.JSONStrings(m.Ratios), + "resolutions": repo.JSONStrings(m.Resolutions), + "durations": repo.JSONStrings(m.Durations), + "description": m.Name, "stub": false, - }, - { - "id": "firefly-gpt-image-2", - "provider": "adobe", - "kind": "image", - "ratios": []string{"1:1", "5:4", "9:16", "21:9", "16:9", "4:3", "3:2", "4:5", "3:4", "2:3"}, - "resolutions": []string{"1K", "2K", "4K"}, - "description": "Adobe Firefly GPT Image", - "stub": false, - }, - { - "id": "firefly-image-5", - "provider": "adobe", - "kind": "image", - "ratios": []string{"1:1", "16:9", "9:16", "4:3", "3:4"}, - "resolutions": []string{"1K", "2K"}, - "description": "Adobe Firefly Image 5", - "stub": false, - }, - { - "id": "runway-nano-banana-2", - "provider": "runway", - "kind": "image", - "ratios": []string{"1:1", "1:4", "1:8", "2:3", "3:2", "3:4", "4:1", "4:3", "4:5", "5:4", "8:1", "9:16", "16:9", "21:9"}, - "resolutions": []string{"1K", "2K", "4K"}, - "description": "Runway Nano Banana 2", - "stub": false, - }, - { - "id": "runway-nano-banana-pro", - "provider": "runway", - "kind": "image", - "ratios": []string{"1:1", "1:4", "1:8", "2:3", "3:2", "3:4", "4:1", "4:3", "4:5", "5:4", "8:1", "9:16", "16:9", "21:9"}, - "resolutions": []string{"1K", "2K", "4K"}, - "description": "Runway Nano Banana Pro", - "stub": false, - }, - { - "id": "nano-banana-2", - "provider": "adobe", - "kind": "image", - "ratios": []string{"1:1", "1:4", "1:8", "2:3", "3:2", "3:4", "4:1", "4:3", "4:5", "5:4", "8:1", "9:16", "16:9", "21:9"}, - "resolutions": []string{"1K", "2K", "4K"}, - "description": "Nano Banana 2", - "stub": false, - }, - { - "id": "nano-banana-pro", - "provider": "adobe", - "kind": "image", - "ratios": []string{"1:1", "1:4", "1:8", "2:3", "3:2", "3:4", "4:1", "4:3", "4:5", "5:4", "8:1", "9:16", "16:9", "21:9"}, - "resolutions": []string{"1K", "2K", "4K"}, - "description": "Nano Banana Pro", - "stub": false, - }, - { - "id": "gemini-veo3.1-fast", - "provider": "adobe", - "kind": "video", - "ratios": []string{"16:9", "9:16"}, - "resolutions": []string{"720p", "1080p"}, - "description": "Veo 3.1 Fast", - "stub": false, - }, - { - "id": "gemini-veo3.1", - "provider": "adobe", - "kind": "video", - "ratios": []string{"16:9", "9:16"}, - "resolutions": []string{"720p"}, - "description": "Veo 3.1", - "stub": false, - }, - { - "id": "firefly-ray", - "provider": "adobe", - "kind": "video", - "ratios": []string{"21:9", "16:9", "4:3", "1:1", "3:4", "9:16", "9:21"}, - "resolutions": []string{"720p"}, - "description": "Luma Ray video", - "stub": false, - }, - { - "id": "firefly-video", - "provider": "adobe", - "kind": "video", - "ratios": []string{"16:9", "1:1", "9:16"}, - "resolutions": []string{"720p", "1080p"}, - "description": "Adobe Firefly Video", - "stub": false, - }, - { - "id": "adobe-seedance-2.0-fast", - "provider": "adobe", - "kind": "video", - "ratios": []string{"16:9", "9:16"}, - "resolutions": []string{"720p"}, - "description": "Seedance 2.0 Fast", - "stub": false, - }, - { - "id": "adobe-seedance-2.0", - "provider": "adobe", - "kind": "video", - "ratios": []string{"16:9", "9:16"}, - "resolutions": []string{"720p", "1080p"}, - "description": "Seedance 2.0", - "stub": false, - }, - { - "id": "seedance-2.0-fast", - "provider": "creativefabrica", - "kind": "video", - "ratios": []string{"16:9", "9:16"}, - "resolutions": []string{"720p"}, - "description": "Seedance 2.0 Fast (Creative Fabrica)", - "stub": false, - }, - { - "id": "seedance-2.0", - "provider": "creativefabrica", - "kind": "video", - "ratios": []string{"16:9", "9:16"}, - "resolutions": []string{"720p"}, - "description": "Seedance 2.0 (Creative Fabrica)", - "stub": false, - }, - { - "id": "runway-gen4-turbo", - "provider": "runway", - "kind": "video", - "ratios": []string{"16:9", "9:16", "1:1", "4:3", "3:4", "21:9"}, - "resolutions": []string{"720p"}, - "description": "Runway Gen-4 Turbo video", - "stub": false, - }, - { - "id": "grok-video", - "provider": "grok", - "kind": "video", - "ratios": []string{"2:3", "3:2", "1:1", "9:16", "16:9"}, - "resolutions": []string{"720p"}, - "description": "Grok Imagine video", - "stub": false, - }, - { - "id": "grok-image", - "provider": "grok", - "kind": "image", - "ratios": []string{"1:1", "2:3", "3:2", "3:4", "4:3", "9:16", "16:9"}, - "resolutions": []string{"1K", "2K"}, - "description": "Grok Imagine image", - "stub": false, - }, - { - "id": "seedream-4.5", - "provider": "leonardo", - "kind": "image", - "ratios": []string{"2:3", "1:1", "16:9", "4:3", "4:5", "9:16", "2:1"}, - "resolutions": []string{"2K", "4K"}, - "description": "Leonardo Seedream 4.5", - "stub": false, - }, - { - "id": "seedance-2.0-不卡人脸", - "provider": "leonardo", - "kind": "video", - "ratios": []string{"16:9", "9:16"}, - "resolutions": []string{"720p"}, - "description": "Leonardo Seedance 2.0 (私有生成)", - "stub": false, - }, - { - "id": "seedance-2.0-fast-不卡人脸", - "provider": "leonardo", - "kind": "video", - "ratios": []string{"16:9", "9:16"}, - "resolutions": []string{"720p"}, - "description": "Leonardo Seedance 2.0 Fast (私有生成)", - "stub": false, - }, - { - "id": "minimax-h3", - "provider": "leonardo", - "kind": "video", - "ratios": []string{"16:9", "9:16"}, - "resolutions": []string{"1440p"}, - "description": "Leonardo MiniMax H3 (1440p)", - "stub": false, - }, - { - "id": "flux-klein-2", - "provider": "krea", - "kind": "image", - "ratios": []string{"1:1", "4:3", "3:4", "16:9", "9:16"}, - "resolutions": []string{"1K", "2K"}, - "description": "Krea Flux Klein", - "stub": false, - }, - { - "id": "imagine-1.5", - "provider": "imagine", - "kind": "image", - "ratios": []string{"1:3", "9:16", "2:3", "3:4", "1:1", "4:3", "3:2", "16:9", "3:1"}, - "resolutions": []string{"2K"}, - "description": "Imagine 1.5", - "stub": false, - }, - { - "id": "imagine-1.5pro", - "provider": "imagine", - "kind": "image", - "ratios": []string{"1:3", "9:16", "2:3", "3:4", "1:1", "4:3", "3:2", "16:9", "3:1"}, - "resolutions": []string{"4K"}, - "description": "Imagine 1.5 Pro", - "stub": false, - }, + }) } return items, nil } diff --git a/backend/internal/http/handler/v1.go b/backend/internal/http/handler/v1.go index 0b46cb5..9849b6f 100644 --- a/backend/internal/http/handler/v1.go +++ b/backend/internal/http/handler/v1.go @@ -1,19 +1,33 @@ +// Package handler — /v1/* 消费者 HTTP 接口（vivid 协议同款）。 +// 这套接口同时是渠道上游契约：同构实例经相同的四个端点互相对接。 +// +// POST /v1/generate 创建生成任务（NDJSON 流：首帧 jobId → 终态帧；断流携 jobId 重发续等） +// POST /v1/uploads 参考素材预签名直传（返回 key 随 generate 的 refs 提交） +// GET /v1/models 可售模型裸数组（启用 + 有渠道绑定 + 全部档位已配价） +// GET /v1/balance {balance}：余额积分 +// +// 鉴权：Authorization: Bearer vk_...（用户 API Key，哈希落库）。 package handler import ( - "encoding/base64" "encoding/json" "errors" - "fmt" - "io" "net/http" "strconv" "strings" + "time" + "backend/internal/model" "backend/internal/service" "github.com/gin-gonic/gin" ) +// v1MaxWait 同步等待终态的上限。视频单渠道预算 30 分钟、最多换几个渠道， +// 留足余量取 2 小时（与 vivid 一致）。 +const v1MaxWait = 2 * time.Hour + +const v1PollInterval = time.Second + type V1Handler struct { v1 *service.V1Service } @@ -22,385 +36,213 @@ func NewV1Handler(v1 *service.V1Service) *V1Handler { return &V1Handler{v1: v1} } -func (h *V1Handler) Models(c *gin.Context) { +func (h *V1Handler) auth(c *gin.Context) (*service.APIPrincipal, bool) { principal, err := h.v1.Authenticate(c.Request.Context(), c.GetHeader("Authorization")) if err != nil { - h.writeAuthError(c, err) - return + v1Err(c, http.StatusUnauthorized, "4011", "无效的 API Key") + return nil, false } - _ = principal + return principal, true +} - items, err := h.v1.ListModels(c.Request.Context()) - if err != nil { - c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load models"}) - return - } - c.JSON(http.StatusOK, gin.H{ - "object": "list", - "data": items, +// v1Err writes the {"error":{"code","message"}} envelope — the wire shape every +// vivid-protocol caller parses. +func v1Err(c *gin.Context, httpStatus int, code, msg string) { + c.JSON(httpStatus, gin.H{ + "error": gin.H{"code": code, "message": msg}, }) } -// UserBalance — GET /v1/user/balance. 返回 API Key 所属用户的账户级余额 -// （剩余 / 累计已用），与具体令牌无关。 -func (h *V1Handler) UserBalance(c *gin.Context) { - principal, err := h.v1.Authenticate(c.Request.Context(), c.GetHeader("Authorization")) - if err != nil { - h.writeAuthError(c, err) - return +// v1Fail maps a service error to the vivid code/HTTP status pair. +func v1Fail(c *gin.Context, err error) { + switch { + case errors.Is(err, service.ErrMissingAPIKey), errors.Is(err, service.ErrInvalidAPIKey): + v1Err(c, http.StatusUnauthorized, "4011", "无效的 API Key") + case errors.Is(err, service.ErrInsufficientFunds): + v1Err(c, http.StatusBadRequest, "4002", "积分不足，请先充值") + case errors.Is(err, service.ErrConcurrencyFull), errors.Is(err, service.ErrUserConcurrencyFull): + v1Err(c, http.StatusTooManyRequests, "4293", "当前并发已满，请稍后重试") + case errors.Is(err, service.ErrVideoJobNotFound): + v1Err(c, http.StatusBadRequest, "4001", "任务不存在") + case errors.Is(err, service.ErrReferenceTooLarge): + v1Err(c, http.StatusRequestEntityTooLarge, "4001", "参考素材超过大小上限") + case errors.Is(err, service.ErrUnknownModel): + v1Err(c, http.StatusBadRequest, "4001", "模型不存在或未启用") + case errors.Is(err, service.ErrNoProviderAccount): + v1Err(c, http.StatusBadRequest, "4001", "模型无可用渠道") + case errors.Is(err, service.ErrBannedPrompt): + v1Err(c, http.StatusBadRequest, "4001", "提示词包含违禁内容") + case errors.Is(err, service.ErrUnsupportedParams): + v1Err(c, http.StatusBadRequest, "4001", err.Error()) + case errors.Is(err, service.ErrProviderTemporary), errors.Is(err, service.ErrProviderExecution), + errors.Is(err, service.ErrProviderAuth): + v1Err(c, http.StatusInternalServerError, "5001", "生成暂时不可用，积分已退还") + default: + v1Err(c, http.StatusBadRequest, "4001", err.Error()) } - resp, err := h.v1.UserBalance(c.Request.Context(), principal) - if err != nil { - c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to load balance"}) - return - } - c.JSON(http.StatusOK, resp) } -// ImageGenerations — OpenAI POST /v1/images/generations (text-to-image only). -// Accepts exactly OpenAI's fields; size→aspect ratio and quality→resolution tier -// are mapped server-side. Returns {created, data:[{b64_json}]}. -func (h *V1Handler) ImageGenerations(c *gin.Context) { - principal, err := h.v1.Authenticate(c.Request.Context(), c.GetHeader("Authorization")) - if err != nil { - h.writeAuthError(c, err) +// ---- POST /v1/generate ---- + +func (h *V1Handler) Generate(c *gin.Context) { + principal, ok := h.auth(c) + if !ok { return } - - var body struct { - Model string `json:"model"` - Prompt string `json:"prompt"` - N int `json:"n"` - Size string `json:"size"` - Quality string `json:"quality"` - ResponseFormat string `json:"response_format"` - Background string `json:"background"` - OutputFormat string `json:"output_format"` - User string `json:"user"` - } + var body service.V1JobRequest if err := c.ShouldBindJSON(&body); err != nil { - c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"}) + v1Err(c, http.StatusBadRequest, "4001", "请求体不是合法 JSON") return } - - resp, err := h.v1.PrepareImageRequest(c.Request.Context(), principal, service.V1ImageRequest{ - Model: body.Model, - Prompt: body.Prompt, - N: body.N, - Size: body.Size, - BaseURL: requestBaseURL(c), - }) - if err != nil { - h.writeV1Error(c, err, resp) - return - } - c.JSON(http.StatusOK, openaiImageResponse(resp)) -} - -// ImageEdits — OpenAI POST /v1/images/edits (image-to-image). multipart/form-data -// only: image / image[] file uploads (+ optional mask), prompt, model, n, size, -// quality. Files become reference images. Returns {created, data:[{b64_json}]}. -func (h *V1Handler) ImageEdits(c *gin.Context) { - principal, err := h.v1.Authenticate(c.Request.Context(), c.GetHeader("Authorization")) - if err != nil { - h.writeAuthError(c, err) - return - } - if !strings.HasPrefix(c.GetHeader("Content-Type"), "multipart/form-data") { - c.JSON(http.StatusBadRequest, gin.H{"detail": "images/edits requires multipart/form-data"}) - return - } - if err := c.Request.ParseMultipartForm(64 << 20); err != nil { - c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid multipart form"}) - return - } - refs := readMultipartImages(c, "image", "image[]") - if len(refs) == 0 { - c.JSON(http.StatusBadRequest, gin.H{"detail": "images/edits requires at least one image file"}) - return - } - n, _ := strconv.Atoi(strings.TrimSpace(c.PostForm("n"))) - resp, err := h.v1.PrepareImageRequest(c.Request.Context(), principal, service.V1ImageRequest{ - Model: c.PostForm("model"), - Prompt: c.PostForm("prompt"), - N: n, - Size: c.PostForm("size"), - ReferenceImages: refs, - BaseURL: requestBaseURL(c), - }) - if err != nil { - h.writeV1Error(c, err, resp) - return - } - c.JSON(http.StatusOK, openaiImageResponse(resp)) -} - -// CreateVideo — OpenAI POST /v1/videos. Creates an async job and returns the -// video object immediately ({id, status:"queued"}). Accepts JSON {model, prompt, -// seconds, size} or multipart (with an input_reference file). size→ratio+ -// resolution, seconds→duration. -func (h *V1Handler) CreateVideo(c *gin.Context) { - principal, err := h.v1.Authenticate(c.Request.Context(), c.GetHeader("Authorization")) - if err != nil { - h.writeAuthError(c, err) - return - } - var modelID, prompt, seconds, size, referenceMode string - var refs []string - if strings.HasPrefix(c.GetHeader("Content-Type"), "multipart/form-data") { - if err := c.Request.ParseMultipartForm(256 << 20); err != nil { - c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid multipart form"}) - return - } - modelID = c.PostForm("model") - prompt = c.PostForm("prompt") - seconds = c.PostForm("seconds") - size = c.PostForm("size") - referenceMode = c.PostForm("reference_mode") - refs = readMultipartImages(c, "input_reference", "input_reference[]") + var ev *model.EventLog + var err error + // 重连模式：只带 jobId → 不建新任务，回读该任务并等待终态。 + if body.JobID != nil && strings.TrimSpace(*body.JobID) != "" { + ev, err = h.v1.JobForUser(c.Request.Context(), principal, *body.JobID) } else { - var body struct { - Model string `json:"model"` - Prompt string `json:"prompt"` - Seconds json.RawMessage `json:"seconds"` - Size string `json:"size"` - ReferenceMode string `json:"reference_mode"` - InputReference []string `json:"input_reference"` - ReferenceImages []string `json:"reference_images"` + ev, err = h.v1.CreateJob(c.Request.Context(), principal, body) + } + if err != nil { + v1Fail(c, err) + return + } + + wait := 300 * time.Second + if h := strings.TrimSpace(c.GetHeader("X-Wait-Seconds")); h != "" { + if sec, serr := strconv.Atoi(h); serr == nil && sec > 0 && time.Duration(sec)*time.Second <= v1MaxWait { + wait = time.Duration(sec) * time.Second } - if err := c.ShouldBindJSON(&body); err != nil { - c.JSON(http.StatusBadRequest, gin.H{"detail": "invalid request body"}) - return + } + + c.Writer.Header().Set("Content-Type", "application/x-ndjson") + c.Writer.Header().Set("Cache-Control", "no-cache, no-transform") + c.Writer.Header().Set("X-Accel-Buffering", "no") + c.Writer.WriteHeader(http.StatusOK) + flusher, _ := c.Writer.(http.Flusher) + + writeFrame := func(v *model.EventLog) bool { + frame, _ := json.Marshal(v1JobBody(h.v1, v)) + if _, werr := c.Writer.Write(append(frame, '\n')); werr != nil { + return false } - modelID, prompt, size = body.Model, body.Prompt, body.Size - seconds = rawToString(body.Seconds) - referenceMode = strings.TrimSpace(body.ReferenceMode) - refs = append(body.InputReference, body.ReferenceImages...) + if flusher != nil { + flusher.Flush() + } + return true } - duration := strings.TrimSpace(seconds) - if duration != "" && !strings.HasSuffix(duration, "s") { - duration += "s" - } - aspect, resolution := videoSizeToInternal(size) - resp, err := h.v1.StartVideoJob(c.Request.Context(), principal, service.V1VideoRequest{ - Model: modelID, - Prompt: prompt, - Duration: duration, - AspectRatio: aspect, - Resolution: resolution, - ReferenceImages: refs, - ReferenceMode: referenceMode, - BaseURL: requestBaseURL(c), - }) - if err != nil { - h.writeV1Error(c, err, nil) - return - } - c.JSON(http.StatusOK, resp) -} -// GetVideo — OpenAI GET /v1/videos/{id}. Returns the job status object. -func (h *V1Handler) GetVideo(c *gin.Context) { - principal, err := h.v1.Authenticate(c.Request.Context(), c.GetHeader("Authorization")) - if err != nil { - h.writeAuthError(c, err) + // 已终态：直接回终态帧关流。 + if st := service.JobStatus(ev); st == "succeeded" || st == "failed" { + writeFrame(ev) return } - resp, err := h.v1.VideoJob(c.Request.Context(), principal, c.Param("id")) - if err != nil { - h.writeV1Error(c, err, nil) + // 连上即回执 jobId（断流后携它重发的凭据）。 + if !writeFrame(ev) { return } - c.JSON(http.StatusOK, resp) -} -// GetVideoContent — OpenAI GET /v1/videos/{id}/content. Streams the rendered mp4 -// by proxying the stored upstream URL (downloaded on demand, never persisted). -func (h *V1Handler) GetVideoContent(c *gin.Context) { - principal, err := h.v1.Authenticate(c.Request.Context(), c.GetHeader("Authorization")) - if err != nil { - h.writeAuthError(c, err) - return - } - body, contentType, err := h.v1.OpenVideoContent(c.Request.Context(), principal, c.Param("id")) - if err != nil { - h.writeV1Error(c, err, nil) - return - } - defer body.Close() - c.Header("Content-Type", contentType) - c.Status(http.StatusOK) - _, _ = io.Copy(c.Writer, body) -} - -// GetImageContent — GET /v1/images/{id}/content. Streams a no-store image by -// proxying its stored (possibly auth-gated) upstream URL. Never persisted. -func (h *V1Handler) GetImageContent(c *gin.Context) { - principal, err := h.v1.Authenticate(c.Request.Context(), c.GetHeader("Authorization")) - if err != nil { - h.writeAuthError(c, err) - return - } - body, contentType, err := h.v1.OpenImageContent(c.Request.Context(), principal, c.Param("id")) - if err != nil { - h.writeV1Error(c, err, nil) - return - } - defer body.Close() - c.Header("Content-Type", contentType) - c.Status(http.StatusOK) - _, _ = io.Copy(c.Writer, body) -} - -// readMultipartImages reads the given file fields and returns each as base64. -func readMultipartImages(c *gin.Context, keys ...string) []string { - var out []string - form := c.Request.MultipartForm - if form == nil { - return out - } - for _, key := range keys { - for _, fh := range form.File[key] { - f, e := fh.Open() - if e != nil { + deadline := time.Now().Add(wait) + poll := time.NewTicker(v1PollInterval) + heartbeat := time.NewTicker(25 * time.Second) + defer poll.Stop() + defer heartbeat.Stop() + for { + select { + case <-c.Request.Context().Done(): + return // client gone — the job keeps running; caller may reconnect with jobId + case <-poll.C: + latest, lerr := h.v1.JobForUser(c.Request.Context(), principal, ev.ID) + if lerr != nil || latest == nil { continue } - b, _ := io.ReadAll(io.LimitReader(f, 20<<20+1)) - f.Close() - if len(b) > 0 { - out = append(out, base64.StdEncoding.EncodeToString(b)) + ev = latest + if st := service.JobStatus(ev); st == "succeeded" || st == "failed" { + writeFrame(ev) + return } + case <-heartbeat.C: + if _, werr := c.Writer.WriteString("\n"); werr != nil { + return + } + if flusher != nil { + flusher.Flush() + } + case <-time.After(time.Until(deadline)): + writeFrame(ev) // snapshot; client reconnects with jobId + return } } - return out } -// rawToString accepts OpenAI's `seconds` whether sent as a JSON string or number. -func rawToString(raw json.RawMessage) string { - if len(raw) == 0 { - return "" +// v1JobBody assembles one NDJSON frame (running receipt and terminal frame +// share the shape; succeeded adds result+creditsCharged, failed adds error). +func v1JobBody(v1 *service.V1Service, ev *model.EventLog) map[string]any { + body := map[string]any{"jobId": ev.ID, "status": service.JobStatus(ev)} + switch service.JobStatus(ev) { + case "succeeded": + body["result"] = v1.JobResultURL(ev) + body["creditsCharged"] = ev.Cost + case "failed": + body["error"] = ev.Error + body["refunded"] = true } - var s string - if json.Unmarshal(raw, &s) == nil { - return s - } - var n json.Number - if json.Unmarshal(raw, &n) == nil { - return n.String() - } - return strings.Trim(string(raw), `"`) + return body } -// videoSizeToInternal maps OpenAI's "WxH" size to our aspect ratio + resolution -// tier. An absent/unparsable size leaves the resolution empty so the caller can -// fall back to whatever tier the model actually prices — hardcoding 720p here -// rejects models that only offer 1440p. -func videoSizeToInternal(size string) (ratio, resolution string) { - var w, h int - if s := strings.TrimSpace(strings.ToLower(size)); s != "" { - _, _ = fmt.Sscanf(s, "%dx%d", &w, &h) +// ---- POST /v1/uploads ---- + +func (h *V1Handler) Uploads(c *gin.Context) { + principal, ok := h.auth(c) + if !ok { + return } - if w == 0 || h == 0 { - return "16:9", "" + var body struct { + Kind string `json:"kind"` + Mime string `json:"mime"` } - // The "p" resolution is the SHORT edge (720p = 1280×720, 1080p = 1920×1080), - // so a standard 1280×720 must read as 720p — not 1080p off the long edge. - resolution = "720p" - switch { - case min(w, h) >= 1440: - resolution = "1440p" - case min(w, h) >= 1080: - resolution = "1080p" + if err := c.ShouldBindJSON(&body); err != nil { + v1Err(c, http.StatusBadRequest, "4001", "请求体不是合法 JSON") + return } - return guessRatioWH(w, h), resolution + view, err := h.v1.PresignRefUpload(c.Request.Context(), principal, body.Kind, body.Mime) + if err != nil { + v1Fail(c, err) + return + } + c.JSON(http.StatusOK, view) } -func guessRatioWH(w, h int) string { - if w == h { - return "1:1" +// ---- GET /v1/models ---- + +func (h *V1Handler) Models(c *gin.Context) { + if _, ok := h.auth(c); !ok { + return } - r := float64(w) / float64(h) - cands := []struct { - name string - v float64 - }{{"16:9", 16.0 / 9}, {"9:16", 9.0 / 16}, {"4:3", 4.0 / 3}, {"3:4", 3.0 / 4}, {"1:1", 1}} - best, bestD := "16:9", 1e9 - for _, cd := range cands { - d := r - cd.v - if d < 0 { - d = -d - } - if d < bestD { - best, bestD = cd.name, d - } + items, err := h.v1.SellableModels(c.Request.Context()) + if err != nil { + v1Err(c, http.StatusInternalServerError, "5001", "模型加载失败") + return } - return best + c.JSON(http.StatusOK, items) } -// openaiImageResponse strips our rich internal map down to OpenAI's image shape. -func openaiImageResponse(m map[string]any) gin.H { - out := gin.H{"created": m["created"]} - if d, ok := m["data"]; ok && d != nil { - out["data"] = d - } else { - out["data"] = []any{} - } - return out -} +// ---- GET /v1/balance ---- -func (h *V1Handler) writeAuthError(c *gin.Context, err error) { - switch { - case errors.Is(err, service.ErrMissingAPIKey): - c.JSON(http.StatusUnauthorized, gin.H{"detail": err.Error()}) - case errors.Is(err, service.ErrInvalidAPIKey): - c.JSON(http.StatusUnauthorized, gin.H{"detail": err.Error()}) - default: - c.JSON(http.StatusInternalServerError, gin.H{"detail": "failed to validate api key"}) +func (h *V1Handler) UserBalance(c *gin.Context) { + principal, ok := h.auth(c) + if !ok { + return } -} - -func (h *V1Handler) writeV1Error(c *gin.Context, err error, payload map[string]any) { - switch { - case errors.Is(err, service.ErrUnknownModel): - c.JSON(http.StatusNotFound, gin.H{"detail": err.Error()}) - case errors.Is(err, service.ErrUnsupportedParams), errors.Is(err, service.ErrBannedPrompt): - c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()}) - case errors.Is(err, service.ErrInsufficientFunds): - c.JSON(http.StatusPaymentRequired, gin.H{"detail": err.Error()}) - case errors.Is(err, service.ErrReferenceTooLarge): - c.JSON(http.StatusRequestEntityTooLarge, gin.H{"detail": err.Error()}) - case errors.Is(err, service.ErrNoProviderAccount): - c.JSON(http.StatusServiceUnavailable, gin.H{"detail": err.Error()}) - case errors.Is(err, service.ErrProviderAuth): - c.JSON(http.StatusServiceUnavailable, gin.H{"detail": err.Error()}) - case errors.Is(err, service.ErrProviderQuota): - // Match the Python contract: provider quota exhaustion maps to 401 - // (QuotaExhaustedError is handled alongside AuthError in routes.py). - c.JSON(http.StatusUnauthorized, gin.H{"detail": err.Error()}) - case errors.Is(err, service.ErrProviderTemporary): - c.JSON(http.StatusServiceUnavailable, gin.H{"detail": err.Error()}) - case errors.Is(err, service.ErrConcurrencyFull), errors.Is(err, service.ErrUserConcurrencyFull): - c.JSON(http.StatusTooManyRequests, gin.H{"detail": err.Error()}) - case errors.Is(err, service.ErrVideoJobNotFound): - c.JSON(http.StatusNotFound, gin.H{"detail": err.Error()}) - case errors.Is(err, service.ErrVideoNotReady): - c.JSON(http.StatusConflict, gin.H{"detail": err.Error()}) - case errors.Is(err, service.ErrProviderUnsupported): - c.JSON(http.StatusNotImplemented, gin.H{"detail": err.Error()}) - case errors.Is(err, service.ErrProviderExecution): - c.JSON(http.StatusBadGateway, gin.H{"detail": err.Error()}) - case errors.Is(err, service.ErrGenerationPending): - c.JSON(http.StatusNotImplemented, payload) - default: - c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()}) + bal, err := h.v1.UserBalance(c.Request.Context(), principal) + if err != nil { + v1Err(c, http.StatusInternalServerError, "5001", "余额查询失败") + return } + c.JSON(http.StatusOK, gin.H{"balance": bal}) } // requestBaseURL derives the scheme+host of the inbound request so the service -// layer can build absolute, directly-downloadable output URLs. Honors -// X-Forwarded-Proto (reverse-proxy / TLS termination) before falling back to -// the connection's TLS state. Returns "" when the host is unknown, which makes -// the service fall back to a relative path. +// layer can build absolute /images URLs for the session path. func requestBaseURL(c *gin.Context) string { host := c.Request.Host if host == "" { diff --git a/backend/internal/http/router/router.go b/backend/internal/http/router/router.go index 8ff660b..6b15e24 100644 --- a/backend/internal/http/router/router.go +++ b/backend/internal/http/router/router.go @@ -11,25 +11,25 @@ import ( ) type Handlers struct { - Health *handler.HealthHandler - Images *handler.ImageHandler - V1 *handler.V1Handler - Site *handler.SiteHandler - Showcase *handler.ShowcaseHandler - Auth *handler.AuthHandler - SiteSettings *handler.SiteSettingsHandler - AppSettings *handler.AppSettingsHandler - AdminRead *handler.AdminReadHandler - AdminWrite *handler.AdminWriteHandler - CDK *handler.CDKHandler - UserTools *handler.UserToolsHandler - UserGen *handler.UserGenerationHandler - ProviderAdmin *handler.ProviderAdminHandler - ConcGroups *handler.ConcurrencyGroupHandler - Announcement *handler.AnnouncementHandler - Payment *handler.PaymentHandler - BannedWords *handler.BannedWordsHandler - Canvas *handler.CanvasHandler + Health *handler.HealthHandler + Images *handler.ImageHandler + V1 *handler.V1Handler + Site *handler.SiteHandler + Showcase *handler.ShowcaseHandler + Auth *handler.AuthHandler + SiteSettings *handler.SiteSettingsHandler + AppSettings *handler.AppSettingsHandler + AdminRead *handler.AdminReadHandler + AdminWrite *handler.AdminWriteHandler + CDK *handler.CDKHandler + UserTools *handler.UserToolsHandler + UserGen *handler.UserGenerationHandler + ChannelAdmin *handler.ChannelAdminHandler + ConcGroups *handler.ConcurrencyGroupHandler + Announcement *handler.AnnouncementHandler + Payment *handler.PaymentHandler + BannedWords *handler.BannedWordsHandler + Canvas *handler.CanvasHandler } func New(cfg *config.Config, auth *service.AuthService, handlers Handlers) *gin.Engine { @@ -59,16 +59,12 @@ func New(cfg *config.Config, auth *service.AuthService, handlers Handlers) *gin. engine.GET("/health", handlers.Health.Handle) engine.GET("/images/:user/:name", handlers.Images.Serve) + // vivid-protocol public API — same four endpoints a vivid instance exposes, + // which is what lets this service act as a channel upstream itself. + engine.POST("/v1/generate", handlers.V1.Generate) + engine.POST("/v1/uploads", handlers.V1.Uploads) engine.GET("/v1/models", handlers.V1.Models) - engine.GET("/v1/user/balance", handlers.V1.UserBalance) - engine.POST("/v1/images/generations", handlers.V1.ImageGenerations) - engine.POST("/v1/images/edits", handlers.V1.ImageEdits) - // OpenAI Sora-style async video: create job → poll → stream content. - engine.POST("/v1/videos", handlers.V1.CreateVideo) - engine.GET("/v1/videos/:id", handlers.V1.GetVideo) - engine.GET("/v1/videos/:id/content", handlers.V1.GetVideoContent) - // No-store image content proxy (auth-gated upstream URLs, e.g. chatgpt). - engine.GET("/v1/images/:id/content", handlers.V1.GetImageContent) + engine.GET("/v1/balance", handlers.V1.UserBalance) publicAdmin := engine.Group("/admin/api") { @@ -79,7 +75,6 @@ func New(cfg *config.Config, auth *service.AuthService, handlers Handlers) *gin. // owner names is admin-only, see /stats/recent below. publicAdmin.GET("/stats", handlers.AdminRead.Stats) publicAdmin.GET("/video-presets", handlers.UserGen.VideoPresets) - publicAdmin.GET("/catalog", handlers.UserGen.Catalog) publicAdmin.GET("/models", handlers.UserGen.Models) publicAdmin.GET("/deai-pricing", handlers.AppSettings.DeAIGet) // 易支付 async notify — called server-to-server by the pay platform, no auth. @@ -147,23 +142,11 @@ func New(cfg *config.Config, auth *service.AuthService, handlers Handlers) *gin. authed.POST("/cdks", handlers.CDK.Create) authed.POST("/cdks/delete-bulk", handlers.CDK.DeleteBulk) authed.DELETE("/cdks/:code", handlers.CDK.Delete) - authed.GET("/tokens", handlers.ProviderAdmin.TokensList) - authed.POST("/tokens", handlers.ProviderAdmin.TokensCreate) - authed.POST("/tokens/import-chatgpt-token", handlers.ProviderAdmin.ImportChatGPTToken) - authed.POST("/tokens/import-adobe-cookie", handlers.ProviderAdmin.ImportAdobeCookie) - authed.POST("/tokens/import-creativefabrica-cookie", handlers.ProviderAdmin.ImportCreativeFabricaCookie) - authed.POST("/tokens/import-runway-token", handlers.ProviderAdmin.ImportRunwayToken) - 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/import-custom-account", handlers.ProviderAdmin.ImportCustomAccount) - authed.POST("/tokens/delete-bulk", handlers.ProviderAdmin.TokenDeleteBulk) - authed.PATCH("/tokens/:pool/:id", handlers.ProviderAdmin.TokenUpdate) - authed.DELETE("/tokens/:pool/:id", handlers.ProviderAdmin.TokenDelete) - authed.GET("/accounts", handlers.ProviderAdmin.AccountsList) - authed.GET("/accounts/:pool/:id/quota", handlers.ProviderAdmin.AccountQuota) - authed.GET("/accounts/:pool/:id/email", handlers.ProviderAdmin.AccountEmail) + authed.GET("/channels", handlers.ChannelAdmin.List) + authed.POST("/channels", handlers.ChannelAdmin.Save) + authed.POST("/channels/delete", handlers.ChannelAdmin.Delete) + authed.POST("/channels/fetch-models", handlers.ChannelAdmin.FetchModels) + authed.POST("/channels/import-models", handlers.ChannelAdmin.ImportModels) authed.GET("/providers", handlers.AdminRead.Providers) authed.GET("/images", handlers.AdminRead.Images) authed.DELETE("/images", handlers.AdminRead.DeleteImage) @@ -172,10 +155,6 @@ func New(cfg *config.Config, auth *service.AuthService, handlers Handlers) *gin. authed.POST("/banned-words/import", handlers.BannedWords.Import) authed.DELETE("/banned-words/:id", handlers.BannedWords.Delete) authed.GET("/banned-word-hits", handlers.BannedWords.Hits) - authed.GET("/refresh/profiles", handlers.ProviderAdmin.RefreshProfiles) - authed.POST("/refresh/profiles/:profile_id/refresh-now", handlers.ProviderAdmin.RefreshNow) - authed.PATCH("/refresh/profiles/:profile_id", handlers.ProviderAdmin.RefreshUpdate) - authed.DELETE("/refresh/profiles/:profile_id", handlers.ProviderAdmin.RefreshDelete) authed.POST("/managed-models", handlers.AdminWrite.CreateModel) authed.PATCH("/managed-models/:model_id", handlers.AdminWrite.UpdateModel) authed.DELETE("/managed-models/:model_id", handlers.AdminWrite.DeleteModel) diff --git a/backend/internal/model/models.go b/backend/internal/model/models.go index ead0a05..172afea 100644 --- a/backend/internal/model/models.go +++ b/backend/internal/model/models.go @@ -97,17 +97,18 @@ type EventLog struct { DeAI bool `gorm:"not null;default:false"` // 去AI特征 was applied (image only) RefFiles datatypes.JSON `gorm:"type:jsonb"` // relative paths of saved reference images, for回显 on reload Source string `gorm:"size:32;index"` - // AccountID is the provider token/account chosen to fulfil this generation, - // stamped when the upstream call begins. Drives the accounts view's live - // in-flight count (pending events per account) and lets an abandoned-event - // purge attribute the failure back to the account it was using. + // AccountID is the channel chosen to fulfil this generation, stamped when + // the upstream call begins. (Field name kept from the pre-channel schema.) AccountID string `gorm:"size:64;index"` - // AccountEmail denormalizes the account's email at stamp time, so log rows - // keep showing which mailbox fulfilled the generation even after the account - // is deleted or re-imported under a different ID. - AccountEmail string `gorm:"size:255"` - UserID string `gorm:"size:32;index"` - Cost float64 `gorm:"not null;default:0"` + // AccountEmail denormalizes the channel's display name at stamp time, so log + // rows keep showing which channel fulfilled the generation even after the + // channel is deleted or renamed. + AccountEmail string `gorm:"size:255"` + // UpstreamJob is the job id assigned by the upstream vivid instance (from + // the NDJSON stream's jobId field) — kept for cross-instance debugging. + UpstreamJob string `gorm:"size:64"` + UserID string `gorm:"size:32;index"` + Cost float64 `gorm:"not null;default:0"` // Refunded marks that this event's up-front charge has already been credited // back, so the normal failure path and the abandoned-purge sweep can never // double-refund the same generation. @@ -138,16 +139,20 @@ type ModelConfig struct { PricesAgent datatypes.JSONMap `gorm:"type:jsonb;column:prices_agent"` DurationPricesAgent datatypes.JSONMap `gorm:"type:jsonb;column:duration_prices_agent"` Durations datatypes.JSON `gorm:"type:jsonb"` - MaxReferenceImages int `gorm:"not null;default:0"` - // Per-kind reference caps for upstreams that accept video/audio refs - // (vivid-protocol custom upstreams). 0 = that kind is not accepted. - MaxReferenceVideos int `gorm:"not null;default:0"` - MaxReferenceAudios int `gorm:"not null;default:0"` + // DurationMin/DurationMax bound a free-form duration (seconds) for models + // whose duration spec is a range rather than discrete options. 0 = unset. + DurationMin int `gorm:"not null;default:0"` + DurationMax int `gorm:"not null;default:0"` + MaxReferenceImages int `gorm:"not null;default:0"` + // Per-kind reference caps for upstreams that accept video/audio refs. + // 0 = that kind is not accepted. + MaxReferenceVideos int `gorm:"not null;default:0"` + MaxReferenceAudios int `gorm:"not null;default:0"` + // MaxRefVideoSeconds/MaxRefAudioSeconds cap the combined duration of + // reference videos/audios. 0 = unlimited. + MaxRefVideoSeconds int `gorm:"not null;default:0"` + MaxRefAudioSeconds int `gorm:"not null;default:0"` ReferenceMode string `gorm:"size:32;not null;default:'none'"` - // Custom-upstream models (provider="custom"): UpstreamModel is the model name - // sent to the upstream OpenAI-compatible API; the base_url + key live on the - // matching custom account (pool="custom", meta.base_url). Empty for built-ins. - UpstreamModel string `gorm:"size:255;not null;default:''"` // Weight controls display order in the model dropdown / admin list: higher // weight floats to the top (matches ShowcaseItem.Weight semantics). Ties fall // back to created_at desc. Default 0. @@ -180,55 +185,37 @@ type CDKCode struct { UpdatedAt time.Time } -type TokenAccount struct { - ID string `gorm:"primaryKey;size:64"` - Pool string `gorm:"size:64;index;not null"` - Value string `gorm:"type:text"` - Status string `gorm:"size:32;index;not null"` - Fails int `gorm:"not null;default:0"` - FailTotal int `gorm:"not null;default:0"` - SuccessTotal int `gorm:"not null;default:0"` - Dead bool `gorm:"not null;default:false"` - Meta datatypes.JSONMap `gorm:"type:jsonb"` - AddedAt *time.Time - LastUsedAt *time.Time - CachedQuotaResetAfter string `gorm:"size:128"` - QuotaRecoverAt *time.Time - // Adobe quota is tracked separately for image vs video. An account only - // enters the shared "quota" waiting status when BOTH are limited; a single - // limit leaves the account usable for the other kind. Recovery time is shared - // (QuotaRecoverAt / CachedQuotaResetAfter) since Adobe resets both at once. - ImageLimited bool `gorm:"not null;default:false"` - VideoLimited bool `gorm:"not null;default:false"` - AccountEmail string `gorm:"size:255"` - AccountDisplayName string `gorm:"size:255"` - // Weight biases scheduling order for ANY account — higher weight is picked - // first within its pool (ties fall back to round-robin). Default 0. - Weight int `gorm:"not null;default:0"` - // Concurrency is the max simultaneous jobs for THIS account. Only custom - // (upstream) accounts honor it; built-in pools use their system default - // (1 per account, grok 10). 0 = use the system default. +// Channel is one upstream vivid-protocol instance this gateway routes +// generation to. Replaces the old per-provider token account pool: a channel +// is just a base_url + api_key plus routing knobs (enabled/weight/concurrency) +// and local usage counters. ApiKey is stored AES-GCM encrypted. +type Channel struct { + ID string `gorm:"primaryKey;size:40"` + Name string `gorm:"size:255;not null"` + BaseURL string `gorm:"size:500;not null"` + APIKeyEnc string `gorm:"type:text;not null"` // AES-GCM encrypted upstream API key + Enabled bool `gorm:"not null;default:true;index"` + // Concurrency caps simultaneous in-flight jobs on this channel. 0 = default. Concurrency int `gorm:"not null;default:0"` - CreatedAt time.Time - UpdatedAt time.Time + // Weight biases weighted-random channel order: weight 2 is picked first at + // roughly 2x the rate of weight 1. + Weight int `gorm:"not null;default:1"` + SuccessTotal int64 `gorm:"not null;default:0"` // locally-tracked successful generations + UsedCredits float64 `gorm:"not null;default:0"` // upstream credits spent (reported by the upstream) + LastUsedAt *time.Time + CreatedAt time.Time + UpdatedAt time.Time } -type RefreshProfile struct { - ID string `gorm:"primaryKey;size:64"` - Name string `gorm:"size:255;not null"` - Pool string `gorm:"size:64;index;not null"` - Kind string `gorm:"size:64;index;not null"` - Cookie string `gorm:"type:text"` - Enabled bool `gorm:"not null;default:true"` - IntervalSeconds int `gorm:"not null;default:54000"` - ImportedAt *time.Time - LastAttemptAt *time.Time - LastSuccessAt *time.Time - LastError string `gorm:"type:text"` - NextRetryAt *time.Time - ConsecutiveFailures int `gorm:"not null;default:0"` - CreatedAt time.Time - UpdatedAt time.Time +// ChannelModel binds a local model_configs row to a channel and records the +// upstream model name to send on the wire (the local row's alias stays the +// user-facing name). +type ChannelModel struct { + ChannelID string `gorm:"primaryKey;size:40"` + ModelID string `gorm:"primaryKey;size:255;index"` + UpstreamModel string `gorm:"size:255;not null"` + CreatedAt time.Time + UpdatedAt time.Time } type SiteSetting struct { @@ -248,8 +235,8 @@ func AutoMigrateModels() []any { &EventLog{}, &ModelConfig{}, &CDKCode{}, - &TokenAccount{}, - &RefreshProfile{}, + &Channel{}, + &ChannelModel{}, &SiteSetting{}, &StatCounter{}, &ConcurrencyGroup{}, diff --git a/backend/internal/provider/adobe/arp_shape_test.go b/backend/internal/provider/adobe/arp_shape_test.go deleted file mode 100644 index 5667b74..0000000 --- a/backend/internal/provider/adobe/arp_shape_test.go +++ /dev/null @@ -1,162 +0,0 @@ -package adobe - -import ( - "encoding/base64" - "encoding/json" - "regexp" - "strings" - "testing" -) - -// captured from live firefly.adobe.com traffic: -// -// ark: 91818c89a54748463.1048135404|r=ap-southeast-1|…|rid=84|ag=101|… -// ftr: 40d7ba7c70ff478e9c212073ec8254c4_1791795834777_27064_UDF43-mnf-a4_31ck_s/Clo68HDL4=-3629-v2_tt -// -// The Arkose slot is deliberately emitted empty rather than synthesized — see -// buildARPSessionID — so the expected ftr ends in "_31ck__tt". -var ( - arkPat = regexp.MustCompile(`^[0-9a-f]{17}\.[1-9][0-9]{9}\|r=ap-southeast-1\|.*\|rid=[0-9]{1,2}\|ag=101\|`) - ftrPat = regexp.MustCompile(`^[0-9a-f]{32}_[0-9]{13}_[0-9]{4,5}_UDF43-mnf-a4_31ck__tt$`) -) - -func TestARPSessionIDShape(t *testing.T) { - raw, err := base64.StdEncoding.DecodeString(buildARPSessionID("tok-a")) - if err != nil { - t.Fatalf("not base64: %v", err) - } - var got struct{ Sid, Ark, Ftr, Bfp string } - if err := json.Unmarshal(raw, &got); err != nil { - t.Fatalf("not json: %v", err) - } - if len(got.Bfp) != 36 { - t.Errorf("bfp not a uuid: %q", got.Bfp) - } - if !arkPat.MatchString(got.Ark) { - t.Errorf("ark shape mismatch:\n%s", got.Ark) - } - if !ftrPat.MatchString(got.Ftr) { - t.Errorf("ftr shape mismatch:\n%s", got.Ftr) - } - if len(got.Sid) != 36 { - t.Errorf("sid not a uuid: %q", got.Sid) - } -} - -// ark must differ per call — a frozen blob is a cross-account correlation key. -func TestARKVariesAcrossCalls(t *testing.T) { - if buildARKBlob() == buildARKBlob() { - t.Error("ark is constant across calls") - } -} - -// pid is stable per token but distinct across tokens. -func TestPIDStablePerToken(t *testing.T) { - defer ReleasePID("tok-x") - defer ReleasePID("tok-y") - if allocPID("tok-x") != allocPID("tok-x") { - t.Error("pid changed for the same token") - } - if allocPID("tok-x") == allocPID("tok-y") { - t.Error("two tokens share a pid") - } -} - -// A real browser keeps sid / devHex / bfp for the whole session; only the ftr -// timestamp advances. Re-randomizing them per call is a bot signal. -func TestIdentityStablePerToken(t *testing.T) { - decode := func(token string) map[string]string { - raw, err := base64.StdEncoding.DecodeString(buildARPSessionID(token)) - if err != nil { - t.Fatalf("not base64: %v", err) - } - var m map[string]string - if err := json.Unmarshal(raw, &m); err != nil { - t.Fatalf("not json: %v", err) - } - return m - } - a, b := decode("tok-s"), decode("tok-s") - if a["sid"] != b["sid"] || a["bfp"] != b["bfp"] { - t.Error("sid/bfp changed across calls for the same token") - } - if a["ftr"][:39] != b["ftr"][:39] { - t.Error("ftr device hex changed across calls for the same token") - } -} - -// CacheARP stores a cookie-derived arp which getARPSessionID prefers over -// synthesis. -func TestCacheARPPreferred(t *testing.T) { - c := NewClient("", "") - CacheARP("tok-c", "eyJjb29raWUiOiJhYXJkZWYa") - if got := c.getARPSessionID("tok-c"); got != "eyJjb29raWUiOiJhYXJkZWYa" { - t.Errorf("cached arp not preferred: %q", got) - } -} - -// refreshARPTimestamp bumps only the ftr timestamp segment; devHex, pid, the -// Arkose slot, and every other payload field pass through untouched. -func TestRefreshARPTimestamp(t *testing.T) { - mk := func(ftr string) string { - b, _ := json.Marshal(map[string]string{ - "sid": "4e54df9a-79cf-4569-b40d-fa793e38b727", - "ftr": ftr, - "ark": "ark_blob", - "bfp": "bf", - }) - return base64.StdEncoding.EncodeToString(b) - } - old := "40d7ba7c70ff478e9c212073ec8254c4_1791795834777_27064_UDF43-mnf-a4_31ck__tt" - raw, err := base64.StdEncoding.DecodeString(refreshARPTimestamp(mk(old))) - if err != nil { - t.Fatalf("refreshed arp not base64: %v", err) - } - var m map[string]string - if err := json.Unmarshal(raw, &m); err != nil { - t.Fatalf("not json: %v", err) - } - parts := strings.Split(m["ftr"], "_") - if len(parts) != 7 { - t.Fatalf("ftr segment count changed: %q", m["ftr"]) - } - if parts[0] != "40d7ba7c70ff478e9c212073ec8254c4" || parts[2] != "27064" || parts[4] != "31ck" || parts[6] != "tt" { - t.Errorf("ftr fields mangled: %q", m["ftr"]) - } - if parts[1] == "1791795834777" { - t.Error("ftr timestamp not bumped") - } - if m["sid"] != "4e54df9a-79cf-4569-b40d-fa793e38b727" || m["ark"] != "ark_blob" || m["bfp"] != "bf" { - t.Errorf("non-ftr fields mangled: %v", m) - } - // malformed inputs pass through unchanged - if got := refreshARPTimestamp("!!!not-base64"); got != "!!!not-base64" { - t.Errorf("garbage input changed: %q", got) - } - if got := refreshARPTimestamp(mk("too_short")); got != mk("too_short") { - t.Error("short ftr changed") - } -} - -// ARPFromCookie returns the sherlockToken verbatim when present, and rebuilds -// the JSON from component cookies otherwise. -func TestARPFromCookie(t *testing.T) { - if got := ARPFromCookie("a=1; sherlockToken=eyJhIjoxfQ==; b=2"); got != "eyJhIjoxfQ==" { - t.Errorf("sherlockToken not returned verbatim: %q", got) - } - if got := ARPFromCookie("a=1"); got != "" { - t.Errorf("expected empty for cookie without arp pieces, got %q", got) - } - raw, err := base64.StdEncoding.DecodeString(ARPFromCookie( - "ff_session_guid=4e54df9a-79cf-4569-b40d-fa793e38b727; forterToken=ft_x%3D-v2; arkose=ark_blob; bfp=bf")) - if err != nil { - t.Fatalf("reassembled arp not base64: %v", err) - } - var m map[string]string - if err := json.Unmarshal(raw, &m); err != nil { - t.Fatalf("not json: %v", err) - } - if m["sid"] != "4e54df9a-79cf-4569-b40d-fa793e38b727" || m["ftr"] != "ft_x=-v2" || m["ark"] != "ark_blob" || m["bfp"] != "bf" { - t.Errorf("reassembled arp mismatch: %v", m) - } -} diff --git a/backend/internal/provider/adobe/auth.go b/backend/internal/provider/adobe/auth.go deleted file mode 100644 index 2a1ab77..0000000 --- a/backend/internal/provider/adobe/auth.go +++ /dev/null @@ -1,44 +0,0 @@ -package adobe - -import ( - "context" - "errors" - "net/http" - "strings" -) - -const ( - refreshURL = "https://adobeid-na1.services.adobe.com/ims/check/v6/token?jslVersion=v2-v0.48.0-1-g1e322cb" - clientID = "projectx_webapp" - scopeValue = "AdobeID,firefly_api,openid" - // clioClientID / clioScopeValue mint the firefly.adobe.com token that is - // entitled to the Firefly-native seedance models (creative_production scope). - // Unlike the Express token this exchange requires user_id and no guest_allowed. - clioClientID = "clio-playground-web" - clioScopeValue = "AdobeID,firefly_api,openid,pps.read,pps.write,additional_info.projectedProductContext,additional_info.ownerOrg,uds_read,uds_write,ab.manage,read_organizations,additional_info.roles,account_cluster.read,creative_production,tk_platform,tk_platform_sync,profile" -) - -var ErrAdobeCookieEmpty = errors.New("cookie is empty") - -type CookieExchangeResult struct { - AccessToken string - ExpiresIn int - Raw map[string]any -} - -func ExchangeCookieToAccessToken(ctx context.Context, client *http.Client, cookie string) (*CookieExchangeResult, error) { - _ = client - tlsClient, err := NewClient(clientID, "").newTLSClient() - if err != nil { - return nil, err - } - return exchangeCookieWithTLSClient(ctx, tlsClient, cookie) -} - -func normalizeCookie(v string) string { - v = strings.TrimSpace(v) - if strings.HasPrefix(strings.ToLower(v), "cookie:") { - v = strings.TrimSpace(v[len("cookie:"):]) - } - return v -} diff --git a/backend/internal/provider/adobe/client.go b/backend/internal/provider/adobe/client.go deleted file mode 100644 index 2242129..0000000 --- a/backend/internal/provider/adobe/client.go +++ /dev/null @@ -1,1246 +0,0 @@ -package adobe - -import ( - "bytes" - "context" - "crypto/sha256" - "encoding/hex" - "encoding/json" - "errors" - "fmt" - "io" - "math/rand" - "net/url" - "strings" - "time" - - http "github.com/bogdanfinn/fhttp" - tlsclient "github.com/bogdanfinn/tls-client" - "github.com/bogdanfinn/tls-client/profiles" -) - -const ( - submitURL = "https://firefly-3p.ff.adobe.io/v2/3p-images/generate-async" - image5SubmitURL = "https://image-v5.ff.adobe.io/v1/images/generate-async" - videoSubmitURL = "https://firefly-3p.ff.adobe.io/v2/3p-videos/generate-async" - // Firefly-native video model (project id "firefly-video"): distinct host, - // submit path and storage host from the 3p (veo/luma) video flow. - fireflyVideoSubmitURL = "https://video-v1.ff.adobe.io/v2/videos/generate" - fireflyVideoUploadURL = "https://video-v1.ff.adobe.io/v2/storage/image" - uploadURL = "https://firefly-3p.ff.adobe.io/v2/storage/image" - // Seedance reference video/audio go to their own storage endpoints; posting - // them to /v2/storage/image is rejected. - uploadVideoURL = "https://firefly-3p.ff.adobe.io/v2/storage/video" - uploadAudioURL = "https://firefly-3p.ff.adobe.io/v2/storage/audio" - creditsURL = "https://firefly.adobe.io/v1/credits/balance" - creditsAPIKey = "SunbreakWebUI1" -) - -var ( - ErrAuth = errors.New("adobe auth failed") - ErrQuotaExhausted = errors.New("adobe quota exhausted") - ErrTemporaryUpstream = errors.New("adobe upstream temporary error") - ErrDeadUpstream = errors.New("adobe upstream fatal error") - ErrRateLimited = errors.New("adobe rate limited") - // ErrContentRejected is Adobe's content-safety filter refusing the prompt or - // the generated image (HTTP 451 image_unsafe). It is the prompt's fault, not - // the account's — every account rejects the same content — so the caller must - // surface it as-is without penalizing/killing the account or failing over. - ErrContentRejected = errors.New("adobe content rejected") - // ErrNotEntitled is a 403 user_not_entitled on submit: the account has no - // Firefly entitlement at all. It wraps ErrAuth so existing auth handling still - // applies, but callers can single it out — unlike an expired access token this - // cannot be fixed by refreshing from the cookie, so the account is done. - ErrNotEntitled = fmt.Errorf("%w: user not entitled", ErrAuth) - // errTransport marks a request that never got a response back (EOF / 连接被切 / - // 超时)。上游没收到就没开始生成，所以原地换条连接重试是安全的。 - errTransport = errors.New("transport failure") -) - -// videoSubmitMaxRetries is how many extra in-place attempts a video submit gets -// when the connection dies before any response arrives (EOF/reset/timeout). -const videoSubmitMaxRetries = 3 - -// isContentRejection reports whether an Adobe response (status + body) is a -// content-safety refusal rather than a genuine upstream/account failure. Adobe -// returns HTTP 451 with an "*_unsafe" error_code when moderation blocks the -// prompt or the produced image, and "reference_image_privacy_error" when a -// reference image contains a real person's face. Both are the request's fault — -// every account rejects them identically — so neither may penalize the account. -func isContentRejection(status int, body string) bool { - if status != 451 { - return false - } - return strings.Contains(body, "unsafe") || strings.Contains(body, "privacy_error") -} - -// isClassifierGlitch reports whether a 400 comes from Adobe's media moderation -// backend failing to READ the media it was handed ("Gemini classification API -// returned HTTP 400 ... Cannot fetch content from the provided URL") rather than -// from the request being wrong. The blobs were just uploaded and the identical -// request succeeds on a retry, so this is a transient upstream fault: retrying -// (with freshly uploaded blobs) is the fix, and the account must not be blamed. -func isClassifierGlitch(status int, body string) bool { - if status != 400 { - return false - } - return strings.Contains(body, "media_classification_client_error") || - strings.Contains(body, "classification API returned") -} - -var profileURLs = []string{ - "https://ims-na1.adobelogin.com/ims/profile/v1", - "https://adobeid-na1.services.adobe.com/ims/profile/v1", -} - -type Client struct { - apiKey string - proxy string -} - -func NewClient(apiKey, proxy string) *Client { - return &Client{ - apiKey: defaultString(apiKey, clientID), - proxy: strings.TrimSpace(proxy), - } -} - -// getARPSessionID returns the x-arp-session-id for this token. When a real, -// cookie-derived value was captured during ExchangeCookie (CacheARP) it is -// reused verbatim — it carries a genuine Arkose payload which the 3p endpoints -// require. Otherwise a shaped value is synthesized (see buildARPSessionID). -func (c *Client) getARPSessionID(token string) string { - if id := arpIdentity(token); strings.TrimSpace(id.arp) != "" { - return refreshARPTimestamp(id.arp) - } - return buildARPSessionID(token) -} - -func (c *Client) SetProxy(proxy string) { - c.proxy = strings.TrimSpace(proxy) -} - -func (c *Client) ExchangeCookie(ctx context.Context, cookie string) (*CookieExchangeResult, error) { - sess, err := c.newDirectTLSClient() - if err != nil { - return nil, err - } - return exchangeCookieWithTLSClient(ctx, sess, cookie) -} - -// uploadMaxRetries is how many extra in-place attempts a transient upload -// failure (transport error / timeout, 429/451/5xx) gets on a fresh connection -// before the error is surfaced. -const uploadMaxRetries = 5 - -// UploadImage stores a reference image and returns its blob id. Transient -// failures are retried in place (uploadMaxRetries times); the final error keeps -// its original (non-temporary) classification so the account is not penalized -// for a network blip. -func (c *Client) UploadImage(ctx context.Context, token string, content []byte, contentType, engine string) (string, error) { - // Reference-image upload runs on the local IP (not the proxy). - body, err, retryable := c.uploadImageOnce(ctx, token, content, contentType, engine) - for attempt := 0; err != nil && retryable && attempt < uploadMaxRetries && ctx.Err() == nil; attempt++ { - body, err, retryable = c.uploadImageOnce(ctx, token, content, contentType, engine) - } - if err != nil { - return "", err - } - - var payload map[string]any - if err := json.Unmarshal(body, &payload); err != nil { - return "", fmt.Errorf("adobe upload bad response: %w (%s)", err, clip(body, 300)) - } - if id := blobIDFromUploadPayload(payload); id != "" { - return id, nil - } - return "", fmt.Errorf("adobe upload missing blob id: %s", clip(body, 300)) -} - -// blobIDFromUploadPayload digs the blob id out of a storage response. The image -// endpoint answers {"images":[{"id":...}]}, video/audio use their own key -// ("videos"/"audios"/"assets"/"files"), and some responses put the id at the top -// level — accept any of them instead of only "images". -func blobIDFromUploadPayload(payload map[string]any) string { - if id := strings.TrimSpace(stringValue(payload["id"])); id != "" { - return id - } - for _, key := range []string{"images", "videos", "audios", "audio", "assets", "files"} { - items, ok := payload[key].([]any) - if !ok || len(items) == 0 { - continue - } - first, ok := items[0].(map[string]any) - if !ok { - continue - } - if id := strings.TrimSpace(stringValue(first["id"])); id != "" { - return id - } - } - return "" -} - -// uploadImageOnce performs a single upload attempt and returns the raw response -// body plus whether a failure is retryable (transport error / 429/451/5xx). -// Auth failures (401/403) and other non-200s are not retryable. -func (c *Client) uploadImageOnce(ctx context.Context, token string, content []byte, contentType, engine string) ([]byte, error, bool) { - sess, err := c.newDirectTLSClient() - if err != nil { - return nil, err, false - } - - endpoint := uploadURL - switch { - case engine == "firefly-video": - endpoint = fireflyVideoUploadURL - case strings.HasPrefix(contentType, "video/"): - endpoint = uploadVideoURL - case strings.HasPrefix(contentType, "audio/"): - endpoint = uploadAudioURL - } - // Seedance is a Firefly-native model: its uploads need the clio-playground-web - // key, the express key is rejected with 403 access_error. - apiKey := c.apiKey - if strings.HasPrefix(engine, "seedance-") { - apiKey = clioClientID - } - req, err := http.NewRequest(http.MethodPost, endpoint, bytes.NewReader(content)) - if err != nil { - return nil, err, false - } - req = req.WithContext(ctx) - req.Header = http.Header{ - "authorization": {"Bearer " + strings.TrimSpace(token)}, - "x-api-key": {apiKey}, - "content-type": {defaultString(contentType, "image/png")}, - "accept": {"*/*"}, - "user-agent": {sess.fp.userAgent}, - http.HeaderOrderKey: { - "authorization", - "x-api-key", - "content-type", - "accept", - "user-agent", - }, - } - - resp, err := sess.client.Do(req) - if err != nil { - // Transport error (incl. Client.Timeout on the storage endpoint) — a - // network blip, retryable on a fresh connection. - return nil, fmt.Errorf("adobe upload request: %w", err), true - } - defer resp.Body.Close() - - body, err := io.ReadAll(resp.Body) - if err != nil { - return nil, err, true - } - if resp.StatusCode == 401 || resp.StatusCode == 403 { - return nil, fmt.Errorf("%w (upload %d %s: %s)", ErrAuth, resp.StatusCode, resp.Header.Get("x-access-error"), clip(body, 300)), false - } - if resp.StatusCode == 429 { - return nil, fmt.Errorf("%w (upload %d): %s", ErrRateLimited, resp.StatusCode, clip(body, 300)), false - } - if resp.StatusCode == 451 || resp.StatusCode >= 500 { - return nil, fmt.Errorf("adobe upload failed: %d %s", resp.StatusCode, clip(body, 300)), true - } - if resp.StatusCode != 200 { - return nil, fmt.Errorf("adobe upload failed: %d %s", resp.StatusCode, clip(body, 300)), false - } - return body, nil, true -} - -func (c *Client) GenerateImage(ctx context.Context, token, modelID, prompt, aspectRatio, resolution string, blobIDs []string, downloadResult bool) ([]byte, map[string]any, error) { - // Only the generate submit goes through the proxy; polling + download run on - // the local IP. - submitSess, err := c.newTLSClient() - if err != nil { - return nil, nil, err - } - pollSess, err := c.newDirectTLSClient() - if err != nil { - return nil, nil, err - } - - var lastBody []byte - var lastErr error - // Firefly Image 5 uses a different endpoint + request schema (modelVersion - // "image5", resolutionLevel, top-level aspectRatio label, no modelId/size). - endpoint := submitURL - var candidates []map[string]any - if modelID == "firefly-image-5" { - endpoint = image5SubmitURL - candidates = []map[string]any{buildImage5Payload(prompt, aspectRatio, resolution, blobIDs)} - } else { - candidates = BuildImagePayloadCandidates(modelID, prompt, aspectRatio, resolution, blobIDs) - } - for _, payload := range candidates { - respBody, pollURL, err := c.submitImage(ctx, submitSess, token, prompt, endpoint, payload) - if err == nil { - meta, data, pollErr := c.pollImage(ctx, pollSess, token, pollURL, downloadResult) - if pollErr != nil { - return nil, nil, pollErr - } - return data, meta, nil - } - lastBody = respBody - lastErr = err - if errors.Is(err, ErrAuth) || errors.Is(err, ErrQuotaExhausted) || errors.Is(err, ErrContentRejected) { - return nil, nil, err - } - } - // Content-safety refusal: the prompt/image is blocked, retrying other payloads - // or accounts is pointless — surface it as-is so the pool doesn't fail over. - if errors.Is(lastErr, ErrContentRejected) { - return nil, nil, fmt.Errorf("%w: adobe submit: %s", ErrContentRejected, submitDetail(lastErr, lastBody)) - } - // Preserve the temporary classification so the pool retries (overload / 5xx / - // rate-limit) instead of failing the request outright. - if errors.Is(lastErr, ErrDeadUpstream) { - return nil, nil, fmt.Errorf("%w: adobe submit: %s", ErrDeadUpstream, submitDetail(lastErr, lastBody)) - } - if errors.Is(lastErr, ErrTemporaryUpstream) { - return nil, nil, fmt.Errorf("%w: adobe submit: %s", ErrTemporaryUpstream, submitDetail(lastErr, lastBody)) - } - return nil, nil, fmt.Errorf("adobe submit failed: %s", submitDetail(lastErr, lastBody)) -} - -// submitDetail 拼接上游细节：优先用响应体，响应体为空时（传输层报错，例如超时、 -// 连接被切）退回内层错误文本，避免错误只剩一句 "adobe submit:" 无从排查。 -func submitDetail(err error, body []byte) string { - if s := clip(body, 300); strings.TrimSpace(s) != "" { - return s - } - if err != nil { - return err.Error() - } - return "(empty response)" -} - -// GenerateVideo renders the clip and (when downloadResult) downloads the MP4. -// With downloadResult=false it returns nil bytes and the upstream presigned URL -// in meta["video_url"] — used by the async /v1/videos job, which proxies that URL -// on /content instead of persisting the file. -func (c *Client) GenerateVideo(ctx context.Context, token, engine, prompt, aspectRatio string, durationSeconds int, resolution, referenceMode, upstreamModel string, blobIDs, videoBlobIDs, audioBlobIDs []string, downloadResult bool) ([]byte, map[string]any, error) { - // Only the submit goes through the proxy; polling + download run on the local IP. - submitSess, err := c.newTLSClient() - if err != nil { - return nil, nil, err - } - pollSess, err := c.newDirectTLSClient() - if err != nil { - return nil, nil, err - } - - payload := BuildVideoPayload(engine, prompt, aspectRatio, durationSeconds, resolution, referenceMode, upstreamModel, blobIDs, videoBlobIDs, audioBlobIDs) - endpoint := videoSubmitURL - if engine == "firefly-video" { - endpoint = fireflyVideoSubmitURL - } - // 连接在拿到响应前就断掉（EOF / reset / 超时）说明上游根本没收到这单， - // 换一条新连接原地重试，最多 videoSubmitMaxRetries 次。 - respBody, pollURL, err := c.submitVideo(ctx, submitSess, token, endpoint, payload) - for attempt := 0; err != nil && errors.Is(err, errTransport) && attempt < videoSubmitMaxRetries && ctx.Err() == nil; attempt++ { - select { - case <-ctx.Done(): - return nil, nil, ctx.Err() - case <-time.After(time.Duration(attempt+1) * 2 * time.Second): - } - retrySess, sessErr := c.newTLSClient() - if sessErr != nil { - break - } - submitSess = retrySess - respBody, pollURL, err = c.submitVideo(ctx, submitSess, token, endpoint, payload) - } - if err != nil { - return nil, nil, err - } - _ = respBody - meta, data, pollErr := c.pollVideo(ctx, pollSess, token, pollURL, downloadResult) - if pollErr != nil { - return nil, nil, pollErr - } - return data, meta, nil -} - -func (c *Client) FetchAccountProfile(ctx context.Context, token string) (map[string]any, error) { - token = strings.TrimSpace(token) - if token == "" { - return map[string]any{}, nil - } - sess, err := c.newDirectTLSClient() - if err != nil { - return nil, err - } - - for _, rawURL := range profileURLs { - req, err := http.NewRequest(http.MethodGet, rawURL, nil) - if err != nil { - return nil, err - } - req = req.WithContext(ctx) - req.Header = http.Header{ - "authorization": {"Bearer " + token}, - "accept": {"application/json"}, - "user-agent": {sess.fp.userAgent}, - http.HeaderOrderKey: { - "authorization", - "accept", - "user-agent", - }, - } - - resp, err := sess.client.Do(req) - if err != nil { - continue - } - body, readErr := io.ReadAll(resp.Body) - resp.Body.Close() - if readErr != nil || resp.StatusCode != 200 { - continue - } - - var payload map[string]any - if err := json.Unmarshal(body, &payload); err != nil { - continue - } - email := strings.TrimSpace(stringValue(payload["email"])) - displayName := strings.TrimSpace(stringValue(payload["displayName"])) - if displayName == "" { - displayName = strings.TrimSpace(stringValue(payload["name"])) - } - if displayName == "" { - displayName = strings.TrimSpace(stringValue(payload["fullName"])) - } - userID := strings.TrimSpace(stringValue(payload["userId"])) - if userID == "" { - userID = strings.TrimSpace(stringValue(payload["authId"])) - } - if email != "" || displayName != "" || userID != "" { - return map[string]any{ - "email": emptyStringNil(email), - "display_name": emptyStringNil(displayName), - "user_id": emptyStringNil(userID), - }, nil - } - } - - return map[string]any{}, nil -} - -func (c *Client) FetchCreditsBalance(ctx context.Context, token string) (map[string]any, error) { - token = strings.TrimSpace(token) - if token == "" { - return map[string]any{ - "remaining": nil, - "used": nil, - "total": nil, - "available_until": nil, - "unknown": true, - "error": "empty token", - }, nil - } - - accountID := ExtractAccountID(token) - if accountID == "" { - return map[string]any{ - "remaining": nil, - "used": nil, - "total": nil, - "available_until": nil, - "unknown": true, - "error": "no account id", - }, nil - } - - sess, err := c.newTLSClient() - if err != nil { - return nil, err - } - req, err := http.NewRequest(http.MethodGet, creditsURL, nil) - if err != nil { - return nil, err - } - req = req.WithContext(ctx) - req.Header = http.Header{ - "authorization": {"Bearer " + token}, - "x-api-key": {creditsAPIKey}, - "x-account-id": {accountID}, - "accept": {"application/json"}, - "content-type": {"application/json"}, - "user-agent": {sess.fp.userAgent}, - http.HeaderOrderKey: { - "authorization", - "x-api-key", - "x-account-id", - "accept", - "content-type", - "user-agent", - }, - } - - resp, err := sess.client.Do(req) - if err != nil { - return map[string]any{ - "remaining": nil, - "used": nil, - "total": nil, - "available_until": nil, - "unknown": true, - "error": "network: " + err.Error(), - }, nil - } - defer resp.Body.Close() - - body, err := io.ReadAll(resp.Body) - if err != nil { - return nil, err - } - if resp.StatusCode == 401 { - return nil, ErrAuth - } - if resp.StatusCode != 200 { - return map[string]any{ - "remaining": nil, - "used": nil, - "total": nil, - "available_until": nil, - "unknown": true, - "error": fmt.Sprintf("http %d: %s", resp.StatusCode, clip(body, 160)), - }, nil - } - - var payload map[string]any - if err := json.Unmarshal(body, &payload); err != nil { - return map[string]any{ - "remaining": nil, - "used": nil, - "total": nil, - "available_until": nil, - "unknown": true, - "error": "non-json", - }, nil - } - - totalInfo, _ := payload["total"].(map[string]any) - quota, _ := totalInfo["quota"].(map[string]any) - return map[string]any{ - "remaining": intOrNil(quota["available"]), - "used": intOrNil(quota["used"]), - "total": intOrNil(quota["total"]), - "available_until": emptyStringNil(strings.TrimSpace(stringValue(totalInfo["availableUntil"]))), - "plan": strings.TrimSpace(stringValue(totalInfo["planCap"])), - "unknown": false, - "error": nil, - }, nil -} - -func (c *Client) submitImage(ctx context.Context, sess *tlsSession, token, prompt, endpoint string, payload map[string]any) ([]byte, string, error) { - body, err := json.Marshal(payload) - if err != nil { - return nil, "", err - } - req, err := http.NewRequest(http.MethodPost, endpoint, bytes.NewReader(body)) - if err != nil { - return nil, "", err - } - req = req.WithContext(ctx) - req.Header = http.Header{ - "authorization": {"Bearer " + strings.TrimSpace(token)}, - "x-api-key": {c.apiKey}, - "x-arp-session-id": {c.getARPSessionID(token)}, - "content-type": {"application/json"}, - "accept": {"*/*"}, - "origin": {"https://new.express.adobe.com"}, - "referer": {"https://new.express.adobe.com/"}, - "accept-language": {"en-US,en;q=0.9"}, - "sec-ch-ua": {sess.fp.secCHUA}, - "sec-ch-ua-mobile": {"?0"}, - "sec-ch-ua-platform": {sess.fp.platform}, - "sec-fetch-site": {"cross-site"}, - "sec-fetch-mode": {"cors"}, - "sec-fetch-dest": {"empty"}, - "user-agent": {sess.fp.userAgent}, - http.HeaderOrderKey: { - "authorization", - "x-api-key", - "x-arp-session-id", - "content-type", - "accept", - "origin", - "referer", - "accept-language", - "sec-ch-ua", - "sec-ch-ua-mobile", - "sec-ch-ua-platform", - "sec-fetch-site", - "sec-fetch-mode", - "sec-fetch-dest", - "user-agent", - "x-nonce", - }, - } - if nonce := buildSubmitNonce(token, prompt); nonce != "" { - req.Header.Set("x-nonce", nonce) - } - - resp, err := sess.client.Do(req) - if err != nil { - return nil, "", fmt.Errorf("%w: %v", ErrTemporaryUpstream, err) - } - defer resp.Body.Close() - respBody, err := io.ReadAll(resp.Body) - if err != nil { - return nil, "", err - } - if resp.StatusCode == 401 || resp.StatusCode == 403 { - accessErr := resp.Header.Get("x-access-error") - if strings.EqualFold(accessErr, "taste_exhausted") { - return respBody, "", ErrQuotaExhausted - } - if strings.EqualFold(accessErr, "user_not_entitled") { - return respBody, "", fmt.Errorf("%w (submit %d %s: %s)", ErrNotEntitled, resp.StatusCode, accessErr, clip(respBody, 300)) - } - return respBody, "", fmt.Errorf("%w (submit %d %s: %s)", ErrAuth, resp.StatusCode, accessErr, clip(respBody, 300)) - } - // "system under load" / timeout_error = adobe rate-limit/overload (can come on a - // non-5xx) — treat as temporary so the pool retries instead of failing. - if b := string(respBody); strings.Contains(b, "system under load") || strings.Contains(b, "timeout_error") { - return respBody, "", ErrTemporaryUpstream - } - if isContentRejection(resp.StatusCode, string(respBody)) { - return respBody, "", ErrContentRejected - } - if resp.StatusCode == 429 || resp.StatusCode == 451 || resp.StatusCode >= 500 { - return respBody, "", fmt.Errorf("%w (submit %d: %s)", ErrDeadUpstream, resp.StatusCode, clip(respBody, 300)) - } - if resp.StatusCode != 200 { - return respBody, "", fmt.Errorf("submit rejected: %d %s", resp.StatusCode, clip(respBody, 300)) - } - - var payloadResp map[string]any - if err := json.Unmarshal(respBody, &payloadResp); err != nil { - return respBody, "", err - } - if override := strings.TrimSpace(resp.Header.Get("x-override-status-link")); override != "" { - return respBody, normalizePollURL(override), nil - } - if links, ok := payloadResp["links"].(map[string]any); ok { - if result, ok := links["result"].(map[string]any); ok { - if href := strings.TrimSpace(stringValue(result["href"])); href != "" { - return respBody, normalizePollURL(href), nil - } - } - if href := strings.TrimSpace(stringValue(links["result"])); href != "" { - return respBody, normalizePollURL(href), nil - } - } - return respBody, "", errors.New("submit ok but no poll url") -} - -func (c *Client) pollImage(ctx context.Context, sess *tlsSession, token, pollURL string, downloadResult bool) (map[string]any, []byte, error) { - for { - if err := ctx.Err(); err != nil { - return nil, nil, fmt.Errorf("adobe generation timed out: %w", err) - } - - req, err := http.NewRequest(http.MethodGet, pollURL, nil) - if err != nil { - return nil, nil, err - } - req = req.WithContext(ctx) - req.Header = http.Header{ - "authorization": {"Bearer " + strings.TrimSpace(token)}, - "accept": {"*/*"}, - "origin": {"https://new.express.adobe.com"}, - "referer": {"https://new.express.adobe.com/"}, - "user-agent": {sess.fp.userAgent}, - http.HeaderOrderKey: { - "authorization", - "accept", - "origin", - "referer", - "user-agent", - }, - } - - resp, err := sess.client.Do(req) - if err != nil { - return nil, nil, fmt.Errorf("%w: %v", ErrTemporaryUpstream, err) - } - body, readErr := io.ReadAll(resp.Body) - resp.Body.Close() - if readErr != nil { - return nil, nil, readErr - } - if b := string(body); strings.Contains(b, "system under load") || strings.Contains(b, "timeout_error") { - return nil, nil, ErrTemporaryUpstream - } - if isContentRejection(resp.StatusCode, string(body)) { - return nil, nil, fmt.Errorf("%w: %s", ErrContentRejected, clip(body, 300)) - } - if resp.StatusCode == 429 || resp.StatusCode == 451 || resp.StatusCode >= 500 { - return nil, nil, fmt.Errorf("%w (poll %d: %s)", ErrDeadUpstream, resp.StatusCode, clip(body, 300)) - } - if isClassifierGlitch(resp.StatusCode, string(body)) { - return nil, nil, fmt.Errorf("%w (poll classifier: %s)", ErrTemporaryUpstream, clip(body, 300)) - } - if resp.StatusCode != 200 { - return nil, nil, fmt.Errorf("adobe poll failed: %d %s", resp.StatusCode, clip(body, 300)) - } - - var payload map[string]any - if err := json.Unmarshal(body, &payload); err != nil { - return nil, nil, err - } - if outputs, ok := payload["outputs"].([]any); ok && len(outputs) > 0 { - if first, ok := outputs[0].(map[string]any); ok { - if image, ok := first["image"].(map[string]any); ok { - if url := strings.TrimSpace(stringValue(image["presignedUrl"])); url != "" { - // Expose the presigned URL for callers that want it directly. - payload["image_url"] = url - if !downloadResult { - return payload, nil, nil - } - data, err := c.download(ctx, sess, url) - if err != nil { - return nil, nil, err - } - return payload, data, nil - } - } - } - } - - status := strings.ToUpper(strings.TrimSpace(stringValue(payload["status"]))) - if status == "FAILED" || status == "CANCELLED" || status == "ERROR" { - return nil, nil, fmt.Errorf("adobe job failed: %s", clip(body, 300)) - } - time.Sleep(3 * time.Second) - } -} - -func (c *Client) submitVideo(ctx context.Context, sess *tlsSession, token, endpoint string, payload map[string]any) ([]byte, string, error) { - body, err := json.Marshal(payload) - if err != nil { - return nil, "", err - } - req, err := http.NewRequest(http.MethodPost, endpoint, bytes.NewReader(body)) - if err != nil { - return nil, "", err - } - req = req.WithContext(ctx) - // Seedance is a Firefly-native model and needs the clio-playground-web key; - // the express key is rejected with 403 access_error. The origin stays on the - // Express surface for every engine — the upstream only gates on the key and - // x-arp-session-id, not on origin. - apiKey := c.apiKey - origin := "https://new.express.adobe.com" - if strings.EqualFold(stringValue(payload["modelId"]), "seedance") { - apiKey = clioClientID - } - req.Header = http.Header{ - "authorization": {"Bearer " + strings.TrimSpace(token)}, - "x-api-key": {apiKey}, - "x-arp-session-id": {c.getARPSessionID(token)}, - "content-type": {"application/json"}, - "accept": {"*/*"}, - "origin": {origin}, - "referer": {origin + "/"}, - "accept-language": {"en-US,en;q=0.9"}, - "sec-ch-ua": {sess.fp.secCHUA}, - "sec-ch-ua-mobile": {"?0"}, - "sec-ch-ua-platform": {sess.fp.platform}, - "sec-fetch-site": {"cross-site"}, - "sec-fetch-mode": {"cors"}, - "sec-fetch-dest": {"empty"}, - "user-agent": {sess.fp.userAgent}, - http.HeaderOrderKey: { - "authorization", - "x-api-key", - "x-arp-session-id", - "content-type", - "accept", - "origin", - "referer", - "accept-language", - "sec-ch-ua", - "sec-ch-ua-mobile", - "sec-ch-ua-platform", - "sec-fetch-site", - "sec-fetch-mode", - "sec-fetch-dest", - "user-agent", - "x-nonce", - }, - } - // The working video submit (HAR) carries x-nonce just like the image submit. - if prompt, _ := payload["prompt"].(string); prompt != "" { - if nonce := buildSubmitNonce(token, prompt); nonce != "" { - req.Header.Set("x-nonce", nonce) - } - } - - resp, err := sess.client.Do(req) - if err != nil { - return nil, "", fmt.Errorf("%w: %w: %v", ErrTemporaryUpstream, errTransport, err) - } - defer resp.Body.Close() - - respBody, err := io.ReadAll(resp.Body) - if err != nil { - return nil, "", err - } - if resp.StatusCode == 401 || resp.StatusCode == 403 { - accessErr := resp.Header.Get("x-access-error") - if strings.EqualFold(accessErr, "taste_exhausted") { - return respBody, "", ErrQuotaExhausted - } - if strings.EqualFold(accessErr, "user_not_entitled") { - return respBody, "", fmt.Errorf("%w (%d %s: %s)", ErrNotEntitled, resp.StatusCode, accessErr, clip(respBody, 300)) - } - // Surface Adobe's response body — "adobe auth failed" alone hides whether - // it's a bad token, a missing scope, or a WAF/fingerprint block. - return respBody, "", fmt.Errorf("%w (%d %s: %s)", ErrAuth, resp.StatusCode, accessErr, clip(respBody, 300)) - } - if isContentRejection(resp.StatusCode, string(respBody)) { - return respBody, "", ErrContentRejected - } - if resp.StatusCode == 408 || resp.StatusCode == 429 || resp.StatusCode == 451 || resp.StatusCode >= 500 { - return respBody, "", fmt.Errorf("%w (video submit %d: %s)", ErrDeadUpstream, resp.StatusCode, clip(respBody, 300)) - } - // "system under load" / timeout_error = adobe overload — treat as a temporary - // error so the tempFailover policy moves to the next account (same as the image path). - if b := string(respBody); strings.Contains(b, "system under load") || strings.Contains(b, "timeout_error") { - return respBody, "", ErrTemporaryUpstream - } - if resp.StatusCode != 200 { - return respBody, "", fmt.Errorf("video submit rejected: %d %s", resp.StatusCode, clip(respBody, 300)) - } - - var payloadResp map[string]any - if err := json.Unmarshal(respBody, &payloadResp); err != nil { - return respBody, "", err - } - if override := strings.TrimSpace(resp.Header.Get("x-override-status-link")); override != "" { - return respBody, normalizePollURL(override), nil - } - if links, ok := payloadResp["links"].(map[string]any); ok { - if result, ok := links["result"].(map[string]any); ok { - if href := strings.TrimSpace(stringValue(result["href"])); href != "" { - return respBody, normalizePollURL(href), nil - } - } - if href := strings.TrimSpace(stringValue(links["result"])); href != "" { - return respBody, normalizePollURL(href), nil - } - } - return respBody, "", errors.New("video submit ok but no poll url") -} - -func (c *Client) pollVideo(ctx context.Context, sess *tlsSession, token, pollURL string, downloadResult bool) (map[string]any, []byte, error) { - for { - if err := ctx.Err(); err != nil { - return nil, nil, fmt.Errorf("adobe video generation timed out: %w", err) - } - - req, err := http.NewRequest(http.MethodGet, pollURL, nil) - if err != nil { - return nil, nil, err - } - req = req.WithContext(ctx) - req.Header = http.Header{ - "authorization": {"Bearer " + strings.TrimSpace(token)}, - "accept": {"*/*"}, - "origin": {"https://new.express.adobe.com"}, - "referer": {"https://new.express.adobe.com/"}, - "user-agent": {sess.fp.userAgent}, - http.HeaderOrderKey: { - "authorization", - "accept", - "origin", - "referer", - "user-agent", - }, - } - - resp, err := sess.client.Do(req) - if err != nil { - return nil, nil, fmt.Errorf("%w: %v", ErrTemporaryUpstream, err) - } - body, readErr := io.ReadAll(resp.Body) - resp.Body.Close() - if readErr != nil { - return nil, nil, readErr - } - if resp.StatusCode == 401 || resp.StatusCode == 403 { - return nil, nil, fmt.Errorf("%w (%d %s: %s)", ErrAuth, resp.StatusCode, resp.Header.Get("x-access-error"), clip(body, 300)) - } - if b := string(body); strings.Contains(b, "system under load") || strings.Contains(b, "timeout_error") { - return nil, nil, ErrTemporaryUpstream - } - if isContentRejection(resp.StatusCode, string(body)) { - return nil, nil, fmt.Errorf("%w: %s", ErrContentRejected, clip(body, 300)) - } - if resp.StatusCode == 429 || resp.StatusCode == 451 || resp.StatusCode >= 500 { - return nil, nil, fmt.Errorf("%w (video poll %d: %s)", ErrDeadUpstream, resp.StatusCode, clip(body, 300)) - } - if isClassifierGlitch(resp.StatusCode, string(body)) { - return nil, nil, fmt.Errorf("%w (video poll classifier: %s)", ErrTemporaryUpstream, clip(body, 300)) - } - if resp.StatusCode != 200 { - return nil, nil, fmt.Errorf("adobe video poll failed: %d %s", resp.StatusCode, clip(body, 300)) - } - - var payload map[string]any - if err := json.Unmarshal(body, &payload); err != nil { - return nil, nil, err - } - if outputs, ok := payload["outputs"].([]any); ok && len(outputs) > 0 { - if first, ok := outputs[0].(map[string]any); ok { - if video, ok := first["video"].(map[string]any); ok { - if raw := strings.TrimSpace(stringValue(video["presignedUrl"])); raw != "" { - payload["video_url"] = raw - if !downloadResult { - return payload, nil, nil - } - data, err := c.download(ctx, sess, raw) - if err != nil { - return nil, nil, err - } - return payload, data, nil - } - } - } - } - - status := strings.ToUpper(strings.TrimSpace(stringValue(payload["status"]))) - if status == "FAILED" || status == "CANCELLED" || status == "ERROR" { - return nil, nil, fmt.Errorf("adobe video job failed: %s", clip(body, 300)) - } - time.Sleep(3 * time.Second) - } -} - -// download fetches the generated artifact from Adobe's presigned S3 URL. The -// asset is already produced at this point, so a transient network hiccup -// (connection EOF/reset mid-body, S3 accelerate 5xx) must not fail the whole -// generation — retry with backoff before giving up. -const downloadTimeout = 3 * time.Minute - -func (c *Client) download(parent context.Context, sess *tlsSession, url string) ([]byte, error) { - // The artifact is already produced; detach from the (often nearly-elapsed) - // generation ctx deadline/cancel and give the download its own budget so a - // slow CDN read isn't killed mid-body and the backoff-retries can actually run. - ctx, cancel := context.WithTimeout(context.WithoutCancel(parent), downloadTimeout) - defer cancel() - var data []byte - var err error - for _, wait := range []time.Duration{0, 1 * time.Second, 2 * time.Second, 5 * time.Second, 10 * time.Second} { - if wait > 0 { - select { - case <-ctx.Done(): - return nil, ctx.Err() - case <-time.After(wait): - } - } - data, err = c.downloadOnce(ctx, sess, url) - if err == nil || !errors.Is(err, ErrTemporaryUpstream) { - break - } - } - return data, err -} - -func (c *Client) downloadOnce(ctx context.Context, sess *tlsSession, 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{ - "accept": {"*/*"}, - "user-agent": {sess.fp.userAgent}, - http.HeaderOrderKey: { - "accept", - "user-agent", - }, - } - resp, err := sess.client.Do(req) - if err != nil { - // Network-level failure (EOF, reset, timeout) — transient, worth a retry. - return nil, fmt.Errorf("%w: %v", ErrTemporaryUpstream, err) - } - defer resp.Body.Close() - if resp.StatusCode != 200 { - body, _ := io.ReadAll(resp.Body) - if resp.StatusCode >= 500 { - return nil, fmt.Errorf("%w: adobe download failed: %d %s", ErrTemporaryUpstream, resp.StatusCode, clip(body, 200)) - } - return nil, fmt.Errorf("adobe download failed: %d %s", resp.StatusCode, clip(body, 200)) - } - data, err := io.ReadAll(resp.Body) - if err != nil { - // Body truncated mid-read (EOF) — transient, retry. - return nil, fmt.Errorf("%w: read body: %v", ErrTemporaryUpstream, err) - } - return data, nil -} - -// fingerprint bundles a TLS ClientProfile with the browser headers that match -// it, so the JA3/JA4 handshake and the advertised User-Agent / client hints -// stay internally consistent within a single request. -type fingerprint struct { - profile profiles.ClientProfile - userAgent string - secCHUA string - platform string // quoted sec-ch-ua-platform value, e.g. `"Windows"` -} - -const ( - osWindows = iota - osMac -) - -// newChromeFingerprint pairs a TLS ClientProfile with headers for the same -// Chrome major. secCHUA is passed in verbatim rather than templated: the GREASE -// brand entry ("Not(A:Brand", "Not_A Brand", …), its version, and its position -// in the list all rotate per Chrome release, so there is no formula to derive -// it from the major. -func newChromeFingerprint(profile profiles.ClientProfile, major, osKind int, secCHUA string) fingerprint { - ua := fmt.Sprintf("Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/%d.0.0.0 Safari/537.36", major) - platform := `"Windows"` - if osKind == osMac { - ua = fmt.Sprintf("Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/%d.0.0.0 Safari/537.36", major) - platform = `"macOS"` - } - return fingerprint{ - profile: profile, - userAgent: ua, - secCHUA: secCHUA, - platform: platform, - } -} - -// adobeFingerprints is a small pool of recent desktop Chrome builds across -// Windows/macOS. Each request draws one at random so the outbound TLS -// fingerprint and headers vary from call to call instead of presenting a single -// static signature to Adobe's anti-bot layer. -// -// The majors here are capped by what tls-client ships a TLS profile for -// (Chrome_146 is the newest in v1.15.1). Live Adobe traffic is on Chromium -// 150+, so the UA is deliberately NOT set higher than the pool: advertising a -// newer UA than the JA3/JA4 handshake backs up is a sharper signal than being a -// few releases behind. tls-client has no Edge profile at any version, so the -// pool stays pure Chrome rather than pairing an Edge UA with a Chrome -// handshake. -var adobeFingerprints = []fingerprint{ - newChromeFingerprint(profiles.Chrome_146, 146, osWindows, `"Chromium";v="146", "Google Chrome";v="146", "Not?A_Brand";v="24"`), - newChromeFingerprint(profiles.Chrome_146, 146, osMac, `"Chromium";v="146", "Google Chrome";v="146", "Not?A_Brand";v="24"`), - newChromeFingerprint(profiles.Chrome_144, 144, osWindows, `"Chromium";v="144", "Google Chrome";v="144", "Not?A_Brand";v="24"`), - newChromeFingerprint(profiles.Chrome_144, 144, osMac, `"Chromium";v="144", "Google Chrome";v="144", "Not?A_Brand";v="24"`), - newChromeFingerprint(profiles.Chrome_133, 133, osWindows, `"Not(A:Brand";v="99", "Google Chrome";v="133", "Chromium";v="133"`), - newChromeFingerprint(profiles.Chrome_131, 131, osMac, `"Google Chrome";v="131", "Chromium";v="131", "Not_A Brand";v="24"`), -} - -func randomFingerprint() fingerprint { - return adobeFingerprints[rand.Intn(len(adobeFingerprints))] -} - -// tlsSession pairs a configured TLS client with the fingerprint it was built -// for, so downstream header builders advertise the matching UA / client hints. -type tlsSession struct { - client tlsclient.HttpClient - fp fingerprint -} - -// newTLSClient builds a session that routes through the configured proxy (when -// set). newDirectTLSClient builds one that always uses the local IP. Only the -// image-generation submit goes through the proxy; reference-image upload, -// polling and result download run on the local IP. -func (c *Client) newTLSClient() (*tlsSession, error) { - return c.newTLSSession(randomFingerprint(), true) -} - -func (c *Client) newDirectTLSClient() (*tlsSession, error) { - return c.newTLSSession(randomFingerprint(), false) -} - -func (c *Client) newTLSSession(fp fingerprint, useProxy bool) (*tlsSession, error) { - options := []tlsclient.HttpClientOption{ - tlsclient.WithTimeoutSeconds(60), - tlsclient.WithClientProfile(fp.profile), - tlsclient.WithNotFollowRedirects(), - tlsclient.WithRandomTLSExtensionOrder(), - } - if useProxy && c.proxy != "" { - options = append(options, tlsclient.WithProxyUrl(c.proxy)) - } - client, err := tlsclient.NewHttpClient(tlsclient.NewNoopLogger(), options...) - if err != nil { - return nil, err - } - return &tlsSession{client: client, fp: fp}, nil -} - -func exchangeCookieWithTLSClient(ctx context.Context, sess *tlsSession, cookie string) (*CookieExchangeResult, error) { - cookie = normalizeCookie(cookie) - if cookie == "" { - return nil, ErrAdobeCookieEmpty - } - - body := "client_id=" + clientID + "&guest_allowed=true&scope=" + strings.ReplaceAll(scopeValue, ",", "%2C") - payload, err := doCookieExchange(ctx, sess, cookie, body, "https://new.express.adobe.com") - if err != nil { - return nil, err - } - token := strings.TrimSpace(stringValue(payload["access_token"])) - if token == "" { - return nil, errors.New("adobe cookie exchange missing access_token") - } - // The Express token (projectx_webapp) covers image/veo/luma/sora but is not - // entitled to the Firefly-native seedance models — those return 403 - // model_not_authorized. firefly.adobe.com mints a second token for - // client_id=clio-playground-web with the creative_production scope, keyed by - // user_id, which IS entitled. Mint that here too (best-effort: if it fails we - // keep the Express token so the other models still work). - // A cookie without the IMS session (ims_sid) mints an anonymous guest - // token (no user_id) — generation rejects it with 401, so fail the - // exchange here instead of landing a dead account in the pool. - uid := ExtractAccountID(token) - if uid == "" { - return nil, errors.New("adobe cookie has no IMS session (ims_sid); minted token is anonymous") - } - arp := ARPFromCookie(cookie) - CacheARP(token, arp) - if uid != "" { - clioBody := "client_id=" + clioClientID + "&scope=" + strings.ReplaceAll(clioScopeValue, ",", "%2C") + "&user_id=" + url.QueryEscape(uid) - if clioPayload, err := doCookieExchange(ctx, sess, cookie, clioBody, "https://firefly.adobe.com"); err == nil { - if clioToken := strings.TrimSpace(stringValue(clioPayload["access_token"])); clioToken != "" { - CacheARP(clioToken, arp) - return &CookieExchangeResult{ - AccessToken: clioToken, - ExpiresIn: intValue(clioPayload["expires_in"]), - Raw: clioPayload, - }, nil - } - } - } - return &CookieExchangeResult{ - AccessToken: token, - ExpiresIn: intValue(payload["expires_in"]), - Raw: payload, - }, nil -} - -func doCookieExchange(ctx context.Context, sess *tlsSession, cookie, body, origin string) (map[string]any, error) { - req, err := http.NewRequest(http.MethodPost, refreshURL, strings.NewReader(body)) - if err != nil { - return nil, err - } - req = req.WithContext(ctx) - req.Header = http.Header{ - "accept": {"*/*"}, - "accept-language": {"en-US,en;q=0.9"}, - "content-type": {"application/x-www-form-urlencoded;charset=UTF-8"}, - "cookie": {cookie}, - "origin": {origin}, - "referer": {origin + "/"}, - "user-agent": {sess.fp.userAgent}, - http.HeaderOrderKey: { - "accept", "accept-language", "content-type", "cookie", - "origin", "referer", "user-agent", - }, - } - resp, err := sess.client.Do(req) - if err != nil { - return nil, fmt.Errorf("adobe cookie exchange network error: %w", err) - } - defer resp.Body.Close() - - respBody, err := io.ReadAll(resp.Body) - if err != nil { - return nil, err - } - if resp.StatusCode != 200 { - return nil, fmt.Errorf("adobe cookie exchange upstream %d: %s", resp.StatusCode, clip(respBody, 200)) - } - var payload map[string]any - if err := json.Unmarshal(respBody, &payload); err != nil { - return nil, fmt.Errorf("adobe cookie exchange invalid json: %w", err) - } - return payload, nil -} - -func buildSubmitNonce(token, prompt string) string { - claims := decodeJWTPayload(token) - userID := strings.TrimSpace(stringValue(claims["user_id"])) - if userID == "" { - userID = strings.TrimSpace(stringValue(claims["aa_id"])) - } - if userID == "" { - userID = strings.TrimSpace(stringValue(claims["sub"])) - } - prompt = strings.TrimSpace(prompt) - if userID == "" || prompt == "" { - return "" - } - if len(prompt) > 256 { - prompt = prompt[:256] - } - sum := sha256.Sum256([]byte(userID + "-" + prompt)) - return hex.EncodeToString(sum[:]) -} - -func ExtractAccountID(token string) string { - claims := decodeJWTPayload(token) - userID := strings.TrimSpace(stringValue(claims["user_id"])) - if userID == "" { - userID = strings.TrimSpace(stringValue(claims["aa_id"])) - } - if userID == "" { - userID = strings.TrimSpace(stringValue(claims["sub"])) - } - return userID -} - -func normalizePollURL(raw string) string { - if strings.TrimSpace(raw) == "" { - return raw - } - parsed, err := url.Parse(raw) - if err != nil { - return raw - } - host := parsed.Hostname() - if !strings.HasPrefix(host, "firefly-epo") { - return raw - } - parts := strings.Split(strings.Trim(parsed.Path, "/"), "/") - if len(parts) == 0 { - return raw - } - jobID := strings.TrimSpace(parts[len(parts)-1]) - hostSuffix := strings.TrimPrefix(host, "firefly-epo") - hostSuffix = strings.SplitN(hostSuffix, ".", 2)[0] - if len(hostSuffix) != 4 { - return raw - } - for _, ch := range hostSuffix { - if ch < '0' || ch > '9' { - return raw - } - } - return "https://bks-epo" + hostSuffix + ".adobe.io/v2/jobs/result/" + jobID + "?host=" + parsed.Host + "/" -} - -func clip(v []byte, n int) string { - s := strings.TrimSpace(string(v)) - if len(s) <= n { - return s - } - return s[:n] -} diff --git a/backend/internal/provider/adobe/payloads.go b/backend/internal/provider/adobe/payloads.go deleted file mode 100644 index 844661a..0000000 --- a/backend/internal/provider/adobe/payloads.go +++ /dev/null @@ -1,517 +0,0 @@ -package adobe - -import ( - "encoding/json" - "math/rand" - "strings" -) - -type modelSpec struct { - UpstreamModelID string - UpstreamModelVersion string -} - -var lumaSize = map[string]map[string][2]int{ - "720p": { - "21:9": {1280, 548}, "16:9": {1280, 720}, "4:3": {960, 720}, - "1:1": {720, 720}, "3:4": {720, 960}, "9:16": {720, 1280}, "9:21": {548, 1280}, - }, - "1080p": { - "21:9": {1920, 822}, "16:9": {1920, 1080}, "4:3": {1440, 1080}, - "1:1": {1080, 1080}, "3:4": {1080, 1440}, "9:16": {1080, 1920}, "9:21": {822, 1920}, - }, - "4k": { - "21:9": {3840, 1646}, "16:9": {3840, 2160}, "4:3": {2880, 2160}, - "1:1": {2160, 2160}, "3:4": {2160, 2880}, "9:16": {2160, 3840}, "9:21": {1646, 3840}, - }, -} - -var gptImageSize = map[string]map[string][2]int{ - "1K": {"1:1": {1024, 1024}, "5:4": {1120, 896}, "9:16": {720, 1280}, "21:9": {1456, 624}, "16:9": {1280, 720}, "4:3": {1152, 864}, "3:2": {1248, 832}, "4:5": {896, 1120}, "3:4": {864, 1152}, "2:3": {832, 1248}}, - "2K": {"1:1": {2048, 2048}, "5:4": {2240, 1792}, "9:16": {1440, 2560}, "21:9": {3024, 1296}, "16:9": {2560, 1440}, "4:3": {2304, 1728}, "3:2": {2496, 1664}, "4:5": {1792, 2240}, "3:4": {1728, 2304}, "2:3": {1664, 2496}}, - "4K": {"1:1": {2880, 2880}, "5:4": {3200, 2560}, "9:16": {2160, 3840}, "21:9": {3696, 1584}, "16:9": {3840, 2160}, "4:3": {3264, 2448}, "3:2": {3504, 2336}, "4:5": {2560, 3200}, "3:4": {2448, 3264}, "2:3": {2336, 3504}}, -} - -var fluxSize = map[string][2]int{ - "1:1": {1024, 1024}, - "16:9": {1408, 768}, - "9:16": {768, 1408}, - "4:3": {1280, 896}, - "3:4": {896, 1280}, -} - -var defaultSize = map[string]map[string][2]int{ - "1K": {"1:1": {1024, 1024}, "1:8": {384, 3072}, "1:4": {512, 2048}, "16:9": {1360, 768}, "9:16": {768, 1360}, "4:1": {2048, 512}, "4:3": {1152, 864}, "3:4": {864, 1152}, "8:1": {3072, 384}}, - "2K": {"1:1": {2048, 2048}, "1:8": {768, 6144}, "1:4": {1024, 4096}, "16:9": {2752, 1536}, "9:16": {1536, 2752}, "4:1": {4096, 1024}, "4:3": {2048, 1536}, "3:4": {1536, 2048}, "8:1": {6144, 768}}, - "4K": {"1:1": {4096, 4096}, "1:8": {1536, 12288}, "1:4": {2048, 8192}, "16:9": {5504, 3072}, "9:16": {3072, 5504}, "4:1": {8192, 2048}, "4:3": {4096, 3072}, "3:4": {3072, 4096}, "8:1": {12288, 1536}}, -} - -func ResolveModelSpec(modelID string) modelSpec { - switch modelID { - case "firefly-gpt-image", "firefly-gpt-image-2": - return modelSpec{UpstreamModelID: "gpt-image", UpstreamModelVersion: "2"} - case "flux-kontext-max": - return modelSpec{UpstreamModelID: "flux", UpstreamModelVersion: "fluxKontextMax"} - case "nano-banana-2": - return modelSpec{UpstreamModelID: "gemini-flash", UpstreamModelVersion: "nano-banana-3"} - case "nano-banana-pro": - return modelSpec{UpstreamModelID: "gemini-flash", UpstreamModelVersion: "nano-banana-2"} - default: - return modelSpec{UpstreamModelID: "gemini-flash", UpstreamModelVersion: "gemini-3-pro-image-preview"} - } -} - -// buildImage5Payload builds the Adobe Firefly Image 5 request. It uses a distinct -// schema from the firefly-3p models: NO modelId/size, a top-level aspectRatio -// string label and a resolutionLevel (1K→1MP, 2K→4MP). Mirrors a captured -// working image-v5.ff.adobe.io request. -func buildImage5Payload(prompt, aspectRatio, resolution string, blobIDs []string) map[string]any { - p := map[string]any{ - "n": 1, - "seeds": []int{rand.Intn(999999)}, - "output": map[string]any{"storeInputs": true}, - "prompt": prompt, - "referenceBlobs": []any{}, - "modelSpecificPayload": map[string]any{"locale": "en-US", "prompt_reasoner": "quality"}, - "modelVersion": "image5", - "resolutionLevel": image5ResolutionLevel(resolution), - "generationMetadata": map[string]any{"module": "text2image", "submodule": "ff-image-generate"}, - } - if len(blobIDs) > 0 { - // Instruct-edit: aspect ratio is derived from the reference image; sending - // aspectRatio is rejected with a validation_error. - p["referenceBlobs"] = blobRefs(blobIDs, "general") - } else { - p["aspectRatio"] = defaultString(aspectRatio, "1:1") - } - return p -} - -// image5ResolutionLevel maps the UI resolution tier to Image 5's megapixel level. -func image5ResolutionLevel(resolution string) string { - switch strings.ToUpper(strings.TrimSpace(resolution)) { - case "1K": - return "1MP" - case "2K": - return "4MP" - default: - return "4MP" - } -} - -func BuildImagePayloadCandidates(modelID, prompt, aspectRatio, outputResolution string, blobIDs []string) []map[string]any { - spec := ResolveModelSpec(modelID) - ratio := defaultString(aspectRatio, "1:1") - resolution := defaultString(outputResolution, "2K") - - switch spec.UpstreamModelID { - case "gpt-image": - return buildGPTImagePayloads(spec, prompt, ratio, resolution, blobIDs) - case "flux": - return buildFluxPayloads(spec, prompt, ratio, blobIDs) - default: - return buildDefaultPayloads(spec, prompt, ratio, resolution, blobIDs) - } -} - -func buildGPTImagePayloads(spec modelSpec, prompt, ratio, resolution string, blobIDs []string) []map[string]any { - size := getSize(gptImageSize, resolution, ratio, "1:1") - // Mirrors the captured working gpt-image request shape: modelSpecificPayload.size, - // generationSettings.detailLevel 3, and NO top-level size / outputResolution - // (sending those got 403). Keeps the chosen size via modelSpecificPayload.size - // ("WxH") rather than "auto". - base := map[string]any{ - "modelId": spec.UpstreamModelID, - "modelVersion": spec.UpstreamModelVersion, - "n": 1, - "prompt": prompt, - "seeds": []int{rand.Intn(999999)}, - "output": map[string]any{"storeInputs": true}, - "referenceBlobs": []any{}, - "generationMetadata": map[string]any{"module": "text2image", "submodule": "ff-image-generate"}, - "modelSpecificPayload": map[string]any{"size": sizeString(size)}, - "generationSettings": map[string]any{"detailLevel": 3}, - } - if len(blobIDs) == 0 { - return []map[string]any{base} - } - subject := cloneMap(base) - subject["referenceBlobs"] = blobRefs(blobIDs, "subject") - return []map[string]any{subject} -} - -func buildFluxPayloads(spec modelSpec, prompt, ratio string, blobIDs []string) []map[string]any { - size := fluxSize[ratio] - if size == [2]int{} { - size = fluxSize["1:1"] - } - base := map[string]any{ - "modelId": spec.UpstreamModelID, - "modelVersion": spec.UpstreamModelVersion, - "n": 1, - "prompt": prompt, - "size": map[string]any{"width": size[0], "height": size[1]}, - "seeds": []int{rand.Intn(999999)}, - "output": map[string]any{"storeInputs": true}, - "referenceBlobs": []any{}, - "modelSpecificPayload": map[string]any{ - "prompt_upsampling": true, - "safety_tolerance": 2, - "aspect_ratio": ratio, - }, - "generationMetadata": map[string]any{"module": "text2image", "submodule": "ff-image-generate"}, - } - if len(blobIDs) == 0 { - return []map[string]any{base} - } - edited := cloneMap(base) - edited["generationMetadata"] = map[string]any{"module": "image2image", "submodule": "ff-image-generate"} - edited["referenceBlobs"] = blobRefs(blobIDs, "general") - return []map[string]any{edited} -} - -func buildDefaultPayloads(spec modelSpec, prompt, ratio, resolution string, blobIDs []string) []map[string]any { - size := getSize(defaultSize, resolution, ratio, "16:9") - // Shape mirrors a captured working firefly.adobe.com request exactly: top-level - // size object, modelSpecificPayload only {parameters:{addWatermark:false}}, - // groundSearch:false, module "text2image" (even with a reference blob). NO - // skipCai and NO modelSpecificPayload.aspectRatio — sending those got 403. - base := map[string]any{ - "modelId": spec.UpstreamModelID, - "modelVersion": spec.UpstreamModelVersion, - "n": 1, - "prompt": prompt, - "size": map[string]any{"width": size[0], "height": size[1]}, - "seeds": []int{rand.Intn(999999)}, - "groundSearch": false, - "output": map[string]any{"storeInputs": true}, - "generationMetadata": map[string]any{ - "module": "text2image", - "submodule": "ff-image-generate", - }, - "modelSpecificPayload": map[string]any{ - "parameters": map[string]any{"addWatermark": false}, - }, - } - if len(blobIDs) == 0 { - base["referenceBlobs"] = []any{} - return []map[string]any{base} - } - edited := cloneMap(base) - edited["referenceBlobs"] = blobRefs(blobIDs, "general") - return []map[string]any{edited} -} - -func getSize(table map[string]map[string][2]int, resolution, ratio, fallbackRatio string) [2]int { - level := defaultString(resolution, "2K") - levelTable, ok := table[level] - if !ok { - levelTable = table["2K"] - } - size, ok := levelTable[ratio] - if !ok { - size = levelTable[fallbackRatio] - } - return size -} - -func sizeString(size [2]int) string { - return itoa(size[0]) + "x" + itoa(size[1]) -} - -func blobRefs(ids []string, usage string) []any { - out := make([]any, 0, len(ids)) - for _, id := range ids { - out = append(out, map[string]any{"id": id, "usage": usage}) - } - return out -} - -func referenceImagesByID(ids []string) []any { - out := make([]any, 0, len(ids)) - for _, id := range ids { - out = append(out, map[string]any{"id": id}) - } - return out -} - -func referenceImagesByLocal(ids []string) []any { - out := make([]any, 0, len(ids)) - for _, id := range ids { - out = append(out, map[string]any{"localBlobRef": id}) - } - return out -} - -func cloneMap(in map[string]any) map[string]any { - out := make(map[string]any, len(in)) - for k, v := range in { - out[k] = v - } - return out -} - -func BuildVideoPayload(engine, prompt, aspectRatio string, durationSeconds int, resolution, referenceMode, upstreamModel string, blobIDs, videoBlobIDs, audioBlobIDs []string) map[string]any { - seedVal := rand.Intn(999999) - engine = defaultString(engine, "sora2") - resolution = defaultString(resolution, "720p") - aspectRatio = defaultString(aspectRatio, "16:9") - if durationSeconds <= 0 { - durationSeconds = 5 - } - - switch engine { - case "firefly-video": - // Firefly-native video model — a distinct schema (mirrors a captured - // working video-v1.ff.adobe.io request): sizes[] carries width/height + - // numFrames (numFrames encodes duration, ~25.6fps so 5s = 128), and - // reference frames go under image.conditions with placement.start - // (0 = first frame / 首帧, 1 = last frame / 末帧). NO modelId / version / - // engine / duration / referenceBlobs fields. - w, h, frames := fireflyVideoSize(aspectRatio, resolution, durationSeconds) - payload := map[string]any{ - "addOnTransparentBackground": false, - "prompt": prompt, - "seeds": []int{seedVal}, - "sizes": []any{map[string]any{"width": w, "height": h, "numFrames": frames}}, - "videoSettings": map[string]any{}, - "locale": "en-US", - "generationMetadata": map[string]any{"module": "text2video", "submodule": "ff-video-generate"}, - "output": map[string]any{"storeInputs": true}, - } - if len(blobIDs) > 0 { - conds := make([]any, 0, 2) - conds = append(conds, map[string]any{ - "source": map[string]any{"id": blobIDs[0]}, - "placement": map[string]any{"start": 0}, - }) - if len(blobIDs) > 1 { - conds = append(conds, map[string]any{ - "source": map[string]any{"id": blobIDs[1]}, - "placement": map[string]any{"start": 1}, - }) - } - payload["image"] = map[string]any{"conditions": conds} - } - return payload - case "veo31-fast", "veo31-standard": - modelVersion := "3.1-fast-generate" - if engine == "veo31-standard" { - modelVersion = "3.1-generate" - } - payload := map[string]any{ - "modelId": "veo", - "modelVersion": modelVersion, - "size": videoSize(aspectRatio, resolution), - "seeds": []int{seedVal}, - "prompt": prompt, - "negativePrompt": "", - "duration": durationSeconds, - "generateAudio": false, - "generationMetadata": map[string]any{ - "module": "text2video", - "submodule": "ff-video-generate", - }, - "output": map[string]any{"storeInputs": true}, - "referenceBlobs": []any{}, - } - if len(blobIDs) > 0 { - payload["generationMetadata"] = map[string]any{"module": "image2video", "submodule": "ff-video-generate"} - refs := make([]any, 0, min(len(blobIDs), 3)) - if referenceMode == "frame" { - for idx, id := range blobIDs[:min(len(blobIDs), 3)] { - refs = append(refs, map[string]any{"id": id, "usage": "frame", "order": idx + 1}) - } - } else { - for _, id := range blobIDs[:min(len(blobIDs), 3)] { - refs = append(refs, map[string]any{"id": id, "usage": "asset"}) - } - } - payload["referenceBlobs"] = refs - } - return payload - case "seedance-2.0-fast", "seedance-2.0": - modelVersion := "seedance_2.0" - if engine == "seedance-2.0-fast" { - modelVersion = "seedance_2.0_fast" - } - payload := map[string]any{ - "modelId": "seedance", - "modelVersion": modelVersion, - "size": videoSize(aspectRatio, resolution), - "seeds": []int{seedVal}, - "prompt": prompt, - "negativePrompt": "", - "duration": durationSeconds, - "generateAudio": true, - "generationSettings": map[string]any{"aspectRatio": aspectRatio}, - "generationMetadata": map[string]any{"module": "text2video", "submodule": "ff-video-generate"}, - "output": map[string]any{"storeInputs": true}, - "referenceBlobs": []any{}, - } - if len(blobIDs) > 0 { - // Seedance keeps module "text2video" even with reference blobs - // (HAR-confirmed) — unlike veo/luma it does not switch to image2video. - if referenceMode == "frame" { - refs := make([]any, 0, min(len(blobIDs), 2)) - for idx, id := range blobIDs[:min(len(blobIDs), 2)] { - refs = append(refs, map[string]any{"id": id, "usage": "frame", "order": idx + 1}) - } - payload["referenceBlobs"] = refs - } else { - refs := make([]any, 0, min(len(blobIDs), 9)) - for _, id := range blobIDs[:min(len(blobIDs), 9)] { - refs = append(refs, map[string]any{"id": id, "usage": "style"}) - } - payload["referenceBlobs"] = refs - } - } - // Video/audio source refs (seedance) - for _, id := range videoBlobIDs[:min(len(videoBlobIDs), 3)] { - payload["referenceBlobs"] = append(payload["referenceBlobs"].([]any), map[string]any{"id": id, "usage": "source"}) - } - for _, id := range audioBlobIDs[:min(len(audioBlobIDs), 3)] { - payload["referenceBlobs"] = append(payload["referenceBlobs"].([]any), map[string]any{"id": id, "usage": "source"}) - } - return payload - case "luma": - payload := map[string]any{ - "modelId": "luma", - "modelVersion": "3.14-ray", - "size": lumaVideoSize(aspectRatio, resolution), - "mode": "flex_2", - "prompt": prompt, - "negativePrompt": "", - "duration": durationSeconds, - "generationMetadata": map[string]any{ - "module": "text2video", - "submodule": "ff-video-generate", - }, - "modelSpecificPayload": map[string]any{ - "resolution": strings.ToLower(resolution), - "aspect_ratio": aspectRatio, - }, - "output": map[string]any{"storeInputs": true}, - } - if len(blobIDs) > 0 { - payload["generationMetadata"] = map[string]any{ - "module": "image2video", - "submodule": "ff-video-generate", - } - refs := make([]any, 0, min(len(blobIDs), 2)) - for idx, id := range blobIDs[:min(len(blobIDs), 2)] { - refs = append(refs, map[string]any{"id": id, "usage": "frame", "order": idx + 1}) - } - payload["referenceBlobs"] = refs - } - return payload - default: - upstream := defaultString(upstreamModel, "openai:firefly:colligo:sora2") - payload := map[string]any{ - "n": 1, - "seeds": []int{seedVal}, - "modelId": "sora", - "modelVersion": "sora-2", - "size": videoSize(aspectRatio, resolution), - "duration": durationSeconds, - "fps": 24, - "prompt": buildVideoPromptJSON(prompt, durationSeconds), - "generationMetadata": map[string]any{"module": "text2video"}, - "model": upstream, - "generateAudio": true, - "generateLoop": false, - "transparentBackground": false, - "seed": itoa(seedVal), - "locale": "en-US", - "camera": map[string]any{"angle": "none", "shotSize": "none", "motion": nil, "promptStyle": nil}, - "negativePrompt": "", - "jobMode": "standard", - "debugGenerationEndpoint": "", - "referenceBlobs": []any{}, - "referenceFrames": []any{}, - "referenceVideo": nil, - "cameraMotionReferenceVideo": nil, - "characterReference": nil, - "editReferenceVideo": nil, - "output": map[string]any{"storeInputs": true}, - } - if len(blobIDs) > 0 { - firstID := blobIDs[0] - payload["generationMetadata"] = map[string]any{"module": "image2video"} - payload["referenceBlobs"] = []any{ - map[string]any{"id": firstID, "usage": "general", "promptReference": 1}, - } - payload["referenceFrames"] = []any{map[string]any{"localBlobRef": firstID}, nil} - } - return payload - } -} - -// fireflyVideoSizeTable maps the firefly-video resolution tier + aspect ratio to -// pixel dimensions. Only 1080p 9:16 (1080x1920) is HAR-confirmed; the rest follow -// the standard 540p/720p/1080p grid for each ratio. -var fireflyVideoSizeTable = map[string]map[string][2]int{ - "540p": {"16:9": {960, 540}, "1:1": {540, 540}, "9:16": {540, 960}}, - "720p": {"16:9": {1280, 720}, "1:1": {720, 720}, "9:16": {720, 1280}}, - "1080p": {"16:9": {1920, 1080}, "1:1": {1080, 1080}, "9:16": {1080, 1920}}, -} - -// fireflyVideoSize returns width, height and numFrames. numFrames encodes the -// clip length (~25.6fps; 5s = 128 frames, HAR-confirmed). -func fireflyVideoSize(aspectRatio, resolution string, durationSeconds int) (int, int, int) { - table, ok := fireflyVideoSizeTable[strings.ToLower(defaultString(resolution, "1080p"))] - if !ok { - table = fireflyVideoSizeTable["1080p"] - } - wh, ok := table[defaultString(aspectRatio, "9:16")] - if !ok { - wh = table["9:16"] - } - frames := durationSeconds * 128 / 5 - if frames <= 0 { - frames = 128 - } - return wh[0], wh[1], frames -} - -func videoSize(aspectRatio, resolution string) map[string]any { - if strings.EqualFold(resolution, "1080p") { - if aspectRatio == "16:9" { - return map[string]any{"width": 1920, "height": 1080} - } - return map[string]any{"width": 1080, "height": 1920} - } - if aspectRatio == "16:9" { - return map[string]any{"width": 1280, "height": 720} - } - return map[string]any{"width": 720, "height": 1280} -} - -func lumaVideoSize(aspectRatio, resolution string) map[string]any { - table, ok := lumaSize[strings.ToLower(defaultString(resolution, "720p"))] - if !ok { - table = lumaSize["720p"] - } - size, ok := table[defaultString(aspectRatio, "16:9")] - if !ok { - size = table["16:9"] - } - return map[string]any{"width": size[0], "height": size[1]} -} - -func buildVideoPromptJSON(prompt string, durationSeconds int) string { - payload := map[string]any{ - "id": 1, - "duration_sec": durationSeconds, - "prompt_text": prompt, - } - b, _ := json.Marshal(payload) - return string(b) -} - -func min(a, b int) int { - if a < b { - return a - } - return b -} diff --git a/backend/internal/provider/adobe/util.go b/backend/internal/provider/adobe/util.go deleted file mode 100644 index db01867..0000000 --- a/backend/internal/provider/adobe/util.go +++ /dev/null @@ -1,386 +0,0 @@ -package adobe - -import ( - "crypto/rand" - "encoding/base64" - "encoding/hex" - "encoding/json" - "math/big" - "net/url" - "regexp" - "strconv" - "strings" - "sync" - "time" - - "github.com/google/uuid" -) - -// adobeUserIDPat matches Adobe IMS user IDs embedded in cookies (e.g. -// "4BDA81F069FC6DA40A495FAB@AdobeID"). -var adobeUserIDPat = regexp.MustCompile(`[A-Fa-f0-9]{20,}@AdobeID`) - -// PID pool: one PID per token, reused across requests for the same account. -var ( - arpPIDMu sync.Mutex - arpTokenPID = map[string]int{} // token → pid - arpPIDToken = map[int]string{} // pid → token -) - -func extractUserIDFromCookie(cookie string) string { - decoded, err := url.QueryUnescape(cookie) - if err != nil { - decoded = cookie - } - return adobeUserIDPat.FindString(decoded) -} - -func stringValue(v any) string { - switch x := v.(type) { - case string: - return x - case nil: - return "" - default: - return strings.TrimSpace(strings.ReplaceAll(toJSONScalar(x), "\n", " ")) - } -} - -func toJSONScalar(v any) string { - b, err := json.Marshal(v) - if err != nil { - return "" - } - return string(b) -} - -func intValue(v any) int { - switch x := v.(type) { - case int: - return x - case int64: - return int(x) - case float64: - return int(x) - case float32: - return int(x) - case json.Number: - n, _ := x.Int64() - return int(n) - case string: - n, _ := strconv.Atoi(strings.TrimSpace(x)) - return n - default: - return 0 - } -} - -func defaultString(v, fallback string) string { - v = strings.TrimSpace(v) - if v == "" { - return fallback - } - return v -} - -func itoa(v int) string { - return strconv.Itoa(v) -} - -func decodeJWTPayload(token string) map[string]any { - parts := strings.Split(strings.TrimSpace(token), ".") - 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 -} - -// arkHead and arkTail are the site-constant halves of the ark blob, as captured -// from firefly.adobe.com. The leading "<17hex>.<10digits>" pair and the rid -// field sit between them and vary per session, so buildARKBlob randomizes those -// instead of shipping one frozen blob for every account. -const ( - arkHead = "|r=ap-southeast-1|meta=3|metabgclr=transparent|metaiconclr=%23757575|guitextcolor=%23000000|pk=BBCC314C-4937-4CCD-B0A3-FDF0F0F7603C|at=40|sup=1" - arkTail = "|ag=101|cdn_url=https%3A%2F%2Farks-client.adobe.com%2Fcdn%2Ffc|surl=https%3A%2F%2Farks-client.adobe.com|smurl=https%3A%2F%2Farks-client.adobe.com%2Fcdn%2Ffc%2Fassets%2Fstyle-manager" -) - -func buildARKBlob() string { - return randomHex(9)[:17] + "." + randomDigits(10) + - arkHead + "|rid=" + strconv.Itoa(randomInt(1, 99)) + arkTail -} - -// Per-token stable session identity: a real browser keeps sid / the ftr device -// hex / bfp constant for the whole session, only the ftr timestamp advances. -// Re-randomizing them every call is a bot signal, so they're cached per token -// alongside the pid. -var ( - arpIdentMu sync.Mutex - arpIdent = map[string]arpSessionIdent{} -) - -type arpSessionIdent struct { - sid string - devHex string - bfp string - arp string // cookie-derived value (sherlockToken), preferred over synthesis -} - -// cookieValue extracts one cookie value from a raw cookie blob. -func cookieValue(cookie, name string) string { - decoded, err := url.QueryUnescape(cookie) - if err != nil { - decoded = cookie - } - for _, part := range strings.Split(decoded, ";") { - part = strings.TrimSpace(part) - if k, v, ok := strings.Cut(part, "="); ok && k == name { - return v - } - } - return "" -} - -// ARPFromCookie rebuilds x-arp-session-id from a browser cookie blob. The -// sherlockToken cookie IS the arp value verbatim; when it's absent the same -// {sid,ark,ftr,bfp} JSON is reassembled from the component cookies — -// forterToken is the complete ftr (with the real Arkose payload), arkose the -// ark blob, ff_session_guid the sid, bfp the fingerprint id. Returns "" when -// no usable pieces are found. -func ARPFromCookie(cookie string) string { - if st := cookieValue(cookie, "sherlockToken"); st != "" { - return st - } - ftr, ark := cookieValue(cookie, "forterToken"), cookieValue(cookie, "arkose") - if ftr == "" { - ftr = cookieValue(cookie, "forter") - } - if ftr == "" || ark == "" { - return "" - } - payload := map[string]any{"ftr": ftr, "ark": ark} - if v := cookieValue(cookie, "ff_session_guid"); v != "" { - payload["sid"] = v - } - if v := cookieValue(cookie, "bfp"); v != "" { - payload["bfp"] = v - } - b, _ := json.Marshal(payload) - return base64.StdEncoding.EncodeToString(b) -} - -// refreshARPTimestamp decodes a cookie-derived arp payload and bumps the ftr -// timestamp field to now before re-encoding. ftr layout is -// ____31ck__tt — a stale timestamp reads as -// a dead session to the gateway; every other field stays verbatim so the -// Forter/Arkose linkage is preserved. Returns arp unchanged on any shape -// mismatch. -func refreshARPTimestamp(arp string) string { - raw, err := base64.StdEncoding.DecodeString(arp) - if err != nil { - return arp - } - var m map[string]any - if err := json.Unmarshal(raw, &m); err != nil { - return arp - } - ftr, _ := m["ftr"].(string) - parts := strings.Split(ftr, "_") - if len(parts) < 3 { - return arp - } - parts[1] = strconv.FormatInt(time.Now().UnixMilli(), 10) - m["ftr"] = strings.Join(parts, "_") - b, err := json.Marshal(m) - if err != nil { - return arp - } - return base64.StdEncoding.EncodeToString(b) -} - -// CacheARP stores a cookie-derived arp value for token so getARPSessionID -// prefers the real (Arkose-bearing) value over the synthesized fallback. -func CacheARP(token, arp string) { - token, arp = strings.TrimSpace(token), strings.TrimSpace(arp) - if token == "" || arp == "" { - return - } - arpIdentMu.Lock() - defer arpIdentMu.Unlock() - id := arpIdent[token] - id.arp = arp - arpIdent[token] = id -} - -// buildARPSessionID returns the x-arp-session-id value: base64 of -// {"sid","ark","ftr","bfp"}, shaped to match live firefly.adobe.com traffic: -// -// ark: <17hex>.<10digits>|…|rid=|… -// ftr: <32hex>___UDF43-mnf-a4_31ck__tt -// -// pid, sid, the ftr device hex and bfp stay stable per token — a real browser -// keeps them for the whole session and re-randomizing is a bot signal. -// -// The Arkose segment is left empty on purpose. In a real browser that slot -// carries a "=-<4digits>-v2" payload emitted by Adobe's Arkose bot-manager -// script; a synthesized one has the right shape but fails Adobe's server-side -// validation, and 3p-images/generate-async rejects it with a misleading -// {"error_code":"timeout_error","message":"system under load"}. Sending the -// slot empty reads as "no Arkose payload yet" and is accepted by the native -// endpoints; 3p models may require the cookie-derived value (see CacheARP). -func buildARPSessionID(token string) string { - ftr := strings.Join([]string{ - arpIdentity(token).devHex, - strconv.FormatInt(time.Now().UnixMilli(), 10), - strconv.Itoa(allocPID(token)), - "UDF43-mnf-a4", - "31ck", - "", - "tt", - }, "_") - - id := arpIdentity(token) - b, _ := json.Marshal(map[string]any{ - "sid": id.sid, - "ark": buildARKBlob(), - "ftr": ftr, - "bfp": id.bfp, - }) - return base64.StdEncoding.EncodeToString(b) -} - -// arpIdentity returns (creating on first use) the stable per-token identity -// pieces, preferring a cached cookie-derived arp value when one exists. -func arpIdentity(token string) arpSessionIdent { - arpIdentMu.Lock() - defer arpIdentMu.Unlock() - id, ok := arpIdent[token] - if !ok || id.sid == "" { - id.sid = uuid.NewString() - } - if !ok || id.devHex == "" { - id.devHex = randomHex(16) - } - if !ok || id.bfp == "" { - id.bfp = uuid.NewString() - } - arpIdent[token] = id - return id -} - -// randomDigits returns n decimal digits with a non-zero leading digit, so the -// rendered value keeps a constant width. -func randomDigits(n int) string { - if n <= 0 { - return "" - } - var b strings.Builder - b.Grow(n) - b.WriteByte(byte('1' + randomInt(0, 8))) - for i := 1; i < n; i++ { - b.WriteByte(byte('0' + randomInt(0, 9))) - } - return b.String() -} - -func randomBase64(n int) string { - buf := make([]byte, n) - rand.Read(buf) - return base64.RawURLEncoding.EncodeToString(buf) -} - -func allocPID(token string) int { - arpPIDMu.Lock() - defer arpPIDMu.Unlock() - if pid, ok := arpTokenPID[token]; ok { - return pid - } - for { - pid := randomInt(1000, 99999) - if _, used := arpPIDToken[pid]; !used { - arpPIDToken[pid] = token - arpTokenPID[token] = pid - return pid - } - } -} - -// ReleasePID releases the PID bound to token so it can be reused by another -// account. -func ReleasePID(token string) { - arpPIDMu.Lock() - defer arpPIDMu.Unlock() - if pid, ok := arpTokenPID[token]; ok { - delete(arpPIDToken, pid) - delete(arpTokenPID, token) - } -} - -func randomHex(n int) string { - if n <= 0 { - return "" - } - buf := make([]byte, n) - if _, err := rand.Read(buf); err != nil { - now := time.Now().UnixNano() - for i := range buf { - buf[i] = byte(now >> ((i % 8) * 8)) - } - } - return hex.EncodeToString(buf) -} - -func randomInt(min, max int) int { - if max <= min { - return min - } - n, err := rand.Int(rand.Reader, big.NewInt(int64(max-min+1))) - if err != nil { - return min - } - return min + int(n.Int64()) -} - -func intOrNil(v any) any { - switch x := v.(type) { - case nil: - return nil - case int: - return x - case int64: - return int(x) - case float64: - return int(x) - case float32: - return int(x) - case json.Number: - n, err := x.Int64() - if err != nil { - return nil - } - return int(n) - case string: - n, err := strconv.Atoi(strings.TrimSpace(x)) - if err != nil { - return nil - } - return n - default: - return nil - } -} - -func emptyStringNil(v string) any { - v = strings.TrimSpace(v) - if v == "" { - return nil - } - return v -} diff --git a/backend/internal/provider/chatgpt/client.go b/backend/internal/provider/chatgpt/client.go deleted file mode 100644 index a6fc18b..0000000 --- a/backend/internal/provider/chatgpt/client.go +++ /dev/null @@ -1,1396 +0,0 @@ -package chatgpt - -import ( - "bufio" - "bytes" - "context" - "encoding/json" - "errors" - "fmt" - "image" - _ "image/gif" - _ "image/jpeg" - _ "image/png" - "io" - stdhttp "net/http" - "net/url" - "os" - "path/filepath" - "strconv" - "strings" - "time" - - http "github.com/bogdanfinn/fhttp" - tlsclient "github.com/bogdanfinn/tls-client" - "github.com/bogdanfinn/tls-client/profiles" -) - -var ( - ErrAuth = errors.New("chatgpt auth failed") - ErrQuotaExhausted = errors.New("chatgpt quota exhausted") - ErrTemporaryUpstream = errors.New("chatgpt upstream temporary error") - // ErrContentPolicy marks a prompt rejected by ChatGPT's content audit. It is - // terminal and NOT retryable: the same prompt fails on every account, so the - // caller must fail fast rather than poll or fail over. - ErrContentPolicy = errors.New("chatgpt content policy rejection") -) - -type Client struct { - proxy string - deviceID string - sessionID string -} - -type fileEntry struct { - FileID string - UploadURL string -} - -type uploadedReference struct { - FileID string - LibraryFileID string - FileName string - MimeType string - SizeBytes int - Width int - Height int -} - -func NewClient(proxy string) *Client { - return &Client{ - proxy: strings.TrimSpace(proxy), - deviceID: newUUID(), - sessionID: newUUID(), - } -} - -func (c *Client) SetProxy(proxy string) { - c.proxy = strings.TrimSpace(proxy) -} - -func (c *Client) GenerateImage(ctx context.Context, accessToken, prompt, model, aspectRatio, resolution string, refs [][]byte, downloadResult bool) ([]byte, map[string]any, error) { - // Everything except the generation submit egresses on the local IP. Only - // startImageGeneration (the /backend-api/f/conversation POST) goes through - // the proxy; the bootstrap / chat-requirements / reference upload / prepare - // handshake and the poll / resolve / download run direct. - session, err := c.newDirectSession(accessToken) - if err != nil { - return nil, nil, err - } - - // Fail over immediately when the account's image_gen allowance is spent: - // submitting anyway just burns the whole poll budget and surfaces as - // "image poll timeout". Unknown quota (init failed) proceeds as before. - if quota, qErr := c.fetchImageQuota(ctx, session, accessToken); qErr == nil && quota["unknown"] == false { - if remaining, ok := quota["remaining"].(int); ok && remaining <= 0 { - return nil, nil, fmt.Errorf("%w: image_gen remaining 0 (resets %s)", ErrQuotaExhausted, stringValue(quota["reset_after"])) - } - } - - scriptSources, dataBuild, err := c.bootstrap(ctx, session) - if err != nil { - return nil, nil, err - } - reqs, err := c.getChatRequirements(ctx, session, accessToken, scriptSources, dataBuild) - if err != nil { - return nil, nil, err - } - effectivePrompt := injectSizeHint(prompt, aspectRatio, resolution) - uploadedRefs, err := c.uploadReferenceImages(ctx, session, accessToken, refs) - if err != nil { - return nil, nil, err - } - conduitToken, err := c.prepareImageConversation(ctx, session, accessToken, effectivePrompt, reqs, model, uploadedRefs) - if err != nil { - return nil, nil, err - } - // The generation submit is the only request that egresses via the proxy. A - // bare connection error here (EOF / connection reset) is almost always the - // proxy hop flaking rather than the account, so retry the submit on a fresh - // connection a few times before surfacing it — otherwise a momentary network - // blip fails over accounts (burning their quota) or fails the request. - var conversationID string - var fileIDs, sedimentIDs []string - for attempt := 0; ; attempt++ { - submitSession, sErr := c.newSession(accessToken) - if sErr != nil { - return nil, nil, sErr - } - conversationID, fileIDs, sedimentIDs, err = c.startImageGeneration(ctx, submitSession, accessToken, effectivePrompt, reqs, conduitToken, model, uploadedRefs) - if err == nil || attempt >= submitConnectRetries || ctx.Err() != nil || !isConnResetErr(err) { - break - } - time.Sleep(submitRetryBackoff) - } - if err != nil { - return nil, nil, err - } - // The SSE stream and conversation JSON echo the user's uploaded reference - // assets; treating those ids as "the generated image" would return the - // reference itself. Drop them from every id set we collect. - refIDs := uploadedRefIDSet(uploadedRefs) - fileIDs = dropIDs(fileIDs, refIDs) - sedimentIDs = dropIDs(sedimentIDs, refIDs) - fileIDs, sedimentIDs, err = c.pollForImage(ctx, session, accessToken, conversationID, fileIDs, sedimentIDs, refIDs, pollBudget(ctx)) - if err != nil { - return nil, nil, err - } - urls, err := c.resolveImageURLs(ctx, session, accessToken, conversationID, fileIDs, sedimentIDs) - if err != nil { - return nil, nil, err - } - if len(urls) == 0 { - return nil, nil, errors.New("no image urls resolved") - } - meta := map[string]any{ - "provider": "chatgpt", - "model": model, - "conversation_id": conversationID, - "image_url": urls[0], // auth-gated (files.oaiusercontent.com) — needs the account token to fetch - } - // downloadResult=false: skip the (auth-gated) download and return just the URL; - // the caller proxies it via OpenAsset with the account token. - if !downloadResult { - return nil, meta, nil - } - images, err := c.downloadBytes(ctx, session, accessToken, urls) - if err != nil { - return nil, nil, err - } - if len(images) == 0 { - return nil, nil, errors.New("download produced no bytes") - } - return images[0], meta, nil -} - -// OpenAsset streams an auth-gated ChatGPT image URL (files.oaiusercontent.com) -// using the generating account's token — a plain GET 403s. Mirrors downloadBytes -// but returns a live stream instead of buffering. -func (c *Client) OpenAsset(ctx context.Context, accessToken, rawURL string) (io.ReadCloser, string, error) { - session, err := c.newDirectSession(accessToken) - if err != nil { - return nil, "", err - } - req, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil) - if err != nil { - return nil, "", err - } - req.Header = c.baseHeaders(accessToken) - req.Header.Set("accept", "*/*") - resp, err := session.Do(req) - if err != nil { - return nil, "", fmt.Errorf("%w: %v", ErrTemporaryUpstream, err) - } - if resp.StatusCode != 200 { - resp.Body.Close() - return nil, "", fmt.Errorf("%w: chatgpt asset status %d", ErrTemporaryUpstream, resp.StatusCode) - } - ct := strings.TrimSpace(resp.Header.Get("Content-Type")) - if ct == "" { - ct = "image/png" - } - return resp.Body, ct, nil -} - -func ExtractAccountInfo(token string) map[string]any { - claims := decodeJWTPayload(token) - profile, _ := claims["https://api.openai.com/profile"].(map[string]any) - auth, _ := claims["https://api.openai.com/auth"].(map[string]any) - return map[string]any{ - "email": emptyStringNil(strings.TrimSpace(stringValue(profile["email"]))), - "email_verified": profile["email_verified"] == true, - "plan_type": emptyStringNil(strings.TrimSpace(stringValue(auth["chatgpt_plan_type"]))), - "user_id": emptyStringNil(strings.TrimSpace(stringValue(auth["chatgpt_user_id"]))), - "issued_at": claims["iat"], - "expires_at": claims["exp"], - } -} - -func (c *Client) FetchImageQuota(ctx context.Context, accessToken string) (map[string]any, error) { - session, err := c.newDirectSession(accessToken) - if err != nil { - return nil, err - } - return c.fetchImageQuota(ctx, session, accessToken) -} - -func (c *Client) fetchImageQuota(ctx context.Context, session tlsclient.HttpClient, accessToken string) (map[string]any, error) { - path := "/backend-api/conversation/init" - body, _ := json.Marshal(map[string]any{ - "gizmo_id": nil, - "requested_default_model": nil, - "conversation_id": nil, - }) - req, err := http.NewRequest(http.MethodPost, baseURL+path, bytes.NewReader(body)) - if err != nil { - return nil, err - } - req = req.WithContext(ctx) - req.Header = c.headers(accessToken, path, map[string]string{ - "accept": "application/json", - "content-type": "application/json", - }) - resp, err := session.Do(req) - if err != nil { - return map[string]any{"remaining": nil, "reset_after": nil, "unknown": true, "error": "network: " + err.Error()}, nil - } - defer resp.Body.Close() - respBody, err := io.ReadAll(resp.Body) - if err != nil { - return nil, err - } - if resp.StatusCode == 401 { - return map[string]any{"remaining": nil, "reset_after": nil, "unknown": true, "error": "token invalid", "auth_failed": true}, nil - } - if resp.StatusCode != 200 { - return map[string]any{"remaining": nil, "reset_after": nil, "unknown": true, "error": fmt.Sprintf("http %d: %s", resp.StatusCode, clip(respBody, 160))}, nil - } - var payload map[string]any - if err := json.Unmarshal(respBody, &payload); err != nil { - return map[string]any{"remaining": nil, "reset_after": nil, "unknown": true, "error": "non-json response"}, nil - } - limits, _ := payload["limits_progress"].([]any) - for _, raw := range limits { - item, _ := raw.(map[string]any) - if strings.TrimSpace(stringValue(item["feature_name"])) != "image_gen" { - continue - } - return map[string]any{ - "remaining": intOrNil(item["remaining"]), - "reset_after": emptyStringNil(strings.TrimSpace(stringValue(item["reset_after"]))), - "unknown": false, - "error": nil, - }, nil - } - return map[string]any{"remaining": nil, "reset_after": nil, "unknown": true, "error": nil}, nil -} - -type chatRequirements struct { - Token string - ProofToken string - TurnstileToken string -} - -func (c *Client) newSession(accessToken string) (tlsclient.HttpClient, error) { - return c.newSessionP(accessToken, true) -} - -// newDirectSession egresses on the local IP (never the proxy). Used for -// everything except the generation submit (the /backend-api/f/conversation -// POST), which is the only request that uses the proxy. -func (c *Client) newDirectSession(accessToken string) (tlsclient.HttpClient, error) { - return c.newSessionP(accessToken, false) -} - -func (c *Client) newSessionP(accessToken string, useProxy bool) (tlsclient.HttpClient, error) { - options := []tlsclient.HttpClientOption{ - tlsclient.WithTimeoutSeconds(600), - // Match the Python reference (curl_cffi impersonate="chrome110"): the - // Chrome_133 JA3/JA4 was tripping Cloudflare on the bootstrap GET (403). - tlsclient.WithClientProfile(profiles.Chrome_110), - tlsclient.WithRandomTLSExtensionOrder(), - } - if useProxy && c.proxy != "" { - options = append(options, tlsclient.WithProxyUrl(c.proxy)) - } - client, err := tlsclient.NewHttpClient(tlsclient.NewNoopLogger(), options...) - if err != nil { - return nil, err - } - client.SetCookies(&url.URL{Scheme: "https", Host: "chatgpt.com"}, nil) - return client, nil -} - -func (c *Client) baseHeaders(accessToken string) http.Header { - return http.Header{ - "accept-language": {"zh-CN,zh;q=0.9,en;q=0.8,en-GB;q=0.7,en-US;q=0.6"}, - "oai-client-build-number": {defaultClientBuildNumber}, - "oai-client-version": {defaultClientVersion}, - "oai-device-id": {c.deviceID}, - "oai-language": {"zh-CN"}, - "oai-session-id": {c.sessionID}, - "origin": {baseURL}, - "priority": {"u=1, i"}, - "referer": {baseURL + "/"}, - "sec-ch-ua": {`"Microsoft Edge";v="149", "Chromium";v="149", "Not)A;Brand";v="24"`}, - "sec-ch-ua-arch": {`"x86"`}, - "sec-ch-ua-bitness": {`"64"`}, - "sec-ch-ua-full-version": {`"149.0.4022.69"`}, - "sec-ch-ua-full-version-list": {`"Microsoft Edge";v="149.0.4022.69", "Chromium";v="149.0.7827.115", "Not)A;Brand";v="24.0.0.0"`}, - "sec-ch-ua-mobile": {"?0"}, - "sec-ch-ua-model": {`""`}, - "sec-ch-ua-platform": {`"Windows"`}, - "sec-ch-ua-platform-version": {`"19.0.0"`}, - "sec-fetch-dest": {"empty"}, - "sec-fetch-mode": {"cors"}, - "sec-fetch-site": {"same-origin"}, - "user-agent": {defaultUserAgent}, - "authorization": {"Bearer " + strings.TrimSpace(accessToken)}, - } -} - -func (c *Client) headers(accessToken, path string, extra map[string]string) http.Header { - h := http.Header{ - "x-openai-target-path": {path}, - "x-openai-target-route": {path}, - } - for k, values := range c.baseHeaders(accessToken) { - h[k] = append([]string{}, values...) - } - for k, v := range extra { - h.Set(k, v) - } - return h -} - -func (c *Client) bootstrap(ctx context.Context, session tlsclient.HttpClient) ([]string, string, error) { - req, err := http.NewRequest(http.MethodGet, baseURL+"/", nil) - if err != nil { - return nil, "", err - } - req = req.WithContext(ctx) - req.Header = http.Header{ - "user-agent": {defaultUserAgent}, - "accept": {"text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8"}, - "accept-language": {"zh-CN,zh;q=0.9,en;q=0.8"}, - "sec-ch-ua": {`"Microsoft Edge";v="143", "Chromium";v="143", "Not A(Brand";v="24"`}, - "sec-ch-ua-mobile": {"?0"}, - "sec-ch-ua-platform": {`"Windows"`}, - "sec-fetch-dest": {"document"}, - "sec-fetch-mode": {"navigate"}, - "sec-fetch-site": {"none"}, - "upgrade-insecure-requests": {"1"}, - } - resp, err := session.Do(req) - if err != nil { - return nil, "", fmt.Errorf("%w: %v", ErrTemporaryUpstream, err) - } - defer resp.Body.Close() - body, err := io.ReadAll(resp.Body) - if err != nil { - return nil, "", err - } - if err := ensureOK(resp.StatusCode, body, "bootstrap"); err != nil { - return nil, "", err - } - sources, dataBuild := parsePOWResources(string(body)) - return sources, dataBuild, nil -} - -func (c *Client) getChatRequirements(ctx context.Context, session tlsclient.HttpClient, accessToken string, scriptSources []string, dataBuild string) (*chatRequirements, error) { - pToken := buildLegacyRequirementsToken(defaultUserAgent, scriptSources, dataBuild) - path := "/backend-api/sentinel/chat-requirements/prepare" - reqBody, _ := json.Marshal(map[string]any{"p": pToken}) - req, err := http.NewRequest(http.MethodPost, baseURL+path, bytes.NewReader(reqBody)) - if err != nil { - return nil, err - } - req = req.WithContext(ctx) - req.Header = c.headers(accessToken, path, map[string]string{ - "content-type": "application/json", - }) - resp, err := session.Do(req) - if err != nil { - return nil, fmt.Errorf("%w: %v", ErrTemporaryUpstream, err) - } - body, err := io.ReadAll(resp.Body) - resp.Body.Close() - if err != nil { - return nil, err - } - if err := ensureOK(resp.StatusCode, body, "chat_requirements_prepare"); err != nil { - return nil, err - } - var prepare map[string]any - if err := json.Unmarshal(body, &prepare); err != nil { - return nil, err - } - if arkose, _ := prepare["arkose"].(map[string]any); arkose["required"] == true { - return nil, errors.New("chat-requirements requires arkose token") - } - proofToken := "" - if powInfo, _ := prepare["proofofwork"].(map[string]any); powInfo["required"] == true { - proofToken, err = buildProofToken(strings.TrimSpace(stringValue(powInfo["seed"])), strings.TrimSpace(stringValue(powInfo["difficulty"])), defaultUserAgent, scriptSources, dataBuild) - if err != nil { - return nil, err - } - } - turnstileToken := "" - if tsInfo, _ := prepare["turnstile"].(map[string]any); tsInfo["required"] == true { - turnstileToken = solveTurnstileToken(strings.TrimSpace(stringValue(tsInfo["dx"])), pToken) - } - - path = "/backend-api/sentinel/chat-requirements/finalize" - finalizeBody, _ := json.Marshal(map[string]any{ - "prepare_token": stringValue(prepare["prepare_token"]), - "proof_token": proofToken, - "turnstile_token": turnstileToken, - }) - req, err = http.NewRequest(http.MethodPost, baseURL+path, bytes.NewReader(finalizeBody)) - if err != nil { - return nil, err - } - req = req.WithContext(ctx) - req.Header = c.headers(accessToken, path, map[string]string{ - "content-type": "application/json", - }) - resp, err = session.Do(req) - if err != nil { - return nil, fmt.Errorf("%w: %v", ErrTemporaryUpstream, err) - } - body, err = io.ReadAll(resp.Body) - resp.Body.Close() - if err != nil { - return nil, err - } - if err := ensureOK(resp.StatusCode, body, "chat_requirements_finalize"); err != nil { - return nil, err - } - var data map[string]any - if err := json.Unmarshal(body, &data); err != nil { - return nil, err - } - token := strings.TrimSpace(stringValue(data["token"])) - if token == "" { - return nil, errors.New("chat-requirements missing token") - } - return &chatRequirements{Token: token, ProofToken: proofToken, TurnstileToken: turnstileToken}, nil -} - -func (c *Client) imageHeaders(accessToken, path string, reqs *chatRequirements, conduitToken, accept string) http.Header { - h := c.headers(accessToken, path, map[string]string{ - "content-type": "application/json", - "accept": accept, - "openai-sentinel-chat-requirements-token": reqs.Token, - }) - if reqs.ProofToken != "" { - h.Set("openai-sentinel-proof-token", reqs.ProofToken) - } - // The real browser also sends the turnstile token on the conversation call - // (HAR confirms openai-sentinel-turnstile-token). We already compute it in - // chat-requirements/finalize but were dropping it here — send it so the - // request matches the browser and isn't extra-challenged by sentinel. - if reqs.TurnstileToken != "" { - h.Set("openai-sentinel-turnstile-token", reqs.TurnstileToken) - } - if conduitToken != "" { - h.Set("x-conduit-token", conduitToken) - } - if accept == "text/event-stream" { - h.Set("x-oai-turn-trace-id", newUUID()) - h.Set("oai-telemetry", "[1,null]") - } - return h -} - -func (c *Client) uploadReferenceImages(ctx context.Context, session tlsclient.HttpClient, accessToken string, refs [][]byte) ([]uploadedReference, error) { - if len(refs) == 0 { - return nil, nil - } - out := make([]uploadedReference, 0, len(refs)) - for i, ref := range refs { - meta, err := inspectReferenceImage(ref, i) - if err != nil { - return nil, err - } - entry, err := c.createFileEntry(ctx, session, accessToken, meta) - if err != nil { - return nil, err - } - if err := c.uploadRawFile(ctx, session, entry.UploadURL, meta.MimeType, ref); err != nil { - return nil, err - } - libraryFileID, err := c.processUploadStream(ctx, session, accessToken, entry.FileID, meta.FileName) - if err != nil { - return nil, err - } - meta.FileID = entry.FileID - meta.LibraryFileID = libraryFileID - out = append(out, meta) - } - return out, nil -} - -func imageModelSlug(model string) string { - if strings.EqualFold(strings.TrimSpace(model), "gpt-image-2") { - return "gpt-5-5-thinking" - } - return "auto" -} - -func inspectReferenceImage(data []byte, index int) (uploadedReference, error) { - if len(data) == 0 { - return uploadedReference{}, errors.New("empty reference image") - } - mimeType := normalizeImageMime(stdhttp.DetectContentType(data)) - if mimeType == "" { - return uploadedReference{}, errors.New("unsupported reference image type") - } - cfg, _, err := image.DecodeConfig(bytes.NewReader(data)) - if err != nil { - return uploadedReference{}, errors.New("failed to decode reference image") - } - ext := extensionForMime(mimeType) - fileName := fmt.Sprintf("reference_%s_%02d%s", time.Now().UTC().Format("20060102_150405"), index+1, ext) - return uploadedReference{ - FileName: fileName, - MimeType: mimeType, - SizeBytes: len(data), - Width: cfg.Width, - Height: cfg.Height, - }, nil -} - -func normalizeImageMime(v string) string { - v = strings.ToLower(strings.TrimSpace(strings.Split(v, ";")[0])) - switch v { - case "image/jpeg", "image/jpg": - return "image/jpeg" - case "image/png": - return "image/png" - case "image/gif": - return "image/gif" - default: - return "" - } -} - -func extensionForMime(mimeType string) string { - switch strings.ToLower(strings.TrimSpace(mimeType)) { - case "image/jpeg": - return ".jpg" - case "image/png": - return ".png" - case "image/gif": - return ".gif" - default: - return filepath.Ext(mimeType) - } -} - -func (c *Client) createFileEntry(ctx context.Context, session tlsclient.HttpClient, accessToken string, meta uploadedReference) (*fileEntry, error) { - path := "/backend-api/files" - payload := map[string]any{ - "file_name": meta.FileName, - "file_size": meta.SizeBytes, - "use_case": "multimodal", - "timezone_offset_min": -480, - "reset_rate_limits": false, - "mime_type": meta.MimeType, - "entry_surface": "chat_composer", - "selection_method": "file_picker", - "client_resolved_mime_type": meta.MimeType, - "mime_resolution_source": "filename_extension", - "store_in_library": true, - "library_persistence_mode": "opportunistic", - } - body, _ := json.Marshal(payload) - req, err := http.NewRequest(http.MethodPost, baseURL+path, bytes.NewReader(body)) - if err != nil { - return nil, err - } - req = req.WithContext(ctx) - req.Header = c.headers(accessToken, path, map[string]string{ - "accept": "application/json", - "content-type": "application/json", - }) - resp, err := session.Do(req) - if err != nil { - return nil, fmt.Errorf("%w: %v", ErrTemporaryUpstream, err) - } - respBody, err := io.ReadAll(resp.Body) - resp.Body.Close() - if err != nil { - return nil, err - } - if err := ensureOK(resp.StatusCode, respBody, "file_create"); err != nil { - return nil, err - } - var data map[string]any - if err := json.Unmarshal(respBody, &data); err != nil { - return nil, err - } - entry := &fileEntry{ - FileID: strings.TrimSpace(stringValue(data["file_id"])), - UploadURL: strings.TrimSpace(stringValue(data["upload_url"])), - } - if entry.FileID == "" || entry.UploadURL == "" { - return nil, errors.New("file_create missing upload payload") - } - return entry, nil -} - -func (c *Client) uploadRawFile(ctx context.Context, session tlsclient.HttpClient, uploadURL, mimeType string, data []byte) error { - req, err := http.NewRequest(http.MethodPut, uploadURL, bytes.NewReader(data)) - if err != nil { - return err - } - req = req.WithContext(ctx) - req.Header = http.Header{ - "accept": {"application/json, text/plain, */*"}, - "accept-language": {"zh-CN,zh;q=0.9,en;q=0.8"}, - "content-type": {mimeType}, - "origin": {baseURL}, - "referer": {baseURL + "/"}, - "sec-ch-ua": {`"Microsoft Edge";v="143", "Chromium";v="143", "Not A(Brand";v="24"`}, - "sec-ch-ua-mobile": {"?0"}, - "sec-ch-ua-platform": {`"Windows"`}, - "sec-fetch-dest": {"empty"}, - "sec-fetch-mode": {"cors"}, - "sec-fetch-site": {"cross-site"}, - "user-agent": {defaultUserAgent}, - "x-ms-blob-type": {"BlockBlob"}, - "x-ms-version": {"2020-04-08"}, - } - resp, err := session.Do(req) - if err != nil { - return fmt.Errorf("%w: %v", ErrTemporaryUpstream, err) - } - body, readErr := io.ReadAll(resp.Body) - resp.Body.Close() - if readErr != nil { - return readErr - } - if err := ensureOK(resp.StatusCode, body, "file_upload"); err != nil { - return err - } - return nil -} - -func (c *Client) processUploadStream(ctx context.Context, session tlsclient.HttpClient, accessToken, fileID, fileName string) (string, error) { - path := "/backend-api/files/process_upload_stream" - payload := map[string]any{ - "file_id": fileID, - "use_case": "multimodal", - "index_for_retrieval": false, - "file_name": fileName, - "library_persistence_mode": "opportunistic", - "entry_surface": "chat_composer", - "metadata": map[string]any{ - "store_in_library": true, - "is_temporary_chat": false, - "library_eligibility_reason": "eligible", - "is_project_thread": false, - }, - } - body, _ := json.Marshal(payload) - req, err := http.NewRequest(http.MethodPost, baseURL+path, bytes.NewReader(body)) - if err != nil { - return "", err - } - req = req.WithContext(ctx) - req.Header = c.headers(accessToken, path, map[string]string{ - "accept": "text/event-stream", - "content-type": "application/json", - }) - resp, err := session.Do(req) - if err != nil { - return "", fmt.Errorf("%w: %v", ErrTemporaryUpstream, err) - } - if resp.StatusCode < 200 || resp.StatusCode >= 300 { - respBody, _ := io.ReadAll(resp.Body) - resp.Body.Close() - if err := ensureOK(resp.StatusCode, respBody, "file_process_upload"); err != nil { - return "", err - } - } - defer resp.Body.Close() - scanner := bufio.NewScanner(resp.Body) - scanner.Buffer(make([]byte, 0, 1024*32), 1024*1024) - libraryFileID := "" - for scanner.Scan() { - line := strings.TrimSpace(scanner.Text()) - if line == "" { - continue - } - var item map[string]any - if err := json.Unmarshal([]byte(line), &item); err != nil { - continue - } - if extra, _ := item["extra"].(map[string]any); extra != nil { - if v := strings.TrimSpace(stringValue(extra["metadata_object_id"])); v != "" { - libraryFileID = v - } - } - if strings.TrimSpace(stringValue(item["event"])) == "file.processing.completed" { - break - } - } - if err := scanner.Err(); err != nil { - return "", err - } - return libraryFileID, nil -} - -func attachmentMimeTypes(refs []uploadedReference) []string { - if len(refs) == 0 { - return nil - } - out := make([]string, 0, len(refs)) - for _, ref := range refs { - if ref.MimeType != "" && !containsString(out, ref.MimeType) { - out = append(out, ref.MimeType) - } - } - return out -} - -func buildAttachmentMetadata(refs []uploadedReference) []map[string]any { - out := make([]map[string]any, 0, len(refs)) - for _, ref := range refs { - out = append(out, map[string]any{ - "id": ref.FileID, - "size": ref.SizeBytes, - "name": ref.FileName, - "mime_type": ref.MimeType, - "width": ref.Width, - "height": ref.Height, - "source": "local", - "library_file_id": emptyStringNil(ref.LibraryFileID), - "is_big_paste": false, - }) - } - return out -} - -func buildMultimodalParts(refs []uploadedReference, prompt string) []any { - out := make([]any, 0, len(refs)+1) - for _, ref := range refs { - out = append(out, map[string]any{ - "content_type": "image_asset_pointer", - "asset_pointer": "sediment://" + ref.FileID, - "size_bytes": ref.SizeBytes, - "width": ref.Width, - "height": ref.Height, - }) - } - out = append(out, prompt) - return out -} - -func (c *Client) prepareImageConversation(ctx context.Context, session tlsclient.HttpClient, accessToken, prompt string, reqs *chatRequirements, model string, refs []uploadedReference) (string, error) { - path := "/backend-api/f/conversation/prepare" - payload := map[string]any{ - "action": "next", - "parent_message_id": "client-created-root", - "model": imageModelSlug(model), - "timezone_offset_min": -480, - "timezone": "Asia/Shanghai", - "conversation_mode": map[string]any{"kind": "primary_assistant"}, - "system_hints": []string{"picture_v2"}, - "supports_buffering": true, - "supported_encodings": []string{"v1"}, - "client_contextual_info": map[string]any{"app_name": "chatgpt.com"}, - } - if len(refs) > 0 { - payload["client_prepare_state"] = "none" - payload["attachment_mime_types"] = attachmentMimeTypes(refs) - } else { - payload["client_prepare_state"] = "success" - payload["partial_query"] = map[string]any{ - "id": newUUID(), - "author": map[string]any{"role": "user"}, - "content": map[string]any{"content_type": "text", "parts": []string{prompt}}, - } - } - body, _ := json.Marshal(payload) - req, err := http.NewRequest(http.MethodPost, baseURL+path, bytes.NewReader(body)) - if err != nil { - return "", err - } - req = req.WithContext(ctx) - req.Header = c.imageHeaders(accessToken, path, reqs, "no-token", "application/json") - resp, err := session.Do(req) - if err != nil { - return "", fmt.Errorf("%w: %v", ErrTemporaryUpstream, err) - } - respBody, err := io.ReadAll(resp.Body) - resp.Body.Close() - if err != nil { - return "", err - } - if err := ensureOK(resp.StatusCode, respBody, "image_prepare"); err != nil { - return "", err - } - var data map[string]any - if err := json.Unmarshal(respBody, &data); err != nil { - return "", err - } - return strings.TrimSpace(stringValue(data["conduit_token"])), nil -} - -func (c *Client) startImageGeneration(ctx context.Context, session tlsclient.HttpClient, accessToken, prompt string, reqs *chatRequirements, conduitToken, model string, refs []uploadedReference) (string, []string, []string, error) { - path := "/backend-api/f/conversation" - // Match the web app: a picture_v2 prompt is sent as the "@创建图片" ecosystem - // mention followed by the prompt, with the mention marked via a - // custom_symbol_offset so the server routes the turn to image generation. - content := map[string]any{"content_type": "text", "parts": []string{pictureV2Command + " " + prompt}} - metadata := map[string]any{ - "system_hints": []string{"picture_v2"}, - "serialization_metadata": map[string]any{ - "custom_symbol_offsets": []any{ - map[string]any{ - "id": "picture_v2", - "symbol": "ecosystemMention", - "startIndex": 0, - "endIndex": pictureV2MentionEnd, - }, - }, - }, - } - if len(refs) > 0 { - // Reference-image turns carry image parts and no leading mention, so keep - // the multimodal shape and drop the ecosystemMention offset. - content = map[string]any{ - "content_type": "multimodal_text", - "parts": buildMultimodalParts(refs, prompt), - } - metadata["attachments"] = buildAttachmentMetadata(refs) - metadata["serialization_metadata"] = map[string]any{"custom_symbol_offsets": []any{}} - } - payload := map[string]any{ - "action": "next", - "messages": []map[string]any{{ - "id": newUUID(), - "author": map[string]any{"role": "user"}, - "create_time": float64(timeMillis()) / 1000.0, - "content": content, - "metadata": metadata, - }}, - "parent_message_id": "client-created-root", - "model": imageModelSlug(model), - "client_prepare_state": "success", - "timezone_offset_min": -480, - "timezone": "Asia/Shanghai", - "conversation_mode": map[string]any{"kind": "primary_assistant"}, - "enable_message_followups": true, - "system_hints": []string{"picture_v2"}, - "supports_buffering": true, - "supported_encodings": []string{"v1"}, - "client_contextual_info": map[string]any{"is_dark_mode": false, "time_since_loaded": 1200, "page_height": 1072, "page_width": 1724, "pixel_ratio": 1.2, "screen_height": 1440, "screen_width": 2560, "app_name": "chatgpt.com", "has_web_push_capabilities": true, "web_push_notification_permission": "default"}, - "paragen_cot_summary_display_override": "allow", - "force_parallel_switch": "auto", - "local_function_names": []string{"local.continue_in_work"}, - } - body, _ := json.Marshal(payload) - req, err := http.NewRequest(http.MethodPost, baseURL+path, bytes.NewReader(body)) - if err != nil { - return "", nil, nil, err - } - req = req.WithContext(ctx) - req.Header = c.imageHeaders(accessToken, path, reqs, conduitToken, "text/event-stream") - resp, err := session.Do(req) - if err != nil { - return "", nil, nil, fmt.Errorf("%w: %v", ErrTemporaryUpstream, err) - } - if resp.StatusCode < 200 || resp.StatusCode >= 300 { - respBody, _ := io.ReadAll(resp.Body) - resp.Body.Close() - if err := ensureOK(resp.StatusCode, respBody, "image_start"); err != nil { - return "", nil, nil, err - } - } - defer resp.Body.Close() - - conversationID := "" - asyncStarted := false - var fileIDs, sedimentIDs []string - scanner := bufio.NewScanner(resp.Body) - scanner.Buffer(make([]byte, 0, 1024*1024), 8*1024*1024) - var chunks []string - refIDs := uploadedRefIDSet(refs) - // Watchdog: scanner.Scan() blocks on a read while the SSE is silent, which - // would hold the whole ctx budget. ChatGPT's inline image pipeline keeps the - // connection open with periodic ": ping" keepalives (~15s apart) while the - // image renders, then streams the asset (a role:"tool" multimodal_text turn - // carrying image_asset_pointer / sediment://) in the same SSE. So this is an - // *idle* timeout reset on every received line — a fixed grace from the - // conversation id would abort mid-render before the asset arrives. It only - // closes the body (unblocking the read) once the stream is genuinely silent. - activity := make(chan struct{}, 1) - watchdogDone := make(chan struct{}) - defer close(watchdogDone) - go func() { - timer := time.NewTimer(sseIdleGrace) - defer timer.Stop() - for { - select { - case <-activity: - if !timer.Stop() { - <-timer.C - } - timer.Reset(sseIdleGrace) - case <-timer.C: - resp.Body.Close() - return - case <-watchdogDone: - return - } - } - }() - for scanner.Scan() { - select { - case activity <- struct{}{}: - default: - } - line := scanner.Text() - if !strings.HasPrefix(line, "data:") { - continue - } - payload := strings.TrimSpace(line[5:]) - if payload == "" { - continue - } - if payload == "[DONE]" { - break - } - chunks = append(chunks, payload) - if conversationID == "" { - if match := conversationIDRE.FindStringSubmatch(payload); len(match) >= 2 { - conversationID = match[1] - } - } - newFiles, newSeds := scanForIDs(payload) - fileIDs = mergeStrings(fileIDs, newFiles) - sedimentIDs = mergeStrings(sedimentIDs, newSeds) - if !asyncStarted && containsAsyncMarker(payload) { - asyncStarted = true - } - // Stop once the image pipeline is confirmed engaged (async/tool marker) or - // the generated asset id has streamed inline. The stream echoes back the - // user's uploaded reference images too, so exclude those ids before deciding - // — otherwise an edit request would stop on its own upload. The idle watchdog - // above bounds a silent stream; pollForImage backstops asset retrieval when - // nothing streamed inline. - haveAsset := len(dropIDs(append([]string(nil), fileIDs...), refIDs)) > 0 || - len(dropIDs(append([]string(nil), sedimentIDs...), refIDs)) > 0 - if conversationID != "" && (asyncStarted || haveAsset) { - break - } - } - if conversationID == "" { - joined := strings.Join(chunks, "\n") - if match := conversationIDRE.FindStringSubmatch(joined); len(match) >= 2 { - conversationID = match[1] - } - } - if conversationID == "" { - return "", nil, nil, errors.New("chatgpt SSE closed without conversation_id") - } - // The stream closed with a conversation id but no sign the image pipeline - // engaged — no async/tool marker and no asset streamed inline. Polling such a - // conversation only burns the whole budget and surfaces as the non-retryable - // "image poll timeout", so treat the attempt as a transient upstream failure. - // That is retryable: a fresh submission reliably engages the pipeline, so the - // pool retries the same account a few times and then fails over to another - // account (换号重试) instead of failing the request. - if !asyncStarted && len(fileIDs) == 0 && len(sedimentIDs) == 0 { - return "", nil, nil, fmt.Errorf("%w: image generation did not start (no async marker)", ErrTemporaryUpstream) - } - return conversationID, fileIDs, sedimentIDs, nil -} - -func (c *Client) getConversation(ctx context.Context, session tlsclient.HttpClient, accessToken, conversationID string) (map[string]any, error) { - path := "/backend-api/conversation/" + conversationID - req, err := http.NewRequest(http.MethodGet, baseURL+path, nil) - if err != nil { - return nil, err - } - req = req.WithContext(ctx) - req.Header = c.headers(accessToken, path, map[string]string{"accept": "application/json"}) - resp, err := session.Do(req) - if err != nil { - return nil, fmt.Errorf("%w: %v", ErrTemporaryUpstream, err) - } - body, err := io.ReadAll(resp.Body) - resp.Body.Close() - if err != nil { - return nil, err - } - if err := ensureOK(resp.StatusCode, body, "conversation_get"); err != nil { - return nil, err - } - if dir := os.Getenv("CHATGPT_DEBUG_DUMP"); dir != "" { - _ = os.WriteFile(filepath.Join(dir, fmt.Sprintf("conv-%s-%d.json", conversationID, time.Now().UnixMilli())), body, 0o644) - } - var payload map[string]any - if err := json.Unmarshal(body, &payload); err != nil { - return nil, err - } - return payload, nil -} - -// pollBudget derives how long to poll for the async image from the caller's -// remaining context budget, leaving headroom to resolve+download the asset -// before the outer deadline (genCtx, 8min) fires. Async image generation under -// load routinely exceeds the old hard-coded 180s, which surfaced as -// "image poll timeout"; tying the budget to the deadline lets slow gens finish -// while the context still backstops a truly stuck request. -func pollBudget(ctx context.Context) time.Duration { - const ( - maxBudget = 6 * time.Minute - headroom = 25 * time.Second - ) - deadline, ok := ctx.Deadline() - if !ok { - return 3 * time.Minute - } - budget := time.Until(deadline) - headroom - if budget < 0 { - budget = 0 - } - if budget > maxBudget { - budget = maxBudget - } - return budget -} - -func (c *Client) pollForImage(ctx context.Context, session tlsclient.HttpClient, accessToken, conversationID string, initialFileIDs, initialSedimentIDs []string, refIDs map[string]bool, timeout time.Duration) ([]string, []string, error) { - start := time.Now() - fileIDs := dropIDs(append([]string{}, initialFileIDs...), refIDs) - sedimentIDs := dropIDs(append([]string{}, initialSedimentIDs...), refIDs) - if len(fileIDs) == 0 { - time.Sleep(8 * time.Second) - } else { - time.Sleep(2 * time.Second) - } - attempt := 0 - for time.Since(start) < timeout { - // Bail out immediately if the caller's context is already done — without - // this, a cancelled request spins here re-issuing doomed upstream calls - // (each fails instantly with "operation was canceled") until `timeout`. - if err := ctx.Err(); err != nil { - return nil, nil, err - } - attempt++ - conv, err := c.getConversation(ctx, session, accessToken, conversationID) - if err != nil { - if errors.Is(err, ErrTemporaryUpstream) { - time.Sleep(time.Duration(minInt(1< 0 || len(sedimentIDs) > 0 { - time.Sleep(2 * time.Second) - conv, err = c.getConversation(ctx, session, accessToken, conversationID) - if err == nil { - finalFiles, finalSeds := extractImageIDs(conv) - fileIDs = mergeStrings(fileIDs, dropIDs(finalFiles, refIDs)) - sedimentIDs = mergeStrings(sedimentIDs, dropIDs(finalSeds, refIDs)) - } - return fileIDs, sedimentIDs, nil - } - time.Sleep(5 * time.Second) - } - return nil, nil, errors.New("image poll timeout") -} - -// uploadedRefIDSet collects every id belonging to the user's uploaded -// reference images so they can be excluded from generated-asset extraction. -func uploadedRefIDSet(refs []uploadedReference) map[string]bool { - ids := make(map[string]bool, len(refs)*2) - for _, ref := range refs { - if ref.FileID != "" { - ids[ref.FileID] = true - } - if ref.LibraryFileID != "" { - ids[ref.LibraryFileID] = true - } - } - return ids -} - -func dropIDs(ids []string, exclude map[string]bool) []string { - if len(exclude) == 0 || len(ids) == 0 { - return ids - } - out := ids[:0] - for _, id := range ids { - if !exclude[id] { - out = append(out, id) - } - } - return out -} - -func (c *Client) getFileDownloadURL(ctx context.Context, session tlsclient.HttpClient, accessToken, conversationID, fileID string, inline bool) (string, error) { - // Current web client form: GET /backend-api/files/download/{id} - // ?conversation_id=...&inline=false → {"status":"success","download_url":...}. - // Falls back to the legacy /files/{id}/download form if the new one fails. - paths := []string{ - "/backend-api/files/download/" + fileID + "?conversation_id=" + conversationID + "&inline=" + strconv.FormatBool(inline), - "/backend-api/files/" + fileID + "/download", - } - var lastErr error - for _, path := range paths { - rawURL, err := c.fetchDownloadURL(ctx, session, accessToken, path) - if err == nil && rawURL != "" { - return rawURL, nil - } - if err != nil { - lastErr = err - } - } - return "", lastErr -} - -func (c *Client) fetchDownloadURL(ctx context.Context, session tlsclient.HttpClient, accessToken, path string) (string, error) { - req, err := http.NewRequest(http.MethodGet, baseURL+path, nil) - if err != nil { - return "", err - } - req = req.WithContext(ctx) - req.Header = c.headers(accessToken, path, map[string]string{"accept": "application/json"}) - resp, err := session.Do(req) - if err != nil { - return "", fmt.Errorf("%w: %v", ErrTemporaryUpstream, err) - } - body, err := io.ReadAll(resp.Body) - resp.Body.Close() - if err != nil { - return "", err - } - if err := ensureOK(resp.StatusCode, body, "file_download_url"); err != nil { - return "", err - } - var payload map[string]any - if err := json.Unmarshal(body, &payload); err != nil { - return "", err - } - rawURL := strings.TrimSpace(stringValue(payload["download_url"])) - if rawURL == "" { - rawURL = strings.TrimSpace(stringValue(payload["url"])) - } - return rawURL, nil -} - -func (c *Client) resolveImageURLs(ctx context.Context, session tlsclient.HttpClient, accessToken, conversationID string, fileIDs, sedimentIDs []string) ([]string, error) { - var urls []string - for _, fileID := range fileIDs { - if fileID == "file_upload" { - continue - } - rawURL, err := c.getFileDownloadURL(ctx, session, accessToken, conversationID, fileID, false) - if err != nil { - continue - } - if rawURL != "" && !containsString(urls, rawURL) { - urls = append(urls, rawURL) - } - } - for _, sedimentID := range sedimentIDs { - path := "/backend-api/conversation/" + conversationID + "/attachment/" + sedimentID + "/download" - req, err := http.NewRequest(http.MethodGet, baseURL+path, nil) - if err != nil { - continue - } - req = req.WithContext(ctx) - req.Header = c.headers(accessToken, path, map[string]string{"accept": "application/json"}) - resp, err := session.Do(req) - if err != nil { - continue - } - body, readErr := io.ReadAll(resp.Body) - resp.Body.Close() - if readErr != nil || resp.StatusCode < 200 || resp.StatusCode >= 300 { - continue - } - var payload map[string]any - if err := json.Unmarshal(body, &payload); err != nil { - continue - } - rawURL := strings.TrimSpace(stringValue(payload["download_url"])) - if rawURL == "" { - rawURL = strings.TrimSpace(stringValue(payload["url"])) - } - if rawURL != "" && !containsString(urls, rawURL) { - urls = append(urls, rawURL) - } - } - return urls, nil -} - -func (c *Client) downloadBytes(ctx context.Context, session tlsclient.HttpClient, accessToken string, urls []string) ([][]byte, error) { - out := make([][]byte, 0, len(urls)) - for _, rawURL := range urls { - req, err := http.NewRequest(http.MethodGet, rawURL, nil) - if err != nil { - return nil, err - } - req = req.WithContext(ctx) - // Mirror Python's session.get(url): the resolved download_url is a - // backend-api stream that requires the same default headers as every - // other call — crucially Authorization. Without it the fetch 403s with - // {"detail":"File stream access denied."}. - req.Header = c.baseHeaders(accessToken) - req.Header.Set("accept", "*/*") - resp, err := session.Do(req) - if err != nil { - return nil, fmt.Errorf("%w: %v", ErrTemporaryUpstream, err) - } - body, readErr := io.ReadAll(resp.Body) - resp.Body.Close() - if readErr != nil { - return nil, readErr - } - if err := ensureOK(resp.StatusCode, body, "image_download"); err != nil { - return nil, err - } - if len(body) > 0 { - out = append(out, body) - } - } - return out, nil -} - -func ensureOK(statusCode int, body []byte, context string) error { - if statusCode >= 200 && statusCode < 300 { - return nil - } - switch statusCode { - case 401, 403: - return fmt.Errorf("%w: %s %d %s", ErrAuth, context, statusCode, clip(body, 400)) - case 429: - return fmt.Errorf("%w: %s 429 %s", ErrQuotaExhausted, context, clip(body, 400)) - case 500, 502, 503, 504: - return fmt.Errorf("%w: %s %d %s", ErrTemporaryUpstream, context, statusCode, clip(body, 400)) - default: - return fmt.Errorf("%s: %d %s", context, statusCode, clip(body, 400)) - } -} - -func injectSizeHint(prompt, aspectRatio, resolution string) string { - _ = resolution - ratio := strings.TrimSpace(aspectRatio) - if ratio == "" || strings.EqualFold(ratio, "auto") { - return strings.TrimSpace(prompt) - } - return strings.TrimSpace(prompt) + "\n\n将宽高比设为 " + ratio -} - -func scanForIDs(text string) ([]string, []string) { - var fileIDs []string - for _, id := range fileServiceIDPattern.FindAllStringSubmatch(text, -1) { - if len(id) >= 2 && id[1] != "file_upload" && !containsString(fileIDs, id[1]) { - fileIDs = append(fileIDs, id[1]) - } - } - for _, id := range realImageIDPattern.FindAllString(text, -1) { - if !containsString(fileIDs, id) { - fileIDs = append(fileIDs, id) - } - } - var sedimentIDs []string - for _, id := range sedimentIDPattern.FindAllStringSubmatch(text, -1) { - if len(id) >= 2 && !containsString(sedimentIDs, id[1]) { - sedimentIDs = append(sedimentIDs, id[1]) - } - } - return fileIDs, sedimentIDs -} - -func extractImageIDs(conversation map[string]any) ([]string, []string) { - var fileIDs, sedimentIDs []string - mapping, _ := conversation["mapping"].(map[string]any) - for _, rawNode := range mapping { - node, _ := rawNode.(map[string]any) - message, _ := node["message"].(map[string]any) - author, _ := message["author"].(map[string]any) - role := strings.ToLower(strings.TrimSpace(stringValue(author["role"]))) - if role != "tool" && role != "assistant" { - continue - } - walkForIDs(message["content"], &fileIDs, &sedimentIDs) - walkForIDs(message["metadata"], &fileIDs, &sedimentIDs) - } - return fileIDs, sedimentIDs -} - -func walkForIDs(value any, fileIDs, sedimentIDs *[]string) { - switch x := value.(type) { - case string: - newFiles, newSeds := scanForIDs(x) - *fileIDs = mergeStrings(*fileIDs, newFiles) - *sedimentIDs = mergeStrings(*sedimentIDs, newSeds) - case map[string]any: - for _, item := range x { - walkForIDs(item, fileIDs, sedimentIDs) - } - case []any: - for _, item := range x { - walkForIDs(item, fileIDs, sedimentIDs) - } - } -} - -func mergeStrings(dst, src []string) []string { - for _, item := range src { - if !containsString(dst, item) { - dst = append(dst, item) - } - } - return dst -} - -func containsString(items []string, value string) bool { - for _, item := range items { - if item == value { - return true - } - } - return false -} - -func intOrNil(v any) any { - switch x := v.(type) { - case nil: - return nil - case int: - return x - case int64: - return int(x) - case float64: - return int(x) - case float32: - return int(x) - case json.Number: - n, err := x.Int64() - if err != nil { - return nil - } - return int(n) - case string: - n, err := strconv.Atoi(strings.TrimSpace(x)) - if err != nil { - return nil - } - return n - default: - return nil - } -} - -func emptyStringNil(v string) any { - v = strings.TrimSpace(v) - if v == "" { - return nil - } - return v -} - -func minInt(a, b int) int { - if a < b { - return a - } - return b -} diff --git a/backend/internal/provider/chatgpt/pow.go b/backend/internal/provider/chatgpt/pow.go deleted file mode 100644 index 0f387f7..0000000 --- a/backend/internal/provider/chatgpt/pow.go +++ /dev/null @@ -1,151 +0,0 @@ -package chatgpt - -import ( - "encoding/base64" - "encoding/hex" - "encoding/json" - "errors" - "math/rand" - "strconv" - "strings" - "time" - - "golang.org/x/crypto/sha3" -) - -var ( - cores = []int{8, 16, 24, 32} - documentKeys = []string{"__reactContainer$fzelfjyxej8", "_reactListening5dehydibo78", "location"} - screenResolutions = [][2]int{{1920, 1080}, {1440, 900}, {2560, 1440}, {3840, 2160}} - navKeys = []string{ - "registerProtocolHandler−function registerProtocolHandler() { [native code] }", - "storage−[object StorageManager]", - "locks−[object LockManager]", - "appCodeName−Mozilla", - "permissions−[object Permissions]", - "share−function share() { [native code] }", - "webdriver−false", - "vendor−Google Inc.", - "mediaDevices−[object MediaDevices]", - "cookieEnabled−true", - "onLine−true", - "mimeTypes−[object MimeTypeArray]", - "credentials−[object CredentialsContainer]", - "serviceWorker−[object ServiceWorkerContainer]", - "keyboard−[object Keyboard]", - "gpu−[object GPU]", - "doNotTrack", - "language−zh-CN", - "geolocation−[object Geolocation]", - "hardwareConcurrency−32", - } - winKeys = []string{ - "0", "window", "self", "document", "name", "location", "history", - "navigation", "innerWidth", "innerHeight", "screen", "chrome", - "navigator", "performance", "crypto", "indexedDB", "sessionStorage", - "localStorage", "fetch", "matchMedia", "postMessage", "setTimeout", - "caches", "__NEXT_DATA__", - } -) - -func buildLegacyRequirementsToken(userAgent string, scriptSources []string, dataBuild string) string { - cfg := buildPOWConfig(userAgent, scriptSources, dataBuild) - body, _ := json.Marshal(cfg) - return "gAAAAAC" + base64.StdEncoding.EncodeToString(body) -} - -func buildProofToken(seed, difficulty, userAgent string, scriptSources []string, dataBuild string) (string, error) { - cfg := buildPOWConfig(userAgent, scriptSources, dataBuild) - answer, solved := powGenerate(seed, difficulty, cfg, 500000) - if !solved { - return "", errors.New("failed to solve proof token") - } - return "gAAAAAB" + answer, nil -} - -func buildPOWConfig(userAgent string, scriptSources []string, dataBuild string) []any { - // scriptSources/dataBuild are no longer part of the sentinel config array - // (the current chatgpt.com client dropped them); kept in the signature for - // call-site compatibility. - _ = scriptSources - _ = dataBuild - rng := rand.New(rand.NewSource(time.Now().UnixNano())) - screen := screenResolutions[rng.Intn(len(screenResolutions))] - loc := time.FixedZone("GMT+0800", 8*3600) - nowLocal := time.Now().In(loc).Format("Mon Jan 02 2006 15:04:05") + " GMT+0800 (中国标准时间)" - perf := float64(time.Now().UnixNano()%1_000_000_000) / 1_000_000 - return []any{ - screen[0] + screen[1], // [0] - nowLocal, // [1] local time, JS Date.toString() shape - 4395630592, // [2] - 1, // [3] overwritten by powGenerate counter - userAgent, // [4] - nil, // [5] (was script source; now null) - defaultClientVersion, // [6] oai-client-version, must match header - "zh-CN", // [7] matches oai-language - "zh-CN,en,en-GB,en-US", // [8] - rng.Float64(), // [9] overwritten by powGenerate counter - navKeys[rng.Intn(len(navKeys))], - documentKeys[rng.Intn(len(documentKeys))], - winKeys[rng.Intn(len(winKeys))], - perf, // [13] - newUUID(), // [14] - "", // [15] - cores[rng.Intn(len(cores))], // [16] - float64(timeMillis()) - perf, // [17] - 0, 0, 0, 0, 0, 0, - 0, - } -} - -func powGenerate(seed, difficulty string, cfg []any, limit int) (string, bool) { - target, err := hex.DecodeString(strings.TrimSpace(difficulty)) - if err != nil { - return "", false - } - diffLen := len(strings.TrimSpace(difficulty)) / 2 - seedBytes := []byte(seed) - head1, _ := json.Marshal(cfg[:3]) - head2, _ := json.Marshal(cfg[4:9]) - head3, _ := json.Marshal(cfg[10:]) - static1 := []byte(string(head1[:len(head1)-1]) + ",") - static2 := []byte("," + string(head2[1:len(head2)-1]) + ",") - static3 := []byte("," + string(head3[1:])) - - for i := 0; i < limit; i++ { - finalJSON := append([]byte{}, static1...) - finalJSON = append(finalJSON, []byte(strconvItoa(i))...) - finalJSON = append(finalJSON, static2...) - finalJSON = append(finalJSON, []byte(strconvItoa(i>>1))...) - finalJSON = append(finalJSON, static3...) - encoded := base64.StdEncoding.EncodeToString(finalJSON) - sum := sha3.Sum512(append(seedBytes, []byte(encoded)...)) - if bytesCompare(sum[:diffLen], target) <= 0 { - return encoded, true - } - } - fallback := "wQ8Lk5FbGpA2NcR9dShT6gYjU7VxZ4D" + base64.StdEncoding.EncodeToString([]byte(`"`+seed+`"`)) - return fallback, false -} - -func bytesCompare(a, b []byte) int { - for i := 0; i < len(a) && i < len(b); i++ { - if a[i] < b[i] { - return -1 - } - if a[i] > b[i] { - return 1 - } - } - if len(a) < len(b) { - return -1 - } - if len(a) > len(b) { - return 1 - } - return 0 -} - -func strconvItoa(v int) string { - return strconv.Itoa(v) -} diff --git a/backend/internal/provider/chatgpt/turnstile.go b/backend/internal/provider/chatgpt/turnstile.go deleted file mode 100644 index db16e3e..0000000 --- a/backend/internal/provider/chatgpt/turnstile.go +++ /dev/null @@ -1,182 +0,0 @@ -package chatgpt - -import ( - "encoding/base64" - "encoding/json" - "math/rand" - "strings" - "time" -) - -type orderedMap struct { - keys []string - values map[string]any -} - -func newOrderedMap() *orderedMap { - return &orderedMap{values: map[string]any{}} -} - -func (m *orderedMap) add(key string, value any) { - if _, ok := m.values[key]; !ok { - m.keys = append(m.keys, key) - } - m.values[key] = value -} - -func solveTurnstileToken(dx, p string) string { - decoded, err := base64.StdEncoding.DecodeString(dx) - if err != nil { - return "" - } - var tokenList [][]any - if err := json.Unmarshal([]byte(xorString(string(decoded), p)), &tokenList); err != nil { - return "" - } - - processMap := map[int]any{16: p} - start := time.Now() - result := "" - rng := rand.New(rand.NewSource(time.Now().UnixNano())) - - toStr := func(value any) string { - if value == nil { - return "undefined" - } - if s, ok := value.(string); ok { - special := map[string]string{ - "window.Math": "[object Math]", - "window.Reflect": "[object Reflect]", - "window.performance": "[object Performance]", - "window.localStorage": "[object Storage]", - "window.Object": "function Object() { [native code] }", - "window.Reflect.set": "function set() { [native code] }", - "window.performance.now": "function () { [native code] }", - "window.Object.create": "function create() { [native code] }", - "window.Object.keys": "function keys() { [native code] }", - "window.Math.random": "function random() { [native code] }", - } - if specialValue, ok := special[s]; ok { - return specialValue - } - return s - } - if list, ok := value.([]string); ok { - return strings.Join(list, ",") - } - return stringValue(value) - } - - for _, token := range tokenList { - if len(token) == 0 { - continue - } - op := intValue(token[0]) - switch op { - case 2: - if len(token) >= 3 { - processMap[intValue(token[1])] = token[2] - } - case 3: - if len(token) >= 2 { - result = base64.StdEncoding.EncodeToString([]byte(toStr(processMap[intValue(token[1])]))) - } - case 5: - if len(token) >= 3 { - e := intValue(token[1]) - t := intValue(token[2]) - cur := processMap[e] - inc := processMap[t] - if list, ok := cur.([]any); ok { - processMap[e] = append(list, inc) - } else if _, ok := cur.(string); ok { - processMap[e] = toStr(cur) + toStr(inc) - } else { - processMap[e] = "NaN" - } - } - case 6, 24: - if len(token) >= 4 { - e := intValue(token[1]) - t := toStr(processMap[intValue(token[2])]) - n := toStr(processMap[intValue(token[3])]) - v := t + "." + n - if op == 6 && v == "window.document.location" { - v = "https://chatgpt.com/" - } - processMap[e] = v - } - case 8: - if len(token) >= 3 { - processMap[intValue(token[1])] = processMap[intValue(token[2])] - } - case 14: - if len(token) >= 3 { - var parsed any - if err := json.Unmarshal([]byte(toStr(processMap[intValue(token[2])])), &parsed); err == nil { - processMap[intValue(token[1])] = parsed - } - } - case 15: - if len(token) >= 3 { - b, _ := json.Marshal(processMap[intValue(token[2])]) - processMap[intValue(token[1])] = string(b) - } - case 17: - if len(token) >= 3 { - e := intValue(token[1]) - target := toStr(processMap[intValue(token[2])]) - switch target { - case "window.performance.now": - processMap[e] = float64(time.Since(start).Nanoseconds())/1e6 + rng.Float64() - case "window.Object.create": - processMap[e] = newOrderedMap() - case "window.Object.keys": - processMap[e] = []string{ - "STATSIG_LOCAL_STORAGE_INTERNAL_STORE_V4", - "STATSIG_LOCAL_STORAGE_STABLE_ID", - "client-correlated-secret", - "oai/apps/capExpiresAt", - "oai-did", - "STATSIG_LOCAL_STORAGE_LOGGING_REQUEST", - "UiState.isNavigationCollapsed.1", - } - case "window.Math.random": - processMap[e] = rng.Float64() - } - } - case 18: - if len(token) >= 2 { - raw, err := base64.StdEncoding.DecodeString(toStr(processMap[intValue(token[1])])) - if err == nil { - processMap[intValue(token[1])] = string(raw) - } - } - case 19: - if len(token) >= 2 { - processMap[intValue(token[1])] = base64.StdEncoding.EncodeToString([]byte(toStr(processMap[intValue(token[1])]))) - } - case 20: - if len(token) >= 4 { - if toStr(processMap[intValue(token[1])]) == toStr(processMap[intValue(token[2])]) { - if intValue(token[3]) == 3 && len(token) >= 5 { - result = base64.StdEncoding.EncodeToString([]byte(toStr(processMap[intValue(token[4])]))) - } - } - } - } - } - return result -} - -func xorString(text, key string) string { - if key == "" { - return text - } - out := make([]rune, 0, len(text)) - keyRunes := []rune(key) - for i, ch := range text { - out = append(out, ch^keyRunes[i%len(keyRunes)]) - } - return string(out) -} diff --git a/backend/internal/provider/chatgpt/util.go b/backend/internal/provider/chatgpt/util.go deleted file mode 100644 index 4ac2fd0..0000000 --- a/backend/internal/provider/chatgpt/util.go +++ /dev/null @@ -1,325 +0,0 @@ -package chatgpt - -import ( - "encoding/base64" - "encoding/json" - "fmt" - "regexp" - "strconv" - "strings" - "time" - - "github.com/google/uuid" -) - -const ( - baseURL = "https://chatgpt.com" - defaultUserAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/149.0.0.0 Safari/537.36 Edg/149.0.0.0" - defaultClientVersion = "prod-db390ebea64862bf1899c420a4c736e0cf639747" - defaultClientBuildNumber = "7904904" - defaultPOWScript = "https://chatgpt.com/backend-api/sentinel/sdk.js" - // sseIdleGrace bounds how long startImageGeneration keeps reading the image - // SSE with no activity at all before giving up. ChatGPT's inline image pipeline - // holds the stream open with periodic ": ping" keepalives (~15s apart) while - // the image renders, so this is an *idle* timeout reset on every received line - // (see the watchdog in startImageGeneration): it only fires when the stream is - // genuinely silent, not on the normal render wait. - sseIdleGrace = 45 * time.Second - - // pictureV2Command is the "@创建图片" ecosystem mention the ChatGPT web app - // prepends to an image (picture_v2) prompt; the server identifies the image - // command from this mention plus its custom_symbol_offset, so we send it to - // match the browser exactly. - pictureV2Command = "@创建图片" - // pictureV2MentionEnd is the UTF-16 length of pictureV2Command — the - // ecosystemMention custom_symbol_offset endIndex the web app reports. - pictureV2MentionEnd = 5 - - // submitConnectRetries / submitRetryBackoff bound how many times the image - // submit (the proxied /backend-api/f/conversation POST) is retried on a bare - // connection error (EOF / reset) before surfacing it. The proxy hop - // occasionally drops the connection; a fresh retry usually succeeds without - // failing over to another account. - submitConnectRetries = 2 - submitRetryBackoff = 500 * time.Millisecond -) - -var ( - fileServiceIDPattern = regexp.MustCompile(`file-service://([A-Za-z0-9_-]+)`) - sedimentIDPattern = regexp.MustCompile(`sediment://([A-Za-z0-9_-]+)`) - realImageIDPattern = regexp.MustCompile(`\bfile_00000000[a-f0-9]{24}\b`) - conversationIDRE = regexp.MustCompile(`"conversation_id"\s*:\s*"([^"]+)"`) - scriptSrcRE = regexp.MustCompile(`]+src="([^"]+)"`) - dataBuildPathRE = regexp.MustCompile(`c/[^/]*/_`) - htmlDataBuildRE = regexp.MustCompile(`]*data-build="([^"]*)"`) - - // asyncMarkers signal that ChatGPT accepted the prompt and switched to the - // async image pipeline (image is delivered later via conversation polling - // rather than inline in the SSE stream). Their presence means "generating — - // keep polling", NOT failure. - asyncMarkers = []string{"image_gen_async", "image_gen_task_id", "trigger_async_ux", "ImageGenToolTemporal", "image_gen_title"} - - // contentPolicyMarkers are stable substrings of ChatGPT's content-audit - // refusal message. When one appears in an assistant turn the prompt was - // rejected upstream — no image will ever arrive, so we must fail fast - // instead of polling until timeout. - contentPolicyMarkers = []string{ - "this request may violate our content polic", - "this prompt may violate our content polic", - "may violate our content policies", - "i can't help with", - "i can\u2019t help with", - "i cannot help with", - "i can't assist with", - "i can\u2019t assist with", - "i cannot assist with", - "i'm unable to help with", - "i\u2019m unable to help with", - "can't create images", - "can\u2019t create images", - "cannot create images", - "can't generate images", - "can\u2019t generate images", - "cannot generate images", - "unable to create images", - "unable to generate images", - } -) - -// containsAsyncMarker reports whether the SSE payload indicates the async image -// pipeline was engaged. -func containsAsyncMarker(text string) bool { - for _, m := range asyncMarkers { - if strings.Contains(text, m) { - return true - } - } - return false -} - -// detectContentPolicyRejection reports whether text contains a ChatGPT content -// audit refusal. Matching is case-insensitive for the English variants. -func detectContentPolicyRejection(text string) bool { - if text == "" { - return false - } - lower := strings.ToLower(text) - for _, m := range contentPolicyMarkers { - if strings.Contains(text, m) || strings.Contains(lower, m) { - return true - } - } - return false -} - -// collectText concatenates every string found under value (recursively) into sb. -func collectText(value any, sb *strings.Builder) { - switch x := value.(type) { - case string: - sb.WriteString(x) - sb.WriteByte('\n') - case map[string]any: - for _, item := range x { - collectText(item, sb) - } - case []any: - for _, item := range x { - collectText(item, sb) - } - } -} - -// conversationRejected scans assistant turns of a fetched conversation for a -// content-policy refusal. -func conversationRejected(conversation map[string]any) bool { - mapping, _ := conversation["mapping"].(map[string]any) - for _, rawNode := range mapping { - node, _ := rawNode.(map[string]any) - message, _ := node["message"].(map[string]any) - if message == nil { - continue - } - author, _ := message["author"].(map[string]any) - role := strings.ToLower(strings.TrimSpace(stringValue(author["role"]))) - if role != "assistant" { - continue - } - var sb strings.Builder - collectText(message["content"], &sb) - if detectContentPolicyRejection(sb.String()) { - return true - } - } - return false -} - -// conversationEndedWithoutImage reports whether the model finished its turn -// with a plain-text answer and never engaged the async image pipeline (no tool -// turn, no image_gen task). ChatGPT localizes audit refusals, so this catches -// rejections in ANY language structurally: a closed text-only turn can never -// produce an image, and polling further only burns the budget. -func conversationEndedWithoutImage(conversation map[string]any) bool { - mapping, _ := conversation["mapping"].(map[string]any) - if len(mapping) == 0 { - return false - } - for _, rawNode := range mapping { - node, _ := rawNode.(map[string]any) - message, _ := node["message"].(map[string]any) - if message == nil { - continue - } - if metadata, _ := message["metadata"].(map[string]any); metadata != nil { - if stringValue(metadata["image_gen_task_id"]) != "" || metadata["image_gen_async"] == true { - return false - } - } - author, _ := message["author"].(map[string]any) - role := strings.ToLower(strings.TrimSpace(stringValue(author["role"]))) - if role == "tool" { - return false - } - if role == "assistant" { - if content, _ := message["content"].(map[string]any); content != nil { - // Only tool-invoking content types mean generation is underway. - // Context nodes (model_editable_context, thoughts, …) also appear - // on refused turns and must NOT suppress the detection. - switch stringValue(content["content_type"]) { - case "code", "multimodal_text": - return false - } - } - } - } - for _, rawNode := range mapping { - node, _ := rawNode.(map[string]any) - message, _ := node["message"].(map[string]any) - if message == nil { - continue - } - author, _ := message["author"].(map[string]any) - if strings.ToLower(strings.TrimSpace(stringValue(author["role"]))) != "assistant" { - continue - } - if message["end_turn"] != true || stringValue(message["status"]) != "finished_successfully" { - continue - } - var sb strings.Builder - collectText(message["content"], &sb) - if strings.TrimSpace(sb.String()) != "" { - return true - } - } - return false -} - -// isConnResetErr reports whether err is a bare connection-level failure (the -// peer/proxy closed the connection before responding) rather than an HTTP-level -// error. Such failures are safe to retry: no response was received, so no -// conversation was created upstream. -func isConnResetErr(err error) bool { - if err == nil { - return false - } - msg := strings.ToLower(err.Error()) - for _, m := range []string{"eof", "connection reset", "reset by peer", "broken pipe", "connection refused"} { - if strings.Contains(msg, m) { - return true - } - } - return false -} - -func stringValue(v any) string { - switch x := v.(type) { - case string: - return x - case nil: - return "" - default: - return fmt.Sprint(v) - } -} - -func intValue(v any) int { - switch x := v.(type) { - case int: - return x - case int64: - return int(x) - case float64: - return int(x) - case float32: - return int(x) - case json.Number: - n, _ := x.Int64() - return int(n) - case string: - n, _ := strconv.Atoi(strings.TrimSpace(x)) - return n - default: - return 0 - } -} - -func decodeJWTPayload(token string) map[string]any { - parts := strings.Split(strings.TrimSpace(token), ".") - 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 newUUID() string { - return uuid.NewString() -} - -func clip(v []byte, n int) string { - s := strings.TrimSpace(string(v)) - if len(s) <= n { - return s - } - return s[:n] -} - -func parsePOWResources(html string) ([]string, string) { - matches := scriptSrcRE.FindAllStringSubmatch(html, -1) - sources := make([]string, 0, len(matches)) - dataBuild := "" - for _, match := range matches { - if len(match) < 2 { - continue - } - src := strings.TrimSpace(match[1]) - if src == "" { - continue - } - sources = append(sources, src) - if dataBuild == "" { - if path := dataBuildPathRE.FindString(src); path != "" { - dataBuild = path - } - } - } - if dataBuild == "" { - if match := htmlDataBuildRE.FindStringSubmatch(html); len(match) >= 2 { - dataBuild = strings.TrimSpace(match[1]) - } - } - if len(sources) == 0 { - sources = []string{defaultPOWScript} - } - return sources, dataBuild -} - -func timeMillis() int64 { - return time.Now().UnixMilli() -} diff --git a/backend/internal/provider/creativefabrica/client.go b/backend/internal/provider/creativefabrica/client.go deleted file mode 100644 index 1a1e685..0000000 --- a/backend/internal/provider/creativefabrica/client.go +++ /dev/null @@ -1,629 +0,0 @@ -// Package creativefabrica implements the Creative Fabrica Studio -// (studio.creativefabrica.com) video-generation upstream. -// -// One-shot accounts: every account's coins are just enough for exactly one -// generation, so a successful render kills the account. The credential is a -// .creativefabrica.com session cookie; a short-lived JWT is minted from it on -// demand via GraphQL /query/userAuth, and every model request authenticates -// with that JWT (the cookie is sent along as a fallback). -package creativefabrica - -import ( - "bytes" - "context" - "crypto/rand" - "encoding/base64" - "encoding/json" - "errors" - "fmt" - "io" - "math/big" - "net/http" - "strings" - "time" - - fhttp "github.com/bogdanfinn/fhttp" - tlsclient "github.com/bogdanfinn/tls-client" - "github.com/bogdanfinn/tls-client/profiles" -) - -var ( - ErrAuth = errors.New("creativefabrica auth failed") - ErrQuotaExhausted = errors.New("creativefabrica quota exhausted") - ErrTemporaryUpstream = errors.New("creativefabrica upstream temporary error") - ErrDeadUpstream = errors.New("creativefabrica upstream fatal error") - ErrRateLimited = errors.New("creativefabrica rate limited") - // ErrPaymentRequired marks an account whose payment intent is in a failed - // state — it can never generate (the studio answers 400 failed_precondition - // "payment required ... COIN_PAYMENT_INTENT_STATUS_FAILED"). It wraps ErrAuth - // so the pool kills the account and fails over instead of burning retries. - ErrPaymentRequired = fmt.Errorf("%w: payment required", ErrAuth) -) - -// isPaymentRequired reports whether a non-200 body is the account-level -// "payment required" rejection rather than a request-level parameter error. -// -// Connect unary errors don't always carry the marker in plaintext: the studio -// answers failed_precondition with the real detail base64-protobuf-encoded in -// details[].value ("payment required. payment status: COIN_PAYMENT_INTENT_..."). -// Decode those values and scan the decoded bytes, so a failed coin intent still -// kills the account instead of being misread as a request-level 400. -func isPaymentRequired(status int, body string) bool { - if status != 400 { - return false - } - b := strings.ToLower(body) - if strings.Contains(b, "payment required") || strings.Contains(b, "coin_payment_intent") { - return true - } - var env struct { - Details []struct { - Value string `json:"value"` - } `json:"details"` - } - if err := json.Unmarshal([]byte(body), &env); err != nil { - return false - } - for _, d := range env.Details { - raw, err := base64.StdEncoding.DecodeString(d.Value) - if err != nil { - continue - } - low := strings.ToLower(string(raw)) - if strings.Contains(low, "payment required") || - strings.Contains(low, "coin_payment_intent") || - strings.Contains(low, "coin_error_code_payment_required") { - return true - } - } - return false -} - -const ( - graphQLHost = "https://graphql-gw.creativefabrica.com" - mediaMatrixHost = "https://studio-media-matrix.creativefabrica.com" - userAuthPath = "/query/userAuth" - userBalancePath = "/query/userBalance" - userPath = "/query/user" - initiatePath = "/creativefabrica.studiomediamatrix.v1.StudioMediaMatrixService/InitiateSession" - listSessionsPath = "/creativefabrica.studiomediamatrix.v1.StudioMediaMatrixService/ListSessions" - origin = "https://studio.creativefabrica.com" - pollInterval = 5 * time.Second - pollTimeout = 16 * time.Minute - downloadTimeout = 3 * time.Minute - videoServiceType = "SERVICE_TYPE_VIDEO_GENERATOR" - videoFrameRef = "VIDEO_FRAME_TYPE_REFERENCE" - visibilityPrivate = "SESSION_VISIBILITY_PRIVATE" -) - -// Model is one Creative Fabrica video model: the local catalog id, the upstream -// enum, the fixed duration in seconds, and the upstream resolution label. -type Model struct { - ID string // local model_configs id, e.g. "seedance-2.0" - Enum string // upstream enum, e.g. VIDEO_GENERATOR_MODEL_BYTEDANCE_SEEDDREAM_2 - Duration int // fixed seconds (account plan is fixed-length) - Resolution string // upstream resolution label, e.g. 720p -} - -// Models returns the two Creative Fabrica seedance models. The upstream enum -// differs between the two (SEEDANCE_2_FAST vs SEEDDREAM_2), matching the -// studio frontend's InitiateSession payloads. -func Models() map[string]Model { - return map[string]Model{ - "seedance-2.0-fast": { - ID: "seedance-2.0-fast", - Enum: "VIDEO_GENERATOR_MODEL_BYTEDANCE_SEEDANCE_2_FAST", - Duration: 14, - Resolution: "720p", - }, - "seedance-2.0": { - ID: "seedance-2.0", - Enum: "VIDEO_GENERATOR_MODEL_BYTEDANCE_SEEDDREAM_2", - Duration: 10, - Resolution: "720p", - }, - } -} - -// LookupModel resolves a local model id to its upstream config. Returns ok=false -// when the id isn't a Creative Fabrica model. -func LookupModel(modelID string) (Model, bool) { - m, ok := Models()[modelID] - return m, ok -} - -// Client talks to the Creative Fabrica Studio API through a Chrome-fingerprinted -// TLS client so the Cloudflare-protected Connect endpoints don't reject us. -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) -} - -// ExchangeToken mints the short-lived JWT from the account cookie via GraphQL -// /query/userAuth. Returns the token and the user id. A null / missing me means -// the cookie no longer authenticates → ErrAuth. -func (c *Client) ExchangeToken(ctx context.Context, cookie string) (token, userID string, err error) { - query := `{"query":"\n query userAuth {\n me {\n token\n user {\n id\n isTemporary\n }\n }\n}\n "}` - var payload struct { - Data struct { - Me *struct { - Token string `json:"token"` - User struct { - ID string `json:"id"` - } `json:"user"` - } `json:"me"` - } `json:"data"` - } - body, err := c.postGraphQL(ctx, cookie, "", userAuthPath, query) - if err != nil { - return "", "", err - } - if err := json.Unmarshal(body, &payload); err != nil { - return "", "", fmt.Errorf("%w: bad token response: %s", ErrAuth, clip(body, 300)) - } - if payload.Data.Me == nil || strings.TrimSpace(payload.Data.Me.Token) == "" { - return "", "", fmt.Errorf("%w: cookie did not authenticate", ErrAuth) - } - return payload.Data.Me.Token, payload.Data.Me.User.ID, nil -} - -// FetchBalance reads the coin balance via GraphQL /query/userBalance (the -// request authenticates with the cookie alone). Negative value on error. -func (c *Client) FetchBalance(ctx context.Context, cookie string) (int64, error) { - query := `{"query":"\nquery userBalance {\n userBalance {\n balance\n }\n}\n\n"}` - var payload struct { - Data struct { - UserBalance *struct { - Balance json.Number `json:"balance"` - } `json:"userBalance"` - } `json:"data"` - } - body, err := c.postGraphQL(ctx, cookie, "", userBalancePath, query) - if err != nil { - return -1, err - } - if err := json.Unmarshal(body, &payload); err != nil { - return -1, fmt.Errorf("bad balance response: %s", clip(body, 300)) - } - if payload.Data.UserBalance == nil { - return -1, ErrAuth - } - b, _ := payload.Data.UserBalance.Balance.Int64() - return b, nil -} - -// FetchUser reads the profile (email, name) via GraphQL /query/user, which the -// studio browser hits on every page load. The request authenticates with the -// cookie alone. Returns the email (empty on error); a null me means the cookie -// no longer authenticates → ErrAuth. -func (c *Client) FetchUser(ctx context.Context, cookie string) (string, error) { - query := `{"query":"\n query user {\n me {\n token\n user {\n id\n email\n }\n }\n}\n "}` - var payload struct { - Data struct { - Me *struct { - Token string `json:"token"` - User struct { - ID string `json:"id"` - Email string `json:"email"` - } `json:"user"` - } `json:"me"` - } `json:"data"` - } - body, err := c.postGraphQL(ctx, cookie, "", userPath, query) - if err != nil { - return "", err - } - if err := json.Unmarshal(body, &payload); err != nil { - return "", fmt.Errorf("%w: bad user response: %s", ErrAuth, clip(body, 300)) - } - if payload.Data.Me == nil || payload.Data.Me.User.ID == "" { - return "", ErrAuth - } - return strings.TrimSpace(payload.Data.Me.User.Email), nil -} - -// GenerateVideo runs the full generation: InitiateSession → PUT reference -// images to the presigned S3 URLs → poll ListSessions until COMPLETED → (when -// downloadResult) fetch the MP4. Returns the bytes (nil when url-only) and the -// previewMediaUrl. durationSeconds is ignored: the plan fixes the length per -// model (LookupModel.Duration). -func (c *Client) GenerateVideo(ctx context.Context, cookie, token, modelID, prompt, aspectRatio string, refs [][]byte, downloadResult bool) ([]byte, string, error) { - m, ok := LookupModel(modelID) - if !ok { - return nil, "", fmt.Errorf("creativefabrica: unknown model %q", modelID) - } - sessionID, uploads, err := c.initiateSession(ctx, cookie, token, m, prompt, aspectRatio, refs) - if err != nil { - return nil, "", err - } - // Each ref uploads to the presigned URL the session returned for it. - for i, u := range uploads { - if err := c.putS3(ctx, u, refs[i]); err != nil { - return nil, "", err - } - } - videoURL, err := c.pollSession(ctx, cookie, token, sessionID) - if err != nil { - return nil, "", err - } - if !downloadResult { - return nil, videoURL, nil - } - data, err := c.download(ctx, videoURL) - if err != nil { - return nil, "", err - } - return data, videoURL, nil -} - -// initiateSession submits the generation and returns the session id plus the -// presigned upload URLs (one per reference image). -func (c *Client) initiateSession(ctx context.Context, cookie, token string, m Model, prompt, aspectRatio string, refs [][]byte) (string, []string, error) { - frames := make([]any, 0, len(refs)) - refPrompt := strings.TrimSpace(prompt) - for i := range refs { - ref := fmt.Sprintf("img%d", i+1) - frames = append(frames, map[string]any{ - "type": videoFrameRef, - "fileSize": fmt.Sprintf("%d", len(refs[i])), - "fileName": randomFileName(i + 1), - "ref": ref, - }) - if !strings.Contains(refPrompt, "["+ref+"]") { - refPrompt += " Use [" + ref + "]" - } - } - reqBody := map[string]any{ - "visibility": visibilityPrivate, - "sessionRequestPromptToVideoGeneratorContent": map[string]any{ - "serviceType": videoServiceType, - "promptContent": map[string]any{"prompt": refPrompt}, - "resolution": resolutionEnum(m.Resolution), - "model": m.Enum, - "frames": frames, - "aspectRatio": aspectRatioEnum(aspectRatio), - "directorConfig": map[string]any{ - "filmStock": map[string]any{"color": "DIRECTOR_FILM_STOCK_COLOR_FULL_COLOR"}, - }, - "videoDuration": map[string]any{"inSeconds": m.Duration}, - }, - } - body, err := json.Marshal(reqBody) - if err != nil { - return "", nil, err - } - resp, err := c.postConnect(ctx, cookie, token, initiatePath, body, true) - if err != nil { - return "", nil, err - } - var payload struct { - Session struct { - ID string `json:"id"` - Status string `json:"status"` - // The presigned upload URLs are nested under promptToVideoGeneratorContent.frames - // (each frame echoes the request frame + the S3 presigned `url`), not under a - // top-level session.frames. Reading the wrong level yields 0 uploads. - PromptToVideoGeneratorContent struct { - Frames []struct { - URL string `json:"url"` - } `json:"frames"` - } `json:"promptToVideoGeneratorContent"` - } `json:"session"` - } - if err := json.Unmarshal(resp, &payload); err != nil { - return "", nil, fmt.Errorf("creativefabrica bad initiate response: %s", clip(resp, 300)) - } - if strings.TrimSpace(payload.Session.ID) == "" { - return "", nil, fmt.Errorf("creativefabrica initiate missing session: %s", clip(resp, 300)) - } - respFrames := payload.Session.PromptToVideoGeneratorContent.Frames - uploads := make([]string, 0, len(respFrames)) - for _, f := range respFrames { - if u := strings.TrimSpace(f.URL); u != "" { - uploads = append(uploads, u) - } - } - if len(uploads) != len(refs) { - return "", nil, fmt.Errorf("creativefabrica initiate returned %d upload urls for %d refs", len(uploads), len(refs)) - } - return payload.Session.ID, uploads, nil -} - -// pollSession polls ListSessions until the session is COMPLETED / FAILED and -// returns previewMediaUrl on success. -func (c *Client) pollSession(ctx context.Context, cookie, token, sessionID string) (string, error) { - reqBody, _ := json.Marshal(map[string]any{ - "serviceType": videoServiceType, - "pagination": map[string]any{"take": 100}, - "surface": "SURFACE_STUDIO", - }) - deadline := time.Now().Add(pollTimeout) - for { - if err := ctx.Err(); err != nil { - return "", err - } - if time.Now().After(deadline) { - return "", fmt.Errorf("creativefabrica generation timed out after %v", pollTimeout) - } - resp, err := c.postConnect(ctx, cookie, token, listSessionsPath, reqBody, false) - if err != nil { - return "", err - } - var payload struct { - Sessions []map[string]any `json:"sessions"` - } - if err := json.Unmarshal(resp, &payload); err != nil { - return "", fmt.Errorf("creativefabrica bad list response: %s", clip(resp, 300)) - } - for _, s := range payload.Sessions { - id, _ := s["id"].(string) - if id != sessionID { - continue - } - status, _ := s["status"].(string) - switch status { - case "SESSION_STATUS_COMPLETED": - if u := strings.TrimSpace(stringValue(s["previewMediaUrl"])); u != "" { - return u, nil - } - return "", fmt.Errorf("creativefabrica session completed without preview url") - case "SESSION_STATUS_FAILED", "SESSION_STATUS_CANCELLED", "SESSION_STATUS_ERROR": - detail := "" - for _, k := range []string{"errorMessage", "failureReason", "error", "message"} { - if v, ok := s[k]; ok { - if d := strings.TrimSpace(stringValue(v)); d != "" { - detail = d - break - } - } - } - return "", fmt.Errorf("creativefabrica session %s%s", status, withDetail(detail)) - } - } - select { - case <-ctx.Done(): - return "", ctx.Err() - case <-time.After(pollInterval): - } - } -} - -// putS3 uploads a reference image to the presigned URL. S3 doesn't care about -// TLS fingerprinting, so a plain client is fine here. -func (c *Client) putS3(ctx context.Context, presigned string, data []byte) error { - req, err := http.NewRequestWithContext(ctx, http.MethodPut, presigned, bytes.NewReader(data)) - if err != nil { - return err - } - req.Header.Set("Content-Type", "image/png") - req.Header.Set("Origin", origin) - req.Header.Set("Referer", origin+"/") - req.Header.Set("User-Agent", defaultUA()) - resp, err := http.DefaultClient.Do(req) - if err != nil { - return fmt.Errorf("creativefabrica s3 upload: %w", err) - } - defer resp.Body.Close() - if resp.StatusCode != 200 { - body, _ := io.ReadAll(resp.Body) - return fmt.Errorf("creativefabrica s3 upload failed: %d %s", resp.StatusCode, clip(body, 200)) - } - return nil -} - -// download fetches the finished MP4 from the public video-v2 URL. -func (c *Client) download(parent context.Context, videoURL string) ([]byte, error) { - ctx, cancel := context.WithTimeout(context.WithoutCancel(parent), downloadTimeout) - defer cancel() - req, err := http.NewRequestWithContext(ctx, http.MethodGet, videoURL, nil) - if err != nil { - return nil, err - } - req.Header.Set("Accept", "*/*") - req.Header.Set("User-Agent", defaultUA()) - resp, err := http.DefaultClient.Do(req) - if err != nil { - return nil, fmt.Errorf("%w: download: %v", ErrTemporaryUpstream, err) - } - defer resp.Body.Close() - if resp.StatusCode != 200 { - return nil, fmt.Errorf("creativefabrica download failed: %d", resp.StatusCode) - } - data, err := io.ReadAll(resp.Body) - if err != nil { - return nil, fmt.Errorf("%w: read body: %v", ErrTemporaryUpstream, err) - } - return data, nil -} - -// postGraphQL calls a graphql-gw endpoint with the raw JSON body. Authenticated -// by cookie (balance) and optionally the JWT too. Direct connection (no proxy). -func (c *Client) postGraphQL(ctx context.Context, cookie, token, path, body string) ([]byte, error) { - return c.postJSON(ctx, graphQLHost+path, token, cookie, []byte(body), false, true, false) -} - -// postConnect calls a Connect unary endpoint on the media-matrix service. The -// Connect-Protocol-Version header marks the POST as a Connect RPC (distinct from -// a plain JSON REST POST) — the studio browser always sends it. Only the -// InitiateSession call (下单) egresses through the proxy; polling runs direct. -func (c *Client) postConnect(ctx context.Context, cookie, token, path string, body []byte, useProxy bool) ([]byte, error) { - return c.postJSON(ctx, mediaMatrixHost+path, token, cookie, body, true, false, useProxy) -} - -func (c *Client) postJSON(ctx context.Context, url, token, cookie string, body []byte, connect, graphql, useProxy bool) ([]byte, error) { - sess, err := c.newTLSClient(useProxy) - if err != nil { - return nil, err - } - req, err := fhttp.NewRequest(fhttp.MethodPost, url, bytes.NewReader(body)) - if err != nil { - return nil, err - } - req = req.WithContext(ctx) - req.Header = fhttp.Header{ - "content-type": {"application/json"}, - "accept": {"*/*"}, - "origin": {origin}, - "referer": {origin + "/"}, - "user-agent": {defaultUA()}, - } - if graphql { - req.Header.Set("accept", "application/json, multipart/mixed") - } - if connect { - req.Header.Set("connect-protocol-version", "1") - } - if cookie != "" { - req.Header.Set("cookie", cookie) - } - if token != "" { - req.Header.Set("authorization", "Bearer "+token) - } - resp, err := sess.client.Do(req) - if err != nil { - return nil, fmt.Errorf("creativefabrica request: %w", err) - } - defer resp.Body.Close() - data, err := io.ReadAll(resp.Body) - if err != nil { - return nil, err - } - switch { - case resp.StatusCode == 401 || resp.StatusCode == 403: - return nil, fmt.Errorf("%w (%d: %s)", ErrAuth, resp.StatusCode, clip(data, 300)) - case isPaymentRequired(resp.StatusCode, string(data)): - return nil, fmt.Errorf("%w (%d: %s)", ErrPaymentRequired, resp.StatusCode, clip(data, 300)) - case resp.StatusCode == 429: - return nil, fmt.Errorf("%w (429)", ErrRateLimited) - case resp.StatusCode >= 500: - return nil, fmt.Errorf("%w (%d: %s)", ErrDeadUpstream, resp.StatusCode, clip(data, 300)) - case resp.StatusCode != 200: - return nil, fmt.Errorf("creativefabrica %d: %s", resp.StatusCode, clip(data, 300)) - } - return data, nil -} - -func resolutionEnum(res string) string { - switch strings.ToLower(strings.TrimSpace(res)) { - case "720p": - return "PROMPT_TO_VIDEO_GENERATOR_CONTENT_RESOLUTION_720P" - case "1080p": - return "PROMPT_TO_VIDEO_GENERATOR_CONTENT_RESOLUTION_1080P" - default: - return "PROMPT_TO_VIDEO_GENERATOR_CONTENT_RESOLUTION_720P" - } -} - -func aspectRatioEnum(ratio string) string { - switch strings.ReplaceAll(strings.TrimSpace(ratio), " ", "") { - case "16:9": - return "PROMPT_TO_VIDEO_GENERATOR_CONTENT_ASPECT_RATIO_16_9" - case "9:16": - return "PROMPT_TO_VIDEO_GENERATOR_CONTENT_ASPECT_RATIO_9_16" - case "1:1": - return "PROMPT_TO_VIDEO_GENERATOR_CONTENT_ASPECT_RATIO_1_1" - case "4:3": - return "PROMPT_TO_VIDEO_GENERATOR_CONTENT_ASPECT_RATIO_4_3" - case "3:4": - return "PROMPT_TO_VIDEO_GENERATOR_CONTENT_ASPECT_RATIO_3_4" - default: - return "PROMPT_TO_VIDEO_GENERATOR_CONTENT_ASPECT_RATIO_16_9" - } -} - -// randomFileName mimics the studio's client-generated reference name -// ("__.png"). The value only needs to be unique per upload. -func randomFileName(n int) string { - return "_" + randomID() + "_" + fmt.Sprintf("%d", n) + ".png" -} - -func randomID() string { - const alphabet = "abcdefghijklmnopqrstuvwxyz0123456789" - b := make([]byte, 26) - for i := range b { - v, _ := rand.Int(rand.Reader, big.NewInt(int64(len(alphabet)))) - b[i] = alphabet[v.Int64()] - } - return string(b) -} - -// TLS client plumbing: a Chrome-fingerprinted client so the Cloudflare front of -// the Connect endpoints sees a plausible browser handshake. -var fingerprints = []profiles.ClientProfile{ - profiles.Chrome_146, - profiles.Chrome_144, - profiles.Chrome_133, - profiles.Chrome_131, -} - -type tlsSession struct { - client tlsclient.HttpClient -} - -func (c *Client) newTLSClient(useProxy bool) (*tlsSession, error) { - idx, _ := rand.Int(rand.Reader, big.NewInt(int64(len(fingerprints)))) - options := []tlsclient.HttpClientOption{ - tlsclient.WithTimeoutSeconds(60), - tlsclient.WithClientProfile(fingerprints[idx.Int64()]), - tlsclient.WithNotFollowRedirects(), - tlsclient.WithRandomTLSExtensionOrder(), - } - if useProxy && c.proxy != "" { - options = append(options, tlsclient.WithProxyUrl(c.proxy)) - } - client, err := tlsclient.NewHttpClient(tlsclient.NewNoopLogger(), options...) - if err != nil { - return nil, err - } - return &tlsSession{client: client}, nil -} - -func defaultUA() string { - return "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/146.0.0.0 Safari/537.36" -} - -func clip(v []byte, n int) string { - s := strings.TrimSpace(string(v)) - if len(s) <= n { - return s - } - return s[:n] -} - -// stringValue coerces a decoded JSON value to its string form. -func stringValue(v any) string { - switch x := v.(type) { - case string: - return x - case float64: - return fmt.Sprintf("%.0f", x) - case json.Number: - return x.String() - case bool: - if x { - return "true" - } - return "false" - case nil: - return "" - default: - return fmt.Sprintf("%v", x) - } -} - -// withDetail appends a failure detail to an error message when present. -func withDetail(detail string) string { - if detail == "" { - return "" - } - return " (" + detail + ")" -} - - diff --git a/backend/internal/provider/creativefabrica/gen_test.go b/backend/internal/provider/creativefabrica/gen_test.go deleted file mode 100644 index 558b614..0000000 --- a/backend/internal/provider/creativefabrica/gen_test.go +++ /dev/null @@ -1,27 +0,0 @@ -package creativefabrica - -import ( - "context" - "os" - "testing" - "time" -) - -func TestBalance(t *testing.T) { - cookie := os.Getenv("CF_COOKIE") - if cookie == "" { - t.Skip("CF_COOKIE not set") - } - c := NewClient("") - ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) - defer cancel() - tok, uid, err := c.ExchangeToken(ctx, cookie) - if err != nil { - t.Fatalf("exchange: %v", err) - } - bal, err := c.FetchBalance(ctx, cookie) - if err != nil { - t.Fatalf("balance: %v", err) - } - t.Logf("user=%s token=%d bal=%d", uid, len(tok), bal) -} diff --git a/backend/internal/provider/custom/client.go b/backend/internal/provider/custom/client.go deleted file mode 100644 index e6cdc5b..0000000 --- a/backend/internal/provider/custom/client.go +++ /dev/null @@ -1,721 +0,0 @@ -// Package custom implements a generic OpenAI-compatible upstream client. A -// "custom" model forwards generation to any OpenAI-compatible API: the upstream -// base_url + api_key live on a custom account (pool="custom"), the upstream model -// name on the model config (UpstreamModel). Calls go DIRECT (no tls-client, no -// proxy) — the upstream is a normal API with no anti-bot. -package custom - -import ( - "bufio" - "bytes" - "context" - "encoding/json" - "errors" - "fmt" - "io" - "mime/multipart" - "net/http" - "net/url" - "strings" - "time" -) - -var ( - ErrAuth = errors.New("custom upstream auth failed") - ErrQuotaExhausted = errors.New("custom upstream quota exhausted") - ErrTemporaryUpstream = errors.New("custom upstream temporary error") -) - -type Client struct{} - -func NewClient() *Client { return &Client{} } - -// sanitizeErr strips the upstream URL/host from a network error so a user's -// private upstream URL never leaks into the event log / API response. -func sanitizeErr(err error) string { - if err == nil { - return "" - } - s := err.Error() - switch { - case strings.Contains(s, "context deadline exceeded"), strings.Contains(s, "Client.Timeout"), strings.Contains(s, "timeout"): - return "request timeout" - case strings.Contains(s, "connection refused"): - return "connection refused" - case strings.Contains(s, "no such host"), strings.Contains(s, "dial tcp"), strings.Contains(s, "lookup "): - return "cannot reach upstream" - case strings.Contains(s, "tls"), strings.Contains(s, "TLS"), strings.Contains(s, "certificate"): - return "TLS error" - case strings.Contains(s, "EOF"), strings.Contains(s, "reset by peer"), strings.Contains(s, "broken pipe"): - return "connection reset" - } - var ue *url.Error - if errors.As(err, &ue) { - return strings.ToLower(ue.Op) + " upstream failed" - } - return "upstream request failed" -} - -func httpClient() *http.Client { return &http.Client{Timeout: 30 * time.Minute} } - -// GenerateImage calls the upstream OpenAI image API. With reference images it -// uses /v1/images/edits (multipart); otherwise /v1/images/generations. Returns -// the raw image bytes (decoded from b64_json, or downloaded from url). -func (c *Client) GenerateImage(ctx context.Context, baseURL, apiKey, model, prompt, size, quality string, refs [][]byte, downloadResult bool) ([]byte, string, error) { - baseURL = strings.TrimRight(strings.TrimSpace(baseURL), "/") - if baseURL == "" || apiKey == "" { - return nil, "", ErrAuth - } - var req *http.Request - var err error - if len(refs) > 0 { - body := &bytes.Buffer{} - w := multipart.NewWriter(body) - _ = w.WriteField("model", model) - _ = w.WriteField("prompt", prompt) - // Ask the upstream for a URL (not base64) so we can pass it through directly. - _ = w.WriteField("response_format", "url") - if size != "" { - _ = w.WriteField("size", size) - } - for i, r := range refs { - fw, e := w.CreateFormFile("image[]", fmt.Sprintf("ref_%d.png", i+1)) - if e != nil { - return nil, "", e - } - _, _ = fw.Write(r) - } - _ = w.Close() - req, err = http.NewRequest(http.MethodPost, baseURL+"/v1/images/edits", body) - if err != nil { - return nil, "", err - } - req.Header.Set("Content-Type", w.FormDataContentType()) - } else { - payload := map[string]any{"model": model, "prompt": prompt, "n": 1, "response_format": "url"} - if size != "" { - payload["size"] = size - } - if quality != "" { - payload["quality"] = quality - } - raw, _ := json.Marshal(payload) - req, err = http.NewRequest(http.MethodPost, baseURL+"/v1/images/generations", bytes.NewReader(raw)) - if err != nil { - return nil, "", err - } - req.Header.Set("Content-Type", "application/json") - } - req = req.WithContext(ctx) - req.Header.Set("Authorization", "Bearer "+apiKey) - - resp, err := httpClient().Do(req) - if err != nil { - return nil, "", fmt.Errorf("%w: %s", ErrTemporaryUpstream, sanitizeErr(err)) - } - defer resp.Body.Close() - body, _ := io.ReadAll(resp.Body) - if e := mapStatus(resp.StatusCode, body); e != nil { - return nil, "", e - } - return imageFromResponse(ctx, body, downloadResult) -} - -// ── Vivid 协议 ────────────────────────────────────────────────────────────── -// Vivid 上游用自家 API:POST /v1/generate 返回 application/x-ndjson 流 —— 首行 -// {"jobId":...,"status":"running"},等待期间每 25s 一个空行心跳,末行为终态 -// (succeeded 带 result 预签名 URL / failed 带 error)。参考素材走两步: -// POST /v1/uploads {kind,mime} → PUT 直传预签名 URL → 拿 key 放进 refs。 -// 断流可携 jobId 重发重连。错误体为 {"error":{"code":"4011|4001|4002|4293|5001"}}。 - -// GenerateVividImage submits an image job to a Vivid-protocol upstream. -// quality is the tier string the upstream expects ("1K"/"2K"/"4K"), ratio the -// "W:H" aspect. refs are pre-uploaded via /v1/uploads into key strings. -func (c *Client) GenerateVividImage(ctx context.Context, baseURL, apiKey, model, prompt, ratio, quality string, refs [][]byte, downloadResult bool) ([]byte, string, error) { - payload, err := c.vividPayload(ctx, baseURL, apiKey, model, prompt, ratio, quality, 0, refs) - if err != nil { - return nil, "", err - } - return c.vividGenerate(ctx, baseURL, apiKey, payload, downloadResult) -} - -// GenerateVividVideo submits a video job (duration required by ranged models) -// to a Vivid-protocol upstream. -func (c *Client) GenerateVividVideo(ctx context.Context, baseURL, apiKey, model, prompt, ratio, quality string, durationSeconds int, refs [][]byte, downloadResult bool) ([]byte, string, error) { - payload, err := c.vividPayload(ctx, baseURL, apiKey, model, prompt, ratio, quality, durationSeconds, refs) - if err != nil { - return nil, "", err - } - return c.vividGenerate(ctx, baseURL, apiKey, payload, downloadResult) -} - -func (c *Client) vividPayload(ctx context.Context, baseURL, apiKey, model, prompt, ratio, quality string, durationSeconds int, refs [][]byte) (map[string]any, error) { - baseURL = strings.TrimRight(strings.TrimSpace(baseURL), "/") - if baseURL == "" || apiKey == "" { - return nil, ErrAuth - } - payload := map[string]any{"model": model, "prompt": prompt} - if quality != "" { - payload["quality"] = quality - } - if ratio != "" { - payload["ratio"] = ratio - } - if durationSeconds > 0 { - payload["duration"] = durationSeconds - } - if len(refs) > 0 { - keys := make([]string, 0, len(refs)) - for _, ref := range refs { - key, err := c.vividUploadRef(ctx, baseURL, apiKey, ref) - if err != nil { - return nil, err - } - keys = append(keys, key) - } - payload["refs"] = keys - } - return payload, nil -} - -// vividUploadRef runs the two-step upload: POST /v1/uploads {kind,mime} → PUT -// the bytes to the presigned url → returns the key for refs[]. kind is image / -// video / audio, derived from the file magic (same classes the upstream -// accepts: png/jpeg/webp · mp4/webm/mov · mp3/wav/m4a/ogg). -func (c *Client) vividUploadRef(ctx context.Context, baseURL, apiKey string, data []byte) (string, error) { - kind, mime := vividRefKindMime(data) - raw, _ := json.Marshal(map[string]any{"kind": kind, "mime": mime}) - up, err := c.doJSON(ctx, http.MethodPost, baseURL+"/v1/uploads", apiKey, raw) - if err != nil { - return "", err - } - putURL, key := vividUploadSession(up) - if putURL == "" || key == "" { - raw, _ := json.Marshal(up) - return "", fmt.Errorf("%w: upload session missing url/key (got %s)", ErrTemporaryUpstream, clip(raw, 200)) - } - req, err := http.NewRequest(http.MethodPut, putURL, bytes.NewReader(data)) - if err != nil { - return "", err - } - req = req.WithContext(ctx) - req.Header.Set("Content-Type", mime) // must match the signed mime exactly - resp, err := httpClient().Do(req) - if err != nil { - return "", fmt.Errorf("%w: %s", ErrTemporaryUpstream, sanitizeErr(err)) - } - defer resp.Body.Close() - body, _ := io.ReadAll(resp.Body) - if resp.StatusCode < 200 || resp.StatusCode >= 300 { - return "", fmt.Errorf("%w: ref upload %d %s", ErrTemporaryUpstream, resp.StatusCode, clip(body, 120)) - } - return key, nil -} - -// vividUploadSession extracts the presigned PUT url + object key from a -// /v1/uploads response. Vivid-protocol implementations vary: some return the -// pair at top level, others wrap it in {"data":{...}} / {"result":{...}}, and -// field spellings differ (upload_url / presigned_url / put_url; key / -// file_key / object_key). Search top level first, then one layer deep. -func vividUploadSession(up map[string]any) (putURL, key string) { - urlKeys := []string{"url", "upload_url", "uploadUrl", "presigned_url", "presignedUrl", "put_url", "putUrl", "signed_url", "signedUrl"} - keyKeys := []string{"key", "file_key", "fileKey", "object_key", "objectKey", "upload_key", "uploadKey"} - pick := func(m map[string]any) (string, string) { - var u, k string - for _, f := range urlKeys { - if u = strings.TrimSpace(stringValue(m[f])); u != "" { - break - } - } - for _, f := range keyKeys { - if k = strings.TrimSpace(stringValue(m[f])); k != "" { - break - } - } - return u, k - } - levels := []map[string]any{up} - for _, nest := range []string{"data", "result", "upload", "session", "upload_session"} { - if sub, ok := up[nest].(map[string]any); ok { - levels = append(levels, sub) - } - } - // url 和 key 允许跨层出现(比如 url 在顶层、key 在 data 里)。 - for _, m := range levels { - u, k := pick(m) - if putURL == "" { - putURL = u - } - if key == "" { - key = k - } - } - return putURL, key -} - -// vividRefKindMime classifies a reference blob into the vivid "kind" (image / -// video / audio) plus a concrete mime. Magic bytes first (DetectContentType -// misses mov/m4a/ogg-audio), then sniffing, then a png default. -func vividRefKindMime(data []byte) (kind, mime string) { - n := len(data) - // ISOBMFF box: ftyp — mp4/mov/m4a share the container; the brand decides. - if n >= 12 && string(data[4:8]) == "ftyp" { - brand := string(data[8:12]) - switch { - case strings.Contains(brand, "M4A"), strings.Contains(brand, "M4B"): - return "audio", "audio/mp4" - case strings.Contains(brand, "qt"): - return "video", "video/quicktime" - default: - return "video", "video/mp4" - } - } - // WebM/MKV EBML - if n >= 4 && data[0] == 0x1A && data[1] == 0x45 && data[2] == 0xDF && data[3] == 0xA3 { - return "video", "video/webm" - } - // MP3: ID3 tag or MPEG sync - if n >= 3 && string(data[:3]) == "ID3" { - return "audio", "audio/mpeg" - } - if n >= 2 && data[0] == 0xFF && (data[1]&0xE0) == 0xE0 { - return "audio", "audio/mpeg" - } - // WAV / OGG / FLAC - if n >= 4 && string(data[:4]) == "RIFF" { - return "audio", "audio/wav" - } - if n >= 4 && string(data[:4]) == "OggS" { - return "audio", "audio/ogg" - } - if n >= 4 && string(data[:4]) == "fLaC" { - return "audio", "audio/x-flac" - } - // Sniffed type (images land here; also catches anything odd above missed). - if m := http.DetectContentType(data); m != "" && m != "application/octet-stream" { - switch { - case strings.HasPrefix(m, "video/"): - return "video", m - case strings.HasPrefix(m, "audio/"), m == "application/ogg": - return "audio", m - case strings.HasPrefix(m, "image/"): - return "image", m - } - } - return "image", "image/png" -} - -// vividGenerate POSTs /v1/generate and reads the NDJSON stream to the terminal -// frame. If the stream drops before terminal, it reconnects once with the -// jobId (same-request resume — no new task, no double charge). -func (c *Client) vividGenerate(ctx context.Context, baseURL, apiKey string, payload map[string]any, downloadResult bool) ([]byte, string, error) { - var last map[string]any - for attempt := 0; attempt < 3; attempt++ { - if attempt > 0 { - if sleepCtx(ctx, 2*time.Second) != nil { - return nil, "", ctx.Err() - } - } - term, err := c.vividStreamOnce(ctx, baseURL, apiKey, payload) - if err == nil && term != nil { - last = term - break - } - if err != nil && !errors.Is(err, ErrTemporaryUpstream) { - return nil, "", err - } - // Temporary error mid-stream: if we hold a jobId, reconnect mode retries - // without creating a new task. - if job := strings.TrimSpace(stringValue(payload["jobId"])); job != "" { - payload = map[string]any{"jobId": job} - continue - } - if attempt == 2 { - return nil, "", err - } - } - if last == nil { - return nil, "", fmt.Errorf("%w: no terminal frame", ErrTemporaryUpstream) - } - status := strings.ToLower(strings.TrimSpace(stringValue(last["status"]))) - switch status { - case "succeeded", "completed", "success": - resultURL := strings.TrimSpace(stringValue(last["result"])) - if resultURL == "" { - return nil, "", fmt.Errorf("%w: succeeded frame missing result", ErrTemporaryUpstream) - } - if !downloadResult { - return nil, resultURL, nil - } - data, err := c.vividDownload(ctx, resultURL) - if err != nil { - return nil, "", err - } - return data, resultURL, nil - default: - reason := strings.TrimSpace(stringValue(last["error"])) - if reason == "" { - reason = "upstream job " + status - } - if isCreditError(reason) || strings.Contains(reason, "积分不足") || strings.Contains(reason, "余额") { - return nil, "", fmt.Errorf("%w: %s", ErrQuotaExhausted, clip([]byte(reason), 160)) - } - return nil, "", errors.New("custom: " + clip([]byte(reason), 160)) - } -} - -// vividStreamOnce performs one /v1/generate request and scans the NDJSON body -// to the last frame. Returns the terminal frame, or a temporary error on -// transport failure (caller's reconnect logic decides what to do with jobId). -func (c *Client) vividStreamOnce(ctx context.Context, baseURL, apiKey string, payload map[string]any) (map[string]any, error) { - raw, _ := json.Marshal(payload) - req, err := http.NewRequest(http.MethodPost, baseURL+"/v1/generate", bytes.NewReader(raw)) - if err != nil { - return nil, err - } - req = req.WithContext(ctx) - req.Header.Set("Authorization", "Bearer "+apiKey) - req.Header.Set("Content-Type", "application/json") - req.Header.Set("X-Wait-Seconds", "7200") - resp, err := httpClient().Do(req) - if err != nil { - return nil, fmt.Errorf("%w: %s", ErrTemporaryUpstream, sanitizeErr(err)) - } - defer resp.Body.Close() - if resp.StatusCode < 200 || resp.StatusCode >= 300 { - body, _ := io.ReadAll(resp.Body) - return nil, vividError(resp.StatusCode, body) - } - var last map[string]any - sc := bufio.NewScanner(resp.Body) - sc.Buffer(make([]byte, 0, 64*1024), 4*1024*1024) - for sc.Scan() { - line := bytes.TrimSpace(sc.Bytes()) - if len(line) == 0 { - continue // 25s heartbeat empty line - } - var frame map[string]any - if json.Unmarshal(line, &frame) != nil { - continue - } - last = frame - // Remember the jobId for potential reconnect. - if job := strings.TrimSpace(stringValue(frame["jobId"])); job != "" { - payload["jobId"] = job - } - } - if err := sc.Err(); err != nil { - return nil, fmt.Errorf("%w: %s", ErrTemporaryUpstream, sanitizeErr(err)) - } - return last, nil -} - -// vividError maps the {"error":{"code":"NNNN"}} envelope to our error classes. -func vividError(status int, body []byte) error { - var e struct { - Error struct { - Code string `json:"code"` - Message string `json:"message"` - } `json:"error"` - } - code := "" - msg := clip(body, 160) - if json.Unmarshal(body, &e) == nil { - code = strings.TrimSpace(e.Error.Code) - if m := strings.TrimSpace(e.Error.Message); m != "" { - msg = clip([]byte(m), 160) - } - } - switch code { - case "4011": - return fmt.Errorf("%w: %s", ErrAuth, msg) - case "4002": - return fmt.Errorf("%w: %s", ErrQuotaExhausted, msg) - case "4293", "5001": - return fmt.Errorf("%w: %s", ErrTemporaryUpstream, msg) - case "4001": - return errors.New("custom: " + msg) // 参数/业务拒绝 → 永久错,不换号 - } - return mapStatus(status, body) -} - -// vividDownload fetches the presigned result URL — no auth header (it's signed). -func (c *Client) vividDownload(ctx context.Context, url string) ([]byte, error) { - req, _ := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) - resp, err := httpClient().Do(req) - if err != nil { - return nil, fmt.Errorf("%w: %s", ErrTemporaryUpstream, sanitizeErr(err)) - } - defer resp.Body.Close() - if resp.StatusCode < 200 || resp.StatusCode >= 300 { - return nil, fmt.Errorf("%w: result 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 result", ErrTemporaryUpstream) - } - return data, nil -} - -// GenerateVideo drives the upstream Sora-style async video API: -// POST /v1/videos → poll GET /v1/videos/{id} → GET /v1/videos/{id}/content. -// Reference frames (image-to-video / first-last frames) are sent as multipart -// input_reference[] files, matching the OpenAI videos API. When downloadResult -// is false it returns the upstream content URL instead. -func (c *Client) GenerateVideo(ctx context.Context, baseURL, apiKey, model, prompt, size string, seconds int, frames [][]byte, downloadResult bool) ([]byte, string, error) { - baseURL = strings.TrimRight(strings.TrimSpace(baseURL), "/") - if baseURL == "" || apiKey == "" { - return nil, "", ErrAuth - } - var created map[string]any - var err error - if len(frames) > 0 { - body := &bytes.Buffer{} - w := multipart.NewWriter(body) - _ = w.WriteField("model", model) - _ = w.WriteField("prompt", prompt) - if size != "" { - _ = w.WriteField("size", size) - } - if seconds > 0 { - _ = w.WriteField("seconds", fmt.Sprintf("%d", seconds)) - } - for i, f := range frames { - fw, e := w.CreateFormFile("input_reference[]", fmt.Sprintf("frame_%d.png", i+1)) - if e != nil { - return nil, "", e - } - _, _ = fw.Write(f) - } - _ = w.Close() - created, err = c.doMultipart(ctx, baseURL+"/v1/videos", apiKey, body, w.FormDataContentType()) - } else { - payload := map[string]any{"model": model, "prompt": prompt} - if size != "" { - payload["size"] = size - } - if seconds > 0 { - payload["seconds"] = fmt.Sprintf("%d", seconds) - } - raw, _ := json.Marshal(payload) - created, err = c.doJSON(ctx, http.MethodPost, baseURL+"/v1/videos", apiKey, raw) - } - if err != nil { - return nil, "", err - } - jobID := strings.TrimSpace(stringValue(created["id"])) - if jobID == "" { - return nil, "", fmt.Errorf("%w: video create missing id", ErrTemporaryUpstream) - } - // Poll until terminal. - for { - if err := ctx.Err(); err != nil { - return nil, "", err - } - job, err := c.doJSON(ctx, http.MethodGet, baseURL+"/v1/videos/"+jobID, apiKey, nil) - if err != nil { - if errors.Is(err, ErrTemporaryUpstream) { - if sleepCtx(ctx, 5*time.Second) != nil { - return nil, "", ctx.Err() - } - continue - } - return nil, "", err - } - switch strings.ToLower(strings.TrimSpace(stringValue(job["status"]))) { - case "completed", "succeeded", "success": - contentURL := baseURL + "/v1/videos/" + jobID + "/content" - if !downloadResult { - return nil, contentURL, nil - } - data, err := c.download(ctx, contentURL, apiKey) - if err != nil { - return nil, "", err - } - return data, contentURL, nil - case "failed", "error", "canceled", "cancelled": - reason := stringValue(job["error"]) - if isCreditError(reason) { - return nil, "", fmt.Errorf("%w: %s", ErrTemporaryUpstream, clip([]byte(reason), 160)) - } - return nil, "", fmt.Errorf("custom: video %s", clip([]byte(reason), 160)) - } - if sleepCtx(ctx, 5*time.Second) != nil { - return nil, "", ctx.Err() - } - } -} - -func (c *Client) doJSON(ctx context.Context, method, url, apiKey string, body []byte) (map[string]any, error) { - var reader io.Reader - if body != nil { - reader = bytes.NewReader(body) - } - req, err := http.NewRequest(method, url, reader) - if err != nil { - return nil, err - } - req = req.WithContext(ctx) - req.Header.Set("Authorization", "Bearer "+apiKey) - if body != nil { - req.Header.Set("Content-Type", "application/json") - } - resp, err := httpClient().Do(req) - if err != nil { - return nil, fmt.Errorf("%w: %s", ErrTemporaryUpstream, sanitizeErr(err)) - } - defer resp.Body.Close() - raw, _ := io.ReadAll(resp.Body) - if e := mapStatus(resp.StatusCode, 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: non-json: %s", ErrTemporaryUpstream, clip(raw, 120)) - } - return out, nil -} - -func (c *Client) doMultipart(ctx context.Context, url, apiKey string, body io.Reader, contentType string) (map[string]any, error) { - req, err := http.NewRequest(http.MethodPost, url, body) - if err != nil { - return nil, err - } - req = req.WithContext(ctx) - req.Header.Set("Authorization", "Bearer "+apiKey) - req.Header.Set("Content-Type", contentType) - resp, err := httpClient().Do(req) - if err != nil { - return nil, fmt.Errorf("%w: %s", ErrTemporaryUpstream, sanitizeErr(err)) - } - defer resp.Body.Close() - raw, _ := io.ReadAll(resp.Body) - if e := mapStatus(resp.StatusCode, 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: non-json: %s", ErrTemporaryUpstream, clip(raw, 120)) - } - return out, nil -} - -func (c *Client) download(ctx context.Context, url, apiKey string) ([]byte, error) { - req, _ := http.NewRequest(http.MethodGet, url, nil) - req = req.WithContext(ctx) - req.Header.Set("Authorization", "Bearer "+apiKey) - resp, err := httpClient().Do(req) - if err != nil { - return nil, fmt.Errorf("%w: %s", ErrTemporaryUpstream, sanitizeErr(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 download", ErrTemporaryUpstream) - } - return data, nil -} - -// imageBytesFromResponse extracts image bytes from an OpenAI images response: -// data[0].b64_json (preferred) or data[0].url (downloaded). -// imageFromResponse parses an OpenAI image response and returns the upstream URL. -// We always request response_format=url, so the response must carry a URL — a -// base64-only response is treated as an error (no base64 pass-through). With -// downloadResult=false the URL is returned directly (no download). -func imageFromResponse(ctx context.Context, body []byte, downloadResult bool) ([]byte, string, error) { - var out struct { - Data []struct { - URL string `json:"url"` - } `json:"data"` - } - if err := json.Unmarshal(body, &out); err != nil || len(out.Data) == 0 { - return nil, "", fmt.Errorf("%w: bad image response: %s", ErrTemporaryUpstream, clip(body, 160)) - } - url := strings.TrimSpace(out.Data[0].URL) - if url == "" { - return nil, "", fmt.Errorf("%w: image response had no url (upstream ignored response_format=url)", ErrTemporaryUpstream) - } - if !downloadResult { - return nil, url, nil - } - req, _ := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) - resp, err := httpClient().Do(req) - if err != nil { - return nil, "", fmt.Errorf("%w: %s", ErrTemporaryUpstream, sanitizeErr(err)) - } - defer resp.Body.Close() - raw, err := io.ReadAll(resp.Body) - return raw, url, err -} - -func mapStatus(status int, body []byte) error { - switch { - case status >= 200 && status < 300: - return nil - case status == 401 || status == 403: - return fmt.Errorf("%w: %d %s", ErrAuth, status, clip(body, 160)) - case status == 429: - // Custom upstreams have NO "quota exhausted" lock — a 429 is just rate - // limiting, treated as a temporary error (fail over, account stays active). - return fmt.Errorf("%w: 429 %s", ErrTemporaryUpstream, clip(body, 160)) - case status >= 500: - return fmt.Errorf("%w: %d %s", ErrTemporaryUpstream, status, clip(body, 160)) - default: - if isCreditError(string(body)) { - return fmt.Errorf("%w: %s", ErrTemporaryUpstream, clip(body, 160)) - } - return fmt.Errorf("custom: %d %s", status, clip(body, 160)) - } -} - -func isCreditError(s string) bool { - s = strings.ToLower(s) - return strings.Contains(s, "insufficient") || strings.Contains(s, "quota") || - strings.Contains(s, "credit") || strings.Contains(s, "balance") -} - -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 clip(b []byte, n int) string { - s := strings.TrimSpace(string(b)) - if len(s) > n { - return s[:n] - } - return s -} - -func sleepCtx(ctx context.Context, d time.Duration) error { - t := time.NewTimer(d) - defer t.Stop() - select { - case <-ctx.Done(): - return ctx.Err() - case <-t.C: - return nil - } -} diff --git a/backend/internal/provider/custom/client_vivid_test.go b/backend/internal/provider/custom/client_vivid_test.go deleted file mode 100644 index 53d112b..0000000 --- a/backend/internal/provider/custom/client_vivid_test.go +++ /dev/null @@ -1,229 +0,0 @@ -package custom - -import ( - "context" - "encoding/json" - "errors" - "fmt" - "io" - "net/http" - "net/http/httptest" - "strings" - "testing" -) - -// vividServer serves /v1/uploads, /v1/generate (ndjson) and a fake presigned -// result URL, all under one httptest server. -func vividServer(t *testing.T, frames []string, httpStatus int, errCode string) *httptest.Server { - t.Helper() - return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - switch r.URL.Path { - case "/v1/uploads": - var req struct { - Kind string `json:"kind"` - Mime string `json:"mime"` - } - _ = json.NewDecoder(r.Body).Decode(&req) - if req.Kind != "image" || req.Mime == "" { - t.Errorf("bad upload request: %+v", req) - } - w.Header().Set("Content-Type", "application/json") - fmt.Fprintf(w, `{"uuid":"u1","key":"u/refs/image/u1.png","url":"%s","expiresAt":9999999999}`, - // presigned URL points back at this server's /put endpoint - "http://"+r.Host+"/put") - case "/put": - if r.Method != http.MethodPut { - t.Errorf("expected PUT, got %s", r.Method) - } - if ct := r.Header.Get("Content-Type"); !strings.HasPrefix(ct, "image/") { - t.Errorf("bad content-type %q", ct) - } - io.Copy(io.Discard, r.Body) - w.WriteHeader(http.StatusOK) - case "/v1/generate": - if httpStatus != 0 { - w.Header().Set("Content-Type", "application/json") - w.WriteHeader(httpStatus) - fmt.Fprintf(w, `{"error":{"code":%q,"message":"boom"}}`, errCode) - return - } - w.Header().Set("Content-Type", "application/x-ndjson") - for _, f := range frames { - fmt.Fprintln(w, f) - fmt.Fprintln(w) // heartbeat line - } - case "/result.png": - w.Write([]byte("PNGDATA")) - default: - http.NotFound(w, r) - } - })) -} - -func TestVividGenerateImageSuccess(t *testing.T) { - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - switch r.URL.Path { - case "/v1/generate": - w.Header().Set("Content-Type", "application/x-ndjson") - fmt.Fprintln(w, `{"jobId":"j1","status":"running"}`) - fmt.Fprintln(w) // heartbeat - fmt.Fprintf(w, `{"jobId":"j1","status":"succeeded","result":"http://%s/result.png","creditsCharged":10}`+"\n", r.Host) - case "/result.png": - w.Write([]byte("PNGDATA")) - default: - http.NotFound(w, r) - } - })) - defer srv.Close() - - data, url, err := NewClient().GenerateVividImage(context.Background(), srv.URL, "vk_test", "m1", "hi", "1:1", "1K", nil, true) - if err != nil { - t.Fatalf("generate: %v", err) - } - if string(data) != "PNGDATA" { - t.Fatalf("data=%q", data) - } - if url == "" { - t.Fatal("empty url") - } -} - -func TestVividGenerateImageURLPassthrough(t *testing.T) { - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Content-Type", "application/x-ndjson") - fmt.Fprintln(w, `{"jobId":"j1","status":"running"}`) - fmt.Fprintln(w, `{"jobId":"j1","status":"succeeded","result":"https://fs.example.com/x.png"}`) - })) - defer srv.Close() - data, url, err := NewClient().GenerateVividImage(context.Background(), srv.URL, "vk_test", "m1", "hi", "1:1", "1K", nil, false) - if err != nil { - t.Fatalf("generate: %v", err) - } - if data != nil { - t.Fatalf("expected no bytes, got %d", len(data)) - } - if url != "https://fs.example.com/x.png" { - t.Fatalf("url=%q", url) - } -} - -func TestVividGenerateFailedFrame(t *testing.T) { - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Content-Type", "application/x-ndjson") - fmt.Fprintln(w, `{"jobId":"j1","status":"running"}`) - fmt.Fprintln(w, `{"jobId":"j1","status":"failed","error":"积分不足","refunded":true}`) - })) - defer srv.Close() - _, _, err := NewClient().GenerateVividImage(context.Background(), srv.URL, "vk_test", "m1", "hi", "1:1", "1K", nil, false) - if !errors.Is(err, ErrQuotaExhausted) { - t.Fatalf("want quota, got %v", err) - } -} - -func TestVividErrorCodes(t *testing.T) { - cases := []struct { - status int - code string - want error - }{ - {401, "4011", ErrAuth}, - {402, "4002", ErrQuotaExhausted}, - {429, "4293", ErrTemporaryUpstream}, - {500, "5001", ErrTemporaryUpstream}, - } - for _, tc := range cases { - srv := vividServer(t, nil, tc.status, tc.code) - _, _, err := NewClient().GenerateVividImage(context.Background(), srv.URL, "vk_test", "m1", "hi", "1:1", "1K", nil, false) - srv.Close() - if !errors.Is(err, tc.want) { - t.Fatalf("code %s: want %v, got %v", tc.code, tc.want, err) - } - } -} - -func TestVividRefKindMime(t *testing.T) { - cases := []struct { - name string - data []byte - kind string - }{ - {"png", []byte{0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, 1, 2, 3, 4}, "image"}, - {"jpeg", []byte{0xFF, 0xD8, 0xFF, 0xE0, 1, 2, 3, 4, 5, 6}, "image"}, - {"mp4", []byte{0, 0, 0, 0, 'f', 't', 'y', 'p', 'i', 's', 'o', 'm'}, "video"}, - {"mov", []byte{0, 0, 0, 0, 'f', 't', 'y', 'p', 'q', 't', ' ', ' '}, "video"}, - {"webm", []byte{0x1A, 0x45, 0xDF, 0xA3, 1, 2, 3, 4}, "video"}, - {"mp3-id3", []byte{'I', 'D', '3', 4, 0, 0, 1, 2}, "audio"}, - {"mp3-sync", []byte{0xFF, 0xFB, 1, 2, 3, 4}, "audio"}, - {"wav", []byte{'R', 'I', 'F', 'F', 1, 2, 3, 4, 'W', 'A', 'V', 'E'}, "audio"}, - {"ogg", []byte{'O', 'g', 'g', 'S', 1, 2, 3, 4}, "audio"}, - {"m4a", []byte{0, 0, 0, 0, 'f', 't', 'y', 'p', 'M', '4', 'A', ' '}, "audio"}, - {"empty", []byte{}, "image"}, - } - for _, tc := range cases { - kind, mime := vividRefKindMime(tc.data) - if kind != tc.kind { - t.Errorf("%s: kind=%s mime=%s", tc.name, kind, mime) - } - } -} - -func TestVividUploadRefKindVideo(t *testing.T) { - var gotKind, gotMime string - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.URL.Path == "/v1/uploads" { - var req map[string]any - _ = json.NewDecoder(r.Body).Decode(&req) - gotKind, _ = req["kind"].(string) - gotMime, _ = req["mime"].(string) - fmt.Fprintf(w, `{"key":"k1","url":"http://%s/up"}`, r.Host) - return - } - w.WriteHeader(http.StatusOK) - })) - defer srv.Close() - webm := []byte{0x1A, 0x45, 0xDF, 0xA3, 1, 2, 3, 4, 5, 6, 7, 8} - key, err := NewClient().vividUploadRef(context.Background(), srv.URL, "vk", webm) - if err != nil { - t.Fatalf("upload: %v", err) - } - if key != "k1" || gotKind != "video" || gotMime != "video/webm" { - t.Fatalf("key=%q kind=%q mime=%q", key, gotKind, gotMime) - } -} - -func TestVividUploadRefFlow(t *testing.T) { - var putOK bool - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - switch r.URL.Path { - case "/v1/uploads": - w.Header().Set("Content-Type", "application/json") - fmt.Fprintf(w, `{"key":"k123","url":"http://%s/up"}`, r.Host) - case "/up": - if r.Method == http.MethodPut { - putOK = true - } - w.WriteHeader(http.StatusOK) - case "/v1/generate": - var req map[string]any - _ = json.NewDecoder(r.Body).Decode(&req) - refs, _ := req["refs"].([]any) - if len(refs) != 1 || refs[0] != "k123" { - t.Errorf("refs=%v", req["refs"]) - } - w.Header().Set("Content-Type", "application/x-ndjson") - fmt.Fprintln(w, `{"jobId":"j1","status":"succeeded","result":"https://x/1.png"}`) - } - })) - defer srv.Close() - png := []byte{0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A} // fake png magic - _, url, err := NewClient().GenerateVividImage(context.Background(), srv.URL, "vk_test", "m1", "hi", "1:1", "1K", [][]byte{png}, false) - if err != nil { - t.Fatalf("generate: %v", err) - } - if !putOK { - t.Fatal("PUT never happened") - } - if url == "" { - t.Fatal("no url") - } -} diff --git a/backend/internal/provider/grok/client.go b/backend/internal/provider/grok/client.go deleted file mode 100644 index 07a8de6..0000000 --- a/backend/internal/provider/grok/client.go +++ /dev/null @@ -1,935 +0,0 @@ -// 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; its value -// is a 70-byte anti-bot record — header[49] | counter_le32 | -// sha256("METHOD!path!counter"+suffix)[:16] | trailer — XOR-masked with one -// random byte and base64-encoded; we reproduce it exactly in statsigID (the -// build-specific header/suffix/trailer are env-overridable when grok rotates -// them). Uses tls-client so the JA3/JA4 fingerprint matches Chrome. -package grok - -import ( - "context" - "crypto/sha256" - "encoding/base64" - "encoding/binary" - "encoding/hex" - "encoding/json" - "errors" - "fmt" - "io" - "log" - "math" - "math/rand/v2" - "os" - "regexp" - "strconv" - "strings" - "sync" - "time" - - 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.newDirectTLSClient() - if err != nil { - return nil, err - } - c.ensureChallenge(ctx, client, token) - // 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. - // Field #5 carries the credits' own reset timestamp (weekly grant refill) — - // this is the 恢复时间 we surface, NOT the subscription's billing-period end. - used, resetUnix, ok := parseCreditsConfig(raw) - if !ok { - return unknownBalance("unparsable credits config"), nil - } - if used < 0 { - used = 0 - } - if used > fullCredits { - used = fullCredits - } - remaining := fullCredits - used - - // 恢复时间: the credits-config weekly reset (when the free grant refills). - reset := strings.TrimSpace(resetUnix) - return map[string]any{ - "remaining": remaining, - "used": used, - "total": fullCredits, - "reset_after": emptyStringNil(reset), - "unknown": false, - "error": nil, - }, nil -} - -// Subscription is the membership view parsed from GET /rest/subscriptions. -type Subscription struct { - Member bool // a subscription with status ACTIVE exists (INACTIVE entries don't count) - Tier string // e.g. SUBSCRIPTION_TIER_GROK_PRO ("" for free) - Status string // e.g. SUBSCRIPTION_STATUS_ACTIVE - BillingPeriodEnd string // RFC3339; when the plan renews / credits reset - FreeTrial bool // currently in a free-trial offer -} - -// FetchSubscription reads GET /rest/subscriptions and reports the account's -// membership. Member is true only when an entry with SUBSCRIPTION_STATUS_ACTIVE -// exists: a lapsed membership keeps its entry but flips to -// SUBSCRIPTION_STATUS_INACTIVE, and an empty array means never subscribed — -// both read as Member=false (the entry's tier/status are still surfaced). -// A 401/403 maps to ErrAuth; other transport/HTTP errors are returned so callers -// can treat them as best-effort (they already have the credit balance). -func (c *Client) FetchSubscription(ctx context.Context, token string) (*Subscription, error) { - token = strings.TrimSpace(strings.TrimPrefix(token, "Bearer ")) - if token == "" { - return nil, ErrAuth - } - client, err := c.newDirectTLSClient() - if err != nil { - return nil, err - } - c.ensureChallenge(ctx, client, token) - req, err := http.NewRequest(http.MethodGet, apiBase+"/rest/subscriptions", nil) - if err != nil { - return nil, err - } - req = req.WithContext(ctx) - c.applyHeaders(req, token, nil) - resp, err := client.Do(req) - if err != nil { - return nil, fmt.Errorf("%w: %v", ErrTemporaryUpstream, err) - } - defer resp.Body.Close() - raw, _ := io.ReadAll(resp.Body) - if resp.StatusCode == 401 || resp.StatusCode == 403 { - return nil, ErrAuth - } - if resp.StatusCode != 200 { - return nil, fmt.Errorf("%w: subscriptions http %d", ErrTemporaryUpstream, resp.StatusCode) - } - var body struct { - Subscriptions []struct { - Tier string `json:"tier"` - Status string `json:"status"` - BillingPeriodEnd string `json:"billingPeriodEnd"` - ActiveOffer struct { - FreeTrial *struct { - TrialDays int `json:"trialDays"` - } `json:"freeTrial"` - } `json:"activeOffer"` - } `json:"subscriptions"` - } - if err := json.Unmarshal(raw, &body); err != nil { - return nil, fmt.Errorf("%w: subscriptions non-json", ErrTemporaryUpstream) - } - out := &Subscription{} - // Surface the ACTIVE subscription if any (falling back to the first entry - // for tier/status info), but only an ACTIVE one sets Member. - for i, s := range body.Subscriptions { - active := strings.EqualFold(s.Status, "SUBSCRIPTION_STATUS_ACTIVE") - if i == 0 || active { - out.Member = active - out.Tier = strings.TrimSpace(s.Tier) - out.Status = strings.TrimSpace(s.Status) - out.BillingPeriodEnd = strings.TrimSpace(s.BillingPeriodEnd) - out.FreeTrial = s.ActiveOffer.FreeTrial != nil - if active { - break - } - } - } - return out, 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.newDirectTLSClient() - if err != nil { - return "", "", err - } - c.ensureChallenge(ctx, client, token) - 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 { - Status string `json:"status"` - 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) - } - // A dead sso cookie answers 200 {"status":"unauthenticated"} rather than 401. - if !strings.EqualFold(body.Status, "authenticated") { - return "", "", ErrAuth - } - return strings.TrimSpace(body.Session.Email), strings.TrimSpace(body.Session.UserID), nil -} - -// grok's x-statsig-id is a per-session anti-bot token. Its 3-byte "F" tail is a -// browser fingerprint the server recomputes from the homepage seed + curve set, -// and the byte-indexing that derives it ROTATES on every grok web reship — so any -// hand-ported algorithm goes stale within a day (403 anti-bot). The durable path -// therefore runs grok's OWN signer in goja (statsig_engine.go): we fetch the seed -// + curves from the homepage and let grok's code do all the (rotating) indexing. -// statsigID prefers that engine; the hand-ported computeStatsigTail below and the -// static env-overridable defaults are only a last-resort fallback. See the package -// doc and statsig_engine.go for the full picture. -// statsigEpoch is the challenge epoch (2023-05-01 00:00 UTC). -const ( - statsigEpoch = 1682924400 - defaultStatsigHeader = "00a1adb5012bd32f844f4426c62680d91c6129361eb9459a759710e179a888a99b21678e1f0b1e8952de6a6b3ca019f74b" - defaultStatsigSuffix = "obfiowerehiringd244100f5c28f5c28f5c047ae147ae147b047ae147ae147b0f5c28f5c28f5c00" - defaultStatsigTrailer = 3 - statsigSaltPrefix = "obfiowerehiring" - statsigAnimDuration = 4096 - statsigTTL = 5 * time.Minute -) - -var ( - statsigHeader = resolveStatsigHeader() - statsigSuffix = envOr("GROK_STATSIG_SUFFIX", defaultStatsigSuffix) - statsigTrailer = resolveStatsigTrailer() - - statsigMetaRe = regexp.MustCompile(`name="grok[^"]*verification"[^>]*content="([^"]+)"`) - - statsigMu sync.Mutex - statsigCache = map[string]statsigChallenge{} // keyed by sso token -) - -// statsigChallenge is a resolved, self-consistent (header, salt) pair for one -// grok session, derived from the homepage seed + curves. -type statsigChallenge struct { - header []byte - suffix string - trailer byte - fetchedAt time.Time - - // Inputs for the durable goja signer (statsig_js.go): the raw seed - // content and the curves JSON. Empty when parsing failed (engine then skipped). - seedB64 string - curvesJSON string -} - -// statsigCurve is one entry of the per-load curve set injected via the Next.js -// RSC stream; the server uses it (with the seed) to recompute F. -type statsigCurve struct { - Color []int `json:"color"` - Deg int `json:"deg"` - Bezier []int `json:"bezier"` -} - -func resolveStatsigHeader() []byte { - h := envOr("GROK_STATSIG_HEADER_HEX", defaultStatsigHeader) - b, err := hex.DecodeString(h) - if err != nil || len(b) != 49 { - b, _ = hex.DecodeString(defaultStatsigHeader) - } - return b -} - -func resolveStatsigTrailer() byte { - if v := strings.TrimSpace(os.Getenv("GROK_STATSIG_TRAILER")); v != "" { - if n, err := strconv.Atoi(v); err == nil && n >= 0 && n <= 255 { - return byte(n) - } - } - return defaultStatsigTrailer -} - -func envOr(key, def string) string { - if v := strings.TrimSpace(os.Getenv(key)); v != "" { - return v - } - return def -} - -// ensureChallenge refreshes the cached (header, salt) for the session if missing -// or stale. Any failure is non-fatal: statsigID then falls back to the static -// defaults. An explicit env override disables dynamic fetching entirely. -func (c *Client) ensureChallenge(ctx context.Context, client tlsclient.HttpClient, token string) { - if token == "" || client == nil { - return - } - if os.Getenv("GROK_STATSIG_HEADER_HEX") != "" || os.Getenv("GROK_STATSIG_SUFFIX") != "" { - return - } - statsigMu.Lock() - cur, ok := statsigCache[token] - fresh := ok && time.Since(cur.fetchedAt) < statsigTTL - statsigMu.Unlock() - if fresh { - return - } - ch, err := fetchStatsigChallenge(ctx, client) - if err != nil { - // Silent fallback to static defaults is the #1 cause of a recurring - // "403 anti-bot": the homepage structure changed and we never notice. - // Surface it so the failure mode (fetch/parse broke vs. offsets rotated) - // is diagnosable from logs instead of guessing. - log.Printf("grok statsig: self-heal failed, using stale static defaults (403 likely): %v", err) - return - } - log.Printf("grok statsig: self-heal ok header[:6]=%x suffix=%s", ch.header[:6], ch.suffix) - statsigMu.Lock() - statsigCache[token] = ch - statsigMu.Unlock() -} - -// fetchStatsigChallenge does a browser-free homepage GET and derives a -// self-consistent (header, salt) pair: header = 0x00 + seed, salt = prefix + F. -// -// The request is ANONYMOUS (no sso cookie): grok redirects a request carrying an -// sso cookie to the sign-in page, whose HTML omits the anti-bot seed and -// the curves array — so a cookied fetch silently loses the challenge inputs and -// statsigID falls back to the stale static recipe (403 anti-bot). The seed/curves -// are per-build and account-independent (the seed is transmitted inside the -// x-statsig-id header the server re-verifies), so the anonymous landing page is -// the reliable source for every session. -func fetchStatsigChallenge(ctx context.Context, client tlsclient.HttpClient) (statsigChallenge, error) { - req, err := http.NewRequest(http.MethodGet, apiBase+"/", nil) - if err != nil { - return statsigChallenge{}, err - } - req = req.WithContext(ctx) - req.Header = http.Header{ - "accept": {"text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8"}, - "accept-language": {"en-US,en;q=0.9"}, - "user-agent": {userAgent}, - http.HeaderOrderKey: {"accept", "accept-language", "user-agent"}, - } - resp, err := client.Do(req) - if err != nil { - return statsigChallenge{}, err - } - defer resp.Body.Close() - raw, err := io.ReadAll(resp.Body) - if err != nil { - return statsigChallenge{}, err - } - html := string(raw) - - mm := statsigMetaRe.FindStringSubmatch(html) - if mm == nil { - return statsigChallenge{}, errors.New("statsig: seed meta not found") - } - curves, err := parseStatsigCurves(html) - if err != nil { - return statsigChallenge{}, err - } - curvesJSON, err := json.Marshal(curves) - if err != nil { - return statsigChallenge{}, err - } - - // Primary path is the goja signer (statsig_js.go); set it up / refresh it for - // this build. Static (header, salt) below is only a last-resort fallback. - ensureEngine(ctx, client, html) - - ch := statsigChallenge{ - header: statsigHeader, - suffix: statsigSuffix, - trailer: statsigTrailer, - fetchedAt: time.Now(), - seedB64: mm[1], - curvesJSON: string(curvesJSON), - } - // NOTE: the old hand-ported per-session derivation (0x00+seed header + - // computeStatsigTail salt) is intentionally NOT applied here. Verified - // 2026-07-13 against /rest/app-chat/conversations/new: that dynamic token is - // rejected ("Request rejected by anti-bot rules.", 403), while the static - // (header, salt) defaults are accepted (200) — so we keep the static values - // as the fallback. When the goja signer (ensureEngine above) locates and - // verifies grok's own chunk it still takes over via seedB64/curvesJSON in - // statsigID; computeStatsigTail is retained only for reference/tests. - return ch, nil -} - -func decodeStatsigSeed(s string) ([]byte, error) { - if b, err := base64.StdEncoding.DecodeString(s); err == nil && len(b) == 48 { - return b, nil - } - b, err := base64.RawStdEncoding.DecodeString(strings.TrimRight(s, "=")) - if err != nil { - return nil, fmt.Errorf("statsig: seed decode: %w", err) - } - if len(b) != 48 { - return nil, fmt.Errorf("statsig: seed len %d", len(b)) - } - return b, nil -} - -// parseStatsigCurves extracts the [[{color,deg,bezier}...]...] array from the -// RSC-escaped homepage HTML (the one immediately followed by color/bezier keys). -func parseStatsigCurves(html string) ([][]statsigCurve, error) { - marker := -1 - for from := 0; ; { - i := strings.Index(html[from:], "curves") - if i < 0 { - break - } - i += from - end := i + 160 - if end > len(html) { - end = len(html) - } - w := html[i:end] - if strings.Contains(w, "color") && strings.Contains(w, "bezier") { - marker = i - break - } - from = i + 6 - } - if marker < 0 { - return nil, errors.New("statsig: curves not found") - } - rel := strings.IndexByte(html[marker:], '[') - if rel < 0 { - return nil, errors.New("statsig: curves array start not found") - } - start := marker + rel - depth, stop := 0, -1 - for k := start; k < len(html); k++ { - switch html[k] { - case '[': - depth++ - case ']': - depth-- - if depth == 0 { - stop = k + 1 - } - } - if stop > 0 { - break - } - } - if stop < 0 { - return nil, errors.New("statsig: curves array end not found") - } - sub := strings.ReplaceAll(html[start:stop], `\`, "") - var out [][]statsigCurve - if err := json.Unmarshal([]byte(sub), &out); err != nil { - return nil, fmt.Errorf("statsig: curves json: %w", err) - } - return out, nil -} - -// computeStatsigTail reproduces the browser's 3-byte color + 6-number transform -// matrix "F" tail: it selects a curve by the seed, samples the curve's keyframe -// animation (color lerp + rotate) at a seed-derived paused currentTime, and -// serializes getComputedStyle(color)+getComputedStyle(transform) exactly as the -// signer does (each number -> Number(v.toFixed(2)).toString(16), '.'/'-' stripped). -func computeStatsigTail(seed []byte, curves [][]statsigCurve) (string, error) { - if len(seed) < 46 { - return "", errors.New("statsig: short seed") - } - if len(curves) == 0 { - return "", errors.New("statsig: no curves") - } - group := int(seed[5]) % len(curves) - if len(curves[group]) == 0 { - return "", errors.New("statsig: empty curve group") - } - idx := int(seed[45]) % len(curves[group]) - cv := curves[group][idx] - if len(cv.Color) < 6 || len(cv.Bezier) < 4 { - return "", errors.New("statsig: malformed curve") - } - - n := (int(seed[1]) % 16) * (int(seed[6]) % 16) * (int(seed[13]) % 16) - currentTime := jsRound(float64(n)/10) * 10 - progress := float64(currentTime) / statsigAnimDuration - - x1 := toFixed2(float64(cv.Bezier[0]) / 255) - y1 := toFixed2(float64(cv.Bezier[1])*2/255 - 1) - x2 := toFixed2(float64(cv.Bezier[2]) / 255) - y2 := toFixed2(float64(cv.Bezier[3])*2/255 - 1) - eased := cubicBezierEase(x1, y1, x2, y2, progress) - - nums := make([]float64, 0, 9) - for k := 0; k < 3; k++ { - v := jsRound(float64(cv.Color[k]) + (float64(cv.Color[k+3])-float64(cv.Color[k]))*eased) - if v < 0 { - v = 0 - } - if v > 255 { - v = 255 - } - nums = append(nums, float64(v)) - } - theta := int(math.Floor(float64(cv.Deg)*300/255 + 60)) - rad := float64(theta) * eased * math.Pi / 180 - cos, sin := math.Cos(rad), math.Sin(rad) - nums = append(nums, cos, sin, -sin, cos, 0, 0) - - var b strings.Builder - for _, v := range nums { - b.WriteString(jsHex(v)) - } - out := strings.NewReplacer(".", "", "-", "").Replace(b.String()) - return out, nil -} - -// jsRound matches JavaScript Math.round (round half up toward +Inf). -func jsRound(x float64) int { - return int(math.Floor(x + 0.5)) -} - -// toFixed2 matches JavaScript Number(v.toFixed(2)). -func toFixed2(v float64) float64 { - f, _ := strconv.ParseFloat(strconv.FormatFloat(v, 'f', 2, 64), 64) - return f -} - -// jsHex matches JavaScript Number(v.toFixed(2)).toString(16). -func jsHex(v float64) string { - v = toFixed2(v) - neg := "" - if v < 0 { - neg = "-" - v = -v - } - ip := int64(math.Floor(v)) - frac := v - float64(ip) - s := neg + strconv.FormatInt(ip, 16) - if frac == 0 { - return s - } - const digits = "0123456789abcdef" - var b strings.Builder - b.WriteString(s) - b.WriteByte('.') - for i := 0; i < 20 && frac != 0; i++ { - frac *= 16 - d := int(frac) - b.WriteByte(digits[d]) - frac -= float64(d) - } - return b.String() -} - -// cubicBezierEase evaluates a CSS cubic-bezier(x1,y1,x2,y2) easing at input -// fraction p: solve X(t)=p for t (bisection), then return Y(t). -func cubicBezierEase(x1, y1, x2, y2, p float64) float64 { - bez := func(t, a, b float64) float64 { - mt := 1 - t - return 3*a*mt*mt*t + 3*b*mt*t*t + t*t*t - } - lo, hi := 0.0, 1.0 - for i := 0; i < 100; i++ { - mid := (lo + hi) / 2 - if bez(mid, x1, x2) < p { - lo = mid - } else { - hi = mid - } - } - return bez((lo+hi)/2, y1, y2) -} - -// statsigID reproduces grok.com's x-statsig-id anti-bot token for a request. The -// token binds to the request METHOD and URL path and to a coarse timestamp, so -// it must be regenerated per request. See the package doc for the layout. It uses -// the session's self-healed (header, salt) when available, else static defaults. -func statsigID(path, method, token string) string { - header, suffix, trailer := statsigHeader, statsigSuffix, statsigTrailer - statsigMu.Lock() - ch, ok := statsigCache[token] - statsigMu.Unlock() - if ok { - header, suffix, trailer = ch.header, ch.suffix, ch.trailer - // Primary: run grok's own signer in goja (durable across reships). Falls - // through to the static computation below on any failure. - if ch.seedB64 != "" && ch.curvesJSON != "" { - if id, err := signWithEngine(ch.seedB64, ch.curvesJSON, path, method); err == nil { - return id - } else if !errors.Is(err, errEngineNotReady) { - log.Printf("grok statsig: js signer failed, using static fallback: %v", err) - } - } - } - - counter := uint32(time.Now().Unix() - statsigEpoch) - sig := fmt.Sprintf("%s!%s!%d%s", method, path, counter, suffix) - hash := sha256.Sum256([]byte(sig)) - - raw := make([]byte, 0, 70) - raw = append(raw, header...) - raw = binary.LittleEndian.AppendUint32(raw, counter) - raw = append(raw, hash[:16]...) - raw = append(raw, trailer) - - key := byte(rand.IntN(256)) - for i := range raw { - raw[i] ^= key - } - return base64.RawStdEncoding.EncodeToString(raw) -} - -// 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(req.URL.Path, req.Method, token)}, - "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) { return c.newTLSClientP(true) } - -// newDirectTLSClient egresses on the local IP (never the proxy). Used for -// reference-frame upload and result (video) download; only the generate submit -// (/rest/app-chat/conversations/new) uses the proxy. -func (c *Client) newDirectTLSClient() (tlsclient.HttpClient, error) { return c.newTLSClientP(false) } - -func (c *Client) newTLSClientP(useProxy bool) (tlsclient.HttpClient, error) { - options := []tlsclient.HttpClientOption{ - // Video generation streams inline until progress=100; a 15s clip can take - // several minutes, so allow up to 10m (caller's genCtx caps at 12m). - tlsclient.WithTimeoutSeconds(600), - tlsclient.WithClientProfile(profiles.Chrome_133), - tlsclient.WithRandomTLSExtensionOrder(), - } - if useProxy && 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) (used int, resetUnix string, ok bool) { - var usedF float32 - seen := false - for len(msg) > 0 { - fn, wt, val, rest, good := readField(msg) - if !good { - break - } - msg = rest - seen = true - switch { - case fn == 1 && wt == 5: // float32 credits USED this period - usedF = float32FromLE(val) - case fn == 5 && wt == 2: // reset timestamp message { #1 varint=seconds } - if sec, sok := firstVarint(val); sok { - resetUnix = strconv.FormatInt(sec, 10) - } - } - } - // A valid config message may OMIT field #1 when used == 0 (proto3 drops zero - // scalars) — a full-quota account. So as long as the message had any field, - // treat it as parsed with used defaulting to 0 (= 100 remaining). - return int(usedF), resetUnix, seen -} - -// 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 -} diff --git a/backend/internal/provider/grok/console.go b/backend/internal/provider/grok/console.go deleted file mode 100644 index 1f9e7b7..0000000 --- a/backend/internal/provider/grok/console.go +++ /dev/null @@ -1,609 +0,0 @@ -package grok - -// Grok Console (console.x.ai) is xAI's own product console. It serves the same -// grok-imagine media models as grok.com, but over a plain JSON API that is NOT -// anti-bot gated: no x-statsig-id, no media-post + streaming-conversation dance. -// Auth reuses the very same website "sso" cookie — POST /v1/dpop/token exchanges -// it for a short-lived access token bound to a locally generated P-256 key, and -// every call then carries "Authorization: DPoP " plus a per-request ES256 -// DPoP proof (RFC 9449) over (method, url, token hash). - -import ( - "context" - "crypto/ecdsa" - "crypto/elliptic" - "crypto/rand" - "crypto/sha256" - "encoding/base64" - "encoding/json" - "errors" - "fmt" - "io" - "net/url" - "strings" - "sync" - "time" - - http "github.com/bogdanfinn/fhttp" - tlsclient "github.com/bogdanfinn/tls-client" - "github.com/google/uuid" -) - -const ( - consoleBase = "https://console.x.ai" - consoleOrigin = "https://console.x.ai" - - // Upstream model ids of the Console media catalog. The -quality image model - // is what the reference-image (edit) path uses. - ConsoleVideoModel = "grok-imagine-video" - ConsoleImageModel = "grok-imagine-image" - ConsoleImageQualityModel = "grok-imagine-image-quality" - - // consoleTokenSkew refreshes a DPoP token slightly before it expires. - consoleTokenSkew = 20 * time.Second - // consoleMaxEditImages is the upstream cap for image-to-image references - // (video takes a single 首图). - consoleMaxEditImages = 3 - - consoleVideoPollEvery = 2 * time.Second - consoleVideoDeadline = 30 * time.Minute -) - -// consoleSession is one minted DPoP token plus the key it is bound to. skew is -// (server time − local time) learned from the mint response Date header; the -// proof iat is shifted by it so a drifting host clock doesn't fail every proof. -type consoleSession struct { - accessToken string - key *ecdsa.PrivateKey - jwk consoleJWK - expiresAt time.Time - skew time.Duration -} - -type consoleJWK struct { - Crv string `json:"crv"` - Kty string `json:"kty"` - X string `json:"x"` - Y string `json:"y"` -} - -var ( - consoleMu sync.Mutex - consoleSessions = map[string]consoleSession{} // keyed by sso token hash -) - -// GenerateConsoleVideo runs Console's video pipeline: POST /v1/videos/generations -// returns a request_id, GET /v1/videos/{id} is polled until the clip is rendered -// and reports its vidgen.x.ai URL. frames is optional (image-to-video, one 首图 -// max) and is inlined as a data URL — Console takes the image in the request, so -// there is no separate upload step. When downloadResult is false, returns nil -// bytes and the artifact URL in meta["video_url"]; otherwise downloads the mp4. -func (c *Client) GenerateConsoleVideo(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 - } - prompt = strings.TrimSpace(prompt) - if prompt == "" && len(frames) == 0 { - return nil, nil, fmt.Errorf("grok console: prompt required") - } - if seconds < 1 || seconds > 15 { - seconds = 10 - } - submitClient, err := c.newTLSClient() - if err != nil { - return nil, nil, err - } - directClient, err := c.newDirectTLSClient() - if err != nil { - return nil, nil, err - } - - payload := map[string]any{"model": ConsoleVideoModel, "duration": seconds} - if prompt != "" { - payload["prompt"] = prompt - } - if ratio := strings.TrimSpace(aspectRatio); ratio != "" { - payload["aspect_ratio"] = ratio - } - if res := consoleVideoResolution(resolution); res != "" { - payload["resolution"] = res - } - for _, f := range frames { - if len(f) == 0 { - continue - } - payload["image"] = map[string]any{"url": dataURL(f)} - break // 上游只吃 1 张首图 - } - - created, err := c.consoleJSON(ctx, submitClient, token, http.MethodPost, "/v1/videos/generations", payload) - if err != nil { - return nil, nil, err - } - requestID := strings.TrimSpace(stringValue(created["request_id"])) - if requestID == "" { - return nil, nil, fmt.Errorf("%w: video create missing request_id", ErrTemporaryUpstream) - } - - deadline := time.Now().Add(consoleVideoDeadline) - videoURL := "" - for { - status, pollErr := c.consoleJSON(ctx, submitClient, token, http.MethodGet, "/v1/videos/"+url.PathEscape(requestID), nil) - if pollErr != nil { - return nil, nil, pollErr - } - url, done, sErr := parseConsoleVideoStatus(status) - if sErr != nil { - return nil, nil, sErr - } - if done { - videoURL = url - break - } - if time.Now().After(deadline) { - return nil, nil, fmt.Errorf("%w: video render did not complete in time", ErrTemporaryUpstream) - } - select { - case <-ctx.Done(): - return nil, nil, ctx.Err() - case <-time.After(consoleVideoPollEvery): - } - } - - meta := map[string]any{ - "provider": "grok", - "request_id": requestID, - "video_url": videoURL, - } - if !downloadResult { - return nil, meta, nil - } - data, err := c.download(ctx, directClient, token, videoURL) - if err != nil { - return nil, nil, err - } - return data, meta, nil -} - -// GenerateConsoleImage runs Console's image pipeline: /v1/images/generations on -// grok-imagine-image for text-to-image, and /v1/images/edits on the -// grok-imagine-image-quality model (up to 3 inlined reference images) as soon as -// references are supplied — 带参考图自动走 quality 那个上游模型. Returns the -// artifact URL in meta["image_url"]; the bytes are downloaded unless urlOnly. -func (c *Client) GenerateConsoleImage(ctx context.Context, token, prompt, aspectRatio, resolution string, refs [][]byte, urlOnly 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 console: prompt required") - } - submitClient, err := c.newTLSClient() - if err != nil { - return nil, nil, err - } - directClient, err := c.newDirectTLSClient() - if err != nil { - return nil, nil, err - } - - payload := map[string]any{ - "model": ConsoleImageModel, "prompt": strings.TrimSpace(prompt), - "n": 1, "response_format": "url", - } - - if ratio := consoleImageAspectRatio(aspectRatio); ratio != "" { - payload["aspect_ratio"] = ratio - } - if res := consoleImageResolution(resolution); res != "" { - payload["resolution"] = res - } - path := "/v1/images/generations" - var images []map[string]any - for _, r := range refs { - if len(r) == 0 || len(images) >= consoleMaxEditImages { - continue - } - images = append(images, map[string]any{"type": "image_url", "url": dataURL(r)}) - } - if len(images) > 0 { - path = "/v1/images/edits" - payload["model"] = ConsoleImageQualityModel - if len(images) == 1 { - payload["image"] = images[0] - } else { - payload["images"] = images - } - } - - res, err := c.consoleJSON(ctx, submitClient, token, http.MethodPost, path, payload) - if err != nil { - return nil, nil, err - } - items, _ := res["data"].([]any) - imageURL := "" - for _, item := range items { - entry, _ := item.(map[string]any) - if entry == nil { - continue - } - if v := strings.TrimSpace(stringValue(entry["url"])); v != "" { - imageURL = v - break - } - } - if imageURL == "" { - return nil, nil, fmt.Errorf("%w: image response missing url", ErrTemporaryUpstream) - } - meta := map[string]any{"provider": "grok", "image_url": imageURL} - if urlOnly { - return nil, meta, nil - } - data, err := c.download(ctx, directClient, token, imageURL) - if err != nil { - return nil, nil, err - } - return data, meta, nil -} - -// consoleJSON does an authed Console JSON call and parses the response object. A -// 401 means the DPoP token went stale (or the sso died): the cached session is -// dropped and the call retried once with a freshly minted token. -func (c *Client) consoleJSON(ctx context.Context, client tlsclient.HttpClient, token, method, path string, body any) (map[string]any, error) { - var payload []byte - if body != nil { - encoded, err := json.Marshal(body) - if err != nil { - return nil, err - } - payload = encoded - } - for attempt := 0; attempt < 2; attempt++ { - session, err := c.consoleSession(ctx, client, token) - if err != nil { - return nil, err - } - var reader io.Reader - if len(payload) > 0 { - reader = strings.NewReader(string(payload)) - } - req, err := http.NewRequest(method, consoleBase+path, reader) - if err != nil { - return nil, err - } - req = req.WithContext(ctx) - if err := c.applyConsoleHeaders(req, token, session); err != nil { - return nil, err - } - resp, err := client.Do(req) - if err != nil { - return nil, fmt.Errorf("%w: %v", ErrTemporaryUpstream, err) - } - raw, readErr := io.ReadAll(resp.Body) - resp.Body.Close() - if resp.StatusCode == 401 && attempt == 0 { - consoleForget(token, session.accessToken) - continue - } - if readErr != nil { - return nil, fmt.Errorf("%w: %v", ErrTemporaryUpstream, readErr) - } - if e := consoleMapStatus(path, resp.StatusCode, raw); e != nil { - return nil, e - } - out := map[string]any{} - if len(raw) == 0 { - return out, 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 - } - return nil, fmt.Errorf("%w: console retry state invalid", ErrTemporaryUpstream) -} - -// consoleSession returns a live DPoP session for the sso token, minting one when -// the cache is empty or the token is about to expire. -func (c *Client) consoleSession(ctx context.Context, client tlsclient.HttpClient, token string) (consoleSession, error) { - key := consoleCacheKey(token) - consoleMu.Lock() - cached, ok := consoleSessions[key] - consoleMu.Unlock() - if ok && cached.expiresAt.After(time.Now().Add(consoleTokenSkew)) { - return cached, nil - } - session, err := c.mintConsoleSession(ctx, client, token) - if err != nil { - return consoleSession{}, err - } - consoleMu.Lock() - consoleSessions[key] = session - consoleMu.Unlock() - return session, nil -} - -// consoleForget drops a cached session (only if it is still the one that failed, -// so a concurrent refresh isn't thrown away). -func consoleForget(token, accessToken string) { - key := consoleCacheKey(token) - consoleMu.Lock() - if current, ok := consoleSessions[key]; ok && current.accessToken == accessToken { - delete(consoleSessions, key) - } - consoleMu.Unlock() -} - -func consoleCacheKey(token string) string { - sum := sha256.Sum256([]byte(token)) - return fmt.Sprintf("%x", sum[:]) -} - -// mintConsoleSession posts the freshly generated public JWK to /v1/dpop/token -// with the sso cookie and keeps the returned DPoP-bound access token. -func (c *Client) mintConsoleSession(ctx context.Context, client tlsclient.HttpClient, token string) (consoleSession, error) { - privateKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) - if err != nil { - return consoleSession{}, err - } - jwk := consoleJWK{ - Crv: "P-256", Kty: "EC", - X: base64.RawURLEncoding.EncodeToString(privateKey.PublicKey.X.FillBytes(make([]byte, 32))), - Y: base64.RawURLEncoding.EncodeToString(privateKey.PublicKey.Y.FillBytes(make([]byte, 32))), - } - body, err := json.Marshal(map[string]any{"jwk": jwk}) - if err != nil { - return consoleSession{}, err - } - req, err := http.NewRequest(http.MethodPost, consoleBase+"/v1/dpop/token", strings.NewReader(string(body))) - if err != nil { - return consoleSession{}, err - } - req = req.WithContext(ctx) - c.applyConsoleBrowserHeaders(req, token, nil) - before := time.Now().UTC() - resp, err := client.Do(req) - if err != nil { - return consoleSession{}, fmt.Errorf("%w: %v", ErrTemporaryUpstream, err) - } - after := time.Now().UTC() - raw, err := io.ReadAll(resp.Body) - resp.Body.Close() - if err != nil { - return consoleSession{}, fmt.Errorf("%w: %v", ErrTemporaryUpstream, err) - } - if e := consoleMapStatus("/v1/dpop/token", resp.StatusCode, raw); e != nil { - return consoleSession{}, e - } - var out struct { - AccessToken string `json:"access_token"` - TokenType string `json:"token_type"` - ExpiresIn int `json:"expires_in"` - } - if err := json.Unmarshal(raw, &out); err != nil { - return consoleSession{}, fmt.Errorf("%w: dpop token non-json: %s", ErrTemporaryUpstream, clip(raw, 120)) - } - if strings.TrimSpace(out.AccessToken) == "" || !strings.EqualFold(strings.TrimSpace(out.TokenType), "DPoP") { - return consoleSession{}, fmt.Errorf("%w: dpop token response invalid", ErrTemporaryUpstream) - } - if out.ExpiresIn <= 0 || out.ExpiresIn > 3600 { - out.ExpiresIn = 300 - } - return consoleSession{ - accessToken: out.AccessToken, - key: privateKey, - jwk: jwk, - expiresAt: time.Now().Add(time.Duration(out.ExpiresIn) * time.Second), - skew: consoleClockSkew(resp.Header.Get("Date"), before, after), - }, nil -} - -// applyConsoleHeaders sets the browser-like header set plus the DPoP -// Authorization header and its per-request proof. -func (c *Client) applyConsoleHeaders(req *http.Request, token string, session consoleSession) error { - proof, err := consoleProof(session, req) - if err != nil { - return err - } - c.applyConsoleBrowserHeaders(req, token, map[string]string{ - "authorization": "DPoP " + session.accessToken, - "dpop": proof, - }) - return nil -} - -func (c *Client) applyConsoleBrowserHeaders(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": {consoleOrigin}, - "referer": {consoleOrigin + "/"}, - "user-agent": {userAgent}, - "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", "authorization", "content-type", "dpop", - "origin", "referer", "user-agent", "sec-ch-ua", "sec-ch-ua-mobile", - "sec-ch-ua-platform", "sec-fetch-dest", "sec-fetch-mode", "sec-fetch-site", "cookie", - } - req.Header = h -} - -// consoleProof builds the ES256 DPoP proof JWT for one request: it binds the -// method + URL (without query) and the access token hash (ath), signed with the -// key the token was issued against. -func consoleProof(session consoleSession, req *http.Request) (string, error) { - if session.key == nil || strings.TrimSpace(session.accessToken) == "" || req == nil || req.URL == nil { - return "", errors.New("grok console: dpop session invalid") - } - path := req.URL.EscapedPath() - if path == "" { - path = "/" - } - ath := sha256.Sum256([]byte(session.accessToken)) - header := map[string]any{"typ": "dpop+jwt", "alg": "ES256", "jwk": session.jwk} - claims := map[string]any{ - "jti": uuid.NewString(), - "htm": strings.ToUpper(req.Method), - "htu": req.URL.Scheme + "://" + req.URL.Host + path, - "iat": time.Now().Add(session.skew).UTC().Unix(), - "ath": base64.RawURLEncoding.EncodeToString(ath[:]), - } - headerJSON, err := json.Marshal(header) - if err != nil { - return "", err - } - claimsJSON, err := json.Marshal(claims) - if err != nil { - return "", err - } - signingInput := base64.RawURLEncoding.EncodeToString(headerJSON) + "." + base64.RawURLEncoding.EncodeToString(claimsJSON) - digest := sha256.Sum256([]byte(signingInput)) - r, s, err := ecdsa.Sign(rand.Reader, session.key, digest[:]) - if err != nil { - return "", err - } - signature := make([]byte, 64) - r.FillBytes(signature[:32]) - s.FillBytes(signature[32:]) - return signingInput + "." + base64.RawURLEncoding.EncodeToString(signature), nil -} - -// consoleClockSkew mirrors the console.x.ai frontend: proof iat is local time -// corrected by (server Date − local time) so host clock drift doesn't break -// every proof. before/after bound the RTT around the mint response. -func consoleClockSkew(dateHeader string, before, after time.Time) time.Duration { - dateHeader = strings.TrimSpace(dateHeader) - if dateHeader == "" { - return 0 - } - serverTime, err := http.ParseTime(dateHeader) - if err != nil { - return 0 - } - if after.Before(before) { - after = before - } - mid := before.Add(after.Sub(before) / 2) - return serverTime.UTC().Sub(mid.UTC()).Round(time.Second) -} - -// parseConsoleVideoStatus reads one poll response: (url, done, error). -func parseConsoleVideoStatus(res map[string]any) (string, bool, error) { - status := strings.ToLower(strings.TrimSpace(stringValue(res["status"]))) - switch status { - case "done", "completed", "succeeded", "success", "ready": - video, _ := res["video"].(map[string]any) - url := "" - if video != nil { - url = strings.TrimSpace(stringValue(video["url"])) - } - if url == "" { - return "", false, fmt.Errorf("%w: video done without content url", ErrTemporaryUpstream) - } - return url, true, nil - case "pending", "processing", "in_progress", "queued", "": - return "", false, nil - default: - message := consoleErrorMessage(res["error"]) - if isCreditError(message) { - return "", false, fmt.Errorf("%w: %s", ErrQuotaExhausted, message) - } - if message == "" { - message = status - } - return "", false, fmt.Errorf("grok console: video generation failed: %s", message) - } -} - -// consoleMapStatus maps a Console HTTP status to the shared provider sentinels. -// Console is not anti-bot gated, so a 403 is a real rejection (dead sso / no -// access) unless the body is a Cloudflare interstitial. -func consoleMapStatus(path string, status int, raw []byte) error { - body := string(raw) - switch { - case status >= 200 && status < 300: - return nil - case status == 401: - return fmt.Errorf("%w: %s 401 %s", ErrAuth, path, clip(raw, 160)) - case status == 403: - if isBotChallenge(body) { - return fmt.Errorf("%w: %s 403 %s", ErrTemporaryUpstream, path, clip(raw, 160)) - } - return fmt.Errorf("%w: %s 403 %s", ErrAuth, path, clip(raw, 160)) - case status == 429: - if isCreditError(body) { - return fmt.Errorf("%w: %s 429 %s", ErrQuotaExhausted, path, clip(raw, 160)) - } - return fmt.Errorf("%w: %s 429 %s", ErrTemporaryUpstream, path, clip(raw, 160)) - case status >= 500: - return fmt.Errorf("%w: %s %d %s", ErrTemporaryUpstream, path, status, clip(raw, 160)) - default: - if isCreditError(body) { - return fmt.Errorf("%w: %s", ErrQuotaExhausted, clip(raw, 160)) - } - return fmt.Errorf("grok console: %s %d %s", path, status, clip(raw, 160)) - } -} - -func consoleErrorMessage(value any) string { - switch v := value.(type) { - case string: - return strings.TrimSpace(v) - case map[string]any: - for _, key := range []string{"message", "code", "type"} { - if text := strings.TrimSpace(stringValue(v[key])); text != "" { - return text - } - } - } - return "" -} - -// consoleImageAspectRatio keeps only the ratios Console's image models accept; -// anything else is dropped so the upstream picks its default instead of 400ing. -func consoleImageAspectRatio(value string) string { - switch strings.TrimSpace(value) { - case "1:1", "16:9", "9:16", "4:3", "3:4", "3:2", "2:3", "2:1", "1:2": - return strings.TrimSpace(value) - } - return "" -} - -// consoleImageResolution maps our tier to the Console enum (only 1k / 2k exist). -func consoleImageResolution(value string) string { - switch strings.ToUpper(strings.TrimSpace(value)) { - case "1K": - return "1k" - case "2K", "4K": - return "2k" - } - return "" -} - -// consoleVideoResolution clamps to the two tiers Console accepts. -func consoleVideoResolution(value string) string { - switch strings.ToLower(strings.TrimSpace(value)) { - case "480p": - return "480p" - default: - return "720p" - } -} - -// dataURL inlines reference image bytes; Console takes images in the request -// body, so no upload endpoint is involved. -func dataURL(img []byte) string { - return "data:image/png;base64," + base64.StdEncoding.EncodeToString(img) -} diff --git a/backend/internal/provider/grok/statsig_engine.go b/backend/internal/provider/grok/statsig_engine.go deleted file mode 100644 index eb00451..0000000 --- a/backend/internal/provider/grok/statsig_engine.go +++ /dev/null @@ -1,398 +0,0 @@ -package grok - -// This file makes x-statsig-id durable across grok web reships by executing grok's -// OWN obfuscated signer (a Turbopack chunk) inside an embedded JS engine (goja), -// under a synthesized DOM + Web-Animations getComputedStyle shim (statsig_shim.js). -// grok's code does all the per-build byte-indexing / curve-selection; we only supply -// the stable browser primitives. This replaces the brittle hand-ported byte-offset -// algorithm in computeStatsigTail (kept as a last-resort fallback). See the package -// doc and the grok-statsig-signer memory for the reverse-engineering details. - -import ( - "context" - "crypto/sha256" - "crypto/sha512" - _ "embed" - "encoding/base64" - "encoding/json" - "errors" - "fmt" - "io" - "log" - "regexp" - "strconv" - "strings" - "sync" - - http "github.com/bogdanfinn/fhttp" - tlsclient "github.com/bogdanfinn/tls-client" - "github.com/dop251/goja" -) - -//go:embed statsig_shim.js -var statsigShimJS string - -const sigPoolSize = 4 - -// errEngineNotReady means the durable signer engine has not been built yet (no -// homepage fetched, or chunk location failed). Callers fall back to the static path. -var errEngineNotReady = errors.New("statsig engine not ready") - -// locateConcurrency bounds parallel chunk fetches during signer discovery. -const locateConcurrency = 24 - -var ( - // chunkPathRe matches chunk URLs in the homepage; allChunkRefRe additionally - // matches the loader-manifest's lazy refs (which drop the /_next/ prefix). - // The Turbopack content-hash alphabet includes '~' (e.g. 0~4v1h2zpw7n0.js) - - // the signer chunk is frequently one of these, so the class MUST list it or - // the signer is never fetched/located (403 anti-bot, static fallback goes stale). - chunkPathRe = regexp.MustCompile(`/_next/static/chunks/[a-zA-Z0-9_.~\-/]+\.js`) - allChunkRefRe = regexp.MustCompile(`(?:/_next/)?static/chunks/[a-zA-Z0-9_.~\-/]+\.js`) - // goja's parser tries to fetch //# sourceMappingURL=... from disk and errors. - sourceMapRe = regexp.MustCompile(`(?m)//[#@]\s*sourceMappingURL=\S*`) - - sigMgrMu sync.Mutex - sigBuildKey string // hash of the homepage chunk list; changes on reship - sigChunkSrc string // current signer chunk source - sigPool chan *sigEngine // pool of ready engines for sigChunkSrc -) - -// sigEngine wraps one goja runtime with grok's signer chunk loaded. A goja runtime -// is not safe for concurrent use; the pool hands each engine to one goroutine at a -// time so no per-engine locking is needed. -type sigEngine struct { - rt *goja.Runtime - fire goja.Callable // __grokSignInto -} - -func newSigEngine(chunkSrc string) (*sigEngine, error) { - rt := goja.New() - // SHA-256 bridge for crypto.subtle.digest. - if err := rt.Set("__goSha256", func(call goja.FunctionCall) goja.Value { - data := jsBytes(rt, call.Argument(0)) - sum := sha256.Sum256(data) - return rt.ToValue(rt.NewArrayBuffer(sum[:])) - }); err != nil { - return nil, err - } - if _, err := rt.RunString(statsigShimJS); err != nil { - return nil, fmt.Errorf("shim: %w", err) - } - if _, err := rt.RunString(sourceMapRe.ReplaceAllString(chunkSrc, "")); err != nil { - return nil, fmt.Errorf("chunk eval: %w", err) - } - if _, err := rt.RunString("__grokBootstrap()"); err != nil { - return nil, fmt.Errorf("bootstrap: %w", err) - } - fire, ok := goja.AssertFunction(rt.Get("__grokSignInto")) - if !ok { - return nil, errors.New("statsig js: __grokSignInto missing") - } - return &sigEngine{rt: rt, fire: fire}, nil -} - -// statsigID runs grok's signer for one request. seedB64 is the raw content; -// curvesJSON is [[{color,deg,bezier}...]...]. -func (e *sigEngine) statsigID(seedB64, curvesJSON, path, method string) (string, error) { - _ = e.rt.Set("__SEED", seedB64) - _ = e.rt.Set("__CURVES", curvesJSON) - _ = e.rt.Set("__PATH", path) - _ = e.rt.Set("__METHOD", method) - // RunString drains goja's microtask queue, settling the async signer's promise. - if _, err := e.fire(goja.Undefined()); err != nil { - return "", err - } - if errv := e.rt.Get("__grokErr"); errv != nil && !goja.IsNull(errv) && !goja.IsUndefined(errv) { - return "", fmt.Errorf("statsig js: %s", errv.String()) - } - res := e.rt.Get("__grokResult") - if res == nil || goja.IsNull(res) || goja.IsUndefined(res) { - return "", errors.New("statsig js: promise did not settle") - } - id := res.String() - if id == "" { - return "", errors.New("statsig js: empty id") - } - return id, nil -} - -// jsBytes extracts the byte contents of a JS Uint8Array / ArrayBuffer value. -func jsBytes(rt *goja.Runtime, v goja.Value) []byte { - if ab, ok := v.Export().(goja.ArrayBuffer); ok { - return ab.Bytes() - } - obj := v.ToObject(rt) - if buf := obj.Get("buffer"); buf != nil { - if ab, ok := buf.Export().(goja.ArrayBuffer); ok { - return ab.Bytes() - } - } - n := int(obj.Get("length").ToInteger()) - out := make([]byte, n) - for i := 0; i < n; i++ { - out[i] = byte(obj.Get(strconv.Itoa(i)).ToInteger()) - } - return out -} - -// signWithEngine borrows an engine from the pool (building one on demand), signs, -// and returns it. Returns an error if the engine subsystem is not ready. -func signWithEngine(seedB64, curvesJSON, path, method string) (string, error) { - sigMgrMu.Lock() - src, pool := sigChunkSrc, sigPool - sigMgrMu.Unlock() - if src == "" || pool == nil { - return "", errEngineNotReady - } - var eng *sigEngine - select { - case eng = <-pool: - default: - var err error - if eng, err = newSigEngine(src); err != nil { - return "", err - } - } - id, err := eng.statsigID(seedB64, curvesJSON, path, method) - select { - case pool <- eng: - default: - } - return id, err -} - -// ensureEngine refreshes the global engine pool when the homepage's chunk set -// changes (i.e. grok reshipped). It locates the signer chunk build-agnostically and -// rebuilds the pool. Cheap no-op when the build is unchanged. -func ensureEngine(ctx context.Context, client tlsclient.HttpClient, homeHTML string) { - paths := chunkPathRe.FindAllString(homeHTML, -1) - if len(paths) == 0 { - return - } - key := hashStrings(paths) - - sigMgrMu.Lock() - unchanged := key == sigBuildKey && sigPool != nil - sigMgrMu.Unlock() - if unchanged { - return - } - - // Inputs for build-agnostic behavioral verification of candidate chunks. - mm := statsigMetaRe.FindStringSubmatch(homeHTML) - if mm == nil { - log.Printf("grok statsig: no seed meta in homepage; cannot locate signer") - return - } - seed, err := decodeStatsigSeed(mm[1]) - if err != nil { - log.Printf("grok statsig: seed decode failed: %v", err) - return - } - curves, err := parseStatsigCurves(homeHTML) - if err != nil { - log.Printf("grok statsig: curves parse failed: %v", err) - return - } - cj, err := json.Marshal(curves) - if err != nil { - return - } - - src, err := locateSignerChunk(ctx, client, dedupe(paths), mm[1], string(cj), seed) - if err != nil { - log.Printf("grok statsig: locate signer chunk failed (will use static fallback): %v", err) - return - } - // smoke-test: a build must produce a loadable engine before we commit to it. - eng, err := newSigEngine(src) - if err != nil { - log.Printf("grok statsig: signer chunk did not load in goja (static fallback): %v", err) - return - } - pool := make(chan *sigEngine, sigPoolSize) - pool <- eng // reuse the smoke-test engine instead of discarding it - sigMgrMu.Lock() - sigBuildKey = key - sigChunkSrc = src - sigPool = pool - sigMgrMu.Unlock() - log.Printf("grok statsig: self-heal engine ready (build %s..)", key[:8]) -} - -// locateSignerChunk finds grok's obfuscated anti-bot signer chunk build-agnostically, -// with no dependency on rotating literals (header name, class names, module ids). -// It screens every reachable chunk with a cheap obfuscator.io fingerprint, then -// confirms the true signer by RUNNING it in goja and checking its output embeds the -// homepage seed (signerEmbedsSeed). The signer is lazily loaded (not referenced in -// the HTML directly), so candidates also include the chunk paths named inside the -// homepage's Turbopack loader manifest. -func locateSignerChunk(ctx context.Context, client tlsclient.HttpClient, homeChunks []string, seedB64, curvesJSON string, seed []byte) (string, error) { - seen := map[string]bool{} - var lazy []string - for _, p := range homeChunks { - seen[p] = true - } - - // Pass 1 (sequential): fetch the homepage chunks, verify them directly, and - // harvest every chunk path they reference (the loader manifest lists the signer). - for _, p := range homeChunks { - body, err := fetchChunk(ctx, client, p) - if err != nil { - continue - } - if src, ok := verifySignerChunk(body, seedB64, curvesJSON, seed); ok { - return src, nil - } - for _, ref := range allChunkRefRe.FindAllString(body, -1) { - np := normalizeChunkPath(ref) - if !seen[np] { - seen[np] = true - lazy = append(lazy, np) - } - } - } - - // Pass 2 (concurrent): fetch the lazily-referenced chunks, cheap-fingerprint each, - // behaviorally verify the matches, and stop at the first chunk that round-trips seed. - ctx2, cancel := context.WithCancel(ctx) - defer cancel() - found := make(chan string, 1) - sem := make(chan struct{}, locateConcurrency) - var wg sync.WaitGroup - for _, p := range lazy { - if ctx2.Err() != nil { - break - } - sem <- struct{}{} - wg.Add(1) - go func(p string) { - defer wg.Done() - defer func() { <-sem }() - if ctx2.Err() != nil { - return - } - body, err := fetchChunk(ctx2, client, p) - if err != nil { - return - } - if src, ok := verifySignerChunk(body, seedB64, curvesJSON, seed); ok { - select { - case found <- src: - cancel() - default: - } - } - }(p) - } - go func() { wg.Wait(); close(found) }() - if src, ok := <-found; ok { - return src, nil - } - return "", fmt.Errorf("statsig signer chunk not found among %d candidates", len(homeChunks)+len(lazy)) -} - -// verifySignerChunk returns the goja-ready source if body is grok's signer: it must -// carry the obfuscator.io fingerprint AND, when executed, produce an x-statsig-id -// whose decoded record embeds the homepage seed. -func verifySignerChunk(body, seedB64, curvesJSON string, seed []byte) (string, bool) { - if !isObfuscatedSigner(body) { - return "", false - } - clean := sourceMapRe.ReplaceAllString(body, "") - eng, err := newSigEngine(clean) - if err != nil { - return "", false - } - id, err := eng.statsigID(seedB64, curvesJSON, "/rest/app-chat/conversations/new", "POST") - if err != nil { - return "", false - } - if !signerEmbedsSeed(id, seed) { - return "", false - } - return clean, true -} - -// normalizeChunkPath turns a bare or /_next/-prefixed chunk ref into a fetch path. -func normalizeChunkPath(ref string) string { - return "/_next/" + strings.TrimPrefix(ref, "/_next/") -} - -// isObfuscatedSigner cheaply screens a chunk for the obfuscator.io string-array -// decoder that grok applies ONLY to its anti-bot signer (the rest of the app is -// plain Turbopack output). The RC4-style byte decoder (`...^...%256`) is the stable -// tell; it matches a handful of chunks, which behavioral verification then narrows -// to exactly one. This is deliberately build-agnostic (no rotating class/id/header). -func isObfuscatedSigner(src string) bool { - return strings.Contains(src, "%256") && - strings.Contains(src, "String.fromCharCode") && - strings.Contains(src, "charCodeAt") -} - -// signerEmbedsSeed decodes a candidate x-statsig-id, strips the per-call XOR mask -// (plaintext[0] is 0x00, so masked[0] is the key), and checks the plaintext record -// begins with 0x00 + the exact homepage seed. Only grok's real signer round-trips -// OUR seed, so this uniquely identifies the signer chunk regardless of obfuscation. -func signerEmbedsSeed(id string, seed []byte) bool { - raw, err := base64.RawStdEncoding.DecodeString(id) - if err != nil || len(raw) < 1+len(seed) { - return false - } - key := raw[0] // plaintext[0] is 0x00, so the mask key == masked byte 0 - for i := 0; i < len(seed); i++ { - if raw[1+i]^key != seed[i] { - return false - } - } - return true -} - -func fetchChunk(ctx context.Context, client tlsclient.HttpClient, path string) (string, error) { - if !strings.HasPrefix(path, "http") { - path = apiBase + path - } - req, err := http.NewRequest(http.MethodGet, path, nil) - if err != nil { - return "", err - } - req = req.WithContext(ctx) - req.Header = http.Header{ - "accept": {"*/*"}, - "user-agent": {userAgent}, - http.HeaderOrderKey: {"accept", "user-agent"}, - } - resp, err := client.Do(req) - if err != nil { - return "", err - } - defer resp.Body.Close() - if resp.StatusCode != 200 { - return "", fmt.Errorf("chunk http %d", resp.StatusCode) - } - b, err := io.ReadAll(resp.Body) - return string(b), err -} - -func hashStrings(ss []string) string { - uniq := dedupe(ss) - h := sha512.New() - for _, s := range uniq { - _, _ = io.WriteString(h, s) - _, _ = io.WriteString(h, "\n") - } - return fmt.Sprintf("%x", h.Sum(nil)) -} - -func dedupe(ss []string) []string { - seen := map[string]bool{} - out := ss[:0:0] - for _, s := range ss { - if !seen[s] { - seen[s] = true - out = append(out, s) - } - } - return out -} diff --git a/backend/internal/provider/grok/statsig_refresh.go b/backend/internal/provider/grok/statsig_refresh.go deleted file mode 100644 index d63c6c9..0000000 --- a/backend/internal/provider/grok/statsig_refresh.go +++ /dev/null @@ -1,300 +0,0 @@ -package grok - -import ( - "context" - "encoding/base64" - "encoding/hex" - "errors" - "log" - "os" - "os/exec" - "strings" - "sync" - "time" - - "github.com/chromedp/cdproto/page" - "github.com/chromedp/chromedp" -) - -// The x-statsig-id recipe (49-byte header + salt suffix) encodes grok's per-load -// browser fingerprint "F". grok recomputes F server-side from the homepage -// seed+curves and the byte-indexing that derives it ROTATES on every web reship, -// so any hand-ported or baked-in recipe goes stale (403 "Request rejected by -// anti-bot rules.") within days. The durable fix is to let grok's OWN signer -// produce the recipe: we drive a headless Chrome to grok.com, read the exact -// string grok feeds to crypto.subtle.digest (reveals the current suffix) and one -// of the x-statsig-id headers its client attaches (reveals the current 49-byte -// header + trailer), then feed those into statsigID's existing per-request -// computation. Because we read grok's live output we never reverse-engineer the -// (rotating) obfuscation; a reship just needs a re-capture, which the auto-refresh -// loop does on a timer and whenever a live request hits an anti-bot 403. -// -// The captured recipe is persisted (site settings) so a restart keeps the last -// good values, and every failure is non-fatal — statsigID keeps using whatever -// recipe is currently live (persisted → last capture → static defaults). - -// hookJS is installed before grok's own scripts run. It records (a) every -// crypto.subtle.digest input that contains the statsig salt marker — the plaintext -// "METHOD!/path!counter" — and (b) every 70-byte x-statsig-id header grok -// attaches to its API calls. Both are only produced once grok's signer succeeds -// (after its keyframe animation settles), so a short wait yields valid values. -const statsigHookJS = ` -(() => { - window.__statsigDigests = []; window.__statsigCaps = []; - try { - const td = new TextDecoder('utf-8', {fatal:false}); - const orig = crypto.subtle.digest.bind(crypto.subtle); - crypto.subtle.digest = function(alg, data) { - try { - let buf = data; if (data && data.buffer) buf = data.buffer.slice(data.byteOffset, data.byteOffset + data.byteLength); - const txt = td.decode(new Uint8Array(buf.slice ? buf.slice(0) : buf)); - if (txt.indexOf('` + statsigSaltPrefix + `') >= 0) window.__statsigDigests.push(txt); - } catch (e) {} - return orig(alg, data); - }; - const of = window.fetch; - window.fetch = function(input, init) { - try { - let hdrs = (init && init.headers); let sid = null; - if (hdrs) { if (typeof hdrs.get === 'function') sid = hdrs.get('x-statsig-id'); else for (const k in hdrs) if (k.toLowerCase() === 'x-statsig-id') sid = hdrs[k]; } - if (sid && sid.length === 94) window.__statsigCaps.push(sid); - } catch (e) {} - return of.apply(this, arguments); - }; - } catch (e) {} -})(); -` - -var ( - statsigRefreshTrigger = make(chan struct{}, 1) - statsigRefreshStarted sync.Once - statsigLastCapture time.Time - statsigLastCaptureMu sync.Mutex -) - -// SetStatsigRecipe atomically replaces the live (header, suffix, trailer) used by -// statsigID and drops the per-session challenge cache so the next request adopts -// the new values. headerHex must decode to 49 bytes; suffix must be non-empty. -func SetStatsigRecipe(headerHex, suffix string, trailer int) error { - h, err := hex.DecodeString(strings.TrimSpace(headerHex)) - if err != nil { - return err - } - if len(h) != 49 { - return errors.New("statsig: header must be 49 bytes") - } - if strings.TrimSpace(suffix) == "" { - return errors.New("statsig: empty suffix") - } - if trailer < 0 || trailer > 255 { - return errors.New("statsig: trailer out of range") - } - statsigMu.Lock() - statsigHeader = h - statsigSuffix = suffix - statsigTrailer = byte(trailer) - statsigCache = map[string]statsigChallenge{} - statsigMu.Unlock() - return nil -} - -// chromeExecPath resolves the Chrome/Chromium binary for headless capture, from -// GROK_STATSIG_CHROME or the usual names on PATH. Empty means none available. -func chromeExecPath() string { - if p := strings.TrimSpace(os.Getenv("GROK_STATSIG_CHROME")); p != "" { - return p - } - for _, name := range []string{"google-chrome-stable", "google-chrome", "chromium", "chromium-browser"} { - if p, err := exec.LookPath(name); err == nil { - return p - } - } - return "" -} - -type statsigCapResult struct { - Digs []string `json:"digs"` - Caps []string `json:"caps"` -} - -// CaptureStatsigRecipe drives a headless Chrome to grok.com and reads grok's own -// signer output, returning the current (headerHex, suffix, trailer). It never -// mutates global state. Returns an error if Chrome is unavailable or grok's -// signer output could not be observed (e.g. blocked before the signer ran). -func CaptureStatsigRecipe(ctx context.Context) (headerHex, suffix string, trailer int, err error) { - chromePath := chromeExecPath() - if chromePath == "" { - return "", "", 0, errors.New("statsig: no chrome binary (set GROK_STATSIG_CHROME)") - } - opts := append(chromedp.DefaultExecAllocatorOptions[:], - chromedp.ExecPath(chromePath), - chromedp.Flag("headless", "new"), - chromedp.Flag("no-sandbox", true), - chromedp.Flag("disable-gpu", true), - chromedp.Flag("disable-dev-shm-usage", true), - chromedp.UserAgent(userAgent), - ) - if proxy := strings.TrimSpace(os.Getenv("GROK_STATSIG_PROXY")); proxy != "" { - opts = append(opts, chromedp.ProxyServer(proxy)) - } - allocCtx, cancel := chromedp.NewExecAllocator(ctx, opts...) - defer cancel() - bctx, cancel2 := chromedp.NewContext(allocCtx) - defer cancel2() - - var res statsigCapResult - runErr := chromedp.Run(bctx, - chromedp.ActionFunc(func(ctx context.Context) error { - _, err := page.AddScriptToEvaluateOnNewDocument(statsigHookJS).Do(ctx) - return err - }), - chromedp.Navigate(apiBase+"/"), - chromedp.Sleep(14*time.Second), - chromedp.Evaluate(`({digs: window.__statsigDigests || [], caps: window.__statsigCaps || []})`, &res), - ) - if runErr != nil { - return "", "", 0, runErr - } - return parseStatsigCapture(res) -} - -// parseStatsigCapture extracts (headerHex, suffix, trailer) from the observed -// digest plaintexts and x-statsig-id headers of one page load. Both must be -// present and are self-consistent (same per-load fingerprint F). -func parseStatsigCapture(res statsigCapResult) (headerHex, suffix string, trailer int, err error) { - for _, d := range res.Digs { - if i := strings.Index(d, statsigSaltPrefix); i >= 0 { - suffix = d[i:] - break - } - } - if suffix == "" { - return "", "", 0, errors.New("statsig: no signer digest observed (blocked before signer ran?)") - } - var header []byte - for _, sid := range res.Caps { - raw, derr := base64.RawStdEncoding.DecodeString(sid) - if derr != nil || len(raw) != 70 { - continue - } - key := raw[0] // header[0] is 0x00, so the XOR mask key == byte 0 - r := make([]byte, 70) - for i := range raw { - r[i] = raw[i] ^ key - } - if r[0] != 0x00 { - continue - } - header = r[0:49] - trailer = int(r[69]) - break - } - if header == nil { - return "", "", 0, errors.New("statsig: no valid x-statsig-id header observed") - } - return hex.EncodeToString(header), suffix, trailer, nil -} - -// TriggerStatsigRefresh asks the auto-refresh loop to re-capture soon (coalesced). -// Called when a live request hits an anti-bot 403 — the recipe likely went stale. -func TriggerStatsigRefresh() { - select { - case statsigRefreshTrigger <- struct{}{}: - default: - } -} - -// StartStatsigAutoRefresh launches the headless capture loop once. It is -// event-driven, NOT a polling loop: it seeds the live recipe from persisted -// values (load), captures once at startup so the process starts fresh, and then -// re-captures ONLY when a live request hits an anti-bot 403 (TriggerStatsigRefresh) -// — i.e. exactly when the recipe actually went stale on a grok reship. The recipe -// has no fixed clock expiry, so there is nothing useful to poll on. An optional -// safetyInterval>0 adds a slow backstop re-capture; pass 0 to disable it (default). -// Captured recipes are handed to `save` for persistence. No-op when a manual -// GROK_STATSIG_* override is set or no Chrome binary is available. -func StartStatsigAutoRefresh(ctx context.Context, safetyInterval time.Duration, - load func(context.Context) (headerHex, suffix string, trailer int, ok bool), - save func(ctx context.Context, headerHex, suffix string, trailer int)) { - - statsigRefreshStarted.Do(func() { - if os.Getenv("GROK_STATSIG_HEADER_HEX") != "" || os.Getenv("GROK_STATSIG_SUFFIX") != "" { - log.Printf("grok statsig: manual GROK_STATSIG_* override set — headless auto-refresh disabled") - return - } - if chromeExecPath() == "" { - log.Printf("grok statsig: no chrome binary — headless auto-refresh disabled, using static recipe (set GROK_STATSIG_CHROME to enable)") - return - } - if load != nil { - if h, s, t, ok := load(ctx); ok { - if err := SetStatsigRecipe(h, s, t); err != nil { - log.Printf("grok statsig: persisted recipe invalid, ignoring: %v", err) - } else { - log.Printf("grok statsig: loaded persisted recipe header[:6]=%s", safePrefix(h, 12)) - } - } - } - go statsigRefreshLoop(ctx, safetyInterval, save) - }) -} - -func statsigRefreshLoop(ctx context.Context, safetyInterval time.Duration, save func(context.Context, string, string, int)) { - refresh := func(reason string) { - statsigLastCaptureMu.Lock() - if time.Since(statsigLastCapture) < 30*time.Second { - statsigLastCaptureMu.Unlock() - return // debounce bursts (e.g. a 403 storm across concurrent requests) - } - statsigLastCapture = time.Now() - statsigLastCaptureMu.Unlock() - - cctx, cancel := context.WithTimeout(ctx, 90*time.Second) - defer cancel() - h, s, t, err := CaptureStatsigRecipe(cctx) - if err != nil { - log.Printf("grok statsig: headless capture (%s) failed, keeping current recipe: %v", reason, err) - return - } - if err := SetStatsigRecipe(h, s, t); err != nil { - log.Printf("grok statsig: captured recipe invalid: %v", err) - return - } - if save != nil { - save(ctx, h, s, t) - } - log.Printf("grok statsig: headless recipe refreshed (%s) header[:6]=%s suffix[:24]=%s", reason, safePrefix(h, 12), safePrefix(s, 24)) - } - - initial := time.NewTimer(10 * time.Second) - defer initial.Stop() - - // Optional slow backstop only; disabled (nil channel blocks forever) when - // safetyInterval<=0 so the loop stays purely event-driven. - var tickC <-chan time.Time - if safetyInterval > 0 { - ticker := time.NewTicker(safetyInterval) - defer ticker.Stop() - tickC = ticker.C - } - - for { - select { - case <-ctx.Done(): - return - case <-initial.C: - refresh("startup") - case <-tickC: - refresh("safety-interval") - case <-statsigRefreshTrigger: - refresh("anti-bot 403") - } - } -} - -func safePrefix(s string, n int) string { - if len(s) < n { - return s - } - return s[:n] -} diff --git a/backend/internal/provider/grok/statsig_selfheal_test.go b/backend/internal/provider/grok/statsig_selfheal_test.go deleted file mode 100644 index be27b28..0000000 --- a/backend/internal/provider/grok/statsig_selfheal_test.go +++ /dev/null @@ -1,175 +0,0 @@ -package grok - -import ( - "context" - "encoding/hex" - "encoding/json" - "errors" - "os" - "strings" - "sync" - "testing" - "time" -) - -// Ground-truth (seed, curves, F) triples. The first row is server-verified -// (captured from a real grok browser token that returned 200); the remainder -// exercise every curve group / currentTime branch and are checked against the -// reference implementation the algorithm was reverse-engineered against. -var statsigTriples = []struct { - seedHex string - wantM string -}{ - {seedHex: "732c9800d181e47c9b2a2705713306f42a51a10d69c6a6c32e53b26d1599b0b75399035a9953fedf05903cb9eb70a989", wantM: "ff75470a3d70a3d70a3d80c51eb851eb8520c51eb851eb8520a3d70a3d70a3d800"}, - {seedHex: "732c9800d1802530bb1d6d057133064d2aa5a1ca69c6a6c32e53b26d1599b0b75399035a9953fe1805903cb9eb70a989", wantM: "27fa10fd70a3d70a3d7028f5c28f5c28f6028f5c28f5c28f60fd70a3d70a3d700"}, - {seedHex: "732c9800d180237b2ed91e057133062c2a13a1de69c6a6c32e53b26d1599b0b75399035a9953fed605903cb9eb70a989", wantM: "079ff0d1eb851eb851e80947ae147ae14780947ae147ae14780d1eb851eb851e800"}, - {seedHex: "732c9800d180197117449405713306722a3fa11f69c6a6c32e53b26d1599b0b75399035a9953fecb05903cb9eb70a989", wantM: "195fe0fae147ae147ae035c28f5c28f5c2035c28f5c28f5c20fae147ae147ae00"}, - {seedHex: "732c9800d1815c3460be3105713306492ad6a13c69c6a6c32e53b26d1599b0b75399035a9953fe9d05903cb9eb70a989", wantM: "ff6ab5035c28f5c28f5c20fae147ae147ae0fae147ae147ae035c28f5c28f5c200"}, - {seedHex: "732c9800d181daa0eee8b9057133061e2a20a16969c6a6c32e53b26d1599b0b75399035a9953fefe05903cb9eb70a989", wantM: "3da0d70e3d70a3d70a3d8075c28f5c28f5c4075c28f5c28f5c40e3d70a3d70a3d800"}, - {seedHex: "732c9800d1812999fdafe5057133067f2a99a15c69c6a6c32e53b26d1599b0b75399035a9953fe7c05903cb9eb70a989", wantM: "1cac4c0451eb851eb8520f5c28f5c28f5c0f5c28f5c28f5c0451eb851eb85200"}, - {seedHex: "732c9800d18254af4dfad705713306252a93a13c69c6a6c32e53b26d1599b0b75399035a9953fed605903cb9eb70a989", wantM: "b40480d70a3d70a3d70808a3d70a3d70a408a3d70a3d70a40d70a3d70a3d70800"}, - {seedHex: "732c9800d182b3fee9232f05713306272a14a1a069c6a6c32e53b26d1599b0b75399035a9953feae05903cb9eb70a989", wantM: "13757a100100"}, - {seedHex: "732c9800d1829ee491c5b105713306f22a8aa12169c6a6c32e53b26d1599b0b75399035a9953fe1f05903cb9eb70a989", wantM: "eb4fb9100a3d70a3d70a3d800a3d70a3d70a3d8100"}, - {seedHex: "732c9800d1833bfc1e6f9305713306ec2a0ba1b569c6a6c32e53b26d1599b0b75399035a9953fe5605903cb9eb70a989", wantM: "8a56650fae147ae147ae02b851eb851eb8602b851eb851eb860fae147ae147ae00"}, - {seedHex: "732c9800d183fe2955e5cd057133067e2a42a1cb69c6a6c32e53b26d1599b0b75399035a9953fec805903cb9eb70a989", wantM: "b242950deb851eb851eb807d70a3d70a3d707d70a3d70a3d70deb851eb851eb800"}, - {seedHex: "732c9800d183d4b7c2764d05713306462a8ea1dc69c6a6c32e53b26d1599b0b75399035a9953fe8e05903cb9eb70a989", wantM: "1e70d8070a3d70a3d70a40e66666666666680e6666666666668070a3d70a3d70a400"}} - -const statsigTestCurves = `[[{"color":[48,44,6,37,198,15],"deg":192,"bezier":[118,76,158,16]},{"color":[224,216,196,111,43,97],"deg":119,"bezier":[67,167,95,219]},{"color":[90,235,250,5,223,64],"deg":104,"bezier":[100,57,106,204]},{"color":[6,109,253,44,29,224],"deg":151,"bezier":[204,60,142,122]},{"color":[81,0,121,208,228,133],"deg":181,"bezier":[182,89,105,123]},{"color":[31,3,160,181,226,184],"deg":98,"bezier":[207,150,215,136]},{"color":[231,243,81,28,109,131],"deg":175,"bezier":[17,103,7,81]},{"color":[222,250,130,169,55,247],"deg":141,"bezier":[21,238,12,84]},{"color":[74,62,116,145,209,185],"deg":109,"bezier":[0,74,58,233]},{"color":[204,168,203,138,107,125],"deg":214,"bezier":[41,13,239,45]},{"color":[246,116,162,162,84,234],"deg":127,"bezier":[160,202,52,76]},{"color":[138,202,210,36,61,195],"deg":234,"bezier":[31,50,177,229]},{"color":[47,46,196,176,79,255],"deg":221,"bezier":[39,14,9,76]},{"color":[245,164,227,71,201,110],"deg":190,"bezier":[193,208,125,9]},{"color":[199,253,44,33,130,240],"deg":191,"bezier":[168,247,61,37]},{"color":[19,91,5,180,202,243],"deg":216,"bezier":[96,152,42,123]}],[{"color":[42,154,230,170,197,128],"deg":108,"bezier":[160,52,34,184]},{"color":[224,132,83,0,231,162],"deg":238,"bezier":[227,37,122,102]},{"color":[24,82,237,199,100,74],"deg":130,"bezier":[186,106,188,209]},{"color":[47,50,169,65,124,44],"deg":228,"bezier":[198,22,146,225]},{"color":[111,131,95,195,131,73],"deg":207,"bezier":[134,146,202,172]},{"color":[192,10,166,28,236,3],"deg":219,"bezier":[85,77,157,235]},{"color":[236,139,199,72,233,250],"deg":197,"bezier":[180,9,79,147]},{"color":[146,195,78,16,231,211],"deg":93,"bezier":[123,18,87,27]},{"color":[168,211,49,42,124,18],"deg":172,"bezier":[232,171,120,118]},{"color":[43,146,96,11,203,53],"deg":146,"bezier":[11,4,83,198]},{"color":[183,97,101,38,115,41],"deg":159,"bezier":[69,223,197,236]},{"color":[78,73,207,132,109,134],"deg":49,"bezier":[59,18,133,168]},{"color":[239,153,225,207,127,157],"deg":194,"bezier":[13,206,154,182]},{"color":[19,146,127,191,68,199],"deg":165,"bezier":[148,212,7,16]},{"color":[93,225,249,144,88,255],"deg":215,"bezier":[62,30,77,69]},{"color":[203,79,164,145,3,20],"deg":81,"bezier":[228,28,93,130]}],[{"color":[171,53,4,125,232,43],"deg":122,"bezier":[136,9,128,97]},{"color":[143,97,205,57,63,69],"deg":192,"bezier":[232,20,219,189]},{"color":[28,173,18,107,158,44],"deg":255,"bezier":[93,187,184,198]},{"color":[170,17,62,142,200,32],"deg":29,"bezier":[14,31,12,97]},{"color":[19,117,122,173,239,66],"deg":74,"bezier":[117,176,139,212]},{"color":[213,151,230,112,224,255],"deg":25,"bezier":[1,223,72,233]},{"color":[153,131,51,105,69,47],"deg":108,"bezier":[123,177,126,140]},{"color":[120,114,44,151,88,83],"deg":165,"bezier":[16,104,134,75]},{"color":[245,145,194,75,120,26],"deg":142,"bezier":[79,235,38,43]},{"color":[147,63,50,255,239,106],"deg":190,"bezier":[122,143,160,150]},{"color":[233,78,184,130,25,123],"deg":54,"bezier":[87,18,184,226]},{"color":[162,180,233,70,57,249],"deg":87,"bezier":[123,238,61,124]},{"color":[146,95,56,171,38,240],"deg":239,"bezier":[241,134,228,44]},{"color":[111,11,149,62,208,177],"deg":70,"bezier":[103,149,4,37]},{"color":[159,128,118,21,197,153],"deg":175,"bezier":[246,215,172,236]},{"color":[194,131,68,247,215,108],"deg":30,"bezier":[23,91,151,231]}],[{"color":[239,129,141,243,85,208],"deg":38,"bezier":[252,248,245,195]},{"color":[205,56,138,49,126,99],"deg":107,"bezier":[72,85,228,91]},{"color":[185,54,148,122,170,158],"deg":192,"bezier":[33,88,51,136]},{"color":[7,14,75,26,23,41],"deg":93,"bezier":[153,21,55,147]},{"color":[245,211,213,64,5,253],"deg":49,"bezier":[253,19,106,155]},{"color":[144,232,165,21,114,130],"deg":200,"bezier":[193,179,133,226]},{"color":[34,196,100,42,114,0],"deg":52,"bezier":[4,4,71,65]},{"color":[195,8,130,102,201,141],"deg":210,"bezier":[8,173,23,33]},{"color":[86,136,44,95,223,62],"deg":249,"bezier":[220,98,68,113]},{"color":[205,48,9,247,236,71],"deg":75,"bezier":[163,240,28,25]},{"color":[43,190,29,239,55,135],"deg":146,"bezier":[109,245,34,188]},{"color":[146,91,92,2,3,251],"deg":97,"bezier":[183,188,95,157]},{"color":[40,213,196,70,81,174],"deg":120,"bezier":[153,197,61,201]},{"color":[131,92,180,68,131,214],"deg":251,"bezier":[94,191,198,89]},{"color":[47,113,219,96,115,228],"deg":238,"bezier":[22,35,60,63]},{"color":[246,244,203,196,78,136],"deg":44,"bezier":[88,23,205,184]}]]` - -// TestComputeStatsigTail was the offline regression for the hand-ported F -// derivation (computeStatsigTail). That algorithm is now only a last-resort -// fallback: the durable path runs grok's own signer in goja (statsig_engine.go), -// which self-heals across web reships. grok rotates the F algorithm on every -// reship, so these browser-captured (seed,F) triples go stale and CANNOT be -// regenerated offline. TestSelfHealStatsigE2E (live, via goja) is the real -// regression now. Kept skipped so the reference vectors aren't lost. -func TestComputeStatsigTail(t *testing.T) { - t.Skip("static F derivation is a deprecated fallback; see TestSelfHealStatsigE2E (goja)") - var curves [][]statsigCurve - if err := json.Unmarshal([]byte(statsigTestCurves), &curves); err != nil { - t.Fatalf("curves: %v", err) - } - for i, tc := range statsigTriples { - seed, err := hex.DecodeString(tc.seedHex) - if err != nil { - t.Fatalf("[%d] seed hex: %v", i, err) - } - got, err := computeStatsigTail(seed, curves) - if err != nil { - t.Fatalf("[%d] computeStatsigTail: %v", i, err) - } - if got != tc.wantM { - t.Errorf("[%d] group=%d\n got=%s\nwant=%s", i, int(seed[5])%len(curves), got, tc.wantM) - } - } -} - -// TestSelfHealStatsigE2E exercises the full browser-free self-healing path: -// fetch the homepage, derive seed+F, cache the challenge, then hit the -// anti-bot-gated conversations/new endpoint. Requires a live GROK_TOK and no -// GROK_STATSIG_* env overrides. -func TestSelfHealStatsigE2E(t *testing.T) { - token := strings.TrimSpace(os.Getenv("GROK_TOK")) - if token == "" { - t.Skip("no GROK_TOK") - } - c := NewClient("") - client, err := c.newTLSClient() - if err != nil { - t.Fatal(err) - } - ctx, cancel := context.WithTimeout(context.Background(), 45*time.Second) - defer cancel() - - c.ensureChallenge(ctx, client, token) - statsigMu.Lock() - ch, ok := statsigCache[token] - statsigMu.Unlock() - if !ok { - t.Fatal("challenge not cached (homepage fetch/derive failed)") - } - t.Logf("dynamic header[:6]=%x suffix=%s", ch.header[:6], ch.suffix) - - body, err := c.postStream(ctx, client, token, "/rest/app-chat/conversations/new", map[string]any{ - "temporary": true, - "modelName": "grok-3", - "message": "hi", - }) - if err != nil { - t.Fatalf("conversations/new: %v", err) - } - t.Logf("OK bytes=%d head=%.80s", len(body), strings.ReplaceAll(body, "\n", " ")) -} - -// TestStatsigEngineConcurrent fires many concurrent goja signs through the pool to -// catch data races / engine cross-talk (run with -race). Requires a live GROK_TOK -// so the engine is built and the challenge cached. -func TestStatsigEngineConcurrent(t *testing.T) { - token := strings.TrimSpace(os.Getenv("GROK_TOK")) - if token == "" { - t.Skip("no GROK_TOK") - } - c := NewClient("") - client, err := c.newTLSClient() - if err != nil { - t.Fatal(err) - } - ctx, cancel := context.WithTimeout(context.Background(), 45*time.Second) - defer cancel() - c.ensureChallenge(ctx, client, token) - statsigMu.Lock() - ch, ok := statsigCache[token] - statsigMu.Unlock() - if !ok || ch.seedB64 == "" || ch.curvesJSON == "" { - t.Fatal("engine inputs not cached") - } - - const n = 40 - var wg sync.WaitGroup - errs := make(chan error, n) - for i := 0; i < n; i++ { - wg.Add(1) - go func(i int) { - defer wg.Done() - id, err := signWithEngine(ch.seedB64, ch.curvesJSON, "/rest/app-chat/conversations/new", "POST") - if err != nil { - errs <- err - return - } - if len(id) < 40 { - errs <- errTooShort - } - }(i) - } - wg.Wait() - close(errs) - for e := range errs { - t.Fatalf("concurrent sign: %v", e) - } -} - -var errTooShort = errors.New("statsig id too short") - -// TestGenerateVideoE2E generates a real grok video using only the dynamic -// self-healed statsig (no env overrides). Requires a live GROK_TOK. -func TestGenerateVideoE2E(t *testing.T) { - token := strings.TrimSpace(os.Getenv("GROK_TOK")) - if token == "" { - t.Skip("no GROK_TOK") - } - c := NewClient("") - ctx, cancel := context.WithTimeout(context.Background(), 12*time.Minute) - defer cancel() - - data, meta, err := c.GenerateVideo(ctx, token, "a cat playing piano", "16:9", "720p", 6, nil, true) - if err != nil { - t.Fatalf("GenerateVideo: %v", err) - } - t.Logf("video bytes=%d meta=%v", len(data), meta) - if len(data) < 1<<20 { - t.Fatalf("video too small: %d bytes", len(data)) - } - if !strings.Contains(string(data[:16]), "ftyp") { - t.Fatalf("not an mp4: % x", data[:16]) - } -} diff --git a/backend/internal/provider/grok/statsig_shim.js b/backend/internal/provider/grok/statsig_shim.js deleted file mode 100644 index 99c167c..0000000 --- a/backend/internal/provider/grok/statsig_shim.js +++ /dev/null @@ -1,216 +0,0 @@ -// statsig_shim.js — browser-environment shim that lets grok.com's own obfuscated -// x-statsig-id signer run inside goja. grok's code does all the (per-build, -// rotating) byte-indexing/curve-selection; we only provide the STABLE browser -// primitives it reads from: the seed , the .r-aufz1o SVG curve DOM, and a -// standard Web-Animations getComputedStyle sampler. Inputs arrive via globals set -// by Go before each sign: __SEED (base64 str), __CURVES (JSON [[{color,deg,bezier}]]), -// __PATH, __METHOD. Go also injects __goSha256(Uint8Array)->ArrayBuffer. -(function () { - 'use strict'; - var g = globalThis; - - // ---- base64 (goja has no atob/btoa) ---- - var B64 = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/'; - g.atob = function (s) { - s = String(s).replace(/=+$/, ''); - var out = '', bits = 0, val = 0; - for (var i = 0; i < s.length; i++) { - var c = B64.indexOf(s.charAt(i)); - if (c < 0) continue; - val = (val << 6) | c; bits += 6; - if (bits >= 8) { bits -= 8; out += String.fromCharCode((val >> bits) & 0xff); } - } - return out; - }; - g.btoa = function (s) { - s = String(s); var out = ''; - for (var i = 0; i < s.length; i += 3) { - var b0 = s.charCodeAt(i), b1 = s.charCodeAt(i + 1), b2 = s.charCodeAt(i + 2); - var h0 = b0 >> 2, h1 = ((b0 & 3) << 4) | (b1 >> 4); - var h2 = ((b1 & 15) << 2) | (b2 >> 6), h3 = b2 & 63; - out += B64[h0] + B64[h1]; - out += isNaN(b1) ? '=' : B64[h2]; - out += isNaN(b2) ? '=' : B64[h3]; - } - return out; - }; - - // ---- TextEncoder (goja has no TextEncoder) ---- - if (typeof g.TextEncoder === 'undefined') { - g.TextEncoder = function () {}; - g.TextEncoder.prototype.encode = function (str) { - str = String(str); - var bytes = []; - for (var i = 0; i < str.length; i++) { - var c = str.charCodeAt(i); - if (c < 0x80) bytes.push(c); - else if (c < 0x800) { bytes.push(0xc0 | (c >> 6), 0x80 | (c & 0x3f)); } - else if (c >= 0xd800 && c <= 0xdbff) { // surrogate pair - var c2 = str.charCodeAt(++i); - var cp = 0x10000 + ((c & 0x3ff) << 10) + (c2 & 0x3ff); - bytes.push(0xf0 | (cp >> 18), 0x80 | ((cp >> 12) & 0x3f), 0x80 | ((cp >> 6) & 0x3f), 0x80 | (cp & 0x3f)); - } else { bytes.push(0xe0 | (c >> 12), 0x80 | ((c >> 6) & 0x3f), 0x80 | (c & 0x3f)); } - } - return Uint8Array.from(bytes); - }; - } - - // ---- crypto.subtle.digest, backed by Go SHA-256 ---- - g.crypto = g.crypto || {}; - g.crypto.subtle = g.crypto.subtle || {}; - g.crypto.subtle.digest = function (algo, data) { - // grok only ever asks for sha-256; g.__goSha256 returns an ArrayBuffer - var bytes = data instanceof Uint8Array ? data : new Uint8Array(data); - return Promise.resolve(g.__goSha256(bytes)); - }; - - // ---- Web Animations getComputedStyle sampler (the only real math we own) ---- - var K = 4096; - function cubicBezier(x1, y1, x2, y2, p) { - if (p <= 0) return 0; if (p >= 1) return 1; - function bez(t, a, b) { var mt = 1 - t; return 3 * a * mt * mt * t + 3 * b * mt * t * t + t * t * t; } - var lo = 0, hi = 1; - for (var i = 0; i < 100; i++) { var m = (lo + hi) / 2; if (bez(m, x1, x2) < p) lo = m; else hi = m; } - return bez((lo + hi) / 2, y1, y2); - } - function hexToRgb(h) { h = h.replace('#', ''); return [parseInt(h.slice(0, 2), 16), parseInt(h.slice(2, 4), 16), parseInt(h.slice(4, 6), 16)]; } - function sample(anim) { - var kf = anim.keyframes, dur = anim.duration || K; - var frac = anim.currentTime / dur; if (frac < 0) frac = 0; if (frac > 1) frac = 1; - var bm = /cubic-bezier\(([^)]+)\)/.exec(kf.easing || ''); - var eased = frac; - if (bm) { var p = bm[1].split(',').map(Number); eased = cubicBezier(p[0], p[1], p[2], p[3], frac); } - var c0 = hexToRgb(kf.color[0]), c1 = hexToRgb(kf.color[1]); - var col = [0, 1, 2].map(function (i) { return Math.round(c0[i] + (c1[i] - c0[i]) * eased); }); - var d0 = parseFloat(/rotate\(([-\d.]+)deg\)/.exec(kf.transform[0])[1]); - var d1 = parseFloat(/rotate\(([-\d.]+)deg\)/.exec(kf.transform[1])[1]); - var ang = (d0 + (d1 - d0) * eased) * Math.PI / 180; - var cos = Math.cos(ang), sin = Math.sin(ang); - return { color: 'rgb(' + col[0] + ', ' + col[1] + ', ' + col[2] + ')', - transform: 'matrix(' + cos + ', ' + sin + ', ' + (-sin) + ', ' + cos + ', 0, 0)' }; - } - - // ---- minimal DOM ---- - function makeEl(props) { - var e = { - nodeName: 'DIV', style: {}, childNodes: [], attrs: {}, _anim: null, _parent: null, - setAttribute: function (k, v) { e.attrs[k] = v; }, - getAttribute: function (k) { return (k in e.attrs) ? e.attrs[k] : (props && props.attrs && k in props.attrs ? props.attrs[k] : null); }, - appendChild: function (c) { e.childNodes.push(c); return c; }, - append: function (c) { e.childNodes.push(c); return c; }, - removeChild: function (c) { return c; }, - remove: function () {}, - animate: function (keyframes, opts) { - var anim = { - keyframes: keyframes, - duration: (opts && opts.duration) || (typeof opts === 'number' ? opts : K), - currentTime: 0, pause: function () {}, play: function () {}, cancel: function () {}, - effect: { getKeyframes: function () { return Array.isArray(keyframes) ? keyframes : [keyframes]; } }, - }; - e._anim = anim; return anim; - }, - getAnimations: function () { return e._anim ? [e._anim] : []; }, - }; - if (props) for (var k in props) if (k !== 'attrs') e[k] = props[k]; - Object.defineProperty(e, 'parentElement', { get: function () { return e._parent; } }); - Object.defineProperty(e, 'innerHTML', { set: function (v) { e._html = v; }, get: function () { return e._html; } }); - Object.defineProperty(e, 'textContent', { set: function (v) { e._text = v; }, get: function () { return e._text; } }); - return e; - } - - // .r-aufz1o group: g.childNodes[0].childNodes[1].getAttribute('d') = svg path whose - // numbers (after substring(9), split('C')) decode back to each curve [c0..c5,deg,b0..b3]. - function groupEl(flatCurves) { - var d = '_________' + flatCurves.map(function (c) { return c.join(' '); }).join('C'); - var path = makeEl({ attrs: { d: d } }); - var inner = makeEl(); inner.childNodes = [makeEl(), path]; - var outer = makeEl(); outer.childNodes = [inner]; outer._parent = makeEl(); - return outer; - } - - var docBody = makeEl(); - g.document = { - currentScript: null, body: docBody, head: makeEl(), - createElement: function (tag) { return makeEl({ nodeName: String(tag || 'div').toUpperCase() }); }, - querySelectorAll: function (sel) { - sel = String(sel); - // The seed is selected by a name/verification attribute selector - // (e.g. [name^=gr] or [name*=verification]); the curve group container is - // selected by a per-build hashed CLASS selector (e.g. .r-aufz1o, .r-3nqkqc) - // which rotates on every reship — so match by shape, not literal class. - if (/verification|name/i.test(sel)) { - var seed = g.__SEED; - return [{ nodeName: 'META', getAttribute: function (a) { return a === 'content' ? seed : null; }, - get content() { return seed; } }]; - } - if (/(^|\s|,)\./.test(sel)) { - var curves = JSON.parse(g.__CURVES); - return curves.map(function (grp) { - return groupEl(grp.map(function (cv) { return cv.color.concat([cv.deg], cv.bezier); })); - }); - } - return []; - }, - querySelector: function (sel) { var r = this.querySelectorAll(sel); return r[0] || null; }, - }; - g.window = g; - g.self = g; - g.getComputedStyle = function (el) { return el && el._anim ? sample(el._anim) : { color: 'rgb(0, 0, 0)', transform: 'none' }; }; - g.navigator = g.navigator || { userAgent: 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/133.0.0.0 Safari/537.36' }; - g.location = g.location || { href: 'https://grok.com/', origin: 'https://grok.com', pathname: '/' }; - - // ---- Turbopack capture + bootstrap ---- - var TP = []; TP.push = function (entry) { TP._entry = entry; return 0; }; - g.TURBOPACK = TP; - - // Called by Go AFTER the signer chunk is eval'd: run the module factory, grab default. - g.__grokBootstrap = function () { - var entry = TP._entry; - if (!entry) throw new Error('turbopack entry not registered'); - var factory = entry[2]; - var exports = {}; - var ctx = { - s: function () { - var flat = Array.prototype.slice.call(arguments).flat(Infinity); - var name = null; - for (var i = 0; i < flat.length; i++) { - var x = flat[i]; - if (typeof x === 'string') name = x; - else if (typeof x === 'function' && name != null) { - (function (nm, getter) { Object.defineProperty(exports, nm, { get: getter, configurable: true, enumerable: true }); })(name, x); - name = null; - } - } - }, - }; - factory(ctx); - // exports.default is a GETTER that invokes the module factory on every read, - // returning a FRESH async signer (path,method)=>Promise with a fresh internal - // curve cache. Keep the exports object; read .default fresh per sign so different - // sessions' curves never leak through the closure cache. - g.__grokExports = exports; - if (typeof exports.default !== 'function') throw new Error('no default export'); - }; - - // Called by Go per sign. Fresh signer each time via the re-invoking getter. - g.__grokSign = function () { - var signer = g.__grokExports.default; // fresh async signer - return signer(g.__PATH, g.__METHOD); // returns Promise - }; - - // Synchronous bridge: Go calls this via RunString (which drains goja's job queue), - // then reads __grokResult / __grokErr. Works because crypto.subtle.digest resolves - // synchronously (Promise.resolve over a Go SHA-256), so the whole await chain settles - // within the microtask drain. - g.__grokResult = null; - g.__grokErr = null; - g.__grokSignInto = function () { - g.__grokResult = null; g.__grokErr = null; - try { - g.__grokSign().then( - function (r) { g.__grokResult = r; }, - function (e) { g.__grokErr = (e && e.stack) ? String(e.stack) : String(e); } - ); - } catch (e) { g.__grokErr = (e && e.stack) ? String(e.stack) : String(e); } - }; -})(); diff --git a/backend/internal/provider/grok/video.go b/backend/internal/provider/grok/video.go deleted file mode 100644 index 2b2061d..0000000 --- a/backend/internal/provider/grok/video.go +++ /dev/null @@ -1,528 +0,0 @@ -package grok - -import ( - "bytes" - "context" - "encoding/json" - "errors" - "fmt" - "io" - "mime/multipart" - "net/textproto" - "os" - "regexp" - "strings" - "time" - - http "github.com/bogdanfinn/fhttp" - tlsclient "github.com/bogdanfinn/tls-client" - "github.com/google/uuid" -) - -// assetBase is where generated media artifacts live (the stream returns a -// path like "users//generated//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 != 15 { - seconds = 10 - } - - // Only the conversations/new submit egresses via the proxy; reference-frame - // upload and the mp4 download run on the local IP. - submitClient, err := c.newTLSClient() - if err != nil { - return nil, nil, err - } - directClient, err := c.newDirectTLSClient() - if err != nil { - return nil, nil, err - } - - // Self-heal the x-statsig-id anti-bot challenge for THIS session before the - // gated submit. Without it the header falls back to the static defaults, - // which go stale on a new grok web build and make conversations/new answer - // 403 (anti-bot). Fetch on the local IP (the proxy is rotating, so pinning - // the seed to a proxy egress IP is pointless). Failure is non-fatal - // (statsigID then uses defaults). - c.ensureChallenge(ctx, directClient, token) - - // 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, directClient, token, f) - if upErr != nil { - return nil, nil, upErr - } - imageRefs = append(imageRefs, url) - } - - // grok streams progress inline, but the stream is fragile: it regularly dies - // mid-body ("unexpected EOF") after only the conversation object / an early - // progress event — while grok keeps rendering the clip server-side. The - // rendered clip always lands at the deterministic URL - // assets.grok.com/users/{userId}/generated/{postId}/generated_video.mp4, so a - // dead stream is NOT a failed generation: poll that URL until the clip is - // ready instead of failing the job (which would also fail over and burn - // another account's credits for a video that already exists). - postID, userID, err := c.createPost(ctx, submitClient, 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 may be partial when psErr != nil — still worth scanning. - body, psErr := c.postStream(ctx, submitClient, token, "/rest/app-chat/conversations/new", payload) - if dir := strings.TrimSpace(os.Getenv("GROK_DUMP")); dir != "" { - _ = os.WriteFile(dir+"/grok_body_"+postID+".txt", []byte(body), 0o644) - } - // Out-of-credits surfaces as a stream error (HTTP is still 200) — not retryable. - if strings.Contains(body, "usagePoolExhausted") || strings.Contains(body, "media generation credits") { - return nil, nil, fmt.Errorf("%w: media generation credits exhausted", ErrQuotaExhausted) - } - if psErr != nil && (errors.Is(psErr, ErrAuth) || errors.Is(psErr, ErrQuotaExhausted)) { - return nil, nil, psErr - } - // A fatal stream error means grok definitively rejected this generation - // (content moderation, an unsupported parameter, etc.). Retrying — on this - // or any other account — fails identically and only burns the pool, so fail - // fast with a non-temporary error the pool won't fail over on. - if strings.Contains(body, "STREAM_ERROR_SEVERITY_FATAL") { - return nil, nil, fmt.Errorf("grok: video generation rejected by upstream (fatal stream error): %s", clip([]byte(body), 200)) - } - // Fast path: the artifact appears as "videoUrl":"users/.../generated_video.mp4". - artifact := "" - for _, m := range videoURLRe.FindAllStringSubmatch(body, -1) { - if v := strings.TrimSpace(m[1]); v != "" { - artifact = v // keep the last (progress=100) one - } - } - if artifact == "" && userID != "" { - // Stream ended without the artifact (dropped connection or early close): - // the render continues upstream — wait for the deterministic asset URL. - url := fmt.Sprintf("%susers/%s/generated/%s/generated_video.mp4", assetBase, userID, postID) - if werr := c.waitForAsset(ctx, directClient, token, url); werr == nil { - artifact = url - } else if errors.Is(werr, ErrAuth) || errors.Is(werr, context.Canceled) || errors.Is(werr, context.DeadlineExceeded) { - return nil, nil, werr - } - } - if artifact == "" { - if psErr != nil { - return nil, nil, psErr - } - 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, directClient, token, fullURL) - if err != nil { - return nil, nil, err - } - return data, meta, nil -} - -// uploadImage uploads one reference frame via /http/upload-file-v2/direct and -// returns its asset content URL for imageReferences. Cloudflare's bot score is -// per-request, so retry transient failures with backoff instead of failing the -// whole task. -func (c *Client) uploadImage(ctx context.Context, client tlsclient.HttpClient, token string, img []byte) (string, error) { - var res map[string]any - var err error - backoffs := []time.Duration{0, 2 * time.Second, 5 * time.Second, 10 * time.Second} - for _, wait := range backoffs { - if wait > 0 { - select { - case <-ctx.Done(): - return "", ctx.Err() - case <-time.After(wait): - } - } - res, err = c.uploadFileV2(ctx, client, token, img) - if err == nil || !errors.Is(err, ErrTemporaryUpstream) { - break - } - } - if err != nil { - return "", err - } - meta, _ := res["fileMetadata"].(map[string]any) - fileURI := "" - if meta != nil { - fileURI = strings.TrimSpace(stringValue(meta["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 -} - -func (c *Client) uploadFileV2(ctx context.Context, client tlsclient.HttpClient, token string, img []byte) (map[string]any, error) { - var buf bytes.Buffer - mw := multipart.NewWriter(&buf) - partHeader := textproto.MIMEHeader{} - partHeader.Set("Content-Disposition", fmt.Sprintf(`form-data; name="file"; filename="%s"`, uuid.NewString()+".png")) - partHeader.Set("Content-Type", "image/png") - part, err := mw.CreatePart(partHeader) - if err != nil { - return nil, err - } - if _, err := part.Write(img); err != nil { - return nil, err - } - if err := mw.WriteField("file_source", "IMAGINE_SELF_UPLOAD_FILE_SOURCE"); err != nil { - return nil, err - } - if err := mw.Close(); err != nil { - return nil, err - } - - req, err := http.NewRequest(http.MethodPost, apiBase+"/http/upload-file-v2/direct", bytes.NewReader(buf.Bytes())) - if err != nil { - return nil, err - } - req = req.WithContext(ctx) - c.applyHeaders(req, token, map[string]string{"content-type": mw.FormDataContentType()}) - - resp, err := client.Do(req) - if err != nil { - return nil, fmt.Errorf("%w: %v", ErrTemporaryUpstream, err) - } - defer resp.Body.Close() - raw, err := io.ReadAll(resp.Body) - if err != nil { - return nil, fmt.Errorf("%w: %v", ErrTemporaryUpstream, err) - } - if e := mapStatus("/http/upload-file-v2/direct", resp.StatusCode, raw); e != nil { - return nil, e - } - var out map[string]any - if err := json.Unmarshal(raw, &out); err != nil { - return nil, fmt.Errorf("%w: upload non-json: %s", ErrTemporaryUpstream, clip(raw, 120)) - } - return out, nil -} - -// createPost registers a video media post and returns its id (parentPostId). -// createPost creates the parent media post. grok uses its id as the videoId, -// and the rendered clip is stored under users/{userId}/generated/{postId}/. -func (c *Client) createPost(ctx context.Context, client tlsclient.HttpClient, token, prompt string) (postID, userID string, err 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) - postID = strings.TrimSpace(stringValue(post["id"])) - userID = strings.TrimSpace(stringValue(post["userId"])) - if postID == "" { - return "", "", fmt.Errorf("%w: media post missing id", ErrTemporaryUpstream) - } - return postID, userID, nil -} - -// waitForAsset polls a generated-asset URL (ranged GET on the local IP — -// assets.grok.com is not anti-bot gated) until grok finishes rendering it. -// 404 means "still rendering"; auth failures abort. -func (c *Client) waitForAsset(ctx context.Context, client tlsclient.HttpClient, token, url string) error { - deadline := time.Now().Add(30 * time.Minute) - for { - if err := ctx.Err(); err != nil { - return err - } - req, err := http.NewRequest(http.MethodGet, url, nil) - if err != nil { - return err - } - req = req.WithContext(ctx) - req.Header = http.Header{ - "user-agent": {userAgent}, - "referer": {origin + "/"}, - "cookie": {"sso=" + token + "; sso-rw=" + token}, - "range": {"bytes=0-0"}, - } - if resp, derr := client.Do(req); derr == nil { - code := resp.StatusCode - _, _ = io.Copy(io.Discard, resp.Body) - resp.Body.Close() - switch { - case code == http.StatusOK || code == http.StatusPartialContent: - return nil - case code == http.StatusUnauthorized || code == http.StatusForbidden: - return fmt.Errorf("%w: asset %d", ErrAuth, code) - } - } - if time.Now().After(deadline) { - return fmt.Errorf("%w: video render did not complete in time", ErrTemporaryUpstream) - } - select { - case <-ctx.Done(): - return ctx.Err() - case <-time.After(4 * time.Second): - } - } -} - -// 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 string(raw), err - } - if e := mapStatus(path, status, raw); e != nil { - return string(raw), 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 { - // Mid-body stream drops ("unexpected EOF", HTTP/2 resets) are common on - // long renders — return what was received so callers can still use it. - return raw, resp.StatusCode, fmt.Errorf("%w: %v", ErrTemporaryUpstream, err) - } - return raw, resp.StatusCode, nil -} - -// OpenAsset streams a grok asset (e.g. a generated video) authenticated with the -// account token — used by the async /v1/videos /content proxy. The caller MUST -// close the returned ReadCloser. -func (c *Client) OpenAsset(ctx context.Context, token, url string) (io.ReadCloser, string, error) { - token = strings.TrimSpace(strings.TrimPrefix(token, "Bearer ")) - if token == "" { - return nil, "", ErrAuth - } - // Asset streaming is pure bandwidth (no anti-bot) — egress on the local IP. - client, err := c.newDirectTLSClient() - if err != nil { - return nil, "", err - } - 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) - } - if resp.StatusCode < 200 || resp.StatusCode >= 300 { - resp.Body.Close() - if resp.StatusCode == 401 || resp.StatusCode == 403 { - return nil, "", fmt.Errorf("%w: asset %d", ErrAuth, resp.StatusCode) - } - return nil, "", fmt.Errorf("%w: asset %d", ErrTemporaryUpstream, resp.StatusCode) - } - ct := resp.Header.Get("Content-Type") - if ct == "" { - ct = "video/mp4" - } - return resp.Body, ct, nil -} - -// download fetches the rendered artifact. The clip is already generated at this -// point, so a transient failure here (HTTP/2 stream reset, CF hiccup) must not -// fail the whole task — retry with backoff. -func (c *Client) download(ctx context.Context, client tlsclient.HttpClient, token, url string) ([]byte, error) { - var data []byte - var err error - for _, wait := range []time.Duration{0, 2 * time.Second, 5 * time.Second, 10 * time.Second} { - if wait > 0 { - select { - case <-ctx.Done(): - return nil, ctx.Err() - case <-time.After(wait): - } - } - data, err = c.downloadOnce(ctx, client, token, url) - if err == nil || !errors.Is(err, ErrTemporaryUpstream) { - break - } - } - return data, err -} - -func (c *Client) downloadOnce(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, fmt.Errorf("%w: %v", ErrTemporaryUpstream, 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 == 403 && isBotChallenge(string(raw)): - // grok bot-detection or a Cloudflare challenge page ("Just a moment…"), - // NOT a dead token — transient, so a good account isn't killed by an - // IP/anti-bot hiccup. A statsig anti-bot rejection means the recipe likely - // went stale on a reship; kick the headless refresher to re-capture. - TriggerStatsigRefresh() - return fmt.Errorf("%w: %s 403 %s", ErrTemporaryUpstream, path, clip(raw, 160)) - case status == 401 || status == 403: - return fmt.Errorf("%w: %s %d %s", ErrAuth, path, status, clip(raw, 160)) - case status == 429: - // 429 is grok RATE-LIMITING ("Too many requests") — a transient error that - // must NOT kill the account. Only a body that names a credit/usage-pool - // exhaustion is a real quota wall. - if isCreditError(string(raw)) { - return fmt.Errorf("%w: %s 429 %s", ErrQuotaExhausted, path, clip(raw, 160)) - } - return fmt.Errorf("%w: %s 429 %s", ErrTemporaryUpstream, 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)) - } -} - -// isBotChallenge reports whether a 403 body is an anti-bot interstitial rather -// than a real auth rejection: grok's own "anti-bot" marker or a Cloudflare -// challenge page ("Just a moment…" / cf-chl / challenge-platform). -func isBotChallenge(s string) bool { - s = strings.ToLower(s) - return strings.Contains(s, "anti-bot") || - strings.Contains(s, "just a moment") || - strings.Contains(s, "cf-chl") || - strings.Contains(s, "challenge-platform") || - strings.Contains(s, "cf_chl") -} - -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") -} diff --git a/backend/internal/provider/grok/zz_probe_test.go b/backend/internal/provider/grok/zz_probe_test.go deleted file mode 100644 index 98fef07..0000000 --- a/backend/internal/provider/grok/zz_probe_test.go +++ /dev/null @@ -1,110 +0,0 @@ -package grok - -import ( - "context" - "crypto/sha1" - "fmt" - "io" - "os" - "path/filepath" - "regexp" - "strings" - "testing" - "time" - - http "github.com/bogdanfinn/fhttp" -) - -var innerChunkRe = regexp.MustCompile(`static/chunks/[a-zA-Z0-9_.~\-/]+\.js`) - -func hnew(ctx context.Context, url, token string) (*http.Request, error) { - req, err := http.NewRequest(http.MethodGet, url, nil) - if err != nil { - return nil, err - } - req = req.WithContext(ctx) - req.Header = http.Header{ - "accept": {"text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8"}, - "accept-language": {"en-US,en;q=0.9"}, - "user-agent": {userAgent}, - "cookie": {"sso=" + token + "; sso-rw=" + token}, - http.HeaderOrderKey: {"accept", "accept-language", "user-agent", "cookie"}, - } - return req, nil -} - -func TestGrokCacheChunks(t *testing.T) { - token := os.Getenv("DIAG_SSO") - if token == "" { - t.Skip("no DIAG_SSO") - } - dir := os.Getenv("DIAG_CACHE") - if dir == "" { - dir = `f:\ai-gateway\vivid-ai\backend\_chunks` - } - os.MkdirAll(dir, 0o755) - logf, _ := os.Create(filepath.Join(dir, "_log.txt")) - defer logf.Close() - lg := func(f string, a ...any) { - fmt.Fprintf(logf, f+"\n", a...) - logf.Sync() - } - - c := NewClient(os.Getenv("DIAG_PROXY")) - ctx, cancel := context.WithTimeout(context.Background(), 3000*time.Second) - defer cancel() - client, err := c.newTLSClient() - if err != nil { - t.Fatal(err) - } - - hreq, _ := hnew(ctx, apiBase+"/", token) - hresp, err := client.Do(hreq) - if err != nil { - t.Fatal(err) - } - body, _ := io.ReadAll(hresp.Body) - hresp.Body.Close() - html := string(body) - os.WriteFile(filepath.Join(dir, "_home.html"), body, 0o644) - - seen := map[string]bool{} - queue := dedupe(chunkPathRe.FindAllString(html, -1)) - lg("home chunk refs: %d", len(queue)) - - saved := 0 - for len(queue) > 0 && saved < 3000 { - p := queue[0] - queue = queue[1:] - norm := p - if !strings.HasPrefix(norm, "/_next/") { - norm = "/_next/" + strings.TrimPrefix(norm, "/") - } - if seen[norm] { - continue - } - seen[norm] = true - src, err := fetchChunk(ctx, client, norm) - if err != nil { - continue - } - saved++ - base := norm[strings.LastIndex(norm, "/")+1:] - fn := fmt.Sprintf("%x_%s", sha1.Sum([]byte(norm)), base) - if len(fn) > 120 { - fn = fn[:120] - } - os.WriteFile(filepath.Join(dir, fn), []byte(src), 0o644) - for _, m := range innerChunkRe.FindAllString(src, -1) { - nn := "/_next/" + m - if !seen[nn] { - queue = append(queue, nn) - } - } - if saved%50 == 0 { - lg("saved=%d queue=%d", saved, len(queue)) - } - } - lg("CRAWL DONE saved=%d queue_left=%d", saved, len(queue)) - fmt.Printf("CRAWL DONE saved=%d\n", saved) -} diff --git a/backend/internal/provider/imagine/client.go b/backend/internal/provider/imagine/client.go deleted file mode 100644 index f01a738..0000000 --- a/backend/internal/provider/imagine/client.go +++ /dev/null @@ -1,452 +0,0 @@ -// Package imagine implements the Imagine.art (vyro.ai) provider client. The -// durable credential is a JSON blob {"token","refreshToken"}: `token` is a ~6h -// access JWT used as Authorization: Bearer for the API, and `refreshToken` is a -// ~7d JWT that mints a fresh pair via /apis/v1/auth/other/refresh/web when the -// access token expires. Both rotate on refresh, so the new pair MUST be saved. -// tls-client gives a Chrome JA3/JA4 so vyro's edge doesn't flag the requests. -package imagine - -import ( - "context" - "crypto/rand" - "encoding/base64" - "encoding/json" - "errors" - "fmt" - "io" - "strconv" - "strings" - "sync" - "time" - - http "github.com/bogdanfinn/fhttp" - tlsclient "github.com/bogdanfinn/tls-client" - "github.com/bogdanfinn/tls-client/profiles" -) - -const ( - apiBase = "https://imagine.vyro.ai" - teamsBase = "https://teams-imagine.vyro.ai" - authBase = "https://auth.vyro.ai" - webOrigin = "https://www.imagine.art" - userAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36" -) - -var ( - ErrAuth = errors.New("imagine auth failed") - ErrQuotaExhausted = errors.New("imagine quota exhausted") - ErrTemporaryUpstream = errors.New("imagine upstream temporary error") -) - -// refreshLeadSeconds renews the access token this many seconds BEFORE it expires -// (proactive, not lazy at expiry) — the maintenance sweep keeps tokens fresh so a -// dormant account's rotating refreshToken never lapses. -const refreshLeadSeconds = 600 // 10 minutes - -type Client struct { - proxy string - // freshest credential per account (key: user id) + a per-account refresh lock, - // so concurrent callers don't each spend the rotating refresh_token — the first - // refreshes, the rest reuse the cached fresh credential. - mu sync.Mutex - creds map[string]string - locks map[string]*sync.Mutex -} - -func NewClient(proxy string) *Client { - return &Client{proxy: strings.TrimSpace(proxy), creds: map[string]string{}, locks: map[string]*sync.Mutex{}} -} - -func (c *Client) SetProxy(proxy string) { - c.proxy = strings.TrimSpace(proxy) -} - -func (c *Client) userLock(userID string) *sync.Mutex { - c.mu.Lock() - defer c.mu.Unlock() - m, ok := c.locks[userID] - if !ok { - m = &sync.Mutex{} - c.locks[userID] = m - } - return m -} - -// --------------------------------------------------------------------------- -// Credential helpers -// --------------------------------------------------------------------------- - -type credential struct { - Token string `json:"token"` - RefreshToken string `json:"refreshToken"` - // Email is the real account email — supplied at import, used for display and - // (pool,email) dedup. It is NOT in the JWT (which only carries userId), so it - // must be carried across refreshes (the refresh response omits it). - Email string `json:"email,omitempty"` - // ParentID is a canvas node the account OWNS, used as the generation's - // parent_id. Imagine rejects any parent the account doesn't own ("user does - // not have access to parent asset") and silently orphans a parent-less - // generation (charged but never produced) — so it's supplied at import and - // carried across refreshes. - ParentID string `json:"parentId,omitempty"` -} - -func parseCred(s string) (credential, bool) { - var cr credential - if json.Unmarshal([]byte(strings.TrimSpace(s)), &cr) != nil { - return cr, false - } - if strings.TrimSpace(cr.Token) == "" || strings.TrimSpace(cr.RefreshToken) == "" { - return cr, false - } - return cr, true -} - -func buildCred(token, refresh, email, parentID string) string { - b, _ := json.Marshal(credential{ - Token: strings.TrimSpace(token), - RefreshToken: strings.TrimSpace(refresh), - Email: strings.TrimSpace(email), - ParentID: strings.TrimSpace(parentID), - }) - return string(b) -} - -// ParentIDFromCred returns the canvas parent node id supplied at import. -func ParentIDFromCred(cred string) string { - cr, ok := parseCred(cred) - if !ok { - return "" - } - return strings.TrimSpace(cr.ParentID) -} - -func looksLikeJWT(s string) bool { - return len(strings.Split(strings.TrimSpace(s), ".")) == 3 -} - -// IsImagineToken reports whether a pasted credential is an Imagine.art account: -// a JSON object carrying a non-empty token + refreshToken that both look like -// JWTs. Distinguishes it from adobe/leonardo/krea cookies. -func IsImagineToken(value string) bool { - cr, ok := parseCred(value) - if !ok { - return false - } - return looksLikeJWT(cr.Token) && looksLikeJWT(cr.RefreshToken) -} - -// jwtClaims base64url-decodes the JWT payload (segment 1) into a claims map. -func jwtClaims(token string) map[string]any { - parts := strings.Split(strings.TrimSpace(token), ".") - if len(parts) < 2 { - return nil - } - raw, err := base64.RawURLEncoding.DecodeString(parts[1]) - if err != nil { - // tolerate padded variants - if raw, err = base64.URLEncoding.DecodeString(parts[1]); err != nil { - return nil - } - } - var m map[string]any - if json.Unmarshal(raw, &m) != nil { - return nil - } - return m -} - -func userIDFromToken(token string) string { - claims := jwtClaims(token) - if claims == nil { - return "" - } - if v := strings.TrimSpace(stringValue(claims["userId"])); v != "" { - return v - } - return strings.TrimSpace(stringValue(claims["sub"])) -} - -func tokenExp(token string) int64 { - claims := jwtClaims(token) - if claims == nil { - return 0 - } - return toInt64(claims["exp"]) -} - -// EmailFromCred returns the real account email supplied at import; if absent it -// falls back to the JWT userId so (pool,email) dedup still has a stable key. -func EmailFromCred(cred string) string { - cr, ok := parseCred(cred) - if !ok { - return "" - } - if e := strings.TrimSpace(cr.Email); e != "" { - return e - } - return userIDFromToken(cr.Token) -} - -// UserIDFromCred returns the JWT userId (== org_id used for credit/generation). -func UserIDFromCred(cred string) string { - cr, ok := parseCred(cred) - if !ok { - return "" - } - return userIDFromToken(cr.Token) -} - -// --------------------------------------------------------------------------- -// Refresh -// --------------------------------------------------------------------------- - -// RefreshIfNeeded returns a credential whose access token is still valid: if the -// stored one is (near) expired it spends the refreshToken to mint a fresh pair -// and rebuilds the credential. Returns (cred, changed, err); changed=true means -// the caller must persist the new credential (both tokens rotate). ErrAuth means -// the refreshToken is dead → the account is gone. -func (c *Client) RefreshIfNeeded(ctx context.Context, cred string) (string, bool, error) { - cr, ok := parseCred(cred) - if !ok { - return cred, false, nil // unparseable — let the downstream call surface the error - } - userID := userIDFromToken(cr.Token) - now := time.Now().Unix() - - lk := c.userLock(userID) - lk.Lock() - defer lk.Unlock() - - // A concurrent caller may already have refreshed this account. - if userID != "" { - c.mu.Lock() - cached := c.creds[userID] - c.mu.Unlock() - if cc, ok := parseCred(cached); ok && tokenExp(cc.Token)-refreshLeadSeconds > now { - return cached, cached != cred, nil - } - } - if tokenExp(cr.Token)-refreshLeadSeconds > now { - return cred, false, nil // still valid - } - - respBody, status, err := c.refreshPost(ctx, cr.RefreshToken) - if err != nil { - return "", false, fmt.Errorf("%w: refresh: %s", ErrTemporaryUpstream, err.Error()) - } - if status == 400 || status == 401 || status == 403 { - return "", false, ErrAuth - } - if status != 200 { - return "", false, fmt.Errorf("%w: refresh http %d: %s", ErrTemporaryUpstream, status, clip(respBody, 120)) - } - var rb struct { - Result struct { - SessionToken string `json:"sessionToken"` - RefreshToken string `json:"refreshToken"` - } `json:"result"` - } - if json.Unmarshal(respBody, &rb) != nil || strings.TrimSpace(rb.Result.SessionToken) == "" { - return "", false, ErrAuth - } - newRefresh := rb.Result.RefreshToken - if strings.TrimSpace(newRefresh) == "" { - newRefresh = cr.RefreshToken // some responses may omit it — keep the old one - } - newCred := buildCred(rb.Result.SessionToken, newRefresh, cr.Email, cr.ParentID) - if userID != "" { - c.mu.Lock() - c.creds[userID] = newCred - c.mu.Unlock() - } - return newCred, true, nil -} - -func (c *Client) refreshPost(ctx context.Context, refreshToken string) ([]byte, int, error) { - client, err := c.newDirectTLSClient() - if err != nil { - return nil, 0, err - } - req, err := http.NewRequest(http.MethodPost, authBase+"/apis/v1/auth/other/refresh/web", nil) - if err != nil { - return nil, 0, err - } - req = req.WithContext(ctx) - req.Header = http.Header{ - "accept": {"application/json, text/plain, */*"}, - "authorization": {"Bearer " + refreshToken}, - "origin": {webOrigin}, - "referer": {webOrigin + "/"}, - "user-agent": {userAgent}, - http.HeaderOrderKey: { - "accept", "authorization", "origin", "referer", "user-agent", - }, - } - resp, err := client.Do(req) - if err != nil { - return nil, 0, err - } - defer resp.Body.Close() - b, err := io.ReadAll(resp.Body) - return b, resp.StatusCode, err -} - -// --------------------------------------------------------------------------- -// Credits -// --------------------------------------------------------------------------- - -// FetchCreditsBalance reads the account's credit balance via /v1/credit. -// remaining is the `total` field. 401/403 → ErrAuth (token dead). Returns the -// normalized map shared by all providers. -func (c *Client) FetchCreditsBalance(ctx context.Context, cred string) (map[string]any, error) { - cr, ok := parseCred(cred) - if !ok { - return unknownBalance("bad credential"), nil - } - userID := userIDFromToken(cr.Token) - body, status, err := c.apiGetP(ctx, cr.Token, apiBase+"/v1/credit?org_id="+userID, false) - if err != nil { - return unknownBalance("network: " + err.Error()), nil - } - if status == 401 || status == 403 { - return nil, ErrAuth - } - if status != 200 { - return unknownBalance(fmt.Sprintf("http %d: %s", status, clip(body, 160))), nil - } - var cb struct { - Status string `json:"status"` - Total int `json:"total"` - } - if err := json.Unmarshal(body, &cb); err != nil { - return unknownBalance("non-json"), nil - } - return map[string]any{ - "remaining": cb.Total, - "used": nil, - "total": nil, - "unknown": false, - "error": nil, - "email": emptyStringNil(EmailFromCred(cred)), - }, nil -} - -// --------------------------------------------------------------------------- -// HTTP helpers -// --------------------------------------------------------------------------- - -func (c *Client) apiGet(ctx context.Context, token, url string) ([]byte, int, error) { - return c.apiGetP(ctx, token, url, true) -} - -// apiGetP picks the egress: polling runs direct (local IP). -func (c *Client) apiGetP(ctx context.Context, token, url string, useProxy bool) ([]byte, int, error) { - client, err := c.newTLSClientP(useProxy) - if err != nil { - return nil, 0, err - } - req, err := http.NewRequest(http.MethodGet, url, nil) - if err != nil { - return nil, 0, err - } - req = req.WithContext(ctx) - req.Header = http.Header{ - "accept": {"application/json, text/plain, */*"}, - "authorization": {"Bearer " + token}, - "origin": {webOrigin}, - "referer": {webOrigin + "/"}, - "user-agent": {userAgent}, - http.HeaderOrderKey: { - "accept", "authorization", "origin", "referer", "user-agent", - }, - } - resp, err := client.Do(req) - if err != nil { - return nil, 0, err - } - defer resp.Body.Close() - b, err := io.ReadAll(resp.Body) - return b, resp.StatusCode, err -} - -func (c *Client) newTLSClient() (tlsclient.HttpClient, error) { return c.newTLSClientP(true) } - -// newDirectTLSClient egresses on the local IP (never the proxy). Used for -// polling and result download. -func (c *Client) newDirectTLSClient() (tlsclient.HttpClient, error) { return c.newTLSClientP(false) } - -func (c *Client) newTLSClientP(useProxy bool) (tlsclient.HttpClient, error) { - options := []tlsclient.HttpClientOption{ - tlsclient.WithTimeoutSeconds(60), - tlsclient.WithClientProfile(profiles.Chrome_120), - } - if useProxy && c.proxy != "" { - options = append(options, tlsclient.WithProxyUrl(c.proxy)) - } - return tlsclient.NewHttpClient(tlsclient.NewNoopLogger(), options...) -} - -// --------------------------------------------------------------------------- -// Small util -// --------------------------------------------------------------------------- - -func uuid4() string { - var b [16]byte - _, _ = rand.Read(b[:]) - b[6] = (b[6] & 0x0f) | 0x40 - b[8] = (b[8] & 0x3f) | 0x80 - return fmt.Sprintf("%x-%x-%x-%x-%x", b[0:4], b[4:6], b[6:8], b[8:10], b[10:16]) -} - -func unknownBalance(reason string) map[string]any { - return map[string]any{ - "remaining": nil, "used": nil, "total": nil, "unknown": true, "error": reason, - } -} - -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 toInt64(v any) int64 { - switch x := v.(type) { - case float64: - return int64(x) - case int64: - return x - case int: - return int64(x) - case json.Number: - n, _ := x.Int64() - return n - case string: - n, _ := strconv.ParseInt(strings.TrimSpace(x), 10, 64) - return n - default: - return 0 - } -} - -func emptyStringNil(v string) any { - if strings.TrimSpace(v) == "" { - return nil - } - return v -} - -func clip(b []byte, n int) string { - s := strings.TrimSpace(string(b)) - if len(s) > n { - return s[:n] - } - return s -} diff --git a/backend/internal/provider/imagine/image.go b/backend/internal/provider/imagine/image.go deleted file mode 100644 index e539002..0000000 --- a/backend/internal/provider/imagine/image.go +++ /dev/null @@ -1,251 +0,0 @@ -package imagine - -import ( - "bytes" - "context" - "encoding/json" - "fmt" - "io" - "mime/multipart" - "strconv" - "strings" - "time" - - http "github.com/bogdanfinn/fhttp" -) - -// GenerateImage runs the full Imagine.art pipeline: submit the txt2img job, poll -// the org objects feed until the batch finishes, then download the produced -// image. styleID picks the model (41001 = 1.5 / 2K, 41004 = 1.5pro / 4K). -// HTTP 402 → ErrQuotaExhausted. These models are pure text2img (no refs). -func (c *Client) GenerateImage(ctx context.Context, cred string, styleID int, resolution, aspectRatio, prompt string, downloadResult bool) ([]byte, map[string]any, error) { - cr, ok := parseCred(cred) - if !ok { - return nil, nil, ErrAuth - } - userID := userIDFromToken(cr.Token) - - metadata, _ := json.Marshal(map[string]any{ - "placeholderUuid": uuid4(), - "promptWithoutManipulation": prompt, - "modeId": 0, - }) - // parent_id MUST be a canvas node this account owns — the server rejects a - // foreign id ("user does not have access to parent asset") and silently - // orphans a parent-less generation (charged but never produced). It's supplied - // at import (credential.parentId). - fields := map[string]string{ - "style_id": strconv.Itoa(styleID), - "aspect_ratio": aspectRatio, - "resolution": resolution, - "variation": "txt2img", - "prompt": prompt, - "is_enhance": "0", - "count": "1", - "clientVersion": "1", - "org_id": userID, - "use_plugin": "false", - "metadata": string(metadata), - } - if pid := strings.TrimSpace(cr.ParentID); pid != "" { - fields["parent_id"] = pid - } - - var buf bytes.Buffer - w := multipart.NewWriter(&buf) - for k, v := range fields { - _ = w.WriteField(k, v) - } - _ = w.Close() - - body, status, err := c.apiPost(ctx, cr.Token, apiBase+"/v1/image/generations/upload", w.FormDataContentType(), buf.Bytes()) - if err != nil { - return nil, nil, fmt.Errorf("%w: submit: %s", ErrTemporaryUpstream, err.Error()) - } - if status == 401 || status == 403 { - return nil, nil, ErrAuth - } - if status == 402 { - return nil, nil, ErrQuotaExhausted - } - if status != 200 && status != 201 { - return nil, nil, fmt.Errorf("%w: submit http %d: %s", ErrTemporaryUpstream, status, clip(body, 200)) - } - var jobs []struct { - BatchID string `json:"batchId"` - ID string `json:"id"` - Status string `json:"status"` - } - if err := json.Unmarshal(body, &jobs); err != nil || len(jobs) == 0 || jobs[0].BatchID == "" { - return nil, nil, fmt.Errorf("%w: no batch id: %s", ErrTemporaryUpstream, clip(body, 200)) - } - batchID := jobs[0].BatchID - - imageURL, err := c.pollImage(ctx, cr.Token, userID, batchID) - if err != nil { - return nil, nil, err - } - meta := map[string]any{"batch_id": batchID, "image_url": imageURL, "org_id": userID} - if !downloadResult { - return nil, meta, nil - } - data, err := c.download(ctx, imageURL) - if err != nil { - return nil, nil, err - } - return data, meta, nil -} - -// pollImage polls the org objects feed until the entry for our batch finishes, -// then extracts its asset URL (image_url is a JSON-encoded array string). -func (c *Client) pollImage(ctx context.Context, token, userID, batchID string) (string, error) { - ticker := time.NewTicker(3 * time.Second) - defer ticker.Stop() - // Poll for the full generation budget (caller's genCtx), leaving headroom for - // the download, instead of a shorter hardcoded cap that killed slow jobs early. - deadline := time.Now().Add(4 * time.Minute) - if dl, ok := ctx.Deadline(); ok { - deadline = dl.Add(-60 * time.Second) - } - - url := teamsBase + "/v1/org/" + userID + "/objects?batch=true&limit=50&service=image,chat-image" - for { - body, status, err := c.apiGetP(ctx, token, url, false) - if err == nil && status == 200 { - var resp struct { - Data []struct { - BatchID string `json:"batch_id"` - Status string `json:"status"` - Code int `json:"code"` - // Current shape: the produced asset lives at url.generation[0]. - URL struct { - Generation []string `json:"generation"` - } `json:"url"` - ImageURL string `json:"image_url"` // legacy fallback - } `json:"data"` - } - if json.Unmarshal(body, &resp) == nil { - for _, o := range resp.Data { - if o.BatchID != batchID { - continue - } - st := strings.ToLower(strings.TrimSpace(o.Status)) - switch { - case st == "finished" || o.Code == 2: - if u := firstNonEmpty(o.URL.Generation); u != "" { - return u, nil - } - if u := firstImageURL(o.ImageURL); u != "" { - return u, nil - } - case st == "failed" || st == "error": - return "", fmt.Errorf("%w: job %s", ErrTemporaryUpstream, st) - } - } - } - } else if status == 401 || status == 403 { - return "", ErrAuth - } - if time.Now().After(deadline) { - return "", fmt.Errorf("%w: generation timed out", ErrTemporaryUpstream) - } - select { - case <-ctx.Done(): - return "", ctx.Err() - case <-ticker.C: - } - } -} - -// firstNonEmpty returns the first non-blank string in a slice. -func firstNonEmpty(ss []string) string { - for _, s := range ss { - if strings.TrimSpace(s) != "" { - return strings.TrimSpace(s) - } - } - return "" -} - -// firstImageURL parses the image_url field — a JSON-encoded array of URLs — and -// returns the first one. Tolerates a bare string too. -func firstImageURL(raw string) string { - raw = strings.TrimSpace(raw) - if raw == "" { - return "" - } - var urls []string - if json.Unmarshal([]byte(raw), &urls) == nil { - for _, u := range urls { - if strings.TrimSpace(u) != "" { - return strings.TrimSpace(u) - } - } - return "" - } - if strings.HasPrefix(raw, "http") { - return raw - } - return "" -} - -func (c *Client) download(ctx context.Context, url string) ([]byte, error) { - client, err := c.newDirectTLSClient() - if err != nil { - return nil, err - } - req, err := http.NewRequest(http.MethodGet, url, nil) - if err != nil { - return nil, err - } - req = req.WithContext(ctx) - req.Header = http.Header{ - "accept": {"image/avif,image/webp,image/png,image/*,*/*;q=0.8"}, - "user-agent": {userAgent}, - "referer": {webOrigin + "/"}, - } - resp, err := client.Do(req) - if err != nil { - return nil, err - } - defer resp.Body.Close() - b, err := io.ReadAll(resp.Body) - if err != nil { - return nil, err - } - if resp.StatusCode != 200 { - return nil, fmt.Errorf("%w: image download http %d", ErrTemporaryUpstream, resp.StatusCode) - } - return b, nil -} - -// apiPost issues a POST with a raw body + content-type, carrying the bearer token. -func (c *Client) apiPost(ctx context.Context, token, url, contentType string, body []byte) ([]byte, int, error) { - client, err := c.newTLSClient() - if err != nil { - return nil, 0, err - } - req, err := http.NewRequest(http.MethodPost, url, bytes.NewReader(body)) - if err != nil { - return nil, 0, err - } - req = req.WithContext(ctx) - req.Header = http.Header{ - "accept": {"application/json, text/plain, */*"}, - "authorization": {"Bearer " + token}, - "content-type": {contentType}, - "origin": {webOrigin}, - "referer": {webOrigin + "/"}, - "user-agent": {userAgent}, - http.HeaderOrderKey: { - "accept", "authorization", "content-type", "origin", "referer", "user-agent", - }, - } - resp, err := client.Do(req) - if err != nil { - return nil, 0, err - } - defer resp.Body.Close() - b, err := io.ReadAll(resp.Body) - return b, resp.StatusCode, err -} diff --git a/backend/internal/provider/krea/client.go b/backend/internal/provider/krea/client.go deleted file mode 100644 index 72c0373..0000000 --- a/backend/internal/provider/krea/client.go +++ /dev/null @@ -1,499 +0,0 @@ -// Package krea implements the Krea.ai (krea.ai) provider client. The durable -// credential is the browser cookie (Supabase "sb-superb-auth-token"); Krea's own -// Next.js backend reads it directly, so quota and generation just forward the -// cookie — there's no separate token-exchange step. tls-client gives a Chrome -// JA3/JA4 so Krea's Cloudflare edge doesn't flag the requests. -package krea - -import ( - "bytes" - "context" - "encoding/base64" - "encoding/json" - "errors" - "fmt" - "io" - "strconv" - "strings" - "sync" - "time" - - http "github.com/bogdanfinn/fhttp" - tlsclient "github.com/bogdanfinn/tls-client" - "github.com/bogdanfinn/tls-client/profiles" -) - -// kreaAnonKey is Krea's public Supabase anon key (fixed, embedded in their -// frontend) — required as the apikey/bearer when refreshing a session token. -const kreaAnonKey = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJyb2xlIjoiYW5vbiIsImlzcyI6InN1cGFiYXNlIiwiaWF0IjoxNzc1Mjc4ODU3LCJleHAiOjE5MzI5NTg4NTd9.NUiqEOd__QsCCMjo3D1zrCAda5dLV2F5p6Kf584sZKc" - -const ( - apiBase = "https://www.krea.ai" - userAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36" -) - -var ( - ErrAuth = errors.New("krea auth failed") - ErrQuotaExhausted = errors.New("krea quota exhausted") - ErrTemporaryUpstream = errors.New("krea upstream temporary error") -) - -// refreshLeadSeconds renews the access token this many seconds BEFORE it expires -// (not lazily at expiry). The maintenance sweep refreshes proactively, so a -// dormant account's token never lapses — once the rotating refresh_token is gone -// (expired/consumed) the account can't recover, so we keep it perpetually fresh. -const refreshLeadSeconds = 600 // 10 minutes - -type Client struct { - proxy string - // freshest cookie per account (key: user id) + a per-account refresh lock, so - // concurrent callers don't each spend the single-use (rotating) refresh_token — - // the first refreshes, the rest reuse the cached fresh cookie. - mu sync.Mutex - cookies map[string]string - locks map[string]*sync.Mutex - // actAt = last /app activation time per account (key: user id); actLocks gives - // a per-account lock so concurrent generations wait for the first to finish the - // (once-per-daily-reset) activation instead of each loading /app. - actAt map[string]int64 - actLocks map[string]*sync.Mutex -} - -func NewClient(proxy string) *Client { - return &Client{ - proxy: strings.TrimSpace(proxy), - cookies: map[string]string{}, - locks: map[string]*sync.Mutex{}, - actAt: map[string]int64{}, - actLocks: map[string]*sync.Mutex{}, - } -} - -func (c *Client) userLock(userID string) *sync.Mutex { - c.mu.Lock() - defer c.mu.Unlock() - m, ok := c.locks[userID] - if !ok { - m = &sync.Mutex{} - c.locks[userID] = m - } - return m -} - -// RefreshIfNeeded returns a cookie whose access_token is still valid: if the -// stored one is (near) expired it spends the refresh_token to mint a new session -// and rebuilds the cookie. Returns (cookie, changed, err); changed=true means the -// caller must persist the new cookie (the refresh_token rotated). ErrAuth means -// the refresh_token is dead → the account is gone. -func (c *Client) RefreshIfNeeded(ctx context.Context, cookie string) (string, bool, error) { - authVal := authCookieValue(cookie) - sess, ok := decodeSession(authVal) - if !ok { - return cookie, false, nil // unparseable — let the downstream call surface the error - } - userID := nestedStr(sess, "user", "id") - now := time.Now().Unix() - - lk := c.userLock(userID) - lk.Lock() - defer lk.Unlock() - - // A concurrent caller may already have refreshed this account. - if userID != "" { - c.mu.Lock() - cached := c.cookies[userID] - c.mu.Unlock() - if cs, ok := decodeSession(authCookieValue(cached)); ok && toInt64(cs["expires_at"])-refreshLeadSeconds > now { - return cached, cached != cookie, nil - } - } - if toInt64(sess["expires_at"])-refreshLeadSeconds > now { - return cookie, false, nil // still valid - } - refreshTok := strings.TrimSpace(stringValue(sess["refresh_token"])) - if refreshTok == "" { - return "", false, ErrAuth - } - - respBody, status, err := c.refreshPost(ctx, refreshTok) - if err != nil { - return "", false, fmt.Errorf("%w: refresh: %s", ErrTemporaryUpstream, err.Error()) - } - if status == 400 || status == 401 || status == 403 { - return "", false, ErrAuth - } - if status != 200 { - return "", false, fmt.Errorf("%w: refresh http %d: %s", ErrTemporaryUpstream, status, clip(respBody, 120)) - } - var ns map[string]any - if json.Unmarshal(respBody, &ns) != nil || strings.TrimSpace(stringValue(ns["access_token"])) == "" { - return "", false, ErrAuth - } - newCookie := replaceAuthCookie(cookie, "base64-"+base64.StdEncoding.EncodeToString(respBody)) - if userID != "" { - c.mu.Lock() - c.cookies[userID] = newCookie - c.mu.Unlock() - } - return newCookie, true, nil -} - -func (c *Client) refreshPost(ctx context.Context, refreshToken string) ([]byte, int, error) { - client, err := c.newDirectTLSClient() - if err != nil { - return nil, 0, err - } - body, _ := json.Marshal(map[string]string{"refresh_token": refreshToken}) - req, err := http.NewRequest(http.MethodPost, apiBase+"/auth/v1/token?grant_type=refresh_token", bytes.NewReader(body)) - if err != nil { - return nil, 0, err - } - req = req.WithContext(ctx) - req.Header = http.Header{ - "accept": {"*/*"}, - "content-type": {"application/json;charset=UTF-8"}, - "apikey": {kreaAnonKey}, - "authorization": {"Bearer " + kreaAnonKey}, - "x-client-info": {"supabase-ssr/0.6.1 createBrowserClient"}, - "x-supabase-api-version": {"2024-01-01"}, - "origin": {apiBase}, - "referer": {apiBase + "/"}, - "user-agent": {userAgent}, - http.HeaderOrderKey: { - "accept", "content-type", "apikey", "authorization", "x-client-info", - "x-supabase-api-version", "origin", "referer", "user-agent", - }, - } - resp, err := client.Do(req) - if err != nil { - return nil, 0, err - } - defer resp.Body.Close() - b, err := io.ReadAll(resp.Body) - return b, resp.StatusCode, err -} - -// chunkSize is supabase-ssr's per-cookie chunk limit; larger sessions (e.g. -// Google-OAuth accounts) are split into sb-superb-auth-token.0/.1/... -const chunkSize = 3600 - -func cookieVal(cookie, name string) string { - for _, p := range strings.Split(cookie, ";") { - p = strings.TrimSpace(p) - if v, ok := strings.CutPrefix(p, name+"="); ok { - return v - } - } - return "" -} - -// authCookieValue returns the full auth value, transparently reassembling a -// chunked cookie (sb-superb-auth-token.0 + .1 + ...) or returning the single one. -func authCookieValue(cookie string) string { - if v := cookieVal(cookie, "sb-superb-auth-token"); v != "" { - return v - } - var b strings.Builder - for i := 0; ; i++ { - v := cookieVal(cookie, fmt.Sprintf("sb-superb-auth-token.%d", i)) - if v == "" { - break - } - b.WriteString(v) - } - return b.String() -} - -// replaceAuthCookie drops every sb-superb-auth-token[.N] cookie and re-adds the -// new value (chunked the same way supabase-ssr would if it's large), preserving -// all other cookies (krea-workspace-id, etc.). -func replaceAuthCookie(cookie, newValue string) string { - var out []string - for _, p := range strings.Split(cookie, ";") { - t := strings.TrimSpace(p) - if t == "" || strings.HasPrefix(t, "sb-superb-auth-token=") || strings.HasPrefix(t, "sb-superb-auth-token.") { - continue - } - out = append(out, t) - } - if len(newValue) <= chunkSize { - out = append(out, "sb-superb-auth-token="+newValue) - } else { - for i, off := 0, 0; off < len(newValue); i++ { - end := off + chunkSize - if end > len(newValue) { - end = len(newValue) - } - out = append(out, fmt.Sprintf("sb-superb-auth-token.%d=%s", i, newValue[off:end])) - off = end - } - } - return strings.Join(out, "; ") -} - -// decodeSession base64-decodes the auth cookie value into the session JSON map. -func decodeSession(authValue string) (map[string]any, bool) { - v := strings.TrimPrefix(strings.TrimSpace(authValue), "base64-") - if v == "" { - return nil, false - } - raw, err := base64.StdEncoding.DecodeString(v) - if err != nil { - raw, err = base64.RawURLEncoding.DecodeString(v) - if err != nil { - return nil, false - } - } - var m map[string]any - if json.Unmarshal(raw, &m) != nil { - return nil, false - } - return m, true -} - -func nestedStr(m map[string]any, k1, k2 string) string { - if sub, ok := m[k1].(map[string]any); ok { - return strings.TrimSpace(stringValue(sub[k2])) - } - return "" -} - -func toInt64(v any) int64 { - switch x := v.(type) { - case float64: - return int64(x) - case int64: - return x - case int: - return int64(x) - case json.Number: - n, _ := x.Int64() - return n - case string: - n, _ := strconv.ParseInt(strings.TrimSpace(x), 10, 64) - return n - default: - return 0 - } -} - -func (c *Client) SetProxy(proxy string) { - c.proxy = strings.TrimSpace(proxy) -} - -// IsKreaCookie reports whether a pasted credential is a Krea cookie: it carries -// the Supabase auth cookie. Distinguishes it from adobe/leonardo cookies. -func IsKreaCookie(value string) bool { - return strings.Contains(value, "sb-superb-auth-token") -} - -// EmailFromCookie decodes the account email straight out of the cookie's embedded -// Supabase session (no network), handling chunked cookies too. -func EmailFromCookie(cookie string) string { - sess, ok := decodeSession(authCookieValue(cookie)) - if !ok { - return "" - } - return nestedStr(sess, "user", "email") -} - -// FetchCreditsBalance reads the account's free-credit balance via /api/billing-data. -// remaining is the integer floor of balance.free (per spec: 17.94 → 17). 401 → -// ErrAuth (cookie dead). Returns the normalized map shared by all providers. -func (c *Client) FetchCreditsBalance(ctx context.Context, cookie string) (map[string]any, error) { - // Detach from the request ctx so a page refresh can't cancel the probe - // mid-flight (which left accounts stuck at "—"). - probeCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 30*time.Second) - defer cancel() - // NOTE: no /app here — the heavy SSR activation is done separately (on recovery - // and at generation via Activate). This probe just reads the current balance. - body, status, err := c.apiGetP(probeCtx, cookie, "/api/billing-data", false) - if err != nil { - return unknownBalance("network: " + err.Error()), nil - } - if status == 401 || status == 403 { - return nil, ErrAuth - } - if status != 200 { - return unknownBalance(fmt.Sprintf("http %d: %s", status, clip(body, 160))), nil - } - // 余额在 balance.free(真实剩余,小数,随用量递减)。krea 的这个字段一直都在, - // 只是排在很长的 entitlements 之后 —— 只读它,不用套餐配额兜底。 - var bd struct { - Balance struct { - Free float64 `json:"free"` - Total float64 `json:"total"` - } `json:"balance"` - } - if err := json.Unmarshal(body, &bd); err != nil { - return unknownBalance("non-json"), nil - } - remaining := int(bd.Balance.Free) // floor - return map[string]any{ - "remaining": remaining, - "used": nil, - "total": int(bd.Balance.Total), - "unknown": false, - "error": nil, - "email": emptyStringNil(EmailFromCookie(cookie)), - }, nil -} - -// Activate loads the authenticated SSR app page (/app), which is what makes krea -// grant the account's DAILY free balance — a cold API-only call (billing-data / -// generate) otherwise sees balance.free=0 and 402s. Done at most ONCE per account -// per daily reset, under a per-account lock: the first caller loads /app while -// concurrent callers wait, then everyone proceeds (no redundant /app). Called -// before each generation and by the daily activation sweep. Best-effort. -func (c *Client) Activate(ctx context.Context, cookie string) { - key := accountKey(cookie) - lastReset := (time.Now().Unix() / 86400) * 86400 - c.mu.Lock() - doneToday := key != "" && c.actAt[key] >= lastReset - c.mu.Unlock() - if doneToday { - return - } - lk := c.actLock(key) - lk.Lock() - defer lk.Unlock() - // Re-check after acquiring the lock — another caller may have just activated. - c.mu.Lock() - doneToday = key != "" && c.actAt[key] >= lastReset - c.mu.Unlock() - if doneToday { - return - } - actCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 60*time.Second) - defer cancel() - _, _, _ = c.apiGetP(actCtx, cookie, "/app", false) - c.mu.Lock() - c.actAt[key] = time.Now().Unix() - c.mu.Unlock() -} - -func (c *Client) actLock(key string) *sync.Mutex { - c.mu.Lock() - defer c.mu.Unlock() - m, ok := c.actLocks[key] - if !ok { - m = &sync.Mutex{} - c.actLocks[key] = m - } - return m -} - -// accountKey is the stable per-account id (the Supabase user id from the session -// cookie) used to key activation state — survives cookie rotation. -func accountKey(cookie string) string { - if sess, ok := decodeSession(authCookieValue(cookie)); ok { - return nestedStr(sess, "user", "id") - } - return "" -} - -// apiGet issues a GET to a krea.ai API path carrying the account cookie. -func (c *Client) apiGet(ctx context.Context, cookie, path string) ([]byte, int, error) { - return c.apiGetP(ctx, cookie, path, true) -} - -// apiGetP picks the egress: polling / asset resolution run direct (local IP). -func (c *Client) apiGetP(ctx context.Context, cookie, path string, useProxy bool) ([]byte, int, error) { - client, err := c.newTLSClientP(useProxy) - if err != nil { - return nil, 0, err - } - req, err := http.NewRequest(http.MethodGet, apiBase+path, nil) - if err != nil { - return nil, 0, err - } - req = req.WithContext(ctx) - req.Header = http.Header{ - "accept": {"*/*"}, - "accept-language": {"en-US,en;q=0.9"}, - "cookie": {cookie}, - "referer": {apiBase + "/"}, - "user-agent": {userAgent}, - "sec-fetch-dest": {"empty"}, - "sec-fetch-mode": {"cors"}, - "sec-fetch-site": {"same-origin"}, - http.HeaderOrderKey: { - "accept", "accept-language", "cookie", "referer", "user-agent", - "sec-fetch-dest", "sec-fetch-mode", "sec-fetch-site", - }, - } - resp, err := client.Do(req) - if err != nil { - return nil, 0, err - } - defer resp.Body.Close() - b, err := io.ReadAll(resp.Body) - return b, resp.StatusCode, err -} - -func (c *Client) newTLSClient() (tlsclient.HttpClient, error) { return c.newTLSClientP(true) } - -// newDirectTLSClient egresses on the local IP (never the proxy). Used for -// reference-image upload, polling and result download. -func (c *Client) newDirectTLSClient() (tlsclient.HttpClient, error) { return c.newTLSClientP(false) } - -func (c *Client) newTLSClientP(useProxy bool) (tlsclient.HttpClient, error) { - options := []tlsclient.HttpClientOption{ - tlsclient.WithTimeoutSeconds(60), - tlsclient.WithClientProfile(profiles.Chrome_120), - } - if useProxy && c.proxy != "" { - options = append(options, tlsclient.WithProxyUrl(c.proxy)) - } - return tlsclient.NewHttpClient(tlsclient.NewNoopLogger(), options...) -} - -func unknownBalance(reason string) map[string]any { - return map[string]any{ - "remaining": nil, "used": nil, "total": nil, "unknown": true, "error": reason, - } -} - -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 intValue(v any) int { - switch x := v.(type) { - case int: - return x - case float64: - return int(x) - case json.Number: - n, _ := x.Int64() - return int(n) - case string: - n, _ := strconv.Atoi(strings.TrimSpace(x)) - return n - default: - return 0 - } -} - -func emptyStringNil(v string) any { - if strings.TrimSpace(v) == "" { - return nil - } - return v -} - -func clip(b []byte, n int) string { - s := strings.TrimSpace(string(b)) - if len(s) > n { - return s[:n] - } - return s -} diff --git a/backend/internal/provider/krea/image.go b/backend/internal/provider/krea/image.go deleted file mode 100644 index fe2c374..0000000 --- a/backend/internal/provider/krea/image.go +++ /dev/null @@ -1,324 +0,0 @@ -package krea - -import ( - "bytes" - "context" - "encoding/json" - "fmt" - "io" - "mime/multipart" - "strings" - "time" - - http "github.com/bogdanfinn/fhttp" -) - -const ( - genModel = "flux2-klein4b" // the single exposed Krea model - genEndpoint = "/api/jobs/v2/new/fluxKlein4b" - refStrength = 0.4 -) - -// ensureProject returns a flux project id for the account: the first existing -// project, or a freshly created one. Generation requires a project. -func (c *Client) ensureProject(ctx context.Context, cookie string) (string, error) { - body, status, err := c.apiGetP(ctx, cookie, "/api/flux-projects", false) - if err != nil { - return "", fmt.Errorf("%w: list projects: %s", ErrTemporaryUpstream, err.Error()) - } - if status == 401 || status == 403 { - return "", ErrAuth - } - if status == 200 { - var projs []struct { - ID string `json:"id"` - } - if json.Unmarshal(body, &projs) == nil { - for _, p := range projs { - if strings.TrimSpace(p.ID) != "" { - return p.ID, nil - } - } - } - } - // No project on this account (e.g. brand-new). Try to create one; if that - // fails, fall back to generating WITHOUT a project (Krea assigns a default). - cb, cs, cerr := c.apiPostJSON(ctx, cookie, "/api/flux-projects", map[string]any{"title": "vivid"}) - if cerr == nil && (cs == 200 || cs == 201) { - var pr struct { - ID string `json:"id"` - } - if json.Unmarshal(cb, &pr) == nil && pr.ID != "" { - return pr.ID, nil - } - } - return "", nil // generate without an explicit project -} - -// uploadImage uploads a reference image (i2i) and returns its app-uploads URL. -func (c *Client) uploadImage(ctx context.Context, cookie string, img []byte) (string, error) { - var buf bytes.Buffer - w := multipart.NewWriter(&buf) - fw, err := w.CreateFormFile("file", "image.png") - if err != nil { - return "", err - } - if _, err := fw.Write(img); err != nil { - return "", err - } - _ = w.Close() - body, status, err := c.apiPostP(ctx, cookie, "/api/upload?", w.FormDataContentType(), buf.Bytes(), false) - if err != nil { - return "", fmt.Errorf("%w: upload: %s", ErrTemporaryUpstream, err.Error()) - } - if status == 401 || status == 403 { - return "", ErrAuth - } - if status != 200 { - return "", fmt.Errorf("%w: upload http %d: %s", ErrTemporaryUpstream, status, clip(body, 160)) - } - var ur struct { - ImageURL string `json:"imageUrl"` - } - if json.Unmarshal(body, &ur) != nil || ur.ImageURL == "" { - return "", fmt.Errorf("%w: no imageUrl", ErrTemporaryUpstream) - } - return ur.ImageURL, nil -} - -// GenerateImage runs the full Krea image pipeline: ensure a project, (for i2i) -// upload reference images, submit the job, poll until done, then resolve and -// download the produced image. 402 INSUFFICIENT_BALANCE → ErrQuotaExhausted. -func (c *Client) GenerateImage(ctx context.Context, cookie, prompt string, width, height int, refImages [][]byte, downloadResult bool) ([]byte, map[string]any, error) { - // Ensure the daily free balance is granted (load /app) before generating, so a - // not-yet-activated account doesn't 402 INSUFFICIENT_BALANCE. Lock-guarded and - // once-per-daily-reset — concurrent gens wait for the first activation, already - // activated ones skip straight through. - c.Activate(ctx, cookie) - - projectID, err := c.ensureProject(ctx, cookie) - if err != nil { - return nil, nil, err - } - - var styleImages []map[string]any - for _, img := range refImages { - if len(img) == 0 { - continue - } - url, upErr := c.uploadImage(ctx, cookie, img) - if upErr != nil { - return nil, nil, upErr - } - styleImages = append(styleImages, map[string]any{"url": url, "strength": refStrength, "source": "upload"}) - } - - payload := map[string]any{ - "provider": genModel, - "prompt": prompt, - "width": width, - "height": height, - "strength": 1, - "steps": 28, - "guidance_scale_flux": 3.5, - "presetStyles": []any{}, - "batchSize": 2, - "guidance": 3.5, - } - if projectID != "" { - payload["project"] = projectID - } - if len(styleImages) > 0 { - payload["styleImages"] = styleImages - } - payloadJSON, _ := json.Marshal(payload) - - // Submit (multipart with a single "payload" field). - var buf bytes.Buffer - w := multipart.NewWriter(&buf) - _ = w.WriteField("payload", string(payloadJSON)) - _ = w.Close() - body, status, err := c.apiPost(ctx, cookie, genEndpoint, w.FormDataContentType(), buf.Bytes()) - if err != nil { - return nil, nil, fmt.Errorf("%w: submit: %s", ErrTemporaryUpstream, err.Error()) - } - if status == 401 || status == 403 { - return nil, nil, ErrAuth - } - if status == 402 || strings.Contains(string(body), "INSUFFICIENT_BALANCE") { - return nil, nil, ErrQuotaExhausted - } - if status != 200 && status != 201 { - return nil, nil, fmt.Errorf("%w: submit http %d: %s", ErrTemporaryUpstream, status, clip(body, 200)) - } - var jobs []struct { - JobID string `json:"job_id"` - } - if err := json.Unmarshal(body, &jobs); err != nil || len(jobs) == 0 || jobs[0].JobID == "" { - return nil, nil, fmt.Errorf("%w: no job id: %s", ErrTemporaryUpstream, clip(body, 200)) - } - // batchSize=2 returns two jobs; keep the SECOND image and discard the first - // (per spec). Fall back to the first if only one came back. - jobID := jobs[0].JobID - if len(jobs) >= 2 && jobs[1].JobID != "" { - jobID = jobs[1].JobID - } - - // Poll until terminal, then resolve the produced image. - imageURL, err := c.pollImage(ctx, cookie, jobID) - if err != nil { - return nil, nil, err - } - meta := map[string]any{"job_id": jobID, "image_url": imageURL, "project": projectID} - if !downloadResult { - return nil, meta, nil - } - data, err := c.download(ctx, imageURL) - if err != nil { - return nil, nil, err - } - return data, meta, nil -} - -// pollImage polls job-status until the job leaves the queue, then matches the -// produced asset by generation_job_id and returns its image URL. -func (c *Client) pollImage(ctx context.Context, cookie, jobID string) (string, error) { - ticker := time.NewTicker(3 * time.Second) - defer ticker.Stop() - // Poll for the full generation budget (caller's genCtx), leaving headroom for - // the download, instead of a shorter hardcoded cap that killed slow jobs early. - deadline := time.Now().Add(4 * time.Minute) - if dl, ok := ctx.Deadline(); ok { - deadline = dl.Add(-60 * time.Second) - } - - for { - body, status, err := c.apiGetP(ctx, cookie, "/api/job-status?id="+jobID, false) - if err == nil && status == 200 { - var js struct { - Status string `json:"status"` - } - if json.Unmarshal(body, &js) == nil { - switch strings.ToLower(js.Status) { - case "complete", "completed", "succeeded", "success", "done", "finished": - if url, e := c.assetForJob(ctx, cookie, jobID); e == nil && url != "" { - return url, nil - } - case "failed", "error", "cancelled", "canceled": - return "", fmt.Errorf("%w: job %s", ErrTemporaryUpstream, js.Status) - } - } - } else if status == 401 || status == 403 { - return "", ErrAuth - } - if time.Now().After(deadline) { - return "", fmt.Errorf("%w: generation timed out", ErrTemporaryUpstream) - } - select { - case <-ctx.Done(): - return "", ctx.Err() - case <-ticker.C: - } - } -} - -// assetForJob finds the generated asset produced by a job and returns its URL. -func (c *Client) assetForJob(ctx context.Context, cookie, jobID string) (string, error) { - body, status, err := c.apiGetP(ctx, cookie, "/api/assets?filter=generated&offset=0", false) - if err != nil || status != 200 { - return "", fmt.Errorf("assets http %d", status) - } - var assets []struct { - ImageURL string `json:"image_url"` - Metadata struct { - GenerationJobID string `json:"generation_job_id"` - } `json:"metadata"` - } - if json.Unmarshal(body, &assets) != nil { - return "", fmt.Errorf("assets non-json") - } - for _, a := range assets { - if a.Metadata.GenerationJobID == jobID && strings.TrimSpace(a.ImageURL) != "" { - return a.ImageURL, nil - } - } - return "", fmt.Errorf("asset not found yet") -} - -func (c *Client) download(ctx context.Context, url string) ([]byte, error) { - client, err := c.newDirectTLSClient() - if err != nil { - return nil, err - } - req, err := http.NewRequest(http.MethodGet, url, nil) - if err != nil { - return nil, err - } - req = req.WithContext(ctx) - req.Header = http.Header{ - "accept": {"image/avif,image/webp,image/png,image/*,*/*;q=0.8"}, - "user-agent": {userAgent}, - "referer": {apiBase + "/"}, - } - resp, err := client.Do(req) - if err != nil { - return nil, err - } - defer resp.Body.Close() - b, err := io.ReadAll(resp.Body) - if err != nil { - return nil, err - } - if resp.StatusCode != 200 { - return nil, fmt.Errorf("%w: image download http %d", ErrTemporaryUpstream, resp.StatusCode) - } - return b, nil -} - -// apiPost issues a POST with a raw body + content-type, carrying the cookie. -func (c *Client) apiPost(ctx context.Context, cookie, path, contentType string, body []byte) ([]byte, int, error) { - return c.apiPostP(ctx, cookie, path, contentType, body, true) -} - -// apiPostP picks the egress: reference-image upload runs direct (local IP), the -// generate submit uses the proxy. -func (c *Client) apiPostP(ctx context.Context, cookie, path, contentType string, body []byte, useProxy bool) ([]byte, int, error) { - client, err := c.newTLSClientP(useProxy) - if err != nil { - return nil, 0, err - } - req, err := http.NewRequest(http.MethodPost, apiBase+path, bytes.NewReader(body)) - if err != nil { - return nil, 0, err - } - req = req.WithContext(ctx) - req.Header = http.Header{ - "accept": {"*/*"}, - "accept-language": {"en-US,en;q=0.9"}, - "content-type": {contentType}, - "cookie": {cookie}, - "origin": {apiBase}, - "referer": {apiBase + "/"}, - "user-agent": {userAgent}, - "sec-fetch-dest": {"empty"}, - "sec-fetch-mode": {"cors"}, - "sec-fetch-site": {"same-origin"}, - http.HeaderOrderKey: { - "accept", "accept-language", "content-type", "cookie", "origin", - "referer", "user-agent", "sec-fetch-dest", "sec-fetch-mode", "sec-fetch-site", - }, - } - resp, err := client.Do(req) - if err != nil { - return nil, 0, err - } - defer resp.Body.Close() - b, err := io.ReadAll(resp.Body) - return b, resp.StatusCode, err -} - -func (c *Client) apiPostJSON(ctx context.Context, cookie, path string, payload any) ([]byte, int, error) { - // project bootstrap runs on the local IP; only the generate submit uses the proxy. - b, _ := json.Marshal(payload) - return c.apiPostP(ctx, cookie, path, "application/json", b, false) -} diff --git a/backend/internal/provider/leonardo/client.go b/backend/internal/provider/leonardo/client.go deleted file mode 100644 index bd60499..0000000 --- a/backend/internal/provider/leonardo/client.go +++ /dev/null @@ -1,732 +0,0 @@ -// Package leonardo implements the Leonardo.ai (app.leonardo.ai) provider client. -// Unlike chatgpt/runway (whose JWT IS the stored credential), Leonardo's durable -// credential is the browser COOKIE (better-auth session): the bearer access token -// it mints lives only ~1h. So every call here takes the cookie and derives a -// fresh JWT on the fly via /api/auth/get-session — there is no long-lived token to -// store or a separate refresh profile to maintain. tls-client gives a Chrome -// JA3/JA4 fingerprint so the requests aren't flagged. -package leonardo - -import ( - "bytes" - "context" - "encoding/json" - "errors" - "fmt" - "io" - "net/url" - "strconv" - "strings" - "sync" - "time" - - http "github.com/bogdanfinn/fhttp" - tlsclient "github.com/bogdanfinn/tls-client" - "github.com/bogdanfinn/tls-client/profiles" -) - -const ( - appBase = "https://app.leonardo.ai" - graphqlURL = "https://api.leonardo.ai/v1/graphql" - schemaVersion = "1.255.2" - userAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/133.0.0.0 Safari/537.36" - // sec-ch-ua must agree with userAgent's major version — a mismatch is itself a - // bot signal. - secChUA = `"Not(A:Brand";v="99", "Google Chrome";v="133", "Chromium";v="133"` - // get-session only answers with a token once better-auth's cookie cache has - // been refreshed by a cross-origin-cookie call; called cold it returns 200 - // null, which looks exactly like a dead cookie. It also sits behind Vercel's - // checkpoint, which 429s machine requests from some exit IPs — so a failed - // attempt is retried, falling back to the proxy for a different exit IP. - getSessionAttempts = 10 - getSessionRetryDelay = 2 * time.Second - // A warmed 200 null usually means the session is gone server-side, but the - // same answer also comes back when the checkpoint silently serves a session- - // less page to an exit IP — so retry it a few times (later attempts through - // the proxy) before calling the account dead. Fewer attempts than the 429 - // budget: each one costs two requests and a truly dead cookie never recovers. - getSessionNullAttempts = 3 -) - -var ( - ErrAuth = errors.New("leonardo auth failed") - ErrQuotaExhausted = errors.New("leonardo quota exhausted") - ErrTemporaryUpstream = errors.New("leonardo upstream temporary error") -) - -type Client struct { - proxy string - // sessions caches the short-lived access token per cookie so we don't hit - // /api/auth/get-session on every call — Leonardo rate-limits that endpoint - // (429) hard, so re-using the ~1h JWT is essential. - mu sync.Mutex - sessions map[string]*Session - // rotated maps a stored cookie to the freshest value Leonardo handed back via - // Set-Cookie (better-auth rotates its session_data cookie cache). The service - // persists it; keeping it here means an unpersisted rotation still works for - // the rest of the process's life. - rotated map[string]string - // refreshing serialises get-session per cookie. Two concurrent refreshes hand - // Cognito the same refresh token twice and its reuse detection revokes the - // whole session — the account then answers 401 forever. - refreshing map[string]*sync.Mutex -} - -func NewClient(proxy string) *Client { - return &Client{proxy: strings.TrimSpace(proxy), sessions: map[string]*Session{}, rotated: map[string]string{}, refreshing: map[string]*sync.Mutex{}} -} - -func (c *Client) SetProxy(proxy string) { - c.proxy = strings.TrimSpace(proxy) -} - -// IsLeonardoCookie reports whether a pasted credential is a Leonardo cookie: it -// carries the better-auth session cookie name. This is what disambiguates it from -// an Adobe cookie at import time. -func IsLeonardoCookie(value string) bool { - return strings.Contains(value, "__Secure-better-auth.session_token") || - strings.Contains(value, "better-auth.session_data") -} - -// HasSessionData reports whether the cookie carries better-auth's session_data -// cache. Leonardo authenticates get-session off THAT cookie: session_token alone -// answers 200 null (no bearer), which looks exactly like a dead account — so a -// cookie without it must be rejected at import instead of dying later. -func HasSessionData(value string) bool { - return strings.Contains(value, "better-auth.session_data") -} - -// RotatedCookie returns the freshest value for a stored cookie when Leonardo -// rotated its session_data cache, so the caller can persist it. -func (c *Client) RotatedCookie(cookie string) (string, bool) { - key := strings.TrimSpace(cookie) - c.mu.Lock() - defer c.mu.Unlock() - fresh, ok := c.rotated[key] - return fresh, ok && fresh != key -} - -// mergeCookies applies a response's Set-Cookie pairs onto a request cookie -// string, keeping the original order and appending new names. -func mergeCookies(cookie string, setCookies []string) string { - if len(setCookies) == 0 { - return cookie - } - updates := map[string]string{} - order := []string{} - for _, sc := range setCookies { - pair := strings.TrimSpace(strings.Split(sc, ";")[0]) - name, value, ok := strings.Cut(pair, "=") - name = strings.TrimSpace(name) - if !ok || name == "" { - continue - } - if _, seen := updates[name]; !seen { - order = append(order, name) - } - updates[name] = value - } - if len(updates) == 0 { - return cookie - } - var out []string - used := map[string]bool{} - for _, part := range strings.Split(cookie, ";") { - part = strings.TrimSpace(part) - if part == "" { - continue - } - name, _, _ := strings.Cut(part, "=") - name = strings.TrimSpace(name) - if v, ok := updates[name]; ok { - used[name] = true - if v == "" { // a cleared cookie drops out - continue - } - out = append(out, name+"="+v) - continue - } - out = append(out, part) - } - for _, name := range order { - if used[name] || updates[name] == "" { - continue - } - out = append(out, name+"="+updates[name]) - } - return strings.Join(out, "; ") -} - -// keepsSession reports whether a merged cookie still carries BOTH components -// get-session needs: the session token and better-auth's session_data cache. -// A merge that loses either one (a Set-Cookie clearing a cache chunk) must be -// discarded — sending it would answer 200 null, i.e. look like a dead account. -func keepsSession(cookie string) bool { - return strings.Contains(cookie, "__Secure-better-auth.session_token") && HasSessionData(cookie) -} - -// Session is the result of /api/auth/get-session: the short-lived bearer plus the -// ids the GraphQL API needs (cognitoSub for the quota query, userId for the feed -// and the CDN image path) and the human-facing account fields. -type Session struct { - AccessToken string - CognitoSub string - UserID string - Email string - Name string - ExpiresAt int64 - // Cookie is the cookie that produced this session, with any Set-Cookie - // rotation applied — persist it so the account keeps authenticating. - Cookie string -} - -// GetSession exchanges the cookie for a fresh access token + account ids. Only a -// 401 or a 200 carrying no access token means the session is dead → ErrAuth; -// everything else (notably the 403/429 人机校验 page) is a temporary error. -func (c *Client) GetSession(ctx context.Context, cookie string) (*Session, error) { - cookie = strings.TrimSpace(cookie) - if cookie == "" { - return nil, ErrAuth - } - // Re-use a cached, still-valid access token (keep a 60s safety margin) instead - // of hitting the heavily rate-limited get-session endpoint again. - c.mu.Lock() - if cs, ok := c.sessions[cookie]; ok && cs.ExpiresAt-60 > time.Now().Unix() { - c.mu.Unlock() - return cs, nil - } - c.mu.Unlock() - - // Only one refresh per cookie at a time; the others wait and then re-use the - // token it minted. - gate := c.refreshGate(cookie) - gate.Lock() - defer gate.Unlock() - c.mu.Lock() - if cs, ok := c.sessions[cookie]; ok && cs.ExpiresAt-60 > time.Now().Unix() { - c.mu.Unlock() - return cs, nil - } - c.mu.Unlock() - - // Use the freshest known value (an earlier response may have rotated the - // better-auth cookie cache) rather than the possibly stale stored cookie. - send := cookie - c.mu.Lock() - if fresh, ok := c.rotated[cookie]; ok && fresh != "" { - send = fresh - } - c.mu.Unlock() - - var body []byte - var status int - var setCookies []string - var err error - var warmed bool - // 401 from cross-origin-cookie means Leonardo已经作废了这个 session（不是人机校验）， - // 记下来好把日志写成"会话被吊销"而不是含糊的 get-session null。 - warmStatus := 0 - for attempt := 0; attempt < getSessionAttempts; attempt++ { - if attempt > 0 { - select { - case <-ctx.Done(): - return nil, ctx.Err() - case <-time.After(getSessionRetryDelay): - } - } - // The local IP normally passes the checkpoint; only once it doesn't is the - // proxy worth its latency (and its own share of 429s). - useProxy := attempt > 0 && c.proxy != "" - client, cerr := c.newTLSClientP(useProxy) - if cerr != nil { - err = cerr - continue - } - // Warm up on the same connection: cross-origin-cookie hands back a fresh - // better-auth cookie cache (and CF_Access_Token). Without it get-session - // answers 200 null even for a perfectly healthy cookie. - warmed = false - warmCookies, wstatus, werr := c.warmSession(ctx, client, send) - warmStatus = wstatus - if werr == nil { - warmed = true - if merged := mergeCookies(send, warmCookies); merged != send && keepsSession(merged) { - send = merged - } - } - status, body, setCookies, err = c.fetchSession(ctx, client, send) - if err == nil && warmed && status != 429 && status != 403 { - if status == 200 && sessionAccessToken(body) == "" && attempt < getSessionNullAttempts-1 { - continue // retry a null session on another exit IP before giving up - } - break - } - } - if err != nil { - return nil, fmt.Errorf("%w: %s", ErrTemporaryUpstream, err.Error()) - } - // Only a real app answer may rotate the stored cookie. The 403/429 人机校验 页 - // also sends Set-Cookie (often CLEARING better-auth cookies), and persisting - // that would strip the session_data cache — after which get-session answers - // 200 null and a perfectly healthy account looks dead. - if status == 200 { - if merged := mergeCookies(send, setCookies); keepsSession(merged) { - send = merged - } - if send != cookie { - c.mu.Lock() - c.rotated[cookie] = send - c.mu.Unlock() - } - } - if status == 401 { - return nil, fmt.Errorf("%w: get-session http 401: %s", ErrAuth, clip(body, 160)) - } - if status != 200 { - // 403 / 429 here is the Vercel / Cloudflare 人机校验 页，不是 cookie 失效 — - // 当成临时错误，否则健康的号会被误判死。 - return nil, fmt.Errorf("%w: get-session http %d: %s", ErrTemporaryUpstream, status, clip(body, 160)) - } - var raw struct { - Session struct { - AccessToken string `json:"accessToken"` - CognitoSub string `json:"cognitoSub"` - UserID string `json:"userId"` - HasuraUserID string `json:"hasuraUserId"` - TokenExpiry int64 `json:"accessTokenExpiry"` - } `json:"session"` - User struct { - ID string `json:"id"` - Email string `json:"email"` - Name string `json:"name"` - } `json:"user"` - } - if err := json.Unmarshal(body, &raw); err != nil { - return nil, fmt.Errorf("%w: get-session non-json", ErrTemporaryUpstream) - } - if strings.TrimSpace(raw.Session.AccessToken) == "" { - if warmStatus == 401 { - // Leonardo 明确拒了这份 session_token：服务端已把会话吊销，重新导入 cookie 才能恢复。 - return nil, fmt.Errorf("%w: session revoked (cross-origin-cookie 401)", ErrAuth) - } - if !warmed { - // A cold get-session (the cookie cache was never refreshed) answers null - // for healthy accounts too — temporary, never a reason to kill the account. - return nil, fmt.Errorf("%w: get-session 200 without accessToken (no warmup): %s", ErrTemporaryUpstream, clip(body, 160)) - } - // No bearer despite 200 → the cookie no longer authenticates. Carry the body - // so the log says WHICH shape it was (null session vs a session without a - // token) instead of a bare "auth failed". - return nil, fmt.Errorf("%w: get-session 200 without accessToken: %s", ErrAuth, clip(body, 160)) - } - uid := raw.Session.UserID - if uid == "" { - uid = raw.Session.HasuraUserID - } - if uid == "" { - uid = raw.User.ID - } - sess := &Session{ - AccessToken: raw.Session.AccessToken, - CognitoSub: raw.Session.CognitoSub, - UserID: uid, - Email: strings.TrimSpace(raw.User.Email), - Name: strings.TrimSpace(raw.User.Name), - ExpiresAt: raw.Session.TokenExpiry, - Cookie: send, - } - if sess.ExpiresAt > time.Now().Unix() { - c.mu.Lock() - c.sessions[cookie] = sess - c.mu.Unlock() - } - return sess, nil -} - -// refreshGate returns the per-cookie lock that serialises get-session refreshes. -func (c *Client) refreshGate(cookie string) *sync.Mutex { - c.mu.Lock() - defer c.mu.Unlock() - gate, ok := c.refreshing[cookie] - if !ok { - gate = &sync.Mutex{} - c.refreshing[cookie] = gate - } - return gate -} - -// warmSession calls cross-origin-cookie, whose response refreshes better-auth's -// cookie cache and CF_Access_Token. It returns the Set-Cookie headers; a -// checkpoint answer (403/429) is an error, since get-session would then be cold. -func (c *Client) warmSession(ctx context.Context, client tlsclient.HttpClient, cookie string) ([]string, int, error) { - req, err := http.NewRequest(http.MethodGet, appBase+"/api/auth/cross-origin-cookie", nil) - if err != nil { - return nil, 0, err - } - req = req.WithContext(ctx) - req.Header = sessionHeader(cookie) - resp, err := client.Do(req) - if err != nil { - return nil, 0, err - } - defer resp.Body.Close() - _, _ = io.Copy(io.Discard, resp.Body) - if resp.StatusCode == 403 || resp.StatusCode == 429 { - return nil, resp.StatusCode, fmt.Errorf("cross-origin-cookie http %d", resp.StatusCode) - } - return resp.Header["Set-Cookie"], resp.StatusCode, nil -} - -// fetchSession performs one get-session call and returns its status, body and -// Set-Cookie headers. -func (c *Client) fetchSession(ctx context.Context, client tlsclient.HttpClient, cookie string) (int, []byte, []string, error) { - req, err := http.NewRequest(http.MethodGet, appBase+"/api/auth/get-session", nil) - if err != nil { - return 0, nil, nil, err - } - req = req.WithContext(ctx) - req.Header = sessionHeader(cookie) - resp, err := client.Do(req) - if err != nil { - return 0, nil, nil, err - } - defer resp.Body.Close() - body, _ := io.ReadAll(resp.Body) - return resp.StatusCode, body, resp.Header["Set-Cookie"], nil -} - -// sessionAccessToken pulls the bearer out of a get-session body; "" means the -// answer carried none (a null session, or a session without a token). -func sessionAccessToken(body []byte) string { - var raw struct { - Session struct { - AccessToken string `json:"accessToken"` - } `json:"session"` - } - if err := json.Unmarshal(body, &raw); err != nil { - return "" - } - return strings.TrimSpace(raw.Session.AccessToken) -} - -// sessionHeader is the auth endpoints' request shape, copied from a real -// browser's call (HAR): a same-origin GET carries NO origin header and DOES -// carry the ua client hints + priority — sending origin while omitting the -// hints is exactly the shape Vercel's checkpoint 429s. -func sessionHeader(cookie string) http.Header { - return http.Header{ - "accept": {"*/*"}, - "accept-language": {"en-US,en;q=0.9"}, - "cache-control": {"no-cache"}, - "cookie": {cookie}, - "pragma": {"no-cache"}, - "priority": {"u=1, i"}, - "referer": {appBase + "/"}, - "sec-ch-ua": {secChUA}, - "sec-ch-ua-mobile": {"?0"}, - "sec-ch-ua-platform": {`"Windows"`}, - "sec-fetch-dest": {"empty"}, - "sec-fetch-mode": {"cors"}, - "sec-fetch-site": {"same-origin"}, - "user-agent": {userAgent}, - http.HeaderOrderKey: { - "accept", "accept-language", "cache-control", "cookie", "pragma", - "priority", "referer", "sec-ch-ua", "sec-ch-ua-mobile", - "sec-ch-ua-platform", "sec-fetch-dest", "sec-fetch-mode", - "sec-fetch-site", "user-agent", - }, - } -} - -// session returns the cookie's access token, optionally forcing a fresh mint -// (dropping the cache) — used when the upstream rejected the current bearer. -func (c *Client) session(ctx context.Context, cookie string, force bool) (*Session, error) { - if force { - c.mu.Lock() - delete(c.sessions, strings.TrimSpace(cookie)) - c.mu.Unlock() - } - return c.GetSession(ctx, cookie) -} - -// ProbeSession force-mints a session from the cookie, bypassing the cached -// bearer. Callers use it to double-check an auth failure before killing an -// account: a rejected bearer (rotation race / expired token) still yields a -// working cookie here, only a genuinely dead cookie returns ErrAuth. -func (c *Client) ProbeSession(ctx context.Context, cookie string) (*Session, error) { - return c.session(ctx, cookie, true) -} - -// callGraphQL runs one GraphQL call for an account cookie. The bearer only lives -// ~1h, so a rejected token (401/403 or a JWTExpired GraphQL error) is re-minted -// from the cookie and the call retried once. Only a cookie that itself stops -// authenticating yields ErrAuth — an upstream bearer rejection stays temporary so -// the account is never killed for it. -func (c *Client) callGraphQL(ctx context.Context, cookie string, payload []byte, useProxy bool, label string) ([]byte, error) { - var lastStatus int - var lastBody []byte - for attempt := 0; attempt < 2; attempt++ { - sess, err := c.session(ctx, cookie, attempt > 0) - if err != nil { - return nil, err - } - body, status, err := c.graphqlP(ctx, sess.AccessToken, payload, useProxy) - if err != nil { - return nil, fmt.Errorf("%w: %s: %s", ErrTemporaryUpstream, label, err.Error()) - } - lastStatus, lastBody = status, body - stale := status == 401 || status == 403 - var gqlErr error - if !stale && status == 200 { - gqlErr = graphqlError(body) - stale = errors.Is(gqlErr, ErrAuth) - } - if stale { - if attempt == 0 { - continue - } - break - } - if status != 200 { - return nil, fmt.Errorf("%w: %s http %d: %s", ErrTemporaryUpstream, label, status, clip(body, 160)) - } - if gqlErr != nil { - return nil, gqlErr - } - return body, nil - } - return nil, fmt.Errorf("%w: %s rejected a freshly minted bearer (http %d): %s", - ErrTemporaryUpstream, label, lastStatus, clip(lastBody, 160)) -} - -const qGetTokens = `query GetUserTokensFromSub($sub: String) { - user_details(where: {cognitoId: {_eq: $sub}}) { - id - plan - subscriptionTokens - paidTokens - rolloverTokens - tokenRenewalDate - __typename - } -}` - -// FetchCreditsBalance derives a JWT from the cookie then reads the account's image -// token balance. Returns a normalized map mirroring the other providers so the -// TokenService quota plumbing is uniform. remaining = subscription+paid+rollover -// (the spendable image tokens); available_until carries the daily renewal time so -// the maintenance sweep can auto-recover a 限额 account. -func (c *Client) FetchCreditsBalance(ctx context.Context, cookie string) (map[string]any, error) { - sess, err := c.GetSession(ctx, cookie) - if err != nil { - if errors.Is(err, ErrAuth) { - return nil, ErrAuth - } - return unknownBalance(err.Error()), nil - } - if sess.CognitoSub == "" { - return unknownBalance("no cognitoSub"), nil - } - - payload, _ := json.Marshal(map[string]any{ - "operationName": "GetUserTokensFromSub", - "variables": map[string]any{"sub": sess.CognitoSub}, - "query": qGetTokens, - }) - body, err := c.callGraphQL(ctx, cookie, payload, false, "credits") - if err != nil { - if errors.Is(err, ErrAuth) { - return nil, ErrAuth - } - return unknownBalance(err.Error()), nil - } - var result struct { - Data struct { - UserDetails []struct { - Plan string `json:"plan"` - SubscriptionTokens int `json:"subscriptionTokens"` - PaidTokens int `json:"paidTokens"` - RolloverTokens int `json:"rolloverTokens"` - TokenRenewalDate string `json:"tokenRenewalDate"` - } `json:"user_details"` - } `json:"data"` - } - if err := json.Unmarshal(body, &result); err != nil { - return unknownBalance("non-json"), nil - } - if len(result.Data.UserDetails) == 0 { - return unknownBalance("no user_details"), nil - } - ud := result.Data.UserDetails[0] - remaining := ud.SubscriptionTokens + ud.PaidTokens + ud.RolloverTokens - return map[string]any{ - "remaining": remaining, - "used": nil, - "total": nil, - "unknown": false, - "error": nil, - "plan": ud.Plan, - "available_until": strings.TrimSpace(ud.TokenRenewalDate), - "email": emptyStringNil(sess.Email), - "display_name": emptyStringNil(sess.Name), - "user_id": emptyStringNil(sess.UserID), - }, nil -} - -// graphqlP runs a GraphQL call; callers pick the egress: only the generate submit -// uses the proxy; reference-image upload and polling run direct (local IP). -func (c *Client) graphqlP(ctx context.Context, accessToken string, payload []byte, useProxy bool) ([]byte, int, error) { - client, err := c.newTLSClientP(useProxy) - if err != nil { - return nil, 0, err - } - req, err := http.NewRequest(http.MethodPost, graphqlURL, bytes.NewReader(payload)) - if err != nil { - return nil, 0, err - } - req = req.WithContext(ctx) - req.Header = http.Header{ - "content-type": {"application/json"}, - "accept": {"*/*"}, - "accept-language": {"en-US,en;q=0.9"}, - "origin": {appBase}, - "priority": {"u=1, i"}, - "referer": {appBase + "/"}, - "sec-ch-ua": {secChUA}, - "sec-ch-ua-mobile": {"?0"}, - "sec-ch-ua-platform": {`"Windows"`}, - "user-agent": {userAgent}, - "authorization": {"Bearer " + accessToken}, - "x-leo-schema-version": {schemaVersion}, - "sec-fetch-dest": {"empty"}, - "sec-fetch-mode": {"cors"}, - "sec-fetch-site": {"same-site"}, - http.HeaderOrderKey: { - "content-type", "accept", "accept-language", "origin", "priority", - "referer", "sec-ch-ua", "sec-ch-ua-mobile", "sec-ch-ua-platform", - "user-agent", "authorization", "x-leo-schema-version", - "sec-fetch-dest", "sec-fetch-mode", "sec-fetch-site", - }, - } - resp, err := client.Do(req) - if err != nil { - return nil, 0, err - } - defer resp.Body.Close() - body, err := io.ReadAll(resp.Body) - if err != nil { - return nil, resp.StatusCode, err - } - return body, resp.StatusCode, nil -} - -func unknownBalance(reason string) map[string]any { - return map[string]any{ - "remaining": nil, - "used": nil, - "total": nil, - "unknown": true, - "error": reason, - } -} - -func (c *Client) newTLSClient() (tlsclient.HttpClient, error) { return c.newTLSClientP(true) } - -// newDirectTLSClient egresses on the local IP (never the proxy). Used for -// reference-image upload, polling and result download. -func (c *Client) newDirectTLSClient() (tlsclient.HttpClient, error) { return c.newTLSClientP(false) } - -func (c *Client) newTLSClientP(useProxy bool) (tlsclient.HttpClient, error) { - // Match the fingerprint proven to work against Leonardo's Cloudflare edge: - // Chrome_120, fixed extension order. A randomized JA3 (Chrome_133 + - // WithRandomTLSExtensionOrder) gets flagged and 429'd at get-session. - options := []tlsclient.HttpClientOption{ - tlsclient.WithTimeoutSeconds(60), - tlsclient.WithClientProfile(profiles.Chrome_120), - } - if useProxy && c.proxy != "" { - options = append(options, tlsclient.WithProxyUrl(c.proxy)) - } - return tlsclient.NewHttpClient(tlsclient.NewNoopLogger(), options...) -} - -// downloadImage fetches a generated image (cdn.leonardo.ai) and returns the bytes. -func (c *Client) downloadImage(ctx context.Context, imageURL string) ([]byte, error) { - if _, err := url.Parse(imageURL); err != nil { - return nil, err - } - client, err := c.newDirectTLSClient() - if err != nil { - return nil, err - } - req, err := http.NewRequest(http.MethodGet, imageURL, nil) - if err != nil { - return nil, err - } - req = req.WithContext(ctx) - req.Header = http.Header{ - "accept": {"image/avif,image/webp,image/png,image/*,*/*;q=0.8"}, - "user-agent": {userAgent}, - "referer": {appBase + "/"}, - } - resp, err := client.Do(req) - if err != nil { - return nil, err - } - defer resp.Body.Close() - body, err := io.ReadAll(resp.Body) - if err != nil { - return nil, err - } - if resp.StatusCode != 200 { - return nil, fmt.Errorf("%w: image download http %d", ErrTemporaryUpstream, resp.StatusCode) - } - return body, nil -} - -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 intValue(v any) int { - switch x := v.(type) { - case int: - return x - case int64: - return int(x) - case float64: - return int(x) - case json.Number: - n, _ := x.Int64() - return int(n) - case string: - n, _ := strconv.Atoi(strings.TrimSpace(x)) - return n - default: - return 0 - } -} - -func emptyStringNil(v string) any { - v = strings.TrimSpace(v) - if v == "" { - return nil - } - return v -} - -func clip(b []byte, n int) string { - s := strings.TrimSpace(string(b)) - if len(s) > n { - return s[:n] - } - return s -} diff --git a/backend/internal/provider/leonardo/image.go b/backend/internal/provider/leonardo/image.go deleted file mode 100644 index a4bb671..0000000 --- a/backend/internal/provider/leonardo/image.go +++ /dev/null @@ -1,401 +0,0 @@ -package leonardo - -import ( - "bytes" - "context" - "encoding/json" - "errors" - "fmt" - "mime/multipart" - "strings" - "time" - - http "github.com/bogdanfinn/fhttp" -) - -// defaultStyleID is the "Dynamic" style applied when the caller doesn't specify -// one — Leonardo's Generate mutation expects a style_ids entry. -const defaultStyleID = "111dc692-d470-4eec-b791-3475abac4c46" - -const mGenerate = `mutation Generate($request: CreateGenerationRequest!) { - generate(request: $request) { - apiCreditCost - generationId - __typename - } -}` - -// qGenerationImages polls one generation's status AND its produced images in a -// single round-trip (where: id _in [genId]). -const qGenerationImages = `query GenerationImages($where: generations_bool_exp = {}) { - generations(where: $where) { - id - status - generated_images { - id - url - __typename - } - __typename - } -}` - -const mUploadImage = `mutation UploadImage($uploadImageInput: UploadImageInput!) { - uploadImage(arg1: $uploadImageInput) { - uploadId - url - fields - __typename - } -}` - -const mUploadInitImage = `mutation UploadInitImage($arg1: InitImageUploadInput!) { - uploadInitImage(arg1: $arg1) { - id - url - fields - __typename - } -}` - -// initImageExtension narrows a sniffed extension to what uploadInitImage accepts -// (png / jpg / jpeg / webp); anything else is sent as png. -func initImageExtension(extension string) string { - switch strings.TrimPrefix(strings.ToLower(strings.TrimSpace(extension)), ".") { - case "jpg": - return "jpg" - case "jpeg": - return "jpeg" - case "webp": - return "webp" - default: - return "png" - } -} - -// uploadInitImage uploads a reference image and returns the init image id to put -// in a Generate request's image_reference guidance. It has to go through -// uploadInitImage (permanent init-image bucket): the uploadImage mutation only -// hands out temporary-bucket ids, which the generation service can't resolve. -func (c *Client) uploadInitImage(ctx context.Context, cookie, extension string, img []byte) (string, error) { - extension = initImageExtension(extension) - payload, _ := json.Marshal(map[string]any{ - "operationName": "UploadInitImage", - "query": mUploadInitImage, - "variables": map[string]any{"arg1": map[string]any{"extension": extension}}, - }) - body, err := c.callGraphQL(ctx, cookie, payload, false, "upload-init-image") - if err != nil { - return "", err - } - var ur struct { - Data struct { - UploadInitImage struct { - ID string `json:"id"` - URL string `json:"url"` - Fields string `json:"fields"` - } `json:"uploadInitImage"` - } `json:"data"` - } - if err := json.Unmarshal(body, &ur); err != nil { - return "", fmt.Errorf("%w: upload-init-image non-json", ErrTemporaryUpstream) - } - up := ur.Data.UploadInitImage - if up.ID == "" || up.URL == "" { - return "", fmt.Errorf("%w: no upload url", ErrTemporaryUpstream) - } - if err := c.putPresigned(ctx, up.URL, up.Fields, "asset."+extension, img); err != nil { - return "", err - } - return up.ID, nil -} - -// putPresigned performs the presigned S3 POST: all policy fields first, the file -// part LAST. -func (c *Client) putPresigned(ctx context.Context, url, fieldsJSON, filename string, asset []byte) error { - var fields map[string]string - if err := json.Unmarshal([]byte(fieldsJSON), &fields); err != nil { - return fmt.Errorf("%w: bad upload fields", ErrTemporaryUpstream) - } - var buf bytes.Buffer - w := multipart.NewWriter(&buf) - for k, v := range fields { - _ = w.WriteField(k, v) - } - fw, err := w.CreateFormFile("file", filename) - if err != nil { - return err - } - if _, err := fw.Write(asset); err != nil { - return err - } - _ = w.Close() - - client, err := c.newDirectTLSClient() - if err != nil { - return err - } - req, err := http.NewRequest(http.MethodPost, url, &buf) - if err != nil { - return err - } - req = req.WithContext(ctx) - req.Header = http.Header{ - "content-type": {w.FormDataContentType()}, - "user-agent": {userAgent}, - "origin": {appBase}, - "referer": {appBase + "/"}, - } - resp, err := client.Do(req) - if err != nil { - return fmt.Errorf("%w: s3 upload: %s", ErrTemporaryUpstream, err.Error()) - } - defer resp.Body.Close() - if resp.StatusCode != 204 && resp.StatusCode != 200 && resp.StatusCode != 201 { - return fmt.Errorf("%w: s3 upload http %d", ErrTemporaryUpstream, resp.StatusCode) - } - return nil -} - -// uploadAsset uploads one reference asset (extension png / mp3 / mp4 …) through -// the same UploadImage presigned-S3 flow images use, and returns its upload id. -func (c *Client) uploadAsset(ctx context.Context, cookie, extension string, asset []byte) (string, error) { - extension = strings.TrimPrefix(strings.ToLower(strings.TrimSpace(extension)), ".") - if extension == "" { - extension = "png" - } - payload, _ := json.Marshal(map[string]any{ - "operationName": "UploadImage", - "query": mUploadImage, - // originalFilename is mandatory for audio uploads and harmless otherwise. - "variables": map[string]any{"uploadImageInput": map[string]any{ - "uploadType": "INIT", - "extension": extension, - "originalFilename": "asset." + extension, - }}, - }) - body, err := c.callGraphQL(ctx, cookie, payload, false, "upload-init") - if err != nil { - return "", err - } - var ur struct { - Data struct { - UploadImage struct { - UploadID string `json:"uploadId"` - URL string `json:"url"` - Fields string `json:"fields"` - } `json:"uploadImage"` - } `json:"data"` - } - if err := json.Unmarshal(body, &ur); err != nil { - return "", fmt.Errorf("%w: upload-init non-json", ErrTemporaryUpstream) - } - up := ur.Data.UploadImage - if up.UploadID == "" || up.URL == "" { - return "", fmt.Errorf("%w: no upload url", ErrTemporaryUpstream) - } - if err := c.putPresigned(ctx, up.URL, up.Fields, "asset."+extension, asset); err != nil { - return "", err - } - return up.UploadID, nil -} - -// GenerateImage runs the full Leonardo image pipeline against one account cookie: -// mint a JWT, (for image-to-image) upload each reference image, submit the -// Generate mutation, poll until COMPLETE, then download the first produced image. -// Returns the image bytes, an info map, and a classified error. -func (c *Client) GenerateImage(ctx context.Context, cookie, model, prompt string, width, height int, styleIDs []string, refImages [][]byte, downloadResult bool) ([]byte, map[string]any, error) { - sess, err := c.GetSession(ctx, cookie) - if err != nil { - return nil, nil, err - } - if len(styleIDs) == 0 { - styleIDs = []string{defaultStyleID} - } - if strings.TrimSpace(model) == "" { - model = "seedream-4.5" - } - - // Image-to-image: upload each reference and collect its guidance entry. - var imageRefs []map[string]any - for _, img := range refImages { - if len(img) == 0 { - continue - } - uploadID, upErr := c.uploadInitImage(ctx, cookie, assetExtension(img, "png"), img) - if upErr != nil { - return nil, nil, upErr - } - imageRefs = append(imageRefs, map[string]any{ - "image": map[string]any{"id": uploadID, "type": "UPLOADED"}, - "strength": "MID", - }) - } - - promptEnhance := "AUTO" - parameters := map[string]any{ - "height": height, - "width": width, - "prompt_enhance": promptEnhance, - "quantity": 1, - "style_ids": styleIDs, - "prompt": prompt, - } - if len(imageRefs) > 0 { - // Preserve the reference when image-guided (matches the web app). - parameters["prompt_enhance"] = "OFF" - parameters["guidances"] = map[string]any{"image_reference": imageRefs} - } - - // 1. submit - genReq := map[string]any{ - "operationName": "Generate", - "query": mGenerate, - "variables": map[string]any{ - "request": map[string]any{ - "model": model, - "public": true, - "parameters": parameters, - }, - }, - } - payload, _ := json.Marshal(genReq) - body, err := c.callGraphQL(ctx, cookie, payload, true, "generate") - if err != nil { - return nil, nil, err - } - var genResp struct { - Data struct { - Generate struct { - GenerationID string `json:"generationId"` - } `json:"generate"` - } `json:"data"` - } - if err := json.Unmarshal(body, &genResp); err != nil { - return nil, nil, fmt.Errorf("%w: generate non-json", ErrTemporaryUpstream) - } - genID := strings.TrimSpace(genResp.Data.Generate.GenerationID) - if genID == "" { - return nil, nil, fmt.Errorf("%w: no generationId: %s", ErrTemporaryUpstream, clip(body, 200)) - } - - // 2. poll until COMPLETE, then read the image url. - imageURL, err := c.pollImage(ctx, cookie, genID) - if err != nil { - return nil, nil, err - } - - info := map[string]any{ - "generation_id": genID, - "image_url": imageURL, - "user_id": sess.UserID, - } - if !downloadResult { - return nil, info, nil - } - // 3. download bytes - data, err := c.downloadImage(ctx, imageURL) - if err != nil { - return nil, nil, err - } - return data, info, nil -} - -// pollImage polls one generation until it reports COMPLETE (returning the first -// image url) or FAILED (error). Honors ctx cancellation / deadline. -func (c *Client) pollImage(ctx context.Context, cookie, genID string) (string, error) { - payload, _ := json.Marshal(map[string]any{ - "operationName": "GenerationImages", - "query": qGenerationImages, - "variables": map[string]any{ - "where": map[string]any{"id": map[string]any{"_in": []string{genID}}}, - }, - }) - - ticker := time.NewTicker(3 * time.Second) - defer ticker.Stop() - // Poll for the full generation budget (caller's genCtx), leaving headroom for - // the download, instead of a shorter hardcoded cap that killed slow jobs early. - // ctx already bounds the wait, so a stuck job still can't hang indefinitely. - deadline := time.Now().Add(5 * time.Minute) - if dl, ok := ctx.Deadline(); ok { - deadline = dl.Add(-60 * time.Second) - } - - for { - body, err := c.callGraphQL(ctx, cookie, payload, false, "poll") - if errors.Is(err, ErrAuth) { - return "", err - } - // 其它错误（含上游临时抖动）不中断轮询，等 deadline 再判超时。 - if err == nil { - var pr struct { - Data struct { - Generations []struct { - Status string `json:"status"` - GeneratedImages []struct { - URL string `json:"url"` - } `json:"generated_images"` - } `json:"generations"` - } `json:"data"` - } - if err := json.Unmarshal(body, &pr); err == nil && len(pr.Data.Generations) > 0 { - g := pr.Data.Generations[0] - switch strings.ToUpper(g.Status) { - case "COMPLETE": - for _, img := range g.GeneratedImages { - if u := strings.TrimSpace(img.URL); u != "" { - return u, nil - } - } - return "", fmt.Errorf("%w: complete but no image url", ErrTemporaryUpstream) - case "FAILED": - return "", fmt.Errorf("%w: generation failed", ErrTemporaryUpstream) - } - } - } - - if time.Now().After(deadline) { - return "", fmt.Errorf("%w: generation timed out", ErrTemporaryUpstream) - } - select { - case <-ctx.Done(): - return "", ctx.Err() - case <-ticker.C: - } - } -} - -// graphqlError inspects a GraphQL response body for an "errors" array and maps the -// first message to a classified sentinel (auth / quota / temporary). Returns nil -// when there are no errors. -func graphqlError(body []byte) error { - var env struct { - Errors []struct { - Message string `json:"message"` - Extensions struct { - Code string `json:"code"` - Details struct { - Message string `json:"message"` - } `json:"details"` - } `json:"extensions"` - } `json:"errors"` - } - if err := json.Unmarshal(body, &env); err != nil || len(env.Errors) == 0 { - return nil - } - msg := strings.TrimSpace(env.Errors[0].Message) - // The generic "An error occurred." hides the real reason in extensions. - if detail := strings.TrimSpace(env.Errors[0].Extensions.Details.Message); detail != "" && detail != msg { - msg = msg + " (" + detail + ")" - } - low := strings.ToLower(msg) - switch { - case strings.Contains(low, "unauthor") || strings.Contains(low, "jwt") || strings.Contains(low, "token is") || strings.Contains(low, "forbidden"): - return ErrAuth - case strings.Contains(low, "token") || strings.Contains(low, "credit") || strings.Contains(low, "quota") || strings.Contains(low, "insufficient") || strings.Contains(low, "not enough"): - return ErrQuotaExhausted - default: - return fmt.Errorf("leonardo: %s", clip([]byte(msg), 200)) - } -} diff --git a/backend/internal/provider/leonardo/video.go b/backend/internal/provider/leonardo/video.go deleted file mode 100644 index 7b563f4..0000000 --- a/backend/internal/provider/leonardo/video.go +++ /dev/null @@ -1,412 +0,0 @@ -package leonardo - -import ( - "context" - "encoding/binary" - "encoding/json" - "errors" - "fmt" - "strings" - "time" -) - -// qGenerationVideos polls one video generation's status AND its produced clip in -// a single round-trip (motionMP4URL carries the mp4). -const qGenerationVideos = `query GenerationVideos($where: generations_bool_exp = {}) { - generations(where: $where) { - id - status - generated_images { - id - motionMP4URL - __typename - } - __typename - } -}` - -// VideoAssets are the decoded reference assets a video request can carry: -// image references (strength MID), one audio track and video references. The -// caller enforces the per-type caps; durations are derived from the bytes here -// because Leonardo requires them for audio/video guidances. -type VideoAssets struct { - Images [][]byte - Audios [][]byte - Videos [][]byte -} - -// GenerateVideo runs the Leonardo video pipeline (seedance-2.0 / -fast, hailuo) against -// one account cookie: upload every reference asset, submit the Generate mutation -// as a PRIVATE generation, poll until COMPLETE, then optionally download the mp4. -func (c *Client) GenerateVideo(ctx context.Context, cookie, model, prompt string, width, height, durationSeconds int, refs VideoAssets, downloadResult bool) ([]byte, map[string]any, error) { - sess, err := c.GetSession(ctx, cookie) - if err != nil { - return nil, nil, err - } - if strings.TrimSpace(model) == "" { - model = "seedance-2.0-fast" - } - - guidances := map[string]any{} - var imageRefs []map[string]any - for _, img := range refs.Images { - if len(img) == 0 { - continue - } - uploadID, upErr := c.uploadInitImage(ctx, cookie, assetExtension(img, "png"), img) - if upErr != nil { - return nil, nil, upErr - } - imageRefs = append(imageRefs, map[string]any{ - "image": map[string]any{"id": uploadID, "type": "UPLOADED"}, - "strength": "MID", - }) - } - if len(imageRefs) > 0 { - guidances["image_reference"] = imageRefs - } - var videoRefs []map[string]any - for _, vid := range refs.Videos { - if len(vid) == 0 { - continue - } - uploadID, upErr := c.uploadAsset(ctx, cookie, assetExtension(vid, "mp4"), vid) - if upErr != nil { - return nil, nil, upErr - } - video := map[string]any{"id": uploadID, "type": "UPLOADED"} - if secs := MediaDurationSeconds(vid); secs > 0 { - video["duration"] = secs - } - videoRefs = append(videoRefs, map[string]any{"video": video}) - } - if len(videoRefs) > 0 { - guidances["video_reference_base"] = videoRefs - } - // Leonardo rejects an audio reference that isn't paired with an image or - // video reference (audio_reference_only_compatible_with_image_or_video_reference). - var audioRefs []map[string]any - for _, aud := range refs.Audios { - if len(aud) == 0 { - continue - } - if len(imageRefs) == 0 && len(videoRefs) == 0 { - return nil, nil, errors.New("leonardo: audio reference requires an image or video reference") - } - uploadID, upErr := c.uploadAsset(ctx, cookie, assetExtension(aud, "mp3"), aud) - if upErr != nil { - return nil, nil, upErr - } - audio := map[string]any{"id": uploadID, "type": "UPLOADED"} - if secs := MediaDurationSeconds(aud); secs > 0 { - audio["duration"] = secs - } - audioRefs = append(audioRefs, map[string]any{"audio": audio}) - } - if len(audioRefs) > 0 { - guidances["audio_reference"] = audioRefs - } - - parameters := map[string]any{ - "height": height, - "width": width, - "duration": durationSeconds, - "motion_has_audio": true, - "quantity": 1, - "prompt": prompt, - "guidances": guidances, - } - // seedance 走随机种子;hailuo 的生成页 seedEnabled=false,请求里不带 seed。 - if !strings.HasPrefix(model, "hailuo") { - parameters["seed"] = -1 - } - genReq := map[string]any{ - "operationName": "Generate", - "query": mGenerate, - "variables": map[string]any{ - "request": map[string]any{ - "model": model, - // 私有生成：不进公开 feed。 - "public": false, - "parameters": parameters, - }, - }, - } - payload, _ := json.Marshal(genReq) - body, err := c.callGraphQL(ctx, cookie, payload, true, "generate-video") - if err != nil { - return nil, nil, err - } - var genResp struct { - Data struct { - Generate struct { - GenerationID string `json:"generationId"` - } `json:"generate"` - } `json:"data"` - } - if err := json.Unmarshal(body, &genResp); err != nil { - return nil, nil, fmt.Errorf("%w: generate non-json", ErrTemporaryUpstream) - } - genID := strings.TrimSpace(genResp.Data.Generate.GenerationID) - if genID == "" { - return nil, nil, fmt.Errorf("%w: no generationId: %s", ErrTemporaryUpstream, clip(body, 200)) - } - - videoURL, err := c.pollVideo(ctx, cookie, genID) - if err != nil { - return nil, nil, err - } - info := map[string]any{ - "generation_id": genID, - "video_url": videoURL, - "user_id": sess.UserID, - } - if !downloadResult { - return nil, info, nil - } - data, err := c.downloadImage(ctx, videoURL) - if err != nil { - return nil, nil, err - } - return data, info, nil -} - -// pollVideo polls one generation until COMPLETE (returning the mp4 url) or -// FAILED. Temporary upstream hiccups don't abort the wait — only ctx/deadline do. -func (c *Client) pollVideo(ctx context.Context, cookie, genID string) (string, error) { - payload, _ := json.Marshal(map[string]any{ - "operationName": "GenerationVideos", - "query": qGenerationVideos, - "variables": map[string]any{ - "where": map[string]any{"id": map[string]any{"_in": []string{genID}}}, - }, - }) - - ticker := time.NewTicker(5 * time.Second) - defer ticker.Stop() - deadline := time.Now().Add(10 * time.Minute) - if dl, ok := ctx.Deadline(); ok { - deadline = dl.Add(-60 * time.Second) - } - - for { - body, err := c.callGraphQL(ctx, cookie, payload, false, "poll-video") - if errors.Is(err, ErrAuth) { - return "", err - } - if err == nil { - var pr struct { - Data struct { - Generations []struct { - Status string `json:"status"` - GeneratedImages []struct { - MotionMP4URL string `json:"motionMP4URL"` - } `json:"generated_images"` - } `json:"generations"` - } `json:"data"` - } - if err := json.Unmarshal(body, &pr); err == nil && len(pr.Data.Generations) > 0 { - g := pr.Data.Generations[0] - switch strings.ToUpper(g.Status) { - case "COMPLETE": - for _, img := range g.GeneratedImages { - if u := strings.TrimSpace(img.MotionMP4URL); u != "" { - return u, nil - } - } - return "", fmt.Errorf("%w: complete but no video url", ErrTemporaryUpstream) - case "FAILED": - return "", fmt.Errorf("%w: generation failed", ErrTemporaryUpstream) - } - } - } - - if time.Now().After(deadline) { - return "", fmt.Errorf("%w: generation timed out", ErrTemporaryUpstream) - } - select { - case <-ctx.Done(): - return "", ctx.Err() - case <-ticker.C: - } - } -} - -// assetExtension sniffs the upload extension Leonardo expects for a reference -// asset; fallback is used when the bytes aren't recognized. -func assetExtension(data []byte, fallback string) string { - n := len(data) - switch { - case n >= 8 && string(data[1:4]) == "PNG": - return "png" - case n >= 3 && data[0] == 0xFF && data[1] == 0xD8 && data[2] == 0xFF: - return "jpg" - case n >= 12 && string(data[0:4]) == "RIFF" && string(data[8:12]) == "WEBP": - return "webp" - case n >= 12 && string(data[4:8]) == "ftyp": - if n >= 12 && strings.HasPrefix(string(data[8:12]), "qt") { - return "mov" - } - if n >= 12 && strings.HasPrefix(string(data[8:12]), "M4A") { - return "m4a" - } - return "mp4" - case n >= 4 && data[0] == 0x1A && data[1] == 0x45 && data[2] == 0xDF && data[3] == 0xA3: - return "webm" - case n >= 3 && string(data[0:3]) == "ID3": - return "mp3" - case n >= 2 && data[0] == 0xFF && (data[1]&0xE0) == 0xE0: - return "mp3" - case n >= 12 && string(data[0:4]) == "RIFF" && string(data[8:12]) == "WAVE": - return "wav" - } - return fallback -} - -// MediaDurationSeconds returns a media asset's duration in seconds (0 when it -// can't be determined). Leonardo's audio_reference / video_reference_base -// guidances carry the duration, and the caller validates video lengths with it, -// so it's derived from the uploaded bytes: mp4/mov via the mvhd box, wav via the -// fmt byte rate, mp3 from the first frame's bitrate. -func MediaDurationSeconds(data []byte) float64 { - if secs := mp4Duration(data); secs > 0 { - return secs - } - if secs := wavDuration(data); secs > 0 { - return secs - } - return mp3Duration(data) -} - -// mp4Duration walks the ISOBMFF box tree to moov/mvhd and reads timescale+duration. -func mp4Duration(data []byte) float64 { - if len(data) < 12 || string(data[4:8]) != "ftyp" { - return 0 - } - moov := findBox(data, "moov") - if moov == nil { - return 0 - } - mvhd := findBox(moov, "mvhd") - if len(mvhd) < 20 { - return 0 - } - version := mvhd[0] - if version == 1 { - if len(mvhd) < 32 { - return 0 - } - timescale := binary.BigEndian.Uint32(mvhd[20:24]) - duration := binary.BigEndian.Uint64(mvhd[24:32]) - if timescale == 0 { - return 0 - } - return float64(duration) / float64(timescale) - } - timescale := binary.BigEndian.Uint32(mvhd[12:16]) - duration := binary.BigEndian.Uint32(mvhd[16:20]) - if timescale == 0 { - return 0 - } - return float64(duration) / float64(timescale) -} - -// findBox returns the payload of the first box named typ among data's boxes -// (callers descend one level at a time by passing a parent's payload back in). -func findBox(data []byte, typ string) []byte { - for off := 0; off+8 <= len(data); { - size := int(binary.BigEndian.Uint32(data[off : off+4])) - name := string(data[off+4 : off+8]) - header := 8 - if size == 1 { // 64-bit size - if off+16 > len(data) { - return nil - } - size = int(binary.BigEndian.Uint64(data[off+8 : off+16])) - header = 16 - } else if size == 0 { - size = len(data) - off - } - if size < header || off+size > len(data) { - return nil - } - if name == typ { - return data[off+header : off+size] - } - off += size - } - return nil -} - -// wavDuration reads the RIFF fmt chunk's byte rate and the data chunk size. -func wavDuration(data []byte) float64 { - if len(data) < 44 || string(data[0:4]) != "RIFF" || string(data[8:12]) != "WAVE" { - return 0 - } - byteRate := 0 - dataSize := 0 - for off := 12; off+8 <= len(data); { - name := string(data[off : off+4]) - size := int(binary.LittleEndian.Uint32(data[off+4 : off+8])) - body := off + 8 - if size < 0 || body > len(data) { - break - } - switch name { - case "fmt ": - if body+16 <= len(data) { - byteRate = int(binary.LittleEndian.Uint32(data[body+8 : body+12])) - } - case "data": - dataSize = size - if body+size > len(data) { - dataSize = len(data) - body - } - } - off = body + size + size%2 - } - if byteRate <= 0 || dataSize <= 0 { - return 0 - } - return float64(dataSize) / float64(byteRate) -} - -// mp3Bitrates are the Layer III bitrate tables (kbps) indexed by the frame -// header's bitrate index: 1 = MPEG1, 2 = MPEG2/2.5. -var mp3Bitrates = map[int][]int{ - 1: {0, 32, 40, 48, 56, 64, 80, 96, 112, 128, 160, 192, 224, 256, 320, 0}, - 2: {0, 8, 16, 24, 32, 40, 48, 56, 64, 80, 96, 112, 128, 144, 160, 0}, -} - -// mp3Duration estimates the duration from the first frame header's bitrate -// (constant-bitrate assumption — good enough for the guidance duration field). -func mp3Duration(data []byte) float64 { - off := 0 - if len(data) >= 10 && string(data[0:3]) == "ID3" { - // syncsafe int: 7 bits per byte - tagSize := int(data[6]&0x7F)<<21 | int(data[7]&0x7F)<<14 | int(data[8]&0x7F)<<7 | int(data[9]&0x7F) - off = 10 + tagSize - } - for ; off+4 <= len(data); off++ { - if data[off] != 0xFF || data[off+1]&0xE0 != 0xE0 { - continue - } - versionBits := (data[off+1] >> 3) & 0x03 - layerBits := (data[off+1] >> 1) & 0x03 - if layerBits != 0x01 { // Layer III only - continue - } - table := 1 - if versionBits != 0x03 { // MPEG2 / 2.5 - table = 2 - } - idx := int((data[off+2] >> 4) & 0x0F) - rates := mp3Bitrates[table] - if idx <= 0 || idx >= len(rates) || rates[idx] == 0 { - continue - } - kbps := rates[idx] - return float64(len(data)-off) * 8 / float64(kbps*1000) - } - return 0 -} diff --git a/backend/internal/provider/runway/client.go b/backend/internal/provider/runway/client.go deleted file mode 100644 index df23a7f..0000000 --- a/backend/internal/provider/runway/client.go +++ /dev/null @@ -1,373 +0,0 @@ -// Package runway implements the Runway (runwayml.com) provider client. For now -// it only covers account management — JWT detection, workspace/team id -// extraction and credit-balance probing — mirroring the curl_cffi reference in -// query_credits.py with tls-client so the JA3/JA4 fingerprint matches Chrome. -package runway - -import ( - "context" - crand "crypto/rand" - "encoding/base64" - "encoding/hex" - "encoding/json" - "errors" - "fmt" - "io" - "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://api.runwayml.com" - origin = "https://app.runwayml.com" - - // userAgent / secChUA mirror the real browser that produced a successful - // registration + generation in the reference HAR (Edge 150 on Windows). - userAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/150.0.0.0 Safari/537.36 Edg/150.0.0.0" - secChUA = `"Not;A=Brand";v="8", "Chromium";v="150", "Microsoft Edge";v="150"` - - // sourceApp / buildHash are the x-runway-source-application[-version] the real - // web app stamps on EVERY authed API call (319/336 requests in the reference - // HAR). buildHash is the web bundle's git sha; it also appears as the - // `sentry-release` segment of the baggage header. Runway anti-abuse keys on - // these being present + client-id (see clientIDFromToken): a write/upload that - // lacks them is judged a non-web client and the account's free credits are - // zeroed on the FIRST upload ("上传清零积分"). buildHash tracks Runway's web - // releases — refresh it from a current HAR if uploads start getting flagged. - sourceApp = "web" - buildHash = "3e96b2f0f85b8c0cafb7c0dcd7a6878305aaa0f0" -) - -// clientIDOverride, when non-empty, forces the x-runway-client-id (used to pin an -// account to the exact persistent device id its real browser session used). -var clientIDOverride string - -// clientIDUUIDNamespace is a fixed namespace so clientIDFromToken is stable -// across processes/restarts — the same account always derives the same -// x-runway-client-id, exactly like the browser's localStorage-persisted id. -var clientIDUUIDNamespace = uuid.MustParse("6ba7b811-9dad-11d1-80b4-00c04fd430c8") - -// clientIDFromToken derives the persistent per-account x-runway-client-id. The -// real SPA generates this UUID once and stores it in localStorage; every request -// from that browser reuses the SAME value (5b89c786-… appears on all 319 authed -// calls in the HAR). We reproduce that stability by deterministically deriving a -// v5 UUID from the account's JWT "id" claim, so one account == one client-id for -// its whole lifetime without needing a DB column. Falls back to the raw token if -// the id claim is missing. -func clientIDFromToken(token string) string { - if clientIDOverride != "" { - return clientIDOverride - } - seed := TeamIDFromToken(token) - if seed == "" { - seed = strings.TrimPrefix(token, "Bearer ") - } - return uuid.NewSHA1(clientIDUUIDNamespace, []byte("runway-client-id:"+seed)).String() -} - -// randHex returns n random bytes hex-encoded (2n chars), for sentry trace ids. -func randHex(n int) string { - b := make([]byte, n) - if _, err := crand.Read(b); err != nil { - return strings.Repeat("0", n*2) - } - return hex.EncodeToString(b) -} - -// browserHeaders builds the full Chrome/Edge header set the Runway web app sends -// on its JSON API calls. tls-client only spoofs the TLS/JA3 fingerprint — it does -// NOT inject a User-Agent or the client-hint / fetch-metadata headers. Without -// them a request looks like a headless bot and Runway's anti-abuse revokes the -// account's free credits within minutes ("秒死"). teamID is optional. -func browserHeaders(token, teamID string) http.Header { - traceID := randHex(16) - spanID := randHex(8) - h := http.Header{ - "accept": {"application/json"}, - "accept-language": {"zh-CN,zh;q=0.9,en;q=0.8,en-GB;q=0.7,en-US;q=0.6"}, - "authorization": {"Bearer " + strings.TrimPrefix(token, "Bearer ")}, - "baggage": {"sentry-environment=production,sentry-release=" + buildHash + ",sentry-public_key=8ea832c064ed4bbcb4b8952c02ba119a,sentry-trace_id=" + traceID}, - "content-type": {"application/json"}, - "origin": {origin}, - "priority": {"u=1, i"}, - "referer": {origin + "/"}, - "sec-ch-ua": {secChUA}, - "sec-ch-ua-mobile": {"?0"}, - "sec-ch-ua-platform": {`"Windows"`}, - "sec-fetch-dest": {"empty"}, - "sec-fetch-mode": {"cors"}, - "sec-fetch-site": {"same-site"}, - "sentry-trace": {traceID + "-" + spanID}, - "user-agent": {userAgent}, - // x-runway-* identify this as the real web app to Runway's anti-abuse. - // client-id must be persistent per account (see clientIDFromToken); - // omitting these is what zeroes an account's credits on first upload. - "x-runway-client-id": {clientIDFromToken(token)}, - "x-runway-source-application": {sourceApp}, - "x-runway-source-application-version": {buildHash}, - "x-runway-workspace": {teamID}, - http.HeaderOrderKey: { - "accept", "accept-language", "authorization", "baggage", - "content-type", "origin", "priority", "referer", - "sec-ch-ua", "sec-ch-ua-mobile", "sec-ch-ua-platform", - "sec-fetch-dest", "sec-fetch-mode", "sec-fetch-site", - "sentry-trace", "user-agent", - "x-runway-client-id", "x-runway-source-application", - "x-runway-source-application-version", "x-runway-workspace", - }, - } - if strings.TrimSpace(teamID) == "" { - h.Del("x-runway-workspace") - } - return h -} - -// browserAssetHeaders is the lighter header set the browser sends on cross-site -// asset transfers (presigned S3 upload / CloudFront artifact download): a -// User-Agent and client hints, but no Runway authorization. -func browserAssetHeaders() http.Header { - return http.Header{ - "accept": {"*/*"}, - "accept-language": {"zh-CN,zh;q=0.9,en;q=0.8,en-GB;q=0.7,en-US;q=0.6"}, - "origin": {origin}, - "referer": {origin + "/"}, - "sec-ch-ua": {secChUA}, - "sec-ch-ua-mobile": {"?0"}, - "sec-ch-ua-platform": {`"Windows"`}, - "sec-fetch-dest": {"empty"}, - "sec-fetch-mode": {"cors"}, - "sec-fetch-site": {"cross-site"}, - "user-agent": {userAgent}, - } -} - -var ( - ErrAuth = errors.New("runway auth failed") - ErrQuotaExhausted = errors.New("runway quota exhausted") - ErrTemporaryUpstream = errors.New("runway 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) -} - -// IsRunwayToken reports whether a JWT looks like a Runway access token: a -// top-level numeric "id" plus an "sso" claim, and crucially NO OpenAI -// (https://api.openai.com/*) claims — that's what disambiguates it from a -// ChatGPT token, which is otherwise also an opaque three-part JWT. -func IsRunwayToken(token string) bool { - claims := decodeJWTPayload(token) - if len(claims) == 0 { - return false - } - for k := range claims { - if strings.HasPrefix(k, "https://api.openai.com/") { - return false - } - } - _, hasSSO := claims["sso"] - return hasSSO && claims["id"] != nil -} - -// TeamIDFromToken returns the Runway workspace/team id, which equals the JWT -// "id" claim (query_credits.py / gen_video.py both derive teamId this way). -func TeamIDFromToken(token string) string { - claims := decodeJWTPayload(token) - switch v := claims["id"].(type) { - case float64: - return strconv.FormatInt(int64(v), 10) - case json.Number: - return v.String() - case string: - return strings.TrimSpace(v) - default: - return "" - } -} - -// ExtractAccountInfo decodes the free (no-network) JWT claims for the accounts -// view: email, team id and expiry. -func ExtractAccountInfo(token string) map[string]any { - claims := decodeJWTPayload(token) - return map[string]any{ - "email": emptyStringNil(strings.TrimSpace(stringValue(claims["email"]))), - "team_id": emptyStringNil(TeamIDFromToken(token)), - "expires_at": claims["exp"], - } -} - -// FetchCreditsBalance probes the account's plan credits via /v1/profile/features -// (query_credits.py). Returns a normalized map mirroring the Adobe client so the -// TokenService quota plumbing can treat all providers uniformly. A 401/403 maps -// to ErrAuth (token dead); any other failure is reported as unknown without -// killing the account. -func (c *Client) FetchCreditsBalance(ctx context.Context, token string) (map[string]any, error) { - token = strings.TrimSpace(token) - if token == "" { - return unknownBalance("empty token"), nil - } - teamID := TeamIDFromToken(token) - if teamID == "" { - return unknownBalance("no team id"), nil - } - - client, err := c.newDirectTLSClient() - if err != nil { - return nil, err - } - url := apiBase + "/v1/profile/features?asTeamId=" + teamID - req, err := http.NewRequest(http.MethodGet, url, nil) - if err != nil { - return nil, err - } - req = req.WithContext(ctx) - req.Header = browserHeaders(token, teamID) - - resp, err := client.Do(req) - if err != nil { - return unknownBalance("network: " + err.Error()), nil - } - defer resp.Body.Close() - - body, err := io.ReadAll(resp.Body) - 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 - } - if resp.StatusCode != 200 { - return unknownBalance(fmt.Sprintf("http %d: %s", resp.StatusCode, clip(body, 160))), nil - } - - var payload map[string]any - if err := json.Unmarshal(body, &payload); err != nil { - return unknownBalance("non-json"), nil - } - features, _ := payload["features"].(map[string]any) - permitted, _ := features["permitted"].(map[string]any) - used, _ := features["used"].(map[string]any) - total := intValue(permitted["numPlanCredits"]) - spent := intValue(used["numPlanCredits"]) - remaining := total - spent - if remaining < 0 { - remaining = 0 - } - return map[string]any{ - "remaining": remaining, - "used": spent, - "total": total, - "unknown": false, - "error": nil, - }, nil -} - -func unknownBalance(reason string) map[string]any { - return map[string]any{ - "remaining": nil, - "used": nil, - "total": nil, - "unknown": true, - "error": reason, - } -} - -func (c *Client) newTLSClient() (tlsclient.HttpClient, error) { return c.newTLSClientP(true) } - -// newDirectTLSClient egresses on the local IP (never the proxy). Used for -// reference-image upload, polling and result download. -func (c *Client) newDirectTLSClient() (tlsclient.HttpClient, error) { return c.newTLSClientP(false) } - -func (c *Client) newTLSClientP(useProxy bool) (tlsclient.HttpClient, error) { - options := []tlsclient.HttpClientOption{ - // Whole-request timeout, incl. reading the response body. 30s was too - // tight for downloading a rendered image/video artifact off a slow CDN, - // surfacing as "request canceled (Client.Timeout ... while reading body)". - // The caller's genCtx still bounds the overall generation. - tlsclient.WithTimeoutSeconds(300), - tlsclient.WithClientProfile(profiles.Chrome_133), - tlsclient.WithRandomTLSExtensionOrder(), - } - if useProxy && c.proxy != "" { - options = append(options, tlsclient.WithProxyUrl(c.proxy)) - } - return tlsclient.NewHttpClient(tlsclient.NewNoopLogger(), options...) -} - -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 intValue(v any) int { - switch x := v.(type) { - case int: - return x - case int64: - return int(x) - case float64: - return int(x) - case json.Number: - n, _ := x.Int64() - return int(n) - case string: - n, _ := strconv.Atoi(strings.TrimSpace(x)) - return n - default: - return 0 - } -} - -func emptyStringNil(v string) any { - v = strings.TrimSpace(v) - if v == "" { - return nil - } - return v -} - -func clip(b []byte, n int) string { - s := strings.TrimSpace(string(b)) - if len(s) > n { - return s[:n] - } - return s -} diff --git a/backend/internal/provider/runway/image.go b/backend/internal/provider/runway/image.go deleted file mode 100644 index 0407566..0000000 --- a/backend/internal/provider/runway/image.go +++ /dev/null @@ -1,181 +0,0 @@ -package runway - -import ( - "bytes" - "context" - "errors" - "fmt" - "image" - "strings" - "time" - - tlsclient "github.com/bogdanfinn/tls-client" -) - -// GenerateImage runs a Runway gemini 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. modelID selects the variant: "runway-nano-banana-2" -// → gemini_3_1_flash_image / gemini-3.1-flash-image-preview (with aspect_ratio), -// anything else → workflow_gemini_image / gemini-3-pro-image-preview. imageSize -// is the "1K"/"2K"/"4K" tier. teamID is the workspace id; if empty it's derived -// from the token. refs may be empty (pure text-to-image). -func (c *Client) GenerateImage(ctx context.Context, token, teamID, modelID, prompt, aspectRatio, imageSize string, refs [][]byte, downloadResult bool) ([]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(imageSize) == "" { - imageSize = "1K" - } - - // Only the task-create (generate submit) egresses via the proxy; reference - // upload, polling and download run on the local IP. - submitClient, err := c.newTLSClient() - if err != nil { - return nil, nil, err - } - directClient, err := c.newDirectTLSClient() - if err != nil { - return nil, nil, err - } - - var refImages []map[string]any - var refAssetIDs []string - 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, directClient, 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, - }) - refAssetIDs = append(refAssetIDs, assetID) - } - - // Browser order: create the real session FIRST, attach the references and its - // own asset group, THEN submit the task into it (see createSession). The task - // carries this session's assetGroupId, NOT the account-wide "Generations" - // group. - sessionID, err := c.createSession(ctx, submitClient, token, teamID) - if err != nil { - return nil, nil, err - } - for _, aid := range refAssetIDs { - c.attachReference(ctx, submitClient, token, teamID, sessionID, aid) - } - assetGroupID, _ := c.sessionAssetGroup(ctx, submitClient, token, teamID, sessionID) // best-effort - - taskID, err := c.createImageTask(ctx, submitClient, token, teamID, modelID, prompt, aspectRatio, imageSize, sessionID, assetGroupID, refImages) - if err != nil { - return nil, nil, err - } - artifactURL, err := c.pollTask(ctx, directClient, token, teamID, taskID) - if err != nil { - return nil, nil, err - } - meta := map[string]any{ - "provider": "runway", - "task_id": taskID, - "team_id": teamID, - "image_url": artifactURL, - } - // downloadResult=false (API-key url-only mode): return the upstream artifact - // URL without downloading the PNG. - if !downloadResult { - return nil, meta, nil - } - data, err := c.download(ctx, directClient, artifactURL) - if err != nil { - return nil, nil, err - } - 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 submits a gemini image task (workflow_gemini_image for Pro, -// gemini_3_1_flash_image for Nano Banana 2) INTO the pre-created session and -// returns its id. sessionID must be a real server-side session (see -// createSession) and assetGroupID must be that session's own group. -func (c *Client) createImageTask(ctx context.Context, client tlsclient.HttpClient, token, teamID, modelID, prompt, aspectRatio, imageSize, sessionID, assetGroupID string, refImages []map[string]any) (string, error) { - if strings.TrimSpace(aspectRatio) == "" { - aspectRatio = "16:9" - } - taskType := "workflow_gemini_image" - opts := map[string]any{ - "name": "Nano Banana Pro - " + prompt, - "text_prompt": prompt, - "aspect_ratio": aspectRatio, - "num_images": 1, - "image_size": imageSize, - "model": "gemini-3-pro-image-preview", - "exploreMode": false, - "creationSource": "tool-mode", - } - if modelID == "runway-nano-banana-2" { - taskType = "gemini_3_1_flash_image" - opts["name"] = "Nano Banana 2 - " + prompt - opts["model"] = "gemini-3.1-flash-image-preview" - } - if assetGroupID != "" { - opts["assetGroupId"] = assetGroupID - } - if len(refImages) > 0 { - opts["reference_images"] = refImages - } - // Pre-flight cost estimate, exactly as the web app does before every spend. - c.estimateCost(ctx, client, token, teamID, "gemini_image", opts) - - res, err := c.submitTask(ctx, client, token, teamID, map[string]any{ - "taskType": taskType, - "options": opts, - "asTeamId": jsonNumberOrString(teamID), - "sessionId": sessionID, - }) - 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) - } - // Post-submit: generation record + session play (session already exists). - c.recordGeneration(ctx, client, token, teamID, id, sessionID, prompt, opts) - return id, nil -} diff --git a/backend/internal/provider/runway/video.go b/backend/internal/provider/runway/video.go deleted file mode 100644 index c617c65..0000000 --- a/backend/internal/provider/runway/video.go +++ /dev/null @@ -1,516 +0,0 @@ -package runway - -import ( - "bytes" - "context" - "encoding/json" - "errors" - "fmt" - "image" - _ "image/gif" - _ "image/jpeg" - _ "image/png" - "io" - mrand "math/rand/v2" - "strings" - "time" - - http "github.com/bogdanfinn/fhttp" - tlsclient "github.com/bogdanfinn/tls-client" -) - -// ratioDimensions maps an aspect ratio to the Gen-4 Turbo native output size. -// These are the only dimensions gen4_turbo accepts; "2K" is a UI label over this -// native tier (see runway-video-gen-spec). Unknown ratios fall back to 16:9. -func ratioDimensions(aspectRatio string) (int, int) { - switch strings.TrimSpace(strings.ReplaceAll(aspectRatio, "x", ":")) { - case "16:9": - return 1280, 720 - case "9:16": - return 720, 1280 - case "1:1": - return 960, 960 - case "4:3": - return 1104, 832 - case "3:4": - return 832, 1104 - case "21:9": - return 1584, 672 - default: - return 1280, 720 - } -} - -// GenerateVideo runs the full i2v pipeline (gen_video.py): upload the first-frame -// image (preview + dataset), create a dataset, create a gen4_turbo task and poll -// it to completion, then download the rendered MP4. teamID is the workspace id -// (meta["team_id"]); if empty it's derived from the token. seconds must be 5 or -// 10; aspectRatio picks the native output size. -// GenerateVideo renders the clip and (when downloadResult) downloads the MP4. -// With downloadResult=false it returns nil bytes and the upstream artifact URL in -// meta["video_url"] — used by the async /v1/videos job, which proxies that URL on -// /content instead of persisting the file. -func (c *Client) GenerateVideo(ctx context.Context, token, teamID, prompt, aspectRatio string, seconds int, frame []byte, downloadResult bool) ([]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 len(frame) == 0 { - return nil, nil, errors.New("runway: first-frame image required") - } - cfg, _, err := image.DecodeConfig(bytes.NewReader(frame)) - if err != nil { - return nil, nil, errors.New("runway: failed to decode first-frame image") - } - - // Only the task-create (generate submit) egresses via the proxy; reference - // upload, dataset create, polling and download run on the local IP (matches - // the image pipeline). - submitClient, err := c.newTLSClient() - if err != nil { - return nil, nil, err - } - directClient, err := c.newDirectTLSClient() - if err != nil { - return nil, nil, err - } - - filename := "frame_" + time.Now().UTC().Format("20060102_150405") + ".png" - previewUploadID, _, err := c.uploadFile(ctx, directClient, token, teamID, filename, "DATASET_PREVIEW", frame) - if err != nil { - return nil, nil, err - } - datasetUploadID, _, err := c.uploadFile(ctx, directClient, token, teamID, filename, "DATASET", frame) - if err != nil { - return nil, nil, err - } - - assetID, imageURL, err := c.createDataset(ctx, directClient, token, teamID, filename, datasetUploadID, previewUploadID, cfg.Width, cfg.Height) - if err != nil { - return nil, nil, err - } - // Browser order: real session first, attach the first-frame asset + its own - // asset group, THEN submit the task into it (see createSession). - sessionID, err := c.createSession(ctx, submitClient, token, teamID) - if err != nil { - return nil, nil, err - } - c.attachReference(ctx, submitClient, token, teamID, sessionID, assetID) - assetGroupID, _ := c.sessionAssetGroup(ctx, submitClient, token, teamID, sessionID) // best-effort - - taskID, err := c.createTask(ctx, submitClient, token, teamID, prompt, imageURL, assetID, assetGroupID, sessionID, aspectRatio, seconds) - if err != nil { - return nil, nil, err - } - - artifactURL, err := c.pollTask(ctx, directClient, token, teamID, taskID) - if err != nil { - return nil, nil, err - } - meta := map[string]any{ - "provider": "runway", - "task_id": taskID, - "team_id": teamID, - "video_url": artifactURL, - } - if !downloadResult { - return nil, meta, nil - } - data, err := c.download(ctx, directClient, artifactURL) - if err != nil { - return nil, nil, err - } - return data, meta, nil -} - -// uploadFile mirrors gen_video.upload_file: register the upload, PUT the bytes to -// the returned S3 URL, then complete. Returns the upload id and final url. -func (c *Client) uploadFile(ctx context.Context, client tlsclient.HttpClient, token, teamID, filename, uploadType string, data []byte) (string, string, error) { - info, err := c.apiJSON(ctx, client, token, teamID, http.MethodPost, "/v1/uploads", map[string]any{ - "filename": filename, - "numberOfParts": 1, - "type": uploadType, - }) - if err != nil { - return "", "", err - } - uploadID := strings.TrimSpace(stringValue(info["id"])) - urls, _ := info["uploadUrls"].([]any) - if uploadID == "" || len(urls) == 0 { - return "", "", fmt.Errorf("%w: upload register missing fields", ErrTemporaryUpstream) - } - putURL := strings.TrimSpace(stringValue(urls[0])) - contentType := "application/octet-stream" - if hdrs, ok := info["uploadHeaders"].(map[string]any); ok { - if ct := strings.TrimSpace(stringValue(hdrs["Content-Type"])); ct != "" { - contentType = ct - } - } - - etag, err := c.putBytes(ctx, client, putURL, contentType, data) - if err != nil { - return "", "", err - } - - res, err := c.apiJSON(ctx, client, token, teamID, http.MethodPost, "/v1/uploads/"+uploadID+"/complete", map[string]any{ - "parts": []map[string]any{{"PartNumber": 1, "ETag": etag}}, - }) - if err != nil { - return "", "", err - } - return uploadID, strings.TrimSpace(stringValue(res["url"])), nil -} - -func (c *Client) createDataset(ctx context.Context, client tlsclient.HttpClient, token, teamID, filename, datasetUploadID, previewUploadID string, w, h int) (string, string, error) { - teamIDNum := jsonNumberOrString(teamID) - res, err := c.apiJSON(ctx, client, token, teamID, http.MethodPost, "/v1/datasets", map[string]any{ - "fileCount": 1, - "name": filename, - "uploadId": datasetUploadID, - "previewUploadIds": []string{previewUploadID}, - "metadata": map[string]any{"size": map[string]any{"width": w, "height": h}}, - "type": map[string]any{"name": "image", "type": "image", "isDirectory": false}, - "asTeamId": teamIDNum, - "privateInTeam": true, - }) - if err != nil { - return "", "", err - } - ds, _ := res["dataset"].(map[string]any) - id := strings.TrimSpace(stringValue(ds["id"])) - url := strings.TrimSpace(stringValue(ds["url"])) - if id == "" || url == "" { - return "", "", fmt.Errorf("%w: dataset missing fields", ErrTemporaryUpstream) - } - return id, url, nil -} - -func (c *Client) assetGroupID(ctx context.Context, client tlsclient.HttpClient, token, teamID string) (string, error) { - res, err := c.apiJSON(ctx, client, token, teamID, http.MethodGet, - "/v1/asset_groups/by_name?name=Generations&asTeamId="+teamID+"&privateInTeam=true", nil) - if err != nil { - return "", err - } - ag, _ := res["assetGroup"].(map[string]any) - return strings.TrimSpace(stringValue(ag["id"])), nil -} - -func (c *Client) createTask(ctx context.Context, client tlsclient.HttpClient, token, teamID, prompt, imageURL, assetID, assetGroupID, sessionID, aspectRatio string, seconds int) (string, error) { - w, h := ratioDimensions(aspectRatio) - opts := map[string]any{ - "route": "i2v", - "name": "Gen-4 Turbo - " + prompt, - "text_prompt": prompt, - "seconds": seconds, - "width": w, - "height": h, - "init_image": imageURL, - "imageAssetId": assetID, - "exploreMode": false, - "creationSource": "tool-mode", - "seed": mrand.IntN(999999999) + 1, - "watermark": true, - } - if assetGroupID != "" { - opts["assetGroupId"] = assetGroupID - } - res, err := c.submitTask(ctx, client, token, teamID, map[string]any{ - "taskType": "gen4_turbo", - "options": opts, - "asTeamId": jsonNumberOrString(teamID), - "sessionId": sessionID, - }) - if err != nil { - return "", err - } - task, _ := res["task"].(map[string]any) - id := strings.TrimSpace(stringValue(task["id"])) - if id == "" { - return "", fmt.Errorf("%w: task missing id", ErrTemporaryUpstream) - } - // Post-submit: generation record + session play (session already exists). - c.recordGeneration(ctx, client, token, teamID, id, sessionID, prompt, opts) - return id, nil -} - -func (c *Client) pollTask(ctx context.Context, client tlsclient.HttpClient, token, teamID, taskID string) (string, error) { - for { - if err := ctx.Err(); err != nil { - return "", err - } - res, err := c.apiJSON(ctx, client, token, teamID, http.MethodGet, "/v1/tasks/"+taskID+"?asTeamId="+teamID, nil) - if err != nil { - // A transient blip shouldn't kill a render that may still succeed. - if errors.Is(err, ErrTemporaryUpstream) { - if sleepCtx(ctx, 5*time.Second) != nil { - return "", ctx.Err() - } - continue - } - return "", err - } - task, _ := res["task"].(map[string]any) - status := strings.ToUpper(strings.TrimSpace(stringValue(task["status"]))) - switch status { - case "SUCCEEDED": - arts, _ := task["artifacts"].([]any) - for _, raw := range arts { - art, _ := raw.(map[string]any) - if url := strings.TrimSpace(stringValue(art["url"])); url != "" { - return url, nil - } - } - return "", errors.New("runway: task succeeded with no artifact url") - case "FAILED", "CANCELED": - reason := strings.TrimSpace(stringValue(task["error"])) - if isCreditError(reason) { - return "", fmt.Errorf("%w: %s", ErrQuotaExhausted, reason) - } - return "", fmt.Errorf("runway: task %s: %s", status, reason) - } - if sleepCtx(ctx, 5*time.Second) != nil { - return "", ctx.Err() - } - } -} - -// estimateCost fires the /v1/billing/estimate_feature_cost_credits pre-flight the -// web app ALWAYS does before a spend (52 times in the reference HAR, always with -// the exact taskOptions right before /v1/tasks). It returns no reservation token, -// so its purpose is purely to mark the spend as coming from the real UI flow; -// submitting a task with no preceding estimate for that spec is a scripted-abuse -// tell. Best-effort — the estimate itself costs nothing. -func (c *Client) estimateCost(ctx context.Context, client tlsclient.HttpClient, token, teamID, feature string, taskOptions map[string]any) { - _, _ = c.apiJSON(ctx, client, token, teamID, http.MethodPost, "/v1/billing/estimate_feature_cost_credits", map[string]any{ - "feature": feature, - "count": 1, - "asTeamId": jsonNumberOrString(teamID), - "taskOptions": taskOptions, - }) -} - -// createSession creates an EMPTY tool-mode session (POST /v1/sessions with an -// empty taskIds array) and returns its real server-side id. This is the crux of -// the whole flow: the browser creates the session BEFORE submitting the first -// task, then submits the task INTO that real session. Submitting /v1/tasks with a -// sessionId that was never created server-side is a forgery tell that makes -// Runway revoke the account's free credits (permitted numPlanCredits 500 -> 0) -// the instant the task lands ("提交生图就清零"). Note taskIds must be [] — passing -// a bogus id yields "Invalid task IDs". -func (c *Client) createSession(ctx context.Context, client tlsclient.HttpClient, token, teamID string) (string, error) { - res, err := c.apiJSON(ctx, client, token, teamID, http.MethodPost, "/v1/sessions", map[string]any{ - "asTeamId": jsonNumberOrString(teamID), - "taskIds": []string{}, - }) - if err != nil { - return "", err - } - sess, _ := res["session"].(map[string]any) - id := strings.TrimSpace(stringValue(sess["id"])) - if id == "" { - return "", fmt.Errorf("%w: session missing id", ErrTemporaryUpstream) - } - return id, nil -} - -// attachReference attaches an uploaded asset to the session (the browser does -// this before the task, so the task's reference_images belong to a real session). -func (c *Client) attachReference(ctx context.Context, client tlsclient.HttpClient, token, teamID, sessionID, assetID string) { - if strings.TrimSpace(assetID) == "" { - return - } - _, _ = c.apiJSON(ctx, client, token, teamID, http.MethodPost, - "/v1/sessions/"+sessionID+"/references", map[string]any{ - "assetId": assetID, - "asTeamId": jsonNumberOrString(teamID), - }) -} - -// sessionAssetGroup creates the session's OWN asset group (POST -// /v1/sessions/{id}/assetGroup) and returns its id — this, NOT the account-wide -// "Generations" group, is the assetGroupId the browser puts on the task. -func (c *Client) sessionAssetGroup(ctx context.Context, client tlsclient.HttpClient, token, teamID, sessionID string) (string, error) { - res, err := c.apiJSON(ctx, client, token, teamID, http.MethodPost, - "/v1/sessions/"+sessionID+"/assetGroup", map[string]any{"asTeamId": jsonNumberOrString(teamID)}) - if err != nil { - return "", err - } - ag, _ := res["assetGroup"].(map[string]any) - return strings.TrimSpace(stringValue(ag["id"])), nil -} - -// recordGeneration is the post-submit lifecycle the browser fires AFTER the task -// lands: the generation record (POST /v1/generations, recordingEnabled=true) and -// a session /play. The session itself already exists (createSession ran before -// the task), so this no longer creates it. Best-effort. -func (c *Client) recordGeneration(ctx context.Context, client tlsclient.HttpClient, token, teamID, taskID, sessionID, prompt string, taskOptions map[string]any) { - settings := map[string]any{"taskId": taskID, "recordingEnabled": true} - for k, v := range taskOptions { - if k == "exploreMode" { // not present on the browser's generation record - continue - } - settings[k] = v - } - // The generation record carries the model id WITHOUT the "-preview" suffix - // the task uses (gemini-3-pro-image-preview -> gemini-3-pro-image in the HAR). - if m, ok := settings["model"].(string); ok { - settings["model"] = strings.TrimSuffix(m, "-preview") - } - _, _ = c.apiJSON(ctx, client, token, teamID, http.MethodPost, "/v1/generations", map[string]any{ - "toolId": "generate", - "prompt": prompt, - "outputs": map[string]any{"outputUrls": []any{}}, - "settings": settings, - }) - _, _ = c.apiJSON(ctx, client, token, teamID, http.MethodPost, "/v1/sessions/"+sessionID+"/play", - map[string]any{"asTeamId": jsonNumberOrString(teamID), "taskId": taskID}) -} - -// submitTask POSTs a /v1/tasks create. Transient (network / 5xx) failures -// surface as ErrTemporaryUpstream so the pool fails over to the NEXT account -// (换号重试) instead of retrying this one. -func (c *Client) submitTask(ctx context.Context, client tlsclient.HttpClient, token, teamID string, body map[string]any) (map[string]any, error) { - return c.apiJSON(ctx, client, token, teamID, http.MethodPost, "/v1/tasks", body) -} - -// apiJSON performs an authed JSON request against the Runway API and returns the -// parsed body, mapping status codes to the shared provider error sentinels. -func (c *Client) apiJSON(ctx context.Context, client tlsclient.HttpClient, token, teamID, method, path string, body any) (map[string]any, error) { - var reader io.Reader - if body != nil { - raw, _ := json.Marshal(body) - reader = bytes.NewReader(raw) - } - req, err := http.NewRequest(method, apiBase+path, reader) - if err != nil { - return nil, err - } - req = req.WithContext(ctx) - req.Header = browserHeaders(token, teamID) - // The generate-submit endpoint is the ONE authed call the real web bundle - // sends WITHOUT x-runway-client-id / source-application[-version] — every - // other endpoint carries them, /v1/tasks carries only x-runway-workspace - // (verified across the whole reference HAR). Leaving them on /v1/tasks is a - // bot tell that gets the account's free credits revoked (permitted - // numPlanCredits 500 -> 0) the instant a task is submitted. Strip them here so - // the submit matches the browser exactly. - if method == http.MethodPost && path == "/v1/tasks" { - // NOTE: raw map delete, NOT req.Header.Del() — fhttp canonicalizes the key - // in Del() ("X-Runway-Client-Id") but our headers are stored lowercase (so - // they go on the h2 wire lowercase like Chrome), so Del() silently no-ops. - delete(req.Header, "x-runway-client-id") - delete(req.Header, "x-runway-source-application") - delete(req.Header, "x-runway-source-application-version") - } - resp, err := client.Do(req) - if err != nil { - return nil, fmt.Errorf("%w: %v", ErrTemporaryUpstream, err) - } - defer resp.Body.Close() - raw, err := io.ReadAll(resp.Body) - if err != nil { - return nil, err - } - 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)) - case resp.StatusCode >= 500: - return nil, fmt.Errorf("%w: %s %d %s", ErrTemporaryUpstream, path, resp.StatusCode, clip(raw, 200)) - case resp.StatusCode < 200 || resp.StatusCode >= 300: - if isCreditError(string(raw)) { - return nil, fmt.Errorf("%w: %s", ErrQuotaExhausted, clip(raw, 200)) - } - return nil, fmt.Errorf("runway: %s %d %s", path, resp.StatusCode, clip(raw, 200)) - } - 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 -} - -// putBytes uploads raw bytes to a presigned S3 URL (no auth) and returns the -// ETag, mirroring the plain requests.Session().put in gen_video.py. -func (c *Client) putBytes(ctx context.Context, client tlsclient.HttpClient, url, contentType string, data []byte) (string, error) { - req, err := http.NewRequest(http.MethodPut, url, bytes.NewReader(data)) - if err != nil { - return "", err - } - req = req.WithContext(ctx) - req.Header = browserAssetHeaders() - req.Header.Set("content-type", contentType) - resp, err := client.Do(req) - if err != nil { - return "", fmt.Errorf("%w: %v", ErrTemporaryUpstream, err) - } - defer resp.Body.Close() - _, _ = io.Copy(io.Discard, resp.Body) - if resp.StatusCode < 200 || resp.StatusCode >= 300 { - return "", fmt.Errorf("%w: s3 put %d", ErrTemporaryUpstream, resp.StatusCode) - } - return strings.Trim(resp.Header.Get("ETag"), `"`), nil -} - -func (c *Client) download(ctx context.Context, client tlsclient.HttpClient, url string) ([]byte, error) { - req, err := http.NewRequest(http.MethodGet, url, nil) - if err != nil { - return nil, err - } - req = req.WithContext(ctx) - req.Header = browserAssetHeaders() - req.Header.Set("accept", "image/avif,image/webp,image/apng,image/svg+xml,image/*,*/*;q=0.8") - req.Header.Set("sec-fetch-dest", "image") - req.Header.Set("sec-fetch-mode", "no-cors") - 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, errors.New("runway: empty artifact download") - } - return data, nil -} - -// jsonNumberOrString returns the team id as a JSON number when it's purely -// numeric (Runway's asTeamId is an integer in the reference payloads), else the -// raw string. -func jsonNumberOrString(teamID string) any { - return json.Number(strings.TrimSpace(teamID)) -} - -func isCreditError(s string) bool { - s = strings.ToLower(s) - return strings.Contains(s, "credit") || strings.Contains(s, "insufficient") || strings.Contains(s, "quota") -} - -// sleepCtx sleeps for d or until ctx is done; returns ctx.Err() if cancelled. -func sleepCtx(ctx context.Context, d time.Duration) error { - t := time.NewTimer(d) - defer t.Stop() - select { - case <-ctx.Done(): - return ctx.Err() - case <-t.C: - return nil - } -} diff --git a/backend/internal/repo/channel_repo.go b/backend/internal/repo/channel_repo.go new file mode 100644 index 0000000..c196bca --- /dev/null +++ b/backend/internal/repo/channel_repo.go @@ -0,0 +1,190 @@ +package repo + +import ( + "context" + "time" + + "backend/internal/model" + "gorm.io/gorm" + "gorm.io/gorm/clause" +) + +type ChannelRepository struct { + db *gorm.DB +} + +func NewChannelRepository(db *gorm.DB) *ChannelRepository { + return &ChannelRepository{db: db} +} + +// ChannelWithUpstream pairs a channel row with the upstream model name a +// binding carries — the shape dispatch needs in one fetch. +type ChannelWithUpstream struct { + Channel model.Channel + UpstreamModel string +} + +func (r *ChannelRepository) List(ctx context.Context) ([]model.Channel, error) { + var items []model.Channel + if err := r.db.WithContext(ctx).Order("created_at asc").Find(&items).Error; err != nil { + return nil, err + } + return items, nil +} + +func (r *ChannelRepository) Get(ctx context.Context, id string) (*model.Channel, error) { + var item model.Channel + if err := r.db.WithContext(ctx).First(&item, "id = ?", id).Error; err != nil { + return nil, err + } + return &item, nil +} + +func (r *ChannelRepository) Create(ctx context.Context, item *model.Channel) error { + return r.db.WithContext(ctx).Create(item).Error +} + +func (r *ChannelRepository) Update(ctx context.Context, id string, patch map[string]any) (*model.Channel, error) { + patch["updated_at"] = time.Now() + if err := r.db.WithContext(ctx).Model(&model.Channel{}).Where("id = ?", id).Updates(patch).Error; err != nil { + return nil, err + } + return r.Get(ctx, id) +} + +func (r *ChannelRepository) Delete(ctx context.Context, id string) (int64, error) { + res := r.db.WithContext(ctx).Delete(&model.Channel{}, "id = ?", id) + return res.RowsAffected, res.Error +} + +// Bindings returns all channel_models rows for a channel. +func (r *ChannelRepository) Bindings(ctx context.Context, channelID string) ([]model.ChannelModel, error) { + var items []model.ChannelModel + if err := r.db.WithContext(ctx).Where("channel_id = ?", channelID).Find(&items).Error; err != nil { + return nil, err + } + return items, nil +} + +// Bind upserts one channel↔model binding (upstream model name is refreshed on +// conflict so re-import can rename it). +func (r *ChannelRepository) Bind(ctx context.Context, channelID, modelID, upstreamModel string) error { + row := model.ChannelModel{ChannelID: channelID, ModelID: modelID, UpstreamModel: upstreamModel} + return r.db.WithContext(ctx).Clauses(clause.OnConflict{ + Columns: []clause.Column{{Name: "channel_id"}, {Name: "model_id"}}, + DoUpdates: clause.AssignmentColumns([]string{"upstream_model", "updated_at"}), + }).Create(&row).Error +} + +func (r *ChannelRepository) Unbind(ctx context.Context, channelID, modelID string) error { + return r.db.WithContext(ctx).Delete(&model.ChannelModel{}, "channel_id = ? AND model_id = ?", channelID, modelID).Error +} + +func (r *ChannelRepository) DeleteBindingsByChannel(ctx context.Context, channelID string) error { + return r.db.WithContext(ctx).Delete(&model.ChannelModel{}, "channel_id = ?", channelID).Error +} + +func (r *ChannelRepository) DeleteBindingsByModel(ctx context.Context, modelID string) error { + return r.db.WithContext(ctx).Delete(&model.ChannelModel{}, "model_id = ?", modelID).Error +} + +// ForModel returns enabled channels bound to modelID (channel rows + the +// upstream model name each binding carries), stable-ordered for dispatch. +func (r *ChannelRepository) ForModel(ctx context.Context, modelID string) ([]ChannelWithUpstream, error) { + var rows []struct { + model.Channel + UpstreamModel string `gorm:"column:upstream_model"` + } + err := r.db.WithContext(ctx). + Table("channels"). + Select("channels.*, channel_models.upstream_model"). + Joins("JOIN channel_models ON channel_models.channel_id = channels.id"). + Where("channel_models.model_id = ? AND channels.enabled = ?", modelID, true). + Order("channels.id asc"). + Scan(&rows).Error + if err != nil { + return nil, err + } + out := make([]ChannelWithUpstream, 0, len(rows)) + for _, row := range rows { + out = append(out, ChannelWithUpstream{Channel: row.Channel, UpstreamModel: row.UpstreamModel}) + } + return out, nil +} + +// ForModelAnyState is ForModel without the enabled filter — used by admin +// channel-pinned tests which deliberately target one channel even when it's +// disabled for traffic. +func (r *ChannelRepository) ForModelAnyState(ctx context.Context, modelID string) ([]ChannelWithUpstream, error) { + var rows []struct { + model.Channel + UpstreamModel string `gorm:"column:upstream_model"` + } + err := r.db.WithContext(ctx). + Table("channels"). + Select("channels.*, channel_models.upstream_model"). + Joins("JOIN channel_models ON channel_models.channel_id = channels.id"). + Where("channel_models.model_id = ?", modelID). + Order("channels.id asc"). + Scan(&rows).Error + if err != nil { + return nil, err + } + out := make([]ChannelWithUpstream, 0, len(rows)) + for _, row := range rows { + out = append(out, ChannelWithUpstream{Channel: row.Channel, UpstreamModel: row.UpstreamModel}) + } + return out, nil +} + +// ModelIDsByChannel lists local model ids bound to a channel. +func (r *ChannelRepository) ModelIDsByChannel(ctx context.Context, channelID string) ([]string, error) { + var ids []string + err := r.db.WithContext(ctx).Model(&model.ChannelModel{}). + Where("channel_id = ?", channelID).Pluck("model_id", &ids).Error + return ids, err +} + +// UpstreamModelMap returns model_id → upstream_model for a channel. +func (r *ChannelRepository) UpstreamModelMap(ctx context.Context, channelID string) (map[string]string, error) { + var rows []model.ChannelModel + if err := r.db.WithContext(ctx).Where("channel_id = ?", channelID).Find(&rows).Error; err != nil { + return nil, err + } + out := make(map[string]string, len(rows)) + for _, row := range rows { + out[row.ModelID] = row.UpstreamModel + } + return out, nil +} + +// ChannelIDsByModel lists channel ids bound to a model (any enabled state). +func (r *ChannelRepository) ChannelIDsByModel(ctx context.Context, modelID string) ([]string, error) { + var ids []string + err := r.db.WithContext(ctx).Model(&model.ChannelModel{}). + Where("model_id = ?", modelID).Pluck("channel_id", &ids).Error + return ids, err +} + +// OrphanModelIDs returns local model ids that no channel binds anymore. +func (r *ChannelRepository) OrphanModelIDs(ctx context.Context) ([]string, error) { + var ids []string + err := r.db.WithContext(ctx). + Table("model_configs"). + Select("model_configs.id"). + Joins("LEFT JOIN channel_models ON channel_models.model_id = model_configs.id"). + Where("channel_models.model_id IS NULL"). + Scan(&ids).Error + return ids, err +} + +// BumpSuccess records one successful upstream call (count + reported credits +// spent + last-used timestamp). Best-effort. +func (r *ChannelRepository) BumpSuccess(ctx context.Context, channelID string, credits float64) error { + now := time.Now() + return r.db.WithContext(ctx).Model(&model.Channel{}).Where("id = ?", channelID).Updates(map[string]any{ + "success_total": gorm.Expr("success_total + 1"), + "used_credits": gorm.Expr("used_credits + ?", credits), + "last_used_at": now, + }).Error +} diff --git a/backend/internal/repo/event_repo.go b/backend/internal/repo/event_repo.go index 3587bb5..1bdc9b4 100644 --- a/backend/internal/repo/event_repo.go +++ b/backend/internal/repo/event_repo.go @@ -24,8 +24,8 @@ type EventListFilter struct { UserID string UserIDs []string // when set, keep ONLY rows whose user_id is in this list (admin 用户搜索) Query string // free-text search over prompt / model / error (server-side, 跨页) - ExcludeSource string // when set, omit rows with this source (e.g. hide API-key "v1" usage from the customer logs page) - Source string // when set, keep ONLY rows with this source (admin 来源 filter): "v1" (API key) / "user" (前台) / "admin" (测试模型) + ExcludeSource string // when set, omit rows with this source (e.g. hide API-key "v1" usage from the customer logs page) + Source string // when set, keep ONLY rows with this source (admin 来源 filter): "v1" (API key) / "user" (前台) / "admin" (测试模型) HasFile bool // when true, keep ONLY rows with a non-empty file (the 创作记录 gallery — paginates over real media) ExcludeFiles []string // when set, omit rows whose file is in this list (e.g. hide homepage showcase media from user galleries) MediaOnly bool // when true, keep only rows that are pending or have a stored file — the 画图台 grid, so deleted works don't eat a slot @@ -618,9 +618,16 @@ func (r *EventRepository) MarkRefunded(ctx context.Context, eventID string) (boo return res.RowsAffected == 1, nil } -// SetAccount stamps which provider account is fulfilling an in-flight event. -// Called when generation commits to a token, so the accounts view can count -// pending events per account and an abandoned-event purge can attribute back. +// SetUpstreamJob records the job id the upstream assigned (NDJSON jobId frame) +// for cross-instance debugging. Best-effort. +func (r *EventRepository) SetUpstreamJob(ctx context.Context, eventID, jobID string) error { + return r.db.WithContext(ctx). + Model(&model.EventLog{}). + Where("id = ?", eventID). + Update("upstream_job", jobID).Error +} + +// SetAccount stamps which channel is fulfilling an in-flight event. func (r *EventRepository) SetAccount(ctx context.Context, eventID, accountID, accountEmail string) error { return r.db.WithContext(ctx). Model(&model.EventLog{}). diff --git a/backend/internal/repo/refresh_profile_repo.go b/backend/internal/repo/refresh_profile_repo.go deleted file mode 100644 index 532b9cd..0000000 --- a/backend/internal/repo/refresh_profile_repo.go +++ /dev/null @@ -1,73 +0,0 @@ -package repo - -import ( - "context" - "time" - - "backend/internal/model" - "gorm.io/gorm" -) - -type RefreshProfileRepository struct { - db *gorm.DB -} - -func NewRefreshProfileRepository(db *gorm.DB) *RefreshProfileRepository { - return &RefreshProfileRepository{db: db} -} - -func (r *RefreshProfileRepository) List(ctx context.Context) ([]model.RefreshProfile, error) { - var items []model.RefreshProfile - if err := r.db.WithContext(ctx). - Order("created_at desc"). - Find(&items).Error; err != nil { - return nil, err - } - return items, nil -} - -func (r *RefreshProfileRepository) Get(ctx context.Context, id string) (*model.RefreshProfile, error) { - var item model.RefreshProfile - if err := r.db.WithContext(ctx).First(&item, "id = ?", id).Error; err != nil { - return nil, err - } - return &item, nil -} - -func (r *RefreshProfileRepository) Create(ctx context.Context, item *model.RefreshProfile) error { - return r.db.WithContext(ctx).Create(item).Error -} - -func (r *RefreshProfileRepository) Update(ctx context.Context, id string, patch map[string]any) (*model.RefreshProfile, error) { - patch["updated_at"] = time.Now() - if err := r.db.WithContext(ctx). - Model(&model.RefreshProfile{}). - Where("id = ?", id). - Updates(patch).Error; err != nil { - return nil, err - } - return r.Get(ctx, id) -} - -func (r *RefreshProfileRepository) Delete(ctx context.Context, id string) error { - return r.db.WithContext(ctx).Delete(&model.RefreshProfile{}, "id = ?", id).Error -} - -func (r *RefreshProfileRepository) DeleteByIDs(ctx context.Context, ids []string) error { - if len(ids) == 0 { - return nil - } - return r.db.WithContext(ctx).Delete(&model.RefreshProfile{}, "id IN ?", ids).Error -} - -// ListDue returns enabled profiles whose next_retry_at has passed (or is unset, -// e.g. freshly imported). The background maintenance loop refreshes these. -func (r *RefreshProfileRepository) ListDue(ctx context.Context, now time.Time) ([]model.RefreshProfile, error) { - var items []model.RefreshProfile - if err := r.db.WithContext(ctx). - Where("enabled = ? AND (next_retry_at IS NULL OR next_retry_at <= ?)", true, now). - Find(&items).Error; err != nil { - return nil, err - } - return items, nil -} diff --git a/backend/internal/repo/token_repo.go b/backend/internal/repo/token_repo.go deleted file mode 100644 index 9b2d9e9..0000000 --- a/backend/internal/repo/token_repo.go +++ /dev/null @@ -1,529 +0,0 @@ -package repo - -import ( - "context" - "encoding/json" - "errors" - "strconv" - "strings" - "time" - - "backend/internal/model" - "gorm.io/datatypes" - "gorm.io/gorm" - "gorm.io/gorm/clause" -) - -type TokenRepository struct { - db *gorm.DB -} - -func NewTokenRepository(db *gorm.DB) *TokenRepository { - return &TokenRepository{db: db} -} - -func (r *TokenRepository) List(ctx context.Context) ([]model.TokenAccount, error) { - var items []model.TokenAccount - if err := r.db.WithContext(ctx). - Order("pool asc, created_at desc"). - Find(&items).Error; err != nil { - return nil, err - } - return items, nil -} - -func (r *TokenRepository) ListByPool(ctx context.Context, pool string) ([]model.TokenAccount, error) { - var items []model.TokenAccount - if err := r.db.WithContext(ctx). - Where("pool = ?", pool). - Order("created_at desc"). - Find(&items).Error; err != nil { - return nil, err - } - return items, nil -} - -func (r *TokenRepository) Get(ctx context.Context, pool, id string) (*model.TokenAccount, error) { - var item model.TokenAccount - if err := r.db.WithContext(ctx). - First(&item, "pool = ? AND id = ?", pool, id).Error; err != nil { - return nil, err - } - return &item, nil -} - -// GetByPoolEmail finds an account in a pool by its account_email (the logical -// identity for import dedup). Returns (nil, nil) when none / email is blank. -func (r *TokenRepository) GetByPoolEmail(ctx context.Context, pool, email string) (*model.TokenAccount, error) { - email = strings.TrimSpace(email) - if email == "" { - return nil, nil - } - var item model.TokenAccount - err := r.db.WithContext(ctx). - Where("pool = ? AND account_email = ?", pool, email). - First(&item).Error - if errors.Is(err, gorm.ErrRecordNotFound) { - return nil, nil - } - if err != nil { - return nil, err - } - return &item, nil -} - -func (r *TokenRepository) Create(ctx context.Context, item *model.TokenAccount) error { - return r.db.WithContext(ctx).Create(item).Error -} - -func (r *TokenRepository) Update(ctx context.Context, pool, id string, patch map[string]any) (*model.TokenAccount, error) { - patch["updated_at"] = time.Now() - if err := r.db.WithContext(ctx). - Model(&model.TokenAccount{}). - Where("pool = ? AND id = ?", pool, id). - Updates(patch).Error; err != nil { - return nil, err - } - return r.Get(ctx, pool, id) -} - -// SwapValue replaces an account's credential only while the stored one is still -// the value the caller started from. A rotating cookie is minted from whatever -// was in the row, so a goroutine that has been holding an older copy (a long -// render, a slow quota probe) must NOT be allowed to write it back over a newer -// rotation — that older copy no longer authenticates, and the account then looks -// dead. Reports whether the row was updated. -func (r *TokenRepository) SwapValue(ctx context.Context, pool, id, from, to string) (bool, error) { - res := r.db.WithContext(ctx). - Model(&model.TokenAccount{}). - Where("pool = ? AND id = ? AND value = ?", pool, id, from). - Updates(map[string]any{"value": to, "updated_at": time.Now()}) - if res.Error != nil { - return false, res.Error - } - return res.RowsAffected > 0, nil -} - -// ReserveQuota atomically pre-deducts `amount` from an account's cached image -// token balance under a row lock, so concurrent picks of the same near-empty -// account can never over-commit it. Returns: -// - allowed=true, deducted=true: balance was known and ≥ amount → decremented. -// - allowed=true, deducted=false: balance unknown → allowed without a hold -// (benefit of the doubt; a post-render reconcile writes the real value). -// - allowed=false: balance known and < amount → caller should fail over. -// RefundQuota releases a hold made with deducted=true when the render fails. -func (r *TokenRepository) ReserveQuota(ctx context.Context, pool, id string, amount int) (allowed, deducted bool, err error) { - err = r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { - var item model.TokenAccount - if e := tx.Clauses(clause.Locking{Strength: "UPDATE"}). - First(&item, "pool = ? AND id = ?", pool, id).Error; e != nil { - return e - } - rem, known := metaInt(item.Meta, "cached_quota_remaining") - if !known { - allowed, deducted = true, false - return nil - } - if rem < amount { - allowed, deducted = false, false - return nil - } - meta := cloneMeta(item.Meta) - meta["cached_quota_remaining"] = rem - amount - if e := tx.Model(&model.TokenAccount{}). - Where("pool = ? AND id = ?", pool, id). - Updates(map[string]any{"meta": meta, "updated_at": time.Now()}).Error; e != nil { - return e - } - allowed, deducted = true, true - return nil - }) - return allowed, deducted, err -} - -// Grok accounts carry a forced local quota instead of an upstream balance: -// Console 没有额度接口，所以导入时写死 图 5 / 视频 2，下单预扣、失败退回，两个都归零直接判死。 -const ( - GrokImageQuotaKey = "grok_image_remaining" - GrokVideoQuotaKey = "grok_video_remaining" - GrokImageQuota = 5 - GrokVideoQuota = 2 -) - -// grokQuotas reads an account's local per-kind counters, falling back to the -// forced defaults for accounts imported before the counters existed. -func grokQuotas(meta datatypes.JSONMap) (images, videos int) { - images, known := metaInt(meta, GrokImageQuotaKey) - if !known { - images = GrokImageQuota - } - videos, known = metaInt(meta, GrokVideoQuotaKey) - if !known { - videos = GrokVideoQuota - } - return images, videos -} - -// ReserveGrokQuota pre-deducts one unit of a grok account's local per-kind quota -// under a row lock, so concurrent orders on the same near-empty account can't -// over-commit it（下单即扣，不会超扣）。allowed=false 表示这一份已经用光，调用方换号。 -// 归零的那一份打上 image_limited / video_limited 让调度跳过；判死留给 -// FinalizeGrokQuota（生成真的成功了才锁号），失败时用 RefundGrokQuota 退回。 -func (r *TokenRepository) ReserveGrokQuota(ctx context.Context, id, kind string) (allowed bool, err error) { - err = r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { - var item model.TokenAccount - if e := tx.Clauses(clause.Locking{Strength: "UPDATE"}). - First(&item, "pool = ? AND id = ?", "grok", id).Error; e != nil { - return e - } - images, videos := grokQuotas(item.Meta) - if kind == "video" { - if videos <= 0 { - return nil - } - videos-- - } else { - if images <= 0 { - return nil - } - images-- - } - meta := cloneMeta(item.Meta) - meta[GrokImageQuotaKey] = images - meta[GrokVideoQuotaKey] = videos - if e := tx.Model(&model.TokenAccount{}). - Where("pool = ? AND id = ?", "grok", id). - Updates(map[string]any{ - "meta": meta, - "image_limited": images <= 0, - "video_limited": videos <= 0, - "updated_at": time.Now(), - }).Error; e != nil { - return e - } - allowed = true - return nil - }) - return allowed, err -} - -// RefundGrokQuota gives back a unit reserved by ReserveGrokQuota when the render -// failed, clearing that kind's limited flag again. Capped at the forced default so -// repeated refunds can't inflate the quota. -func (r *TokenRepository) RefundGrokQuota(ctx context.Context, id, kind string) error { - return r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { - var item model.TokenAccount - if e := tx.Clauses(clause.Locking{Strength: "UPDATE"}). - First(&item, "pool = ? AND id = ?", "grok", id).Error; e != nil { - return e - } - images, videos := grokQuotas(item.Meta) - if kind == "video" { - videos = min(videos+1, GrokVideoQuota) - } else { - images = min(images+1, GrokImageQuota) - } - meta := cloneMeta(item.Meta) - meta[GrokImageQuotaKey] = images - meta[GrokVideoQuotaKey] = videos - return tx.Model(&model.TokenAccount{}). - Where("pool = ? AND id = ?", "grok", id). - Updates(map[string]any{ - "meta": meta, - "image_limited": images <= 0, - "video_limited": videos <= 0, - "updated_at": time.Now(), - }).Error - }) -} - -// FinalizeGrokQuota locks an account down once a successful generation left both -// local counters at zero (no reset time — 用完就废). -func (r *TokenRepository) FinalizeGrokQuota(ctx context.Context, id string) error { - return r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { - var item model.TokenAccount - if e := tx.Clauses(clause.Locking{Strength: "UPDATE"}). - First(&item, "pool = ? AND id = ?", "grok", id).Error; e != nil { - return e - } - images, videos := grokQuotas(item.Meta) - if images > 0 || videos > 0 { - return nil - } - return tx.Model(&model.TokenAccount{}). - Where("pool = ? AND id = ?", "grok", id). - Updates(map[string]any{ - "status": "disabled", - "dead": true, - "updated_at": time.Now(), - }).Error - }) -} - -// RefundQuota atomically adds `amount` back to cached_quota_remaining (releasing a -// hold from a reservation whose render then failed). No-op if the balance is -// unknown. Row-locked like ReserveQuota. -func (r *TokenRepository) RefundQuota(ctx context.Context, pool, id string, amount int) error { - return r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { - var item model.TokenAccount - if e := tx.Clauses(clause.Locking{Strength: "UPDATE"}). - First(&item, "pool = ? AND id = ?", pool, id).Error; e != nil { - return e - } - rem, known := metaInt(item.Meta, "cached_quota_remaining") - if !known { - return nil - } - meta := cloneMeta(item.Meta) - meta["cached_quota_remaining"] = rem + amount - return tx.Model(&model.TokenAccount{}). - Where("pool = ? AND id = ?", pool, id). - Updates(map[string]any{"meta": meta, "updated_at": time.Now()}).Error - }) -} - -func cloneMeta(m datatypes.JSONMap) datatypes.JSONMap { - out := datatypes.JSONMap{} - for k, v := range m { - out[k] = v - } - return out -} - -func metaInt(m datatypes.JSONMap, key string) (int, bool) { - if m == nil { - return 0, false - } - v, ok := m[key] - if !ok || v == nil { - return 0, false - } - switch x := v.(type) { - case int: - return x, true - case int64: - return int(x), true - case float64: - return int(x), true - case json.Number: - n, e := x.Int64() - if e != nil { - return 0, false - } - return int(n), true - case string: - n, e := strconv.Atoi(strings.TrimSpace(x)) - if e != nil { - return 0, false - } - return n, true - default: - return 0, false - } -} - -// TouchLastUsed stamps last_used_at at the moment a token is SELECTED, so the -// accounts view reflects an accurate "last used" time. Rotation order is driven -// by the in-memory strict round-robin cursor in the service layer (see -// V1Service.rotateRoundRobin), not by this timestamp. -func (r *TokenRepository) TouchLastUsed(ctx context.Context, id string) error { - return r.db.WithContext(ctx). - Model(&model.TokenAccount{}). - Where("id = ?", id). - Update("last_used_at", time.Now()).Error -} - -// IncrementFail bumps an account's failure counters by one. Used to attribute -// an abandoned (purged) generation's failure back to the account it was using, -// since that generation never reached the normal markTokenFailure path. -func (r *TokenRepository) IncrementFail(ctx context.Context, id string) error { - return r.db.WithContext(ctx). - Model(&model.TokenAccount{}). - Where("id = ?", id). - Updates(map[string]any{ - "fail_total": gorm.Expr("fail_total + 1"), - "fails": gorm.Expr("fails + 1"), - "updated_at": time.Now(), - }).Error -} - -func (r *TokenRepository) Delete(ctx context.Context, pool, id string) (int64, error) { - res := r.db.WithContext(ctx). - Delete(&model.TokenAccount{}, "pool = ? AND id = ?", pool, id) - return res.RowsAffected, res.Error -} - -// DeleteByIDs removes accounts by id across pools (ids are globally unique), -// for bulk delete. Returns the number of rows removed. -func (r *TokenRepository) DeleteByIDs(ctx context.Context, ids []string) (int64, error) { - if len(ids) == 0 { - return 0, nil - } - res := r.db.WithContext(ctx).Delete(&model.TokenAccount{}, "id IN ?", ids) - return res.RowsAffected, res.Error -} - -// leonardoDailyTokens is the free-tier daily allowance restored at each reset. -// A paid account's true balance is reconciled on its next successful render. -const leonardoDailyTokens = 150 - -// RecoverQuota reactivates quota-exhausted tokens whose reset time has passed. -// Reset source: cached_quota_reset_after (upstream marker) first, else the -// quota_recover_at fallback stamped when the token was marked quota-exhausted. -// Mirrors Python TokenPool.recover_quota; returns the count reactivated. -// RecoverQuota reactivates quota-exhausted tokens whose reset time has passed and -// returns the accounts it recovered, so the caller can re-sync their real balance -// (the providers only sync quota when accessed). -func (r *TokenRepository) RecoverQuota(ctx context.Context) ([]model.TokenAccount, error) { - // Also pick up accounts that are only single-kind limited (image_limited / - // video_limited) — those keep status "active" and would otherwise never have - // their per-kind flag cleared. This is legacy support; Adobe no longer sets - // per-kind limits. - var items []model.TokenAccount - if err := r.db.WithContext(ctx). - Where("status = ? OR image_limited = ? OR video_limited = ?", "quota", true, true). - Find(&items).Error; err != nil { - return nil, err - } - now := time.Now() - var recovered []model.TokenAccount - for i := range items { - t := &items[i] - // Runway's reset marker is the JWT expiry, not a quota-refresh time, and - // there's no way to refresh a bare JWT — so a runway account is never - // "recovered"; it's expired-to-dead by ExpireByReset instead. - if t.Pool == "runway" { - continue - } - reset := parseResetMarker(t.CachedQuotaResetAfter) - if reset == nil { - reset = t.QuotaRecoverAt - } - if reset == nil || now.Before(*reset) { - continue - } - patch := map[string]any{ - "fails": 0, - "quota_recover_at": nil, - } - // For Adobe, also clear any legacy per-kind limit flags. - if t.Pool == "adobe" { - patch["image_limited"] = false - patch["video_limited"] = false - } - // Only flip status back to active if it was sunk to "quota" (both kinds - // limited); a single-kind limit left status untouched. - if t.Status == "quota" { - patch["status"] = "active" - } - // Leonardo's free tokens fully renew at each daily reset — restore the - // balance and advance the reset marker to the next 08:00 Beijing (== next - // UTC midnight), so the account is immediately usable instead of stuck at a - // stale 0. A paid account's real balance is corrected on its next render. - if t.Pool == "leonardo" || t.Pool == "krea" || t.Pool == "imagine" { - meta := cloneMeta(t.Meta) - if t.Pool == "leonardo" { - meta["cached_quota_remaining"] = leonardoDailyTokens - } else { - // Krea/Imagine balances re-sync from upstream (billing-data / v1/credit) - // on next probe — drop the stale value so the account isn't shown as - // empty after reset. - delete(meta, "cached_quota_remaining") - } - meta["cached_quota_at"] = int(now.Unix()) - patch["meta"] = meta - patch["cached_quota_reset_after"] = time.Unix((now.Unix()/86400+1)*86400, 0).UTC().Format(time.RFC3339) - } - if _, err := r.Update(ctx, t.Pool, t.ID, patch); err != nil { - return recovered, err - } - recovered = append(recovered, *t) - } - return recovered, nil -} - -// RollResetMarkers advances a stale (past) daily-reset marker to its next future -// occurrence — same time-of-day, +N whole days — for ACTIVE accounts of the given -// daily-reset pools, so the 恢复时间 column always shows the upcoming reset rather -// than yesterday's. Only active accounts are rolled: a 限额 account must keep its -// past marker so RecoverQuota can recover it (rolling it forward early would -// prevent recovery). Returns the number advanced. -func (r *TokenRepository) RollResetMarkers(ctx context.Context, pools []string) (int, error) { - var items []model.TokenAccount - if err := r.db.WithContext(ctx). - Where("pool IN ? AND dead = ? AND status = ? AND image_limited = ? AND video_limited = ? AND cached_quota_reset_after <> ''", - pools, false, "active", false, false). - Find(&items).Error; err != nil { - return 0, err - } - now := time.Now() - n := 0 - for i := range items { - t := &items[i] - reset := parseResetMarker(t.CachedQuotaResetAfter) - if reset == nil || !reset.Before(now) { - continue // unparseable or already in the future - } - next := *reset - for !next.After(now) { - next = next.Add(24 * time.Hour) - } - if _, err := r.Update(ctx, t.Pool, t.ID, map[string]any{ - "cached_quota_reset_after": next.UTC().Format(time.RFC3339), - }); err != nil { - return n, err - } - n++ - } - return n, nil -} - -// ExpireByReset marks accounts of a pool dead once their reset marker has passed. -// For runway the marker IS the JWT expiry and there's no refresh, so an expired -// token can only 401 — we proactively flip it to disabled+dead (the same end -// state a 401 would produce) instead of leaving a doomed account "active". -func (r *TokenRepository) ExpireByReset(ctx context.Context, pool string) (int, error) { - var items []model.TokenAccount - if err := r.db.WithContext(ctx). - Where("pool = ? AND dead = ?", pool, false). - Find(&items).Error; err != nil { - return 0, err - } - now := time.Now() - expired := 0 - for i := range items { - t := &items[i] - reset := parseResetMarker(t.CachedQuotaResetAfter) - if reset == nil || now.Before(*reset) { - continue - } - if _, err := r.Update(ctx, t.Pool, t.ID, map[string]any{ - "status": "disabled", - "dead": true, - }); err != nil { - return expired, err - } - expired++ - } - return expired, nil -} - -// parseResetMarker best-effort parses a quota reset marker into a time. Accepts -// epoch seconds (numeric string) or ISO-8601 (e.g. Adobe's available_until -// "2026-06-16T23:59:59.999Z"). Returns nil if unparseable. -func parseResetMarker(v string) *time.Time { - v = strings.TrimSpace(v) - if v == "" { - return nil - } - if f, err := strconv.ParseFloat(v, 64); err == nil && f > 946684800 { - t := time.Unix(int64(f), 0) - return &t - } - for _, layout := range []string{time.RFC3339, "2006-01-02T15:04:05.999Z07:00", "2006-01-02T15:04:05Z07:00"} { - if t, err := time.Parse(layout, v); err == nil { - return &t - } - } - return nil -} diff --git a/backend/internal/service/admin_read.go b/backend/internal/service/admin_read.go index 9e8c34f..37ed753 100644 --- a/backend/internal/service/admin_read.go +++ b/backend/internal/service/admin_read.go @@ -19,20 +19,20 @@ type AdminReadService struct { models *repo.ModelRepository events *repo.EventRepository settings *repo.SiteSettingRepository - tokens *repo.TokenRepository + channels *repo.ChannelRepository cdks *repo.CDKRepository store *storage.Client showcase *repo.ShowcaseRepository } -func NewAdminReadService(cfg *config.Config, users *repo.UserRepository, models *repo.ModelRepository, events *repo.EventRepository, settings *repo.SiteSettingRepository, tokens *repo.TokenRepository, cdks *repo.CDKRepository, store *storage.Client, showcase *repo.ShowcaseRepository) *AdminReadService { +func NewAdminReadService(cfg *config.Config, users *repo.UserRepository, models *repo.ModelRepository, events *repo.EventRepository, settings *repo.SiteSettingRepository, channels *repo.ChannelRepository, cdks *repo.CDKRepository, store *storage.Client, showcase *repo.ShowcaseRepository) *AdminReadService { return &AdminReadService{ cfg: cfg, users: users, models: models, events: events, settings: settings, - tokens: tokens, + channels: channels, cdks: cdks, store: store, showcase: showcase, @@ -104,6 +104,10 @@ func (s *AdminReadService) ModelsView(ctx context.Context) ([]map[string]any, er "max_reference_images": item.MaxReferenceImages, "max_reference_videos": item.MaxReferenceVideos, "max_reference_audios": item.MaxReferenceAudios, + "duration_min": item.DurationMin, + "duration_max": item.DurationMax, + "max_ref_video_seconds": item.MaxRefVideoSeconds, + "max_ref_audio_seconds": item.MaxRefAudioSeconds, "reference_mode": item.ReferenceMode, "weight": item.Weight, "generation_count": item.GenerationCount, @@ -196,24 +200,24 @@ func (s *AdminReadService) UserNameMap(ctx context.Context) (map[string]string, return out, nil } -// AccountNameMap builds a token-account id -> display label lookup (account -// email, else display name, else id) used to annotate log rows with which -// provider account fulfilled each generation (event_logs.account_id). +// AccountNameMap builds a channel id -> display name lookup used to annotate +// log rows with which channel fulfilled each generation (event_logs.account_id +// now stores the channel id; account_email carries a denormalized name stamp). func (s *AdminReadService) AccountNameMap(ctx context.Context) (map[string]string, error) { - accounts, err := s.tokens.List(ctx) + if s.channels == nil { + return map[string]string{}, nil + } + channels, err := s.channels.List(ctx) if err != nil { return nil, err } - out := make(map[string]string, len(accounts)) - for _, a := range accounts { - label := strings.TrimSpace(a.AccountEmail) + out := make(map[string]string, len(channels)) + for _, ch := range channels { + label := strings.TrimSpace(ch.Name) if label == "" { - label = strings.TrimSpace(a.AccountDisplayName) + label = ch.ID } - if label == "" { - label = a.ID - } - out[a.ID] = label + out[ch.ID] = label } return out, nil } @@ -422,66 +426,30 @@ func (s *AdminReadService) Invites(ctx context.Context) ([]repo.InviteRecord, *r return s.users.AllInvites(ctx, reward) } +// Providers is now a channel summary (the account pools are gone): one row per +// channel with its bound-model count and local usage stats. func (s *AdminReadService) Providers(ctx context.Context) ([]map[string]any, error) { - models, err := s.models.List(ctx) + if s.channels == nil { + return []map[string]any{}, nil + } + channels, err := s.channels.List(ctx) if err != nil { return nil, err } - tokens, err := s.tokens.List(ctx) - if err != nil { - return nil, err - } - modelCounts := map[string]int{} - for _, item := range models { - modelCounts[item.Provider]++ - } - type aggregate struct { - active int - disabled int - quota int - } - tokenCounts := map[string]*aggregate{} - for _, item := range tokens { - if _, ok := tokenCounts[item.Pool]; !ok { - tokenCounts[item.Pool] = &aggregate{} - } - switch item.Status { - case "active": - tokenCounts[item.Pool].active++ - case "quota": - tokenCounts[item.Pool].quota++ - default: - tokenCounts[item.Pool].disabled++ - } - } - providers := []struct { - Name string - Pool string - Type string - }{ - {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 { - count := tokenCounts[item.Pool] - if count == nil { - count = &aggregate{} - } + out := make([]map[string]any, 0, len(channels)) + for _, ch := range channels { + modelIDs, _ := s.channels.ModelIDsByChannel(ctx, ch.ID) out = append(out, map[string]any{ - "name": item.Name, - "token_pool": item.Pool, - "type": item.Type, - "model_count": modelCounts[item.Name], - "tokens_total": count.active + count.disabled + count.quota, - "tokens_active": count.active, - "tokens_disabled": count.disabled, - "tokens_quota": count.quota, + "id": ch.ID, + "name": ch.Name, + "base_url": ch.BaseURL, + "enabled": ch.Enabled, + "concurrency": ch.Concurrency, + "weight": ch.Weight, + "model_count": len(modelIDs), + "success_total": ch.SuccessTotal, + "used_credits": ch.UsedCredits, + "type": "vivid", }) } return out, nil diff --git a/backend/internal/service/admin_write.go b/backend/internal/service/admin_write.go index 7f0ab72..dbb4335 100644 --- a/backend/internal/service/admin_write.go +++ b/backend/internal/service/admin_write.go @@ -28,18 +28,18 @@ type AdminWriteService struct { models *repo.ModelRepository events *repo.EventRepository apiKeys *repo.APIKeyRepository - tokens *repo.TokenRepository + channels *repo.ChannelRepository orders *repo.OrderRepository } -func NewAdminWriteService(users *repo.UserRepository, showcase *repo.ShowcaseRepository, models *repo.ModelRepository, events *repo.EventRepository, apiKeys *repo.APIKeyRepository, tokens *repo.TokenRepository, orders *repo.OrderRepository) *AdminWriteService { +func NewAdminWriteService(users *repo.UserRepository, showcase *repo.ShowcaseRepository, models *repo.ModelRepository, events *repo.EventRepository, apiKeys *repo.APIKeyRepository, channels *repo.ChannelRepository, orders *repo.OrderRepository) *AdminWriteService { return &AdminWriteService{ users: users, showcase: showcase, models: models, events: events, apiKeys: apiKeys, - tokens: tokens, + channels: channels, orders: orders, } } @@ -371,7 +371,7 @@ func (s *AdminWriteService) CreateModel(ctx context.Context, body map[string]any return nil, errors.New("type must be image or video") } if provider == "" { - return nil, errors.New("provider required") + provider = "channel" } if err := s.validateModelNameSpace(ctx, "", modelID, alias); err != nil { return nil, err @@ -436,10 +436,9 @@ func (s *AdminWriteService) UpdateModel(ctx context.Context, modelID string, bod } if _, ok := body["provider"]; ok { provider := strings.TrimSpace(stringValue(body["provider"])) - if provider == "" { - return nil, errors.New("provider required") + if provider != "" { + patch["provider"] = provider } - patch["provider"] = provider } // Only touch `enabled` when the caller explicitly sends a non-null value; // mirrors Python models_store.update ("enabled" in fields and is not None). @@ -488,6 +487,18 @@ func (s *AdminWriteService) UpdateModel(ctx context.Context, modelID string, bod if _, ok := body["reference_mode"]; ok { patch["reference_mode"] = defaultString(strings.TrimSpace(stringValue(body["reference_mode"])), "none") } + if _, ok := body["duration_min"]; ok { + patch["duration_min"] = intValue(body["duration_min"]) + } + if _, ok := body["duration_max"]; ok { + patch["duration_max"] = intValue(body["duration_max"]) + } + if _, ok := body["max_ref_video_seconds"]; ok { + patch["max_ref_video_seconds"] = intValue(body["max_ref_video_seconds"]) + } + if _, ok := body["max_ref_audio_seconds"]; ok { + patch["max_ref_audio_seconds"] = intValue(body["max_ref_audio_seconds"]) + } if _, ok := body["weight"]; ok { patch["weight"] = intValue(body["weight"]) } @@ -532,52 +543,14 @@ func (s *AdminWriteService) DeleteModel(ctx context.Context, modelID string) err if rows == 0 { return ErrNotFound } - // Clean up: strip this model id from every custom upstream account's - // supported-models list so no orphan reference is left behind. - s.removeModelFromUpstreams(ctx, modelID) + // Clean up: drop every channel binding pointing at the deleted model so no + // orphan reference is left behind. + if s.channels != nil { + _ = s.channels.DeleteBindingsByModel(ctx, modelID) + } return nil } -// removeModelFromUpstreams drops modelID from the CSV in each custom account's -// Meta["models"]. Best-effort — a failure here doesn't undo the model delete. -func (s *AdminWriteService) removeModelFromUpstreams(ctx context.Context, modelID string) { - if s.tokens == nil { - return - } - items, err := s.tokens.ListByPool(ctx, "custom") - if err != nil { - return - } - for _, it := range items { - raw, _ := it.Meta["models"].(string) - if strings.TrimSpace(raw) == "" { - continue - } - kept := make([]string, 0, len(strings.Split(raw, ","))) - changed := false - for _, p := range strings.Split(raw, ",") { - p = strings.TrimSpace(p) - if p == "" { - continue - } - if p == modelID { - changed = true - continue - } - kept = append(kept, p) - } - if !changed { - continue - } - meta := datatypes.JSONMap{} - for k, v := range it.Meta { - meta[k] = v - } - meta["models"] = strings.Join(kept, ",") - _, _ = s.tokens.Update(ctx, "custom", it.ID, map[string]any{"meta": meta}) - } -} - func (s *AdminWriteService) ClearLogs(ctx context.Context) (int64, error) { return s.events.DeleteAll(ctx) } diff --git a/backend/internal/service/api_key.go b/backend/internal/service/api_key.go index 1912e5b..650eb51 100644 --- a/backend/internal/service/api_key.go +++ b/backend/internal/service/api_key.go @@ -2,6 +2,7 @@ package service import ( "context" + "crypto/rand" "crypto/sha256" "encoding/hex" "errors" @@ -109,8 +110,15 @@ func (s *APIKeyService) DeleteOne(ctx context.Context, userID, keyID string) err return s.keys.DeleteByID(ctx, userID, keyID) } +// generatePlainAPIKey mints a vivid-shaped key: vk_ + 48 lowercase hex chars +// (upstream channel api_keys use the same shape, so one key format rules +// everything — no per-direction prefixes). func generatePlainAPIKey() (string, error) { - return "sk-" + randomUpper(38), nil + b := make([]byte, 24) + if _, err := rand.Read(b); err != nil { + return "", err + } + return "vk_" + hex.EncodeToString(b), nil } func previewAPIKey(plain string) string { diff --git a/backend/internal/service/channel.go b/backend/internal/service/channel.go new file mode 100644 index 0000000..054723b --- /dev/null +++ b/backend/internal/service/channel.go @@ -0,0 +1,450 @@ +package service + +import ( + "context" + "crypto/aes" + "crypto/cipher" + "crypto/rand" + "crypto/sha256" + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "io" + "strings" + "sync" + "time" + + "backend/internal/model" + "backend/internal/repo" + "backend/internal/upstream" + + "gorm.io/datatypes" + "gorm.io/gorm" +) + +// ChannelService manages upstream vivid-protocol channels: CRUD, upstream key +// encryption at rest, model fetching/import/binding, and live balance probes. +// The actual request dispatch lives on V1Service (it owns concurrency + events). +type ChannelService struct { + channels *repo.ChannelRepository + models *repo.ModelRepository + up *upstream.Client + secret []byte +} + +func NewChannelService(channels *repo.ChannelRepository, models *repo.ModelRepository, up *upstream.Client, secret string) *ChannelService { + sum := sha256.Sum256([]byte("image2api-channel-keys|" + secret)) + return &ChannelService{channels: channels, models: models, up: up, secret: sum[:]} +} + +// ChannelView is the admin list row: local stats + bound model names + a live +// upstream balance probe (nil when the upstream doesn't implement /v1/balance +// or the probe failed). +type ChannelView struct { + ID string `json:"id"` + Name string `json:"name"` + BaseURL string `json:"base_url"` + Enabled bool `json:"enabled"` + Concurrency int `json:"concurrency"` + Weight int `json:"weight"` + Models []string `json:"models"` + SuccessTotal int64 `json:"success_total"` + UsedCredits float64 `json:"used_credits"` + UpstreamBalance *int64 `json:"upstream_balance"` + LastUsedAt *int64 `json:"last_used_at"` + CreatedAt int64 `json:"created_at"` + UpdatedAt int64 `json:"updated_at"` +} + +func (s *ChannelService) List(ctx context.Context) ([]ChannelView, error) { + items, err := s.channels.List(ctx) + if err != nil { + return nil, err + } + out := make([]ChannelView, len(items)) + var wg sync.WaitGroup + for i, ch := range items { + out[i] = ChannelView{ + ID: ch.ID, + Name: ch.Name, + BaseURL: ch.BaseURL, + Enabled: ch.Enabled, + Concurrency: ch.Concurrency, + Weight: ch.Weight, + SuccessTotal: ch.SuccessTotal, + UsedCredits: ch.UsedCredits, + CreatedAt: ch.CreatedAt.Unix(), + UpdatedAt: ch.UpdatedAt.Unix(), + } + if ch.LastUsedAt != nil { + t := ch.LastUsedAt.Unix() + out[i].LastUsedAt = &t + } + if ids, err := s.channels.ModelIDsByChannel(ctx, ch.ID); err == nil { + names := make([]string, 0, len(ids)) + for _, id := range ids { + if m, err := s.models.Get(ctx, id); err == nil { + names = append(names, m.EffectiveName()) + } else { + names = append(names, id) + } + } + out[i].Models = names + } + if key := s.decrypt(ch.APIKeyEnc); key != "" { + wg.Add(1) + go func(i int, ch model.Channel, key string) { + defer wg.Done() + bctx, cancel := context.WithTimeout(context.Background(), 8*time.Second) + defer cancel() + if bal, err := s.up.Balance(bctx, ch.BaseURL, key); err == nil { + out[i].UpstreamBalance = bal + } + }(i, ch, key) + } + } + wg.Wait() + return out, nil +} + +// SaveChannelDto is the admin create/update payload. ApiKey empty or the mask +// placeholder keeps the existing key on update; required on create. +type SaveChannelDto struct { + ID string `json:"id"` + Name string `json:"name"` + BaseURL string `json:"base_url"` + APIKey string `json:"api_key"` + Enabled bool `json:"enabled"` + Concurrency int `json:"concurrency"` + Weight int `json:"weight"` +} + +func (s *ChannelService) Save(ctx context.Context, d SaveChannelDto) (*model.Channel, error) { + name := strings.TrimSpace(d.Name) + base := strings.TrimRight(strings.TrimSpace(d.BaseURL), "/") + if base == "" { + return nil, errors.New("base_url required") + } + if !strings.HasPrefix(base, "http://") && !strings.HasPrefix(base, "https://") { + return nil, errors.New("base_url must start with http:// or https://") + } + weight := d.Weight + if weight < 1 { + weight = 1 + } + if d.ID == "" { + if strings.TrimSpace(d.APIKey) == "" { + return nil, errors.New("api_key required") + } + enc, err := s.encrypt(d.APIKey) + if err != nil { + return nil, err + } + ch := &model.Channel{ + ID: "ch-" + randomUpper(10), + Name: name, + BaseURL: base, + APIKeyEnc: enc, + Enabled: d.Enabled, + Concurrency: d.Concurrency, + Weight: weight, + } + if ch.Name == "" { + ch.Name = base + } + return ch, s.channels.Create(ctx, ch) + } + patch := map[string]any{ + "name": name, + "base_url": base, + "enabled": d.Enabled, + "concurrency": d.Concurrency, + "weight": weight, + } + if key := strings.TrimSpace(d.APIKey); key != "" && !strings.Contains(key, "•") && !strings.Contains(key, "*") { + enc, err := s.encrypt(key) + if err != nil { + return nil, err + } + patch["api_key_enc"] = enc + } + return s.channels.Update(ctx, d.ID, patch) +} + +// Delete removes a channel, its bindings, and any local model that ends up +// bound to no channel at all (same semantics as vivid). +func (s *ChannelService) Delete(ctx context.Context, id string) error { + modelIDs, _ := s.channels.ModelIDsByChannel(ctx, id) + if err := s.channels.DeleteBindingsByChannel(ctx, id); err != nil { + return err + } + res := 0 + if n, err := s.channels.Delete(ctx, id); err != nil { + return err + } else { + res = int(n) + } + if res == 0 { + return gorm.ErrRecordNotFound + } + for _, mid := range modelIDs { + rest, _ := s.channels.ChannelIDsByModel(ctx, mid) + if len(rest) == 0 { + _, _ = s.models.Delete(ctx, mid) + } + } + return nil +} + +// FetchModelsDto resolves credentials for a model fetch: Id uses the stored +// channel's key; otherwise BaseURL+ApiKey preview an unsaved channel. +type FetchModelsDto struct { + ID string `json:"id"` + BaseURL string `json:"base_url"` + APIKey string `json:"api_key"` +} + +// FetchModels lists the upstream's sellable models and flags which are already +// imported locally. +func (s *ChannelService) FetchModels(ctx context.Context, d FetchModelsDto) ([]upstream.RemoteModel, error) { + base, key, err := s.resolveCreds(ctx, d.ID, d.BaseURL, d.APIKey) + if err != nil { + return nil, err + } + items, err := s.up.Models(ctx, base, key) + if err != nil { + return nil, err + } + return items, nil +} + +// ChannelModelSpec is one model the admin ticks for import (same shape as the +// upstream /v1/models item plus an optional local alias). +type ChannelModelSpec struct { + ID string `json:"id"` + Alias string `json:"alias"` + Kind string `json:"kind"` + Qualities []string `json:"qualities"` + Duration *upstream.RemoteDuration `json:"duration,omitempty"` + RefImages int64 `json:"refImages"` + RefVideos int64 `json:"refVideos"` + RefAudios int64 `json:"refAudios"` + RefVideoSeconds int64 `json:"refVideoSeconds,omitempty"` + RefAudioSeconds int64 `json:"refAudioSeconds,omitempty"` +} + +func (s *ChannelService) resolveCreds(ctx context.Context, id, base, key string) (string, string, error) { + if strings.TrimSpace(id) != "" { + ch, err := s.channels.Get(ctx, id) + if err != nil { + return "", "", err + } + return strings.TrimRight(ch.BaseURL, "/"), s.decrypt(ch.APIKeyEnc), nil + } + base = strings.TrimRight(strings.TrimSpace(base), "/") + if base == "" || strings.TrimSpace(key) == "" { + return "", "", errors.New("id or base_url+api_key required") + } + return base, strings.TrimSpace(key), nil +} + +// ImportModels creates/updates local model rows for the ticked specs and binds +// them to the channel. Pricing stays whatever the local row already has +// (empty for new rows — admin prices them in 模型管理). +func (s *ChannelService) ImportModels(ctx context.Context, channelID string, specs []ChannelModelSpec) error { + ch, err := s.channels.Get(ctx, channelID) + if err != nil { + return err + } + for _, spec := range specs { + upID := strings.TrimSpace(spec.ID) + if upID == "" { + continue + } + localID := strings.TrimSpace(spec.Alias) + if localID == "" { + localID = upID + } + existing, gerr := s.models.Get(ctx, localID) + if gerr != nil && !errors.Is(gerr, gorm.ErrRecordNotFound) { + existing = nil + } + var local *model.ModelConfig + if existing != nil && existing.ID == localID { + local = existing + } + patch := modelPatchFromSpec(spec, local) + if local == nil { + local = &model.ModelConfig{ + ID: localID, + Name: localID, + Type: spec.Kind, + Provider: "channel", + Enabled: true, + Ratios: jsonArray([]string{"16:9"}), + Prices: datatypes.JSONMap{}, + CreatedAt: time.Now(), + } + applySpec(local, patch) + if err := s.models.Create(ctx, local); err != nil { + return fmt.Errorf("model %s: %w", localID, err) + } + } else { + if _, err := s.models.Update(ctx, localID, patch); err != nil { + return fmt.Errorf("model %s: %w", localID, err) + } + } + if err := s.channels.Bind(ctx, ch.ID, localID, upID); err != nil { + return err + } + } + return nil +} + +// modelPatchFromSpec maps an upstream model spec onto model_configs columns. +// Prices/enabled/alias are intentionally untouched on update so admin pricing +// survives re-import. +func modelPatchFromSpec(spec ChannelModelSpec, existing *model.ModelConfig) map[string]any { + patch := map[string]any{ + "type": spec.Kind, + "resolutions": jsonArray(spec.Qualities), + "max_reference_images": spec.RefImages, + "max_reference_videos": spec.RefVideos, + "max_reference_audios": spec.RefAudios, + "max_ref_video_seconds": spec.RefVideoSeconds, + "max_ref_audio_seconds": spec.RefAudioSeconds, + } + if spec.RefImages+spec.RefVideos+spec.RefAudios > 0 { + patch["reference_mode"] = "asset" + patch["image_to_image"] = spec.RefImages > 0 + } + if d := spec.Duration; d != nil { + switch d.Mode { + case "fixed": + patch["durations"] = jsonArray([]string{fmt.Sprintf("%ds", d.Seconds)}) + patch["duration_min"], patch["duration_max"] = int64(0), int64(0) + case "options": + opts := make([]string, 0, len(d.Options)) + for _, s := range d.Options { + opts = append(opts, fmt.Sprintf("%ds", s)) + } + patch["durations"] = jsonArray(opts) + patch["duration_min"], patch["duration_max"] = int64(0), int64(0) + case "range": + patch["duration_min"], patch["duration_max"] = d.Min, d.Max + } + } + return patch +} + +func applySpec(m *model.ModelConfig, patch map[string]any) { + if v, ok := patch["resolutions"]; ok { + m.Resolutions = datatypes.JSON(mustJSON(v)) + } + if v, ok := patch["durations"]; ok { + m.Durations = datatypes.JSON(mustJSON(v)) + } + if v, ok := patch["duration_min"]; ok { + m.DurationMin = int(v.(int64)) + } + if v, ok := patch["duration_max"]; ok { + m.DurationMax = int(v.(int64)) + } + m.MaxReferenceImages = int(specInt(patch, "max_reference_images")) + m.MaxReferenceVideos = int(specInt(patch, "max_reference_videos")) + m.MaxReferenceAudios = int(specInt(patch, "max_reference_audios")) + m.MaxRefVideoSeconds = int(specInt(patch, "max_ref_video_seconds")) + m.MaxRefAudioSeconds = int(specInt(patch, "max_ref_audio_seconds")) + if v, ok := patch["reference_mode"].(string); ok { + m.ReferenceMode = v + } + if v, ok := patch["image_to_image"].(bool); ok { + m.ImageToImage = v + } +} + +func specInt(patch map[string]any, key string) int64 { + switch v := patch[key].(type) { + case int64: + return v + case int: + return int64(v) + default: + return 0 + } +} + +func mustJSON(v any) []byte { + b, _ := json.Marshal(v) + return b +} + +// ChannelsForModel returns enabled channels bound to the model with their +// upstream model names — dispatch order is decided by the caller. +func (s *ChannelService) ChannelsForModel(ctx context.Context, modelID string) ([]repo.ChannelWithUpstream, error) { + return s.channels.ForModel(ctx, modelID) +} + +// ChannelsForModelAnyState includes disabled channels (admin pin tests). +func (s *ChannelService) ChannelsForModelAnyState(ctx context.Context, modelID string) ([]repo.ChannelWithUpstream, error) { + return s.channels.ForModelAnyState(ctx, modelID) +} + +// Decrypt exposes the upstream key for dispatch. +func (s *ChannelService) Decrypt(ch model.Channel) string { return s.decrypt(ch.APIKeyEnc) } + +// Upstream returns the shared protocol client (dispatch needs it for +// UploadRef/Generate/Download; admin handlers for balance probes). +func (s *ChannelService) Upstream() *upstream.Client { return s.up } + +// BumpSuccess forwards local stats updates to the repo. +func (s *ChannelService) BumpSuccess(ctx context.Context, channelID string, credits float64) { + _ = s.channels.BumpSuccess(ctx, channelID, credits) +} + +// --- AES-GCM key encryption --- + +const encPrefix = "enc:v1:" + +func (s *ChannelService) encrypt(plain string) (string, error) { + block, err := aes.NewCipher(s.secret) + if err != nil { + return "", err + } + gcm, err := cipher.NewGCM(block) + if err != nil { + return "", err + } + nonce := make([]byte, gcm.NonceSize()) + if _, err := io.ReadFull(rand.Reader, nonce); err != nil { + return "", err + } + return encPrefix + base64.StdEncoding.EncodeToString(gcm.Seal(nonce, nonce, []byte(plain), nil)), nil +} + +func (s *ChannelService) decrypt(stored string) string { + if !strings.HasPrefix(stored, encPrefix) { + return stored // plaintext rows (e.g. seeded by hand) + } + raw, err := base64.StdEncoding.DecodeString(strings.TrimPrefix(stored, encPrefix)) + if err != nil { + return "" + } + block, err := aes.NewCipher(s.secret) + if err != nil { + return "" + } + gcm, err := cipher.NewGCM(block) + if err != nil { + return "" + } + if len(raw) < gcm.NonceSize() { + return "" + } + plain, err := gcm.Open(nil, raw[:gcm.NonceSize()], raw[gcm.NonceSize():], nil) + if err != nil { + return "" + } + return string(plain) +} diff --git a/backend/internal/service/maintenance.go b/backend/internal/service/maintenance.go index 10c825a..22ae379 100644 --- a/backend/internal/service/maintenance.go +++ b/backend/internal/service/maintenance.go @@ -5,25 +5,19 @@ import ( "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. +// MaintenanceService runs the periodic self-healing sweep. Without it stale +// pending events permanently block a user's generation gate, and old media / +// logs / temp objects accumulate unbounded. Token-pool and cookie-refresh +// upkeep is gone — channels are stateless upstream credentials. type MaintenanceService struct { - tokens *repo.TokenRepository - tokenSvc *TokenService events *repo.EventRepository users *repo.UserRepository - refresh *RefreshProfileService settings *repo.SiteSettingRepository store *storage.Client inflight *InflightRegistry @@ -36,20 +30,17 @@ type MaintenanceService struct { 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, orders *repo.OrderRepository) *MaintenanceService { +func NewMaintenanceService(events *repo.EventRepository, users *repo.UserRepository, settings *repo.SiteSettingRepository, store *storage.Client, inflight *InflightRegistry, showcase *repo.ShowcaseRepository, orders *repo.OrderRepository) *MaintenanceService { return &MaintenanceService{ - tokens: tokens, - tokenSvc: tokenSvc, - events: events, - users: users, - refresh: refresh, - settings: settings, - store: store, - inflight: inflight, - showcase: showcase, - orders: orders, - interval: 60 * time.Second, - stalePending: 600 * time.Second, + events: events, + users: users, + settings: settings, + store: store, + inflight: inflight, + showcase: showcase, + orders: orders, + interval: 60 * time.Second, + stalePending: 600 * time.Second, // 视频比图片慢得多，单条执行预算就是 30 分钟（videoGenBudget），用图片的 // 10 分钟去扫会在它出片前判死并退款。 stalePendingVid: 30 * time.Minute, @@ -63,10 +54,6 @@ func NewMaintenanceService(tokens *repo.TokenRepository, tokenSvc *TokenService, func (m *MaintenanceService) Run(ctx context.Context) { ticker := time.NewTicker(m.interval) defer ticker.Stop() - // The leonardo session keep-alive gets its own loop: one tick() can take tens - // of minutes (219 adobe cookie profiles alone), which would stretch a 5-minute - // keep-alive to the tick's real duration. - go m.runLeonardoKeepalive(ctx) m.tick(ctx) for { select { @@ -78,55 +65,6 @@ func (m *MaintenanceService) Run(ctx context.Context) { } } -// runLeonardoKeepalive re-checks every minute which leonardo accounts are due for -// a session renewal (the 5-minute due-ness itself is read per account from the DB, -// so a restart can't skip one). -func (m *MaintenanceService) runLeonardoKeepalive(ctx context.Context) { - if m.tokenSvc == nil { - return - } - ticker := time.NewTicker(time.Minute) - defer ticker.Stop() - for { - m.tokenSvc.RefreshLeonardoSessions(ctx) - select { - case <-ctx.Done(): - return - case <-ticker.C: - } - } -} - -// 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) { // 0. Auto-cancel unpaid recharge orders past their 30-min TTL. if m.orders != nil { @@ -137,74 +75,7 @@ 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 marker passes the - // token can only 401, so flip it to disabled+dead proactively instead of - // leaving a doomed account "active". Grok is intentionally NOT swept here: - // its reset marker is billingPeriodEnd (a credits-renewal date), not a - // death deadline — a grok sso keeps working past billingPeriodEnd, so - // expiring on it kills live accounts. Grok death is caught for real by the - // import-time FetchSession check and by marking dead on a 401 at use. - for _, pool := range []string{"runway"} { - if n, err := m.tokens.ExpireByReset(ctx, pool); err != nil { - log.Printf("maintenance: expire_%s: %v", pool, err) - } else if n > 0 { - log.Printf("maintenance: expired %d %s token(s)", n, pool) - } - } - - // 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) - // 1e. Re-validate grok accounts: an empty /rest/subscriptions (or 401) - // means the membership lapsed → disable+dead; otherwise re-sync the - // credits balance and 恢复时间 (from the credits' weekly reset). - m.tokenSvc.RefreshGrokLiveness(ctx) - } - - // 2. Auto-renew Adobe cookies whose refresh interval has elapsed. - if m.refresh != nil { - if m.settings != nil { - if proxy, err := m.settings.GetValue(ctx, "proxy.url"); err == nil && proxy != "" { - m.refresh.SetProxy(proxy) - } - } - 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 + // 1. 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, m.stalePendingVid); err != nil { @@ -219,13 +90,6 @@ func (m *MaintenanceService) tick(ctx context.Context) { 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 } @@ -248,10 +112,11 @@ func (m *MaintenanceService) tick(ctx context.Context) { 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. + // 2. Enforce the admin-configured log retention window. m.pruneLogs(ctx) - // 5. Enforce the media retention window. Runs every 60s like the log prune; + // 3. Enforce the media retention window + sweep expired tmp/ (v1 API job + // results) and stale refs/ objects. 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) @@ -274,10 +139,11 @@ func (m *MaintenanceService) pruneMedia(ctx context.Context) { return } days := m.retentionDays(ctx, "media.retention_days") - if days <= 0 { - return - } cutoff := time.Now().Add(-time.Duration(days) * 24 * time.Hour) + // tmp/ holds /v1 job results: presigned for 30 min, pruned shortly after. + tmpCutoff := time.Now().Add(-40 * time.Minute) + // /refs/ holds uploaded reference material: live for the presign TTL. + refCutoff := time.Now().Add(-2 * PresignTTL) objs, err := m.store.List(ctx, "") if err != nil { log.Printf("maintenance: list media: %v", err) @@ -303,10 +169,20 @@ func (m *MaintenanceService) pruneMedia(ctx context.Context) { removed, skipped := 0, 0 var clearedKeys []string for _, o := range objs { - if !o.LastModified.Before(cutoff) { + key := strings.TrimLeft(o.Key, "/") + expired := false + switch { + case strings.HasPrefix(key, "tmp/"): + expired = o.LastModified.Before(tmpCutoff) + case strings.Contains(key, "/refs/"): + expired = o.LastModified.Before(refCutoff) + case days > 0: + expired = o.LastModified.Before(cutoff) + } + if !expired { continue } - if _, ok := pinned[strings.TrimLeft(o.Key, "/")]; ok { + if _, ok := pinned[key]; ok { skipped++ continue } diff --git a/backend/internal/service/refresh_profiles.go b/backend/internal/service/refresh_profiles.go deleted file mode 100644 index e0b88bd..0000000 --- a/backend/internal/service/refresh_profiles.go +++ /dev/null @@ -1,235 +0,0 @@ -package service - -import ( - "context" - "errors" - "strings" - "time" - - "backend/internal/model" - "backend/internal/provider/adobe" - "backend/internal/repo" - "gorm.io/datatypes" -) - -type RefreshProfileService struct { - profiles *repo.RefreshProfileRepository - tokens *repo.TokenRepository - adobe *adobe.Client -} - -func NewRefreshProfileService(profiles *repo.RefreshProfileRepository, tokens *repo.TokenRepository, adobeClient *adobe.Client) *RefreshProfileService { - return &RefreshProfileService{ - profiles: profiles, - tokens: tokens, - adobe: adobeClient, - } -} - -func (s *RefreshProfileService) SetProxy(proxy string) { - if s.adobe != nil { - s.adobe.SetProxy(proxy) - } -} - -func (s *RefreshProfileService) List(ctx context.Context) ([]model.RefreshProfile, error) { - return s.profiles.List(ctx) -} - -func (s *RefreshProfileService) RefreshNow(ctx context.Context, id string) error { - if s.adobe == nil || s.tokens == nil { - return errors.New("refresh client not configured") - } - profile, err := s.profiles.Get(ctx, id) - if err != nil { - return err - } - if profile.Pool != "adobe" || profile.Kind != "adobe_cookie" { - return errors.New("unsupported refresh profile") - } - - now := time.Now() - _, _ = s.profiles.Update(ctx, id, map[string]any{ - "last_attempt_at": now, - }) - - result, err := s.adobe.ExchangeCookie(ctx, profile.Cookie) - if err != nil { - failures := profile.ConsecutiveFailures + 1 - // Exponential backoff: 60s per consecutive failure, capped at 1h. - secs := 60 * failures - if secs > 3600 { - secs = 3600 - } - msg := err.Error() - if len(msg) > 300 { - msg = msg[:300] - } - _, _ = s.profiles.Update(ctx, id, map[string]any{ - "last_error": msg, - "consecutive_failures": failures, - "next_retry_at": now.Add(time.Duration(secs) * time.Second), - }) - // Mark the token dead when the cookie exchange keeps failing: - // - 5 consecutive failures → cookie is likely expired/revoked - // - ride_AdobeID_acct_actreq (Adobe requires account action) → permanently broken - if failures >= 5 || strings.Contains(msg, "ride_AdobeID_acct_actreq") { - _, _ = s.tokens.Update(ctx, "adobe", id, map[string]any{ - "status": "disabled", - "dead": true, - }) - } - return err - } - - tokenPatch := map[string]any{ - "value": result.AccessToken, - "status": "active", - "dead": false, - "fails": 0, - "updated_at": now, - } - email, _ := parseJWTEmailExpiry(result.AccessToken) - if email != "" { - tokenPatch["account_email"] = email - } - if profile.IntervalSeconds > 0 { - tokenPatch["cached_quota_reset_after"] = now.Add(time.Duration(profile.IntervalSeconds) * time.Second).Format(time.RFC3339) - } else { - tokenPatch["cached_quota_reset_after"] = now.Add(54000 * time.Second).Format(time.RFC3339) - } - if profileData, profileErr := s.adobe.FetchAccountProfile(ctx, result.AccessToken); profileErr == nil { - if email := strings.TrimSpace(stringValue(profileData["email"])); email != "" { - tokenPatch["account_email"] = email - } - if displayName := strings.TrimSpace(stringValue(profileData["display_name"])); displayName != "" { - tokenPatch["account_display_name"] = displayName - } - } - planKnown := false - if quotaData, quotaErr := s.adobe.FetchCreditsBalance(ctx, result.AccessToken); quotaErr == nil { - meta := datatypes.JSONMap{ - "cached_quota_at": int(time.Now().Unix()), - } - if remaining, ok := quotaData["remaining"].(int); ok { - meta["cached_quota_remaining"] = remaining - } - if used, ok := quotaData["used"].(int); ok { - meta["cached_quota_used"] = used - } - if total, ok := quotaData["total"].(int); ok { - meta["cached_quota_total"] = total - } - tokenPatch["meta"] = meta - - planCap := strings.ToLower(strings.TrimSpace(stringValue(quotaData["plan"]))) - isVIP := planCap != "" && !strings.EqualFold(planCap, "free") - planKnown = planCap != "" - // 额度打到 0 的 VIP(母号/子号) 视为普号：写死 plan=free，额度恢复(>0)后下次刷新自动还原真实身份。 - if isVIP { - if rem, ok := quotaData["remaining"].(int); ok && rem <= 0 { - isVIP = false - planCap = "free" - } - } - meta["plan"] = planCap - - // VIP: use Adobe's reset time; set concurrency 5. - // 母号/子号 身份固定，不随积分动态变化——只有降级普号才刷新身份。 - if isVIP { - if resetAfter := strings.TrimSpace(stringValue(quotaData["available_until"])); resetAfter != "" { - tokenPatch["cached_quota_reset_after"] = resetAfter - } - tokenPatch["concurrency"] = 5 - // 身份固定：这里重建了整个 meta，若不带上 is_sub_account，刷新会把导入时 - // 确定的 母号/子号 标记冲掉。身份只在导入时确定，刷新只沿用、绝不靠积分猜： - // 已有标记就带过来；取不到（没 meta / 缺字段）说明身份未知，直接置死号。 - if existing, gerr := s.tokens.Get(ctx, "adobe", id); gerr == nil { - if v, ok := existing.Meta["is_sub_account"]; ok { - meta["is_sub_account"] = v - } - } - if _, ok := meta["is_sub_account"]; !ok { - tokenPatch["status"] = "disabled" - tokenPatch["dead"] = true - } - } else { - // Free account: concurrency 1, limit image+video - tokenPatch["concurrency"] = 1 - tokenPatch["image_limited"] = true - tokenPatch["video_limited"] = true - meta["is_sub_account"] = false - } - } - // 探测不到会员身份（credits 接口失败或没返回 plan）的号既不算普号也不算会员号， - // 留在池里只会在需要会员的请求上撞 403 —— 直接置死号。 - if !planKnown { - tokenPatch["status"] = "disabled" - tokenPatch["dead"] = true - } - if _, err := s.tokens.Update(ctx, "adobe", id, tokenPatch); err != nil { - return err - } - - interval := profile.IntervalSeconds - if interval <= 0 { - interval = 54000 - } - _, err = s.profiles.Update(ctx, id, map[string]any{ - "last_success_at": now, - "next_retry_at": now.Add(time.Duration(interval) * time.Second), - "last_error": "", - "consecutive_failures": 0, - }) - return err -} - -// RefreshDue refreshes every enabled profile whose next_retry_at has passed. -// Driven by the background maintenance loop so Adobe cookies auto-renew without -// an admin clicking "refresh". Individual failures are recorded on the profile -// (backoff + dead escalation) and don't abort the sweep. -func (s *RefreshProfileService) RefreshDue(ctx context.Context) (int, error) { - if s.adobe == nil || s.tokens == nil { - return 0, nil - } - due, err := s.profiles.ListDue(ctx, time.Now()) - if err != nil { - return 0, err - } - refreshed := 0 - for _, p := range due { - if p.Pool != "adobe" || p.Kind != "adobe_cookie" { - continue - } - if err := s.RefreshNow(ctx, p.ID); err != nil { - continue - } - refreshed++ - } - return refreshed, nil -} - -func (s *RefreshProfileService) Update(ctx context.Context, id string, body map[string]any) (*model.RefreshProfile, error) { - patch := map[string]any{} - if raw, ok := body["enabled"]; ok { - patch["enabled"] = boolValueWithDefault(raw, false) - } - if raw, ok := body["name"]; ok { - patch["name"] = stringValue(raw) - } - if raw, ok := body["interval_seconds"]; ok { - n := intValue(raw) - if n <= 0 { - return nil, errors.New("interval_seconds must be positive") - } - patch["interval_seconds"] = n - } - if len(patch) == 0 { - return s.profiles.Get(ctx, id) - } - return s.profiles.Update(ctx, id, patch) -} - -func (s *RefreshProfileService) Delete(ctx context.Context, id string) error { - return s.profiles.Delete(ctx, id) -} diff --git a/backend/internal/service/tokens.go b/backend/internal/service/tokens.go deleted file mode 100644 index 29c0791..0000000 --- a/backend/internal/service/tokens.go +++ /dev/null @@ -1,2239 +0,0 @@ -package service - -import ( - "context" - "encoding/base64" - "encoding/json" - "errors" - "fmt" - "log" - "strings" - "sync" - "sync/atomic" - "time" - - "backend/internal/model" - "backend/internal/provider/adobe" - "backend/internal/provider/chatgpt" - "backend/internal/provider/creativefabrica" - "backend/internal/provider/grok" - "backend/internal/provider/imagine" - "backend/internal/provider/krea" - "backend/internal/provider/leonardo" - "backend/internal/provider/runway" - "backend/internal/repo" - - "gorm.io/datatypes" - "gorm.io/gorm" -) - -var validTokenPools = map[string]string{ - "chatgpt": "openai", - "adobe": "adobe", - "runway": "runway", - "leonardo": "leonardo", - "krea": "krea", - "imagine": "imagine", - "grok": "grok", - "custom": "custom", - "creativefabrica": "creativefabrica", -} - -type TokenService struct { - tokens *repo.TokenRepository - refresh *repo.RefreshProfileRepository - events *repo.EventRepository - settings *repo.SiteSettingRepository - adobe *adobe.Client - chatgpt *chatgpt.Client - runway *runway.Client - leonardo *leonardo.Client - krea *krea.Client - imagine *imagine.Client - grok *grok.Client - cf *creativefabrica.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. - sem chan struct{} - // kreaActivating guards the once-per-day krea /app activation sweep so the 60s - // maintenance tick can't pile up overlapping sweeps. - kreaActivating atomic.Bool - // leonardoKeeping guards the leonardo session keep-alive sweep so the 60s - // maintenance tick can't pile up overlapping sweeps. - leonardoKeeping 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, grokClient *grok.Client, cfClient *creativefabrica.Client) *TokenService { - return &TokenService{ - tokens: tokens, - refresh: refresh, - events: events, - settings: settings, - adobe: adobeClient, - chatgpt: chatGPTClient, - runway: runwayClient, - leonardo: leonardoClient, - krea: kreaClient, - imagine: imagineClient, - grok: grokClient, - cf: cfClient, - sem: make(chan struct{}, 10), - } -} - -// applyProxy snapshots the configured outbound proxy onto a provider client right -// before an upstream call. Adobe/ChatGPT/Runway/Leonardo all egress through it; -// without this the import/quota probes would dial from the bare server IP (which -// Leonardo rate-limits with a 429). -func (s *TokenService) applyProxy(ctx context.Context) { - if s.settings == nil { - return - } - proxy, err := s.settings.GetValue(ctx, "proxy.url") - if err != nil { - return - } - if s.leonardo != nil { - s.leonardo.SetProxy(proxy) - } - if s.krea != nil { - s.krea.SetProxy(proxy) - } - if s.imagine != nil { - s.imagine.SetProxy(proxy) - } - if s.grok != nil { - s.grok.SetProxy(proxy) - } - if s.cf != nil { - s.cf.SetProxy(proxy) - } -} - -// RefreshExpiringTokens proactively renews krea/imagine sessions ~10min before -// the access token expires (the providers' refreshLeadSeconds gate), so a dormant -// account's rotating refresh_token never lapses. A dead token can't be recovered -// and — for krea — also means the daily free-credit meter can't be re-created, so -// keeping it perpetually fresh is what lets额度 auto-recover each day. Called by -// the maintenance sweep; refresh only hits the network for near-expiry accounts. -func (s *TokenService) RefreshExpiringTokens(ctx context.Context) { - items, err := s.tokens.List(ctx) - if err != nil { - return - } - s.applyProxy(ctx) - for i := range items { - it := items[i] - if it.Dead || it.Status == "disabled" { - continue - } - switch it.Pool { - case "krea": - if s.krea == nil { - continue - } - if _, rerr := kreaRefreshAndPersist(ctx, s.krea, s.tokens, it.ID, it.Value); rerr != nil && errors.Is(rerr, krea.ErrAuth) { - _, _ = s.tokens.Update(ctx, "krea", it.ID, map[string]any{"status": "disabled", "dead": true}) - } - case "imagine": - if s.imagine == nil { - continue - } - if _, rerr := imagineRefreshAndPersist(ctx, s.imagine, s.tokens, it.ID, it.Value); rerr != nil && errors.Is(rerr, imagine.ErrAuth) { - _, _ = s.tokens.Update(ctx, "imagine", it.ID, map[string]any{"status": "disabled", "dead": true}) - } - } - } -} - -// leonardoKeepaliveEvery is how long a leonardo cookie may go unrenewed. The due -// check reads each account's own meta["session_kept_at"], so the cadence survives -// a restart (no in-memory timer to lose) — after a restart every account whose -// stamp is older than this is simply due again. -const leonardoKeepaliveEvery = 5 * time.Minute - -// leonardoKeptAtKey stamps the last successful session keep-alive. -const leonardoKeptAtKey = "session_kept_at" - -// RefreshLeonardoSessions re-mints a session for every live leonardo account whose -// last keep-alive is older than leonardoKeepaliveEvery, instead of waiting for the -// ~1h bearer cache to lapse or for the account to be used. get-session rolls the -// better-auth session (expiresAt is pushed out) and may rotate CF_Access_Token / -// session_data, so the rotated cookie is written back. It never disables an -// account: a dead cookie is left to the quota-refresh strike logic, which -// double-checks before killing. -func (s *TokenService) RefreshLeonardoSessions(ctx context.Context) { - if s.leonardo == nil { - return - } - if !s.leonardoKeeping.CompareAndSwap(false, true) { - return // a sweep is already running - } - bg := context.WithoutCancel(ctx) - go func() { - defer s.leonardoKeeping.Store(false) - items, err := s.tokens.ListByPool(bg, "leonardo") - if err != nil { - return - } - due := time.Now().Add(-leonardoKeepaliveEvery).Unix() - proxied := false - renewed, failed := 0, 0 - for i := range items { - a := items[i] - if a.Dead || a.Status == "disabled" || strings.TrimSpace(a.Value) == "" { - continue - } - if at, ok := jsonMapInt(a.Meta, leonardoKeptAtKey); ok && int64(at) > due { - continue // renewed less than leonardoKeepaliveEvery ago - } - if !proxied { - s.applyProxy(bg) - proxied = true - } - callCtx, cancel := context.WithTimeout(bg, 90*time.Second) - _, err := s.leonardo.ProbeSession(callCtx, a.Value) - if err != nil { - cancel() - failed++ - if errors.Is(err, leonardo.ErrAuth) { - log.Printf("leonardo %s: session keepalive auth failure (%v)", a.ID, err) - } - continue // leave the stamp alone so the next tick retries - } - if fresh, ok := s.leonardo.RotatedCookie(a.Value); ok && strings.TrimSpace(fresh) != "" { - _, _ = s.tokens.SwapValue(callCtx, "leonardo", a.ID, a.Value, fresh) - } - fields := map[string]any{} - meta := cloneJSONMap(a.Meta) - meta[leonardoKeptAtKey] = int(time.Now().Unix()) - fields["meta"] = meta - _, _ = s.tokens.Update(callCtx, "leonardo", a.ID, fields) - cancel() - renewed++ - } - if renewed > 0 || failed > 0 { - log.Printf("leonardo: session keepalive renewed %d, failed %d", renewed, failed) - } - }() -} - -// ActivateKreaDue loads /app (Activate) for each krea account that hasn't been -// synced since the most recent daily reset, then re-syncs its balance. Krea only -// grants the daily free balance after the SSR app page loads, so without this an -// always-active account (one that never went 限额, hence never recovered) would -// read 0 / 402 after the reset. Runs in the background off the maintenance tick, -// guarded so sweeps never overlap, bounded concurrency to avoid a reset burst. -func (s *TokenService) ActivateKreaDue(ctx context.Context) { - if s.krea == nil { - return - } - if !s.kreaActivating.CompareAndSwap(false, true) { - return // a sweep is already running - } - bg := context.WithoutCancel(ctx) - go func() { - defer s.kreaActivating.Store(false) - items, err := s.tokens.ListByPool(bg, "krea") - if err != nil { - return - } - s.applyProxy(bg) - // Most recent UTC midnight (== last Beijing-08:00 reset). An account whose - // last sync (cached_quota_at) predates this hasn't been activated today. - lastReset := (time.Now().Unix() / 86400) * 86400 - sem := make(chan struct{}, 4) - var wg sync.WaitGroup - for i := range items { - it := items[i] - if it.Dead || it.Status == "disabled" || strings.TrimSpace(it.Value) == "" { - continue - } - if at, ok := jsonMapInt(it.Meta, "cached_quota_at"); ok && int64(at) >= lastReset { - continue // already activated/synced since the last reset - } - wg.Add(1) - sem <- struct{}{} - go func(it model.TokenAccount) { - defer wg.Done() - defer func() { <-sem }() - actx, cancel := context.WithTimeout(bg, 90*time.Second) - defer cancel() - s.krea.Activate(actx, it.Value) // load /app → grant daily balance - _, _ = s.Quota(actx, "krea", it.ID) - }(it) - } - wg.Wait() - }() -} - -func (s *TokenService) List(ctx context.Context) (map[string][]ginToken, error) { - items, err := s.tokens.List(ctx) - if err != nil { - return nil, err - } - out := map[string][]ginToken{} - for _, item := range items { - out[item.Pool] = append(out[item.Pool], ginToken{ - ID: item.ID, - ValuePreview: previewSecret(item.Value), - Status: item.Status, - Fails: item.Fails, - AddedAt: item.AddedAt, - }) - } - return out, nil -} - -func (s *TokenService) Add(ctx context.Context, pool, value, tokenID string) (*model.TokenAccount, error) { - pool = normalizePool(pool) - if pool == "" { - return nil, errors.New("unknown pool") - } - value = strings.TrimSpace(value) - if value == "" { - return nil, errors.New("pool and value required") - } - if tokenID == "" { - tokenID = newTokenID(pool) - } - return s.createToken(ctx, pool, tokenID, value, "active", nil) -} - -func (s *TokenService) ImportChatGPTToken(ctx context.Context, accessToken, tokenID string) (*model.TokenAccount, error) { - accessToken = strings.TrimSpace(accessToken) - if accessToken == "" { - return nil, errors.New("access_token required") - } - // Land as pending and return instantly; a background worker probes quota and - // flips the row active/dead (Python import_chatgpt_token). pending tokens are - // not schedulable — the pool only hands out status=="active". - meta := datatypes.JSONMap{"pending_check": true} - info := chatgpt.ExtractAccountInfo(accessToken) - _, exp := parseJWTEmailExpiry(accessToken) - // Identity is (pool, email): reuse the existing row for this email, else mint a - // fresh id — never trust the caller's id (it can collide with an unrelated row - // → a spurious 23505/400 for a brand-new account). - email := strings.TrimSpace(stringValue(info["email"])) - if existing, _ := s.tokens.GetByPoolEmail(ctx, "chatgpt", email); existing != nil { - tokenID = existing.ID - } else if email != "" || tokenID == "" { - tokenID = newTokenID("chatgpt") - } - item, err := s.createToken(ctx, "chatgpt", tokenID, accessToken, "pending", meta) - if err != nil { - if errors.Is(err, gorm.ErrDuplicatedKey) { - if item, err = s.tokens.Update(ctx, "chatgpt", tokenID, map[string]any{ - "value": accessToken, "status": "pending", "meta": meta, - }); err != nil { - return nil, err - } - } else { - return nil, err - } - } - // JWT-derived fields are free (no network) — hydrate them up front. OpenAI - // tokens carry the email under the nested "https://api.openai.com/profile" - // claim, so read it from ExtractAccountInfo — parseJWTEmailExpiry only sees - // top-level claims and returns "" for ChatGPT tokens. - patch := map[string]any{} - if email := strings.TrimSpace(stringValue(info["email"])); email != "" { - patch["account_email"] = email - } - if display := strings.TrimSpace(stringValue(info["plan_type"])); display != "" { - patch["account_display_name"] = display - } - if exp != nil { - patch["cached_quota_reset_after"] = exp.Format(time.RFC3339) - } - if len(patch) > 0 { - if updated, uerr := s.tokens.Update(ctx, "chatgpt", tokenID, patch); uerr == nil { - item = updated - } - } - go s.checkPendingChatGPT(tokenID, accessToken) - return item, nil -} - -// ImportRunwayToken lands a Runway JWT as a pending account and probes its -// credit balance off-thread (mirrors ImportChatGPTToken). The workspace/team id -// (= the JWT "id" claim) is stashed in meta["team_id"] so generation can send it -// as x-runway-workspace later. Recovery time == the JWT expiry. -func (s *TokenService) ImportRunwayToken(ctx context.Context, accessToken, tokenID string) (*model.TokenAccount, error) { - accessToken = strings.TrimSpace(strings.TrimPrefix(accessToken, "Bearer ")) - if accessToken == "" { - return nil, errors.New("access_token required") - } - if !runway.IsRunwayToken(accessToken) { - return nil, errors.New("not a runway token") - } - teamID := runway.TeamIDFromToken(accessToken) - // JWT-derived fields are free (no network); email + exp are top-level claims. - email, exp := parseJWTEmailExpiry(accessToken) - // Identity is (pool, email): reuse the existing row for this email, else mint a - // fresh id — never trust the caller's id (it can collide → spurious 400). - if existing, _ := s.tokens.GetByPoolEmail(ctx, "runway", email); existing != nil { - tokenID = existing.ID - } else if email != "" || tokenID == "" { - tokenID = newTokenID("runway") - } - meta := datatypes.JSONMap{"pending_check": true} - if teamID != "" { - meta["team_id"] = teamID - } - item, err := s.createToken(ctx, "runway", tokenID, accessToken, "pending", meta) - if err != nil { - if errors.Is(err, gorm.ErrDuplicatedKey) { - if item, err = s.tokens.Update(ctx, "runway", tokenID, map[string]any{ - "value": accessToken, "status": "pending", "meta": meta, - }); err != nil { - return nil, err - } - } else { - return nil, err - } - } - patch := map[string]any{} - if email != "" { - patch["account_email"] = email - } - if exp != nil { - patch["cached_quota_reset_after"] = exp.Format(time.RFC3339) - } - if len(patch) > 0 { - if updated, uerr := s.tokens.Update(ctx, "runway", tokenID, patch); uerr == nil { - item = updated - } - } - go s.checkPendingRunway(tokenID, accessToken) - return item, nil -} - -// ImportLeonardoCookie imports a Leonardo account. Unlike Adobe (which keeps a -// refresh profile to re-mint a bearer), Leonardo's stored credential IS the -// cookie — the bearer is derived on demand at generation time via get-session — -// so there's no refresh profile. Lands a pending row, then the worker validates -// the cookie + hydrates email/quota off-thread. -func (s *TokenService) ImportLeonardoCookie(ctx context.Context, cookie, tokenID string) (*model.TokenAccount, error) { - cookie = cleanAdobeCookie(cookie) // same paste-cleanup (JSON / "Cookie:" prefix) - if cookie == "" { - return nil, errors.New("cookie required") - } - if !leonardo.IsLeonardoCookie(cookie) { - return nil, errors.New("not a leonardo cookie") - } - // get-session 认的是 better-auth 的 session_data 缓存 cookie：只有 session_token - // 时它返回 200 null（换不到 accessToken），和死号一模一样，所以导入就拦掉。 - if !leonardo.HasSessionData(cookie) { - return nil, errors.New("leonardo cookie 缺少 __Secure-better-auth.session_data，请复制完整 cookie") - } - if tokenID == "" { - tokenID = newTokenID("leonardo") - } - meta := datatypes.JSONMap{"pending_check": true} - item, err := s.createToken(ctx, "leonardo", tokenID, cookie, "pending", meta) - if err != nil { - if errors.Is(err, gorm.ErrDuplicatedKey) { - item, err = s.tokens.Update(ctx, "leonardo", tokenID, map[string]any{ - "value": cookie, - "status": "pending", - "meta": meta, - }) - if err != nil { - return nil, err - } - } else { - return nil, err - } - } - // Fill 恢复时间 synchronously at import (next 08:00 Beijing) so it appears - // alongside 创建时间 immediately — not only after the background quota probe. - if updated, uerr := s.tokens.Update(ctx, "leonardo", tokenID, map[string]any{ - "cached_quota_reset_after": leonardoResetAfter(""), - }); uerr == nil { - item = updated - } - go s.checkPendingLeonardo(tokenID, cookie) - return item, nil -} - -// persistLeonardoCookie writes back the account cookie when Leonardo rotated its -// better-auth session_data cache — that cache is what authenticates get-session, -// so keeping a stale copy would eventually look like a dead account. -func (s *TokenService) persistLeonardoCookie(ctx context.Context, tokenID, cookie string) { - if s.leonardo == nil { - return - } - if fresh, ok := s.leonardo.RotatedCookie(cookie); ok && strings.TrimSpace(fresh) != "" { - _, _ = s.tokens.SwapValue(ctx, "leonardo", tokenID, cookie, fresh) - } -} - -// checkPendingLeonardo validates a freshly imported Leonardo cookie off-thread: -// get-session must succeed (else the cookie is dead → disabled), then it hydrates -// email/display-name + the token balance and the daily renewal time (so the -// maintenance sweep can auto-recover a 限额 account). -func (s *TokenService) checkPendingLeonardo(tokenID, cookie string) { - defer func() { - if r := recover(); r != nil { - log.Printf("token import: leonardo 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.leonardo == nil { - s.finishPending(ctx, "leonardo", tokenID, "active", false, nil) - return - } - s.applyProxy(ctx) - data, err := s.leonardo.FetchCreditsBalance(ctx, cookie) - s.persistLeonardoCookie(ctx, tokenID, cookie) - if err != nil { - if errors.Is(err, leonardo.ErrAuth) { - s.finishPending(ctx, "leonardo", tokenID, "disabled", true, nil) - return - } - // network/proxy blip — benefit of the doubt, activate. - s.finishPending(ctx, "leonardo", tokenID, "active", false, nil) - return - } - seed := map[string]any{} - if em := strings.TrimSpace(stringValue(data["email"])); em != "" { - seed["account_email"] = em - } - if dn := strings.TrimSpace(stringValue(data["display_name"])); dn != "" { - seed["account_display_name"] = dn - } - // Always fill 恢复时间: use upstream's renewal time if present, else the next - // daily reset (08:00 Beijing == next UTC midnight). - seed["cached_quota_reset_after"] = leonardoResetAfter(stringValue(data["available_until"])) - if len(seed) > 0 { - _, _ = s.tokens.Update(ctx, "leonardo", tokenID, seed) - } - 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 uid := strings.TrimSpace(stringValue(data["user_id"])); uid != "" { - quotaMeta["user_id"] = uid - } - s.finishPending(ctx, "leonardo", tokenID, "active", false, quotaMeta) -} - -// ImportKreaCookie imports a Krea account. Like Leonardo the stored credential -// IS the cookie (Supabase session); quota/generation forward it directly. -func (s *TokenService) ImportKreaCookie(ctx context.Context, cookie, tokenID string) (*model.TokenAccount, error) { - cookie = cleanAdobeCookie(cookie) // same paste-cleanup (JSON / "Cookie:" prefix) - if cookie == "" { - return nil, errors.New("cookie required") - } - if !krea.IsKreaCookie(cookie) { - return nil, errors.New("not a krea cookie") - } - // Identity is (pool, email), NOT the caller-supplied id — a colliding id from - // the upstream importer would otherwise raise a spurious 23505/400 for a brand - // new account. Reuse the existing row for this email; else mint a fresh unique - // id (never trust the caller's id for a new row). - email := krea.EmailFromCookie(cookie) - if existing, _ := s.tokens.GetByPoolEmail(ctx, "krea", email); existing != nil { - tokenID = existing.ID - } else { - tokenID = newTokenID("krea") - } - meta := datatypes.JSONMap{"pending_check": true} - item, err := s.createToken(ctx, "krea", tokenID, cookie, "pending", meta) - if err != nil { - if errors.Is(err, gorm.ErrDuplicatedKey) { - item, err = s.tokens.Update(ctx, "krea", tokenID, map[string]any{ - "value": cookie, - "status": "pending", - "meta": meta, - }) - if err != nil { - return nil, err - } - } else { - return nil, err - } - } - // Email is in the cookie (no network) — hydrate up front, plus 恢复时间. - seed := map[string]any{"cached_quota_reset_after": leonardoResetAfter("")} - if email != "" { - seed["account_email"] = email - } - if updated, uerr := s.tokens.Update(ctx, "krea", tokenID, seed); uerr == nil { - item = updated - } - go s.checkPendingKrea(tokenID, cookie) - return item, nil -} - -// checkPendingKrea validates a freshly imported Krea cookie off-thread and -// hydrates the credit balance. A 401 from billing-data → the cookie is dead. -func (s *TokenService) checkPendingKrea(tokenID, cookie string) { - defer func() { - if r := recover(); r != nil { - log.Printf("token import: krea 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.krea == nil { - s.finishPending(ctx, "krea", tokenID, "active", false, nil) - return - } - s.applyProxy(ctx) - cookie, rerr := kreaRefreshAndPersist(ctx, s.krea, s.tokens, tokenID, cookie) - if rerr != nil { - if errors.Is(rerr, krea.ErrAuth) { - s.finishPending(ctx, "krea", tokenID, "disabled", true, nil) - return - } - s.finishPending(ctx, "krea", tokenID, "active", false, nil) - return - } - data, err := s.krea.FetchCreditsBalance(ctx, cookie) - if err != nil { - if errors.Is(err, krea.ErrAuth) { - s.finishPending(ctx, "krea", tokenID, "disabled", true, nil) - return - } - s.finishPending(ctx, "krea", tokenID, "active", false, nil) - return - } - if em := strings.TrimSpace(stringValue(data["email"])); em != "" { - _, _ = s.tokens.Update(ctx, "krea", tokenID, map[string]any{"account_email": em}) - } - quotaMeta := map[string]any{} - if rem, ok := data["remaining"].(int); ok { - quotaMeta["cached_quota_remaining"] = rem - quotaMeta["cached_quota_at"] = int(time.Now().Unix()) - } - s.finishPending(ctx, "krea", tokenID, "active", false, quotaMeta) -} - -// ImportImagineToken imports an Imagine.art account. The stored credential IS the -// JSON {"token","refreshToken"}; quota/generation forward it (refreshing the -// access token from the refreshToken when expired). -func (s *TokenService) ImportImagineToken(ctx context.Context, cred, tokenID string) (*model.TokenAccount, error) { - cred = strings.TrimSpace(cred) - if cred == "" { - return nil, errors.New("credential required") - } - if !imagine.IsImagineToken(cred) { - return nil, errors.New("not an imagine token") - } - // Identity is (pool, email=userId), NOT the caller-supplied id — reuse the - // existing row for this account; else mint a fresh unique id. - email := imagine.EmailFromCred(cred) - if existing, _ := s.tokens.GetByPoolEmail(ctx, "imagine", email); existing != nil { - tokenID = existing.ID - } else { - tokenID = newTokenID("imagine") - } - meta := datatypes.JSONMap{"pending_check": true} - item, err := s.createToken(ctx, "imagine", tokenID, cred, "pending", meta) - if err != nil { - if errors.Is(err, gorm.ErrDuplicatedKey) { - item, err = s.tokens.Update(ctx, "imagine", tokenID, map[string]any{ - "value": cred, - "status": "pending", - "meta": meta, - }) - if err != nil { - return nil, err - } - } else { - return nil, err - } - } - // userId is in the token (no network) — hydrate up front, plus 恢复时间. - // Imagine free credits renew daily like Krea (next UTC midnight = 08:00 北京). - seed := map[string]any{"cached_quota_reset_after": leonardoResetAfter("")} - if email != "" { - seed["account_email"] = email - } - if updated, uerr := s.tokens.Update(ctx, "imagine", tokenID, seed); uerr == nil { - item = updated - } - go s.checkPendingImagine(tokenID, cred) - return item, nil -} - -// checkPendingImagine validates a freshly imported Imagine token off-thread and -// hydrates the credit balance. A 401 from /v1/credit → the token is dead. -func (s *TokenService) checkPendingImagine(tokenID, cred string) { - defer func() { - if r := recover(); r != nil { - log.Printf("token import: imagine 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.imagine == nil { - s.finishPending(ctx, "imagine", tokenID, "active", false, nil) - return - } - s.applyProxy(ctx) - cred, rerr := imagineRefreshAndPersist(ctx, s.imagine, s.tokens, tokenID, cred) - if rerr != nil { - if errors.Is(rerr, imagine.ErrAuth) { - s.finishPending(ctx, "imagine", tokenID, "disabled", true, nil) - return - } - s.finishPending(ctx, "imagine", tokenID, "active", false, nil) - return - } - data, err := s.imagine.FetchCreditsBalance(ctx, cred) - if err != nil { - if errors.Is(err, imagine.ErrAuth) { - s.finishPending(ctx, "imagine", tokenID, "disabled", true, nil) - return - } - s.finishPending(ctx, "imagine", tokenID, "active", false, nil) - return - } - if em := strings.TrimSpace(stringValue(data["email"])); em != "" { - _, _ = s.tokens.Update(ctx, "imagine", tokenID, map[string]any{"account_email": em}) - } - quotaMeta := map[string]any{} - if rem, ok := data["remaining"].(int); ok { - quotaMeta["cached_quota_remaining"] = rem - quotaMeta["cached_quota_at"] = int(time.Now().Unix()) - } - s.finishPending(ctx, "imagine", tokenID, "active", false, quotaMeta) -} - -func (s *TokenService) ImportAdobeCookie(ctx context.Context, cookie, tokenID string) (*model.TokenAccount, *model.RefreshProfile, error) { - cookie = cleanAdobeCookie(cookie) - if cookie == "" { - return nil, nil, errors.New("cookie required") - } - if tokenID == "" { - tokenID = newTokenID("adobe") - } - now := time.Now() - // Register the cookie refresh profile up front. Push next_retry_at out a full - // interval so the maintenance loop doesn't race the import worker on the first - // exchange — the worker below owns the initial hydrate. - nextRetry := now.Add(54000 * time.Second) - profile := &model.RefreshProfile{ - ID: tokenID, - Name: tokenID, - Pool: "adobe", - Kind: "adobe_cookie", - Cookie: cookie, - Enabled: true, - IntervalSeconds: 54000, - ImportedAt: &now, - NextRetryAt: &nextRetry, - CreatedAt: now, - UpdatedAt: now, - } - if err := s.refresh.Create(ctx, profile); err != nil && !errors.Is(err, gorm.ErrDuplicatedKey) { - return nil, nil, err - } - // Land a placeholder pending token (value filled in by the worker). NOT - // schedulable — the pool only hands out status=="active". The import returns - // instantly; the row flips active/dead once the worker finishes the three - // Adobe round-trips. Mirrors Python import_adobe_cookie. - meta := datatypes.JSONMap{"pending_check": true} - item, err := s.createToken(ctx, "adobe", tokenID, "", "pending", meta) - if err != nil { - if errors.Is(err, gorm.ErrDuplicatedKey) { - item, err = s.tokens.Update(ctx, "adobe", tokenID, map[string]any{ - "status": "pending", - "meta": meta, - }) - if err != nil { - return nil, nil, err - } - } else { - return nil, nil, err - } - } - go s.checkPendingAdobe(tokenID, cookie) - return item, profile, nil -} - -// checkPendingAdobe runs the three Adobe round-trips off-thread for a freshly -// imported cookie so the import request returns instantly (Python -// _check_pending_adobe). Step 1 (exchange) is authoritative — a bad/expired -// cookie can't mint a token, so failure marks the row dead. Steps 2-3 (credits / -// profile) are best-effort hydration. -func (s *TokenService) checkPendingAdobe(tokenID, cookie string) { - defer func() { - if r := recover(); r != nil { - log.Printf("token import: adobe pending check panicked for %s: %v", tokenID, r) - } - }() - s.sem <- struct{}{} - defer func() { <-s.sem }() - ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second) - defer cancel() - - if s.adobe == nil { - s.finishPending(ctx, "adobe", tokenID, "disabled", true, nil) - return - } - result, err := s.adobe.ExchangeCookie(ctx, cookie) - if err != nil { - s.finishPending(ctx, "adobe", tokenID, "disabled", true, nil) - _, _ = s.refresh.Update(ctx, tokenID, map[string]any{ - "last_attempt_at": time.Now(), - "last_error": err.Error(), - "consecutive_failures": 1, - }) - return - } - // Seed the real access token, then activate so the pool can schedule it. - seed := map[string]any{"value": result.AccessToken} - email, exp := parseJWTEmailExpiry(result.AccessToken) - if email != "" { - seed["account_email"] = email - } - if exp != nil { - seed["cached_quota_reset_after"] = exp.Format(time.RFC3339) - } - _, _ = s.tokens.Update(ctx, "adobe", tokenID, seed) - - quotaMeta := map[string]any{} - planKnown := false - if cb, e := s.adobe.FetchCreditsBalance(ctx, result.AccessToken); e == nil { - quotaMeta["cached_quota_at"] = int(time.Now().Unix()) - planCap := strings.ToLower(strings.TrimSpace(stringValue(cb["plan"]))) - isVIP := planCap != "" && planCap != "free" - planKnown = planCap != "" - - rem, remOK := cb["remaining"].(int) - if remOK { - quotaMeta["cached_quota_remaining"] = rem - } - // 额度打到 0 的 VIP(母号/子号) 视为普号：写死 plan=free，额度恢复(>0)后下次刷新自动还原真实身份。 - if isVIP && remOK && rem <= 0 { - isVIP = false - planCap = "free" - } - quotaMeta["plan"] = planCap - // VIP 号导入即固定 母号/子号 身份并显式写入 is_sub_account：积分 1..4000 视为子号， - // 其余（含 >4000）为母号(false)。母号必须显式写 false，否则 isVipMotherAccount - // 因字段缺失把它判为非母号，导致 Seedance 无号可调度。 - if isVIP { - quotaMeta["is_sub_account"] = rem > 0 && rem <= 4000 - } - - // Set concurrency + limits based on plan - concurrency := 1 - imageLimit := true - videoLimit := true - if isVIP { - concurrency = 5 - imageLimit = false - videoLimit = false - if ra := strings.TrimSpace(stringValue(cb["available_until"])); ra != "" { - _, _ = s.tokens.Update(ctx, "adobe", tokenID, map[string]any{"cached_quota_reset_after": ra}) - } - } - patch := map[string]any{ - "concurrency": concurrency, - "image_limited": imageLimit, - "video_limited": videoLimit, - "meta": quotaMeta, - } - if planKnown { - patch["status"] = "active" - patch["dead"] = false - } - _, _ = s.tokens.Update(ctx, "adobe", tokenID, patch) - } - if prof, e := s.adobe.FetchAccountProfile(ctx, result.AccessToken); e == nil { - p := map[string]any{} - if em := strings.TrimSpace(stringValue(prof["email"])); em != "" { - p["account_email"] = em - } - if dn := strings.TrimSpace(stringValue(prof["display_name"])); dn != "" { - p["account_display_name"] = dn - } - if len(p) > 0 { - _, _ = s.tokens.Update(ctx, "adobe", tokenID, p) - } - } - - // 探测不到会员身份（credits 接口失败或没返回 plan）的号既不算普号也不算会员号， - // 留在池里只会在需要会员的请求上撞 403 —— 直接置死号，等重新探测到 plan 再启用。 - if !planKnown { - log.Printf("token import: adobe %s plan unknown, marking dead", tokenID) - s.finishPending(ctx, "adobe", tokenID, "disabled", true, quotaMeta) - return - } - s.finishPending(ctx, "adobe", tokenID, "active", false, quotaMeta) - _, _ = s.refresh.Update(ctx, tokenID, map[string]any{ - "last_attempt_at": time.Now(), - "last_success_at": time.Now(), - "last_error": "", - "consecutive_failures": 0, - "next_retry_at": time.Now().Add(54000 * time.Second), - }) -} - -// checkPendingChatGPT probes a freshly imported ChatGPT token's quota off-thread -// (Python _check_pending_chatgpt). 401 → dead; a non-auth error gets the benefit -// of the doubt and activates so a transient blip can't sideline a good account. -func (s *TokenService) checkPendingChatGPT(tokenID, accessToken string) { - defer func() { - if r := recover(); r != nil { - log.Printf("token import: chatgpt 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.chatgpt == nil { - s.finishPending(ctx, "chatgpt", tokenID, "active", false, nil) - return - } - data, err := s.chatgpt.FetchImageQuota(ctx, accessToken) - if err != nil { - // network/proxy error — benefit of the doubt, activate. - s.finishPending(ctx, "chatgpt", tokenID, "active", false, nil) - return - } - if boolValueWithDefault(data["auth_failed"], false) { - s.finishPending(ctx, "chatgpt", tokenID, "disabled", true, nil) - return - } - rem, exhausted := chatgptRemaining(data) - quotaMeta := map[string]any{ - "cached_quota_remaining": rem, - "cached_quota_at": int(time.Now().Unix()), - } - // reset 时间:优先用 OpenAI 的 reset_after,缺失则默认次日重置,保证限额号能被 - // RecoverQuota 到点自动复活、重新探测,而不会永久搁置。 - reset := strings.TrimSpace(stringValue(data["reset_after"])) - if reset == "" { - reset = leonardoResetAfter("") - } - _, _ = s.tokens.Update(ctx, "chatgpt", tokenID, map[string]any{"cached_quota_reset_after": reset}) - // remaining<=0(0 / 负数 / 未知)→ 置「限额」,池子不再调度,到点自动恢复。 - status := "active" - if exhausted { - status = "quota" - } - s.finishPending(ctx, "chatgpt", tokenID, status, false, quotaMeta) -} - -// chatgptRemaining normalizes OpenAI's image_gen remaining: the raw rate-limit -// counter can go NEGATIVE on over-used accounts, and "—"(absent)means unknown — -// both clamp to 0, and 0 counts as exhausted (→ 限额). Returns (remaining≥0, -// exhausted). -func chatgptRemaining(data map[string]any) (int, bool) { - raw, ok := data["remaining"] - if !ok || raw == nil { - return 0, true - } - rem := intValue(raw) - if rem < 0 { - rem = 0 - } - return rem, rem <= 0 -} - -// checkPendingRunway probes a freshly imported Runway token's credit balance -// off-thread (mirrors checkPendingChatGPT). ErrAuth → dead; any other error -// gets the benefit of the doubt and activates so a transient blip can't sideline -// a good account. -func (s *TokenService) checkPendingRunway(tokenID, accessToken string) { - defer func() { - if r := recover(); r != nil { - log.Printf("token import: runway 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.runway == nil { - s.finishPending(ctx, "runway", tokenID, "active", false, nil) - return - } - data, err := s.runway.FetchCreditsBalance(ctx, accessToken) - if err != nil { - if errors.Is(err, runway.ErrAuth) { - s.finishPending(ctx, "runway", tokenID, "disabled", true, nil) - return - } - // network/proxy error — benefit of the doubt, activate. - s.finishPending(ctx, "runway", 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 - } - 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) - // Fully async, no dedup: every import mints a fresh row (a passed-in tokenID is - // an explicit edit → update). We do NOT look up an existing account by - // email/session_id, and we leave account_email empty — email, quota and - // recovery time are all filled off-thread by checkPendingGrok, which also - // disables the account if the sso session is dead. - if strings.TrimSpace(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 - } - } - 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) - // Validate the session first: a dead sso answers 200 {"status":"unauthenticated"}, - // which FetchSession now maps to ErrAuth → disable. Also backfill the email if - // import couldn't resolve it (empty account_email currently shows the session id). - if email, _, serr := s.grok.FetchSession(ctx, ssoToken); serr != nil { - if errors.Is(serr, grok.ErrAuth) { - s.finishPending(ctx, "grok", tokenID, "disabled", true, nil) - return - } - } else if strings.TrimSpace(email) != "" { - _, _ = s.tokens.Update(ctx, "grok", tokenID, map[string]any{"account_email": strings.TrimSpace(email)}) - } - // Console 没有额度接口，所以额度不查上游：导入即写死 图 5 / 视频 2，生成时各扣各的， - // 两个都归零就判死；没有恢复时间。 - s.finishPending(ctx, "grok", tokenID, "active", false, map[string]any{ - repo.GrokImageQuotaKey: repo.GrokImageQuota, - repo.GrokVideoQuotaKey: repo.GrokVideoQuota, - "cached_quota_at": int(time.Now().Unix()), - }) -} - -// RefreshGrokLiveness re-validates every live grok account each maintenance tick. -// Grok sso can't be renewed and has no reset-based death deadline (billingPeriodEnd -// is only a credits-renewal date — the sso keeps working past it), so liveness is -// probed directly: GET /rest/subscriptions. No ACTIVE entry — a lapsed membership -// flips to SUBSCRIPTION_STATUS_INACTIVE (empty array / 401 also count) — means the -// paid membership is gone → the account is disabled+dead. Otherwise the credits -// balance is re-synced and 恢复时间 is refreshed from the credits' own weekly -// reset (NOT the subscription's billing-period end). -func (s *TokenService) RefreshGrokLiveness(ctx context.Context) { - if s.grok == nil { - return - } - items, err := s.tokens.List(ctx) - if err != nil { - return - } - s.applyProxy(ctx) - for i := range items { - it := items[i] - if it.Pool != "grok" || it.Dead || it.Status == "disabled" || strings.TrimSpace(it.Value) == "" { - continue - } - // A grok sso can momentarily 401 (upstream blip / proxy hiccup / anti-bot) - // while still being fully valid, so a single auth failure must never kill a - // live account. Retry the subscription probe up to 3 times on 401/403; only - // if all 3 attempts still return ErrAuth do we treat the session as dead. - var sub *grok.Subscription - var serr error - for attempt := 1; attempt <= 3; attempt++ { - sub, serr = s.grok.FetchSubscription(ctx, it.Value) - if serr == nil || !errors.Is(serr, grok.ErrAuth) { - break - } - if attempt < 3 { - time.Sleep(2 * time.Second) - } - } - if serr != nil { - if errors.Is(serr, grok.ErrAuth) { - // 3 consecutive 401/403 → the sso session is genuinely dead. - _, _ = s.tokens.Update(ctx, "grok", it.ID, map[string]any{"status": "disabled", "dead": true}) - } - continue // other transient upstream error → leave as-is, retry next tick - } - if sub == nil || !sub.Member { - // no ACTIVE subscription (INACTIVE / empty) → membership lapsed → dead. - _, _ = s.tokens.Update(ctx, "grok", it.ID, map[string]any{"status": "disabled", "dead": true}) - continue - } - // 额度不查上游（Console 无额度接口）：只给老号补齐写死的 图 5 / 视频 2。 - if _, ok := jsonMapInt(it.Meta, repo.GrokImageQuotaKey); ok { - continue - } - meta := cloneJSONMap(it.Meta) - meta[repo.GrokImageQuotaKey] = repo.GrokImageQuota - meta[repo.GrokVideoQuotaKey] = repo.GrokVideoQuota - meta["cached_quota_at"] = int(time.Now().Unix()) - _, _ = s.tokens.Update(ctx, "grok", it.ID, map[string]any{"meta": meta}) - } -} - -// ImportCreativeFabricaCookie imports a Creative Fabrica session cookie the same -// way Adobe does (paste the whole Cookie header; JSON array/object accepted). -// The cookie IS the credential — a fresh short-lived JWT is minted on demand via -// /query/userAuth, so no RefreshProfile is registered (there is nothing to -// refresh). Accounts are one-shot: their coins buy exactly one generation, so a -// successful render disables the account (see generateCreativeFabricaVideo). -func (s *TokenService) ImportCreativeFabricaCookie(ctx context.Context, cookie, tokenID string) (*model.TokenAccount, error) { - cookie = cleanAdobeCookie(cookie) - if cookie == "" { - return nil, errors.New("cookie required") - } - if tokenID == "" { - tokenID = newTokenID("creativefabrica") - } - meta := datatypes.JSONMap{"pending_check": true} - item, err := s.createToken(ctx, "creativefabrica", tokenID, cookie, "pending", meta) - if err != nil { - if errors.Is(err, gorm.ErrDuplicatedKey) { - item, err = s.tokens.Update(ctx, "creativefabrica", tokenID, map[string]any{ - "status": "pending", - "meta": meta, - }) - if err != nil { - return nil, err - } - } else { - return nil, err - } - } - go s.checkPendingCreativeFabrica(tokenID, cookie) - return item, nil -} - -// checkPendingCreativeFabrica probes a freshly imported cookie off-thread: -// /query/userAuth must mint a token (else the cookie is dead), then best-effort -// balance hydration. Balance of exactly 0 means the account can't generate → -// dead. -func (s *TokenService) checkPendingCreativeFabrica(tokenID, cookie string) { - defer func() { - if r := recover(); r != nil { - log.Printf("token import: creativefabrica pending check panicked for %s: %v", tokenID, r) - } - }() - s.sem <- struct{}{} - defer func() { <-s.sem }() - ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second) - defer cancel() - - if s.cf == nil { - s.finishPending(ctx, "creativefabrica", tokenID, "disabled", true, nil) - return - } - if _, _, err := s.cf.ExchangeToken(ctx, cookie); err != nil { - log.Printf("token import: creativefabrica %s cookie failed to authenticate, marking dead: %v", tokenID, err) - s.finishPending(ctx, "creativefabrica", tokenID, "disabled", true, nil) - return - } - meta := map[string]any{} - // Best-effort profile hydration: the studio browser hits /query/user on - // every page load, so the email is free to grab here (avoids 邮箱列 showing - // the raw id). - if email, e := s.cf.FetchUser(ctx, cookie); e == nil && strings.TrimSpace(email) != "" { - if _, uerr := s.tokens.Update(ctx, "creativefabrica", tokenID, map[string]any{"account_email": strings.TrimSpace(email)}); uerr != nil { - log.Printf("token import: creativefabrica %s email write failed: %v", tokenID, uerr) - } - } - if bal, e := s.cf.FetchBalance(ctx, cookie); e == nil { - meta["cached_quota_remaining"] = bal - meta["cached_quota_at"] = int(time.Now().Unix()) - // One-shot accounts with zero coins left can't generate at all. - if bal <= 0 { - log.Printf("token import: creativefabrica %s has 0 coins, marking dead", tokenID) - s.finishPending(ctx, "creativefabrica", tokenID, "disabled", true, meta) - return - } - } - s.finishPending(ctx, "creativefabrica", tokenID, "active", false, meta) -} - -// ImportCustomAccount adds an upstream as a custom account: base_url + key, the -// csv list of model ids it serves (empty = all), plus optional weight and -// per-account concurrency. No probe — the account goes active immediately and is -// matched to custom models by id at generation time. Calls go direct (no proxy). -func (s *TokenService) ImportCustomAccount(ctx context.Context, baseURL, apiKey, models, name, protocol string, weight, concurrency int, tokenID string) (*model.TokenAccount, error) { - baseURL = strings.TrimRight(strings.TrimSpace(baseURL), "/") - apiKey = strings.TrimSpace(apiKey) - protocol = strings.ToLower(strings.TrimSpace(protocol)) - // Edit mode: tokenID points at an existing custom account. base_url required; - // a blank key keeps the stored one. - if strings.TrimSpace(tokenID) != "" { - existing, gerr := s.tokens.Get(ctx, "custom", tokenID) - if gerr != nil { - return nil, gerr - } - if baseURL == "" { - return nil, errors.New("base_url required") - } - proto := protocol - if proto == "" { - proto = strings.ToLower(strings.TrimSpace(stringValue(existing.Meta["protocol"]))) - } - if proto != "vivid" { - proto = "openai" - } - meta := datatypes.JSONMap{"base_url": baseURL, "protocol": proto} - if m := strings.TrimSpace(models); m != "" { - meta["models"] = m - } - patch := map[string]any{"meta": meta, "weight": weight, "concurrency": concurrency, "account_email": strings.TrimSpace(name)} - if apiKey != "" { - patch["value"] = apiKey - } - item, uerr := s.tokens.Update(ctx, "custom", tokenID, patch) - if uerr != nil { - return nil, uerr - } - _ = existing - return item, nil - } - if baseURL == "" || apiKey == "" { - return nil, errors.New("base_url and key required") - } - if protocol != "vivid" { - protocol = "openai" - } - meta := datatypes.JSONMap{"base_url": baseURL, "protocol": protocol} - if m := strings.TrimSpace(models); m != "" { - meta["models"] = m - } - tokenID = newTokenID("custom") - item, err := s.createToken(ctx, "custom", tokenID, apiKey, "active", meta) - if err != nil { - return nil, err - } - patch := map[string]any{} - if strings.TrimSpace(name) != "" { - patch["account_email"] = strings.TrimSpace(name) - } - if weight != 0 { - patch["weight"] = weight - } - if concurrency > 0 { - patch["concurrency"] = concurrency - } - if len(patch) > 0 { - if updated, uerr := s.tokens.Update(ctx, "custom", tokenID, patch); uerr == nil { - item = updated - } - } - return item, nil -} - -// leonardoAuthStrikeLimit 是 leonardo 号被判死前允许的连续鉴权失败次数。上游偶发 -// 返回 200 null / 401(cookie 轮换竞态、人机校验)时一次就判死会误杀健康号，所以要 -// 连续失败到这个次数才判死；任何一次成功都会清零。 -const leonardoAuthStrikeLimit = 3 - -// leonardoAuthStrike 记一次鉴权失败：返回要写回的 meta、当前连续失败次数,以及是否 -// 该判死。 -func leonardoAuthStrike(item *model.TokenAccount, reason string) (datatypes.JSONMap, int, bool) { - meta := cloneJSONMap(item.Meta) - strikes := 1 - if n, ok := jsonMapInt(item.Meta, "auth_fails"); ok { - strikes = n + 1 - } - meta["auth_fails"] = strikes - meta["last_auth_error"] = reason - return meta, strikes, strikes >= leonardoAuthStrikeLimit -} - -// 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) { - item, err := s.tokens.Get(ctx, pool, id) - if err != nil { - return - } - meta := cloneJSONMap(item.Meta) - meta["pending_check"] = false - for k, v := range quotaMeta { - meta[k] = v - } - patch := map[string]any{"status": status, "meta": meta} - if dead { - patch["dead"] = true - } - _, _ = s.tokens.Update(ctx, pool, id, patch) -} - -func (s *TokenService) Update(ctx context.Context, pool, id string, body map[string]any) (*model.TokenAccount, error) { - pool = normalizePool(pool) - if pool == "" { - return nil, errors.New("unknown pool") - } - patch := map[string]any{} - if raw, ok := body["status"]; ok { - status := normalizeTokenStatus(stringValue(raw)) - if status == "" { - return nil, errors.New("invalid status") - } - patch["status"] = status - if status == "active" { - patch["dead"] = false - if _, hasFails := body["fails"]; !hasFails { - patch["fails"] = 0 - } - } - } - if raw, ok := body["value"]; ok { - value := strings.TrimSpace(stringValue(raw)) - if value == "" { - return nil, errors.New("value cannot be empty") - } - patch["value"] = value - patch["dead"] = false - } - if raw, ok := body["fails"]; ok { - patch["fails"] = intValue(raw) - } - if len(patch) == 0 { - return s.tokens.Get(ctx, pool, id) - } - return s.tokens.Update(ctx, pool, id, patch) -} - -func (s *TokenService) Delete(ctx context.Context, pool, id string) error { - pool = normalizePool(pool) - if pool == "" { - return errors.New("unknown pool") - } - rows, err := s.tokens.Delete(ctx, pool, id) - if err != nil { - return err - } - // Also drop the matching cookie refresh profile (token id == profile id), - // otherwise the background refresher re-creates the token. Track whether a - // profile existed so we can mirror Python's 404-when-nothing-removed. - profileRemoved := false - if _, getErr := s.refresh.Get(ctx, id); getErr == nil { - profileRemoved = true - } - _ = s.refresh.Delete(ctx, id) - if rows == 0 && !profileRemoved { - return ErrNotFound - } - return nil -} - -// DeleteBulk removes many accounts by id (across pools) plus their cookie -// refresh profiles. Returns how many account rows were removed. -func (s *TokenService) DeleteBulk(ctx context.Context, ids []string) (int, error) { - seen := make(map[string]struct{}, len(ids)) - clean := make([]string, 0, len(ids)) - for _, id := range ids { - id = strings.TrimSpace(id) - if id == "" { - continue - } - if _, ok := seen[id]; ok { - continue - } - seen[id] = struct{}{} - clean = append(clean, id) - } - if len(clean) == 0 { - return 0, nil - } - rows, err := s.tokens.DeleteByIDs(ctx, clean) - if err != nil { - return 0, err - } - // Drop matching cookie refresh profiles (id == token id) so the background - // refresher doesn't re-create the tokens. - _ = s.refresh.DeleteByIDs(ctx, clean) - return int(rows), nil -} - -func (s *TokenService) Accounts(ctx context.Context) ([]map[string]any, error) { - items, err := s.tokens.List(ctx) - if err != nil { - return nil, err - } - inFlight, err := s.events.InFlightByAccount(ctx) - if err != nil { - return nil, err - } - out := make([]map[string]any, 0, len(items)) - for _, item := range items { - out = append(out, accountRow(item, inFlight[item.ID])) - } - return out, nil -} - -func (s *TokenService) Quota(ctx context.Context, pool, id string) (map[string]any, error) { - item, err := s.tokens.Get(ctx, normalizePool(pool), id) - if err != nil { - return nil, err - } - if poolToType(item.Pool) == "openai" && s.chatgpt != nil { - data, err := s.chatgpt.FetchImageQuota(ctx, item.Value) - if err != nil { - return nil, err - } - authFailed := boolValueWithDefault(data["auth_failed"], false) - if authFailed { - _, _ = s.tokens.Update(ctx, item.Pool, item.ID, map[string]any{ - "status": "disabled", - "dead": true, - "fails": gorm.Expr("fails + 1"), - }) - } - patch := map[string]any{} - meta := cloneJSONMap(item.Meta) - meta["cached_quota_at"] = int(time.Now().Unix()) - rem, exhausted := chatgptRemaining(data) - meta["cached_quota_remaining"] = rem - patch["meta"] = meta - resetAfter := strings.TrimSpace(stringValue(data["reset_after"])) - if resetAfter == "" { - resetAfter = leonardoResetAfter("") - } - patch["cached_quota_reset_after"] = resetAfter - item.CachedQuotaResetAfter = resetAfter - // remaining<=0(负数/未知/0)→ 限额;>0 且当前是 quota → 恢复 active。 - // auth 失效已在上面置死,这里不再改它的状态。 - if !authFailed { - if exhausted { - patch["status"] = "quota" - } else if item.Status == "quota" { - patch["status"] = "active" - } - } - if updated, updateErr := s.tokens.Update(ctx, item.Pool, item.ID, patch); updateErr == nil { - item = updated - } - return map[string]any{ - "supported": true, - "remaining": rem, - "total": nil, - "reset_after": emptyToNil(item.CachedQuotaResetAfter), - "quota_cached_at": meta["cached_quota_at"], - "unchanged": false, - "unknown": boolValueWithDefault(data["unknown"], false), - "error": data["error"], - }, nil - } - if poolToType(item.Pool) == "adobe" && s.adobe != nil { - data, err := s.adobe.FetchCreditsBalance(ctx, item.Value) - if err != nil { - if errors.Is(err, adobe.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()) - plan := strings.ToLower(strings.TrimSpace(stringValue(data["plan"]))) - if plan != "" { - meta["plan"] = plan - } - if remaining, ok := data["remaining"].(int); ok { - meta["cached_quota_remaining"] = remaining - // 额度打到 0 的 VIP(母号/子号) 视为普号：写死 plan=free，额度恢复(>0)后下次扫描自动还原真实身份。 - if plan != "" && plan != "free" && remaining <= 0 { - plan = "free" - meta["plan"] = "free" - } - // is_sub_account 仅在导入时写入，刷新配额时不覆盖 - if _, hasFlag := meta["is_sub_account"]; !hasFlag && plan != "free" { - meta["is_sub_account"] = remaining > 0 && remaining <= 4000 - } - } - // 子号必然是会员号：plan 查回 free 说明这个号没有会员，不能再挂子号标记 - // （否则同一行会同时显示 子号 和 普号）。 - if plan == "free" { - meta["is_sub_account"] = false - } - 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 - // 探测不到会员身份（新旧 plan 都为空）的号既不算普号也不算会员号， - // 留在池里只会在需要会员的请求上撞 403 —— 直接置死号。 - if strings.TrimSpace(stringValue(meta["plan"])) == "" { - patch["status"] = "disabled" - patch["dead"] = true - } - if resetAfter := strings.TrimSpace(stringValue(data["available_until"])); resetAfter != "" { - patch["cached_quota_reset_after"] = resetAfter - item.CachedQuotaResetAfter = resetAfter - } - if len(patch) > 0 { - if updated, updateErr := s.tokens.Update(ctx, item.Pool, item.ID, patch); updateErr == nil { - item = updated - } - } - subAccount, _ := jsonMapBool(meta, "is_sub_account") - 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"], - // 会员身份跟着额度一起返回，前端查一次额度就能刷新 子号/母号/普号 徽章。 - "plan": emptyToNil(strings.ToLower(strings.TrimSpace(stringValue(meta["plan"])))), - "sub_account": subAccount, - }, nil - } - if poolToType(item.Pool) == "krea" && s.krea != nil { - s.applyProxy(ctx) - cookie, rerr := kreaRefreshAndPersist(ctx, s.krea, s.tokens, item.ID, item.Value) - if rerr != nil { - if errors.Is(rerr, krea.ErrAuth) { - _, _ = s.tokens.Update(ctx, item.Pool, item.ID, map[string]any{ - "status": "disabled", "dead": true, "fails": gorm.Expr("fails + 1"), - }) - } - return nil, rerr - } - data, err := s.krea.FetchCreditsBalance(ctx, cookie) - if err != nil { - if errors.Is(err, krea.ErrAuth) { - _, _ = s.tokens.Update(ctx, item.Pool, item.ID, map[string]any{ - "status": "disabled", "dead": true, "fails": gorm.Expr("fails + 1"), - }) - } - return nil, err - } - meta := cloneJSONMap(item.Meta) - meta["cached_quota_at"] = int(time.Now().Unix()) - rem, hasRem := data["remaining"].(int) - if hasRem { - meta["cached_quota_remaining"] = rem - } - // 每日免费额度在 UTC 零点(北京 08:00)重置 —— 恢复时间始终重算为"下一个零点", - // 不保留已过期的旧值(否则过了 08:00 还一直显示今天 08:00,不会变明天)。 - resetAfter := leonardoResetAfter("") - patch := map[string]any{"meta": meta, "cached_quota_reset_after": resetAfter} - // Krea 限额由生成时的 402 判定;一旦余额恢复(galactus 触发刷新后 >0),把之前 - // 沉下去的 quota 翻回 active,避免有余额却卡在限额。 - if item.Status == "quota" && hasRem && rem > 0 { - patch["status"] = "active" - } - 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(resetAfter), - "quota_cached_at": meta["cached_quota_at"], - "unchanged": false, - "unknown": boolValueWithDefault(data["unknown"], false), - "error": data["error"], - }, nil - } - if poolToType(item.Pool) == "imagine" && s.imagine != nil { - s.applyProxy(ctx) - cred, rerr := imagineRefreshAndPersist(ctx, s.imagine, s.tokens, item.ID, item.Value) - if rerr != nil { - if errors.Is(rerr, imagine.ErrAuth) { - _, _ = s.tokens.Update(ctx, item.Pool, item.ID, map[string]any{ - "status": "disabled", "dead": true, "fails": gorm.Expr("fails + 1"), - }) - } - return nil, rerr - } - data, err := s.imagine.FetchCreditsBalance(ctx, cred) - if err != nil { - if errors.Is(err, imagine.ErrAuth) { - _, _ = s.tokens.Update(ctx, item.Pool, item.ID, map[string]any{ - "status": "disabled", "dead": true, "fails": gorm.Expr("fails + 1"), - }) - } - return nil, err - } - meta := cloneJSONMap(item.Meta) - meta["cached_quota_at"] = int(time.Now().Unix()) - rem, hasRem := data["remaining"].(int) - if hasRem { - meta["cached_quota_remaining"] = rem - } - // 每日免费额度在 UTC 零点(北京 08:00)重置 —— 恢复时间始终重算为"下一个零点", - // 不保留已过期的旧值(否则过了 08:00 还一直显示今天 08:00,不会变明天)。 - resetAfter := leonardoResetAfter("") - patch := map[string]any{"meta": meta, "cached_quota_reset_after": resetAfter} - // 余额恢复(>0)→ 把之前因 402 沉下去的 quota 翻回 active。 - if item.Status == "quota" && hasRem && rem > 0 { - patch["status"] = "active" - } - 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(resetAfter), - "quota_cached_at": meta["cached_quota_at"], - "unchanged": false, - "unknown": boolValueWithDefault(data["unknown"], false), - "error": data["error"], - }, nil - } - if poolToType(item.Pool) == "leonardo" && s.leonardo != nil { - s.applyProxy(ctx) - data, err := s.leonardo.FetchCreditsBalance(ctx, item.Value) - s.persistLeonardoCookie(ctx, item.ID, item.Value) - if err != nil { - if errors.Is(err, leonardo.ErrAuth) { - meta, strikes, kill := leonardoAuthStrike(item, "quota refresh: "+err.Error()) - patch := map[string]any{"meta": meta, "fails": gorm.Expr("fails + 1")} - if kill { - patch["status"] = "disabled" - patch["dead"] = true - log.Printf("account leonardo/%s disabled after %d consecutive auth failures: %v", item.ID, strikes, err) - } else { - log.Printf("leonardo %s: auth failure %d/%d on quota refresh (%v) — kept active", item.ID, strikes, leonardoAuthStrikeLimit, err) - } - _, _ = s.tokens.Update(ctx, item.Pool, item.ID, patch) - } - return nil, err - } - patch := map[string]any{} - meta := cloneJSONMap(item.Meta) - meta["cached_quota_at"] = int(time.Now().Unix()) - meta["auth_fails"] = 0 // cookie 还能换 token，连续失败计数清零 - if remaining, ok := data["remaining"].(int); ok { - meta["cached_quota_remaining"] = remaining - // Below the per-generation floor → sink to "限额" so it stops being - // scheduled. The daily renewal time (below) lets the sweep auto-recover. - if remaining < leonardoMinCredits && item.Status == "active" { - patch["status"] = "quota" - } - } - if uid := strings.TrimSpace(stringValue(data["user_id"])); uid != "" { - meta["user_id"] = uid - } - patch["meta"] = meta - // Daily reset (08:00 Beijing == next UTC midnight) unless upstream gives an - // explicit renewal time. Drives RecoverQuota. - resetAfter := leonardoResetAfter(stringValue(data["available_until"])) - patch["cached_quota_reset_after"] = resetAfter - 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(resetAfter), - "quota_cached_at": meta["cached_quota_at"], - "unchanged": false, - "unknown": boolValueWithDefault(data["unknown"], false), - "error": data["error"], - }, nil - } - if poolToType(item.Pool) == "runway" && s.runway != nil { - data, err := s.runway.FetchCreditsBalance(ctx, item.Value) - if err != nil { - if errors.Is(err, runway.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 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 - } - 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 updated, updateErr := s.tokens.Update(ctx, item.Pool, item.ID, patch); updateErr == nil { - item = updated - } - // Recovery time stays the JWT expiry (cached at import) — Runway credits - // reset monthly, so the credits endpoint carries no reset timestamp. - 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 - } - if poolToType(item.Pool) == "grok" { - // 本地写死的额度（图 5 / 视频 2），没有上游接口可查，所以刷新只是回读本地计数。 - images, ok := jsonMapInt(item.Meta, repo.GrokImageQuotaKey) - if !ok { - images = repo.GrokImageQuota - } - videos, ok := jsonMapInt(item.Meta, repo.GrokVideoQuotaKey) - if !ok { - videos = repo.GrokVideoQuota - } - return map[string]any{ - "supported": true, - "remaining": images + videos, - "image_remaining": images, - "video_remaining": videos, - "used": nil, - "total": nil, - "reset_after": nil, - "quota_cached_at": item.Meta["cached_quota_at"], - "unchanged": true, - "unknown": false, - "error": nil, - }, 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" || 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), - "unchanged": true, - "unknown": false, - "error": nil, - }, nil -} - -func (s *TokenService) Email(ctx context.Context, pool, id string) (map[string]any, error) { - item, err := s.tokens.Get(ctx, normalizePool(pool), id) - if err != nil { - return nil, err - } - if poolToType(item.Pool) == "openai" { - email := strings.TrimSpace(item.AccountEmail) - if email != "" { - return map[string]any{"email": email, "cached": true}, nil - } - info := chatgpt.ExtractAccountInfo(item.Value) - if extracted := strings.TrimSpace(stringValue(info["email"])); extracted != "" { - _, _ = s.tokens.Update(ctx, item.Pool, item.ID, map[string]any{"account_email": extracted}) - return map[string]any{"email": extracted, "cached": false}, nil - } - return map[string]any{"email": nil, "cached": false}, nil - } - if poolToType(item.Pool) == "runway" { - email := strings.TrimSpace(item.AccountEmail) - if email != "" { - return map[string]any{"email": email, "cached": true}, nil - } - // Runway email is a top-level JWT claim — decode it (no network). - if extracted, _ := parseJWTEmailExpiry(item.Value); extracted != "" { - _, _ = s.tokens.Update(ctx, item.Pool, item.ID, map[string]any{"account_email": extracted}) - return map[string]any{"email": extracted, "cached": false}, nil - } - return map[string]any{"email": nil, "cached": false}, nil - } - if poolToType(item.Pool) != "adobe" && poolToType(item.Pool) != "creativefabrica" { - return map[string]any{"email": nil}, nil - } - if poolToType(item.Pool) == "creativefabrica" { - email := strings.TrimSpace(item.AccountEmail) - if email != "" { - return map[string]any{"email": email, "cached": true}, nil - } - if s.cf == nil { - return map[string]any{"email": nil, "cached": false}, nil - } - fetched, err := s.cf.FetchUser(ctx, item.Value) - if err != nil { - if errors.Is(err, creativefabrica.ErrAuth) { - _, _ = s.tokens.Update(ctx, item.Pool, item.ID, map[string]any{ - "status": "disabled", - "dead": true, - "fails": gorm.Expr("fails + 1"), - }) - } - return nil, err - } - if fetched = strings.TrimSpace(fetched); fetched != "" { - _, _ = s.tokens.Update(ctx, item.Pool, item.ID, map[string]any{"account_email": fetched}) - email = fetched - } - return map[string]any{"email": emptyToNil(email), "cached": false}, nil - } - email := strings.TrimSpace(item.AccountEmail) - if email == "" { - if s.adobe == nil { - return map[string]any{"email": nil, "cached": false}, nil - } - profile, err := s.adobe.FetchAccountProfile(ctx, item.Value) - if err != nil { - if errors.Is(err, adobe.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{} - if profileEmail := strings.TrimSpace(stringValue(profile["email"])); profileEmail != "" { - patch["account_email"] = profileEmail - email = profileEmail - } - if displayName := strings.TrimSpace(stringValue(profile["display_name"])); displayName != "" { - patch["account_display_name"] = displayName - } - if len(patch) > 0 { - _, _ = s.tokens.Update(ctx, item.Pool, item.ID, patch) - } - return map[string]any{"email": emptyToNil(email), "cached": false}, nil - } - return map[string]any{"email": email, "cached": true}, nil -} - -func (s *TokenService) createToken(ctx context.Context, pool, tokenID, value, status string, meta datatypes.JSONMap) (*model.TokenAccount, error) { - now := time.Now() - item := &model.TokenAccount{ - ID: tokenID, - Pool: pool, - Value: value, - Status: status, - AddedAt: &now, - Meta: meta, - CreatedAt: now, - UpdatedAt: now, - } - if err := s.tokens.Create(ctx, item); err != nil { - return nil, err - } - return item, nil -} - -func accountRow(item model.TokenAccount, inFlight int64) map[string]any { - remaining, hasRemaining := jsonMapInt(item.Meta, "cached_quota_remaining") - quotaAt, _ := jsonMapInt(item.Meta, "cached_quota_at") - pending, _ := jsonMapBool(item.Meta, "pending_check") - subAccount, _ := jsonMapBool(item.Meta, "is_sub_account") - plan, _ := item.Meta["plan"] - // OpenAI email lives in the token's JWT (nested profile claim). Decode it at - // render time like the Python reference (_account_row) so accounts imported - // before the email was persisted still show a name; fall back to the cached - // field for adobe (whose email comes from a network profile fetch). - email := item.AccountEmail - if poolToType(item.Pool) == "openai" { - if decoded := strings.TrimSpace(stringValue(chatgpt.ExtractAccountInfo(item.Value)["email"])); decoded != "" { - email = decoded - } - } - typeLabel := poolToType(item.Pool) - teamID := "" - if item.Meta != nil { - teamID = strings.TrimSpace(stringValue(item.Meta["team_id"])) - } - // grok 额度是本地写死的两个计数（图 / 视频），前台单独一列展示成 "5/2"。 - var grokImages, grokVideos any - if typeLabel == "grok" { - images, ok := jsonMapInt(item.Meta, repo.GrokImageQuotaKey) - if !ok { - images = repo.GrokImageQuota - } - videos, ok := jsonMapInt(item.Meta, repo.GrokVideoQuotaKey) - if !ok { - videos = repo.GrokVideoQuota - } - grokImages, grokVideos = images, videos - } - hasQuota := typeLabel == "openai" || typeLabel == "adobe" || typeLabel == "runway" || typeLabel == "leonardo" || typeLabel == "krea" || typeLabel == "imagine" || typeLabel == "grok" || typeLabel == "creativefabrica" - return map[string]any{ - "id": item.ID, - "pool": item.Pool, - "type": typeLabel, - "email": emptyToNil(email), - "team_id": emptyToNil(teamID), - "remaining": valueOrNil(hasQuota && hasRemaining, remaining), - "image_remaining": grokImages, - "video_remaining": grokVideos, - "reset_after": emptyToNil(item.CachedQuotaResetAfter), - "quota_cached_at": valueOrNil(quotaAt != 0, quotaAt), - "created_at": unixOrNil(item.AddedAt), - "last_used_at": unixOrNil(item.LastUsedAt), - "expires_at": jwtExpiryUnix(item.Value), - "in_flight": inFlight, - "success_total": item.SuccessTotal, - "fail_total": item.FailTotal, - "fails_streak": item.Fails, - "status": item.Status, - "dead": item.Dead, - "image_limited": item.ImageLimited, - "video_limited": item.VideoLimited, - "pending": pending, - "sub_account": subAccount, - "plan": emptyToNil(strings.ToLower(strings.TrimSpace(stringValue(plan)))), - "quota_supported": hasQuota, - "needs_reset_fetch": typeLabel == "adobe" && item.Status == "active" && strings.TrimSpace(item.CachedQuotaResetAfter) == "", - "weight": item.Weight, - "concurrency": item.Concurrency, - "base_url": emptyToNil(strings.TrimSpace(stringValue(item.Meta["base_url"]))), - "protocol": emptyToNil(strings.TrimSpace(stringValue(item.Meta["protocol"]))), - "models": strings.TrimSpace(stringValue(item.Meta["models"])), - } -} - -// jwtExpiryUnix returns the access token's exp claim (epoch seconds) for the -// accounts UI, or nil when the token is absent/opaque (e.g. a pending row whose -// value hasn't been minted yet). -func jwtExpiryUnix(token string) any { - if strings.TrimSpace(token) == "" { - return nil - } - _, exp := parseJWTEmailExpiry(token) - if exp == nil { - return nil - } - return exp.Unix() -} - -// cleanAdobeCookie mirrors the Python admin import preprocessing -// (api/admin.py import_adobe_cookie): tolerate JSON array/object pastes, -// unwrap a one-level {"cookie": "..."} wrapper, strip a leading "Cookie:" -// prefix and collapse stray whitespace/newlines. -func cleanAdobeCookie(cookie string) string { - cookieStr := strings.TrimSpace(cookie) - - // ① A JSON array/object paste -> turn into a cookie string up front. - if strings.HasPrefix(cookieStr, "[") { - if converted := cookieStringFromInput(cookieStr); converted != "" { - cookieStr = converted - } - } - - // ② Tolerate the whole JSON object `{"cookie": "..."}` pasted into the - // textarea. Unwrap one level of JSON before validating. - if strings.HasPrefix(cookieStr, "{") && strings.HasSuffix(cookieStr, "}") { - var parsed map[string]any - if err := json.Unmarshal([]byte(cookieStr), &parsed); err == nil { - inner, ok := parsed["cookie"] - if !ok { - inner = parsed["value"] - } - switch v := inner.(type) { - case string: - cookieStr = strings.TrimSpace(v) - case []any, map[string]any: - if converted := cookieStringFromInputValue(v); converted != "" { - cookieStr = converted - } - } - } - } - - // ③ Strip a leading "Cookie: " prefix (case-insensitive). - if len(cookieStr) >= 7 && strings.EqualFold(cookieStr[:7], "cookie:") { - cookieStr = strings.TrimSpace(cookieStr[7:]) - } - - // ④ Collapse stray newlines and excess whitespace. - cookieStr = strings.Join(strings.Fields(cookieStr), " ") - return cookieStr -} - -// cookieStringFromInput parses a JSON string (array or object) describing -// browser cookies into a "name=value; name=value" cookie string, mirroring -// providers/adobe/_auth.py _cookie_string_from_input. -func cookieStringFromInput(raw string) string { - raw = strings.TrimSpace(raw) - var parsed any - if err := json.Unmarshal([]byte(raw), &parsed); err != nil { - return "" - } - return cookieStringFromInputValue(parsed) -} - -func cookieStringFromInputValue(raw any) string { - switch v := raw.(type) { - case string: - text := strings.TrimSpace(v) - if len(text) >= 7 && strings.EqualFold(text[:7], "cookie:") { - text = strings.TrimSpace(text[7:]) - } - return text - case []any: - parts := make([]string, 0, len(v)) - for _, item := range v { - m, ok := item.(map[string]any) - if !ok { - continue - } - name := strings.TrimSpace(stringValue(m["name"])) - value := stringValue(m["value"]) - if name != "" { - parts = append(parts, name+"="+value) - } - } - return strings.Join(parts, "; ") - case map[string]any: - if cookies, ok := v["cookies"].([]any); ok { - return cookieStringFromInputValue(cookies) - } - if inner, ok := v["cookie"]; ok { - switch inner.(type) { - case string, []any: - return cookieStringFromInputValue(inner) - } - } - return "" - default: - return "" - } -} - -func previewSecret(v string) string { - v = strings.TrimSpace(v) - if v == "" { - return "" - } - if len(v) <= 16 { - return "***" - } - return v[:6] + "…" + v[len(v)-4:] -} - -func normalizePool(pool string) string { - pool = strings.ToLower(strings.TrimSpace(pool)) - if _, ok := validTokenPools[pool]; ok { - return pool - } - return "" -} - -func normalizeTokenStatus(v string) string { - switch strings.ToLower(strings.TrimSpace(v)) { - case "active", "disabled", "quota", "pending": - return strings.ToLower(strings.TrimSpace(v)) - default: - return "" - } -} - -func newTokenID(pool string) string { - prefix := "TK" - if pool == "adobe" { - prefix = "AD" - } - if pool == "chatgpt" { - prefix = "OA" - } - if pool == "runway" { - prefix = "RW" - } - if pool == "leonardo" { - prefix = "LN" - } - if pool == "krea" { - prefix = "KR" - } - if pool == "imagine" { - prefix = "IM" - } - if pool == "creativefabrica" { - prefix = "CF" - } - return prefix + randomUpper(10) -} - -func poolToType(pool string) string { - if mapped, ok := validTokenPools[pool]; ok { - return mapped - } - return pool -} - -func parseJWTEmailExpiry(token string) (string, *time.Time) { - parts := strings.Split(strings.TrimSpace(token), ".") - if len(parts) < 2 { - return "", nil - } - payload := parts[1] - raw, err := base64.RawURLEncoding.DecodeString(payload) - if err != nil { - return "", nil - } - var claims map[string]any - if err := json.Unmarshal(raw, &claims); err != nil { - return "", nil - } - email := strings.TrimSpace(stringValue(claims["email"])) - switch v := claims["exp"].(type) { - case float64: - t := time.Unix(int64(v), 0) - return email, &t - case json.Number: - n, err := v.Int64() - if err == nil { - t := time.Unix(n, 0) - return email, &t - } - } - return email, nil -} - -func jsonMapInt(m datatypes.JSONMap, key string) (int, bool) { - if m == nil { - return 0, false - } - v, ok := m[key] - if !ok || v == nil { - return 0, false - } - return intValue(v), true -} - -func jsonMapBool(m datatypes.JSONMap, key string) (bool, bool) { - if m == nil { - return false, false - } - v, ok := m[key] - if !ok || v == nil { - return false, false - } - switch x := v.(type) { - case bool: - return x, true - default: - return boolValueWithDefault(x, false), true - } -} - -func cloneJSONMap(in datatypes.JSONMap) datatypes.JSONMap { - out := datatypes.JSONMap{} - for k, v := range in { - out[k] = v - } - return out -} - -func timeOrNil(t *time.Time) any { - if t == nil { - return nil - } - return *t -} - -func unixOrNil(t *time.Time) any { - if t == nil { - return nil - } - return t.Unix() -} - -func emptyToNil(v string) any { - if strings.TrimSpace(v) == "" { - return nil - } - return strings.TrimSpace(v) -} - -func valueOrNil(ok bool, v any) any { - if !ok { - return nil - } - return v -} - -type ginToken struct { - ID string `json:"id"` - ValuePreview string `json:"value_preview"` - Status string `json:"status"` - Fails int `json:"fails"` - AddedAt *time.Time `json:"added_at"` -} - -var _ = fmt.Sprint diff --git a/backend/internal/service/user_generation.go b/backend/internal/service/user_generation.go index 1cb1141..8133a8b 100644 --- a/backend/internal/service/user_generation.go +++ b/backend/internal/service/user_generation.go @@ -40,8 +40,8 @@ type UserGenerateRequest struct { // DeAI applies 去AI特征 post-processing to the generated image (image only) // and charges the per-tier surcharge on top of the model price. DeAI bool - // AccountID pins an admin test to one specific provider account (账号生图测试). - AccountID string + // ChannelID pins an admin test to one specific channel (渠道生图测试). + ChannelID string } func (s *UserGenerationService) Generate(ctx context.Context, user *model.User, in UserGenerateRequest) (map[string]any, error) { @@ -138,7 +138,7 @@ func (s *UserGenerationService) AdminTest(ctx context.Context, user *model.User, AspectRatio: in.Ratio, Resolution: in.Resolution, ReferenceImages: in.ReferenceImages, - AccountID: in.AccountID, + ChannelID: in.ChannelID, }) default: return s.v1.prepareAdminTestImage(ctx, principal, V1ImageRequest{ @@ -147,7 +147,7 @@ func (s *UserGenerationService) AdminTest(ctx context.Context, user *model.User, AspectRatio: in.Ratio, Resolution: in.Resolution, ReferenceImages: in.ReferenceImages, - AccountID: in.AccountID, + ChannelID: in.ChannelID, }) } } diff --git a/backend/internal/service/v1.go b/backend/internal/service/v1.go index bb7ee09..552dedf 100644 --- a/backend/internal/service/v1.go +++ b/backend/internal/service/v1.go @@ -8,122 +8,109 @@ import ( "fmt" "io" "log" - "net/http" + "math/rand" "path/filepath" "runtime/debug" - "sort" "strings" "sync" - "sync/atomic" "time" - "strconv" - "backend/internal/config" "backend/internal/model" - "backend/internal/provider/adobe" - "backend/internal/provider/chatgpt" - "backend/internal/provider/creativefabrica" - "backend/internal/provider/custom" - "backend/internal/provider/grok" - "backend/internal/provider/imagine" - "backend/internal/provider/krea" - "backend/internal/provider/leonardo" - "backend/internal/provider/runway" "backend/internal/repo" "backend/internal/storage" + "backend/internal/upstream" + "backend/internal/util" "gorm.io/gorm" ) var ( - ErrMissingAPIKey = errors.New("missing api key") - ErrInvalidAPIKey = errors.New("invalid api key") - ErrUnknownModel = errors.New("unknown model") - ErrUnsupportedParams = errors.New("unsupported or unpriced parameters for this model") - ErrBannedPrompt = errors.New("prompt contains banned content") - ErrInsufficientFunds = errors.New("insufficient credits") - ErrGenerationPending = errors.New("generation executor not implemented yet") - ErrProviderAuth = errors.New("provider token invalid or expired") - ErrNoProviderAccount = errors.New("no provider account available, please ask an admin to configure one") + ErrMissingAPIKey = errors.New("missing api key") + ErrInvalidAPIKey = errors.New("invalid api key") + ErrUnknownModel = errors.New("unknown model") + ErrUnsupportedParams = errors.New("unsupported or unpriced parameters for this model") + ErrBannedPrompt = errors.New("prompt contains banned content") + ErrInsufficientFunds = errors.New("insufficient credits") + ErrGenerationPending = errors.New("generation executor not implemented yet") + ErrProviderAuth = errors.New("upstream channel credentials rejected") + // ErrNoProviderAccount = no enabled channel is bound to the model. + ErrNoProviderAccount = errors.New("no upstream channel available, please ask an admin to configure one") + // ErrProviderQuota — kept for handler compatibility; channel failover no + // longer produces it (quota errors just rotate to the next channel). ErrProviderQuota = errors.New("provider quota exhausted") - ErrProviderTemporary = errors.New("provider temporary unavailable") - ErrProviderExecution = errors.New("provider request failed") + ErrProviderTemporary = errors.New("upstream temporary unavailable") + ErrProviderExecution = errors.New("upstream request failed") ErrProviderUnsupported = errors.New("provider not implemented") ErrReferenceTooLarge = errors.New("reference image too large") - // ErrConcurrencyFull — every eligible account is busy (each account runs at - // most ONE generation at a time). English message: surfaced to API / UI. - ErrConcurrencyFull = errors.New("all accounts are busy (1 concurrent job each), please try again shortly") - // ErrUserConcurrencyFull — the caller already has their concurrency-group's max - // generations in flight (画图台 + API key combined). 0 = unlimited. + // ErrConcurrencyFull — every eligible channel is busy right now. + ErrConcurrencyFull = errors.New("all upstream channels are busy, please try again shortly") + // ErrUserConcurrencyFull — the caller already has their concurrency-group's + // max generations in flight (画图台 + API key combined). 0 = unlimited. ErrUserConcurrencyFull = errors.New("too many generations in progress, please wait for one to finish") - // ErrVideoJobNotFound / ErrVideoNotReady — /v1/videos async job lookups. - ErrVideoJobNotFound = errors.New("video job not found") - ErrVideoNotReady = errors.New("video is not ready yet") + // ErrVideoJobNotFound / ErrVideoNotReady — /v1 job lookups (kept names). + ErrVideoJobNotFound = errors.New("job not found") + ErrVideoNotReady = errors.New("job result is not ready yet") ) -// maxReferenceImageBytes bounds a single decoded reference image. 20 MB -// comfortably covers real photos/screenshots; anything larger is almost -// certainly abuse or a mistake. Mirrors Python core/refs.py. +// maxReferenceImageBytes bounds a single decoded reference blob. 20 MB +// comfortably covers real photos/screenshots/clips; anything larger is almost +// certainly abuse or a mistake. const maxReferenceImageBytes = 20 * 1024 * 1024 +// videoGenBudget caps one video render end-to-end (submit + stream + download). +const videoGenBudget = 30 * time.Minute + +// imageGenBudget caps one image generation end-to-end. +const imageGenBudget = 8 * time.Minute + +// defaultChannelConcurrency applies when a channel's Concurrency is 0. +const defaultChannelConcurrency = 5 + +// maxChannelFailover caps how many channels one request may burn on +// retryable/temporary errors before giving up, so an upstream-wide blip can't +// fan a single request out across every channel. +const maxChannelFailover = 6 + +// PresignTTL is the lifetime of /v1/uploads PUT urls and result GET urls. +const PresignTTL = 30 * time.Minute + type V1Service struct { cfg *config.Config models *repo.ModelRepository users *repo.UserRepository events *repo.EventRepository - tokens *repo.TokenRepository + channels *ChannelService settings *repo.SiteSettingRepository cgroups *repo.ConcurrencyGroupRepository - adobe *adobe.Client - chatgpt *chatgpt.Client - runway *runway.Client - leonardo *leonardo.Client - krea *krea.Client - imagine *imagine.Client - grok *grok.Client - custom *custom.Client - cf *creativefabrica.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 - // init cycle). nil for deployments without cookie refresh. - refresh *RefreshProfileService + up *upstream.Client // banned is the admin-managed prompt blocklist (set via SetBannedWords). // nil disables the check. banned *repo.BannedWordRepository - // tokenCursors holds one strict round-robin cursor per pool (key: pool name, - // value: *uint64). Each pick advances the pool's cursor by one so accounts - // are used in a fixed, even rotation (acct1→acct2→acct3→acct1…) independent - // of fails/last_used. The atomic counter also serializes concurrent picks so - // two simultaneous requests never start on the same account. - tokenCursors sync.Map - // inflight maps an in-progress event ID → the cancel func of its generation // work context, so the maintenance sweep can stop a stuck generation the - // moment it abandons the row (instead of letting an orphaned goroutine run on - // for minutes and surface a late "success" on an already-abandoned event). + // moment it abandons the row. inflight *InflightRegistry - // conc is the Redis-backed concurrency limiter for BOTH the per-account - // upstream gate (1+ jobs per account) and the per-user gate (画图台 + API key, - // capped by the user's concurrency group). Self-healing + fail-open. + // conc is the Redis-backed concurrency limiter for BOTH the per-channel + // upstream gate and the per-user gate (画图台 + API key, capped by the + // user's concurrency group). Self-healing + fail-open. conc *ConcurrencyService } -// acctAcquire takes one per-account upstream slot (capped at max; 0/1 = single), -// tagged with the generation's eventID (unique per job; a generation only ever -// holds one slot on a given account at a time, so failover reuses it cleanly). -func (s *V1Service) acctAcquire(ctx context.Context, accountID, eventID string, max int) bool { +// chanAcquire takes one per-channel slot (capped at max; 0 = default), tagged +// with the generation's eventID. +func (s *V1Service) chanAcquire(ctx context.Context, channelID, eventID string, max int) bool { if max < 1 { - max = 1 + max = defaultChannelConcurrency } - return s.conc.Acquire(ctx, "conc:a:"+accountID, max, eventID) + return s.conc.Acquire(ctx, "conc:c:"+channelID, max, eventID) } -func (s *V1Service) acctRelease(ctx context.Context, accountID, eventID string) { - s.conc.Release(ctx, "conc:a:"+accountID, eventID) +func (s *V1Service) chanRelease(ctx context.Context, channelID, eventID string) { + s.conc.Release(ctx, "conc:c:"+channelID, eventID) } // userAcquire takes one per-user generation slot, capped by the user's @@ -194,11 +181,8 @@ type V1ImageRequest struct { Model string Prompt string Size string - // Quality is OpenAI's image quality (low|medium|high|auto). For our tiered - // models it selects the resolution (low→1K, medium→2K, high→4K, auto→default), - // clamped to whatever tiers the model actually prices. Only used when - // Resolution is left blank (the strict /v1 OpenAI path); the playground passes - // Resolution directly and ignores this. + // Quality is the vivid-protocol quality label (= resolution tier verbatim, + // e.g. "1K"/"2K"/"4K"). Only used when Resolution is blank. Quality string AspectRatio string Resolution string @@ -206,15 +190,14 @@ type V1ImageRequest struct { ReferenceImages []string // DeAI applies 去AI特征 post-processing (crop / noise / tone jitter + // re-encode) to the output and charges the per-tier surcharge on top of - // the model price. Playground-only; the /v1 OpenAI path never sets it. + // the model price. Playground-only; the /v1 path never sets it. DeAI bool - // BaseURL is the scheme+host of the inbound request (e.g. "https://host"), - // used to build absolute, directly-downloadable output URLs. Empty falls - // back to a relative "/images/..." path. + // ChannelID pins an admin test to one specific channel (渠道测试). Ignored + // on the normal user paths. + ChannelID string + // BaseURL is the scheme+host of the inbound request, used to build absolute + // /images URLs. Empty falls back to a relative path. BaseURL string - // AccountID pins the generation to one specific provider account (admin - // account-test). Empty keeps the normal pool selection with failover. - AccountID string } type V1VideoRequest struct { @@ -225,32 +208,24 @@ type V1VideoRequest struct { Resolution string ReferenceImages []string ReferenceMode string // "frame", "asset" or "media", overrides model default + // ChannelID — see V1ImageRequest.ChannelID. + ChannelID string // BaseURL — see V1ImageRequest.BaseURL. BaseURL string - // AccountID — see V1ImageRequest.AccountID. - AccountID string } -func NewV1Service(cfg *config.Config, models *repo.ModelRepository, users *repo.UserRepository, events *repo.EventRepository, tokens *repo.TokenRepository, settings *repo.SiteSettingRepository, cgroups *repo.ConcurrencyGroupRepository, conc *ConcurrencyService, adobeClient *adobe.Client, chatGPTClient *chatgpt.Client, runwayClient *runway.Client, leonardoClient *leonardo.Client, kreaClient *krea.Client, imagineClient *imagine.Client, grokClient *grok.Client, customClient *custom.Client, cfClient *creativefabrica.Client, store *storage.Client) *V1Service { +func NewV1Service(cfg *config.Config, models *repo.ModelRepository, users *repo.UserRepository, events *repo.EventRepository, channels *ChannelService, settings *repo.SiteSettingRepository, cgroups *repo.ConcurrencyGroupRepository, conc *ConcurrencyService, store *storage.Client) *V1Service { return &V1Service{ cfg: cfg, models: models, users: users, events: events, - tokens: tokens, + channels: channels, settings: settings, cgroups: cgroups, conc: conc, - adobe: adobeClient, - chatgpt: chatGPTClient, - runway: runwayClient, - leonardo: leonardoClient, - krea: kreaClient, - imagine: imagineClient, - grok: grokClient, - custom: customClient, - cf: cfClient, store: store, + up: channels.Upstream(), inflight: &InflightRegistry{}, } } @@ -259,11 +234,6 @@ func NewV1Service(cfg *config.Config, models *repo.ModelRepository, users *repo. // generation when it abandons that event. func (s *V1Service) Inflight() *InflightRegistry { return s.inflight } -// SetRefresh wires the Adobe cookie-refresh service in after construction -// (RefreshProfileService is built later in bootstrap, so it can't be a ctor arg -// without reordering). Enables refresh-then-retry on a mid-request 401. -func (s *V1Service) SetRefresh(r *RefreshProfileService) { s.refresh = r } - // SetBannedWords wires the prompt blocklist in after construction. func (s *V1Service) SetBannedWords(r *repo.BannedWordRepository) { s.banned = r } @@ -324,37 +294,11 @@ func (s *V1Service) logRejectedEvent(ctx context.Context, kind, modelID string, _ = s.events.Create(ctx, event) } -// refreshAdobeToken re-mints an Adobe account's access token from its cookie -// (RefreshNow) and returns the updated row. Used to retry a 401 with a fresh -// token instead of replaying the stale one. Returns false if refresh is -// unavailable or the cookie can no longer mint a token (genuinely dead). -func (s *V1Service) refreshAdobeToken(ctx context.Context, tokenID string) (model.TokenAccount, bool) { - if s.refresh == nil { - return model.TokenAccount{}, false - } - if s.settings != nil { - if proxy, err := s.settings.GetValue(ctx, "proxy.url"); err == nil && proxy != "" { - s.refresh.SetProxy(proxy) - } - } - if err := s.refresh.RefreshNow(ctx, tokenID); err != nil { - return model.TokenAccount{}, false - } - t, err := s.tokens.Get(ctx, "adobe", tokenID) - if err != nil || t == nil { - return model.TokenAccount{}, false - } - return *t, true -} - func (s *V1Service) Authenticate(ctx context.Context, authHeader string) (*APIPrincipal, error) { token := ParseBearer(authHeader) if token == "" { return nil, ErrMissingAPIKey } - - // Only per-user API keys (hashed in the DB) authenticate to /v1. The old - // global/shared API_KEY backdoor has been removed. user, err := s.users.GetByAPIKeyHash(ctx, HashAPIKey(token)) if err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { @@ -372,49 +316,342 @@ func (s *V1Service) Authenticate(ctx context.Context, authHeader string) (*APIPr }, nil } -func (s *V1Service) ListModels(ctx context.Context) ([]map[string]any, error) { +// ===== /v1 public protocol (vivid-compatible) ===== + +// V1ModelItem is one GET /v1/models entry — identical shape to a vivid +// upstream's, which is what makes channel model-import work symmetrically. +type V1ModelItem struct { + ID string `json:"id"` + Kind string `json:"kind"` + Qualities []string `json:"qualities"` + Duration *V1ModelDur `json:"duration,omitempty"` + RefImages int64 `json:"refImages"` + RefVideos int64 `json:"refVideos"` + RefAudios int64 `json:"refAudios"` + RefVideoSeconds int64 `json:"refVideoSeconds,omitempty"` + RefAudioSeconds int64 `json:"refAudioSeconds,omitempty"` +} + +// V1ModelDur exposes a video model's duration capability: fixed=locked seconds, +// options=discrete list, range=continuous interval. +type V1ModelDur struct { + Mode string `json:"mode"` + Seconds int64 `json:"seconds,omitempty"` + Min int64 `json:"min,omitempty"` + Max int64 `json:"max,omitempty"` + Options []int64 `json:"options,omitempty"` +} + +// SellableModels — GET /v1/models. Only enabled models bound to at least one +// enabled channel AND fully priced are sellable (a model with an unpriced tier +// would charge-reject at submit time — hide it instead). +func (s *V1Service) SellableModels(ctx context.Context) ([]V1ModelItem, error) { items, err := s.models.List(ctx) if err != nil { return nil, err } - now := time.Now().Unix() - out := make([]map[string]any, 0, len(items)) + out := make([]V1ModelItem, 0, len(items)) for _, item := range items { if !item.Enabled { continue } - out = append(out, map[string]any{ - "id": item.EffectiveName(), - "object": "model", - "created": now, - "owned_by": item.Provider, - "kind": item.Type, - "supported_ratios": repo.JSONStrings(item.Ratios), - "supported_resolutions": repo.JSONStrings(item.Resolutions), - }) + if !s.modelSellable(ctx, &item) { + continue + } + qualities := repo.JSONStrings(item.Resolutions) + if len(qualities) == 0 { + // Resolution list unset → advertise the priced tiers. + for k := range item.Prices { + qualities = append(qualities, k) + } + } + entry := V1ModelItem{ + ID: item.EffectiveName(), + Kind: item.Type, + Qualities: qualities, + RefImages: int64(item.MaxReferenceImages), + RefVideos: int64(item.MaxReferenceVideos), + RefAudios: int64(item.MaxReferenceAudios), + RefVideoSeconds: int64(item.MaxRefVideoSeconds), + RefAudioSeconds: int64(item.MaxRefAudioSeconds), + } + if entry.Qualities == nil { + entry.Qualities = []string{} + } + if item.Type == "video" { + entry.Duration = modelDurationSpec(&item) + } + out = append(out, entry) } return out, nil } -// UserBalance — GET /v1/user/balance 的数据。重新读一次用户行保证实时 -// （principal 里那份是鉴权时读的，可能已经过期）。 -func (s *V1Service) UserBalance(ctx context.Context, principal *APIPrincipal) (map[string]any, error) { +// modelSellable checks binding + pricing so /v1/models only advertises what a +// generate call can actually fulfil. +func (s *V1Service) modelSellable(ctx context.Context, item *model.ModelConfig) bool { + channels, err := s.channels.ChannelsForModel(ctx, item.ID) + if err != nil || len(channels) == 0 { + return false + } + for _, res := range repo.JSONStrings(item.Resolutions) { + if _, ok := jsonMapFloat(item.Prices, res); !ok { + return false + } + } + if item.Type == "video" { + durs := repo.JSONStrings(item.Durations) + if len(durs) > 0 { + for _, d := range durs { + if _, ok := jsonMapFloat(item.DurationPrices, d); !ok { + return false + } + } + } else if item.DurationMin > 0 || item.DurationMax > 0 { + if _, ok := jsonMapFloat(item.DurationPrices, "per_second"); !ok { + return false + } + } + } + return len(repo.JSONStrings(item.Resolutions)) > 0 || len(item.Prices) > 0 +} + +func modelDurationSpec(item *model.ModelConfig) *V1ModelDur { + if item.DurationMin > 0 || item.DurationMax > 0 { + return &V1ModelDur{Mode: "range", Min: int64(item.DurationMin), Max: int64(item.DurationMax)} + } + opts := repo.JSONStrings(item.Durations) + switch len(opts) { + case 0: + return nil + case 1: + return &V1ModelDur{Mode: "fixed", Seconds: int64(parseDurationSeconds(opts[0]))} + default: + vals := make([]int64, 0, len(opts)) + for _, d := range opts { + if n := parseDurationSeconds(d); n > 0 { + vals = append(vals, int64(n)) + } + } + return &V1ModelDur{Mode: "options", Options: vals} + } +} + +// UserBalance — GET /v1/balance. Re-reads the user row so the number is live. +func (s *V1Service) UserBalance(ctx context.Context, principal *APIPrincipal) (float64, error) { user := principal.User if fresh, err := s.users.GetByID(ctx, user.ID); err == nil && fresh != nil { user = fresh } - return map[string]any{ - "object": "user.balance", - "balance": user.Credits, - "used": user.CreditsUsed, - "total": user.Credits + user.CreditsUsed, + return user.Credits, nil +} + +// ===== channel dispatch ===== + +// generationParams bundles the wire fields one generate call carries. +type generationParams struct { + Kind string // "image" | "video" + Prompt string + Ratio string + Quality string // resolution tier verbatim + DurationSecs int64 + Refs [][]byte + // PinChannelID forces this one channel (admin test) — bypasses the enabled + // filter, weighted order and failover. + PinChannelID string +} + +// genResult is what a successful dispatch produces. +type genResult struct { + Data []byte + ContentType string + ChannelID string + ChannelName string + UpstreamJob string + Credits float64 // upstream-reported spend +} + +// dispatchGenerate routes one generation through the channel pool bound to the +// model: weighted-random order, per-channel concurrency slot, refs re-uploaded +// per attempt (upload sessions are channel-local), failover on auth/temporary +// errors. Deterministic errors (invalid params / failed job) are terminal. +func (s *V1Service) dispatchGenerate(ctx context.Context, eventID string, modelItem *model.ModelConfig, p generationParams) (*genResult, error) { + var channels []repo.ChannelWithUpstream + var err error + if pin := strings.TrimSpace(p.PinChannelID); pin != "" { + // Admin test pin: single channel, disabled included, no failover. + channels, err = s.channels.ChannelsForModelAnyState(ctx, modelItem.ID) + if err != nil { + return nil, err + } + kept := channels[:0] + for _, c := range channels { + if c.Channel.ID == pin { + kept = append(kept, c) + } + } + channels = kept + } else { + channels, err = s.channels.ChannelsForModel(ctx, modelItem.ID) + if err != nil { + return nil, err + } + } + if len(channels) == 0 { + return nil, ErrNoProviderAccount + } + var lastErr error + failovers := 0 + for _, cand := range weightedChannelOrder(channels) { + if ctx.Err() != nil { + return nil, ctx.Err() + } + ch := cand.Channel + if !s.chanAcquire(ctx, ch.ID, eventID, ch.Concurrency) { + continue // channel at capacity — try the next one + } + res, attemptErr := s.generateOnChannel(ctx, eventID, ch, cand.UpstreamModel, p) + s.chanRelease(ctx, ch.ID, eventID) + if attemptErr == nil { + return res, nil + } + lastErr = attemptErr + if !upstream.Failoverable(attemptErr) { + return nil, upstreamUserError(attemptErr) + } + if errors.Is(attemptErr, upstream.ErrAuth) { + continue // bad credentials don't burn the failover budget + } + failovers++ + if failovers >= maxChannelFailover { + break + } + } + if lastErr == nil { + return nil, ErrConcurrencyFull + } + return nil, fmt.Errorf("%w: %v", ErrProviderTemporary, upstreamUserError(lastErr)) +} + +// generateOnChannel performs one full attempt on a single channel: upload refs +// → POST /v1/generate (NDJSON, jobId reconnect) → download result bytes. +func (s *V1Service) generateOnChannel(ctx context.Context, eventID string, ch model.Channel, upstreamModel string, p generationParams) (*genResult, error) { + apiKey := s.channels.Decrypt(ch) + if strings.TrimSpace(apiKey) == "" { + return nil, fmt.Errorf("%w: empty key", upstream.ErrAuth) + } + base := strings.TrimRight(ch.BaseURL, "/") + refKeys := make([]string, 0, len(p.Refs)) + for _, blob := range p.Refs { + key, err := s.up.UploadRef(ctx, base, apiKey, blob) + if err != nil { + return nil, err + } + refKeys = append(refKeys, key) + } + wait := int64(imageGenBudget / time.Second) + if p.Kind == "video" { + wait = int64(videoGenBudget / time.Second) + } + var upstreamJob string + resultURL, credits, err := s.up.Generate(ctx, base, apiKey, upstream.GenParams{ + Model: upstreamModel, + Prompt: p.Prompt, + Quality: p.Quality, + Ratio: p.Ratio, + Duration: p.DurationSecs, + Refs: refKeys, + }, wait, func(job string) { + if upstreamJob == "" { + upstreamJob = job + _ = s.events.SetUpstreamJob(ctx, eventID, job) + } + }) + if err != nil { + return nil, err + } + data, ct, err := s.up.Download(ctx, resultURL) + if err != nil { + return nil, err + } + s.channels.BumpSuccess(ctx, ch.ID, float64(credits)) + _ = s.events.SetAccount(ctx, eventID, ch.ID, ch.Name) + return &genResult{ + Data: data, + ContentType: ct, + ChannelID: ch.ID, + ChannelName: ch.Name, + UpstreamJob: upstreamJob, + Credits: float64(credits), }, nil } -func (s *V1Service) PrepareImageRequest(ctx context.Context, principal *APIPrincipal, in V1ImageRequest) (map[string]any, error) { - return s.prepareImageExecution(ctx, principal, in, "v1", true) +// weightedChannelOrder builds a weighted-random permutation: each pick draws a +// channel with probability weight/remainingTotal (weight 2 → ~2x first-pick +// rate of weight 1). Ported from vivid's job executor. +func weightedChannelOrder(channels []repo.ChannelWithUpstream) []repo.ChannelWithUpstream { + if len(channels) <= 1 { + return channels + } + rest := append([]repo.ChannelWithUpstream(nil), channels...) + out := make([]repo.ChannelWithUpstream, 0, len(rest)) + for len(rest) > 0 { + total := 0 + for _, c := range rest { + w := c.Channel.Weight + if w < 1 { + w = 1 + } + total += w + } + r := rand.Intn(total) + for i, c := range rest { + w := c.Channel.Weight + if w < 1 { + w = 1 + } + if r < w { + out = append(out, c) + rest = append(rest[:i], rest[i+1:]...) + break + } + r -= w + } + } + return out } +// upstreamUserError maps an upstream error to what the caller sees. failed 帧 +// 文案是上游给用户看的中文，透传；4001 参数错同样透传；其余归一为通用文案。 +func upstreamUserError(err error) error { + if err == nil { + return ErrProviderTemporary + } + msg := err.Error() + for _, prefix := range []string{ + upstream.ErrJobFailed.Error() + ": ", + upstream.ErrInvalid.Error() + ": ", + upstream.ErrAuth.Error() + ": ", + upstream.ErrRetryable.Error() + ": ", + upstream.ErrAmbiguous.Error() + ": ", + } { + msg = strings.TrimSpace(strings.TrimPrefix(msg, prefix)) + } + switch upstream.KindOf(err) { + case upstream.KindJobFailed, upstream.KindInvalid: + if msg == "" { + return ErrProviderExecution + } + return fmt.Errorf("%w: %s", ErrProviderExecution, msg) + default: + if msg == "" { + return ErrProviderTemporary + } + return fmt.Errorf("%w: %s", ErrProviderTemporary, msg) + } +} + +// ===== session / admin-test generation (unchanged response shapes) ===== + func (s *V1Service) prepareSessionImage(ctx context.Context, principal *APIPrincipal, in V1ImageRequest) (map[string]any, error) { return s.prepareImageExecution(ctx, principal, in, "user", true) } @@ -425,19 +662,8 @@ func (s *V1Service) prepareAdminTestImage(ctx context.Context, principal *APIPri func (s *V1Service) prepareImageExecution(ctx context.Context, principal *APIPrincipal, in V1ImageRequest, source string, charge bool) (map[string]any, error) { // Detach the whole execution from the request lifecycle. The frontend tracks - // progress by polling /jobs/mine, so a client disconnect — or an nginx/CDN - // gateway timeout on the slow synchronous response — must NOT cancel an - // in-flight generation. Binding to the request ctx meant a cancelled request - // (a) spun uselessly in the upstream poll until its 180s timeout and - // (b) silently dropped the refund + final status write, leaving the row stuck - // pending until the maintenance sweep mislabeled it "abandoned". - // - // `ctx` (WithoutCancel) is durable and used for ALL bookkeeping (status / - // refund / cleanup) so those always land. `genCtx` is the cancellable WORK - // context: an 8-min backstop, AND registered in s.inflight so the maintenance - // sweep can cancel it the instant it abandons the row — stopping a stuck - // generation from running on for minutes and surfacing a late "success" on an - // already-abandoned event. + // progress by polling /jobs/mine, so a client disconnect must NOT cancel an + // in-flight generation or drop the refund/final-status bookkeeping. ctx = context.WithoutCancel(ctx) if source != "admin" { if err := s.checkBannedPrompt(ctx, principal, in.Prompt); err != nil { @@ -445,16 +671,12 @@ func (s *V1Service) prepareImageExecution(ctx context.Context, principal *APIPri return nil, err } } - // 去AI特征 is gated by a system-settings switch (default off) — drop the - // flag when disabled so no surcharge is charged and no processing runs. if in.DeAI && !s.deaiEnabled(ctx) { in.DeAI = false } - genCtx, cancel := context.WithTimeout(ctx, 8*time.Minute) + genCtx, cancel := context.WithTimeout(ctx, imageGenBudget) defer cancel() - // Per-user concurrency gate (画图台 + API key combined). Admin model-tests are - // exempt. Held for the whole generation; released on return. if source != "admin" && principal != nil && principal.User != nil { slot := randomUpper(12) if !s.userAcquire(ctx, principal.User, slot) { @@ -470,204 +692,43 @@ func (s *V1Service) prepareImageExecution(ctx context.Context, principal *APIPri return nil, err } refCount := len(in.ReferenceImages) - // API-key (source "v1") requests don't persist the output: we return the image - // as base64 inline (OpenAI gpt-image-1 also returns only b64_json) and never - // upload to RustFS, so there's no URL. The event is still logged (empty file) - // for usage; the customer logs page hides source="v1" rows. - noStore := source == "v1" - var fileURL, relativePath string - if !noStore { - fileURL, relativePath = s.allocateOutput(principal, "png", in.BaseURL) - } - // upstreamURL is the provider's original artifact URL. For API-key (source - // "v1") requests we return it instead of base64. When gatedURL is true the URL - // is auth-gated (chatgpt files.oaiusercontent.com — a plain GET 403s), so we - // store it on the event and hand the caller a proxy URL - // ({base}/v1/images/{eventID}/content) that re-fetches with the account token. - var upstreamURL string - var gatedURL bool + fileURL, relativePath := s.allocateOutput(principal, "png", in.BaseURL) eventID, err := s.logPendingEvent(ctx, "image", modelItem, principal, in.Prompt, aspectRatio, resolution, "", refCount, price, relativePath, source, nil, in.DeAI) if err != nil { return nil, err } - // Register so the maintenance sweep can cancel this generation if it abandons - // the row; deregister on return. s.inflight.Add(eventID, cancel) defer s.inflight.Done(eventID) startedAt := time.Now() - var imageBytes []byte - switch s.effectiveProvider(genCtx, modelItem) { - case "adobe": - b, u, execErr := s.generateAdobeImage(genCtx, eventID, modelItem, in, aspectRatio, resolution, noStore) - if execErr != nil { - _ = s.refundIfNeeded(ctx, principal, eventID, price) - _ = s.events.UpdateStatus(ctx, eventID, "failed", execErr.Error(), 0) - switch { - case errors.Is(execErr, adobe.ErrAuth): - return nil, ErrProviderAuth - case errors.Is(execErr, adobe.ErrQuotaExhausted): - return nil, ErrProviderQuota - case errors.Is(execErr, adobe.ErrTemporaryUpstream): - return nil, ErrProviderTemporary - default: - return nil, fmt.Errorf("%w: %v", ErrProviderExecution, execErr) - } - } - imageBytes = b - upstreamURL = u - case "chatgpt": - b, u, execErr := s.generateChatGPTImage(genCtx, eventID, modelItem, in, aspectRatio, resolution, noStore) - if execErr != nil { - _ = s.refundIfNeeded(ctx, principal, eventID, price) - _ = s.events.UpdateStatus(ctx, eventID, "failed", execErr.Error(), 0) - switch { - case errors.Is(execErr, chatgpt.ErrAuth): - return nil, ErrProviderAuth - case errors.Is(execErr, chatgpt.ErrQuotaExhausted): - return nil, ErrProviderQuota - case errors.Is(execErr, chatgpt.ErrTemporaryUpstream): - return nil, ErrProviderTemporary - default: - return nil, fmt.Errorf("%w: %v", ErrProviderExecution, execErr) - } - } - imageBytes = b - upstreamURL = u - gatedURL = true // chatgpt URL needs the account token → proxy it - case "leonardo": - b, u, execErr := s.generateLeonardoImage(genCtx, eventID, modelItem, in, aspectRatio, resolution, noStore) - if execErr != nil { - _ = s.refundIfNeeded(ctx, principal, eventID, price) - _ = s.events.UpdateStatus(ctx, eventID, "failed", execErr.Error(), 0) - switch { - case errors.Is(execErr, leonardo.ErrAuth): - return nil, ErrProviderAuth - case errors.Is(execErr, leonardo.ErrQuotaExhausted): - return nil, ErrProviderQuota - case errors.Is(execErr, leonardo.ErrTemporaryUpstream): - return nil, ErrProviderTemporary - default: - return nil, fmt.Errorf("%w: %v", ErrProviderExecution, execErr) - } - } - imageBytes = b - upstreamURL = u - case "krea": - b, u, execErr := s.generateKreaImage(genCtx, eventID, modelItem, in, aspectRatio, resolution, noStore) - if execErr != nil { - _ = s.refundIfNeeded(ctx, principal, eventID, price) - _ = s.events.UpdateStatus(ctx, eventID, "failed", execErr.Error(), 0) - switch { - case errors.Is(execErr, krea.ErrAuth): - return nil, ErrProviderAuth - case errors.Is(execErr, krea.ErrQuotaExhausted): - return nil, ErrProviderQuota - case errors.Is(execErr, krea.ErrTemporaryUpstream): - return nil, ErrProviderTemporary - default: - return nil, fmt.Errorf("%w: %v", ErrProviderExecution, execErr) - } - } - imageBytes = b - upstreamURL = u - case "imagine": - b, u, execErr := s.generateImagineImage(genCtx, eventID, modelItem, in, aspectRatio, resolution, noStore) - if execErr != nil { - _ = s.refundIfNeeded(ctx, principal, eventID, price) - _ = s.events.UpdateStatus(ctx, eventID, "failed", execErr.Error(), 0) - switch { - case errors.Is(execErr, imagine.ErrAuth): - return nil, ErrProviderAuth - case errors.Is(execErr, imagine.ErrQuotaExhausted): - return nil, ErrProviderQuota - case errors.Is(execErr, imagine.ErrTemporaryUpstream): - return nil, ErrProviderTemporary - default: - return nil, fmt.Errorf("%w: %v", ErrProviderExecution, execErr) - } - } - imageBytes = b - upstreamURL = u - case "grok": - b, u, execErr := s.generateGrokImage(genCtx, eventID, modelItem, in, aspectRatio, resolution, noStore) - if execErr != nil { - _ = s.refundIfNeeded(ctx, principal, eventID, price) - _ = s.events.UpdateStatus(ctx, eventID, "failed", execErr.Error(), 0) - switch { - case errors.Is(execErr, grok.ErrAuth): - return nil, ErrProviderAuth - case errors.Is(execErr, grok.ErrQuotaExhausted): - return nil, ErrProviderQuota - case errors.Is(execErr, grok.ErrTemporaryUpstream): - return nil, ErrProviderTemporary - default: - return nil, fmt.Errorf("%w: %v", ErrProviderExecution, execErr) - } - } - imageBytes = b - upstreamURL = u - case "runway": - b, u, execErr := s.generateRunwayImage(genCtx, eventID, modelItem, in, aspectRatio, resolution, noStore) - 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 - upstreamURL = u - case "custom": - b, u, execErr := s.generateCustomImage(genCtx, eventID, modelItem, in, aspectRatio, resolution, noStore) - if execErr != nil { - _ = s.refundIfNeeded(ctx, principal, eventID, price) - _ = s.events.UpdateStatus(ctx, eventID, "failed", execErr.Error(), 0) - switch { - case errors.Is(execErr, custom.ErrAuth): - return nil, ErrProviderAuth - case errors.Is(execErr, custom.ErrQuotaExhausted): - return nil, ErrProviderQuota - case errors.Is(execErr, custom.ErrTemporaryUpstream): - return nil, ErrProviderTemporary - default: - return nil, fmt.Errorf("%w: %v", ErrProviderExecution, execErr) - } - } - imageBytes = b - upstreamURL = u - default: + refs, err := decodeReferenceImages(in.ReferenceImages, max(totalRefLimit(modelItem), 1)) + if err != nil { _ = s.refundIfNeeded(ctx, principal, eventID, price) - _ = s.events.UpdateStatus(ctx, eventID, "failed", "provider not implemented", 0) - return nil, fmt.Errorf("%w: %s", ErrProviderUnsupported, modelItem.Provider) + _ = s.events.UpdateStatus(ctx, eventID, "failed", err.Error(), 0) + return nil, err } - // 去AI特征: post-process before storing/returning. Best-effort — a decode - // failure keeps the original bytes rather than failing a paid generation. + res, execErr := s.dispatchGenerate(genCtx, eventID, modelItem, generationParams{ + Kind: "image", Prompt: in.Prompt, Ratio: aspectRatio, Quality: resolution, Refs: refs, + PinChannelID: in.ChannelID, + }) + if execErr != nil { + _ = s.refundIfNeeded(ctx, principal, eventID, price) + _ = s.events.UpdateStatus(ctx, eventID, "failed", execErr.Error(), 0) + return nil, execErr + } + imageBytes := res.Data if in.DeAI { if processed, derr := applyDeAI(imageBytes); derr == nil { imageBytes = processed } } - if !noStore { - // Upload to RustFS. On failure the generation fails and credits are - // refunded — we never fall back to local disk. - if err := s.store.Put(genCtx, relativePath, imageBytes, "image/png"); err != nil { - _ = s.refundIfNeeded(ctx, principal, eventID, price) - _ = s.events.UpdateStatus(ctx, eventID, "failed", "storage upload failed: "+err.Error(), 0) - return nil, fmt.Errorf("%w: %v", ErrProviderExecution, err) - } - // Best-effort thumbnail for list views; the image serving route falls - // back to the original when the thumb object is missing. - if thumb, terr := makeThumbnail(imageBytes); terr == nil { - _ = s.store.Put(genCtx, ThumbKey(relativePath), thumb, "image/jpeg") - } + if err := s.store.Put(genCtx, relativePath, imageBytes, "image/png"); err != nil { + _ = s.refundIfNeeded(ctx, principal, eventID, price) + _ = s.events.UpdateStatus(ctx, eventID, "failed", "storage upload failed: "+err.Error(), 0) + return nil, fmt.Errorf("%w: %v", ErrProviderExecution, err) + } + if thumb, terr := makeThumbnail(imageBytes); terr == nil { + _ = s.store.Put(genCtx, ThumbKey(relativePath), thumb, "image/jpeg") } elapsedMS := int(time.Since(startedAt).Milliseconds()) if err := s.events.UpdateStatus(ctx, eventID, "success", "", elapsedMS); err != nil { @@ -680,51 +741,11 @@ func (s *V1Service) prepareImageExecution(ctx context.Context, principal *APIPri if charge { _ = s.maybeGrantInviteReward(ctx, principal) } - if noStore { - // Prefer the provider's original URL — return it directly, no base64. - // (API-key requests don't support DeAI, so there's no post-processing that - // would invalidate the upstream URL.) - if strings.TrimSpace(upstreamURL) != "" { - outURL := upstreamURL - if gatedURL { - // Auth-gated URL (chatgpt): store it on the event and return a proxy - // URL that re-fetches with the account token (see OpenImageContent). - _ = s.events.SetFile(ctx, eventID, upstreamURL) - if base := strings.TrimRight(strings.TrimSpace(in.BaseURL), "/"); base != "" { - outURL = base + "/v1/images/" + eventID + "/content" - } - } - return map[string]any{ - "created": time.Now().Unix(), - "data": []map[string]any{{"url": outURL}}, - "model": modelItem.EffectiveName(), - "provider": modelItem.Provider, - "kind": "image", - "url": outURL, - "elapsed_ms": elapsedMS, - "charged": price, - "credits": principalCredits(principal), - }, nil - } - // Fallback: providers without an upstream URL still return base64. - b64 := base64.StdEncoding.EncodeToString(imageBytes) - return map[string]any{ - "created": time.Now().Unix(), - "data": []map[string]any{{"b64_json": b64}}, - "model": modelItem.EffectiveName(), - "provider": modelItem.Provider, - "kind": "image", - "b64_json": b64, - "elapsed_ms": elapsedMS, - "charged": price, - "credits": principalCredits(principal), - }, nil - } return map[string]any{ "created": time.Now().Unix(), "data": []map[string]any{{"url": fileURL, "b64_json": nil}}, "model": modelItem.EffectiveName(), - "provider": modelItem.Provider, + "provider": res.ChannelName, "kind": "image", "url": fileURL, "elapsed_ms": elapsedMS, @@ -733,10 +754,6 @@ func (s *V1Service) prepareImageExecution(ctx context.Context, principal *APIPri }, nil } -func (s *V1Service) PrepareVideoRequest(ctx context.Context, principal *APIPrincipal, in V1VideoRequest) (map[string]any, error) { - return s.prepareVideoExecution(ctx, principal, in, "v1", true) -} - func (s *V1Service) prepareSessionVideo(ctx context.Context, principal *APIPrincipal, in V1VideoRequest) (map[string]any, error) { return s.prepareVideoExecution(ctx, principal, in, "user", true) } @@ -746,10 +763,6 @@ func (s *V1Service) prepareAdminTestVideo(ctx context.Context, principal *APIPri } func (s *V1Service) prepareVideoExecution(ctx context.Context, principal *APIPrincipal, in V1VideoRequest, source string, charge bool) (map[string]any, error) { - // Detach from the request lifecycle — see prepareImageExecution. `ctx` - // (WithoutCancel) carries all bookkeeping; `genCtx` is the cancellable work - // context (12-min backstop — video polls up to 10 min — and registered so the - // maintenance sweep can cancel a stuck render when it abandons the row). ctx = context.WithoutCancel(ctx) if source != "admin" { if err := s.checkBannedPrompt(ctx, principal, in.Prompt); err != nil { @@ -760,7 +773,6 @@ func (s *V1Service) prepareVideoExecution(ctx context.Context, principal *APIPri genCtx, cancel := context.WithTimeout(ctx, videoGenBudget) defer cancel() - // Per-user concurrency gate (画图台 + API key combined); admin tests exempt. if source != "admin" && principal != nil && principal.User != nil { slot := randomUpper(12) if !s.userAcquire(ctx, principal.User, slot) { @@ -776,82 +788,43 @@ func (s *V1Service) prepareVideoExecution(ctx context.Context, principal *APIPri return nil, err } refCount := len(in.ReferenceImages) - // API-key (source "v1") requests return base64 inline and never persist a - // file — see prepareImageExecution for the rationale. - noStore := source == "v1" - var fileURL, relativePath string - if !noStore { - fileURL, relativePath = s.allocateOutput(principal, "mp4", in.BaseURL) - } + fileURL, relativePath := s.allocateOutput(principal, "mp4", in.BaseURL) eventID, err := s.logPendingEvent(ctx, "video", modelItem, principal, in.Prompt, aspectRatio, resolution, duration, refCount, price, relativePath, source, nil, false) if err != nil { return nil, err } - // Register so the maintenance sweep can cancel this render if it abandons the - // row; deregister on return. s.inflight.Add(eventID, cancel) defer s.inflight.Done(eventID) startedAt := time.Now() - // API-key (noStore) requests return the upstream video URL directly. - // downloadResult=false skips the download. grok asset URLs are auth-gated - // (a plain GET 403s) → gatedVideoURL routes them through the /content proxy. - prov := s.effectiveProvider(genCtx, modelItem) - urlOnly := noStore - gatedVideoURL := prov == "grok" - var videoBytes []byte - var videoURL string - var execErr error - switch prov { - case "adobe": - videoBytes, videoURL, execErr = s.generateAdobeVideo(genCtx, eventID, modelItem, in, aspectRatio, resolution, parseDurationSeconds(duration), !urlOnly) - case "runway": - videoBytes, videoURL, execErr = s.generateRunwayVideo(genCtx, eventID, modelItem, in, aspectRatio, parseDurationSeconds(duration), !urlOnly) - case "grok": - videoBytes, videoURL, execErr = s.generateGrokVideo(genCtx, eventID, modelItem, in, aspectRatio, resolution, parseDurationSeconds(duration), !urlOnly) - case "leonardo": - videoBytes, videoURL, execErr = s.generateLeonardoVideo(genCtx, eventID, modelItem, in, aspectRatio, parseDurationSeconds(duration), !urlOnly) - case "custom": - videoBytes, videoURL, execErr = s.generateCustomVideo(genCtx, eventID, modelItem, in, aspectRatio, resolution, parseDurationSeconds(duration), !urlOnly) - case "creativefabrica": - videoBytes, videoURL, execErr = s.generateCreativeFabricaVideo(genCtx, eventID, modelItem, in, aspectRatio, !urlOnly) - default: + refs, err := decodeReferenceImages(in.ReferenceImages, max(totalRefLimit(modelItem), 1)) + if err != nil { _ = s.refundIfNeeded(ctx, principal, eventID, price) - _ = s.events.UpdateStatus(ctx, eventID, "failed", "provider not implemented", 0) - return nil, fmt.Errorf("%w: %s", ErrProviderUnsupported, modelItem.Provider) + _ = s.events.UpdateStatus(ctx, eventID, "failed", err.Error(), 0) + return nil, err } + res, execErr := s.dispatchGenerate(genCtx, eventID, modelItem, generationParams{ + Kind: "video", Prompt: in.Prompt, Ratio: aspectRatio, Quality: resolution, + DurationSecs: int64(parseDurationSeconds(duration)), Refs: refs, + PinChannelID: in.ChannelID, + }) if execErr != nil { _ = s.refundIfNeeded(ctx, principal, eventID, price) _ = s.events.UpdateStatus(ctx, eventID, "failed", execErr.Error(), 0) - switch { - case errors.Is(execErr, ErrNoProviderAccount): - return nil, ErrNoProviderAccount - case errors.Is(execErr, adobe.ErrAuth), errors.Is(execErr, runway.ErrAuth), errors.Is(execErr, grok.ErrAuth), errors.Is(execErr, leonardo.ErrAuth), errors.Is(execErr, custom.ErrAuth), errors.Is(execErr, creativefabrica.ErrAuth): - return nil, ErrProviderAuth - case errors.Is(execErr, adobe.ErrQuotaExhausted), errors.Is(execErr, runway.ErrQuotaExhausted), errors.Is(execErr, grok.ErrQuotaExhausted), errors.Is(execErr, leonardo.ErrQuotaExhausted), errors.Is(execErr, custom.ErrQuotaExhausted), errors.Is(execErr, creativefabrica.ErrQuotaExhausted): - return nil, ErrProviderQuota - case errors.Is(execErr, adobe.ErrTemporaryUpstream), errors.Is(execErr, runway.ErrTemporaryUpstream), errors.Is(execErr, grok.ErrTemporaryUpstream), errors.Is(execErr, leonardo.ErrTemporaryUpstream), errors.Is(execErr, custom.ErrTemporaryUpstream), errors.Is(execErr, creativefabrica.ErrTemporaryUpstream): - return nil, ErrProviderTemporary - default: - return nil, fmt.Errorf("%w: %v", ErrProviderExecution, execErr) - } + return nil, execErr } - if !noStore { - if err := s.store.Put(genCtx, relativePath, videoBytes, "video/mp4"); err != nil { - _ = s.refundIfNeeded(ctx, principal, eventID, price) - _ = s.events.UpdateStatus(ctx, eventID, "failed", "storage upload failed: "+err.Error(), 0) - return nil, fmt.Errorf("%w: %v", ErrProviderExecution, err) + videoBytes := res.Data + if err := s.store.Put(genCtx, relativePath, videoBytes, "video/mp4"); err != nil { + _ = s.refundIfNeeded(ctx, principal, eventID, price) + _ = s.events.UpdateStatus(ctx, eventID, "failed", "storage upload failed: "+err.Error(), 0) + return nil, fmt.Errorf("%w: %v", ErrProviderExecution, err) + } + if thumb, last, terr := extractVideoFrames(genCtx, videoBytes); terr == nil { + if len(thumb) > 0 { + _ = s.store.Put(genCtx, ThumbKey(relativePath), thumb, "image/jpeg") } - // Best-effort stills: first frame (downscaled) for list thumbnails and - // the full-res last frame for 首尾帧 continuation. Missing objects fall - // back to the video itself at serve time. - if thumb, last, terr := extractVideoFrames(genCtx, videoBytes); terr == nil { - if len(thumb) > 0 { - _ = s.store.Put(genCtx, ThumbKey(relativePath), thumb, "image/jpeg") - } - if len(last) > 0 { - _ = s.store.Put(genCtx, LastFrameKey(relativePath), last, "image/jpeg") - } + if len(last) > 0 { + _ = s.store.Put(genCtx, LastFrameKey(relativePath), last, "image/jpeg") } } elapsedMS := int(time.Since(startedAt).Milliseconds()) @@ -865,47 +838,11 @@ func (s *V1Service) prepareVideoExecution(ctx context.Context, principal *APIPri if charge { _ = s.maybeGrantInviteReward(ctx, principal) } - if noStore && strings.TrimSpace(videoURL) != "" { - // Return the upstream video URL. grok URLs are auth-gated → store on the - // event and hand back the /content proxy (re-fetches with the account token). - outURL := videoURL - if gatedVideoURL { - _ = s.events.SetFile(ctx, eventID, videoURL) - if base := strings.TrimRight(strings.TrimSpace(in.BaseURL), "/"); base != "" { - outURL = base + "/v1/videos/" + eventID + "/content" - } - } - return map[string]any{ - "created": time.Now().Unix(), - "data": []map[string]any{{"url": outURL}}, - "model": modelItem.EffectiveName(), - "provider": modelItem.Provider, - "kind": "video", - "url": outURL, - "elapsed_ms": elapsedMS, - "charged": price, - "credits": principalCredits(principal), - }, nil - } - if noStore { - b64 := base64.StdEncoding.EncodeToString(videoBytes) - return map[string]any{ - "created": time.Now().Unix(), - "data": []map[string]any{{"b64_json": b64}}, - "model": modelItem.EffectiveName(), - "provider": modelItem.Provider, - "kind": "video", - "b64_json": b64, - "elapsed_ms": elapsedMS, - "charged": price, - "credits": principalCredits(principal), - }, nil - } return map[string]any{ "created": time.Now().Unix(), "data": []map[string]any{{"url": fileURL}}, "model": modelItem.EffectiveName(), - "provider": modelItem.Provider, + "provider": res.ChannelName, "kind": "video", "url": fileURL, "elapsed_ms": elapsedMS, @@ -914,107 +851,198 @@ func (s *V1Service) prepareVideoExecution(ctx context.Context, principal *APIPri }, nil } -// ===== /v1/videos — OpenAI Sora-style async jobs ===== -// POST /v1/videos charges + creates a pending event and renders in the -// background; the render captures only the UPSTREAM video URL (no download, no -// RustFS). GET /v1/videos/{id} polls status; /content proxies the upstream URL. +// ===== /v1/generate jobs (vivid protocol) ===== -// StartVideoJob validates+charges, creates the job event, kicks the render off in -// the background, and returns the OpenAI video object (status "queued"). -func (s *V1Service) StartVideoJob(ctx context.Context, principal *APIPrincipal, in V1VideoRequest) (map[string]any, error) { +// V1JobRequest is the POST /v1/generate body. +type V1JobRequest struct { + Model string `json:"model"` + Prompt string `json:"prompt"` + Quality string `json:"quality"` + Ratio *string `json:"ratio"` + Duration *int64 `json:"duration"` + // Refs are object keys issued by /v1/uploads. + Refs []string `json:"refs"` + // JobId reconnect mode: re-attach to an existing job, never creates one. + JobID *string `json:"jobId"` +} + +// CreateJob validates + charges + creates the pending event and kicks the +// background render; returns the event row (its ID is the jobId). +func (s *V1Service) CreateJob(ctx context.Context, principal *APIPrincipal, in V1JobRequest) (*model.EventLog, error) { ctx = context.WithoutCancel(ctx) if err := s.checkBannedPrompt(ctx, principal, in.Prompt); err != nil { - s.logRejectedEvent(ctx, "video", in.Model, principal, in.Prompt, "v1", err.Error()) + s.logRejectedEvent(ctx, "", in.Model, principal, in.Prompt, "v1", err.Error()) return nil, err } - // Validate reference_mode against model capabilities and reference count - // BEFORE charging — a bad override must fail fast with no debit, never - // charge-then-reject (which would silently eat the user's credits). modelItem, err := s.models.Get(ctx, strings.TrimSpace(in.Model)) if err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { - s.logRejectedEvent(ctx, "video", in.Model, principal, in.Prompt, "v1", ErrUnknownModel.Error()) return nil, ErrUnknownModel } return nil, err } - if rm := strings.TrimSpace(in.ReferenceMode); rm != "" { - if err := validateReferenceMode(rm, modelItem, len(in.ReferenceImages)); err != nil { - s.logRejectedEvent(ctx, "video", in.Model, principal, in.Prompt, "v1", err.Error()) + if !modelItem.Enabled { + return nil, ErrUnknownModel + } + kind := modelItem.Type + if kind != "image" && kind != "video" { + return nil, ErrUnknownModel + } + if strings.TrimSpace(in.Prompt) == "" { + return nil, fmt.Errorf("%w: prompt required", ErrUnsupportedParams) + } + quality := strings.TrimSpace(in.Quality) + if quality == "" { + return nil, fmt.Errorf("%w: quality required", ErrUnsupportedParams) + } + if res := repo.JSONStrings(modelItem.Resolutions); len(res) > 0 && !strSliceContains(res, quality) { + return nil, fmt.Errorf("%w: quality %q not in %v", ErrUnsupportedParams, quality, res) + } + ratio := "" + if in.Ratio != nil { + ratio = snapRatio(strings.TrimSpace(strings.ReplaceAll(*in.Ratio, "x", ":")), repo.JSONStrings(modelItem.Ratios)) + } + var durSecs int64 + duration := "" + if kind == "video" { + if in.Duration == nil || *in.Duration <= 0 { + return nil, fmt.Errorf("%w: duration required", ErrUnsupportedParams) + } + durSecs = *in.Duration + if err := validateModelDuration(modelItem, durSecs); err != nil { return nil, err } - if rm == modelItem.ReferenceMode { - in.ReferenceMode = "" // same as default, don't override - } + duration = fmt.Sprintf("%ds", durSecs) } - modelItem, resolution, aspectRatio, duration, price, err := s.prepareVideo(ctx, principal, in, true) + // Refs: keys must live under the caller's refs/ dir; bytes are fetched now + // (existence + caps + duration checks) and re-uploaded to whatever channel + // wins the dispatch. + refBlobs, refKeys, err := s.resolveJobRefs(ctx, principal, in.Refs, modelItem) if err != nil { - s.logRejectedEvent(ctx, "video", in.Model, principal, in.Prompt, "v1", err.Error()) + s.logRejectedEvent(ctx, kind, modelItem.ID, principal, in.Prompt, "v1", err.Error()) return nil, err } - // Source "v1": no output file is allocated — the result is the upstream URL, - // stored on the event when the render completes. - eventID, err := s.logPendingEvent(ctx, "video", modelItem, principal, in.Prompt, aspectRatio, resolution, duration, len(in.ReferenceImages), price, "", "v1", nil, false) + // The job runs async — acquire the per-user slot now and release it when the + // job settles, so the slot survives past this HTTP request. + slot := randomUpper(12) + if principal != nil && principal.User != nil && !s.userAcquire(ctx, principal.User, slot) { + s.logRejectedEvent(ctx, kind, modelItem.ID, principal, in.Prompt, "v1", ErrUserConcurrencyFull.Error()) + return nil, ErrUserConcurrencyFull + } + price, err := s.chargeForModel(ctx, principal, modelItem, kind, quality, duration, 0, true) if err != nil { + s.userRelease(ctx, principal.User.ID, slot) + s.logRejectedEvent(ctx, kind, modelItem.ID, principal, in.Prompt, "v1", err.Error()) return nil, err } - go s.runVideoJob(ctx, principal, in, modelItem, eventID, aspectRatio, resolution, duration, price) - return videoJobObject(eventID, modelItem.EffectiveName(), "queued", 0, duration, sizeFromRatioRes(aspectRatio, resolution), time.Now().Unix(), 0, ""), nil + eventID, err := s.logPendingEvent(ctx, kind, modelItem, principal, in.Prompt, ratio, quality, duration, len(refBlobs), price, "", "v1", nil, false) + if err != nil { + s.userRelease(ctx, principal.User.ID, slot) + _, _ = s.users.RefundCredits(ctx, principal.User.ID, price) + return nil, err + } + go s.runV1Job(ctx, principal, eventID, modelItem, generationParams{ + Kind: kind, Prompt: in.Prompt, Ratio: ratio, Quality: quality, + DurationSecs: durSecs, Refs: refBlobs, + }, refKeys, slot, price) + return s.events.GetByID(ctx, eventID) } -// runVideoJob renders the clip in the background, capturing the upstream URL -// (downloadResult=false → no bytes, no RustFS) and storing it on the event. -func (s *V1Service) runVideoJob(ctx context.Context, principal *APIPrincipal, in V1VideoRequest, modelItem *model.ModelConfig, eventID, aspectRatio, resolution, duration string, price float64) { - genCtx, cancel := context.WithTimeout(ctx, videoGenBudget) +// JobForUser reloads a job scoped to the caller (reconnect mode + wait loop). +func (s *V1Service) JobForUser(ctx context.Context, principal *APIPrincipal, jobID string) (*model.EventLog, error) { + ev, err := s.events.GetByID(ctx, strings.TrimSpace(jobID)) + if err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, ErrVideoJobNotFound + } + return nil, err + } + if ev == nil || (principal != nil && principal.User != nil && ev.UserID != principal.User.ID) { + return nil, ErrVideoJobNotFound + } + return ev, nil +} + +// JobStatus maps an event to the vivid wire status. +func JobStatus(ev *model.EventLog) string { + switch ev.Status { + case "success": + return "succeeded" + case "failed": + return "failed" + default: + if strings.TrimSpace(ev.AccountID) != "" { + return "running" + } + return "queued" + } +} + +// JobResultURL presigns the stored result for GET (30 min). API-key results +// expire: the URL is only minted within 30 minutes of completion. +func (s *V1Service) JobResultURL(ev *model.EventLog) string { + if ev == nil || ev.Status != "success" || strings.TrimSpace(ev.File) == "" { + return "" + } + if s.store == nil || !s.store.Configured() { + return "" + } + if time.Since(ev.UpdatedAt) > PresignTTL { + return "" + } + url, err := s.store.PresignGet(ev.File, PresignTTL) + if err != nil { + return "" + } + return url +} + +// runV1Job executes a /v1 job in the background: dispatch → store under tmp/ → +// final event state. Caller polls the row; the NDJSON stream just re-reads it. +// `slot` is the per-user concurrency token acquired in CreateJob. +func (s *V1Service) runV1Job(ctx context.Context, principal *APIPrincipal, eventID string, modelItem *model.ModelConfig, p generationParams, refKeys []string, slot string, price float64) { + budget := imageGenBudget + if p.Kind == "video" { + budget = videoGenBudget + } + genCtx, cancel := context.WithTimeout(ctx, budget) defer cancel() s.inflight.Add(eventID, cancel) defer s.inflight.Done(eventID) + if principal != nil && principal.User != nil { + defer s.userRelease(ctx, principal.User.ID, slot) + } + // Reference uploads are session-scoped on OUR side too — drop them once the + // job reaches a terminal state. + defer func() { + for _, k := range refKeys { + _ = s.store.Delete(context.Background(), k) + } + }() startedAt := time.Now() - - // A panic in provider fan-out must fail the event (and refund), not - // take down the whole API process with it. defer func() { if r := recover(); r != nil { - log.Printf("[v1] panic in runVideoJob event=%s model=%s: %v\n%s", eventID, modelItem.Name, r, debug.Stack()) + log.Printf("[v1] panic in runV1Job event=%s model=%s: %v\n%s", eventID, modelItem.Name, r, debug.Stack()) _ = s.refundIfNeeded(ctx, principal, eventID, price) _ = s.events.UpdateStatus(ctx, eventID, "failed", fmt.Sprintf("internal error: %v", r), 0) } }() - // No-store: capture only the UPSTREAM video URL. /content streams it on demand - // (grok URLs are auth-gated → fetched with the generating account's token). - var videoURL string - var execErr error - switch s.effectiveProvider(genCtx, modelItem) { - case "adobe": - _, 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) - case "leonardo": - _, videoURL, execErr = s.generateLeonardoVideo(genCtx, eventID, modelItem, in, aspectRatio, parseDurationSeconds(duration), false) - case "custom": - _, videoURL, execErr = s.generateCustomVideo(genCtx, eventID, modelItem, in, aspectRatio, resolution, parseDurationSeconds(duration), false) - case "creativefabrica": - _, videoURL, execErr = s.generateCreativeFabricaVideo(genCtx, eventID, modelItem, in, aspectRatio, false) - default: - _ = s.refundIfNeeded(ctx, principal, eventID, price) - _ = s.events.UpdateStatus(ctx, eventID, "failed", "provider not implemented", 0) - return - } + res, execErr := s.dispatchGenerate(genCtx, eventID, modelItem, p) if execErr != nil { _ = s.refundIfNeeded(ctx, principal, eventID, price) _ = s.events.UpdateStatus(ctx, eventID, "failed", execErr.Error(), 0) return } - if strings.TrimSpace(videoURL) == "" { + ext := resultExt(res.ContentType, p.Kind) + rel := filepath.ToSlash(filepath.Join("tmp", eventID+"."+ext)) + if err := s.store.Put(genCtx, rel, res.Data, res.ContentType); err != nil { _ = s.refundIfNeeded(ctx, principal, eventID, price) - _ = s.events.UpdateStatus(ctx, eventID, "failed", "upstream returned no video url", 0) + _ = s.events.UpdateStatus(ctx, eventID, "failed", "storage upload failed: "+err.Error(), 0) return } - // Store the upstream URL as the event's "file"; /content fetches it on demand. - if err := s.events.MarkVideoReady(ctx, eventID, videoURL, int(time.Since(startedAt).Milliseconds())); err != nil { + _ = s.events.SetFile(ctx, eventID, rel) + if err := s.events.UpdateStatus(ctx, eventID, "success", "", int(time.Since(startedAt).Milliseconds())); err != nil { return } _ = s.models.IncrementGenerationCount(ctx, modelItem.ID) @@ -1024,226 +1052,241 @@ func (s *V1Service) runVideoJob(ctx context.Context, principal *APIPrincipal, in _ = s.maybeGrantInviteReward(ctx, principal) } -// VideoJob returns the OpenAI video object for a job, scoped to the caller. -func (s *V1Service) VideoJob(ctx context.Context, principal *APIPrincipal, id string) (map[string]any, error) { - ev, err := s.videoEventForUser(ctx, principal, id) - if err != nil { - return nil, err - } - status, progress := videoJobStatus(ev) - completedAt := int64(0) - if ev.Status == "success" || ev.Status == "failed" { - completedAt = ev.UpdatedAt.Unix() - } - errMsg := "" - if ev.Status == "failed" { - errMsg = ev.Error - } - modelName := ev.Model - if nameByID, nerr := s.models.NameMap(ctx); nerr == nil { - if name, ok := nameByID[ev.Model]; ok && strings.TrimSpace(name) != "" { - modelName = name - } - } - return videoJobObject(ev.ID, modelName, status, progress, ev.Duration, sizeFromRatioRes(ev.Ratio, ev.Resolution), ev.TS.Unix(), completedAt, errMsg), nil -} - -// OpenVideoContent streams a completed job's video by proxying the stored -// upstream URL (downloaded on demand — never persisted). -func (s *V1Service) OpenVideoContent(ctx context.Context, principal *APIPrincipal, id string) (io.ReadCloser, string, error) { - ev, err := s.videoEventForUser(ctx, principal, id) - if err != nil { - return nil, "", err - } - if ev.Status != "success" || strings.TrimSpace(ev.File) == "" { - return nil, "", ErrVideoNotReady - } - // grok asset URLs (assets.grok.com) are auth-gated — a plain GET 403s. Stream - // them through the SAME account that generated the clip, using its token. If - // that account is gone (grok pools churn often), the clip is unrecoverable. - if ev.Provider == "grok" && s.grok != nil { - if s.settings != nil { - if proxy, perr := s.settings.GetValue(ctx, "proxy.url"); perr == nil { - s.grok.SetProxy(proxy) - } - } - acct, _ := s.tokens.Get(ctx, "grok", ev.AccountID) - if acct == nil || strings.TrimSpace(acct.Value) == "" { - return nil, "", fmt.Errorf("%w: grok account no longer available for this video", ErrProviderTemporary) - } - return s.grok.OpenAsset(ctx, acct.Value, ev.File) - } - // Other providers return publicly-fetchable URLs — proxy directly. - req, err := http.NewRequestWithContext(ctx, http.MethodGet, ev.File, nil) - if err != nil { - return nil, "", err - } - resp, err := (&http.Client{Timeout: 5 * time.Minute}).Do(req) - if err != nil { - return nil, "", fmt.Errorf("%w: fetch upstream video: %v", ErrProviderTemporary, err) - } - if resp.StatusCode != http.StatusOK { - resp.Body.Close() - return nil, "", fmt.Errorf("%w: upstream video status %d", ErrProviderTemporary, resp.StatusCode) - } - ct := strings.TrimSpace(resp.Header.Get("Content-Type")) - if ct == "" { - ct = "video/mp4" - } - return resp.Body, ct, nil -} - -// OpenImageContent streams a no-store image by proxying the stored upstream URL. -// chatgpt URLs are auth-gated (files.oaiusercontent.com — a plain GET 403s), so -// they're fetched through the generating account's token; other providers' -// URLs are public and proxied directly. Never persisted. -func (s *V1Service) OpenImageContent(ctx context.Context, principal *APIPrincipal, id string) (io.ReadCloser, string, error) { - ev, err := s.events.GetByID(ctx, strings.TrimSpace(id)) - if err != nil { - return nil, "", err - } - if ev == nil || ev.Kind != "image" { - return nil, "", ErrVideoJobNotFound - } - if principal != nil && principal.User != nil && ev.UserID != principal.User.ID { - return nil, "", ErrVideoJobNotFound - } - if ev.Status != "success" || strings.TrimSpace(ev.File) == "" { - return nil, "", ErrVideoNotReady - } - if ev.Provider == "chatgpt" && s.chatgpt != nil { - if s.settings != nil { - if proxy, perr := s.settings.GetValue(ctx, "proxy.url"); perr == nil { - s.chatgpt.SetProxy(proxy) - } - } - acct, _ := s.tokens.Get(ctx, "chatgpt", ev.AccountID) - if acct == nil || strings.TrimSpace(acct.Value) == "" { - return nil, "", fmt.Errorf("%w: chatgpt account no longer available for this image", ErrProviderTemporary) - } - return s.chatgpt.OpenAsset(ctx, acct.Value, ev.File) - } - req, err := http.NewRequestWithContext(ctx, http.MethodGet, ev.File, nil) - if err != nil { - return nil, "", err - } - resp, err := (&http.Client{Timeout: 5 * time.Minute}).Do(req) - if err != nil { - return nil, "", fmt.Errorf("%w: fetch upstream image: %v", ErrProviderTemporary, err) - } - if resp.StatusCode != http.StatusOK { - resp.Body.Close() - return nil, "", fmt.Errorf("%w: upstream image status %d", ErrProviderTemporary, resp.StatusCode) - } - ct := strings.TrimSpace(resp.Header.Get("Content-Type")) - if ct == "" { - ct = "image/png" - } - return resp.Body, ct, nil -} - -func (s *V1Service) videoEventForUser(ctx context.Context, principal *APIPrincipal, id string) (*model.EventLog, error) { - ev, err := s.events.GetByID(ctx, strings.TrimSpace(id)) - if err != nil { - return nil, err - } - if ev == nil || ev.Kind != "video" { - return nil, ErrVideoJobNotFound - } - if principal != nil && principal.User != nil && ev.UserID != principal.User.ID { - return nil, ErrVideoJobNotFound - } - return ev, nil -} - -// videoJobStatus maps our event status → OpenAI's (queued|in_progress|completed| -// failed) plus a coarse progress. -func videoJobStatus(ev *model.EventLog) (string, int) { - switch ev.Status { - case "success": - return "completed", 100 - case "failed": - return "failed", 0 - default: - if strings.TrimSpace(ev.AccountID) != "" { - return "in_progress", 50 - } - return "queued", 0 - } -} - -func videoJobObject(id, modelID, status string, progress int, seconds, size string, createdAt, completedAt int64, errMsg string) map[string]any { - obj := map[string]any{ - "id": id, - "object": "video", - "model": modelID, - "status": status, - "progress": progress, - "created_at": createdAt, - "size": size, - "seconds": strings.TrimSuffix(strings.TrimSpace(seconds), "s"), - } - if completedAt > 0 { - obj["completed_at"] = completedAt - } else { - obj["completed_at"] = nil - } - if errMsg != "" { - obj["error"] = map[string]any{"message": errMsg} - } else { - obj["error"] = nil - } - return obj -} - -// sizeFromRatioRes reconstructs an OpenAI-style "WxH" label from our stored ratio -// + resolution tier (best-effort; only for display in the job object). -func sizeFromRatioRes(ratio, resolution string) string { - long := 720 - res := strings.ToUpper(resolution) +func resultExt(contentType, kind string) string { switch { - case strings.Contains(res, "1080") || strings.Contains(res, "2K"): - long = 1080 - case strings.Contains(res, "4K") || strings.Contains(res, "2160"): - long = 2160 + case strings.Contains(contentType, "webp"): + return "webp" + case strings.Contains(contentType, "jpeg"), strings.Contains(contentType, "jpg"): + return "jpg" + case strings.Contains(contentType, "gif"): + return "gif" + case strings.Contains(contentType, "webm"): + return "webm" + case strings.Contains(contentType, "mp4"): + return "mp4" + case strings.Contains(contentType, "quicktime"): + return "mov" } - w, h := long, long - switch strings.TrimSpace(ratio) { - case "16:9": - w, h = long, long*9/16 - case "9:16": - w, h = long*9/16, long - case "4:3": - w, h = long, long*3/4 - case "3:4": - w, h = long*3/4, long - case "1:1": - w, h = long, long - default: - w, h = long, long*9/16 + if kind == "video" { + return "mp4" } - return fmt.Sprintf("%dx%d", w, h) + return "png" } -// hasActiveProviderToken reports whether the provider pool holds at least one -// usable token for this kind of generation — mirrors the selection filter in -// the generate* paths. Used to fail fast (before charging / creating a job) -// with a clear "no account" error instead of dialing upstream with no token. -func (s *V1Service) hasActiveProviderToken(ctx context.Context, provider, kind string) (bool, error) { - items, err := s.tokens.ListByPool(ctx, provider) - if err != nil { - return false, err - } - for _, item := range items { - if item.Status != "active" || item.Dead || strings.TrimSpace(item.Value) == "" { - continue +// validateModelDuration enforces the model's duration spec (options list or +// min/max range) on a /v1 request. +func validateModelDuration(item *model.ModelConfig, secs int64) error { + if item.DurationMin > 0 || item.DurationMax > 0 { + if item.DurationMin > 0 && secs < int64(item.DurationMin) || + item.DurationMax > 0 && secs > int64(item.DurationMax) { + return fmt.Errorf("%w: duration %ds out of range [%d-%d]", ErrUnsupportedParams, secs, item.DurationMin, item.DurationMax) } - // Adobe accounts are credit-based (积分号) — no per-kind quota locks. - return true, nil + return nil } - return false, nil + if durs := repo.JSONStrings(item.Durations); len(durs) > 0 { + for _, d := range durs { + if int64(parseDurationSeconds(d)) == secs { + return nil + } + } + return fmt.Errorf("%w: duration %ds not in %v", ErrUnsupportedParams, secs, durs) + } + return nil } +// resolveJobRefs fetches + validates /v1 ref keys. Returns blobs (for channel +// upload) and the raw keys (for cleanup on job completion). +func (s *V1Service) resolveJobRefs(ctx context.Context, principal *APIPrincipal, keys []string, modelItem *model.ModelConfig) ([][]byte, []string, error) { + if len(keys) == 0 { + return nil, nil, nil + } + if principal == nil || principal.User == nil { + return nil, nil, ErrInvalidAPIKey + } + uid := principal.User.ID + total := len(keys) + if cap := totalRefLimit(modelItem); cap > 0 && total > cap { + return nil, nil, fmt.Errorf("%w: too many reference files", ErrUnsupportedParams) + } + blobs := make([][]byte, 0, total) + counts := map[string]int{} + secSum := map[string]float64{} + for _, key := range keys { + key = strings.TrimSpace(key) + if !validRefKey(uid, key) { + return nil, nil, fmt.Errorf("%w: 参考素材无效", ErrUnsupportedParams) + } + data, err := s.getObjectBytes(ctx, key) + if err != nil { + return nil, nil, fmt.Errorf("%w: 参考素材已过期或不存在", ErrUnsupportedParams) + } + kind, mime := upstream.RefKindMime(data) + if kind == "" { + return nil, nil, fmt.Errorf("%w: 参考素材格式不支持", ErrUnsupportedParams) + } + counts[kind]++ + limit := refLimitFor(modelItem, kind) + if counts[kind] > limit { + return nil, nil, fmt.Errorf("%w: 参考%s最多 %d 个", ErrUnsupportedParams, refKindLabel(kind), limit) + } + if kind == "video" || kind == "audio" { + if sec, ok := util.MediaDurationSeconds(mime, data[:minInt(len(data), 256<<10)]); ok { + secSum[kind] += sec + } + } + blobs = append(blobs, data) + } + if lim := modelItem.MaxRefVideoSeconds; lim > 0 && secSum["video"] > float64(lim) { + return nil, nil, fmt.Errorf("%w: 参考视频总时长 %.0f 秒超过上限 %d 秒", ErrUnsupportedParams, secSum["video"], lim) + } + if lim := modelItem.MaxRefAudioSeconds; lim > 0 && secSum["audio"] > float64(lim) { + return nil, nil, fmt.Errorf("%w: 参考音频总时长 %.0f 秒超过上限 %d 秒", ErrUnsupportedParams, secSum["audio"], lim) + } + return blobs, keys, nil +} + +func refLimitFor(item *model.ModelConfig, kind string) int { + switch kind { + case "video": + return item.MaxReferenceVideos + case "audio": + return item.MaxReferenceAudios + default: + return item.MaxReferenceImages + } +} + +func refKindLabel(kind string) string { + switch kind { + case "video": + return "视频" + case "audio": + return "音频" + default: + return "图" + } +} + +// validRefKey checks ownership + shape: /refs//.. +func validRefKey(uid, key string) bool { + if uid == "" || !strings.HasPrefix(key, uid+"/refs/") { + return false + } + rest := strings.TrimPrefix(key, uid+"/refs/") + parts := strings.Split(rest, "/") + if len(parts) != 2 { + return false + } + switch parts[0] { + case "image", "video", "audio": + default: + return false + } + name := parts[1] + dot := strings.LastIndexByte(name, '.') + if dot <= 0 || dot == len(name)-1 { + return false + } + switch name[dot+1:] { + case "png", "jpg", "webp", "mp4", "webm", "mov", "mp3", "wav", "m4a", "ogg": + return true + } + return false +} + +func (s *V1Service) getObjectBytes(ctx context.Context, key string) ([]byte, error) { + resp, err := s.store.Get(ctx, key, "") + if err != nil { + return nil, err + } + defer resp.Body.Close() + if resp.StatusCode/100 != 2 { + return nil, fmt.Errorf("status %d", resp.StatusCode) + } + return io.ReadAll(io.LimitReader(resp.Body, maxReferenceImageBytes+1)) +} + +// PresignRefUpload — POST /v1/uploads. Returns the object key + presigned PUT +// URL the caller uploads bytes to within PresignTTL. +type V1UploadView struct { + UUID string `json:"uuid"` + Key string `json:"key"` + URL string `json:"url"` + ExpiresAt int64 `json:"expiresAt"` +} + +func (s *V1Service) PresignRefUpload(ctx context.Context, principal *APIPrincipal, kind, mime string) (*V1UploadView, error) { + if principal == nil || principal.User == nil { + return nil, ErrInvalidAPIKey + } + kind = strings.ToLower(strings.TrimSpace(kind)) + mime = strings.ToLower(strings.TrimSpace(mime)) + if !presignRefMimeOK(kind, mime) { + return nil, fmt.Errorf("%w: 参考素材类型或格式不支持", ErrUnsupportedParams) + } + if s.store == nil || !s.store.Configured() { + return nil, fmt.Errorf("%w: 对象存储不可用", ErrProviderTemporary) + } + id := "r" + strings.ToLower(randomUpper(16)) + key := fmt.Sprintf("%s/refs/%s/%s.%s", principal.User.ID, kind, id, refExtOfMime(mime)) + url, err := s.store.PresignPut(key, mime, PresignTTL) + if err != nil { + return nil, fmt.Errorf("%w: 预签名签发失败", ErrProviderTemporary) + } + return &V1UploadView{UUID: id, Key: key, URL: url, ExpiresAt: time.Now().Add(PresignTTL).Unix()}, nil +} + +func presignRefMimeOK(kind, mime string) bool { + if !strings.HasPrefix(mime, kind+"/") { + return false + } + switch mime { + case "image/png", "image/jpeg", "image/webp", + "video/mp4", "video/webm", "video/quicktime", + "audio/mpeg", "audio/mp3", "audio/wav", "audio/x-wav", "audio/mp4", "audio/m4a", "audio/ogg": + return true + } + return false +} + +func refExtOfMime(mime string) string { + switch mime { + case "image/png": + return "png" + case "image/jpeg": + return "jpg" + case "image/webp": + return "webp" + case "video/mp4": + return "mp4" + case "video/webm": + return "webm" + case "video/quicktime": + return "mov" + case "audio/mpeg", "audio/mp3": + return "mp3" + case "audio/wav", "audio/x-wav": + return "wav" + case "audio/mp4", "audio/m4a": + return "m4a" + case "audio/ogg": + return "ogg" + default: + return "bin" + } +} + +func strSliceContains(list []string, v string) bool { + for _, item := range list { + if strings.TrimSpace(item) == v { + return true + } + } + return false +} + +// ===== pricing / charge / refund (unchanged semantics) ===== + func (s *V1Service) prepareImage(ctx context.Context, principal *APIPrincipal, in V1ImageRequest, charge bool) (*model.ModelConfig, string, string, float64, error) { modelID := strings.TrimSpace(in.Model) prompt := strings.TrimSpace(in.Prompt) @@ -1260,16 +1303,11 @@ func (s *V1Service) prepareImage(ctx context.Context, principal *APIPrincipal, i if !modelItem.Enabled || modelItem.Type != "image" { return nil, "", "", 0, ErrUnknownModel } - // Fail fast before charging if the provider has no usable account. Use the - // effective provider: a custom upstream serving this model id routes to - // "custom" (effectiveProvider only returns it when such an account exists, so - // the precheck is satisfied); otherwise check the native provider pool. - if eff := s.effectiveProvider(ctx, modelItem); eff != "custom" { - if ok, err := s.hasActiveProviderToken(ctx, eff, "image"); err != nil { - return nil, "", "", 0, err - } else if !ok { - return nil, "", "", 0, ErrNoProviderAccount - } + // Fail fast before charging if no channel serves this model. + if ok, err := s.hasChannelForModel(ctx, modelItem.ID); err != nil { + return nil, "", "", 0, err + } else if !ok { + return nil, "", "", 0, ErrNoProviderAccount } refLimit := 0 if modelItem.ImageToImage { @@ -1281,23 +1319,15 @@ func (s *V1Service) prepareImage(ctx context.Context, principal *APIPrincipal, i if len(in.ReferenceImages) > refLimit { return nil, "", "", 0, errors.New("too many reference images") } - // Reject oversized reference images before charging (all providers, all paths). if err := ensureReferenceSizes(in.ReferenceImages); err != nil { return nil, "", "", 0, err } - // `size` (WxH) drives BOTH the aspect ratio AND the resolution tier — its long - // edge maps to a tier (<1800→1K, 1800–3499→2K, ≥3500→4K). The web path passes - // an explicit resolution; the OpenAI /v1 path derives it from size. There is no - // `quality` param — size is the single source of truth for resolution. aspectRatio, resolution := parseImageSize(in.Size, in.AspectRatio, in.Resolution) - // Snap to the nearest ratio the model actually supports — a `size`-derived - // ratio (e.g. 1:3) must never be passed through to an upstream that rejects - // it (Runway 400s on ratios outside its list). + // An explicit quality label (vivid callers) wins over a blank resolution. + if resolution == "" && strings.TrimSpace(in.Quality) != "" { + resolution = strings.TrimSpace(in.Quality) + } aspectRatio = snapRatio(aspectRatio, repo.JSONStrings(modelItem.Ratios)) - // parseImageSize defaults a blank resolution to "2K" (OpenAI-size parity). - // For a model that doesn't price that tier — e.g. gpt-image-2 is 1K-only — - // fall back to its first supported tier so a missing/stale resolution from - // the client doesn't get rejected as "unsupported or unpriced". if _, ok := modelPrice(modelItem, "image", resolution, "", false); !ok { if fb := firstPricedResolution(modelItem); fb != "" { resolution = fb @@ -1334,74 +1364,36 @@ func (s *V1Service) prepareVideo(ctx context.Context, principal *APIPrincipal, i if !modelItem.Enabled || modelItem.Type != "video" { return nil, "", "", "", 0, ErrUnknownModel } - // Validate duration against model's supported range (from Durations JSON array). - if secs := parseDurationSeconds(duration); secs > 0 { - if durList := repo.JSONStrings(modelItem.Durations); len(durList) > 0 { - minSecs, maxSecs := 9999, 0 - for _, d := range durList { - n := parseDurationSeconds(d) - if n > 0 { - if n < minSecs { - minSecs = n - } - if n > maxSecs { - maxSecs = n - } - } - } - if secs < minSecs || secs > maxSecs { - return nil, "", "", "", 0, fmt.Errorf("duration %ds out of range [%d-%d] for model %s", secs, minSecs, maxSecs, modelItem.EffectiveName()) - } + // Validate duration against the model's spec (options list or min/max range). + if secs := int64(parseDurationSeconds(duration)); secs > 0 { + if err := validateModelDuration(modelItem, secs); err != nil { + return nil, "", "", "", 0, err } } - // Fail fast before charging — effective provider (custom upstream by id, else native). - if eff := s.effectiveProvider(ctx, modelItem); eff == "custom" { - // custom serves this id (effectiveProvider guaranteed it) — precheck ok - } else if ok, err := s.hasActiveProviderToken(ctx, eff, "video"); err != nil { + if ok, err := s.hasChannelForModel(ctx, modelItem.ID); err != nil { return nil, "", "", "", 0, err } else if !ok { return nil, "", "", "", 0, ErrNoProviderAccount } refLimit := totalRefLimit(modelItem) + if refLimit <= 0 && len(in.ReferenceImages) > 0 { + return nil, "", "", "", 0, errors.New("this model does not accept references") + } if refLimit <= 0 { refLimit = 10 } if len(in.ReferenceImages) > refLimit { return nil, "", "", "", 0, errors.New("too many reference images") } - // Reject oversized reference images before charging (all providers, all paths). if err := ensureReferenceSizes(in.ReferenceImages); err != nil { return nil, "", "", "", 0, err } - // Runway i2v strictly requires exactly one first-frame image. Enforce it here, - // BEFORE charging, so a missing/extra frame fails fast instead of charge → - // upstream reject → refund. generateRunwayVideo keeps its own guard too. - if modelItem.Provider == "runway" { - n := 0 - for _, r := range in.ReferenceImages { - if strings.TrimSpace(r) != "" { - n++ - } - } - if n != 1 { - return nil, "", "", "", 0, errors.New("runway 图生视频需要且仅需 1 张首帧图") - } - } - // Leonardo seedance 的参考资产分三类且各有上限/时长限制,在扣费前拦掉, - // 不让请求带着非法参考走到上游。 - if modelItem.Provider == "leonardo" { - if _, err := classifyLeonardoVideoRefs(in.ReferenceImages, leonardoVideoSpecOf(modelItem.ID)); err != nil { - return nil, "", "", "", 0, err - } - } aspectRatio := strings.TrimSpace(strings.ReplaceAll(in.AspectRatio, "x", ":")) if aspectRatio == "" { aspectRatio = "16:9" } resolution := strings.TrimSpace(in.Resolution) if resolution == "" { - // 调用方没指定档位时用模型自己配的第一档,而不是假定 720p —— - // 只卖 1440p 的模型会被 720p 判成"没定价"。 if resList := repo.JSONStrings(modelItem.Resolutions); len(resList) > 0 { resolution = strings.TrimSpace(resList[0]) } @@ -1413,17 +1405,24 @@ func (s *V1Service) prepareVideo(ctx context.Context, principal *APIPrincipal, i if err != nil { return nil, "", "", "", 0, err } - // 规范化 duration 字段：前端 per_second 计费模式可能发来 "per_second" 字符串， - // 统一转为 "Xs" 格式（如 "4s"）存库，避免日志显示原始键名。 if n := parseDurationSeconds(duration); n > 0 { duration = fmt.Sprintf("%ds", n) } return modelItem, resolution, aspectRatio, duration, price, nil } +// hasChannelForModel reports whether at least one enabled channel serves the +// model — the channel-era replacement for the token-pool precheck. +func (s *V1Service) hasChannelForModel(ctx context.Context, modelID string) (bool, error) { + channels, err := s.channels.ChannelsForModel(ctx, modelID) + if err != nil { + return false, err + } + return len(channels) > 0, nil +} + func (s *V1Service) chargeForModel(ctx context.Context, principal *APIPrincipal, modelItem *model.ModelConfig, kind, resolution, duration string, surcharge float64, charge bool) (float64, error) { - // 代理用户走代理价(某档未设代理价则回退普通价)。principal.User 即将被扣费的 - // 用户,无论画图台还是 key 调用都从这里取,所以一处即覆盖所有路径。 + // 代理用户走代理价(某档未设代理价则回退普通价)。 agent := principal != nil && principal.User != nil && principal.User.Role == "agent" price, ok := modelPrice(modelItem, kind, resolution, duration, agent) if !ok { @@ -1447,2028 +1446,6 @@ func (s *V1Service) chargeForModel(ctx context.Context, principal *APIPrincipal, return price, nil } -func (s *V1Service) userDir(principal *APIPrincipal) string { - if principal == nil { - return "anon" - } - return OwnerDir(principal.User) -} - -// OwnerDir is the storage directory (= /images// segment) a user's outputs -// live under: sanitized name → sanitized email-local → id → "anon". -func OwnerDir(user *model.User) string { - if user != nil { - if d := sanitizeOwnerName(user.Name); d != "" { - return d - } - if d := sanitizeOwnerName(strings.Split(user.Email, "@")[0]); d != "" { - return d - } - if user.ID != "" { - return user.ID - } - } - return "anon" -} - -// contentTypeForExt maps a file extension to a MIME type for storage uploads. -func contentTypeForExt(ext string) string { - switch strings.ToLower(strings.TrimPrefix(ext, ".")) { - case "png": - return "image/png" - case "jpg", "jpeg": - return "image/jpeg" - case "webp": - return "image/webp" - case "gif": - return "image/gif" - case "mp4": - return "video/mp4" - case "webm": - return "video/webm" - case "mov": - return "video/quicktime" - default: - return "application/octet-stream" - } -} - -// imageExtFromBytes sniffs a sensible file extension from the magic bytes so the -// saved reference keeps its real type (the /images handler types by extension). -func imageExtFromBytes(b []byte) string { - switch { - case len(b) >= 3 && b[0] == 0xFF && b[1] == 0xD8 && b[2] == 0xFF: - return "jpg" - case len(b) >= 6 && string(b[0:6]) == "GIF89a", len(b) >= 6 && string(b[0:6]) == "GIF87a": - return "gif" - case len(b) >= 12 && string(b[0:4]) == "RIFF" && string(b[8:12]) == "WEBP": - return "webp" - default: - return "png" - } -} - -// allocateOutput builds the object key (= relative path, user-scoped) and the -// directly-downloadable URL pointing at this site's /images proxy. Nothing is -// written here — the bytes are uploaded to RustFS by the caller. -func (s *V1Service) allocateOutput(principal *APIPrincipal, ext, baseURL string) (string, string) { - userDir := s.userDir(principal) - filename := time.Now().Format("20060102-150405") + "-" + randomUpper(8) + "." + strings.TrimPrefix(ext, ".") - relativePath := filepath.ToSlash(filepath.Join(userDir, filename)) - // OpenAI-style clients need a directly-downloadable absolute URL. When the - // inbound request's base URL is known, build "{scheme}://{host}/images/..."; - // otherwise fall back to the relative path for backward compatibility. - if base := strings.TrimRight(strings.TrimSpace(baseURL), "/"); base != "" { - return base + "/images/" + relativePath, relativePath - } - return "/images/" + relativePath, relativePath -} - -func (s *V1Service) logPendingEvent(ctx context.Context, kind string, modelItem *model.ModelConfig, principal *APIPrincipal, prompt, ratio, resolution, duration string, refs int, cost float64, file, source string, refFiles []string, deai bool) (string, error) { - event := &model.EventLog{ - ID: "evt-" + randomUpper(12), - TS: time.Now(), - Kind: kind, - Status: "pending", - Model: modelItem.ID, - Provider: modelItem.Provider, - Prompt: prompt, - Ratio: ratio, - Resolution: resolution, - Duration: duration, - Refs: refs, - DeAI: deai, - Source: source, - Cost: cost, - File: file, - CreatedAt: time.Now(), - UpdatedAt: time.Now(), - } - if len(refFiles) > 0 { - event.RefFiles = jsonArray(refFiles) - } - if principal != nil && principal.User != nil { - event.UserID = principal.User.ID - } - if err := s.events.Create(ctx, event); err != nil { - return "", err - } - return event.ID, nil -} - -func (s *V1Service) finishUnimplementedEvent(ctx context.Context, eventID string) error { - return s.events.UpdateStatus(ctx, eventID, "failed", "generation executor not implemented yet", 0) -} - -// videoGenBudget caps one video render end-to-end (submit + poll + download). -// 上游慢的时候（seedance 长镜头）12 分钟不够，统一给 30 分钟。 -const videoGenBudget = 30 * time.Minute - -// 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 = fail over" -// policy may burn per request before giving up, so an upstream-wide blip -// ("system under load") can't fan a single request out across the whole pool. -// After this many accounts fail this way, the request fails. -const maxTempDeadAccounts = 10 - -// 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: -// - 额度耗尽 quota → mark the account and FAIL OVER to the next account -// immediately (same-account retry can't help). Repeats until one succeeds or -// the pool is exhausted. -// - 认证失效 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 → record the failure (no disable/dead) and FAIL OVER to -// the next account immediately, capped at maxTempDeadAccounts accounts so a -// pool-wide blip can't fan a single request out across everything. -// - 参数错 / request-level (anything else) → return immediately, no retry, no -// account penalty (the account isn't at fault). -// -// Returns the actual upstream error (never a synthetic "retry failed"). On -// success it stamps success_total/fails=0 on the winning account. classify maps -// a provider error to (isAuth, isQuota, isTemporary). refreshOnAuth (nil for -// providers whose token IS the credential) re-mints the account's token so an -// auth retry uses a FRESH token instead of replaying the stale one. -func (s *V1Service) runPoolWithFailover(ctx context.Context, eventID, pool string, active []model.TokenAccount, kind string, - attempt func(token model.TokenAccount) ([]byte, error), - classify func(error) (isAuth, isQuota, isTemporary, isDead bool), - refreshOnAuth func(tokenID string) (model.TokenAccount, bool), - tempFailover bool, -) ([]byte, error) { - var lastErr error - busy := 0 - tempDeadCount := 0 - for _, token := range active { - // Per-account concurrency gate (defaults to 1 for built-in pools). - if !s.acctAcquire(ctx, token.ID, eventID, accountConcurrency(token)) { - busy++ - continue - } - // release via defer so a panic in tryAccount can't leak the 1-job slot. - data, err, failover, tempDead := func() ([]byte, error, bool, bool) { - defer s.acctRelease(ctx, token.ID, eventID) - return s.tryAccount(ctx, eventID, pool, token, kind, attempt, classify, refreshOnAuth, tempFailover) - }() - if err == nil { - return data, nil - } - lastErr = err - if tempDead { - // temp-failover policy: this account hit a temporary upstream error. - // Cap how many accounts one request may burn before we stop, so an - // upstream-wide blip doesn't fan out across the whole pool. - tempDeadCount++ - if tempDeadCount >= maxTempDeadAccounts { - return nil, lastErr - } - } - if failover { - continue - } - // temporary exhausted or request-level error → surface it, no fan-out. - return nil, lastErr - } - // Nothing ran. If accounts were skipped ONLY because they were all busy - // (no real failure), tell the caller the pool is at its concurrency cap. - if lastErr == nil { - if busy > 0 { - return nil, ErrConcurrencyFull - } - return nil, ErrProviderExecution - } - return nil, lastErr -} - -// tryAccount runs one account's attempt with the pool's retry policy: -// 额度耗尽/认证失效 → mark + failover; 上游临时 → record failure + failover (capped -// via the tempDead return); 参数错 → fail fast. Returns (data, err, failover, -// tempDead) — failover=true means move on to the next account. The per-account -// concurrency gate is held by the caller. -func (s *V1Service) tryAccount(ctx context.Context, eventID, pool string, token model.TokenAccount, kind string, - attempt func(token model.TokenAccount) ([]byte, error), - classify func(error) (isAuth, isQuota, isTemporary, isDead bool), - refreshOnAuth func(tokenID string) (model.TokenAccount, bool), - tempFailover bool, -) ([]byte, error, bool, bool) { - _ = s.events.SetAccount(ctx, eventID, token.ID, token.AccountEmail) - _ = s.tokens.TouchLastUsed(ctx, token.ID) - authRefreshed := false - if strings.TrimSpace(token.Value) == "" && refreshOnAuth != nil { - if refreshed, ok := refreshOnAuth(token.ID); ok { - token = refreshed - authRefreshed = true - } else { - s.markTokenDead(ctx, pool, token, kind) - return nil, ErrProviderExecution, true, true - } - } - for { - data, err := attempt(token) - if err == nil { - _, _ = s.tokens.Update(ctx, pool, token.ID, map[string]any{ - "last_used_at": time.Now(), - "success_total": gorm.Expr("success_total + 1"), - "fails": 0, - }) - return data, nil, false, false - } - isAuth, isQuota, isTemp, isDead := classify(err) - if isQuota { - s.markTokenFailure(ctx, pool, token, kind, false, true) - return nil, err, true, false - } - if isAuth { - // A 403 user_not_entitled means the account has no Firefly entitlement - // — refreshing the access token can't grant one, so kill it now instead - // of leaving it in rotation to burn every future request. - if errors.Is(err, adobe.ErrNotEntitled) { - s.markTokenDead(ctx, pool, token, kind) - return nil, err, true, true - } - // Refresh from cookie and retry ONCE; otherwise the credential is dead. - if refreshOnAuth != nil && !authRefreshed { - if refreshed, ok := refreshOnAuth(token.ID); ok { - token = refreshed - authRefreshed = true - continue - } - } - s.markTokenFailure(ctx, pool, token, kind, true, false) - return nil, err, true, false - } - // Fatal / temporary-under-failover-policy upstream error. - if isDead || (isTemp && tempFailover) { - if tempFailover { - // Ops policy (adobe): NEVER kill on these upstream errors — a - // genuinely bad account and a transient Adobe blip (429/5xx/ - // overload) look the same, and killing wipes healthy accounts. - // Record the failure and fail over to the next account (no - // disable/dead). The 4th return value caps how many accounts one - // request may burn this way (maxTempDeadAccounts) so a pool-wide - // blip can't fan a single request across the whole pool. - s.markTokenFailure(ctx, pool, token, kind, false, false) - return nil, err, true, true - } - s.markTokenDead(ctx, pool, token, kind) - return nil, err, true, true - } - if isTemp { - // Temporary upstream error → record the failure (no disable/dead) and - // fail over to the NEXT account, capped via the tempDead return so a - // pool-wide blip can't fan one request across the whole pool. - s.markTokenFailure(ctx, pool, token, kind, false, false) - return nil, err, true, true - } - return nil, err, false, false // 参数错 / request-level - } -} - -func adobeErrClass(e error) (bool, bool, bool, bool) { - return errors.Is(e, adobe.ErrAuth), errors.Is(e, adobe.ErrQuotaExhausted), errors.Is(e, adobe.ErrTemporaryUpstream) || errors.Is(e, adobe.ErrRateLimited), errors.Is(e, adobe.ErrDeadUpstream) -} - -// creativefabricaErrClass maps a creativefabrica upstream error onto the pool's -// (auth, quota, temporary, dead) classification. -func creativefabricaErrClass(e error) (bool, bool, bool, bool) { - return errors.Is(e, creativefabrica.ErrAuth), - errors.Is(e, creativefabrica.ErrQuotaExhausted), - errors.Is(e, creativefabrica.ErrTemporaryUpstream) || errors.Is(e, creativefabrica.ErrRateLimited), - errors.Is(e, creativefabrica.ErrDeadUpstream) -} - -// noStore url-only mode: adobe returns a presigned image URL (meta["image_url"]); -// skip the download and return it directly. -func (s *V1Service) generateAdobeImage(ctx context.Context, eventID string, modelItem *model.ModelConfig, in V1ImageRequest, aspectRatio, resolution string, noStore bool) ([]byte, string, error) { - urlOnly := noStore - if s.adobe == nil { - return nil, "", errors.New("adobe client not configured") - } - if s.settings != nil { - if proxy, err := s.settings.GetValue(ctx, "proxy.url"); err == nil { - s.adobe.SetProxy(proxy) - } - } - - items, err := s.tokens.ListByPool(ctx, "adobe") - if err != nil { - return nil, "", err - } - var active []model.TokenAccount - for _, item := range items { - // Adobe accounts are credit-based (积分号) — no per-kind quota locks. - // Only skip accounts that are dead or disabled. - if item.Status != "active" || item.Dead { - continue - } - // plan 未探测到的号既不算普号也不算会员号，置死号、不参与调度 - if planUnknown(item.Meta) { - s.markPlanUnknownDead(ctx, "adobe", item.ID) - continue - } - // 普号(free)只能调度 free_allowed 的模型（香蕉2 仅 1K） - if !freeAccountsAllowed(modelItem, resolution) && isFreeAccount(item.Meta) { - continue - } - active = append(active, item) - } - active = pinTestAccount(items, active, in.AccountID) - if len(active) == 0 { - return nil, "", ErrNoProviderAccount - } - s.rotateRoundRobin("adobe", active) - // 非 seedance 图片生成：普号 → 子号 → 母号 - active = prioritizeSubAccounts(active) - - refs, err := decodeReferenceImages(in.ReferenceImages, max(1, modelItem.MaxReferenceImages)) - if err != nil { - return nil, "", err - } - - // Round-robin order. Adobe uses tempFailover=true: a temporary upstream error - // ("system under load") fails over to the next account without penalizing the - // current one, capped at maxTempDeadAccounts; auth/quota also fail over - // (see runPoolWithFailover). imageURL is captured from the successful attempt. - var imageURL string - data, err := s.runPoolWithFailover(ctx, eventID, "adobe", active, "image", func(token model.TokenAccount) ([]byte, error) { - var blobIDs []string - for _, ref := range refs { - id, upErr := s.adobe.UploadImage(ctx, token.Value, ref, "image/png", "") - if upErr != nil { - if errors.Is(upErr, adobe.ErrRateLimited) { - recoverAt := time.Now().Add(4 * time.Hour) - s.tokens.Update(ctx, "adobe", token.ID, map[string]any{ - "status": "quota", - "quota_recover_at": &recoverAt, - }) - } - return nil, upErr - } - blobIDs = append(blobIDs, id) - } - d, meta, genErr := s.adobe.GenerateImage(ctx, token.Value, modelItem.ID, in.Prompt, aspectRatio, resolution, blobIDs, !urlOnly) - if genErr == nil { - imageURL = strings.TrimSpace(stringValue(meta["image_url"])) - } - return d, genErr - }, adobeErrClass, func(id string) (model.TokenAccount, bool) { - return s.refreshAdobeToken(ctx, id) - }, true) - return data, imageURL, err -} - -func (s *V1Service) generateAdobeVideo(ctx context.Context, eventID string, modelItem *model.ModelConfig, in V1VideoRequest, aspectRatio, resolution string, durationSeconds int, downloadResult bool) ([]byte, string, error) { - if s.adobe == nil { - return nil, "", errors.New("adobe client not configured") - } - if s.settings != nil { - if proxy, err := s.settings.GetValue(ctx, "proxy.url"); err == nil { - s.adobe.SetProxy(proxy) - } - } - - items, err := s.tokens.ListByPool(ctx, "adobe") - if err != nil { - return nil, "", err - } - var active []model.TokenAccount - for _, item := range items { - if item.Status != "active" || item.Dead { - continue - } - // plan 未探测到的号既不算普号也不算会员号，置死号、不参与调度 - if planUnknown(item.Meta) { - s.markPlanUnknownDead(ctx, "adobe", item.ID) - continue - } - // Seedance 模型只允许 VIP 母号：必须正向识别（plan 非 free、非子号、 - // 积分 >4000），plan/额度未探测的账号一律不参与调度 - if isSeedanceModel(modelItem.ID) && !isVipMotherAccount(item.Meta) { - continue - } - // 普号(free)只能调度 free_allowed 的模型 - if !freeAccountsAllowed(modelItem, resolution) && isFreeAccount(item.Meta) { - continue - } - active = append(active, item) - } - active = pinTestAccount(items, active, in.AccountID) - if len(active) == 0 { - return nil, "", ErrNoProviderAccount - } - s.rotateRoundRobin("adobe", active) - // 非 seedance 视频生成：普号 → 子号 → 母号 - // （seedance 已在上面过滤掉子号，此处无需额外处理） - if !isSeedanceModel(modelItem.ID) { - active = prioritizeSubAccounts(active) - } - - refLimit := modelItem.MaxReferenceImages - if refLimit <= 0 { - refLimit = 10 - } - refs, err := decodeReferenceImages(in.ReferenceImages, refLimit) - if err != nil { - return nil, "", err - } - // Classify refs for seedance: images (usage:style), videos, audio (usage:source). - var imgRefs, vidRefs, audRefs [][]byte - for _, r := range refs { - switch detectMediaType(r) { - case "video": - vidRefs = append(vidRefs, r) - case "audio": - audRefs = append(audRefs, r) - default: - imgRefs = append(imgRefs, r) - } - } - prompt := in.Prompt - - engine, upstreamModel := resolveAdobeVideoEngine(modelItem.ID) - referenceMode := defaultString(strings.TrimSpace(modelItem.ReferenceMode), "frame") - if rm := strings.TrimSpace(in.ReferenceMode); rm != "" { - referenceMode = rm - } - - // Round-robin order; fail over to the next account on auth/quota; temporary - // upstream errors fail over too without penalizing the account (tempFailover, - // 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 - for _, ref := range imgRefs { - id, upErr := s.adobe.UploadImage(ctx, token.Value, ref, "image/png", engine) - if upErr != nil { - if errors.Is(upErr, adobe.ErrRateLimited) { - recoverAt := time.Now().Add(4 * time.Hour) - s.tokens.Update(ctx, "adobe", token.ID, map[string]any{ - "status": "quota", - "quota_recover_at": &recoverAt, - }) - } - return nil, upErr - } - blobIDs = append(blobIDs, id) - } - var videoBlobIDs []string - for _, ref := range vidRefs { - id, upErr := s.adobe.UploadImage(ctx, token.Value, ref, "video/mp4", engine) - if upErr != nil { - return nil, upErr - } - videoBlobIDs = append(videoBlobIDs, id) - } - var audioBlobIDs []string - for _, ref := range audRefs { - id, upErr := s.adobe.UploadImage(ctx, token.Value, ref, "audio/mp3", engine) - if upErr != nil { - return nil, upErr - } - audioBlobIDs = append(audioBlobIDs, id) - } - bytes, meta, genErr := s.adobe.GenerateVideo(ctx, token.Value, engine, prompt, aspectRatio, durationSeconds, resolution, referenceMode, upstreamModel, blobIDs, videoBlobIDs, audioBlobIDs, downloadResult) - if genErr == nil { - videoURL = strings.TrimSpace(stringValue(meta["video_url"])) - } - return bytes, genErr - }, adobeErrClass, func(id string) (model.TokenAccount, bool) { - return s.refreshAdobeToken(ctx, id) - }, true) - return data, videoURL, err -} - -// maxCreativeFabricaRefs caps how many reference images a Creative Fabrica -// generation may carry (the studio UI allows up to 9). -const maxCreativeFabricaRefs = 9 - -// generateCreativeFabricaVideo renders a video through the Creative Fabrica -// Studio upstream. Accounts are ONE-SHOT: the coins buy exactly one generation, -// so a successful render immediately disables the account. A fresh short-lived -// JWT is minted from the stored cookie for every attempt (there is no -// long-lived token to cache). Only image reference frames are supported. -func (s *V1Service) generateCreativeFabricaVideo(ctx context.Context, eventID string, modelItem *model.ModelConfig, in V1VideoRequest, aspectRatio string, downloadResult bool) ([]byte, string, error) { - if s.cf == nil { - return nil, "", errors.New("creativefabrica client not configured") - } - if s.settings != nil { - if proxy, err := s.settings.GetValue(ctx, "proxy.url"); err == nil { - s.cf.SetProxy(proxy) - } - } - - items, err := s.tokens.ListByPool(ctx, "creativefabrica") - 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 - } - active = append(active, item) - } - active = pinTestAccount(items, active, in.AccountID) - if len(active) == 0 { - return nil, "", ErrNoProviderAccount - } - s.rotateRoundRobin("creativefabrica", active) - - refLimit := modelItem.MaxReferenceImages - if refLimit <= 0 { - refLimit = maxCreativeFabricaRefs - } - refs, err := decodeReferenceImages(in.ReferenceImages, refLimit) - if err != nil { - return nil, "", err - } - // The studio only accepts image reference frames — reject video/audio refs. - for _, r := range refs { - if detectMediaType(r) != "image" { - return nil, "", errors.New("creativefabrica only supports image reference frames") - } - } - - var videoURL string - data, err := s.runPoolWithFailover(ctx, eventID, "creativefabrica", active, "video", func(token model.TokenAccount) ([]byte, error) { - jwt, _, terr := s.cf.ExchangeToken(ctx, token.Value) - if terr != nil { - s.markTokenDead(ctx, "creativefabrica", token, "video") - return nil, terr - } - bytes, url, gerr := s.cf.GenerateVideo(ctx, token.Value, jwt, modelItem.ID, in.Prompt, aspectRatio, refs, downloadResult) - if gerr == nil { - videoURL = url - // One-shot: the account's coins paid for exactly this generation. - s.markTokenDead(ctx, "creativefabrica", token, "video") - } - return bytes, gerr - }, creativefabricaErrClass, nil, true) - return data, videoURL, err -} - -// 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. -const leonardoMinCredits = 30 - -func (s *V1Service) generateRunwayVideo(ctx context.Context, eventID string, modelItem *model.ModelConfig, in V1VideoRequest, aspectRatio string, durationSeconds int, downloadResult bool) ([]byte, string, error) { - if s.runway == nil { - return nil, "", errors.New("runway client not configured") - } - if s.settings != nil { - if proxy, err := s.settings.GetValue(ctx, "proxy.url"); err == nil { - s.runway.SetProxy(proxy) - } - } - - // Runway i2v strictly requires exactly one first-frame image. - refs, err := decodeReferenceImages(in.ReferenceImages, 1) - if err != nil { - return nil, "", err - } - if len(refs) != 1 { - return nil, "", errors.New("runway 图生视频需要且仅需 1 张首帧图") - } - frame := refs[0] - - items, err := s.tokens.ListByPool(ctx, "runway") - 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 - } - // 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) - } - active = pinTestAccount(items, active, in.AccountID) - if len(active) == 0 { - return nil, "", ErrNoProviderAccount - } - s.rotateRoundRobin("runway", active) - - var lastErr error - var videoURL string - busy := 0 - for _, token := range active { - // Per-account concurrency gate (defaults to 1 for built-in pools). - if !s.acctAcquire(ctx, token.ID, eventID, accountConcurrency(token)) { - busy++ - continue - } - var data []byte - done, failover := func() (bool, bool) { - defer s.acctRelease(ctx, token.ID, eventID) - _ = s.events.SetAccount(ctx, eventID, token.ID, token.AccountEmail) - _ = s.tokens.TouchLastUsed(ctx, token.ID) - teamID := "" - if token.Meta != nil { - teamID = strings.TrimSpace(stringValue(token.Meta["team_id"])) - } - d, meta, genErr := s.runway.GenerateVideo(ctx, token.Value, teamID, in.Prompt, aspectRatio, durationSeconds, frame, downloadResult) - 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 - videoURL = strings.TrimSpace(stringValue(meta["video_url"])) - 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, "video", true, false) - return false, true - case errors.Is(genErr, runway.ErrTemporaryUpstream): - // 上游临时错误 → 直接换下一个号。 - return false, true - default: - // 参数级错误(如 prompt 未过审)→ 直接失败,不换号。 - 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 -} - -// totalRefLimit is the aggregate reference-file cap across all media kinds. -// Models with per-kind caps (max_reference_videos/audios) total them; legacy -// models only have max_reference_images, so the sum degenerates correctly. -func totalRefLimit(m *model.ModelConfig) int { - return m.MaxReferenceImages + m.MaxReferenceVideos + m.MaxReferenceAudios -} - -// refKinds counts decoded refs by media class (image / video / audio) using the -// same magic-byte detection the leonardo upload path relies on. -func refKinds(refs [][]byte) (images, videos, audios int) { - for _, r := range refs { - switch detectMediaType(r) { - case "video": - videos++ - case "audio": - audios++ - default: - images++ - } - } - return -} - -// checkCustomRefCaps validates decoded refs against a custom model's per-kind -// caps (Chinese messages mirror the leonardo spec errors). -func checkCustomRefCaps(modelItem *model.ModelConfig, refs [][]byte) (images, videos, audios int, err error) { - images, videos, audios = refKinds(refs) - switch { - case videos > modelItem.MaxReferenceVideos: - err = fmt.Errorf("最多 %d 段视频参考", modelItem.MaxReferenceVideos) - case audios > modelItem.MaxReferenceAudios: - err = fmt.Errorf("最多 %d 段音频参考", modelItem.MaxReferenceAudios) - case images > modelItem.MaxReferenceImages: - err = fmt.Errorf("最多 %d 张图片参考", modelItem.MaxReferenceImages) - } - return -} - -// customAccountServes reports whether a custom (upstream) account is usable for a -// given model id: active, not dead, has a base_url, and its meta.models list (csv -// of model ids it serves) contains the id. An empty models list serves ALL ids. -func customAccountServes(item model.TokenAccount, modelID string) bool { - if item.Status != "active" || item.Dead || strings.TrimSpace(item.Value) == "" { - return false - } - if item.Meta == nil || strings.TrimSpace(stringValue(item.Meta["base_url"])) == "" { - return false - } - list := strings.TrimSpace(stringValue(item.Meta["models"])) - if list == "" { - return true - } - for _, m := range strings.Split(list, ",") { - if strings.EqualFold(strings.TrimSpace(m), modelID) { - return true - } - } - return false -} - -// customActive returns the custom accounts that serve modelID, ordered by weight -// (higher first; ties by id) so heavier upstreams are preferred. -func (s *V1Service) customActive(ctx context.Context, modelID string) ([]model.TokenAccount, error) { - items, err := s.tokens.ListByPool(ctx, "custom") - if err != nil { - return nil, err - } - var active []model.TokenAccount - for _, item := range items { - if customAccountServes(item, modelID) { - active = append(active, item) - } - } - s.rotateRoundRobin("custom", active) // weight priority + round-robin within ties - return active, nil -} - -// accountConcurrency is the per-account simultaneous-job cap. Custom accounts use -// their configured Concurrency (default 1); built-in pools use the system value. -func accountConcurrency(item model.TokenAccount) int { - if item.Pool == "adobe" { - if isFreeAccount(item.Meta) { - return 1 // FREE 普号 / 降级号限制为 1 并发 - } - if item.Concurrency > 0 { - return item.Concurrency - } - return 5 // VIP 会员号默认 5 并发 - } - if item.Concurrency > 0 { - return item.Concurrency - } - if item.Pool == "grok" { - return grokConcurrencyPerAccount // 10 - } - return 1 -} - -// effectiveProvider routes a model to the "custom" upstream whenever a custom -// account declares it serves that model id (id-based override of the model's -// native provider) — so an upstream can take over any model by matching its id. -// Otherwise the model's own provider is used. -func (s *V1Service) effectiveProvider(ctx context.Context, modelItem *model.ModelConfig) string { - if s.custom != nil { - if active, err := s.customActive(ctx, modelItem.ID); err == nil && len(active) > 0 { - return "custom" - } - } - return modelItem.Provider -} - -// generateCustomImage forwards an image generation to an OpenAI-compatible -// upstream. The upstream (custom account) is matched by model id; calls go direct -// (no proxy). Billing uses the local model price. -func (s *V1Service) generateCustomImage(ctx context.Context, eventID string, modelItem *model.ModelConfig, in V1ImageRequest, aspectRatio, resolution string, noStore bool) ([]byte, string, error) { - urlOnly := noStore - if s.custom == nil { - return nil, "", errors.New("custom client not configured") - } - refs, err := decodeReferenceImages(in.ReferenceImages, max(1, totalRefLimit(modelItem))) - if err != nil { - return nil, "", err - } - _, vids, auds, err := checkCustomRefCaps(modelItem, refs) - if err != nil { - return nil, "", err - } - active, err := s.customActive(ctx, modelItem.ID) - if err != nil { - return nil, "", err - } - active = pinTestAccount(active, active, in.AccountID) - if len(active) == 0 { - return nil, "", ErrNoProviderAccount - } - size := upstreamSize(aspectRatio, resolution) - quality := upstreamQuality(resolution) - var lastErr error - busy := 0 - for _, token := range active { - if !s.acctAcquire(ctx, token.ID, eventID, accountConcurrency(token)) { - busy++ - continue - } - var data []byte - var imgURL string - done, failover := func() (bool, bool) { - defer s.acctRelease(ctx, token.ID, eventID) - _ = s.events.SetAccount(ctx, eventID, token.ID, token.AccountEmail) - _ = s.tokens.TouchLastUsed(ctx, token.ID) - baseURL := stringValue(token.Meta["base_url"]) - vivid := strings.EqualFold(stringValue(token.Meta["protocol"]), "vivid") - if !vivid && (vids > 0 || auds > 0) { - // OpenAI-protocol upstreams only accept image refs — skip this - // account; another vivid-protocol account may still serve. - lastErr = errors.New("upstream does not accept video/audio refs") - return false, true - } - var d []byte - var u string - var genErr error - if vivid { - d, u, genErr = s.custom.GenerateVividImage(ctx, baseURL, token.Value, modelItem.ID, in.Prompt, aspectRatio, vividQuality(resolution), refs, !urlOnly) - } else { - d, u, genErr = s.custom.GenerateImage(ctx, baseURL, token.Value, modelItem.ID, in.Prompt, size, quality, refs, !urlOnly) - } - if genErr == nil { - _, _ = s.tokens.Update(ctx, "custom", token.ID, map[string]any{ - "last_used_at": time.Now(), "success_total": gorm.Expr("success_total + 1"), "fails": 0, - }) - data = d - imgURL = u - return true, false - } - lastErr = genErr - switch { - case errors.Is(genErr, custom.ErrAuth): - s.markTokenFailure(ctx, "custom", token, "image", true, false) - return false, true - case errors.Is(genErr, custom.ErrQuotaExhausted): - s.markTokenFailure(ctx, "custom", token, "image", false, true) - return false, true - case errors.Is(genErr, custom.ErrTemporaryUpstream): - return false, true - default: - return false, false - } - }() - if done { - return data, imgURL, nil - } - if failover { - continue - } - return nil, "", lastErr - } - if lastErr == nil { - if busy > 0 { - return nil, "", ErrConcurrencyFull - } - lastErr = ErrProviderExecution - } - return nil, "", lastErr -} - -// generateCustomVideo forwards a video generation to an OpenAI-compatible -// (Sora-style) upstream, matched by model id. No proxy; local-price billing. -func (s *V1Service) generateCustomVideo(ctx context.Context, eventID string, modelItem *model.ModelConfig, in V1VideoRequest, aspectRatio, resolution string, durationSeconds int, downloadResult bool) ([]byte, string, error) { - if s.custom == nil { - return nil, "", errors.New("custom client not configured") - } - active, err := s.customActive(ctx, modelItem.ID) - if err != nil { - return nil, "", err - } - active = pinTestAccount(active, active, in.AccountID) - if len(active) == 0 { - return nil, "", ErrNoProviderAccount - } - size := upstreamVideoSize(aspectRatio, resolution) - // Optional reference files (images / videos / audios for vivid upstreams; - // images only for OpenAI-protocol ones). - frames, err := decodeReferenceImages(in.ReferenceImages, max(1, totalRefLimit(modelItem))) - if err != nil { - return nil, "", err - } - _, vids, auds, err := checkCustomRefCaps(modelItem, frames) - if err != nil { - return nil, "", err - } - var lastErr error - var videoURL string - busy := 0 - for _, token := range active { - if !s.acctAcquire(ctx, token.ID, eventID, accountConcurrency(token)) { - busy++ - continue - } - var data []byte - done, failover := func() (bool, bool) { - defer s.acctRelease(ctx, token.ID, eventID) - _ = s.events.SetAccount(ctx, eventID, token.ID, token.AccountEmail) - _ = s.tokens.TouchLastUsed(ctx, token.ID) - baseURL := stringValue(token.Meta["base_url"]) - vivid := strings.EqualFold(stringValue(token.Meta["protocol"]), "vivid") - if !vivid && (vids > 0 || auds > 0) { - lastErr = errors.New("upstream does not accept video/audio refs") - return false, true - } - var d []byte - var url string - var genErr error - if vivid { - d, url, genErr = s.custom.GenerateVividVideo(ctx, baseURL, token.Value, modelItem.ID, in.Prompt, aspectRatio, vividQuality(resolution), durationSeconds, frames, downloadResult) - } else { - d, url, genErr = s.custom.GenerateVideo(ctx, baseURL, token.Value, modelItem.ID, in.Prompt, size, durationSeconds, frames, downloadResult) - } - if genErr == nil { - _, _ = s.tokens.Update(ctx, "custom", token.ID, map[string]any{ - "last_used_at": time.Now(), "success_total": gorm.Expr("success_total + 1"), "fails": 0, - }) - data = d - videoURL = url - return true, false - } - lastErr = genErr - switch { - case errors.Is(genErr, custom.ErrAuth): - s.markTokenFailure(ctx, "custom", token, "video", true, false) - return false, true - case errors.Is(genErr, custom.ErrQuotaExhausted): - s.markTokenFailure(ctx, "custom", token, "video", false, true) - return false, true - case errors.Is(genErr, custom.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 -} - -// upstreamSize maps our (ratio, resolution) to an OpenAI-style "WxH" size string -// for the upstream. The pixel base scales with the tier (1K/2K/4K); the ratio -// sets the shape. Upstreams that key off ratio (our own /v1) read it fine. -func upstreamSize(aspectRatio, resolution string) string { - base := 1024 - switch strings.ToUpper(strings.TrimSpace(resolution)) { - case "2K": - base = 2048 - case "4K": - base = 4096 - } - w, h := 1, 1 - parts := strings.Split(strings.ReplaceAll(strings.TrimSpace(aspectRatio), "x", ":"), ":") - if len(parts) == 2 { - if a, e1 := strconv.Atoi(strings.TrimSpace(parts[0])); e1 == nil && a > 0 { - if b, e2 := strconv.Atoi(strings.TrimSpace(parts[1])); e2 == nil && b > 0 { - w, h = a, b - } - } - } - if w >= h { - return fmt.Sprintf("%dx%d", base, base*h/w) - } - return fmt.Sprintf("%dx%d", base*w/h, base) -} - -// upstreamVideoSize maps our (ratio, resolution) to a "WxH" size for video -// upstreams. Video "Np" tiers set the SHORT edge in pixels (like grok: -// 720p 1:1 → 720x720, 720p 16:9 → 1280x720); 2K/4K fall back to the -// long-edge mapping shared with images. -func upstreamVideoSize(aspectRatio, resolution string) string { - short := 0 - switch res := strings.ToLower(strings.TrimSpace(resolution)); res { - case "540p": - short = 540 - case "720p", "": - short = 720 - case "1080p": - short = 1080 - } - if short == 0 { - return upstreamSize(aspectRatio, resolution) - } - w, h := 1, 1 - parts := strings.Split(strings.ReplaceAll(strings.TrimSpace(aspectRatio), "x", ":"), ":") - if len(parts) == 2 { - if a, e1 := strconv.Atoi(strings.TrimSpace(parts[0])); e1 == nil && a > 0 { - if b, e2 := strconv.Atoi(strings.TrimSpace(parts[1])); e2 == nil && b > 0 { - w, h = a, b - } - } - } - if w >= h { - return fmt.Sprintf("%dx%d", short*w/h, short) - } - return fmt.Sprintf("%dx%d", short, short*h/w) -} - -// vividQuality maps a resolution tier to the quality string a Vivid-protocol -// upstream expects. Vivid passes the tier through verbatim — "1K"/"2K"/"4K" -// for images, "720p"/"1080p"/"540p" for videos — so don't re-map to the OpenAI -// enum; only default a blank to "1K" (the field is required). -func vividQuality(resolution string) string { - if q := strings.TrimSpace(resolution); q != "" { - return q - } - return "1K" -} - -// upstreamQuality maps a resolution tier to the OpenAI quality enum. -func upstreamQuality(resolution string) string { - switch strings.ToUpper(strings.TrimSpace(resolution)) { - case "2K": - return "medium" - case "4K": - return "high" - case "1K": - return "low" - } - return "" -} - -// generateGrokVideo runs grok's imagine video pipeline across the grok pool, -// via Grok Console (console.x.ai) — the same sso account, but the clean JSON -// media API instead of the anti-bot gated grok.com website flow. -// 额度是本地写死的（每号 图 5 / 视频 2）：视频计数归零的号不再调度，下单先预扣一个、 -// 失败退回（并发不超扣），图/视频都归零直接判死；auth / 额度错误同样判死换号 -// （grok sso 不续期，失效就失效）。 -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, repo.GrokVideoQuotaKey); ok && rem <= 0 { - continue - } - active = append(active, item) - } - active = pinTestAccount(items, active, in.AccountID) - 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 { - // Per-account concurrency gate (defaults to 1 for built-in pools). - if !s.acctAcquire(ctx, token.ID, eventID, accountConcurrency(token)) { - busy++ - continue - } - var data []byte - done, failover := func() (bool, bool) { - defer s.acctRelease(ctx, token.ID, eventID) - _ = s.events.SetAccount(ctx, eventID, token.ID, token.AccountEmail) - _ = s.tokens.TouchLastUsed(ctx, token.ID) - // 下单先预扣本地额度（并发时不会超扣），失败再退回。 - allowed, reserveErr := s.tokens.ReserveGrokQuota(ctx, token.ID, "video") - if reserveErr != nil || !allowed { - return false, true - } - d, meta, genErr := s.grok.GenerateConsoleVideo(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, - }) - // 图/视频都归零时账号直接判死。 - _ = s.tokens.FinalizeGrokQuota(ctx, token.ID) - data = d - videoURL = strings.TrimSpace(stringValue(meta["video_url"])) - return true, false - } - _ = s.tokens.RefundGrokQuota(ctx, token.ID, "video") - 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 -} - -// generateGrokImage runs Grok Console's image pipeline (grok-imagine-image) -// across the grok pool. 额度策略同视频路径，只是扣的是图片那份计数。带参考图时 -// （最多 3 张，内联在请求里）自动走 /images/edits 的 quality 上游 — 图生图。 -func (s *V1Service) generateGrokImage(ctx context.Context, eventID string, modelItem *model.ModelConfig, in V1ImageRequest, aspectRatio, resolution string, noStore bool) ([]byte, string, error) { - // API-key (noStore) requests skip the download and return the upstream URL. - urlOnly := noStore - 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) - } - } - - refs, 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, repo.GrokImageQuotaKey); ok && rem <= 0 { - continue - } - active = append(active, item) - } - active = pinTestAccount(items, active, in.AccountID) - if len(active) == 0 { - return nil, "", ErrNoProviderAccount - } - s.rotateRoundRobin("grok", active) - - var lastErr error - busy := 0 - for _, token := range active { - // Per-account concurrency gate. - if !s.acctAcquire(ctx, token.ID, eventID, accountConcurrency(token)) { - busy++ - continue - } - var data []byte - var artURL string - done, failover := func() (bool, bool) { - defer s.acctRelease(ctx, token.ID, eventID) - _ = s.events.SetAccount(ctx, eventID, token.ID, token.AccountEmail) - _ = s.tokens.TouchLastUsed(ctx, token.ID) - // 下单先预扣本地额度（并发时不会超扣），失败再退回。 - allowed, reserveErr := s.tokens.ReserveGrokQuota(ctx, token.ID, "image") - if reserveErr != nil || !allowed { - return false, true - } - d, meta, genErr := s.grok.GenerateConsoleImage(ctx, token.Value, in.Prompt, aspectRatio, resolution, refs, urlOnly) - 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, - }) - // 图/视频都归零时账号直接判死。 - _ = s.tokens.FinalizeGrokQuota(ctx, token.ID) - data = d - artURL = strings.TrimSpace(stringValue(meta["image_url"])) - return true, false - } - _ = s.tokens.RefundGrokQuota(ctx, token.ID, "image") - lastErr = genErr - switch { - case errors.Is(genErr, grok.ErrAuth), errors.Is(genErr, grok.ErrQuotaExhausted): - // 失效 / 额度没了 → 当 401 判死(不续期),换号。 - s.markTokenFailure(ctx, "grok", token, "image", true, false) - return false, true - case errors.Is(genErr, grok.ErrTemporaryUpstream): - return false, true - default: - return false, false - } - }() - if done { - return data, artURL, 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 gemini image pipeline (Nano Banana Pro or -// Nano Banana 2, selected by the model id) 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. -// noStore url-only mode (API-key requests without DeAI): skip the artifact -// download and return the upstream image URL directly, no bytes. -func (s *V1Service) generateRunwayImage(ctx context.Context, eventID string, modelItem *model.ModelConfig, in V1ImageRequest, aspectRatio, resolution string, noStore bool) ([]byte, string, error) { - // API-key (noStore) requests don't support DeAI (only the web drawing board - // does), so url-only mode == noStore — skip the download, return the URL. - urlOnly := noStore - if s.runway == nil { - return nil, "", errors.New("runway client not configured") - } - 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 nil, "", err - } - - items, err := s.tokens.ListByPool(ctx, "runway") - 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 - } - // 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) - } - active = pinTestAccount(items, active, in.AccountID) - if len(active) == 0 { - return nil, "", ErrNoProviderAccount - } - s.rotateRoundRobin("runway", active) - - imageSize := strings.TrimSpace(resolution) - if imageSize == "" { - imageSize = "1K" - } - var lastErr error - busy := 0 - for _, token := range active { - // Per-account concurrency gate (defaults to 1 for built-in pools). - if !s.acctAcquire(ctx, token.ID, eventID, accountConcurrency(token)) { - busy++ - continue - } - var data []byte - var artURL string - done, failover := func() (bool, bool) { - defer s.acctRelease(ctx, token.ID, eventID) - _ = s.events.SetAccount(ctx, eventID, token.ID, token.AccountEmail) - _ = s.tokens.TouchLastUsed(ctx, token.ID) - teamID := "" - if token.Meta != nil { - teamID = strings.TrimSpace(stringValue(token.Meta["team_id"])) - } - // downloadResult=false in url-only mode → skip the artifact download and - // just return meta["image_url"]. - d, meta, genErr := s.runway.GenerateImage(ctx, token.Value, teamID, modelItem.ID, in.Prompt, aspectRatio, imageSize, refs, !urlOnly) - 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 - artURL = strings.TrimSpace(stringValue(meta["image_url"])) - 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, artURL, 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 -// successful generation and writes it back (negative / unknown clamp to 0), -// flipping the account to 限额 when it hits 0 — so accounts limit one-by-one as -// they're used, not all at once on a later batch probe. Runs while the -// per-account concurrency gate is still held. Best-effort (never fails the render). -func (s *V1Service) reconcileChatGPTQuota(ctx context.Context, tokenID, accessToken string) { - if s.chatgpt == nil { - return - } - data, err := s.chatgpt.FetchImageQuota(ctx, accessToken) - if err != nil || boolValueWithDefault(data["auth_failed"], false) { - return - } - rem, exhausted := chatgptRemaining(data) - item, err := s.tokens.Get(ctx, "chatgpt", tokenID) - if err != nil { - return - } - meta := cloneJSONMap(item.Meta) - meta["cached_quota_remaining"] = rem - meta["cached_quota_at"] = int(time.Now().Unix()) - patch := map[string]any{"meta": meta} - if reset := strings.TrimSpace(stringValue(data["reset_after"])); reset != "" { - patch["cached_quota_reset_after"] = reset - } else if strings.TrimSpace(item.CachedQuotaResetAfter) == "" { - patch["cached_quota_reset_after"] = leonardoResetAfter("") - } - if exhausted && item.Status == "active" { - patch["status"] = "quota" - } - _, _ = s.tokens.Update(ctx, "chatgpt", tokenID, patch) -} - -// chatgpt image URLs are auth-gated (files.oaiusercontent.com — a plain GET -// 403s), so url-only mode returns the URL for the caller to proxy via -// OpenImageContent using the generating account's token. -func (s *V1Service) generateChatGPTImage(ctx context.Context, eventID string, modelItem *model.ModelConfig, in V1ImageRequest, aspectRatio, resolution string, noStore bool) ([]byte, string, error) { - urlOnly := noStore - if s.chatgpt == nil { - return nil, "", errors.New("chatgpt client not configured") - } - if s.settings != nil { - if proxy, err := s.settings.GetValue(ctx, "proxy.url"); err == nil { - s.chatgpt.SetProxy(proxy) - } - } - - items, err := s.tokens.ListByPool(ctx, "chatgpt") - if err != nil { - return nil, "", err - } - var active []model.TokenAccount - for _, item := range items { - if item.Status == "active" && !item.Dead && strings.TrimSpace(item.Value) != "" { - active = append(active, item) - } - } - active = pinTestAccount(items, active, in.AccountID) - if len(active) == 0 { - return nil, "", ErrNoProviderAccount - } - s.rotateRoundRobin("chatgpt", active) - - refLimit := modelItem.MaxReferenceImages - if refLimit <= 0 { - refLimit = 1 - } - refs, err := decodeReferenceImages(in.ReferenceImages, refLimit) - if err != nil { - return nil, "", err - } - - // Round-robin order; on a transient upstream error (e.g. "image generation - // did not start (no async marker)") FAIL OVER to the next account - // (tempFailover=true, capped at maxTempDeadAccounts) — never mark the - // account dead. Auth/quota fail over immediately (see runPoolWithFailover). - var imageURL string - data, err := s.runPoolWithFailover(ctx, eventID, "chatgpt", active, "image", func(token model.TokenAccount) ([]byte, error) { - d, meta, genErr := s.chatgpt.GenerateImage(ctx, token.Value, in.Prompt, modelItem.ID, aspectRatio, resolution, refs, !urlOnly) - if genErr == nil { - imageURL = strings.TrimSpace(stringValue(meta["image_url"])) - // Sync the real OpenAI quota BEFORE the concurrency gate releases, so the - // freshly-decremented remaining (and 限额 flip at 0) gates the next pick. - s.reconcileChatGPTQuota(ctx, token.ID, token.Value) - } - return d, genErr - }, func(e error) (bool, bool, bool, bool) { - return errors.Is(e, chatgpt.ErrAuth), errors.Is(e, chatgpt.ErrQuotaExhausted), errors.Is(e, chatgpt.ErrTemporaryUpstream), false - }, nil, true) // chatgpt token IS the credential — no cookie to refresh; switch accounts on transient errors - return data, imageURL, err -} - -// leonardoResetAfter returns when a Leonardo account's daily free tokens renew. -// Leonardo resets at 08:00 Beijing == 00:00 UTC, so when the upstream gives no -// explicit renewal time we deterministically use the next UTC midnight — this is -// filled at import so 恢复时间 is always populated, not left blank. -func leonardoResetAfter(availableUntil string) string { - if v := strings.TrimSpace(availableUntil); v != "" { - return v - } - return time.Unix((time.Now().Unix()/86400+1)*86400, 0).UTC().Format(time.RFC3339) -} - -// leonardoDimensions maps the catalog's resolution+ratio to Leonardo pixel sizes. -func leonardoDimensions(resolution, aspectRatio string) (int, int) { - res := strings.ToUpper(strings.TrimSpace(resolution)) - ar := strings.TrimSpace(aspectRatio) - if res == "4K" { - switch ar { - case "2:3": - return 2000, 3000 - case "16:9": - return 4096, 2304 - case "4:3": - return 4096, 3072 - case "4:5": - return 3264, 4080 - case "9:16": - return 2160, 3840 - case "2:1": - return 4096, 2048 - default: // 1:1 - return 4096, 4096 - } - } - switch ar { // 2K (default) - case "2:3": - return 1664, 2496 - case "16:9": - return 2560, 1440 - case "4:3": - return 2304, 1728 - case "4:5": - return 2432, 3040 - case "9:16": - return 1440, 2560 - case "2:1": - return 3232, 1616 - default: // 1:1 - return 2048, 2048 - } -} - -func (s *V1Service) generateLeonardoImage(ctx context.Context, eventID string, modelItem *model.ModelConfig, in V1ImageRequest, aspectRatio, resolution string, noStore bool) ([]byte, string, error) { - urlOnly := noStore - if s.leonardo == nil { - return nil, "", errors.New("leonardo client not configured") - } - if s.settings != nil { - if proxy, err := s.settings.GetValue(ctx, "proxy.url"); err == nil { - s.leonardo.SetProxy(proxy) - } - } - - items, err := s.tokens.ListByPool(ctx, "leonardo") - 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 - } - // Skip accounts under the per-generation floor (treated as 限额). Unknown - // balance gets the benefit of the doubt (upstream rejects if truly empty). - if rem, ok := jsonMapInt(item.Meta, "cached_quota_remaining"); ok && rem < leonardoMinCredits { - continue - } - active = append(active, item) - } - active = pinTestAccount(items, active, in.AccountID) - if len(active) == 0 { - return nil, "", ErrNoProviderAccount - } - s.rotateRoundRobin("leonardo", active) - - width, height := leonardoDimensions(resolution, aspectRatio) - // The catalog model id is the upstream Leonardo model name (e.g. seedream-4.5). - upstreamModel := strings.TrimSpace(modelItem.ID) - - // Optional image-to-image: decode the reference image once up front (Leonardo - // seedream takes at most one). - refLimit := modelItem.MaxReferenceImages - if refLimit <= 0 { - refLimit = 1 - } - refs, err := decodeReferenceImages(in.ReferenceImages, refLimit) - if err != nil { - return nil, "", err - } - - // token.Value is the cookie; GenerateImage mints a fresh JWT each attempt (and - // re-mints it internally when the bearer is rejected), so an auth failure means - // the cookie itself no longer authenticates — no refresher (nil). - var imageURL string - data, err := s.runPoolWithFailover(ctx, eventID, "leonardo", active, "image", func(token model.TokenAccount) ([]byte, error) { - // Atomically pre-deduct the per-generation cost so concurrent picks of the - // same near-empty account can't over-commit it. A known-insufficient - // balance surfaces as quota → the driver fails over to the next account. - allowed, deducted, rerr := s.tokens.ReserveQuota(ctx, "leonardo", token.ID, leonardoMinCredits) - if rerr != nil { - return nil, fmt.Errorf("%w: reserve: %v", leonardo.ErrTemporaryUpstream, rerr) - } - if !allowed { - return nil, leonardo.ErrQuotaExhausted - } - data, meta, genErr := s.leonardo.GenerateImage(ctx, token.Value, upstreamModel, in.Prompt, width, height, nil, refs, !urlOnly) - cookie := s.leonardoPersistCookie(ctx, token.ID, token.Value) - if genErr != nil { - // Release the hold so a failed render doesn't burn credits. - if deducted { - _ = s.tokens.RefundQuota(ctx, "leonardo", token.ID, leonardoMinCredits) - } - return nil, genErr - } - imageURL = strings.TrimSpace(stringValue(meta["image_url"])) - // Success → overwrite the held value with the REAL upstream balance and - // sink to 限额 if below the floor (best-effort; never fails a done render). - s.reconcileLeonardoCredits(ctx, token.ID, cookie) - return data, nil - }, func(e error) (bool, bool, bool, bool) { - return errors.Is(e, leonardo.ErrAuth), errors.Is(e, leonardo.ErrQuotaExhausted), errors.Is(e, leonardo.ErrTemporaryUpstream), false - }, nil, true) - return data, imageURL, err -} - -// leonardoPrivateSuffix 标记目录里 Leonardo 私有视频模型(public:false 生成)的 -// 后缀,上游 slug 就是去掉它之后的 id。 -const leonardoPrivateSuffix = "-不卡人脸" - -// leonardoVideoSpec 描述一个 Leonardo 视频模型的上游 slug、输出尺寸和参考资产 -// 限制(各类上限 + 时长约束),校验在扣费前跑,免得非法参考白扣积分。 -type leonardoVideoSpec struct { - upstream string - // long/short 是长边/短边像素,按比例组合成 16:9 或 9:16。 - long int - short int - maxImages int - maxAudios int - maxVideos int - // 视频参考单个时长区间与总时长上限(秒),0 表示不限。 - videoMinSeconds float64 - videoMaxSeconds float64 - videoTotalSeconds float64 - // 音频参考总时长上限(秒),0 表示不限。 - audioTotalSeconds float64 -} - -var leonardoVideoSpecs = map[string]leonardoVideoSpec{ - "seedance-2.0" + leonardoPrivateSuffix: { - upstream: "seedance-2.0", long: 1280, short: 720, - maxImages: 4, maxAudios: 1, maxVideos: 3, - videoMinSeconds: 3, videoMaxSeconds: 10, videoTotalSeconds: 15, - }, - "seedance-2.0-fast" + leonardoPrivateSuffix: { - upstream: "seedance-2.0-fast", long: 1280, short: 720, - maxImages: 4, maxAudios: 1, maxVideos: 3, - videoMinSeconds: 3, videoMaxSeconds: 10, videoTotalSeconds: 15, - }, - "minimax-h3": { - upstream: "hailuo-03", long: 2560, short: 1440, - maxImages: 5, maxAudios: 3, maxVideos: 0, - audioTotalSeconds: 15, - }, -} - -func leonardoVideoSpecOf(modelID string) leonardoVideoSpec { - id := strings.TrimSpace(modelID) - if spec, ok := leonardoVideoSpecs[id]; ok { - return spec - } - // 目录里新增的同族模型退化成 seedance 规格,上游 slug 取去掉私有后缀的 id。 - return leonardoVideoSpec{ - upstream: strings.TrimSuffix(id, leonardoPrivateSuffix), long: 1280, short: 720, - maxImages: 4, maxAudios: 1, maxVideos: 3, - videoMinSeconds: 3, videoMaxSeconds: 10, videoTotalSeconds: 15, - } -} - -// dimensions 按比例返回像素尺寸,只支持 16:9 / 9:16。 -func (spec leonardoVideoSpec) dimensions(aspectRatio string) (int, int) { - if strings.TrimSpace(aspectRatio) == "9:16" { - return spec.short, spec.long - } - return spec.long, spec.short -} - -// classifyLeonardoVideoRefs decodes the mixed reference payload (画图台把图片/ -// 音频/视频一起塞进 reference_images)并按类型分流,同时按模型规格校验各类上限和 -// 时长。 -func classifyLeonardoVideoRefs(inputs []string, spec leonardoVideoSpec) (leonardo.VideoAssets, error) { - var refs leonardo.VideoAssets - decoded, err := decodeReferenceImages(inputs, spec.maxImages+spec.maxAudios+spec.maxVideos) - if err != nil { - return refs, err - } - var videoTotal, audioTotal float64 - for _, ref := range decoded { - switch detectMediaType(ref) { - case "video": - if spec.maxVideos == 0 { - return refs, errors.New("该模型不支持视频参考") - } - secs := leonardo.MediaDurationSeconds(ref) - if secs <= 0 { - return refs, errors.New("无法解析视频参考的时长,请换一个 mp4") - } - if secs < spec.videoMinSeconds || secs > spec.videoMaxSeconds { - return refs, fmt.Errorf("单个视频参考需 %.0f-%.0f 秒,当前 %.1f 秒", - spec.videoMinSeconds, spec.videoMaxSeconds, secs) - } - videoTotal += secs - if spec.videoTotalSeconds > 0 && videoTotal >= spec.videoTotalSeconds { - return refs, fmt.Errorf("视频参考总时长 %.1f 秒,需短于 %.0f 秒", - videoTotal, spec.videoTotalSeconds) - } - refs.Videos = append(refs.Videos, ref) - case "audio": - if spec.maxAudios == 0 { - return refs, errors.New("该模型不支持音频参考") - } - secs := leonardo.MediaDurationSeconds(ref) - if spec.audioTotalSeconds > 0 { - if secs <= 0 { - return refs, errors.New("无法解析音频参考的时长,请换一个 mp3") - } - audioTotal += secs - if audioTotal > spec.audioTotalSeconds { - return refs, fmt.Errorf("音频参考总时长 %.1f 秒,最多 %.0f 秒", - audioTotal, spec.audioTotalSeconds) - } - } - refs.Audios = append(refs.Audios, ref) - default: - refs.Images = append(refs.Images, ref) - } - } - switch { - case len(refs.Images) > spec.maxImages: - return refs, fmt.Errorf("最多 %d 张参考图", spec.maxImages) - case len(refs.Audios) > spec.maxAudios: - return refs, fmt.Errorf("最多 %d 段音频参考", spec.maxAudios) - case len(refs.Videos) > spec.maxVideos: - return refs, fmt.Errorf("最多 %d 段视频参考", spec.maxVideos) - } - return refs, nil -} - -// generateLeonardoVideo renders a Leonardo video across the leonardo pool. Mirrors -// generateLeonardoImage (cookie → JWT, quota reserve, cookie 轮换持久化), only the -// upstream call differs: 私有生成 + 三类参考资产 + motionMP4URL。 -func (s *V1Service) generateLeonardoVideo(ctx context.Context, eventID string, modelItem *model.ModelConfig, in V1VideoRequest, aspectRatio string, durationSeconds int, downloadResult bool) ([]byte, string, error) { - if s.leonardo == nil { - return nil, "", errors.New("leonardo client not configured") - } - if s.settings != nil { - if proxy, err := s.settings.GetValue(ctx, "proxy.url"); err == nil { - s.leonardo.SetProxy(proxy) - } - } - spec := leonardoVideoSpecOf(modelItem.ID) - refs, err := classifyLeonardoVideoRefs(in.ReferenceImages, spec) - if err != nil { - return nil, "", err - } - - items, err := s.tokens.ListByPool(ctx, "leonardo") - 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 < leonardoMinCredits { - continue - } - active = append(active, item) - } - active = pinTestAccount(items, active, in.AccountID) - if len(active) == 0 { - return nil, "", ErrNoProviderAccount - } - s.rotateRoundRobin("leonardo", active) - - width, height := spec.dimensions(aspectRatio) - - var videoURL string - data, err := s.runPoolWithFailover(ctx, eventID, "leonardo", active, "video", func(token model.TokenAccount) ([]byte, error) { - allowed, deducted, rerr := s.tokens.ReserveQuota(ctx, "leonardo", token.ID, leonardoMinCredits) - if rerr != nil { - return nil, fmt.Errorf("%w: reserve: %v", leonardo.ErrTemporaryUpstream, rerr) - } - if !allowed { - return nil, leonardo.ErrQuotaExhausted - } - data, meta, genErr := s.leonardo.GenerateVideo(ctx, token.Value, spec.upstream, in.Prompt, width, height, durationSeconds, refs, downloadResult) - cookie := s.leonardoPersistCookie(ctx, token.ID, token.Value) - if genErr != nil { - if deducted { - _ = s.tokens.RefundQuota(ctx, "leonardo", token.ID, leonardoMinCredits) - } - return nil, genErr - } - videoURL = strings.TrimSpace(stringValue(meta["video_url"])) - s.reconcileLeonardoCredits(ctx, token.ID, cookie) - return data, nil - }, func(e error) (bool, bool, bool, bool) { - return errors.Is(e, leonardo.ErrAuth), errors.Is(e, leonardo.ErrQuotaExhausted), errors.Is(e, leonardo.ErrTemporaryUpstream), false - }, nil, true) - return data, videoURL, err -} - -// leonardoPersistCookie writes back the account cookie when Leonardo rotated its -// better-auth session_data cache (that cache is what actually authenticates -// get-session, so a stale copy would eventually look like a dead account). -// Returns the value now stored. -func (s *V1Service) leonardoPersistCookie(ctx context.Context, tokenID, cookie string) string { - if s.leonardo == nil { - return cookie - } - fresh, ok := s.leonardo.RotatedCookie(cookie) - if !ok || strings.TrimSpace(fresh) == "" { - return cookie - } - _, _ = s.tokens.SwapValue(ctx, "leonardo", tokenID, cookie, fresh) - return fresh -} - -// reconcileLeonardoCredits re-fetches an account's real token balance after a -// render and writes it back, flipping the account to 限额 when below the per-gen -// floor. Stores the daily renewal time so RecoverQuota can auto-recover it. -func (s *V1Service) reconcileLeonardoCredits(ctx context.Context, tokenID, cookie string) { - if s.leonardo == nil { - return - } - data, err := s.leonardo.FetchCreditsBalance(ctx, cookie) - s.leonardoPersistCookie(ctx, tokenID, cookie) - if err != nil { - return - } - rem, ok := data["remaining"].(int) - if !ok { - return - } - item, err := s.tokens.Get(ctx, "leonardo", tokenID) - if err != nil { - return - } - meta := cloneJSONMap(item.Meta) - meta["cached_quota_remaining"] = rem - meta["cached_quota_at"] = int(time.Now().Unix()) - patch := map[string]any{"meta": meta} - patch["cached_quota_reset_after"] = leonardoResetAfter(stringValue(data["available_until"])) - if rem < leonardoMinCredits && item.Status == "active" { - patch["status"] = "quota" - } - _, _ = s.tokens.Update(ctx, "leonardo", tokenID, patch) -} - -// kreaRefreshAndPersist ensures the account's Krea cookie has a valid access token -// (refreshing via the rotating refresh_token when expired) and persists the new -// cookie — the refresh_token is single-use, so the rotated value MUST be saved. -func kreaRefreshAndPersist(ctx context.Context, client *krea.Client, tokens *repo.TokenRepository, tokenID, cookie string) (string, error) { - if client == nil { - return cookie, nil - } - fresh, changed, err := client.RefreshIfNeeded(ctx, cookie) - if err != nil { - return "", err - } - if changed && tokenID != "" { - _, _ = tokens.Update(ctx, "krea", tokenID, map[string]any{"value": fresh}) - } - return fresh, nil -} - -// kreaDimensions maps the catalog's resolution+ratio to Krea pixel sizes. -func kreaDimensions(resolution, aspectRatio string) (int, int) { - res := strings.ToUpper(strings.TrimSpace(resolution)) - ar := strings.TrimSpace(aspectRatio) - if res == "2K" { - switch ar { - case "4:3": - return 2048, 1536 - case "3:4": - return 1536, 2048 - case "16:9": - return 2048, 1152 - case "9:16": - return 1152, 2048 - default: // 1:1 - return 2048, 2048 - } - } - switch ar { // 1K (default) - case "4:3": - return 1024, 768 - case "3:4": - return 768, 1024 - case "16:9": - return 1024, 576 - case "9:16": - return 576, 1024 - default: // 1:1 - return 1024, 1024 - } -} - -func (s *V1Service) generateKreaImage(ctx context.Context, eventID string, modelItem *model.ModelConfig, in V1ImageRequest, aspectRatio, resolution string, noStore bool) ([]byte, string, error) { - urlOnly := noStore - if s.krea == nil { - return nil, "", errors.New("krea client not configured") - } - if s.settings != nil { - if proxy, err := s.settings.GetValue(ctx, "proxy.url"); err == nil { - s.krea.SetProxy(proxy) - } - } - - items, err := s.tokens.ListByPool(ctx, "krea") - if err != nil { - return nil, "", err - } - var active []model.TokenAccount - for _, item := range items { - // No numeric floor — Krea signals 限额 with a 402 at generation time, which - // the failover driver turns into mark-quota + next account. - if item.Status == "active" && !item.Dead && strings.TrimSpace(item.Value) != "" { - active = append(active, item) - } - } - active = pinTestAccount(items, active, in.AccountID) - if len(active) == 0 { - return nil, "", ErrNoProviderAccount - } - s.rotateRoundRobin("krea", active) - - width, height := kreaDimensions(resolution, aspectRatio) - refLimit := modelItem.MaxReferenceImages - if refLimit <= 0 { - refLimit = 1 - } - refs, err := decodeReferenceImages(in.ReferenceImages, refLimit) - if err != nil { - return nil, "", err - } - - var imageURL string - data, err := s.runPoolWithFailover(ctx, eventID, "krea", active, "image", func(token model.TokenAccount) ([]byte, error) { - // Refresh the (rotating) Supabase token if expired and persist the new - // cookie, then generate with the fresh cookie. - cookie, rerr := kreaRefreshAndPersist(ctx, s.krea, s.tokens, token.ID, token.Value) - if rerr != nil { - return nil, rerr - } - data, meta, genErr := s.krea.GenerateImage(ctx, cookie, in.Prompt, width, height, refs, !urlOnly) - if genErr == nil { - imageURL = strings.TrimSpace(stringValue(meta["image_url"])) - } - return data, genErr - }, func(e error) (bool, bool, bool, bool) { - return errors.Is(e, krea.ErrAuth), errors.Is(e, krea.ErrQuotaExhausted), errors.Is(e, krea.ErrTemporaryUpstream), false - }, nil, true) - return data, imageURL, err -} - -// imagineRefreshAndPersist ensures the account's Imagine credential has a valid -// access token (refreshing via the rotating refreshToken when expired) and -// persists the new credential — both tokens rotate, so the value MUST be saved. -func imagineRefreshAndPersist(ctx context.Context, client *imagine.Client, tokens *repo.TokenRepository, tokenID, cred string) (string, error) { - if client == nil { - return cred, nil - } - fresh, changed, err := client.RefreshIfNeeded(ctx, cred) - if err != nil { - return "", err - } - if changed && tokenID != "" { - _, _ = tokens.Update(ctx, "imagine", tokenID, map[string]any{"value": fresh}) - } - return fresh, nil -} - -// imagineStyle maps the catalog model id to its upstream style_id + resolution. -func imagineStyle(modelID string) (int, string) { - if strings.TrimSpace(modelID) == "imagine-1.5pro" { - return 41004, "4K" - } - return 41001, "2K" -} - -func (s *V1Service) generateImagineImage(ctx context.Context, eventID string, modelItem *model.ModelConfig, in V1ImageRequest, aspectRatio, resolution string, noStore bool) ([]byte, string, error) { - urlOnly := noStore - if s.imagine == nil { - return nil, "", errors.New("imagine client not configured") - } - if s.settings != nil { - if proxy, err := s.settings.GetValue(ctx, "proxy.url"); err == nil { - s.imagine.SetProxy(proxy) - } - } - - items, err := s.tokens.ListByPool(ctx, "imagine") - if err != nil { - return nil, "", err - } - var active []model.TokenAccount - for _, item := range items { - // No numeric floor — Imagine signals 限额 with a 402 at generation time, - // which the failover driver turns into mark-quota + next account. - if item.Status == "active" && !item.Dead && strings.TrimSpace(item.Value) != "" { - active = append(active, item) - } - } - active = pinTestAccount(items, active, in.AccountID) - if len(active) == 0 { - return nil, "", ErrNoProviderAccount - } - s.rotateRoundRobin("imagine", active) - - // Each model supports exactly one resolution (2K / 4K) — force it per model. - styleID, res := imagineStyle(modelItem.ID) - - var imageURL string - data, err := s.runPoolWithFailover(ctx, eventID, "imagine", active, "image", func(token model.TokenAccount) ([]byte, error) { - // Refresh the (rotating) access token if expired and persist the new - // credential, then generate with the fresh token. - cred, rerr := imagineRefreshAndPersist(ctx, s.imagine, s.tokens, token.ID, token.Value) - if rerr != nil { - return nil, rerr - } - data, meta, genErr := s.imagine.GenerateImage(ctx, cred, styleID, res, aspectRatio, in.Prompt, !urlOnly) - if genErr != nil { - return nil, genErr - } - imageURL = strings.TrimSpace(stringValue(meta["image_url"])) - return data, nil - }, func(e error) (bool, bool, bool, bool) { - return errors.Is(e, imagine.ErrAuth), errors.Is(e, imagine.ErrQuotaExhausted), errors.Is(e, imagine.ErrTemporaryUpstream), false - }, nil, true) - return data, imageURL, err -} - func (s *V1Service) refundIfNeeded(ctx context.Context, principal *APIPrincipal, eventID string, price float64) error { if principal == nil || principal.User == nil || price <= 0 { return nil @@ -3509,11 +1486,76 @@ func (s *V1Service) maybeGrantInviteReward(ctx context.Context, principal *APIPr return err } -// ensureReferenceSizes rejects any reference image over the byte cap BEFORE -// charging, so an oversized image fails fast (no charge, no pending-log churn) -// across every entry path — session /generate, API-key /v1, and admin /test. -// decodeReferenceImages re-checks at decode time as a backstop; this mirrors its -// base64 length pre-check (decoded ≈ len(b64)*3/4). +func (s *V1Service) userDir(principal *APIPrincipal) string { + if principal == nil { + return "anon" + } + return OwnerDir(principal.User) +} + +// OwnerDir is the storage directory (= /images// segment) a user's outputs +// live under: sanitized name → sanitized email-local → id → "anon". +func OwnerDir(user *model.User) string { + if user != nil { + if d := sanitizeOwnerName(user.Name); d != "" { + return d + } + if d := sanitizeOwnerName(strings.Split(user.Email, "@")[0]); d != "" { + return d + } + if user.ID != "" { + return user.ID + } + } + return "anon" +} + +// allocateOutput builds the object key (= relative path, user-scoped) and the +// /images proxy URL. Nothing is written here — the caller uploads the bytes. +func (s *V1Service) allocateOutput(principal *APIPrincipal, ext, baseURL string) (string, string) { + userDir := s.userDir(principal) + filename := time.Now().Format("20060102-150405") + "-" + randomUpper(8) + "." + strings.TrimPrefix(ext, ".") + relativePath := filepath.ToSlash(filepath.Join(userDir, filename)) + if base := strings.TrimRight(strings.TrimSpace(baseURL), "/"); base != "" { + return base + "/images/" + relativePath, relativePath + } + return "/images/" + relativePath, relativePath +} + +func (s *V1Service) logPendingEvent(ctx context.Context, kind string, modelItem *model.ModelConfig, principal *APIPrincipal, prompt, ratio, resolution, duration string, refs int, cost float64, file, source string, refFiles []string, deai bool) (string, error) { + event := &model.EventLog{ + ID: "evt-" + randomUpper(12), + TS: time.Now(), + Kind: kind, + Status: "pending", + Model: modelItem.ID, + Provider: modelItem.Provider, + Prompt: prompt, + Ratio: ratio, + Resolution: resolution, + Duration: duration, + Refs: refs, + DeAI: deai, + Source: source, + Cost: cost, + File: file, + CreatedAt: time.Now(), + UpdatedAt: time.Now(), + } + if len(refFiles) > 0 { + event.RefFiles = jsonArray(refFiles) + } + if principal != nil && principal.User != nil { + event.UserID = principal.User.ID + } + if err := s.events.Create(ctx, event); err != nil { + return "", err + } + return event.ID, nil +} + +// ===== small helpers ===== + func ensureReferenceSizes(inputs []string) error { for _, raw := range inputs { v := strings.TrimSpace(raw) @@ -3540,10 +1582,7 @@ func decodeReferenceImages(inputs []string, limit int) ([][]byte, error) { if v == "" { continue } - // Only raw base64 is accepted (no "data:...;base64," URL prefix). A data - // URL now fails to decode rather than being silently stripped. - // decoded size ≈ len(b64) * 3 / 4 — reject oversized payloads up front, - // before allocating the decoded buffer. + // Only raw base64 is accepted (no "data:...;base64," URL prefix). if (len(v)*3)/4 > maxReferenceImageBytes { return nil, ErrReferenceTooLarge } @@ -3565,38 +1604,16 @@ func decodeReferenceImages(inputs []string, limit int) ([][]byte, error) { return out, nil } -// detectMediaType inspects the first bytes of a decoded reference to classify it -// as "video", "audio", or "image". Used to route refs to the correct upload MIME -// and the correct referenceBlobs usage for seedance. -func detectMediaType(data []byte) string { - n := len(data) - if n < 8 { - return "image" +// totalRefLimit is the combined per-kind cap for a model. +func totalRefLimit(item *model.ModelConfig) int { + if item == nil { + return 0 } - // MP4 / ISOBMFF - if n >= 12 && string(data[4:8]) == "ftyp" { - return "video" + n := item.MaxReferenceImages + item.MaxReferenceVideos + item.MaxReferenceAudios + if item.ImageToImage && n == 0 { + n = 1 } - // WebM - if n >= 4 && data[0] == 0x1A && data[1] == 0x45 && data[2] == 0xDF && data[3] == 0xA3 { - return "video" - } - // MP3: ID3 header or sync word 0xFFFx - if n >= 3 && data[0] == 0x49 && data[1] == 0x44 && data[2] == 0x33 { - return "audio" - } - if n >= 2 && data[0] == 0xFF && (data[1]&0xE0) == 0xE0 { - return "audio" - } - // WAV - if n >= 4 && data[0] == 0x52 && data[1] == 0x49 && data[2] == 0x46 && data[3] == 0x46 { - return "audio" - } - // OGG - if n >= 4 && data[0] == 0x4F && data[1] == 0x67 && data[2] == 0x67 && data[3] == 0x53 { - return "audio" - } - return "image" + return n } func parseImageSize(size, aspectRatio, resolution string) (string, string) { @@ -3635,8 +1652,6 @@ func parseImageSize(size, aspectRatio, resolution string) (string, string) { } // snapRatio returns the entry in supported closest in value to ar ("W:H"). -// ar is returned as-is when it's already supported, unparsable, or the model -// has no ratio list. func snapRatio(ar string, supported []string) string { parse := func(s string) (float64, bool) { var w, h int @@ -3673,11 +1688,6 @@ func guessRatio(w, h int) string { W int H int } - // The 17 ratios actually used across our models. Must stay in sync with the - // custom-model picker (CustomModelModal RATIO_OPTS) and the docs 对照表, so a - // /v1 `size` maps to exactly one of them. 9:21 is intentionally absent — - // no image provider accepts it (Runway 400s on it). snapRatio then clamps - // the guess to the target model's own supported list. candidates := []candidate{ {1, 1}, {5, 4}, {4, 3}, {3, 2}, {16, 9}, {2, 1}, {21, 9}, {3, 1}, {4, 1}, {8, 1}, // 横 @@ -3695,11 +1705,8 @@ func guessRatio(w, h int) string { return fmt.Sprintf("%d:%d", best.W, best.H) } -// firstPricedResolution returns the model's lowest priced image tier (1K/2K/4K -// order), or "" if none is priced. Used to rescue a request whose resolution -// the model doesn't support. // deaiEnabled reports whether the 去AI特征 feature is switched on in system -// settings (default off). When off, an incoming deai flag is ignored entirely. +// settings (default off). func (s *V1Service) deaiEnabled(ctx context.Context) bool { if s.settings == nil { return false @@ -3744,59 +1751,22 @@ func firstPricedResolution(item *model.ModelConfig) string { return r } } + // Fall back to whatever priced tier exists (video-style labels like 720p). + for _, r := range repo.JSONStrings(item.Resolutions) { + if _, ok := jsonMapFloat(item.Prices, r); ok { + return r + } + } return "" } -// resolutionForQuality maps OpenAI's `quality` to one of the model's priced -// resolution tiers: low→1K, medium→2K, high→4K, auto/blank→the model's lowest -// priced tier. The desired tier is clamped to the nearest tier the model -// actually prices (e.g. seedream is 2K/4K only: low→2K, high→4K). -func resolutionForQuality(item *model.ModelConfig, quality string) string { - order := []string{"1K", "2K", "4K"} - var priced []string - for _, r := range order { - if _, ok := jsonMapFloat(item.Prices, r); ok { - priced = append(priced, r) - } - } - if len(priced) == 0 { - return firstPricedResolution(item) - } - rank := map[string]int{"low": 0, "medium": 1, "high": 2} - want, ok := rank[strings.ToLower(strings.TrimSpace(quality))] - if !ok { - return priced[0] // auto / unknown → model default (lowest priced) - } - idxOf := func(r string) int { - for i, v := range order { - if v == r { - return i - } - } - return 0 - } - best, bestDist := priced[0], 99 - for _, r := range priced { - d := idxOf(r) - want - if d < 0 { - d = -d - } - if d < bestDist { - best, bestDist = r, d - } - } - return best -} - // modelPrice returns the charge for (kind, resolution, duration). The set of -// supported tiers is always driven by the NORMAL prices; `agent` only overrides -// the amount with the agent price when one is set for that tier (else it falls -// back to the normal price). +// supported tiers is driven by the NORMAL prices; `agent` only overrides the +// amount with the agent price when one is set for that tier. func modelPrice(item *model.ModelConfig, kind, resolution, duration string, agent bool) (float64, bool) { if item == nil { return 0, false } - // tierPrice: normal price gates support; agent price (if present) overrides. tierPrice := func(normal, agentMap map[string]any, key string) (float64, bool) { nv, ok := jsonMapFloat(normal, key) if !ok { @@ -3852,8 +1822,7 @@ func jsonMapFloat(m map[string]any, key string) (float64, bool) { return float64(x), true case json.Number: // datatypes.JSONMap.Scan decodes with UseNumber(), so values loaded from - // the DB arrive as json.Number — NOT float64. Without this case every - // price read back from Postgres looked "unpriced". + // the DB arrive as json.Number — NOT float64. if f, err := x.Float64(); err == nil { return f, true } @@ -3895,25 +1864,6 @@ func parseDurationSeconds(raw string) int { return n } -func resolveAdobeVideoEngine(modelID string) (string, string) { - switch strings.ToLower(strings.TrimSpace(modelID)) { - case "gemini-veo3.1-fast", "gemini-veo31", "firefly-veo31": - return "veo31-fast", "" - case "gemini-veo3.1": - return "veo31-standard", "" - case "adobe-seedance-2.0-fast": - return "seedance-2.0-fast", "" - case "adobe-seedance-2.0": - return "seedance-2.0", "" - case "firefly-ray": - return "luma", "" - case "firefly-video": - return "firefly-video", "" - default: - return "sora2", "" - } -} - func absFloat(v float64) float64 { if v < 0 { return -v @@ -3928,277 +1878,16 @@ func max(a, b int) int { return b } +func minInt(a, b int) int { + if a < b { + return a + } + return b +} + func principalCredits(principal *APIPrincipal) float64 { if principal == nil || principal.User == nil { return 0 } return principal.User.Credits } - -// markTokenFailure applies Python mark_bad semantics for a failed generation -// attempt against a pool token. It always bumps fail counters; the status side -// effects depend on the failure reason and the provider/pool. -// -// - quota: status="quota"; when no cached_quota_reset_after is present, set -// quota_recover_at to next UTC midnight so the maintenance loop can revive it. -// - auth on chatgpt: status="disabled" + dead=true (the access token IS the -// credential; a 401 means it's dead). -// - auth on adobe: NOT disabled/dead — the access token auto-refreshes from the -// cookie, so rotate for this request and let the refresh loop mint a new one. -// - other (non-auth/non-quota): NEITHER pool is auto-disabled — accounts stay -// active/green and fails is tracked only for rotation ordering. -func (s *V1Service) markTokenFailure(ctx context.Context, pool string, token model.TokenAccount, kind string, isAuth, isQuota bool) { - patch := map[string]any{ - "last_used_at": time.Now(), - "fail_total": gorm.Expr("fail_total + 1"), - "fails": gorm.Expr("fails + 1"), - } - switch { - case isQuota: - patch["status"] = "quota" - if strings.TrimSpace(token.CachedQuotaResetAfter) == "" { - recoverAt := time.Unix((time.Now().Unix()/86400+1)*86400, 0).UTC() - patch["quota_recover_at"] = &recoverAt - } - case isAuth: - // Adobe auth failures are NOT disabling: the access token refreshes from - // the cookie. chatgpt/runway/leonardo auth means the stored credential is - // dead — a raw JWT (chatgpt/runway) or a cookie whose session no longer - // authenticates (leonardo) — there's nothing left to refresh from. - // grok is intentionally excluded: a grok sso can momentarily 401 while - // still valid (upstream blip / proxy / anti-bot), so an auth failure just - // fails over for this request without permanently killing the account. - disable := pool == "chatgpt" || pool == "runway" || pool == "leonardo" || pool == "krea" || pool == "imagine" || pool == "creativefabrica" - if disable && pool == "leonardo" { - // 两道保险：先重新 get-session 复核（单次失败常是 bearer 轮换竞态），复核 - // 也不过就只记一次连续失败，连续到上限才判死。 - if s.leonardoCookieAlive(ctx, token) { - log.Printf("leonardo %s: auth failure on %s but cookie still authenticates — kept active", token.ID, kind) - disable = false - } else { - meta, strikes, kill := leonardoAuthStrike(&token, "auth failure on "+kind) - patch["meta"] = meta - disable = kill - if kill { - log.Printf("account leonardo/%s disabled after %d consecutive auth failures: %s", token.ID, strikes, kind) - } else { - log.Printf("leonardo %s: auth failure %d/%d on %s — kept active", token.ID, strikes, leonardoAuthStrikeLimit, kind) - } - } - } - if disable { - patch["status"] = "disabled" - patch["dead"] = true - if pool != "leonardo" { - log.Printf("account %s/%s disabled: auth failure on %s", pool, token.ID, kind) - } - } - default: - // Neither pool is auto-disabled on generic (non-auth / non-quota) failures - // — the account usually still works, so it stays active (green). fails is - // only tracked for rotation ordering. (A chatgpt *auth* failure still marks - // the token dead in the isAuth case above; that is a genuinely dead token.) - } - _, _ = s.tokens.Update(ctx, pool, token.ID, patch) -} - -// leonardoCookieAlive re-checks a Leonardo cookie after an auth failure by -// force-minting a session (bypassing the cached bearer). Only a cookie that -// still fails to authenticate counts as dead; a temporary upstream answer -// (403/429 人机校验) also keeps the account alive. -func (s *V1Service) leonardoCookieAlive(ctx context.Context, token model.TokenAccount) bool { - if s.leonardo == nil || strings.TrimSpace(token.Value) == "" { - return false - } - probeCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 30*time.Second) - defer cancel() - sess, err := s.leonardo.ProbeSession(probeCtx, token.Value) - if err == nil && sess != nil && strings.TrimSpace(sess.AccessToken) != "" { - s.leonardoPersistCookie(probeCtx, token.ID, token.Value) - return true - } - return !errors.Is(err, leonardo.ErrAuth) -} - -// markTokenDead disables an account and marks it dead on a fatal upstream error -// (a non-overload temporary Adobe failure that ops policy treats as account death). -func (s *V1Service) markTokenDead(ctx context.Context, pool string, token model.TokenAccount, kind string) { - _, _ = s.tokens.Update(ctx, pool, token.ID, map[string]any{ - "last_used_at": time.Now(), - "fail_total": gorm.Expr("fail_total + 1"), - "fails": gorm.Expr("fails + 1"), - "status": "disabled", - "dead": true, - }) -} - -// nextCursor returns the pool's current round-robin position and atomically -// advances it by one. Concurrent callers each get a distinct value, so parallel -// picks land on different accounts instead of racing onto the same one. The -// counter is in-memory (per process): it resets on restart, which only shifts -// the rotation's starting point — distribution stays even. -func (s *V1Service) nextCursor(pool string) uint64 { - v, _ := s.tokenCursors.LoadOrStore(pool, new(uint64)) - return atomic.AddUint64(v.(*uint64), 1) - 1 -} - -// rotateRoundRobin orders the active accounts by a stable key (ID) and rotates -// the slice in place so iteration begins at the pool's current cursor position, -// then advances the cursor. This is strict round-robin: account selection -// cycles in fixed order regardless of fails or last_used. The fall-through -// retry chain is preserved — on failure the caller's loop simply continues to -// the next account in rotation order. -// pinTestAccount narrows account selection to the single account requested by -// an admin 账号生图测试. The pinned account is taken from the pool's full list -// (bypassing active/dead/limited filters) so a limited or disabled account can -// still be probed. Returns nil when the account isn't in this pool. -func pinTestAccount(items, active []model.TokenAccount, accountID string) []model.TokenAccount { - id := strings.TrimSpace(accountID) - if id == "" { - return active - } - for _, item := range items { - if item.ID == id && strings.TrimSpace(item.Value) != "" { - return []model.TokenAccount{item} - } - } - return nil -} - -func (s *V1Service) rotateRoundRobin(pool string, items []model.TokenAccount) { - if len(items) <= 1 { - return - } - // Weight = priority: higher-weight accounts come first, so the scheduler tries - // them before lower-weight ones (and only falls through when they're at their - // concurrency cap). Within the SAME weight all accounts are equal, so they're - // rotated by the pool cursor for even distribution. - sort.SliceStable(items, func(i, j int) bool { - if items[i].Weight != items[j].Weight { - return items[i].Weight > items[j].Weight - } - return items[i].ID < items[j].ID - }) - start := int(s.nextCursor(pool)) - for i := 0; i < len(items); { - j := i + 1 - for j < len(items) && items[j].Weight == items[i].Weight { - j++ - } - if g := j - i; g > 1 { - off := start % g - if off != 0 { - grp := items[i:j] - rot := make([]model.TokenAccount, 0, g) - rot = append(rot, grp[off:]...) - rot = append(rot, grp[:off]...) - copy(grp, rot) - } - } - i = j - } -} - -// freeOnly1KModelID is the one free-allowed model 普号 may only serve at 1K -// (香蕉2 的 2K/4K 需要会员号) — see freeAccountsAllowed. -const freeOnly1KModelID = "nano-banana-2" - -func isSeedanceModel(modelID string) bool { - return modelID == "adobe-seedance-2.0-fast" || modelID == "adobe-seedance-2.0" -} - -// freeAccountsAllowed reports whether 普号(free) may serve this request: the model -// must be marked free_allowed. 香蕉2 另外只允许 1K 档，它的 2K/4K 只走会员号。 -func freeAccountsAllowed(modelItem *model.ModelConfig, resolution string) bool { - if modelItem == nil || !modelItem.FreeAllowed { - return false - } - if modelItem.ID == freeOnly1KModelID { - return strings.EqualFold(strings.TrimSpace(resolution), "1K") - } - return true -} - -func isFreeAccount(meta map[string]interface{}) bool { - if meta == nil { - return false - } - plan := strings.ToLower(strings.TrimSpace(stringValue(meta["plan"]))) - return plan == "free" -} - -// planUnknown 报告账号的会员身份还没探测出来（meta.plan 缺失或为空）。这类号 -// 既不能当普号也不能当会员号用：当会员号派出去会在需要会员的模型上撞 403 -// user_not_entitled。 -func planUnknown(meta map[string]interface{}) bool { - if meta == nil { - return true - } - return strings.TrimSpace(stringValue(meta["plan"])) == "" -} - -// markPlanUnknownDead 把选号时遇到的 plan 未探测账号置为死号，等重新探测到 -// plan 后再由额度刷新恢复。 -func (s *V1Service) markPlanUnknownDead(ctx context.Context, pool, id string) { - s.tokens.Update(ctx, pool, id, map[string]any{"status": "disabled", "dead": true}) -} - -// prioritizeSubAccounts 对非 Seedance 模型按 普号 → 子号 → 母号 的顺序排序： -// 先消耗普号，普号不可用再用低积分子号，最后才动 vip 母号。 -func prioritizeSubAccounts(active []model.TokenAccount) []model.TokenAccount { - var frees, subs, mothers []model.TokenAccount - for _, a := range active { - switch { - case isFreeAccount(a.Meta): - frees = append(frees, a) - case isLowCredits(a.Meta): - subs = append(subs, a) - default: - mothers = append(mothers, a) - } - } - return append(append(frees, subs...), mothers...) -} - -// isVipMotherAccount 正向识别 VIP 母号：plan 已知且非 free，且 is_sub_account -// 显式为 false。只看身份不看积分余额（低积分母号也可用）；plan 未探测或 -// is_sub_account 缺失的账号返回 false，等刷新补齐后才可被 Seedance 调度。 -func isVipMotherAccount(meta map[string]interface{}) bool { - if meta == nil { - return false - } - plan := strings.ToLower(strings.TrimSpace(stringValue(meta["plan"]))) - if plan == "" || plan == "free" { - return false - } - v, ok := meta["is_sub_account"] - if !ok { - return false - } - switch val := v.(type) { - case bool: - return !val - case float64: - return val == 0 - } - return false -} - -func isLowCredits(meta map[string]interface{}) bool { - if meta == nil { - return false - } - if v, ok := meta["is_sub_account"]; ok { - switch val := v.(type) { - case bool: - return val - case float64: - return val != 0 - } - } - // 兼容存量账号：is_sub_account 字段不存在时，用积分余额判断（>0 且 ≤4000 视为子号） - if rem, ok := jsonMapInt(meta, "cached_quota_remaining"); ok { - return rem > 0 && rem <= 4000 - } - return false -} diff --git a/backend/internal/storage/presign.go b/backend/internal/storage/presign.go new file mode 100644 index 0000000..56beefe --- /dev/null +++ b/backend/internal/storage/presign.go @@ -0,0 +1,83 @@ +package storage + +import ( + "encoding/hex" + "fmt" + "net/http" + "net/url" + "strconv" + "strings" + "time" +) + +// PresignGet issues a query-string SigV4 URL that lets anyone holding it GET +// the object for ttl. Used for /v1/generate result URLs and reference +// read-back. +func (c *Client) PresignGet(key string, ttl time.Duration) (string, error) { + return c.presign(http.MethodGet, key, "", ttl) +} + +// PresignPut issues a query-string SigV4 URL bound to PUT with the exact +// contentType signed in — the uploader must send that Content-Type verbatim +// (same contract vivid upstreams expect on /v1/uploads sessions). +func (c *Client) PresignPut(key, contentType string, ttl time.Duration) (string, error) { + if strings.TrimSpace(contentType) == "" { + return "", fmt.Errorf("presign put requires contentType") + } + return c.presign(http.MethodPut, key, contentType, ttl) +} + +// presign builds an AWS SigV4 presigned URL (UNSIGNED-PAYLOAD). Signed headers +// are host (+content-type for PUT). Validity is capped at 7 days by SigV4. +func (c *Client) presign(method, key, contentType string, ttl time.Duration) (string, error) { + if !c.Configured() { + return "", fmt.Errorf("rustfs not configured") + } + if ttl <= 0 { + ttl = 30 * time.Minute + } + if ttl > 7*24*time.Hour { + ttl = 7 * 24 * time.Hour + } + now := time.Now().UTC() + amzDate := now.Format("20060102T150405Z") + dateStamp := now.Format("20060102") + scope := dateStamp + "/" + region + "/" + service + "/aws4_request" + + signedHeaders := "host" + if contentType != "" { + signedHeaders = "content-type;host" + } + + q := map[string]string{ + "X-Amz-Algorithm": "AWS4-HMAC-SHA256", + "X-Amz-Credential": c.ak + "/" + scope, + "X-Amz-Date": amzDate, + "X-Amz-Expires": strconv.FormatInt(int64(ttl.Seconds()), 10), + "X-Amz-SignedHeaders": signedHeaders, + } + canonicalURI := "/" + uriEncode(c.bucket+"/"+strings.TrimPrefix(key, "/"), true) + canonicalQuery := canonicalQueryString(q) + + var canonHeaders strings.Builder + if contentType != "" { + canonHeaders.WriteString("content-type:" + contentType + "\n") + } + canonHeaders.WriteString("host:" + c.host + "\n") + + canonicalRequest := strings.Join([]string{ + method, canonicalURI, canonicalQuery, canonHeaders.String(), signedHeaders, "UNSIGNED-PAYLOAD", + }, "\n") + stringToSign := strings.Join([]string{ + "AWS4-HMAC-SHA256", amzDate, scope, hexSHA256([]byte(canonicalRequest)), + }, "\n") + sig := hex.EncodeToString(hmacSHA256(signingKey(c.sk, dateStamp), []byte(stringToSign))) + + // Rebuild the query with url.Values so the signature chars get escaped + // correctly, then append X-Amz-Signature verbatim (hex, safe). + vals := url.Values{} + for k, v := range q { + vals.Set(k, v) + } + return c.endpoint + canonicalURI + "?" + vals.Encode() + "&X-Amz-Signature=" + sig, nil +} diff --git a/backend/internal/upstream/client.go b/backend/internal/upstream/client.go new file mode 100644 index 0000000..04d1c47 --- /dev/null +++ b/backend/internal/upstream/client.go @@ -0,0 +1,631 @@ +// Package upstream is the vivid-protocol client used to talk to channel +// upstreams (other same-protocol instances). It covers the full public +// surface: POST /v1/generate (NDJSON stream with jobId reconnect), POST +// /v1/uploads (presigned reference upload), GET /v1/models, GET /v1/balance, +// plus presigned-URL result download. +package upstream + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "strings" + "time" +) + +// Error classes. Dispatch decisions (retry / fail over / fail the job) key off +// these sentinels via errors.Is — KindOf maps them for log messages. +var ( + // ErrAuth = channel credentials rejected (4011 / HTTP 401 / 403). The + // channel is skipped without consuming a retry budget. + ErrAuth = errors.New("upstream auth rejected") + // ErrInvalid = upstream rejected the request itself (4001 / HTTP 4xx). + // Retrying on another channel would fail identically — terminal. + ErrInvalid = errors.New("upstream rejected request") + // ErrJobFailed = the upstream ran the job and returned a failed frame. + // The frame's error text is user-safe and surfaced as-is. + ErrJobFailed = errors.New("upstream job failed") + // ErrRetryable = transient upstream failure (4293 / 5001 / 5xx / quota + // / timeout / transport). Fail over to the next channel. + ErrRetryable = errors.New("upstream temporary failure") + // ErrAmbiguous = the generate request was sent but the stream dropped + // before any jobId arrived — the job may or may not exist upstream. + // Treated as retryable for dispatch purposes. + ErrAmbiguous = errors.New("upstream submission ambiguous") +) + +// Kind is the coarse dispatch classification of an upstream error. +type Kind int + +const ( + KindRetryable Kind = iota + KindAuth + KindInvalid + KindJobFailed +) + +// KindOf classifies an error returned by this package. +func KindOf(err error) Kind { + switch { + case errors.Is(err, ErrAuth): + return KindAuth + case errors.Is(err, ErrInvalid): + return KindInvalid + case errors.Is(err, ErrJobFailed): + return KindJobFailed + default: + return KindRetryable + } +} + +// Failoverable reports whether the request should be re-tried on another +// channel. ErrInvalid and ErrJobFailed are deterministic — same request, same +// result — so they are terminal. +func Failoverable(err error) bool { + return !errors.Is(err, ErrInvalid) && !errors.Is(err, ErrJobFailed) +} + +// GenParams is the POST /v1/generate request body. +type GenParams struct { + Model string `json:"model"` + Prompt string `json:"prompt"` + Quality string `json:"quality,omitempty"` + Ratio string `json:"ratio,omitempty"` + Duration int64 `json:"duration,omitempty"` + Refs []string `json:"refs,omitempty"` + JobID string `json:"jobId,omitempty"` +} + +// RemoteModel mirrors one item of a vivid GET /v1/models response. +type RemoteModel struct { + ID string `json:"id"` + Kind string `json:"kind"` + Qualities []string `json:"qualities"` + Duration *RemoteDuration `json:"duration,omitempty"` + // Per-kind reference caps; 0 = that kind is not accepted. + RefImages int64 `json:"refImages"` + RefVideos int64 `json:"refVideos"` + RefAudios int64 `json:"refAudios"` + // Combined reference duration caps in seconds; 0 = unlimited. + RefVideoSeconds int64 `json:"refVideoSeconds,omitempty"` + RefAudioSeconds int64 `json:"refAudioSeconds,omitempty"` +} + +// RemoteDuration mirrors the duration spec of a vivid model item. +type RemoteDuration struct { + Mode string `json:"mode"` // fixed | options | range + Seconds int64 `json:"seconds,omitempty"` + Min int64 `json:"min,omitempty"` + Max int64 `json:"max,omitempty"` + Options []int64 `json:"options,omitempty"` +} + +// Client is stateless; per-call baseURL/key come from the channel row. +type Client struct { + HTTP *http.Client +} + +func NewClient() *Client { + return &Client{HTTP: &http.Client{Timeout: 30 * time.Minute}} +} + +func (c *Client) http() *http.Client { + if c != nil && c.HTTP != nil { + return c.HTTP + } + return NewClient().HTTP +} + +// maxStreamFailures caps reconnect attempts after a jobId is known — each +// reconnect re-opens /v1/generate with {jobId} and never double-charges. +const maxStreamFailures = 3 + +// Generate submits a vivid job and reads the NDJSON stream to a terminal +// frame, reconnecting with the reported jobId when the stream drops mid-run +// (upstream heartbeat keeps us attached; transport blips don't lose the job). +// onJobID, when non-nil, fires once the upstream assigns a job id. Returns the +// presigned result URL and the credits the upstream reports charging us. +func (c *Client) Generate(ctx context.Context, baseURL, apiKey string, p GenParams, waitSeconds int64, onJobID func(string)) (resultURL string, credits int64, err error) { + baseURL = strings.TrimRight(strings.TrimSpace(baseURL), "/") + if baseURL == "" || apiKey == "" { + return "", 0, ErrAuth + } + var last map[string]any + failures := 0 + for { + term, ferr := c.streamOnce(ctx, baseURL, apiKey, &p, waitSeconds, onJobID) + if ctx.Err() != nil { + return "", 0, ctx.Err() + } + if ferr != nil && !errors.Is(ferr, ErrRetryable) && !errors.Is(ferr, ErrAmbiguous) { + return "", 0, ferr + } + if ferr != nil { + failures++ + } else { + failures = 0 + } + if ferr == nil && term != nil { + status := strings.ToLower(strings.TrimSpace(str(term["status"]))) + if status != "running" && status != "" { + last = term + break + } + } + if p.JobID == "" { + if ferr != nil { + return "", 0, ferr + } + return "", 0, ErrAmbiguous + } + if failures >= maxStreamFailures { + if ferr != nil { + return "", 0, ferr + } + return "", 0, fmt.Errorf("%w: stream dropped repeatedly", ErrRetryable) + } + if sleepCtx(ctx, 2*time.Second) != nil { + return "", 0, ctx.Err() + } + } + status := strings.ToLower(strings.TrimSpace(str(last["status"]))) + switch status { + case "succeeded", "completed", "success": + resultURL = strings.TrimSpace(str(last["result"])) + if resultURL == "" { + return "", 0, fmt.Errorf("%w: succeeded frame missing result", ErrRetryable) + } + return resultURL, int64(num(last["creditsCharged"])), nil + default: + reason := strings.TrimSpace(str(last["error"])) + if reason == "" { + reason = "upstream job " + status + } + if isCreditError(reason) { + return "", 0, fmt.Errorf("%w: %s", ErrRetryable, clip([]byte(reason), 160)) + } + // failed 帧的 error 是上游给用户看的中文文案 — 透传 + return "", 0, fmt.Errorf("%w: %s", ErrJobFailed, clip([]byte(reason), 160)) + } +} + +// streamOnce performs one /v1/generate call and scans the NDJSON body. It +// records the upstream jobId into p.JobID as soon as a frame carries it. +func (c *Client) streamOnce(ctx context.Context, baseURL, apiKey string, p *GenParams, waitSeconds int64, onJobID func(string)) (map[string]any, error) { + raw, _ := json.Marshal(p) + req, err := http.NewRequestWithContext(ctx, http.MethodPost, baseURL+"/v1/generate", bytes.NewReader(raw)) + if err != nil { + return nil, err + } + req.Header.Set("Authorization", "Bearer "+apiKey) + req.Header.Set("Content-Type", "application/json") + if waitSeconds <= 0 { + waitSeconds = 300 + } + req.Header.Set("X-Wait-Seconds", fmt.Sprintf("%d", waitSeconds)) + resp, err := c.http().Do(req) + if err != nil { + if p.JobID == "" { + return nil, fmt.Errorf("%w: %s", ErrAmbiguous, sanitizeErr(err)) + } + return nil, fmt.Errorf("%w: %s", ErrRetryable, sanitizeErr(err)) + } + defer resp.Body.Close() + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + body, _ := io.ReadAll(resp.Body) + return nil, vividError(resp.StatusCode, body) + } + var last map[string]any + sc := bufio.NewScanner(resp.Body) + sc.Buffer(make([]byte, 0, 64*1024), 4*1024*1024) + for sc.Scan() { + line := bytes.TrimSpace(sc.Bytes()) + if len(line) == 0 { + continue // 25s heartbeat empty line + } + var frame map[string]any + if json.Unmarshal(line, &frame) != nil { + continue + } + last = frame + if job := strings.TrimSpace(str(frame["jobId"])); job != "" && job != p.JobID { + p.JobID = job + if onJobID != nil { + onJobID(job) + } + } + if status := strings.ToLower(strings.TrimSpace(str(frame["status"]))); status == "succeeded" || status == "failed" { + return frame, nil + } + } + if err := sc.Err(); err != nil { + if p.JobID == "" { + return nil, fmt.Errorf("%w: %s", ErrAmbiguous, sanitizeErr(err)) + } + return nil, fmt.Errorf("%w: %s", ErrRetryable, sanitizeErr(err)) + } + return last, nil +} + +// UploadRef runs the two-step reference upload: POST /v1/uploads {kind,mime} +// then PUT the bytes to the presigned url with the exact signed Content-Type. +// Returns the object key for refs[]. +func (c *Client) UploadRef(ctx context.Context, baseURL, apiKey string, data []byte) (string, error) { + baseURL = strings.TrimRight(strings.TrimSpace(baseURL), "/") + if baseURL == "" || apiKey == "" { + return "", ErrAuth + } + kind, mime := RefKindMime(data) + if kind == "" { + return "", fmt.Errorf("%w: 参考素材格式不支持", ErrInvalid) + } + raw, _ := json.Marshal(map[string]any{"kind": kind, "mime": mime}) + up, err := c.postJSON(ctx, baseURL+"/v1/uploads", apiKey, raw) + if err != nil { + return "", err + } + putURL, key := UploadSession(up) + if putURL == "" || key == "" { + rawUp, _ := json.Marshal(up) + return "", fmt.Errorf("%w: upload session missing url/key (got %s)", ErrRetryable, clip(rawUp, 200)) + } + req, err := http.NewRequestWithContext(ctx, http.MethodPut, putURL, bytes.NewReader(data)) + if err != nil { + return "", err + } + req.Header.Set("Content-Type", mime) // must match the signed mime exactly + resp, err := c.http().Do(req) + if err != nil { + return "", fmt.Errorf("%w: %s", ErrRetryable, sanitizeErr(err)) + } + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return "", fmt.Errorf("%w: ref upload %d %s", ErrRetryable, resp.StatusCode, clip(body, 120)) + } + return key, nil +} + +// UploadSession extracts the presigned PUT url + object key from a +// /v1/uploads response. Implementations vary: some return the pair at top +// level, others wrap it in {"data":{...}} / {"result":{...}}, and field +// spellings differ (upload_url / presigned_url / put_url; key / file_key / +// object_key). Search top level first, then one layer deep. +func UploadSession(up map[string]any) (putURL, key string) { + // Go map 查找区分大小写，而上游实现可能返回 PascalCase 字段；按去分隔符+ + // 小写归一化匹配。 + urlKeys := map[string]bool{"url": true, "uploadurl": true, "presignedurl": true, "puturl": true, "signedurl": true} + keyKeys := map[string]bool{"key": true, "filekey": true, "objectkey": true, "uploadkey": true} + norm := func(s string) string { return strings.ToLower(strings.NewReplacer("_", "", "-", "").Replace(s)) } + pick := func(m map[string]any) (string, string) { + var u, k string + for name, v := range m { + s := strings.TrimSpace(str(v)) + if s == "" { + continue + } + switch { + case u == "" && urlKeys[norm(name)]: + u = s + case k == "" && keyKeys[norm(name)]: + k = s + } + } + return u, k + } + levels := []map[string]any{up} + for _, nest := range []string{"data", "result", "upload", "session", "upload_session"} { + if sub, ok := up[nest].(map[string]any); ok { + levels = append(levels, sub) + } + } + for _, m := range levels { + u, k := pick(m) + if putURL == "" { + putURL = u + } + if key == "" { + key = k + } + } + return putURL, key +} + +// Models fetches the upstream's sellable model list (GET /v1/models → bare +// array of RemoteModel). +func (c *Client) Models(ctx context.Context, baseURL, apiKey string) ([]RemoteModel, error) { + baseURL = strings.TrimRight(strings.TrimSpace(baseURL), "/") + if baseURL == "" || apiKey == "" { + return nil, ErrAuth + } + req, err := http.NewRequestWithContext(ctx, http.MethodGet, baseURL+"/v1/models", nil) + if err != nil { + return nil, err + } + req.Header.Set("Authorization", "Bearer "+apiKey) + resp, err := c.http().Do(req) + if err != nil { + return nil, fmt.Errorf("%w: %s", ErrRetryable, sanitizeErr(err)) + } + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return nil, vividError(resp.StatusCode, body) + } + var items []RemoteModel + if err := json.Unmarshal(body, &items); err != nil { + return nil, fmt.Errorf("%w: bad /v1/models body", ErrRetryable) + } + return items, nil +} + +// Balance fetches the upstream balance (GET /v1/balance → {"balance":N}). +// Returns nil when the upstream does not implement the endpoint. +func (c *Client) Balance(ctx context.Context, baseURL, apiKey string) (*int64, error) { + baseURL = strings.TrimRight(strings.TrimSpace(baseURL), "/") + if baseURL == "" || apiKey == "" { + return nil, ErrAuth + } + req, err := http.NewRequestWithContext(ctx, http.MethodGet, baseURL+"/v1/balance", nil) + if err != nil { + return nil, err + } + req.Header.Set("Authorization", "Bearer "+apiKey) + resp, err := c.http().Do(req) + if err != nil { + return nil, fmt.Errorf("%w: %s", ErrRetryable, sanitizeErr(err)) + } + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + if resp.StatusCode == http.StatusNotFound { + return nil, nil + } + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return nil, vividError(resp.StatusCode, body) + } + var env struct { + Balance *int64 `json:"balance"` + } + if json.Unmarshal(body, &env) != nil || env.Balance == nil { + return nil, nil + } + return env.Balance, nil +} + +// Download fetches a presigned result URL — no Authorization header. +func (c *Client) Download(ctx context.Context, url string) ([]byte, string, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if err != nil { + return nil, "", err + } + resp, err := c.http().Do(req) + if err != nil { + return nil, "", fmt.Errorf("%w: %s", ErrRetryable, sanitizeErr(err)) + } + defer resp.Body.Close() + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return nil, "", fmt.Errorf("%w: result download %d", ErrRetryable, resp.StatusCode) + } + data, err := io.ReadAll(resp.Body) + if err != nil { + return nil, "", err + } + if len(data) == 0 { + return nil, "", fmt.Errorf("%w: empty result", ErrRetryable) + } + ct := strings.TrimSpace(resp.Header.Get("Content-Type")) + if ct == "" || ct == "application/octet-stream" { + ct = http.DetectContentType(data) + } + return data, ct, nil +} + +func (c *Client) postJSON(ctx context.Context, url, apiKey string, body []byte) (map[string]any, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body)) + if err != nil { + return nil, err + } + req.Header.Set("Authorization", "Bearer "+apiKey) + req.Header.Set("Content-Type", "application/json") + resp, err := c.http().Do(req) + if err != nil { + return nil, fmt.Errorf("%w: %s", ErrRetryable, sanitizeErr(err)) + } + defer resp.Body.Close() + raw, _ := io.ReadAll(resp.Body) + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return nil, vividError(resp.StatusCode, raw) + } + var out map[string]any + if err := json.Unmarshal(raw, &out); err != nil { + return nil, fmt.Errorf("%w: bad response %s", ErrRetryable, clip(raw, 120)) + } + return out, nil +} + +// vividError maps the {"error":{"code":"NNNN","message":...}} envelope to our +// error classes; falls back to plain HTTP status mapping. +func vividError(status int, body []byte) error { + var e struct { + Error struct { + Code string `json:"code"` + Message string `json:"message"` + } `json:"error"` + } + code := "" + msg := clip(body, 160) + if json.Unmarshal(body, &e) == nil { + code = strings.TrimSpace(e.Error.Code) + if m := strings.TrimSpace(e.Error.Message); m != "" { + msg = clip([]byte(m), 160) + } + } + switch code { + case "4011": + return fmt.Errorf("%w: %s", ErrAuth, msg) + case "4002": + return fmt.Errorf("%w: 积分不足 %s", ErrRetryable, msg) + case "4293", "5001": + return fmt.Errorf("%w: %s", ErrRetryable, msg) + case "4001": + return fmt.Errorf("%w: %s", ErrInvalid, msg) + } + switch { + case status == http.StatusUnauthorized || status == http.StatusForbidden: + return fmt.Errorf("%w: http %d", ErrAuth, status) + case status == http.StatusTooManyRequests || status >= 500: + return fmt.Errorf("%w: http %d %s", ErrRetryable, status, msg) + default: + return fmt.Errorf("%w: http %d %s", ErrInvalid, status, msg) + } +} + +// RefKindMime classifies a reference blob into the vivid "kind" (image / +// video / audio) plus a concrete mime. Magic bytes first (DetectContentType +// misses mov/m4a/ogg-audio), then sniffing, then a png default. +func RefKindMime(data []byte) (kind, mime string) { + n := len(data) + // ISOBMFF box: ftyp — mp4/mov/m4a share the container; the brand decides. + if n >= 12 && string(data[4:8]) == "ftyp" { + brand := string(data[8:12]) + switch { + case strings.HasPrefix(brand, "M4A"), strings.HasPrefix(brand, "M4B"), strings.HasPrefix(brand, "F4A"), hasBMFFHandler(data, "soun") && !hasBMFFHandler(data, "vide"): + return "audio", "audio/mp4" + case brand == "qt ": + return "video", "video/quicktime" + default: + return "video", "video/mp4" + } + } + // WebM/MKV EBML + if n >= 4 && data[0] == 0x1A && data[1] == 0x45 && data[2] == 0xDF && data[3] == 0xA3 && bytes.Contains(bytes.ToLower(data[:minInt(n, 4096)]), []byte("webm")) { + return "video", "video/webm" + } + // MP3: ID3 tag or MPEG sync + if n >= 3 && string(data[:3]) == "ID3" { + return "audio", "audio/mpeg" + } + if n >= 4 && data[0] == 0xFF && (data[1]&0xE0) == 0xE0 && + (data[1]>>3)&3 != 1 && (data[1]>>1)&3 != 0 && data[2]>>4 != 0 && data[2]>>4 != 15 && (data[2]>>2)&3 != 3 { + return "audio", "audio/mpeg" + } + // WAV / OGG + if n >= 12 && string(data[:4]) == "RIFF" && string(data[8:12]) == "WAVE" { + return "audio", "audio/wav" + } + if n >= 27 && string(data[:4]) == "OggS" && (bytes.Contains(data[:minInt(n, 4096)], []byte("OpusHead")) || + bytes.Contains(data[:minInt(n, 4096)], []byte("vorbis")) || bytes.Contains(data[:minInt(n, 4096)], []byte("Speex ")) || + bytes.Contains(data[:minInt(n, 4096)], []byte("fLaC"))) { + return "audio", "audio/ogg" + } + if m := http.DetectContentType(data); m == "image/png" || m == "image/jpeg" || m == "image/webp" { + return "image", m + } + return "", "" +} + +func hasBMFFHandler(data []byte, handler string) bool { + data = data[:minInt(len(data), 4096)] + for offset := 0; offset < len(data); { + i := bytes.Index(data[offset:], []byte("hdlr")) + if i < 0 { + return false + } + i += offset + if i+16 <= len(data) && string(data[i+12:i+16]) == handler { + return true + } + offset = i + 4 + } + return false +} + +func isCreditError(s string) bool { + m := strings.ToLower(s) + return strings.Contains(m, "credit") || strings.Contains(m, "quota") || + strings.Contains(m, "insufficient") || strings.Contains(s, "积分不足") || + strings.Contains(s, "余额") +} + +func str(v any) string { + switch t := v.(type) { + case string: + return t + case fmt.Stringer: + return t.String() + default: + return "" + } +} + +func num(v any) float64 { + switch t := v.(type) { + case float64: + return t + case int64: + return float64(t) + case json.Number: + f, _ := t.Float64() + return f + default: + return 0 + } +} + +func clip(b []byte, n int) string { + s := strings.TrimSpace(string(b)) + if len(s) <= n { + return s + } + return s[:n] + "…" +} + +func minInt(a, b int) int { + if a < b { + return a + } + return b +} + +// sanitizeErr strips the upstream URL/host from a network error so a channel's +// private upstream URL never leaks into the event log / API response. +func sanitizeErr(err error) string { + if err == nil { + return "" + } + s := err.Error() + switch { + case strings.Contains(s, "context deadline exceeded"), strings.Contains(s, "Client.Timeout"), strings.Contains(s, "timeout"): + return "request timeout" + case strings.Contains(s, "connection refused"): + return "connection refused" + case strings.Contains(s, "no such host"), strings.Contains(s, "dial tcp"), strings.Contains(s, "lookup "): + return "cannot reach upstream" + case strings.Contains(s, "tls"), strings.Contains(s, "TLS"), strings.Contains(s, "certificate"): + return "TLS error" + case strings.Contains(s, "EOF"), strings.Contains(s, "reset by peer"), strings.Contains(s, "broken pipe"): + return "connection reset" + } + var ue *url.Error + if errors.As(err, &ue) { + return strings.ToLower(ue.Op) + " upstream failed" + } + return "upstream request failed" +} + +func sleepCtx(ctx context.Context, d time.Duration) error { + t := time.NewTimer(d) + defer t.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-t.C: + return nil + } +} diff --git a/backend/internal/util/mediadur.go b/backend/internal/util/mediadur.go new file mode 100644 index 0000000..00a6b67 --- /dev/null +++ b/backend/internal/util/mediadur.go @@ -0,0 +1,138 @@ +package util + +import ( + "bytes" + "encoding/binary" + "math" +) + +// maxMediaDurationSec 时长 sanity 上限：解析出超过 24h 视为垃圾数据丢弃。 +const maxMediaDurationSec = 24 * 3600 + +// MediaDurationSeconds 从已读片段尽力解析媒体时长（秒）。 +// mp4/mov/m4a 扫描 mvhd 盒子；webm 解 EBML Duration×TimecodeScale；wav 按 +// byteRate 换算。解析不出返回 false——调用方可用文件尾片段再试一次 +// （未 faststart 的 mp4 moov 盒子常在文件尾）。mp3/ogg 无廉价可靠头长度，跳过。 +func MediaDurationSeconds(mime string, buf []byte) (float64, bool) { + if len(buf) == 0 { + return 0, false + } + var sec float64 + var ok bool + switch mime { + case "video/mp4", "video/quicktime", "audio/mp4": + sec, ok = mvhdDuration(buf) + case "video/webm": + sec, ok = webmDuration(buf) + case "audio/wav": + sec, ok = wavDuration(buf) + } + if !ok || sec <= 0 || sec > maxMediaDurationSec { + return 0, false + } + return sec, true +} + +// mvhdDuration 在缓冲里扫 "mvhd" 盒子并解析 duration/timescale。 +// 盒子布局：[size4][type4=mvhd][version1][flags3] +// v0: creation4 mod4 timescale4 duration4 → timescale@idx+16, duration@idx+20 +// v1: creation8 mod8 timescale4 duration8 → timescale@idx+24, duration@idx+28 +func mvhdDuration(buf []byte) (float64, bool) { + for off := 0; ; { + i := bytes.Index(buf[off:], []byte("mvhd")) + if i < 0 { + return 0, false + } + i += off + if ver := buf[i+4]; ver == 0 && i+24 <= len(buf) { + ts := binary.BigEndian.Uint32(buf[i+16 : i+20]) + dur := binary.BigEndian.Uint32(buf[i+20 : i+24]) + if ts > 0 && dur > 0 { + return float64(dur) / float64(ts), true + } + } else if ver == 1 && i+36 <= len(buf) { + ts := binary.BigEndian.Uint32(buf[i+24 : i+28]) + dur := binary.BigEndian.Uint64(buf[i+28 : i+36]) + if ts > 0 && dur > 0 { + return float64(dur) / float64(ts), true + } + } + off = i + 4 + } +} + +// webmDuration 扫 EBML Duration 元素（0x4489，float）和 TimecodeScale +// （0x2AD7B1，uint，缺省 1,000,000ns）。秒 = duration × scale / 1e9。 +func webmDuration(buf []byte) (float64, bool) { + scale := uint64(1_000_000) + if i := bytes.Index(buf, []byte{0x2a, 0xd7, 0xb1}); i >= 0 { + if v, n := ebmlVint(buf[i+3:]); n > 0 && i+3+n+v <= len(buf) { + scale = 0 + for _, b := range buf[i+3+n : i+3+n+v] { + scale = scale<<8 | uint64(b) + } + if scale == 0 { + scale = 1_000_000 + } + } + } + i := bytes.Index(buf, []byte{0x44, 0x89}) + if i < 0 { + return 0, false + } + v, n := ebmlVint(buf[i+2:]) + if n == 0 || i+2+n+v > len(buf) { + return 0, false + } + f := buf[i+2+n : i+2+n+v] + var dur float64 + switch v { + case 4: + dur = float64(math.Float32frombits(binary.BigEndian.Uint32(f))) + case 8: + dur = math.Float64frombits(binary.BigEndian.Uint64(f)) + default: + return 0, false + } + return dur * float64(scale) / 1e9, true +} + +// ebmlVint 读 EBML 变长整数的长度值：返回 (负载长度, vint 本身字节数)。 +func ebmlVint(buf []byte) (int, int) { + if len(buf) == 0 { + return 0, 0 + } + for n := 1; n <= 8; n++ { + mask := byte(0x80 >> (n - 1)) + if buf[0]&mask != 0 { + if len(buf) < n { + return 0, 0 + } + v := int(buf[0] &^ mask) + for i := 1; i < n; i++ { + v = v<<8 | int(buf[i]) + } + return v, n + } + } + return 0, 0 +} + +// wavDuration 解析 RIFF/WAVE：byteRate@28(LE u32) + "data" 块大小换算秒数。 +func wavDuration(buf []byte) (float64, bool) { + if len(buf) < 44 || string(buf[:4]) != "RIFF" || string(buf[8:12]) != "WAVE" { + return 0, false + } + rate := binary.LittleEndian.Uint32(buf[28:32]) + if rate == 0 { + return 0, false + } + for off := 12; off+8 <= len(buf); { + size := int(binary.LittleEndian.Uint32(buf[off+4 : off+8])) + if string(buf[off:off+4]) == "data" { + return float64(size) / float64(rate), true + } + off += 8 + size + size&1 // 块按 2 字节对齐 + } + return 0, false +} diff --git a/frontend/src/components/AccountEditModal.vue b/frontend/src/components/AccountEditModal.vue deleted file mode 100644 index 58c5d97..0000000 --- a/frontend/src/components/AccountEditModal.vue +++ /dev/null @@ -1,99 +0,0 @@ - - - - - diff --git a/frontend/src/components/AccountTestModal.vue b/frontend/src/components/AccountTestModal.vue deleted file mode 100644 index bc4b473..0000000 --- a/frontend/src/components/AccountTestModal.vue +++ /dev/null @@ -1,168 +0,0 @@ - - - - - diff --git a/frontend/src/components/CustomModelModal.vue b/frontend/src/components/CustomModelModal.vue deleted file mode 100644 index 3e6dd06..0000000 --- a/frontend/src/components/CustomModelModal.vue +++ /dev/null @@ -1,215 +0,0 @@ - - - diff --git a/frontend/src/components/ImportModal.vue b/frontend/src/components/ImportModal.vue deleted file mode 100644 index b8058d0..0000000 --- a/frontend/src/components/ImportModal.vue +++ /dev/null @@ -1,166 +0,0 @@ - - - diff --git a/frontend/src/components/ModelFormModal.vue b/frontend/src/components/ModelFormModal.vue index 09a8109..bf7e45d 100644 --- a/frontend/src/components/ModelFormModal.vue +++ b/frontend/src/components/ModelFormModal.vue @@ -61,11 +61,12 @@ const resolutions = computed(() => { // offerable palette — for custom models the palette is just the input grid. const storedResolutions = computed(() => sortResolutions(entry.value?.resolutions || [])) // 按秒计价的视频模型（一个 /s 价，实付 = 分辨率价 + 每秒价 × 秒数）。目录用 -// per_second 声明；已入库的模型则看它是否已存了 per_second 价。 +// per_second 声明；已入库的模型则看它是否已存了 per_second 价。渠道导入的 +// duration range 模型（duration_min/max，无固定档）也只能走 per_second。 const isPerSecond = computed(() => { const e = entry.value if (!e) return false - return !!e.per_second || e.duration_prices?.per_second != null + return !!e.per_second || e.duration_prices?.per_second != null || (e.type === 'video' && (e.duration_max || 0) > 0 && !(e.durations || []).length) }) const durationTiers = computed(() => { @@ -221,7 +222,7 @@ async function save() { placeholder="选择一个支持的模型" />
{{ entry?.id }}

- 所有支持的模型都已添加。 + 模型经「渠道管理 → 拉取上游模型」导入。

@@ -259,7 +260,8 @@ async function save() { 时长
{{ d }} - — + {{ entry.duration_min }}–{{ entry.duration_max }}s（区间） + —
参考图
diff --git a/frontend/src/components/TestModal.vue b/frontend/src/components/TestModal.vue index 44573ba..cd71af2 100644 --- a/frontend/src/components/TestModal.vue +++ b/frontend/src/components/TestModal.vue @@ -16,11 +16,16 @@ const isVideo = props.model.type === 'video' // (core/video_models.FAMILIES), not the user-edited managed-model record // which may have stale max_reference_images / reference_mode values. const familyPreset = ref(null) +const channels = ref([]) // 渠道 pin 候选（空 = 自动调度） +const channelId = ref('') onMounted(async () => { - if (!isVideo) return - const r = await api('/video-presets') - const list = r.data?.data || [] - familyPreset.value = list.find((p) => p.key === props.model.id) || null + if (isVideo) { + const r = await api('/video-presets') + const list = r.data?.data || [] + familyPreset.value = list.find((p) => p.key === props.model.id) || null + } + const cr = await api('/channels') + channels.value = (cr.data?.data || []).filter((c) => c.enabled) }) const ratios = (props.model.ratios && props.model.ratios.length) @@ -137,6 +142,7 @@ async function run() { // Backend accepts raw base64 only — strip the "data:...;base64," prefix. payload.reference_images = refImages.value.map((r) => r.dataUrl.replace(/^data:[^,]*,/, '')) } + if (channelId.value) payload.channel_id = channelId.value recoverSubmitTs = Date.now() const r = await api('/test', jsonBody('POST', payload)) if (r.ok && r.data?.url) { @@ -239,6 +245,15 @@ async function recover() {
+ +
+ + +
+