From 0e675e259cdd5988d8d32a71e62b3b00e96d97ce Mon Sep 17 00:00:00 2001 From: wangbo Date: Sun, 28 Jun 2026 20:50:23 +0800 Subject: [PATCH] feat: add runner failover rules and cache affinity --- apps/api/internal/clients/clients_test.go | 33 +- apps/api/internal/clients/gemini.go | 4 +- apps/api/internal/clients/helpers.go | 65 +- apps/api/internal/clients/simulation.go | 58 +- apps/api/internal/clients/types.go | 9 +- .../httpapi/core_flow_integration_test.go | 349 ++++++- apps/api/internal/httpapi/streaming.go | 1 + apps/api/internal/runner/cache_affinity.go | 69 ++ .../internal/runner/cache_affinity_test.go | 42 + apps/api/internal/runner/limits_test.go | 29 + apps/api/internal/runner/pricing.go | 2 +- apps/api/internal/runner/recording.go | 16 +- apps/api/internal/runner/recording_test.go | 38 + apps/api/internal/runner/retry_decision.go | 155 ++- .../internal/runner/retry_decision_test.go | 116 +++ apps/api/internal/runner/runtime_policy.go | 52 +- apps/api/internal/runner/service.go | 33 +- apps/api/internal/runner/trace.go | 19 + apps/api/internal/store/cache_affinity.go | 73 ++ apps/api/internal/store/candidates.go | 273 +++++- apps/api/internal/store/candidates_test.go | 336 +++++++ apps/api/internal/store/postgres.go | 1 + apps/api/internal/store/runner_policies.go | 283 +++++- apps/api/internal/store/runtime_policies.go | 61 +- apps/api/internal/store/runtime_types.go | 16 + .../api/migrations/0007_pricing_rule_sets.sql | 13 + .../0055_cache_affinity_routing.sql | 46 + .../0056_runner_failover_action_rules.sql | 169 ++++ .../src/pages/admin/RuntimePoliciesPanel.tsx | 899 +++++++++++------- apps/web/src/styles/pages.css | 147 ++- packages/contracts/src/index.ts | 2 + 31 files changed, 2939 insertions(+), 470 deletions(-) create mode 100644 apps/api/internal/runner/cache_affinity.go create mode 100644 apps/api/internal/runner/cache_affinity_test.go create mode 100644 apps/api/internal/runner/recording_test.go create mode 100644 apps/api/internal/store/cache_affinity.go create mode 100644 apps/api/migrations/0055_cache_affinity_routing.sql create mode 100644 apps/api/migrations/0056_runner_failover_action_rules.sql diff --git a/apps/api/internal/clients/clients_test.go b/apps/api/internal/clients/clients_test.go index 74456f1..85cf910 100644 --- a/apps/api/internal/clients/clients_test.go +++ b/apps/api/internal/clients/clients_test.go @@ -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 diff --git a/apps/api/internal/clients/gemini.go b/apps/api/internal/clients/gemini.go index 3c70d12..88f5990 100644 --- a/apps/api/internal/clients/gemini.go +++ b/apps/api/internal/clients/gemini.go @@ -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 { diff --git a/apps/api/internal/clients/helpers.go b/apps/api/internal/clients/helpers.go index 20f32a8..98e35bb 100644 --- a/apps/api/internal/clients/helpers.go +++ b/apps/api/internal/clients/helpers.go @@ -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 { diff --git a/apps/api/internal/clients/simulation.go b/apps/api/internal/clients/simulation.go index a07ee5b..a0b60dc 100644 --- a/apps/api/internal/clients/simulation.go +++ b/apps/api/internal/clients/simulation.go @@ -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 == "" { diff --git a/apps/api/internal/clients/types.go b/apps/api/internal/clients/types.go index 87a4dab..b017eaf 100644 --- a/apps/api/internal/clients/types.go +++ b/apps/api/internal/clients/types.go @@ -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 { diff --git a/apps/api/internal/httpapi/core_flow_integration_test.go b/apps/api/internal/httpapi/core_flow_integration_test.go index 1981e56..321dd67 100644 --- a/apps/api/internal/httpapi/core_flow_integration_test.go +++ b/apps/api/internal/httpapi/core_flow_integration_test.go @@ -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, ®isterResponse) + + 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: diff --git a/apps/api/internal/httpapi/streaming.go b/apps/api/internal/httpapi/streaming.go index be58ef3..e82e895 100644 --- a/apps/api/internal/httpapi/streaming.go +++ b/apps/api/internal/httpapi/streaming.go @@ -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 } diff --git a/apps/api/internal/runner/cache_affinity.go b/apps/api/internal/runner/cache_affinity.go new file mode 100644 index 0000000..a82910b --- /dev/null +++ b/apps/api/internal/runner/cache_affinity.go @@ -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[:]) +} diff --git a/apps/api/internal/runner/cache_affinity_test.go b/apps/api/internal/runner/cache_affinity_test.go new file mode 100644 index 0000000..babfb47 --- /dev/null +++ b/apps/api/internal/runner/cache_affinity_test.go @@ -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) + } +} diff --git a/apps/api/internal/runner/limits_test.go b/apps/api/internal/runner/limits_test.go index 6155439..31490b3 100644 --- a/apps/api/internal/runner/limits_test.go +++ b/apps/api/internal/runner/limits_test.go @@ -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{ diff --git a/apps/api/internal/runner/pricing.go b/apps/api/internal/runner/pricing.go index 3ae7092..f5adbd5 100644 --- a/apps/api/internal/runner/pricing.go +++ b/apps/api/internal/runner/pricing.go @@ -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{} diff --git a/apps/api/internal/runner/recording.go b/apps/api/internal/runner/recording.go index 6c975ed..c9b4f8e 100644 --- a/apps/api/internal/runner/recording.go +++ b/apps/api/internal/runner/recording.go @@ -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 diff --git a/apps/api/internal/runner/recording_test.go b/apps/api/internal/runner/recording_test.go new file mode 100644 index 0000000..a024a52 --- /dev/null +++ b/apps/api/internal/runner/recording_test.go @@ -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) + } +} diff --git a/apps/api/internal/runner/retry_decision.go b/apps/api/internal/runner/retry_decision.go index ecc06bd..30f47ee 100644 --- a/apps/api/internal/runner/retry_decision.go +++ b/apps/api/internal/runner/retry_decision.go @@ -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 diff --git a/apps/api/internal/runner/retry_decision_test.go b/apps/api/internal/runner/retry_decision_test.go index 0a2e493..aebc1f1 100644 --- a/apps/api/internal/runner/retry_decision_test.go +++ b/apps/api/internal/runner/retry_decision_test.go @@ -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) + } +} diff --git a/apps/api/internal/runner/runtime_policy.go b/apps/api/internal/runner/runtime_policy.go index a80375c..131a90e 100644 --- a/apps/api/internal/runner/runtime_policy.go +++ b/apps/api/internal/runner/runtime_policy.go @@ -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), diff --git a/apps/api/internal/runner/service.go b/apps/api/internal/runner/service.go index bcd5983..bec0f01 100644 --- a/apps/api/internal/runner/service.go +++ b/apps/api/internal/runner/service.go @@ -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 } diff --git a/apps/api/internal/runner/trace.go b/apps/api/internal/runner/trace.go index 83f4de5..d585bcf 100644 --- a/apps/api/internal/runner/trace.go +++ b/apps/api/internal/runner/trace.go @@ -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 { diff --git a/apps/api/internal/store/cache_affinity.go b/apps/api/internal/store/cache_affinity.go new file mode 100644 index 0000000..6d949e9 --- /dev/null +++ b/apps/api/internal/store/cache_affinity.go @@ -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 +} diff --git a/apps/api/internal/store/candidates.go b/apps/api/internal/store/candidates.go index c628244..872d749 100644 --- a/apps/api/internal/store/candidates.go +++ b/apps/api/internal/store/candidates.go @@ -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 diff --git a/apps/api/internal/store/candidates_test.go b/apps/api/internal/store/candidates_test.go index b5ce0f6..b5e09fd 100644 --- a/apps/api/internal/store/candidates_test.go +++ b/apps/api/internal/store/candidates_test.go @@ -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 } diff --git a/apps/api/internal/store/postgres.go b/apps/api/internal/store/postgres.go index 82b11b1..e847db4 100644 --- a/apps/api/internal/store/postgres.go +++ b/apps/api/internal/store/postgres.go @@ -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"` diff --git a/apps/api/internal/store/runner_policies.go b/apps/api/internal/store/runner_policies.go index 57e73fd..071358e 100644 --- a/apps/api/internal/store/runner_policies.go +++ b/apps/api/internal/store/runner_policies.go @@ -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{}, diff --git a/apps/api/internal/store/runtime_policies.go b/apps/api/internal/store/runtime_policies.go index 6ff7ceb..3b04730 100644 --- a/apps/api/internal/store/runtime_policies.go +++ b/apps/api/internal/store/runtime_policies.go @@ -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 } diff --git a/apps/api/internal/store/runtime_types.go b/apps/api/internal/store/runtime_types.go index f614bee..a8d301b 100644 --- a/apps/api/internal/store/runtime_types.go +++ b/apps/api/internal/store/runtime_types.go @@ -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 diff --git a/apps/api/migrations/0007_pricing_rule_sets.sql b/apps/api/migrations/0007_pricing_rule_sets.sql index 18bae32..a5e249b 100644 --- a/apps/api/migrations/0007_pricing_rule_sets.sql +++ b/apps/api/migrations/0007_pricing_rule_sets.sql @@ -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', diff --git a/apps/api/migrations/0055_cache_affinity_routing.sql b/apps/api/migrations/0055_cache_affinity_routing.sql new file mode 100644 index 0000000..11d364d --- /dev/null +++ b/apps/api/migrations/0055_cache_affinity_routing.sql @@ -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); diff --git a/apps/api/migrations/0056_runner_failover_action_rules.sql b/apps/api/migrations/0056_runner_failover_action_rules.sql new file mode 100644 index 0000000..f832f44 --- /dev/null +++ b/apps/api/migrations/0056_runner_failover_action_rules.sql @@ -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'; diff --git a/apps/web/src/pages/admin/RuntimePoliciesPanel.tsx b/apps/web/src/pages/admin/RuntimePoliciesPanel.tsx index 7fca483..d90ac79 100644 --- a/apps/web/src/pages/admin/RuntimePoliciesPanel.tsx +++ b/apps/web/src/pages/admin/RuntimePoliciesPanel.tsx @@ -6,7 +6,20 @@ import { Badge, Button, Card, CardContent, CardHeader, CardTitle, ConfirmDialog, import type { LoadState } from '../../types'; type RuntimePanelTab = 'model' | 'runner'; -type RunnerPolicyStrategy = 'failover' | 'hardStop' | 'priorityDemote' | 'singleSource'; +type RunnerFailoverAction = 'next' | 'cooldown_and_next' | 'demote_and_next' | 'disable_and_next' | 'stop'; +type RunnerFailoverTarget = 'platform' | 'model'; +type RunnerActionRule = { + id: string; + enabled: boolean; + categories: string[]; + codes: string[]; + statusCodes: string[]; + keywords: string[]; + action: RunnerFailoverAction; + target: RunnerFailoverTarget; + demoteSteps: string; + cooldownSeconds: string; +}; type RuntimePolicyForm = { policyKey: string; @@ -32,52 +45,35 @@ type RuntimePolicyForm = { type RunnerPolicyForm = { name: string; description: string; - failoverEnabled: boolean; maxPlatforms: string; maxDurationSeconds: string; - allowCategories: string[]; - denyCategories: string[]; - allowCodes: string[]; - denyCodes: string[]; - allowKeywords: string[]; - denyKeywords: string[]; - allowStatusCodes: string[]; - denyStatusCodes: string[]; - failoverActions: Record; - hardStopEnabled: boolean; - hardStopCategories: string[]; - hardStopCodes: string[]; - hardStopStatusCodes: string[]; - hardStopKeywords: string[]; - priorityDemoteEnabled: boolean; - priorityDemoteCategories: string[]; - priorityDemoteCodes: string[]; - priorityDemoteStatusCodes: string[]; - priorityDemoteKeywords: string[]; singleSourceProtectionEnabled: boolean; + actionRules: RunnerActionRule[]; + cacheAffinityPolicyJson: string; metadataJson: string; status: string; }; -const failoverActionDefinitions = [ - { - value: 'next', - title: '仅轮转', - description: '这些错误分类只尝试下一个平台，不修改当前平台状态。', - }, - { - value: 'disable_and_next', - title: '禁用后轮转', - description: '这些错误分类会先禁用当前平台，再尝试下一个平台。', - }, - { - value: 'cooldown_and_next', - title: '冷却后轮转', - description: '这些错误分类会先让当前平台模型进入冷却，再尝试下一个平台。', - }, -] as const; +const failoverActionOptions: Array<{ label: string; value: RunnerFailoverAction }> = [ + { label: '仅重试下一个', value: 'next' }, + { label: '冷却后重试', value: 'cooldown_and_next' }, + { label: '降级后重试', value: 'demote_and_next' }, + { label: '禁用后重试', value: 'disable_and_next' }, + { label: '不再重试', value: 'stop' }, +]; -const failoverActionValues = new Set(failoverActionDefinitions.map((item) => item.value)); +const failoverActionText: Record = { + next: '仅重试下一个', + cooldown_and_next: '冷却后重试', + demote_and_next: '降级后重试', + disable_and_next: '禁用后重试', + stop: '不再重试', +}; + +const failoverTargetOptions: Array<{ label: string; value: RunnerFailoverTarget }> = [ + { label: '平台', value: 'platform' }, + { label: '模型', value: 'model' }, +]; const failoverCategoryOptions = [ 'network', @@ -93,6 +89,8 @@ const failoverCategoryOptions = [ 'insufficient_balance', 'client_error', ].map((item) => ({ label: item, value: item })); +const legacyRunnerFailoverFields = ['allowCategories', 'denyCategories', 'allowCodes', 'denyCodes', 'allowKeywords', 'denyKeywords', 'allowStatusCodes', 'denyStatusCodes', 'actions']; +let runnerActionRuleId = 0; export function RuntimePoliciesPanel(props: { message: string; @@ -330,39 +328,39 @@ function RunnerPolicyEditor(props: { onChange: (form: RunnerPolicyForm) => void; onSubmit: (event: FormEvent) => void; }) { - const [activeStrategy, setActiveStrategy] = useState('failover'); + const [activeRuleId, setActiveRuleId] = useState(''); const patch = (next: Partial) => props.onChange({ ...props.form, ...next }); - const strategyDefinitions = [ - { - value: 'failover' as const, - title: '平台间故障切换', - description: '当前平台内部重试耗尽后是否尝试下一个候选平台', - enabled: props.form.failoverEnabled, - icon: , - }, - { - value: 'hardStop' as const, - title: '硬拒绝规则', - description: '参数、余额、权限等错误直接失败，不受模型覆盖影响', - enabled: props.form.hardStopEnabled, - icon: , - }, - { - value: 'priorityDemote' as const, - title: '优先级降级', - description: '命中后将失败平台的动态优先级调整到当前模型队尾', - enabled: props.form.priorityDemoteEnabled, - icon: , - }, - { - value: 'singleSource' as const, - title: '单一源保护', - description: '只有一个可用源时跳过自动禁用和冷却', - enabled: props.form.singleSourceProtectionEnabled, - icon: , - }, - ]; - const activeDefinition = strategyDefinitions.find((item) => item.value === activeStrategy) ?? strategyDefinitions[0]; + const activeRule = props.form.actionRules.find((rule) => rule.id === activeRuleId) ?? props.form.actionRules[0] ?? null; + const activeRuleIndex = activeRule ? props.form.actionRules.findIndex((rule) => rule.id === activeRule.id) : -1; + const patchRule = (ruleId: string, next: Partial) => { + patch({ actionRules: props.form.actionRules.map((rule) => rule.id === ruleId ? { ...rule, ...next } : rule) }); + }; + const addRule = () => { + const rule = createRunnerActionRule({ categories: [], action: 'next' }); + patch({ actionRules: [...props.form.actionRules, rule] }); + setActiveRuleId(rule.id); + }; + const removeRule = (ruleId: string) => { + if (props.form.actionRules.length <= 1) return; + const removedIndex = props.form.actionRules.findIndex((rule) => rule.id === ruleId); + const nextRules = props.form.actionRules.filter((rule) => rule.id !== ruleId); + patch({ actionRules: nextRules }); + if (activeRuleId === ruleId) { + setActiveRuleId(nextRules[Math.max(0, removedIndex - 1)]?.id ?? nextRules[0]?.id ?? ''); + } + }; + + useEffect(() => { + if (!props.form.actionRules.length) { + setActiveRuleId(''); + return; + } + if (!props.form.actionRules.some((rule) => rule.id === activeRuleId)) { + setActiveRuleId(props.form.actionRules[0].id); + } + }, [activeRuleId, props.form.actionRules]); + + const formId = 'runner-global-policy-form'; return ( @@ -371,129 +369,147 @@ function RunnerPolicyEditor(props: { {props.form.name}

