localchat/internal/web/server_test.go

package web

import (
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"net/http/cookiejar"
	"net/http/httptest"
	"net/url"
	"regexp"
	"strings"
	"testing"
	"time"

	"git.b0b.be/bdeb/localchat/internal/config"
)

// fakeLLM streams back "echo: <last user message>" and records what it got.
func fakeLLM(t *testing.T, seen *[][]map[string]string) *httptest.Server {
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/v1/models" {
			fmt.Fprint(w, `{"data":[{"id":"fake"}]}`)
			return
		}
		var req struct {
			Messages []map[string]string `json:"messages"`
		}
		json.NewDecoder(r.Body).Decode(&req)
		*seen = append(*seen, req.Messages)
		last := req.Messages[len(req.Messages)-1]["content"]
		w.Header().Set("Content-Type", "text/event-stream")
		for _, part := range []string{"echo: ", last, "\n\n**bold** <script>x</script>"} {
			b, _ := json.Marshal(part)
			fmt.Fprintf(w, "data: {\"choices\":[{\"delta\":{\"content\":%s}}]}\n\n", b)
		}
		fmt.Fprint(w, "data: [DONE]\n\n")
	}))
}

func newTestApp(t *testing.T, llmURL string) (*httptest.Server, *http.Client) {
	t.Helper()
	cfg := config.Config{
		BaseURL:      llmURL + "/v1",
		SystemPrompt: "be nice",
		Title:        "test",
		MaxHistory:   20,
		Timeout:      10 * time.Second,
	}
	app := httptest.NewServer(New(cfg, slog.New(slog.NewTextHandler(io.Discard, nil))).Handler())
	t.Cleanup(app.Close)
	jar, _ := cookiejar.New(nil)
	return app, &http.Client{Jar: jar}
}

var streamURL = regexp.MustCompile(`sse-connect="(/chat/stream/[0-9a-f]+)"`)

func get(t *testing.T, c *http.Client, u string) string {
	t.Helper()
	resp, err := c.Get(u)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	b, _ := io.ReadAll(resp.Body)
	return string(b)
}

// ask posts a message, follows the SSE stream and returns the raw events.
func ask(t *testing.T, app *httptest.Server, c *http.Client, msg string) string {
	t.Helper()
	resp, err := c.PostForm(app.URL+"/chat", url.Values{"message": {msg}})
	if err != nil {
		t.Fatal(err)
	}
	body, _ := io.ReadAll(resp.Body)
	resp.Body.Close()
	m := streamURL.FindStringSubmatch(string(body))
	if m == nil {
		t.Fatalf("no stream url in %s", body)
	}
	return get(t, c, app.URL+m[1])
}

func TestChatFlow(t *testing.T) {
	var seen [][]map[string]string
	llm := fakeLLM(t, &seen)
	defer llm.Close()
	app, c := newTestApp(t, llm.URL)

	if page := get(t, c, app.URL+"/"); !strings.Contains(page, `id="composer"`) {
		t.Fatal("index page missing composer")
	}
	if h := get(t, c, app.URL+"/health"); !strings.Contains(h, "dot ok") || !strings.Contains(h, "fake") {
		t.Fatalf("health: %s", h)
	}

	events := ask(t, app, c, "hi <b>there</b>")
	for _, want := range []string{
		"event: token\ndata: <span>echo: </span>",
		"hi &lt;b&gt;there&lt;/b&gt;", // user text is escaped in tokens
		"event: done",
		"<strong>bold</strong>",     // final reply is rendered Markdown
		"<!-- raw HTML omitted -->", // ... with raw HTML from the model dropped
	} {
		if !strings.Contains(events, want) {
			t.Errorf("stream missing %q:\n%s", want, events)
		}
	}
	if strings.Contains(events, "<script>") {
		t.Error("raw <script> leaked into the page")
	}

	// second turn sends system prompt + full history
	ask(t, app, c, "again")
	last := seen[len(seen)-1]
	roles := []string{}
	for _, m := range last {
		roles = append(roles, m["role"])
	}
	if got := strings.Join(roles, ","); got != "system,user,assistant,user" {
		t.Fatalf("history roles = %s", got)
	}

	// page reload shows the history, reset clears it
	if page := get(t, c, app.URL+"/"); strings.Count(page, `class="msg `) != 4 {
		t.Fatal("expected 4 messages after reload")
	}
	resp, _ := c.Post(app.URL+"/chat/reset", "", nil)
	resp.Body.Close()
	if page := get(t, c, app.URL+"/"); strings.Contains(page, `class="msg `) {
		t.Fatal("reset did not clear history")
	}
}

func TestStreamReconnectDoesNotRegenerate(t *testing.T) {
	var seen [][]map[string]string
	llm := fakeLLM(t, &seen)
	defer llm.Close()
	app, c := newTestApp(t, llm.URL)

	resp, _ := c.PostForm(app.URL+"/chat", url.Values{"message": {"once"}})
	body, _ := io.ReadAll(resp.Body)
	resp.Body.Close()
	path := streamURL.FindStringSubmatch(string(body))[1]

	get(t, c, app.URL+path)
	again := get(t, c, app.URL+path)
	if len(seen) != 1 {
		t.Fatalf("model called %d times, want 1", len(seen))
	}
	if strings.Contains(again, "event: token") || !strings.Contains(again, "echo: once") {
		t.Fatalf("reconnect should replay the final message only:\n%s", again)
	}
}

func TestSessionsAreIsolated(t *testing.T) {
	var seen [][]map[string]string
	llm := fakeLLM(t, &seen)
	defer llm.Close()
	app, alice := newTestApp(t, llm.URL)
	jar, _ := cookiejar.New(nil)
	bob := &http.Client{Jar: jar}

	ask(t, app, alice, "alice secret")
	if strings.Contains(get(t, bob, app.URL+"/"), "alice secret") {
		t.Fatal("bob can see alice's chat")
	}
	resp, _ := bob.PostForm(app.URL+"/chat", url.Values{"message": {"x"}})
	body, _ := io.ReadAll(resp.Body)
	resp.Body.Close()
	path := streamURL.FindStringSubmatch(string(body))[1]
	if out := get(t, alice, app.URL+path); !strings.Contains(out, "lost") {
		t.Fatalf("alice should not be able to stream bob's reply:\n%s", out)
	}
}

func TestEmptyMessageRejected(t *testing.T) {
	var seen [][]map[string]string
	llm := fakeLLM(t, &seen)
	defer llm.Close()
	app, c := newTestApp(t, llm.URL)

	resp, _ := c.PostForm(app.URL+"/chat", url.Values{"message": {"   "}})
	resp.Body.Close()
	if resp.StatusCode != http.StatusBadRequest {
		t.Fatalf("status %d", resp.StatusCode)
	}
}

func TestModelOffline(t *testing.T) {
	app, c := newTestApp(t, "http://127.0.0.1:1") // nothing listens here
	if h := get(t, c, app.URL+"/health"); !strings.Contains(h, "model offline") {
		t.Fatalf("health: %s", h)
	}
	events := ask(t, app, c, "hello?")
	if !strings.Contains(events, "event: done") || !strings.Contains(events, `class="error"`) {
		t.Fatalf("want an error bubble:\n%s", events)
	}
}