feat: use message LCP cache affinity keys

apps/api/internal/httpapi/request_preparation.go

@@ -63,12 +63,14 @@ func (s *Server) prepareTaskRequest(ctx context.Context, r *http.Request, user *
 		return preparedTaskRequest{}, err
 	}
 	inputs := make([]store.ConversationMessageInput, 0, len(messages))
+	messageHashes := make([]any, 0, len(messages))
 	for _, rawMessage := range messages {
 		message, _ := rawMessage.(map[string]any)
 		if message == nil {
 			message = map[string]any{"content": rawMessage}
 		}
 		hash, assetHashes := canonicalConversationMessageHash(message)
+		messageHashes = append(messageHashes, hash)
 		inputs = append(inputs, store.ConversationMessageInput{
 			Hash:         hash,
 			Role:         stringFromRequestAny(message["role"]),
@@ -83,6 +85,7 @@ func (s *Server) prepareTaskRequest(ctx context.Context, r *http.Request, user *
 	preparedBody["conversationId"] = conversationKey
 	preparedBody["conversationRecordId"] = conversationID
 	preparedBody["messageRefs"] = messageRefsForRequest(refs)
+	preparedBody["messageHashes"] = messageHashes
 	preparedBody["newMessageCount"] = newCount
 	delete(preparedBody, "messages")
 	result.ConversationID = conversationID

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 {
+	if len(basePayload) <= 2 {
-		return ""
+		return cacheAffinityKeys{}
 	}
-	raw, err := json.Marshal(payload)
+	key := "prompt:" + stableJSONHash(basePayload)
-	if err != nil || len(raw) == 0 {
+	return cacheAffinityKeys{Primary: key, Lookup: []string{key}, Record: []string{key}}
-		return ""
-	}
-	return "prompt:" + sha256Text(string(raw))
 }
 
 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[:])

apps/api/internal/runner/cache_affinity_test.go

@@ -1,6 +1,9 @@
 package runner
 
-import "testing"
+import (
+	"strings"
+	"testing"
+)
 
 func TestBuildCacheAffinityKeyUsesExplicitKeyFirst(t *testing.T) {
 	body := map[string]any{
@@ -10,7 +13,7 @@ func TestBuildCacheAffinityKeyUsesExplicitKeyFirst(t *testing.T) {
 	}
 
 	got := buildCacheAffinityKey("chat.completions", "text_generate", body)
-	if got == "" || got[:9] != "explicit:" {
+	if !strings.HasPrefix(got, "explicit:") {
 		t.Fatalf("expected explicit cache affinity key, got %q", got)
 	}
 }
@@ -20,7 +23,7 @@ func TestBuildCacheAffinityKeyUsesSessionContinuity(t *testing.T) {
 		"session_id": "conversation-1",
 	})
 
-	if got == "" || got[:11] != "session_id:" {
+	if !strings.HasPrefix(got, "session_id:") {
 		t.Fatalf("expected session cache affinity key, got %q", got)
 	}
 }
@@ -36,7 +39,35 @@ func TestBuildCacheAffinityKeyFallsBackToPromptHash(t *testing.T) {
 
 	first := buildCacheAffinityKey("chat.completions", "text_generate", body)
 	second := buildCacheAffinityKey("chat.completions", "text_generate", body)
-	if first == "" || first != second || first[:7] != "prompt:" {
+	if first == "" || first != second || !strings.HasPrefix(first, "prompt_lcp:") {
-		t.Fatalf("expected stable prompt hash cache affinity key, first=%q second=%q", first, second)
+		t.Fatalf("expected stable prompt LCP cache affinity key, first=%q second=%q", first, second)
+	}
+}
+
+func TestBuildCacheAffinityKeysUsesMessagePrefixLCPKeys(t *testing.T) {
+	first := buildCacheAffinityKeys("chat.completions", "text_generate", map[string]any{
+		"messageHashes": []any{"system-hash", "user-hash"},
+		"tools":         []any{map[string]any{"type": "function", "function": map[string]any{"name": "lookup"}}},
+	})
+	next := buildCacheAffinityKeys("chat.completions", "text_generate", map[string]any{
+		"messageHashes": []any{"system-hash", "user-hash", "assistant-hash", "tool-hash"},
+		"tools":         []any{map[string]any{"type": "function", "function": map[string]any{"name": "lookup"}}},
+	})
+
+	if !strings.HasPrefix(first.Primary, "prompt_lcp:") {
+		t.Fatalf("expected LCP prompt key, got %+v", first)
+	}
+	if len(first.Record) != 2 || len(next.Lookup) != 4 {
+		t.Fatalf("expected per-message prefix keys, first=%+v next=%+v", first, next)
+	}
+	found := false
+	for _, key := range next.Lookup {
+		if key == first.Primary {
+			found = true
+			break
+		}
+	}
+	if !found {
+		t.Fatalf("next lookup keys should include first request prefix key, first=%+v next=%+v", first, next)
 	}
 }

apps/api/internal/runner/service.go

@@ -149,9 +149,10 @@ func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *aut
 	if err != nil {
 		return Result{}, err
 	}
-	cacheAffinityKey := buildCacheAffinityKey(task.Kind, modelType, body)
+	cacheAffinityKeys := buildCacheAffinityKeys(task.Kind, modelType, body)
 	candidates, err := s.store.ListModelCandidates(ctx, task.Model, modelType, user, store.ListModelCandidatesOptions{
-		CacheAffinityKey:    cacheAffinityKey,
+		CacheAffinityKey:    cacheAffinityKeys.Primary,
+		CacheAffinityKeys:   cacheAffinityKeys.Lookup,
 		CacheAffinityPolicy: runnerPolicy.CacheAffinityPolicy,
 	})
 	if err != nil {
@@ -324,7 +325,7 @@ candidatesLoop:
 				break candidatesLoop
 			}
 			candidateBody := preprocessing.Body
-			response, err := s.runCandidate(ctx, task, user, candidateBody, preprocessing.Log, candidate, nextAttemptNo, onDelta, singleSourceProtected, runnerPolicy.CacheAffinityPolicy)
+			response, err := s.runCandidate(ctx, task, user, candidateBody, preprocessing.Log, candidate, nextAttemptNo, onDelta, singleSourceProtected, runnerPolicy.CacheAffinityPolicy, cacheAffinityKeys.Record)
 			if err == nil {
 				attemptNo = nextAttemptNo
 				billings := s.billings(ctx, user, task.Kind, candidateBody, candidate, response, isSimulation(task, candidate))
@@ -535,7 +536,7 @@ candidatesLoop:
 	return Result{Task: failed, Output: failed.Result}, lastErr
 }
 
-func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user *auth.User, body map[string]any, preprocessing parameterPreprocessingLog, candidate store.RuntimeModelCandidate, attemptNo int, onDelta clients.StreamDelta, singleSourceProtected bool, cacheAffinityPolicy map[string]any) (clients.Response, error) {
+func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user *auth.User, body map[string]any, preprocessing parameterPreprocessingLog, candidate store.RuntimeModelCandidate, attemptNo int, onDelta clients.StreamDelta, singleSourceProtected bool, cacheAffinityPolicy map[string]any, cacheAffinityRecordKeys []string) (clients.Response, error) {
 	simulated := isSimulation(task, candidate)
 	baseAttemptMetrics := mergeMetrics(attemptMetrics(candidate, attemptNo, simulated), parameterPreprocessingMetrics(preprocessing))
 	reservations := s.rateLimitReservations(ctx, user, candidate, body)
@@ -769,6 +770,7 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
 	}
 	if err := s.store.RecordCacheAffinityObservation(context.WithoutCancel(ctx), store.CacheAffinityObservationInput{
 		CacheAffinityKey:       candidate.CacheAffinity.Key,
+		CacheAffinityKeys:      cacheAffinityRecordKeys,
 		CacheAffinityPolicy:    cacheAffinityPolicy,
 		PlatformID:             candidate.PlatformID,
 		PlatformModelID:        candidate.PlatformModelID,

apps/api/internal/store/cache_affinity.go

@@ -7,6 +7,7 @@ import (
 
 type CacheAffinityObservationInput struct {
 	CacheAffinityKey       string
+	CacheAffinityKeys      []string
 	CacheAffinityPolicy    map[string]any
 	PlatformID             string
 	PlatformModelID        string
@@ -19,9 +20,10 @@ type CacheAffinityObservationInput struct {
 
 func (s *Store) RecordCacheAffinityObservation(ctx context.Context, input CacheAffinityObservationInput) error {
 	input.CacheAffinityKey = strings.TrimSpace(input.CacheAffinityKey)
+	input.CacheAffinityKeys = normalizedCacheAffinityKeys(input.CacheAffinityKey, input.CacheAffinityKeys)
 	input.ClientID = strings.TrimSpace(input.ClientID)
 	input.ModelType = strings.TrimSpace(input.ModelType)
-	if input.CacheAffinityKey == "" || input.ClientID == "" || input.ModelType == "" {
+	if len(input.CacheAffinityKeys) == 0 || input.ClientID == "" || input.ModelType == "" {
 		return nil
 	}
 	if !input.CachedInputTokensKnown || !cacheAffinityPolicyEnabled(input.CacheAffinityPolicy, input.ModelType) {
@@ -39,6 +41,15 @@ func (s *Store) RecordCacheAffinityObservation(ctx context.Context, input CacheA
 	}
 	alpha := cacheAffinityEMAAlpha(input.CacheAffinityPolicy)
 	hitRatio := float64(input.CachedInputTokens) / float64(input.InputTokens)
+	for _, key := range input.CacheAffinityKeys {
+		if err := s.recordCacheAffinityObservationKey(ctx, input, key, hitRatio, alpha); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func (s *Store) recordCacheAffinityObservationKey(ctx context.Context, input CacheAffinityObservationInput, key string, hitRatio float64, alpha float64) error {
 	_, err := s.pool.Exec(ctx, `
 	INSERT INTO gateway_cache_affinity_stats (
 	  client_id, cache_affinity_key, platform_id, platform_model_id, model_type,
@@ -60,7 +71,7 @@ func (s *Store) RecordCacheAffinityObservation(ctx context.Context, input CacheA
 	    last_observed_at = now(),
 	    updated_at = now()`,
 		input.ClientID,
-		input.CacheAffinityKey,
+		key,
 		input.PlatformID,
 		input.PlatformModelID,
 		input.ModelType,

apps/api/internal/store/candidates.go

@@ -13,6 +13,7 @@ import (
 
 type ListModelCandidatesOptions struct {
 	CacheAffinityKey    string
+	CacheAffinityKeys   []string
 	CacheAffinityPolicy map[string]any
 }
 
@@ -46,6 +47,7 @@ func (s *Store) ListModelCandidates(ctx context.Context, model string, modelType
 	       COALESCE(s.waiting_count, 0)::float8,
 	       COALESCE(s.limiter_ratio, 0)::float8,
 	       COALESCE(EXTRACT(EPOCH FROM s.last_assigned_at), 0)::float8,
+	       COALESCE(ca.cache_affinity_key, '')::text,
 	       COALESCE(ca.request_count, 0)::float8,
 	       COALESCE(ca.input_tokens, 0)::float8,
 	       COALESCE(ca.cached_input_tokens, 0)::float8,
@@ -59,10 +61,15 @@ func (s *Store) ListModelCandidates(ctx context.Context, model string, modelType
 	LEFT JOIN model_runtime_policy_sets rp ON rp.id = COALESCE(m.runtime_policy_set_id, b.runtime_policy_set_id)
 	LEFT JOIN runtime_client_states s
 	  ON s.client_id = p.platform_key || ':' || $2::text || ':' || COALESCE(NULLIF(m.provider_model_name, ''), m.model_name)
-	LEFT JOIN gateway_cache_affinity_stats ca
+	LEFT JOIN LATERAL (
-	  ON ca.client_id = p.platform_key || ':' || $2::text || ':' || COALESCE(NULLIF(m.provider_model_name, ''), m.model_name)
+	  SELECT ca.cache_affinity_key, ca.request_count, ca.input_tokens, ca.cached_input_tokens, ca.ema_hit_ratio, ca.last_hit_ratio, ca.last_observed_at
-	  AND ca.cache_affinity_key = NULLIF($4::text, '')
+	  FROM unnest($4::text[]) WITH ORDINALITY AS affinity_keys(cache_affinity_key, affinity_rank)
-	  AND ($5::int <= 0 OR ca.last_observed_at >= now() - ($5::int * interval '1 second'))
+	  JOIN gateway_cache_affinity_stats ca ON ca.cache_affinity_key = affinity_keys.cache_affinity_key
+	  WHERE ca.client_id = p.platform_key || ':' || $2::text || ':' || COALESCE(NULLIF(m.provider_model_name, ''), m.model_name)
+	    AND ($5::int <= 0 OR ca.last_observed_at >= now() - ($5::int * interval '1 second'))
+	  ORDER BY affinity_keys.affinity_rank ASC, ca.cached_input_tokens DESC, ca.ema_hit_ratio DESC, ca.last_observed_at DESC
+	  LIMIT 1
+	) ca ON TRUE
 LEFT JOIN (
 	  SELECT scope_key, SUM(lease_value) AS active
 	  FROM gateway_concurrency_leases
@@ -175,7 +182,7 @@ WHERE p.status = 'enabled'
 	         COALESCE(s.running_count, 0) ASC,
 	         COALESCE(s.waiting_count, 0) ASC,
 	         COALESCE(s.last_assigned_at, to_timestamp(0)) ASC,
-	         m.created_at ASC`, exactModel, modelType, modelMatchKey, listOptions.CacheAffinityKey, cacheAffinityStaleAfterSeconds(listOptions.CacheAffinityPolicy))
+	         m.created_at ASC`, exactModel, modelType, modelMatchKey, listOptions.CacheAffinityKeys, cacheAffinityStaleAfterSeconds(listOptions.CacheAffinityPolicy))
 	if err != nil {
 		return nil, err
 	}
@@ -211,6 +218,7 @@ WHERE p.status = 'enabled'
 		var stateWaitingCount float64
 		var stateLimiterRatio float64
 		var lastAssignedUnix float64
+		var cacheAffinityKey string
 		var cacheRequestCount float64
 		var cacheInputTokens float64
 		var cacheCachedInputTokens float64
@@ -269,6 +277,7 @@ WHERE p.status = 'enabled'
 			&stateWaitingCount,
 			&stateLimiterRatio,
 			&lastAssignedUnix,
+			&cacheAffinityKey,
 			&cacheRequestCount,
 			&cacheInputTokens,
 			&cacheCachedInputTokens,
@@ -302,6 +311,7 @@ WHERE p.status = 'enabled'
 		item.WaitingCount = maxFloat(queuedWaiting, stateWaitingCount)
 		item.LastAssignedUnix = lastAssignedUnix
 		applyRuntimeCandidateCacheAffinity(&item, listOptions, runtimeCandidateCacheAffinityInput{
+			Key:               cacheAffinityKey,
 			RequestCount:      int(cacheRequestCount),
 			InputTokens:       int(cacheInputTokens),
 			CachedInputTokens: int(cacheCachedInputTokens),
@@ -421,6 +431,7 @@ type runtimeCandidateLoadInput struct {
 }
 
 type runtimeCandidateCacheAffinityInput struct {
+	Key               string
 	RequestCount      int
 	InputTokens       int
 	CachedInputTokens int
@@ -435,15 +446,38 @@ func normalizeListModelCandidatesOptions(modelType string, options ...ListModelC
 	}
 	out := options[0]
 	out.CacheAffinityKey = strings.TrimSpace(out.CacheAffinityKey)
-	if out.CacheAffinityKey == "" || !cacheAffinityPolicyEnabled(out.CacheAffinityPolicy, modelType) {
+	out.CacheAffinityKeys = normalizedCacheAffinityKeys(out.CacheAffinityKey, out.CacheAffinityKeys)
+	if len(out.CacheAffinityKeys) > 0 {
+		out.CacheAffinityKey = out.CacheAffinityKeys[0]
+	}
+	if len(out.CacheAffinityKeys) == 0 || !cacheAffinityPolicyEnabled(out.CacheAffinityPolicy, modelType) {
 		out.CacheAffinityKey = ""
+		out.CacheAffinityKeys = nil
+	}
+	return out
+}
+
+func normalizedCacheAffinityKeys(primary string, keys []string) []string {
+	seen := map[string]bool{}
+	out := make([]string, 0, len(keys)+1)
+	for _, key := range append([]string{primary}, keys...) {
+		key = strings.TrimSpace(key)
+		if key == "" || seen[key] {
+			continue
+		}
+		seen[key] = true
+		out = append(out, key)
 	}
 	return out
 }
 
 func applyRuntimeCandidateCacheAffinity(candidate *RuntimeModelCandidate, options ListModelCandidatesOptions, input runtimeCandidateCacheAffinityInput) {
+	key := strings.TrimSpace(input.Key)
+	if key == "" {
+		key = options.CacheAffinityKey
+	}
 	affinity := RuntimeCandidateCacheAffinity{
-		Key:               options.CacheAffinityKey,
+		Key:               key,
 		RequestCount:      input.RequestCount,
 		InputTokens:       input.InputTokens,
 		CachedInputTokens: input.CachedInputTokens,
@@ -455,7 +489,7 @@ func applyRuntimeCandidateCacheAffinity(candidate *RuntimeModelCandidate, option
 	minSamples := cacheAffinityMinSamples(options.CacheAffinityPolicy)
 	hasObservedCachedHit := affinity.CachedInputTokens > 0 || affinity.LastHitRatio > 0
 	hasSampledAffinity := input.RequestCount >= minSamples && affinity.EMAHitRatio > 0
-	if options.CacheAffinityKey == "" || input.RequestCount <= 0 || (!hasObservedCachedHit && !hasSampledAffinity) {
+	if key == "" || input.RequestCount <= 0 || (!hasObservedCachedHit && !hasSampledAffinity) {
 		candidate.CacheAffinity = affinity
 		return
 	}