feat(volces): 接入任务兼容查询与上游取消

2026-07-22 00:25:22 +08:00
parent ddd68cfebd
commit d7951cfdd2
11 changed files with 885 additions and 55 deletions
@@ -0,0 +1,43 @@
+package store
+
+import (
+	"context"
+	"strings"
+)
+
+// GetRuntimeModelCandidateForRemoteTask restores the exact platform model used
+// to submit an asynchronous provider task. It deliberately ignores enabled
+// state so a task can still be cancelled after its platform is disabled.
+func (s *Store) GetRuntimeModelCandidateForRemoteTask(ctx context.Context, platformModelID string, platformID string) (RuntimeModelCandidate, bool, error) {
+	platformModelID = strings.TrimSpace(platformModelID)
+	platformID = strings.TrimSpace(platformID)
+	if platformModelID == "" || platformID == "" {
+		return RuntimeModelCandidate{}, false, nil
+	}
+	var candidate RuntimeModelCandidate
+	var credentials, config []byte
+	err := s.pool.QueryRow(ctx, `
+SELECT p.id::text, p.platform_key, p.name, p.provider,
+       COALESCE(NULLIF(p.config->>'specType', ''), p.provider), COALESCE(p.base_url, ''), p.auth_type,
+       p.credentials, p.config, m.id::text, COALESCE(NULLIF(m.provider_model_name, ''), m.model_name),
+       m.model_name, COALESCE(m.model_alias, ''),
+       COALESCE((m.model_type->>0), 'video_generate')
+FROM platform_models m
+JOIN integration_platforms p ON p.id = m.platform_id
+WHERE m.id = $1::uuid AND p.id = $2::uuid AND p.deleted_at IS NULL`, platformModelID, platformID).Scan(
+		&candidate.PlatformID, &candidate.PlatformKey, &candidate.PlatformName, &candidate.Provider,
+		&candidate.SpecType, &candidate.BaseURL, &candidate.AuthType, &credentials, &config,
+		&candidate.PlatformModelID, &candidate.ProviderModelName, &candidate.ModelName, &candidate.ModelAlias, &candidate.ModelType,
+	)
+	if IsNotFound(err) {
+		return RuntimeModelCandidate{}, false, nil
+	}
+	if err != nil {
+		return RuntimeModelCandidate{}, false, err
+	}
+	candidate.Credentials = decodeObject(credentials)
+	candidate.PlatformConfig = decodeObject(config)
+	candidate.ClientID = candidate.PlatformKey + ":" + candidate.ModelType + ":" + firstNonEmpty(candidate.ProviderModelName, candidate.ModelName)
+	candidate.QueueKey = candidate.ClientID
+	return candidate, true, nil
+}
@@ -530,7 +530,8 @@ WHERE id = $1::uuid
 UPDATE gateway_task_attempts
 SET remote_task_id = NULLIF($2::text, ''),
    response_snapshot = COALESCE(response_snapshot, '{}'::jsonb) || jsonb_build_object('remote_task_payload', $3::jsonb)
-WHERE id = $1::uuid`,
+WHERE id = $1::uuid
+  AND status = 'running'`,
 			attemptID,
 			remoteTaskID,
 			string(payloadJSON),
@@ -585,6 +586,72 @@ WHERE id = $1::uuid
 	return task, true, nil
 }