{props.form.description}

- +
+ + +
-
-
- - -
-
- {activeDefinition.icon} -
- {activeDefinition.title} - {activeDefinition.description} -
- {activeDefinition.enabled ? '启用' : '关闭'} -
- - {activeStrategy === 'failover' && ( -
- patch({ failoverEnabled: checked })} /> - - - patch(value)} - /> -
- 补充触发条件（可自定义扩展） - 这些条件只决定是否进入平台间切换；分类命中后的处理方式以上面的分组为准。 +
+
- )} + {props.form.actionRules.map((rule, index) => ( + + ))} + - {activeStrategy === 'hardStop' && ( -
- patch({ hardStopEnabled: checked })} /> - patch({ hardStopCategories: value })} /> - patch({ hardStopCodes: value })} /> - patch({ hardStopStatusCodes: value })} /> - 建议只用错误码、分类和关键词做硬拒绝；401/403 这类平台鉴权错误默认走平台间切换。 - patch({ hardStopKeywords: value })} /> -
- )} - - {activeStrategy === 'priorityDemote' && ( -
- patch({ priorityDemoteEnabled: checked })} /> - 命中降级规则后，失败平台会自动调整到提供当前模型的其它客户端优先级队尾。 - patch({ priorityDemoteCategories: value })} /> - patch({ priorityDemoteCodes: value })} /> - patch({ priorityDemoteStatusCodes: value })} /> - patch({ priorityDemoteKeywords: value })} /> -
- )} - - {activeStrategy === 'singleSource' && ( -
- patch({ singleSourceProtectionEnabled: checked })} - /> - 本次调度只有一个可用源时，命中自动禁用、运行策略冷却或故障切换冷却动作都只记录保护事件，不修改平台或模型状态。 -
- )} + {activeRule && ( +
+
+
+ #{activeRuleIndex + 1} {failoverActionText[activeRule.action]} + {runnerActionRuleConditionText(activeRule)} +
+
+ patchRule(activeRule.id, { enabled: checked })} /> + +
+
+
+ patchRule(activeRule.id, { categories: value })} /> + patchRule(activeRule.id, { codes: value })} /> + patchRule(activeRule.id, { statusCodes: value })} /> + patchRule(activeRule.id, { keywords: value })} /> + + {actionNeedsTarget(activeRule.action) && ( + + )} + {activeRule.action === 'cooldown_and_next' && ( + + )} + {activeRule.action === 'demote_and_next' && ( + + )} +
+
+ )} +
-
+