easyai-ai-gateway/apps/api/internal/runner/cache_affinity.go

package runner

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"strings"
)

type cacheAffinityKeys struct {
	Primary string
	Lookup  []string
	Record  []string
}

const cacheAffinityMinimumStableMessages = 2

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 cacheAffinityKeys{}
	}
	for _, key := range []string{"cacheAffinityKey", "cache_affinity_key"} {
		if value := strings.TrimSpace(stringFromAny(body[key])); 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 != "" {
			key := key + ":" + sha256Text(value)
			return cacheAffinityKeys{Primary: key, Lookup: []string{key}, Record: []string{key}}
		}
	}
	if !cacheAffinityPromptHashSupported(kind, modelType) {
		return cacheAffinityKeys{}
	}
	basePayload := map[string]any{
		"version": "v2",
	}
	if tools, ok := body["tools"]; ok {
		basePayload["tools"] = tools
	}
	messageHashes := cacheAffinityMessageHashes(body)
	if len(messageHashes) >= cacheAffinityMinimumStableMessages {
		baseHash := stableJSONHash(basePayload)
		keys := cacheAffinityPrefixKeys("prompt_lcp_v2:", baseHash, messageHashes)
		if len(keys) == 0 {
			return cacheAffinityKeys{}
		}
		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,
		}
	}
	return cacheAffinityKeys{}
}

func cacheAffinityPromptHashSupported(kind string, modelType string) bool {
	switch strings.TrimSpace(kind) {
	case "chat.completions", "responses":
		return true
	}
	switch strings.TrimSpace(modelType) {
	case "chat", "text_generate", "responses":
		return true
	default:
		return false
	}
}

func 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
	}
	if hashes := stringListFromAny(body["message_hashes"]); len(hashes) > 0 {
		return hashes
	}
	messages, ok := body["messages"].([]any)
	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 {
			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 {
		return ""
	}
	return sha256Text(string(raw))
}

func sha256Text(value string) string {
	sum := sha256.Sum256([]byte(value))
	return hex.EncodeToString(sum[:])
}