llmgw/internal/router/cache_test.go

package router

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/prometheus/client_golang/prometheus/testutil"

	"gitea.mynodes.xyz/ilker/llmgw/internal/guard"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router/proxy"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// cachingRouter is pricedRouter with its project reusing answers for an
// hour, reproducible ones only unless anyAnswer.
func cachingRouter(t *testing.T, engineURL string, anyAnswer bool) (*Router, string) {
	t.Helper()
	rt, secret := pricedRouter(t, engineURL)
	if err := rt.store.SetAnswerCache(t.Context(), mustProjectOf(t, rt, "key-1"), store.AnswerCache{TTL: time.Hour, Any: anyAnswer}); err != nil {
		t.Fatal(err)
	}
	return rt, secret
}

func ask(temperature any, text string) map[string]any {
	b := map[string]any{"model": "q14", "messages": []any{map[string]any{"role": "user", "content": text}}}
	if temperature != nil {
		b["temperature"] = temperature
	}
	return b
}

// usageOf waits for the record of one request.
func usageOf(t *testing.T, rt *Router, id string) store.Usage {
	t.Helper()
	deadline := time.Now().Add(5 * time.Second)
	for time.Now().Before(deadline) {
		if u, ok, err := rt.store.UsageByRequest(t.Context(), id); err == nil && ok {
			return u
		}
		time.Sleep(20 * time.Millisecond)
	}
	t.Fatalf("no usage was recorded for %s", id)
	return store.Usage{}
}

func TestAProjectThatReusesAnswersAnswersTheSameRequestWithoutTheModel(t *testing.T) {
	e := newFakeEngine(t, "qwen2.5-14b-instruct")
	rt, secret := cachingRouter(t, e.srv.URL, false)
	h := rt.Handler()
	auth := func(id string) map[string]string {
		return map[string]string{"Authorization": "Bearer " + secret, "X-Request-Id": id}
	}

	first := post(t, h, ask(0, "what is 2+2"), auth("first"))
	if first.Code != 200 || first.Header().Get(cacheHeader) != "miss" {
		t.Fatalf("first: %d %q %s", first.Code, first.Header().Get(cacheHeader), first.Body.String())
	}
	// The same request, said differently: keys in another order, another
	// end user, the temperature written as a float, by the canonical name.
	again := map[string]any{"messages": []any{map[string]any{"content": "what is 2+2", "role": "user"}},
		"temperature": 0.0, "model": "qwen2.5-14b-instruct", "user": "someone-else"}
	second := post(t, h, again, auth("second"))
	if second.Code != 200 || second.Header().Get(cacheHeader) != "hit" {
		t.Fatalf("second: %d %q %s", second.Code, second.Header().Get(cacheHeader), second.Body.String())
	}
	if e.calls.Load() != 1 {
		t.Errorf("the engine was asked %d times; the second answer should have been reused", e.calls.Load())
	}
	if second.Body.String() != first.Body.String() {
		t.Errorf("reused %s, first was %s", second.Body.String(), first.Body.String())
	}
	if second.Header().Get("Age") == "" {
		t.Error("a reused answer did not say its age")
	}

	// Recorded, as a request that generated nothing.
	miss, hit := usageOf(t, rt, "first"), usageOf(t, rt, "second")
	if miss.CacheHit || miss.CostMicros == 0 || miss.PromptTokens == 0 {
		t.Errorf("the answer the model gave: %+v", miss)
	}
	if !hit.CacheHit || hit.CostMicros != 0 || hit.PromptTokens != 0 || hit.CompletionTokens != 0 ||
		hit.Status != 200 || hit.ServedModel != "qwen2.5-14b-instruct" {
		t.Errorf("the reused answer: %+v", hit)
	}

	// A sampled answer is asked for each time, and says so.
	for i := range 2 {
		rec := post(t, h, ask(0.7, "write a poem"), auth(fmt.Sprintf("poem-%d", i)))
		if rec.Header().Get(cacheHeader) != "bypass" {
			t.Errorf("a sampled request: %q", rec.Header().Get(cacheHeader))
		}
	}
	if e.calls.Load() != 3 {
		t.Errorf("the engine was asked %d times, expected each sampled request to reach it", e.calls.Load())
	}
	// A seed makes it reproducible.
	seeded := ask(0.7, "write a poem")
	seeded["seed"] = 42
	post(t, h, seeded, auth("seed-1"))
	if rec := post(t, h, seeded, auth("seed-2")); rec.Header().Get(cacheHeader) != "hit" {
		t.Errorf("a seeded request: %q", rec.Header().Get(cacheHeader))
	}
	// A client that asks for a fresh answer gets one.
	hdr := auth("fresh")
	hdr["Cache-Control"] = "no-cache"
	if rec := post(t, h, ask(0, "what is 2+2"), hdr); rec.Header().Get(cacheHeader) != "bypass" {
		t.Errorf("Cache-Control: no-cache: %q", rec.Header().Get(cacheHeader))
	}
	hdr = auth("fresh-2")
	hdr[cacheHeader] = "no"
	if rec := post(t, h, ask(0, "what is 2+2"), hdr); rec.Header().Get(cacheHeader) != "bypass" {
		t.Errorf("x-llmgw-cache: no: %q", rec.Header().Get(cacheHeader))
	}

	if n := testutil.ToFloat64(rt.m.cacheRequests.WithLabelValues("qwen2.5-14b-instruct", "hit")); n != 2 {
		t.Errorf("%v hits counted", n)
	}
	if n := testutil.ToFloat64(rt.m.cacheRequests.WithLabelValues("qwen2.5-14b-instruct", "miss")); n != 2 {
		t.Errorf("%v misses counted", n)
	}
	if n := testutil.ToFloat64(rt.m.cacheBytes); n <= 0 {
		t.Errorf("the cache holds %v bytes", n)
	}

	// Turned off, nothing is reused, and nothing is said about it.
	_ = rt.store.SetAnswerCache(t.Context(), mustProjectOf(t, rt, "key-1"), store.AnswerCache{})
	rt.reqLogs.by = nil
	calls := e.calls.Load()
	if rec := post(t, h, ask(0, "what is 2+2"), auth("off")); rec.Header().Get(cacheHeader) != "" || e.calls.Load() != calls+1 {
		t.Errorf("a project that reuses nothing: %q, %d calls", rec.Header().Get(cacheHeader), e.calls.Load()-calls)
	}
}

