fix gateway loopback validation chains
This commit is contained in:
@@ -3,9 +3,67 @@ package store
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgconn"
|
||||
)
|
||||
|
||||
func (s *Store) ListTasks(ctx context.Context, user *auth.User, limit int) ([]GatewayTask, error) {
|
||||
if limit <= 0 {
|
||||
limit = 50
|
||||
}
|
||||
if limit > 100 {
|
||||
limit = 100
|
||||
}
|
||||
gatewayUserID := localGatewayUserID(user)
|
||||
apiKeyID := ""
|
||||
userID := ""
|
||||
if user != nil {
|
||||
apiKeyID = strings.TrimSpace(user.APIKeyID)
|
||||
userID = strings.TrimSpace(user.ID)
|
||||
}
|
||||
if gatewayUserID == "" && userID == "" {
|
||||
return nil, ErrLocalUserRequired
|
||||
}
|
||||
rows, err := s.pool.Query(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
|
||||
)
|
||||
ORDER BY created_at DESC
|
||||
LIMIT $4`, gatewayUserID, userID, apiKeyID, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
items := make([]GatewayTask, 0)
|
||||
for rows.Next() {
|
||||
task, err := scanGatewayTask(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, task)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
|
||||
func (s *Store) MarkTaskRunning(ctx context.Context, taskID string, modelType string, normalizedRequest map[string]any) error {
|
||||
normalizedJSON, _ := json.Marshal(emptyObjectIfNil(normalizedRequest))
|
||||
_, err := s.pool.Exec(ctx, `
|
||||
@@ -140,6 +198,103 @@ WHERE id = $1::uuid`,
|
||||
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"
|
||||
}
|
||||
metadata, _ := json.Marshal(map[string]any{
|
||||
"taskId": task.ID,
|
||||
"kind": task.Kind,
|
||||
"model": task.Model,
|
||||
"resolvedModel": task.ResolvedModel,
|
||||
"billings": task.Billings,
|
||||
"billingSummary": task.BillingSummary,
|
||||
})
|
||||
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
|
||||
}
|
||||
var exists bool
|
||||
if err := tx.QueryRow(ctx, `
|
||||
SELECT EXISTS (
|
||||
SELECT 1
|
||||
FROM gateway_wallet_transactions t
|
||||
JOIN gateway_wallet_accounts a ON a.id = t.account_id
|
||||
WHERE a.gateway_user_id = $1::uuid
|
||||
AND a.currency = $2
|
||||
AND t.idempotency_key = $3
|
||||
)`, task.GatewayUserID, currency, billingIdempotencyKey(task.ID)).Scan(&exists); err != nil {
|
||||
return err
|
||||
}
|
||||
if exists {
|
||||
return nil
|
||||
}
|
||||
var accountID string
|
||||
var balanceBefore float64
|
||||
var gatewayTenantID string
|
||||
if err := tx.QueryRow(ctx, `
|
||||
SELECT id::text, 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, &gatewayTenantID); err != nil {
|
||||
return err
|
||||
}
|
||||
amount := roundMoney(task.FinalChargeAmount)
|
||||
balanceAfter := roundMoney(balanceBefore - amount)
|
||||
if _, err := tx.Exec(ctx, `
|
||||
UPDATE gateway_wallet_accounts
|
||||
SET balance = $2,
|
||||
total_spent = total_spent + $3,
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid`, accountID, balanceAfter, amount); err != nil {
|
||||
return err
|
||||
}
|
||||
_, 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))
|
||||
if pgErr, ok := err.(*pgconn.PgError); ok && pgErr.Code == "23505" {
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
})
|
||||
}
|
||||
|
||||
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 text
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (s *Store) FinishTaskFailure(ctx context.Context, input FinishTaskFailureInput) (GatewayTask, error) {
|
||||
metricsJSON, _ := json.Marshal(emptyObjectIfNil(input.Metrics))
|
||||
if _, err := s.pool.Exec(ctx, `
|
||||
|
||||
Reference in New Issue
Block a user