llmgw/internal/router/threads_test.go

package router

import (
	"encoding/json"
	"strings"
	"testing"

	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// splicedRoles reads back the roles and contents of a spliced body, which is
// what the engine will actually see.
func splicedRoles(t *testing.T, body []byte) []string {
	t.Helper()
	var b struct {
		Messages []struct {
			Role    string `json:"role"`
			Content string `json:"content"`
		} `json:"messages"`
	}
	if err := json.Unmarshal(body, &b); err != nil {
		t.Fatalf("the spliced body is not valid JSON: %v\n%s", err, body)
	}
	out := make([]string, 0, len(b.Messages))
	for _, m := range b.Messages {
		out = append(out, m.Role+":"+m.Content)
	}
	return out
}

func TestHistoryGoesBeforeWhatTheCallerJustSaid(t *testing.T) {
	// The whole point of a thread: the client sends the new turn, not the
	// conversation again. Order is what makes the reply make sense.
	body := []byte(`{"model":"m","messages":[{"role":"user","content":"third"}]}`)
	history := []store.Message{
		{Role: "user", Content: "first"},
		{Role: "assistant", Content: "second"},
	}

	out, dropped, err := spliceThread(body, history, 0)
	if err != nil {
		t.Fatal(err)
	}
	if dropped != 0 {
		t.Errorf("nothing should have been dropped, %d was", dropped)
	}
	got := splicedRoles(t, out)
	want := []string{"user:first", "assistant:second", "user:third"}
	if strings.Join(got, "|") != strings.Join(want, "|") {
		t.Errorf("messages came out as %v, want %v", got, want)
	}
}

func TestEverythingElseInTheBodySurvives(t *testing.T) {
	// The router forwards fields it does not understand — that is how engine
	// specific options reach an engine at all. Splicing a thread must not be
	// the thing that starts dropping them.
	body := []byte(`{"model":"m","messages":[{"role":"user","content":"hi"}],` +
		`"temperature":0.2,"stream":true,"some_engine_option":{"nested":[1,2]}}`)

	out, _, err := spliceThread(body, []store.Message{{Role: "user", Content: "old"}}, 0)
	if err != nil {
		t.Fatal(err)
	}
	var got map[string]json.RawMessage
	if err := json.Unmarshal(out, &got); err != nil {
		t.Fatal(err)
	}
	for _, k := range []string{"model", "temperature", "stream", "some_engine_option"} {
		if _, ok := got[k]; !ok {
			t.Errorf("%q was dropped by splicing", k)
		}
	}
	if string(got["some_engine_option"]) != `{"nested":[1,2]}` {
		t.Errorf("a field llmgw does not understand was rewritten: %s", got["some_engine_option"])
	}
}

func TestTheOldestTurnsAreDroppedFirst(t *testing.T) {
	// A conversation that outgrew the model it is now being served by — after
	// a budget fallback, that is not a choice the caller made, so the thread
	// has to stay usable.
	history := []store.Message{
		{Role: "user", Content: strings.Repeat("a", 100)},
		{Role: "assistant", Content: strings.Repeat("b", 100)},
		{Role: "user", Content: strings.Repeat("c", 100)},
		{Role: "assistant", Content: "keep me"},
	}
	body := []byte(`{"model":"m","messages":[{"role":"user","content":"now"}]}`)

	out, dropped, err := spliceThread(body, history, 200)
	if err != nil {
		t.Fatal(err)
	}
	if dropped == 0 {
		t.Fatal("nothing was dropped although the history was over budget")
	}
	got := strings.Join(splicedRoles(t, out), "|")
	if !strings.Contains(got, "keep me") {
		t.Errorf("the most recent stored turn was dropped: %s", got)
	}
	if !strings.Contains(got, "user:now") {
		t.Errorf("the caller's own turn was dropped: %s", got)
	}
	if strings.Contains(got, strings.Repeat("a", 100)) {
		t.Errorf("the oldest turn survived while newer ones were dropped: %s", got)
	}
}

func TestASystemMessageIsNeverDropped(t *testing.T) {
	// It is usually the instruction that shapes every answer. Losing it
	// silently changes what the model is for, which is a stranger failure than
	// a forgotten turn — and much harder to spot, because the model keeps
	// answering.
	history := []store.Message{
		{Role: "system", Content: "Answer only in the user's own language."},
		{Role: "user", Content: strings.Repeat("x", 500)},
		{Role: "assistant", Content: strings.Repeat("y", 500)},
	}
	body := []byte(`{"model":"m","messages":[{"role":"user","content":"now"}]}`)

	out, dropped, err := spliceThread(body, history, 100)
	if err != nil {
		t.Fatal(err)
	}
	if dropped != 2 {
		t.Errorf("dropped %d turns, expected both non-system ones", dropped)
	}
	if got := strings.Join(splicedRoles(t, out), "|"); !strings.Contains(got, "user's own language") {
		t.Errorf("the system message was dropped: %s", got)
	}
}

func TestAModelWithNoDeclaredContextTrimsNothing(t *testing.T) {
	// Guessing a limit for a model that never stated one would drop turns for
	// no reason, and the caller would never learn why the model forgot.
	history := make([]store.Message, 50)
	for i := range history {
		history[i] = store.Message{Role: "user", Content: strings.Repeat("z", 1000)}
	}
	body := []byte(`{"model":"m","messages":[{"role":"user","content":"now"}]}`)

	_, dropped, err := spliceThread(body, history, 0)
	if err != nil {
		t.Fatal(err)
	}
	if dropped != 0 {
		t.Errorf("%d turns were dropped for a model that declares no context", dropped)
	}
}

func TestNoHistoryLeavesTheBodyByteForByte(t *testing.T) {
	// A request without a thread, and the first request of a new one, must
	// take exactly the path they took before this existed.
	body := []byte(`{"model":"m","messages":[{"role":"user","content":"hi"}]}`)
	out, dropped, err := spliceThread(body, nil, 4096)
	if err != nil {
		t.Fatal(err)
	}
	if dropped != 0 || string(out) != string(body) {
		t.Errorf("an empty history rewrote the body:\n%s", out)
	}
}

func TestStructuredContentIsStoredAsItArrived(t *testing.T) {
	// An image or an array of parts. Flattening it to a string would replay
	// something the client never sent.
	body := []byte(`{"model":"m","messages":[{"role":"user","content":[{"type":"text","text":"hi"}]}]}`)
	got := ownMessages(body)
	if len(got) != 1 {
		t.Fatalf("read %d messages", len(got))
	}
	if !strings.HasPrefix(got[0].Content, `[{"type":"text"`) {
		t.Errorf("structured content was flattened: %q", got[0].Content)
	}
}