func TestAProjectThatReusesAnyAnswerReusesSampledOnes(t *testing.T) {
	e := newFakeEngine(t, "qwen2.5-14b-instruct")
	rt, secret := cachingRouter(t, e.srv.URL, true)
	auth := map[string]string{"Authorization": "Bearer " + secret}
	post(t, rt.Handler(), ask(nil, "write a poem"), auth)
	if rec := post(t, rt.Handler(), ask(nil, "write a poem"), auth); rec.Header().Get(cacheHeader) != "hit" {
		t.Errorf("%q", rec.Header().Get(cacheHeader))
	}
}

func TestKeysOfOneProjectNeverSeeAnotherProjectsAnswers(t *testing.T) {
	e := newFakeEngine(t, "qwen2.5-14b-instruct")
	rt, secret := cachingRouter(t, e.srv.URL, false)
	team, err := rt.store.CreateTeam(t.Context(), mustOrgOf(t, rt), "other")
	if err != nil {
		t.Fatal(err)
	}
	other, err := rt.store.CreateProject(t.Context(), team, "elsewhere")
	if err != nil {
		t.Fatal(err)
	}
	_ = rt.store.SetAnswerCache(t.Context(), other, store.AnswerCache{TTL: time.Hour, Any: true})
	secret2 := store.KeyPrefix + strings.Repeat("b", 43)
	sum := sha256.Sum256([]byte(secret2))
	if err := rt.store.CreateKey(t.Context(), store.Key{ID: "key-2", Hash: sum[:], Prefix: "bbbb", Name: "other",
		ProjectID: other, CreatedAt: time.Now()}); err != nil {
		t.Fatal(err)
	}

	post(t, rt.Handler(), ask(0, "the secret plan"), map[string]string{"Authorization": "Bearer " + secret})
	rec := post(t, rt.Handler(), ask(0, "the secret plan"), map[string]string{"Authorization": "Bearer " + secret2})
	if rec.Header().Get(cacheHeader) != "miss" || e.calls.Load() != 2 {
		t.Errorf("another project's key: %q, the engine asked %d times", rec.Header().Get(cacheHeader), e.calls.Load())
	}
}

func mustOrgOf(t *testing.T, rt *Router) string {
	t.Helper()
	orgs, err := rt.store.Tree(t.Context())
	if err != nil || len(orgs) == 0 {
		t.Fatalf("no organization: %v", err)
	}
	return orgs[0].ID
}

func TestAFailedAnswerIsNotReused(t *testing.T) {
	e := newFakeEngine(t, "qwen2.5-14b-instruct")
	rt, secret := cachingRouter(t, e.srv.URL, false)
	auth := map[string]string{"Authorization": "Bearer " + secret}
	e.status.Store(http.StatusBadRequest)
	if rec := post(t, rt.Handler(), ask(0, "hello"), auth); rec.Code == 200 {
		t.Fatalf("%d", rec.Code)
	}
	e.status.Store(0)
	if rec := post(t, rt.Handler(), ask(0, "hello"), auth); rec.Code != 200 || rec.Header().Get(cacheHeader) != "miss" {
		t.Errorf("after a failure: %d %q", rec.Code, rec.Header().Get(cacheHeader))
	}
}

// switchGuard blocks every request, or puts every one under no_store, when
// told to.
type switchGuard struct {
	block, noStore atomic.Bool
	// answers, when set, stops answers saying "hi".
	answers atomic.Bool
}

func (g *switchGuard) Input(context.Context, guard.Subject, []byte) (guard.Verdict, error) {
	if g.block.Load() {
		return guard.Verdict{Action: guard.Block, Policy: "nothing"}, nil
	}
	return guard.Verdict{}, nil
}

func (g *switchGuard) Output(context.Context, guard.Subject) guard.OutputFilter {
	if !g.answers.Load() {
		return nil
	}
	return &stopHi{}
}

func (g *switchGuard) Directives(guard.Subject) guard.Directives {
	return guard.Directives{NoStore: g.noStore.Load()}
}

func TestTheGuardrailsStillDecideAboutAnAnswerFromTheCache(t *testing.T) {
	e := newFakeEngine(t, "qwen2.5-14b-instruct")
	rt, secret := cachingRouter(t, e.srv.URL, false)
	g := &switchGuard{}
	rt.SetGuard(g)
	auth := map[string]string{"Authorization": "Bearer " + secret}
	if rec := post(t, rt.Handler(), ask(0, "hello"), auth); rec.Header().Get(cacheHeader) != "miss" {
		t.Fatalf("%q", rec.Header().Get(cacheHeader))
	}

	// A policy that now refuses the request refuses it, answer kept or not.
	g.block.Store(true)
	if rec := post(t, rt.Handler(), ask(0, "hello"), auth); rec.Code != http.StatusBadRequest || rec.Header().Get(cacheHeader) != "" {
		t.Errorf("a blocked request: %d %q %s", rec.Code, rec.Header().Get(cacheHeader), rec.Body.String())
	}
	g.block.Store(false)

	// A policy that keeps nothing neither reads the cache nor fills it.
	g.noStore.Store(true)
	if rec := post(t, rt.Handler(), ask(0, "hello"), auth); rec.Code != 200 || rec.Header().Get(cacheHeader) != "bypass" {
		t.Errorf("under no_store: %d %q", rec.Code, rec.Header().Get(cacheHeader))
	}
	if rec := post(t, rt.Handler(), ask(0, "under no_store"), auth); rec.Header().Get(cacheHeader) != "bypass" {
		t.Errorf("under no_store: %q", rec.Header().Get(cacheHeader))
	}
	g.noStore.Store(false)
	if rec := post(t, rt.Handler(), ask(0, "under no_store"), auth); rec.Header().Get(cacheHeader) != "miss" {
		t.Errorf("an answer given under no_store was kept: %q", rec.Header().Get(cacheHeader))
	}

	// An output policy added since judges the reused answer too.
	g.answers.Store(true)
	rec := post(t, rt.Handler(), ask(0, "hello"), auth)
	if rec.Header().Get(cacheHeader) != "hit" || !strings.Contains(rec.Body.String(), "content_filter") ||
		rec.Header().Get(guardHeader) != "blocked" {
		t.Errorf("an answer the output policy stops: %q %q %s", rec.Header().Get(cacheHeader),
			rec.Header().Get(guardHeader), rec.Body.String())
	}
	// And an answer it stopped is not kept.
	if rec := post(t, rt.Handler(), ask(0, "new question"), auth); rec.Header().Get(cacheHeader) != "miss" {
		t.Fatalf("%q", rec.Header().Get(cacheHeader))
	}
	g.answers.Store(false)
	if rec := post(t, rt.Handler(), ask(0, "new question"), auth); rec.Header().Get(cacheHeader) != "miss" {
		t.Errorf("an answer a guard stopped was kept: %q", rec.Header().Get(cacheHeader))
	}
}

// greetingEngine streams "Hello world" to a request that asks for a
// stream, with usage when asked, and answers whole otherwise.
func greetingEngine(t *testing.T) (*httptest.Server, *atomic.Int64) {
	t.Helper()
	var calls atomic.Int64
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/metrics" {
			return
		}
		calls.Add(1)
		var b struct {
			Stream        bool `json:"stream"`
			StreamOptions struct {
				IncludeUsage bool `json:"include_usage"`
			} `json:"stream_options"`
		}
		body, _ := io.ReadAll(r.Body)
		_ = json.Unmarshal(body, &b)
		if !b.Stream {
			w.Header().Set("Content-Type", "application/json")
			_, _ = io.WriteString(w, `{"id":"c1","object":"chat.completion","created":1,"model":"qwen2.5-14b-instruct",`+
				`"choices":[{"index":0,"message":{"role":"assistant","content":"Hello world"},"finish_reason":"stop"}],`+
				`"usage":{"prompt_tokens":5,"completion_tokens":2,"total_tokens":7}}`)
			return
		}
		w.Header().Set("Content-Type", "text/event-stream")
		head := `{"id":"c1","object":"chat.completion.chunk","created":1,"model":"qwen2.5-14b-instruct","choices":`
		for _, c := range []string{
			`[{"index":0,"delta":{"role":"assistant","content":""},"finish_reason":null}]}`,
			`[{"index":0,"delta":{"content":"Hello"},"finish_reason":null}]}`,
			`[{"index":0,"delta":{"content":" world"},"finish_reason":null}]}`,
			`[{"index":0,"delta":{},"finish_reason":"stop"}]}`,
		} {
			fmt.Fprintf(w, "data: %s%s\n\n", head, c)
		}
		if b.StreamOptions.IncludeUsage {
			fmt.Fprintf(w, "data: %s[],\"usage\":{\"prompt_tokens\":5,\"completion_tokens\":2,\"total_tokens\":7}}\n\n", head)
		}
		_, _ = io.WriteString(w, "data: [DONE]\n\n")
	}))
	t.Cleanup(srv.Close)
	return srv, &calls
}

// streamedText is the content of a streamed chat answer, and whether it
// carried usage and ended with [DONE].
func streamedText(t *testing.T, sse string) (text string, usage, done bool) {
	t.Helper()
	for _, line := range strings.Split(sse, "\n") {
		data, ok := strings.CutPrefix(line, "data: ")
		if !ok {
			continue
		}
		if data == "[DONE]" {
			done = true
			continue
		}
		var c struct {
			Object  string `json:"object"`
			Choices []struct {
				Delta struct {
					Content string `json:"content"`
				} `json:"delta"`
			} `json:"choices"`
			Usage *struct {
				CompletionTokens int `json:"completion_tokens"`
			} `json:"usage"`
		}
		if err := json.Unmarshal([]byte(data), &c); err != nil || c.Object != "chat.completion.chunk" {
			t.Fatalf("not a chunk: %q", data)
		}
		for _, ch := range c.Choices {
			text += ch.Delta.Content
		}
		if c.Usage != nil && c.Usage.CompletionTokens == 2 {
			usage = true
		}
	}
	return text, usage, done
}

func TestAStreamedAnswerIsReusedForClientsThatStreamAndThatDoNot(t *testing.T) {
	srv, calls := greetingEngine(t)
	rt, secret := cachingRouter(t, srv.URL, false)
	auth := map[string]string{"Authorization": "Bearer " + secret}
	streamed := func(includeUsage bool) map[string]any {
		b := ask(0, "greet me")
		b["stream"] = true
		if includeUsage {
			b["stream_options"] = map[string]any{"include_usage": true}
		}
		return b
	}

	first := post(t, rt.Handler(), streamed(false), auth)
	if first.Header().Get(cacheHeader) != "miss" {
		t.Fatalf("%q %s", first.Header().Get(cacheHeader), first.Body.String())
	}
	if text, _, _ := streamedText(t, first.Body.String()); text != "Hello world" {
		t.Fatalf("streamed %q", text)
	}

	// Whole, to a client that does not stream.
	whole := post(t, rt.Handler(), ask(0, "greet me"), auth)
	var c struct {
		Object  string `json:"object"`
		Choices []struct {
			Message struct {
				Role, Content string
			} `json:"message"`
			FinishReason string `json:"finish_reason"`
		} `json:"choices"`
		Usage struct {
			PromptTokens, CompletionTokens int
		} `json:"usage"`
	}
	if err := json.Unmarshal(whole.Body.Bytes(), &c); err != nil || whole.Header().Get(cacheHeader) != "hit" {
		t.Fatalf("%q %s %v", whole.Header().Get(cacheHeader), whole.Body.String(), err)
	}
	if c.Object != "chat.completion" || len(c.Choices) != 1 || c.Choices[0].Message.Content != "Hello world" ||
		c.Choices[0].Message.Role != "assistant" || c.Choices[0].FinishReason != "stop" {
		t.Errorf("put back together as %s", whole.Body.String())
	}

	// Streamed again, with the usage only to a client that asks for it.
	for _, withUsage := range []bool{false, true} {
		rec := post(t, rt.Handler(), streamed(withUsage), auth)
		text, usage, done := streamedText(t, rec.Body.String())
		if rec.Header().Get(cacheHeader) != "hit" || text != "Hello world" || !done || usage != withUsage ||
			!strings.HasPrefix(rec.Header().Get("Content-Type"), "text/event-stream") {
			t.Errorf("streamed from the cache (usage %v): %q %q usage=%v done=%v\n%s", withUsage,
				rec.Header().Get(cacheHeader), text, usage, done, rec.Body.String())
		}
	}
	if calls.Load() != 1 {
		t.Errorf("the engine was asked %d times", calls.Load())
	}
}

func TestManyClientsReadingAndFillingTheCacheAtOnce(t *testing.T) {
	e := newFakeEngine(t, "qwen2.5-14b-instruct")
	rt, secret := cachingRouter(t, e.srv.URL, false)
	h := rt.Handler()
	auth := map[string]string{"Authorization": "Bearer " + secret}
	var wg sync.WaitGroup
	for i := range 40 {
		wg.Go(func() {
			rec := post(t, h, ask(0, fmt.Sprintf("question %d", i%5)), auth)
			if rec.Code != 200 {
				t.Errorf("%d %s", rec.Code, rec.Body.String())
			}
		})
	}
	wg.Wait()
	for i := range 5 {
		if rec := post(t, h, ask(0, fmt.Sprintf("question %d", i)), auth); rec.Header().Get(cacheHeader) != "hit" {
			t.Errorf("question %d: %q", i, rec.Header().Get(cacheHeader))
		}
	}
}

