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

chore(git): 同步主分支发布与业务变更

Browse Source 
# Conflicts:
#	apps/api/go.mod
#	apps/api/go.sum
This commit is contained in:
chengcheng
2026-07-28 10:47:41 +08:00
parent 9dfd1aafa5 7c142f5960
commit 0b02e62c72
238 changed files with 25370 additions and 2079 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/internal/store/admin_task_redaction.go
+94
View File
@@ -0,0 +1,94 @@
package store
import "strings"
const maskedAdminTaskValue = "***"
func MaskAdminGatewayTask(task AdminGatewayTask) AdminGatewayTask {
task.Request = maskSensitiveObject(task.Request)
task.Result = maskSensitiveObject(task.Result)
task.Billings = maskSensitiveArray(task.Billings)
task.Usage = maskSensitiveObject(task.Usage)
task.Metrics = maskSensitiveObject(task.Metrics)
task.BillingSummary = maskSensitiveObject(task.BillingSummary)
task.PricingSnapshot = maskSensitiveObject(task.PricingSnapshot)
task.CompatibilitySubmitHeaders = maskSensitiveObject(task.CompatibilitySubmitHeaders)
task.CompatibilitySubmitBody = maskSensitiveObject(task.CompatibilitySubmitBody)
task.Attempts = append([]TaskAttempt(nil), task.Attempts...)
for index := range task.Attempts {
task.Attempts[index].Usage = maskSensitiveObject(task.Attempts[index].Usage)
task.Attempts[index].Metrics = maskSensitiveObject(task.Attempts[index].Metrics)
task.Attempts[index].RequestSnapshot = maskSensitiveObject(task.Attempts[index].RequestSnapshot)
task.Attempts[index].ResponseSnapshot = maskSensitiveObject(task.Attempts[index].ResponseSnapshot)
task.Attempts[index].PricingSnapshot = maskSensitiveObject(task.Attempts[index].PricingSnapshot)
}
return task
}
func MaskTaskParamPreprocessingLogs(items []TaskParamPreprocessingLog) []TaskParamPreprocessingLog {
masked := make([]TaskParamPreprocessingLog, len(items))
for index, item := range items {
item.ActualInput = maskSensitiveObject(item.ActualInput)
item.ConvertedOutput = maskSensitiveObject(item.ConvertedOutput)
item.Changes = maskSensitiveArray(item.Changes)
item.ModelSnapshot = maskSensitiveObject(item.ModelSnapshot)
masked[index] = item
}
return masked
}
func maskSensitiveObject(value map[string]any) map[string]any {
if value == nil {
return nil
}
masked, _ := maskSensitiveValue(value).(map[string]any)
return masked
}
func maskSensitiveArray(value []any) []any {
if value == nil {
return nil
}
masked, _ := maskSensitiveValue(value).([]any)
return masked
}
func maskSensitiveValue(value any) any {
switch typed := value.(type) {
case map[string]any:
masked := make(map[string]any, len(typed))
for key, item := range typed {
if adminTaskSensitiveKey(key) {
masked[key] = maskedAdminTaskValue
continue
}
masked[key] = maskSensitiveValue(item)
}
return masked
case []any:
masked := make([]any, len(typed))
for index, item := range typed {
masked[index] = maskSensitiveValue(item)
}
return masked
default:
return value
}
}
func adminTaskSensitiveKey(key string) bool {
normalized := strings.NewReplacer("_", "", "-", "", ".", "", " ", "").Replace(strings.ToLower(strings.TrimSpace(key)))
switch normalized {
case "authorization", "proxyauthorization", "apikey", "xapikey", "accesstoken", "refreshtoken",
"idtoken", "token", "secret", "clientsecret", "password", "passwd", "cookie", "setcookie",
"credentials", "credential", "privatekey":
return true
default:
return strings.HasSuffix(normalized, "apikey") ||
strings.HasSuffix(normalized, "accesstoken") ||
strings.HasSuffix(normalized, "refreshtoken") ||
strings.HasSuffix(normalized, "clientsecret") ||
strings.HasSuffix(normalized, "password") ||
strings.HasSuffix(normalized, "privatekey")
}
}
apps/api/internal/store/admin_task_redaction_test.go
+87
View File
@@ -0,0 +1,87 @@
package store
import (
"strings"
"testing"
)
func TestMaskAdminGatewayTaskRecursivelyMasksSecretsWithoutMutatingSource(t *testing.T) {
source := AdminGatewayTask{
GatewayTask: GatewayTask{
Request: map[string]any{
"model": "example",
"headers": map[string]any{
"Authorization": "Bearer private",
"X-Api-Key": "secret-key",
},
"usage": map[string]any{"input_tokens": float64(12)},
},
Result: map[string]any{
"nested": []any{map[string]any{"password": "private-password"}},
},
Attempts: []TaskAttempt{{
RequestSnapshot: map[string]any{"client_secret": "private-client-secret"},
}},
},
}
masked := MaskAdminGatewayTask(source)
headers := masked.Request["headers"].(map[string]any)
if headers["Authorization"] != maskedAdminTaskValue || headers["X-Api-Key"] != maskedAdminTaskValue {
t.Fatalf("sensitive headers were not masked: %#v", headers)
}
if masked.Request["usage"].(map[string]any)["input_tokens"] != float64(12) {
t.Fatalf("token usage must remain visible: %#v", masked.Request)
}
nested := masked.Result["nested"].([]any)[0].(map[string]any)
if nested["password"] != maskedAdminTaskValue {
t.Fatalf("nested password was not masked: %#v", nested)
}
if masked.Attempts[0].RequestSnapshot["client_secret"] != maskedAdminTaskValue {
t.Fatalf("attempt snapshot secret was not masked: %#v", masked.Attempts[0])
}
if source.Request["headers"].(map[string]any)["Authorization"] != "Bearer private" {
t.Fatal("masking mutated the source request")
}
if source.Attempts[0].RequestSnapshot["client_secret"] != "private-client-secret" {
t.Fatal("masking mutated the source attempts")
}
}
func TestBuildAdminTaskWhereUsesSharedPlaceholdersAndAllFilters(t *testing.T) {
where, args := buildAdminTaskWhere(AdminTaskListFilter{
Query: "needle",
GatewayTenant: "00000000-0000-0000-0000-000000000001",
GatewayUser: "00000000-0000-0000-0000-000000000002",
UserGroup: "00000000-0000-0000-0000-000000000003",
Status: "failed",
Platform: "00000000-0000-0000-0000-000000000004",
Model: "model-a",
ModelType: "image_generate",
RunMode: "production",
BillingStatus: "settled",
APIKey: "ops",
})
sql := strings.Join(where, "\n")
if strings.Contains(sql, "%!") || strings.Contains(sql, "$%d") {
t.Fatalf("SQL contains an unresolved placeholder: %s", sql)
}
if len(args) != 11 {
t.Fatalf("argument count=%d, want 11", len(args))
}
if strings.Count(sql, "$1") < 10 {
t.Fatalf("keyword search should reuse one placeholder, got: %s", sql)
}
for _, fragment := range []string{
"t.gateway_tenant_id",
"t.gateway_user_id",
"t.user_group_id",
"platform_attempt.platform_id",
"t.billing_status",
"t.api_key_prefix",
} {
if !strings.Contains(sql, fragment) {
t.Fatalf("SQL is missing %q: %s", fragment, sql)
}
}
}
apps/api/internal/store/admin_tasks.go
+385
View File
@@ -0,0 +1,385 @@
package store
import (
"context"
"fmt"
"strings"
"time"
)
type AdminTaskListFilter struct {
Query string
GatewayTenant string
GatewayUser string
UserGroup string
Status string
Platform string
Model string
ModelType string
RunMode string
BillingStatus string
APIKey string
CreatedFrom *time.Time
CreatedTo *time.Time
Page int
PageSize int
}
type AdminTaskUserSummary struct {
ID string `json:"id"`
Username string `json:"username,omitempty"`
DisplayName string `json:"displayName,omitempty"`
Email string `json:"email,omitempty"`
Source string `json:"source,omitempty"`
}
type AdminTaskTenantSummary struct {
ID string `json:"id"`
TenantKey string `json:"tenantKey,omitempty"`
Name string `json:"name,omitempty"`
}
type AdminTaskPlatformSummary struct {
ID string `json:"id"`
Name string `json:"name,omitempty"`
Provider string `json:"provider,omitempty"`
}
type AdminTaskContext struct {
User *AdminTaskUserSummary `json:"user,omitempty"`
Tenant *AdminTaskTenantSummary `json:"tenant,omitempty"`
LatestPlatform *AdminTaskPlatformSummary `json:"latestPlatform,omitempty"`
}
type AdminGatewayTask struct {
GatewayTask
AdminContext AdminTaskContext `json:"adminContext"`
}
type AdminTaskListResult struct {
Items []AdminGatewayTask
Total int
Page int
PageSize int
}
func (s *Store) ListAdminTasks(ctx context.Context, filter AdminTaskListFilter) (AdminTaskListResult, error) {
page := filter.Page
if page <= 0 {
page = 1
}
pageSize := filter.PageSize
if pageSize <= 0 {
pageSize = 50
}
if pageSize > 100 {
pageSize = 100
}
where, args := buildAdminTaskWhere(filter)
var total int
if err := s.pool.QueryRow(ctx, `
SELECT count(*)
FROM gateway_tasks t
WHERE `+strings.Join(where, "\n AND "), args...).Scan(&total); err != nil {
return AdminTaskListResult{}, err
}
queryArgs := append([]any{}, args...)
queryArgs = append(queryArgs, pageSize, (page-1)*pageSize)
rows, err := s.pool.Query(ctx, `
SELECT t.id::text
FROM gateway_tasks t
WHERE `+strings.Join(where, "\n AND ")+`
ORDER BY t.created_at DESC, t.id DESC
LIMIT $`+fmt.Sprint(len(args)+1)+` OFFSET $`+fmt.Sprint(len(args)+2), queryArgs...)
if err != nil {
return AdminTaskListResult{}, err
}
defer rows.Close()
taskIDs := make([]string, 0, pageSize)
for rows.Next() {
var taskID string
if err := rows.Scan(&taskID); err != nil {
return AdminTaskListResult{}, err
}
taskIDs = append(taskIDs, taskID)
}
if err := rows.Err(); err != nil {
return AdminTaskListResult{}, err
}
items, err := s.loadAdminTasksByIDs(ctx, taskIDs)
if err != nil {
return AdminTaskListResult{}, err
}
return AdminTaskListResult{
Items: items,
Total: total,
Page: page,
PageSize: pageSize,
}, nil
}
func (s *Store) GetAdminTask(ctx context.Context, taskID string) (AdminGatewayTask, error) {
task, err := s.GetTask(ctx, taskID)
if err != nil {
return AdminGatewayTask{}, err
}
contexts, err := s.loadAdminTaskContexts(ctx, []string{task.ID})
if err != nil {
return AdminGatewayTask{}, err
}
return adminTaskWithContext(task, contexts[task.ID]), nil
}
func buildAdminTaskWhere(filter AdminTaskListFilter) ([]string, []any) {
where := []string{"TRUE"}
args := make([]any, 0, 14)
add := func(clause string, value any) {
args = append(args, value)
placeholder := fmt.Sprintf("$%d", len(args))
where = append(where, strings.ReplaceAll(clause, "$%d", placeholder))
}
if query := strings.TrimSpace(filter.Query); query != "" {
add(`(
t.id::text ILIKE $%d
OR COALESCE(t.request_id, '') ILIKE $%d
OR t.kind ILIKE $%d
OR t.model ILIKE $%d
OR COALESCE(t.requested_model, '') ILIKE $%d
OR COALESCE(t.resolved_model, '') ILIKE $%d
OR COALESCE(t.api_key_id, '') ILIKE $%d
OR COALESCE(t.api_key_name, '') ILIKE $%d
OR COALESCE(t.api_key_prefix, '') ILIKE $%d
OR COALESCE(t.model_type, '') ILIKE $%d
OR EXISTS (
SELECT 1
FROM gateway_users admin_user
WHERE admin_user.id = t.gateway_user_id
AND (
admin_user.username ILIKE $%d
OR COALESCE(admin_user.display_name, '') ILIKE $%d
OR COALESCE(admin_user.email, '') ILIKE $%d
OR admin_user.user_key ILIKE $%d
)
)
OR EXISTS (
SELECT 1
FROM gateway_tenants admin_tenant
WHERE admin_tenant.id = t.gateway_tenant_id
AND (
admin_tenant.name ILIKE $%d
OR admin_tenant.tenant_key ILIKE $%d
)
)
OR EXISTS (
SELECT 1
FROM gateway_task_attempts admin_attempt
LEFT JOIN integration_platforms admin_platform ON admin_platform.id = admin_attempt.platform_id
LEFT JOIN platform_models admin_model ON admin_model.id = admin_attempt.platform_model_id
WHERE admin_attempt.task_id = t.id
AND (
COALESCE(admin_platform.name, '') ILIKE $%d
OR COALESCE(admin_platform.internal_name, '') ILIKE $%d
OR COALESCE(admin_platform.provider, '') ILIKE $%d
OR COALESCE(admin_model.model_name, '') ILIKE $%d
OR COALESCE(admin_model.provider_model_name, '') ILIKE $%d
OR COALESCE(admin_model.model_alias, '') ILIKE $%d
)
)
)`, "%"+query+"%")
}
if value := strings.TrimSpace(filter.GatewayTenant); value != "" {
add("t.gateway_tenant_id = $%d::uuid", value)
}
if value := strings.TrimSpace(filter.GatewayUser); value != "" {
add("t.gateway_user_id = $%d::uuid", value)
}
if value := strings.TrimSpace(filter.UserGroup); value != "" {
add("t.user_group_id = $%d::uuid", value)
}
if value := strings.TrimSpace(filter.Status); value != "" {
add("t.status = $%d", value)
}
if value := strings.TrimSpace(filter.Platform); value != "" {
add(`EXISTS (
SELECT 1
FROM gateway_task_attempts platform_attempt
WHERE platform_attempt.task_id = t.id
AND platform_attempt.platform_id = $%d::uuid
)`, value)
}
if value := strings.TrimSpace(filter.Model); value != "" {
add(`(
t.model = $%d
OR COALESCE(t.requested_model, '') = $%d
OR COALESCE(t.resolved_model, '') = $%d
OR EXISTS (
SELECT 1
FROM gateway_task_attempts model_attempt
LEFT JOIN platform_models selected_model ON selected_model.id = model_attempt.platform_model_id
WHERE model_attempt.task_id = t.id
AND (
COALESCE(selected_model.model_name, '') = $%d
OR COALESCE(selected_model.provider_model_name, '') = $%d
OR COALESCE(selected_model.model_alias, '') = $%d
)
)
)`, value)
}
if value := strings.TrimSpace(filter.ModelType); value != "" {
add("COALESCE(t.model_type, '') = $%d", value)
}
if value := strings.TrimSpace(filter.RunMode); value != "" {
add("t.run_mode = $%d", value)
}
if value := strings.TrimSpace(filter.BillingStatus); value != "" {
add("t.billing_status = $%d", value)
}
if value := strings.TrimSpace(filter.APIKey); value != "" {
add(`(
COALESCE(t.api_key_id, '') ILIKE $%d
OR COALESCE(t.api_key_name, '') ILIKE $%d
OR COALESCE(t.api_key_prefix, '') ILIKE $%d
)`, "%"+value+"%")
}
if filter.CreatedFrom != nil {
add("t.created_at >= $%d::timestamptz", *filter.CreatedFrom)
}
if filter.CreatedTo != nil {
add("t.created_at <= $%d::timestamptz", *filter.CreatedTo)
}
return where, args
}
func (s *Store) loadAdminTasksByIDs(ctx context.Context, taskIDs []string) ([]AdminGatewayTask, error) {
if len(taskIDs) == 0 {
return []AdminGatewayTask{}, nil
}
rows, err := s.pool.Query(ctx, `
SELECT `+gatewayTaskColumns+`
FROM gateway_tasks
WHERE id::text = ANY($1)`, taskIDs)
if err != nil {
return nil, err
}
defer rows.Close()
tasksByID := make(map[string]GatewayTask, len(taskIDs))
tasks := make([]GatewayTask, 0, len(taskIDs))
for rows.Next() {
task, err := scanGatewayTask(rows)
if err != nil {
return nil, err
}
tasks = append(tasks, task)
}
if err := rows.Err(); err != nil {
return nil, err
}
tasks, err = s.attachTaskAttempts(ctx, tasks)
if err != nil {
return nil, err
}
for _, task := range tasks {
tasksByID[task.ID] = task
}
contexts, err := s.loadAdminTaskContexts(ctx, taskIDs)
if err != nil {
return nil, err
}
items := make([]AdminGatewayTask, 0, len(taskIDs))
for _, taskID := range taskIDs {
task, ok := tasksByID[taskID]
if !ok {
continue
}
items = append(items, adminTaskWithContext(task, contexts[taskID]))
}
return items, nil
}
func (s *Store) loadAdminTaskContexts(ctx context.Context, taskIDs []string) (map[string]AdminTaskContext, error) {
contexts := make(map[string]AdminTaskContext, len(taskIDs))
if len(taskIDs) == 0 {
return contexts, nil
}
rows, err := s.pool.Query(ctx, `
SELECT t.id::text,
COALESCE(u.id::text, ''), COALESCE(u.username, ''), COALESCE(u.display_name, ''),
COALESCE(u.email, ''), COALESCE(u.source, ''),
COALESCE(tenant.id::text, ''), COALESCE(tenant.tenant_key, ''), COALESCE(tenant.name, '')
FROM gateway_tasks t
LEFT JOIN gateway_users u ON u.id = t.gateway_user_id
LEFT JOIN gateway_tenants tenant ON tenant.id = t.gateway_tenant_id
WHERE t.id::text = ANY($1)`, taskIDs)
if err != nil {
return nil, err
}
defer rows.Close()
for rows.Next() {
var taskID string
var user AdminTaskUserSummary
var tenant AdminTaskTenantSummary
if err := rows.Scan(
&taskID,
&user.ID,
&user.Username,
&user.DisplayName,
&user.Email,
&user.Source,
&tenant.ID,
&tenant.TenantKey,
&tenant.Name,
); err != nil {
return nil, err
}
context := AdminTaskContext{}
if user.ID != "" {
context.User = &user
}
if tenant.ID != "" {
context.Tenant = &tenant
}
contexts[taskID] = context
}
return contexts, rows.Err()
}
func adminTaskWithContext(task GatewayTask, context AdminTaskContext) AdminGatewayTask {
if context.User == nil {
userID := strings.TrimSpace(task.GatewayUserID)
if userID == "" {
userID = strings.TrimSpace(task.UserID)
}
if userID != "" {
context.User = &AdminTaskUserSummary{ID: userID, Source: task.UserSource}
}
}
if context.Tenant == nil {
tenantID := strings.TrimSpace(task.GatewayTenantID)
if tenantID == "" {
tenantID = strings.TrimSpace(task.TenantID)
}
if tenantID != "" || strings.TrimSpace(task.TenantKey) != "" {
context.Tenant = &AdminTaskTenantSummary{ID: tenantID, TenantKey: task.TenantKey}
}
}
for index := len(task.Attempts) - 1; index >= 0; index-- {
attempt := task.Attempts[index]
if strings.TrimSpace(attempt.PlatformID) == "" {
continue
}
context.LatestPlatform = &AdminTaskPlatformSummary{
ID: attempt.PlatformID,
Name: attempt.PlatformName,
Provider: attempt.Provider,
}
break
}
return AdminGatewayTask{GatewayTask: task, AdminContext: context}
}
apps/api/internal/store/async_worker_capacity.go
+152
View File
@@ -0,0 +1,152 @@
package store
import (
"context"
"fmt"
)
type AsyncWorkerCapacitySnapshot struct {
Capacity int
Desired int
HardLimit int
Capped bool
EnabledModels int
UnlimitedModels int
EnabledGroups int
UnlimitedGroups int
ModelDesired int
GroupDesired int
}
func (s *Store) AsyncWorkerCapacity(ctx context.Context, hardLimit int) (AsyncWorkerCapacitySnapshot, error) {
if hardLimit < 1 {
return AsyncWorkerCapacitySnapshot{}, fmt.Errorf("async worker hard limit must be positive")
}
rows, err := s.pool.Query(ctx, `
SELECT COALESCE(b.default_rate_limit_policy, '{}'::jsonb),
p.rate_limit_policy,
COALESCE(rp.rate_limit_policy, '{}'::jsonb),
(m.runtime_policy_set_id IS NOT NULL),
COALESCE(m.runtime_policy_override, '{}'::jsonb),
m.rate_limit_policy,
m.rate_limit_policy_mode
FROM platform_models m
JOIN integration_platforms p ON p.id = m.platform_id
LEFT JOIN base_model_catalog b ON b.id = m.base_model_id
LEFT JOIN model_runtime_policy_sets rp ON rp.id = COALESCE(m.runtime_policy_set_id, b.runtime_policy_set_id)
WHERE p.status = 'enabled'
AND p.deleted_at IS NULL
AND m.enabled = true`)
if err != nil {
return AsyncWorkerCapacitySnapshot{}, err
}
defer rows.Close()
modelPolicies := make([]map[string]any, 0)
for rows.Next() {
var basePolicyBytes, platformPolicyBytes, runtimePolicyBytes []byte
var runtimeOverrideBytes, modelPolicyBytes []byte
var runtimeExplicit bool
var modelPolicyMode string
if err := rows.Scan(
&basePolicyBytes,
&platformPolicyBytes,
&runtimePolicyBytes,
&runtimeExplicit,
&runtimeOverrideBytes,
&modelPolicyBytes,
&modelPolicyMode,
); err != nil {
return AsyncWorkerCapacitySnapshot{}, err
}
modelPolicies = append(modelPolicies, EffectiveRateLimitPolicy(EffectiveRateLimitPolicyInput{
BasePolicy: decodeObject(basePolicyBytes),
PlatformPolicy: decodeObject(platformPolicyBytes),
RuntimePolicy: decodeObject(runtimePolicyBytes),
RuntimePolicyExplicit: runtimeExplicit,
RuntimePolicyOverride: decodeObject(runtimeOverrideBytes),
ModelPolicy: decodeObject(modelPolicyBytes),
ModelPolicyMode: modelPolicyMode,
}))
}
if err := rows.Err(); err != nil {
return AsyncWorkerCapacitySnapshot{}, err
}
groupRows, err := s.pool.Query(ctx, `
SELECT rate_limit_policy
FROM gateway_user_groups
WHERE status = 'active'`)
if err != nil {
return AsyncWorkerCapacitySnapshot{}, err
}
defer groupRows.Close()
groupPolicies := make([]map[string]any, 0)
for groupRows.Next() {
var policyBytes []byte
if err := groupRows.Scan(&policyBytes); err != nil {
return AsyncWorkerCapacitySnapshot{}, err
}
groupPolicies = append(groupPolicies, NormalizeRateLimitPolicy(decodeObject(policyBytes)))
}
if err := groupRows.Err(); err != nil {
return AsyncWorkerCapacitySnapshot{}, err
}
return asyncWorkerCapacityFromPolicySets(modelPolicies, groupPolicies, hardLimit), nil
}
func asyncWorkerCapacityFromPolicies(policies []map[string]any, hardLimit int) AsyncWorkerCapacitySnapshot {
return asyncWorkerCapacityFromPolicySets(policies, nil, hardLimit)
}
func asyncWorkerCapacityFromPolicySets(modelPolicies []map[string]any, groupPolicies []map[string]any, hardLimit int) AsyncWorkerCapacitySnapshot {
snapshot := AsyncWorkerCapacitySnapshot{
HardLimit: hardLimit,
EnabledModels: len(modelPolicies),
EnabledGroups: len(groupPolicies),
}
modelDesired, modelFinite, unlimitedModels := concurrentPolicySetCapacity(modelPolicies)
groupDesired, groupFinite, unlimitedGroups := concurrentPolicySetCapacity(groupPolicies)
snapshot.ModelDesired = modelDesired
snapshot.GroupDesired = groupDesired
snapshot.UnlimitedModels = unlimitedModels
snapshot.UnlimitedGroups = unlimitedGroups
desired := hardLimit
switch {
case snapshot.EnabledModels == 0:
desired = 1
case modelFinite && groupFinite:
desired = min(modelDesired, groupDesired)
case modelFinite:
desired = modelDesired
case groupFinite:
desired = groupDesired
}
if desired < 1 {
desired = 1
}
snapshot.Desired = desired
snapshot.Capacity = desired
if snapshot.Capacity > hardLimit {
snapshot.Capacity = hardLimit
snapshot.Capped = true
}
return snapshot
}
func concurrentPolicySetCapacity(policies []map[string]any) (total int, finite bool, unlimited int) {
if len(policies) == 0 {
return 0, false, 0
}
finite = true
for _, policy := range policies {
capacity, policyFinite := ConcurrentPolicyCapacity(policy)
if !policyFinite {
unlimited++
finite = false
continue
}
total += capacity
}
return total, finite, unlimited
}
apps/api/internal/store/base_models.go
+194 -49
View File
@@ -3,25 +3,36 @@ package store
import (
"context"
"encoding/json"
"fmt"
"strings"
"github.com/jackc/pgx/v5"
)
const baseModelColumns = `
id::text, provider_key, canonical_model_key, provider_model_name, model_type, display_name,
id::text, provider_key, canonical_model_key, invocation_name, provider_model_name, model_type, display_name,
capabilities, base_billing_config, default_rate_limit_policy, COALESCE(pricing_rule_set_id::text, ''),
COALESCE(runtime_policy_set_id::text, ''), runtime_policy_override, metadata,
catalog_type, COALESCE(default_snapshot, '{}'::jsonb), COALESCE(customized_at::text, ''),
pricing_version, status, created_at, updated_at`
pricing_version, status, created_at, updated_at,
COALESCE((
SELECT jsonb_agg(DISTINCT compatibility_alias.alias ORDER BY compatibility_alias.alias)
FROM model_compatibility_aliases compatibility_alias
WHERE compatibility_alias.base_model_id = base_model_catalog.id
AND compatibility_alias.active = true
AND (compatibility_alias.expires_at IS NULL OR compatibility_alias.expires_at > now())
), '[]'::jsonb),
(SELECT count(*)::int FROM platform_models platform_model WHERE platform_model.base_model_id = base_model_catalog.id)`
type BaseModelInput struct {
ProviderKey string `json:"providerKey"`
CanonicalModelKey string `json:"canonicalModelKey"`
InvocationName string `json:"invocationName"`
ProviderModelName string `json:"providerModelName"`
ModelType StringList `json:"modelType"`
ModelAlias string `json:"modelAlias"`
DisplayName string `json:"displayName"`
LegacyAliases StringList `json:"legacyAliases"`
Capabilities map[string]any `json:"capabilities"`
BaseBillingConfig map[string]any `json:"baseBillingConfig"`
DefaultRateLimitPolicy map[string]any `json:"defaultRateLimitPolicy"`
@@ -86,25 +97,32 @@ func (s *Store) CreateBaseModel(ctx context.Context, input BaseModelInput) (Base
defaultSnapshot, _ := json.Marshal(emptyObjectIfNil(input.DefaultSnapshot))
modelType, _ := json.Marshal(input.ModelType)
return scanBaseModel(s.pool.QueryRow(ctx, `
tx, err := s.pool.Begin(ctx)
if err != nil {
return BaseModel{}, err
}
defer tx.Rollback(ctx)
item, err := scanBaseModel(tx.QueryRow(ctx, `
INSERT INTO base_model_catalog (
provider_id, provider_key, canonical_model_key, provider_model_name, model_type, display_name,
provider_id, provider_key, canonical_model_key, invocation_name, provider_model_name, model_type, display_name,
capabilities, base_billing_config, default_rate_limit_policy, pricing_rule_set_id, runtime_policy_set_id, runtime_policy_override,
metadata, catalog_type, default_snapshot, pricing_version, status
)
VALUES (
(SELECT id FROM model_catalog_providers WHERE provider_key = $1 OR provider_code = $1 LIMIT 1),
$1, $2, $3, $4::jsonb, $5, $6, $7, $8,
COALESCE(NULLIF($9, '')::uuid, (SELECT id FROM model_pricing_rule_sets WHERE rule_set_key = 'default-multimodal-v1' LIMIT 1)),
COALESCE(NULLIF($10, '')::uuid, (SELECT id FROM model_runtime_policy_sets WHERE policy_key = 'default-runtime-v1' LIMIT 1)),
$11, $12, NULLIF($13, ''), NULLIF($14::jsonb, '{}'::jsonb), $15, $16
$1, $2, $3, $4, $5::jsonb, $6, $7, $8, $9,
COALESCE(NULLIF($10, '')::uuid, (SELECT id FROM model_pricing_rule_sets WHERE rule_set_key = 'default-multimodal-v1' LIMIT 1)),
COALESCE(NULLIF($11, '')::uuid, (SELECT id FROM model_runtime_policy_sets WHERE policy_key = 'default-runtime-v1' LIMIT 1)),
$12, $13, NULLIF($14, ''), NULLIF($15::jsonb, '{}'::jsonb), $16, $17
)
RETURNING `+baseModelColumns,
input.ProviderKey,
input.CanonicalModelKey,
input.InvocationName,
input.ProviderModelName,
string(modelType),
input.ModelAlias,
input.DisplayName,
capabilities,
billingConfig,
rateLimitPolicy,
@@ -117,6 +135,17 @@ RETURNING `+baseModelColumns,
input.PricingVersion,
input.Status,
))
if err != nil {
return BaseModel{}, err
}
if err := replaceBaseModelAliases(ctx, tx, item.ID, input); err != nil {
return BaseModel{}, err
}
if err := tx.Commit(ctx); err != nil {
return BaseModel{}, err
}
item.LegacyAliases = input.LegacyAliases
return item, nil
}
func (s *Store) UpdateBaseModel(ctx context.Context, id string, input BaseModelInput) (BaseModel, error) {
@@ -129,35 +158,43 @@ func (s *Store) UpdateBaseModel(ctx context.Context, id string, input BaseModelI
defaultSnapshot, _ := json.Marshal(emptyObjectIfNil(input.DefaultSnapshot))
modelType, _ := json.Marshal(input.ModelType)
return scanBaseModel(s.pool.QueryRow(ctx, `
tx, err := s.pool.Begin(ctx)
if err != nil {
return BaseModel{}, err
}
defer tx.Rollback(ctx)
item, err := scanBaseModel(tx.QueryRow(ctx, `
UPDATE base_model_catalog
SET provider_id = (SELECT id FROM model_catalog_providers WHERE provider_key = $2 OR provider_code = $2 LIMIT 1),
provider_key = $2,
canonical_model_key = $3,
provider_model_name = $4,
model_type = $5::jsonb,
display_name = $6,
capabilities = $7,
base_billing_config = $8,
default_rate_limit_policy = $9,
pricing_rule_set_id = COALESCE(NULLIF($10, '')::uuid, (SELECT id FROM model_pricing_rule_sets WHERE rule_set_key = 'default-multimodal-v1' LIMIT 1)),
runtime_policy_set_id = COALESCE(NULLIF($11, '')::uuid, (SELECT id FROM model_runtime_policy_sets WHERE policy_key = 'default-runtime-v1' LIMIT 1)),
runtime_policy_override = $12,
metadata = $13,
catalog_type = NULLIF($14, ''),
default_snapshot = COALESCE(NULLIF($15::jsonb, '{}'::jsonb), default_snapshot),
customized_at = CASE WHEN NULLIF($14, '') = 'system' THEN now() ELSE NULL END,
pricing_version = $16,
status = $17,
invocation_name = $4,
provider_model_name = $5,
model_type = $6::jsonb,
display_name = $7,
capabilities = $8,
base_billing_config = $9,
default_rate_limit_policy = $10,
pricing_rule_set_id = COALESCE(NULLIF($11, '')::uuid, (SELECT id FROM model_pricing_rule_sets WHERE rule_set_key = 'default-multimodal-v1' LIMIT 1)),
runtime_policy_set_id = COALESCE(NULLIF($12, '')::uuid, (SELECT id FROM model_runtime_policy_sets WHERE policy_key = 'default-runtime-v1' LIMIT 1)),
runtime_policy_override = $13,
metadata = $14,
catalog_type = NULLIF($15, ''),
default_snapshot = COALESCE(NULLIF($16::jsonb, '{}'::jsonb), default_snapshot),
customized_at = CASE WHEN NULLIF($15, '') = 'system' THEN now() ELSE NULL END,
pricing_version = $17,
status = $18,
updated_at = now()
WHERE id = $1::uuid
RETURNING `+baseModelColumns,
id,
input.ProviderKey,
input.CanonicalModelKey,
input.InvocationName,
input.ProviderModelName,
string(modelType),
input.ModelAlias,
input.DisplayName,
capabilities,
billingConfig,
rateLimitPolicy,
@@ -170,6 +207,17 @@ RETURNING `+baseModelColumns,
input.PricingVersion,
input.Status,
))
if err != nil {
return BaseModel{}, err
}
if err := replaceBaseModelAliases(ctx, tx, item.ID, input); err != nil {
return BaseModel{}, err
}
if err := tx.Commit(ctx); err != nil {
return BaseModel{}, err
}
item.LegacyAliases = input.LegacyAliases
return item, nil
}
func (s *Store) ResetBaseModelToDefault(ctx context.Context, id string) (BaseModel, error) {
@@ -194,18 +242,19 @@ UPDATE base_model_catalog
SET provider_id = (SELECT id FROM model_catalog_providers WHERE provider_key = COALESCE($2::text, provider_key) OR provider_code = COALESCE($2::text, provider_key) LIMIT 1),
provider_key = COALESCE($2::text, provider_key),
canonical_model_key = COALESCE($3::text, canonical_model_key),
provider_model_name = COALESCE($4::text, provider_model_name),
model_type = COALESCE($5::jsonb, model_type),
display_name = COALESCE($6::text, display_name),
capabilities = COALESCE($7::jsonb, capabilities),
base_billing_config = COALESCE($8::jsonb, base_billing_config),
default_rate_limit_policy = COALESCE($9::jsonb, default_rate_limit_policy),
pricing_rule_set_id = COALESCE(NULLIF($10::text, '')::uuid, pricing_rule_set_id),
runtime_policy_set_id = COALESCE(NULLIF($11::text, '')::uuid, runtime_policy_set_id),
runtime_policy_override = COALESCE($12::jsonb, runtime_policy_override),
metadata = COALESCE($13::jsonb, metadata),
pricing_version = COALESCE($14::integer, pricing_version),
status = COALESCE($15::text, status),
invocation_name = COALESCE($4::text, invocation_name),
provider_model_name = COALESCE($5::text, provider_model_name),
model_type = COALESCE($6::jsonb, model_type),
display_name = COALESCE($7::text, display_name),
capabilities = COALESCE($8::jsonb, capabilities),
base_billing_config = COALESCE($9::jsonb, base_billing_config),
default_rate_limit_policy = COALESCE($10::jsonb, default_rate_limit_policy),
pricing_rule_set_id = COALESCE(NULLIF($11::text, '')::uuid, pricing_rule_set_id),
runtime_policy_set_id = COALESCE(NULLIF($12::text, '')::uuid, runtime_policy_set_id),
runtime_policy_override = COALESCE($13::jsonb, runtime_policy_override),
metadata = COALESCE($14::jsonb, metadata),
pricing_version = COALESCE($15::integer, pricing_version),
status = COALESCE($16::text, status),
customized_at = NULL,
updated_at = now()
WHERE id = $1::uuid
@@ -213,9 +262,10 @@ RETURNING `+baseModelColumns,
id,
stringFromSnapshot(snapshot, "providerKey"),
stringFromSnapshot(snapshot, "canonicalModelKey"),
stringFromSnapshot(snapshot, "invocationName", "modelAlias", "providerModelName"),
stringFromSnapshot(snapshot, "providerModelName"),
jsonStringListFromSnapshot(snapshot, "modelType"),
stringFromSnapshot(snapshot, "modelAlias", "displayName"),
stringFromSnapshot(snapshot, "displayName", "modelAlias", "providerModelName"),
jsonFromSnapshot(snapshot, "capabilities"),
jsonFromSnapshot(snapshot, "baseBillingConfig"),
jsonFromSnapshot(snapshot, "defaultRateLimitPolicy"),
@@ -240,13 +290,14 @@ SET provider_id = (
),
provider_key = COALESCE(NULLIF(default_snapshot->>'providerKey', ''), provider_key),
canonical_model_key = COALESCE(NULLIF(default_snapshot->>'canonicalModelKey', ''), canonical_model_key),
invocation_name = COALESCE(NULLIF(default_snapshot->>'invocationName', ''), NULLIF(default_snapshot->>'modelAlias', ''), invocation_name),
provider_model_name = COALESCE(NULLIF(default_snapshot->>'providerModelName', ''), provider_model_name),
model_type = COALESCE(NULLIF(CASE
WHEN jsonb_typeof(default_snapshot->'modelType') = 'array' THEN default_snapshot->'modelType'
WHEN COALESCE(default_snapshot->>'modelType', '') <> '' THEN jsonb_build_array(default_snapshot->>'modelType')
ELSE NULL
END, '[]'::jsonb), model_type),
display_name = COALESCE(NULLIF(COALESCE(default_snapshot->>'modelAlias', default_snapshot->>'displayName'), ''), display_name),
display_name = COALESCE(NULLIF(COALESCE(default_snapshot->>'displayName', default_snapshot->>'modelAlias'), ''), display_name),
capabilities = COALESCE(default_snapshot->'capabilities', capabilities),
base_billing_config = COALESCE(default_snapshot->'baseBillingConfig', base_billing_config),
default_rate_limit_policy = COALESCE(default_snapshot->'defaultRateLimitPolicy', default_rate_limit_policy),
@@ -269,11 +320,23 @@ RETURNING `+baseModelColumns)
}
func (s *Store) DeleteBaseModel(ctx context.Context, id string) error {
result, err := s.pool.Exec(ctx, `DELETE FROM base_model_catalog WHERE id = $1::uuid`, id)
result, err := s.pool.Exec(ctx, `
DELETE FROM base_model_catalog base_model
WHERE base_model.id = $1::uuid
AND NOT EXISTS (
SELECT 1 FROM platform_models platform_model WHERE platform_model.base_model_id = base_model.id
)`, id)
if err != nil {
return err
}
if result.RowsAffected() == 0 {
var exists bool
if err := s.pool.QueryRow(ctx, `SELECT EXISTS (SELECT 1 FROM base_model_catalog WHERE id = $1::uuid)`, id).Scan(&exists); err != nil {
return err
}
if exists {
return ErrBaseModelInUse
}
return pgx.ErrNoRows
}
return nil
@@ -294,7 +357,7 @@ func scanBaseModelRows(rows pgx.Rows) ([]BaseModel, error) {
func scanBaseModel(scanner baseModelScanner) (BaseModel, error) {
var item BaseModel
var modelType []byte
var modelAlias string
var legacyAliases []byte
var capabilities []byte
var billingConfig []byte
var rateLimitPolicy []byte
@@ -305,9 +368,10 @@ func scanBaseModel(scanner baseModelScanner) (BaseModel, error) {
&item.ID,
&item.ProviderKey,
&item.CanonicalModelKey,
&item.InvocationName,
&item.ProviderModelName,
&modelType,
&modelAlias,
&item.DisplayName,
&capabilities,
&billingConfig,
&rateLimitPolicy,
@@ -322,6 +386,8 @@ func scanBaseModel(scanner baseModelScanner) (BaseModel, error) {
&item.Status,
&item.CreatedAt,
&item.UpdatedAt,
&legacyAliases,
&item.ReferenceCount,
); err != nil {
return BaseModel{}, err
}
@@ -332,18 +398,20 @@ func scanBaseModel(scanner baseModelScanner) (BaseModel, error) {
item.Metadata = decodeObject(metadata)
item.DefaultSnapshot = decodeObject(defaultSnapshot)
item.ModelType = decodeStringArray(modelType)
item.ModelAlias = modelAlias
item.DisplayName = modelAlias
item.LegacyAliases = decodeStringArray(legacyAliases)
item.ModelAlias = item.InvocationName
return item, nil
}
func normalizeBaseModelInput(input BaseModelInput) BaseModelInput {
input.ProviderKey = strings.TrimSpace(input.ProviderKey)
input.CanonicalModelKey = strings.TrimSpace(input.CanonicalModelKey)
input.InvocationName = strings.TrimSpace(input.InvocationName)
input.ProviderModelName = strings.TrimSpace(input.ProviderModelName)
input.ModelType = uniqueStringList(input.ModelType)
input.ModelAlias = strings.TrimSpace(input.ModelAlias)
input.DisplayName = strings.TrimSpace(input.DisplayName)
input.LegacyAliases = normalizeCompatibilityAliases(input.LegacyAliases)
input.PricingRuleSetID = strings.TrimSpace(input.PricingRuleSetID)
input.RuntimePolicySetID = strings.TrimSpace(input.RuntimePolicySetID)
input.CatalogType = strings.TrimSpace(input.CatalogType)
@@ -351,12 +419,17 @@ func normalizeBaseModelInput(input BaseModelInput) BaseModelInput {
if input.CanonicalModelKey == "" && input.ProviderKey != "" && input.ProviderModelName != "" {
input.CanonicalModelKey = input.ProviderKey + ":" + input.ProviderModelName
}
if input.ModelAlias == "" {
input.ModelAlias = input.DisplayName
if input.InvocationName == "" {
input.InvocationName = input.ModelAlias
}
if input.ModelAlias == "" {
input.ModelAlias = input.ProviderModelName
if input.InvocationName == "" {
input.InvocationName = input.ProviderModelName
}
if input.DisplayName == "" {
input.DisplayName = input.InvocationName
}
input.ModelAlias = input.InvocationName
input.LegacyAliases = withoutStrings(input.LegacyAliases, input.InvocationName)
if len(input.ModelType) == 0 {
input.ModelType = StringList{"text_generate"}
}
@@ -372,6 +445,78 @@ func normalizeBaseModelInput(input BaseModelInput) BaseModelInput {
return input
}
func replaceBaseModelAliases(ctx context.Context, tx pgx.Tx, baseModelID string, input BaseModelInput) error {
for _, alias := range input.LegacyAliases {
for _, modelType := range input.ModelType {
var conflictingCanonicalKey string
if err := tx.QueryRow(ctx, `
SELECT COALESCE((
SELECT other.canonical_model_key
FROM base_model_catalog other
WHERE other.id <> $1::uuid
AND other.invocation_name <> $4::text
AND other.model_type ? $3::text
AND (
other.invocation_name = $2::text
OR EXISTS (
SELECT 1
FROM model_compatibility_aliases other_alias
WHERE other_alias.base_model_id = other.id
AND other_alias.alias = $2::text
AND other_alias.model_type = $3::text
AND other_alias.active = true
AND (other_alias.expires_at IS NULL OR other_alias.expires_at > now())
)
)
LIMIT 1
), '')`, baseModelID, alias, modelType, input.InvocationName).Scan(&conflictingCanonicalKey); err != nil {
return err
}
if conflictingCanonicalKey != "" {
return fmt.Errorf("%w: %q (%s) conflicts with %s", ErrModelAliasConflict, alias, modelType, conflictingCanonicalKey)
}
}
}
if _, err := tx.Exec(ctx, `DELETE FROM model_compatibility_aliases WHERE base_model_id = $1::uuid`, baseModelID); err != nil {
return err
}
for _, alias := range input.LegacyAliases {
for _, modelType := range input.ModelType {
if _, err := tx.Exec(ctx, `
INSERT INTO model_compatibility_aliases (base_model_id, alias, model_type, expires_at)
VALUES ($1::uuid, $2, $3, now() + interval '14 days')`, baseModelID, alias, modelType); err != nil {
return err
}
}
}
return nil
}
func normalizeCompatibilityAliases(values []string) StringList {
out := make(StringList, 0, len(values))
seen := map[string]bool{}
for _, value := range values {
value = strings.TrimSpace(value)
if value == "" || seen[value] {
continue
}
seen[value] = true
out = append(out, value)
}
return out
}
func withoutStrings(values []string, excluded string) StringList {
out := make(StringList, 0, len(values))
for _, value := range values {
if value != excluded {
out = append(out, value)
}
}
return out
}
func stringFromSnapshot(snapshot map[string]any, keys ...string) any {
for _, key := range keys {
value, ok := snapshot[key]
apps/api/internal/store/billing_settlements.go
+5 -5
View File
@@ -432,13 +432,13 @@ SELECT task.id,
COALESCE((SELECT MAX(event.seq) + 1 FROM gateway_task_events event WHERE event.task_id = task.id), 1),
CASE WHEN $2 = 'settled' THEN 'task.billing.settled' ELSE 'task.billing.released' END,
task.status,
'billing',
1,
CASE WHEN $2 = 'settled' THEN 'task billing settled' ELSE 'task billing reservation released' END,
jsonb_build_object('settlementId', $3::text, 'billingStatus', $2::text),
NULL,
0,
NULL,
'{}'::jsonb,
task.run_mode = 'simulation'
FROM gateway_tasks task
WHERE task.id = $1::uuid`, taskID, billingStatus, settlementID); err != nil {
WHERE task.id = $1::uuid`, taskID, billingStatus); err != nil {
return err
}
tag, err := tx.Exec(ctx, `
apps/api/internal/store/binary_result_backfill.go
+87
View File
@@ -0,0 +1,87 @@
package store
import (
"context"
"encoding/json"
"time"
"github.com/jackc/pgx/v5"
)
type TaskBinaryResultBackfillItem struct {
ID string
Result map[string]any
FinishedAt time.Time
}
func (s *Store) ListTaskBinaryResultBackfillBatch(ctx context.Context, afterID string, batchSize int) ([]TaskBinaryResultBackfillItem, error) {
if batchSize < 1 || batchSize > 100 {
batchSize = 100
}
var items []TaskBinaryResultBackfillItem
err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
if _, err := tx.Exec(ctx, `SET LOCAL lock_timeout = '5s'; SET LOCAL statement_timeout = '5s'`); err != nil {
return err
}
rows, err := tx.Query(ctx, `
SELECT id::text, result, COALESCE(finished_at, updated_at)
FROM gateway_tasks
WHERE status = 'succeeded'
AND result <> '{}'::jsonb
AND (NULLIF($1::text, '') IS NULL OR id > NULLIF($1::text, '')::uuid)
ORDER BY id
LIMIT $2`, afterID, batchSize)
if err != nil {
return err
}
defer rows.Close()
items = make([]TaskBinaryResultBackfillItem, 0, batchSize)
for rows.Next() {
var item TaskBinaryResultBackfillItem
var resultJSON []byte
if err := rows.Scan(&item.ID, &resultJSON, &item.FinishedAt); err != nil {
return err
}
item.Result = decodeObject(resultJSON)
items = append(items, item)
}
return rows.Err()
})
return items, err
}
func (s *Store) UpdateTaskBinaryResultBackfill(ctx context.Context, taskID string, result map[string]any) (bool, error) {
report := sanitizeJSONForStorageWithReport(minimalTaskResult(result))
if report.BinaryCount > 0 {
return false, &taskPayloadBinaryError{
target: ErrTaskResultBinaryNotMaterialized,
code: "result_binary_not_materialized",
count: report.BinaryCount,
}
}
resultJSON, err := json.Marshal(report.Value)
if err != nil {
return false, err
}
updated := false
err = pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
if _, err := tx.Exec(ctx, `SET LOCAL lock_timeout = '5s'; SET LOCAL statement_timeout = '5s'`); err != nil {
return err
}
tag, err := tx.Exec(ctx, `
UPDATE gateway_tasks
SET result = $2::jsonb,
updated_at = now()
WHERE id = $1::uuid
AND status = 'succeeded'`,
taskID,
string(resultJSON),
)
if err != nil {
return err
}
updated = tag.RowsAffected() == 1
return nil
})
return updated, err
}
apps/api/internal/store/candidates.go
+194 -194
View File
@@ -2,13 +2,15 @@ package store
import (
"context"
"errors"
"fmt"
"math"
"sort"
"strings"
"unicode"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
"github.com/jackc/pgx/v5"
)
type ListModelCandidatesOptions struct {
@@ -19,7 +21,6 @@ type ListModelCandidatesOptions struct {
func (s *Store) ListModelCandidates(ctx context.Context, model string, modelType string, user *auth.User, options ...ListModelCandidatesOptions) ([]RuntimeModelCandidate, error) {
exactModel := strings.TrimSpace(model)
modelMatchKey := normalizeModelMatchKey(exactModel)
listOptions := normalizeListModelCandidatesOptions(modelType, options...)
rows, err := s.pool.Query(ctx, `
SELECT p.id::text, p.platform_key, p.name, p.provider,
@@ -27,27 +28,40 @@ func (s *Store) ListModelCandidates(ctx context.Context, model string, modelType
COALESCE(p.base_url, ''),
p.auth_type, p.credentials, p.config, p.default_pricing_mode,
p.default_discount_factor::float8, COALESCE(p.pricing_rule_set_id::text, ''),
p.retry_policy, p.rate_limit_policy,
p.retry_policy, p.rate_limit_policy, COALESCE(b.default_rate_limit_policy, '{}'::jsonb),
COALESCE(p.dynamic_priority, p.priority) AS effective_priority,
m.id::text, COALESCE(m.base_model_id::text, ''), COALESCE(b.canonical_model_key, ''),
COALESCE(NULLIF(m.provider_model_name, ''), m.model_name), m.model_name, COALESCE(m.model_alias, ''),
$2::text AS requested_model_type, m.display_name,
EXISTS (
SELECT 1
FROM model_compatibility_aliases compatibility_alias
JOIN base_model_catalog alias_base ON alias_base.id = compatibility_alias.base_model_id
WHERE compatibility_alias.alias = $1::text
AND compatibility_alias.model_type = $2::text
AND compatibility_alias.active = true
AND (compatibility_alias.expires_at IS NULL OR compatibility_alias.expires_at > now())
AND (compatibility_alias.base_model_id = b.id OR alias_base.invocation_name = b.invocation_name)
) AS legacy_alias_used,
CASE
WHEN b.capabilities #> '{text_generate,supportedApiProtocols}' IS NOT NULL
THEN jsonb_set(
COALESCE(m.capabilities, '{}'::jsonb),
COALESCE(b.capabilities, '{}'::jsonb) || COALESCE(m.capabilities, '{}'::jsonb),
'{text_generate,supportedApiProtocols}',
b.capabilities #> '{text_generate,supportedApiProtocols}',
true
)
ELSE m.capabilities
ELSE COALESCE(b.capabilities, '{}'::jsonb) || COALESCE(m.capabilities, '{}'::jsonb)
END AS effective_capabilities,
m.capability_override,
COALESCE(b.base_billing_config, '{}'::jsonb), m.billing_config, m.billing_config_override,
m.pricing_mode, COALESCE(m.discount_factor, 0)::float8, COALESCE(m.pricing_rule_set_id::text, ''),
COALESCE(b.pricing_rule_set_id::text, ''),
m.permission_config, m.retry_policy, m.rate_limit_policy, COALESCE(m.runtime_policy_set_id::text, COALESCE(b.runtime_policy_set_id::text, '')),
m.permission_config, m.retry_policy, m.rate_limit_policy, m.rate_limit_policy_mode,
COALESCE(m.runtime_policy_set_id::text, COALESCE(b.runtime_policy_set_id::text, '')),
(m.runtime_policy_set_id IS NOT NULL),
COALESCE(NULLIF(m.runtime_policy_override, '{}'::jsonb), b.runtime_policy_override, '{}'::jsonb),
COALESCE(m.runtime_policy_override, '{}'::jsonb),
COALESCE(rp.retry_policy, '{}'::jsonb), COALESCE(rp.rate_limit_policy, '{}'::jsonb),
COALESCE(rp.auto_disable_policy, '{}'::jsonb), COALESCE(rp.degrade_policy, '{}'::jsonb),
COALESCE(con.active, 0)::float8,
@@ -74,10 +88,10 @@ func (s *Store) ListModelCandidates(ctx context.Context, model string, modelType
ON s.client_id = p.platform_key || ':' || $2::text || ':' || COALESCE(NULLIF(m.provider_model_name, ''), m.model_name)
LEFT JOIN LATERAL (
SELECT ca.cache_affinity_key, ca.request_count, ca.input_tokens, ca.cached_input_tokens, ca.ema_hit_ratio, ca.last_hit_ratio, ca.last_observed_at
FROM unnest($4::text[]) WITH ORDINALITY AS affinity_keys(cache_affinity_key, affinity_rank)
FROM unnest($3::text[]) WITH ORDINALITY AS affinity_keys(cache_affinity_key, affinity_rank)
JOIN gateway_cache_affinity_stats ca ON ca.cache_affinity_key = affinity_keys.cache_affinity_key
WHERE ca.client_id = p.platform_key || ':' || $2::text || ':' || COALESCE(NULLIF(m.provider_model_name, ''), m.model_name)
AND ($5::int <= 0 OR ca.last_observed_at >= now() - ($5::int * interval '1 second'))
AND ($4::int <= 0 OR ca.last_observed_at >= now() - ($4::int * interval '1 second'))
ORDER BY affinity_keys.affinity_rank ASC, ca.cached_input_tokens DESC, ca.ema_hit_ratio DESC, ca.last_observed_at DESC
LIMIT 1
) ca ON TRUE
@@ -138,61 +152,25 @@ WHERE p.status = 'enabled'
AND m.model_type @> jsonb_build_array($2::text)
AND (p.cooldown_until IS NULL OR p.cooldown_until <= now())
AND (m.cooldown_until IS NULL OR m.cooldown_until <= now())
AND (
(
$2::text IN ('audio_generate', 'text_to_speech', 'voice_clone')
AND (
m.model_alias = $1::text
OR m.model_name = $1::text
OR b.canonical_model_key = $1::text
OR b.provider_model_name = $1::text
OR (
NULLIF($3::text, '') IS NOT NULL
AND (
regexp_replace(COALESCE(m.model_alias, ''), '[[:space:]]+', '', 'g') = $3::text
OR regexp_replace(COALESCE(m.model_name, ''), '[[:space:]]+', '', 'g') = $3::text
OR regexp_replace(COALESCE(b.canonical_model_key, ''), '[[:space:]]+', '', 'g') = $3::text
OR regexp_replace(COALESCE(b.provider_model_name, ''), '[[:space:]]+', '', 'g') = $3::text
)
)
)
)
OR (
$2::text NOT IN ('audio_generate', 'text_to_speech', 'voice_clone')
AND (
(
COALESCE(m.model_alias, '') <> ''
AND (
m.model_alias = $1::text
OR (
NULLIF($3::text, '') IS NOT NULL
AND regexp_replace(COALESCE(m.model_alias, ''), '[[:space:]]+', '', 'g') = $3::text
)
)
)
OR (
m.model_name = $1::text
OR COALESCE(NULLIF(m.provider_model_name, ''), m.model_name) = $1::text
OR b.canonical_model_key = $1::text
OR b.provider_model_name = $1::text
OR (
NULLIF($3::text, '') IS NOT NULL
AND (
regexp_replace(COALESCE(m.model_name, ''), '[[:space:]]+', '', 'g') = $3::text
OR regexp_replace(COALESCE(NULLIF(m.provider_model_name, ''), m.model_name), '[[:space:]]+', '', 'g') = $3::text
OR regexp_replace(COALESCE(b.canonical_model_key, ''), '[[:space:]]+', '', 'g') = $3::text
OR regexp_replace(COALESCE(b.provider_model_name, ''), '[[:space:]]+', '', 'g') = $3::text
)
)
)
)
)
)
AND (
b.invocation_name = $1::text
OR (b.id IS NULL AND m.model_name = $1::text)
OR EXISTS (
SELECT 1
FROM model_compatibility_aliases compatibility_alias
JOIN base_model_catalog alias_base ON alias_base.id = compatibility_alias.base_model_id
WHERE compatibility_alias.alias = $1::text
AND compatibility_alias.model_type = $2::text
AND compatibility_alias.active = true
AND (compatibility_alias.expires_at IS NULL OR compatibility_alias.expires_at > now())
AND (compatibility_alias.base_model_id = b.id OR alias_base.invocation_name = b.invocation_name)
)
)
ORDER BY effective_priority ASC,
COALESCE(s.running_count, 0) ASC,
COALESCE(s.waiting_count, 0) ASC,
COALESCE(s.last_assigned_at, to_timestamp(0)) ASC,
m.created_at ASC`, exactModel, modelType, modelMatchKey, listOptions.CacheAffinityKeys, cacheAffinityStaleAfterSeconds(listOptions.CacheAffinityPolicy))
m.created_at ASC`, exactModel, modelType, listOptions.CacheAffinityKeys, cacheAffinityStaleAfterSeconds(listOptions.CacheAffinityPolicy))
if err != nil {
return nil, err
}
@@ -205,6 +183,7 @@ WHERE p.status = 'enabled'
var platformConfig []byte
var platformRetryPolicy []byte
var platformRateLimitPolicy []byte
var baseRateLimitPolicy []byte
var capabilities []byte
var capabilityOverride []byte
var baseBilling []byte
@@ -214,6 +193,7 @@ WHERE p.status = 'enabled'
var modelRetryPolicy []byte
var modelRateLimitPolicy []byte
var runtimePolicyOverride []byte
var rateLimitRuntimeOverride []byte
var runtimeRetryPolicy []byte
var runtimeRateLimitPolicy []byte
var autoDisablePolicy []byte
@@ -250,6 +230,7 @@ WHERE p.status = 'enabled'
&item.PlatformPricingRuleSetID,
&platformRetryPolicy,
&platformRateLimitPolicy,
&baseRateLimitPolicy,
&item.PlatformPriority,
&item.PlatformModelID,
&item.BaseModelID,
@@ -259,6 +240,7 @@ WHERE p.status = 'enabled'
&item.ModelAlias,
&item.ModelType,
&item.DisplayName,
&item.LegacyAliasUsed,
&capabilities,
&capabilityOverride,
&baseBilling,
@@ -271,8 +253,11 @@ WHERE p.status = 'enabled'
&permissionConfig,
&modelRetryPolicy,
&modelRateLimitPolicy,
&item.ModelRateLimitPolicyMode,
&item.RuntimePolicySetID,
&item.RuntimePolicyExplicit,
&runtimePolicyOverride,
&rateLimitRuntimeOverride,
&runtimeRetryPolicy,
&runtimeRateLimitPolicy,
&autoDisablePolicy,
@@ -301,6 +286,7 @@ WHERE p.status = 'enabled'
item.PlatformConfig = decodeObject(platformConfig)
item.PlatformRetryPolicy = decodeObject(platformRetryPolicy)
item.PlatformRateLimitPolicy = decodeObject(platformRateLimitPolicy)
item.BaseRateLimitPolicy = decodeObject(baseRateLimitPolicy)
item.Capabilities = decodeObject(capabilities)
item.CapabilityOverride = decodeObject(capabilityOverride)
item.BaseBillingConfig = decodeObject(baseBilling)
@@ -310,6 +296,7 @@ WHERE p.status = 'enabled'
item.ModelRetryPolicy = decodeObject(modelRetryPolicy)
item.ModelRateLimitPolicy = decodeObject(modelRateLimitPolicy)
item.RuntimePolicyOverride = decodeObject(runtimePolicyOverride)
item.RateLimitRuntimeOverride = decodeObject(rateLimitRuntimeOverride)
item.RuntimeRetryPolicy = decodeObject(runtimeRetryPolicy)
item.RuntimeRateLimitPolicy = decodeObject(runtimeRateLimitPolicy)
item.AutoDisablePolicy = decodeObject(autoDisablePolicy)
@@ -330,7 +317,15 @@ WHERE p.status = 'enabled'
LastObservedUnix: cacheLastObservedUnix,
})
applyRuntimeCandidateLoad(&item, runtimeCandidateLoadInput{
Policy: effectiveModelRateLimitPolicy(item.PlatformRateLimitPolicy, item.RuntimeRateLimitPolicy, item.RuntimePolicySetID, item.RuntimePolicyOverride, item.ModelRateLimitPolicy),
Policy: EffectiveRateLimitPolicy(EffectiveRateLimitPolicyInput{
BasePolicy: item.BaseRateLimitPolicy,
PlatformPolicy: item.PlatformRateLimitPolicy,
RuntimePolicy: item.RuntimeRateLimitPolicy,
RuntimePolicyExplicit: item.RuntimePolicyExplicit,
RuntimePolicyOverride: item.RateLimitRuntimeOverride,
ModelPolicy: item.ModelRateLimitPolicy,
ModelPolicyMode: item.ModelRateLimitPolicyMode,
}),
ConcurrentActive: concurrentActive,
QueuedWaiting: queuedWaiting,
RPMUsed: rpmUsed,
@@ -361,10 +356,31 @@ WHERE p.status = 'enabled'
if len(items) == 0 {
return nil, ErrNoModelCandidate
}
s.recordLegacyAliasUsage(ctx, exactModel, items)
sortRuntimeModelCandidates(items)
return items, nil
}
func (s *Store) recordLegacyAliasUsage(ctx context.Context, alias string, items []RuntimeModelCandidate) {
seen := map[string]bool{}
for _, item := range items {
if !item.LegacyAliasUsed || item.CanonicalModelKey == "" {
continue
}
key := item.CanonicalModelKey + "\x00" + item.ModelType
if seen[key] {
continue
}
seen[key] = true
_, _ = s.pool.Exec(ctx, `
INSERT INTO model_alias_usage_metrics (alias, canonical_model_key, model_type, hit_count)
VALUES ($1, $2, $3, 1)
ON CONFLICT (alias, canonical_model_key, model_type) DO UPDATE
SET hit_count = model_alias_usage_metrics.hit_count + 1,
last_used_at = now()`, alias, item.CanonicalModelKey, item.ModelType)
}
}
func (s *Store) GetRuntimeModelCandidateForVoiceCloneDeletion(ctx context.Context, platformModelID string, platformID string) (RuntimeModelCandidate, bool, error) {
platformModelID = strings.TrimSpace(platformModelID)
platformID = strings.TrimSpace(platformID)
@@ -487,14 +503,15 @@ func applyRuntimeCandidateCacheAffinity(candidate *RuntimeModelCandidate, option
key = options.CacheAffinityKey
}
affinity := RuntimeCandidateCacheAffinity{
Key: key,
RequestCount: input.RequestCount,
InputTokens: input.InputTokens,
CachedInputTokens: input.CachedInputTokens,
EMAHitRatio: boundedRatio(input.EMAHitRatio),
LastHitRatio: boundedRatio(input.LastHitRatio),
LastObservedUnix: input.LastObservedUnix,
AdjustedPriority: float64(candidate.PlatformPriority),
Key: key,
RequestCount: input.RequestCount,
InputTokens: input.InputTokens,
CachedInputTokens: input.CachedInputTokens,
EMAHitRatio: boundedRatio(input.EMAHitRatio),
LastHitRatio: boundedRatio(input.LastHitRatio),
LastObservedUnix: input.LastObservedUnix,
AdjustedPriority: float64(candidate.PlatformPriority),
MatchedPrefixDepth: cacheAffinityMatchedPrefixDepth(key, options.CacheAffinityKeys),
}
minSamples := cacheAffinityMinSamples(options.CacheAffinityPolicy)
hasObservedCachedHit := affinity.CachedInputTokens > 0 || affinity.LastHitRatio > 0
@@ -525,6 +542,32 @@ func applyRuntimeCandidateCacheAffinity(candidate *RuntimeModelCandidate, option
candidate.CacheAffinity = affinity
}
func cacheAffinityMatchedPrefixDepth(key string, lookup []string) int {
prefix := ""
switch {
case strings.HasPrefix(key, "prompt_lcp_v2:"):
prefix = "prompt_lcp_v2:"
case strings.HasPrefix(key, "prompt_lcp:"):
prefix = "prompt_lcp:"
default:
return 0
}
matchingKeys := make([]string, 0, len(lookup))
for _, candidateKey := range lookup {
if strings.HasPrefix(candidateKey, prefix) {
matchingKeys = append(matchingKeys, candidateKey)
}
}
for index, candidateKey := range matchingKeys {
if candidateKey == key {
return len(matchingKeys) - index + cacheAffinityMinimumPrefixDepth - 1
}
}
return 0
}
const cacheAffinityMinimumPrefixDepth = 2
func cacheAffinityPolicyEnabled(policy map[string]any, modelType string) bool {
if enabled, ok := policy["enabled"].(bool); ok && !enabled {
return false
@@ -699,6 +742,9 @@ func sortRuntimeModelCandidates(items []RuntimeModelCandidate) {
if aFull != bFull {
return !aFull
}
if items[i].PlatformPriority != items[j].PlatformPriority {
return items[i].PlatformPriority < items[j].PlatformPriority
}
if items[i].CacheAffinity.Applied != items[j].CacheAffinity.Applied {
return items[i].CacheAffinity.Applied
}
@@ -733,6 +779,32 @@ func sortRuntimeModelCandidates(items []RuntimeModelCandidate) {
}
return false
})
appliedCount := 0
for index := range items {
if items[index].CacheAffinity.Applied {
appliedCount++
}
}
for index := range items {
items[index].CacheAffinity.CandidateCount = appliedCount
}
if len(items) == 0 || items[0].CacheAffinity.Applied {
return
}
for index := 1; index < len(items); index++ {
if !items[index].CacheAffinity.Applied {
continue
}
switch {
case runtimeCandidateFull(items[index]) && !runtimeCandidateFull(items[0]):
items[0].CacheAffinity.OverrideReason = "capacity_tier_unavailable"
case items[index].PlatformPriority != items[0].PlatformPriority:
items[0].CacheAffinity.OverrideReason = "different_effective_priority_tier"
}
if items[0].CacheAffinity.OverrideReason != "" {
break
}
}
}
func runtimeCandidateFull(candidate RuntimeModelCandidate) bool {
@@ -741,14 +813,20 @@ func runtimeCandidateFull(candidate RuntimeModelCandidate) bool {
func (s *Store) modelCandidateCooldownError(ctx context.Context, model string, modelType string) (error, error) {
exactModel := strings.TrimSpace(model)
modelMatchKey := normalizeModelMatchKey(exactModel)
rows, err := s.pool.Query(ctx, `
SELECT p.name,
COALESCE(NULLIF(m.display_name, ''), NULLIF(m.model_alias, ''), m.model_name),
COALESCE(to_char(p.cooldown_until AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'), ''),
GREATEST(COALESCE(EXTRACT(EPOCH FROM p.cooldown_until - now()), 0), 0)::float8,
COALESCE(to_char(m.cooldown_until AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'), ''),
GREATEST(COALESCE(EXTRACT(EPOCH FROM m.cooldown_until - now()), 0), 0)::float8
var code string
var recoveryAt time.Time
var remainingSeconds float64
err := s.pool.QueryRow(ctx, `
SELECT CASE
WHEN COALESCE(m.cooldown_until, to_timestamp(0)) >= COALESCE(p.cooldown_until, to_timestamp(0))
THEN 'model_cooling_down'
ELSE 'platform_cooling_down'
END,
GREATEST(COALESCE(p.cooldown_until, to_timestamp(0)), COALESCE(m.cooldown_until, to_timestamp(0))) AS recovery_at,
GREATEST(
EXTRACT(EPOCH FROM GREATEST(COALESCE(p.cooldown_until, to_timestamp(0)), COALESCE(m.cooldown_until, to_timestamp(0))) - now()),
0
)::float8
FROM platform_models m
JOIN integration_platforms p ON p.id = m.platform_id
LEFT JOIN base_model_catalog b ON b.id = m.base_model_id
@@ -756,129 +834,51 @@ WHERE p.status = 'enabled'
AND p.deleted_at IS NULL
AND m.enabled = true
AND m.model_type @> jsonb_build_array($2::text)
AND (
b.invocation_name = $1::text
OR (b.id IS NULL AND m.model_name = $1::text)
OR EXISTS (
SELECT 1
FROM model_compatibility_aliases compatibility_alias
JOIN base_model_catalog alias_base ON alias_base.id = compatibility_alias.base_model_id
WHERE compatibility_alias.alias = $1::text
AND compatibility_alias.model_type = $2::text
AND compatibility_alias.active = true
AND (compatibility_alias.expires_at IS NULL OR compatibility_alias.expires_at > now())
AND (compatibility_alias.base_model_id = b.id OR alias_base.invocation_name = b.invocation_name)
)
)
AND (
(
$2::text IN ('audio_generate', 'text_to_speech', 'voice_clone')
AND (
m.model_alias = $1::text
OR m.model_name = $1::text
OR b.canonical_model_key = $1::text
OR b.provider_model_name = $1::text
OR (
NULLIF($3::text, '') IS NOT NULL
AND (
regexp_replace(COALESCE(m.model_alias, ''), '[[:space:]]+', '', 'g') = $3::text
OR regexp_replace(COALESCE(m.model_name, ''), '[[:space:]]+', '', 'g') = $3::text
OR regexp_replace(COALESCE(b.canonical_model_key, ''), '[[:space:]]+', '', 'g') = $3::text
OR regexp_replace(COALESCE(b.provider_model_name, ''), '[[:space:]]+', '', 'g') = $3::text
)
)
)
)
OR (
$2::text NOT IN ('audio_generate', 'text_to_speech', 'voice_clone')
AND (
(
COALESCE(m.model_alias, '') <> ''
AND (
m.model_alias = $1::text
OR (
NULLIF($3::text, '') IS NOT NULL
AND regexp_replace(COALESCE(m.model_alias, ''), '[[:space:]]+', '', 'g') = $3::text
)
)
)
OR (
COALESCE(m.model_alias, '') = ''
AND (
m.model_name = $1::text
OR b.canonical_model_key = $1::text
OR b.provider_model_name = $1::text
OR (
NULLIF($3::text, '') IS NOT NULL
AND (
regexp_replace(COALESCE(m.model_name, ''), '[[:space:]]+', '', 'g') = $3::text
OR regexp_replace(COALESCE(b.canonical_model_key, ''), '[[:space:]]+', '', 'g') = $3::text
OR regexp_replace(COALESCE(b.provider_model_name, ''), '[[:space:]]+', '', 'g') = $3::text
)
)
)
)
)
)
p.cooldown_until > now()
OR m.cooldown_until > now()
)
ORDER BY GREATEST(COALESCE(p.cooldown_until, to_timestamp(0)), COALESCE(m.cooldown_until, to_timestamp(0))) DESC,
p.priority ASC,
m.created_at ASC`, exactModel, modelType, modelMatchKey)
ORDER BY GREATEST(COALESCE(p.cooldown_until, to_timestamp(0)), COALESCE(m.cooldown_until, to_timestamp(0))) ASC,
p.priority ASC,
m.created_at ASC
LIMIT 1`, exactModel, modelType).Scan(&code, &recoveryAt, &remainingSeconds)
if errors.Is(err, pgx.ErrNoRows) {
return nil, nil
}
if err != nil {
return nil, err
}
defer rows.Close()
for rows.Next() {
var platformName string
var displayName string
var platformCooldownUntil string
var platformRemainingSeconds float64
var modelCooldownUntil string
var modelRemainingSeconds float64
if err := rows.Scan(
&platformName,
&displayName,
&platformCooldownUntil,
&platformRemainingSeconds,
&modelCooldownUntil,
&modelRemainingSeconds,
); err != nil {
return nil, err
}
if modelRemainingSeconds > 0 {
return &ModelCandidateUnavailableError{
Code: "model_cooling_down",
Message: cooldownErrorMessage("模型", displayName, modelRemainingSeconds, modelCooldownUntil),
}, nil
}
if platformRemainingSeconds > 0 {
return &ModelCandidateUnavailableError{
Code: "platform_cooling_down",
Message: cooldownErrorMessage("平台", platformName, platformRemainingSeconds, platformCooldownUntil),
}, nil
}
retryAfterSeconds := int(math.Ceil(remainingSeconds))
if retryAfterSeconds < 1 {
retryAfterSeconds = 1
}
if err := rows.Err(); err != nil {
return nil, err
message := "请求的模型暂时不可用，请稍后重试"
if code == "platform_cooling_down" {
message = "可用平台暂时不可用，请稍后重试"
}
return nil, nil
}
func cooldownErrorMessage(scope string, name string, remainingSeconds float64, cooldownUntil string) string {
name = strings.TrimSpace(name)
if name == "" {
name = "候选"
}
remainingMinutes := remainingSeconds / 60
if remainingMinutes < 0.1 {
remainingMinutes = 0.1
}
message := fmt.Sprintf("%s %s 冷却中，剩余 %.1f 分钟", scope, name, remainingMinutes)
if strings.TrimSpace(cooldownUntil) != "" {
message += "，预计恢复时间 " + cooldownUntil
}
return message
}
func normalizeModelMatchKey(value string) string {
value = strings.TrimSpace(value)
if value == "" {
return ""
}
var builder strings.Builder
builder.Grow(len(value))
for _, char := range value {
if unicode.IsSpace(char) {
continue
}
builder.WriteRune(char)
}
return builder.String()
recoveryAt = recoveryAt.UTC()
return &ModelCandidateUnavailableError{
Code: code,
Message: message,
RetryAfter: time.Duration(retryAfterSeconds) * time.Second,
RecoveryAt: recoveryAt,
Details: map[string]any{
"retryAfterSeconds": retryAfterSeconds,
"recoveryAt": recoveryAt.Format(time.RFC3339),
},
}, nil
}
apps/api/internal/store/candidates_test.go
+35 -16
View File
@@ -2,13 +2,6 @@ package store
import "testing"
func TestNormalizeModelMatchKeyRemovesWhitespace(t *testing.T) {
got := normalizeModelMatchKey(" doubao-5.0 图像\t编辑\n")
if got != "doubao-5.0图像编辑" {
t.Fatalf("expected whitespace-insensitive model key, got %q", got)
}
}
func TestNormalizeModelTypeListExpandsOmniVideoBaseCapabilities(t *testing.T) {
got := normalizeModelTypeList([]string{"omni_video"})
want := StringList{"video_generate", "image_to_video", "omni_video"}
@@ -127,7 +120,7 @@ func TestRuntimeCandidateSortingPrefersCacheAffinityWithinAvailableCandidates(t
},
{
PlatformID: "cache-affinity",
PlatformPriority: 20,
PlatformPriority: 10,
},
}
applyRuntimeCandidateCacheAffinity(&candidates[0], ListModelCandidatesOptions{
@@ -155,7 +148,7 @@ func TestRuntimeCandidateSortingPrefersCacheAffinityWithinAvailableCandidates(t
}
}
func TestRuntimeCandidateSortingPromotesObservedCacheHitAbovePriority(t *testing.T) {
func TestRuntimeCandidateSortingKeepsEffectivePriorityAheadOfCacheAffinity(t *testing.T) {
candidates := []RuntimeModelCandidate{
{
PlatformID: "volces-priority",
@@ -184,14 +177,17 @@ func TestRuntimeCandidateSortingPromotesObservedCacheHitAbovePriority(t *testing
sortRuntimeModelCandidates(candidates)
if candidates[0].PlatformID != "deepseek-cache-hit" || !candidates[0].CacheAffinity.Applied {
t.Fatalf("expected observed cache hit to outrank platform priority, got %+v", candidates)
if candidates[0].PlatformID != "volces-priority" || candidates[0].CacheAffinity.Applied {
t.Fatalf("expected effective priority tier to remain first, got %+v", candidates)
}
if candidates[0].CacheAffinity.Score <= 0 || candidates[0].CacheAffinity.Boost <= 0 {
t.Fatalf("expected observed cache hit to carry score and boost, got %+v", candidates[0].CacheAffinity)
if candidates[0].CacheAffinity.OverrideReason != "different_effective_priority_tier" {
t.Fatalf("expected priority override reason on selected candidate, got %+v", candidates[0].CacheAffinity)
}
if candidates[1].PlatformID != "volces-priority" || candidates[1].CacheAffinity.Applied {
t.Fatalf("expected priority-only candidate second, got %+v", candidates)
if candidates[1].PlatformID != "deepseek-cache-hit" || !candidates[1].CacheAffinity.Applied {
t.Fatalf("expected cache affinity candidate to remain visible as fallback, got %+v", candidates)
}
if candidates[1].CacheAffinity.Score <= 0 || candidates[1].CacheAffinity.Boost <= 0 {
t.Fatalf("expected observed cache hit to retain score and boost diagnostics, got %+v", candidates[1].CacheAffinity)
}
}
@@ -203,7 +199,7 @@ func TestRuntimeCandidateSortingOrdersMultipleCacheAffinityCandidatesByScore(t *
},
{
PlatformID: "higher-cache-score",
PlatformPriority: 50,
PlatformPriority: 1,
},
}
policy := map[string]any{"enabled": true, "minSamples": 1, "maxPriorityBoost": 20, "modelTypes": []any{"text_generate"}}
@@ -267,6 +263,29 @@ func TestRuntimeCandidateSortingKeepsFullCacheAffinityCandidateAvoided(t *testin
if candidates[1].PlatformID != "cache-affinity-full" || !candidates[1].LoadAvoided {
t.Fatalf("expected full cache-affinity candidate to remain avoided fallback, got %+v", candidates)
}
if candidates[0].CacheAffinity.OverrideReason != "capacity_tier_unavailable" {
t.Fatalf("expected capacity override reason on selected candidate, got %+v", candidates[0].CacheAffinity)
}
}
func TestCacheAffinityMatchedPrefixDepthIgnoresLegacyCompatibilityKeys(t *testing.T) {
lookup := []string{
"prompt_lcp_v2:depth-four",
"prompt_lcp_v2:depth-three",
"prompt_lcp_v2:depth-two",
"prompt_lcp:depth-four",
"prompt_lcp:depth-three",
"prompt_lcp:depth-two",
}
if depth := cacheAffinityMatchedPrefixDepth("prompt_lcp_v2:depth-three", lookup); depth != 3 {
t.Fatalf("V2 matched prefix depth=%d want=3", depth)
}
if depth := cacheAffinityMatchedPrefixDepth("prompt_lcp:depth-two", lookup); depth != 2 {
t.Fatalf("legacy matched prefix depth=%d want=2", depth)
}
if depth := cacheAffinityMatchedPrefixDepth("explicit:key", lookup); depth != 0 {
t.Fatalf("explicit key prefix depth=%d want=0", depth)
}
}
func TestDefaultRunnerPriorityDemotePolicyUsesAutoMode(t *testing.T) {
apps/api/internal/store/compatibility_tasks.go
+11 -14
View File
@@ -2,7 +2,6 @@ package store
import (
"context"
"encoding/json"
"strings"
)
@@ -16,25 +15,19 @@ type CompatibilitySubmission struct {
}
func (s *Store) SetTaskCompatibilitySubmission(ctx context.Context, taskID string, submission CompatibilitySubmission) error {
headers, _ := json.Marshal(emptyObjectIfNil(submission.Headers))
body, _ := json.Marshal(emptyObjectIfNil(submission.Body))
_, err := s.pool.Exec(ctx, `
UPDATE gateway_tasks
SET compatibility_protocol = NULLIF($2, ''),
compatibility_public_id = NULLIF($3, ''),
compatibility_source_protocol = NULLIF($4, ''),
compatibility_submit_http_status = NULLIF($5, 0),
compatibility_submit_headers = $6::jsonb,
compatibility_submit_body = $7::jsonb,
SET compatibility_protocol = COALESCE(NULLIF($2, ''), compatibility_protocol),
compatibility_public_id = NULL,
compatibility_source_protocol = COALESCE(NULLIF($3, ''), compatibility_source_protocol),
compatibility_submit_http_status = NULL,
compatibility_submit_headers = '{}'::jsonb,
compatibility_submit_body = '{}'::jsonb,
updated_at = now()
WHERE id = $1::uuid`,
taskID,
strings.TrimSpace(submission.TargetProtocol),
strings.TrimSpace(submission.PublicID),
strings.TrimSpace(submission.SourceProtocol),
submission.HTTPStatus,
string(headers),
string(body),
)
return err
}
@@ -44,7 +37,11 @@ func (s *Store) GetCompatibilityTask(ctx context.Context, protocol string, publi
SELECT `+gatewayTaskColumns+`
FROM gateway_tasks
WHERE compatibility_protocol = $1
AND compatibility_public_id = $2
AND (
id::text = $2
OR remote_task_id = $2
OR compatibility_public_id = $2
)
LIMIT 1`, strings.TrimSpace(protocol), strings.TrimSpace(publicID)))
if err != nil {
return GatewayTask{}, err
apps/api/internal/store/conversations.go
+10 -2
View File
@@ -41,7 +41,7 @@ func (s *Store) EnsureConversation(ctx context.Context, user *auth.User, convers
if userID == "" {
userID = "anonymous"
}
metadataJSON, _ := json.Marshal(emptyObjectIfNil(metadata))
metadataJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(metadata)))
var conversationID string
err := s.pool.QueryRow(ctx, `
INSERT INTO gateway_conversations (user_id, gateway_user_id, conversation_key, metadata)
@@ -67,7 +67,15 @@ func (s *Store) UpsertConversationMessages(ctx context.Context, conversationID s
refs := make([]TaskMessageRefInput, 0, len(messages))
newCount := 0
for index, message := range messages {
snapshotJSON, _ := json.Marshal(emptyObjectIfNil(message.Snapshot))
snapshotReport := sanitizeJSONForStorageWithReport(emptyObjectIfNil(message.Snapshot))
if snapshotReport.BinaryCount > 0 {
return nil, 0, &taskPayloadBinaryError{
target: ErrTaskRequestBinaryNotMaterialized,
code: "request_binary_not_materialized",
count: snapshotReport.BinaryCount,
}
}
snapshotJSON, _ := json.Marshal(snapshotReport.Value)
var messageID string
var inserted bool
if err := tx.QueryRow(ctx, `
apps/api/internal/store/file_storage_channels.go
+1 -1
View File
@@ -578,7 +578,7 @@ func normalizeFileStorageScene(scene string) string {
}
func defaultFileStorageScenes() []string {
return []string{FileStorageSceneUpload, FileStorageSceneImageResult}
return []string{FileStorageSceneUpload, FileStorageSceneImageResult, FileStorageSceneRequestAsset}
}
func defaultFileStorageRetryPolicyIfEmpty(policy map[string]any) map[string]any {
apps/api/internal/store/file_storage_channels_test.go
+17
View File
@@ -0,0 +1,17 @@
package store
import (
"reflect"
"testing"
)
func TestDefaultFileStorageScenesIncludeCrossNodeRequestAssets(t *testing.T) {
want := []string{
FileStorageSceneUpload,
FileStorageSceneImageResult,
FileStorageSceneRequestAsset,
}
if got := defaultFileStorageScenes(); !reflect.DeepEqual(got, want) {
t.Fatalf("default file storage scenes = %#v, want %#v", got, want)
}
}
apps/api/internal/store/json_storage_guard.go
+344
View File
@@ -0,0 +1,344 @@
package store
import (
"crypto/sha256"
"encoding/base64"
"encoding/json"
"errors"
"fmt"
"strconv"
"strings"
)
const (
storageBinaryPrefixChars = 16
storageGenericBase64MinLength = 4096
storageInvalidBinaryMinLength = 512
storageJSONSanitizerMaxDepth = 64
storagePlaceholderPrefix = "[GatewayBinary:v1;"
storageBufferObjectType = "buffer"
storageDataURLPrefix = "data:"
storageDataURLBase64Marker = ";base64"
storageDataURLMaxContentType = 64
)
var (
ErrTaskRequestBinaryNotMaterialized = errors.New("task request binary payload was not materialized")
ErrTaskResultBinaryNotMaterialized = errors.New("task result binary payload was not materialized")
)
type taskPayloadBinaryError struct {
target error
code string
count int
}
func (e *taskPayloadBinaryError) Error() string {
return fmt.Sprintf("%s: detected %d inline binary value(s)", e.code, e.count)
}
func (e *taskPayloadBinaryError) ErrorCode() string {
return e.code
}
func (e *taskPayloadBinaryError) Is(target error) bool {
return target == e.target
}
type storageSanitizeReport struct {
Value any
BinaryCount int
}
// sanitizeJSONForStorage is the final task-domain persistence guard. It always
// returns a detached JSON-compatible value and replaces inline binary payloads
// with a bounded, deterministic placeholder.
func sanitizeJSONForStorage(value any) any {
return sanitizeJSONForStorageWithReport(value).Value
}
func sanitizeJSONForStorageWithReport(value any) storageSanitizeReport {
next, count := sanitizeJSONStorageValue(value, nil, 0)
return storageSanitizeReport{Value: next, BinaryCount: count}
}
func sanitizeJSONStorageValue(value any, path []string, depth int) (any, int) {
if depth >= storageJSONSanitizerMaxDepth {
return "[JSON,max-depth]", 0
}
switch typed := value.(type) {
case map[string]any:
if payload, contentType, ok := storageBufferObjectBytes(typed); ok {
return storageBinaryPlaceholder(payload, contentType, "buffer"), 1
}
next := make(map[string]any, len(typed))
count := 0
for key, child := range typed {
sanitized, childCount := sanitizeJSONStorageValue(child, appendStoragePath(path, key), depth+1)
next[key] = sanitized
count += childCount
}
return next, count
case []any:
if storagePathIsBinary(path) {
if payload, ok := storageNumberArrayBytes(typed); ok {
return storageBinaryPlaceholder(payload, "", "buffer"), 1
}
}
next := make([]any, len(typed))
count := 0
for index, child := range typed {
sanitized, childCount := sanitizeJSONStorageValue(child, path, depth+1)
next[index] = sanitized
count += childCount
}
return next, count
case []byte:
return storageBinaryPlaceholder(typed, "", "buffer"), 1
case string:
if placeholder, ok := storageStringPlaceholder(typed, path); ok {
return placeholder, 1
}
return typed, 0
default:
return value, 0
}
}
func appendStoragePath(path []string, key string) []string {
next := make([]string, len(path)+1)
copy(next, path)
next[len(path)] = key
return next
}
func storageStringPlaceholder(value string, path []string) (string, bool) {
raw := strings.TrimSpace(value)
if raw == "" || strings.HasPrefix(raw, storagePlaceholderPrefix) {
return "", false
}
encoded, contentType, dataURL := storageBase64StringParts(raw)
strict := dataURL || storagePathIsBinary(path)
if !strict && len(encoded) < storageGenericBase64MinLength {
return "", false
}
payload, ok := storageDecodeBase64(encoded)
if ok {
encoding := "raw"
if dataURL {
encoding = "data-uri"
}
return storageBinaryPlaceholder(payload, contentType, encoding), true
}
if strict && len(raw) >= storageInvalidBinaryMinLength {
return storageBinaryPlaceholder([]byte(raw), contentType, "raw"), true
}
return "", false
}
func storageBase64StringParts(value string) (encoded string, contentType string, dataURL bool) {
if !strings.HasPrefix(strings.ToLower(value), storageDataURLPrefix) {
return value, "", false
}
prefix, payload, ok := strings.Cut(value, ",")
if !ok || !strings.Contains(strings.ToLower(prefix), storageDataURLBase64Marker) {
return value, "", false
}
mediaType := strings.TrimSpace(prefix[len(storageDataURLPrefix):])
if before, _, found := strings.Cut(mediaType, ";"); found {
mediaType = before
}
if len(mediaType) > storageDataURLMaxContentType || !storageSafeContentType(mediaType) {
mediaType = ""
}
return payload, strings.ToLower(mediaType), true
}
func storageSafeContentType(value string) bool {
if value == "" {
return true
}
for _, char := range value {
switch {
case char >= 'a' && char <= 'z':
case char >= 'A' && char <= 'Z':
case char >= '0' && char <= '9':
case char == '/', char == '.', char == '+', char == '-':
default:
return false
}
}
return true
}
func storageDecodeBase64(value string) ([]byte, bool) {
normalized := removeStorageASCIIWhitespace(value)
if normalized == "" {
return nil, false
}
for _, encoding := range []*base64.Encoding{
base64.StdEncoding,
base64.RawStdEncoding,
base64.URLEncoding,
base64.RawURLEncoding,
} {
payload, err := encoding.DecodeString(normalized)
if err != nil || len(payload) == 0 {
continue
}
canonical := encoding.EncodeToString(payload)
if strings.TrimRight(normalized, "=") == strings.TrimRight(canonical, "=") {
return payload, true
}
}
return nil, false
}
func removeStorageASCIIWhitespace(value string) string {
return strings.Map(func(char rune) rune {
switch char {
case ' ', '\n', '\r', '\t':
return -1
default:
return char
}
}, value)
}
func storagePathIsBinary(path []string) bool {
if len(path) == 0 {
return false
}
key := normalizeStorageBinaryKey(path[len(path)-1])
if storageBinaryKey(key) {
return true
}
if len(path) < 2 {
return false
}
parent := normalizeStorageBinaryKey(path[len(path)-2])
return (parent == "inlinedata" || parent == "binary" || parent == "media") &&
(key == "data" || key == "content")
}
func normalizeStorageBinaryKey(value string) string {
return strings.Map(func(char rune) rune {
switch {
case char >= 'a' && char <= 'z':
return char
case char >= 'A' && char <= 'Z':
return char + ('a' - 'A')
case char >= '0' && char <= '9':
return char
default:
return -1
}
}, value)
}
func storageBinaryKey(key string) bool {
return key == "b64" ||
key == "b64json" ||
key == "base64" ||
key == "buffer" ||
key == "bytes" ||
strings.Contains(key, "base64") ||
strings.Contains(key, "buffer") ||
strings.Contains(key, "binary") ||
strings.HasSuffix(key, "b64") ||
strings.HasSuffix(key, "bytes")
}
func storageBufferObjectBytes(value map[string]any) ([]byte, string, bool) {
kind, _ := value["type"].(string)
if normalizeStorageBinaryKey(kind) != storageBufferObjectType {
return nil, "", false
}
contentType := firstNonEmpty(
stringFromAny(value["mime_type"]),
stringFromAny(value["mimeType"]),
stringFromAny(value["contentType"]),
)
switch data := value["data"].(type) {
case []byte:
if len(data) == 0 {
return nil, "", false
}
return append([]byte(nil), data...), contentType, true
case []any:
payload, ok := storageNumberArrayBytes(data)
return payload, contentType, ok
default:
return nil, "", false
}
}
func storageNumberArrayBytes(values []any) ([]byte, bool) {
if len(values) == 0 {
return nil, false
}
payload := make([]byte, len(values))
for index, value := range values {
next, ok := storageByteFromAny(value)
if !ok {
return nil, false
}
payload[index] = next
}
return payload, true
}
func storageByteFromAny(value any) (byte, bool) {
switch typed := value.(type) {
case byte:
return typed, true
case int:
if typed >= 0 && typed <= 255 {
return byte(typed), true
}
case int32:
if typed >= 0 && typed <= 255 {
return byte(typed), true
}
case int64:
if typed >= 0 && typed <= 255 {
return byte(typed), true
}
case float64:
asInt := int(typed)
if typed == float64(asInt) && asInt >= 0 && asInt <= 255 {
return byte(asInt), true
}
case json.Number:
asInt, err := strconv.ParseInt(string(typed), 10, 16)
if err == nil && asInt >= 0 && asInt <= 255 {
return byte(asInt), true
}
}
return 0, false
}
func storageBinaryPlaceholder(payload []byte, contentType string, encoding string) string {
digest := sha256.Sum256(payload)
prefix := base64.StdEncoding.EncodeToString(payload)
if len(prefix) > storageBinaryPrefixChars {
prefix = prefix[:storageBinaryPrefixChars]
}
contentType = strings.ToLower(strings.TrimSpace(contentType))
if contentType == "" || len(contentType) > 32 || !storageSafeContentType(contentType) {
contentType = "application/octet-stream"
}
switch encoding {
case "data-uri", "buffer":
default:
encoding = "raw"
}
return fmt.Sprintf(
"[GatewayBinary:v1;prefix=%s;sha256=%x;bytes=%d;mime=%s;encoding=%s]",
prefix,
digest,
len(payload),
contentType,
encoding,
)
}
apps/api/internal/store/json_storage_guard_test.go
+103
View File
@@ -0,0 +1,103 @@
package store
import (
"encoding/base64"
"encoding/json"
"strings"
"testing"
)
func TestSanitizeJSONForStorageReplacesBinaryWithoutMutatingInput(t *testing.T) {
payload := []byte("shared binary payload")
encoded := base64.StdEncoding.EncodeToString(payload)
input := map[string]any{
"data": []any{
map[string]any{"b64_json": encoded},
},
"buffer": map[string]any{
"type": "Buffer",
"data": []any{float64(1), float64(2), float64(3)},
},
"text": "ordinary text",
}
report := sanitizeJSONForStorageWithReport(input)
if report.BinaryCount != 2 {
t.Fatalf("binary count = %d, want 2", report.BinaryCount)
}
next := report.Value.(map[string]any)
data := next["data"].([]any)
placeholder := data[0].(map[string]any)["b64_json"].(string)
if !strings.HasPrefix(placeholder, storagePlaceholderPrefix) {
t.Fatalf("unexpected placeholder: %q", placeholder)
}
if len(placeholder) > 200 {
t.Fatalf("placeholder exceeds 200 bytes: %d", len(placeholder))
}
if !strings.Contains(placeholder, ";prefix="+encoded[:16]+";") {
t.Fatalf("placeholder should retain the bounded Base64 prefix: %q", placeholder)
}
if got := input["data"].([]any)[0].(map[string]any)["b64_json"]; got != encoded {
t.Fatalf("sanitizer mutated input: %v", got)
}
if next["text"] != "ordinary text" {
t.Fatalf("ordinary text changed: %+v", next)
}
}
func TestSanitizeJSONForStorageUsesDecodedBytesForEquivalentRepresentations(t *testing.T) {
payload := []byte("equivalent payload")
encoded := base64.StdEncoding.EncodeToString(payload)
report := sanitizeJSONForStorageWithReport(map[string]any{
"rawBase64": encoded,
"dataURL": "data:image/png;base64," + encoded,
"bytes": []any{float64('e'), float64('q'), float64('u'), float64('i'), float64('v'), float64('a'), float64('l'), float64('e'), float64('n'), float64('t'), float64(' '), float64('p'), float64('a'), float64('y'), float64('l'), float64('o'), float64('a'), float64('d')},
})
if report.BinaryCount != 3 {
t.Fatalf("binary count = %d, want 3", report.BinaryCount)
}
next := report.Value.(map[string]any)
hashes := map[string]struct{}{}
for _, key := range []string{"rawBase64", "dataURL", "bytes"} {
value := next[key].(string)
hashStart := strings.Index(value, ";sha256=")
hashEnd := strings.Index(value[hashStart+1:], ";bytes=")
if hashStart < 0 || hashEnd < 0 {
t.Fatalf("missing hash in %s placeholder: %q", key, value)
}
hashes[value[hashStart+8:hashStart+1+hashEnd]] = struct{}{}
}
if len(hashes) != 1 {
t.Fatalf("equivalent binary values produced different hashes: %+v", next)
}
}
func TestSanitizeJSONForStorageAvoidsShortGenericTextFalsePositive(t *testing.T) {
input := map[string]any{"message": "YWJjZA==", "count": json.Number("12")}
report := sanitizeJSONForStorageWithReport(input)
if report.BinaryCount != 0 {
t.Fatalf("short generic Base64-like text should not be sanitized: %+v", report)
}
}
func TestSanitizeJSONForStorageRecognizesExistingPlaceholder(t *testing.T) {
value := storageBinaryPlaceholder([]byte("payload"), "image/png", "raw")
report := sanitizeJSONForStorageWithReport(map[string]any{"b64_json": value})
if report.BinaryCount != 0 {
t.Fatalf("existing placeholder should pass the final guard: %+v", report)
}
}
func TestStorageBinaryPlaceholderIsAlwaysBounded(t *testing.T) {
value := storageBinaryPlaceholder(
[]byte("payload"),
"application/vnd.a-very-long-provider-specific-generated-binary-result+json",
"data-uri",
)
if len(value) > 200 {
t.Fatalf("placeholder exceeds 200 bytes: %d %q", len(value), value)
}
if !strings.Contains(value, ";mime=application/octet-stream;") {
t.Fatalf("oversized MIME should use the bounded fallback: %q", value)
}
}
apps/api/internal/store/model_identity_integration_test.go
+119
View File
@@ -0,0 +1,119 @@
package store
import (
"context"
"errors"
"os"
"strings"
"testing"
"time"
)
func TestModelIdentityRoutingAndDeleteProtection(t *testing.T) {
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
if databaseURL == "" {
t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run model identity PostgreSQL integration tests")
}
ctx := context.Background()
db, err := Connect(ctx, databaseURL)
if err != nil {
t.Fatalf("connect test database: %v", err)
}
defer db.Close()
suffix := time.Now().UTC().Format("20060102150405.000000000")
invocationName := "official-model-" + suffix
legacyAlias := "Legacy Model " + suffix
displayName := "Model Card " + suffix
baseModel, err := db.CreateBaseModel(ctx, BaseModelInput{
ProviderKey: "openai",
CanonicalModelKey: "openai:" + invocationName,
InvocationName: invocationName,
ProviderModelName: "provider/" + invocationName,
ModelType: StringList{"text_generate"},
DisplayName: displayName,
LegacyAliases: StringList{legacyAlias},
Status: "active",
})
if err != nil {
t.Fatalf("create base model: %v", err)
}
defer func() {
_, _ = db.pool.Exec(ctx, `DELETE FROM platform_models WHERE base_model_id = $1::uuid`, baseModel.ID)
_, _ = db.pool.Exec(ctx, `DELETE FROM base_model_catalog WHERE id = $1::uuid`, baseModel.ID)
}()
var platformID string
if err := db.pool.QueryRow(ctx, `
SELECT id::text
FROM integration_platforms
WHERE status = 'enabled' AND deleted_at IS NULL
ORDER BY created_at ASC
LIMIT 1`).Scan(&platformID); err != nil {
t.Fatalf("find enabled platform: %v", err)
}
var platformModelID string
if err := db.pool.QueryRow(ctx, `
INSERT INTO platform_models (
platform_id, base_model_id, model_name, provider_model_name, model_alias,
model_type, display_name, enabled
)
VALUES ($1::uuid, $2::uuid, $3, $4, $3, '["text_generate"]'::jsonb, $5, true)
RETURNING id::text`, platformID, baseModel.ID, invocationName, "provider/"+invocationName, displayName).Scan(&platformModelID); err != nil {
t.Fatalf("bind platform model: %v", err)
}
if err := db.DeleteBaseModel(ctx, baseModel.ID); !errors.Is(err, ErrBaseModelInUse) {
t.Fatalf("bound base model deletion should be protected: %v", err)
}
officialCandidates, err := db.ListModelCandidates(ctx, invocationName, "text_generate", nil)
if err != nil || len(officialCandidates) != 1 || officialCandidates[0].PlatformModelID != platformModelID {
t.Fatalf("official invocation should resolve the bound source: candidates=%+v err=%v", officialCandidates, err)
}
legacyCandidates, err := db.ListModelCandidates(ctx, legacyAlias, "text_generate", nil)
if err != nil || len(legacyCandidates) != 1 || !legacyCandidates[0].LegacyAliasUsed {
t.Fatalf("registered legacy alias should resolve and be marked: candidates=%+v err=%v", legacyCandidates, err)
}
if _, err := db.ListModelCandidates(ctx, displayName, "text_generate", nil); !errors.Is(err, ErrNoModelCandidate) {
t.Fatalf("display name must not become an invocation alias: %v", err)
}
var legacyHits int64
if err := db.pool.QueryRow(ctx, `
SELECT hit_count
FROM model_alias_usage_metrics
WHERE alias = $1 AND canonical_model_key = $2 AND model_type = 'text_generate'`, legacyAlias, baseModel.CanonicalModelKey).Scan(&legacyHits); err != nil {
t.Fatalf("read legacy alias metric: %v", err)
}
if legacyHits != 1 {
t.Fatalf("expected one legacy alias hit, got %d", legacyHits)
}
deprecatedInvocation := "deprecated-model-" + suffix
deprecatedBase, err := db.CreateBaseModel(ctx, BaseModelInput{
ProviderKey: "openai",
CanonicalModelKey: "openai:" + deprecatedInvocation,
InvocationName: deprecatedInvocation,
ProviderModelName: deprecatedInvocation,
ModelType: StringList{"text_generate"},
DisplayName: "Deprecated Model",
Status: "deprecated",
})
if err != nil {
t.Fatalf("create deprecated base model: %v", err)
}
defer func() {
_, _ = db.pool.Exec(ctx, `DELETE FROM base_model_catalog WHERE id = $1::uuid`, deprecatedBase.ID)
}()
if _, err := db.CreatePlatformModel(ctx, CreatePlatformModelInput{
PlatformID: platformID,
BaseModelID: deprecatedBase.ID,
ModelName: deprecatedInvocation,
ProviderModelName: deprecatedInvocation,
ModelType: StringList{"text_generate"},
}); !errors.Is(err, ErrInvalidPlatformModelConfiguration) {
t.Fatalf("deprecated base model must reject new bindings: %v", err)
}
}
apps/api/internal/store/oidc_users_integration_test.go
+16
View File
@@ -307,6 +307,22 @@ func applyOIDCJITTestMigrations(t *testing.T, ctx context.Context, databaseURL s
if err != nil {
t.Fatalf("read migration %s: %v", version, err)
}
migrationSQL := string(migration)
const noTransactionMarker = "-- easyai:migration:no-transaction"
if strings.HasPrefix(strings.TrimSpace(migrationSQL), noTransactionMarker) {
migrationSQL = strings.TrimSpace(strings.TrimPrefix(strings.TrimSpace(migrationSQL), noTransactionMarker))
for _, statement := range strings.Split(migrationSQL, "-- easyai:migration:statement") {
if statement = strings.TrimSpace(statement); statement != "" {
if _, err := pool.Exec(ctx, statement); err != nil {
t.Fatalf("apply non-transaction migration %s: %v", version, err)
}
}
}
if _, err := pool.Exec(ctx, `INSERT INTO schema_migrations(version) VALUES($1)`, version); err != nil {
t.Fatalf("record non-transaction migration %s: %v", version, err)
}
continue
}
tx, err := pool.Begin(ctx)
if err != nil {
t.Fatalf("begin migration %s: %v", version, err)
apps/api/internal/store/platform_models.go
+33 -24
View File
@@ -18,6 +18,7 @@ type modelCatalogSnapshot struct {
ID string
ProviderKey string
CanonicalModelKey string
InvocationName string
ProviderModelName string
ModelType StringList
DisplayName string
@@ -27,6 +28,7 @@ type modelCatalogSnapshot struct {
DefaultRateLimitPolicy map[string]any
RuntimePolicySetID string
RuntimePolicyOverride map[string]any
Status string
}
func (s *Store) CreatePlatformModel(ctx context.Context, input CreatePlatformModelInput) (PlatformModel, error) {
@@ -95,6 +97,18 @@ func (s *Store) createPlatformModel(ctx context.Context, q platformModelQuerier,
if err != nil && !IsNotFound(err) {
return PlatformModel{}, err
}
if base.ID != "" && base.Status != "" && base.Status != "active" {
var existingBinding bool
if err := q.QueryRow(ctx, `
SELECT EXISTS (
SELECT 1 FROM platform_models WHERE platform_id = $1::uuid AND base_model_id = $2::uuid
)`, input.PlatformID, base.ID).Scan(&existingBinding); err != nil {
return PlatformModel{}, err
}
if !existingBinding {
return PlatformModel{}, fmt.Errorf("%w: base model %q is %s and cannot be newly bound", ErrInvalidPlatformModelConfiguration, base.InvocationName, base.Status)
}
}
if len(input.ModelType) == 0 {
input.ModelType = base.ModelType
}
@@ -102,7 +116,9 @@ func (s *Store) createPlatformModel(ctx context.Context, q platformModelQuerier,
if len(input.ModelType) == 0 {
input.ModelType = StringList{"text_generate"}
}
if input.ModelName == "" {
if base.InvocationName != "" {
input.ModelName = base.InvocationName
} else if input.ModelName == "" {
input.ModelName = base.ProviderModelName
}
if input.ProviderModelName == "" {
@@ -127,18 +143,13 @@ func (s *Store) createPlatformModel(ctx context.Context, q platformModelQuerier,
// soon as the base pricing rule changes and can mask the authoritative rule.
billingConfig := input.BillingConfig
explicitRuntimePolicySetID := strings.TrimSpace(input.RuntimePolicySetID)
rateLimitPolicyMode := NormalizeRateLimitPolicyMode(input.RateLimitPolicyMode, input.RateLimitPolicy != nil)
rateLimitPolicy := input.RateLimitPolicy
if len(rateLimitPolicy) == 0 && explicitRuntimePolicySetID == "" {
rateLimitPolicy = base.DefaultRateLimitPolicy
if rateLimitPolicyMode == RateLimitPolicyModeInherit {
rateLimitPolicy = nil
}
runtimePolicySetID := explicitRuntimePolicySetID
if runtimePolicySetID == "" {
runtimePolicySetID = base.RuntimePolicySetID
}
runtimePolicyOverride := input.RuntimePolicyOverride
if len(runtimePolicyOverride) == 0 {
runtimePolicyOverride = base.RuntimePolicyOverride
}
capabilityOverrideJSON, _ := json.Marshal(emptyObjectIfNil(input.CapabilityOverride))
capabilitiesJSON, _ := json.Marshal(emptyObjectIfNil(capabilities))
@@ -173,14 +184,14 @@ func (s *Store) createPlatformModel(ctx context.Context, q platformModelQuerier,
INSERT INTO platform_models (
platform_id, base_model_id, model_name, provider_model_name, model_alias, model_type, display_name,
capability_override, capabilities, pricing_mode, discount_factor,
pricing_rule_set_id, billing_config_override, billing_config, permission_config, retry_policy, rate_limit_policy,
pricing_rule_set_id, billing_config_override, billing_config, permission_config, retry_policy, rate_limit_policy, rate_limit_policy_mode,
runtime_policy_set_id, runtime_policy_override, enabled
)
VALUES (
$1::uuid, $2::uuid, $3, NULLIF($4, ''), NULLIF($5, ''), $6::jsonb, $7,
$8::jsonb, $9::jsonb, $10, $11::numeric,
NULLIF($12, '')::uuid, $13::jsonb, $14::jsonb, $15::jsonb, $16::jsonb, $17::jsonb,
NULLIF($18, '')::uuid, $19::jsonb, true
NULLIF($12, '')::uuid, $13::jsonb, $14::jsonb, $15::jsonb, $16::jsonb, $17::jsonb, $18,
NULLIF($19, '')::uuid, $20::jsonb, true
)
ON CONFLICT (platform_id, model_name) DO UPDATE
SET base_model_id = EXCLUDED.base_model_id,
@@ -197,6 +208,7 @@ SET base_model_id = EXCLUDED.base_model_id,
permission_config = EXCLUDED.permission_config,
retry_policy = EXCLUDED.retry_policy,
rate_limit_policy = EXCLUDED.rate_limit_policy,
rate_limit_policy_mode = EXCLUDED.rate_limit_policy_mode,
runtime_policy_set_id = EXCLUDED.runtime_policy_set_id,
runtime_policy_override = EXCLUDED.runtime_policy_override,
enabled = true,
@@ -205,7 +217,7 @@ RETURNING id::text, platform_id::text, COALESCE(base_model_id::text, ''), model_
COALESCE(NULLIF(provider_model_name, ''), model_name), COALESCE(model_alias, ''), model_type, display_name, capability_override,
capabilities, pricing_mode, COALESCE(discount_factor, 0)::float8,
COALESCE(pricing_rule_set_id::text, ''), billing_config_override, billing_config,
permission_config, retry_policy, rate_limit_policy, COALESCE(runtime_policy_set_id::text, ''), runtime_policy_override,
permission_config, retry_policy, rate_limit_policy, rate_limit_policy_mode, COALESCE(runtime_policy_set_id::text, ''), runtime_policy_override,
COALESCE(to_char(cooldown_until AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'), ''),
enabled, created_at, updated_at`,
input.PlatformID,
@@ -225,6 +237,7 @@ RETURNING id::text, platform_id::text, COALESCE(base_model_id::text, ''), model_
string(permissionJSON),
string(retryJSON),
string(rateLimitJSON),
rateLimitPolicyMode,
runtimePolicySetID,
string(runtimePolicyOverrideJSON),
).Scan(
@@ -246,6 +259,7 @@ RETURNING id::text, platform_id::text, COALESCE(base_model_id::text, ''), model_
&permissionBytes,
&retryPolicyBytes,
&rateLimitPolicyBytes,
&model.RateLimitPolicyMode,
&model.RuntimePolicySetID,
&runtimePolicyOverrideBytes,
&model.CooldownUntil,
@@ -370,9 +384,9 @@ func (s *Store) lookupBaseModel(ctx context.Context, q platformModelQuerier, id
var runtimePolicyOverride []byte
var modelTypeBytes []byte
err := q.QueryRow(ctx, `
SELECT id::text, provider_key, canonical_model_key, provider_model_name, model_type, display_name,
SELECT id::text, provider_key, canonical_model_key, invocation_name, provider_model_name, model_type, display_name,
capabilities, base_billing_config, COALESCE(pricing_rule_set_id::text, ''), default_rate_limit_policy,
COALESCE(runtime_policy_set_id::text, ''), runtime_policy_override
COALESCE(runtime_policy_set_id::text, ''), runtime_policy_override, status
FROM base_model_catalog
WHERE ($1 <> '' AND id = NULLIF($1, '')::uuid)
OR ($2 <> '' AND canonical_model_key = $2)
@@ -382,6 +396,7 @@ LIMIT 1`, strings.TrimSpace(id), strings.TrimSpace(canonicalKey), strings.TrimSp
&item.ID,
&item.ProviderKey,
&item.CanonicalModelKey,
&item.InvocationName,
&item.ProviderModelName,
&modelTypeBytes,
&item.DisplayName,
@@ -391,6 +406,7 @@ LIMIT 1`, strings.TrimSpace(id), strings.TrimSpace(canonicalKey), strings.TrimSp
&rateLimitPolicy,
&item.RuntimePolicySetID,
&runtimePolicyOverride,
&item.Status,
)
if err != nil {
if err == pgx.ErrNoRows {
@@ -407,15 +423,8 @@ LIMIT 1`, strings.TrimSpace(id), strings.TrimSpace(canonicalKey), strings.TrimSp
}
func normalizePlatformModelAlias(alias string, base modelCatalogSnapshot) string {
alias = strings.TrimSpace(alias)
if alias == "" {
alias = firstNonEmpty(base.ProviderModelName, base.DisplayName, base.CanonicalModelKey)
}
if base.ProviderKey != "" {
alias = strings.TrimPrefix(alias, base.ProviderKey+":")
}
if alias == base.CanonicalModelKey {
alias = stripAliasPrefix(alias)
if base.InvocationName != "" {
return base.InvocationName
}
return strings.TrimSpace(alias)
}
apps/api/internal/store/postgres.go
+78 -34
View File
@@ -69,6 +69,8 @@ var (
ErrTaskExecutionFinished = errors.New("task execution already finished")
ErrTaskExecutionManualReview = errors.New("task execution requires manual review")
ErrProtectedDefault = errors.New("protected default resource cannot be deleted")
ErrBaseModelInUse = errors.New("base model is still referenced by platform models")
ErrModelAliasConflict = errors.New("model compatibility alias conflicts with another model identity")
ErrUserAlreadyExists = errors.New("user already exists")
ErrWeakPassword = errors.New("password must be at least 8 characters")
)
@@ -217,20 +219,24 @@ type CreatedAPIKey struct {
}
type PlatformModel struct {
ID string `json:"id"`
PlatformID string `json:"platformId"`
BaseModelID string `json:"baseModelId,omitempty"`
Provider string `json:"provider,omitempty"`
PlatformName string `json:"platformName,omitempty"`
ModelName string `json:"modelName"`
ProviderModelName string `json:"providerModelName,omitempty"`
ID string `json:"id"`
PlatformID string `json:"platformId"`
BaseModelID string `json:"baseModelId,omitempty"`
Provider string `json:"provider,omitempty"`
PlatformName string `json:"platformName,omitempty"`
ModelName string `json:"modelName"`
ProviderModelName string `json:"providerModelName,omitempty"`
// Deprecated: use ModelName. This field mirrors the canonical invocation name during the compatibility window.
ModelAlias string `json:"modelAlias,omitempty"`
LegacyAliases StringList `json:"legacyAliases,omitempty"`
ModelType StringList `json:"modelType"`
DisplayName string `json:"displayName"`
CapabilityOverride map[string]any `json:"capabilityOverride,omitempty"`
Capabilities map[string]any `json:"capabilities,omitempty"`
BaseCapabilities map[string]any `json:"-"`
BaseBillingConfig map[string]any `json:"-"`
BaseRateLimitPolicy map[string]any `json:"-"`
BaseRuntimePolicySetID string `json:"-"`
BasePricingRuleSetID string `json:"-"`
PlatformPricingRuleSetID string `json:"-"`
PricingMode string `json:"pricingMode"`
@@ -241,6 +247,7 @@ type PlatformModel struct {
PermissionConfig map[string]any `json:"permissionConfig,omitempty"`
RetryPolicy map[string]any `json:"retryPolicy,omitempty"`
RateLimitPolicy map[string]any `json:"rateLimitPolicy,omitempty"`
RateLimitPolicyMode string `json:"rateLimitPolicyMode" enums:"inherit,override"`
RuntimePolicySetID string `json:"runtimePolicySetId,omitempty"`
RuntimePolicyOverride map[string]any `json:"runtimePolicyOverride,omitempty"`
CooldownUntil string `json:"cooldownUntil,omitempty"`
@@ -284,13 +291,17 @@ type CatalogProvider struct {
}
type BaseModel struct {
ID string `json:"id"`
ProviderKey string `json:"providerKey"`
CanonicalModelKey string `json:"canonicalModelKey"`
ProviderModelName string `json:"providerModelName"`
ModelType StringList `json:"modelType"`
ID string `json:"id"`
ProviderKey string `json:"providerKey"`
CanonicalModelKey string `json:"canonicalModelKey"`
InvocationName string `json:"invocationName"`
ProviderModelName string `json:"providerModelName"`
ModelType StringList `json:"modelType"`
// Deprecated: use InvocationName.
ModelAlias string `json:"modelAlias"`
DisplayName string `json:"-"`
DisplayName string `json:"displayName"`
LegacyAliases StringList `json:"legacyAliases,omitempty"`
ReferenceCount int `json:"referenceCount"`
Capabilities map[string]any `json:"capabilities,omitempty"`
BaseBillingConfig map[string]any `json:"baseBillingConfig,omitempty"`
DefaultRateLimitPolicy map[string]any `json:"defaultRateLimitPolicy,omitempty"`
@@ -576,17 +587,19 @@ COALESCE(compatibility_submit_headers, '{}'::jsonb), COALESCE(compatibility_subm
created_at, updated_at, COALESCE(finished_at::text, '')`
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"`
ID string `json:"id"`
TaskID string `json:"taskId"`
PlatformID string `json:"-"`
SkippedReason string `json:"-"`
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"`
}
type TaskAttempt struct {
@@ -938,18 +951,34 @@ func (s *Store) listModels(ctx context.Context, platformID string) ([]PlatformMo
rows, err := s.pool.Query(ctx, `
SELECT m.id::text, m.platform_id::text, COALESCE(m.base_model_id::text, ''), p.provider, p.name,
m.model_name, COALESCE(NULLIF(m.provider_model_name, ''), m.model_name), COALESCE(m.model_alias, ''), m.model_type, m.display_name,
COALESCE(NULLIF(b.invocation_name, ''), m.model_name),
COALESCE(NULLIF(m.provider_model_name, ''), m.model_name),
COALESCE(NULLIF(b.invocation_name, ''), NULLIF(m.model_alias, ''), m.model_name),
COALESCE(b.model_type, m.model_type),
COALESCE(NULLIF(b.display_name, ''), NULLIF(m.display_name, ''), COALESCE(NULLIF(b.invocation_name, ''), m.model_name)),
COALESCE(b.legacy_aliases, '[]'::jsonb),
m.capability_override, m.capabilities, COALESCE(b.capabilities, '{}'::jsonb),
COALESCE(b.base_billing_config, '{}'::jsonb), COALESCE(b.pricing_rule_set_id::text, ''),
COALESCE(p.pricing_rule_set_id::text, ''), m.pricing_mode, COALESCE(m.discount_factor, 0)::float8,
COALESCE(m.pricing_rule_set_id::text, ''), m.billing_config_override, m.billing_config,
m.permission_config, m.retry_policy, m.rate_limit_policy, COALESCE(m.runtime_policy_set_id::text, ''), m.runtime_policy_override,
m.permission_config, m.retry_policy, m.rate_limit_policy, m.rate_limit_policy_mode,
COALESCE(m.runtime_policy_set_id::text, ''), m.runtime_policy_override,
COALESCE(b.default_rate_limit_policy, '{}'::jsonb), COALESCE(b.runtime_policy_set_id::text, ''),
COALESCE(to_char(m.cooldown_until AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'), ''),
m.enabled, m.created_at, m.updated_at
FROM platform_models m
JOIN integration_platforms p ON p.id = m.platform_id
LEFT JOIN LATERAL (
SELECT catalog.capabilities, catalog.base_billing_config, catalog.pricing_rule_set_id
SELECT catalog.invocation_name, catalog.display_name, catalog.model_type,
catalog.capabilities, catalog.base_billing_config, catalog.pricing_rule_set_id,
catalog.default_rate_limit_policy, catalog.runtime_policy_set_id,
COALESCE((
SELECT jsonb_agg(DISTINCT compatibility_alias.alias ORDER BY compatibility_alias.alias)
FROM model_compatibility_aliases compatibility_alias
WHERE compatibility_alias.base_model_id = catalog.id
AND compatibility_alias.active = true
AND (compatibility_alias.expires_at IS NULL OR compatibility_alias.expires_at > now())
), '[]'::jsonb) AS legacy_aliases
FROM base_model_catalog catalog
WHERE (m.base_model_id IS NOT NULL AND catalog.id = m.base_model_id)
OR (
@@ -983,7 +1012,9 @@ ORDER BY m.model_type ASC, m.model_name ASC`, args...)
var retryPolicy []byte
var rateLimitPolicy []byte
var runtimePolicyOverride []byte
var baseRateLimitPolicy []byte
var modelTypeBytes []byte
var legacyAliasesBytes []byte
if err := rows.Scan(
&model.ID,
&model.PlatformID,
@@ -995,6 +1026,7 @@ ORDER BY m.model_type ASC, m.model_name ASC`, args...)
&model.ModelAlias,
&modelTypeBytes,
&model.DisplayName,
&legacyAliasesBytes,
&capabilityOverride,
&capabilities,
&baseCapabilities,
@@ -1009,8 +1041,11 @@ ORDER BY m.model_type ASC, m.model_name ASC`, args...)
&permissionConfig,
&retryPolicy,
&rateLimitPolicy,
&model.RateLimitPolicyMode,
&model.RuntimePolicySetID,
&runtimePolicyOverride,
&baseRateLimitPolicy,
&model.BaseRuntimePolicySetID,
&model.CooldownUntil,
&model.Enabled,
&model.CreatedAt,
@@ -1023,12 +1058,14 @@ ORDER BY m.model_type ASC, m.model_name ASC`, args...)
model.BaseCapabilities = decodeObject(baseCapabilities)
model.BaseBillingConfig = decodeObject(baseBillingConfig)
model.ModelType = decodeStringArray(modelTypeBytes)
model.LegacyAliases = decodeStringArray(legacyAliasesBytes)
model.BillingConfigOverride = decodeObject(billingConfigOverride)
model.BillingConfig = decodeObject(billingConfig)
model.PermissionConfig = decodeObject(permissionConfig)
model.RetryPolicy = decodeObject(retryPolicy)
model.RateLimitPolicy = decodeObject(rateLimitPolicy)
model.RuntimePolicyOverride = decodeObject(runtimePolicyOverride)
model.BaseRateLimitPolicy = decodeObject(baseRateLimitPolicy)
models = append(models, model)
}
return models, rows.Err()
@@ -1956,7 +1993,15 @@ func (s *Store) CreateTask(ctx context.Context, input CreateTaskInput, user *aut
}
func (s *Store) CreateTaskIdempotent(ctx context.Context, input CreateTaskInput, user *auth.User) (CreateTaskResult, error) {
requestBody, _ := json.Marshal(input.Request)
requestReport := sanitizeJSONForStorageWithReport(input.Request)
if requestReport.BinaryCount > 0 {
return CreateTaskResult{}, &taskPayloadBinaryError{
target: ErrTaskRequestBinaryNotMaterialized,
code: "request_binary_not_materialized",
count: requestReport.BinaryCount,
}
}
requestBody, _ := json.Marshal(requestReport.Value)
runMode := normalizeRunMode(input.RunMode, input.Request)
status := "queued"
resultBody, _ := json.Marshal(map[string]any(nil))
@@ -2011,7 +2056,7 @@ WHERE user_id = $1 AND idempotency_key_hash = $2`, user.ID, strings.TrimSpace(in
}
events := taskEventsForCreate(task.ID, runMode, status, nil)
for _, event := range events {
payload, _ := json.Marshal(event.Payload)
payload, _ := json.Marshal(sanitizeJSONForStorage(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::text, NULLIF($4::text, ''), NULLIF($5::text, ''), $6, NULLIF($7::text, ''), $8::jsonb, $9)`,
@@ -2146,7 +2191,8 @@ func scanGatewayTask(scanner taskScanner) (GatewayTask, error) {
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
COALESCE(progress, 0)::float8, COALESCE(message, ''), payload, simulated, created_at,
COALESCE(platform_id::text, '')
FROM gateway_task_events
WHERE task_id = $1::uuid
ORDER BY seq ASC`, taskID)
@@ -2171,6 +2217,7 @@ ORDER BY seq ASC`, taskID)
&payload,
&item.Simulated,
&item.CreatedAt,
&item.PlatformID,
); err != nil {
return nil, err
}
@@ -2386,10 +2433,7 @@ func taskEventsForCreate(taskID string, runMode string, status string, result ma
Seq: 1,
EventType: "task.accepted",
Status: "queued",
Phase: "queued",
Progress: 0,
Message: "task accepted",
Payload: map[string]any{"taskId": taskID},
Payload: map[string]any{},
Simulated: runMode == "simulation",
}}
}
apps/api/internal/store/rate_limit_policy.go
+167
View File
@@ -0,0 +1,167 @@
package store
import (
"math"
"strings"
)
const (
RateLimitPolicyModeInherit = "inherit"
RateLimitPolicyModeOverride = "override"
)
type EffectiveRateLimitPolicyInput struct {
BasePolicy map[string]any
PlatformPolicy map[string]any
RuntimePolicy map[string]any
RuntimePolicyExplicit bool
RuntimePolicyOverride map[string]any
ModelPolicy map[string]any
ModelPolicyMode string
}
// EffectiveRateLimitPolicy resolves one authoritative policy. Platform
// policies are defaults for every bound model; explicit model/runtime settings
// replace the complete policy instead of merging individual metrics.
func EffectiveRateLimitPolicy(input EffectiveRateLimitPolicyInput) map[string]any {
policy := input.BasePolicy
if policySpecified(input.PlatformPolicy) {
policy = input.PlatformPolicy
}
if input.RuntimePolicyExplicit {
policy = input.RuntimePolicy
}
if raw, ok := input.RuntimePolicyOverride["rateLimitPolicy"]; ok {
policy, _ = raw.(map[string]any)
}
mode := NormalizeRateLimitPolicyMode(input.ModelPolicyMode, input.ModelPolicy != nil)
if mode == RateLimitPolicyModeOverride {
policy = input.ModelPolicy
}
return NormalizeRateLimitPolicy(policy)
}
func NormalizeRateLimitPolicyMode(mode string, policyProvided bool) string {
switch strings.ToLower(strings.TrimSpace(mode)) {
case RateLimitPolicyModeOverride:
return RateLimitPolicyModeOverride
case RateLimitPolicyModeInherit:
return RateLimitPolicyModeInherit
default:
if policyProvided {
return RateLimitPolicyModeOverride
}
return RateLimitPolicyModeInherit
}
}
func RateLimitPolicyMetric(policy map[string]any, metric string) (float64, bool) {
rules, _ := NormalizeRateLimitPolicy(policy)["rules"].([]any)
for _, rawRule := range rules {
rule, _ := rawRule.(map[string]any)
if strings.TrimSpace(stringValue(rule["metric"])) != metric {
continue
}
value := floatValue(rule["limit"])
return value, true
}
return 0, false
}
// NormalizeRateLimitPolicy keeps the canonical rules contract while accepting
// policies imported from server-main before that contract existed. Runtime
// enforcement and worker sizing must agree on these legacy limits; otherwise a
// configured max_concurrent_requests silently becomes unlimited.
func NormalizeRateLimitPolicy(policy map[string]any) map[string]any {
if policy == nil {
return nil
}
out := clonePolicy(policy)
rules, _ := out["rules"].([]any)
if len(rules) > 0 {
return out
}
legacyScopes := []map[string]any{policy}
for _, key := range []string{"platformLimits", "modelLimits", "platform_limits", "model_limits"} {
if nested, ok := policy[key].(map[string]any); ok {
legacyScopes = append(legacyScopes, nested)
}
}
if limit, ok := lowestPositiveLegacyLimit(legacyScopes,
"max_concurrent_requests", "maxConcurrentRequests", "concurrent"); ok {
rules = append(rules, map[string]any{
"metric": "concurrent",
"limit": limit,
"leaseTtlSeconds": 120,
})
}
if limit, ok := lowestPositiveLegacyLimit(legacyScopes,
"max_request_per_minute", "maxRequestsPerMinute", "rpm"); ok {
rules = append(rules, map[string]any{
"metric": "rpm",
"limit": limit,
"windowSeconds": 60,
})
}
if limit, ok := lowestPositiveLegacyLimit(legacyScopes,
"max_tokens_per_minute", "maxTokensPerMinute", "tpm_total"); ok {
rules = append(rules, map[string]any{
"metric": "tpm_total",
"limit": limit,
"windowSeconds": 60,
})
}
if len(rules) > 0 {
out["rules"] = rules
}
return out
}
func ConcurrentPolicyCapacity(policy map[string]any) (int, bool) {
limit, ok := RateLimitPolicyMetric(policy, "concurrent")
if !ok || limit <= 0 {
return 0, false
}
if limit < 1 {
return 1, true
}
if limit >= float64(math.MaxInt) {
return math.MaxInt, true
}
return int(math.Floor(limit)), true
}
func policySpecified(policy map[string]any) bool {
rules, _ := NormalizeRateLimitPolicy(policy)["rules"].([]any)
return len(rules) > 0
}
func lowestPositiveLegacyLimit(scopes []map[string]any, keys ...string) (float64, bool) {
limit := 0.0
found := false
for _, scope := range scopes {
for _, key := range keys {
value := floatValue(scope[key])
if value <= 0 {
continue
}
if !found || value < limit {
limit = value
found = true
}
}
}
return limit, found
}
func clonePolicy(policy map[string]any) map[string]any {
if policy == nil {
return nil
}
out := make(map[string]any, len(policy))
for key, value := range policy {
out[key] = value
}
return out
}
apps/api/internal/store/rate_limit_policy_migration_integration_test.go
+240
View File
@@ -0,0 +1,240 @@
package store
import (
"context"
"fmt"
"net/url"
"os"
"path/filepath"
"runtime"
"strings"
"testing"
"time"
"github.com/jackc/pgx/v5/pgxpool"
)
func TestRateLimitPolicyModeMigrationClassifiesExistingRows(t *testing.T) {
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
if databaseURL == "" {
t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run rate-limit migration PostgreSQL integration tests")
}
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
defer cancel()
schema := fmt.Sprintf("rate_limit_mode_%d", time.Now().UnixNano())
adminPool, err := pgxpool.New(ctx, databaseURL)
if err != nil {
t.Fatalf("connect migration test database: %v", err)
}
defer adminPool.Close()
if _, err := adminPool.Exec(ctx, `CREATE SCHEMA `+schema); err != nil {
t.Fatalf("create migration test schema: %v", err)
}
defer adminPool.Exec(context.Background(), `DROP SCHEMA IF EXISTS `+schema+` CASCADE`)
schemaURL, err := url.Parse(databaseURL)
if err != nil {
t.Fatalf("parse test database url: %v", err)
}
query := schemaURL.Query()
query.Set("search_path", schema)
schemaURL.RawQuery = query.Encode()
pool, err := pgxpool.New(ctx, schemaURL.String())
if err != nil {
t.Fatalf("connect migration test schema: %v", err)
}
defer pool.Close()
if _, err := pool.Exec(ctx, `
CREATE TABLE base_model_catalog (
id uuid PRIMARY KEY,
default_rate_limit_policy jsonb NOT NULL DEFAULT '{}'::jsonb,
runtime_policy_set_id uuid,
runtime_policy_override jsonb NOT NULL DEFAULT '{}'::jsonb
);
CREATE TABLE integration_platforms (
id uuid PRIMARY KEY,
rate_limit_policy jsonb NOT NULL DEFAULT '{}'::jsonb,
updated_at timestamptz NOT NULL DEFAULT now()
);
CREATE TABLE platform_models (
id uuid PRIMARY KEY,
base_model_id uuid,
rate_limit_policy jsonb NOT NULL DEFAULT '{}'::jsonb,
runtime_policy_set_id uuid,
runtime_policy_override jsonb NOT NULL DEFAULT '{}'::jsonb,
updated_at timestamptz NOT NULL DEFAULT now()
);
INSERT INTO integration_platforms (id, rate_limit_policy)
VALUES ('20000000-0000-0000-0000-000000000001', '{"rules":[]}');
INSERT INTO base_model_catalog (id, default_rate_limit_policy)
VALUES ('00000000-0000-0000-0000-000000000001', '{"rules":[{"metric":"concurrent","limit":4}]}');
INSERT INTO platform_models (id, base_model_id, rate_limit_policy, runtime_policy_override) VALUES
('10000000-0000-0000-0000-000000000001', '00000000-0000-0000-0000-000000000001', '{"rules":[{"metric":"concurrent","limit":4}]}', '{}'),
('10000000-0000-0000-0000-000000000002', '00000000-0000-0000-0000-000000000001', '{"rules":[{"metric":"concurrent","limit":8}]}', '{}'),
('10000000-0000-0000-0000-000000000003', '00000000-0000-0000-0000-000000000001', '{"rules":[{"metric":"concurrent","limit":4}]}', '{"rateLimitPolicy":{"rules":[]}}'),
('10000000-0000-0000-0000-000000000004', NULL, '{}', '{}');`); err != nil {
t.Fatalf("seed pre-migration rows: %v", err)
}
_, currentFile, _, _ := runtime.Caller(0)
migrationPath := filepath.Join(filepath.Dir(currentFile), "..", "..", "migrations", "0080_platform_model_rate_limit_policy_mode.sql")
migration, err := os.ReadFile(migrationPath)
if err != nil {
t.Fatalf("read migration: %v", err)
}
if _, err := pool.Exec(ctx, string(migration)); err != nil {
t.Fatalf("apply migration: %v", err)
}
rows, err := pool.Query(ctx, `SELECT id::text, rate_limit_policy_mode FROM platform_models ORDER BY id`)
if err != nil {
t.Fatalf("read classified rows: %v", err)
}
defer rows.Close()
want := []string{"inherit", "override", "override", "inherit"}
index := 0
for rows.Next() {
var id, mode string
if err := rows.Scan(&id, &mode); err != nil {
t.Fatalf("scan classified row: %v", err)
}
if index >= len(want) || mode != want[index] {
t.Fatalf("row %s mode=%s, want=%s", id, mode, want[index])
}
index++
}
if err := rows.Err(); err != nil {
t.Fatalf("read classified rows: %v", err)
}
if index != len(want) {
t.Fatalf("classified %d rows, want %d", index, len(want))
}
var normalizedPlatformPolicy string
if err := pool.QueryRow(ctx, `
SELECT rate_limit_policy::text
FROM integration_platforms
WHERE id = '20000000-0000-0000-0000-000000000001'`).Scan(&normalizedPlatformPolicy); err != nil {
t.Fatalf("read normalized platform policy: %v", err)
}
if normalizedPlatformPolicy != "{}" {
t.Fatalf("normalized platform policy=%s, want={}", normalizedPlatformPolicy)
}
}
func TestRemoveUnconfiguredGPTImageConcurrencyLimitMigration(t *testing.T) {
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
if databaseURL == "" {
t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run rate-limit migration PostgreSQL integration tests")
}
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
defer cancel()
schema := fmt.Sprintf("remove_gpt_image_limit_%d", time.Now().UnixNano())
adminPool, err := pgxpool.New(ctx, databaseURL)
if err != nil {
t.Fatalf("connect migration test database: %v", err)
}
defer adminPool.Close()
if _, err := adminPool.Exec(ctx, `CREATE SCHEMA `+schema); err != nil {
t.Fatalf("create migration test schema: %v", err)
}
defer adminPool.Exec(context.Background(), `DROP SCHEMA IF EXISTS `+schema+` CASCADE`)
schemaURL, err := url.Parse(databaseURL)
if err != nil {
t.Fatalf("parse test database url: %v", err)
}
query := schemaURL.Query()
query.Set("search_path", schema)
schemaURL.RawQuery = query.Encode()
pool, err := pgxpool.New(ctx, schemaURL.String())
if err != nil {
t.Fatalf("connect migration test schema: %v", err)
}
defer pool.Close()
if _, err := pool.Exec(ctx, `
CREATE TABLE base_model_catalog (
id uuid PRIMARY KEY,
invocation_name text NOT NULL,
default_rate_limit_policy jsonb NOT NULL DEFAULT '{}'::jsonb,
metadata jsonb NOT NULL DEFAULT '{}'::jsonb,
updated_at timestamptz NOT NULL DEFAULT now()
);
INSERT INTO base_model_catalog (id, invocation_name, default_rate_limit_policy, metadata) VALUES
(
'00000000-0000-0000-0000-000000000001',
'gpt-image-2',
'{"rules":[{"metric":"concurrent","limit":5,"leaseTtlSeconds":300}]}',
'{"rateLimitSource":"image-gateway-migration","rateLimitUpdatedAt":"2026-07-24"}'
),
(
'00000000-0000-0000-0000-000000000002',
'gpt-image-2',
'{"rules":[{"metric":"concurrent","limit":6,"leaseTtlSeconds":300}]}',
'{"rateLimitSource":"image-gateway-migration"}'
),
(
'00000000-0000-0000-0000-000000000003',
'gpt-image-2',
'{"rules":[{"metric":"concurrent","limit":5,"leaseTtlSeconds":300}]}',
'{"rateLimitSource":"operator"}'
),
(
'00000000-0000-0000-0000-000000000004',
'gemini-3-pro-image',
'{"rules":[{"metric":"concurrent","limit":10,"leaseTtlSeconds":600}]}',
'{"rateLimitSource":"image-gateway-migration"}'
);`); err != nil {
t.Fatalf("seed pre-migration rows: %v", err)
}
_, currentFile, _, _ := runtime.Caller(0)
migrationPath := filepath.Join(filepath.Dir(currentFile), "..", "..", "migrations", "0082_remove_unconfigured_gpt_image_concurrency_limit.sql")
migration, err := os.ReadFile(migrationPath)
if err != nil {
t.Fatalf("read migration: %v", err)
}
if _, err := pool.Exec(ctx, string(migration)); err != nil {
t.Fatalf("apply migration: %v", err)
}
rows, err := pool.Query(ctx, `
SELECT id::text, default_rate_limit_policy::text, metadata->>'rateLimitSource',
metadata->>'rateLimitRemovalReason'
FROM base_model_catalog
ORDER BY id`)
if err != nil {
t.Fatalf("read migrated rows: %v", err)
}
defer rows.Close()
type result struct {
policy string
source *string
removalReason *string
}
got := make([]result, 0, 4)
for rows.Next() {
var id string
var item result
if err := rows.Scan(&id, &item.policy, &item.source, &item.removalReason); err != nil {
t.Fatalf("scan migrated row: %v", err)
}
got = append(got, item)
}
if err := rows.Err(); err != nil {
t.Fatalf("read migrated rows: %v", err)
}
if len(got) != 4 {
t.Fatalf("migrated %d rows, want 4", len(got))
}
if got[0].policy != "{}" || got[0].source != nil ||
got[0].removalReason == nil || *got[0].removalReason != "upstream-limit-unconfigured" {
t.Fatalf("migration-owned GPT policy not removed: %+v", got[0])
}
if got[1].policy == "{}" || got[1].source == nil || *got[1].source != "image-gateway-migration" {
t.Fatalf("custom GPT policy was changed: %+v", got[1])
}
if got[2].policy == "{}" || got[2].source == nil || *got[2].source != "operator" {
t.Fatalf("operator GPT policy was changed: %+v", got[2])
}
if got[3].policy == "{}" || got[3].source == nil || *got[3].source != "image-gateway-migration" {
t.Fatalf("Gemini policy was changed: %+v", got[3])
}
}
apps/api/internal/store/rate_limit_policy_test.go
+169
View File
@@ -0,0 +1,169 @@
package store
import "testing"
func TestEffectiveRateLimitPolicyPrecedence(t *testing.T) {
policy := func(limit float64) map[string]any {
return map[string]any{"rules": []any{map[string]any{"metric": "concurrent", "limit": limit}}}
}
tests := []struct {
name string
input EffectiveRateLimitPolicyInput
want float64
ok bool
}{
{name: "base", input: EffectiveRateLimitPolicyInput{BasePolicy: policy(2)}, want: 2, ok: true},
{name: "platform", input: EffectiveRateLimitPolicyInput{BasePolicy: policy(2), PlatformPolicy: policy(4)}, want: 4, ok: true},
{name: "empty platform inherits base", input: EffectiveRateLimitPolicyInput{BasePolicy: policy(2), PlatformPolicy: map[string]any{"rules": []any{}}}, want: 2, ok: true},
{name: "runtime", input: EffectiveRateLimitPolicyInput{PlatformPolicy: policy(4), RuntimePolicy: policy(8), RuntimePolicyExplicit: true}, want: 8, ok: true},
{name: "runtime override", input: EffectiveRateLimitPolicyInput{PlatformPolicy: policy(4), RuntimePolicyOverride: map[string]any{"rateLimitPolicy": policy(16)}}, want: 16, ok: true},
{name: "model override", input: EffectiveRateLimitPolicyInput{PlatformPolicy: policy(4), ModelPolicy: policy(32), ModelPolicyMode: "override"}, want: 32, ok: true},
{name: "model explicit unlimited", input: EffectiveRateLimitPolicyInput{PlatformPolicy: policy(4), ModelPolicy: map[string]any{}, ModelPolicyMode: "override"}, ok: false},
{name: "model inherit", input: EffectiveRateLimitPolicyInput{PlatformPolicy: policy(4), ModelPolicy: policy(32), ModelPolicyMode: "inherit"}, want: 4, ok: true},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got, ok := ConcurrentPolicyCapacity(EffectiveRateLimitPolicy(tt.input))
if ok != tt.ok || (ok && got != int(tt.want)) {
t.Fatalf("capacity = (%d, %v), want (%d, %v)", got, ok, int(tt.want), tt.ok)
}
})
}
}
func TestNormalizeRateLimitPolicyLegacyShapes(t *testing.T) {
tests := []struct {
name string
policy map[string]any
metric string
want float64
}{
{
name: "platform concurrent",
policy: map[string]any{"platformLimits": map[string]any{"max_concurrent_requests": 5.0}},
metric: "concurrent",
want: 5,
},
{
name: "model concurrent camel case",
policy: map[string]any{"modelLimits": map[string]any{"maxConcurrentRequests": 10.0}},
metric: "concurrent",
want: 10,
},
{
name: "stricter duplicate wins",
policy: map[string]any{
"platformLimits": map[string]any{"max_concurrent_requests": 8.0},
"modelLimits": map[string]any{"max_concurrent_requests": 3.0},
},
metric: "concurrent",
want: 3,
},
{
name: "requests per minute",
policy: map[string]any{"model_limits": map[string]any{"max_request_per_minute": 60.0}},
metric: "rpm",
want: 60,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got, ok := RateLimitPolicyMetric(tt.policy, tt.metric)
if !ok || got != tt.want {
t.Fatalf("metric %s = (%v, %v), want (%v, true)", tt.metric, got, ok, tt.want)
}
})
}
}
func TestConcurrentPolicyCapacityRoundsDown(t *testing.T) {
policy := map[string]any{"rules": []any{map[string]any{"metric": "concurrent", "limit": 96.9}}}
got, ok := ConcurrentPolicyCapacity(policy)
if !ok || got != 96 {
t.Fatalf("capacity = (%d, %v), want (96, true)", got, ok)
}
}
func TestAsyncWorkerCapacityAggregation(t *testing.T) {
policy := func(limit float64) map[string]any {
return map[string]any{"rules": []any{map[string]any{"metric": "concurrent", "limit": limit}}}
}
tests := []struct {
name string
policies []map[string]any
hardLimit int
wantCapacity int
wantDesired int
wantCapped bool
}{
{name: "no enabled models", hardLimit: 2048, wantCapacity: 1, wantDesired: 1},
{name: "finite sum", policies: []map[string]any{policy(64), policy(32)}, hardLimit: 2048, wantCapacity: 96, wantDesired: 96},
{name: "unlimited model", policies: []map[string]any{policy(64), {}}, hardLimit: 2048, wantCapacity: 2048, wantDesired: 2048},
{name: "hard limit cap", policies: []map[string]any{policy(80), policy(80)}, hardLimit: 96, wantCapacity: 96, wantDesired: 160, wantCapped: true},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := asyncWorkerCapacityFromPolicies(tt.policies, tt.hardLimit)
if got.Capacity != tt.wantCapacity || got.Desired != tt.wantDesired || got.Capped != tt.wantCapped {
t.Fatalf("snapshot=%+v, want capacity=%d desired=%d capped=%v", got, tt.wantCapacity, tt.wantDesired, tt.wantCapped)
}
})
}
}
func TestAsyncWorkerCapacityRespectsUserGroupCeiling(t *testing.T) {
policy := func(limit float64) map[string]any {
return map[string]any{"rules": []any{map[string]any{"metric": "concurrent", "limit": limit}}}
}
tests := []struct {
name string
models []map[string]any
groups []map[string]any
hardLimit int
wantCapacity int
wantDesired int
wantCapped bool
}{
{
name: "group ceiling prevents worker oversubscription",
models: []map[string]any{{}},
groups: []map[string]any{policy(3), policy(10)},
hardLimit: 2048,
wantCapacity: 13,
wantDesired: 13,
},
{
name: "model ceiling is stricter",
models: []map[string]any{policy(5), policy(7)},
groups: []map[string]any{policy(300)},
hardLimit: 2048,
wantCapacity: 12,
wantDesired: 12,
},
{
name: "both policy sets unlimited use hard limit",
models: []map[string]any{{}},
groups: []map[string]any{{}},
hardLimit: 256,
wantCapacity: 256,
wantDesired: 256,
},
{
name: "finite group desired still reports hard cap",
models: []map[string]any{{}},
groups: []map[string]any{policy(500)},
hardLimit: 256,
wantCapacity: 256,
wantDesired: 500,
wantCapped: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := asyncWorkerCapacityFromPolicySets(tt.models, tt.groups, tt.hardLimit)
if got.Capacity != tt.wantCapacity || got.Desired != tt.wantDesired || got.Capped != tt.wantCapped {
t.Fatalf("snapshot=%+v, want capacity=%d desired=%d capped=%v", got, tt.wantCapacity, tt.wantDesired, tt.wantCapped)
}
})
}
}
apps/api/internal/store/rate_limit_retry_test.go
+19
View File
@@ -0,0 +1,19 @@
package store
import (
"testing"
"time"
)
func TestConcurrencyRetryAfterBoundsPolling(t *testing.T) {
if got := concurrencyRetryAfter(time.Time{}); got != 5*time.Second {
t.Fatalf("zero expiry retry=%s, want=5s", got)
}
if got := concurrencyRetryAfter(time.Now().Add(30 * time.Second)); got != 5*time.Second {
t.Fatalf("distant expiry retry=%s, want=5s", got)
}
got := concurrencyRetryAfter(time.Now().Add(2500 * time.Millisecond))
if got < 2*time.Second || got > 2500*time.Millisecond {
t.Fatalf("near expiry retry=%s, want within [2s,2.5s]", got)
}
}
apps/api/internal/store/rate_limit_status.go
+45 -63
View File
@@ -143,8 +143,10 @@ func (s *Store) ListModelRateLimitStatuses(ctx context.Context) ([]ModelRateLimi
p.priority, p.dynamic_priority, COALESCE(p.dynamic_priority, p.priority),
m.model_name, COALESCE(NULLIF(m.provider_model_name, ''), m.model_name), COALESCE(m.model_alias, ''),
m.model_type, m.display_name, m.enabled,
p.rate_limit_policy, COALESCE(rp.rate_limit_policy, '{}'::jsonb), COALESCE(m.runtime_policy_set_id::text, b.runtime_policy_set_id::text, ''),
COALESCE(NULLIF(m.runtime_policy_override, '{}'::jsonb), b.runtime_policy_override, '{}'::jsonb), m.rate_limit_policy,
COALESCE(b.default_rate_limit_policy, '{}'::jsonb), p.rate_limit_policy,
COALESCE(rp.rate_limit_policy, '{}'::jsonb),
(m.runtime_policy_set_id IS NOT NULL),
COALESCE(m.runtime_policy_override, '{}'::jsonb), m.rate_limit_policy, m.rate_limit_policy_mode,
COALESCE(to_char(p.cooldown_until AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'), ''),
COALESCE(to_char(m.cooldown_until AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'), ''),
COALESCE(con.active, 0)::float8,
@@ -217,11 +219,13 @@ ORDER BY p.priority ASC, m.model_name ASC`)
for rows.Next() {
var item ModelRateLimitStatus
var modelTypeBytes []byte
var basePolicyBytes []byte
var platformPolicyBytes []byte
var runtimePolicyBytes []byte
var runtimePolicySetID string
var runtimePolicyExplicit bool
var runtimeOverrideBytes []byte
var modelPolicyBytes []byte
var modelPolicyMode string
var platformDynamicPriority sql.NullInt64
var platformCooldownUntil string
var modelCooldownUntil string
@@ -248,11 +252,13 @@ ORDER BY p.priority ASC, m.model_name ASC`)
&modelTypeBytes,
&item.DisplayName,
&item.Enabled,
&basePolicyBytes,
&platformPolicyBytes,
&runtimePolicyBytes,
&runtimePolicySetID,
&runtimePolicyExplicit,
&runtimeOverrideBytes,
&modelPolicyBytes,
&modelPolicyMode,
&platformCooldownUntil,
&modelCooldownUntil,
&concurrentCurrent,
@@ -268,13 +274,15 @@ ORDER BY p.priority ASC, m.model_name ASC`)
}
item.PlatformDynamicPriority = intPointerFromNull(platformDynamicPriority)
item.ModelType = decodeStringArray(modelTypeBytes)
policy := effectiveModelRateLimitPolicy(
decodeObject(platformPolicyBytes),
decodeObject(runtimePolicyBytes),
runtimePolicySetID,
decodeObject(runtimeOverrideBytes),
decodeObject(modelPolicyBytes),
)
policy := EffectiveRateLimitPolicy(EffectiveRateLimitPolicyInput{
BasePolicy: decodeObject(basePolicyBytes),
PlatformPolicy: decodeObject(platformPolicyBytes),
RuntimePolicy: decodeObject(runtimePolicyBytes),
RuntimePolicyExplicit: runtimePolicyExplicit,
RuntimePolicyOverride: decodeObject(runtimeOverrideBytes),
ModelPolicy: decodeObject(modelPolicyBytes),
ModelPolicyMode: modelPolicyMode,
})
item.PlatformCooldownUntil = platformCooldownUntil
item.ModelCooldownUntil = modelCooldownUntil
item.RateLimitPolicy = policy
@@ -330,19 +338,20 @@ func (s *Store) listRecentPriorityDemotionsByPlatform(ctx context.Context, statu
return out, nil
}
rows, err := s.pool.Query(ctx, `
SELECT id::text, task_id::text, COALESCE(message, ''), payload, created_at
SELECT id::text, task_id::text, COALESCE(message, ''), payload, created_at, event_platform_id
FROM (
SELECT e.*,
COALESCE(e.platform_id::text, e.payload->>'platformId') AS event_platform_id,
row_number() OVER (
PARTITION BY e.payload->>'platformId'
PARTITION BY COALESCE(e.platform_id::text, e.payload->>'platformId')
ORDER BY e.created_at DESC, e.seq DESC
) AS demotion_rank
FROM gateway_task_events e
WHERE e.event_type = 'task.policy.priority_demoted'
AND e.payload->>'platformId' = ANY($1::text[])
AND COALESCE(e.platform_id::text, e.payload->>'platformId') = ANY($1::text[])
) ranked
WHERE demotion_rank <= $2
ORDER BY payload->>'platformId' ASC, created_at DESC`, platformIDs, limit)
ORDER BY event_platform_id ASC, created_at DESC`, platformIDs, limit)
if err != nil {
return nil, err
}
@@ -353,10 +362,16 @@ func (s *Store) listRecentPriorityDemotionsByPlatform(ctx context.Context, statu
var message string
var payloadBytes []byte
var createdAt time.Time
if err := rows.Scan(&id, &taskID, &message, &payloadBytes, &createdAt); err != nil {
var platformID string
if err := rows.Scan(&id, &taskID, &message, &payloadBytes, &createdAt, &platformID); err != nil {
return nil, err
}
record := priorityDemotionRecordFromEventPayload(id, taskID, message, decodeObject(payloadBytes), createdAt)
payload := decodeObject(payloadBytes)
if payload == nil {
payload = map[string]any{}
}
payload["platformId"] = platformID
record := priorityDemotionRecordFromEventPayload(id, taskID, message, payload, createdAt)
if record.PlatformID == "" {
continue
}
@@ -412,12 +427,13 @@ func (s *Store) listLatestPlatformDisabledReasons(ctx context.Context, statuses
return out, nil
}
rows, err := s.pool.Query(ctx, `
SELECT id::text, task_id::text, event_type, COALESCE(message, ''), payload, COALESCE(attempt_error_message, ''), created_at
SELECT id::text, task_id::text, event_type, COALESCE(message, ''), payload, COALESCE(attempt_error_message, ''), created_at, event_platform_id
FROM (
SELECT e.*,
COALESCE(e.platform_id::text, e.payload->>'platformId') AS event_platform_id,
a.error_message AS attempt_error_message,
row_number() OVER (
PARTITION BY e.payload->>'platformId'
PARTITION BY COALESCE(e.platform_id::text, e.payload->>'platformId')
ORDER BY e.created_at DESC, e.seq DESC
) AS disabled_rank
FROM gateway_task_events e
@@ -425,12 +441,12 @@ FROM (
SELECT error_message
FROM gateway_task_attempts attempt
WHERE attempt.task_id = e.task_id
AND attempt.platform_id::text = e.payload->>'platformId'
AND attempt.platform_id::text = COALESCE(e.platform_id::text, e.payload->>'platformId')
ORDER BY attempt.attempt_no DESC, attempt.started_at DESC
LIMIT 1
) a ON TRUE
WHERE e.event_type IN ('task.policy.failover_disabled', 'task.policy.auto_disabled')
AND e.payload->>'platformId' = ANY($1::text[])
AND COALESCE(e.platform_id::text, e.payload->>'platformId') = ANY($1::text[])
) ranked
WHERE disabled_rank = 1`, platformIDs)
if err != nil {
@@ -445,10 +461,16 @@ WHERE disabled_rank = 1`, platformIDs)
var payloadBytes []byte
var attemptErrorMessage string
var createdAt time.Time
if err := rows.Scan(&id, &taskID, &eventType, &message, &payloadBytes, &attemptErrorMessage, &createdAt); err != nil {
var platformID string
if err := rows.Scan(&id, &taskID, &eventType, &message, &payloadBytes, &attemptErrorMessage, &createdAt, &platformID); err != nil {
return nil, err
}
record := platformPolicyEventFromPayload(id, taskID, eventType, message, attemptErrorMessage, decodeObject(payloadBytes), createdAt)
payload := decodeObject(payloadBytes)
if payload == nil {
payload = map[string]any{}
}
payload["platformId"] = platformID
record := platformPolicyEventFromPayload(id, taskID, eventType, message, attemptErrorMessage, payload, createdAt)
if record.PlatformID == "" {
continue
}
@@ -491,46 +513,6 @@ func platformPolicyEventFromPayload(id string, taskID string, eventType string,
}
}
func effectiveModelRateLimitPolicy(platformPolicy map[string]any, runtimePolicy map[string]any, runtimePolicySetID string, runtimeOverride map[string]any, modelPolicy map[string]any) map[string]any {
policy := platformPolicy
if strings.TrimSpace(runtimePolicySetID) != "" {
policy = runtimePolicy
} else if hasRateLimitRules(runtimePolicy) {
policy = shallowMergeMap(policy, runtimePolicy)
}
if _, hasOverride := runtimeOverride["rateLimitPolicy"]; hasOverride {
nested, _ := runtimeOverride["rateLimitPolicy"].(map[string]any)
if len(nested) == 0 {
policy = nil
} else {
policy = shallowMergeMap(policy, nested)
}
}
if hasRateLimitRules(modelPolicy) {
policy = shallowMergeMap(policy, modelPolicy)
}
if hasRateLimitRules(policy) {
return policy
}
return nil
}
func hasRateLimitRules(policy map[string]any) bool {
rules, _ := policy["rules"].([]any)
return len(rules) > 0
}
func shallowMergeMap(base map[string]any, override map[string]any) map[string]any {
out := map[string]any{}
for key, value := range base {
out[key] = value
}
for key, value := range override {
out[key] = value
}
return out
}
func rateLimitForMetric(policy map[string]any, metric string) float64 {
rules, _ := policy["rules"].([]any)
for _, rawRule := range rules {
apps/api/internal/store/rate_limit_status_test.go
+19 -21
View File
@@ -6,23 +6,23 @@ import (
)
func TestEffectiveModelRateLimitPolicyTreatsModelRulesAsAuthoritative(t *testing.T) {
policy := effectiveModelRateLimitPolicy(
map[string]any{"rules": []any{
policy := EffectiveRateLimitPolicy(EffectiveRateLimitPolicyInput{
PlatformPolicy: map[string]any{"rules": []any{
map[string]any{"metric": "rpm", "limit": 500},
map[string]any{"metric": "tpm_total", "limit": 100000},
}},
map[string]any{"rules": []any{
RuntimePolicy: map[string]any{"rules": []any{
map[string]any{"metric": "rpm", "limit": 120},
map[string]any{"metric": "tpm_total", "limit": 240000},
map[string]any{"metric": "concurrent", "limit": 6},
}},
"runtime-policy-1",
map[string]any{},
map[string]any{"rules": []any{
RuntimePolicyExplicit: true,
ModelPolicy: map[string]any{"rules": []any{
map[string]any{"metric": "rpm", "limit": 30},
map[string]any{"metric": "concurrent", "limit": 2},
}},
)
ModelPolicyMode: "override",
})
if got := rateLimitForMetric(policy, "rpm"); got != 30 {
t.Fatalf("expected model rpm limit to win, got %v", got)
@@ -36,16 +36,15 @@ func TestEffectiveModelRateLimitPolicyTreatsModelRulesAsAuthoritative(t *testing
}
func TestEffectiveModelRateLimitPolicyTreatsEmptyRuntimePolicyAsUnlimited(t *testing.T) {
policy := effectiveModelRateLimitPolicy(
map[string]any{"rules": []any{
policy := EffectiveRateLimitPolicy(EffectiveRateLimitPolicyInput{
PlatformPolicy: map[string]any{"rules": []any{
map[string]any{"metric": "rpm", "limit": 500},
map[string]any{"metric": "tpm_total", "limit": 100000},
}},
map[string]any{"rules": []any{}},
"runtime-policy-1",
map[string]any{},
map[string]any{},
)
RuntimePolicy: map[string]any{"rules": []any{}},
RuntimePolicyExplicit: true,
ModelPolicyMode: "inherit",
})
if got := rateLimitForMetric(policy, "rpm"); got != 0 {
t.Fatalf("expected empty runtime policy rpm to mean unlimited, got %v", got)
@@ -56,17 +55,16 @@ func TestEffectiveModelRateLimitPolicyTreatsEmptyRuntimePolicyAsUnlimited(t *tes
}
func TestEffectiveModelRateLimitPolicyTreatsNegativeLimitAsUnlimited(t *testing.T) {
policy := effectiveModelRateLimitPolicy(
map[string]any{"rules": []any{
policy := EffectiveRateLimitPolicy(EffectiveRateLimitPolicyInput{
PlatformPolicy: map[string]any{"rules": []any{
map[string]any{"metric": "rpm", "limit": 500},
}},
map[string]any{"rules": []any{
RuntimePolicy: map[string]any{"rules": []any{
map[string]any{"metric": "rpm", "limit": -1},
}},
"runtime-policy-1",
map[string]any{},
map[string]any{},
)
RuntimePolicyExplicit: true,
ModelPolicyMode: "inherit",
})
if got := rateLimitForMetric(policy, "rpm"); got != -1 {
t.Fatalf("expected negative runtime rpm marker to be preserved, got %v", got)
apps/api/internal/store/rate_limits.go
+122 -33
View File
@@ -4,6 +4,7 @@ import (
"context"
"errors"
"fmt"
"sort"
"time"
"github.com/jackc/pgx/v5"
@@ -17,6 +18,8 @@ type RuntimeRecoveryResult struct {
RequeuedAsyncTasks int64 `json:"requeuedAsyncTasks"`
}
var ErrConcurrencyLeaseLost = errors.New("concurrency lease lost")
func (s *Store) ReserveRateLimits(ctx context.Context, taskID string, attemptID string, reservations []RateLimitReservation) (RateLimitResult, error) {
tx, err := s.pool.Begin(ctx)
if err != nil {
@@ -24,6 +27,26 @@ func (s *Store) ReserveRateLimits(ctx context.Context, taskID string, attemptID
}
defer tx.Rollback(ctx)
lockKeys := make([]string, 0)
lockKeySet := make(map[string]struct{})
for _, reservation := range reservations {
if reservation.Metric != "concurrent" || reservation.Limit <= 0 || reservation.Amount <= 0 {
continue
}
key := fmt.Sprintf("%d:%s%d:%s", len(reservation.ScopeType), reservation.ScopeType, len(reservation.ScopeKey), reservation.ScopeKey)
if _, exists := lockKeySet[key]; exists {
continue
}
lockKeySet[key] = struct{}{}
lockKeys = append(lockKeys, key)
}
sort.Strings(lockKeys)
for _, key := range lockKeys {
if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`, key); err != nil {
return RateLimitResult{}, err
}
}
result := RateLimitResult{}
for _, reservation := range reservations {
if reservation.Limit <= 0 || reservation.Amount <= 0 {
@@ -49,11 +72,11 @@ func (s *Store) ReserveRateLimits(ctx context.Context, taskID string, attemptID
reservation.WindowSeconds = 60
}
if reservation.Metric == "concurrent" {
leaseID, err := reserveConcurrencyLease(ctx, tx, taskID, attemptID, reservation)
lease, err := reserveConcurrencyLease(ctx, tx, taskID, attemptID, reservation)
if err != nil {
return RateLimitResult{}, err
}
result.LeaseIDs = append(result.LeaseIDs, leaseID)
result.Leases = append(result.Leases, lease)
continue
}
normalized, err := reserveCounterWindow(ctx, tx, taskID, attemptID, reservation)
@@ -65,7 +88,7 @@ func (s *Store) ReserveRateLimits(ctx context.Context, taskID string, attemptID
return result, tx.Commit(ctx)
}
func reserveConcurrencyLease(ctx context.Context, tx pgx.Tx, taskID string, attemptID string, reservation RateLimitReservation) (string, error) {
func reserveConcurrencyLease(ctx context.Context, tx pgx.Tx, taskID string, attemptID string, reservation RateLimitReservation) (ConcurrencyLease, error) {
if reservation.LeaseTTLSeconds <= 0 {
reservation.LeaseTTLSeconds = 120
}
@@ -83,10 +106,10 @@ WHERE scope_type = $1
reservation.ScopeKey,
reservation.LeaseTTLSeconds,
).Scan(&active, &nextAvailableAt); err != nil {
return "", err
return ConcurrencyLease{}, err
}
if active+reservation.Amount > reservation.Limit {
return "", &RateLimitExceededError{
return ConcurrencyLease{}, &RateLimitExceededError{
ScopeType: reservation.ScopeType,
ScopeKey: reservation.ScopeKey,
ScopeName: reservation.ScopeName,
@@ -117,9 +140,9 @@ RETURNING id::text`,
reservation.Amount,
reservation.LeaseTTLSeconds,
).Scan(&leaseID); err != nil {
return "", err
return ConcurrencyLease{}, err
}
return leaseID, nil
return ConcurrencyLease{ID: leaseID, TTL: time.Duration(reservation.LeaseTTLSeconds) * time.Second}, nil
}
func reserveCounterWindow(ctx context.Context, tx pgx.Tx, taskID string, attemptID string, reservation RateLimitReservation) (RateLimitReservation, error) {
@@ -232,13 +255,16 @@ func retryAfterUntil(when time.Time) time.Duration {
func concurrencyRetryAfter(leaseExpiresAt time.Time) time.Duration {
if leaseExpiresAt.IsZero() {
return time.Second
return 5 * time.Second
}
duration := time.Until(leaseExpiresAt)
if duration <= time.Second {
return time.Second
}
return time.Second
if duration > 5*time.Second {
return 5 * time.Second
}
return duration
}
func (s *Store) CommitRateLimitReservations(ctx context.Context, reservations []RateLimitReservation, actualByMetric map[string]float64) error {
@@ -249,26 +275,57 @@ func (s *Store) ReleaseRateLimitReservations(ctx context.Context, reservations [
return s.finishRateLimitReservations(ctx, reservations, nil, "released", reason)
}
func (s *Store) ReleaseConcurrencyLeases(ctx context.Context, leaseIDs []string) error {
func (s *Store) ReleaseConcurrencyLeases(ctx context.Context, leases []ConcurrencyLease) error {
leaseIDs := concurrencyLeaseIDs(leases)
if len(leaseIDs) == 0 {
return nil
}
for _, leaseID := range leaseIDs {
if leaseID == "" {
continue
}
if _, err := s.pool.Exec(ctx, `
_, err := s.pool.Exec(ctx, `
UPDATE gateway_concurrency_leases
SET released_at = now()
WHERE id = $1::uuid AND released_at IS NULL`, leaseID); err != nil && !errors.Is(err, ErrRateLimited) {
return err
WHERE id = ANY($1::uuid[]) AND released_at IS NULL`, leaseIDs)
return err
}
func (s *Store) RenewConcurrencyLeases(ctx context.Context, leases []ConcurrencyLease) error {
leaseIDs := make([]string, 0, len(leases))
ttlSeconds := make([]int32, 0, len(leases))
for _, lease := range leases {
if lease.ID == "" {
continue
}
ttl := lease.TTL
if ttl <= 0 {
ttl = 120 * time.Second
}
seconds := int32(ttl / time.Second)
if seconds < 1 {
seconds = 1
}
leaseIDs = append(leaseIDs, lease.ID)
ttlSeconds = append(ttlSeconds, seconds)
}
if len(leaseIDs) == 0 {
return nil
}
tag, err := s.pool.Exec(ctx, `
UPDATE gateway_concurrency_leases lease
SET expires_at = now() + (renewal.ttl_seconds * interval '1 second')
FROM unnest($1::uuid[], $2::int[]) AS renewal(id, ttl_seconds)
WHERE lease.id = renewal.id
AND lease.released_at IS NULL
AND lease.expires_at > now()`, leaseIDs, ttlSeconds)
if err != nil {
return err
}
if tag.RowsAffected() != int64(len(leaseIDs)) {
return fmt.Errorf("%w: renewed %d of %d leases", ErrConcurrencyLeaseLost, tag.RowsAffected(), len(leaseIDs))
}
return nil
}
func (s *Store) AttachRateLimitResultToAttempt(ctx context.Context, attemptID string, result RateLimitResult) error {
if attemptID == "" || (len(result.Reservations) == 0 && len(result.LeaseIDs) == 0) {
if attemptID == "" || (len(result.Reservations) == 0 && len(result.Leases) == 0) {
return nil
}
tx, err := s.pool.Begin(ctx)
@@ -289,7 +346,8 @@ WHERE id = $1::uuid`, reservation.ReservationID, attemptID); err != nil {
return err
}
}
for _, leaseID := range result.LeaseIDs {
for _, lease := range result.Leases {
leaseID := lease.ID
if leaseID == "" {
continue
}
@@ -303,6 +361,16 @@ WHERE id = $1::uuid`, leaseID, attemptID); err != nil {
return tx.Commit(ctx)
}
func concurrencyLeaseIDs(leases []ConcurrencyLease) []string {
ids := make([]string, 0, len(leases))
for _, lease := range leases {
if lease.ID != "" {
ids = append(ids, lease.ID)
}
}
return ids
}
func (s *Store) RecoverInterruptedRuntimeState(ctx context.Context) (RuntimeRecoveryResult, error) {
tx, err := s.pool.Begin(ctx)
if err != nil {
@@ -401,17 +469,37 @@ RETURNING id::text`)
for _, taskID := range asyncTaskIDs {
if _, err := tx.Exec(ctx, `
INSERT INTO gateway_task_events (task_id, seq, event_type, status, phase, progress, message, payload, simulated)
VALUES (
$1::uuid,
SELECT
task.id,
COALESCE((SELECT MAX(seq) + 1 FROM gateway_task_events WHERE task_id = $1::uuid), 1),
'task.recovered',
'task.queued',
'queued',
'recovered',
0.2,
'async task recovered after service restart',
'{"code":"server_restarted"}'::jsonb,
NULL,
0,
NULL,
'{}'::jsonb,
false
)`, taskID); err != nil {
FROM gateway_tasks task
WHERE task.id = $1::uuid
AND (
SELECT count(*)
FROM gateway_task_events event
WHERE event.task_id = task.id
AND event.event_type NOT IN (
'task.completed', 'task.failed', 'task.cancelled',
'task.billing.settled', 'task.billing.released', 'task.billing.review'
)
) < 16
AND NOT EXISTS (
SELECT 1
FROM gateway_task_events event
WHERE event.task_id = task.id
AND event.seq = (SELECT MAX(last_event.seq) FROM gateway_task_events last_event WHERE last_event.task_id = task.id)
AND event.event_type = 'task.queued'
AND COALESCE(event.status, '') = 'queued'
AND event.platform_id IS NULL
AND event.simulated = false
)`, taskID); err != nil {
return RuntimeRecoveryResult{}, err
}
}
@@ -420,9 +508,10 @@ VALUES (
taskRows, err := tx.Query(ctx, `
UPDATE gateway_tasks
SET status = 'failed',
error = 'task interrupted by service restart',
error = NULL,
error_code = 'server_restarted',
error_message = 'task interrupted by service restart',
remote_task_payload = '{}'::jsonb,
finished_at = now(),
updated_at = now()
WHERE async_mode = false
@@ -452,12 +541,12 @@ INSERT INTO gateway_task_events (task_id, seq, event_type, status, phase, progre
VALUES (
$1::uuid,
COALESCE((SELECT MAX(seq) + 1 FROM gateway_task_events WHERE task_id = $1::uuid), 1),
'task.recovered',
'task.failed',
'failed',
'recovered',
1,
'task interrupted by service restart',
'{"code":"server_restarted"}'::jsonb,
NULL,
0,
NULL,
'{}'::jsonb,
false
)`, taskID); err != nil {
return RuntimeRecoveryResult{}, err
apps/api/internal/store/rate_limits_integration_test.go
+198
View File
@@ -0,0 +1,198 @@
package store
import (
"context"
"errors"
"os"
"strings"
"sync"
"sync/atomic"
"testing"
"time"
)
func TestConcurrencyLeaseReservationIsAtomicAcrossPools(t *testing.T) {
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
if databaseURL == "" {
t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run concurrency lease PostgreSQL integration tests")
}
ctx, cancel := context.WithTimeout(context.Background(), 45*time.Second)
defer cancel()
first, err := Connect(ctx, databaseURL)
if err != nil {
t.Fatalf("connect first store: %v", err)
}
defer first.Close()
second, err := Connect(ctx, databaseURL)
if err != nil {
t.Fatalf("connect second store: %v", err)
}
defer second.Close()
scopeKey := "atomic-" + time.Now().UTC().Format("20060102150405.000000000")
taskIDs := createLeaseTestTasks(t, ctx, first, 256, scopeKey)
defer deleteLeaseTestTasks(t, first, taskIDs)
var successes atomic.Int64
var peak atomic.Int64
monitorCtx, stopMonitor := context.WithCancel(ctx)
var monitorWG sync.WaitGroup
monitorWG.Add(1)
go func() {
defer monitorWG.Done()
ticker := time.NewTicker(2 * time.Millisecond)
defer ticker.Stop()
for {
select {
case <-monitorCtx.Done():
return
case <-ticker.C:
var active int64
if err := first.Pool().QueryRow(monitorCtx, `
SELECT COUNT(*)
FROM gateway_concurrency_leases
WHERE scope_type = 'platform_model'
AND scope_key = $1
AND released_at IS NULL
AND expires_at > now()`, scopeKey).Scan(&active); err == nil {
for active > peak.Load() && !peak.CompareAndSwap(peak.Load(), active) {
}
}
}
}
}()
var wg sync.WaitGroup
errs := make(chan error, len(taskIDs))
for index, taskID := range taskIDs {
wg.Add(1)
go func(index int, taskID string) {
defer wg.Done()
target := first
if index%2 == 1 {
target = second
}
_, err := target.ReserveRateLimits(ctx, taskID, "", []RateLimitReservation{{
ScopeType: "platform_model",
ScopeKey: scopeKey,
Metric: "concurrent",
Limit: 64,
Amount: 1,
LeaseTTLSeconds: 30,
}})
if err == nil {
successes.Add(1)
return
}
if !errors.Is(err, ErrRateLimited) {
errs <- err
}
}(index, taskID)
}
wg.Wait()
stopMonitor()
monitorWG.Wait()
close(errs)
for err := range errs {
t.Fatalf("unexpected reservation error: %v", err)
}
var active int64
if err := first.Pool().QueryRow(ctx, `
SELECT COUNT(*)
FROM gateway_concurrency_leases
WHERE scope_type = 'platform_model'
AND scope_key = $1
AND released_at IS NULL
AND expires_at > now()`, scopeKey).Scan(&active); err != nil {
t.Fatalf("count active leases: %v", err)
}
if successes.Load() != 64 || active != 64 {
t.Fatalf("successful reservations=%d active leases=%d, want exactly 64", successes.Load(), active)
}
if peak.Load() > 64 {
t.Fatalf("active lease peak=%d, want <=64", peak.Load())
}
}
func TestConcurrencyLeaseRenewalExtendsAndReleases(t *testing.T) {
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
if databaseURL == "" {
t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run concurrency lease PostgreSQL integration tests")
}
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
defer cancel()
db, err := Connect(ctx, databaseURL)
if err != nil {
t.Fatalf("connect store: %v", err)
}
defer db.Close()
scopeKey := "renew-" + time.Now().UTC().Format("20060102150405.000000000")
taskIDs := createLeaseTestTasks(t, ctx, db, 1, scopeKey)
defer deleteLeaseTestTasks(t, db, taskIDs)
result, err := db.ReserveRateLimits(ctx, taskIDs[0], "", []RateLimitReservation{{
ScopeType: "platform_model",
ScopeKey: scopeKey,
Metric: "concurrent",
Limit: 1,
Amount: 1,
LeaseTTLSeconds: 2,
}})
if err != nil {
t.Fatalf("reserve short lease: %v", err)
}
time.Sleep(time.Second)
if err := db.RenewConcurrencyLeases(ctx, result.Leases); err != nil {
t.Fatalf("renew short lease: %v", err)
}
time.Sleep(1500 * time.Millisecond)
var active bool
if err := db.Pool().QueryRow(ctx, `
SELECT EXISTS (
SELECT 1 FROM gateway_concurrency_leases
WHERE id = $1::uuid AND released_at IS NULL AND expires_at > now()
)`, result.Leases[0].ID).Scan(&active); err != nil {
t.Fatalf("read renewed lease: %v", err)
}
if !active {
t.Fatal("renewed lease expired at its original TTL")
}
if err := db.ReleaseConcurrencyLeases(ctx, result.Leases); err != nil {
t.Fatalf("release renewed lease: %v", err)
}
}
func createLeaseTestTasks(t *testing.T, ctx context.Context, db *Store, count int, marker string) []string {
t.Helper()
rows, err := db.Pool().Query(ctx, `
INSERT INTO gateway_tasks (kind, run_mode, user_id, model, model_type, request, status, queue_key)
SELECT 'lease-test', 'simulation', $2, 'lease-test', 'text_generate', '{}'::jsonb, 'queued', $2
FROM generate_series(1, $1)
RETURNING id::text`, count, marker)
if err != nil {
t.Fatalf("create lease test tasks: %v", err)
}
defer rows.Close()
ids := make([]string, 0, count)
for rows.Next() {
var id string
if err := rows.Scan(&id); err != nil {
t.Fatalf("scan lease test task: %v", err)
}
ids = append(ids, id)
}
if err := rows.Err(); err != nil {
t.Fatalf("create lease test tasks: %v", err)
}
return ids
}
func deleteLeaseTestTasks(t *testing.T, db *Store, taskIDs []string) {
t.Helper()
cleanupCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
if _, err := db.Pool().Exec(cleanupCtx, `DELETE FROM gateway_tasks WHERE id = ANY($1::uuid[])`, taskIDs); err != nil {
t.Errorf("delete lease test tasks: %v", err)
}
}
apps/api/internal/store/response_chains.go
+3 -3
View File
@@ -45,15 +45,15 @@ type CreateResponseChainInput struct {
}
func (s *Store) CreateResponseChain(ctx context.Context, input CreateResponseChainInput) error {
requestJSON, err := json.Marshal(input.RequestSnapshot)
requestJSON, err := json.Marshal(sanitizeJSONForStorage(input.RequestSnapshot))
if err != nil {
return err
}
responseJSON, err := json.Marshal(input.ResponseSnapshot)
responseJSON, err := json.Marshal(sanitizeJSONForStorage(input.ResponseSnapshot))
if err != nil {
return err
}
internalJSON, err := json.Marshal(input.InternalSnapshot)
internalJSON, err := json.Marshal(sanitizeJSONForStorage(input.InternalSnapshot))
if err != nil {
return err
}
apps/api/internal/store/runtime_policies.go
+177
View File
@@ -8,6 +8,7 @@ import (
"time"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgconn"
)
const runtimePolicyColumns = `
@@ -38,6 +39,24 @@ type PlatformDynamicPriorityState struct {
UpdatedAt time.Time `json:"updatedAt"`
}
type CandidateFailureEffectInput struct {
Effect string
Target string
PlatformID string
PlatformModelID string
RequestedModel string
ModelType string
CooldownSeconds int
DemoteSteps int
SingleSourceProtection bool
}
type CandidateFailureEffectResult struct {
Applied bool
GuardReason string
DynamicPriority int
}
func (s *Store) ListRuntimePolicySets(ctx context.Context) ([]RuntimePolicySet, error) {
rows, err := s.pool.Query(ctx, `SELECT `+runtimePolicyColumns+` FROM model_runtime_policy_sets ORDER BY policy_key ASC`)
if err != nil {
@@ -171,6 +190,164 @@ WHERE id = $1::uuid`, platformModelID, cooldownSeconds)
return err
}
func (s *Store) RuntimeCandidateAvailable(ctx context.Context, platformID string, platformModelID string) (bool, error) {
if strings.TrimSpace(platformID) == "" || strings.TrimSpace(platformModelID) == "" {
return false, nil
}
var available bool
err := s.pool.QueryRow(ctx, `
SELECT EXISTS (
SELECT 1
FROM platform_models model
JOIN integration_platforms platform ON platform.id = model.platform_id
WHERE platform.id = $1::uuid
AND model.id = $2::uuid
AND platform.status = 'enabled'
AND platform.deleted_at IS NULL
AND (platform.cooldown_until IS NULL OR platform.cooldown_until <= now())
AND model.enabled = true
AND (model.cooldown_until IS NULL OR model.cooldown_until <= now())
)`, platformID, platformModelID).Scan(&available)
return available, err
}
func (s *Store) ApplyCandidateFailureEffect(ctx context.Context, input CandidateFailureEffectInput) (CandidateFailureEffectResult, error) {
if input.Effect == "" || input.Effect == "none" {
return CandidateFailureEffectResult{}, nil
}
if input.Effect == "demote" {
steps := input.DemoteSteps
if steps <= 0 {
steps = 1
}
priority, err := s.DemoteCandidatePlatformPriorityBySteps(
ctx,
input.PlatformID,
input.PlatformModelID,
input.RequestedModel,
input.ModelType,
steps,
)
return CandidateFailureEffectResult{Applied: err == nil, DynamicPriority: priority}, err
}
if input.Effect != "cooldown" && input.Effect != "disable" {
return CandidateFailureEffectResult{}, nil
}
tx, err := s.pool.Begin(ctx)
if err != nil {
return CandidateFailureEffectResult{}, err
}
defer func() {
_ = tx.Rollback(ctx)
}()
lockKey := strings.TrimSpace(input.RequestedModel) + "\x1f" + strings.TrimSpace(input.ModelType)
if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1::text, 0))`, lockKey); err != nil {
return CandidateFailureEffectResult{}, err
}
if input.SingleSourceProtection {
alternativeAvailable, err := candidateAlternativeAvailable(ctx, tx, input)
if err != nil {
return CandidateFailureEffectResult{}, err
}
if !alternativeAvailable {
if err := tx.Commit(ctx); err != nil {
return CandidateFailureEffectResult{}, err
}
return CandidateFailureEffectResult{GuardReason: "single_source"}, nil
}
}
var commandTag pgconn.CommandTag
switch {
case input.Effect == "disable" && input.Target == "model":
commandTag, err = tx.Exec(ctx, `
UPDATE platform_models
SET enabled = false,
updated_at = now()
WHERE id = $1::uuid
AND enabled = true`, input.PlatformModelID)
case input.Effect == "disable":
commandTag, err = tx.Exec(ctx, `
UPDATE integration_platforms
SET status = 'disabled',
updated_at = now()
WHERE id = $1::uuid
AND status = 'enabled'
AND deleted_at IS NULL`, input.PlatformID)
case input.Effect == "cooldown" && input.Target == "platform":
commandTag, err = tx.Exec(ctx, `
UPDATE integration_platforms
SET cooldown_until = GREATEST(
COALESCE(cooldown_until, to_timestamp(0)),
now() + ($2::int * interval '1 second')
),
updated_at = now()
WHERE id = $1::uuid
AND status = 'enabled'
AND deleted_at IS NULL`, input.PlatformID, positiveCooldownSeconds(input.CooldownSeconds))
default:
commandTag, err = tx.Exec(ctx, `
UPDATE platform_models
SET cooldown_until = GREATEST(
COALESCE(cooldown_until, to_timestamp(0)),
now() + ($2::int * interval '1 second')
),
updated_at = now()
WHERE id = $1::uuid
AND enabled = true`, input.PlatformModelID, positiveCooldownSeconds(input.CooldownSeconds))
}
if err != nil {
return CandidateFailureEffectResult{}, err
}
applied := commandTag.RowsAffected() > 0
if err := tx.Commit(ctx); err != nil {
return CandidateFailureEffectResult{}, err
}
return CandidateFailureEffectResult{Applied: applied}, nil
}
func candidateAlternativeAvailable(ctx context.Context, tx pgx.Tx, input CandidateFailureEffectInput) (bool, error) {
var available bool
err := tx.QueryRow(ctx, `
SELECT EXISTS (
SELECT 1
FROM platform_models model
JOIN integration_platforms platform ON platform.id = model.platform_id
LEFT JOIN base_model_catalog base ON base.id = model.base_model_id
WHERE platform.id <> $1::uuid
AND platform.status = 'enabled'
AND platform.deleted_at IS NULL
AND (platform.cooldown_until IS NULL OR platform.cooldown_until <= now())
AND model.enabled = true
AND (model.cooldown_until IS NULL OR model.cooldown_until <= now())
AND (NULLIF($3::text, '') IS NULL OR model.model_type @> jsonb_build_array($3::text))
AND (
base.invocation_name = $2::text
OR (base.id IS NULL AND model.model_name = $2::text)
OR EXISTS (
SELECT 1
FROM model_compatibility_aliases compatibility_alias
JOIN base_model_catalog alias_base ON alias_base.id = compatibility_alias.base_model_id
WHERE compatibility_alias.alias = $2::text
AND compatibility_alias.model_type = $3::text
AND compatibility_alias.active = true
AND (compatibility_alias.expires_at IS NULL OR compatibility_alias.expires_at > now())
AND (compatibility_alias.base_model_id = base.id OR alias_base.invocation_name = base.invocation_name)
)
)
)`, input.PlatformID, strings.TrimSpace(input.RequestedModel), strings.TrimSpace(input.ModelType)).Scan(&available)
return available, err
}
func positiveCooldownSeconds(value int) int {
if value <= 0 {
return 300
}
return value
}
func (s *Store) DemoteCandidatePlatformPriority(ctx context.Context, platformID string, platformModelID string, requestedModel string, modelType string) (int, error) {
return s.DemoteCandidatePlatformPriorityBySteps(ctx, platformID, platformModelID, requestedModel, modelType, 99)
}
apps/api/internal/store/runtime_types.go
+49 -16
View File
@@ -13,9 +13,11 @@ var (
)
type ModelCandidateUnavailableError struct {
Code string
Message string
Details map[string]any
Code string
Message string
Details map[string]any
RetryAfter time.Duration
RecoveryAt time.Time
}
func (e *ModelCandidateUnavailableError) Error() string {
@@ -42,6 +44,22 @@ func ModelCandidateErrorDetails(err error) map[string]any {
return nil
}
func ModelCandidateRetryAfter(err error) time.Duration {
var candidateErr *ModelCandidateUnavailableError
if errors.As(err, &candidateErr) && candidateErr.RetryAfter > 0 {
return candidateErr.RetryAfter
}
return 0
}
func ModelCandidateRecoveryAt(err error) time.Time {
var candidateErr *ModelCandidateUnavailableError
if errors.As(err, &candidateErr) {
return candidateErr.RecoveryAt
}
return time.Time{}
}
type RateLimitExceededError struct {
ScopeType string
ScopeKey string
@@ -111,6 +129,7 @@ type CreatePlatformModelInput struct {
PermissionConfig map[string]any `json:"permissionConfig"`
RetryPolicy map[string]any `json:"retryPolicy"`
RateLimitPolicy map[string]any `json:"rateLimitPolicy"`
RateLimitPolicyMode string `json:"rateLimitPolicyMode" enums:"inherit,override"`
RuntimePolicySetID string `json:"runtimePolicySetId"`
RuntimePolicyOverride map[string]any `json:"runtimePolicyOverride"`
Enabled bool `json:"enabled"`
@@ -130,6 +149,7 @@ type RuntimeModelCandidate struct {
DefaultDiscountFactor float64
PlatformRetryPolicy map[string]any
PlatformRateLimitPolicy map[string]any
BaseRateLimitPolicy map[string]any
PlatformPriority int
PlatformModelID string
BaseModelID string
@@ -139,6 +159,7 @@ type RuntimeModelCandidate struct {
ModelAlias string
ModelType string
DisplayName string
LegacyAliasUsed bool
Capabilities map[string]any
CapabilityOverride map[string]any
BaseBillingConfig map[string]any
@@ -152,8 +173,11 @@ type RuntimeModelCandidate struct {
ModelPricingRuleSetID string
ModelRetryPolicy map[string]any
ModelRateLimitPolicy map[string]any
ModelRateLimitPolicyMode string
RuntimePolicySetID string
RuntimePolicyExplicit bool
RuntimePolicyOverride map[string]any
RateLimitRuntimeOverride map[string]any
RuntimeRetryPolicy map[string]any
RuntimeRateLimitPolicy map[string]any
AutoDisablePolicy map[string]any
@@ -172,18 +196,21 @@ type RuntimeModelCandidate struct {
}
type RuntimeCandidateCacheAffinity struct {
Key string
RequestCount int
InputTokens int
CachedInputTokens int
EMAHitRatio float64
LastHitRatio float64
LastObservedUnix float64
Confidence float64
Score float64
Boost float64
AdjustedPriority float64
Applied bool
Key string
RequestCount int
InputTokens int
CachedInputTokens int
EMAHitRatio float64
LastHitRatio float64
LastObservedUnix float64
Confidence float64
Score float64
Boost float64
AdjustedPriority float64
MatchedPrefixDepth int
CandidateCount int
OverrideReason string
Applied bool
}
type RuntimeCandidateLoadMetrics struct {
@@ -218,10 +245,15 @@ type RateLimitReservation struct {
}
type RateLimitResult struct {
LeaseIDs []string
Leases []ConcurrencyLease
Reservations []RateLimitReservation
}
type ConcurrencyLease struct {
ID string
TTL time.Duration
}
type CreateTaskAttemptInput struct {
TaskID string
AttemptNo int
@@ -248,6 +280,7 @@ type FinishTaskAttemptInput struct {
Status string
Retryable bool
RequestID string
StatusCode int
Usage map[string]any
Metrics map[string]any
ResponseSnapshot map[string]any
apps/api/internal/store/task_history.go
+496
View File
@@ -0,0 +1,496 @@
package store
import (
"context"
"time"
"github.com/jackc/pgx/v5"
)
const (
TaskCallbackBatchSize = 100
TaskCallbackLockTTL = 60 * time.Second
)
type TaskCallbackDelivery struct {
ID string
TaskID string
EventID string
Seq int64
CallbackURL string
Status string
Attempts int
LockToken string
EventType string
TaskStatus string
CreatedAt time.Time
}
type TaskHistoryCleanupResult struct {
CompactedTasks int64
CompactedCheckpoints int64
CompactedAttempts int64
CompactedEvents int64
CompactedCallbacks int64
RetiredCallbacks int64
CompactedParamLogs int64
CompactedClonedVoices int64
DeletedCallbacks int64
DeletedParamLogs int64
DeletedEvents int64
DeletedAttempts int64
DeletedTasks int64
}
func (result TaskHistoryCleanupResult) Total() int64 {
return result.CompactedTasks +
result.CompactedCheckpoints +
result.CompactedAttempts +
result.CompactedEvents +
result.CompactedCallbacks +
result.RetiredCallbacks +
result.CompactedParamLogs +
result.CompactedClonedVoices +
result.DeletedCallbacks +
result.DeletedParamLogs +
result.DeletedEvents +
result.DeletedAttempts +
result.DeletedTasks
}
func (s *Store) ClaimTaskCallbacks(ctx context.Context, workerID string, limit int, staleAfter time.Duration) ([]TaskCallbackDelivery, error) {
if limit <= 0 || limit > TaskCallbackBatchSize {
limit = TaskCallbackBatchSize
}
if staleAfter <= 0 {
staleAfter = TaskCallbackLockTTL
}
rows, err := s.pool.Query(ctx, `
WITH picked AS (
SELECT outbox.id, event.event_type, COALESCE(event.status, '') AS task_status, event.created_at
FROM gateway_task_callback_outbox outbox
JOIN gateway_task_events event ON event.id = outbox.event_id
WHERE outbox.created_at >= COALESCE((
SELECT applied_at
FROM schema_migrations
WHERE version = '0083_task_history_minimal_storage'
), now())
AND (
(
outbox.status = 'pending'
AND outbox.next_attempt_at <= now()
)
OR (
outbox.status = 'processing'
AND outbox.locked_at < now() - ($3::int * interval '1 second')
)
)
ORDER BY outbox.next_attempt_at, outbox.created_at
LIMIT $2
FOR UPDATE OF outbox SKIP LOCKED
)
UPDATE gateway_task_callback_outbox outbox
SET status = 'processing',
attempts = outbox.attempts + 1,
locked_by = $1,
lock_token = gen_random_uuid(),
locked_at = now(),
updated_at = now()
FROM picked
WHERE outbox.id = picked.id
RETURNING outbox.id::text, outbox.task_id::text, COALESCE(outbox.event_id::text, ''),
outbox.seq, outbox.callback_url, outbox.status, outbox.attempts,
COALESCE(outbox.lock_token::text, ''), picked.event_type, picked.task_status,
picked.created_at`,
workerID, limit, int(staleAfter/time.Second))
if err != nil {
return nil, err
}
defer rows.Close()
items := make([]TaskCallbackDelivery, 0)
for rows.Next() {
var item TaskCallbackDelivery
if err := rows.Scan(
&item.ID,
&item.TaskID,
&item.EventID,
&item.Seq,
&item.CallbackURL,
&item.Status,
&item.Attempts,
&item.LockToken,
&item.EventType,
&item.TaskStatus,
&item.CreatedAt,
); err != nil {
return nil, err
}
items = append(items, item)
}
return items, rows.Err()
}
func (s *Store) MarkTaskCallbackDelivered(ctx context.Context, item TaskCallbackDelivery) error {
_, err := s.pool.Exec(ctx, `
UPDATE gateway_task_callback_outbox
SET status = 'delivered',
delivered_at = now(),
failed_at = NULL,
last_error = NULL,
locked_by = NULL,
lock_token = NULL,
locked_at = NULL,
updated_at = now()
WHERE id = $1::uuid
AND status = 'processing'
AND lock_token = $2::uuid`, item.ID, item.LockToken)
return err
}
func (s *Store) MarkTaskCallbackFailed(ctx context.Context, item TaskCallbackDelivery, retry bool, nextAttemptAt time.Time, message string) error {
status := "failed"
if retry {
status = "pending"
}
_, err := s.pool.Exec(ctx, `
UPDATE gateway_task_callback_outbox
SET status = $3,
next_attempt_at = CASE WHEN $3 = 'pending' THEN $4::timestamptz ELSE next_attempt_at END,
failed_at = CASE WHEN $3 = 'failed' THEN now() ELSE NULL END,
last_error = NULLIF(left($5, 2048), ''),
locked_by = NULL,
lock_token = NULL,
locked_at = NULL,
updated_at = now()
WHERE id = $1::uuid
AND status = 'processing'
AND lock_token = $2::uuid`,
item.ID, item.LockToken, status, nextAttemptAt, message)
return err
}
func (s *Store) CleanupTaskHistory(ctx context.Context, analysisCutoff time.Time, taskCutoff time.Time, batchSize int) (TaskHistoryCleanupResult, error) {
if batchSize < 1 {
batchSize = 1000
}
var result TaskHistoryCleanupResult
err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
if _, err := tx.Exec(ctx, `SET LOCAL lock_timeout = '5s'; SET LOCAL statement_timeout = '5s'`); err != nil {
return err
}
steps := []struct {
target *int64
sql string
args []any
}{
{&result.CompactedTasks, `
WITH picked AS (
SELECT id
FROM gateway_tasks
WHERE COALESCE(normalized_request, '{}'::jsonb) <> '{}'::jsonb
OR compatibility_public_id IS NOT NULL
OR compatibility_submit_http_status IS NOT NULL
OR COALESCE(compatibility_submit_headers, '{}'::jsonb) <> '{}'::jsonb
OR COALESCE(compatibility_submit_body, '{}'::jsonb) <> '{}'::jsonb
OR result ?| ARRAY[
'raw', 'raw_data', 'rawData', 'provider_response', 'providerResponse',
'submit', 'file_retrieve', 'fileRetrieve', 'upstream_task_id', 'remote_task_id'
]
OR error IS NOT NULL
OR octet_length(COALESCE(error_message, '')) > 2048
ORDER BY updated_at
LIMIT $1
FOR UPDATE SKIP LOCKED
)
UPDATE gateway_tasks task
SET normalized_request = '{}'::jsonb,
compatibility_public_id = NULL,
compatibility_submit_http_status = NULL,
compatibility_submit_headers = '{}'::jsonb,
compatibility_submit_body = '{}'::jsonb,
result = COALESCE(task.result, '{}'::jsonb)
- 'raw' - 'raw_data' - 'rawData'
- 'provider_response' - 'providerResponse'
- 'submit' - 'file_retrieve' - 'fileRetrieve'
- 'upstream_task_id' - 'remote_task_id',
error = NULL,
error_message = CASE
WHEN octet_length(COALESCE(task.error_message, '')) > 2048 THEN left(task.error_message, 512)
ELSE task.error_message
END,
remote_task_payload = CASE
WHEN task.status IN ('succeeded', 'failed', 'cancelled') THEN '{}'::jsonb
ELSE (
SELECT COALESCE(jsonb_object_agg(entry.key, entry.value), '{}'::jsonb)
FROM jsonb_each(COALESCE(task.remote_task_payload, '{}'::jsonb)) entry
WHERE entry.key IN (
'endpoint', 'pollEndpoint', 'phase', 'mode', 'taskType',
'targetResolution', 'imageToken', 'receipt', 'cleanupElementIds'
)
)
END
FROM picked
WHERE task.id = picked.id`, []any{batchSize}},
{&result.CompactedCheckpoints, `
WITH picked AS (
SELECT id
FROM gateway_tasks
WHERE COALESCE(remote_task_payload, '{}'::jsonb) <> '{}'::jsonb
AND (
status IN ('succeeded', 'failed', 'cancelled')
OR EXISTS (
SELECT 1
FROM jsonb_object_keys(COALESCE(remote_task_payload, '{}'::jsonb)) AS key
WHERE key NOT IN (
'endpoint', 'pollEndpoint', 'phase', 'mode', 'taskType',
'targetResolution', 'imageToken', 'receipt', 'cleanupElementIds'
)
)
)
ORDER BY updated_at
LIMIT $1
FOR UPDATE SKIP LOCKED
)
UPDATE gateway_tasks task
SET remote_task_payload = CASE
WHEN task.status IN ('succeeded', 'failed', 'cancelled') THEN '{}'::jsonb
ELSE (
SELECT COALESCE(jsonb_object_agg(entry.key, entry.value), '{}'::jsonb)
FROM jsonb_each(COALESCE(task.remote_task_payload, '{}'::jsonb)) entry
WHERE entry.key IN (
'endpoint', 'pollEndpoint', 'phase', 'mode', 'taskType',
'targetResolution', 'imageToken', 'receipt', 'cleanupElementIds'
)
)
END
FROM picked
WHERE task.id = picked.id`, []any{batchSize}},
{&result.CompactedAttempts, `
WITH picked AS (
SELECT id
FROM gateway_task_attempts
WHERE request_snapshot <> '{}'::jsonb
OR COALESCE(response_snapshot, '{}'::jsonb) <> '{}'::jsonb
OR COALESCE(usage, '{}'::jsonb) <> '{}'::jsonb
OR COALESCE(metrics, '{}'::jsonb) <> '{}'::jsonb
OR COALESCE(pricing_snapshot, '{}'::jsonb) <> '{}'::jsonb
OR request_fingerprint IS NOT NULL
ORDER BY started_at
LIMIT $1
FOR UPDATE SKIP LOCKED
)
UPDATE gateway_task_attempts attempt
SET request_snapshot = '{}'::jsonb,
response_snapshot = '{}'::jsonb,
usage = '{}'::jsonb,
metrics = '{}'::jsonb,
pricing_snapshot = '{}'::jsonb,
request_fingerprint = NULL,
error_message = NULLIF(left(error_message, 2048), '')
FROM picked
WHERE attempt.id = picked.id`, []any{batchSize}},
{&result.CompactedEvents, `
WITH picked AS (
SELECT id
FROM gateway_task_events
WHERE payload <> '{}'::jsonb
OR phase IS NOT NULL
OR COALESCE(progress, 0) <> 0
OR message IS NOT NULL
ORDER BY created_at
LIMIT $1
FOR UPDATE SKIP LOCKED
)
UPDATE gateway_task_events event
SET platform_id = CASE
WHEN event.platform_id IS NULL
AND event.payload->>'platformId' ~* '^[0-9a-f]{8}-[0-9a-f]{4}-[1-5][0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$'
THEN (event.payload->>'platformId')::uuid
ELSE event.platform_id
END,
phase = NULL,
progress = 0,
message = NULL,
payload = '{}'::jsonb
FROM picked
WHERE event.id = picked.id`, []any{batchSize}},
{&result.CompactedCallbacks, `
WITH picked AS (
SELECT id
FROM gateway_task_callback_outbox
WHERE payload <> '{}'::jsonb
ORDER BY created_at
LIMIT $1
FOR UPDATE SKIP LOCKED
)
UPDATE gateway_task_callback_outbox outbox
SET payload = '{}'::jsonb
FROM picked
WHERE outbox.id = picked.id`, []any{batchSize}},
{&result.RetiredCallbacks, `
WITH migration_boundary AS (
SELECT applied_at
FROM schema_migrations
WHERE version = '0083_task_history_minimal_storage'
), picked AS (
SELECT outbox.id
FROM gateway_task_callback_outbox outbox
CROSS JOIN migration_boundary boundary
WHERE outbox.created_at < boundary.applied_at
AND outbox.status IN ('pending', 'processing')
ORDER BY outbox.created_at
LIMIT $1
FOR UPDATE OF outbox SKIP LOCKED
)
UPDATE gateway_task_callback_outbox outbox
SET status = 'failed',
failed_at = outbox.created_at,
last_error = 'legacy_callback_not_replayed',
locked_by = NULL,
lock_token = NULL,
locked_at = NULL,
updated_at = now()
FROM picked
WHERE outbox.id = picked.id`, []any{batchSize}},
{&result.CompactedParamLogs, `
WITH picked AS (
SELECT id
FROM gateway_task_param_preprocessing_logs
WHERE actual_input <> '{}'::jsonb
OR converted_output <> '{}'::jsonb
OR model_snapshot <> '{}'::jsonb
OR octet_length(changes::text) > 1024
ORDER BY created_at
LIMIT $1
FOR UPDATE SKIP LOCKED
)
UPDATE gateway_task_param_preprocessing_logs log
SET actual_input = '{}'::jsonb,
converted_output = '{}'::jsonb,
model_snapshot = '{}'::jsonb,
changes = CASE WHEN octet_length(log.changes::text) > 1024 THEN '[]'::jsonb ELSE log.changes END
FROM picked
WHERE log.id = picked.id`, []any{batchSize}},
{&result.CompactedClonedVoices, `
WITH picked AS (
SELECT id
FROM gateway_cloned_voices
WHERE metadata ?| ARRAY['raw', 'raw_data', 'rawData', 'provider_response', 'providerResponse', 'request']
ORDER BY updated_at
LIMIT $1
FOR UPDATE SKIP LOCKED
)
UPDATE gateway_cloned_voices voice
SET metadata = COALESCE(voice.metadata, '{}'::jsonb)
- 'raw' - 'raw_data' - 'rawData'
- 'provider_response' - 'providerResponse' - 'request',
updated_at = now()
FROM picked
WHERE voice.id = picked.id`, []any{batchSize}},
{&result.DeletedCallbacks, `
WITH picked AS (
SELECT id
FROM gateway_task_callback_outbox
WHERE (status = 'delivered' AND delivered_at < now() - interval '24 hours')
OR (status = 'failed' AND COALESCE(failed_at, updated_at) < $1::timestamptz)
ORDER BY updated_at
LIMIT $2
FOR UPDATE SKIP LOCKED
)
DELETE FROM gateway_task_callback_outbox outbox
USING picked
WHERE outbox.id = picked.id`, []any{analysisCutoff, batchSize}},
{&result.DeletedParamLogs, `
WITH picked AS (
SELECT id
FROM gateway_task_param_preprocessing_logs
WHERE created_at < $1::timestamptz
ORDER BY created_at
LIMIT $2
FOR UPDATE SKIP LOCKED
)
DELETE FROM gateway_task_param_preprocessing_logs log
USING picked
WHERE log.id = picked.id`, []any{analysisCutoff, batchSize}},
{&result.DeletedEvents, `
WITH picked AS (
SELECT event.id
FROM gateway_task_events event
WHERE event.created_at < $1::timestamptz
AND NOT EXISTS (
SELECT 1
FROM gateway_task_callback_outbox outbox
WHERE outbox.event_id = event.id
AND outbox.status IN ('pending', 'processing')
)
ORDER BY event.created_at
LIMIT $2
FOR UPDATE OF event SKIP LOCKED
)
DELETE FROM gateway_task_events event
USING picked
WHERE event.id = picked.id`, []any{analysisCutoff, batchSize}},
{&result.DeletedAttempts, `
WITH picked AS (
SELECT id
FROM gateway_task_attempts
WHERE COALESCE(finished_at, started_at) < $1::timestamptz
AND status <> 'running'
ORDER BY COALESCE(finished_at, started_at)
LIMIT $2
FOR UPDATE SKIP LOCKED
)
DELETE FROM gateway_task_attempts attempt
USING picked
WHERE attempt.id = picked.id`, []any{analysisCutoff, batchSize}},
{&result.DeletedTasks, `
WITH picked AS (
SELECT task.id
FROM gateway_tasks task
WHERE task.finished_at < $1::timestamptz
AND task.status IN ('succeeded', 'failed', 'cancelled')
AND task.billing_status IN ('settled', 'released', 'not_required')
AND (task.execution_lease_expires_at IS NULL OR task.execution_lease_expires_at <= now())
AND NOT EXISTS (
SELECT 1
FROM settlement_outbox settlement
WHERE settlement.task_id = task.id
AND settlement.status IN ('pending', 'processing', 'retryable_failed', 'manual_review')
)
AND NOT EXISTS (
SELECT 1
FROM gateway_task_callback_outbox callback
WHERE callback.task_id = task.id
)
AND NOT EXISTS (
SELECT 1
FROM gateway_concurrency_leases lease
WHERE lease.task_id = task.id
AND lease.released_at IS NULL
AND lease.expires_at > now()
)
AND NOT EXISTS (
SELECT 1
FROM gateway_rate_limit_reservations reservation
WHERE reservation.task_id = task.id
AND reservation.status = 'reserved'
)
ORDER BY task.finished_at
LIMIT $2
FOR UPDATE OF task SKIP LOCKED
)
DELETE FROM gateway_tasks task
USING picked
WHERE task.id = picked.id`, []any{taskCutoff, batchSize}},
}
for _, step := range steps {
tag, err := tx.Exec(ctx, step.sql, step.args...)
if err != nil {
return err
}
*step.target = tag.RowsAffected()
}
return nil
})
return result, err
}
apps/api/internal/store/task_history_test.go
+401
View File
@@ -0,0 +1,401 @@
package store
import (
"context"
"strings"
"testing"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
"github.com/google/uuid"
)
func TestTaskEventAllowlistRejectsSyntheticProgress(t *testing.T) {
if taskEventAllowed("task.progress") || taskEventAllowed("task.attempt.started") || taskEventAllowed("unknown") {
t.Fatal("synthetic or unknown event type was allowed")
}
for _, eventType := range []string{"task.accepted", "task.running", "task.completed", "task.attempt.failed", "task.policy.priority_demoted"} {
if !taskEventAllowed(eventType) {
t.Fatalf("required event type %q was rejected", eventType)
}
}
}
func TestMinimalTaskResultDropsProviderRawCopy(t *testing.T) {
result := minimalTaskResult(map[string]any{
"data": []any{map[string]any{"url": "https://example.invalid/result"}},
"raw": map[string]any{"provider": "duplicate"},
"raw_data": map[string]any{"provider": "duplicate"},
"providerResponse": map[string]any{"provider": "duplicate"},
"upstream_task_id": "remote-task",
"provider_specific": "kept",
})
if result["raw"] != nil || result["raw_data"] != nil || result["providerResponse"] != nil ||
result["upstream_task_id"] != nil || result["data"] == nil || result["provider_specific"] != "kept" {
t.Fatalf("minimal task result=%#v", result)
}
}
func TestTaskEventsAreMinimalAndConsecutiveDuplicatesAreSkipped(t *testing.T) {
db := billingV2IntegrationStore(t)
ctx := context.Background()
user := &auth.User{ID: "event-minimal-" + uuid.NewString()}
task, err := db.CreateTask(ctx, CreateTaskInput{
Kind: "images.generations", Model: "event-minimal", RunMode: "simulation",
Request: map[string]any{"prompt": "minimal"},
}, user)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_tasks WHERE id=$1::uuid`, task.ID)
})
running, err := db.AddTaskEvent(ctx, task.ID, "task.running", "running", "starting", 0.1, "ignored", map[string]any{"large": "ignored"}, true)
if err != nil || running.ID == "" || len(running.Payload) != 0 || running.Phase != "" || running.Progress != 0 || running.Message != "" {
t.Fatalf("running=%+v err=%v", running, err)
}
duplicate, err := db.AddTaskEvent(ctx, task.ID, "task.running", "running", "polling", 0.8, "ignored again", map[string]any{"other": "ignored"}, true)
if err != nil || duplicate.ID != "" || duplicate.SkippedReason != "duplicate" {
t.Fatalf("duplicate=%+v err=%v", duplicate, err)
}
synthetic, err := db.AddTaskEvent(ctx, task.ID, "task.progress", "running", "polling", 0.9, "ignored", nil, true)
if err != nil || synthetic.ID != "" || synthetic.SkippedReason != "unknown_type" {
t.Fatalf("synthetic=%+v err=%v", synthetic, err)
}
completed, err := db.AddTaskEvent(ctx, task.ID, "task.completed", "succeeded", "completed", 1, "ignored", map[string]any{"result": map[string]any{"duplicate": true}}, true)
if err != nil || completed.ID == "" || len(completed.Payload) != 0 {
t.Fatalf("completed=%+v err=%v", completed, err)
}
if err := db.QueueTaskCallback(ctx, completed, "https://callback.invalid/task"); err != nil {
t.Fatal(err)
}
var events int
var nonEmptyPayloads int
if err := db.pool.QueryRow(ctx, `
SELECT count(*), count(*) FILTER (WHERE payload <> '{}'::jsonb)
FROM gateway_task_events
WHERE task_id=$1::uuid`, task.ID).Scan(&events, &nonEmptyPayloads); err != nil {
t.Fatal(err)
}
if events != 3 || nonEmptyPayloads != 0 {
t.Fatalf("events=%d nonEmptyPayloads=%d, want 3/0", events, nonEmptyPayloads)
}
var callbackPayloadBytes int
if err := db.pool.QueryRow(ctx, `
SELECT pg_column_size(payload)
FROM gateway_task_callback_outbox
WHERE task_id=$1::uuid`, task.ID).Scan(&callbackPayloadBytes); err != nil {
t.Fatal(err)
}
if callbackPayloadBytes > 16 {
t.Fatalf("callback payload storage=%d, want an empty json object", callbackPayloadBytes)
}
budgetExceeded := false
for index := 0; index < 20; index++ {
eventType := "task.running"
status := "running"
if index%2 == 1 {
eventType = "task.queued"
status = "queued"
}
event, err := db.AddTaskEvent(ctx, task.ID, eventType, status, "", 0, "", nil, true)
if err != nil {
t.Fatal(err)
}
if event.SkippedReason == "budget_exceeded" {
budgetExceeded = true
break
}
}
if !budgetExceeded {
t.Fatal("non-terminal event budget was not enforced")
}
var nonTerminalEvents int
if err := db.pool.QueryRow(ctx, `
SELECT count(*)
FROM gateway_task_events
WHERE task_id=$1::uuid
AND event_type NOT IN (
'task.completed', 'task.failed', 'task.cancelled',
'task.billing.settled', 'task.billing.released', 'task.billing.review'
)`, task.ID).Scan(&nonTerminalEvents); err != nil {
t.Fatal(err)
}
if nonTerminalEvents != 16 {
t.Fatalf("non-terminal events=%d, want 16", nonTerminalEvents)
}
}
func TestTaskAttemptKeepsOnlyCompactDiagnostics(t *testing.T) {
db := billingV2IntegrationStore(t)
ctx := context.Background()
user := &auth.User{ID: "attempt-minimal-" + uuid.NewString()}
task, err := db.CreateTask(ctx, CreateTaskInput{
Kind: "images.generations", Model: "attempt-minimal", RunMode: "simulation",
Request: map[string]any{"prompt": "canonical"},
}, user)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_tasks WHERE id=$1::uuid`, task.ID)
})
attemptID, err := db.CreateTaskAttempt(ctx, CreateTaskAttemptInput{
TaskID: task.ID, AttemptNo: 1, QueueKey: "test", Status: "running",
RequestSnapshot: map[string]any{"prompt": "duplicate"},
Metrics: map[string]any{"trace": []any{map[string]any{"large": "duplicate"}}},
PricingSnapshot: map[string]any{"price": 1},
RequestFingerprint: "duplicate-fingerprint",
})
if err != nil {
t.Fatal(err)
}
if err := db.FinishTaskAttempt(ctx, FinishTaskAttemptInput{
AttemptID: attemptID,
Status: "failed",
Retryable: true,
StatusCode: 503,
Usage: map[string]any{"output_tokens": 100, "raw": strings.Repeat("duplicate", 100)},
Metrics: map[string]any{"trace": []any{map[string]any{"large": "duplicate"}}, "cacheAffinityMatched": true, "cacheAffinityMatchedPrefixDepth": 3},
ResponseSnapshot: map[string]any{"result": "duplicate"},
ErrorCode: "upstream_error",
ErrorMessage: strings.Repeat("错误", 4096),
}); err != nil {
t.Fatal(err)
}
attempts, err := db.ListTaskAttempts(ctx, task.ID)
if err != nil {
t.Fatal(err)
}
if len(attempts) != 1 {
t.Fatalf("attempts=%d", len(attempts))
}
attempt := attempts[0]
if len(attempt.RequestSnapshot) != 0 || len(attempt.ResponseSnapshot) != 0 || len(attempt.PricingSnapshot) != 0 {
t.Fatalf("attempt contains duplicate snapshots: %+v", attempt)
}
if len(attempt.Usage) != 1 || taskAttemptMetricInt(attempt.Usage, "output_tokens") != 100 {
t.Fatalf("attempt compact usage mismatch: %+v", attempt.Usage)
}
if len(attempt.Metrics) != 2 ||
attempt.Metrics["cacheAffinityMatched"] != true ||
taskAttemptMetricInt(attempt.Metrics, "cacheAffinityMatchedPrefixDepth") != 3 {
t.Fatalf("attempt compact metrics mismatch: %+v", attempt.Metrics)
}
if attempt.StatusCode != 503 || len(attempt.ErrorMessage) > 2048 {
t.Fatalf("attempt status/error not preserved safely: %+v", attempt)
}
}
func TestTaskRunningDoesNotPersistNormalizedRequestCopy(t *testing.T) {
db := billingV2IntegrationStore(t)
ctx := context.Background()
user := &auth.User{ID: "request-minimal-" + uuid.NewString()}
task, err := db.CreateTask(ctx, CreateTaskInput{
Kind: "images.generations", Model: "request-minimal", RunMode: "simulation",
Request: map[string]any{"prompt": "canonical"},
}, user)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_tasks WHERE id=$1::uuid`, task.ID)
})
executionToken := uuid.NewString()
if _, err := db.ClaimTaskExecution(ctx, task.ID, executionToken, time.Minute); err != nil {
t.Fatal(err)
}
if err := db.MarkTaskRunning(ctx, task.ID, executionToken, "image_generate", map[string]any{
"prompt": "duplicate normalized request",
}); err != nil {
t.Fatal(err)
}
var normalizedBytes int
if err := db.pool.QueryRow(ctx, `
SELECT pg_column_size(normalized_request)
FROM gateway_tasks
WHERE id=$1::uuid`, task.ID).Scan(&normalizedBytes); err != nil {
t.Fatal(err)
}
if normalizedBytes > 16 {
t.Fatalf("normalized request storage=%d, want an empty json object", normalizedBytes)
}
}
func TestTaskHistoryCleanupRemovesHistoricalProviderCopies(t *testing.T) {
db := billingV2IntegrationStore(t)
ctx := context.Background()
user := &auth.User{ID: "history-compaction-" + uuid.NewString()}
task, err := db.CreateTask(ctx, CreateTaskInput{
Kind: "images.generations", Model: "history-compaction", RunMode: "simulation",
Request: map[string]any{"prompt": "canonical"},
}, user)
if err != nil {
t.Fatal(err)
}
voiceID := uuid.NewString()
t.Cleanup(func() {
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_cloned_voices WHERE id=$1::uuid`, voiceID)
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_tasks WHERE id=$1::uuid`, task.ID)
})
if _, err := db.pool.Exec(ctx, `
UPDATE gateway_tasks
SET result='{"data":[{"url":"https://example.invalid/result"}],"raw_data":{"provider":"duplicate"},"upstream_task_id":"remote"}'::jsonb
WHERE id=$1::uuid`,
task.ID); err != nil {
t.Fatal(err)
}
if _, err := db.pool.Exec(ctx, `
INSERT INTO gateway_cloned_voices (id, user_id, provider, voice_id, metadata)
VALUES ($1::uuid, 'history-compaction', 'test', $2, '{"request":{"duplicate":true},"rawData":{"provider":"duplicate"},"keep":"value"}'::jsonb)`,
voiceID, "voice-"+voiceID); err != nil {
t.Fatal(err)
}
result, err := db.CleanupTaskHistory(ctx, time.Now().AddDate(0, 0, -7), time.Now().AddDate(0, 0, -30), 100)
if err != nil {
t.Fatal(err)
}
if result.CompactedTasks < 1 || result.CompactedClonedVoices < 1 {
t.Fatalf("cleanup result=%+v", result)
}
var taskResult map[string]any
var voiceMetadata map[string]any
if err := db.pool.QueryRow(ctx, `SELECT result FROM gateway_tasks WHERE id=$1::uuid`, task.ID).Scan(&taskResult); err != nil {
t.Fatal(err)
}
if err := db.pool.QueryRow(ctx, `SELECT metadata FROM gateway_cloned_voices WHERE id=$1::uuid`, voiceID).Scan(&voiceMetadata); err != nil {
t.Fatal(err)
}
if taskResult["raw_data"] != nil || taskResult["upstream_task_id"] != nil || taskResult["data"] == nil {
t.Fatalf("compacted task result=%+v", taskResult)
}
if voiceMetadata["request"] != nil || voiceMetadata["rawData"] != nil || voiceMetadata["keep"] != "value" {
t.Fatalf("compacted voice metadata=%+v", voiceMetadata)
}
}
func TestTaskHistoryCleanupKeepsTaskUntilCallbackRetentionExpires(t *testing.T) {
db := billingV2IntegrationStore(t)
ctx := context.Background()
user := &auth.User{ID: "cleanup-" + uuid.NewString()}
createOldTask := func(label string) GatewayTask {
t.Helper()
task, err := db.CreateTask(ctx, CreateTaskInput{
Kind: "images.generations", Model: label, RunMode: "simulation",
Request: map[string]any{"prompt": label},
}, user)
if err != nil {
t.Fatal(err)
}
if _, err := db.pool.Exec(ctx, `
UPDATE gateway_tasks
SET status='succeeded', billing_status='not_required',
finished_at=now()-interval '40 days', updated_at=now()-interval '40 days'
WHERE id=$1::uuid`, task.ID); err != nil {
t.Fatal(err)
}
if _, err := db.pool.Exec(ctx, `
UPDATE gateway_task_events
SET created_at=now()-interval '10 days'
WHERE task_id=$1::uuid`, task.ID); err != nil {
t.Fatal(err)
}
return task
}
safe := createOldTask("cleanup-safe")
protected := createOldTask("cleanup-protected")
t.Cleanup(func() {
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_tasks WHERE id IN ($1::uuid, $2::uuid)`, safe.ID, protected.ID)
})
completed, err := db.AddTaskEvent(ctx, protected.ID, "task.completed", "succeeded", "", 0, "", nil, true)
if err != nil {
t.Fatal(err)
}
if err := db.QueueTaskCallback(ctx, completed, "https://callback.invalid/task"); err != nil {
t.Fatal(err)
}
if _, err := db.pool.Exec(ctx, `
UPDATE gateway_task_callback_outbox
SET status='delivered', delivered_at=now(), updated_at=now()
WHERE task_id=$1::uuid`, protected.ID); err != nil {
t.Fatal(err)
}
result, err := db.CleanupTaskHistory(ctx, time.Now().AddDate(0, 0, -7), time.Now().AddDate(0, 0, -30), 100)
if err != nil {
t.Fatal(err)
}
if result.DeletedTasks < 1 {
t.Fatalf("deletedTasks=%d, want at least 1", result.DeletedTasks)
}
var safeExists bool
var protectedExists bool
if err := db.pool.QueryRow(ctx, `SELECT EXISTS (SELECT 1 FROM gateway_tasks WHERE id=$1::uuid)`, safe.ID).Scan(&safeExists); err != nil {
t.Fatal(err)
}
if err := db.pool.QueryRow(ctx, `SELECT EXISTS (SELECT 1 FROM gateway_tasks WHERE id=$1::uuid)`, protected.ID).Scan(&protectedExists); err != nil {
t.Fatal(err)
}
if safeExists || !protectedExists {
t.Fatalf("safeExists=%t protectedExists=%t", safeExists, protectedExists)
}
}
func TestLegacyCallbacksAreNotReplayed(t *testing.T) {
db := billingV2IntegrationStore(t)
ctx := context.Background()
user := &auth.User{ID: "legacy-callback-" + uuid.NewString()}
task, err := db.CreateTask(ctx, CreateTaskInput{
Kind: "images.generations", Model: "legacy-callback", RunMode: "simulation",
Request: map[string]any{"prompt": "legacy"},
}, user)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_tasks WHERE id=$1::uuid`, task.ID)
})
event, err := db.AddTaskEvent(ctx, task.ID, "task.running", "running", "", 0, "", nil, true)
if err != nil {
t.Fatal(err)
}
if err := db.QueueTaskCallback(ctx, event, "https://callback.invalid/legacy"); err != nil {
t.Fatal(err)
}
if _, err := db.pool.Exec(ctx, `
UPDATE gateway_task_callback_outbox
SET created_at=(
SELECT applied_at - interval '1 second'
FROM schema_migrations
WHERE version='0083_task_history_minimal_storage'
)
WHERE task_id=$1::uuid`, task.ID); err != nil {
t.Fatal(err)
}
claimed, err := db.ClaimTaskCallbacks(ctx, "legacy-test", 10, time.Minute)
if err != nil {
t.Fatal(err)
}
if len(claimed) != 0 {
t.Fatalf("claimed legacy callbacks=%d, want 0", len(claimed))
}
result, err := db.CleanupTaskHistory(ctx, time.Now().AddDate(0, 0, -7), time.Now().AddDate(0, 0, -30), 100)
if err != nil {
t.Fatal(err)
}
if result.RetiredCallbacks < 1 {
t.Fatalf("retired legacy callbacks=%d, want at least 1", result.RetiredCallbacks)
}
var status string
if err := db.pool.QueryRow(ctx, `
SELECT status
FROM gateway_task_callback_outbox
WHERE task_id=$1::uuid`, task.ID).Scan(&status); err != nil {
t.Fatal(err)
}
if status != "failed" {
t.Fatalf("legacy callback status=%q, want failed", status)
}
}
apps/api/internal/store/tasks_runtime.go
+299 -104
View File
@@ -8,6 +8,7 @@ import (
"strconv"
"strings"
"time"
"unicode/utf8"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
"github.com/jackc/pgx/v5"
@@ -324,7 +325,7 @@ 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 = NULL,
error_code = 'upstream_submission_unknown',
error_message = 'upstream submission result is unknown',
locked_by = NULL,
@@ -332,6 +333,7 @@ SET status = 'failed',
heartbeat_at = NULL,
execution_token = NULL,
execution_lease_expires_at = NULL,
remote_task_payload = '{}'::jsonb,
finished_at = now(),
updated_at = now()
WHERE id = $1::uuid`, taskID, hasGatewayUser); err != nil {
@@ -416,19 +418,18 @@ SELECT COALESCE((SELECT status = 'running' FROM gateway_tasks WHERE id = $1::uui
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))
func (s *Store) MarkTaskRunning(ctx context.Context, taskID string, executionToken string, modelType string, _ map[string]any) error {
tag, err := s.pool.Exec(ctx, `
UPDATE gateway_tasks
SET status = 'running',
model_type = NULLIF($3::text, ''),
normalized_request = $4::jsonb,
normalized_request = '{}'::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))
AND execution_token = $2::uuid`, taskID, executionToken, modelType)
if err != nil {
return err
}
@@ -469,6 +470,13 @@ func (s *Store) RequeueTask(ctx context.Context, taskID string, executionToken s
delay = 10 * time.Minute
}
nextRunAt := time.Now().Add(delay)
return s.RequeueTaskUntil(ctx, taskID, executionToken, nextRunAt, queueKey)
}
func (s *Store) RequeueTaskUntil(ctx context.Context, taskID string, executionToken string, nextRunAt time.Time, queueKey string) (GatewayTask, error) {
if nextRunAt.Before(time.Now().Add(time.Second)) {
nextRunAt = time.Now().Add(time.Second)
}
return scanGatewayTask(s.pool.QueryRow(ctx, `
UPDATE gateway_tasks
SET status = 'queued',
@@ -502,7 +510,7 @@ WHERE id = $1::uuid`, taskID, riverJobID)
}
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))
payloadJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(payload)))
return pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
tag, err := tx.Exec(ctx, `
UPDATE gateway_tasks
@@ -557,12 +565,14 @@ func (s *Store) CancelQueuedTask(ctx context.Context, taskID string, message str
if message == "" {
message = "任务已取消"
}
message = truncateUTF8Bytes(message, 2048)
tag, err := s.pool.Exec(ctx, `
UPDATE gateway_tasks
SET status = 'cancelled',
error = NULLIF($2::text, ''),
error = NULL,
error_code = 'task_cancelled',
error_message = NULLIF($2::text, ''),
remote_task_payload = '{}'::jsonb,
locked_by = NULL,
locked_at = NULL,
heartbeat_at = NULL,
@@ -592,6 +602,7 @@ func (s *Store) CancelSubmittedTask(ctx context.Context, taskID string, executio
if message == "" {
message = "任务已由上游取消"
}
message = truncateUTF8Bytes(message, 2048)
var task GatewayTask
changed := false
err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
@@ -599,7 +610,7 @@ func (s *Store) CancelSubmittedTask(ctx context.Context, taskID string, executio
task, err = scanGatewayTask(tx.QueryRow(ctx, `
UPDATE gateway_tasks
SET status = 'cancelled',
error = NULLIF($2, ''),
error = NULL,
error_code = 'task_cancelled',
error_message = NULLIF($2, ''),
billing_status = CASE
@@ -613,6 +624,7 @@ SET status = 'cancelled',
heartbeat_at = NULL,
execution_token = NULL,
execution_lease_expires_at = NULL,
remote_task_payload = '{}'::jsonb,
finished_at = now(),
updated_at = now()
WHERE id = $1::uuid
@@ -683,9 +695,6 @@ LIMIT $1`, limit)
}
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
@@ -701,7 +710,7 @@ INSERT INTO gateway_task_attempts (
)
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, ''),
$7, $8, '{}'::jsonb, '{}'::jsonb, '{}'::jsonb, NULL,
'not_submitted', now()
)
RETURNING id::text`,
@@ -713,10 +722,6 @@ RETURNING id::text`,
input.QueueKey,
firstNonEmpty(input.Status, "running"),
input.Simulated,
string(requestJSON),
string(metricsJSON),
string(pricingJSON),
input.RequestFingerprint,
).Scan(&attemptID)
if err != nil {
return "", err
@@ -731,13 +736,16 @@ WHERE id = $1::uuid`, input.TaskID, input.AttemptNo); err != nil {
}
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.Changed {
return "", nil
}
changesJSON, _ := json.Marshal(sanitizeJSONForStorage(input.Changes))
if input.Changes == nil {
changesJSON = []byte("[]")
}
modelSnapshotJSON, _ := json.Marshal(emptyObjectIfNil(input.ModelSnapshot))
if len(changesJSON) > 1024 {
changesJSON = []byte("[]")
}
var attemptNo any
if input.AttemptNo > 0 {
attemptNo = input.AttemptNo
@@ -751,7 +759,7 @@ INSERT INTO gateway_task_param_preprocessing_logs (
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
true, $8::int, '{}'::jsonb, '{}'::jsonb, $9::jsonb, '{}'::jsonb
)
RETURNING id::text`,
input.TaskID,
@@ -761,12 +769,8 @@ RETURNING id::text`,
input.PlatformID,
input.PlatformModelID,
input.ClientID,
input.Changed,
input.ChangeCount,
string(actualInputJSON),
string(convertedOutputJSON),
string(changesJSON),
string(modelSnapshotJSON),
).Scan(&id)
return id, err
}
@@ -823,24 +827,30 @@ ORDER BY COALESCE(attempt_no, 0), created_at`, taskID)
}
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, `
if strings.TrimSpace(taskID) == "" || attemptNo <= 0 || len(entry) == 0 {
return nil
}
encoded, err := json.Marshal(entry)
if err != nil {
return err
}
result, 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),
COALESCE(metrics->'trace', '[]'::jsonb) || jsonb_build_array($3::jsonb),
true
)
WHERE task_id = $1::uuid
AND attempt_no = $2::int`, taskID, attemptNo, string(entryJSON))
return err
AND attempt_no = $2::int`, taskID, attemptNo, encoded)
if err != nil {
return err
}
if result.RowsAffected() == 0 {
return pgx.ErrNoRows
}
return nil
}
func (s *Store) listTaskAttemptsByTaskIDs(ctx context.Context, taskIDs []string) (map[string][]TaskAttempt, error) {
@@ -855,7 +865,8 @@ SELECT a.id::text, a.task_id::text, a.attempt_no,
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),
COALESCE(a.request_id, ''), COALESCE(a.status_code, 0),
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, ''),
@@ -908,6 +919,7 @@ func scanTaskAttempt(scanner taskScanner) (TaskAttempt, error) {
&item.Retryable,
&item.Simulated,
&item.RequestID,
&item.StatusCode,
&usageBytes,
&metricsBytes,
&requestBytes,
@@ -980,7 +992,9 @@ func enrichTaskAttemptFromMetrics(item *TaskAttempt) {
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")
if item.StatusCode == 0 {
item.StatusCode = taskAttemptMetricInt(item.Metrics, "statusCode")
}
}
func taskAttemptMetricString(metrics map[string]any, key string) string {
@@ -1008,47 +1022,96 @@ func taskAttemptMetricInt(metrics map[string]any, key string) int {
}
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))
statusCode := input.StatusCode
if statusCode == 0 {
statusCode = taskAttemptMetricInt(input.Metrics, "statusCode")
}
usageJSON, _ := json.Marshal(sanitizeJSONForStorage(minimalTaskAttemptUsage(input.Usage)))
metricsJSON, _ := json.Marshal(sanitizeJSONForStorage(minimalTaskAttemptMetrics(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, ''),
status_code = NULLIF($5::int, 0),
usage = $11::jsonb,
metrics = $12::jsonb,
response_snapshot = '{}'::jsonb,
pricing_snapshot = '{}'::jsonb,
request_fingerprint = NULL,
response_started_at = $6::timestamptz,
response_finished_at = $7::timestamptz,
response_duration_ms = $8,
error_code = NULLIF($9::text, ''),
error_message = NULLIF(left($10::text, 2048), ''),
finished_at = now()
WHERE id = $1::uuid`,
input.AttemptID,
input.Status,
input.Retryable,
input.RequestID,
string(usageJSON),
string(metricsJSON),
string(responseJSON),
statusCode,
nullableTime(input.ResponseStartedAt),
nullableTime(input.ResponseFinishedAt),
input.ResponseDurationMS,
input.ErrorCode,
input.ErrorMessage,
truncateUTF8Bytes(input.ErrorMessage, 2048),
usageJSON,
metricsJSON,
)
return err
}
func minimalTaskAttemptUsage(usage map[string]any) map[string]any {
return whitelistedTaskAttemptMap(usage, []string{
"inputTokens", "promptTokens", "input_tokens", "prompt_tokens",
"cachedInputTokens", "cachedPromptTokens", "cached_input_tokens", "cached_tokens",
"cachedInputTokensKnown",
"outputTokens", "completionTokens", "output_tokens", "completion_tokens",
"totalTokens", "total_tokens",
})
}
func minimalTaskAttemptMetrics(metrics map[string]any) map[string]any {
return whitelistedTaskAttemptMap(metrics, []string{
"platformPriority", "currentPriority", "loadRatio", "loadAvoided",
"cacheAffinityKey", "cacheAdjustedPriority", "cacheAffinitySamples",
"cacheAffinityScore", "cacheAffinityConfidence", "cacheAffinityHitRatio",
"cacheAffinityEMAHitRatio", "cacheAffinityLastHitRatio",
"cacheAffinityCachedInputTokens", "cacheAffinityBoost",
"cacheAffinityApplied", "cacheAffinityMatched",
"cacheAffinityCandidateCount", "cacheAffinityMatchedPrefixDepth",
"cacheAffinityOverrideReason",
})
}
func whitelistedTaskAttemptMap(input map[string]any, keys []string) map[string]any {
out := map[string]any{}
for _, key := range keys {
if value, ok := input[key]; ok {
out[key] = value
}
}
return out
}
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))
resultReport := sanitizeJSONForStorageWithReport(minimalTaskResult(input.Result))
if resultReport.BinaryCount > 0 {
return GatewayTask{}, &taskPayloadBinaryError{
target: ErrTaskResultBinaryNotMaterialized,
code: "result_binary_not_materialized",
count: resultReport.BinaryCount,
}
}
resultJSON, _ := json.Marshal(resultReport.Value)
billingsJSON, _ := json.Marshal(sanitizeJSONForStorage(input.Billings))
usageJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(input.Usage)))
metricsJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(input.Metrics)))
attemptUsageJSON, _ := json.Marshal(sanitizeJSONForStorage(minimalTaskAttemptUsage(input.Usage)))
attemptMetricsJSON, _ := json.Marshal(sanitizeJSONForStorage(minimalTaskAttemptMetrics(input.Metrics)))
billingSummaryJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(input.BillingSummary)))
pricingSnapshotJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(input.PricingSnapshot)))
finalChargeAmount := strings.TrimSpace(input.FinalChargeAmountText)
if finalChargeAmount == "" {
finalChargeAmount = strconv.FormatFloat(input.FinalChargeAmount, 'f', 9, 64)
@@ -1061,23 +1124,24 @@ 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, ''),
status_code = NULL,
usage = $6::jsonb,
metrics = $7::jsonb,
response_snapshot = '{}'::jsonb,
pricing_snapshot = '{}'::jsonb,
request_fingerprint = NULL,
upstream_submission_status = 'response_received',
upstream_submission_updated_at = now(),
response_started_at = $8::timestamptz,
response_finished_at = $9::timestamptz,
response_duration_ms = $10,
response_started_at = $3::timestamptz,
response_finished_at = $4::timestamptz,
response_duration_ms = $5,
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,
input.AttemptID, input.RequestID,
nullableTime(input.ResponseStartedAt), nullableTime(input.ResponseFinishedAt), input.ResponseDurationMS,
attemptUsageJSON, attemptMetricsJSON,
); err != nil {
return err
}
@@ -1113,6 +1177,7 @@ SET status = 'succeeded',
heartbeat_at = NULL,
execution_token = NULL,
execution_lease_expires_at = NULL,
remote_task_payload = '{}'::jsonb,
finished_at = now(),
updated_at = now()
WHERE id = $1::uuid
@@ -1166,8 +1231,21 @@ func (s *Store) FinishTaskManualReview(ctx context.Context, input FinishTaskManu
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))
message := truncateUTF8Bytes(input.Message, 2048)
result := input.Result
if status == "failed" {
result = nil
}
resultReport := sanitizeJSONForStorageWithReport(minimalTaskResult(result))
if resultReport.BinaryCount > 0 {
return GatewayTask{}, &taskPayloadBinaryError{
target: ErrTaskResultBinaryNotMaterialized,
code: "result_binary_not_materialized",
count: resultReport.BinaryCount,
}
}
resultJSON, _ := json.Marshal(resultReport.Value)
pricingSnapshotJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(input.PricingSnapshot)))
err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
if strings.TrimSpace(input.AttemptID) != "" {
attemptStatus := "failed"
@@ -1186,7 +1264,7 @@ SET status = $2,
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,
WHERE id = $1::uuid`, input.AttemptID, attemptStatus, input.RequestID, input.Code, message,
nullableTime(input.ResponseStartedAt), nullableTime(input.ResponseFinishedAt), input.ResponseDurationMS); err != nil {
return err
}
@@ -1196,7 +1274,7 @@ 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 = NULL,
error_code = NULLIF($5, ''),
error_message = NULLIF($6, ''),
billing_status = 'manual_review',
@@ -1211,11 +1289,12 @@ SET status = $2,
heartbeat_at = NULL,
execution_token = NULL,
execution_lease_expires_at = NULL,
remote_task_payload = '{}'::jsonb,
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,
AND execution_token = $12::uuid`, input.TaskID, status, string(resultJSON), input.RequestID, input.Code, message,
string(pricingSnapshotJSON), input.RequestFingerprint, nullableTime(input.ResponseStartedAt),
nullableTime(input.ResponseFinishedAt), input.ResponseDurationMS, input.ExecutionToken)
if err != nil {
@@ -1268,7 +1347,7 @@ func (s *Store) SettleTaskBilling(ctx context.Context, task GatewayTask) error {
"billings": task.Billings,
"billingSummary": task.BillingSummary,
}
metadata, _ := json.Marshal(metadataMap)
metadata, _ := json.Marshal(sanitizeJSONForStorage(metadataMap))
return pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
if _, err := tx.Exec(ctx, `
INSERT INTO gateway_wallet_accounts (
@@ -1360,7 +1439,7 @@ ON CONFLICT (account_id, idempotency_key) WHERE idempotency_key IS NOT NULL DO N
"frozenBefore": roundMoney(frozenBefore),
"frozenAfter": frozenAfter,
})
metadata, _ = json.Marshal(billingMetadata)
metadata, _ = json.Marshal(sanitizeJSONForStorage(billingMetadata))
if _, err := tx.Exec(ctx, `
INSERT INTO gateway_wallet_transactions (
account_id, gateway_tenant_id, gateway_user_id, direction, transaction_type,
@@ -1396,13 +1475,14 @@ func taskBillingString(value any) string {
}
func (s *Store) FinishTaskFailure(ctx context.Context, input FinishTaskFailureInput) (GatewayTask, error) {
metricsJSON, _ := json.Marshal(emptyObjectIfNil(input.Metrics))
resultJSON, _ := json.Marshal(emptyObjectIfNil(input.Result))
metricsJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(input.Metrics)))
resultJSON, _ := json.Marshal(minimalTaskResult(nil))
message := truncateUTF8Bytes(input.Message, 2048)
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 = NULL,
error_code = NULLIF($3::text, ''),
error_message = NULLIF($2::text, ''),
request_id = NULLIF($4::text, ''),
@@ -1422,13 +1502,14 @@ func (s *Store) FinishTaskFailure(ctx context.Context, input FinishTaskFailureIn
heartbeat_at = NULL,
execution_token = NULL,
execution_lease_expires_at = NULL,
remote_task_payload = '{}'::jsonb,
finished_at = now(),
updated_at = now()
WHERE id = $1::uuid
AND status = 'running'
AND execution_token = $10::uuid`,
input.TaskID,
input.Message,
message,
input.Code,
input.RequestID,
string(metricsJSON),
@@ -1472,29 +1553,152 @@ func nullableTime(value time.Time) any {
return value
}
func truncateUTF8Bytes(value string, maxBytes int) string {
if maxBytes <= 0 || len(value) <= maxBytes {
return value
}
end := maxBytes
for end > 0 && !utf8.ValidString(value[:end]) {
end--
}
return value[:end]
}
func minimalTaskResult(input map[string]any) map[string]any {
result := make(map[string]any, len(input))
for key, value := range input {
switch key {
case "raw",
"raw_data",
"rawData",
"provider_response",
"providerResponse",
"submit",
"file_retrieve",
"fileRetrieve",
"upstream_task_id",
"remote_task_id":
continue
}
result[key] = value
}
return result
}
var allowedTaskEventTypes = map[string]struct{}{
"task.accepted": {},
"task.running": {},
"task.queued": {},
"task.completed": {},
"task.failed": {},
"task.cancelled": {},
"task.attempt.failed": {},
"task.billing.settled": {},
"task.billing.released": {},
"task.billing.review": {},
"task.policy.auto_disabled": {},
"task.policy.degraded": {},
"task.policy.failover_disabled": {},
"task.policy.failover_cooled_down": {},
"task.policy.priority_demoted": {},
}
func taskEventAllowed(eventType string) bool {
_, ok := allowedTaskEventTypes[strings.TrimSpace(eventType)]
return ok
}
func taskEventTerminal(eventType string) bool {
switch strings.TrimSpace(eventType) {
case "task.completed", "task.failed", "task.cancelled",
"task.billing.settled", "task.billing.released", "task.billing.review":
return true
default:
return false
}
}
func taskEventPlatformID(payload map[string]any) string {
platformID, _ := payload["platformId"].(string)
return strings.TrimSpace(platformID)
}
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))
eventType = strings.TrimSpace(eventType)
if !taskEventAllowed(eventType) {
return TaskEvent{SkippedReason: "unknown_type"}, nil
}
platformID := taskEventPlatformID(payload)
tx, err := s.pool.Begin(ctx)
if err != nil {
return TaskEvent{}, err
}
defer tx.Rollback(ctx)
if _, err := tx.Exec(ctx, `SELECT 1 FROM gateway_tasks WHERE id = $1::uuid FOR UPDATE`, taskID); err != nil {
return TaskEvent{}, err
}
var previousType string
var previousStatus string
var previousPlatformID string
var previousSimulated bool
err = tx.QueryRow(ctx, `
SELECT event_type, COALESCE(status, ''), COALESCE(platform_id::text, ''), simulated
FROM gateway_task_events
WHERE task_id = $1::uuid
ORDER BY seq DESC
LIMIT 1`, taskID).Scan(&previousType, &previousStatus, &previousPlatformID, &previousSimulated)
if err != nil && !errors.Is(err, pgx.ErrNoRows) {
return TaskEvent{}, err
}
if err == nil &&
previousType == eventType &&
previousStatus == strings.TrimSpace(status) &&
previousPlatformID == platformID &&
previousSimulated == simulated {
if err := tx.Commit(ctx); err != nil {
return TaskEvent{}, err
}
return TaskEvent{SkippedReason: "duplicate"}, nil
}
if !taskEventTerminal(eventType) {
var nonTerminalCount int
if err := tx.QueryRow(ctx, `
SELECT count(*)
FROM gateway_task_events
WHERE task_id = $1::uuid
AND event_type NOT IN (
'task.completed', 'task.failed', 'task.cancelled',
'task.billing.settled', 'task.billing.released', 'task.billing.review'
)`, taskID).Scan(&nonTerminalCount); err != nil {
return TaskEvent{}, err
}
if nonTerminalCount >= 16 {
if err := tx.Commit(ctx); err != nil {
return TaskEvent{}, err
}
return TaskEvent{SkippedReason: "budget_exceeded"}, nil
}
}
var event TaskEvent
var payloadBytes []byte
err := s.pool.QueryRow(ctx, `
err = tx.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
INSERT INTO gateway_task_events (task_id, seq, event_type, status, phase, progress, message, payload, simulated, platform_id)
SELECT $1::uuid, next_seq.seq, $2::text, NULLIF($3::text, ''), NULL, 0, NULL, '{}'::jsonb, $4,
NULLIF($5::text, '')::uuid
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`,
COALESCE(progress, 0)::float8, COALESCE(message, ''), payload, simulated, created_at,
COALESCE(platform_id::text, '')`,
taskID,
eventType,
status,
phase,
progress,
message,
string(payloadJSON),
simulated,
platformID,
).Scan(
&event.ID,
&event.TaskID,
@@ -1507,39 +1711,30 @@ RETURNING id::text, task_id::text, seq, event_type, COALESCE(status, ''), COALES
&payloadBytes,
&event.Simulated,
&event.CreatedAt,
&event.PlatformID,
)
if err != nil {
return TaskEvent{}, err
}
event.Payload = decodeObject(payloadBytes)
if err := tx.Commit(ctx); err != nil {
return TaskEvent{}, err
}
return event, nil
}
func (s *Store) QueueTaskCallback(ctx context.Context, event TaskEvent, callbackURL string) error {
if callbackURL == "" {
if callbackURL == "" || event.ID == "" {
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)
VALUES ($1::uuid, $2::uuid, $3, $4, '{}'::jsonb)
ON CONFLICT (task_id, seq, callback_url) DO NOTHING`,
event.TaskID,
event.ID,
event.Seq,
callbackURL,
string(payloadJSON),
)
return err
}
apps/api/internal/store/wallet.go
+2 -2
View File
@@ -161,7 +161,7 @@ func (s *Store) ReserveTaskBilling(ctx context.Context, task GatewayTask, user *
}
reservations := make([]WalletBillingReservation, 0, len(amounts))
pricingSnapshotJSON, _ := json.Marshal(pricingSnapshot)
pricingSnapshotJSON, _ := json.Marshal(sanitizeJSONForStorage(pricingSnapshot))
requestFingerprint := walletString(pricingSnapshot["requestFingerprint"])
err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
for currency, rawAmount := range amounts {
@@ -306,7 +306,7 @@ func (s *Store) reserveTaskBillingExact(ctx context.Context, task GatewayTask, g
if currency != "resource" {
return nil, fmt.Errorf("unsupported billing currency %q", currency)
}
pricingSnapshotJSON, _ := json.Marshal(pricingSnapshot)
pricingSnapshotJSON, _ := json.Marshal(sanitizeJSONForStorage(pricingSnapshot))
requestFingerprint := walletString(pricingSnapshot["requestFingerprint"])
var reservations []WalletBillingReservation
err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {