fix(queue): 回收失效任务并阻止不确定提交重放

Worker 进程心跳不能证明单个 River job goroutine 仍存活，改以任务执行租约作为回收所有权栅栏。\n\n排队任务若上一次 attempt 在租约中断时处于 submitting 或 response_received，则转入人工复核并生成 release 记录，不再重新提交上游；明确被上游拒绝的响应仍允许重试。\n\n验证：go vet ./...；env -u AI_GATEWAY_TEST_DATABASE_URL go test ./... -count=1；真实 PostgreSQL 下相关 store 与 River 集成测试通过。
2026-07-30 19:15:19 +08:00
parent 9e06ed6162
commit f33d6d64e0
4 changed files with 256 additions and 65 deletions
@@ -254,18 +254,26 @@ WHERE id = $1`, yielded.RiverJobID).Scan(&yieldedState); err != nil {
 	if err != nil {
 		t.Fatalf("recover orphaned River jobs: %v", err)
 	}
-	if recovered != 1 {
-		t.Fatalf("recovered jobs=%d, want 1", recovered)
+	if recovered != 2 {
+		t.Fatalf("recovered jobs=%d, want 2", recovered)
 	}
-	var orphanedState, protectedState string
+	var orphanedState, protectedState, yieldedRecoveredState string
 	if err := db.Pool().QueryRow(ctx, `SELECT state::text FROM river_job WHERE id = $1`, orphaned.RiverJobID).Scan(&orphanedState); err != nil {
 		t.Fatalf("read orphaned River job state: %v", err)
 	}
 	if err := db.Pool().QueryRow(ctx, `SELECT state::text FROM river_job WHERE id = $1`, protected.RiverJobID).Scan(&protectedState); err != nil {
 		t.Fatalf("read protected River job state: %v", err)
 	}
-	if orphanedState != "retryable" || protectedState != "running" {
-		t.Fatalf("River states orphaned=%s protected=%s, want retryable/running", orphanedState, protectedState)
+	if err := db.Pool().QueryRow(ctx, `SELECT state::text FROM river_job WHERE id = $1`, yielded.RiverJobID).Scan(&yieldedRecoveredState); err != nil {
+		t.Fatalf("read yielded River job state: %v", err)
+	}
+	if orphanedState != "retryable" || protectedState != "running" || yieldedRecoveredState != "retryable" {
+		t.Fatalf(
+			"River states orphaned=%s protected=%s yielded=%s, want retryable/running/retryable",
+			orphanedState,
+			protectedState,
+			yieldedRecoveredState,
+		)
 	}
 	recovered, err = db.RecoverOrphanedAsyncRiverJobs(ctx, 30*time.Second, 10)
 	if err != nil || recovered != 0 {
@@ -168,6 +168,127 @@ WHERE task_id=$1::uuid AND event_type='task.billing.review'`, created.ID).Scan(&
 	}
 }

