feat: use message LCP cache affinity keys

apps/api/internal/runner/cache_affinity.go

@@ -7,47 +7,78 @@ import (
 	"strings"
 )
 
+type cacheAffinityKeys struct {
+	Primary string
+	Lookup  []string
+	Record  []string
+}
+
 func buildCacheAffinityKey(kind string, modelType string, body map[string]any) string {
+	return buildCacheAffinityKeys(kind, modelType, body).Primary
+}
+
+func buildCacheAffinityKeys(kind string, modelType string, body map[string]any) cacheAffinityKeys {
 	if len(body) == 0 {
-		return ""
+		return cacheAffinityKeys{}
 	}
 	for _, key := range []string{"cacheAffinityKey", "cache_affinity_key"} {
 		if value := strings.TrimSpace(stringFromAny(body[key])); value != "" {
-			return "explicit:" + sha256Text(value)
+			key := "explicit:" + sha256Text(value)
+			return cacheAffinityKeys{Primary: key, Lookup: []string{key}, Record: []string{key}}
 		}
 	}
 	for _, key := range []string{"sessionId", "session_id", "conversationId", "conversation_id"} {
 		if value := strings.TrimSpace(stringFromAny(body[key])); value != "" {
-			return key + ":" + sha256Text(value)
+			key := key + ":" + sha256Text(value)
+			return cacheAffinityKeys{Primary: key, Lookup: []string{key}, Record: []string{key}}
 		}
 	}
 	if !cacheAffinityPromptHashSupported(kind, modelType) {
-		return ""
+		return cacheAffinityKeys{}
 	}
-	payload := map[string]any{
+	basePayload := 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
+		basePayload["tools"] = tools
 	}
 	if instructions, ok := body["instructions"]; ok {
-		payload["instructions"] = instructions
+		basePayload["instructions"] = instructions
+	}
+	messageHashes := cacheAffinityMessageHashes(body)
+	if len(messageHashes) > 0 {
+		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))
+		}
+		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])
+		}
+		return cacheAffinityKeys{
+			Primary: keys[len(keys)-1],
+			Lookup:  lookup,
+			Record:  keys,
+		}
 	}
 	if input, ok := body["input"]; ok {
-		payload["input"] = input
+		basePayload["input"] = input
 	}
-	if len(payload) <= 2 {
-		return ""
+	if len(basePayload) <= 2 {
+		return cacheAffinityKeys{}
 	}
-	raw, err := json.Marshal(payload)
-	if err != nil || len(raw) == 0 {
-		return ""
-	}
-	return "prompt:" + sha256Text(string(raw))
+	key := "prompt:" + stableJSONHash(basePayload)
+	return cacheAffinityKeys{Primary: key, Lookup: []string{key}, Record: []string{key}}
 }
 
 func cacheAffinityPromptHashSupported(kind string, modelType string) bool {
@@ -63,6 +94,36 @@ 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, ok := body["messages"].([]any)
+	if !ok || len(messages) == 0 {
+		return nil
+	}
+	hashes := make([]string, 0, len(messages))
+	for _, message := range messages {
+		raw, err := json.Marshal(message)
+		if err != nil || len(raw) == 0 {
+			continue
+		}
+		hashes = append(hashes, sha256Text(string(raw)))
+	}
+	return hashes
+}
+
+func stableJSONHash(value any) string {
+	raw, err := json.Marshal(value)
+	if err != nil || len(raw) == 0 {
+		return ""
+	}
+	return sha256Text(string(raw))
+}
+
 func sha256Text(value string) string {
 	sum := sha256.Sum256([]byte(value))
 	return hex.EncodeToString(sum[:])