+// CancelSubmittedTask records a confirmed upstream cancellation. Callers must
+// first complete the provider-side DELETE so local status never claims a remote
+// task was cancelled when the upstream request was not accepted.
+func (s *Store) CancelSubmittedTask(ctx context.Context, taskID string, executionToken string, message string) (GatewayTask, bool, error) {
+	message = strings.TrimSpace(message)
+	if message == "" {
+		message = "任务已由上游取消"
+	}
+	var task GatewayTask
+	changed := false
+	err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
+		var err error
+		task, err = scanGatewayTask(tx.QueryRow(ctx, `
+UPDATE gateway_tasks
+SET status = 'cancelled',
+    error = NULLIF($2, ''),
+    error_code = 'task_cancelled',
+    error_message = NULLIF($2, ''),
+    billing_status = CASE
+      WHEN run_mode <> 'production' OR gateway_user_id IS NULL THEN 'not_required'
+      WHEN reservation_amount > 0 THEN 'pending'
+      ELSE 'released'
+    END,
+    billing_updated_at = now(),
+    locked_by = NULL,
+    locked_at = NULL,
+    heartbeat_at = NULL,
+    execution_token = NULL,
+    execution_lease_expires_at = NULL,
+    finished_at = now(),
+    updated_at = now()
+WHERE id = $1::uuid
+  AND NULLIF(remote_task_id, '') IS NOT NULL
+  AND (
+    (status = 'running' AND execution_token = NULLIF($3, '')::uuid)
+    OR status = 'queued'
+  )
+RETURNING `+gatewayTaskColumns, taskID, message, strings.TrimSpace(executionToken)))
+		if IsNotFound(err) {
+			return nil
+		}
+		if err != nil {
+			return err
+		}
+		changed = true
+		payloadJSON, _ := json.Marshal(map[string]any{"taskId": taskID, "reason": "upstream_cancelled"})
+		_, err = tx.Exec(ctx, `
+INSERT INTO settlement_outbox (
+  task_id, event_type, action, amount, currency, pricing_snapshot, payload, status, next_attempt_at
+)
+SELECT id, 'task.billing.release', 'release', reservation_amount, billing_currency,
+       pricing_snapshot, $2::jsonb, 'pending', now()
+FROM gateway_tasks
+WHERE id = $1::uuid
+  AND run_mode = 'production'
+  AND gateway_user_id IS NOT NULL
+  AND reservation_amount > 0
+ON CONFLICT (task_id, event_type) DO NOTHING`, taskID, string(payloadJSON))
+		return err
+	})
+	if err != nil {
+		return GatewayTask{}, false, err
+	}
+	return task, changed, nil
+}
+
 func (s *Store) ListRecoverableAsyncTasks(ctx context.Context, limit int) ([]AsyncTaskQueueItem, error) {
 	if limit <= 0 {
 		limit = 500
@@ -0,0 +1,111 @@
+package store
+
+import (
+	"context"
+	"strings"
+
+	"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
+)
+
+// VolcesCompatibleTaskListFilter mirrors the supported filters of Ark's
+// ListContentsGenerationsTasks API. Task IDs are the gateway's public task
+// IDs, which are the IDs returned by the compatibility create endpoint.
+type VolcesCompatibleTaskListFilter struct {
+	CompatibilityMarker string
+	Status              string
+	Model               string
+	TaskIDs             []string
+	Page                int
+	PageSize            int
+}
+
+// ListVolcesCompatibleTasks returns only video tasks created through a named
+// compatibility surface. Keeping this query separate from ListTasks avoids
+// broadening the ordinary task-list API's filtering semantics.
+func (s *Store) ListVolcesCompatibleTasks(ctx context.Context, user *auth.User, filter VolcesCompatibleTaskListFilter) (TaskListResult, error) {
+	page := filter.Page
+	if page < 1 {
+		page = 1
+	}
+	if page > 500 {
+		page = 500
+	}
+	pageSize := filter.PageSize
+	if pageSize < 1 {
+		pageSize = 20
+	}
+	if pageSize > 500 {
+		pageSize = 500
+	}
+	gatewayUserID := localGatewayUserID(user)
+	userID, apiKeyID := "", ""
+	if user != nil {
+		userID = strings.TrimSpace(user.ID)
+		apiKeyID = strings.TrimSpace(user.APIKeyID)
+	}
+	if gatewayUserID == "" && userID == "" {
+		return TaskListResult{}, ErrLocalUserRequired
+	}
+	taskIDs := make([]string, 0, len(filter.TaskIDs))
+	seen := make(map[string]bool, len(filter.TaskIDs))
+	for _, taskID := range filter.TaskIDs {
+		taskID = strings.TrimSpace(taskID)
+		if taskID != "" && !seen[taskID] {
+			seen[taskID] = true
+			taskIDs = append(taskIDs, taskID)
+		}
+	}
+	args := []any{
+		gatewayUserID,
+		userID,
+		apiKeyID,
+		strings.TrimSpace(filter.CompatibilityMarker),
+		strings.ToLower(strings.TrimSpace(filter.Status)),
+		strings.TrimSpace(filter.Model),
+		taskIDs,
+	}
+	whereSQL := `
+WHERE (
+    (
+      NULLIF($1, '')::uuid IS NOT NULL
+      AND gateway_user_id = NULLIF($1, '')::uuid
+    )
+    OR (
+      NULLIF($1, '')::uuid IS NULL
+      AND NULLIF($2, '') IS NOT NULL
+      AND user_id = $2
+    )
+  )
+  AND (NULLIF($3, '') IS NULL OR api_key_id = $3)
+  AND kind = 'videos.generations'
+  AND request->>'_gateway_compatibility' = $4
+  AND (NULLIF($5, '') IS NULL OR LOWER(status) = $5)
+  AND (NULLIF($6, '') IS NULL OR model = $6 OR resolved_model = $6)
+  AND (COALESCE(array_length($7::text[], 1), 0) = 0 OR id::text = ANY($7::text[]))`
+	var total int
+	if err := s.pool.QueryRow(ctx, `SELECT count(*) FROM gateway_tasks `+whereSQL, args...).Scan(&total); err != nil {
+		return TaskListResult{}, err
+	}
+	rows, err := s.pool.Query(ctx, `
+SELECT `+gatewayTaskColumns+`
+FROM gateway_tasks
+`+whereSQL+`
+ORDER BY created_at DESC
+LIMIT $8 OFFSET $9`, append(args, pageSize, (page-1)*pageSize)...)
+	if err != nil {
+		return TaskListResult{}, err
+	}
+	defer rows.Close()
+	items := make([]GatewayTask, 0)
+	for rows.Next() {
+		task, err := scanGatewayTask(rows)
+		if err != nil {
+			return TaskListResult{}, err
+		}
+		items = append(items, task)
+	}
+	if err := rows.Err(); err != nil {
+		return TaskListResult{}, err
+	}
+	return TaskListResult{Items: items, Total: total, Page: page, PageSize: pageSize}, nil
+}