fix(runtime): 运行恢复整合 - 合并 PR #253 本地任务恢复改动

2026-08-18 20:27:06 +08:00
parent 0002cdd2cb f62cc5ef3e
commit f7c5b3b257
13 changed files with 736 additions and 67 deletions
--- a/backend/internal/repository/repository.go
+++ b/backend/internal/repository/repository.go
@@ -294,7 +294,7 @@ func (r *Repository) ClaimNextTask(owner string, leaseDuration time.Duration) (*
 	now := time.Now()
 	leaseExpiresAt := now.Add(leaseDuration)
 	err := r.db.Transaction(func(tx *gorm.DB) error {
-		query := tx.Where("status = ? OR (status = ? AND (lease_expires_at IS NULL OR lease_expires_at <= ?))", model.TaskStatusQueued, model.TaskStatusRunning, now).
+		query := tx.Where("(status = ? OR (status = ? AND (lease_expires_at IS NULL OR lease_expires_at <= ?))) AND (next_poll_at IS NULL OR next_poll_at <= ?)", model.TaskStatusQueued, model.TaskStatusRunning, now, now).
 			Order("created_at asc").Limit(1)
 		if r.Dialect() == "postgres" {
 			query = query.Clauses(clause.Locking{Strength: "UPDATE", Options: "SKIP LOCKED"})
@@ -309,7 +309,7 @@ func (r *Repository) ClaimNextTask(owner string, leaseDuration time.Duration) (*
 		}
 		claim := tx.Model(&model.Task{}).Where("id = ?", task.ID)
 		if r.Dialect() != "postgres" {
-			claim = claim.Where("status = ? OR (status = ? AND (lease_expires_at IS NULL OR lease_expires_at <= ?))", model.TaskStatusQueued, model.TaskStatusRunning, now)
+			claim = claim.Where("(status = ? OR (status = ? AND (lease_expires_at IS NULL OR lease_expires_at <= ?))) AND (next_poll_at IS NULL OR next_poll_at <= ?)", model.TaskStatusQueued, model.TaskStatusRunning, now, now)
 		}
 		updated := claim.
 			Updates(map[string]any{
@@ -320,6 +320,7 @@ func (r *Repository) ClaimNextTask(owner string, leaseDuration time.Duration) (*
 				"started_at":       gorm.Expr("COALESCE(started_at, ?)", now),
 				"lease_owner":      owner,
 				"lease_expires_at": leaseExpiresAt,
+				"next_poll_at":     nil,
 				"updated_at":       now,
 			})
 		if updated.Error != nil {
@@ -358,6 +359,23 @@ func (r *Repository) UpdateTaskProviderState(id string, providerRequestID string
 	return r.db.Model(&model.Task{}).Where("id = ?", id).Updates(updates).Error
 }

+func (r *Repository) DeferRunningTaskForProviderPoll(id string, owner string, stage string, delay time.Duration) error {
+	now := time.Now()
+	result := r.db.Model(&model.Task{}).
+		Where("id = ? AND status = ? AND lease_owner = ?", id, model.TaskStatusRunning, owner).
+		Updates(map[string]any{
+			"stage": stage, "error": "", "completed_at": nil, "next_poll_at": now.Add(delay),
+			"lease_owner": "", "lease_expires_at": nil, "updated_at": now,
+		})
+	if result.Error != nil {
+		return result.Error
+	}
+	if result.RowsAffected != 1 {
+		return ErrTaskStateConflict
+	}
+	return nil
+}
+
 // 人工恢复仅锁定失败任务；旧 worker 的租约可覆盖，但未过期的人工恢复租约不能并发抢占。
 func (r *Repository) ClaimFailedTaskProviderRecovery(id string, userID string, owner string, leaseDuration time.Duration) error {
 	now := time.Now()
--- a/backend/internal/repository/task_provider_poll_test.go
+++ b/backend/internal/repository/task_provider_poll_test.go
@@ -0,0 +1,51 @@
+package repository
+
+import (
+	"testing"
+	"time"
+
+	"infinite-canvas/backend/internal/model"
+
+	"gorm.io/driver/sqlite"
+	"gorm.io/gorm"
+)
+
+func TestDeferredProviderPollKeepsOriginalTaskIdentityWithoutImmediateReclaim(t *testing.T) {
+	db, err := gorm.Open(sqlite.Open("file:provider-poll-defer?mode=memory&cache=shared"), &gorm.Config{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := db.AutoMigrate(&model.Task{}); err != nil {
+		t.Fatal(err)
+	}
+	task := model.Task{
+		ID: "task-1", UserID: "user-1", Type: "canvas_video", Status: model.TaskStatusRunning,
+		LeaseOwner: "worker-1", LeaseExpiresAt: ptrTime(time.Now().Add(time.Minute)), ProviderRequestID: "provider-task-1",
+	}
+	if err := db.Create(&task).Error; err != nil {
+		t.Fatal(err)
+	}
+	repo := New(db)
+	if err := repo.DeferRunningTaskForProviderPoll(task.ID, task.LeaseOwner, "后台仍在生成", time.Minute); err != nil {
+		t.Fatal(err)
+	}
+	claimed, err := repo.ClaimNextTask("worker-2", 45*time.Second)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if claimed != nil {
+		t.Fatalf("task reclaimed before next poll: %#v", claimed)
+	}
+	if err := db.Model(&model.Task{}).Where("id = ?", task.ID).Update("next_poll_at", time.Now().Add(-time.Second)).Error; err != nil {
+		t.Fatal(err)
+	}
+	claimed, err = repo.ClaimNextTask("worker-2", 45*time.Second)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if claimed == nil || claimed.ID != task.ID || claimed.ProviderRequestID != task.ProviderRequestID {
+		t.Fatalf("reclaimed task = %#v", claimed)
+	}
+}
+
+func ptrTime(value time.Time) *time.Time { return &value }
--- a/backend/internal/service/provider.go
+++ b/backend/internal/service/provider.go
@@ -1854,7 +1854,10 @@ func runNewAPIChannel2VideoTask(ctx context.Context, input canvasGenerationInput
 			return nil, err
 		}
 	}
-	return nil, fmt.Errorf("NewAPI Video Generations 视频生成超时（任务 %s）", id)
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+	return nil, context.DeadlineExceeded
 }

 // 单次查询只读取既有上游任务，不创建新任务；自动轮询和人工恢复共用这条安全边界。
--- a/backend/internal/service/provider_test.go
+++ b/backend/internal/service/provider_test.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"encoding/json"
+	"errors"
 	"io"
 	"mime/multipart"
 	"net/http"
@@ -12,6 +13,9 @@ import (
 	"strconv"
 	"strings"
 	"testing"
+	"time"
+
+	"infinite-canvas/backend/internal/model"
 )

 const testReferenceImageDataURL = "data:image/png;base64,aGVsbG8="
@@ -1306,6 +1310,59 @@ func TestRunNewAPIChannel2VideoTaskDownloadsTemporaryResult(t *testing.T) {
 	}
 }

+func TestRunNewAPIChannel2VideoTaskResumesOriginalProviderTaskWithoutAnotherPost(t *testing.T) {
+	t.Setenv("CANVAS_ALLOW_PRIVATE_UPSTREAMS", "true")
+	paths := make([]string, 0, 2)
+	var server *httptest.Server
+	server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		paths = append(paths, r.Method+" "+r.URL.Path)
+		switch r.Method + " " + r.URL.Path {
+		case "GET /v1/video/generations/existing-provider-task":
+			w.Header().Set("Content-Type", "application/json")
+			_, _ = w.Write([]byte(`{"code":"success","data":{"task_id":"existing-provider-task","status":"SUCCESS","result_url":"` + server.URL + `/video.mp4"}}`))
+		case "GET /video.mp4":
+			w.Header().Set("Content-Type", "video/mp4")
+			_, _ = w.Write([]byte("video"))
+		default:
+			http.NotFound(w, r)
+		}
+	}))
+	defer server.Close()
+
+	input := canvasGenerationInput{
+		Mode:   "video",
+		Config: providerConfig{BaseURL: server.URL, APIKey: "test-key", Model: "video-model", InterfaceType: string(model.ChannelInterfaceNewAPIChannel2)},
+	}
+	inputJSON, err := json.Marshal(input)
+	if err != nil {
+		t.Fatal(err)
+	}
+	ctx := withProviderAnalytics(context.Background(), nil, model.Task{
+		ID: "task-1", Type: "canvas_video", ProviderRequestID: "existing-provider-task", InputJSON: string(inputJSON),
+	})
+	result, err := runNewAPIChannel2VideoTask(ctx, input)
+	if err != nil {
+		t.Fatalf("runNewAPIChannel2VideoTask() error = %v", err)
+	}
+	if result["video"] == nil {
+		t.Fatalf("result = %#v", result)
+	}
+	want := "GET /v1/video/generations/existing-provider-task,GET /video.mp4"
+	if got := strings.Join(paths, ","); got != want {
+		t.Fatalf("paths = %q, want %q", got, want)
+	}
+}
+
+func TestRunNewAPIChannel2VideoTaskReturnsTypedDeadlineWhenPollingWindowEnds(t *testing.T) {
+	ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(-time.Second))
+	defer cancel()
+	ctx = withProviderAnalytics(ctx, nil, model.Task{ID: "task-1", Type: "canvas_video", ProviderRequestID: "existing-provider-task"})
+	_, err := runNewAPIChannel2VideoTask(ctx, canvasGenerationInput{})
+	if !errors.Is(err, context.DeadlineExceeded) {
+		t.Fatalf("runNewAPIChannel2VideoTask() error = %v, want context deadline exceeded", err)
+	}
+}
+
 func TestRunGeminiVeoVideoTaskUsesLongRunningOperation(t *testing.T) {
 	t.Setenv("CANVAS_ALLOW_PRIVATE_UPSTREAMS", "true")
 	paths := make([]string, 0, 3)
--- a/backend/internal/service/runtime_policy_test.go
+++ b/backend/internal/service/runtime_policy_test.go
@@ -1,7 +1,10 @@
 package service

 import (
+	"context"
+	"encoding/json"
 	"testing"
+	"time"

 	"infinite-canvas/backend/internal/model"
 	"infinite-canvas/backend/internal/repository"
@@ -10,6 +13,76 @@ import (
 	"gorm.io/gorm"
 )

+func TestVideoTaskTimeoutHasFiveMinuteSafetyFloor(t *testing.T) {
+	policy := defaultRuntimePolicy().Task
+	policy.VideoTimeoutMinutes = 1
+	policy.ImageTimeoutMinutes = 1
+	if got := taskExecutionTimeoutWithPolicy("canvas_video", policy); got != 5*time.Minute {
+		t.Fatalf("video timeout = %s, want 5m", got)
+	}
+	if got := taskExecutionTimeoutWithPolicy("canvas_image", policy); got != time.Minute {
+		t.Fatalf("image timeout = %s, want 1m", got)
+	}
+}
+
+func TestOnlyResumableNewAPIChannel2VideoDeadlinesStayRunning(t *testing.T) {
+	svc := &Service{}
+	input, err := json.Marshal(canvasGenerationInput{Mode: "video", Config: providerConfig{BaseURL: "https://example.com", InterfaceType: string(model.ChannelInterfaceNewAPIChannel2)}})
+	if err != nil {
+		t.Fatal(err)
+	}
+	base := model.Task{ID: "task-1", Type: "canvas_video", ProviderRequestID: "provider-task-1"}
+	if !svc.shouldDeferVideoProviderTask(base, string(input), context.DeadlineExceeded) {
+		t.Fatal("resumable NewAPI Channel 2 deadline should remain running")
+	}
+	if svc.shouldDeferVideoProviderTask(model.Task{ID: "task-2", Type: "canvas_video"}, string(input), context.DeadlineExceeded) {
+		t.Fatal("missing provider task id must not be deferred")
+	}
+	if svc.shouldDeferVideoProviderTask(base, string(input), context.Canceled) {
+		t.Fatal("explicit cancellation must not be deferred")
+	}
+	other, err := json.Marshal(canvasGenerationInput{Mode: "video", Config: providerConfig{BaseURL: "https://example.com", InterfaceType: string(model.ChannelInterfaceNewAPIVideo)}})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if svc.shouldDeferVideoProviderTask(base, string(other), context.DeadlineExceeded) {
+		t.Fatal("providers without a verified query-only resume contract must fail closed")
+	}
+}
+
+func TestResumableVideoDeadlineUsesResolvedSystemChannelProtocol(t *testing.T) {
+	db, err := gorm.Open(sqlite.Open("file:"+newID()+"?mode=memory&cache=shared"), &gorm.Config{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := db.AutoMigrate(&model.ModelChannel{}, &model.ChannelModel{}); err != nil {
+		t.Fatal(err)
+	}
+	channel := model.ModelChannel{
+		ID: "channel-video", UserID: "admin", Scope: model.ChannelScopeSystem, Enabled: true, Name: "Video",
+		BaseURL: "https://example.com", APIKey: "test-key", APIFormat: "openai", ModelsJSON: `["video-model"]`,
+	}
+	channelModel := model.ChannelModel{
+		ID: "model-video", ChannelID: channel.ID, ModelKey: "video-model", Capability: "video",
+		Protocol: model.ChannelInterfaceNewAPIChannel2, Enabled: true, PriceConfigured: true,
+	}
+	if err := db.Create(&channel).Error; err != nil {
+		t.Fatal(err)
+	}
+	if err := db.Create(&channelModel).Error; err != nil {
+		t.Fatal(err)
+	}
+	input, err := json.Marshal(canvasGenerationInput{Mode: "video", Config: providerConfig{ChannelID: channel.ID, Model: "video-model"}})
+	if err != nil {
+		t.Fatal(err)
+	}
+	svc := New(repository.New(db), t.TempDir())
+	task := model.Task{ID: "task-system", Type: "canvas_video", ProviderRequestID: "provider-task-1"}
+	if !svc.shouldDeferVideoProviderTask(task, string(input), context.DeadlineExceeded) {
+		t.Fatal("system-channel task must use the resolved model protocol when deciding query-only recovery")
+	}
+}
+
 func TestRuntimePolicyDefaultsAndSelfUseModeValidate(t *testing.T) {
 	if err := validateRuntimePolicy(defaultRuntimePolicy()); err != nil {
 		t.Fatalf("default runtime policy error = %v", err)
--- a/backend/internal/service/service.go
+++ b/backend/internal/service/service.go
@@ -990,6 +990,14 @@ func (s *Service) processClaimedTask(task *model.Task) error {
 			return nil
 		}
 		if errors.Is(err, context.DeadlineExceeded) {
+			decryptedInput, decryptErr := s.decryptTaskInputJSON(task.InputJSON)
+			if decryptErr == nil && s.shouldDeferVideoProviderTask(*task, decryptedInput, err) {
+				if deferErr := s.repo.DeferRunningTaskForProviderPoll(task.ID, task.LeaseOwner, "后台仍在生成", 15*time.Second); deferErr != nil {
+					return deferErr
+				}
+				_ = s.log(task.UserID, task.ID, "info", "前台等待结束，上游视频仍在生成，将继续回查原任务", task.PollStage)
+				return nil
+			}
 			err = errors.New(taskTimeoutMessage(task.Type))
 		}
 		task.Status = model.TaskStatusFailed
@@ -1079,7 +1087,7 @@ func taskExecutionTimeoutWithPolicy(taskType string, policy RuntimeTaskPolicy) t
 	case taskType == "agent_storyboard" || taskType == "agent_storyboard_rows":
 		return time.Duration(policy.StoryboardTimeoutMinutes) * time.Minute
 	case strings.HasPrefix(taskType, "canvas_video") || strings.HasPrefix(taskType, "video_"):
-		return time.Duration(policy.VideoTimeoutMinutes) * time.Minute
+		return max(time.Duration(policy.VideoTimeoutMinutes)*time.Minute, 5*time.Minute)
 	case strings.HasPrefix(taskType, "canvas_image"):
 		return time.Duration(policy.ImageTimeoutMinutes) * time.Minute
 	case strings.HasPrefix(taskType, "canvas_audio"):
@@ -1091,6 +1099,18 @@ func taskExecutionTimeoutWithPolicy(taskType string, policy RuntimeTaskPolicy) t
 	}
 }

+func (s *Service) shouldDeferVideoProviderTask(task model.Task, decryptedInput string, err error) bool {
+	if !errors.Is(err, context.DeadlineExceeded) || strings.TrimSpace(task.ProviderRequestID) == "" || (!strings.HasPrefix(task.Type, "canvas_video") && !strings.HasPrefix(task.Type, "video_")) {
+		return false
+	}
+	var input canvasGenerationInput
+	if json.Unmarshal([]byte(decryptedInput), &input) != nil {
+		return false
+	}
+	resolved, resolveErr := s.resolveProviderConfig(input.Config)
+	return resolveErr == nil && resolved.InterfaceType == string(model.ChannelInterfaceNewAPIChannel2)
+}
+
 func taskTimeoutMessage(taskType string) string {
 	if strings.HasPrefix(taskType, "canvas_video") || strings.HasPrefix(taskType, "video_") {
 		return "视频生成等待超时，请稍后到任务中心查看或重试。"