package store import ( "context" "encoding/json" "errors" "fmt" "strconv" "strings" "time" "github.com/easyai/easyai-ai-gateway/apps/api/internal/auth" "github.com/jackc/pgx/v5" ) type TaskListFilter struct { Query string ModelType string CreatedFrom *time.Time CreatedTo *time.Time Page int PageSize int } type TaskListResult struct { Items []GatewayTask Total int Page int PageSize int } func (s *Store) ListTasks(ctx context.Context, user *auth.User, filter TaskListFilter) (TaskListResult, error) { page := filter.Page if page <= 0 { page = 1 } pageSize := filter.PageSize if pageSize <= 0 { pageSize = 50 } if pageSize > 100 { pageSize = 100 } offset := (page - 1) * pageSize gatewayUserID := localGatewayUserID(user) apiKeyID := "" userID := "" if user != nil { apiKeyID = strings.TrimSpace(user.APIKeyID) userID = strings.TrimSpace(user.ID) } if gatewayUserID == "" && userID == "" { return TaskListResult{}, ErrLocalUserRequired } queryPattern := "" if query := strings.TrimSpace(filter.Query); query != "" { queryPattern = "%" + query + "%" } args := []any{ gatewayUserID, userID, apiKeyID, queryPattern, strings.TrimSpace(filter.ModelType), nullableTaskListTime(filter.CreatedFrom), nullableTaskListTime(filter.CreatedTo), } 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 ( NULLIF($4, '') IS NULL OR id::text ILIKE $4 OR COALESCE(request_id, '') ILIKE $4 OR kind ILIKE $4 OR model ILIKE $4 OR COALESCE(requested_model, '') ILIKE $4 OR COALESCE(resolved_model, '') ILIKE $4 OR COALESCE(api_key_id, '') ILIKE $4 OR COALESCE(api_key_name, '') ILIKE $4 OR COALESCE(api_key_prefix, '') ILIKE $4 OR status ILIKE $4 OR COALESCE(model_type, '') ILIKE $4 ) AND ( NULLIF($5, '') IS NULL OR model_type = $5 ) AND ( $6::timestamptz IS NULL OR created_at >= $6::timestamptz ) AND ( $7::timestamptz IS NULL OR created_at <= $7::timestamptz )` var total int if err := s.pool.QueryRow(ctx, `SELECT count(*) FROM gateway_tasks `+whereSQL, args...).Scan(&total); err != nil { return TaskListResult{}, err } queryArgs := append(args, pageSize, offset) rows, err := s.pool.Query(ctx, ` SELECT `+gatewayTaskColumns+` FROM gateway_tasks `+whereSQL+` ORDER BY created_at DESC LIMIT $8 OFFSET $9`, queryArgs...) 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 } func nullableTaskListTime(value *time.Time) any { if value == nil { return nil } return *value } func (s *Store) ClaimTaskExecution(ctx context.Context, taskID string, executionToken string, leaseTTL time.Duration) (GatewayTask, error) { if leaseTTL <= 0 { leaseTTL = 5 * time.Minute } var task GatewayTask manualReview := false err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error { var queuedReady bool var runningExpired bool var production bool var hasGatewayUser bool if err := tx.QueryRow(ctx, ` SELECT status = 'queued' AND next_run_at <= now(), status = 'running' AND (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, &runningExpired, &production, &hasGatewayUser); err != nil { return err } 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 { 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 = 'upstream submission result is unknown', 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, finished_at = now(), updated_at = now() WHERE id = $1::uuid`, 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 } } var err error task, err = scanGatewayTask(tx.QueryRow(ctx, ` UPDATE gateway_tasks SET status = 'running', execution_token = $2::uuid, execution_lease_expires_at = now() + ($3::int * interval '1 second'), locked_at = now(), heartbeat_at = now(), updated_at = now() WHERE id = $1::uuid AND ( (status = 'queued' AND next_run_at <= now()) OR (status = 'running' AND (execution_lease_expires_at IS NULL OR execution_lease_expires_at <= now())) ) RETURNING `+gatewayTaskColumns, taskID, executionToken, int(leaseTTL/time.Second))) return err }) if err != nil { if errors.Is(err, pgx.ErrNoRows) { return GatewayTask{}, ErrTaskExecutionLeaseUnavailable } return GatewayTask{}, err } if manualReview { return GatewayTask{}, ErrTaskExecutionManualReview } return task, nil } func (s *Store) RenewTaskExecutionLease(ctx context.Context, taskID string, executionToken string, leaseTTL time.Duration) error { if leaseTTL <= 0 { leaseTTL = 5 * time.Minute } tag, err := s.pool.Exec(ctx, ` UPDATE gateway_tasks SET execution_lease_expires_at = now() + ($3::int * interval '1 second'), heartbeat_at = now(), updated_at = now() WHERE id = $1::uuid AND status = 'running' AND execution_token = $2::uuid`, taskID, executionToken, int(leaseTTL/time.Second)) if err != nil { return err } if tag.RowsAffected() != 1 { var stillRunning bool if err := s.pool.QueryRow(ctx, ` SELECT COALESCE((SELECT status = 'running' FROM gateway_tasks WHERE id = $1::uuid), false)`, taskID).Scan(&stillRunning); err != nil { return err } if !stillRunning { return ErrTaskExecutionFinished } return ErrTaskExecutionLeaseLost } return nil } func (s *Store) MarkTaskRunning(ctx context.Context, taskID string, executionToken string, modelType string, normalizedRequest map[string]any) error { normalizedJSON, _ := json.Marshal(emptyObjectIfNil(normalizedRequest)) tag, err := s.pool.Exec(ctx, ` UPDATE gateway_tasks SET status = 'running', model_type = NULLIF($3::text, ''), normalized_request = $4::jsonb, locked_at = now(), heartbeat_at = now(), updated_at = now() WHERE id = $1::uuid AND status = 'running' AND execution_token = $2::uuid`, taskID, executionToken, modelType, string(normalizedJSON)) if err != nil { return err } if tag.RowsAffected() != 1 { return ErrTaskExecutionLeaseLost } return nil } func (s *Store) ClaimAsyncQueuedTask(ctx context.Context, workerID string) (GatewayTask, error) { return scanGatewayTask(s.pool.QueryRow(ctx, ` WITH picked AS ( SELECT id AS task_id FROM gateway_tasks WHERE async_mode = true AND status = 'queued' AND next_run_at <= now() ORDER BY priority ASC, created_at ASC LIMIT 1 FOR UPDATE SKIP LOCKED ) UPDATE gateway_tasks t SET status = 'running', locked_by = NULLIF($1::text, ''), locked_at = now(), heartbeat_at = now(), updated_at = now() FROM picked WHERE t.id = picked.task_id RETURNING `+gatewayTaskColumns, workerID)) } func (s *Store) RequeueTask(ctx context.Context, taskID string, executionToken string, delay time.Duration, queueKey string) (GatewayTask, error) { if delay < time.Second { delay = time.Second } if delay > 10*time.Minute { delay = 10 * time.Minute } nextRunAt := time.Now().Add(delay) return scanGatewayTask(s.pool.QueryRow(ctx, ` UPDATE gateway_tasks SET status = 'queued', locked_by = NULL, locked_at = NULL, heartbeat_at = NULL, execution_token = NULL, execution_lease_expires_at = NULL, next_run_at = $3::timestamptz, queue_key = COALESCE(NULLIF($4::text, ''), queue_key), error = NULL, error_code = NULL, error_message = NULL, updated_at = now() WHERE id = $1::uuid AND status = 'running' AND execution_token = $2::uuid RETURNING `+gatewayTaskColumns, taskID, executionToken, nextRunAt, strings.TrimSpace(queueKey))) } func (s *Store) SetTaskRiverJobID(ctx context.Context, taskID string, riverJobID int64) error { if riverJobID <= 0 { return nil } _, err := s.pool.Exec(ctx, ` UPDATE gateway_tasks SET river_job_id = $2, updated_at = now() WHERE id = $1::uuid`, taskID, riverJobID) return err } func (s *Store) SetTaskRemoteTask(ctx context.Context, taskID string, executionToken string, attemptID string, remoteTaskID string, payload map[string]any) error { payloadJSON, _ := json.Marshal(emptyObjectIfNil(payload)) return pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error { tag, err := tx.Exec(ctx, ` UPDATE gateway_tasks SET remote_task_id = NULLIF($3::text, ''), remote_task_payload = $4::jsonb, updated_at = now() WHERE id = $1::uuid AND status = 'running' AND execution_token = $2::uuid`, taskID, executionToken, remoteTaskID, string(payloadJSON), ) if err != nil { return err } if tag.RowsAffected() != 1 { return ErrTaskExecutionLeaseLost } if strings.TrimSpace(attemptID) == "" { return nil } _, err = tx.Exec(ctx, ` 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`, attemptID, remoteTaskID, string(payloadJSON), ) return err }) } func (s *Store) SetAttemptUpstreamSubmissionStatus(ctx context.Context, attemptID string, status string) error { switch status { case "not_submitted", "submitting", "response_received": default: return fmt.Errorf("invalid upstream submission status %q", status) } _, err := s.pool.Exec(ctx, ` UPDATE gateway_task_attempts SET upstream_submission_status = $2, upstream_submission_updated_at = now() WHERE id = $1::uuid`, attemptID, status) return err } func (s *Store) CancelQueuedTask(ctx context.Context, taskID string, message string) (GatewayTask, bool, error) { message = strings.TrimSpace(message) if message == "" { message = "任务已取消" } tag, err := s.pool.Exec(ctx, ` UPDATE gateway_tasks SET status = 'cancelled', error = NULLIF($2::text, ''), error_code = 'task_cancelled', error_message = NULLIF($2::text, ''), locked_by = NULL, locked_at = NULL, heartbeat_at = NULL, finished_at = now(), updated_at = now() WHERE id = $1::uuid AND status = 'queued' AND COALESCE(remote_task_id, '') = ''`, taskID, message) if err != nil { return GatewayTask{}, false, err } if tag.RowsAffected() == 0 { return GatewayTask{}, false, nil } task, err := s.GetTask(ctx, taskID) if err != nil { return GatewayTask{}, true, err } return task, true, nil } func (s *Store) ListRecoverableAsyncTasks(ctx context.Context, limit int) ([]AsyncTaskQueueItem, error) { if limit <= 0 { limit = 500 } rows, err := s.pool.Query(ctx, ` SELECT id::text, priority, next_run_at FROM gateway_tasks WHERE async_mode = true AND ( status = 'queued' OR (status = 'running' AND (execution_lease_expires_at IS NULL OR execution_lease_expires_at <= now())) ) ORDER BY priority ASC, created_at ASC LIMIT $1`, limit) if err != nil { return nil, err } defer rows.Close() items := make([]AsyncTaskQueueItem, 0) for rows.Next() { var item AsyncTaskQueueItem if err := rows.Scan(&item.ID, &item.Priority, &item.NextRunAt); err != nil { return nil, err } items = append(items, item) } if err := rows.Err(); err != nil { return nil, err } return items, nil } func (s *Store) CreateTaskAttempt(ctx context.Context, input CreateTaskAttemptInput) (string, error) { requestJSON, _ := json.Marshal(emptyObjectIfNil(input.RequestSnapshot)) metricsJSON, _ := json.Marshal(emptyObjectIfNil(input.Metrics)) pricingJSON, _ := json.Marshal(emptyObjectIfNil(input.PricingSnapshot)) tx, err := s.pool.Begin(ctx) if err != nil { return "", err } defer tx.Rollback(ctx) var attemptID string err = tx.QueryRow(ctx, ` INSERT INTO gateway_task_attempts ( task_id, attempt_no, platform_id, platform_model_id, client_id, queue_key, status, simulated, request_snapshot, metrics, pricing_snapshot, request_fingerprint, upstream_submission_status, upstream_submission_updated_at ) VALUES ( $1::uuid, $2::int, NULLIF($3::text, '')::uuid, NULLIF($4::text, '')::uuid, NULLIF($5::text, ''), $6, $7, $8, $9::jsonb, $10::jsonb, $11::jsonb, NULLIF($12, ''), 'not_submitted', now() ) RETURNING id::text`, input.TaskID, input.AttemptNo, input.PlatformID, input.PlatformModelID, input.ClientID, input.QueueKey, firstNonEmpty(input.Status, "running"), input.Simulated, string(requestJSON), string(metricsJSON), string(pricingJSON), input.RequestFingerprint, ).Scan(&attemptID) if err != nil { return "", err } if _, err := tx.Exec(ctx, ` UPDATE gateway_tasks SET attempt_count = GREATEST(attempt_count, $2::int), updated_at = now() WHERE id = $1::uuid`, input.TaskID, input.AttemptNo); err != nil { return "", err } return attemptID, tx.Commit(ctx) } func (s *Store) CreateTaskParamPreprocessingLog(ctx context.Context, input CreateTaskParamPreprocessingLogInput) (string, error) { actualInputJSON, _ := json.Marshal(emptyObjectIfNil(input.ActualInput)) convertedOutputJSON, _ := json.Marshal(emptyObjectIfNil(input.ConvertedOutput)) changesJSON, _ := json.Marshal(input.Changes) if input.Changes == nil { changesJSON = []byte("[]") } modelSnapshotJSON, _ := json.Marshal(emptyObjectIfNil(input.ModelSnapshot)) var attemptNo any if input.AttemptNo > 0 { attemptNo = input.AttemptNo } var id string err := s.pool.QueryRow(ctx, ` INSERT INTO gateway_task_param_preprocessing_logs ( task_id, attempt_id, attempt_no, model_type, platform_id, platform_model_id, client_id, changed, change_count, actual_input, converted_output, changes, model_snapshot ) VALUES ( $1::uuid, NULLIF($2::text, '')::uuid, $3::int, NULLIF($4::text, ''), NULLIF($5::text, '')::uuid, NULLIF($6::text, '')::uuid, NULLIF($7::text, ''), $8, $9::int, $10::jsonb, $11::jsonb, $12::jsonb, $13::jsonb ) RETURNING id::text`, input.TaskID, input.AttemptID, attemptNo, input.ModelType, input.PlatformID, input.PlatformModelID, input.ClientID, input.Changed, input.ChangeCount, string(actualInputJSON), string(convertedOutputJSON), string(changesJSON), string(modelSnapshotJSON), ).Scan(&id) return id, err } func (s *Store) attachTaskAttempts(ctx context.Context, items []GatewayTask) ([]GatewayTask, error) { if len(items) == 0 { return items, nil } taskIDs := make([]string, 0, len(items)) for _, item := range items { taskIDs = append(taskIDs, item.ID) } attemptsByTaskID, err := s.listTaskAttemptsByTaskIDs(ctx, taskIDs) if err != nil { return nil, err } for index := range items { items[index].Attempts = attemptsByTaskID[items[index].ID] } return items, nil } func (s *Store) ListTaskAttempts(ctx context.Context, taskID string) ([]TaskAttempt, error) { attemptsByTaskID, err := s.listTaskAttemptsByTaskIDs(ctx, []string{taskID}) if err != nil { return nil, err } return attemptsByTaskID[taskID], nil } func (s *Store) ListTaskParamPreprocessingLogs(ctx context.Context, taskID string) ([]TaskParamPreprocessingLog, error) { rows, err := s.pool.Query(ctx, ` SELECT id::text, task_id::text, COALESCE(attempt_id::text, ''), COALESCE(attempt_no, 0), COALESCE(model_type, ''), COALESCE(platform_id::text, ''), COALESCE(platform_model_id::text, ''), COALESCE(client_id, ''), changed, change_count, COALESCE(actual_input, '{}'::jsonb), COALESCE(converted_output, '{}'::jsonb), COALESCE(changes, '[]'::jsonb), COALESCE(model_snapshot, '{}'::jsonb), created_at FROM gateway_task_param_preprocessing_logs WHERE task_id = $1::uuid ORDER BY COALESCE(attempt_no, 0), created_at`, taskID) if err != nil { return nil, err } defer rows.Close() items := make([]TaskParamPreprocessingLog, 0) for rows.Next() { item, err := scanTaskParamPreprocessingLog(rows) if err != nil { return nil, err } items = append(items, item) } return items, rows.Err() } func (s *Store) AppendTaskAttemptTrace(ctx context.Context, taskID string, attemptNo int, entry map[string]any) error { entryJSON, _ := json.Marshal(emptyObjectIfNil(entry)) _, err := s.pool.Exec(ctx, ` UPDATE gateway_task_attempts SET metrics = jsonb_set( COALESCE(metrics, '{}'::jsonb), '{trace}', ( CASE WHEN jsonb_typeof(COALESCE(metrics->'trace', '[]'::jsonb)) = 'array' THEN COALESCE(metrics->'trace', '[]'::jsonb) ELSE '[]'::jsonb END ) || jsonb_build_array($3::jsonb), true ) WHERE task_id = $1::uuid AND attempt_no = $2::int`, taskID, attemptNo, string(entryJSON)) return err } func (s *Store) listTaskAttemptsByTaskIDs(ctx context.Context, taskIDs []string) (map[string][]TaskAttempt, error) { itemsByTaskID := map[string][]TaskAttempt{} if len(taskIDs) == 0 { return itemsByTaskID, nil } rows, err := s.pool.Query(ctx, ` SELECT a.id::text, a.task_id::text, a.attempt_no, COALESCE(a.platform_id::text, ''), COALESCE(p.name, ''), COALESCE(p.provider, ''), COALESCE(a.platform_model_id::text, ''), COALESCE(pm.model_name, ''), COALESCE(NULLIF(pm.provider_model_name, ''), pm.model_name, ''), COALESCE(pm.model_alias, ''), COALESCE(a.client_id, ''), a.queue_key, a.status, a.retryable, a.simulated, COALESCE(a.request_id, ''), COALESCE(a.usage, '{}'::jsonb), COALESCE(a.metrics, '{}'::jsonb), a.request_snapshot, COALESCE(a.response_snapshot, '{}'::jsonb), COALESCE(a.response_started_at::text, ''), COALESCE(a.response_finished_at::text, ''), COALESCE(a.response_duration_ms, 0), COALESCE(a.error_code, ''), COALESCE(a.error_message, ''), COALESCE(a.pricing_snapshot, '{}'::jsonb), COALESCE(a.request_fingerprint, ''), a.upstream_submission_status, COALESCE(a.upstream_submission_updated_at::text, ''), a.started_at, COALESCE(a.finished_at::text, '') FROM gateway_task_attempts a LEFT JOIN integration_platforms p ON p.id = a.platform_id LEFT JOIN platform_models pm ON pm.id = a.platform_model_id WHERE a.task_id::text = ANY($1) ORDER BY a.task_id, a.attempt_no`, taskIDs) if err != nil { return nil, err } defer rows.Close() for rows.Next() { item, err := scanTaskAttempt(rows) if err != nil { return nil, err } itemsByTaskID[item.TaskID] = append(itemsByTaskID[item.TaskID], item) } if err := rows.Err(); err != nil { return nil, err } return itemsByTaskID, nil } func scanTaskAttempt(scanner taskScanner) (TaskAttempt, error) { var item TaskAttempt var usageBytes []byte var metricsBytes []byte var requestBytes []byte var responseBytes []byte var pricingSnapshotBytes []byte if err := scanner.Scan( &item.ID, &item.TaskID, &item.AttemptNo, &item.PlatformID, &item.PlatformName, &item.Provider, &item.PlatformModelID, &item.ModelName, &item.ProviderModelName, &item.ModelAlias, &item.ClientID, &item.QueueKey, &item.Status, &item.Retryable, &item.Simulated, &item.RequestID, &usageBytes, &metricsBytes, &requestBytes, &responseBytes, &item.ResponseStartedAt, &item.ResponseFinishedAt, &item.ResponseDurationMS, &item.ErrorCode, &item.ErrorMessage, &pricingSnapshotBytes, &item.RequestFingerprint, &item.UpstreamSubmissionStatus, &item.UpstreamSubmissionUpdatedAt, &item.StartedAt, &item.FinishedAt, ); err != nil { return TaskAttempt{}, err } item.Usage = decodeObject(usageBytes) item.Metrics = decodeObject(metricsBytes) item.RequestSnapshot = decodeObject(requestBytes) item.ResponseSnapshot = decodeObject(responseBytes) item.PricingSnapshot = decodeObject(pricingSnapshotBytes) enrichTaskAttemptFromMetrics(&item) return item, nil } func scanTaskParamPreprocessingLog(scanner taskScanner) (TaskParamPreprocessingLog, error) { var item TaskParamPreprocessingLog var actualInputBytes []byte var convertedOutputBytes []byte var changesBytes []byte var modelSnapshotBytes []byte if err := scanner.Scan( &item.ID, &item.TaskID, &item.AttemptID, &item.AttemptNo, &item.ModelType, &item.PlatformID, &item.PlatformModelID, &item.ClientID, &item.Changed, &item.ChangeCount, &actualInputBytes, &convertedOutputBytes, &changesBytes, &modelSnapshotBytes, &item.CreatedAt, ); err != nil { return TaskParamPreprocessingLog{}, err } item.ActualInput = decodeObject(actualInputBytes) item.ConvertedOutput = decodeObject(convertedOutputBytes) item.Changes = decodeArray(changesBytes) item.ModelSnapshot = decodeObject(modelSnapshotBytes) return item, nil } func enrichTaskAttemptFromMetrics(item *TaskAttempt) { if item == nil || len(item.Metrics) == 0 { return } item.PlatformID = firstNonEmpty(item.PlatformID, taskAttemptMetricString(item.Metrics, "platformId")) item.PlatformName = firstNonEmpty(item.PlatformName, taskAttemptMetricString(item.Metrics, "platformName")) item.Provider = firstNonEmpty(item.Provider, taskAttemptMetricString(item.Metrics, "provider")) item.PlatformModelID = firstNonEmpty(item.PlatformModelID, taskAttemptMetricString(item.Metrics, "platformModelId")) item.ModelName = firstNonEmpty(item.ModelName, taskAttemptMetricString(item.Metrics, "resolvedModel"), taskAttemptMetricString(item.Metrics, "modelName")) item.ProviderModelName = firstNonEmpty(item.ProviderModelName, taskAttemptMetricString(item.Metrics, "providerModel")) item.ModelAlias = firstNonEmpty(item.ModelAlias, taskAttemptMetricString(item.Metrics, "modelAlias")) item.ModelType = firstNonEmpty(item.ModelType, taskAttemptMetricString(item.Metrics, "modelType")) item.ClientID = firstNonEmpty(item.ClientID, taskAttemptMetricString(item.Metrics, "clientId")) item.StatusCode = taskAttemptMetricInt(item.Metrics, "statusCode") } func taskAttemptMetricString(metrics map[string]any, key string) string { value, _ := metrics[key].(string) return strings.TrimSpace(value) } func taskAttemptMetricInt(metrics map[string]any, key string) int { switch value := metrics[key].(type) { case int: return value case int64: return int(value) case float64: return int(value) case json.Number: next, _ := value.Int64() return int(next) case string: next, _ := strconv.Atoi(strings.TrimSpace(value)) return next default: return 0 } } func (s *Store) FinishTaskAttempt(ctx context.Context, input FinishTaskAttemptInput) error { responseJSON, _ := json.Marshal(emptyObjectIfNil(input.ResponseSnapshot)) usageJSON, _ := json.Marshal(emptyObjectIfNil(input.Usage)) metricsJSON, _ := json.Marshal(emptyObjectIfNil(input.Metrics)) _, err := s.pool.Exec(ctx, ` UPDATE gateway_task_attempts SET status = $2::text, retryable = $3, request_id = NULLIF($4::text, ''), usage = $5::jsonb, metrics = $6::jsonb, response_snapshot = $7::jsonb, response_started_at = $8::timestamptz, response_finished_at = $9::timestamptz, response_duration_ms = $10, error_code = NULLIF($11::text, ''), error_message = NULLIF($12::text, ''), finished_at = now() WHERE id = $1::uuid`, input.AttemptID, input.Status, input.Retryable, input.RequestID, string(usageJSON), string(metricsJSON), string(responseJSON), nullableTime(input.ResponseStartedAt), nullableTime(input.ResponseFinishedAt), input.ResponseDurationMS, input.ErrorCode, input.ErrorMessage, ) return err } func (s *Store) FinishTaskSuccess(ctx context.Context, input FinishTaskSuccessInput) (GatewayTask, error) { resultJSON, _ := json.Marshal(emptyObjectIfNil(input.Result)) billingsJSON, _ := json.Marshal(input.Billings) usageJSON, _ := json.Marshal(emptyObjectIfNil(input.Usage)) metricsJSON, _ := json.Marshal(emptyObjectIfNil(input.Metrics)) billingSummaryJSON, _ := json.Marshal(emptyObjectIfNil(input.BillingSummary)) pricingSnapshotJSON, _ := json.Marshal(emptyObjectIfNil(input.PricingSnapshot)) finalChargeAmount := strings.TrimSpace(input.FinalChargeAmountText) if finalChargeAmount == "" { finalChargeAmount = strconv.FormatFloat(input.FinalChargeAmount, 'f', 9, 64) } currency := normalizeWalletCurrency(input.BillingCurrency) err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error { if strings.TrimSpace(input.AttemptID) != "" { if _, err := tx.Exec(ctx, ` UPDATE gateway_task_attempts SET status = 'succeeded', retryable = false, request_id = NULLIF($2, ''), usage = $3::jsonb, metrics = $4::jsonb, response_snapshot = $5::jsonb, pricing_snapshot = $6::jsonb, request_fingerprint = NULLIF($7, ''), upstream_submission_status = 'response_received', upstream_submission_updated_at = now(), response_started_at = $8::timestamptz, response_finished_at = $9::timestamptz, response_duration_ms = $10, error_code = NULL, error_message = NULL, finished_at = now() WHERE id = $1::uuid`, input.AttemptID, input.RequestID, string(usageJSON), string(metricsJSON), string(resultJSON), string(pricingSnapshotJSON), input.RequestFingerprint, nullableTime(input.ResponseStartedAt), nullableTime(input.ResponseFinishedAt), input.ResponseDurationMS, ); err != nil { return err } } tag, err := tx.Exec(ctx, ` UPDATE gateway_tasks SET status = 'succeeded', result = $2::jsonb, billings = $3::jsonb, request_id = NULLIF($4, ''), resolved_model = NULLIF($5, ''), usage = $6::jsonb, metrics = $7::jsonb, billing_summary = $8::jsonb, final_charge_amount = $9::numeric, billing_version = 'effective-pricing-v2', billing_status = CASE WHEN run_mode = 'production' AND gateway_user_id IS NOT NULL THEN 'pending' ELSE 'not_required' END, billing_currency = $10, pricing_snapshot = $11::jsonb, request_fingerprint = NULLIF($12, ''), billing_updated_at = now(), response_started_at = $13::timestamptz, response_finished_at = $14::timestamptz, response_duration_ms = $15, error = NULL, error_code = NULL, error_message = NULL, 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 status = 'running' AND execution_token = $16::uuid`, input.TaskID, string(resultJSON), string(billingsJSON), input.RequestID, input.ResolvedModel, string(usageJSON), string(metricsJSON), string(billingSummaryJSON), finalChargeAmount, currency, string(pricingSnapshotJSON), input.RequestFingerprint, nullableTime(input.ResponseStartedAt), nullableTime(input.ResponseFinishedAt), input.ResponseDurationMS, input.ExecutionToken, ) if err != nil { return err } if tag.RowsAffected() != 1 { return ErrTaskExecutionLeaseLost } payloadJSON, _ := json.Marshal(map[string]any{"taskId": input.TaskID, "pricingVersion": "effective-pricing-v2"}) _, 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.settle', 'settle', $2::numeric, $3, $4::jsonb, $5::jsonb, 'pending', now() FROM gateway_tasks WHERE id = $1::uuid AND run_mode = 'production' AND gateway_user_id IS NOT NULL ON CONFLICT (task_id, event_type) DO NOTHING`, input.TaskID, finalChargeAmount, currency, string(pricingSnapshotJSON), string(payloadJSON)) return err }) if err != nil { return GatewayTask{}, err } return s.GetTask(ctx, input.TaskID) } func (s *Store) FinishTaskManualReview(ctx context.Context, input FinishTaskManualReviewInput) (GatewayTask, error) { status := strings.TrimSpace(input.TaskStatus) if status != "succeeded" && status != "failed" { return GatewayTask{}, fmt.Errorf("manual review task status must be succeeded or failed") } resultJSON, _ := json.Marshal(emptyObjectIfNil(input.Result)) pricingSnapshotJSON, _ := json.Marshal(emptyObjectIfNil(input.PricingSnapshot)) err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error { if strings.TrimSpace(input.AttemptID) != "" { attemptStatus := "failed" if status == "succeeded" { attemptStatus = "succeeded" } if _, err := tx.Exec(ctx, ` UPDATE gateway_task_attempts SET status = $2, request_id = NULLIF($3, ''), error_code = CASE WHEN $2 = 'failed' THEN NULLIF($4, '') ELSE NULL END, error_message = CASE WHEN $2 = 'failed' THEN NULLIF($5, '') ELSE NULL END, upstream_submission_status = CASE WHEN $2 = 'succeeded' THEN 'response_received' ELSE upstream_submission_status END, upstream_submission_updated_at = now(), response_started_at = $6::timestamptz, response_finished_at = $7::timestamptz, response_duration_ms = $8, finished_at = now() WHERE id = $1::uuid`, input.AttemptID, attemptStatus, input.RequestID, input.Code, input.Message, nullableTime(input.ResponseStartedAt), nullableTime(input.ResponseFinishedAt), input.ResponseDurationMS); err != nil { return err } } tag, err := tx.Exec(ctx, ` UPDATE gateway_tasks SET status = $2, result = $3::jsonb, request_id = NULLIF($4, ''), error = CASE WHEN $2 = 'failed' THEN NULLIF($6, '') ELSE NULL END, error_code = NULLIF($5, ''), error_message = NULLIF($6, ''), billing_status = 'manual_review', pricing_snapshot = CASE WHEN $7::jsonb = '{}'::jsonb THEN pricing_snapshot ELSE $7::jsonb END, request_fingerprint = COALESCE(NULLIF($8, ''), request_fingerprint), billing_updated_at = now(), response_started_at = $9::timestamptz, response_finished_at = $10::timestamptz, response_duration_ms = $11, 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 status = 'running' AND execution_token = $12::uuid`, input.TaskID, status, string(resultJSON), input.RequestID, input.Code, input.Message, string(pricingSnapshotJSON), input.RequestFingerprint, nullableTime(input.ResponseStartedAt), nullableTime(input.ResponseFinishedAt), input.ResponseDurationMS, input.ExecutionToken) if err != nil { return err } if tag.RowsAffected() != 1 { return ErrTaskExecutionLeaseLost } payloadJSON, _ := json.Marshal(map[string]any{ "taskId": input.TaskID, "classification": input.Code, }) _, 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(), $3 FROM gateway_tasks WHERE id = $1::uuid AND run_mode = 'production' AND gateway_user_id IS NOT NULL ON CONFLICT (task_id, event_type) DO NOTHING`, input.TaskID, string(payloadJSON), input.Code) return err }) if err != nil { return GatewayTask{}, err } return s.GetTask(ctx, input.TaskID) } func (s *Store) SettleTaskBilling(ctx context.Context, task GatewayTask) error { if task.FinalChargeAmount <= 0 || strings.TrimSpace(task.GatewayUserID) == "" { return nil } currency := strings.TrimSpace(taskBillingString(task.BillingSummary["currency"])) if currency == "" || currency == "mixed" { currency = "resource" } modelIdentity := taskBillingModelIdentity(task) metadataMap := map[string]any{ "taskId": task.ID, "kind": task.Kind, "model": modelIdentity.Model, "requestedModel": modelIdentity.RequestedModel, "resolvedModel": modelIdentity.ResolvedModel, "modelName": modelIdentity.ModelName, "modelAlias": modelIdentity.ModelAlias, "providerModel": modelIdentity.ProviderModelName, "billings": task.Billings, "billingSummary": task.BillingSummary, } metadata, _ := json.Marshal(metadataMap) return pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error { if _, err := tx.Exec(ctx, ` INSERT INTO gateway_wallet_accounts ( gateway_tenant_id, gateway_user_id, tenant_id, tenant_key, user_id, currency ) VALUES (NULLIF($1, '')::uuid, $2::uuid, NULLIF($3, ''), NULLIF($4, ''), NULLIF($5, ''), $6) ON CONFLICT (gateway_user_id, currency) DO NOTHING`, task.GatewayTenantID, task.GatewayUserID, task.TenantID, task.TenantKey, task.UserID, currency); err != nil { return err } if err := ensureWalletAccountAuditGuard(ctx, tx, task.GatewayUserID, currency); err != nil { return err } var exists bool var accountID string var balanceBefore float64 var frozenBefore float64 var gatewayTenantID string if err := tx.QueryRow(ctx, ` SELECT id::text, balance::float8, frozen_balance::float8, COALESCE(gateway_tenant_id::text, '') FROM gateway_wallet_accounts WHERE gateway_user_id = $1::uuid AND currency = $2 FOR UPDATE`, task.GatewayUserID, currency).Scan(&accountID, &balanceBefore, &frozenBefore, &gatewayTenantID); err != nil { return err } if err := tx.QueryRow(ctx, ` SELECT EXISTS ( SELECT 1 FROM gateway_wallet_transactions WHERE account_id = $1::uuid AND idempotency_key = $2 )`, accountID, billingIdempotencyKey(task.ID)).Scan(&exists); err != nil { return err } if exists { return nil } amount := roundMoney(task.FinalChargeAmount) reservationKey, reservedAmount, err := activeWalletReservation(ctx, tx, accountID, task.ID) if err != nil { return err } reservedAmount = roundMoney(reservedAmount) spendableForTask := roundMoney(balanceBefore - frozenBefore + reservedAmount) if spendableForTask+0.000001 < amount { return fmt.Errorf("%w: required %.6f %s, available %.6f", ErrInsufficientWalletBalance, amount, currency, spendableForTask) } balanceAfter := roundMoney(balanceBefore - amount) frozenAfter := roundMoney(frozenBefore - reservedAmount) if frozenAfter < 0 { frozenAfter = 0 } if _, err := tx.Exec(ctx, ` UPDATE gateway_wallet_accounts SET balance = $2, total_spent = total_spent + $3, frozen_balance = $4, updated_at = now() WHERE id = $1::uuid`, accountID, balanceAfter, amount, frozenAfter); err != nil { return err } if reservedAmount > 0 { releaseMetadata, _ := json.Marshal(map[string]any{ "taskId": task.ID, "reason": "task_billing_settled", "reserved": reservedAmount, "frozenBefore": roundMoney(frozenBefore), "frozenAfter": frozenAfter, }) if _, err := tx.Exec(ctx, ` INSERT INTO gateway_wallet_transactions ( account_id, gateway_tenant_id, gateway_user_id, direction, transaction_type, amount, balance_before, balance_after, idempotency_key, reference_type, reference_id, metadata ) VALUES ( $1::uuid, NULLIF($2, '')::uuid, $3::uuid, 'credit', 'release', $4, $5, $6, $7, 'gateway_task', $8, $9::jsonb ) ON CONFLICT (account_id, idempotency_key) WHERE idempotency_key IS NOT NULL DO NOTHING`, accountID, firstNonEmpty(gatewayTenantID, task.GatewayTenantID), task.GatewayUserID, reservedAmount, roundMoney(balanceBefore), roundMoney(balanceBefore), billingReservationReleaseIdempotencyKey(reservationKey), task.ID, string(releaseMetadata)); err != nil { return err } } billingMetadata := mergeObjects(metadataMap, map[string]any{ "reservedAmount": reservedAmount, "frozenBefore": roundMoney(frozenBefore), "frozenAfter": frozenAfter, }) metadata, _ = json.Marshal(billingMetadata) if _, err := tx.Exec(ctx, ` INSERT INTO gateway_wallet_transactions ( account_id, gateway_tenant_id, gateway_user_id, direction, transaction_type, amount, balance_before, balance_after, idempotency_key, reference_type, reference_id, metadata ) VALUES ( $1::uuid, NULLIF($2, '')::uuid, $3::uuid, 'debit', 'task_billing', $4, $5, $6, $7, 'gateway_task', $8, $9::jsonb )`, accountID, firstNonEmpty(gatewayTenantID, task.GatewayTenantID), task.GatewayUserID, amount, roundMoney(balanceBefore), balanceAfter, billingIdempotencyKey(task.ID), task.ID, string(metadata)); err != nil { return err } return nil }) } func billingIdempotencyKey(taskID string) string { return "task:" + taskID + ":billing" } func roundMoney(value float64) float64 { if value < 0 { return -roundMoney(-value) } return float64(int64(value*1000000+0.5)) / 1000000 } func taskBillingString(value any) string { if text, ok := value.(string); ok { return strings.TrimSpace(text) } return "" } func (s *Store) FinishTaskFailure(ctx context.Context, input FinishTaskFailureInput) (GatewayTask, error) { metricsJSON, _ := json.Marshal(emptyObjectIfNil(input.Metrics)) resultJSON, _ := json.Marshal(emptyObjectIfNil(input.Result)) err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error { tag, err := tx.Exec(ctx, ` UPDATE gateway_tasks SET status = 'failed', error = NULLIF($2::text, ''), error_code = NULLIF($3::text, ''), error_message = NULLIF($2::text, ''), request_id = NULLIF($4::text, ''), metrics = $5::jsonb, response_started_at = $6::timestamptz, response_finished_at = $7::timestamptz, response_duration_ms = $8, result = $9::jsonb, 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 status = 'running' AND execution_token = $10::uuid`, input.TaskID, input.Message, input.Code, input.RequestID, string(metricsJSON), nullableTime(input.ResponseStartedAt), nullableTime(input.ResponseFinishedAt), input.ResponseDurationMS, string(resultJSON), input.ExecutionToken, ) if err != nil { return err } if tag.RowsAffected() != 1 { return ErrTaskExecutionLeaseLost } payloadJSON, _ := json.Marshal(map[string]any{"taskId": input.TaskID, "reason": input.Code}) _, 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`, input.TaskID, string(payloadJSON)) return err }) if err != nil { return GatewayTask{}, err } return s.GetTask(ctx, input.TaskID) } func nullableTime(value time.Time) any { if value.IsZero() { return nil } return value } func (s *Store) AddTaskEvent(ctx context.Context, taskID string, eventType string, status string, phase string, progress float64, message string, payload map[string]any, simulated bool) (TaskEvent, error) { payloadJSON, _ := json.Marshal(emptyObjectIfNil(payload)) var event TaskEvent var payloadBytes []byte err := s.pool.QueryRow(ctx, ` WITH next_seq AS ( SELECT COALESCE(MAX(seq), 0) + 1 AS seq FROM gateway_task_events WHERE task_id = $1::uuid ) INSERT INTO gateway_task_events (task_id, seq, event_type, status, phase, progress, message, payload, simulated) SELECT $1::uuid, next_seq.seq, $2::text, NULLIF($3::text, ''), NULLIF($4::text, ''), $5, NULLIF($6::text, ''), $7::jsonb, $8 FROM next_seq RETURNING id::text, task_id::text, seq, event_type, COALESCE(status, ''), COALESCE(phase, ''), COALESCE(progress, 0)::float8, COALESCE(message, ''), payload, simulated, created_at`, taskID, eventType, status, phase, progress, message, string(payloadJSON), simulated, ).Scan( &event.ID, &event.TaskID, &event.Seq, &event.EventType, &event.Status, &event.Phase, &event.Progress, &event.Message, &payloadBytes, &event.Simulated, &event.CreatedAt, ) if err != nil { return TaskEvent{}, err } event.Payload = decodeObject(payloadBytes) return event, nil } func (s *Store) QueueTaskCallback(ctx context.Context, event TaskEvent, callbackURL string) error { if callbackURL == "" { return nil } payloadJSON, _ := json.Marshal(map[string]any{ "taskId": event.TaskID, "seq": event.Seq, "eventType": event.EventType, "status": event.Status, "phase": event.Phase, "progress": event.Progress, "message": event.Message, "payload": event.Payload, "simulated": event.Simulated, "createdAt": event.CreatedAt, }) _, err := s.pool.Exec(ctx, ` INSERT INTO gateway_task_callback_outbox (task_id, event_id, seq, callback_url, payload) VALUES ($1::uuid, $2::uuid, $3, $4, $5::jsonb) ON CONFLICT (task_id, seq, callback_url) DO NOTHING`, event.TaskID, event.ID, event.Seq, callbackURL, string(payloadJSON), ) return err } func (s *Store) RecordClientAssignment(ctx context.Context, candidate RuntimeModelCandidate) error { _, err := s.pool.Exec(ctx, ` INSERT INTO runtime_client_states ( client_id, platform_id, provider, method_name, queue_key, running_count, last_assigned_at ) VALUES ($1, $2::uuid, $3, $4, $5, 1, now()) ON CONFLICT (client_id) DO UPDATE SET running_count = runtime_client_states.running_count + 1, last_assigned_at = now(), updated_at = now()`, candidate.ClientID, candidate.PlatformID, candidate.Provider, candidate.ModelType, candidate.QueueKey, ) return err } func (s *Store) RecordClientRelease(ctx context.Context, clientID string, lastError string) error { _, err := s.pool.Exec(ctx, ` UPDATE runtime_client_states SET running_count = GREATEST(running_count - 1, 0), last_error = NULLIF($2::text, ''), updated_at = now() WHERE client_id = $1`, clientID, lastError) return err }