func TestTheAnswerCacheDropsWhatWasUsedLongestAgoPastItsBudget(t *testing.T) {
	var held int64
	c := newAnswerCache(1000, 400, func(b int64, _ int) { held = b })
	now := time.Now()
	put := func(key string) bool {
		return c.put(&cachedAnswer{key: key, project: "p", body: bytes.Repeat([]byte("x"), 100),
			at: now, expires: now.Add(time.Hour)}, now)
	}
	for _, k := range []string{"a", "b", "c"} {
		put(k)
	}
	c.get("a", time.Hour, now) // a is now the most recently used
	put("d")
	put("e")
	if c.get("b", time.Hour, now) != nil {
		t.Error("b, used longest ago, was kept past the budget")
	}
	if c.get("a", time.Hour, now) == nil {
		t.Error("a, used recently, was dropped")
	}
	if held > 1000 {
		t.Errorf("holding %d bytes of 1000", held)
	}
	if c.put(&cachedAnswer{key: "big", body: bytes.Repeat([]byte("x"), 500)}, now) {
		t.Error("an answer larger than an entry may be was kept")
	}
	// Past its TTL, or past a TTL shortened since, it is gone.
	if c.get("a", time.Hour, now.Add(2*time.Hour)) != nil {
		t.Error("an expired answer was returned")
	}
	if c.get("e", time.Minute, now.Add(2*time.Minute)) != nil {
		t.Error("an answer older than the project's TTL now was returned")
	}
}

func TestAnAnswerIsLookedUpByWhatChangesIt(t *testing.T) {
	key := func(project, model, body string) string {
		k, ok := answerKey(project, "/v1/chat/completions", model, []byte(body))
		if !ok {
			t.Fatalf("no key for %s", body)
		}
		return k
	}
	base := key("p", "m", `{"model":"m","messages":[{"role":"user","content":"hi"}],"temperature":0}`)
	same := []string{
		`{"messages":[{"content":"hi","role":"user"}],"temperature":0.0,"model":"alias"}`,
		`{"model":"m","messages":[{"role":"user","content":"hi"}],"temperature":0,"stream":true,"stream_options":{"include_usage":true}}`,
		`{"model":"m","messages":[{"role":"user","content":"hi"}],"temperature":0e0,"user":"u","metadata":{"a":"b"}}`,
	}
	for _, b := range same {
		if key("p", "m", b) != base {
			t.Errorf("%s was looked up apart", b)
		}
	}
	different := []string{
		`{"model":"m","messages":[{"role":"user","content":"hi!"}],"temperature":0}`,
		`{"model":"m","messages":[{"role":"user","content":"hi"}],"temperature":0,"max_tokens":5}`,
		`{"model":"m","messages":[{"role":"user","content":"hi"}],"temperature":0,"seed":9007199254740993}`,
	}
	for _, b := range different {
		if key("p", "m", b) == base {
			t.Errorf("%s was looked up as the same", b)
		}
	}
	if key("p", "m", `{"seed":9007199254740993}`) == key("p", "m", `{"seed":9007199254740992}`) {
		t.Error("two large seeds were taken for one")
	}
	if key("other", "m", `{"model":"m","messages":[{"role":"user","content":"hi"}],"temperature":0}`) == base {
		t.Error("another project's request was looked up as the same")
	}
	if key("p", "other", `{"model":"m","messages":[{"role":"user","content":"hi"}],"temperature":0}`) == base {
		t.Error("another model's request was looked up as the same")
	}
	if _, ok := answerKey("p", "/v1/chat/completions", "m", []byte(`not json`)); ok {
		t.Error("a body that is not JSON was given a key")
	}
}

