Merge pull request #4593 from StarryKira/fix/image-storage-env-unreachable

fix: 异步生图开关配了却不生效（环境变量被静默丢弃 + 迁移到后台开关）
2026-07-20 09:20:15 +08:00
parent f3312cf7d9 6102d64274
commit bfabfe60c8
18 changed files with 1324 additions and 35 deletions
@@ -10,12 +10,14 @@ import (
 type BackupHandler struct {
 	backupService *service.BackupService
 	userService   *service.UserService
+	imageStorage  *service.ImageStorageSettingService
 }

-func NewBackupHandler(backupService *service.BackupService, userService *service.UserService) *BackupHandler {
+func NewBackupHandler(backupService *service.BackupService, userService *service.UserService, imageStorage *service.ImageStorageSettingService) *BackupHandler {
 	return &BackupHandler{
 		backupService: backupService,
 		userService:   userService,
+		imageStorage:  imageStorage,
 	}
 }

@@ -203,3 +205,48 @@ func (h *BackupHandler) RestoreBackup(c *gin.Context) {
 	}
 	response.Accepted(c, record)
 }
+
+// ─── 异步生图对象存储配置 ───
+//
+// 与备份共用一套 S3 客户端构造，因此放在同一个页面下：勾选"复用备份 S3"即可直接
+// 借用备份已配置的端点与密钥，只用不同的前缀区分对象（备份走 backups/，图片走 images/）。
+
+func (h *BackupHandler) GetImageStorageConfig(c *gin.Context) {
+	ctx := c.Request.Context()
+	cfg, err := h.imageStorage.Get(ctx)
+	if err != nil {
+		response.ErrorFrom(c, err)
+		return
+	}
+	response.Success(c, gin.H{
+		"config":            cfg,
+		"secret_configured": h.imageStorage.SecretConfigured(ctx),
+	})
+}
+
+func (h *BackupHandler) UpdateImageStorageConfig(c *gin.Context) {
+	var req service.ImageStorageSettings
+	if err := c.ShouldBindJSON(&req); err != nil {
+		response.BadRequest(c, "Invalid request: "+err.Error())
+		return
+	}
+	cfg, err := h.imageStorage.Update(c.Request.Context(), req)
+	if err != nil {
+		response.ErrorFrom(c, err)
+		return
+	}
+	response.Success(c, cfg)
+}
+
+func (h *BackupHandler) TestImageStorageConnection(c *gin.Context) {
+	var req service.ImageStorageSettings
+	if err := c.ShouldBindJSON(&req); err != nil {
+		response.BadRequest(c, "Invalid request: "+err.Error())
+		return
+	}
+	if err := h.imageStorage.TestConnection(c.Request.Context(), req); err != nil {
+		response.Success(c, gin.H{"ok": false, "message": err.Error()})
+		return
+	}
+	response.Success(c, gin.H{"ok": true, "message": "connection successful"})
+}
@@ -0,0 +1,134 @@
+//go:build unit
+
+package handler
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/Wei-Shaw/sub2api/internal/config"
+	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 toggleSettingRepo struct {
+	mu     sync.Mutex
+	values map[string]string
+}
+
+func (r *toggleSettingRepo) Get(context.Context, string) (*service.Setting, error) { return nil, nil }
+func (r *toggleSettingRepo) GetValue(_ context.Context, key string) (string, error) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	return r.values[key], nil
+}
+
+func (r *toggleSettingRepo) Set(_ context.Context, key, value string) error {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.values[key] = value
+	return nil
+}
+
+func (r *toggleSettingRepo) GetMultiple(context.Context, []string) (map[string]string, error) {
+	return map[string]string{}, nil
+}
+func (r *toggleSettingRepo) SetMultiple(context.Context, map[string]string) error { return nil }
+func (r *toggleSettingRepo) GetAll(context.Context) (map[string]string, error) {
+	return map[string]string{}, nil
+}
+func (r *toggleSettingRepo) Delete(context.Context, string) error { return nil }
+
+type passthroughEncryptor struct{}
+
+func (passthroughEncryptor) Encrypt(plaintext string) (string, error)  { return plaintext, nil }
+func (passthroughEncryptor) Decrypt(ciphertext string) (string, error) { return ciphertext, nil }
+
+type noopImageStorage struct{}
+
+func (noopImageStorage) Save(context.Context, string, string, []byte) (string, error) {
+	return "https://cdn.example.test/object.png", nil
+}
+
+// TestAsyncImageEnablesWithoutRestart drives the actual HTTP path for the bug behind
+// #4458 and #4542: with object storage unconfigured the async endpoint 404s, and the
+// only way to turn it on used to be editing config.yaml and restarting the container.
+// Flipping the admin setting must flip the endpoint over in the same process.
+func TestAsyncImageEnablesWithoutRestart(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+
+	repo := &toggleSettingRepo{values: map[string]string{}}
+	backup := service.NewBackupService(repo, &config.Config{}, passthroughEncryptor{}, nil, nil)
+	factory := func(context.Context, *config.ImageStorageConfig) (service.ImageStorage, error) {
+		return noopImageStorage{}, nil
+	}
+	settings := service.NewImageStorageSettingService(repo, passthroughEncryptor{}, backup, factory, config.ImageStorageConfig{})
+
+	store := &asyncImageMemoryStore{tasks: make(map[string]*service.ImageTaskRecord)}
+	tasks := service.NewImageTaskServiceWithResolver(store, settings.Resolver(), time.Hour, time.Minute)
+
+	h := &AsyncImageHandler{tasks: tasks}
+	h.execute = func(_ string, c *gin.Context) {
+		c.JSON(http.StatusOK, gin.H{"created": 1, "data": []gin.H{{"url": "https://upstream.test/i.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)
+
+	submit := func() *httptest.ResponseRecorder {
+		req := httptest.NewRequest(http.MethodPost, "/v1/images/generations/async",
+			strings.NewReader(`{"model":"gpt-image-1","prompt":"a lighthouse"}`))
+		req.Header.Set("Content-Type", "application/json")
+		rec := httptest.NewRecorder()
+		router.ServeHTTP(rec, req)
+		return rec
+	}
+
+	rec := submit()
+	require.Equal(t, http.StatusNotFound, rec.Code, "disabled until an admin configures object storage")
+	require.Contains(t, rec.Body.String(), "async image tasks are not enabled")
+
+	// The admin saves the setting — no restart, same process.
+	_, err := settings.Update(context.Background(), service.ImageStorageSettings{
+		Enabled: true, Bucket: "my-images",
+		Endpoint: "https://acct.r2.cloudflarestorage.com", AccessKeyID: "ak", SecretAccessKey: "sk",
+	})
+	require.NoError(t, err)
+
+	rec = submit()
+	require.Equal(t, http.StatusAccepted, rec.Code, "the endpoint must go live as soon as the setting is saved")
+
+	var accepted struct {
+		TaskID  string `json:"task_id"`
+		PollURL string `json:"poll_url"`
+	}
+	require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &accepted))
+	require.NotEmpty(t, accepted.TaskID)
+
+	// Turning the feature back off must not strand a task that was already accepted.
+	_, err = settings.Update(context.Background(), service.ImageStorageSettings{Enabled: false})
+	require.NoError(t, err)
+
+	require.Equal(t, http.StatusNotFound, submit().Code, "new submissions are refused again")
+
+	pollRec := httptest.NewRecorder()
+	router.ServeHTTP(pollRec, httptest.NewRequest(http.MethodGet, accepted.PollURL, nil))
+	require.Equal(t, http.StatusOK, pollRec.Code, "an already-accepted task stays pollable after the switch is turned off")
+}
@@ -39,6 +39,13 @@ func (h *AsyncImageHandler) enabled() bool {
 	return h != nil && h.tasks != nil && h.tasks.Enabled()
 }

+// pollable reports whether task lookups can be served. It is deliberately weaker
+// than enabled(): results already written to Redis stay readable after the
+// feature is switched off, so an in-flight task is never stranded.
+func (h *AsyncImageHandler) pollable() bool {
+	return h != nil && h.tasks != nil && h.tasks.Pollable()
+}
+
 // 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) {
@@ -155,7 +162,11 @@ func (h *AsyncImageHandler) checkSecurityAuditBeforeSubmit(c *gin.Context, apiKe
 }

 func (h *AsyncImageHandler) Get(c *gin.Context) {
-	if !h.enabled() {
+	// Polling deliberately does not require the feature to be enabled, only that
+	// the task store is reachable. Turning the switch off in the admin UI must not
+	// strand tasks that were already accepted — their results are still in Redis
+	// and their submitters are still polling.
+	if !h.pollable() {
 		imageTaskJSONError(c, http.StatusNotFound, "not_found_error", "async image tasks are not enabled")
 		return
 	}