fix gateway loopback validation chains

This commit is contained in:
2026-05-11 08:48:02 +08:00
parent ff666b1ece
commit ca7e76e815
42 changed files with 1641 additions and 129 deletions
+58 -26
View File
@@ -59,7 +59,7 @@ func (s *Store) ListBaseModels(ctx context.Context) ([]BaseModel, error) {
rows, err := s.pool.Query(ctx, `
SELECT `+baseModelColumns+`
FROM base_model_catalog
ORDER BY provider_key ASC, model_type ASC, canonical_model_key ASC`)
ORDER BY provider_key ASC, canonical_model_key ASC`)
if err != nil {
return nil, err
}
@@ -84,7 +84,7 @@ func (s *Store) CreateBaseModel(ctx context.Context, input BaseModelInput) (Base
runtimePolicyOverride, _ := json.Marshal(emptyObjectIfNil(input.RuntimePolicyOverride))
metadata, _ := json.Marshal(emptyObjectIfNil(input.Metadata))
defaultSnapshot, _ := json.Marshal(emptyObjectIfNil(input.DefaultSnapshot))
modelType := primaryString(input.ModelType, "text_generate")
modelType, _ := json.Marshal(input.ModelType)
return scanBaseModel(s.pool.QueryRow(ctx, `
INSERT INTO base_model_catalog (
@@ -94,7 +94,7 @@ INSERT INTO base_model_catalog (
)
VALUES (
(SELECT id FROM model_catalog_providers WHERE provider_key = $1 OR provider_code = $1 LIMIT 1),
$1, $2, $3, $4, $5, $6, $7, $8,
$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
@@ -103,7 +103,7 @@ RETURNING `+baseModelColumns,
input.ProviderKey,
input.CanonicalModelKey,
input.ProviderModelName,
modelType,
string(modelType),
input.ModelAlias,
capabilities,
billingConfig,
@@ -127,7 +127,7 @@ func (s *Store) UpdateBaseModel(ctx context.Context, id string, input BaseModelI
runtimePolicyOverride, _ := json.Marshal(emptyObjectIfNil(input.RuntimePolicyOverride))
metadata, _ := json.Marshal(emptyObjectIfNil(input.Metadata))
defaultSnapshot, _ := json.Marshal(emptyObjectIfNil(input.DefaultSnapshot))
modelType := primaryString(input.ModelType, "text_generate")
modelType, _ := json.Marshal(input.ModelType)
return scanBaseModel(s.pool.QueryRow(ctx, `
UPDATE base_model_catalog
@@ -135,7 +135,7 @@ SET provider_id = (SELECT id FROM model_catalog_providers WHERE provider_key = $
provider_key = $2,
canonical_model_key = $3,
provider_model_name = $4,
model_type = $5,
model_type = $5::jsonb,
display_name = $6,
capabilities = $7,
base_billing_config = $8,
@@ -156,7 +156,7 @@ RETURNING `+baseModelColumns,
input.ProviderKey,
input.CanonicalModelKey,
input.ProviderModelName,
modelType,
string(modelType),
input.ModelAlias,
capabilities,
billingConfig,
@@ -195,7 +195,7 @@ SET provider_id = (SELECT id FROM model_catalog_providers WHERE provider_key = C
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::text, model_type),
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),
@@ -214,7 +214,7 @@ RETURNING `+baseModelColumns,
stringFromSnapshot(snapshot, "providerKey"),
stringFromSnapshot(snapshot, "canonicalModelKey"),
stringFromSnapshot(snapshot, "providerModelName"),
stringFromSnapshot(snapshot, "modelType"),
jsonStringListFromSnapshot(snapshot, "modelType"),
stringFromSnapshot(snapshot, "modelAlias", "displayName"),
jsonFromSnapshot(snapshot, "capabilities"),
jsonFromSnapshot(snapshot, "baseBillingConfig"),
@@ -242,9 +242,10 @@ SET provider_id = (
canonical_model_key = COALESCE(NULLIF(default_snapshot->>'canonicalModelKey', ''), canonical_model_key),
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'->>0
ELSE default_snapshot->>'modelType'
END, ''), model_type),
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),
capabilities = COALESCE(default_snapshot->'capabilities', capabilities),
base_billing_config = COALESCE(default_snapshot->'baseBillingConfig', base_billing_config),
@@ -292,7 +293,7 @@ func scanBaseModelRows(rows pgx.Rows) ([]BaseModel, error) {
func scanBaseModel(scanner baseModelScanner) (BaseModel, error) {
var item BaseModel
var modelType string
var modelType []byte
var modelAlias string
var capabilities []byte
var billingConfig []byte
@@ -330,7 +331,7 @@ func scanBaseModel(scanner baseModelScanner) (BaseModel, error) {
item.RuntimePolicyOverride = decodeObject(runtimePolicyOverride)
item.Metadata = decodeObject(metadata)
item.DefaultSnapshot = decodeObject(defaultSnapshot)
item.ModelType = baseModelTypes(item.Capabilities, item.Metadata, modelType)
item.ModelType = decodeStringArray(modelType)
item.ModelAlias = modelAlias
item.DisplayName = modelAlias
return item, nil
@@ -420,14 +421,22 @@ func jsonFromSnapshot(snapshot map[string]any, key string) any {
return string(raw)
}
func baseModelTypes(capabilities map[string]any, metadata map[string]any, fallback string) StringList {
values := make([]string, 0)
values = append(values, stringListFromAny(capabilities["originalTypes"])...)
values = append(values, stringListFromAny(metadata["originalTypes"])...)
if fallback != "" {
values = append(values, fallback)
func jsonStringListFromSnapshot(snapshot map[string]any, key string) any {
values := stringListFromAny(snapshot[key])
if len(values) == 0 {
if value, ok := snapshot[key].(string); ok {
values = []string{value}
}
}
return uniqueStringList(values)
normalized := uniqueStringList(values)
if len(normalized) == 0 {
return nil
}
raw, err := json.Marshal(normalized)
if err != nil {
return nil
}
return string(raw)
}
func stringListFromAny(value any) []string {
@@ -451,16 +460,39 @@ func uniqueStringList(values []string) StringList {
out := make([]string, 0, len(values))
seen := map[string]bool{}
for _, value := range values {
value = strings.TrimSpace(value)
if value == "" || seen[value] {
continue
for _, normalized := range modelTypeAliases(value) {
normalized = strings.TrimSpace(normalized)
if normalized == "" || seen[normalized] {
continue
}
seen[normalized] = true
out = append(out, normalized)
}
seen[value] = true
out = append(out, value)
}
return out
}
func normalizeModelTypeList(values []string) StringList {
return uniqueStringList(values)
}
func modelTypeAliases(value string) []string {
switch strings.TrimSpace(value) {
case "chat", "text", "responses":
return []string{"text_generate"}
case "image":
return []string{"image_generate", "image_edit"}
case "images.generations":
return []string{"image_generate"}
case "images.edits":
return []string{"image_edit"}
case "video", "videos.generations":
return []string{"video_generate"}
default:
return []string{value}
}
}
func primaryString(values []string, fallback string) string {
for _, value := range values {
if value = strings.TrimSpace(value); value != "" {
+22 -11
View File
@@ -18,27 +18,25 @@ SELECT p.id::text, p.platform_key, p.name, p.provider,
COALESCE(p.dynamic_priority, p.priority) AS effective_priority,
m.id::text, COALESCE(m.base_model_id::text, ''), COALESCE(b.canonical_model_key, ''),
COALESCE(b.provider_model_name, ''), m.model_name, COALESCE(m.model_alias, ''),
m.model_type, m.display_name, m.capabilities, m.capability_override,
$2 AS requested_model_type, m.display_name, m.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, ''),
m.permission_config, m.retry_policy, m.rate_limit_policy, COALESCE(m.runtime_policy_set_id::text, COALESCE(b.runtime_policy_set_id::text, '')),
COALESCE(NULLIF(m.runtime_policy_override, '{}'::jsonb), b.runtime_policy_override, '{}'::jsonb)
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, '')),
COALESCE(NULLIF(m.runtime_policy_override, '{}'::jsonb), b.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)
FROM platform_models m
JOIN integration_platforms p ON p.id = m.platform_id
LEFT JOIN model_catalog_providers cp ON cp.provider_key = p.provider OR cp.provider_code = p.provider
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)
LEFT JOIN runtime_client_states s
ON s.client_id = p.platform_key || ':' || m.model_type || ':' || m.model_name
ON s.client_id = p.platform_key || ':' || $2 || ':' || m.model_name
WHERE p.status = 'enabled'
AND p.deleted_at IS NULL
AND m.enabled = true
AND (
m.model_type = $2
OR ($2 = 'text_generate' AND m.model_type IN ('chat', 'responses', 'text'))
OR ($2 = 'image_generate' AND m.model_type IN ('image', 'images.generations'))
OR ($2 = 'image_edit' AND m.model_type IN ('images.edits'))
OR ($2 = 'video_generate' AND m.model_type IN ('video', 'videos.generations', 'video_generate', 'text_to_video', 'image_to_video', 'omni_video', 'video_edit', 'video_reference', 'video_first_last_frame'))
)
AND m.model_type @> jsonb_build_array($2)
AND (p.cooldown_until IS NULL OR p.cooldown_until <= now())
AND (
m.model_name = $1
@@ -73,6 +71,10 @@ ORDER BY effective_priority ASC,
var modelRetryPolicy []byte
var modelRateLimitPolicy []byte
var runtimePolicyOverride []byte
var runtimeRetryPolicy []byte
var runtimeRateLimitPolicy []byte
var autoDisablePolicy []byte
var degradePolicy []byte
if err := rows.Scan(
&item.PlatformID,
&item.PlatformKey,
@@ -105,11 +107,16 @@ ORDER BY effective_priority ASC,
&item.PricingMode,
&item.DiscountFactor,
&item.ModelPricingRuleSetID,
&item.BasePricingRuleSetID,
&permissionConfig,
&modelRetryPolicy,
&modelRateLimitPolicy,
&item.RuntimePolicySetID,
&runtimePolicyOverride,
&runtimeRetryPolicy,
&runtimeRateLimitPolicy,
&autoDisablePolicy,
&degradePolicy,
); err != nil {
return nil, err
}
@@ -126,6 +133,10 @@ ORDER BY effective_priority ASC,
item.ModelRetryPolicy = decodeObject(modelRetryPolicy)
item.ModelRateLimitPolicy = decodeObject(modelRateLimitPolicy)
item.RuntimePolicyOverride = decodeObject(runtimePolicyOverride)
item.RuntimeRetryPolicy = decodeObject(runtimeRetryPolicy)
item.RuntimeRateLimitPolicy = decodeObject(runtimeRateLimitPolicy)
item.AutoDisablePolicy = decodeObject(autoDisablePolicy)
item.DegradePolicy = decodeObject(degradePolicy)
item.ClientID = fmt.Sprintf("%s:%s:%s", item.PlatformKey, item.ModelType, item.ModelName)
item.QueueKey = item.ClientID
items = append(items, item)
+19 -9
View File
@@ -17,7 +17,7 @@ type modelCatalogSnapshot struct {
ProviderKey string
CanonicalModelKey string
ProviderModelName string
ModelType string
ModelType StringList
DisplayName string
Capabilities map[string]any
BaseBillingConfig map[string]any
@@ -83,9 +83,13 @@ func (s *Store) createPlatformModel(ctx context.Context, q platformModelQuerier,
if err != nil && !IsNotFound(err) {
return PlatformModel{}, err
}
if input.ModelType == "" {
if len(input.ModelType) == 0 {
input.ModelType = base.ModelType
}
input.ModelType = normalizeModelTypeList(input.ModelType)
if len(input.ModelType) == 0 {
input.ModelType = StringList{"text_generate"}
}
if input.ModelName == "" {
input.ModelName = base.ProviderModelName
}
@@ -104,11 +108,12 @@ func (s *Store) createPlatformModel(ctx context.Context, q platformModelQuerier,
if len(billingConfig) == 0 {
billingConfig = mergeObjects(base.BaseBillingConfig, input.BillingConfigOverride)
}
explicitRuntimePolicySetID := strings.TrimSpace(input.RuntimePolicySetID)
rateLimitPolicy := input.RateLimitPolicy
if len(rateLimitPolicy) == 0 {
if len(rateLimitPolicy) == 0 && explicitRuntimePolicySetID == "" {
rateLimitPolicy = base.DefaultRateLimitPolicy
}
runtimePolicySetID := strings.TrimSpace(input.RuntimePolicySetID)
runtimePolicySetID := explicitRuntimePolicySetID
if runtimePolicySetID == "" {
runtimePolicySetID = base.RuntimePolicySetID
}
@@ -119,6 +124,7 @@ func (s *Store) createPlatformModel(ctx context.Context, q platformModelQuerier,
capabilityOverrideJSON, _ := json.Marshal(emptyObjectIfNil(input.CapabilityOverride))
capabilitiesJSON, _ := json.Marshal(emptyObjectIfNil(capabilities))
modelTypeJSON, _ := json.Marshal(input.ModelType)
billingOverrideJSON, _ := json.Marshal(emptyObjectIfNil(input.BillingConfigOverride))
billingJSON, _ := json.Marshal(emptyObjectIfNil(billingConfig))
permissionJSON, _ := json.Marshal(emptyObjectIfNil(input.PermissionConfig))
@@ -144,6 +150,7 @@ func (s *Store) createPlatformModel(ctx context.Context, q platformModelQuerier,
var retryPolicyBytes []byte
var rateLimitPolicyBytes []byte
var runtimePolicyOverrideBytes []byte
var modelTypeBytes []byte
err = q.QueryRow(ctx, `
INSERT INTO platform_models (
platform_id, base_model_id, model_name, model_alias, model_type, display_name,
@@ -152,12 +159,12 @@ INSERT INTO platform_models (
runtime_policy_set_id, runtime_policy_override, enabled
)
VALUES (
$1::uuid, $2::uuid, $3, NULLIF($4, ''), $5, $6,
$1::uuid, $2::uuid, $3, NULLIF($4, ''), $5::jsonb, $6,
$7::jsonb, $8::jsonb, $9, $10::numeric,
NULLIF($11, '')::uuid, $12::jsonb, $13::jsonb, $14::jsonb, $15::jsonb, $16::jsonb,
NULLIF($17, '')::uuid, $18::jsonb, true
)
ON CONFLICT (platform_id, model_name, model_type) DO UPDATE
ON CONFLICT (platform_id, model_name) DO UPDATE
SET base_model_id = EXCLUDED.base_model_id,
model_alias = EXCLUDED.model_alias,
display_name = EXCLUDED.display_name,
@@ -185,7 +192,7 @@ RETURNING id::text, platform_id::text, COALESCE(base_model_id::text, ''), model_
baseID,
input.ModelName,
input.ModelAlias,
input.ModelType,
string(modelTypeJSON),
input.DisplayName,
string(capabilityOverrideJSON),
string(capabilitiesJSON),
@@ -205,7 +212,7 @@ RETURNING id::text, platform_id::text, COALESCE(base_model_id::text, ''), model_
&model.BaseModelID,
&model.ModelName,
&model.ModelAlias,
&model.ModelType,
&modelTypeBytes,
&model.DisplayName,
&capabilityOverrideBytes,
&capabilitiesBytes,
@@ -228,6 +235,7 @@ RETURNING id::text, platform_id::text, COALESCE(base_model_id::text, ''), model_
}
model.CapabilityOverride = decodeObject(capabilityOverrideBytes)
model.Capabilities = decodeObject(capabilitiesBytes)
model.ModelType = decodeStringArray(modelTypeBytes)
model.BillingConfigOverride = decodeObject(billingOverrideBytes)
model.BillingConfig = decodeObject(billingBytes)
model.PermissionConfig = decodeObject(permissionBytes)
@@ -265,6 +273,7 @@ func (s *Store) lookupBaseModel(ctx context.Context, q platformModelQuerier, id
var billingConfig []byte
var rateLimitPolicy []byte
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,
capabilities, base_billing_config, default_rate_limit_policy,
@@ -279,7 +288,7 @@ LIMIT 1`, strings.TrimSpace(id), strings.TrimSpace(canonicalKey), strings.TrimSp
&item.ProviderKey,
&item.CanonicalModelKey,
&item.ProviderModelName,
&item.ModelType,
&modelTypeBytes,
&item.DisplayName,
&capabilities,
&billingConfig,
@@ -297,6 +306,7 @@ LIMIT 1`, strings.TrimSpace(id), strings.TrimSpace(canonicalKey), strings.TrimSp
item.BaseBillingConfig = decodeObject(billingConfig)
item.DefaultRateLimitPolicy = decodeObject(rateLimitPolicy)
item.RuntimePolicyOverride = decodeObject(runtimePolicyOverride)
item.ModelType = normalizeModelTypeList(decodeStringArray(modelTypeBytes))
return item, nil
}
+13 -4
View File
@@ -134,7 +134,7 @@ type PlatformModel struct {
PlatformName string `json:"platformName,omitempty"`
ModelName string `json:"modelName"`
ModelAlias string `json:"modelAlias,omitempty"`
ModelType string `json:"modelType"`
ModelType StringList `json:"modelType"`
DisplayName string `json:"displayName"`
CapabilityOverride map[string]any `json:"capabilityOverride,omitempty"`
Capabilities map[string]any `json:"capabilities,omitempty"`
@@ -390,6 +390,8 @@ type GatewayTask struct {
ResponseDurationMS int64 `json:"responseDurationMs"`
FinishedAt string `json:"finishedAt,omitempty"`
Error string `json:"error,omitempty"`
ErrorCode string `json:"errorCode,omitempty"`
ErrorMessage string `json:"errorMessage,omitempty"`
CreatedAt time.Time `json:"createdAt"`
UpdatedAt time.Time `json:"updatedAt"`
}
@@ -404,6 +406,7 @@ request, status, COALESCE(result, '{}'::jsonb), COALESCE(billings, '[]'::jsonb),
COALESCE(usage, '{}'::jsonb), COALESCE(metrics, '{}'::jsonb), COALESCE(billing_summary, '{}'::jsonb),
COALESCE(final_charge_amount, 0)::float8, COALESCE(response_started_at::text, ''),
COALESCE(response_finished_at::text, ''), COALESCE(response_duration_ms, 0), COALESCE(error, ''),
COALESCE(error_code, ''), COALESCE(error_message, ''),
created_at, updated_at, COALESCE(finished_at::text, '')`
type TaskEvent struct {
@@ -713,6 +716,7 @@ ORDER BY m.model_type ASC, m.model_name ASC`, args...)
var retryPolicy []byte
var rateLimitPolicy []byte
var runtimePolicyOverride []byte
var modelTypeBytes []byte
if err := rows.Scan(
&model.ID,
&model.PlatformID,
@@ -721,7 +725,7 @@ ORDER BY m.model_type ASC, m.model_name ASC`, args...)
&model.PlatformName,
&model.ModelName,
&model.ModelAlias,
&model.ModelType,
&modelTypeBytes,
&model.DisplayName,
&capabilityOverride,
&capabilities,
@@ -743,6 +747,7 @@ ORDER BY m.model_type ASC, m.model_name ASC`, args...)
}
model.CapabilityOverride = decodeObject(capabilityOverride)
model.Capabilities = decodeObject(capabilities)
model.ModelType = decodeStringArray(modelTypeBytes)
model.BillingConfigOverride = decodeObject(billingConfigOverride)
model.BillingConfig = decodeObject(billingConfig)
model.PermissionConfig = decodeObject(permissionConfig)
@@ -1231,7 +1236,7 @@ func (s *Store) VerifyLocalAPIKey(ctx context.Context, secret string) (*auth.Use
return nil, auth.ErrUnauthorized
}
rows, err := s.pool.Query(ctx, `
SELECT k.id::text, k.key_hash, k.key_prefix, k.name, COALESCE(k.user_group_id::text, u.default_user_group_id::text, ''),
SELECT k.id::text, k.key_hash, k.key_prefix, k.name, k.scopes, COALESCE(k.user_group_id::text, u.default_user_group_id::text, ''),
u.id::text, u.username, u.roles, COALESCE(u.gateway_tenant_id::text, ''),
COALESCE(u.tenant_id, ''), COALESCE(u.tenant_key, '')
FROM gateway_api_keys k
@@ -1252,6 +1257,7 @@ WHERE k.key_prefix = $1
var hash string
var keyPrefix string
var keyName string
var scopesBytes []byte
var userGroupID string
var gatewayUserID string
var username string
@@ -1259,7 +1265,7 @@ WHERE k.key_prefix = $1
var gatewayTenantID string
var tenantID string
var tenantKey string
if err := rows.Scan(&apiKeyID, &hash, &keyPrefix, &keyName, &userGroupID, &gatewayUserID, &username, &rolesBytes, &gatewayTenantID, &tenantID, &tenantKey); err != nil {
if err := rows.Scan(&apiKeyID, &hash, &keyPrefix, &keyName, &scopesBytes, &userGroupID, &gatewayUserID, &username, &rolesBytes, &gatewayTenantID, &tenantID, &tenantKey); err != nil {
return nil, err
}
if bcrypt.CompareHashAndPassword([]byte(hash), []byte(secret)) != nil {
@@ -1281,6 +1287,7 @@ WHERE k.key_prefix = $1
APIKeyID: apiKeyID,
APIKeyName: keyName,
APIKeyPrefix: keyPrefix,
APIKeyScopes: decodeStringArray(scopesBytes),
}, nil
}
if err := rows.Err(); err != nil {
@@ -1661,6 +1668,8 @@ func scanGatewayTask(scanner taskScanner) (GatewayTask, error) {
&task.ResponseFinishedAt,
&task.ResponseDurationMS,
&task.Error,
&task.ErrorCode,
&task.ErrorMessage,
&task.CreatedAt,
&task.UpdatedAt,
&task.FinishedAt,
+57
View File
@@ -176,6 +176,63 @@ func (s *Store) DeletePricingRuleSet(ctx context.Context, id string) error {
return nil
}
func (s *Store) PricingRuleSetBillingConfig(ctx context.Context, id string) (map[string]any, error) {
id = strings.TrimSpace(id)
if id == "" {
return nil, nil
}
rows, err := s.pool.Query(ctx, `
SELECT resource_type, base_price::float8, dynamic_weight
FROM model_pricing_rules
WHERE rule_set_id = $1::uuid
AND status = 'active'
ORDER BY priority ASC, resource_type ASC`, id)
if err != nil {
return nil, err
}
defer rows.Close()
config := map[string]any{}
for rows.Next() {
var resourceType string
var basePrice float64
var dynamicWeightBytes []byte
if err := rows.Scan(&resourceType, &basePrice, &dynamicWeightBytes); err != nil {
return nil, err
}
dynamicWeight := decodeObject(dynamicWeightBytes)
switch resourceType {
case "text_input":
config["textInputPer1k"] = basePrice
case "text_output":
config["textOutputPer1k"] = basePrice
case "image":
config["imageBase"] = basePrice
config["image"] = pricingResourceConfig(basePrice, dynamicWeight)
case "image_edit":
config["editBase"] = basePrice
config["image_edit"] = pricingResourceConfig(basePrice, dynamicWeight)
case "video":
config["videoBase"] = basePrice
config["video"] = pricingResourceConfig(basePrice, dynamicWeight)
default:
config[resourceType] = pricingResourceConfig(basePrice, dynamicWeight)
}
}
if err := rows.Err(); err != nil {
return nil, err
}
return config, nil
}
func pricingResourceConfig(basePrice float64, dynamicWeight map[string]any) map[string]any {
config := map[string]any{"basePrice": basePrice}
if len(dynamicWeight) > 0 {
config["dynamicWeight"] = dynamicWeight
}
return config
}
func insertPricingRules(ctx context.Context, tx pgx.Tx, ruleSetID string, defaultCurrency string, rules []PricingRuleInput) error {
for index, rule := range rules {
rule = normalizePricingRule(rule, index, defaultCurrency)
@@ -107,6 +107,33 @@ func (s *Store) DeleteRuntimePolicySet(ctx context.Context, id string) error {
return nil
}
func (s *Store) DisableCandidatePlatform(ctx context.Context, platformID string) error {
if strings.TrimSpace(platformID) == "" {
return nil
}
_, err := s.pool.Exec(ctx, `
UPDATE integration_platforms
SET status = 'disabled',
updated_at = now()
WHERE id = $1::uuid`, platformID)
return err
}
func (s *Store) CooldownCandidatePlatform(ctx context.Context, platformID string, cooldownSeconds int) error {
if strings.TrimSpace(platformID) == "" {
return nil
}
if cooldownSeconds <= 0 {
cooldownSeconds = 300
}
_, err := s.pool.Exec(ctx, `
UPDATE integration_platforms
SET cooldown_until = now() + ($2::int * interval '1 second'),
updated_at = now()
WHERE id = $1::uuid`, platformID, cooldownSeconds)
return err
}
func scanRuntimePolicySet(scanner runtimePolicyScanner) (RuntimePolicySet, error) {
var item RuntimePolicySet
var rateLimitPolicy []byte
+6 -1
View File
@@ -16,7 +16,7 @@ type CreatePlatformModelInput struct {
CanonicalModelKey string `json:"canonicalModelKey"`
ModelName string `json:"modelName"`
ModelAlias string `json:"modelAlias"`
ModelType string `json:"modelType"`
ModelType StringList `json:"modelType"`
DisplayName string `json:"displayName"`
CapabilityOverride map[string]any `json:"capabilityOverride"`
Capabilities map[string]any `json:"capabilities"`
@@ -64,12 +64,17 @@ type RuntimeModelCandidate struct {
PermissionConfig map[string]any
PricingMode string
DiscountFactor float64
BasePricingRuleSetID string
PlatformPricingRuleSetID string
ModelPricingRuleSetID string
ModelRetryPolicy map[string]any
ModelRateLimitPolicy map[string]any
RuntimePolicySetID string
RuntimePolicyOverride map[string]any
RuntimeRetryPolicy map[string]any
RuntimeRateLimitPolicy map[string]any
AutoDisablePolicy map[string]any
DegradePolicy map[string]any
ClientID string
QueueKey string
}
+155
View File
@@ -3,9 +3,67 @@ package store
import (
"context"
"encoding/json"
"strings"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgconn"
)
func (s *Store) ListTasks(ctx context.Context, user *auth.User, limit int) ([]GatewayTask, error) {
if limit <= 0 {
limit = 50
}
if limit > 100 {
limit = 100
}
gatewayUserID := localGatewayUserID(user)
apiKeyID := ""
userID := ""
if user != nil {
apiKeyID = strings.TrimSpace(user.APIKeyID)
userID = strings.TrimSpace(user.ID)
}
if gatewayUserID == "" && userID == "" {
return nil, ErrLocalUserRequired
}
rows, err := s.pool.Query(ctx, `
SELECT `+gatewayTaskColumns+`
FROM gateway_tasks
WHERE (
(
NULLIF($1, '')::uuid IS NOT NULL
AND gateway_user_id = NULLIF($1, '')::uuid
)
OR (
NULLIF($1, '')::uuid IS NULL
AND NULLIF($2, '') IS NOT NULL
AND user_id = $2
)
)
AND (
NULLIF($3, '') IS NULL
OR api_key_id = $3
)
ORDER BY created_at DESC
LIMIT $4`, gatewayUserID, userID, apiKeyID, limit)
if err != nil {
return nil, err
}
defer rows.Close()
items := make([]GatewayTask, 0)
for rows.Next() {
task, err := scanGatewayTask(rows)
if err != nil {
return nil, err
}
items = append(items, task)
}
return items, rows.Err()
}
func (s *Store) MarkTaskRunning(ctx context.Context, taskID string, modelType string, normalizedRequest map[string]any) error {
normalizedJSON, _ := json.Marshal(emptyObjectIfNil(normalizedRequest))
_, err := s.pool.Exec(ctx, `
@@ -140,6 +198,103 @@ WHERE id = $1::uuid`,
return s.GetTask(ctx, input.TaskID)
}
func (s *Store) SettleTaskBilling(ctx context.Context, task GatewayTask) error {
if task.FinalChargeAmount <= 0 || strings.TrimSpace(task.GatewayUserID) == "" {
return nil
}
currency := strings.TrimSpace(taskBillingString(task.BillingSummary["currency"]))
if currency == "" || currency == "mixed" {
currency = "resource"
}
metadata, _ := json.Marshal(map[string]any{
"taskId": task.ID,
"kind": task.Kind,
"model": task.Model,
"resolvedModel": task.ResolvedModel,
"billings": task.Billings,
"billingSummary": task.BillingSummary,
})
return pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
if _, err := tx.Exec(ctx, `
INSERT INTO gateway_wallet_accounts (
gateway_tenant_id, gateway_user_id, tenant_id, tenant_key, user_id, currency
)
VALUES (NULLIF($1, '')::uuid, $2::uuid, NULLIF($3, ''), NULLIF($4, ''), NULLIF($5, ''), $6)
ON CONFLICT (gateway_user_id, currency) DO NOTHING`,
task.GatewayTenantID, task.GatewayUserID, task.TenantID, task.TenantKey, task.UserID, currency); err != nil {
return err
}
var exists bool
if err := tx.QueryRow(ctx, `
SELECT EXISTS (
SELECT 1
FROM gateway_wallet_transactions t
JOIN gateway_wallet_accounts a ON a.id = t.account_id
WHERE a.gateway_user_id = $1::uuid
AND a.currency = $2
AND t.idempotency_key = $3
)`, task.GatewayUserID, currency, billingIdempotencyKey(task.ID)).Scan(&exists); err != nil {
return err
}
if exists {
return nil
}
var accountID string
var balanceBefore float64
var gatewayTenantID string
if err := tx.QueryRow(ctx, `
SELECT id::text, balance::float8, COALESCE(gateway_tenant_id::text, '')
FROM gateway_wallet_accounts
WHERE gateway_user_id = $1::uuid
AND currency = $2
FOR UPDATE`, task.GatewayUserID, currency).Scan(&accountID, &balanceBefore, &gatewayTenantID); err != nil {
return err
}
amount := roundMoney(task.FinalChargeAmount)
balanceAfter := roundMoney(balanceBefore - amount)
if _, err := tx.Exec(ctx, `
UPDATE gateway_wallet_accounts
SET balance = $2,
total_spent = total_spent + $3,
updated_at = now()
WHERE id = $1::uuid`, accountID, balanceAfter, amount); err != nil {
return err
}
_, err := tx.Exec(ctx, `
INSERT INTO gateway_wallet_transactions (
account_id, gateway_tenant_id, gateway_user_id, direction, transaction_type,
amount, balance_before, balance_after, idempotency_key, reference_type, reference_id, metadata
)
VALUES (
$1::uuid, NULLIF($2, '')::uuid, $3::uuid, 'debit', 'task_billing',
$4, $5, $6, $7, 'gateway_task', $8, $9::jsonb
)`,
accountID, firstNonEmpty(gatewayTenantID, task.GatewayTenantID), task.GatewayUserID, amount, roundMoney(balanceBefore), balanceAfter, billingIdempotencyKey(task.ID), task.ID, string(metadata))
if pgErr, ok := err.(*pgconn.PgError); ok && pgErr.Code == "23505" {
return nil
}
return err
})
}
func billingIdempotencyKey(taskID string) string {
return "task:" + taskID + ":billing"
}
func roundMoney(value float64) float64 {
if value < 0 {
return -roundMoney(-value)
}
return float64(int64(value*1000000+0.5)) / 1000000
}
func taskBillingString(value any) string {
if text, ok := value.(string); ok {
return text
}
return ""
}
func (s *Store) FinishTaskFailure(ctx context.Context, input FinishTaskFailureInput) (GatewayTask, error) {
metricsJSON, _ := json.Marshal(emptyObjectIfNil(input.Metrics))
if _, err := s.pool.Exec(ctx, `