Revert "feat: 支持异步生图任务与结果轮询"

This commit is contained in:
Wesley Liddick
2026-07-16 09:47:27 +08:00
committed by GitHub
parent a3a227c858
commit 502097026f
21 changed files with 5 additions and 961 deletions
-1
View File
@@ -135,7 +135,6 @@ docs/*
!docs/PAYMENT.md
!docs/PAYMENT_CN.md
!docs/ADMIN_PAYMENT_INTEGRATION_API.md
!docs/ASYNC_IMAGE_TASKS.md
!docs/legal/
!docs/legal/*.md
.serena/
-6
View File
@@ -688,12 +688,6 @@ Simple Mode is designed for individual developers or internal teams who want qui
---
## Asynchronous Image Tasks
Long-running OpenAI/Grok image generation and editing can be submitted through `/v1/images/generations/async` or `/v1/images/edits/async`, then polled at `/v1/images/tasks/{task_id}` without holding a CDN connection open. See [Asynchronous Image Tasks](docs/ASYNC_IMAGE_TASKS.md) for request and response examples.
---
## Grok / xAI Support
Sub2API supports both Grok subscription accounts through xAI OAuth and standard xAI API-key accounts. Both account types forward OpenAI-compatible Responses traffic to xAI.
+2 -5
View File
@@ -70,6 +70,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
concurrencyCache := repository.ProvideConcurrencyCache(redisClient, configConfig)
schedulerCache := repository.ProvideSchedulerCache(redisClient, configConfig)
accountRepository := repository.NewAccountRepository(client, db, schedulerCache)
adminAccountRepository := repository.NewAdminAccountRepository(client, db, schedulerCache)
concurrencyService := service.ProvideConcurrencyService(concurrencyCache, accountRepository, configConfig)
apiKeyService := service.ProvideAPIKeyService(apiKeyRepository, userRepository, groupRepository, userSubscriptionRepository, userGroupRateRepository, apiKeyCache, configConfig, billingCacheService, concurrencyService)
apiKeyAuthCacheInvalidator := service.ProvideAPIKeyAuthCacheInvalidator(apiKeyService)
@@ -173,7 +174,6 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
leaderLockCache := repository.NewLeaderLockCache(redisClient)
dashboardAggregationService := service.ProvideDashboardAggregationService(dashboardAggregationRepository, timingWheelService, leaderLockCache, db, configConfig)
dashboardHandler := admin.NewDashboardHandler(dashboardService, dashboardAggregationService)
adminAccountRepository := repository.NewAdminAccountRepository(client, db, schedulerCache)
proxyExitInfoProber := repository.NewProxyExitInfoProber(configConfig)
proxyLatencyCache := repository.NewProxyLatencyCache(redisClient)
adminService := service.NewAdminService(userRepository, groupRepository, adminAccountRepository, proxyRepository, apiKeyRepository, redeemCodeRepository, userGroupRateRepository, userRPMCache, billingCacheService, proxyExitInfoProber, proxyLatencyCache, apiKeyAuthCacheInvalidator, client, settingService, subscriptionService, userSubscriptionRepository, privacyClientFactory, openAIGatewayService, affiliateService)
@@ -261,9 +261,6 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
handlerPaymentHandler := handler.NewPaymentHandler(paymentService, paymentConfigService)
paymentWebhookHandler := handler.NewPaymentWebhookHandler(paymentService, registry)
availableChannelHandler := handler.NewAvailableChannelHandler(channelService, apiKeyService, settingService)
imageTaskStore := repository.NewImageTaskStore(redisClient)
imageTaskService := service.NewImageTaskService(imageTaskStore)
asyncImageHandler := handler.NewAsyncImageHandler(imageTaskService, openAIGatewayHandler)
batchImageRepository := repository.NewBatchImageRepository(db)
batchImageQueue := repository.NewBatchImageQueue(redisClient, configConfig)
batchImageModelPricingResolver := service.ProvideBatchImageModelPricingResolver(modelPricingResolver)
@@ -274,7 +271,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
batchImageHandler := handler.NewBatchImageHandler(batchImagePublicService, batchImageDownloadService, batchImageCleanupService)
idempotencyCoordinator := service.ProvideIdempotencyCoordinator(idempotencyRepository, configConfig)
idempotencyCleanupService := service.ProvideIdempotencyCleanupService(idempotencyRepository, configConfig)
handlers := handler.ProvideHandlers(authHandler, userHandler, apiKeyHandler, usageHandler, redeemHandler, subscriptionHandler, announcementHandler, channelMonitorUserHandler, adminHandlers, gatewayHandler, openAIGatewayHandler, handlerSettingHandler, totpHandler, handlerPaymentHandler, paymentWebhookHandler, availableChannelHandler, asyncImageHandler, batchImageHandler, idempotencyCoordinator, idempotencyCleanupService)
handlers := handler.ProvideHandlers(authHandler, userHandler, apiKeyHandler, usageHandler, redeemHandler, subscriptionHandler, announcementHandler, channelMonitorUserHandler, adminHandlers, gatewayHandler, openAIGatewayHandler, handlerSettingHandler, totpHandler, handlerPaymentHandler, paymentWebhookHandler, availableChannelHandler, batchImageHandler, idempotencyCoordinator, idempotencyCleanupService)
jwtAuthMiddleware := middleware.NewJWTAuthMiddleware(authService, userService)
adminAuthMiddleware := middleware.NewAdminAuthMiddleware(authService, userService, settingService)
apiKeyAuthMiddleware := middleware.NewAPIKeyAuthMiddleware(apiKeyService, subscriptionService, configConfig)
-2
View File
@@ -166,8 +166,6 @@ github.com/google/go-querystring v1.1.0/go.mod h1:Kcdr2DB4koayq7X8pmAG4sNG59So17
github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e h1:ijClszYn+mADRFY17kjQEVQ1XRhq2/JR1M3sGqeJoxs=
github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e/go.mod h1:boTsfXsheKC2y+lKOCMpSfarhxDeIzfZG1jqGcPl3cA=
github.com/google/subcommands v1.2.0 h1:vWQspBTo2nEqTUFita5/KeEWlUL8kQObDFbub/EN9oE=
github.com/google/subcommands v1.2.0/go.mod h1:ZjhPrFU+Olkh9WazFPsl27BQ4UPiG37m3yTrtFlrHVk=
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
github.com/google/wire v0.7.0 h1:JxUKI6+CVBgCO2WToKy/nQk0sS+amI9z9EjVmdaocj4=
-3
View File
@@ -23,7 +23,6 @@ const (
EndpointResponsesCompact = "/v1/responses/compact"
EndpointImagesGenerations = "/v1/images/generations"
EndpointImagesEdits = "/v1/images/edits"
EndpointImageTasks = "/v1/images/tasks"
EndpointVideosGenerations = "/v1/videos/generations"
EndpointVideosEdits = "/v1/videos/edits"
EndpointVideosExtensions = "/v1/videos/extensions"
@@ -89,8 +88,6 @@ func NormalizeInboundEndpoint(path string) string {
return EndpointImagesGenerations
case strings.Contains(path, EndpointImagesEdits) || strings.Contains(path, "/images/edits"):
return EndpointImagesEdits
case strings.Contains(path, EndpointImageTasks) || strings.Contains(path, "/images/tasks/"):
return EndpointImageTasks
case strings.Contains(path, EndpointVideosGenerations) || strings.Contains(path, "/videos/generations"):
return EndpointVideosGenerations
case strings.Contains(path, EndpointVideosEdits) || strings.Contains(path, "/videos/edits"):
@@ -31,7 +31,6 @@ func TestNormalizeInboundEndpoint(t *testing.T) {
{"/v1/responses/compact/detail", EndpointResponsesCompact},
{"/v1/images/generations", EndpointImagesGenerations},
{"/v1/images/edits", EndpointImagesEdits},
{"/v1/images/tasks/imgtask_123", EndpointImageTasks},
{"/v1/videos/generations", EndpointVideosGenerations},
{"/v1/videos/req_123", EndpointVideos},
{"/v1beta/models", EndpointGeminiModels},
@@ -53,7 +52,6 @@ func TestNormalizeInboundEndpoint(t *testing.T) {
{"/responses/compact", EndpointResponsesCompact},
{"/responses/compact/detail", EndpointResponsesCompact},
{"/alpha/search", EndpointAlphaSearch},
{"/images/tasks/imgtask_123", EndpointImageTasks},
// Bare Codex direct alias route — root vs. compact.
{"/backend-api/codex/responses", EndpointResponses},
-1
View File
@@ -58,7 +58,6 @@ type Handlers struct {
Payment *PaymentHandler
PaymentWebhook *PaymentWebhookHandler
AvailableChannel *AvailableChannelHandler
AsyncImage *AsyncImageHandler
BatchImage *BatchImageHandler
}
@@ -1,270 +0,0 @@
package handler
import (
"bytes"
"context"
"encoding/json"
"errors"
"io"
"net/http"
"net/http/httptest"
"strings"
"time"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
pkghttputil "github.com/Wei-Shaw/sub2api/internal/pkg/httputil"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"go.uber.org/zap"
)
type AsyncImageHandler struct {
tasks *service.ImageTaskService
openAI *OpenAIGatewayHandler
execute func(platform string, c *gin.Context)
}
func NewAsyncImageHandler(tasks *service.ImageTaskService, openAI *OpenAIGatewayHandler) *AsyncImageHandler {
h := &AsyncImageHandler{tasks: tasks, openAI: openAI}
h.execute = h.executeWithGateway
return h
}
// Submit accepts the same payload as the synchronous Images endpoint and
// returns before the upstream image generation begins.
func (h *AsyncImageHandler) Submit(c *gin.Context) {
apiKey, ok := middleware2.GetAPIKeyFromContext(c)
if !ok || apiKey == nil || apiKey.UserID <= 0 || apiKey.ID <= 0 {
imageTaskError(c, service.ErrImageTaskForbidden)
return
}
platform := ""
if apiKey.Group != nil {
platform = apiKey.Group.Platform
}
if platform != service.PlatformOpenAI && platform != service.PlatformGrok {
imageTaskJSONError(c, http.StatusNotFound, "not_found_error", "Images API is not supported for this platform")
return
}
if !service.GroupAllowsImageGeneration(apiKey.Group) {
imageTaskJSONError(c, http.StatusForbidden, "permission_error", service.ImageGenerationPermissionMessage())
return
}
if h == nil || h.tasks == nil || h.execute == nil {
imageTaskError(c, service.ErrImageTaskUnavailable)
return
}
body, err := pkghttputil.ReadRequestBodyWithPrealloc(c.Request)
if err != nil {
if maxErr, ok := extractMaxBytesError(err); ok {
imageTaskJSONError(c, http.StatusRequestEntityTooLarge, "invalid_request_error", buildBodyTooLargeMessage(maxErr.Limit))
return
}
imageTaskJSONError(c, http.StatusBadRequest, "invalid_request_error", "Failed to read request body")
return
}
if len(body) == 0 {
imageTaskJSONError(c, http.StatusBadRequest, "invalid_request_error", "Request body is empty")
return
}
if asyncImageRequestStreams(c.GetHeader("Content-Type"), body) {
imageTaskJSONError(c, http.StatusBadRequest, "invalid_request_error", "streaming image requests cannot be submitted as asynchronous tasks")
return
}
if err := h.validateRequest(c, platform, body); err != nil {
imageTaskJSONError(c, http.StatusBadRequest, "invalid_request_error", err.Error())
return
}
taskCtx, recorder, cancel := newAsyncImageContext(c, body, h.tasks.ExecutionTimeout())
task, err := h.tasks.Create(c.Request.Context(), service.ImageTaskOwner{UserID: apiKey.UserID, APIKeyID: apiKey.ID})
if err != nil {
cancel()
imageTaskError(c, err)
return
}
pollURL := imageTaskPollURL(c.Request.URL.Path, task.ID)
c.Header("Cache-Control", "no-store")
c.Header("Location", pollURL)
c.Header("Retry-After", "3")
c.JSON(http.StatusAccepted, gin.H{
"id": task.ID,
"task_id": task.TaskID,
"object": task.Object,
"status": task.Status,
"created_at": task.CreatedAt,
"expires_at": task.ExpiresAt,
"poll_url": pollURL,
})
go h.run(task.ID, platform, taskCtx, recorder, cancel)
}
func (h *AsyncImageHandler) Get(c *gin.Context) {
apiKey, ok := middleware2.GetAPIKeyFromContext(c)
if !ok || apiKey == nil || apiKey.UserID <= 0 || apiKey.ID <= 0 {
imageTaskError(c, service.ErrImageTaskForbidden)
return
}
if h == nil || h.tasks == nil {
imageTaskError(c, service.ErrImageTaskUnavailable)
return
}
task, err := h.tasks.Get(c.Request.Context(), service.ImageTaskOwner{UserID: apiKey.UserID, APIKeyID: apiKey.ID}, c.Param("task_id"))
if err != nil {
imageTaskError(c, err)
return
}
c.Header("Cache-Control", "no-store")
if task.Status == service.ImageTaskStatusProcessing {
c.Header("Retry-After", "3")
}
c.JSON(http.StatusOK, task)
}
func (h *AsyncImageHandler) validateRequest(c *gin.Context, platform string, body []byte) error {
if h.openAI == nil || h.openAI.gatewayService == nil {
return nil
}
if platform == service.PlatformGrok {
parsed := service.ParseGrokMediaRequest(c.GetHeader("Content-Type"), body)
if strings.TrimSpace(parsed.Model) == "" {
return errors.New("model is required")
}
return nil
}
parsed, err := h.openAI.gatewayService.ParseOpenAIImagesRequest(c, body)
if err != nil {
return err
}
if parsed.Stream {
return errors.New("streaming image requests cannot be submitted as asynchronous tasks")
}
return nil
}
func (h *AsyncImageHandler) executeWithGateway(platform string, c *gin.Context) {
if h.openAI == nil {
imageTaskJSONError(c, http.StatusServiceUnavailable, "api_error", "image gateway is unavailable")
return
}
if platform == service.PlatformGrok {
h.openAI.GrokImages(c)
return
}
h.openAI.Images(c)
}
func (h *AsyncImageHandler) run(taskID, platform string, taskCtx *gin.Context, recorder *httptest.ResponseRecorder, cancel context.CancelFunc) {
defer cancel()
defer func() {
if recovered := recover(); recovered != nil {
logger.L().Error("image_task.execution_panicked", zap.String("task_id", taskID), zap.Any("panic", recovered))
h.failTask(taskID, http.StatusInternalServerError, imageTaskErrorPayload("api_error", "image generation task panicked"))
}
}()
h.execute(platform, taskCtx)
body := bytes.TrimSpace(recorder.Body.Bytes())
if err := taskCtx.Request.Context().Err(); err != nil && len(body) == 0 {
h.failTask(taskID, http.StatusGatewayTimeout, imageTaskErrorPayload("timeout_error", "image generation task timed out"))
return
}
statusCode := recorder.Code
if statusCode == 0 {
statusCode = http.StatusOK
}
if statusCode >= http.StatusOK && statusCode < http.StatusMultipleChoices {
if len(body) == 0 || !json.Valid(body) {
h.failTask(taskID, http.StatusBadGateway, imageTaskErrorPayload("api_error", "upstream returned an invalid image response"))
return
}
if err := h.tasks.Complete(context.Background(), taskID, statusCode, json.RawMessage(body)); err != nil {
logger.L().Error("image_task.complete_store_failed", zap.String("task_id", taskID), zap.Error(err))
}
return
}
h.failTask(taskID, statusCode, extractImageTaskError(body))
}
func (h *AsyncImageHandler) failTask(taskID string, statusCode int, taskErr json.RawMessage) {
if err := h.tasks.Fail(context.Background(), taskID, statusCode, taskErr); err != nil {
logger.L().Error("image_task.failure_store_failed", zap.String("task_id", taskID), zap.Error(err))
}
}
func newAsyncImageContext(c *gin.Context, body []byte, timeoutDuration time.Duration) (*gin.Context, *httptest.ResponseRecorder, context.CancelFunc) {
base := context.WithoutCancel(c.Request.Context())
executionCtx, cancel := context.WithTimeout(base, timeoutDuration)
request := c.Request.Clone(executionCtx)
request.Body = io.NopCloser(bytes.NewReader(body))
request.GetBody = func() (io.ReadCloser, error) {
return io.NopCloser(bytes.NewReader(body)), nil
}
request.ContentLength = int64(len(body))
request.URL.Path = strings.TrimSuffix(request.URL.Path, "/async")
taskCtx := c.Copy()
recorder := httptest.NewRecorder()
recorderCtx, _ := gin.CreateTestContext(recorder)
taskCtx.Writer = recorderCtx.Writer
taskCtx.Request = request
return taskCtx, recorder, cancel
}
func asyncImageRequestStreams(contentType string, body []byte) bool {
if isMultipartImagesContentType(contentType) {
return false
}
var envelope struct {
Stream bool `json:"stream"`
}
return json.Unmarshal(body, &envelope) == nil && envelope.Stream
}
func imageTaskPollURL(submitPath, taskID string) string {
if strings.HasPrefix(submitPath, "/v1/") {
return "/v1/images/tasks/" + taskID
}
return "/images/tasks/" + taskID
}
func extractImageTaskError(body []byte) json.RawMessage {
if json.Valid(body) {
var envelope struct {
Error json.RawMessage `json:"error"`
}
if json.Unmarshal(body, &envelope) == nil && len(envelope.Error) > 0 && json.Valid(envelope.Error) {
return envelope.Error
}
return json.RawMessage(body)
}
return imageTaskErrorPayload("api_error", "image generation failed")
}
func imageTaskErrorPayload(errorType, message string) json.RawMessage {
data, _ := json.Marshal(gin.H{"type": errorType, "message": message})
return data
}
func imageTaskError(c *gin.Context, err error) {
status := infraerrors.Code(err)
code := infraerrors.Reason(err)
message := infraerrors.Message(err)
if status <= 0 {
status = http.StatusInternalServerError
}
if strings.TrimSpace(code) == "" {
code = "IMAGE_TASK_ERROR"
}
imageTaskJSONError(c, status, code, message)
}
func imageTaskJSONError(c *gin.Context, status int, code, message string) {
c.Header("Cache-Control", "no-store")
c.JSON(status, gin.H{"error": gin.H{"type": code, "code": code, "message": message}})
}
@@ -1,107 +0,0 @@
package handler
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"strings"
"sync"
"testing"
"time"
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
type asyncImageMemoryStore struct {
mu sync.RWMutex
tasks map[string]*service.ImageTaskRecord
}
func (s *asyncImageMemoryStore) Save(_ context.Context, task *service.ImageTaskRecord, _ time.Duration) error {
s.mu.Lock()
defer s.mu.Unlock()
copy := *task
copy.Result = append(json.RawMessage(nil), task.Result...)
copy.Error = append(json.RawMessage(nil), task.Error...)
s.tasks[task.ID] = &copy
return nil
}
func (s *asyncImageMemoryStore) Get(_ context.Context, id string) (*service.ImageTaskRecord, error) {
s.mu.RLock()
defer s.mu.RUnlock()
task := s.tasks[id]
if task == nil {
return nil, service.ErrImageTaskNotFound
}
copy := *task
copy.Result = append(json.RawMessage(nil), task.Result...)
copy.Error = append(json.RawMessage(nil), task.Error...)
return &copy, nil
}
func TestAsyncImageHandlerSubmitAndPoll(t *testing.T) {
gin.SetMode(gin.TestMode)
store := &asyncImageMemoryStore{tasks: make(map[string]*service.ImageTaskRecord)}
tasks := service.NewImageTaskServiceWithOptions(store, time.Hour, time.Minute)
release := make(chan struct{})
h := &AsyncImageHandler{tasks: tasks}
h.execute = func(_ string, c *gin.Context) {
<-release
c.JSON(http.StatusOK, gin.H{"created": 123, "data": []gin.H{{"url": "https://example.test/image.png"}}})
}
router := gin.New()
router.Use(func(c *gin.Context) {
groupID := int64(3)
c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{
ID: 9,
UserID: 7,
GroupID: &groupID,
Group: &service.Group{ID: groupID, Platform: service.PlatformOpenAI, AllowImageGeneration: true},
})
c.Next()
})
router.POST("/v1/images/generations/async", h.Submit)
router.GET("/v1/images/tasks/:task_id", h.Get)
requestCtx, cancelRequest := context.WithCancel(context.Background())
req := httptest.NewRequest(http.MethodPost, "/v1/images/generations/async", strings.NewReader(`{"model":"gpt-image-1","prompt":"cat"}`)).WithContext(requestCtx)
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.Equal(t, http.StatusAccepted, w.Code)
require.Equal(t, "no-store", w.Header().Get("Cache-Control"))
require.Equal(t, "3", w.Header().Get("Retry-After"))
var accepted struct {
TaskID string `json:"task_id"`
Status string `json:"status"`
PollURL string `json:"poll_url"`
}
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &accepted))
require.Equal(t, service.ImageTaskStatusProcessing, accepted.Status)
require.Equal(t, "/v1/images/tasks/"+accepted.TaskID, accepted.PollURL)
require.Equal(t, accepted.PollURL, w.Header().Get("Location"))
// The detached background request must survive completion of/cancellation
// from the short submission request.
cancelRequest()
close(release)
require.Eventually(t, func() bool {
got, err := tasks.Get(context.Background(), service.ImageTaskOwner{UserID: 7, APIKeyID: 9}, accepted.TaskID)
return err == nil && got.Status == service.ImageTaskStatusCompleted
}, time.Second, 10*time.Millisecond)
pollReq := httptest.NewRequest(http.MethodGet, accepted.PollURL, nil)
pollWriter := httptest.NewRecorder()
router.ServeHTTP(pollWriter, pollReq)
require.Equal(t, http.StatusOK, pollWriter.Code)
require.Equal(t, "no-store", pollWriter.Header().Get("Cache-Control"))
require.Empty(t, pollWriter.Header().Get("Retry-After"))
require.Contains(t, pollWriter.Body.String(), "https://example.test/image.png")
}
-3
View File
@@ -115,7 +115,6 @@ func ProvideHandlers(
paymentHandler *PaymentHandler,
paymentWebhookHandler *PaymentWebhookHandler,
availableChannelHandler *AvailableChannelHandler,
asyncImageHandler *AsyncImageHandler,
batchImageHandler *BatchImageHandler,
_ *service.IdempotencyCoordinator,
_ *service.IdempotencyCleanupService,
@@ -137,7 +136,6 @@ func ProvideHandlers(
Payment: paymentHandler,
PaymentWebhook: paymentWebhookHandler,
AvailableChannel: availableChannelHandler,
AsyncImage: asyncImageHandler,
BatchImage: batchImageHandler,
}
}
@@ -160,7 +158,6 @@ var ProviderSet = wire.NewSet(
NewPaymentHandler,
NewPaymentWebhookHandler,
NewAvailableChannelHandler,
NewAsyncImageHandler,
NewBatchImageHandler,
// Admin handlers
@@ -1,48 +0,0 @@
package repository
import (
"context"
"encoding/json"
"strings"
"time"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/redis/go-redis/v9"
)
const imageTaskKeyPrefix = "image_task:"
type imageTaskStore struct {
rdb *redis.Client
}
func NewImageTaskStore(rdb *redis.Client) service.ImageTaskStore {
return &imageTaskStore{rdb: rdb}
}
func (s *imageTaskStore) Save(ctx context.Context, task *service.ImageTaskRecord, ttl time.Duration) error {
data, err := json.Marshal(task)
if err != nil {
return err
}
return s.rdb.Set(ctx, imageTaskKey(task.ID), data, ttl).Err()
}
func (s *imageTaskStore) Get(ctx context.Context, id string) (*service.ImageTaskRecord, error) {
data, err := s.rdb.Get(ctx, imageTaskKey(id)).Bytes()
if err != nil {
if err == redis.Nil {
return nil, service.ErrImageTaskNotFound
}
return nil, err
}
var task service.ImageTaskRecord
if err := json.Unmarshal(data, &task); err != nil {
return nil, err
}
return &task, nil
}
func imageTaskKey(id string) string {
return imageTaskKeyPrefix + strings.TrimSpace(id)
}
@@ -1,43 +0,0 @@
package repository
import (
"context"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/alicebob/miniredis/v2"
"github.com/redis/go-redis/v9"
"github.com/stretchr/testify/require"
)
func TestImageTaskStoreRoundTripAndTTL(t *testing.T) {
mr := miniredis.RunT(t)
rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
t.Cleanup(func() { _ = rdb.Close() })
store := NewImageTaskStore(rdb)
task := &service.ImageTaskRecord{
ID: "imgtask_123",
UserID: 7,
APIKeyID: 9,
Status: service.ImageTaskStatusProcessing,
CreatedAt: 100,
ExpiresAt: 200,
}
require.NoError(t, store.Save(context.Background(), task, 24*time.Hour))
got, err := store.Get(context.Background(), task.ID)
require.NoError(t, err)
require.Equal(t, task, got)
require.Equal(t, 24*time.Hour, mr.TTL(imageTaskKey(task.ID)))
}
func TestImageTaskStoreMissing(t *testing.T) {
mr := miniredis.RunT(t)
rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
t.Cleanup(func() { _ = rdb.Close() })
store := NewImageTaskStore(rdb)
_, err := store.Get(context.Background(), "imgtask_missing")
require.ErrorIs(t, err, service.ErrImageTaskNotFound)
}
-1
View File
@@ -117,7 +117,6 @@ var ProviderSet = wire.NewSet(
NewRedeemCache,
NewUpdateCache,
NewGeminiTokenCache,
NewImageTaskStore,
NewBatchImageQueue,
NewBatchImageDownloadLimiter,
NewLeaderLockCache,
@@ -4,7 +4,6 @@ import (
"context"
"errors"
"fmt"
"net/http"
"strings"
"github.com/Wei-Shaw/sub2api/internal/config"
@@ -26,7 +25,7 @@ func NewAPIKeyAuthMiddleware(apiKeyService *service.APIKeyService, subscriptionS
// - 鉴权(Authentication):验证 Key 有效性、用户状态、IP 限制 —— 始终执行
// - 计费执行(Billing Enforcement):过期/配额/订阅/余额检查 —— skipBilling 时整块跳过
//
// /v1/usage 与异步生图任务查询只需鉴权,不需要计费执行(允许已耗尽额度的 Key 查询自身结果)。
// /v1/usage 端点只需鉴权,不需要计费执行(允许过期/配额耗尽的 Key 查询自身用量)。
func apiKeyAuthWithSubscription(apiKeyService *service.APIKeyService, subscriptionService *service.SubscriptionService, cfg *config.Config) gin.HandlerFunc {
return func(c *gin.Context) {
// ── 1. 提取 API Key ──────────────────────────────────────────
@@ -147,11 +146,8 @@ func apiKeyAuthWithSubscription(apiKeyService *service.APIKeyService, subscripti
// ── 5. 加载订阅(订阅模式时始终加载) ───────────────────────
// skipBilling: usage and async image task polling only read data that
// already belongs to the authenticated key. In particular, polling must
// remain available after the completed generation consumes the key's
// remaining balance.
skipBilling := c.Request.URL.Path == "/v1/usage" || isAsyncImageTaskRead(c.Request.Method, c.Request.URL.Path)
// skipBilling: /v1/usage 只需鉴权,跳过所有计费执行
skipBilling := c.Request.URL.Path == "/v1/usage"
var subscription *service.UserSubscription
isSubscriptionType := apiKey.Group != nil && apiKey.Group.IsSubscriptionType()
@@ -247,13 +243,6 @@ func apiKeyAuthWithSubscription(apiKeyService *service.APIKeyService, subscripti
}
}
func isAsyncImageTaskRead(method, path string) bool {
if method != http.MethodGet {
return false
}
return strings.HasPrefix(path, "/v1/images/tasks/") || strings.HasPrefix(path, "/images/tasks/")
}
// GetAPIKeyFromContext 从上下文中获取API key
func GetAPIKeyFromContext(c *gin.Context) (*service.APIKey, bool) {
value, exists := c.Get(string(ContextKeyAPIKey))
@@ -1,15 +0,0 @@
package middleware
import (
"net/http"
"testing"
"github.com/stretchr/testify/require"
)
func TestIsAsyncImageTaskRead(t *testing.T) {
require.True(t, isAsyncImageTaskRead(http.MethodGet, "/v1/images/tasks/imgtask_123"))
require.True(t, isAsyncImageTaskRead(http.MethodGet, "/images/tasks/imgtask_123"))
require.False(t, isAsyncImageTaskRead(http.MethodPost, "/v1/images/tasks/imgtask_123"))
require.False(t, isAsyncImageTaskRead(http.MethodGet, "/v1/images/generations"))
}
@@ -191,9 +191,6 @@ func RegisterGatewayRoutes(
})
gateway.POST("/images/generations", imagesHandler)
gateway.POST("/images/edits", imagesHandler)
gateway.POST("/images/generations/async", h.AsyncImage.Submit)
gateway.POST("/images/edits/async", h.AsyncImage.Submit)
gateway.GET("/images/tasks/:task_id", h.AsyncImage.Get)
gateway.POST("/images/batches", h.BatchImage.Submit)
gateway.GET("/images/batches", h.BatchImage.List)
gateway.GET("/images/batches/models", h.BatchImage.Models)
@@ -274,9 +271,6 @@ func RegisterGatewayRoutes(
})
r.POST("/images/generations", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, imagesHandler)
r.POST("/images/edits", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, imagesHandler)
r.POST("/images/generations/async", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, h.AsyncImage.Submit)
r.POST("/images/edits/async", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, h.AsyncImage.Submit)
r.GET("/images/tasks/:task_id", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, h.AsyncImage.Get)
r.POST("/videos/generations", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, videoGenerationHandler)
r.POST("/videos/edits", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, videoEditHandler)
r.POST("/videos/extensions", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, videoExtensionHandler)
@@ -28,7 +28,6 @@ func newGatewayRoutesTestRouter(platform ...string) *gin.Engine {
&handler.Handlers{
Gateway: &handler.GatewayHandler{},
OpenAIGateway: &handler.OpenAIGatewayHandler{},
AsyncImage: handler.NewAsyncImageHandler(nil, nil),
},
servermiddleware.APIKeyAuthMiddleware(func(c *gin.Context) {
groupID := int64(1)
@@ -114,25 +113,6 @@ func TestGatewayRoutesOpenAIImagesPathsAreRegistered(t *testing.T) {
}
}
func TestGatewayRoutesAsyncImagesPathsAreRegistered(t *testing.T) {
router := newGatewayRoutesTestRouter()
registered := make(map[string]bool)
for _, route := range router.Routes() {
registered[route.Method+" "+route.Path] = true
}
for _, route := range []string{
"POST /v1/images/generations/async",
"POST /v1/images/edits/async",
"GET /v1/images/tasks/:task_id",
"POST /images/generations/async",
"POST /images/edits/async",
"GET /images/tasks/:task_id",
} {
require.True(t, registered[route], "%s should be registered", route)
}
}
func TestGatewayRoutesGrokImagesAndVideosPathsAreRegistered(t *testing.T) {
router := newGatewayRoutesTestRouter(service.PlatformGrok)
-210
View File
@@ -1,210 +0,0 @@
package service
import (
"context"
"encoding/json"
"errors"
"net/http"
"strings"
"time"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/google/uuid"
)
const (
ImageTaskStatusProcessing = "processing"
ImageTaskStatusCompleted = "completed"
ImageTaskStatusFailed = "failed"
defaultImageTaskTTL = 24 * time.Hour
defaultImageTaskExecutionTimeout = 30 * time.Minute
)
var (
ErrImageTaskNotFound = infraerrors.New(http.StatusNotFound, "IMAGE_TASK_NOT_FOUND", "image task not found")
ErrImageTaskForbidden = infraerrors.New(http.StatusForbidden, "IMAGE_TASK_FORBIDDEN", "image task does not belong to this API key")
ErrImageTaskUnavailable = infraerrors.New(http.StatusServiceUnavailable, "IMAGE_TASK_UNAVAILABLE", "image task storage is unavailable")
)
// ImageTaskRecord is the private Redis representation of an asynchronous image
// request. Ownership fields are intentionally omitted from the public view.
type ImageTaskRecord struct {
ID string `json:"id"`
UserID int64 `json:"user_id"`
APIKeyID int64 `json:"api_key_id"`
Status string `json:"status"`
HTTPStatus int `json:"http_status,omitempty"`
Result json.RawMessage `json:"result,omitempty"`
Error json.RawMessage `json:"error,omitempty"`
CreatedAt int64 `json:"created_at"`
CompletedAt *int64 `json:"completed_at,omitempty"`
ExpiresAt int64 `json:"expires_at"`
}
// ImageTask is the API-safe task representation returned to callers.
type ImageTask struct {
ID string `json:"id"`
TaskID string `json:"task_id"`
Object string `json:"object"`
Status string `json:"status"`
HTTPStatus int `json:"http_status,omitempty"`
ImageURL string `json:"image_url,omitempty"`
Result json.RawMessage `json:"result,omitempty"`
Error json.RawMessage `json:"error,omitempty"`
CreatedAt int64 `json:"created_at"`
CompletedAt *int64 `json:"completed_at,omitempty"`
ExpiresAt int64 `json:"expires_at"`
}
type ImageTaskOwner struct {
UserID int64
APIKeyID int64
}
type ImageTaskStore interface {
Save(ctx context.Context, task *ImageTaskRecord, ttl time.Duration) error
Get(ctx context.Context, id string) (*ImageTaskRecord, error)
}
type ImageTaskService struct {
store ImageTaskStore
ttl time.Duration
executionTimeout time.Duration
}
func NewImageTaskService(store ImageTaskStore) *ImageTaskService {
return NewImageTaskServiceWithOptions(store, defaultImageTaskTTL, defaultImageTaskExecutionTimeout)
}
func NewImageTaskServiceWithOptions(store ImageTaskStore, ttl, executionTimeout time.Duration) *ImageTaskService {
if ttl <= 0 {
ttl = defaultImageTaskTTL
}
if executionTimeout <= 0 {
executionTimeout = defaultImageTaskExecutionTimeout
}
return &ImageTaskService{store: store, ttl: ttl, executionTimeout: executionTimeout}
}
func (s *ImageTaskService) ExecutionTimeout() time.Duration {
if s == nil || s.executionTimeout <= 0 {
return defaultImageTaskExecutionTimeout
}
return s.executionTimeout
}
func (s *ImageTaskService) Create(ctx context.Context, owner ImageTaskOwner) (*ImageTask, error) {
if s == nil || s.store == nil {
return nil, ErrImageTaskUnavailable
}
now := time.Now().UTC()
task := &ImageTaskRecord{
ID: "imgtask_" + strings.ReplaceAll(uuid.NewString(), "-", ""),
UserID: owner.UserID,
APIKeyID: owner.APIKeyID,
Status: ImageTaskStatusProcessing,
CreatedAt: now.Unix(),
ExpiresAt: now.Add(s.ttl).Unix(),
}
if err := s.store.Save(ctx, task, s.ttl); err != nil {
return nil, ErrImageTaskUnavailable.WithCause(err)
}
return imageTaskToPublic(task), nil
}
func (s *ImageTaskService) Get(ctx context.Context, owner ImageTaskOwner, id string) (*ImageTask, error) {
if s == nil || s.store == nil {
return nil, ErrImageTaskUnavailable
}
task, err := s.store.Get(ctx, strings.TrimSpace(id))
if err != nil {
if errors.Is(err, ErrImageTaskNotFound) {
return nil, ErrImageTaskNotFound
}
return nil, ErrImageTaskUnavailable.WithCause(err)
}
if task.UserID != owner.UserID || task.APIKeyID != owner.APIKeyID {
// Do not reveal whether a random task ID exists for another caller.
return nil, ErrImageTaskNotFound
}
return imageTaskToPublic(task), nil
}
func (s *ImageTaskService) Complete(ctx context.Context, id string, statusCode int, result json.RawMessage) error {
if !json.Valid(result) {
return s.Fail(ctx, id, http.StatusBadGateway, imageTaskErrorJSON("api_error", "upstream returned a non-JSON image response"))
}
return s.finish(ctx, id, ImageTaskStatusCompleted, statusCode, result, nil)
}
func (s *ImageTaskService) Fail(ctx context.Context, id string, statusCode int, taskErr json.RawMessage) error {
if !json.Valid(taskErr) {
taskErr = imageTaskErrorJSON("api_error", "image generation failed")
}
return s.finish(ctx, id, ImageTaskStatusFailed, statusCode, nil, taskErr)
}
func (s *ImageTaskService) finish(ctx context.Context, id, status string, statusCode int, result, taskErr json.RawMessage) error {
if s == nil || s.store == nil {
return ErrImageTaskUnavailable
}
task, err := s.store.Get(ctx, id)
if err != nil {
if errors.Is(err, ErrImageTaskNotFound) {
return ErrImageTaskNotFound
}
return ErrImageTaskUnavailable.WithCause(err)
}
now := time.Now().UTC()
completedAt := now.Unix()
task.Status = status
task.HTTPStatus = statusCode
task.Result = result
task.Error = taskErr
task.CompletedAt = &completedAt
task.ExpiresAt = now.Add(s.ttl).Unix()
if err := s.store.Save(ctx, task, s.ttl); err != nil {
return ErrImageTaskUnavailable.WithCause(err)
}
return nil
}
func imageTaskToPublic(task *ImageTaskRecord) *ImageTask {
if task == nil {
return nil
}
return &ImageTask{
ID: task.ID,
TaskID: task.ID,
Object: "image.generation.task",
Status: task.Status,
HTTPStatus: task.HTTPStatus,
ImageURL: firstImageTaskURL(task.Result),
Result: task.Result,
Error: task.Error,
CreatedAt: task.CreatedAt,
CompletedAt: task.CompletedAt,
ExpiresAt: task.ExpiresAt,
}
}
func firstImageTaskURL(result json.RawMessage) string {
if len(result) == 0 || !json.Valid(result) {
return ""
}
var response struct {
Data []struct {
URL string `json:"url"`
} `json:"data"`
}
if json.Unmarshal(result, &response) != nil || len(response.Data) == 0 {
return ""
}
return strings.TrimSpace(response.Data[0].URL)
}
func imageTaskErrorJSON(errorType, message string) json.RawMessage {
data, _ := json.Marshal(map[string]string{"type": errorType, "message": message})
return data
}
@@ -1,91 +0,0 @@
package service
import (
"context"
"encoding/json"
"errors"
"net/http"
"testing"
"time"
"github.com/stretchr/testify/require"
)
type imageTaskMemoryStore struct {
task *ImageTaskRecord
ttl time.Duration
saveErr error
getErr error
}
func (s *imageTaskMemoryStore) Save(_ context.Context, task *ImageTaskRecord, ttl time.Duration) error {
if s.saveErr != nil {
return s.saveErr
}
copy := *task
s.task = &copy
s.ttl = ttl
return nil
}
func (s *imageTaskMemoryStore) Get(_ context.Context, _ string) (*ImageTaskRecord, error) {
if s.getErr != nil {
return nil, s.getErr
}
if s.task == nil {
return nil, ErrImageTaskNotFound
}
copy := *s.task
return &copy, nil
}
func TestImageTaskServiceLifecycleAndOwnership(t *testing.T) {
store := &imageTaskMemoryStore{}
svc := NewImageTaskServiceWithOptions(store, time.Hour, 10*time.Minute)
owner := ImageTaskOwner{UserID: 7, APIKeyID: 9}
created, err := svc.Create(context.Background(), owner)
require.NoError(t, err)
require.Equal(t, ImageTaskStatusProcessing, created.Status)
require.Equal(t, created.ID, created.TaskID)
require.Equal(t, "image.generation.task", created.Object)
require.Equal(t, time.Hour, store.ttl)
require.Equal(t, owner.UserID, store.task.UserID)
require.Equal(t, owner.APIKeyID, store.task.APIKeyID)
_, err = svc.Get(context.Background(), ImageTaskOwner{UserID: 7, APIKeyID: 10}, created.ID)
require.ErrorIs(t, err, ErrImageTaskNotFound)
result := json.RawMessage(`{"created":123,"data":[{"url":"https://example.test/image.png"}]}`)
require.NoError(t, svc.Complete(context.Background(), created.ID, http.StatusOK, result))
completed, err := svc.Get(context.Background(), owner, created.ID)
require.NoError(t, err)
require.Equal(t, ImageTaskStatusCompleted, completed.Status)
require.Equal(t, http.StatusOK, completed.HTTPStatus)
require.Equal(t, "https://example.test/image.png", completed.ImageURL)
require.JSONEq(t, string(result), string(completed.Result))
require.NotNil(t, completed.CompletedAt)
}
func TestImageTaskServiceInvalidResultBecomesFailed(t *testing.T) {
store := &imageTaskMemoryStore{}
svc := NewImageTaskServiceWithOptions(store, time.Hour, time.Minute)
created, err := svc.Create(context.Background(), ImageTaskOwner{UserID: 1, APIKeyID: 2})
require.NoError(t, err)
require.NoError(t, svc.Complete(context.Background(), created.ID, http.StatusOK, json.RawMessage(`not-json`)))
got, err := svc.Get(context.Background(), ImageTaskOwner{UserID: 1, APIKeyID: 2}, created.ID)
require.NoError(t, err)
require.Equal(t, ImageTaskStatusFailed, got.Status)
require.Equal(t, http.StatusBadGateway, got.HTTPStatus)
require.Contains(t, string(got.Error), "non-JSON")
}
func TestImageTaskServiceMapsStoreFailures(t *testing.T) {
store := &imageTaskMemoryStore{saveErr: errors.New("redis down")}
svc := NewImageTaskService(store)
_, err := svc.Create(context.Background(), ImageTaskOwner{UserID: 1, APIKeyID: 2})
require.ErrorIs(t, err, ErrImageTaskUnavailable)
}
-2
View File
@@ -637,7 +637,6 @@ var ProviderSet = wire.NewSet(
NewAdminService,
NewGatewayService,
NewOpenAIGatewayService,
NewImageTaskService,
ProvideBatchImageModelPricingResolver,
NewBatchImagePublicService,
NewBatchImageDownloadService,
@@ -647,7 +646,6 @@ var ProviderSet = wire.NewSet(
NewOAuthService,
ProvideOpenAIOAuthService,
NewGrokOAuthService,
wire.Bind(new(GrokOAuthTokenService), new(*GrokOAuthService)),
NewGeminiOAuthService,
NewGeminiQuotaService,
NewCompositeTokenCacheInvalidator,
-111
View File
@@ -1,111 +0,0 @@
# Asynchronous Image Tasks
Asynchronous image tasks let clients submit long-running OpenAI-compatible image requests without keeping one HTTP connection open. This avoids proxy/CDN response timeouts such as Cloudflare 524 while preserving the existing image routing, billing, moderation, concurrency, and failover behavior.
## Endpoints
The authenticated gateway exposes both `/v1` paths and their existing no-prefix aliases:
```text
POST /v1/images/generations/async
POST /v1/images/edits/async
GET /v1/images/tasks/{task_id}
```
The aliases are `/images/generations/async`, `/images/edits/async`, and `/images/tasks/{task_id}`.
Only OpenAI and Grok groups are supported. Requests use the same JSON or multipart payload as the corresponding synchronous endpoint. Streaming image requests are rejected because a polled task returns one final JSON result.
## Submit a task
```bash
curl -i https://api.example.com/v1/images/generations/async \
-H 'Authorization: Bearer sk-...' \
-H 'Content-Type: application/json' \
-d '{
"model": "gpt-image-1",
"prompt": "A lighthouse during a winter storm",
"size": "1536x1024"
}'
```
The server stores the initial task in Redis and responds with `202 Accepted`:
```json
{
"id": "imgtask_0123456789abcdef",
"task_id": "imgtask_0123456789abcdef",
"object": "image.generation.task",
"status": "processing",
"created_at": 1784092800,
"expires_at": 1784179200,
"poll_url": "/v1/images/tasks/imgtask_0123456789abcdef"
}
```
`Location` contains the polling path and `Retry-After: 3` provides the recommended polling interval.
## Poll a task
Use the same API key that submitted the task:
```bash
curl https://api.example.com/v1/images/tasks/imgtask_0123456789abcdef \
-H 'Authorization: Bearer sk-...'
```
While work is in progress:
```json
{
"id": "imgtask_0123456789abcdef",
"task_id": "imgtask_0123456789abcdef",
"object": "image.generation.task",
"status": "processing",
"created_at": 1784092800,
"expires_at": 1784179200
}
```
On success, `result` is the unmodified JSON body from the synchronous image API, so URL and base64 response formats both remain supported:
```json
{
"id": "imgtask_0123456789abcdef",
"task_id": "imgtask_0123456789abcdef",
"object": "image.generation.task",
"status": "completed",
"http_status": 200,
"image_url": "https://...",
"result": {
"created": 1784092923,
"data": [{"url": "https://..."}]
},
"created_at": 1784092800,
"completed_at": 1784092923,
"expires_at": 1784179323
}
```
For URL responses, `image_url` mirrors the first `data[].url` for simple clients. On failure, the task reaches `failed` and exposes the original OpenAI-compatible error object where available:
```json
{
"id": "imgtask_0123456789abcdef",
"task_id": "imgtask_0123456789abcdef",
"object": "image.generation.task",
"status": "failed",
"http_status": 502,
"error": {
"type": "api_error",
"message": "Upstream request failed"
},
"created_at": 1784092800,
"completed_at": 1784092923,
"expires_at": 1784179323
}
```
All submit and poll responses include `Cache-Control: no-store`, preventing a CDN from caching the `processing` state. Tasks and results expire 24 hours after their latest state update. A task executes for at most 30 minutes.
Task ownership is scoped to both user and API key. Unknown task IDs and IDs owned by another key both return `404`, avoiding task-existence disclosure. Polling remains available when the completed generation used the key's remaining balance; normal authentication, disabled-key, user, IP, and group checks still apply.