func TestAStreamedAnswerIsPutBackTogetherAndTakenApartAgain(t *testing.T) {
	sse := strings.Join([]string{
		`data: {"id":"c","object":"chat.completion.chunk","created":7,"model":"m","choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}]}`,
		`data: {"id":"c","object":"chat.completion.chunk","created":7,"model":"m","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"get_weather","arguments":""}}]},"finish_reason":null}]}`,
		`data: {"id":"c","object":"chat.completion.chunk","created":7,"model":"m","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\"city\":"}}]},"finish_reason":null}]}`,
		`data: {"id":"c","object":"chat.completion.chunk","created":7,"model":"m","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\"Paris\"}"}}]},"finish_reason":null}]}`,
		`data: {"id":"c","object":"chat.completion.chunk","created":7,"model":"m","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}`,
		`data: [DONE]`,
	}, "\n\n") + "\n\n"
	body, ok := assembleStream("/v1/chat/completions", []byte(sse), &proxy.Usage{PromptTokens: 3, CompletionTokens: 4})
	if !ok {
		t.Fatal("not put back together")
	}
	var c struct {
		Choices []struct {
			Message struct {
				Content   *string    `json:"content"`
				ToolCalls []toolCall `json:"tool_calls"`
			} `json:"message"`
			FinishReason string `json:"finish_reason"`
		} `json:"choices"`
		Usage struct {
			TotalTokens int `json:"total_tokens"`
		} `json:"usage"`
	}
	if err := json.Unmarshal(body, &c); err != nil {
		t.Fatal(err)
	}
	if len(c.Choices) != 1 || c.Choices[0].Message.Content != nil || len(c.Choices[0].Message.ToolCalls) != 1 ||
		c.Choices[0].Message.ToolCalls[0].Function.Arguments != `{"city":"Paris"}` ||
		c.Choices[0].Message.ToolCalls[0].ID != "call_1" || c.Choices[0].FinishReason != "tool_calls" ||
		c.Usage.TotalTokens != 7 {
		t.Errorf("put back together as %s", body)
	}

	// Taken apart, and put back together again, it is the same answer.
	events, err := streamFromAnswer("/v1/chat/completions", body, false)
	if err != nil {
		t.Fatal(err)
	}
	again, ok := assembleStream("/v1/chat/completions",
		[]byte("data: "+strings.Join(append(events, "[DONE]"), "\n\ndata: ")+"\n\n"), &proxy.Usage{PromptTokens: 3, CompletionTokens: 4})
	if !ok || !bytes.Equal(again, body) {
		t.Errorf("round trip:\n%s\n%s", body, again)
	}

	// What it does not understand, and a stream cut short, are not kept.
	for _, bad := range []string{
		strings.Replace(sse, `"delta":{}`, `"delta":{"audio":{"id":"a"}}`, 1),
		strings.Replace(sse, `"finish_reason":null}]}`, `"finish_reason":null,"logprobs":{"content":[]}}]}`, 1),
		strings.Replace(sse, "data: [DONE]", "", 1),
	} {
		if _, ok := assembleStream("/v1/chat/completions", []byte(bad), nil); ok {
			t.Errorf("kept: %s", bad)
		}
	}

	// A completion, too.
	text := `data: {"id":"t","object":"text_completion","created":1,"model":"m","choices":[{"index":0,"text":"Hel","finish_reason":null}]}` + "\n\n" +
		`data: {"id":"t","object":"text_completion","created":1,"model":"m","choices":[{"index":0,"text":"lo","finish_reason":"length"}]}` + "\n\n" +
		"data: [DONE]\n\n"
	tb, ok := assembleStream("/v1/completions", []byte(text), nil)
	if !ok || !strings.Contains(string(tb), `"text":"Hello"`) || !strings.Contains(string(tb), `"object":"text_completion"`) {
		t.Errorf("a completion: %s", tb)
	}
}

func TestAnAnswerFromTheCacheReachesAnthropicClientsInTheirShape(t *testing.T) {
	srv, calls := greetingEngine(t)
	rt, secret := cachingRouter(t, srv.URL, false)
	send := func(stream bool) *httptest.ResponseRecorder {
		b, _ := json.Marshal(map[string]any{"model": "q14", "max_tokens": 50, "temperature": 0, "stream": stream,
			"messages": []any{map[string]any{"role": "user", "content": "greet me"}}})
		req := httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(b))
		req.Header.Set("Content-Type", "application/json")
		req.Header.Set("x-api-key", secret)
		rec := httptest.NewRecorder()
		rt.Handler().ServeHTTP(rec, req)
		return rec
	}
	first := send(true)
	if first.Code != 200 || !strings.Contains(first.Body.String(), "Hello") {
		t.Fatalf("%d %s", first.Code, first.Body.String())
	}
	whole := send(false)
	if whole.Code != 200 || !strings.Contains(whole.Body.String(), `"text":"Hello world"`) {
		t.Errorf("whole: %d %s", whole.Code, whole.Body.String())
	}
	streamed := send(true)
	if !strings.Contains(streamed.Body.String(), "content_block_delta") ||
		!strings.Contains(streamed.Body.String(), "Hello world") ||
		!strings.Contains(streamed.Body.String(), "message_stop") {
		t.Errorf("streamed: %s", streamed.Body.String())
	}
	if calls.Load() != 1 {
		t.Errorf("the engine was asked %d times", calls.Load())
	}
}