easyai-ai-gateway/apps/api/internal/store/volces_compatible_tasks.go

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[])
	  OR remote_task_id = 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
	}
	items, err = s.attachTaskAttempts(ctx, items)
	if err != nil {
		return TaskListResult{}, err
	}
	return TaskListResult{Items: items, Total: total, Page: page, PageSize: pageSize}, nil
}

// GetVolcesCompatibleTask resolves both the gateway UUID used for converted
// tasks and the upstream Ark task ID exposed by native passthrough tasks.
func (s *Store) GetVolcesCompatibleTask(ctx context.Context, user *auth.User, compatibilityMarker string, identifier string) (GatewayTask, error) {
	gatewayUserID := localGatewayUserID(user)
	userID, apiKeyID := "", ""
	if user != nil {
		userID = strings.TrimSpace(user.ID)
		apiKeyID = strings.TrimSpace(user.APIKeyID)
	}
	if gatewayUserID == "" && userID == "" {
		return GatewayTask{}, ErrLocalUserRequired
	}
	task, err := scanGatewayTask(s.pool.QueryRow(ctx, `
SELECT `+gatewayTaskColumns+`
FROM gateway_tasks
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 (id::text = $5 OR remote_task_id = $5)
ORDER BY created_at DESC
LIMIT 1`, gatewayUserID, userID, apiKeyID, strings.TrimSpace(compatibilityMarker), strings.TrimSpace(identifier)))
	if err != nil {
		return GatewayTask{}, err
	}
	attempts, err := s.ListTaskAttempts(ctx, task.ID)
	if err != nil {
		return GatewayTask{}, err
	}
	task.Attempts = attempts
	return task, nil
}