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

2026-07-24 23:51:27 +08:00
parent d7a0ec56f5
commit 152885bbb6
10 changed files with 540 additions and 94 deletions
@@ -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 {