chengcheng
8beb8501fa
ci / verify (pull_request) Failing after 8s
阻止上游提交状态不明的任务被租约接管后重复执行，并为人工复核保留可操作的结算记录。 将生产提交绑定到当前估价签名，统一复用预处理快照，并补强规则形状、定点溢出与历史规则兼容校验。 已通过 PostgreSQL 16 集成测试、Go 全量测试与静态检查、前端测试与构建、OpenAPI、依赖审计、镜像、迁移、流水线和 SemVer 门禁。
1384 lines
43 KiB
Go
1384 lines
43 KiB
Go
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
|
|
}
|