diff --git a/apps/api/internal/httpapi/request_preparation.go b/apps/api/internal/httpapi/request_preparation.go index 662773d..f0c2a36 100644 --- a/apps/api/internal/httpapi/request_preparation.go +++ b/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 diff --git a/apps/api/internal/runner/cache_affinity.go b/apps/api/internal/runner/cache_affinity.go index a82910b..90a2b6c 100644 --- a/apps/api/internal/runner/cache_affinity.go +++ b/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[:]) diff --git a/apps/api/internal/runner/cache_affinity_test.go b/apps/api/internal/runner/cache_affinity_test.go index babfb47..e00280b 100644 --- a/apps/api/internal/runner/cache_affinity_test.go +++ b/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:" { - t.Fatalf("expected stable prompt hash cache affinity key, first=%q second=%q", first, second) + if first == "" || first != second || !strings.HasPrefix(first, "prompt_lcp:") { + 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) } } diff --git a/apps/api/internal/runner/service.go b/apps/api/internal/runner/service.go index bec0f01..2e7255d 100644 --- a/apps/api/internal/runner/service.go +++ b/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, diff --git a/apps/api/internal/store/cache_affinity.go b/apps/api/internal/store/cache_affinity.go index 6d949e9..09a6138 100644 --- a/apps/api/internal/store/cache_affinity.go +++ b/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, diff --git a/apps/api/internal/store/candidates.go b/apps/api/internal/store/candidates.go index 872d749..e15d92f 100644 --- a/apps/api/internal/store/candidates.go +++ b/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 - ON ca.client_id = p.platform_key || ':' || $2::text || ':' || COALESCE(NULLIF(m.provider_model_name, ''), m.model_name) - AND ca.cache_affinity_key = NULLIF($4::text, '') - AND ($5::int <= 0 OR ca.last_observed_at >= now() - ($5::int * interval '1 second')) + LEFT JOIN LATERAL ( + 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 + FROM unnest($4::text[]) WITH ORDINALITY AS affinity_keys(cache_affinity_key, affinity_rank) + 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 }