feat: add runner failover rules and cache affinity

This commit is contained in:
wangbo 2026-06-28 20:50:23 +08:00
parent 41bb3a6525
commit 0e675e259c
31 changed files with 2939 additions and 470 deletions

View File

@ -302,7 +302,7 @@ func TestOpenAIClientChatContract(t *testing.T) {
"choices": []any{map[string]any{
"message": map[string]any{"role": "assistant", "content": "ok"},
}},
"usage": map[string]any{"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5},
"usage": map[string]any{"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5, "prompt_cache_hit_tokens": 2},
})
}))
defer server.Close()
@ -327,6 +327,14 @@ func TestOpenAIClientChatContract(t *testing.T) {
if response.Usage.TotalTokens != 5 || response.Result["id"] != "chatcmpl-test" {
t.Fatalf("unexpected response: %+v", response)
}
if response.Usage.CachedInputTokens != 2 {
t.Fatalf("expected cached input token usage, got %+v", response.Usage)
}
resultUsage, _ := response.Result["usage"].(map[string]any)
promptDetails, _ := resultUsage["prompt_tokens_details"].(map[string]any)
if intFromAny(promptDetails["cached_tokens"]) != 2 {
t.Fatalf("expected normalized cached prompt tokens in result usage, got %+v", resultUsage)
}
if response.RequestID != "req-chat-test" || response.ResponseStartedAt.IsZero() || response.ResponseFinishedAt.IsZero() {
t.Fatalf("response metadata was not captured: %+v", response)
}
@ -335,6 +343,29 @@ func TestOpenAIClientChatContract(t *testing.T) {
}
}
func TestUsageFromOpenAIUsageTracksKnownCachedInputZero(t *testing.T) {
knownZero := usageFromOpenAIUsage(map[string]any{
"prompt_tokens": 100,
"completion_tokens": 10,
"total_tokens": 110,
"prompt_tokens_details": map[string]any{
"cached_tokens": 0,
},
})
if !knownZero.CachedInputTokensKnown || knownZero.CachedInputTokens != 0 {
t.Fatalf("expected known zero cached input tokens, got %+v", knownZero)
}
absent := usageFromOpenAIUsage(map[string]any{
"prompt_tokens": 100,
"completion_tokens": 10,
"total_tokens": 110,
})
if absent.CachedInputTokensKnown {
t.Fatalf("expected missing cached field to remain unknown, got %+v", absent)
}
}
func TestOpenAIClientEmbeddingsContract(t *testing.T) {
var gotPath string
var gotModel string

View File

@ -680,11 +680,11 @@ func geminiUsage(raw map[string]any) Usage {
input := intFromAny(usageMap["prompt_tokens"])
output := intFromAny(usageMap["completion_tokens"])
total := intFromAny(usageMap["total_tokens"])
cachedInput := cachedInputTokensFromOpenAIUsage(usageMap)
cachedInput, cachedInputKnown := cachedInputTokensValueFromOpenAIUsage(usageMap)
if cachedInput > input && input > 0 {
cachedInput = input
}
return Usage{InputTokens: input, OutputTokens: output, CachedInputTokens: cachedInput, TotalTokens: total}
return Usage{InputTokens: input, OutputTokens: output, CachedInputTokens: cachedInput, CachedInputTokensKnown: cachedInputKnown, TotalTokens: total}
}
func geminiUsageMap(raw map[string]any) map[string]any {

View File

@ -259,6 +259,9 @@ func NormalizeChatCompletionResult(result map[string]any) map[string]any {
if result == nil {
return nil
}
if usage, ok := result["usage"].(map[string]any); ok && len(usage) > 0 {
result["usage"] = NormalizeChatCompletionUsage(usage)
}
choices, _ := result["choices"].([]any)
for _, rawChoice := range choices {
choice, _ := rawChoice.(map[string]any)
@ -280,6 +283,9 @@ func NormalizeChatCompletionStreamEvent(event map[string]any) map[string]any {
if event == nil {
return nil
}
if usage, ok := event["usage"].(map[string]any); ok && len(usage) > 0 {
event["usage"] = NormalizeChatCompletionUsage(usage)
}
choices, _ := event["choices"].([]any)
for _, rawChoice := range choices {
choice, _ := rawChoice.(map[string]any)
@ -295,6 +301,33 @@ func NormalizeChatCompletionStreamEvent(event map[string]any) map[string]any {
return event
}
func NormalizeChatCompletionUsage(usage map[string]any) map[string]any {
if len(usage) == 0 {
return usage
}
normalized := usageFromOpenAIUsage(usage)
out := cloneMapAny(usage)
if normalized.InputTokens > 0 && intFromAny(out["prompt_tokens"]) == 0 {
out["prompt_tokens"] = normalized.InputTokens
}
if normalized.OutputTokens > 0 && intFromAny(out["completion_tokens"]) == 0 {
out["completion_tokens"] = normalized.OutputTokens
}
if normalized.TotalTokens > 0 && intFromAny(out["total_tokens"]) == 0 {
out["total_tokens"] = normalized.TotalTokens
}
if normalized.CachedInputTokens > 0 {
promptDetails, _ := firstPresent(out["prompt_tokens_details"], out["promptTokensDetails"]).(map[string]any)
normalizedPromptDetails := cloneMapAny(promptDetails)
if normalizedPromptDetails == nil {
normalizedPromptDetails = map[string]any{}
}
normalizedPromptDetails["cached_tokens"] = normalized.CachedInputTokens
out["prompt_tokens_details"] = normalizedPromptDetails
}
return out
}
func normalizeToolCallsContainer(container map[string]any, stream bool) {
if container == nil {
return
@ -877,27 +910,36 @@ func sortedStreamToolCalls(toolCalls map[int]map[string]any) []any {
func usageFromOpenAI(result map[string]any) Usage {
usage, _ := result["usage"].(map[string]any)
return usageFromOpenAIUsage(usage)
}
func usageFromOpenAIUsage(usage map[string]any) Usage {
input := intFromAny(firstPresent(usage["prompt_tokens"], usage["input_tokens"]))
output := intFromAny(firstPresent(usage["completion_tokens"], usage["output_tokens"]))
total := intFromAny(usage["total_tokens"])
if total == 0 {
total = input + output
}
cachedInput := cachedInputTokensFromOpenAIUsage(usage)
cachedInput, cachedInputKnown := cachedInputTokensValueFromOpenAIUsage(usage)
if cachedInput > input && input > 0 {
cachedInput = input
}
return Usage{InputTokens: input, OutputTokens: output, CachedInputTokens: cachedInput, TotalTokens: total}
return Usage{InputTokens: input, OutputTokens: output, CachedInputTokens: cachedInput, CachedInputTokensKnown: cachedInputKnown, TotalTokens: total}
}
func cachedInputTokensFromOpenAIUsage(usage map[string]any) int {
cachedInput, _ := cachedInputTokensValueFromOpenAIUsage(usage)
return cachedInput
}
func cachedInputTokensValueFromOpenAIUsage(usage map[string]any) (int, bool) {
if len(usage) == 0 {
return 0
return 0, false
}
promptDetails, _ := firstPresent(usage["prompt_tokens_details"], usage["promptTokensDetails"]).(map[string]any)
inputDetails, _ := firstPresent(usage["input_tokens_details"], usage["inputTokensDetails"]).(map[string]any)
usageMetadata, _ := usage["usageMetadata"].(map[string]any)
return intFromAny(firstPresent(
value, ok := firstPresentValue(
promptDetails["cached_tokens"],
promptDetails["cachedTokens"],
promptDetails["cache_read_input_tokens"],
@ -918,7 +960,11 @@ func cachedInputTokensFromOpenAIUsage(usage map[string]any) int {
usage["cachedContentTokenCount"],
usageMetadata["cached_content_token_count"],
usageMetadata["cachedContentTokenCount"],
))
)
if !ok {
return 0, false
}
return intFromAny(value), true
}
func requestIDFromHTTPResponse(resp *http.Response) string {
@ -979,12 +1025,17 @@ func firstNonEmptyString(values ...any) string {
}
func firstPresent(values ...any) any {
value, _ := firstPresentValue(values...)
return value
}
func firstPresentValue(values ...any) (any, bool) {
for _, value := range values {
if value != nil {
return value
return value, true
}
}
return nil
return nil, false
}
func errorMessage(raw []byte, fallback string) string {

View File

@ -107,6 +107,7 @@ func simulationProfile(request Request) string {
}
func simulatedResult(request Request) map[string]any {
usage := simulatedUsage(request)
switch request.Kind {
case "chat.completions":
return map[string]any{
@ -122,7 +123,7 @@ func simulatedResult(request Request) map[string]any {
"content": fmt.Sprintf("simulation response from %s", request.Candidate.Provider),
},
}},
"usage": map[string]any{"prompt_tokens": 12, "completion_tokens": 8, "total_tokens": 20},
"usage": simulatedOpenAIUsageMap(usage),
}
case "responses":
return map[string]any{
@ -131,7 +132,7 @@ func simulatedResult(request Request) map[string]any {
"created_at": nowUnix(),
"model": request.Model,
"output_text": fmt.Sprintf("simulation response from %s", request.Candidate.Provider),
"usage": map[string]any{"input_tokens": 12, "output_tokens": 8, "total_tokens": 20},
"usage": simulatedResponseUsageMap(usage),
}
case "embeddings":
return simulatedEmbeddingResult(request)
@ -373,6 +374,9 @@ func simulatedAudioData(request Request, fallbackPrompt string) []any {
}
func simulatedUsage(request Request) Usage {
if usage, ok := simulationUsageOverride(request); ok {
return usage
}
if request.ModelType == "chat" || request.ModelType == "text_generate" || request.Kind == "responses" {
return Usage{InputTokens: 12, OutputTokens: 8, TotalTokens: 20}
}
@ -382,6 +386,56 @@ func simulatedUsage(request Request) Usage {
return Usage{}
}
func simulationUsageOverride(request Request) (Usage, bool) {
raw, _ := request.Body["simulationUsage"].(map[string]any)
if len(raw) == 0 {
raw, _ = request.Body["testUsage"].(map[string]any)
}
if len(raw) == 0 {
return Usage{}, false
}
usage := Usage{
InputTokens: intFromAny(firstPresent(raw["inputTokens"], raw["input_tokens"], raw["promptTokens"], raw["prompt_tokens"])),
OutputTokens: intFromAny(firstPresent(raw["outputTokens"], raw["output_tokens"], raw["completionTokens"], raw["completion_tokens"])),
CachedInputTokens: intFromAny(firstPresent(raw["cachedInputTokens"], raw["cached_input_tokens"], raw["cachedPromptTokens"], raw["cached_tokens"])),
TotalTokens: intFromAny(firstPresent(raw["totalTokens"], raw["total_tokens"])),
}
if _, ok := firstPresentValue(raw["cachedInputTokens"], raw["cached_input_tokens"], raw["cachedPromptTokens"], raw["cached_tokens"]); ok {
usage.CachedInputTokensKnown = true
}
if usage.TotalTokens == 0 {
usage.TotalTokens = usage.InputTokens + usage.OutputTokens
}
if usage.CachedInputTokens > usage.InputTokens && usage.InputTokens > 0 {
usage.CachedInputTokens = usage.InputTokens
}
return usage, true
}
func simulatedOpenAIUsageMap(usage Usage) map[string]any {
out := map[string]any{
"prompt_tokens": usage.InputTokens,
"completion_tokens": usage.OutputTokens,
"total_tokens": usage.TotalTokens,
}
if usage.CachedInputTokensKnown {
out["prompt_tokens_details"] = map[string]any{"cached_tokens": usage.CachedInputTokens}
}
return out
}
func simulatedResponseUsageMap(usage Usage) map[string]any {
out := map[string]any{
"input_tokens": usage.InputTokens,
"output_tokens": usage.OutputTokens,
"total_tokens": usage.TotalTokens,
}
if usage.CachedInputTokensKnown {
out["input_tokens_details"] = map[string]any{"cached_tokens": usage.CachedInputTokens}
}
return out
}
func simulatedProgress(request Request) []Progress {
provider := request.Candidate.Provider
if provider == "" {

View File

@ -36,10 +36,11 @@ type Response struct {
}
type Usage struct {
InputTokens int
OutputTokens int
CachedInputTokens int
TotalTokens int
InputTokens int
OutputTokens int
CachedInputTokens int
CachedInputTokensKnown bool
TotalTokens int
}
type Progress struct {

View File

@ -275,12 +275,61 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
CanonicalModelKey string `json:"canonicalModelKey"`
ProviderModelName string `json:"providerModelName"`
ModelType []string `json:"modelType"`
ModelAlias string `json:"modelAlias"`
} `json:"items"`
}
doJSON(t, server.URL, http.MethodGet, "/api/admin/catalog/base-models", loginResponse.AccessToken, nil, http.StatusOK, &baseModels)
if len(baseModels.Items) < 300 {
t.Fatalf("server-main seed should include the migrated base model catalog: got %d", len(baseModels.Items))
}
requiredMiniMaxModels := map[string]struct {
providerModelName string
modelType string
alias string
}{
"minimax:speech-2.8-hd": {providerModelName: "speech-2.8-hd", modelType: "text_to_speech", alias: "MiniMax-Speech-2.8-HD"},
"minimax:speech-2.8-turbo": {providerModelName: "speech-2.8-turbo", modelType: "text_to_speech", alias: "MiniMax-Speech-2.8-Turbo"},
"minimax:voice-clone": {providerModelName: "voice_clone", modelType: "voice_clone", alias: "MiniMax-Voice-Clone"},
}
seenMiniMaxModels := map[string]bool{}
for _, model := range baseModels.Items {
required, ok := requiredMiniMaxModels[model.CanonicalModelKey]
if !ok {
continue
}
seenMiniMaxModels[model.CanonicalModelKey] = true
if model.ProviderModelName != required.providerModelName || model.ModelAlias != required.alias || !stringSliceContains(model.ModelType, required.modelType) {
t.Fatalf("MiniMax base model is not aligned with server-main: %+v", model)
}
if strings.Contains(model.ModelAlias, " ") {
t.Fatalf("MiniMax Agent-facing alias should not contain spaces: %+v", model)
}
}
for canonicalKey := range requiredMiniMaxModels {
if !seenMiniMaxModels[canonicalKey] {
t.Fatalf("missing MiniMax base model %s in catalog", canonicalKey)
}
}
var modelCatalog struct {
Items []miniMaxModelCatalogItem `json:"items"`
}
doJSON(t, server.URL, http.MethodGet, "/api/v1/model-catalog", apiKeyResponse.Secret, nil, http.StatusOK, &modelCatalog)
for _, required := range requiredMiniMaxModels {
item := findModelCatalogItem(modelCatalog.Items, required.alias)
if item == nil {
t.Fatalf("missing MiniMax model %s in development model catalog", required.alias)
}
if item.ModelName != required.providerModelName || !stringSliceContains(item.ModelType, required.modelType) {
t.Fatalf("MiniMax model catalog item is not aligned with base model: %+v", *item)
}
if strings.Contains(item.Alias, " ") {
t.Fatalf("MiniMax development catalog alias should not contain spaces: %+v", *item)
}
if !modelCatalogItemHasEnabledSource(item.Sources, required.alias, required.providerModelName, required.modelType) {
t.Fatalf("MiniMax development catalog item has no enabled source: %+v", *item)
}
}
baseModelInput := map[string]any{
"providerKey": "openai",
@ -486,7 +535,7 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
speechMarker := "speech-simulation-" + suffixText
var speechResult map[string]any
doJSON(t, server.URL, http.MethodPost, "/api/v1/speech/generations", apiKeyResponse.Secret, map[string]any{
"model": "speech-2.6-turbo",
"model": "MiniMax-Speech-2.8-HD",
"runMode": "simulation",
"text": "hello gateway speech",
"voice_id": "female-shaonv",
@ -516,6 +565,46 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
t.Fatalf("speech simulation task should succeed with text_to_speech billing: %+v", speechTaskDetail)
}
voiceCloneMarker := "voice-clone-simulation-" + suffixText
voiceID := "voice_clone_" + suffixText
var voiceCloneResult map[string]any
doJSON(t, server.URL, http.MethodPost, "/api/v1/voice_clone", apiKeyResponse.Secret, map[string]any{
"model": "MiniMax-Voice-Clone",
"runMode": "simulation",
"voice_id": voiceID,
"audio_url": "/static/simulation/audio.wav",
"prompt_audio_url": "/static/simulation/audio.wav",
"prompt_text": "hello prompt audio",
"text": "hello voice clone preview",
"preview_model": "speech-2.8-hd",
"simulation": true,
"simulationDurationMs": 5,
"integrationTestMarker": voiceCloneMarker,
}, http.StatusOK, &voiceCloneResult)
if voiceCloneResult["status"] != "success" || voiceCloneResult["voice_id"] != voiceID {
t.Fatalf("unexpected voice clone compatible result: %+v", voiceCloneResult)
}
clonedVoice, _ := voiceCloneResult["cloned_voice"].(map[string]any)
if clonedVoice == nil || clonedVoice["voiceId"] != voiceID || clonedVoice["previewModel"] != "speech-2.8-hd" || clonedVoice["platformId"] == "" {
t.Fatalf("voice clone should persist a platform-bound cloned voice: %+v", voiceCloneResult)
}
var voiceCloneTaskDetail struct {
Status string `json:"status"`
ModelType string `json:"modelType"`
Model string `json:"model"`
Attempts []struct {
PlatformID string `json:"platformId"`
} `json:"attempts"`
}
voiceCloneTaskID := waitForTaskIDByRequestField(t, ctx, testPool, "integrationTestMarker", voiceCloneMarker, 2*time.Second)
doJSON(t, server.URL, http.MethodGet, "/api/v1/tasks/"+voiceCloneTaskID, apiKeyResponse.Secret, nil, http.StatusOK, &voiceCloneTaskDetail)
if voiceCloneTaskDetail.Status != "succeeded" || voiceCloneTaskDetail.ModelType != "voice_clone" || voiceCloneTaskDetail.Model != "MiniMax-Voice-Clone" || len(voiceCloneTaskDetail.Attempts) == 0 {
t.Fatalf("voice clone simulation task should succeed through the development interface: %+v", voiceCloneTaskDetail)
}
if voiceCloneTaskDetail.Attempts[0].PlatformID != clonedVoice["platformId"] {
t.Fatalf("voice clone persisted platform should match task attempt: task=%+v voice=%+v", voiceCloneTaskDetail, clonedVoice)
}
doubaoLiteImageEditModel := "doubao-5.0-lite图像编辑"
var doubaoLitePlatformModel struct {
ID string `json:"id"`
@ -857,7 +946,7 @@ WHERE reference_type = 'gateway_task'
"simulationDurationMs": 5,
"simulationProfile": "non_retryable_failure",
"messages": []map[string]any{{"role": "user", "content": "failed first"}},
}, "rate-limit-failed-first-"+suffixText, http.StatusBadGateway, nil, &rateLimitFailedTask.Task)
}, "rate-limit-failed-first-"+suffixText, http.StatusBadRequest, nil, &rateLimitFailedTask.Task)
if rateLimitFailedTask.Task.Status != "failed" || rateLimitFailedTask.Task.ErrorCode != "bad_request" {
t.Fatalf("failed rate-limited task should fail before consuming rpm: %+v", rateLimitFailedTask.Task)
}
@ -1460,6 +1549,142 @@ WHERE m.platform_id = $1::uuid
}
}
func TestCacheAffinityRoutingFlow(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 the PostgreSQL integration flow")
}
ctx := context.Background()
applyMigration(t, ctx, databaseURL)
db, err := store.Connect(ctx, databaseURL)
if err != nil {
t.Fatalf("connect store: %v", err)
}
defer db.Close()
serverCtx, cancelServer := context.WithCancel(ctx)
defer cancelServer()
server := httptest.NewServer(NewServerWithContext(serverCtx, config.Config{
AppEnv: "test",
HTTPAddr: ":0",
DatabaseURL: databaseURL,
IdentityMode: "hybrid",
JWTSecret: "test-secret",
CORSAllowedOrigin: "*",
}, db, slog.New(slog.NewTextHandler(io.Discard, nil))))
defer server.Close()
suffixText := strconv.FormatInt(time.Now().UnixNano(), 10)
username := "cache_affinity_admin_" + suffixText
password := "password123"
var registerResponse struct {
AccessToken string `json:"accessToken"`
}
doJSON(t, server.URL, http.MethodPost, "/api/v1/auth/register", "", map[string]any{
"username": username,
"email": username + "@example.com",
"password": password,
}, http.StatusCreated, &registerResponse)
testPool, err := pgxpool.New(ctx, databaseURL)
if err != nil {
t.Fatalf("connect test pool: %v", err)
}
defer testPool.Close()
if _, err := testPool.Exec(ctx, `UPDATE gateway_users SET roles = '["admin"]'::jsonb WHERE username = $1`, username); err != nil {
t.Fatalf("promote cache affinity user: %v", err)
}
var loginResponse struct {
AccessToken string `json:"accessToken"`
}
doJSON(t, server.URL, http.MethodPost, "/api/v1/auth/login", "", map[string]any{
"account": username,
"password": password,
}, http.StatusOK, &loginResponse)
var apiKeyResponse struct {
Secret string `json:"secret"`
APIKey struct {
ID string `json:"id"`
} `json:"apiKey"`
}
doJSON(t, server.URL, http.MethodPost, "/api/v1/api-keys", loginResponse.AccessToken, map[string]any{
"name": "cache affinity key",
}, http.StatusCreated, &apiKeyResponse)
model := "cache-affinity-smoke-" + suffixText
cacheKey := "cache-affinity-key-" + suffixText
lowPlatform, lowPlatformModelID := createSimulationTextPlatformModel(t, server.URL, loginResponse.AccessToken, "cache-low-"+suffixText, "Cache Affinity Low Priority", model, 20, map[string]any{
"rules": []map[string]any{
{"metric": "concurrent", "limit": 1, "leaseTtlSeconds": 120},
},
})
for index := 0; index < 3; index++ {
detail := runCacheAffinityChatTask(t, ctx, testPool, server.URL, apiKeyResponse.Secret, model, cacheKey, "prime-"+strconv.Itoa(index)+"-"+suffixText, map[string]any{
"inputTokens": 1000,
"cachedInputTokens": 800,
"outputTokens": 20,
})
if detail.Status != "succeeded" || len(detail.Attempts) != 1 || detail.Attempts[0].PlatformName != "Cache Affinity Low Priority" {
t.Fatalf("priming request should use the only low-priority platform, got %+v", detail)
}
}
highPlatform, _ := createSimulationTextPlatformModel(t, server.URL, loginResponse.AccessToken, "cache-high-"+suffixText, "Cache Affinity High Priority", model, 1, nil)
_ = highPlatform
sameKeyDetail := runCacheAffinityChatTask(t, ctx, testPool, server.URL, apiKeyResponse.Secret, model, cacheKey, "same-key-"+suffixText, map[string]any{
"inputTokens": 1000,
"cachedInputTokens": 700,
"outputTokens": 20,
})
if sameKeyDetail.Status != "succeeded" || len(sameKeyDetail.Attempts) != 1 || sameKeyDetail.Attempts[0].PlatformName != "Cache Affinity Low Priority" {
t.Fatalf("same cache key should prefer cached low-priority platform, got %+v", sameKeyDetail)
}
if !boolFromTestMap(sameKeyDetail.Attempts[0].Metrics, "cacheAffinityApplied") || floatFromTestAny(sameKeyDetail.Attempts[0].Metrics["cacheAffinityBoost"]) <= 0 {
t.Fatalf("same cache key attempt should expose cache affinity metrics: %+v", sameKeyDetail.Attempts[0].Metrics)
}
if intFromTestAny(sameKeyDetail.Usage["inputTokens"]) != 1000 ||
intFromTestAny(sameKeyDetail.Usage["cachedInputTokens"]) != 700 ||
!boolFromTestMap(sameKeyDetail.Usage, "cachedInputTokensKnown") {
t.Fatalf("same cache key task detail should expose simulated usage, got %+v", sameKeyDetail.Usage)
}
newKeyDetail := runCacheAffinityChatTask(t, ctx, testPool, server.URL, apiKeyResponse.Secret, model, "cache-affinity-new-"+suffixText, "new-key-"+suffixText, map[string]any{
"inputTokens": 1000,
"cachedInputTokens": 0,
"outputTokens": 20,
})
if newKeyDetail.Status != "succeeded" || len(newKeyDetail.Attempts) != 1 || newKeyDetail.Attempts[0].PlatformName != "Cache Affinity High Priority" {
t.Fatalf("new cache key should fall back to base priority, got %+v", newKeyDetail)
}
seedQueuedConcurrencyLoad(t, ctx, testPool, lowPlatform.PlatformKey, model, lowPlatformModelID)
fullAvoidedDetail := runCacheAffinityChatTask(t, ctx, testPool, server.URL, apiKeyResponse.Secret, model, cacheKey, "full-avoid-"+suffixText, map[string]any{
"inputTokens": 1000,
"cachedInputTokens": 0,
"outputTokens": 20,
})
if fullAvoidedDetail.Status != "succeeded" || len(fullAvoidedDetail.Attempts) != 1 || fullAvoidedDetail.Attempts[0].PlatformName != "Cache Affinity High Priority" {
t.Fatalf("full cached platform should be avoided before cache affinity boost, got %+v", fullAvoidedDetail)
}
var requestCount int
var inputTokens int
var cachedTokens int
var emaHitRatio float64
if err := testPool.QueryRow(ctx, `
SELECT request_count::int, input_tokens::int, cached_input_tokens::int, ema_hit_ratio::float8
FROM gateway_cache_affinity_stats
WHERE client_id = $1
ORDER BY updated_at DESC
LIMIT 1`, lowPlatform.PlatformKey+":text_generate:"+model).Scan(&requestCount, &inputTokens, &cachedTokens, &emaHitRatio); err != nil {
t.Fatalf("read cache affinity stats: %v", err)
}
if requestCount < 4 || inputTokens < 4000 || cachedTokens < 3100 || emaHitRatio <= 0 {
t.Fatalf("cache affinity stats should aggregate low-priority platform observations, count=%d input=%d cached=%d ema=%v", requestCount, inputTokens, cachedTokens, emaHitRatio)
}
}
func TestOriginAllowedSupportsCommaSeparatedOrigins(t *testing.T) {
allowed := "http://localhost:5178, http://127.0.0.1:5178"
if !originAllowed("http://localhost:5178", allowed) {
@ -1636,6 +1861,108 @@ type taskWaitDetail struct {
} `json:"attempts"`
}
type cacheAffinityPlatformFixture struct {
ID string `json:"id"`
PlatformKey string `json:"platformKey"`
Name string `json:"name"`
}
type cacheAffinityTaskDetail struct {
ID string `json:"id"`
Status string `json:"status"`
Usage map[string]any `json:"usage"`
Metrics map[string]any `json:"metrics"`
Attempts []struct {
AttemptNo int `json:"attemptNo"`
PlatformName string `json:"platformName"`
Status string `json:"status"`
Metrics map[string]any `json:"metrics"`
} `json:"attempts"`
}
func createSimulationTextPlatformModel(t *testing.T, baseURL string, adminToken string, platformKey string, platformName string, model string, priority int, rateLimitPolicy map[string]any) (cacheAffinityPlatformFixture, string) {
t.Helper()
var platform cacheAffinityPlatformFixture
doJSON(t, baseURL, http.MethodPost, "/api/admin/platforms", adminToken, map[string]any{
"provider": "openai",
"platformKey": platformKey,
"name": platformName,
"baseUrl": "https://api.openai.com/v1",
"authType": "bearer",
"credentials": map[string]any{"mode": "simulation"},
"priority": priority,
}, http.StatusCreated, &platform)
if platform.ID == "" || platform.PlatformKey == "" {
t.Fatalf("cache affinity platform was not created: %+v", platform)
}
payload := map[string]any{
"canonicalModelKey": "openai:gpt-4o-mini",
"modelName": model,
"providerModelName": model,
"modelAlias": model,
"modelType": []string{"text_generate"},
"displayName": platformName,
}
if len(rateLimitPolicy) > 0 {
payload["rateLimitPolicy"] = rateLimitPolicy
}
var platformModel struct {
ID string `json:"id"`
}
doJSON(t, baseURL, http.MethodPost, "/api/admin/platforms/"+platform.ID+"/models", adminToken, payload, http.StatusCreated, &platformModel)
if platformModel.ID == "" {
t.Fatalf("cache affinity platform model was not created: %+v", platformModel)
}
return platform, platformModel.ID
}
func runCacheAffinityChatTask(t *testing.T, ctx context.Context, pool *pgxpool.Pool, baseURL string, token string, model string, cacheAffinityKey string, marker string, simulationUsage map[string]any) cacheAffinityTaskDetail {
t.Helper()
var detail cacheAffinityTaskDetail
doAPIV1ChatCompletionAndLoadTask(t, ctx, pool, baseURL, token, map[string]any{
"model": model,
"runMode": "simulation",
"simulation": true,
"simulationDurationMs": 5,
"cacheAffinityKey": cacheAffinityKey,
"simulationUsage": simulationUsage,
"messages": []map[string]any{{"role": "user", "content": "cache affinity route"}},
}, marker, http.StatusOK, nil, &detail)
return detail
}
type miniMaxModelCatalogItem struct {
Alias string `json:"alias"`
ModelName string `json:"modelName"`
ModelType []string `json:"modelType"`
Sources []miniMaxModelCatalogSource `json:"sources"`
}
type miniMaxModelCatalogSource struct {
ModelAlias string `json:"modelAlias"`
ModelName string `json:"modelName"`
ModelType []string `json:"modelType"`
Enabled bool `json:"enabled"`
}
func findModelCatalogItem(items []miniMaxModelCatalogItem, alias string) *miniMaxModelCatalogItem {
for index := range items {
if items[index].Alias == alias {
return &items[index]
}
}
return nil
}
func modelCatalogItemHasEnabledSource(sources []miniMaxModelCatalogSource, alias string, modelName string, modelType string) bool {
for _, source := range sources {
if source.Enabled && source.ModelAlias == alias && source.ModelName == modelName && stringSliceContains(source.ModelType, modelType) {
return true
}
}
return false
}
func waitForTaskStatus(t *testing.T, baseURL string, token string, taskID string, statuses []string, timeout time.Duration) taskWaitDetail {
t.Helper()
wanted := map[string]bool{}
@ -1762,7 +2089,7 @@ func assertLoadAvoidanceSimulatedRetryChain(t *testing.T, ctx context.Context, t
"simulation": true,
"simulationDurationMs": 5,
"messages": []map[string]any{{"role": "user", "content": "load avoidance retry chain"}},
}, "load-avoidance-"+suffixText, http.StatusBadGateway, nil, &taskResponse.Task)
}, "load-avoidance-"+suffixText, http.StatusBadRequest, nil, &taskResponse.Task)
if taskResponse.Task.ID == "" || taskResponse.Task.Status != "failed" || taskResponse.Task.ErrorCode != "bad_request" {
t.Fatalf("load avoidance task should only fail after avoided clients are retried, got %+v", taskResponse.Task)
}
@ -1856,6 +2183,22 @@ func boolFromTestMap(values map[string]any, key string) bool {
return value
}
func intFromTestAny(value any) int {
switch typed := value.(type) {
case int:
return typed
case int64:
return int(typed)
case float64:
return int(typed)
case json.Number:
out, _ := typed.Int64()
return int(out)
default:
return 0
}
}
func floatFromTestAny(value any) float64 {
switch typed := value.(type) {
case float64:

View File

@ -80,6 +80,7 @@ func (s *compatibleStreamWriter) writeDone(w http.ResponseWriter, output map[str
}
if s.includeUsage && !s.sentUsage {
if usage, ok := output["usage"].(map[string]any); ok && len(usage) > 0 {
usage = clients.NormalizeChatCompletionUsage(usage)
s.writeChatData(w, s.chatChunk([]any{}, usage))
s.sentUsage = true
}

View File

@ -0,0 +1,69 @@
package runner
import (
"crypto/sha256"
"encoding/hex"
"encoding/json"
"strings"
)
func buildCacheAffinityKey(kind string, modelType string, body map[string]any) string {
if len(body) == 0 {
return ""
}
for _, key := range []string{"cacheAffinityKey", "cache_affinity_key"} {
if value := strings.TrimSpace(stringFromAny(body[key])); value != "" {
return "explicit:" + sha256Text(value)
}
}
for _, key := range []string{"sessionId", "session_id", "conversationId", "conversation_id"} {
if value := strings.TrimSpace(stringFromAny(body[key])); value != "" {
return key + ":" + sha256Text(value)
}
}
if !cacheAffinityPromptHashSupported(kind, modelType) {
return ""
}
payload := map[string]any{
"kind": kind,
"modelType": modelType,
}
if messages, ok := body["messages"]; ok {
payload["messages"] = messages
}
if tools, ok := body["tools"]; ok {
payload["tools"] = tools
}
if instructions, ok := body["instructions"]; ok {
payload["instructions"] = instructions
}
if input, ok := body["input"]; ok {
payload["input"] = input
}
if len(payload) <= 2 {
return ""
}
raw, err := json.Marshal(payload)
if err != nil || len(raw) == 0 {
return ""
}
return "prompt:" + sha256Text(string(raw))
}
func cacheAffinityPromptHashSupported(kind string, modelType string) bool {
switch strings.TrimSpace(kind) {
case "chat.completions", "responses":
return true
}
switch strings.TrimSpace(modelType) {
case "chat", "text_generate", "responses":
return true
default:
return false
}
}
func sha256Text(value string) string {
sum := sha256.Sum256([]byte(value))
return hex.EncodeToString(sum[:])
}

View File

@ -0,0 +1,42 @@
package runner
import "testing"
func TestBuildCacheAffinityKeyUsesExplicitKeyFirst(t *testing.T) {
body := map[string]any{
"cacheAffinityKey": "session-a",
"sessionId": "session-b",
"messages": []any{map[string]any{"role": "user", "content": "hello"}},
}
got := buildCacheAffinityKey("chat.completions", "text_generate", body)
if got == "" || got[:9] != "explicit:" {
t.Fatalf("expected explicit cache affinity key, got %q", got)
}
}
func TestBuildCacheAffinityKeyUsesSessionContinuity(t *testing.T) {
got := buildCacheAffinityKey("chat.completions", "text_generate", map[string]any{
"session_id": "conversation-1",
})
if got == "" || got[:11] != "session_id:" {
t.Fatalf("expected session cache affinity key, got %q", got)
}
}
func TestBuildCacheAffinityKeyFallsBackToPromptHash(t *testing.T) {
body := map[string]any{
"messages": []any{
map[string]any{"role": "system", "content": "stable"},
map[string]any{"role": "user", "content": "hello"},
},
"tools": []any{map[string]any{"type": "function", "function": map[string]any{"name": "lookup"}}},
}
first := buildCacheAffinityKey("chat.completions", "text_generate", body)
second := buildCacheAffinityKey("chat.completions", "text_generate", body)
if first == "" || first != second || first[:7] != "prompt:" {
t.Fatalf("expected stable prompt hash cache affinity key, first=%q second=%q", first, second)
}
}

View File

@ -68,6 +68,35 @@ func TestBillingsSplitsCachedInputTokens(t *testing.T) {
}
}
func TestBillingsDefaultsCachedInputPriceToOneTenthOfInputPrice(t *testing.T) {
service := &Service{}
candidate := store.RuntimeModelCandidate{
ModelName: "test-model",
BillingConfig: map[string]any{
"textInputPer1k": 1.0,
"textOutputPer1k": 2.0,
},
}
response := clients.Response{
Usage: clients.Usage{
InputTokens: 1000,
CachedInputTokens: 400,
OutputTokens: 500,
TotalTokens: 1500,
},
}
lines := service.billings(context.Background(), nil, "chat.completions", nil, candidate, response, false)
if len(lines) != 3 {
t.Fatalf("expected uncached input, cached input, and output lines, got %+v", lines)
}
cached, _ := lines[1].(map[string]any)
if cached["resourceType"] != "text_cached_input" || cached["quantity"] != 400 || cached["amount"] != 0.04 {
t.Fatalf("unexpected cached input fallback billing line: %+v", cached)
}
}
func TestEffectiveRateLimitPolicyTreatsEmptyRuntimePolicyAsUnlimited(t *testing.T) {
policy := effectiveRateLimitPolicy(store.RuntimeModelCandidate{
PlatformRateLimitPolicy: map[string]any{"rules": []any{

View File

@ -82,7 +82,7 @@ func (s *Service) billings(ctx context.Context, user *auth.User, kind string, bo
inputPrice := resourcePrice(config, "text", "textInputPer1k", "inputTokenPrice", "basePrice")
cachedInputPrice := resourcePrice(config, "text", "textCachedInputPer1k", "cachedInputTokenPrice", "cacheInputTokenPrice", "textInputCacheHitPer1k", "inputCacheHitTokenPrice")
if cachedInputPrice <= 0 {
cachedInputPrice = inputPrice
cachedInputPrice = inputPrice / 10
}
inputAmount := roundPrice(float64(uncachedInputTokens) / 1000 * inputPrice * discount)
lines := []any{}

View File

@ -139,6 +139,19 @@ func attemptMetrics(candidate store.RuntimeModelCandidate, attemptNo int, simula
"queued": candidate.LoadMetrics.QueuedCount,
}
}
if candidate.CacheAffinity.Key != "" {
metrics["cacheAffinityKey"] = candidate.CacheAffinity.Key
metrics["cacheAdjustedPriority"] = candidate.CacheAffinity.AdjustedPriority
metrics["cacheAffinitySamples"] = candidate.CacheAffinity.RequestCount
metrics["cacheAffinityScore"] = candidate.CacheAffinity.Score
metrics["cacheAffinityConfidence"] = candidate.CacheAffinity.Confidence
metrics["cacheAffinityHitRatio"] = candidate.CacheAffinity.EMAHitRatio
metrics["cacheAffinityEMAHitRatio"] = candidate.CacheAffinity.EMAHitRatio
metrics["cacheAffinityLastHitRatio"] = candidate.CacheAffinity.LastHitRatio
metrics["cacheAffinityCachedInputTokens"] = candidate.CacheAffinity.CachedInputTokens
metrics["cacheAffinityBoost"] = candidate.CacheAffinity.Boost
metrics["cacheAffinityApplied"] = candidate.CacheAffinity.Applied
}
if attemptNo > 0 {
metrics["attempt"] = attemptNo
}
@ -151,9 +164,10 @@ func usageToMap(usage clients.Usage) map[string]any {
out["inputTokens"] = usage.InputTokens
out["promptTokens"] = usage.InputTokens
}
if usage.CachedInputTokens > 0 {
if usage.CachedInputTokensKnown || usage.CachedInputTokens > 0 {
out["cachedInputTokens"] = usage.CachedInputTokens
out["cachedPromptTokens"] = usage.CachedInputTokens
out["cachedInputTokensKnown"] = usage.CachedInputTokensKnown
}
if usage.OutputTokens > 0 {
out["outputTokens"] = usage.OutputTokens

View File

@ -0,0 +1,38 @@
package runner
import (
"testing"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestAttemptMetricsIncludesCacheAffinityCoefficient(t *testing.T) {
metrics := attemptMetrics(store.RuntimeModelCandidate{
ModelName: "deepseek-chat",
ModelType: "text_generate",
PlatformID: "deepseek-platform",
PlatformPriority: 100,
CacheAffinity: store.RuntimeCandidateCacheAffinity{
Key: "conversation:test",
RequestCount: 2,
CachedInputTokens: 7296,
EMAHitRatio: 0.91,
LastHitRatio: 0.92,
Confidence: 1,
Score: 1.74,
Boost: 20,
AdjustedPriority: 80,
Applied: true,
},
}, 1, false)
if metrics["cacheAffinityScore"] != 1.74 {
t.Fatalf("expected cache affinity score in attempt metrics, got %+v", metrics)
}
if metrics["cacheAffinityCachedInputTokens"] != 7296 {
t.Fatalf("expected cached input tokens in attempt metrics, got %+v", metrics)
}
if metrics["cacheAffinityHitRatio"] != 0.91 || metrics["cacheAffinityLastHitRatio"] != 0.92 {
t.Fatalf("expected hit ratios in attempt metrics, got %+v", metrics)
}
}

View File

@ -34,17 +34,20 @@ type retryDecision struct {
type failoverDecision struct {
Retry bool
Action string
Target string
Reason string
CooldownSeconds int
DemoteSteps int
Match policyRuleMatch
Info failureInfo
}
type priorityDemoteDecision struct {
Demote bool
Reason string
Match policyRuleMatch
Info failureInfo
Demote bool
Reason string
DemoteSteps int
Match policyRuleMatch
Info failureInfo
}
func shouldRetrySameClient(candidate store.RuntimeModelCandidate, err error) bool {
@ -77,10 +80,13 @@ func failoverDecisionForCandidate(runnerPolicy store.RunnerPolicy, candidate sto
if strings.TrimSpace(runnerPolicy.Status) != "" && runnerPolicy.Status != "active" {
return failoverDecision{Retry: false, Action: "stop", Reason: "runner_policy_disabled", Match: policyRuleMatch{Source: "gateway_runner_policies", Policy: "runnerPolicy", Rule: "status", Value: runnerPolicy.Status}, Info: info}
}
if match, ok := hardStopPolicyMatch(runnerPolicy.HardStopPolicy, info); ok {
return failoverDecision{Retry: false, Action: "stop", Reason: "hard_stop_policy", Match: match, Info: info}
}
overridePolicy := failoverOverridePolicy(candidate.RuntimePolicyOverride)
hasActionRules := len(actionRulesFromPolicy(runnerPolicy.FailoverPolicy)) > 0 || len(actionRulesFromPolicy(overridePolicy)) > 0
if !hasActionRules {
if match, ok := hardStopPolicyMatch(runnerPolicy.HardStopPolicy, info); ok {
return failoverDecision{Retry: false, Action: "stop", Reason: "hard_stop_policy", Match: match, Info: info}
}
}
policy := effectiveFailoverPolicy(runnerPolicy.FailoverPolicy, candidate.RuntimePolicyOverride)
if !boolFromPolicy(policy, "enabled", true) {
source := "gateway_runner_policies.failover_policy"
@ -89,22 +95,26 @@ func failoverDecisionForCandidate(runnerPolicy store.RunnerPolicy, candidate sto
}
return failoverDecision{Retry: false, Action: "stop", Reason: "failover_disabled", Match: policyRuleMatch{Source: source, Policy: "failoverPolicy", Rule: "enabled", Value: "false"}, Info: info}
}
if errors.Is(err, store.ErrRateLimited) {
return failoverDecision{Retry: false, Action: "queue", Reason: "local_rate_limit_wait_queue", Match: policyRuleMatch{Source: "gateway_rate_limits", Policy: "rateLimitPolicy", Rule: "localCapacity", Value: "exceeded"}, Info: info}
}
if ruleDecision, ok := failoverActionRuleDecisionWithSources(runnerPolicy.FailoverPolicy, overridePolicy, info); ok {
return ruleDecision
}
if match, ok := failoverDenyMatchWithSources(runnerPolicy.FailoverPolicy, overridePolicy, info); ok {
return failoverDecision{Retry: false, Action: "stop", Reason: "failover_deny_policy", Match: match, Info: info}
}
action := failoverAction(policy, info)
target := defaultFailoverActionTarget(action)
cooldownSeconds := intFromPolicy(policy, "cooldownSeconds")
if cooldownSeconds <= 0 {
cooldownSeconds = 300
}
if errors.Is(err, store.ErrRateLimited) {
return failoverDecision{Retry: false, Action: "queue", Reason: "local_rate_limit_wait_queue", Match: policyRuleMatch{Source: "gateway_rate_limits", Policy: "rateLimitPolicy", Rule: "localCapacity", Value: "exceeded"}, Info: info}
}
if match, ok := failoverAllowMatchWithSources(runnerPolicy.FailoverPolicy, overridePolicy, info); ok {
return failoverDecision{Retry: true, Action: action, Reason: "failover_allow_policy", CooldownSeconds: cooldownSeconds, Match: match, Info: info}
return failoverDecision{Retry: true, Action: action, Target: target, Reason: "failover_allow_policy", CooldownSeconds: cooldownSeconds, DemoteSteps: 1, Match: match, Info: info}
}
if clients.IsRetryable(err) {
return failoverDecision{Retry: true, Action: action, Reason: "client_retryable", CooldownSeconds: cooldownSeconds, Match: policyRuleMatch{Source: "provider_client", Policy: "ClientError", Rule: "Retryable", Value: "true"}, Info: info}
return failoverDecision{Retry: true, Action: action, Target: target, Reason: "client_retryable", CooldownSeconds: cooldownSeconds, DemoteSteps: 1, Match: policyRuleMatch{Source: "provider_client", Policy: "ClientError", Rule: "Retryable", Value: "true"}, Info: info}
}
return failoverDecision{Retry: false, Action: "stop", Reason: "client_non_retryable", Match: policyRuleMatch{Source: "provider_client", Policy: "ClientError", Rule: "Retryable", Value: "false"}, Info: info}
}
@ -122,12 +132,15 @@ func priorityDemoteDecisionForCandidate(runnerPolicy store.RunnerPolicy, err err
if hardStopPolicyMatches(runnerPolicy.HardStopPolicy, info) {
return priorityDemoteDecision{Demote: false, Reason: "hard_stop_policy", Info: info}
}
if len(actionRulesFromPolicy(runnerPolicy.FailoverPolicy)) > 0 {
return priorityDemoteDecision{Demote: false, Reason: "failover_action_rules_enabled", Info: info}
}
policy := runnerPolicy.PriorityDemotePolicy
if !boolFromPolicy(policy, "enabled", false) {
return priorityDemoteDecision{Demote: false, Reason: "priority_demote_disabled", Info: info}
}
if match, ok := priorityDemotePolicyMatch(policy, info); ok {
return priorityDemoteDecision{Demote: true, Reason: "priority_demote_policy", Match: match, Info: info}
return priorityDemoteDecision{Demote: true, Reason: "priority_demote_policy", DemoteSteps: 99, Match: match, Info: info}
}
return priorityDemoteDecision{Demote: false, Reason: "priority_demote_no_match", Info: info}
}
@ -299,11 +312,118 @@ func priorityDemotePolicyMatch(policy map[string]any, info failureInfo) (policyR
func failoverAction(policy map[string]any, info failureInfo) string {
actions, _ := policy["actions"].(map[string]any)
if action := stringFromAny(actions[info.Category]); action != "" {
return action
return normalizeFailoverAction(action)
}
return "next"
}
func failoverActionRuleDecisionWithSources(base map[string]any, override map[string]any, info failureInfo) (failoverDecision, bool) {
if rules := actionRulesFromPolicy(override); len(rules) > 0 {
return failoverActionRuleDecision(rules, "runtime_policy_override.failoverPolicy", info)
}
return failoverActionRuleDecision(actionRulesFromPolicy(base), "gateway_runner_policies.failover_policy", info)
}
func failoverActionRuleDecision(rules []map[string]any, source string, info failureInfo) (failoverDecision, bool) {
for index, rule := range rules {
if !boolFromPolicy(rule, "enabled", true) {
continue
}
match, ok := failoverActionRuleMatch(rule, info, source, index)
if !ok {
continue
}
action := normalizeFailoverAction(stringFromAny(rule["action"]))
target := normalizeFailoverTarget(stringFromAny(rule["target"]), action)
cooldownSeconds := intFromPolicy(rule, "cooldownSeconds")
if cooldownSeconds <= 0 {
cooldownSeconds = 300
}
demoteSteps := intFromPolicy(rule, "demoteSteps")
if demoteSteps <= 0 {
demoteSteps = 1
}
if action == "stop" {
return failoverDecision{Retry: false, Action: "stop", Target: target, Reason: "failover_action_rule", CooldownSeconds: cooldownSeconds, DemoteSteps: demoteSteps, Match: match, Info: info}, true
}
return failoverDecision{Retry: true, Action: action, Target: target, Reason: "failover_action_rule", CooldownSeconds: cooldownSeconds, DemoteSteps: demoteSteps, Match: match, Info: info}, true
}
return failoverDecision{}, false
}
func actionRulesFromPolicy(policy map[string]any) []map[string]any {
raw, ok := policy["actionRules"].([]any)
if !ok {
return nil
}
rules := make([]map[string]any, 0, len(raw))
for _, item := range raw {
rule, ok := item.(map[string]any)
if !ok {
continue
}
rules = append(rules, rule)
}
return rules
}
func failoverActionRuleMatch(rule map[string]any, info failureInfo, source string, index int) (policyRuleMatch, bool) {
source = fmt.Sprintf("%s.actionRules[%d]", source, index)
if value, ok := matchingStringListValue(rule, "categories", info.Category); ok {
return policyRuleMatch{Source: source, Policy: "failoverPolicy", Rule: "categories", Value: value}, true
}
if value, ok := matchingStringListValue(rule, "codes", info.Code); ok {
return policyRuleMatch{Source: source, Policy: "failoverPolicy", Rule: "codes", Value: value}, true
}
if value, ok := matchingStringListValue(rule, "errorCodes", info.Code); ok {
return policyRuleMatch{Source: source, Policy: "failoverPolicy", Rule: "errorCodes", Value: value}, true
}
if value, ok := matchingStringListValue(rule, "errorCodes", info.Category); ok {
return policyRuleMatch{Source: source, Policy: "failoverPolicy", Rule: "errorCodes", Value: value}, true
}
if value, ok := matchingIntListValue(rule, "statusCodes", info.Status); ok {
return policyRuleMatch{Source: source, Policy: "failoverPolicy", Rule: "statusCodes", Value: fmt.Sprintf("%d", value)}, true
}
if value, ok := matchingKeywordValue(rule, "keywords", info.Target); ok {
return policyRuleMatch{Source: source, Policy: "failoverPolicy", Rule: "keywords", Value: value}, true
}
return policyRuleMatch{}, false
}
func normalizeFailoverAction(action string) string {
switch strings.TrimSpace(action) {
case "rotate":
return "next"
case "cooldown_and_rotate":
return "cooldown_and_next"
case "demote_and_rotate":
return "demote_and_next"
case "disable_and_rotate":
return "disable_and_next"
case "cooldown_and_next", "demote_and_next", "disable_and_next", "stop", "next":
return strings.TrimSpace(action)
default:
return "next"
}
}
func normalizeFailoverTarget(target string, action string) string {
target = strings.ToLower(strings.TrimSpace(target))
if target == "platform" || target == "model" {
return target
}
return defaultFailoverActionTarget(action)
}
func defaultFailoverActionTarget(action string) string {
switch action {
case "cooldown_and_next":
return "model"
default:
return "platform"
}
}
func boolFromPolicy(policy map[string]any, key string, fallback bool) bool {
value, ok := policy[key].(bool)
if !ok {
@ -442,6 +562,13 @@ func intListFromPolicy(policy map[string]any, key string) []int {
out = append(out, typed)
case float64:
out = append(out, int(typed))
case string:
if parsed := strings.TrimSpace(typed); parsed != "" {
var value int
if _, err := fmt.Sscanf(parsed, "%d", &value); err == nil && value > 0 {
out = append(out, value)
}
}
}
}
return out

View File

@ -197,6 +197,97 @@ func TestProviderAuthErrorsFailOverInsteadOfHardStop(t *testing.T) {
}
}
func TestFailoverActionRulesControlPerRuleCooldownAndDemotion(t *testing.T) {
runnerPolicy := store.RunnerPolicy{
Status: "active",
FailoverPolicy: map[string]any{
"enabled": true,
"actionRules": []any{
map[string]any{
"enabled": true,
"categories": []any{"rate_limit"},
"action": "cooldown_and_next",
"target": "platform",
"cooldownSeconds": 45,
},
map[string]any{
"enabled": true,
"keywords": []any{"temporary overload"},
"action": "demote_and_next",
"demoteSteps": 3,
},
},
},
}
decision := failoverDecisionForCandidate(runnerPolicy, store.RuntimeModelCandidate{}, &clients.ClientError{Code: "rate_limit", StatusCode: 429, Retryable: false})
if !decision.Retry || decision.Action != "cooldown_and_next" || decision.Target != "platform" || decision.CooldownSeconds != 45 {
t.Fatalf("rate limit should use per-rule cooldown action, got %+v", decision)
}
decision = failoverDecisionForCandidate(runnerPolicy, store.RuntimeModelCandidate{}, &clients.ClientError{Code: "server_error", Message: "temporary overload from upstream", StatusCode: 503, Retryable: false})
if !decision.Retry || decision.Action != "demote_and_next" || decision.DemoteSteps != 3 {
t.Fatalf("keyword match should use per-rule demote action, got %+v", decision)
}
}
func TestFailoverActionRulesCanStopAndMatchStringStatusCodes(t *testing.T) {
runnerPolicy := store.RunnerPolicy{
Status: "active",
FailoverPolicy: map[string]any{
"enabled": true,
"actionRules": []any{
map[string]any{
"enabled": true,
"statusCodes": []any{"400"},
"keywords": []any{"bad request"},
"action": "stop",
},
map[string]any{
"enabled": true,
"categories": []any{"request_error"},
"action": "next",
},
},
},
}
decision := failoverDecisionForCandidate(runnerPolicy, store.RuntimeModelCandidate{}, &clients.ClientError{Code: "bad_request", Message: "bad request", StatusCode: 400, Retryable: true})
if decision.Retry || decision.Action != "stop" || decision.Reason != "failover_action_rule" || decision.Match.Rule != "statusCodes" {
t.Fatalf("first stop rule should win before later retry rules, got %+v", decision)
}
}
func TestFailoverActionRulesKeepOrderBeforeDerivedHardStopPolicy(t *testing.T) {
runnerPolicy := store.RunnerPolicy{
Status: "active",
FailoverPolicy: map[string]any{
"enabled": true,
"actionRules": []any{
map[string]any{
"enabled": true,
"keywords": []any{"bad request retryable"},
"action": "next",
},
map[string]any{
"enabled": true,
"statusCodes": []any{400},
"action": "stop",
},
},
},
HardStopPolicy: map[string]any{
"enabled": true,
"statusCodes": []any{400},
},
}
decision := failoverDecisionForCandidate(runnerPolicy, store.RuntimeModelCandidate{}, &clients.ClientError{Code: "bad_request", Message: "bad request retryable", StatusCode: 400, Retryable: true})
if !decision.Retry || decision.Action != "next" || decision.Match.Rule != "keywords" {
t.Fatalf("actionRules should keep list order before derived hard stop compatibility fields, got %+v", decision)
}
}
func TestPriorityDemotePolicyIsKeywordGatedAndHardStopSafe(t *testing.T) {
runnerPolicy := store.RunnerPolicy{
Status: "active",
@ -218,3 +309,28 @@ func TestPriorityDemotePolicyIsKeywordGatedAndHardStopSafe(t *testing.T) {
t.Fatal("priority demotion should not run for hard-stop request errors")
}
}
func TestPriorityDemotePolicyIsSkippedWhenFailoverActionRulesExist(t *testing.T) {
runnerPolicy := store.RunnerPolicy{
Status: "active",
FailoverPolicy: map[string]any{
"enabled": true,
"actionRules": []any{
map[string]any{
"enabled": true,
"categories": []any{"rate_limit"},
"action": "demote_and_next",
},
},
},
PriorityDemotePolicy: map[string]any{
"enabled": true,
"keywords": []any{"rate_limit"},
},
}
decision := priorityDemoteDecisionForCandidate(runnerPolicy, &clients.ClientError{Code: "rate_limit", Message: "rate_limit from upstream", Retryable: true})
if decision.Demote || decision.Reason != "failover_action_rules_enabled" {
t.Fatalf("legacy priority demotion should be skipped when actionRules are active, got %+v", decision)
}
}

View File

@ -45,18 +45,29 @@ func (s *Service) applyCandidateFailurePolicies(ctx context.Context, taskID stri
}
}
func (s *Service) applyFailoverAction(ctx context.Context, taskID string, candidate store.RuntimeModelCandidate, decision failoverDecision, simulated bool, singleSourceProtected bool) {
func (s *Service) applyFailoverAction(ctx context.Context, taskID string, candidate store.RuntimeModelCandidate, requestedModel string, decision failoverDecision, simulated bool, singleSourceProtected bool) {
switch decision.Action {
case "disable_and_next":
if singleSourceProtected {
s.emitSingleSourceProtected(ctx, taskID, candidate, decision.Action, decision.Info.Code, simulated)
return
}
if err := s.store.DisableCandidatePlatform(ctx, candidate.PlatformID); err == nil {
target := decision.Target
if target == "" {
target = "platform"
}
var err error
if target == "model" {
err = s.store.DisableCandidatePlatformModel(ctx, candidate.PlatformModelID)
} else {
err = s.store.DisableCandidatePlatform(ctx, candidate.PlatformID)
}
if err == nil {
_ = s.emit(ctx, taskID, "task.policy.failover_disabled", "running", "failover", 0.51, "candidate platform disabled by runner failover policy", addPolicyTracePayload(map[string]any{
"platformId": candidate.PlatformID,
"platformModelId": candidate.PlatformModelID,
"action": decision.Action,
"target": target,
"reason": decision.Reason,
}, decision.Match, decision.Info), simulated)
}
@ -65,15 +76,43 @@ func (s *Service) applyFailoverAction(ctx context.Context, taskID string, candid
s.emitSingleSourceProtected(ctx, taskID, candidate, decision.Action, decision.Info.Code, simulated)
return
}
if err := s.store.CooldownCandidatePlatformModel(ctx, candidate.PlatformModelID, decision.CooldownSeconds); err == nil {
target := decision.Target
if target == "" {
target = "model"
}
var err error
if target == "platform" {
err = s.store.CooldownCandidatePlatform(ctx, candidate.PlatformID, decision.CooldownSeconds)
} else {
err = s.store.CooldownCandidatePlatformModel(ctx, candidate.PlatformModelID, decision.CooldownSeconds)
}
if err == nil {
_ = s.emit(ctx, taskID, "task.policy.failover_cooled_down", "running", "failover", 0.51, "candidate model cooled down by runner failover policy", addPolicyTracePayload(map[string]any{
"platformId": candidate.PlatformID,
"platformModelId": candidate.PlatformModelID,
"cooldownSeconds": decision.CooldownSeconds,
"action": decision.Action,
"target": target,
"reason": decision.Reason,
}, decision.Match, decision.Info), simulated)
}
case "demote_and_next":
demoteSteps := decision.DemoteSteps
if demoteSteps <= 0 {
demoteSteps = 1
}
if dynamicPriority, err := s.store.DemoteCandidatePlatformPriorityBySteps(ctx, candidate.PlatformID, candidate.PlatformModelID, requestedModel, candidate.ModelType, demoteSteps); err == nil {
_ = s.emit(ctx, taskID, "task.policy.priority_demoted", "running", "priority_demote", 0.52, "candidate platform priority demoted by runner failover policy", addPolicyTracePayload(map[string]any{
"platformId": candidate.PlatformID,
"platformModelId": candidate.PlatformModelID,
"dynamicPriority": dynamicPriority,
"demoteSteps": demoteSteps,
"action": decision.Action,
"target": "platform",
"reason": decision.Reason,
"errorMessage": decision.Info.Message,
}, decision.Match, decision.Info), simulated)
}
}
}
@ -94,12 +133,17 @@ func (s *Service) applyPriorityDemotePolicy(ctx context.Context, taskID string,
if !decision.Demote {
return
}
if dynamicPriority, err := s.store.DemoteCandidatePlatformPriority(ctx, candidate.PlatformID, candidate.PlatformModelID, requestedModel, candidate.ModelType); err == nil {
demoteSteps := decision.DemoteSteps
if demoteSteps <= 0 {
demoteSteps = 99
}
if dynamicPriority, err := s.store.DemoteCandidatePlatformPriorityBySteps(ctx, candidate.PlatformID, candidate.PlatformModelID, requestedModel, candidate.ModelType, demoteSteps); err == nil {
s.recordAttemptTrace(ctx, taskID, attemptNo, priorityDemoteTraceEntry(decision, candidate.PlatformID, candidate.PlatformModelID, dynamicPriority))
_ = s.emit(ctx, taskID, "task.policy.priority_demoted", "running", "priority_demote", 0.52, "candidate platform priority demoted by runner policy", addPolicyTracePayload(map[string]any{
"platformId": candidate.PlatformID,
"platformModelId": candidate.PlatformModelID,
"dynamicPriority": dynamicPriority,
"demoteSteps": demoteSteps,
"code": clients.ErrorCode(cause),
"reason": decision.Reason,
"errorMessage": messageFromError(cause),

View File

@ -145,7 +145,15 @@ func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *aut
return Result{}, err
}
}
candidates, err := s.store.ListModelCandidates(ctx, task.Model, modelType, user)
runnerPolicy, err := s.store.GetActiveRunnerPolicy(ctx)
if err != nil {
return Result{}, err
}
cacheAffinityKey := buildCacheAffinityKey(task.Kind, modelType, body)
candidates, err := s.store.ListModelCandidates(ctx, task.Model, modelType, user, store.ListModelCandidatesOptions{
CacheAffinityKey: cacheAffinityKey,
CacheAffinityPolicy: runnerPolicy.CacheAffinityPolicy,
})
if err != nil {
s.recordFailedAttempt(ctx, failedAttemptRecord{
Task: task,
@ -278,10 +286,6 @@ func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *aut
return Result{}, err
}
runnerPolicy, err := s.store.GetActiveRunnerPolicy(ctx)
if err != nil {
return Result{}, err
}
maxPlatforms := maxPlatformsForCandidates(candidates, runnerPolicy)
maxFailoverDuration := maxFailoverDurationForCandidates(candidates, runnerPolicy)
singleSourceProtected := singleSourceProtectionActive(candidates, maxPlatforms, runnerPolicy)
@ -320,7 +324,7 @@ candidatesLoop:
break candidatesLoop
}
candidateBody := preprocessing.Body
response, err := s.runCandidate(ctx, task, user, candidateBody, preprocessing.Log, candidate, nextAttemptNo, onDelta, singleSourceProtected)
response, err := s.runCandidate(ctx, task, user, candidateBody, preprocessing.Log, candidate, nextAttemptNo, onDelta, singleSourceProtected, runnerPolicy.CacheAffinityPolicy)
if err == nil {
attemptNo = nextAttemptNo
billings := s.billings(ctx, user, task.Kind, candidateBody, candidate, response, isSimulation(task, candidate))
@ -490,7 +494,7 @@ candidatesLoop:
if !decision.Retry {
break
}
s.applyFailoverAction(ctx, task.ID, candidate, decision, isSimulation(task, candidate), singleSourceProtected)
s.applyFailoverAction(ctx, task.ID, candidate, task.Model, decision, isSimulation(task, candidate), singleSourceProtected)
if err := s.emit(ctx, task.ID, "task.retrying", "running", "retry", 0.55, "retrying next client", addPolicyTracePayload(map[string]any{
"attempt": attemptNo,
"action": decision.Action,
@ -531,7 +535,7 @@ candidatesLoop:
return Result{Task: failed, Output: failed.Result}, lastErr
}
func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user *auth.User, body map[string]any, preprocessing parameterPreprocessingLog, candidate store.RuntimeModelCandidate, attemptNo int, onDelta clients.StreamDelta, singleSourceProtected bool) (clients.Response, error) {
func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user *auth.User, body map[string]any, preprocessing parameterPreprocessingLog, candidate store.RuntimeModelCandidate, attemptNo int, onDelta clients.StreamDelta, singleSourceProtected bool, cacheAffinityPolicy map[string]any) (clients.Response, error) {
simulated := isSimulation(task, candidate)
baseAttemptMetrics := mergeMetrics(attemptMetrics(candidate, attemptNo, simulated), parameterPreprocessingMetrics(preprocessing))
reservations := s.rateLimitReservations(ctx, user, candidate, body)
@ -763,6 +767,19 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
}); err != nil {
return clients.Response{}, fmt.Errorf("finish task attempt: %w", err)
}
if err := s.store.RecordCacheAffinityObservation(context.WithoutCancel(ctx), store.CacheAffinityObservationInput{
CacheAffinityKey: candidate.CacheAffinity.Key,
CacheAffinityPolicy: cacheAffinityPolicy,
PlatformID: candidate.PlatformID,
PlatformModelID: candidate.PlatformModelID,
ClientID: candidate.ClientID,
ModelType: candidate.ModelType,
InputTokens: response.Usage.InputTokens,
CachedInputTokens: response.Usage.CachedInputTokens,
CachedInputTokensKnown: response.Usage.CachedInputTokensKnown,
}); err != nil {
s.logger.Warn("record cache affinity observation failed", "error", err, "clientId", candidate.ClientID)
}
return response, nil
}

View File

@ -39,12 +39,21 @@ func failoverTraceEntry(decision failoverDecision, candidate store.RuntimeModelC
if decision.CooldownSeconds > 0 {
entry["cooldownSeconds"] = decision.CooldownSeconds
}
if decision.Target != "" {
entry["target"] = decision.Target
}
if decision.DemoteSteps > 0 {
entry["demoteSteps"] = decision.DemoteSteps
}
return entry
}
func priorityDemoteTraceEntry(decision priorityDemoteDecision, platformID string, platformModelID string, dynamicPriority int) map[string]any {
entry := policyTraceEntry("priority_demoted", "priority_demote", "demote", decision.Reason, decision.Match, decision.Info)
entry["demote"] = decision.Demote
if decision.DemoteSteps > 0 {
entry["demoteSteps"] = decision.DemoteSteps
}
if dynamicPriority > 0 {
entry["dynamicPriority"] = dynamicPriority
}
@ -78,6 +87,16 @@ func addCandidatePriorityTraceFields(entry map[string]any, candidate store.Runti
if candidate.PlatformName != "" {
entry["currentPlatformName"] = candidate.PlatformName
}
if candidate.CacheAffinity.Key != "" {
entry["cacheAffinityKey"] = candidate.CacheAffinity.Key
entry["cacheAffinityScore"] = candidate.CacheAffinity.Score
entry["cacheAffinityConfidence"] = candidate.CacheAffinity.Confidence
entry["cacheAffinityHitRatio"] = candidate.CacheAffinity.EMAHitRatio
entry["cacheAffinityLastHitRatio"] = candidate.CacheAffinity.LastHitRatio
entry["cacheAffinityCachedInputTokens"] = candidate.CacheAffinity.CachedInputTokens
entry["cacheAffinityBoost"] = candidate.CacheAffinity.Boost
entry["cacheAffinityApplied"] = candidate.CacheAffinity.Applied
}
}
func policyTraceEntry(event string, scope string, action string, reason string, match policyRuleMatch, info failureInfo) map[string]any {

View File

@ -0,0 +1,73 @@
package store
import (
"context"
"strings"
)
type CacheAffinityObservationInput struct {
CacheAffinityKey string
CacheAffinityPolicy map[string]any
PlatformID string
PlatformModelID string
ClientID string
ModelType string
InputTokens int
CachedInputTokens int
CachedInputTokensKnown bool
}
func (s *Store) RecordCacheAffinityObservation(ctx context.Context, input CacheAffinityObservationInput) error {
input.CacheAffinityKey = strings.TrimSpace(input.CacheAffinityKey)
input.ClientID = strings.TrimSpace(input.ClientID)
input.ModelType = strings.TrimSpace(input.ModelType)
if input.CacheAffinityKey == "" || input.ClientID == "" || input.ModelType == "" {
return nil
}
if !input.CachedInputTokensKnown || !cacheAffinityPolicyEnabled(input.CacheAffinityPolicy, input.ModelType) {
return nil
}
minInputTokens := cacheAffinityMinInputTokens(input.CacheAffinityPolicy)
if input.InputTokens < minInputTokens || input.InputTokens <= 0 {
return nil
}
if input.CachedInputTokens < 0 {
input.CachedInputTokens = 0
}
if input.CachedInputTokens > input.InputTokens {
input.CachedInputTokens = input.InputTokens
}
alpha := cacheAffinityEMAAlpha(input.CacheAffinityPolicy)
hitRatio := float64(input.CachedInputTokens) / float64(input.InputTokens)
_, err := s.pool.Exec(ctx, `
INSERT INTO gateway_cache_affinity_stats (
client_id, cache_affinity_key, platform_id, platform_model_id, model_type,
request_count, input_tokens, cached_input_tokens, ema_hit_ratio, last_hit_ratio, last_observed_at
)
VALUES (
$1, $2, NULLIF($3::text, '')::uuid, NULLIF($4::text, '')::uuid, $5,
1, $6, $7, $8, $8, now()
)
ON CONFLICT (client_id, cache_affinity_key) DO UPDATE
SET platform_id = COALESCE(EXCLUDED.platform_id, gateway_cache_affinity_stats.platform_id),
platform_model_id = COALESCE(EXCLUDED.platform_model_id, gateway_cache_affinity_stats.platform_model_id),
model_type = EXCLUDED.model_type,
request_count = gateway_cache_affinity_stats.request_count + 1,
input_tokens = gateway_cache_affinity_stats.input_tokens + EXCLUDED.input_tokens,
cached_input_tokens = gateway_cache_affinity_stats.cached_input_tokens + EXCLUDED.cached_input_tokens,
ema_hit_ratio = ($9::float8 * EXCLUDED.last_hit_ratio) + ((1 - $9::float8) * gateway_cache_affinity_stats.ema_hit_ratio),
last_hit_ratio = EXCLUDED.last_hit_ratio,
last_observed_at = now(),
updated_at = now()`,
input.ClientID,
input.CacheAffinityKey,
input.PlatformID,
input.PlatformModelID,
input.ModelType,
input.InputTokens,
input.CachedInputTokens,
hitRatio,
alpha,
)
return err
}

View File

@ -3,6 +3,7 @@ package store
import (
"context"
"fmt"
"math"
"sort"
"strings"
"unicode"
@ -10,11 +11,17 @@ import (
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
)
func (s *Store) ListModelCandidates(ctx context.Context, model string, modelType string, user *auth.User) ([]RuntimeModelCandidate, error) {
type ListModelCandidatesOptions struct {
CacheAffinityKey string
CacheAffinityPolicy map[string]any
}
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,
SELECT p.id::text, p.platform_key, p.name, p.provider,
COALESCE(NULLIF(p.config->>'specType', ''), NULLIF(cp.provider_type, ''), NULLIF(p.config->>'sourceSpecType', ''), p.provider) AS spec_type,
COALESCE(p.base_url, ''),
p.auth_type, p.credentials, p.config, p.default_pricing_mode,
@ -35,17 +42,27 @@ SELECT p.id::text, p.platform_key, p.name, p.provider,
COALESCE(queued.waiting, 0)::float8,
COALESCE(rpm.used_value, 0)::float8, COALESCE(rpm.reserved_value, 0)::float8,
COALESCE(tpm.used_value, 0)::float8, COALESCE(tpm.reserved_value, 0)::float8,
COALESCE(s.running_count, 0)::float8,
COALESCE(s.waiting_count, 0)::float8,
COALESCE(s.limiter_ratio, 0)::float8,
COALESCE(EXTRACT(EPOCH FROM s.last_assigned_at), 0)::float8
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 || ':' || $2::text || ':' || COALESCE(NULLIF(m.provider_model_name, ''), m.model_name)
COALESCE(s.running_count, 0)::float8,
COALESCE(s.waiting_count, 0)::float8,
COALESCE(s.limiter_ratio, 0)::float8,
COALESCE(EXTRACT(EPOCH FROM s.last_assigned_at), 0)::float8,
COALESCE(ca.request_count, 0)::float8,
COALESCE(ca.input_tokens, 0)::float8,
COALESCE(ca.cached_input_tokens, 0)::float8,
COALESCE(ca.ema_hit_ratio, 0)::float8,
COALESCE(ca.last_hit_ratio, 0)::float8,
COALESCE(EXTRACT(EPOCH FROM ca.last_observed_at), 0)::float8
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 || ':' || $2::text || ':' || COALESCE(NULLIF(m.provider_model_name, ''), m.model_name)
LEFT JOIN gateway_cache_affinity_stats ca
ON ca.client_id = p.platform_key || ':' || $2::text || ':' || COALESCE(NULLIF(m.provider_model_name, ''), m.model_name)
AND ca.cache_affinity_key = NULLIF($4::text, '')
AND ($5::int <= 0 OR ca.last_observed_at >= now() - ($5::int * interval '1 second'))
LEFT JOIN (
SELECT scope_key, SUM(lease_value) AS active
FROM gateway_concurrency_leases
@ -154,11 +171,11 @@ WHERE p.status = 'enabled'
)
)
)
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)
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.CacheAffinityKey, cacheAffinityStaleAfterSeconds(listOptions.CacheAffinityPolicy))
if err != nil {
return nil, err
}
@ -194,6 +211,12 @@ ORDER BY effective_priority ASC,
var stateWaitingCount float64
var stateLimiterRatio float64
var lastAssignedUnix float64
var cacheRequestCount float64
var cacheInputTokens float64
var cacheCachedInputTokens float64
var cacheEMAHitRatio float64
var cacheLastHitRatio float64
var cacheLastObservedUnix float64
if err := rows.Scan(
&item.PlatformID,
&item.PlatformKey,
@ -246,6 +269,12 @@ ORDER BY effective_priority ASC,
&stateWaitingCount,
&stateLimiterRatio,
&lastAssignedUnix,
&cacheRequestCount,
&cacheInputTokens,
&cacheCachedInputTokens,
&cacheEMAHitRatio,
&cacheLastHitRatio,
&cacheLastObservedUnix,
); err != nil {
return nil, err
}
@ -272,6 +301,14 @@ ORDER BY effective_priority ASC,
item.RunningCount = stateRunningCount
item.WaitingCount = maxFloat(queuedWaiting, stateWaitingCount)
item.LastAssignedUnix = lastAssignedUnix
applyRuntimeCandidateCacheAffinity(&item, listOptions, runtimeCandidateCacheAffinityInput{
RequestCount: int(cacheRequestCount),
InputTokens: int(cacheInputTokens),
CachedInputTokens: int(cacheCachedInputTokens),
EMAHitRatio: cacheEMAHitRatio,
LastHitRatio: cacheLastHitRatio,
LastObservedUnix: cacheLastObservedUnix,
})
applyRuntimeCandidateLoad(&item, runtimeCandidateLoadInput{
Policy: effectiveModelRateLimitPolicy(item.PlatformRateLimitPolicy, item.RuntimeRateLimitPolicy, item.RuntimePolicySetID, item.RuntimePolicyOverride, item.ModelRateLimitPolicy),
ConcurrentActive: concurrentActive,
@ -383,6 +420,182 @@ type runtimeCandidateLoadInput struct {
StateLimiterRatio float64
}
type runtimeCandidateCacheAffinityInput struct {
RequestCount int
InputTokens int
CachedInputTokens int
EMAHitRatio float64
LastHitRatio float64
LastObservedUnix float64
}
func normalizeListModelCandidatesOptions(modelType string, options ...ListModelCandidatesOptions) ListModelCandidatesOptions {
if len(options) == 0 {
return ListModelCandidatesOptions{}
}
out := options[0]
out.CacheAffinityKey = strings.TrimSpace(out.CacheAffinityKey)
if out.CacheAffinityKey == "" || !cacheAffinityPolicyEnabled(out.CacheAffinityPolicy, modelType) {
out.CacheAffinityKey = ""
}
return out
}
func applyRuntimeCandidateCacheAffinity(candidate *RuntimeModelCandidate, options ListModelCandidatesOptions, input runtimeCandidateCacheAffinityInput) {
affinity := RuntimeCandidateCacheAffinity{
Key: options.CacheAffinityKey,
RequestCount: input.RequestCount,
InputTokens: input.InputTokens,
CachedInputTokens: input.CachedInputTokens,
EMAHitRatio: boundedRatio(input.EMAHitRatio),
LastHitRatio: boundedRatio(input.LastHitRatio),
LastObservedUnix: input.LastObservedUnix,
AdjustedPriority: float64(candidate.PlatformPriority),
}
minSamples := cacheAffinityMinSamples(options.CacheAffinityPolicy)
hasObservedCachedHit := affinity.CachedInputTokens > 0 || affinity.LastHitRatio > 0
hasSampledAffinity := input.RequestCount >= minSamples && affinity.EMAHitRatio > 0
if options.CacheAffinityKey == "" || input.RequestCount <= 0 || (!hasObservedCachedHit && !hasSampledAffinity) {
candidate.CacheAffinity = affinity
return
}
affinity.Confidence = float64(input.RequestCount) / float64(minSamples)
if affinity.Confidence > 1 {
affinity.Confidence = 1
}
if hasObservedCachedHit && affinity.Confidence < 1 {
affinity.Confidence = 1
}
affinity.Score = runtimeCandidateCacheAffinityScore(affinity)
maxBoost := cacheAffinityMaxPriorityBoost(options.CacheAffinityPolicy)
affinity.Boost = affinity.Score * maxBoost
if affinity.Boost > maxBoost {
affinity.Boost = maxBoost
}
if hasObservedCachedHit || affinity.Score > 0 {
affinity.Applied = true
}
if affinity.Boost > 0 {
affinity.AdjustedPriority = float64(candidate.PlatformPriority) - affinity.Boost
}
candidate.CacheAffinity = affinity
}
func cacheAffinityPolicyEnabled(policy map[string]any, modelType string) bool {
if enabled, ok := policy["enabled"].(bool); ok && !enabled {
return false
}
modelTypes := stringListFromAny(policy["modelTypes"])
if len(modelTypes) == 0 {
modelTypes = []string{"chat", "text_generate", "responses"}
}
modelType = strings.TrimSpace(modelType)
for _, candidate := range modelTypes {
if strings.EqualFold(strings.TrimSpace(candidate), modelType) {
return true
}
}
return false
}
func cacheAffinityMinSamples(policy map[string]any) int {
value := intFromStorePolicy(policy, "minSamples", 3)
if value <= 0 {
return 1
}
return value
}
func cacheAffinityMinInputTokens(policy map[string]any) int {
value := intFromStorePolicy(policy, "minInputTokens", 512)
if value < 0 {
return 0
}
return value
}
func cacheAffinityStaleAfterSeconds(policy map[string]any) int {
return intFromStorePolicy(policy, "staleAfterSeconds", 86400)
}
func cacheAffinityMaxPriorityBoost(policy map[string]any) float64 {
value := floatFromStorePolicy(policy, "maxPriorityBoost", 20)
if value < 0 {
return 0
}
return value
}
func cacheAffinityEMAAlpha(policy map[string]any) float64 {
value := floatFromStorePolicy(policy, "emaAlpha", 0.35)
if value <= 0 || value > 1 {
return 0.35
}
return value
}
func runtimeCandidateCacheAffinityScore(affinity RuntimeCandidateCacheAffinity) float64 {
hitRatio := affinity.EMAHitRatio
if affinity.LastHitRatio > hitRatio {
hitRatio = affinity.LastHitRatio
}
hitRatio = boundedRatio(hitRatio)
confidence := boundedRatio(affinity.Confidence)
if hitRatio <= 0 || confidence <= 0 {
return 0
}
tokenFactor := 0.0
if affinity.CachedInputTokens > 0 {
tokenFactor = float64(affinity.CachedInputTokens) / 8192
if tokenFactor > 1 {
tokenFactor = 1
}
}
return math.Round(hitRatio*confidence*(1+tokenFactor)*1_000_000) / 1_000_000
}
func intFromStorePolicy(policy map[string]any, key string, fallback int) int {
if policy == nil {
return fallback
}
switch value := policy[key].(type) {
case int:
return value
case int64:
return int(value)
case float64:
return int(value)
default:
return fallback
}
}
func floatFromStorePolicy(policy map[string]any, key string, fallback float64) float64 {
if policy == nil {
return fallback
}
switch value := policy[key].(type) {
case int:
return float64(value)
case int64:
return float64(value)
case float64:
return value
default:
return fallback
}
}
func boundedRatio(value float64) float64 {
if value < 0 {
return 0
}
if value > 1 {
return 1
}
return value
}
func applyRuntimeCandidateLoad(candidate *RuntimeModelCandidate, input runtimeCandidateLoadInput) {
rpmLimit := rateLimitForMetric(input.Policy, "rpm")
tpmLimitValue := tpmLimit(input.Policy)
@ -422,6 +635,9 @@ func sortRuntimeModelCandidates(items []RuntimeModelCandidate) {
hasNonFull := false
for index := range items {
items[index].LoadAvoided = false
if !items[index].CacheAffinity.Applied && items[index].CacheAffinity.AdjustedPriority == 0 {
items[index].CacheAffinity.AdjustedPriority = float64(items[index].PlatformPriority)
}
if runtimeCandidateFull(items[index]) {
hasFull = true
} else {
@ -439,8 +655,25 @@ 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
}
if items[i].CacheAffinity.Applied && items[j].CacheAffinity.Applied {
if items[i].CacheAffinity.Score != items[j].CacheAffinity.Score {
return items[i].CacheAffinity.Score > items[j].CacheAffinity.Score
}
if items[i].CacheAffinity.EMAHitRatio != items[j].CacheAffinity.EMAHitRatio {
return items[i].CacheAffinity.EMAHitRatio > items[j].CacheAffinity.EMAHitRatio
}
if items[i].CacheAffinity.LastHitRatio != items[j].CacheAffinity.LastHitRatio {
return items[i].CacheAffinity.LastHitRatio > items[j].CacheAffinity.LastHitRatio
}
if items[i].CacheAffinity.CachedInputTokens != items[j].CacheAffinity.CachedInputTokens {
return items[i].CacheAffinity.CachedInputTokens > items[j].CacheAffinity.CachedInputTokens
}
}
if items[i].CacheAffinity.AdjustedPriority != items[j].CacheAffinity.AdjustedPriority {
return items[i].CacheAffinity.AdjustedPriority < items[j].CacheAffinity.AdjustedPriority
}
if items[i].LoadRatio != items[j].LoadRatio {
return items[i].LoadRatio < items[j].LoadRatio

View File

@ -106,6 +106,156 @@ func TestRuntimeCandidateSortingAvoidsFullCandidatesButKeepsFallback(t *testing.
}
}
func TestRuntimeCandidateSortingPrefersCacheAffinityWithinAvailableCandidates(t *testing.T) {
candidates := []RuntimeModelCandidate{
{
PlatformID: "base-priority",
PlatformPriority: 10,
},
{
PlatformID: "cache-affinity",
PlatformPriority: 20,
},
}
applyRuntimeCandidateCacheAffinity(&candidates[0], ListModelCandidatesOptions{
CacheAffinityKey: "explicit:test",
CacheAffinityPolicy: map[string]any{"enabled": true, "minSamples": 3, "maxPriorityBoost": 20, "modelTypes": []any{"text_generate"}},
}, runtimeCandidateCacheAffinityInput{
RequestCount: 1,
EMAHitRatio: 1,
})
applyRuntimeCandidateCacheAffinity(&candidates[1], ListModelCandidatesOptions{
CacheAffinityKey: "explicit:test",
CacheAffinityPolicy: map[string]any{"enabled": true, "minSamples": 3, "maxPriorityBoost": 20, "modelTypes": []any{"text_generate"}},
}, runtimeCandidateCacheAffinityInput{
RequestCount: 3,
EMAHitRatio: 0.9,
})
sortRuntimeModelCandidates(candidates)
if candidates[0].PlatformID != "cache-affinity" || !candidates[0].CacheAffinity.Applied {
t.Fatalf("expected cache affinity candidate first, got %+v", candidates)
}
if candidates[1].PlatformID != "base-priority" || candidates[1].CacheAffinity.Applied {
t.Fatalf("expected low sample candidate to keep base priority, got %+v", candidates)
}
}
func TestRuntimeCandidateSortingPromotesObservedCacheHitAbovePriority(t *testing.T) {
candidates := []RuntimeModelCandidate{
{
PlatformID: "volces-priority",
PlatformPriority: 1,
},
{
PlatformID: "deepseek-cache-hit",
PlatformPriority: 100,
},
}
policy := map[string]any{"enabled": true, "minSamples": 3, "maxPriorityBoost": 20, "modelTypes": []any{"text_generate"}}
applyRuntimeCandidateCacheAffinity(&candidates[0], ListModelCandidatesOptions{
CacheAffinityKey: "conversation:test",
CacheAffinityPolicy: policy,
}, runtimeCandidateCacheAffinityInput{})
applyRuntimeCandidateCacheAffinity(&candidates[1], ListModelCandidatesOptions{
CacheAffinityKey: "conversation:test",
CacheAffinityPolicy: policy,
}, runtimeCandidateCacheAffinityInput{
RequestCount: 1,
InputTokens: 7945,
CachedInputTokens: 7296,
EMAHitRatio: 0.918,
LastHitRatio: 0.918,
})
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].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[1].PlatformID != "volces-priority" || candidates[1].CacheAffinity.Applied {
t.Fatalf("expected priority-only candidate second, got %+v", candidates)
}
}
func TestRuntimeCandidateSortingOrdersMultipleCacheAffinityCandidatesByScore(t *testing.T) {
candidates := []RuntimeModelCandidate{
{
PlatformID: "lower-cache-score",
PlatformPriority: 1,
},
{
PlatformID: "higher-cache-score",
PlatformPriority: 50,
},
}
policy := map[string]any{"enabled": true, "minSamples": 1, "maxPriorityBoost": 20, "modelTypes": []any{"text_generate"}}
applyRuntimeCandidateCacheAffinity(&candidates[0], ListModelCandidatesOptions{
CacheAffinityKey: "conversation:test",
CacheAffinityPolicy: policy,
}, runtimeCandidateCacheAffinityInput{
RequestCount: 2,
InputTokens: 8000,
CachedInputTokens: 1000,
EMAHitRatio: 0.5,
LastHitRatio: 0.5,
})
applyRuntimeCandidateCacheAffinity(&candidates[1], ListModelCandidatesOptions{
CacheAffinityKey: "conversation:test",
CacheAffinityPolicy: policy,
}, runtimeCandidateCacheAffinityInput{
RequestCount: 1,
InputTokens: 8000,
CachedInputTokens: 7000,
EMAHitRatio: 0.75,
LastHitRatio: 0.75,
})
sortRuntimeModelCandidates(candidates)
if candidates[0].PlatformID != "higher-cache-score" {
t.Fatalf("expected higher cache score candidate first, got %+v", candidates)
}
if candidates[0].CacheAffinity.Score <= candidates[1].CacheAffinity.Score {
t.Fatalf("expected score to drive cache-affinity order, got %+v", candidates)
}
}
func TestRuntimeCandidateSortingKeepsFullCacheAffinityCandidateAvoided(t *testing.T) {
candidates := []RuntimeModelCandidate{
{
PlatformID: "cache-affinity-full",
PlatformPriority: 50,
LoadLimited: true,
LoadRatio: 1.2,
},
{
PlatformID: "available",
PlatformPriority: 10,
},
}
applyRuntimeCandidateCacheAffinity(&candidates[0], ListModelCandidatesOptions{
CacheAffinityKey: "explicit:test",
CacheAffinityPolicy: map[string]any{"enabled": true, "minSamples": 1, "maxPriorityBoost": 100, "modelTypes": []any{"text_generate"}},
}, runtimeCandidateCacheAffinityInput{
RequestCount: 10,
EMAHitRatio: 1,
})
sortRuntimeModelCandidates(candidates)
if candidates[0].PlatformID != "available" {
t.Fatalf("expected available candidate before full cache-affinity candidate, got %+v", candidates)
}
if candidates[1].PlatformID != "cache-affinity-full" || !candidates[1].LoadAvoided {
t.Fatalf("expected full cache-affinity candidate to remain avoided fallback, got %+v", candidates)
}
}
func TestDefaultRunnerPriorityDemotePolicyUsesAutoMode(t *testing.T) {
policy := defaultRunnerPriorityDemotePolicy()
@ -115,4 +265,190 @@ func TestDefaultRunnerPriorityDemotePolicyUsesAutoMode(t *testing.T) {
if policy["enabled"] != true {
t.Fatalf("expected default priority demotion to stay enabled, got %+v", policy["enabled"])
}
if policy["deprecated"] != true || policy["replacedBy"] != "failoverPolicy.actionRules" {
t.Fatalf("expected legacy priority demotion policy to be marked deprecated, got %+v", policy)
}
}
func TestDefaultRunnerFailoverPolicyCarriesComprehensiveDefaultsInActionRules(t *testing.T) {
policy := defaultRunnerFailoverPolicy()
if policy["legacyFieldsDeprecated"] != true || policy["replacedBy"] != "actionRules" {
t.Fatalf("expected legacy failover summary fields to be marked deprecated, got %+v", policy)
}
rules, ok := policy["actionRules"].([]any)
if !ok || len(rules) < 5 {
t.Fatalf("expected default action rules to be present, got %+v", policy["actionRules"])
}
cooldownRule := defaultActionRuleByAction(t, rules, "cooldown_and_next")
for _, code := range []string{"rate_limit", "too_many_requests", "resource_exhausted", "throttled", "quota_exceeded", "rpm_limit", "tpm_limit"} {
if !actionRuleHasString(cooldownRule, "codes", code) {
t.Fatalf("expected cooldown rule to include rate-limit code %q, got %+v", code, cooldownRule["codes"])
}
}
for _, keyword := range []string{"rate limit", "too many requests", "resource exhausted", "request limit", "requests per minute", "tokens per minute"} {
if !actionRuleHasString(cooldownRule, "keywords", keyword) {
t.Fatalf("expected cooldown rule to include rate-limit keyword %q, got %+v", keyword, cooldownRule["keywords"])
}
}
if !actionRuleHasInt(cooldownRule, "statusCodes", 429) {
t.Fatalf("expected cooldown rule to include 429, got %+v", cooldownRule["statusCodes"])
}
disableRule := defaultActionRuleByAction(t, rules, "disable_and_next")
for _, code := range []string{"auth_failed", "invalid_api_key", "missing_credentials", "account_disabled", "billing_not_active", "permission_denied"} {
if !actionRuleHasString(disableRule, "codes", code) {
t.Fatalf("expected disable rule to include auth/platform code %q, got %+v", code, disableRule["codes"])
}
}
for _, keyword := range []string{"invalid api key", "unauthorized", "credential", "permission denied", "account disabled", "billing_not_active", "余额不足"} {
if !actionRuleHasString(disableRule, "keywords", keyword) {
t.Fatalf("expected disable rule to include auth/platform keyword %q, got %+v", keyword, disableRule["keywords"])
}
}
for _, status := range []int{401, 403} {
if !actionRuleHasInt(disableRule, "statusCodes", status) {
t.Fatalf("expected disable rule to include status %d, got %+v", status, disableRule["statusCodes"])
}
}
demoteRule := defaultActionRuleByAction(t, rules, "demote_and_next")
for _, category := range []string{"timeout", "provider_5xx", "provider_overloaded"} {
if !actionRuleHasString(demoteRule, "categories", category) {
t.Fatalf("expected demote rule to include transient category %q, got %+v", category, demoteRule["categories"])
}
}
for _, code := range []string{"timeout", "server_error", "overloaded", "provider_failed", "invalid_response", "script_timeout", "upload_failed"} {
if !actionRuleHasString(demoteRule, "codes", code) {
t.Fatalf("expected demote rule to include transient code %q, got %+v", code, demoteRule["codes"])
}
}
for _, status := range []int{408, 500, 502, 503, 504, 520, 522, 524, 529, 530} {
if !actionRuleHasInt(demoteRule, "statusCodes", status) {
t.Fatalf("expected demote rule to include transient status %d, got %+v", status, demoteRule["statusCodes"])
}
}
for _, keyword := range []string{"deadline exceeded", "temporarily unavailable", "service unavailable", "bad gateway", "upstream timeout", "try again later"} {
if !actionRuleHasString(demoteRule, "keywords", keyword) {
t.Fatalf("expected demote rule to include transient keyword %q, got %+v", keyword, demoteRule["keywords"])
}
}
nextRule := defaultActionRuleByAction(t, rules, "next")
for _, category := range []string{"network", "stream_error"} {
if !actionRuleHasString(nextRule, "categories", category) {
t.Fatalf("expected next rule to include network category %q, got %+v", category, nextRule["categories"])
}
}
for _, code := range []string{"network", "stream_read_error", "upload_network", "upload_config_failed", "upload_source_fetch_failed"} {
if !actionRuleHasString(nextRule, "codes", code) {
t.Fatalf("expected next rule to include network code %q, got %+v", code, nextRule["codes"])
}
}
for _, keyword := range []string{"socket hang up", "connection reset", "econnreset", "enotfound", "unexpected eof", "stream error", "fetch failed"} {
if !actionRuleHasString(nextRule, "keywords", keyword) {
t.Fatalf("expected next rule to include network keyword %q, got %+v", keyword, nextRule["keywords"])
}
}
stopRule := defaultActionRuleByAction(t, rules, "stop")
for _, code := range []string{"bad_request", "invalid_request", "invalid_parameter", "missing_required", "unsupported_kind", "unsupported_model", "request_asset_input_format_unsupported", "upload_source_too_large", "invalid_multipart_image", "script_error", "insufficient_balance", "permission_denied"} {
if !actionRuleHasString(stopRule, "codes", code) {
t.Fatalf("expected stop rule to include legacy hard-stop code %q, got %+v", code, stopRule["codes"])
}
}
for _, status := range []int{400, 404, 405, 409, 413, 415, 422} {
if !actionRuleHasInt(stopRule, "statusCodes", status) {
t.Fatalf("expected stop rule to include request-error status %d, got %+v", status, stopRule["statusCodes"])
}
}
for _, keyword := range []string{"invalid parameter", "missing required", "not supported", "image is required", "too large", "not found"} {
if !actionRuleHasString(stopRule, "keywords", keyword) {
t.Fatalf("expected stop rule to include hard-stop keyword %q, got %+v", keyword, stopRule["keywords"])
}
}
}
func TestLegacyDefaultRunnerActionRulesAreDetectedForUpgrade(t *testing.T) {
legacy := []any{
map[string]any{
"enabled": true,
"categories": []any{"network", "timeout", "stream_error", "provider_5xx", "provider_overloaded"},
"action": "next",
},
map[string]any{
"enabled": true,
"categories": []any{"rate_limit"},
"codes": []any{"rate_limit"},
"statusCodes": []any{429},
"keywords": []any{"rate_limit", "429", "too many requests"},
"action": "cooldown_and_next",
"target": "model",
"cooldownSeconds": 300,
},
map[string]any{
"enabled": true,
"keywords": []any{"rate limit", "too many requests", "temporarily_unavailable"},
"action": "demote_and_next",
"target": "platform",
"demoteSteps": 1,
},
map[string]any{
"enabled": true,
"categories": []any{"auth_error"},
"codes": []any{"auth_failed", "invalid_api_key", "missing_credentials"},
"statusCodes": []any{401, 403},
"keywords": []any{"invalid api key", "api key is invalid", "unauthorized", "forbidden", "permission denied"},
"action": "disable_and_next",
"target": "platform",
},
map[string]any{
"enabled": true,
"categories": []any{"request_error", "unsupported_model", "user_permission", "insufficient_balance"},
"codes": []any{"bad_request", "invalid_request", "invalid_parameter", "missing_required", "unsupported_kind", "unsupported_model", "insufficient_balance", "permission_denied"},
"keywords": []any{"invalid_parameter", "missing required", "bad request", "insufficient balance"},
"action": "stop",
},
}
if !runnerActionRulesLookLikeLegacyDefaults(legacy) {
t.Fatal("expected old compact default action rules to be detected")
}
if runnerActionRulesLookLikeLegacyDefaults(defaultRunnerFailoverPolicy()["actionRules"]) {
t.Fatal("enhanced defaults should not be treated as legacy")
}
}
func defaultActionRuleByAction(t *testing.T, rules []any, action string) map[string]any {
t.Helper()
for _, raw := range rules {
rule, ok := raw.(map[string]any)
if ok && rule["action"] == action {
return rule
}
}
t.Fatalf("expected action rule %q in %+v", action, rules)
return nil
}
func actionRuleHasString(rule map[string]any, key string, want string) bool {
for _, raw := range ruleSlice(rule, key) {
if raw == want {
return true
}
}
return false
}
func actionRuleHasInt(rule map[string]any, key string, want int) bool {
for _, raw := range ruleSlice(rule, key) {
if raw == want {
return true
}
}
return false
}
func ruleSlice(rule map[string]any, key string) []any {
values, _ := rule[key].([]any)
return values
}

View File

@ -278,6 +278,7 @@ type RunnerPolicy struct {
HardStopPolicy map[string]any `json:"hardStopPolicy,omitempty"`
PriorityDemotePolicy map[string]any `json:"priorityDemotePolicy,omitempty"`
SingleSourcePolicy map[string]any `json:"singleSourcePolicy,omitempty"`
CacheAffinityPolicy map[string]any `json:"cacheAffinityPolicy,omitempty"`
Metadata map[string]any `json:"metadata,omitempty"`
Status string `json:"status"`
CreatedAt time.Time `json:"createdAt"`

View File

@ -3,6 +3,7 @@ package store
import (
"context"
"encoding/json"
"fmt"
"strings"
"time"
@ -11,7 +12,7 @@ import (
const runnerPolicyColumns = `
id::text, policy_key, name, COALESCE(description, ''), failover_policy,
hard_stop_policy, priority_demote_policy, single_source_policy, metadata, status, created_at, updated_at`
hard_stop_policy, priority_demote_policy, single_source_policy, cache_affinity_policy, metadata, status, created_at, updated_at`
type RunnerPolicyInput struct {
PolicyKey string `json:"policyKey"`
@ -21,6 +22,7 @@ type RunnerPolicyInput struct {
HardStopPolicy map[string]any `json:"hardStopPolicy"`
PriorityDemotePolicy map[string]any `json:"priorityDemotePolicy"`
SingleSourcePolicy map[string]any `json:"singleSourcePolicy"`
CacheAffinityPolicy map[string]any `json:"cacheAffinityPolicy"`
Metadata map[string]any `json:"metadata"`
Status string `json:"status"`
}
@ -52,24 +54,26 @@ func (s *Store) UpsertDefaultRunnerPolicy(ctx context.Context, input RunnerPolic
hardStopPolicy, _ := json.Marshal(emptyObjectIfNil(input.HardStopPolicy))
priorityDemotePolicy, _ := json.Marshal(emptyObjectIfNil(input.PriorityDemotePolicy))
singleSourcePolicy, _ := json.Marshal(emptyObjectIfNil(input.SingleSourcePolicy))
cacheAffinityPolicy, _ := json.Marshal(emptyObjectIfNil(input.CacheAffinityPolicy))
metadata, _ := json.Marshal(emptyObjectIfNil(input.Metadata))
return scanRunnerPolicy(s.pool.QueryRow(ctx, `
INSERT INTO gateway_runner_policies (
policy_key, name, description, failover_policy, hard_stop_policy, priority_demote_policy, single_source_policy, metadata, status
)
VALUES ($1, $2, NULLIF($3, ''), $4, $5, $6, $7, $8, $9)
ON CONFLICT (policy_key) DO UPDATE
SET name = EXCLUDED.name,
description = EXCLUDED.description,
failover_policy = EXCLUDED.failover_policy,
hard_stop_policy = EXCLUDED.hard_stop_policy,
priority_demote_policy = EXCLUDED.priority_demote_policy,
single_source_policy = EXCLUDED.single_source_policy,
metadata = EXCLUDED.metadata,
status = EXCLUDED.status,
updated_at = now()
RETURNING `+runnerPolicyColumns,
input.PolicyKey, input.Name, input.Description, failoverPolicy, hardStopPolicy, priorityDemotePolicy, singleSourcePolicy, metadata, input.Status,
INSERT INTO gateway_runner_policies (
policy_key, name, description, failover_policy, hard_stop_policy, priority_demote_policy, single_source_policy, cache_affinity_policy, metadata, status
)
VALUES ($1, $2, NULLIF($3, ''), $4, $5, $6, $7, $8, $9, $10)
ON CONFLICT (policy_key) DO UPDATE
SET name = EXCLUDED.name,
description = EXCLUDED.description,
failover_policy = EXCLUDED.failover_policy,
hard_stop_policy = EXCLUDED.hard_stop_policy,
priority_demote_policy = EXCLUDED.priority_demote_policy,
single_source_policy = EXCLUDED.single_source_policy,
cache_affinity_policy = EXCLUDED.cache_affinity_policy,
metadata = EXCLUDED.metadata,
status = EXCLUDED.status,
updated_at = now()
RETURNING `+runnerPolicyColumns,
input.PolicyKey, input.Name, input.Description, failoverPolicy, hardStopPolicy, priorityDemotePolicy, singleSourcePolicy, cacheAffinityPolicy, metadata, input.Status,
))
}
@ -79,6 +83,7 @@ func scanRunnerPolicy(scanner runnerPolicyScanner) (RunnerPolicy, error) {
var hardStopPolicy []byte
var priorityDemotePolicy []byte
var singleSourcePolicy []byte
var cacheAffinityPolicy []byte
var metadata []byte
if err := scanner.Scan(
&item.ID,
@ -89,6 +94,7 @@ func scanRunnerPolicy(scanner runnerPolicyScanner) (RunnerPolicy, error) {
&hardStopPolicy,
&priorityDemotePolicy,
&singleSourcePolicy,
&cacheAffinityPolicy,
&metadata,
&item.Status,
&item.CreatedAt,
@ -100,7 +106,17 @@ func scanRunnerPolicy(scanner runnerPolicyScanner) (RunnerPolicy, error) {
item.HardStopPolicy = decodeObject(hardStopPolicy)
item.PriorityDemotePolicy = decodeObject(priorityDemotePolicy)
item.SingleSourcePolicy = decodeObject(singleSourcePolicy)
item.CacheAffinityPolicy = decodeObject(cacheAffinityPolicy)
item.Metadata = decodeObject(metadata)
if len(item.CacheAffinityPolicy) == 0 {
item.CacheAffinityPolicy = defaultRunnerCacheAffinityPolicy()
}
item.FailoverPolicy = markRunnerFailoverLegacyFieldsDeprecated(item.FailoverPolicy)
item.HardStopPolicy = markRunnerLegacyPolicyDeprecated(item.HardStopPolicy)
item.PriorityDemotePolicy = markRunnerLegacyPolicyDeprecated(item.PriorityDemotePolicy)
if item.PolicyKey == "default-runner-v1" && runnerActionRulesLookLikeLegacyDefaults(item.FailoverPolicy["actionRules"]) {
item.FailoverPolicy["actionRules"] = defaultRunnerFailoverPolicy()["actionRules"]
}
return item, nil
}
@ -118,9 +134,155 @@ func normalizeRunnerPolicyInput(input RunnerPolicyInput) RunnerPolicyInput {
if input.Status == "" {
input.Status = "active"
}
if len(input.CacheAffinityPolicy) == 0 {
input.CacheAffinityPolicy = defaultRunnerCacheAffinityPolicy()
}
input.FailoverPolicy = markRunnerFailoverLegacyFieldsDeprecated(input.FailoverPolicy)
input.HardStopPolicy = markRunnerLegacyPolicyDeprecated(input.HardStopPolicy)
input.PriorityDemotePolicy = markRunnerLegacyPolicyDeprecated(input.PriorityDemotePolicy)
return input
}
func markRunnerLegacyPolicyDeprecated(policy map[string]any) map[string]any {
if policy == nil {
policy = map[string]any{}
}
policy["deprecated"] = true
policy["replacedBy"] = "failoverPolicy.actionRules"
return policy
}
func markRunnerFailoverLegacyFieldsDeprecated(policy map[string]any) map[string]any {
if policy == nil {
policy = map[string]any{}
}
policy["legacyFieldsDeprecated"] = true
policy["deprecatedFields"] = []any{"allowCategories", "denyCategories", "allowCodes", "denyCodes", "allowKeywords", "denyKeywords", "allowStatusCodes", "denyStatusCodes", "actions"}
policy["replacedBy"] = "actionRules"
return policy
}
func runnerActionRulesLookLikeLegacyDefaults(raw any) bool {
rawRules, ok := raw.([]any)
if !ok || len(rawRules) != 5 {
return false
}
byAction := make(map[string]map[string]any, len(rawRules))
for _, rawRule := range rawRules {
rule, ok := rawRule.(map[string]any)
if !ok {
return false
}
action := strings.TrimSpace(fmt.Sprint(rule["action"]))
if action == "" {
action = "next"
}
if _, exists := byAction[action]; exists {
return false
}
byAction[action] = rule
}
next, nextOK := byAction["next"]
cooldown, cooldownOK := byAction["cooldown_and_next"]
demote, demoteOK := byAction["demote_and_next"]
disable, disableOK := byAction["disable_and_next"]
stop, stopOK := byAction["stop"]
if !nextOK || !cooldownOK || !demoteOK || !disableOK || !stopOK {
return false
}
if runnerRulesHaveEnhancedDefaultTerms(rawRules) {
return false
}
return runnerRuleHasAll(cooldown, "categories", "rate_limit") &&
runnerRuleHasAll(cooldown, "codes", "rate_limit") &&
runnerRuleHasAll(cooldown, "statusCodes", "429") &&
len(runnerRuleStrings(cooldown, "codes")) <= 1 &&
len(runnerRuleStrings(cooldown, "keywords")) <= 3 &&
runnerRuleHasAll(disable, "categories", "auth_error") &&
runnerRuleHasAll(disable, "codes", "auth_failed", "invalid_api_key", "missing_credentials") &&
runnerRuleHasAll(disable, "statusCodes", "401", "403") &&
len(runnerRuleStrings(disable, "codes")) <= 3 &&
len(runnerRuleStrings(disable, "keywords")) <= 8 &&
runnerRuleHasAll(stop, "categories", "request_error", "unsupported_model", "user_permission", "insufficient_balance") &&
len(runnerRuleStrings(stop, "codes")) <= 8 &&
len(runnerRuleStrings(stop, "keywords")) <= 4 &&
len(runnerRuleStrings(next, "categories")) >= 5 &&
len(runnerRuleStrings(next, "codes")) <= 6 &&
len(runnerRuleStrings(next, "keywords")) <= 13 &&
(runnerRuleHasAll(demote, "keywords", "temporarily_unavailable") ||
runnerRuleHasAll(demote, "categories", "provider_5xx") ||
runnerRuleHasAll(demote, "codes", "server_error")) &&
len(runnerRuleStrings(demote, "codes")) <= 6 &&
len(runnerRuleStrings(demote, "keywords")) <= 8
}
func runnerRulesHaveEnhancedDefaultTerms(rawRules []any) bool {
markers := map[string]struct{}{
"resource_exhausted": {},
"rpm_limit": {},
"tpm_limit": {},
"account_disabled": {},
"billing_not_active": {},
"provider_failed": {},
"invalid_response": {},
"request_asset_fetch_failed": {},
"socket hang up": {},
"econnreset": {},
"request_asset_input_format_unsupported": {},
"invalid_multipart_image": {},
}
for _, rawRule := range rawRules {
rule, ok := rawRule.(map[string]any)
if !ok {
continue
}
for _, field := range []string{"categories", "codes", "statusCodes", "keywords"} {
for _, value := range runnerRuleStrings(rule, field) {
if _, ok := markers[value]; ok {
return true
}
}
}
}
return false
}
func runnerRuleHasAll(rule map[string]any, key string, values ...string) bool {
actual := make(map[string]struct{}, len(values))
for _, value := range runnerRuleStrings(rule, key) {
actual[value] = struct{}{}
}
for _, value := range values {
if _, ok := actual[strings.ToLower(strings.TrimSpace(value))]; !ok {
return false
}
}
return true
}
func runnerRuleStrings(rule map[string]any, key string) []string {
raw, ok := rule[key].([]any)
if !ok {
if typed, ok := rule[key].([]string); ok {
out := make([]string, 0, len(typed))
for _, item := range typed {
if value := strings.ToLower(strings.TrimSpace(item)); value != "" {
out = append(out, value)
}
}
return out
}
return nil
}
out := make([]string, 0, len(raw))
for _, item := range raw {
if value := strings.ToLower(strings.TrimSpace(fmt.Sprint(item))); value != "" {
out = append(out, value)
}
}
return out
}
func defaultRunnerPolicy() RunnerPolicy {
now := time.Now()
return RunnerPolicy{
@ -131,6 +293,7 @@ func defaultRunnerPolicy() RunnerPolicy {
HardStopPolicy: defaultRunnerHardStopPolicy(),
PriorityDemotePolicy: defaultRunnerPriorityDemotePolicy(),
SingleSourcePolicy: defaultRunnerSingleSourcePolicy(),
CacheAffinityPolicy: defaultRunnerCacheAffinityPolicy(),
Metadata: map[string]any{"source": "code-default"},
Status: "active",
CreatedAt: now,
@ -141,6 +304,8 @@ func defaultRunnerPolicy() RunnerPolicy {
func defaultRunnerPriorityDemotePolicy() map[string]any {
return map[string]any{
"enabled": true,
"deprecated": true,
"replacedBy": "failoverPolicy.actionRules",
"categories": []any{"network", "timeout", "stream_error", "rate_limit", "provider_5xx", "provider_overloaded"},
"codes": []any{"network", "timeout", "stream_read_error", "rate_limit", "server_error", "overloaded"},
"statusCodes": []any{408, 429, 500, 502, 503, 504},
@ -154,24 +319,88 @@ func defaultRunnerSingleSourcePolicy() map[string]any {
}
}
func defaultRunnerCacheAffinityPolicy() map[string]any {
return map[string]any{
"enabled": true,
"modelTypes": []any{"chat", "text_generate", "responses"},
"minSamples": 3,
"minInputTokens": 512,
"staleAfterSeconds": 86400,
"maxPriorityBoost": 20,
"emaAlpha": 0.35,
}
}
func defaultRunnerFailoverPolicy() map[string]any {
return map[string]any{
"enabled": true,
"maxPlatforms": 99,
"maxDurationSeconds": 600,
"allowCategories": []any{"network", "timeout", "stream_error", "rate_limit", "provider_5xx", "provider_overloaded", "auth_error"},
"denyCategories": []any{"request_error", "unsupported_model", "user_permission", "insufficient_balance"},
"allowCodes": []any{"auth_failed", "invalid_api_key", "missing_credentials"},
"allowKeywords": []any{"timeout", "network", "rate_limit", "overloaded", "temporarily_unavailable", "server_error", "auth_failed", "invalid_api_key", "missing_credentials", "unauthorized", "forbidden", "429", "5xx"},
"denyKeywords": []any{"invalid_parameter", "missing required", "bad request"},
"allowStatusCodes": []any{401, 403, 408, 429, 500, 502, 503, 504},
"denyStatusCodes": []any{},
"enabled": true,
"maxPlatforms": 99,
"maxDurationSeconds": 600,
"legacyFieldsDeprecated": true,
"deprecatedFields": []any{"allowCategories", "denyCategories", "allowCodes", "denyCodes", "allowKeywords", "denyKeywords", "allowStatusCodes", "denyStatusCodes", "actions"},
"replacedBy": "actionRules",
"allowCategories": []any{"network", "timeout", "stream_error", "rate_limit", "provider_5xx", "provider_overloaded", "auth_error"},
"denyCategories": []any{"request_error", "unsupported_model", "user_permission", "insufficient_balance"},
"allowCodes": []any{"auth_failed", "invalid_api_key", "missing_credentials"},
"allowKeywords": []any{"timeout", "network", "rate_limit", "overloaded", "temporarily_unavailable", "server_error", "auth_failed", "invalid_api_key", "missing_credentials", "unauthorized", "forbidden", "429", "5xx"},
"denyKeywords": []any{"invalid_parameter", "missing required", "bad request"},
"allowStatusCodes": []any{401, 403, 408, 429, 500, 502, 503, 504},
"denyStatusCodes": []any{},
"actionRules": []any{
map[string]any{
"enabled": true,
"categories": []any{"rate_limit"},
"codes": []any{"rate_limit", "too_many_requests", "too_many_request", "rate_limit_exceeded", "requests_rate_limited", "requests_limit_exceeded", "resource_exhausted", "throttled", "quota_exceeded", "quota_exhausted", "qps_limit", "rpm_limit", "tpm_limit"},
"statusCodes": []any{429},
"keywords": []any{"rate limit", "rate_limit", "rate-limit", "too many requests", "too_many_requests", "429", "throttle", "throttled", "resource exhausted", "quota exceeded", "quota_exceeded", "rate exceeded", "request limit", "requests per minute", "tokens per minute", "qps", "rpm", "tpm"},
"action": "cooldown_and_next",
"target": "model",
"cooldownSeconds": 300,
},
map[string]any{
"enabled": true,
"categories": []any{"auth_error"},
"codes": []any{"auth_failed", "invalid_api_key", "missing_credentials", "missing_credential", "unauthorized", "forbidden", "account_disabled", "billing_not_active", "permission_denied"},
"statusCodes": []any{401, 403},
"keywords": []any{"invalid api key", "api key is invalid", "invalid_api_key", "apikey invalid", "api key invalid", "unauthorized", "forbidden", "authentication", "auth failed", "credential", "missing credential", "permission denied", "access denied", "invalid signature", "signature mismatch", "secret key", "access key", "account disabled", "account_deactivated", "billing not active", "billing_not_active", "insufficient balance", "quota exceeded", "余额不足", "欠费"},
"action": "disable_and_next",
"target": "platform",
},
map[string]any{
"enabled": true,
"categories": []any{"timeout", "provider_5xx", "provider_overloaded"},
"codes": []any{"timeout", "server_error", "overloaded", "provider_failed", "invalid_response", "script_timeout", "request_asset_fetch_failed", "upload_failed", "upload_read_failed", "upload_source_read_failed"},
"statusCodes": []any{408, 500, 502, 503, 504, 520, 521, 522, 523, 524, 529, 530},
"keywords": []any{"timeout", "timed out", "deadline exceeded", "context deadline exceeded", "overloaded", "overload", "temporarily_unavailable", "temporarily unavailable", "temporary unavailable", "service unavailable", "server error", "internal server error", "bad gateway", "gateway timeout", "upstream timeout", "upstream error", "provider failed", "model is overloaded", "capacity", "try again later", "please try again later", "5xx", "500", "502", "503", "504"},
"action": "demote_and_next",
"target": "platform",
"demoteSteps": 1,
},
map[string]any{
"enabled": true,
"categories": []any{"network", "stream_error"},
"codes": []any{"network", "stream_read_error", "cancelled", "request_asset_fetch_failed", "upload_network", "upload_config_failed", "local_static_store_failed", "upload_source_fetch_failed"},
"statusCodes": []any{},
"keywords": []any{"network", "socket hang up", "socket closed", "connection reset", "connection refused", "connection aborted", "broken pipe", "econnreset", "econnrefused", "eai_again", "enotfound", "no such host", "dns", "tls handshake timeout", "unexpected eof", "connection closed", "stream error", "stream_read_error", "read failed", "fetch failed"},
"action": "next",
},
map[string]any{
"enabled": true,
"categories": []any{"request_error", "unsupported_model", "user_permission", "insufficient_balance"},
"codes": []any{"bad_request", "invalid_request", "invalid_parameter", "missing_required", "unsupported_kind", "unsupported_model", "unsupported_operation", "unsupported_request_resolution", "request_asset_input_format_unsupported", "request_asset_expired", "request_asset_decode_failed", "request_asset_public_url_required", "upload_no_channel", "upload_unsupported_channel", "upload_source_too_large", "upload_source_invalid_url", "upload_source_unsupported_url", "upload_decode_failed", "invalid_proxy", "invalid_json_body", "invalid_multipart_body", "invalid_multipart_file", "invalid_multipart_image", "invalid_multipart_audio", "missing_configuration", "script_error", "cloned_voice_not_found", "cloned_voice_expired", "cloned_voice_unavailable", "insufficient_balance", "permission_denied"},
"statusCodes": []any{400, 404, 405, 406, 409, 410, 411, 413, 415, 422},
"keywords": []any{"invalid_parameter", "invalid parameter", "missing required", "missing_required", "bad request", "unsupported", "not supported", "does not support", "is not supported", "insufficient balance", "permission denied", "no permission", "invalid json", "invalid multipart", "image is required", "audio must", "prompt is required", "model is required", "exceeds", "too large", "expired", "not found", "required"},
"action": "stop",
},
},
}
}
func defaultRunnerHardStopPolicy() map[string]any {
return map[string]any{
"enabled": true,
"deprecated": true,
"replacedBy": "failoverPolicy.actionRules",
"categories": []any{"request_error", "unsupported_model", "user_permission", "insufficient_balance"},
"codes": []any{"bad_request", "invalid_request", "invalid_parameter", "missing_required", "unsupported_kind", "unsupported_model", "insufficient_balance", "permission_denied"},
"statusCodes": []any{},

View File

@ -129,6 +129,18 @@ WHERE id = $1::uuid`, platformID)
return err
}
func (s *Store) DisableCandidatePlatformModel(ctx context.Context, platformModelID string) error {
if strings.TrimSpace(platformModelID) == "" {
return nil
}
_, err := s.pool.Exec(ctx, `
UPDATE platform_models
SET enabled = false,
updated_at = now()
WHERE id = $1::uuid`, platformModelID)
return err
}
func (s *Store) CooldownCandidatePlatform(ctx context.Context, platformID string, cooldownSeconds int) error {
if strings.TrimSpace(platformID) == "" {
return nil
@ -160,9 +172,16 @@ WHERE id = $1::uuid`, platformModelID, cooldownSeconds)
}
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)
}
func (s *Store) DemoteCandidatePlatformPriorityBySteps(ctx context.Context, platformID string, platformModelID string, requestedModel string, modelType string, demoteSteps int) (int, error) {
if strings.TrimSpace(platformID) == "" || strings.TrimSpace(platformModelID) == "" || strings.TrimSpace(requestedModel) == "" {
return 0, nil
}
if demoteSteps <= 0 {
demoteSteps = 1
}
var dynamicPriority int
err := s.pool.QueryRow(ctx, `
WITH current_model AS (
@ -171,10 +190,13 @@ func (s *Store) DemoteCandidatePlatformPriority(ctx context.Context, platformID
WHERE id = $2::uuid
AND platform_id = $1::uuid
),
peer_priority AS (
SELECT MAX(COALESCE(peer.dynamic_priority, peer.priority)) AS max_priority
eligible_raw AS (
SELECT peer.id AS platform_id,
peer.priority,
COALESCE(peer.dynamic_priority, peer.priority) AS effective_priority,
peer.created_at
FROM current_model current
JOIN platform_models peer_model ON peer_model.platform_id <> current.platform_id
JOIN platform_models peer_model ON TRUE
JOIN integration_platforms peer ON peer.id = peer_model.platform_id
LEFT JOIN base_model_catalog peer_base ON peer_base.id = peer_model.base_model_id
WHERE peer.status = 'enabled'
@ -194,14 +216,43 @@ func (s *Store) DemoteCandidatePlatformPriority(ctx context.Context, platformID
)
)
)
),
eligible AS (
SELECT DISTINCT ON (platform_id) platform_id, priority, effective_priority, created_at
FROM eligible_raw
ORDER BY platform_id, effective_priority ASC, priority ASC, created_at ASC
),
ordered AS (
SELECT platform_id, effective_priority,
row_number() OVER (ORDER BY effective_priority ASC, priority ASC, created_at ASC, platform_id ASC) AS rn
FROM eligible
),
current_position AS (
SELECT rn, effective_priority
FROM ordered
WHERE platform_id = $1::uuid
),
next_position AS (
SELECT ordered.effective_priority
FROM ordered, current_position
WHERE ordered.rn >= current_position.rn + $5::int
ORDER BY ordered.rn ASC
LIMIT 1
),
target_priority AS (
SELECT COALESCE(
(SELECT effective_priority FROM next_position),
(SELECT MAX(effective_priority) FROM ordered),
(SELECT effective_priority + $5::int FROM current_position)
) AS value
)
UPDATE integration_platforms target
SET dynamic_priority = COALESCE((SELECT max_priority FROM peer_priority), target.priority) + 1,
SET dynamic_priority = COALESCE((SELECT value FROM target_priority), target.priority) + 1,
updated_at = now()
WHERE target.id = $1::uuid
AND target.deleted_at IS NULL
AND EXISTS (SELECT 1 FROM current_model)
RETURNING dynamic_priority`, platformID, platformModelID, requestedModel, modelType).Scan(&dynamicPriority)
RETURNING dynamic_priority`, platformID, platformModelID, requestedModel, modelType, demoteSteps).Scan(&dynamicPriority)
return dynamicPriority, err
}

View File

@ -163,11 +163,27 @@ type RuntimeModelCandidate struct {
LoadLimited bool
LoadAvoided bool
LoadMetrics RuntimeCandidateLoadMetrics
CacheAffinity RuntimeCandidateCacheAffinity
RunningCount float64
WaitingCount float64
LastAssignedUnix float64
}
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
}
type RuntimeCandidateLoadMetrics struct {
RPMCurrent float64
RPMLimit float64

View File

@ -118,6 +118,19 @@ CROSS JOIN (
'{"formula":"ceil(input_tokens / 1000) * base_price"}'::jsonb,
10
),
(
'text_cached_input_tokens',
'文本缓存输入 Token',
'text_cached_input',
'1k_tokens',
0.001::numeric,
'{"meter":"cached_input_tokens"}'::jsonb,
'{}'::jsonb,
'token_usage',
'{"metrics":["cached_input_tokens"],"unitScale":1000}'::jsonb,
'{"formula":"ceil(cached_input_tokens / 1000) * base_price"}'::jsonb,
15
),
(
'text_output_tokens',
'文本输出 Token',

View File

@ -0,0 +1,46 @@
ALTER TABLE IF EXISTS gateway_runner_policies
ADD COLUMN IF NOT EXISTS cache_affinity_policy jsonb NOT NULL DEFAULT '{
"enabled": true,
"modelTypes": ["chat", "text_generate", "responses"],
"minSamples": 3,
"minInputTokens": 512,
"staleAfterSeconds": 86400,
"maxPriorityBoost": 20,
"emaAlpha": 0.35
}'::jsonb;
UPDATE gateway_runner_policies
SET cache_affinity_policy = COALESCE(NULLIF(cache_affinity_policy, '{}'::jsonb), '{
"enabled": true,
"modelTypes": ["chat", "text_generate", "responses"],
"minSamples": 3,
"minInputTokens": 512,
"staleAfterSeconds": 86400,
"maxPriorityBoost": 20,
"emaAlpha": 0.35
}'::jsonb),
metadata = COALESCE(metadata, '{}'::jsonb) || jsonb_build_object('cacheAffinityPolicy', '0055_cache_affinity_routing'),
updated_at = now();
CREATE TABLE IF NOT EXISTS gateway_cache_affinity_stats (
client_id text NOT NULL,
cache_affinity_key text NOT NULL,
platform_id uuid REFERENCES integration_platforms(id) ON DELETE SET NULL,
platform_model_id uuid REFERENCES platform_models(id) ON DELETE SET NULL,
model_type text NOT NULL,
request_count bigint NOT NULL DEFAULT 0,
input_tokens bigint NOT NULL DEFAULT 0,
cached_input_tokens bigint NOT NULL DEFAULT 0,
ema_hit_ratio numeric NOT NULL DEFAULT 0,
last_hit_ratio numeric NOT NULL DEFAULT 0,
last_observed_at timestamptz NOT NULL DEFAULT now(),
created_at timestamptz NOT NULL DEFAULT now(),
updated_at timestamptz NOT NULL DEFAULT now(),
PRIMARY KEY (client_id, cache_affinity_key)
);
CREATE INDEX IF NOT EXISTS idx_gateway_cache_affinity_key_observed
ON gateway_cache_affinity_stats(cache_affinity_key, last_observed_at DESC);
CREATE INDEX IF NOT EXISTS idx_gateway_cache_affinity_platform_model
ON gateway_cache_affinity_stats(platform_model_id, cache_affinity_key);

View File

@ -0,0 +1,169 @@
WITH defaults AS (
SELECT
'[
{
"enabled": true,
"categories": ["rate_limit"],
"codes": ["rate_limit"],
"statusCodes": [429],
"keywords": ["rate_limit", "429", "too many requests"],
"action": "cooldown_and_next",
"target": "model",
"cooldownSeconds": 300
},
{
"enabled": true,
"categories": ["auth_error"],
"codes": ["auth_failed", "invalid_api_key", "missing_credentials"],
"statusCodes": [401, 403],
"keywords": ["invalid api key", "api key is invalid", "unauthorized", "forbidden", "permission denied"],
"action": "disable_and_next",
"target": "platform"
},
{
"enabled": true,
"categories": ["network", "timeout", "stream_error", "rate_limit", "provider_5xx", "provider_overloaded"],
"codes": ["network", "timeout", "stream_read_error", "rate_limit", "server_error", "overloaded"],
"statusCodes": [408, 429, 500, 502, 503, 504],
"keywords": ["timeout", "network", "rate_limit", "overloaded", "temporarily_unavailable", "server_error", "429", "5xx"],
"action": "demote_and_next",
"target": "platform",
"demoteSteps": 1
},
{
"enabled": true,
"categories": ["network", "timeout", "stream_error", "rate_limit", "provider_5xx", "provider_overloaded", "auth_error"],
"codes": ["auth_failed", "invalid_api_key", "missing_credentials"],
"statusCodes": [401, 403, 408, 429, 500, 502, 503, 504],
"keywords": ["timeout", "network", "rate_limit", "overloaded", "temporarily_unavailable", "server_error", "auth_failed", "invalid_api_key", "missing_credentials", "unauthorized", "forbidden", "429", "5xx"],
"action": "next"
},
{
"enabled": true,
"categories": ["request_error", "unsupported_model", "user_permission", "insufficient_balance"],
"codes": ["bad_request", "invalid_request", "invalid_parameter", "missing_required", "unsupported_kind", "unsupported_model", "insufficient_balance", "permission_denied"],
"keywords": ["invalid_parameter", "missing required", "bad request", "insufficient balance"],
"action": "stop"
}
]'::jsonb AS legacy_action_rules,
'[
{
"enabled": true,
"categories": ["network", "timeout", "stream_error", "provider_5xx", "provider_overloaded"],
"action": "next"
},
{
"enabled": true,
"categories": ["rate_limit"],
"codes": ["rate_limit"],
"statusCodes": [429],
"keywords": ["rate_limit", "429", "too many requests"],
"action": "cooldown_and_next",
"target": "model",
"cooldownSeconds": 300
},
{
"enabled": true,
"keywords": ["rate limit", "too many requests", "temporarily_unavailable"],
"action": "demote_and_next",
"target": "platform",
"demoteSteps": 1
},
{
"enabled": true,
"categories": ["auth_error"],
"codes": ["auth_failed", "invalid_api_key", "missing_credentials"],
"statusCodes": [401, 403],
"keywords": ["invalid api key", "api key is invalid", "unauthorized", "forbidden", "permission denied"],
"action": "disable_and_next",
"target": "platform"
},
{
"enabled": true,
"categories": ["request_error", "unsupported_model", "user_permission", "insufficient_balance"],
"codes": ["bad_request", "invalid_request", "invalid_parameter", "missing_required", "unsupported_kind", "unsupported_model", "insufficient_balance", "permission_denied"],
"keywords": ["invalid_parameter", "missing required", "bad request", "insufficient balance"],
"action": "stop"
}
]'::jsonb AS compact_action_rules,
'[
{
"enabled": true,
"categories": ["rate_limit"],
"codes": ["rate_limit", "too_many_requests", "too_many_request", "rate_limit_exceeded", "requests_rate_limited", "requests_limit_exceeded", "resource_exhausted", "throttled", "quota_exceeded", "quota_exhausted", "qps_limit", "rpm_limit", "tpm_limit"],
"statusCodes": [429],
"keywords": ["rate limit", "rate_limit", "rate-limit", "too many requests", "too_many_requests", "429", "throttle", "throttled", "resource exhausted", "quota exceeded", "quota_exceeded", "rate exceeded", "request limit", "requests per minute", "tokens per minute", "qps", "rpm", "tpm"],
"action": "cooldown_and_next",
"target": "model",
"cooldownSeconds": 300
},
{
"enabled": true,
"categories": ["auth_error"],
"codes": ["auth_failed", "invalid_api_key", "missing_credentials", "missing_credential", "unauthorized", "forbidden", "account_disabled", "billing_not_active", "permission_denied"],
"statusCodes": [401, 403],
"keywords": ["invalid api key", "api key is invalid", "invalid_api_key", "apikey invalid", "api key invalid", "unauthorized", "forbidden", "authentication", "auth failed", "credential", "missing credential", "permission denied", "access denied", "invalid signature", "signature mismatch", "secret key", "access key", "account disabled", "account_deactivated", "billing not active", "billing_not_active", "insufficient balance", "quota exceeded", "余额不足", "欠费"],
"action": "disable_and_next",
"target": "platform"
},
{
"enabled": true,
"categories": ["timeout", "provider_5xx", "provider_overloaded"],
"codes": ["timeout", "server_error", "overloaded", "provider_failed", "invalid_response", "script_timeout", "request_asset_fetch_failed", "upload_failed", "upload_read_failed", "upload_source_read_failed"],
"statusCodes": [408, 500, 502, 503, 504, 520, 521, 522, 523, 524, 529, 530],
"keywords": ["timeout", "timed out", "deadline exceeded", "context deadline exceeded", "overloaded", "overload", "temporarily_unavailable", "temporarily unavailable", "temporary unavailable", "service unavailable", "server error", "internal server error", "bad gateway", "gateway timeout", "upstream timeout", "upstream error", "provider failed", "model is overloaded", "capacity", "try again later", "please try again later", "5xx", "500", "502", "503", "504"],
"action": "demote_and_next",
"target": "platform",
"demoteSteps": 1
},
{
"enabled": true,
"categories": ["network", "stream_error"],
"codes": ["network", "stream_read_error", "cancelled", "request_asset_fetch_failed", "upload_network", "upload_config_failed", "local_static_store_failed", "upload_source_fetch_failed"],
"statusCodes": [],
"keywords": ["network", "socket hang up", "socket closed", "connection reset", "connection refused", "connection aborted", "broken pipe", "econnreset", "econnrefused", "eai_again", "enotfound", "no such host", "dns", "tls handshake timeout", "unexpected eof", "connection closed", "stream error", "stream_read_error", "read failed", "fetch failed"],
"action": "next"
},
{
"enabled": true,
"categories": ["request_error", "unsupported_model", "user_permission", "insufficient_balance"],
"codes": ["bad_request", "invalid_request", "invalid_parameter", "missing_required", "unsupported_kind", "unsupported_model", "unsupported_operation", "unsupported_request_resolution", "request_asset_input_format_unsupported", "request_asset_expired", "request_asset_decode_failed", "request_asset_public_url_required", "upload_no_channel", "upload_unsupported_channel", "upload_source_too_large", "upload_source_invalid_url", "upload_source_unsupported_url", "upload_decode_failed", "invalid_proxy", "invalid_json_body", "invalid_multipart_body", "invalid_multipart_file", "invalid_multipart_image", "invalid_multipart_audio", "missing_configuration", "script_error", "cloned_voice_not_found", "cloned_voice_expired", "cloned_voice_unavailable", "insufficient_balance", "permission_denied"],
"statusCodes": [400, 404, 405, 406, 409, 410, 411, 413, 415, 422],
"keywords": ["invalid_parameter", "invalid parameter", "missing required", "missing_required", "bad request", "unsupported", "not supported", "does not support", "is not supported", "insufficient balance", "permission denied", "no permission", "invalid json", "invalid multipart", "image is required", "audio must", "prompt is required", "model is required", "exceeds", "too large", "expired", "not found", "required"],
"action": "stop"
}
]'::jsonb AS enhanced_action_rules
)
UPDATE gateway_runner_policies
SET failover_policy = jsonb_set(
COALESCE(failover_policy, '{}'::jsonb),
'{actionRules}',
CASE
WHEN failover_policy->'actionRules' IS NULL THEN defaults.enhanced_action_rules
WHEN failover_policy->'actionRules' = defaults.legacy_action_rules THEN defaults.enhanced_action_rules
WHEN failover_policy->'actionRules' = defaults.compact_action_rules THEN defaults.enhanced_action_rules
ELSE failover_policy->'actionRules'
END,
true
) || jsonb_build_object(
'legacyFieldsDeprecated', true,
'deprecatedFields', to_jsonb(ARRAY['allowCategories', 'denyCategories', 'allowCodes', 'denyCodes', 'allowKeywords', 'denyKeywords', 'allowStatusCodes', 'denyStatusCodes', 'actions']::text[]),
'replacedBy', 'actionRules'
),
hard_stop_policy = COALESCE(hard_stop_policy, '{}'::jsonb) || jsonb_build_object(
'deprecated', true,
'replacedBy', 'failoverPolicy.actionRules'
),
priority_demote_policy = jsonb_set(
COALESCE(priority_demote_policy, '{}'::jsonb) || jsonb_build_object(
'deprecated', true,
'replacedBy', 'failoverPolicy.actionRules'
),
'{derivedFromActionRules}',
'true'::jsonb,
true
),
metadata = COALESCE(metadata, '{}'::jsonb) || jsonb_build_object('failoverActionRules', '0056_runner_failover_action_rules_v2'),
updated_at = now()
FROM defaults
WHERE policy_key = 'default-runner-v1';

File diff suppressed because it is too large Load Diff

View File

@ -2136,23 +2136,155 @@
margin: 4px 0 0;
}
.runnerPolicyHeaderActions {
display: flex;
align-items: flex-end;
gap: 10px;
}
.runnerPolicyHeaderStatus {
min-width: 220px;
max-width: 240px;
}
.runnerPolicyWorkbench {
grid-template-columns: 240px minmax(0, 1fr);
.runnerPolicySectionTitle {
display: flex;
min-width: 0;
align-items: center;
gap: 10px;
}
.runnerPolicyDetailPanel {
min-height: 360px;
.runnerPolicySectionTitle > span {
display: grid;
min-width: 0;
gap: 2px;
}
.runnerPolicySectionTitle strong {
color: var(--text-strong);
font-size: var(--font-size-lg);
}
.runnerPolicySectionTitle small {
color: var(--muted-foreground);
font-size: var(--font-size-sm);
line-height: 1.35;
}
.runnerPolicyDetailRows {
padding: 8px 10px;
}
.runnerToggleField {
display: grid;
align-content: start;
gap: 8px;
}
.runnerRuleWorkbench {
display: grid;
grid-template-columns: minmax(220px, 280px) minmax(0, 1fr);
align-items: start;
gap: 12px;
}
.runnerRuleList,
.runnerRuleEditor {
min-width: 0;
border: 1px solid var(--border);
border-radius: 10px;
background: #fff;
}
.runnerRuleList {
display: grid;
align-content: start;
gap: 8px;
padding: 10px;
}
.runnerRuleListHead,
.runnerRuleEditor header,
.runnerRuleTools {
display: flex;
align-items: center;
gap: 10px;
}
.runnerRuleListHead,
.runnerRuleEditor header {
justify-content: space-between;
}
.runnerRuleListHead strong,
.runnerRuleEditor strong {
color: var(--text-normal);
font-size: var(--font-size-sm);
}
.runnerRuleTab {
display: grid;
width: 100%;
min-width: 0;
grid-template-columns: minmax(0, 1fr) auto;
align-items: center;
gap: 8px;
padding: 10px;
border: 1px solid transparent;
border-radius: 8px;
background: transparent;
color: inherit;
cursor: pointer;
text-align: left;
}
.runnerRuleTab[data-active='true'] {
border-color: color-mix(in srgb, var(--primary) 55%, transparent);
background: color-mix(in srgb, var(--primary) 9%, transparent);
}
.runnerRuleTab[data-disabled='true'] {
opacity: 0.62;
}
.runnerRuleTab span,
.runnerRuleEditor header > div:first-child {
display: grid;
min-width: 0;
gap: 3px;
}
.runnerRuleTab strong,
.runnerRuleTab small,
.runnerRuleEditor header span {
overflow: hidden;
text-overflow: ellipsis;
white-space: nowrap;
}
.runnerRuleTab small,
.runnerRuleTab em,
.runnerRuleEditor header span {
color: var(--muted-foreground);
font-size: var(--font-size-xs);
font-style: normal;
font-weight: var(--font-weight-regular);
}
.runnerRuleEditor {
display: grid;
gap: 12px;
padding: 12px;
}
.runnerRuleFields {
align-items: start;
}
.runnerRuleTools .platformToggle {
min-width: auto;
}
.runnerActionMatrix {
display: grid;
gap: 10px;
@ -2242,6 +2374,12 @@
flex-direction: column;
}
.runnerPolicyHeaderActions {
width: 100%;
align-items: stretch;
flex-direction: column;
}
.runnerPolicyHeaderStatus {
width: 100%;
max-width: none;
@ -2284,6 +2422,7 @@
.runtimePolicyFormBody,
.fileStorageDialogBody,
.runtimePolicyRows,
.runnerRuleWorkbench,
.runnerActionGrid,
.accessPermissionGrid,
.accessTreeToolbar,

View File

@ -283,6 +283,7 @@ export interface GatewayRunnerPolicy {
hardStopPolicy?: Record<string, unknown>;
priorityDemotePolicy?: Record<string, unknown>;
singleSourcePolicy?: Record<string, unknown>;
cacheAffinityPolicy?: Record<string, unknown>;
metadata?: Record<string, unknown>;
status: 'active' | 'disabled' | string;
createdAt: string;
@ -297,6 +298,7 @@ export interface GatewayRunnerPolicyUpsertRequest {
hardStopPolicy?: Record<string, unknown>;
priorityDemotePolicy?: Record<string, unknown>;
singleSourcePolicy?: Record<string, unknown>;
cacheAffinityPolicy?: Record<string, unknown>;
metadata?: Record<string, unknown>;
status?: 'active' | 'disabled' | string;
}