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

package runner

import (
	"strings"
	"testing"
)

func TestBuildCacheAffinityKeyUsesExplicitKeyFirst(t *testing.T) {
	body := map[string]any{
		"cacheAffinityKey": "session-a",
		"sessionId":        "session-b",
		"messages":         []any{map[string]any{"role": "user", "content": "hello"}},
	}

	got := buildCacheAffinityKey("chat.completions", "text_generate", body)
	if !strings.HasPrefix(got, "explicit:") {
		t.Fatalf("expected explicit cache affinity key, got %q", got)
	}
}

func TestBuildCacheAffinityKeyUsesSessionContinuity(t *testing.T) {
	got := buildCacheAffinityKey("chat.completions", "text_generate", map[string]any{
		"session_id": "conversation-1",
	})

	if !strings.HasPrefix(got, "session_id:") {
		t.Fatalf("expected session cache affinity key, got %q", got)
	}
}

func TestBuildCacheAffinityKeyFallsBackToPromptHash(t *testing.T) {
	body := map[string]any{
		"messages": []any{
			map[string]any{"role": "system", "content": "stable"},
			map[string]any{"role": "user", "content": "hello"},
		},
		"tools": []any{map[string]any{"type": "function", "function": map[string]any{"name": "lookup"}}},
	}

	first := buildCacheAffinityKey("chat.completions", "text_generate", body)
	second := buildCacheAffinityKey("chat.completions", "text_generate", body)
	if first == "" || first != second || !strings.HasPrefix(first, "prompt_lcp_v2:") {
		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_v2:") {
		t.Fatalf("expected LCP prompt key, got %+v", first)
	}
	if len(first.Record) != 1 || len(next.Record) != 3 || len(next.Lookup) != 6 {
		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)
	}
}

func TestBuildCacheAffinityKeysRequiresTwoStableMessages(t *testing.T) {
	systemOnly := buildCacheAffinityKeys("chat.completions", "text_generate", map[string]any{
		"messages": []any{
			map[string]any{"role": "system", "content": "shared system prompt"},
		},
	})
	firstConversation := buildCacheAffinityKeys("chat.completions", "text_generate", map[string]any{
		"messages": []any{
			map[string]any{"role": "system", "content": "shared system prompt"},
			map[string]any{"role": "user", "content": "conversation A"},
		},
	})
	secondConversation := buildCacheAffinityKeys("chat.completions", "text_generate", map[string]any{
		"messages": []any{
			map[string]any{"role": "system", "content": "shared system prompt"},
			map[string]any{"role": "user", "content": "conversation B"},
		},
	})

	if systemOnly.Primary != "" || len(systemOnly.Lookup) != 0 || len(systemOnly.Record) != 0 {
		t.Fatalf("shared system prompt alone must not create affinity, got %+v", systemOnly)
	}
	if firstConversation.Primary == "" || secondConversation.Primary == "" {
		t.Fatalf("two stable messages should create affinity, first=%+v second=%+v", firstConversation, secondConversation)
	}
	for _, firstKey := range firstConversation.Lookup {
		for _, secondKey := range secondConversation.Lookup {
			if firstKey == secondKey {
				t.Fatalf("different conversations sharing only a system prompt must not intersect: key=%s", firstKey)
			}
		}
	}
}

func TestBuildCacheAffinityKeysNormalizesResponsesToolHistory(t *testing.T) {
	tools := []any{map[string]any{
		"type": "function",
		"function": map[string]any{
			"name": "lookup_weather",
		},
	}}
	chat := buildCacheAffinityKeys("chat.completions", "text_generate", map[string]any{
		"tools": tools,
		"messages": []any{
			map[string]any{"role": "system", "content": "Use tools when helpful."},
			map[string]any{"role": "user", "content": "Weather in Hangzhou?"},
		},
	})
	responses := buildCacheAffinityKeys("responses", "text_generate", map[string]any{
		"tools":        tools,
		"instructions": "Use tools when helpful.",
		"input":        "Weather in Hangzhou?",
	})
	continued := buildCacheAffinityKeys("responses", "text_generate", map[string]any{
		"tools":        tools,
		"instructions": "Use tools when helpful.",
		"input": []any{
			map[string]any{"role": "user", "content": "Weather in Hangzhou?"},
			map[string]any{"type": "function_call", "call_id": "call_weather", "name": "lookup_weather", "arguments": `{"city":"Hangzhou"}`},
			map[string]any{"type": "function_call_output", "call_id": "call_weather", "output": `{"temperature":31}`},
		},
	})

	if chat.Primary == "" || chat.Primary != responses.Primary {
		t.Fatalf("chat and responses should share the V2 protocol-independent prefix, chat=%+v responses=%+v", chat, responses)
	}
	found := false
	for _, key := range continued.Lookup {
		if key == responses.Primary {
			found = true
			break
		}
	}
	if !found {
		t.Fatalf("responses tool continuation should retain its stable prefix, initial=%+v continued=%+v", responses, continued)
	}
}