fix(routing): 对齐缓存亲和力硬规则与审计

This commit is contained in:
2026-07-24 23:51:27 +08:00
parent d7a0ec56f5
commit 152885bbb6
10 changed files with 540 additions and 94 deletions
@@ -1746,6 +1746,7 @@ WHERE m.platform_id = $1::uuid
} }
restartModel := "worker-restart-" + suffixText restartModel := "worker-restart-" + suffixText
createSimulationTextBaseModel(t, server.URL, loginResponse.AccessToken, restartModel)
createSimulationTextPlatformModel( createSimulationTextPlatformModel(
t, t,
server.URL, server.URL,
@@ -1864,6 +1865,7 @@ func TestCacheAffinityRoutingFlow(t *testing.T) {
model := "cache-affinity-smoke-" + suffixText model := "cache-affinity-smoke-" + suffixText
cacheKey := "cache-affinity-key-" + 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{ lowPlatform, lowPlatformModelID := createSimulationTextPlatformModel(t, server.URL, loginResponse.AccessToken, "cache-low-"+suffixText, "Cache Affinity Low Priority", model, 20, map[string]any{
"rules": []map[string]any{ "rules": []map[string]any{
{"metric": "concurrent", "limit": 1, "leaseTtlSeconds": 120}, {"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) tools := []any{map[string]any{
_ = highPlatform "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{ sameKeyDetail := runCacheAffinityChatTask(t, ctx, testPool, server.URL, apiKeyResponse.Secret, model, cacheKey, "same-key-"+suffixText, map[string]any{
"inputTokens": 1000, "inputTokens": 1000,
"cachedInputTokens": 700, "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) 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{ newKeyDetail := runCacheAffinityChatTask(t, ctx, testPool, server.URL, apiKeyResponse.Secret, model, "cache-affinity-new-"+suffixText, "new-key-"+suffixText, map[string]any{
"inputTokens": 1000, "inputTokens": 1000,
"cachedInputTokens": 0, "cachedInputTokens": 0,
"outputTokens": 20, "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) 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, "cachedInputTokens": 0,
"outputTokens": 20, "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) 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 requestCount int
var inputTokens 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 SELECT request_count::int, input_tokens::int, cached_input_tokens::int, ema_hit_ratio::float8
FROM gateway_cache_affinity_stats FROM gateway_cache_affinity_stats
WHERE client_id = $1 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 { LIMIT 1`, lowPlatform.PlatformKey+":text_generate:"+model).Scan(&requestCount, &inputTokens, &cachedTokens, &emaHitRatio); err != nil {
t.Fatalf("read cache affinity stats: %v", err) t.Fatalf("read cache affinity stats: %v", err)
} }
@@ -2214,6 +2266,20 @@ type cacheAffinityTaskDetail struct {
} `json:"attempts"` } `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) { 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() t.Helper()
var platform cacheAffinityPlatformFixture 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) t.Fatalf("cache affinity platform was not created: %+v", platform)
} }
payload := map[string]any{ payload := map[string]any{
"canonicalModelKey": "openai:gpt-4o-mini", "canonicalModelKey": "openai:" + model,
"modelName": model, "modelName": model,
"providerModelName": model, "providerModelName": model,
"modelAlias": 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 { 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() t.Helper()
var detail cacheAffinityTaskDetail var detail cacheAffinityTaskDetail
doAPIV1ChatCompletionAndLoadTask(t, ctx, pool, baseURL, token, map[string]any{ requestBody := map[string]any{
"model": model, "model": model,
"runMode": "simulation", "runMode": "simulation",
"simulation": true, "simulation": true,
"simulationDurationMs": 5, "simulationDurationMs": 5,
"cacheAffinityKey": cacheAffinityKey, }
"simulationUsage": simulationUsage, for key, value := range body {
"messages": []map[string]any{{"role": "user", "content": "cache affinity route"}}, requestBody[key] = value
}, marker, http.StatusOK, nil, &detail) }
doAPIV1ChatCompletionAndLoadTask(t, ctx, pool, baseURL, token, requestBody, marker, http.StatusOK, nil, &detail)
return detail return detail
} }
+173 -30
View File
@@ -13,6 +13,8 @@ type cacheAffinityKeys struct {
Record []string Record []string
} }
const cacheAffinityMinimumStableMessages = 2
func buildCacheAffinityKey(kind string, modelType string, body map[string]any) string { func buildCacheAffinityKey(kind string, modelType string, body map[string]any) string {
return buildCacheAffinityKeys(kind, modelType, body).Primary return buildCacheAffinityKeys(kind, modelType, body).Primary
} }
@@ -37,49 +39,29 @@ func buildCacheAffinityKeys(kind string, modelType string, body map[string]any)
return cacheAffinityKeys{} return cacheAffinityKeys{}
} }
basePayload := map[string]any{ basePayload := map[string]any{
"kind": kind, "version": "v2",
"modelType": modelType,
} }
if tools, ok := body["tools"]; ok { if tools, ok := body["tools"]; ok {
basePayload["tools"] = tools basePayload["tools"] = tools
} }
if instructions, ok := body["instructions"]; ok {
basePayload["instructions"] = instructions
}
messageHashes := cacheAffinityMessageHashes(body) messageHashes := cacheAffinityMessageHashes(body)
if len(messageHashes) > 0 { if len(messageHashes) >= cacheAffinityMinimumStableMessages {
baseHash := stableJSONHash(basePayload) baseHash := stableJSONHash(basePayload)
keys := make([]string, 0, len(messageHashes)) keys := cacheAffinityPrefixKeys("prompt_lcp_v2:", baseHash, 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))
}
if len(keys) == 0 { if len(keys) == 0 {
return cacheAffinityKeys{} return cacheAffinityKeys{}
} }
lookup := make([]string, 0, len(keys)) lookup := reverseCacheAffinityKeys(keys)
for index := len(keys) - 1; index >= 0; index-- { legacyKeys := legacyCacheAffinityPrefixKeys(kind, modelType, body)
lookup = append(lookup, keys[index]) lookup = normalizedCacheAffinityKeyList(append(lookup, reverseCacheAffinityKeys(legacyKeys)...))
}
return cacheAffinityKeys{ return cacheAffinityKeys{
Primary: keys[len(keys)-1], Primary: keys[len(keys)-1],
Lookup: lookup, Lookup: lookup,
Record: keys, Record: keys,
} }
} }
if input, ok := body["input"]; ok {
basePayload["input"] = input
}
if len(basePayload) <= 2 {
return cacheAffinityKeys{} return cacheAffinityKeys{}
} }
key := "prompt:" + stableJSONHash(basePayload)
return cacheAffinityKeys{Primary: key, Lookup: []string{key}, Record: []string{key}}
}
func cacheAffinityPromptHashSupported(kind string, modelType string) bool { func cacheAffinityPromptHashSupported(kind string, modelType string) bool {
switch strings.TrimSpace(kind) { switch strings.TrimSpace(kind) {
@@ -95,6 +77,146 @@ func cacheAffinityPromptHashSupported(kind string, modelType string) bool {
} }
func cacheAffinityMessageHashes(body map[string]any) []string { 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 { if hashes := stringListFromAny(body["messageHashes"]); len(hashes) > 0 {
return hashes return hashes
} }
@@ -102,20 +224,41 @@ func cacheAffinityMessageHashes(body map[string]any) []string {
return hashes return hashes
} }
messages, ok := body["messages"].([]any) messages, ok := body["messages"].([]any)
if !ok || len(messages) == 0 { if !ok {
return nil return nil
} }
hashes := make([]string, 0, len(messages)) hashes := make([]string, 0, len(messages))
for _, message := range messages { for _, message := range messages {
raw, err := json.Marshal(message) raw, err := json.Marshal(message)
if err != nil || len(raw) == 0 { if err == nil && len(raw) > 0 {
continue
}
hashes = append(hashes, sha256Text(string(raw))) hashes = append(hashes, sha256Text(string(raw)))
} }
}
return hashes 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 { func stableJSONHash(value any) string {
raw, err := json.Marshal(value) raw, err := json.Marshal(value)
if err != nil || len(raw) == 0 { if err != nil || len(raw) == 0 {
@@ -39,7 +39,7 @@ func TestBuildCacheAffinityKeyFallsBackToPromptHash(t *testing.T) {
first := buildCacheAffinityKey("chat.completions", "text_generate", body) first := buildCacheAffinityKey("chat.completions", "text_generate", body)
second := 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) 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"}}}, "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) 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) t.Fatalf("expected per-message prefix keys, first=%+v next=%+v", first, next)
} }
found := false 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) 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)
}
}
+8
View File
@@ -170,6 +170,14 @@ func attemptMetrics(candidate store.RuntimeModelCandidate, attemptNo int, simula
metrics["cacheAffinityCachedInputTokens"] = candidate.CacheAffinity.CachedInputTokens metrics["cacheAffinityCachedInputTokens"] = candidate.CacheAffinity.CachedInputTokens
metrics["cacheAffinityBoost"] = candidate.CacheAffinity.Boost metrics["cacheAffinityBoost"] = candidate.CacheAffinity.Boost
metrics["cacheAffinityApplied"] = candidate.CacheAffinity.Applied 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 { if attemptNo > 0 {
metrics["attempt"] = attemptNo metrics["attempt"] = attemptNo
+10 -1
View File
@@ -13,7 +13,7 @@ func TestAttemptMetricsIncludesCacheAffinityCoefficient(t *testing.T) {
PlatformID: "deepseek-platform", PlatformID: "deepseek-platform",
PlatformPriority: 100, PlatformPriority: 100,
CacheAffinity: store.RuntimeCandidateCacheAffinity{ CacheAffinity: store.RuntimeCandidateCacheAffinity{
Key: "conversation:test", Key: "prompt_lcp_v2:test",
RequestCount: 2, RequestCount: 2,
CachedInputTokens: 7296, CachedInputTokens: 7296,
EMAHitRatio: 0.91, EMAHitRatio: 0.91,
@@ -22,6 +22,9 @@ func TestAttemptMetricsIncludesCacheAffinityCoefficient(t *testing.T) {
Score: 1.74, Score: 1.74,
Boost: 20, Boost: 20,
AdjustedPriority: 80, AdjustedPriority: 80,
MatchedPrefixDepth: 4,
CandidateCount: 2,
OverrideReason: "different_effective_priority_tier",
Applied: true, Applied: true,
}, },
}, 1, false) }, 1, false)
@@ -35,4 +38,10 @@ func TestAttemptMetricsIncludesCacheAffinityCoefficient(t *testing.T) {
if metrics["cacheAffinityHitRatio"] != 0.91 || metrics["cacheAffinityLastHitRatio"] != 0.92 { if metrics["cacheAffinityHitRatio"] != 0.91 || metrics["cacheAffinityLastHitRatio"] != 0.92 {
t.Fatalf("expected hit ratios in attempt metrics, got %+v", metrics) 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)
}
} }
+56
View File
@@ -508,6 +508,7 @@ func applyRuntimeCandidateCacheAffinity(candidate *RuntimeModelCandidate, option
LastHitRatio: boundedRatio(input.LastHitRatio), LastHitRatio: boundedRatio(input.LastHitRatio),
LastObservedUnix: input.LastObservedUnix, LastObservedUnix: input.LastObservedUnix,
AdjustedPriority: float64(candidate.PlatformPriority), AdjustedPriority: float64(candidate.PlatformPriority),
MatchedPrefixDepth: cacheAffinityMatchedPrefixDepth(key, options.CacheAffinityKeys),
} }
minSamples := cacheAffinityMinSamples(options.CacheAffinityPolicy) minSamples := cacheAffinityMinSamples(options.CacheAffinityPolicy)
hasObservedCachedHit := affinity.CachedInputTokens > 0 || affinity.LastHitRatio > 0 hasObservedCachedHit := affinity.CachedInputTokens > 0 || affinity.LastHitRatio > 0
@@ -538,6 +539,32 @@ func applyRuntimeCandidateCacheAffinity(candidate *RuntimeModelCandidate, option
candidate.CacheAffinity = affinity 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 { func cacheAffinityPolicyEnabled(policy map[string]any, modelType string) bool {
if enabled, ok := policy["enabled"].(bool); ok && !enabled { if enabled, ok := policy["enabled"].(bool); ok && !enabled {
return false return false
@@ -712,6 +739,9 @@ func sortRuntimeModelCandidates(items []RuntimeModelCandidate) {
if aFull != bFull { if aFull != bFull {
return !aFull 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 { if items[i].CacheAffinity.Applied != items[j].CacheAffinity.Applied {
return items[i].CacheAffinity.Applied return items[i].CacheAffinity.Applied
} }
@@ -746,6 +776,32 @@ func sortRuntimeModelCandidates(items []RuntimeModelCandidate) {
} }
return false 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 { func runtimeCandidateFull(candidate RuntimeModelCandidate) bool {
+35 -9
View File
@@ -120,7 +120,7 @@ func TestRuntimeCandidateSortingPrefersCacheAffinityWithinAvailableCandidates(t
}, },
{ {
PlatformID: "cache-affinity", PlatformID: "cache-affinity",
PlatformPriority: 20, PlatformPriority: 10,
}, },
} }
applyRuntimeCandidateCacheAffinity(&candidates[0], ListModelCandidatesOptions{ applyRuntimeCandidateCacheAffinity(&candidates[0], ListModelCandidatesOptions{
@@ -148,7 +148,7 @@ func TestRuntimeCandidateSortingPrefersCacheAffinityWithinAvailableCandidates(t
} }
} }
func TestRuntimeCandidateSortingPromotesObservedCacheHitAbovePriority(t *testing.T) { func TestRuntimeCandidateSortingKeepsEffectivePriorityAheadOfCacheAffinity(t *testing.T) {
candidates := []RuntimeModelCandidate{ candidates := []RuntimeModelCandidate{
{ {
PlatformID: "volces-priority", PlatformID: "volces-priority",
@@ -177,14 +177,17 @@ func TestRuntimeCandidateSortingPromotesObservedCacheHitAbovePriority(t *testing
sortRuntimeModelCandidates(candidates) sortRuntimeModelCandidates(candidates)
if candidates[0].PlatformID != "deepseek-cache-hit" || !candidates[0].CacheAffinity.Applied { if candidates[0].PlatformID != "volces-priority" || candidates[0].CacheAffinity.Applied {
t.Fatalf("expected observed cache hit to outrank platform priority, got %+v", candidates) t.Fatalf("expected effective priority tier to remain first, got %+v", candidates)
} }
if candidates[0].CacheAffinity.Score <= 0 || candidates[0].CacheAffinity.Boost <= 0 { if candidates[0].CacheAffinity.OverrideReason != "different_effective_priority_tier" {
t.Fatalf("expected observed cache hit to carry score and boost, got %+v", candidates[0].CacheAffinity) t.Fatalf("expected priority override reason on selected candidate, got %+v", candidates[0].CacheAffinity)
} }
if candidates[1].PlatformID != "volces-priority" || candidates[1].CacheAffinity.Applied { if candidates[1].PlatformID != "deepseek-cache-hit" || !candidates[1].CacheAffinity.Applied {
t.Fatalf("expected priority-only candidate second, got %+v", candidates) 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", PlatformID: "higher-cache-score",
PlatformPriority: 50, PlatformPriority: 1,
}, },
} }
policy := map[string]any{"enabled": true, "minSamples": 1, "maxPriorityBoost": 20, "modelTypes": []any{"text_generate"}} 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 { if candidates[1].PlatformID != "cache-affinity-full" || !candidates[1].LoadAvoided {
t.Fatalf("expected full cache-affinity candidate to remain avoided fallback, got %+v", candidates) 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) { func TestDefaultRunnerPriorityDemotePolicyUsesAutoMode(t *testing.T) {
+3
View File
@@ -189,6 +189,9 @@ type RuntimeCandidateCacheAffinity struct {
Score float64 Score float64
Boost float64 Boost float64
AdjustedPriority float64 AdjustedPriority float64
MatchedPrefixDepth int
CandidateCount int
OverrideReason string
Applied bool Applied bool
} }
+13 -6
View File
@@ -128,7 +128,7 @@ WHERE task_id=$1::uuid
} }
} }
func TestTaskAttemptDropsDuplicateJSONSnapshots(t *testing.T) { func TestTaskAttemptKeepsOnlyCompactDiagnostics(t *testing.T) {
db := billingV2IntegrationStore(t) db := billingV2IntegrationStore(t)
ctx := context.Background() ctx := context.Background()
user := &auth.User{ID: "attempt-minimal-" + uuid.NewString()} user := &auth.User{ID: "attempt-minimal-" + uuid.NewString()}
@@ -157,8 +157,8 @@ func TestTaskAttemptDropsDuplicateJSONSnapshots(t *testing.T) {
Status: "failed", Status: "failed",
Retryable: true, Retryable: true,
StatusCode: 503, StatusCode: 503,
Usage: map[string]any{"output_tokens": 100}, Usage: map[string]any{"output_tokens": 100, "raw": strings.Repeat("duplicate", 100)},
Metrics: map[string]any{"trace": []any{map[string]any{"large": "duplicate"}}}, Metrics: map[string]any{"trace": []any{map[string]any{"large": "duplicate"}}, "cacheAffinityMatched": true, "cacheAffinityMatchedPrefixDepth": 3},
ResponseSnapshot: map[string]any{"result": "duplicate"}, ResponseSnapshot: map[string]any{"result": "duplicate"},
ErrorCode: "upstream_error", ErrorCode: "upstream_error",
ErrorMessage: strings.Repeat("错误", 4096), ErrorMessage: strings.Repeat("错误", 4096),
@@ -173,9 +173,16 @@ func TestTaskAttemptDropsDuplicateJSONSnapshots(t *testing.T) {
t.Fatalf("attempts=%d", len(attempts)) t.Fatalf("attempts=%d", len(attempts))
} }
attempt := attempts[0] attempt := attempts[0]
if len(attempt.RequestSnapshot) != 0 || len(attempt.ResponseSnapshot) != 0 || if len(attempt.RequestSnapshot) != 0 || len(attempt.ResponseSnapshot) != 0 || len(attempt.PricingSnapshot) != 0 {
len(attempt.Usage) != 0 || len(attempt.Metrics) != 0 || len(attempt.PricingSnapshot) != 0 { t.Fatalf("attempt contains duplicate snapshots: %+v", attempt)
t.Fatalf("attempt contains duplicate JSON: %+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 { if attempt.StatusCode != 503 || len(attempt.ErrorMessage) > 2048 {
t.Fatalf("attempt status/error not preserved safely: %+v", attempt) t.Fatalf("attempt status/error not preserved safely: %+v", attempt)
+44 -4
View File
@@ -996,14 +996,16 @@ func (s *Store) FinishTaskAttempt(ctx context.Context, input FinishTaskAttemptIn
if statusCode == 0 { if statusCode == 0 {
statusCode = taskAttemptMetricInt(input.Metrics, "statusCode") statusCode = taskAttemptMetricInt(input.Metrics, "statusCode")
} }
usageJSON, _ := json.Marshal(sanitizeJSONForStorage(minimalTaskAttemptUsage(input.Usage)))
metricsJSON, _ := json.Marshal(sanitizeJSONForStorage(minimalTaskAttemptMetrics(input.Metrics)))
_, err := s.pool.Exec(ctx, ` _, err := s.pool.Exec(ctx, `
UPDATE gateway_task_attempts UPDATE gateway_task_attempts
SET status = $2::text, SET status = $2::text,
retryable = $3, retryable = $3,
request_id = NULLIF($4::text, ''), request_id = NULLIF($4::text, ''),
status_code = NULLIF($5::int, 0), status_code = NULLIF($5::int, 0),
usage = '{}'::jsonb, usage = $11::jsonb,
metrics = '{}'::jsonb, metrics = $12::jsonb,
response_snapshot = '{}'::jsonb, response_snapshot = '{}'::jsonb,
pricing_snapshot = '{}'::jsonb, pricing_snapshot = '{}'::jsonb,
request_fingerprint = NULL, request_fingerprint = NULL,
@@ -1024,10 +1026,45 @@ WHERE id = $1::uuid`,
input.ResponseDurationMS, input.ResponseDurationMS,
input.ErrorCode, input.ErrorCode,
truncateUTF8Bytes(input.ErrorMessage, 2048), truncateUTF8Bytes(input.ErrorMessage, 2048),
usageJSON,
metricsJSON,
) )
return err 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) { func (s *Store) FinishTaskSuccess(ctx context.Context, input FinishTaskSuccessInput) (GatewayTask, error) {
resultReport := sanitizeJSONForStorageWithReport(minimalTaskResult(input.Result)) resultReport := sanitizeJSONForStorageWithReport(minimalTaskResult(input.Result))
if resultReport.BinaryCount > 0 { if resultReport.BinaryCount > 0 {
@@ -1041,6 +1078,8 @@ func (s *Store) FinishTaskSuccess(ctx context.Context, input FinishTaskSuccessIn
billingsJSON, _ := json.Marshal(sanitizeJSONForStorage(input.Billings)) billingsJSON, _ := json.Marshal(sanitizeJSONForStorage(input.Billings))
usageJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(input.Usage))) usageJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(input.Usage)))
metricsJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(input.Metrics))) 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))) billingSummaryJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(input.BillingSummary)))
pricingSnapshotJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(input.PricingSnapshot))) pricingSnapshotJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(input.PricingSnapshot)))
finalChargeAmount := strings.TrimSpace(input.FinalChargeAmountText) finalChargeAmount := strings.TrimSpace(input.FinalChargeAmountText)
@@ -1056,8 +1095,8 @@ SET status = 'succeeded',
retryable = false, retryable = false,
request_id = NULLIF($2, ''), request_id = NULLIF($2, ''),
status_code = NULL, status_code = NULL,
usage = '{}'::jsonb, usage = $6::jsonb,
metrics = '{}'::jsonb, metrics = $7::jsonb,
response_snapshot = '{}'::jsonb, response_snapshot = '{}'::jsonb,
pricing_snapshot = '{}'::jsonb, pricing_snapshot = '{}'::jsonb,
request_fingerprint = NULL, request_fingerprint = NULL,
@@ -1072,6 +1111,7 @@ SET status = 'succeeded',
WHERE id = $1::uuid`, WHERE id = $1::uuid`,
input.AttemptID, input.RequestID, input.AttemptID, input.RequestID,
nullableTime(input.ResponseStartedAt), nullableTime(input.ResponseFinishedAt), input.ResponseDurationMS, nullableTime(input.ResponseStartedAt), nullableTime(input.ResponseFinishedAt), input.ResponseDurationMS,
attemptUsageJSON, attemptMetricsJSON,
); err != nil { ); err != nil {
return err return err
} }