+func TestQueuedPreparationDoesNotReplayInterruptedUpstreamSubmission(t *testing.T) {
+	db := billingV2IntegrationStore(t)
+	ctx := context.Background()
+	tenantID, gatewayUserID := seedWalletReservationUser(t, ctx, db)
+
+	for _, submissionStatus := range []string{"submitting", "response_received"} {
+		t.Run(submissionStatus, func(t *testing.T) {
+			user := &auth.User{
+				ID:              "billing-queued-review-" + uuid.NewString(),
+				GatewayUserID:   gatewayUserID,
+				GatewayTenantID: tenantID,
+			}
+			created, err := db.CreateTask(ctx, CreateTaskInput{
+				Kind: "images.generations", Model: "billing-v2-model", RunMode: "production",
+				Request: map[string]any{"model": "billing-v2-model"},
+			}, user)
+			if err != nil {
+				t.Fatal(err)
+			}
+			t.Cleanup(func() {
+				_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_tasks WHERE id=$1::uuid`, created.ID)
+			})
+
+			token := uuid.NewString()
+			if _, err := db.ClaimTaskExecution(ctx, created.ID, token, 5*time.Minute); err != nil {
+				t.Fatal(err)
+			}
+			attemptID, err := db.CreateTaskAttempt(ctx, CreateTaskAttemptInput{
+				TaskID: created.ID, AttemptNo: 1, Status: "running",
+			})
+			if err != nil {
+				t.Fatal(err)
+			}
+			if err := db.SetAttemptUpstreamSubmissionStatus(ctx, attemptID, submissionStatus); err != nil {
+				t.Fatal(err)
+			}
+			if _, err := db.pool.Exec(ctx, `
+UPDATE gateway_task_attempts
+SET status='failed',
+    retryable=true,
+    error_code='execution_lease_expired',
+    error_message='attempt execution lease expired',
+    finished_at=now()
+WHERE id=$1::uuid`, attemptID); err != nil {
+				t.Fatal(err)
+			}
+			if _, err := db.pool.Exec(ctx, `
+UPDATE gateway_tasks
+SET status='queued',
+    execution_token=NULL,
+    execution_lease_expires_at=NULL,
+    heartbeat_at=NULL,
+    locked_at=NULL,
+    updated_at=now()
+WHERE id=$1::uuid`, created.ID); err != nil {
+				t.Fatal(err)
+			}
+
+			if _, err := db.ClaimTaskPreparation(ctx, created.ID, uuid.NewString(), 5*time.Minute); !errors.Is(err, ErrTaskExecutionManualReview) {
+				t.Fatalf("queued ambiguous preparation error=%v", err)
+			}
+			review, err := db.GetTask(ctx, created.ID)
+			if err != nil {
+				t.Fatal(err)
+			}
+			if review.Status != "failed" || review.BillingStatus != "manual_review" || review.ErrorCode != "upstream_submission_unknown" {
+				t.Fatalf("manual review task=%+v", review)
+			}
+		})
+	}
+}
+
+func TestQueuedPreparationCanResumeAfterKnownRejectedResponse(t *testing.T) {
+	db := billingV2IntegrationStore(t)
+	ctx := context.Background()
+	_, gatewayUserID := seedWalletReservationUser(t, ctx, db)
+	user := &auth.User{ID: "billing-queued-known-response-" + uuid.NewString(), GatewayUserID: gatewayUserID}
+	created, err := db.CreateTask(ctx, CreateTaskInput{
+		Kind: "images.generations", Model: "billing-v2-model", RunMode: "production",
+		Request: map[string]any{"model": "billing-v2-model"},
+	}, user)
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() {
+		_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_tasks WHERE id=$1::uuid`, created.ID)
+	})
+
+	if _, err := db.ClaimTaskExecution(ctx, created.ID, uuid.NewString(), 5*time.Minute); err != nil {
+		t.Fatal(err)
+	}
+	attemptID, err := db.CreateTaskAttempt(ctx, CreateTaskAttemptInput{
+		TaskID: created.ID, AttemptNo: 1, Status: "running",
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := db.SetAttemptUpstreamSubmissionStatus(ctx, attemptID, "response_received"); err != nil {
+		t.Fatal(err)
+	}
+	if err := db.FinishTaskAttempt(ctx, FinishTaskAttemptInput{
+		AttemptID: attemptID, Status: "failed", ErrorCode: "upstream_rejected",
+	}); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := db.pool.Exec(ctx, `
+UPDATE gateway_tasks
+SET status='queued',
+    execution_token=NULL,
+    execution_lease_expires_at=NULL,
+    heartbeat_at=NULL,
+    locked_at=NULL,
+    updated_at=now()
+WHERE id=$1::uuid`, created.ID); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := db.ClaimTaskPreparation(ctx, created.ID, uuid.NewString(), 5*time.Minute); err != nil {
+		t.Fatalf("known rejected response should remain retryable: %v", err)
+	}
+}
+
 func TestExpiredExecutionLeaseCanResumeAfterKnownRejectedResponse(t *testing.T) {
 	db := billingV2IntegrationStore(t)
 	ctx := context.Background()
@@ -304,58 +304,15 @@ FOR UPDATE`, taskID).Scan(&queuedReady, &runningExpired, &production, &hasGatewa
 		if !queuedReady && !runningExpired {
 			return ErrTaskExecutionLeaseUnavailable
 		}
-		if runningExpired && production {
-			var submissionAmbiguous bool
-			if err := tx.QueryRow(ctx, `
-SELECT COALESCE((
-  SELECT (
-    (status = 'running' AND upstream_submission_status IN ('submitting', 'response_received'))
-    OR (status = 'failed' AND upstream_submission_status = 'submitting')
-  )
-  FROM gateway_task_attempts
-  WHERE task_id = $1::uuid
-  ORDER BY attempt_no DESC, started_at DESC
-  LIMIT 1
-), false)`, taskID).Scan(&submissionAmbiguous); err != nil {
+		if (queuedReady || runningExpired) && production {
+			submissionAmbiguous, err := taskExecutionRequiresManualReviewTx(ctx, tx, taskID)
+			if err != nil {
 				return err
 			}
 			if submissionAmbiguous {
-				if _, err := tx.Exec(ctx, `
-UPDATE gateway_tasks
-SET status = 'failed',
-    billing_status = CASE WHEN $2 THEN 'manual_review' ELSE 'not_required' END,
-    billing_updated_at = now(),
-    error = NULL,
-    error_code = 'upstream_submission_unknown',
-    error_message = 'upstream submission result is unknown',
-    locked_by = NULL,
-    locked_at = NULL,
-    heartbeat_at = NULL,
-    execution_token = NULL,
-    execution_lease_expires_at = NULL,
-    remote_task_payload = '{}'::jsonb,
-    finished_at = now(),
-    updated_at = now()
-WHERE id = $1::uuid`, taskID, hasGatewayUser); err != nil {
+				if err := markTaskExecutionManualReviewTx(ctx, tx, taskID, hasGatewayUser); err != nil {
 					return err
 				}
-				if hasGatewayUser {
-					payloadJSON, _ := json.Marshal(map[string]any{
-						"taskId": taskID, "classification": "upstream_submission_unknown",
-					})
-					if _, err := tx.Exec(ctx, `
-INSERT INTO settlement_outbox (
-  task_id, event_type, action, amount, currency, pricing_snapshot, payload,
-  status, next_attempt_at, manual_review_reason
-)
-SELECT id, 'task.billing.review', 'release', reservation_amount, billing_currency,
-       pricing_snapshot, $2::jsonb, 'manual_review', now(), 'upstream_submission_unknown'
-FROM gateway_tasks
-WHERE id = $1::uuid
-ON CONFLICT (task_id, event_type) DO NOTHING`, taskID, string(payloadJSON)); err != nil {
-						return err
-					}
-				}
 				manualReview = true
 				return nil
 			}
@@ -389,11 +346,116 @@ RETURNING `+gatewayTaskColumns, taskID, executionToken, int(leaseTTL/time.Second
 	return task, nil
 }

+func taskExecutionRequiresManualReviewTx(ctx context.Context, tx pgx.Tx, taskID string) (bool, error) {
+	var submissionAmbiguous bool
+	err := tx.QueryRow(ctx, `
+SELECT COALESCE((
+  SELECT (
+    (status = 'running' AND upstream_submission_status IN ('submitting', 'response_received'))
+    OR (status = 'failed' AND upstream_submission_status = 'submitting')
+    OR (
+      status = 'failed'
+      AND upstream_submission_status = 'response_received'
+      AND error_code = 'execution_lease_expired'
+    )
+  )
+  FROM gateway_task_attempts
+  WHERE task_id = $1::uuid
+  ORDER BY attempt_no DESC, started_at DESC
+  LIMIT 1
+), false)`, taskID).Scan(&submissionAmbiguous)
+	return submissionAmbiguous, err
+}
+
+func markTaskExecutionManualReviewTx(
+	ctx context.Context,
+	tx pgx.Tx,
+	taskID string,
+	hasGatewayUser bool,
+) error {
+	if _, err := tx.Exec(ctx, `
+UPDATE gateway_tasks
+SET status = 'failed',
+    billing_status = CASE WHEN $2 THEN 'manual_review' ELSE 'not_required' END,
+    billing_updated_at = now(),
+    error = NULL,
+    error_code = 'upstream_submission_unknown',
+    error_message = 'upstream submission result is unknown',
+    locked_by = NULL,
+    locked_at = NULL,
+    heartbeat_at = NULL,
+    execution_token = NULL,
+    execution_lease_expires_at = NULL,
+    remote_task_payload = '{}'::jsonb,
+    finished_at = now(),
+    updated_at = now()
+WHERE id = $1::uuid`, taskID, hasGatewayUser); err != nil {
+		return err
+	}
+	if !hasGatewayUser {
+		return nil
+	}
+	payloadJSON, _ := json.Marshal(map[string]any{
+		"taskId": taskID, "classification": "upstream_submission_unknown",
+	})
+	if _, err := tx.Exec(ctx, `
+INSERT INTO settlement_outbox (
+  task_id, event_type, action, amount, currency, pricing_snapshot, payload,
+  status, next_attempt_at, manual_review_reason
+)
+SELECT id, 'task.billing.review', 'release', reservation_amount, billing_currency,
+       pricing_snapshot, $2::jsonb, 'manual_review', now(), 'upstream_submission_unknown'
+FROM gateway_tasks
+WHERE id = $1::uuid
+ON CONFLICT (task_id, event_type) DO NOTHING`, taskID, string(payloadJSON)); err != nil {
+		return err
+	}
+	return nil
+}
+
 func (s *Store) ClaimTaskPreparation(ctx context.Context, taskID string, executionToken string, leaseTTL time.Duration) (GatewayTask, error) {
 	if leaseTTL <= 0 {
 		leaseTTL = 5 * time.Minute
 	}
-	task, err := scanGatewayTask(s.pool.QueryRow(ctx, `
+	var task GatewayTask
+	manualReview := false
+	err := s.beginTransaction(ctx, func(tx pgx.Tx) error {
+		var queuedReady bool
+		var production bool
+		var hasGatewayUser bool
+		if err := tx.QueryRow(ctx, `
+SELECT status = 'queued'
+         AND next_run_at <= now()
+         AND (
+           execution_token IS NULL
+           OR execution_lease_expires_at IS NULL
+           OR execution_lease_expires_at <= now()
+         ),
+       run_mode = 'production',
+       gateway_user_id IS NOT NULL
+FROM gateway_tasks
+WHERE id = $1::uuid
+FOR UPDATE`, taskID).Scan(&queuedReady, &production, &hasGatewayUser); err != nil {
+			return err
+		}
+		if !queuedReady {
+			return ErrTaskExecutionLeaseUnavailable
+		}
+		if production {
+			submissionAmbiguous, err := taskExecutionRequiresManualReviewTx(ctx, tx, taskID)
+			if err != nil {
+				return err
+			}
+			if submissionAmbiguous {
+				if err := markTaskExecutionManualReviewTx(ctx, tx, taskID, hasGatewayUser); err != nil {
+					return err
+				}
+				manualReview = true
+				return nil
+			}
+		}
+		var err error
+		task, err = scanGatewayTask(tx.QueryRow(ctx, `
 UPDATE gateway_tasks
 SET execution_token = $2::uuid,
    execution_lease_expires_at = now() + ($3::int * interval '1 second'),
@@ -409,9 +471,17 @@ WHERE id = $1::uuid
    OR execution_lease_expires_at <= now()
  )
 RETURNING `+gatewayTaskColumns, taskID, executionToken, int(leaseTTL/time.Second)))
+		return err
+	})
 	if errors.Is(err, pgx.ErrNoRows) {
 		return GatewayTask{}, ErrTaskExecutionLeaseUnavailable
 	}
+	if err != nil {
+		return GatewayTask{}, err
+	}
+	if manualReview {
+		return GatewayTask{}, ErrTaskExecutionManualReview
+	}
 	return task, err
 }

@@ -297,10 +297,10 @@ FROM deleted`, staleAfter.String(), limit).Scan(&yielded); err != nil {
 	return yielded, nil
 }

-// RecoverOrphanedAsyncRiverJobs retries River jobs whose owning worker is no
-// longer alive after the gateway execution lease has already returned the task
-// to queued. Normal long-running jobs remain protected by their running task
-// state and the active worker heartbeat.
+// RecoverOrphanedAsyncRiverJobs retries stale River jobs after the gateway
+// execution lease has already returned the task to queued. A process heartbeat
+// does not prove that an individual River job goroutine is still alive; the
+// task execution lease is the ownership fence.
 func (s *Store) RecoverOrphanedAsyncRiverJobs(
 	ctx context.Context,
 	workerStaleAfter time.Duration,
@@ -335,14 +335,6 @@ WITH orphaned AS MATERIALIZED (
    AND task.async_mode = true
    AND task.status = 'queued'
    AND (task.execution_lease_expires_at IS NULL OR task.execution_lease_expires_at <= now())
-    AND NOT EXISTS (
-      SELECT 1
-      FROM gateway_worker_instances worker
-      WHERE worker.status = 'active'
-        AND worker.heartbeat_at > now() - $1::interval
-        AND COALESCE(job.attempted_by[array_length(job.attempted_by, 1)], '')
-              LIKE worker.instance_id || '-exec-%'
-    )
  ORDER BY job.attempted_at ASC, job.id ASC
  LIMIT $2
  FOR UPDATE OF job SKIP LOCKED