llmgw/test/e2e/threads_test.go

package e2e

import (
	"bytes"
	"context"
	"encoding/json"
	"io"
	"net/http"
	"strings"
	"testing"
	"time"

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

// say sends one turn, optionally as part of a thread, and returns what came
// back along with the headers that explain what llmgw did.
func (h *harness) say(t *testing.T, key, model, thread, text string) (int, http.Header, string) {
	t.Helper()
	body, _ := json.Marshal(map[string]any{
		"model":      model,
		"messages":   []map[string]string{{"role": "user", "content": text}},
		"max_tokens": 16,
	})
	req, err := http.NewRequestWithContext(context.Background(), http.MethodPost,
		h.router.URL+"/v1/chat/completions", bytes.NewReader(body))
	if err != nil {
		t.Fatal(err)
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+key)
	if thread != "" {
		req.Header.Set("x-llmgw-thread", thread)
	}
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	msg, _ := io.ReadAll(resp.Body)
	return resp.StatusCode, resp.Header, string(msg)
}

// authConfig turns on authentication against a provisioned store. The e2e
// tests set this inline elsewhere; threads need it in five places.
func authConfig(path string) func(*config.Server) {
	return func(c *config.Server) {
		c.Auth = config.Auth{
			Enabled: true, Store: path,
			CacheTTL:      50 * time.Millisecond,
			FlushInterval: 20 * time.Millisecond,
			BatchSize:     8,
		}
	}
}

func enableThreads(t *testing.T, path string, hours int) {
	t.Helper()
	withStore(t, path, func(s *store.SQLite) {
		if err := s.SetThreadRetention(context.Background(), "proj", hours); err != nil {
			t.Fatal(err)
		}
	})
}

func storedTurns(t *testing.T, path, thread string) []store.Message {
	t.Helper()
	var out []store.Message
	withStore(t, path, func(s *store.SQLite) {
		msgs, err := s.Thread(context.Background(), "proj", thread)
		if err != nil {
			t.Fatalf("reading thread %q: %v", thread, err)
		}
		out = msgs
	})
	return out
}

func TestAProjectThatNeverAskedStoresNothingAndSaysSo(t *testing.T) {
	// The property being protected: llmgw keeps no conversations unless a
	// project turns it on. Refused rather than ignored, because ignoring the
	// header would send the model a conversation with no history and it would
	// answer confidently about something it was never told.
	path, mint := provision(t)
	key := mint(t, "user")
	h := start(t, 2, authConfig(path))

	code, _, msg := h.say(t, key, testModel, "th1", "hello")
	if code != http.StatusBadRequest {
		t.Fatalf("a thread on a project that keeps none returned %d: %s", code, msg)
	}
	if !strings.Contains(msg, "does not keep conversations") {
		t.Errorf("the refusal does not say why: %s", msg)
	}
}

func TestAThreadRemembersTheEarlierTurns(t *testing.T) {
	path, mint := provision(t)
	key := mint(t, "user")
	enableThreads(t, path, 24)
	h := start(t, 2, authConfig(path))

	if code, _, msg := h.say(t, key, testModel, "th1", "first question"); code != http.StatusOK {
		t.Fatalf("%d %s", code, msg)
	}
	if code, _, msg := h.say(t, key, testModel, "th1", "second question"); code != http.StatusOK {
		t.Fatalf("%d %s", code, msg)
	}

	turns := storedTurns(t, path, "th1")
	if len(turns) != 4 {
		t.Fatalf("stored %d turns, expected two questions and two answers: %+v", len(turns), turns)
	}
	if turns[0].Content != "first question" || turns[2].Content != "second question" {
		t.Errorf("the questions were not stored in order: %+v", turns)
	}
	if turns[1].Role != "assistant" || turns[1].Content == "" {
		t.Errorf("the first answer was not stored: %+v", turns[1])
	}
}

func TestAThreadContinuesOnADifferentModel(t *testing.T) {
	// The whole point. A conversation started on one model is carried to
	// another, and the stored turns say which model produced each answer —
	// without that, a thread that moved reads as one voice, and when the
	// answers disagree that is the first thing anyone asks.
	path, mint := provision(t)
	key := mint(t, "user")
	enableThreads(t, path, 24)
	h := start(t, 2, func(c *config.Server) {
		authConfig(path)(c)
		c.Models = budgetFallbackModels()
		for i := range c.Endpoints {
			c.Endpoints[i].Models = []string{testModel, cheapModel}
		}
	})

	if code, _, msg := h.say(t, key, testModel, "th1", "started here"); code != http.StatusOK {
		t.Fatalf("%d %s", code, msg)
	}
	if code, _, msg := h.say(t, key, cheapModel, "th1", "continued there"); code != http.StatusOK {
		t.Fatalf("%d %s", code, msg)
	}

	turns := storedTurns(t, path, "th1")
	models := map[string]bool{}
	for _, m := range turns {
		if m.Role == "assistant" {
			models[m.Model] = true
		}
	}
	if !models[testModel] || !models[cheapModel] {
		t.Errorf("the thread does not record both models that answered: %+v", turns)
	}
}

func TestOneProjectCannotReadAnothersThread(t *testing.T) {
	// Thread ids are opaque, but an unguessable id is a secret and secrets end
	// up in logs and bug reports. The project scope is what makes a leaked id
	// useless outside it.
	path, mint := provision(t)
	mine := mint(t, "mine")
	enableThreads(t, path, 24)

	// A key in a second project.
	withStore(t, path, func(s *store.SQLite) {
		if err := s.PutProject(context.Background(),
			store.Project{ID: "other", TeamID: "team", Name: "Other"}); err != nil {
			t.Fatal(err)
		}
		if err := s.SetThreadRetention(context.Background(), "other", 24); err != nil {
			t.Fatal(err)
		}
	})
	theirs := mint(t, "theirs", func(k *store.Key) { k.ProjectID = "other" })

	h := start(t, 2, authConfig(path))
	if code, _, msg := h.say(t, mine, testModel, "shared-id", "a private thing"); code != http.StatusOK {
		t.Fatalf("%d %s", code, msg)
	}

	// The other project using the same id must not see it, and must not be
	// able to append to it either.
	code, _, msg := h.say(t, theirs, testModel, "shared-id", "what was said?")
	if code == http.StatusOK {
		if strings.Contains(msg, "a private thing") {
			t.Fatal("one project read another's conversation")
		}
		// Served, but from an empty history of its own — check the store
		// agrees that nothing crossed over.
		if turns := storedTurns(t, path, "shared-id"); len(turns) > 2 {
			for _, m := range turns {
				if m.Content == "a private thing" {
					t.Fatal("a second project appended into another's thread")
				}
			}
		}
	}
}

func TestWithoutAThreadNothingIsStoredAtAll(t *testing.T) {
	// Every request that does not opt in must take exactly the path it took
	// before threads existed.
	path, mint := provision(t)
	key := mint(t, "user")
	enableThreads(t, path, 24)
	h := start(t, 2, authConfig(path))

	if code, _, msg := h.say(t, key, testModel, "", "no thread here"); code != http.StatusOK {
		t.Fatalf("%d %s", code, msg)
	}

	withStore(t, path, func(s *store.SQLite) {
		// Pruning with a cutoff far past every retention window deletes every
		// thread there is, so a count of zero means none were created.
		n, err := s.PruneThreads(context.Background(), time.Now().Add(1000*time.Hour))
		if err != nil {
			t.Fatal(err)
		}
		if n != 0 {
			t.Errorf("%d threads were created by a request that used none", n)
		}
	})
}