BCAI/easyai-ai-gateway
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: add ai gateway local core flow

Browse Source
This commit is contained in:
wangbo
2026-05-09 16:51:28 +08:00
parent 5b20f017eb
commit c0335bd5d0
20 changed files with 2332 additions and 156 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/internal/store/postgres.go
+603 -73
View File
@@ -2,6 +2,8 @@ package store
import (
"context"
"crypto/rand"
"encoding/base64"
"encoding/json"
"errors"
"strings"
@@ -10,6 +12,7 @@ import (
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgconn"
"github.com/jackc/pgx/v5/pgxpool"
"golang.org/x/crypto/bcrypt"
)
@@ -21,6 +24,8 @@ type Store struct {
var (
ErrInvalidCredentials = errors.New("invalid account or password")
ErrInvalidInvitation = errors.New("invalid or expired invitation code")
ErrLocalUserRequired = errors.New("local gateway user is required")
ErrUserAlreadyExists = errors.New("user already exists")
ErrWeakPassword = errors.New("password must be at least 8 characters")
)
@@ -73,6 +78,37 @@ type CreatePlatformInput struct {
Priority int `json:"priority"`
}
type CreateAPIKeyInput struct {
Name string `json:"name"`
Scopes []string `json:"scopes"`
ExpiresAt string `json:"expiresAt"`
}
type APIKey struct {
ID string `json:"id"`
GatewayTenantID string `json:"gatewayTenantId,omitempty"`
GatewayUserID string `json:"gatewayUserId"`
TenantID string `json:"tenantId,omitempty"`
TenantKey string `json:"tenantKey,omitempty"`
UserID string `json:"userId,omitempty"`
KeyPrefix string `json:"keyPrefix"`
Name string `json:"name"`
Scopes []string `json:"scopes,omitempty"`
UserGroupID string `json:"userGroupId,omitempty"`
RateLimitPolicy map[string]any `json:"rateLimitPolicy,omitempty"`
QuotaPolicy map[string]any `json:"quotaPolicy,omitempty"`
Status string `json:"status"`
ExpiresAt string `json:"expiresAt,omitempty"`
LastUsedAt string `json:"lastUsedAt,omitempty"`
CreatedAt time.Time `json:"createdAt"`
UpdatedAt time.Time `json:"updatedAt"`
}
type CreatedAPIKey struct {
APIKey APIKey `json:"apiKey"`
Secret string `json:"secret"`
}
type PlatformModel struct {
ID string `json:"id"`
PlatformID string `json:"platformId"`
@@ -161,8 +197,6 @@ type LocalRegisterInput struct {
Email string `json:"email"`
Password string `json:"password"`
DisplayName string `json:"displayName"`
TenantKey string `json:"tenantKey"`
TenantName string `json:"tenantName"`
InvitationCode string `json:"invitationCode"`
}
@@ -228,12 +262,14 @@ type RateLimitWindow struct {
type CreateTaskInput struct {
Kind string `json:"kind"`
Model string `json:"model"`
RunMode string `json:"runMode"`
Request map[string]any `json:"request"`
}
type GatewayTask struct {
ID string `json:"id"`
Kind string `json:"kind"`
RunMode string `json:"runMode"`
UserID string `json:"userId"`
GatewayUserID string `json:"gatewayUserId,omitempty"`
UserSource string `json:"userSource,omitempty"`
@@ -252,6 +288,20 @@ type GatewayTask struct {
UpdatedAt time.Time `json:"updatedAt"`
}
type TaskEvent struct {
ID string `json:"id"`
TaskID string `json:"taskId"`
Seq int64 `json:"seq"`
EventType string `json:"eventType"`
Status string `json:"status,omitempty"`
Phase string `json:"phase,omitempty"`
Progress float64 `json:"progress,omitempty"`
Message string `json:"message,omitempty"`
Payload map[string]any `json:"payload,omitempty"`
Simulated bool `json:"simulated"`
CreatedAt time.Time `json:"createdAt"`
}
func (s *Store) ListPlatforms(ctx context.Context) ([]Platform, error) {
rows, err := s.pool.Query(ctx, `
SELECT id::text, provider, platform_key, name, COALESCE(base_url, ''), auth_type, status, priority,
@@ -656,6 +706,216 @@ ORDER BY priority ASC, group_key ASC`)
return items, rows.Err()
}
func (s *Store) ListAPIKeys(ctx context.Context, user *auth.User) ([]APIKey, error) {
gatewayUserID := localGatewayUserID(user)
if gatewayUserID == "" {
return []APIKey{}, nil
}
rows, err := s.pool.Query(ctx, `
SELECT id::text, COALESCE(gateway_tenant_id::text, ''), gateway_user_id::text,
COALESCE(tenant_id, ''), COALESCE(tenant_key, ''), COALESCE(user_id, ''),
key_prefix, name, scopes, COALESCE(user_group_id::text, ''),
rate_limit_policy, quota_policy, status, COALESCE(expires_at::text, ''),
COALESCE(last_used_at::text, ''), created_at, updated_at
FROM gateway_api_keys
WHERE gateway_user_id = $1::uuid AND deleted_at IS NULL
ORDER BY created_at DESC`, gatewayUserID)
if err != nil {
return nil, err
}
defer rows.Close()
items := make([]APIKey, 0)
for rows.Next() {
item, err := scanAPIKey(rows)
if err != nil {
return nil, err
}
items = append(items, item)
}
return items, rows.Err()
}
func (s *Store) CreateAPIKey(ctx context.Context, input CreateAPIKeyInput, user *auth.User) (CreatedAPIKey, error) {
gatewayUserID := localGatewayUserID(user)
if gatewayUserID == "" {
return CreatedAPIKey{}, ErrLocalUserRequired
}
name := strings.TrimSpace(input.Name)
if name == "" {
name = "Default API Key"
}
scopes := input.Scopes
if len(scopes) == 0 {
scopes = []string{"chat", "image", "video"}
}
secret, err := generateAPIKeySecret()
if err != nil {
return CreatedAPIKey{}, err
}
keyHash, err := bcrypt.GenerateFromPassword([]byte(secret), bcrypt.DefaultCost)
if err != nil {
return CreatedAPIKey{}, err
}
scopesJSON, err := json.Marshal(scopes)
if err != nil {
return CreatedAPIKey{}, err
}
var item APIKey
var scopesBytes []byte
var rateLimitPolicy []byte
var quotaPolicy []byte
err = s.pool.QueryRow(ctx, `
INSERT INTO gateway_api_keys (
gateway_tenant_id, gateway_user_id, tenant_id, tenant_key, user_id,
key_prefix, key_hash, name, scopes, expires_at
)
VALUES (NULLIF($1, '')::uuid, $2::uuid, NULLIF($3, ''), NULLIF($4, ''), NULLIF($5, ''),
$6, $7, $8, $9::jsonb, NULLIF($10, '')::timestamptz)
RETURNING id::text, COALESCE(gateway_tenant_id::text, ''), gateway_user_id::text,
COALESCE(tenant_id, ''), COALESCE(tenant_key, ''), COALESCE(user_id, ''),
key_prefix, name, scopes, COALESCE(user_group_id::text, ''),
rate_limit_policy, quota_policy, status, COALESCE(expires_at::text, ''),
COALESCE(last_used_at::text, ''), created_at, updated_at`,
user.GatewayTenantID, gatewayUserID, user.TenantID, user.TenantKey, user.ID,
apiKeyPrefix(secret), string(keyHash), name, string(scopesJSON), strings.TrimSpace(input.ExpiresAt),
).Scan(
&item.ID,
&item.GatewayTenantID,
&item.GatewayUserID,
&item.TenantID,
&item.TenantKey,
&item.UserID,
&item.KeyPrefix,
&item.Name,
&scopesBytes,
&item.UserGroupID,
&rateLimitPolicy,
&quotaPolicy,
&item.Status,
&item.ExpiresAt,
&item.LastUsedAt,
&item.CreatedAt,
&item.UpdatedAt,
)
if err != nil {
return CreatedAPIKey{}, err
}
item.Scopes = decodeStringArray(scopesBytes)
item.RateLimitPolicy = decodeObject(rateLimitPolicy)
item.QuotaPolicy = decodeObject(quotaPolicy)
return CreatedAPIKey{APIKey: item, Secret: secret}, nil
}
func (s *Store) DisableAPIKey(ctx context.Context, apiKeyID string, user *auth.User) (APIKey, error) {
gatewayUserID := localGatewayUserID(user)
if gatewayUserID == "" {
return APIKey{}, ErrLocalUserRequired
}
var item APIKey
var scopesBytes []byte
var rateLimitPolicy []byte
var quotaPolicy []byte
err := s.pool.QueryRow(ctx, `
UPDATE gateway_api_keys
SET status = 'disabled', updated_at = now()
WHERE id = $1::uuid AND gateway_user_id = $2::uuid AND deleted_at IS NULL
RETURNING id::text, COALESCE(gateway_tenant_id::text, ''), gateway_user_id::text,
COALESCE(tenant_id, ''), COALESCE(tenant_key, ''), COALESCE(user_id, ''),
key_prefix, name, scopes, COALESCE(user_group_id::text, ''),
rate_limit_policy, quota_policy, status, COALESCE(expires_at::text, ''),
COALESCE(last_used_at::text, ''), created_at, updated_at`,
apiKeyID, gatewayUserID,
).Scan(
&item.ID,
&item.GatewayTenantID,
&item.GatewayUserID,
&item.TenantID,
&item.TenantKey,
&item.UserID,
&item.KeyPrefix,
&item.Name,
&scopesBytes,
&item.UserGroupID,
&rateLimitPolicy,
&quotaPolicy,
&item.Status,
&item.ExpiresAt,
&item.LastUsedAt,
&item.CreatedAt,
&item.UpdatedAt,
)
if err != nil {
return APIKey{}, err
}
item.Scopes = decodeStringArray(scopesBytes)
item.RateLimitPolicy = decodeObject(rateLimitPolicy)
item.QuotaPolicy = decodeObject(quotaPolicy)
return item, nil
}
func (s *Store) VerifyLocalAPIKey(ctx context.Context, secret string) (*auth.User, error) {
prefix := apiKeyPrefix(secret)
if prefix == "" {
return nil, auth.ErrUnauthorized
}
rows, err := s.pool.Query(ctx, `
SELECT k.id::text, k.key_hash, COALESCE(k.user_group_id::text, ''),
u.id::text, u.username, u.roles, COALESCE(u.gateway_tenant_id::text, ''),
COALESCE(u.tenant_id, ''), COALESCE(u.tenant_key, '')
FROM gateway_api_keys k
JOIN gateway_users u ON u.id = k.gateway_user_id
WHERE k.key_prefix = $1
AND k.status = 'active'
AND k.deleted_at IS NULL
AND u.status = 'active'
AND u.deleted_at IS NULL
AND (k.expires_at IS NULL OR k.expires_at > now())`, prefix)
if err != nil {
return nil, err
}
defer rows.Close()
for rows.Next() {
var apiKeyID string
var hash string
var userGroupID string
var gatewayUserID string
var username string
var rolesBytes []byte
var gatewayTenantID string
var tenantID string
var tenantKey string
if err := rows.Scan(&apiKeyID, &hash, &userGroupID, &gatewayUserID, &username, &rolesBytes, &gatewayTenantID, &tenantID, &tenantKey); err != nil {
return nil, err
}
if bcrypt.CompareHashAndPassword([]byte(hash), []byte(secret)) != nil {
continue
}
if _, err := s.pool.Exec(ctx, `UPDATE gateway_api_keys SET last_used_at = now(), updated_at = now() WHERE id = $1::uuid`, apiKeyID); err != nil {
return nil, err
}
return &auth.User{
ID: gatewayUserID,
Username: username,
Roles: decodeStringArray(rolesBytes),
TenantID: tenantID,
GatewayTenantID: gatewayTenantID,
TenantKey: tenantKey,
Source: "gateway",
GatewayUserID: gatewayUserID,
UserGroupID: userGroupID,
APIKeyID: apiKeyID,
APIKeyName: prefix,
}, nil
}
if err := rows.Err(); err != nil {
return nil, err
}
return nil, auth.ErrUnauthorized
}
func (s *Store) RegisterLocalUser(ctx context.Context, input LocalRegisterInput) (GatewayUser, error) {
account := normalizeAccount(firstNonEmpty(input.Username, input.Email))
if account == "" {
@@ -664,14 +924,8 @@ func (s *Store) RegisterLocalUser(ctx context.Context, input LocalRegisterInput)
if len(input.Password) < 8 {
return GatewayUser{}, ErrWeakPassword
}
tenantKey := normalizeKey(input.TenantKey)
if tenantKey == "" {
tenantKey = "personal-" + normalizeKey(account)
}
tenantName := strings.TrimSpace(input.TenantName)
if tenantName == "" {
tenantName = tenantKey
}
tenantKey := "default"
tenantName := "Default Tenant"
displayName := strings.TrimSpace(input.DisplayName)
username := strings.TrimSpace(input.Username)
if username == "" {
@@ -695,61 +949,79 @@ func (s *Store) RegisterLocalUser(ctx context.Context, input LocalRegisterInput)
userGroupID := ""
role := "user"
invitationID := ""
invitedBy := ""
var isBootstrapUser bool
if err := tx.QueryRow(ctx, `
SELECT NOT EXISTS (
SELECT 1 FROM gateway_users WHERE source = 'gateway' AND deleted_at IS NULL
)`).Scan(&isBootstrapUser); err != nil {
return GatewayUser{}, err
}
if isBootstrapUser {
role = "admin"
}
if err := tx.QueryRow(ctx, `
INSERT INTO gateway_tenants (tenant_key, source, external_tenant_id, name)
VALUES ($1, 'gateway', $1, $2)
ON CONFLICT (tenant_key) DO UPDATE SET updated_at=now()
RETURNING id::text`,
tenantKey, tenantName,
).Scan(&tenantID); err != nil {
return GatewayUser{}, err
}
if invitationCode != "" {
if err := tx.QueryRow(ctx, `
SELECT i.id::text,
i.tenant_id::text,
t.tenant_key,
t.name,
COALESCE(i.user_group_id::text, t.default_user_group_id::text, ''),
COALESCE(NULLIF(i.role, ''), 'user')
FROM gateway_tenant_invitations i
JOIN gateway_tenants t ON t.id = i.tenant_id
SELECT i.id::text, COALESCE(i.created_by::text, '')
FROM gateway_invitations i
WHERE lower(i.invite_code) = lower($1)
AND i.status = 'active'
AND t.status = 'active'
AND (i.expires_at IS NULL OR i.expires_at > now())
AND (i.max_uses IS NULL OR i.used_count < i.max_uses)
FOR UPDATE OF i`,
invitationCode,
).Scan(&invitationID, &tenantID, &tenantKey, &tenantName, &userGroupID, &role); err != nil {
).Scan(&invitationID, &invitedBy); err != nil {
if errors.Is(err, pgx.ErrNoRows) {
return GatewayUser{}, ErrInvalidInvitation
}
return GatewayUser{}, err
}
} else if err := tx.QueryRow(ctx, `
INSERT INTO gateway_tenants (tenant_key, source, external_tenant_id, name)
VALUES ($1, 'gateway', $1, $2)
ON CONFLICT (tenant_key) DO UPDATE SET updated_at=now()
RETURNING id::text`,
tenantKey, tenantName,
).Scan(&tenantID); err != nil {
return GatewayUser{}, err
}
rolesJSON, err := json.Marshal([]string{role})
if err != nil {
return GatewayUser{}, err
}
userMetadataJSON := []byte("{}")
if invitationID != "" {
userMetadataJSON, err = json.Marshal(map[string]any{
"registration": map[string]any{
"invitationId": invitationID,
"invitationCode": invitationCode,
"invitedBy": invitedBy,
},
})
if err != nil {
return GatewayUser{}, err
}
}
var user GatewayUser
var roles []byte
var authProfile []byte
var metadata []byte
if err := tx.QueryRow(ctx, `
INSERT INTO gateway_users (
user_key, source, external_user_id, username, display_name, email,
password_hash, gateway_tenant_id, tenant_id, tenant_key, default_user_group_id, roles, status
)
VALUES ($1, 'gateway', $2, $3, NULLIF($4, ''), NULLIF($5, ''), $6, $7::uuid, $8, $8, NULLIF($9, '')::uuid, $10::jsonb, 'active')
RETURNING id::text, user_key, source, COALESCE(external_user_id, ''), username,
COALESCE(display_name, ''), COALESCE(email, ''), COALESCE(phone, ''), COALESCE(avatar_url, ''),
COALESCE(gateway_tenant_id::text, ''), COALESCE(tenant_id, ''), COALESCE(tenant_key, ''),
COALESCE(default_user_group_id::text, ''), roles, auth_profile, metadata,
status, COALESCE(last_login_at::text, ''), COALESCE(synced_at::text, ''), COALESCE(source_updated_at::text, ''),
created_at, updated_at`,
"gateway:"+account, account, username, displayName, email, string(passwordHash), tenantID, tenantKey, userGroupID, string(rolesJSON),
INSERT INTO gateway_users (
user_key, source, external_user_id, username, display_name, email,
password_hash, gateway_tenant_id, tenant_id, tenant_key, default_user_group_id, roles, metadata, status
)
VALUES ($1, 'gateway', $2, $3, NULLIF($4, ''), NULLIF($5, ''), $6, $7::uuid, $8, $8, NULLIF($9, '')::uuid, $10::jsonb, $11::jsonb, 'active')
RETURNING id::text, user_key, source, COALESCE(external_user_id, ''), username,
COALESCE(display_name, ''), COALESCE(email, ''), COALESCE(phone, ''), COALESCE(avatar_url, ''),
COALESCE(gateway_tenant_id::text, ''), COALESCE(tenant_id, ''), COALESCE(tenant_key, ''),
COALESCE(default_user_group_id::text, ''), roles, auth_profile, metadata,
status, COALESCE(last_login_at::text, ''), COALESCE(synced_at::text, ''), COALESCE(source_updated_at::text, ''),
created_at, updated_at`,
"gateway:"+account, account, username, displayName, email, string(passwordHash), tenantID, tenantKey, userGroupID, string(rolesJSON), string(userMetadataJSON),
).Scan(
&user.ID,
&user.UserKey,
@@ -774,33 +1046,28 @@ FOR UPDATE OF i`,
&user.CreatedAt,
&user.UpdatedAt,
); err != nil {
if isUniqueViolation(err) {
return GatewayUser{}, ErrUserAlreadyExists
}
return GatewayUser{}, err
}
if invitationID != "" {
if _, err := tx.Exec(ctx, `
UPDATE gateway_tenant_invitations
UPDATE gateway_invitations
SET used_count = used_count + 1, updated_at = now()
WHERE id = $1::uuid`, invitationID); err != nil {
return GatewayUser{}, err
}
}
if userGroupID != "" {
metadata, err := json.Marshal(map[string]any{
"source": "registration",
"invitationId": invitationID,
})
if err != nil {
return GatewayUser{}, err
}
if _, err := tx.Exec(ctx, `
INSERT INTO gateway_user_group_memberships (group_id, principal_type, principal_id, source, metadata)
VALUES ($1::uuid, 'user', $2, 'gateway', $3::jsonb)
ON CONFLICT (group_id, principal_type, principal_id)
DO UPDATE SET status = 'active', updated_at = now()`,
userGroupID, user.ID, string(metadata),
); err != nil {
return GatewayUser{}, err
}
if _, err := tx.Exec(ctx, `
INSERT INTO gateway_wallet_accounts (
gateway_tenant_id, gateway_user_id, tenant_id, tenant_key, user_id, currency
)
VALUES ($1::uuid, $2::uuid, $3, $3, $4, 'resource')
ON CONFLICT (gateway_user_id, currency) DO NOTHING`,
user.GatewayTenantID, user.ID, user.TenantKey, user.ID,
); err != nil {
return GatewayUser{}, err
}
if err := tx.Commit(ctx); err != nil {
return GatewayUser{}, err
@@ -913,24 +1180,58 @@ ORDER BY window_start DESC, scope_type ASC, scope_key ASC, metric ASC`)
func (s *Store) CreateTask(ctx context.Context, input CreateTaskInput, user *auth.User) (GatewayTask, error) {
requestBody, _ := json.Marshal(input.Request)
runMode := normalizeRunMode(input.RunMode, input.Request)
status := "queued"
result := map[string]any(nil)
billings := []any(nil)
finished := false
if runMode == "simulation" {
status = "succeeded"
result = simulationResult(input.Kind, input.Model)
billings = simulationBillings(input.Kind, input.Model)
finished = true
}
resultBody, _ := json.Marshal(result)
billingsBody, _ := json.Marshal(billings)
tx, err := s.pool.Begin(ctx)
if err != nil {
return GatewayTask{}, err
}
defer tx.Rollback(ctx)
var task GatewayTask
var requestBytes []byte
var resultBytes []byte
var billingsBytes []byte
err := s.pool.QueryRow(ctx, `
INSERT INTO gateway_tasks (
kind, user_id, gateway_user_id, user_source, gateway_tenant_id, tenant_id, tenant_key,
api_key_id, user_group_id, user_group_key, model, request, status
)
VALUES ($1, $2, NULLIF($3, '')::uuid, COALESCE(NULLIF($4, ''), 'gateway'), NULLIF($5, '')::uuid, NULLIF($6, ''), NULLIF($7, ''), NULLIF($8, ''), NULLIF($9, '')::uuid, NULLIF($10, ''), $11, $12, 'queued')
RETURNING id::text, kind, user_id, COALESCE(gateway_user_id::text, ''), user_source,
COALESCE(gateway_tenant_id::text, ''), COALESCE(tenant_id, ''), COALESCE(tenant_key, ''),
COALESCE(user_group_id::text, ''), COALESCE(user_group_key, ''), model, request, status, result, billings, COALESCE(error, ''), created_at, updated_at`,
input.Kind, user.ID, user.GatewayUserID, user.Source, user.GatewayTenantID, user.TenantID, user.TenantKey, user.APIKeyID, user.UserGroupID, user.UserGroupKey, input.Model, requestBody,
).Scan(&task.ID, &task.Kind, &task.UserID, &task.GatewayUserID, &task.UserSource, &task.GatewayTenantID, &task.TenantID, &task.TenantKey, &task.UserGroupID, &task.UserGroupKey, &task.Model, &requestBytes, &task.Status, &resultBytes, &billingsBytes, &task.Error, &task.CreatedAt, &task.UpdatedAt)
err = tx.QueryRow(ctx, `
INSERT INTO gateway_tasks (
kind, run_mode, user_id, gateway_user_id, user_source, gateway_tenant_id, tenant_id, tenant_key,
api_key_id, user_group_id, user_group_key, model, request, status, result, billings, finished_at
)
VALUES ($1, $2, $3, NULLIF($4, '')::uuid, COALESCE(NULLIF($5, ''), 'gateway'), NULLIF($6, '')::uuid, NULLIF($7, ''), NULLIF($8, ''), NULLIF($9, ''), NULLIF($10, '')::uuid, NULLIF($11, ''), $12, $13, $14, $15::jsonb, $16::jsonb, CASE WHEN $17 THEN now() ELSE NULL END)
RETURNING id::text, kind, run_mode, user_id, COALESCE(gateway_user_id::text, ''), user_source,
COALESCE(gateway_tenant_id::text, ''), COALESCE(tenant_id, ''), COALESCE(tenant_key, ''),
COALESCE(user_group_id::text, ''), COALESCE(user_group_key, ''), model, request, status, result, billings, COALESCE(error, ''), created_at, updated_at`,
input.Kind, runMode, user.ID, user.GatewayUserID, user.Source, user.GatewayTenantID, user.TenantID, user.TenantKey, user.APIKeyID, user.UserGroupID, user.UserGroupKey, input.Model, requestBody, status, resultBody, billingsBody, finished,
).Scan(&task.ID, &task.Kind, &task.RunMode, &task.UserID, &task.GatewayUserID, &task.UserSource, &task.GatewayTenantID, &task.TenantID, &task.TenantKey, &task.UserGroupID, &task.UserGroupKey, &task.Model, &requestBytes, &task.Status, &resultBytes, &billingsBytes, &task.Error, &task.CreatedAt, &task.UpdatedAt)
if err != nil {
return GatewayTask{}, err
}
events := taskEventsForCreate(task.ID, runMode, status, result)
for _, event := range events {
payload, _ := json.Marshal(event.Payload)
if _, err := tx.Exec(ctx, `
INSERT INTO gateway_task_events (task_id, seq, event_type, status, phase, progress, message, payload, simulated)
VALUES ($1::uuid, $2, $3, NULLIF($4, ''), NULLIF($5, ''), $6, NULLIF($7, ''), $8::jsonb, $9)`,
task.ID, event.Seq, event.EventType, event.Status, event.Phase, event.Progress, event.Message, string(payload), event.Simulated,
); err != nil {
return GatewayTask{}, err
}
}
if err := tx.Commit(ctx); err != nil {
return GatewayTask{}, err
}
task.Request = decodeObject(requestBytes)
task.Result = decodeObject(resultBytes)
task.Billings = decodeArray(billingsBytes)
@@ -943,12 +1244,12 @@ func (s *Store) GetTask(ctx context.Context, taskID string) (GatewayTask, error)
var resultBytes []byte
var billingsBytes []byte
err := s.pool.QueryRow(ctx, `
SELECT id::text, kind, user_id, COALESCE(gateway_user_id::text, ''), user_source,
SELECT id::text, kind, run_mode, user_id, COALESCE(gateway_user_id::text, ''), user_source,
COALESCE(gateway_tenant_id::text, ''), COALESCE(tenant_id, ''), COALESCE(tenant_key, ''),
COALESCE(user_group_id::text, ''), COALESCE(user_group_key, ''), model, request, status, result, billings, COALESCE(error, ''), created_at, updated_at
FROM gateway_tasks
WHERE id=$1`, taskID,
).Scan(&task.ID, &task.Kind, &task.UserID, &task.GatewayUserID, &task.UserSource, &task.GatewayTenantID, &task.TenantID, &task.TenantKey, &task.UserGroupID, &task.UserGroupKey, &task.Model, &requestBytes, &task.Status, &resultBytes, &billingsBytes, &task.Error, &task.CreatedAt, &task.UpdatedAt)
FROM gateway_tasks
WHERE id=$1`, taskID,
).Scan(&task.ID, &task.Kind, &task.RunMode, &task.UserID, &task.GatewayUserID, &task.UserSource, &task.GatewayTenantID, &task.TenantID, &task.TenantKey, &task.UserGroupID, &task.UserGroupKey, &task.Model, &requestBytes, &task.Status, &resultBytes, &billingsBytes, &task.Error, &task.CreatedAt, &task.UpdatedAt)
if err != nil {
return GatewayTask{}, err
}
@@ -958,10 +1259,239 @@ WHERE id=$1`, taskID,
return task, nil
}
func (s *Store) ListTaskEvents(ctx context.Context, taskID string) ([]TaskEvent, error) {
rows, err := s.pool.Query(ctx, `
SELECT id::text, task_id::text, seq, event_type, COALESCE(status, ''), COALESCE(phase, ''),
COALESCE(progress, 0)::float8, COALESCE(message, ''), payload, simulated, created_at
FROM gateway_task_events
WHERE task_id = $1::uuid
ORDER BY seq ASC`, taskID)
if err != nil {
return nil, err
}
defer rows.Close()
items := make([]TaskEvent, 0)
for rows.Next() {
var item TaskEvent
var payload []byte
if err := rows.Scan(
&item.ID,
&item.TaskID,
&item.Seq,
&item.EventType,
&item.Status,
&item.Phase,
&item.Progress,
&item.Message,
&payload,
&item.Simulated,
&item.CreatedAt,
); err != nil {
return nil, err
}
item.Payload = decodeObject(payload)
items = append(items, item)
}
return items, rows.Err()
}
func IsNotFound(err error) bool {
return err == pgx.ErrNoRows
}
func isUniqueViolation(err error) bool {
var pgErr *pgconn.PgError
return errors.As(err, &pgErr) && pgErr.Code == "23505"
}
type apiKeyScanner interface {
Scan(dest ...any) error
}
func scanAPIKey(scanner apiKeyScanner) (APIKey, error) {
var item APIKey
var scopesBytes []byte
var rateLimitPolicy []byte
var quotaPolicy []byte
if err := scanner.Scan(
&item.ID,
&item.GatewayTenantID,
&item.GatewayUserID,
&item.TenantID,
&item.TenantKey,
&item.UserID,
&item.KeyPrefix,
&item.Name,
&scopesBytes,
&item.UserGroupID,
&rateLimitPolicy,
&quotaPolicy,
&item.Status,
&item.ExpiresAt,
&item.LastUsedAt,
&item.CreatedAt,
&item.UpdatedAt,
); err != nil {
return APIKey{}, err
}
item.Scopes = decodeStringArray(scopesBytes)
item.RateLimitPolicy = decodeObject(rateLimitPolicy)
item.QuotaPolicy = decodeObject(quotaPolicy)
return item, nil
}
func localGatewayUserID(user *auth.User) string {
if user == nil {
return ""
}
if user.GatewayUserID != "" {
return user.GatewayUserID
}
if user.Source == "" || user.Source == "gateway" {
return user.ID
}
return ""
}
func generateAPIKeySecret() (string, error) {
bytes := make([]byte, 32)
if _, err := rand.Read(bytes); err != nil {
return "", err
}
return "sk-gw-" + base64.RawURLEncoding.EncodeToString(bytes), nil
}
func apiKeyPrefix(secret string) string {
secret = strings.TrimSpace(secret)
if !strings.HasPrefix(secret, "sk-gw-") {
return ""
}
if len(secret) <= 18 {
return secret
}
return secret[:18]
}
func normalizeRunMode(input string, request map[string]any) string {
mode := strings.ToLower(strings.TrimSpace(input))
if mode == "" && request != nil {
if raw, ok := request["runMode"].(string); ok {
mode = strings.ToLower(strings.TrimSpace(raw))
}
if raw, ok := request["mode"].(string); ok && mode == "" {
mode = strings.ToLower(strings.TrimSpace(raw))
}
if raw, ok := request["simulation"].(bool); ok && raw {
mode = "simulation"
}
if raw, ok := request["testMode"].(bool); ok && raw {
mode = "simulation"
}
}
if mode == "test" || mode == "dry-run" || mode == "dry_run" {
return "simulation"
}
if mode == "simulation" {
return "simulation"
}
return "production"
}
func simulationResult(kind string, model string) map[string]any {
switch kind {
case "chat.completions":
return map[string]any{
"id": "chatcmpl-simulated",
"object": "chat.completion",
"model": model,
"choices": []any{map[string]any{"index": 0, "finish_reason": "stop", "message": map[string]any{"role": "assistant", "content": "simulation response"}}},
"usage": map[string]any{"prompt_tokens": 12, "completion_tokens": 8, "total_tokens": 20},
}
case "images.generations":
return map[string]any{
"id": "img-simulated",
"model": model,
"data": []any{map[string]any{"url": "/static/simulation/image.png", "revised_prompt": "simulation image"}},
}
case "videos.generations":
return map[string]any{
"id": "video-simulated",
"model": model,
"data": []any{map[string]any{"url": "/static/simulation/video.mp4", "duration": 5}},
}
default:
return map[string]any{"id": "task-simulated", "model": model, "kind": kind, "ok": true}
}
}
func simulationBillings(kind string, model string) []any {
resourceType := "task"
unit := "item"
if kind == "chat.completions" {
resourceType = "text_total"
unit = "1k_tokens"
}
if kind == "images.generations" {
resourceType = "image"
unit = "image"
}
if kind == "videos.generations" {
resourceType = "video"
unit = "second"
}
return []any{map[string]any{
"model": model,
"resourceType": resourceType,
"unit": unit,
"quantity": 1,
"amount": 0,
"currency": "resource",
"simulated": true,
}}
}
func taskEventsForCreate(taskID string, runMode string, status string, result map[string]any) []TaskEvent {
events := []TaskEvent{{
TaskID: taskID,
Seq: 1,
EventType: "task.accepted",
Status: "queued",
Phase: "queued",
Progress: 0,
Message: "task accepted",
Payload: map[string]any{"taskId": taskID},
Simulated: runMode == "simulation",
}}
if runMode != "simulation" {
return events
}
return append(events,
TaskEvent{
TaskID: taskID,
Seq: 2,
EventType: "task.progress",
Status: "running",
Phase: "simulation",
Progress: 0.5,
Message: "simulation client running",
Payload: map[string]any{"taskId": taskID},
Simulated: true,
},
TaskEvent{
TaskID: taskID,
Seq: 3,
EventType: "task.completed",
Status: status,
Phase: "completed",
Progress: 1,
Message: "simulation completed",
Payload: map[string]any{"taskId": taskID, "result": result},
Simulated: true,
},
)
}
func decodeObject(bytes []byte) map[string]any {
if len(bytes) == 0 {
return nil