From 152885bbb61815f55ae0123111ec53d1ccf9aa7f Mon Sep 17 00:00:00 2001 From: wangbo Date: Fri, 24 Jul 2026 23:04:12 +0800 Subject: [PATCH] =?UTF-8?q?fix(routing):=20=E5=AF=B9=E9=BD=90=E7=BC=93?= =?UTF-8?q?=E5=AD=98=E4=BA=B2=E5=92=8C=E5=8A=9B=E7=A1=AC=E8=A7=84=E5=88=99?= =?UTF-8?q?=E4=B8=8E=E5=AE=A1=E8=AE=A1?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../httpapi/core_flow_integration_test.go | 98 ++++++++- apps/api/internal/runner/cache_affinity.go | 205 +++++++++++++++--- .../internal/runner/cache_affinity_test.go | 84 ++++++- apps/api/internal/runner/recording.go | 8 + apps/api/internal/runner/recording_test.go | 29 ++- apps/api/internal/store/candidates.go | 72 +++++- apps/api/internal/store/candidates_test.go | 44 +++- apps/api/internal/store/runtime_types.go | 27 ++- apps/api/internal/store/task_history_test.go | 19 +- apps/api/internal/store/tasks_runtime.go | 48 +++- 10 files changed, 540 insertions(+), 94 deletions(-) diff --git a/apps/api/internal/httpapi/core_flow_integration_test.go b/apps/api/internal/httpapi/core_flow_integration_test.go index 1e72b97..1eb15ac 100644 --- a/apps/api/internal/httpapi/core_flow_integration_test.go +++ b/apps/api/internal/httpapi/core_flow_integration_test.go @@ -1746,6 +1746,7 @@ WHERE m.platform_id = $1::uuid } restartModel := "worker-restart-" + suffixText + createSimulationTextBaseModel(t, server.URL, loginResponse.AccessToken, restartModel) createSimulationTextPlatformModel( t, server.URL, @@ -1864,6 +1865,7 @@ func TestCacheAffinityRoutingFlow(t *testing.T) { model := "cache-affinity-smoke-" + suffixText cacheKey := "cache-affinity-key-" + suffixText + createSimulationTextBaseModel(t, server.URL, loginResponse.AccessToken, model) 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}, @@ -1880,8 +1882,31 @@ func TestCacheAffinityRoutingFlow(t *testing.T) { } } - highPlatform, _ := createSimulationTextPlatformModel(t, server.URL, loginResponse.AccessToken, "cache-high-"+suffixText, "Cache Affinity High Priority", model, 1, nil) - _ = highPlatform + tools := []any{map[string]any{ + "type": "function", + "function": map[string]any{ + "name": "lookup_weather", + "description": "Look up current weather", + "parameters": map[string]any{"type": "object", "properties": map[string]any{"city": map[string]any{"type": "string"}}}, + }, + }} + initialToolMessages := []any{ + map[string]any{"role": "system", "content": "Use tools when helpful."}, + map[string]any{"role": "user", "content": "Weather in Hangzhou?"}, + } + for index := 0; index < 3; index++ { + detail := runCacheAffinityChatTaskWithBody(t, ctx, testPool, server.URL, apiKeyResponse.Secret, model, "tool-prime-"+strconv.Itoa(index)+"-"+suffixText, map[string]any{ + "tools": tools, + "messages": initialToolMessages, + "simulationUsage": map[string]any{"inputTokens": 1000, "cachedInputTokens": 750, "outputTokens": 20}, + }) + if detail.Status != "succeeded" || len(detail.Attempts) != 1 || detail.Attempts[0].PlatformName != "Cache Affinity Low Priority" { + t.Fatalf("tool-history priming request should use the only platform, got %+v", detail) + } + } + + peerPlatform, _ := createSimulationTextPlatformModel(t, server.URL, loginResponse.AccessToken, "cache-peer-"+suffixText, "Cache Affinity Peer", model, 20, nil) + _ = peerPlatform sameKeyDetail := runCacheAffinityChatTask(t, ctx, testPool, server.URL, apiKeyResponse.Secret, model, cacheKey, "same-key-"+suffixText, map[string]any{ "inputTokens": 1000, "cachedInputTokens": 700, @@ -1899,12 +1924,36 @@ func TestCacheAffinityRoutingFlow(t *testing.T) { t.Fatalf("same cache key task detail should expose simulated usage, got %+v", sameKeyDetail.Usage) } + toolHistoryDetail := runCacheAffinityChatTaskWithBody(t, ctx, testPool, server.URL, apiKeyResponse.Secret, model, "tool-history-"+suffixText, map[string]any{ + "tools": tools, + "messages": append(append([]any{}, initialToolMessages...), + map[string]any{"role": "assistant", "tool_calls": []any{map[string]any{ + "id": "call_weather", + "type": "function", + "function": map[string]any{ + "name": "lookup_weather", + "arguments": `{"city":"Hangzhou"}`, + }, + }}}, + map[string]any{"role": "tool", "tool_call_id": "call_weather", "content": `{"temperature":31}`}, + ), + "simulationUsage": map[string]any{"inputTokens": 1400, "cachedInputTokens": 900, "outputTokens": 30}, + }) + if toolHistoryDetail.Status != "succeeded" || len(toolHistoryDetail.Attempts) != 1 || toolHistoryDetail.Attempts[0].PlatformName != "Cache Affinity Low Priority" { + t.Fatalf("tool continuation should retain the cached platform, got %+v", toolHistoryDetail) + } + if !boolFromTestMap(toolHistoryDetail.Attempts[0].Metrics, "cacheAffinityMatched") || + intFromTestAny(toolHistoryDetail.Attempts[0].Metrics["cacheAffinityMatchedPrefixDepth"]) < 2 || + intFromTestAny(toolHistoryDetail.Attempts[0].Metrics["cacheAffinityCandidateCount"]) < 1 { + t.Fatalf("tool continuation should expose prefix and candidate metrics, got %+v", toolHistoryDetail.Attempts[0].Metrics) + } + 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" { + if newKeyDetail.Status != "succeeded" || len(newKeyDetail.Attempts) != 1 || newKeyDetail.Attempts[0].PlatformName != "Cache Affinity Peer" { t.Fatalf("new cache key should fall back to base priority, got %+v", newKeyDetail) } @@ -1914,9 +1963,12 @@ func TestCacheAffinityRoutingFlow(t *testing.T) { "cachedInputTokens": 0, "outputTokens": 20, }) - if fullAvoidedDetail.Status != "succeeded" || len(fullAvoidedDetail.Attempts) != 1 || fullAvoidedDetail.Attempts[0].PlatformName != "Cache Affinity High Priority" { + if fullAvoidedDetail.Status != "succeeded" || len(fullAvoidedDetail.Attempts) != 1 || fullAvoidedDetail.Attempts[0].PlatformName != "Cache Affinity Peer" { t.Fatalf("full cached platform should be avoided before cache affinity boost, got %+v", fullAvoidedDetail) } + if fullAvoidedDetail.Attempts[0].Metrics["cacheAffinityOverrideReason"] != "capacity_tier_unavailable" { + t.Fatalf("full cached platform override reason should be visible, got %+v", fullAvoidedDetail.Attempts[0].Metrics) + } var requestCount int var inputTokens int @@ -1926,7 +1978,7 @@ func TestCacheAffinityRoutingFlow(t *testing.T) { 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 + ORDER BY request_count DESC, cached_input_tokens DESC, 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) } @@ -2214,6 +2266,20 @@ type cacheAffinityTaskDetail struct { } `json:"attempts"` } +func createSimulationTextBaseModel(t *testing.T, baseURL string, adminToken string, model string) { + t.Helper() + doJSON(t, baseURL, http.MethodPost, "/api/admin/catalog/base-models", adminToken, map[string]any{ + "providerKey": "openai", + "canonicalModelKey": "openai:" + model, + "invocationName": model, + "providerModelName": model, + "modelType": []string{"text_generate"}, + "displayName": model, + "capabilities": map[string]any{"originalTypes": []string{"text_generate"}}, + "metadata": map[string]any{"source": "cache-affinity-integration-test"}, + }, http.StatusCreated, nil) +} + 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 @@ -2230,7 +2296,7 @@ func createSimulationTextPlatformModel(t *testing.T, baseURL string, adminToken t.Fatalf("cache affinity platform was not created: %+v", platform) } payload := map[string]any{ - "canonicalModelKey": "openai:gpt-4o-mini", + "canonicalModelKey": "openai:" + model, "modelName": model, "providerModelName": model, "modelAlias": model, @@ -2251,17 +2317,27 @@ func createSimulationTextPlatformModel(t *testing.T, baseURL string, adminToken } 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() + return runCacheAffinityChatTaskWithBody(t, ctx, pool, baseURL, token, model, marker, map[string]any{ + "cacheAffinityKey": cacheAffinityKey, + "simulationUsage": simulationUsage, + "messages": []any{map[string]any{"role": "user", "content": "cache affinity route"}}, + }) +} + +func runCacheAffinityChatTaskWithBody(t *testing.T, ctx context.Context, pool *pgxpool.Pool, baseURL string, token string, model string, marker string, body map[string]any) cacheAffinityTaskDetail { t.Helper() var detail cacheAffinityTaskDetail - doAPIV1ChatCompletionAndLoadTask(t, ctx, pool, baseURL, token, map[string]any{ + requestBody := 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) + } + for key, value := range body { + requestBody[key] = value + } + doAPIV1ChatCompletionAndLoadTask(t, ctx, pool, baseURL, token, requestBody, marker, http.StatusOK, nil, &detail) return detail } diff --git a/apps/api/internal/runner/cache_affinity.go b/apps/api/internal/runner/cache_affinity.go index 90a2b6c..90a96b5 100644 --- a/apps/api/internal/runner/cache_affinity.go +++ b/apps/api/internal/runner/cache_affinity.go @@ -13,6 +13,8 @@ type cacheAffinityKeys struct { Record []string } +const cacheAffinityMinimumStableMessages = 2 + func buildCacheAffinityKey(kind string, modelType string, body map[string]any) string { return buildCacheAffinityKeys(kind, modelType, body).Primary } @@ -37,48 +39,28 @@ func buildCacheAffinityKeys(kind string, modelType string, body map[string]any) return cacheAffinityKeys{} } basePayload := map[string]any{ - "kind": kind, - "modelType": modelType, + "version": "v2", } if tools, ok := body["tools"]; ok { basePayload["tools"] = tools } - if instructions, ok := body["instructions"]; ok { - basePayload["instructions"] = instructions - } messageHashes := cacheAffinityMessageHashes(body) - if len(messageHashes) > 0 { + if len(messageHashes) >= cacheAffinityMinimumStableMessages { baseHash := stableJSONHash(basePayload) - keys := make([]string, 0, len(messageHashes)) - for index := range messageHashes { - keyPayload := map[string]any{ - "base": baseHash, - "messageHashes": messageHashes[:index+1], - "prefixLength": index + 1, - } - keys = append(keys, "prompt_lcp:"+stableJSONHash(keyPayload)) - } + keys := cacheAffinityPrefixKeys("prompt_lcp_v2:", baseHash, messageHashes) if len(keys) == 0 { return cacheAffinityKeys{} } - lookup := make([]string, 0, len(keys)) - for index := len(keys) - 1; index >= 0; index-- { - lookup = append(lookup, keys[index]) - } + lookup := reverseCacheAffinityKeys(keys) + legacyKeys := legacyCacheAffinityPrefixKeys(kind, modelType, body) + lookup = normalizedCacheAffinityKeyList(append(lookup, reverseCacheAffinityKeys(legacyKeys)...)) return cacheAffinityKeys{ Primary: keys[len(keys)-1], Lookup: lookup, Record: keys, } } - if input, ok := body["input"]; ok { - basePayload["input"] = input - } - if len(basePayload) <= 2 { - return cacheAffinityKeys{} - } - key := "prompt:" + stableJSONHash(basePayload) - return cacheAffinityKeys{Primary: key, Lookup: []string{key}, Record: []string{key}} + return cacheAffinityKeys{} } func cacheAffinityPromptHashSupported(kind string, modelType string) bool { @@ -95,6 +77,146 @@ func cacheAffinityPromptHashSupported(kind string, modelType string) bool { } func cacheAffinityMessageHashes(body map[string]any) []string { + if hashes := stringListFromAny(body["messageHashes"]); len(hashes) > 0 { + return hashes + } + if hashes := stringListFromAny(body["message_hashes"]); len(hashes) > 0 { + return hashes + } + messages := cacheAffinityStableMessages(body) + if len(messages) == 0 { + return nil + } + hashes := make([]string, 0, len(messages)) + for _, message := range messages { + hash := stableJSONHash(message) + if hash == "" { + continue + } + hashes = append(hashes, hash) + } + return hashes +} + +func cacheAffinityStableMessages(body map[string]any) []any { + messages := make([]any, 0) + if instructions := strings.TrimSpace(stringFromAny(body["instructions"])); instructions != "" { + messages = append(messages, map[string]any{"role": "system", "content": instructions}) + } + if rawMessages, ok := body["messages"].([]any); ok { + for _, message := range rawMessages { + if normalized := normalizeCacheAffinityMessage(message); normalized != nil { + messages = append(messages, normalized) + } + } + return messages + } + switch input := body["input"].(type) { + case string: + if strings.TrimSpace(input) != "" { + messages = append(messages, map[string]any{"role": "user", "content": input}) + } + case []any: + for _, message := range input { + if normalized := normalizeCacheAffinityMessage(message); normalized != nil { + messages = append(messages, normalized) + } + } + } + return messages +} + +func normalizeCacheAffinityMessage(value any) any { + message, ok := value.(map[string]any) + if !ok { + if text := strings.TrimSpace(stringFromAny(value)); text != "" { + return map[string]any{"role": "user", "content": text} + } + return nil + } + messageType := strings.TrimSpace(stringFromAny(message["type"])) + switch messageType { + case "function_call": + return map[string]any{ + "role": "assistant", + "tool_calls": []any{map[string]any{ + "id": firstNonEmptyCacheAffinityString(message["call_id"], message["id"]), + "type": "function", + "function": map[string]any{ + "name": message["name"], + "arguments": message["arguments"], + }, + }}, + } + case "function_call_output": + return map[string]any{ + "role": "tool", + "tool_call_id": firstNonEmptyCacheAffinityString(message["call_id"], message["id"]), + "content": message["output"], + } + } + normalized := make(map[string]any, len(message)) + for key, item := range message { + switch key { + case "id", "status": + continue + case "type": + if messageType == "message" { + continue + } + } + normalized[key] = item + } + if role := strings.TrimSpace(stringFromAny(normalized["role"])); role != "" { + normalized["role"] = role + } + return normalized +} + +func firstNonEmptyCacheAffinityString(values ...any) string { + for _, value := range values { + if text := strings.TrimSpace(stringFromAny(value)); text != "" { + return text + } + } + return "" +} + +func cacheAffinityPrefixKeys(prefix string, baseHash string, messageHashes []string) []string { + if len(messageHashes) < cacheAffinityMinimumStableMessages { + return nil + } + keys := make([]string, 0, len(messageHashes)-cacheAffinityMinimumStableMessages+1) + for prefixLength := cacheAffinityMinimumStableMessages; prefixLength <= len(messageHashes); prefixLength++ { + keyPayload := map[string]any{ + "base": baseHash, + "messageHashes": messageHashes[:prefixLength], + "prefixLength": prefixLength, + } + keys = append(keys, prefix+stableJSONHash(keyPayload)) + } + return keys +} + +func legacyCacheAffinityPrefixKeys(kind string, modelType string, body map[string]any) []string { + messageHashes := legacyCacheAffinityMessageHashes(body) + if len(messageHashes) < cacheAffinityMinimumStableMessages { + return nil + } + basePayload := map[string]any{ + "kind": kind, + "modelType": modelType, + } + if tools, ok := body["tools"]; ok { + basePayload["tools"] = tools + } + if instructions, ok := body["instructions"]; ok { + basePayload["instructions"] = instructions + } + return cacheAffinityPrefixKeys("prompt_lcp:", stableJSONHash(basePayload), messageHashes) +} + +func legacyCacheAffinityMessageHashes(body map[string]any) []string { if hashes := stringListFromAny(body["messageHashes"]); len(hashes) > 0 { return hashes } @@ -102,20 +224,41 @@ func cacheAffinityMessageHashes(body map[string]any) []string { return hashes } messages, ok := body["messages"].([]any) - if !ok || len(messages) == 0 { + if !ok { return nil } hashes := make([]string, 0, len(messages)) for _, message := range messages { raw, err := json.Marshal(message) - if err != nil || len(raw) == 0 { - continue + if err == nil && len(raw) > 0 { + hashes = append(hashes, sha256Text(string(raw))) } - hashes = append(hashes, sha256Text(string(raw))) } return hashes } +func reverseCacheAffinityKeys(keys []string) []string { + reversed := make([]string, 0, len(keys)) + for index := len(keys) - 1; index >= 0; index-- { + reversed = append(reversed, keys[index]) + } + return reversed +} + +func normalizedCacheAffinityKeyList(keys []string) []string { + seen := map[string]bool{} + out := make([]string, 0, len(keys)) + for _, key := range keys { + key = strings.TrimSpace(key) + if key == "" || seen[key] { + continue + } + seen[key] = true + out = append(out, key) + } + return out +} + func stableJSONHash(value any) string { raw, err := json.Marshal(value) if err != nil || len(raw) == 0 { diff --git a/apps/api/internal/runner/cache_affinity_test.go b/apps/api/internal/runner/cache_affinity_test.go index e00280b..7292f8f 100644 --- a/apps/api/internal/runner/cache_affinity_test.go +++ b/apps/api/internal/runner/cache_affinity_test.go @@ -39,7 +39,7 @@ func TestBuildCacheAffinityKeyFallsBackToPromptHash(t *testing.T) { first := buildCacheAffinityKey("chat.completions", "text_generate", body) second := buildCacheAffinityKey("chat.completions", "text_generate", body) - if first == "" || first != second || !strings.HasPrefix(first, "prompt_lcp:") { + if first == "" || first != second || !strings.HasPrefix(first, "prompt_lcp_v2:") { t.Fatalf("expected stable prompt LCP cache affinity key, first=%q second=%q", first, second) } } @@ -54,10 +54,10 @@ func TestBuildCacheAffinityKeysUsesMessagePrefixLCPKeys(t *testing.T) { "tools": []any{map[string]any{"type": "function", "function": map[string]any{"name": "lookup"}}}, }) - if !strings.HasPrefix(first.Primary, "prompt_lcp:") { + if !strings.HasPrefix(first.Primary, "prompt_lcp_v2:") { t.Fatalf("expected LCP prompt key, got %+v", first) } - if len(first.Record) != 2 || len(next.Lookup) != 4 { + if len(first.Record) != 1 || len(next.Record) != 3 || len(next.Lookup) != 6 { t.Fatalf("expected per-message prefix keys, first=%+v next=%+v", first, next) } found := false @@ -71,3 +71,81 @@ func TestBuildCacheAffinityKeysUsesMessagePrefixLCPKeys(t *testing.T) { t.Fatalf("next lookup keys should include first request prefix key, first=%+v next=%+v", first, next) } } + +func TestBuildCacheAffinityKeysRequiresTwoStableMessages(t *testing.T) { + systemOnly := buildCacheAffinityKeys("chat.completions", "text_generate", map[string]any{ + "messages": []any{ + map[string]any{"role": "system", "content": "shared system prompt"}, + }, + }) + firstConversation := buildCacheAffinityKeys("chat.completions", "text_generate", map[string]any{ + "messages": []any{ + map[string]any{"role": "system", "content": "shared system prompt"}, + map[string]any{"role": "user", "content": "conversation A"}, + }, + }) + secondConversation := buildCacheAffinityKeys("chat.completions", "text_generate", map[string]any{ + "messages": []any{ + map[string]any{"role": "system", "content": "shared system prompt"}, + map[string]any{"role": "user", "content": "conversation B"}, + }, + }) + + if systemOnly.Primary != "" || len(systemOnly.Lookup) != 0 || len(systemOnly.Record) != 0 { + t.Fatalf("shared system prompt alone must not create affinity, got %+v", systemOnly) + } + if firstConversation.Primary == "" || secondConversation.Primary == "" { + t.Fatalf("two stable messages should create affinity, first=%+v second=%+v", firstConversation, secondConversation) + } + for _, firstKey := range firstConversation.Lookup { + for _, secondKey := range secondConversation.Lookup { + if firstKey == secondKey { + t.Fatalf("different conversations sharing only a system prompt must not intersect: key=%s", firstKey) + } + } + } +} + +func TestBuildCacheAffinityKeysNormalizesResponsesToolHistory(t *testing.T) { + tools := []any{map[string]any{ + "type": "function", + "function": map[string]any{ + "name": "lookup_weather", + }, + }} + chat := buildCacheAffinityKeys("chat.completions", "text_generate", map[string]any{ + "tools": tools, + "messages": []any{ + map[string]any{"role": "system", "content": "Use tools when helpful."}, + map[string]any{"role": "user", "content": "Weather in Hangzhou?"}, + }, + }) + responses := buildCacheAffinityKeys("responses", "text_generate", map[string]any{ + "tools": tools, + "instructions": "Use tools when helpful.", + "input": "Weather in Hangzhou?", + }) + continued := buildCacheAffinityKeys("responses", "text_generate", map[string]any{ + "tools": tools, + "instructions": "Use tools when helpful.", + "input": []any{ + map[string]any{"role": "user", "content": "Weather in Hangzhou?"}, + map[string]any{"type": "function_call", "call_id": "call_weather", "name": "lookup_weather", "arguments": `{"city":"Hangzhou"}`}, + map[string]any{"type": "function_call_output", "call_id": "call_weather", "output": `{"temperature":31}`}, + }, + }) + + if chat.Primary == "" || chat.Primary != responses.Primary { + t.Fatalf("chat and responses should share the V2 protocol-independent prefix, chat=%+v responses=%+v", chat, responses) + } + found := false + for _, key := range continued.Lookup { + if key == responses.Primary { + found = true + break + } + } + if !found { + t.Fatalf("responses tool continuation should retain its stable prefix, initial=%+v continued=%+v", responses, continued) + } +} diff --git a/apps/api/internal/runner/recording.go b/apps/api/internal/runner/recording.go index 1022b5c..0754b0d 100644 --- a/apps/api/internal/runner/recording.go +++ b/apps/api/internal/runner/recording.go @@ -170,6 +170,14 @@ func attemptMetrics(candidate store.RuntimeModelCandidate, attemptNo int, simula metrics["cacheAffinityCachedInputTokens"] = candidate.CacheAffinity.CachedInputTokens metrics["cacheAffinityBoost"] = candidate.CacheAffinity.Boost metrics["cacheAffinityApplied"] = candidate.CacheAffinity.Applied + metrics["cacheAffinityMatched"] = candidate.CacheAffinity.Applied + metrics["cacheAffinityCandidateCount"] = candidate.CacheAffinity.CandidateCount + if candidate.CacheAffinity.MatchedPrefixDepth > 0 { + metrics["cacheAffinityMatchedPrefixDepth"] = candidate.CacheAffinity.MatchedPrefixDepth + } + if candidate.CacheAffinity.OverrideReason != "" { + metrics["cacheAffinityOverrideReason"] = candidate.CacheAffinity.OverrideReason + } } if attemptNo > 0 { metrics["attempt"] = attemptNo diff --git a/apps/api/internal/runner/recording_test.go b/apps/api/internal/runner/recording_test.go index a024a52..1fd6f0d 100644 --- a/apps/api/internal/runner/recording_test.go +++ b/apps/api/internal/runner/recording_test.go @@ -13,16 +13,19 @@ func TestAttemptMetricsIncludesCacheAffinityCoefficient(t *testing.T) { 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, + Key: "prompt_lcp_v2:test", + RequestCount: 2, + CachedInputTokens: 7296, + EMAHitRatio: 0.91, + LastHitRatio: 0.92, + Confidence: 1, + Score: 1.74, + Boost: 20, + AdjustedPriority: 80, + MatchedPrefixDepth: 4, + CandidateCount: 2, + OverrideReason: "different_effective_priority_tier", + Applied: true, }, }, 1, false) @@ -35,4 +38,10 @@ func TestAttemptMetricsIncludesCacheAffinityCoefficient(t *testing.T) { if metrics["cacheAffinityHitRatio"] != 0.91 || metrics["cacheAffinityLastHitRatio"] != 0.92 { t.Fatalf("expected hit ratios in attempt metrics, got %+v", metrics) } + if metrics["cacheAffinityMatched"] != true || + metrics["cacheAffinityMatchedPrefixDepth"] != 4 || + metrics["cacheAffinityCandidateCount"] != 2 || + metrics["cacheAffinityOverrideReason"] != "different_effective_priority_tier" { + t.Fatalf("expected cache affinity routing diagnostics in attempt metrics, got %+v", metrics) + } } diff --git a/apps/api/internal/store/candidates.go b/apps/api/internal/store/candidates.go index 1d3e34c..d43e8e1 100644 --- a/apps/api/internal/store/candidates.go +++ b/apps/api/internal/store/candidates.go @@ -500,14 +500,15 @@ func applyRuntimeCandidateCacheAffinity(candidate *RuntimeModelCandidate, option key = options.CacheAffinityKey } affinity := RuntimeCandidateCacheAffinity{ - Key: key, - RequestCount: input.RequestCount, - InputTokens: input.InputTokens, - CachedInputTokens: input.CachedInputTokens, - EMAHitRatio: boundedRatio(input.EMAHitRatio), - LastHitRatio: boundedRatio(input.LastHitRatio), - LastObservedUnix: input.LastObservedUnix, - AdjustedPriority: float64(candidate.PlatformPriority), + Key: key, + RequestCount: input.RequestCount, + InputTokens: input.InputTokens, + CachedInputTokens: input.CachedInputTokens, + EMAHitRatio: boundedRatio(input.EMAHitRatio), + LastHitRatio: boundedRatio(input.LastHitRatio), + LastObservedUnix: input.LastObservedUnix, + AdjustedPriority: float64(candidate.PlatformPriority), + MatchedPrefixDepth: cacheAffinityMatchedPrefixDepth(key, options.CacheAffinityKeys), } minSamples := cacheAffinityMinSamples(options.CacheAffinityPolicy) hasObservedCachedHit := affinity.CachedInputTokens > 0 || affinity.LastHitRatio > 0 @@ -538,6 +539,32 @@ func applyRuntimeCandidateCacheAffinity(candidate *RuntimeModelCandidate, option candidate.CacheAffinity = affinity } +func cacheAffinityMatchedPrefixDepth(key string, lookup []string) int { + prefix := "" + switch { + case strings.HasPrefix(key, "prompt_lcp_v2:"): + prefix = "prompt_lcp_v2:" + case strings.HasPrefix(key, "prompt_lcp:"): + prefix = "prompt_lcp:" + default: + return 0 + } + matchingKeys := make([]string, 0, len(lookup)) + for _, candidateKey := range lookup { + if strings.HasPrefix(candidateKey, prefix) { + matchingKeys = append(matchingKeys, candidateKey) + } + } + for index, candidateKey := range matchingKeys { + if candidateKey == key { + return len(matchingKeys) - index + cacheAffinityMinimumPrefixDepth - 1 + } + } + return 0 +} + +const cacheAffinityMinimumPrefixDepth = 2 + func cacheAffinityPolicyEnabled(policy map[string]any, modelType string) bool { if enabled, ok := policy["enabled"].(bool); ok && !enabled { return false @@ -712,6 +739,9 @@ func sortRuntimeModelCandidates(items []RuntimeModelCandidate) { if aFull != bFull { return !aFull } + if items[i].PlatformPriority != items[j].PlatformPriority { + return items[i].PlatformPriority < items[j].PlatformPriority + } if items[i].CacheAffinity.Applied != items[j].CacheAffinity.Applied { return items[i].CacheAffinity.Applied } @@ -746,6 +776,32 @@ func sortRuntimeModelCandidates(items []RuntimeModelCandidate) { } return false }) + appliedCount := 0 + for index := range items { + if items[index].CacheAffinity.Applied { + appliedCount++ + } + } + for index := range items { + items[index].CacheAffinity.CandidateCount = appliedCount + } + if len(items) == 0 || items[0].CacheAffinity.Applied { + return + } + for index := 1; index < len(items); index++ { + if !items[index].CacheAffinity.Applied { + continue + } + switch { + case runtimeCandidateFull(items[index]) && !runtimeCandidateFull(items[0]): + items[0].CacheAffinity.OverrideReason = "capacity_tier_unavailable" + case items[index].PlatformPriority != items[0].PlatformPriority: + items[0].CacheAffinity.OverrideReason = "different_effective_priority_tier" + } + if items[0].CacheAffinity.OverrideReason != "" { + break + } + } } func runtimeCandidateFull(candidate RuntimeModelCandidate) bool { diff --git a/apps/api/internal/store/candidates_test.go b/apps/api/internal/store/candidates_test.go index 1ef7c05..c11689f 100644 --- a/apps/api/internal/store/candidates_test.go +++ b/apps/api/internal/store/candidates_test.go @@ -120,7 +120,7 @@ func TestRuntimeCandidateSortingPrefersCacheAffinityWithinAvailableCandidates(t }, { PlatformID: "cache-affinity", - PlatformPriority: 20, + PlatformPriority: 10, }, } applyRuntimeCandidateCacheAffinity(&candidates[0], ListModelCandidatesOptions{ @@ -148,7 +148,7 @@ func TestRuntimeCandidateSortingPrefersCacheAffinityWithinAvailableCandidates(t } } -func TestRuntimeCandidateSortingPromotesObservedCacheHitAbovePriority(t *testing.T) { +func TestRuntimeCandidateSortingKeepsEffectivePriorityAheadOfCacheAffinity(t *testing.T) { candidates := []RuntimeModelCandidate{ { PlatformID: "volces-priority", @@ -177,14 +177,17 @@ func TestRuntimeCandidateSortingPromotesObservedCacheHitAbovePriority(t *testing sortRuntimeModelCandidates(candidates) - if candidates[0].PlatformID != "deepseek-cache-hit" || !candidates[0].CacheAffinity.Applied { - t.Fatalf("expected observed cache hit to outrank platform priority, got %+v", candidates) + if candidates[0].PlatformID != "volces-priority" || candidates[0].CacheAffinity.Applied { + t.Fatalf("expected effective priority tier to remain first, got %+v", candidates) } - if candidates[0].CacheAffinity.Score <= 0 || candidates[0].CacheAffinity.Boost <= 0 { - t.Fatalf("expected observed cache hit to carry score and boost, got %+v", candidates[0].CacheAffinity) + if candidates[0].CacheAffinity.OverrideReason != "different_effective_priority_tier" { + t.Fatalf("expected priority override reason on selected candidate, got %+v", candidates[0].CacheAffinity) } - if candidates[1].PlatformID != "volces-priority" || candidates[1].CacheAffinity.Applied { - t.Fatalf("expected priority-only candidate second, got %+v", candidates) + if candidates[1].PlatformID != "deepseek-cache-hit" || !candidates[1].CacheAffinity.Applied { + t.Fatalf("expected cache affinity candidate to remain visible as fallback, got %+v", candidates) + } + if candidates[1].CacheAffinity.Score <= 0 || candidates[1].CacheAffinity.Boost <= 0 { + t.Fatalf("expected observed cache hit to retain score and boost diagnostics, got %+v", candidates[1].CacheAffinity) } } @@ -196,7 +199,7 @@ func TestRuntimeCandidateSortingOrdersMultipleCacheAffinityCandidatesByScore(t * }, { PlatformID: "higher-cache-score", - PlatformPriority: 50, + PlatformPriority: 1, }, } policy := map[string]any{"enabled": true, "minSamples": 1, "maxPriorityBoost": 20, "modelTypes": []any{"text_generate"}} @@ -260,6 +263,29 @@ func TestRuntimeCandidateSortingKeepsFullCacheAffinityCandidateAvoided(t *testin if candidates[1].PlatformID != "cache-affinity-full" || !candidates[1].LoadAvoided { t.Fatalf("expected full cache-affinity candidate to remain avoided fallback, got %+v", candidates) } + if candidates[0].CacheAffinity.OverrideReason != "capacity_tier_unavailable" { + t.Fatalf("expected capacity override reason on selected candidate, got %+v", candidates[0].CacheAffinity) + } +} + +func TestCacheAffinityMatchedPrefixDepthIgnoresLegacyCompatibilityKeys(t *testing.T) { + lookup := []string{ + "prompt_lcp_v2:depth-four", + "prompt_lcp_v2:depth-three", + "prompt_lcp_v2:depth-two", + "prompt_lcp:depth-four", + "prompt_lcp:depth-three", + "prompt_lcp:depth-two", + } + if depth := cacheAffinityMatchedPrefixDepth("prompt_lcp_v2:depth-three", lookup); depth != 3 { + t.Fatalf("V2 matched prefix depth=%d want=3", depth) + } + if depth := cacheAffinityMatchedPrefixDepth("prompt_lcp:depth-two", lookup); depth != 2 { + t.Fatalf("legacy matched prefix depth=%d want=2", depth) + } + if depth := cacheAffinityMatchedPrefixDepth("explicit:key", lookup); depth != 0 { + t.Fatalf("explicit key prefix depth=%d want=0", depth) + } } func TestDefaultRunnerPriorityDemotePolicyUsesAutoMode(t *testing.T) { diff --git a/apps/api/internal/store/runtime_types.go b/apps/api/internal/store/runtime_types.go index e2f5924..6162353 100644 --- a/apps/api/internal/store/runtime_types.go +++ b/apps/api/internal/store/runtime_types.go @@ -178,18 +178,21 @@ type RuntimeModelCandidate struct { } type RuntimeCandidateCacheAffinity struct { - Key string - RequestCount int - InputTokens int - CachedInputTokens int - EMAHitRatio float64 - LastHitRatio float64 - LastObservedUnix float64 - Confidence float64 - Score float64 - Boost float64 - AdjustedPriority float64 - Applied bool + Key string + RequestCount int + InputTokens int + CachedInputTokens int + EMAHitRatio float64 + LastHitRatio float64 + LastObservedUnix float64 + Confidence float64 + Score float64 + Boost float64 + AdjustedPriority float64 + MatchedPrefixDepth int + CandidateCount int + OverrideReason string + Applied bool } type RuntimeCandidateLoadMetrics struct { diff --git a/apps/api/internal/store/task_history_test.go b/apps/api/internal/store/task_history_test.go index e259c4a..a954f6c 100644 --- a/apps/api/internal/store/task_history_test.go +++ b/apps/api/internal/store/task_history_test.go @@ -128,7 +128,7 @@ WHERE task_id=$1::uuid } } -func TestTaskAttemptDropsDuplicateJSONSnapshots(t *testing.T) { +func TestTaskAttemptKeepsOnlyCompactDiagnostics(t *testing.T) { db := billingV2IntegrationStore(t) ctx := context.Background() user := &auth.User{ID: "attempt-minimal-" + uuid.NewString()} @@ -157,8 +157,8 @@ func TestTaskAttemptDropsDuplicateJSONSnapshots(t *testing.T) { Status: "failed", Retryable: true, StatusCode: 503, - Usage: map[string]any{"output_tokens": 100}, - Metrics: map[string]any{"trace": []any{map[string]any{"large": "duplicate"}}}, + Usage: map[string]any{"output_tokens": 100, "raw": strings.Repeat("duplicate", 100)}, + Metrics: map[string]any{"trace": []any{map[string]any{"large": "duplicate"}}, "cacheAffinityMatched": true, "cacheAffinityMatchedPrefixDepth": 3}, ResponseSnapshot: map[string]any{"result": "duplicate"}, ErrorCode: "upstream_error", ErrorMessage: strings.Repeat("错误", 4096), @@ -173,9 +173,16 @@ func TestTaskAttemptDropsDuplicateJSONSnapshots(t *testing.T) { t.Fatalf("attempts=%d", len(attempts)) } attempt := attempts[0] - if len(attempt.RequestSnapshot) != 0 || len(attempt.ResponseSnapshot) != 0 || - len(attempt.Usage) != 0 || len(attempt.Metrics) != 0 || len(attempt.PricingSnapshot) != 0 { - t.Fatalf("attempt contains duplicate JSON: %+v", attempt) + if len(attempt.RequestSnapshot) != 0 || len(attempt.ResponseSnapshot) != 0 || len(attempt.PricingSnapshot) != 0 { + t.Fatalf("attempt contains duplicate snapshots: %+v", attempt) + } + if len(attempt.Usage) != 1 || taskAttemptMetricInt(attempt.Usage, "output_tokens") != 100 { + t.Fatalf("attempt compact usage mismatch: %+v", attempt.Usage) + } + if len(attempt.Metrics) != 2 || + attempt.Metrics["cacheAffinityMatched"] != true || + taskAttemptMetricInt(attempt.Metrics, "cacheAffinityMatchedPrefixDepth") != 3 { + t.Fatalf("attempt compact metrics mismatch: %+v", attempt.Metrics) } if attempt.StatusCode != 503 || len(attempt.ErrorMessage) > 2048 { t.Fatalf("attempt status/error not preserved safely: %+v", attempt) diff --git a/apps/api/internal/store/tasks_runtime.go b/apps/api/internal/store/tasks_runtime.go index 417c292..48dfe38 100644 --- a/apps/api/internal/store/tasks_runtime.go +++ b/apps/api/internal/store/tasks_runtime.go @@ -996,14 +996,16 @@ func (s *Store) FinishTaskAttempt(ctx context.Context, input FinishTaskAttemptIn if statusCode == 0 { statusCode = taskAttemptMetricInt(input.Metrics, "statusCode") } + usageJSON, _ := json.Marshal(sanitizeJSONForStorage(minimalTaskAttemptUsage(input.Usage))) + metricsJSON, _ := json.Marshal(sanitizeJSONForStorage(minimalTaskAttemptMetrics(input.Metrics))) _, err := s.pool.Exec(ctx, ` UPDATE gateway_task_attempts SET status = $2::text, retryable = $3, request_id = NULLIF($4::text, ''), status_code = NULLIF($5::int, 0), - usage = '{}'::jsonb, - metrics = '{}'::jsonb, + usage = $11::jsonb, + metrics = $12::jsonb, response_snapshot = '{}'::jsonb, pricing_snapshot = '{}'::jsonb, request_fingerprint = NULL, @@ -1024,10 +1026,45 @@ WHERE id = $1::uuid`, input.ResponseDurationMS, input.ErrorCode, truncateUTF8Bytes(input.ErrorMessage, 2048), + usageJSON, + metricsJSON, ) return err } +func minimalTaskAttemptUsage(usage map[string]any) map[string]any { + return whitelistedTaskAttemptMap(usage, []string{ + "inputTokens", "promptTokens", "input_tokens", "prompt_tokens", + "cachedInputTokens", "cachedPromptTokens", "cached_input_tokens", "cached_tokens", + "cachedInputTokensKnown", + "outputTokens", "completionTokens", "output_tokens", "completion_tokens", + "totalTokens", "total_tokens", + }) +} + +func minimalTaskAttemptMetrics(metrics map[string]any) map[string]any { + return whitelistedTaskAttemptMap(metrics, []string{ + "platformPriority", "currentPriority", "loadRatio", "loadAvoided", + "cacheAffinityKey", "cacheAdjustedPriority", "cacheAffinitySamples", + "cacheAffinityScore", "cacheAffinityConfidence", "cacheAffinityHitRatio", + "cacheAffinityEMAHitRatio", "cacheAffinityLastHitRatio", + "cacheAffinityCachedInputTokens", "cacheAffinityBoost", + "cacheAffinityApplied", "cacheAffinityMatched", + "cacheAffinityCandidateCount", "cacheAffinityMatchedPrefixDepth", + "cacheAffinityOverrideReason", + }) +} + +func whitelistedTaskAttemptMap(input map[string]any, keys []string) map[string]any { + out := map[string]any{} + for _, key := range keys { + if value, ok := input[key]; ok { + out[key] = value + } + } + return out +} + func (s *Store) FinishTaskSuccess(ctx context.Context, input FinishTaskSuccessInput) (GatewayTask, error) { resultReport := sanitizeJSONForStorageWithReport(minimalTaskResult(input.Result)) if resultReport.BinaryCount > 0 { @@ -1041,6 +1078,8 @@ func (s *Store) FinishTaskSuccess(ctx context.Context, input FinishTaskSuccessIn billingsJSON, _ := json.Marshal(sanitizeJSONForStorage(input.Billings)) usageJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(input.Usage))) metricsJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(input.Metrics))) + attemptUsageJSON, _ := json.Marshal(sanitizeJSONForStorage(minimalTaskAttemptUsage(input.Usage))) + attemptMetricsJSON, _ := json.Marshal(sanitizeJSONForStorage(minimalTaskAttemptMetrics(input.Metrics))) billingSummaryJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(input.BillingSummary))) pricingSnapshotJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(input.PricingSnapshot))) finalChargeAmount := strings.TrimSpace(input.FinalChargeAmountText) @@ -1056,8 +1095,8 @@ SET status = 'succeeded', retryable = false, request_id = NULLIF($2, ''), status_code = NULL, - usage = '{}'::jsonb, - metrics = '{}'::jsonb, + usage = $6::jsonb, + metrics = $7::jsonb, response_snapshot = '{}'::jsonb, pricing_snapshot = '{}'::jsonb, request_fingerprint = NULL, @@ -1072,6 +1111,7 @@ SET status = 'succeeded', WHERE id = $1::uuid`, input.AttemptID, input.RequestID, nullableTime(input.ResponseStartedAt), nullableTime(input.ResponseFinishedAt), input.ResponseDurationMS, + attemptUsageJSON, attemptMetricsJSON, ); err != nil { return err }