atlas-iac/services/hermes/router/session_snapshot_test.go

package main

import (
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"strings"
	"testing"
	"testing/iotest"
)

// webuiSession mirrors the fields of the tenant WebUI `GET /api/session`
// payload that the continuity fallback reads.
type webuiSession struct {
	SessionID    string            `json:"session_id"`
	Messages     []json.RawMessage `json:"messages"`
	MessageCount int               `json:"message_count"`
	Title        string            `json:"title,omitempty"`
	ReadOnly     bool              `json:"read_only,omitempty"`
}

type webuiSessionEnvelope struct {
	Session webuiSession `json:"session"`
}

func snapshotResponse(target, body string) *http.Response {
	request, _ := http.NewRequest(http.MethodGet, target, nil)
	return &http.Response{
		StatusCode: http.StatusOK,
		Header:     make(http.Header),
		Body:       io.NopCloser(strings.NewReader(body)),
		Request:    request,
	}
}

func decodeSnapshot(t *testing.T, response *http.Response) webuiSession {
	t.Helper()
	var envelope webuiSessionEnvelope
	if err := json.NewDecoder(response.Body).Decode(&envelope); err != nil {
		t.Fatal(err)
	}
	return envelope.Session
}

func TestBoundSessionSnapshotKeepsOnlyRecentMessages(t *testing.T) {
	messages := make([]map[string]int, 30)
	for index := range messages {
		messages[index] = map[string]int{"index": index}
	}
	body, _ := json.Marshal(map[string]any{"session": map[string]any{
		"session_id": "resolved", "messages": messages, "message_count": 30,
	}})
	response := snapshotResponse(
		"http://tenant/api/session?session_id=root&messages=1&msg_limit=24&hermes_fallback=1",
		string(body),
	)
	response.Header.Set("ETag", "stale")
	if err := boundSessionSnapshot(response); err != nil {
		t.Fatal(err)
	}
	bounded := decodeSnapshot(t, response)
	if len(bounded.Messages) != sessionSnapshotItems || bounded.MessageCount != 30 {
		t.Fatalf("snapshot was not bounded: %#v", bounded)
	}
	var first map[string]int
	if err := json.Unmarshal(bounded.Messages[0], &first); err != nil || first["index"] != 6 {
		t.Fatalf("snapshot did not retain the recent tail: %#v, %v", first, err)
	}
	if response.Header.Get("Cache-Control") != "no-store" || response.Header.Get("ETag") != "" {
		t.Fatal("bounded snapshot retained cache metadata")
	}
}

func TestBoundSessionSnapshotPreservesLargerReportedTotal(t *testing.T) {
	response := snapshotResponse(
		"http://tenant/api/session?session_id=leaf&hermes_fallback=1",
		`{"session":{"session_id":"leaf","messages":[],"message_count":100}}`,
	)
	if err := boundSessionSnapshot(response); err != nil {
		t.Fatal(err)
	}
	bounded := decodeSnapshot(t, response)
	if bounded.SessionID != "leaf" || bounded.MessageCount != 100 {
		t.Fatalf("reported total was lost: %#v", bounded)
	}
}

// The poller and the WebUI both evolve; bounding the message tail must never
// strip the surrounding session metadata that decides what the banner says.
func TestBoundSessionSnapshotRelaysUnknownSessionMetadata(t *testing.T) {
	response := snapshotResponse(
		"http://tenant/api/session?session_id=root&hermes_fallback=1",
		`{"session":{"session_id":"root","messages":[{"role":"user"}],`+
			`"is_streaming":true,"active_stream_id":"stream-1","read_only":false,`+
			`"future_field":{"kept":true}},"other_envelope_key":7}`,
	)
	if err := boundSessionSnapshot(response); err != nil {
		t.Fatal(err)
	}
	body, err := io.ReadAll(response.Body)
	if err != nil {
		t.Fatal(err)
	}
	for _, expected := range []string{
		`"is_streaming":true`, `"active_stream_id":"stream-1"`,
		`"future_field":{"kept":true}`, `"other_envelope_key":7`,
	} {
		if !strings.Contains(string(body), expected) {
			t.Fatalf("bounded snapshot dropped %s: %s", expected, body)
		}
	}
	if length := response.Header.Get("Content-Length"); length != "" &&
		length != strings.TrimSpace(length) {
		t.Fatal("content length was not rewritten")
	}
}

func TestBoundSessionSnapshotFailsClosedOnMalformedOrOversizedBody(t *testing.T) {
	for name, body := range map[string]string{
		"malformed":      `{`,
		"missing":        `{"error":"Session not found"}`,
		"nonObject":      `{"session":42}`,
		"badMessageList": `{"session":{"messages":"all of them"}}`,
		"oversized":      strings.Repeat("x", sessionSnapshotBytes+1),
	} {
		t.Run(name, func(t *testing.T) {
			response := snapshotResponse(
				"http://tenant/api/session?session_id=root&hermes_fallback=1",
				body,
			)
			if err := boundSessionSnapshot(response); err == nil {
				t.Fatal("unsafe snapshot was accepted")
			}
		})
	}
	response := snapshotResponse(
		"http://tenant/api/session?session_id=root&hermes_fallback=1", `{}`,
	)
	response.Body = io.NopCloser(iotest.ErrReader(errors.New("read failed")))
	if err := boundSessionSnapshot(response); err == nil {
		t.Fatal("snapshot body read failure was ignored")
	}
}

func TestBoundSessionSnapshotLeavesNativeAndErrorResponsesUntouched(t *testing.T) {
	for _, response := range []*http.Response{
		{StatusCode: http.StatusOK},
		// The WebUI's own session reads carry no fallback marker.
		snapshotResponse("http://tenant/api/session?session_id=root", `{}`),
		snapshotResponse("http://tenant/api/sessions", `{}`),
		// The dashboard-only route is not this backend's contract.
		snapshotResponse("http://tenant/api/sessions/root/messages?hermes_fallback=1", `{}`),
	} {
		if err := boundSessionSnapshot(response); err != nil {
			t.Fatal(err)
		}
	}
	errorResponse := snapshotResponse(
		"http://tenant/api/session?session_id=root&hermes_fallback=1", `{}`,
	)
	errorResponse.StatusCode = http.StatusNotFound
	if err := boundSessionSnapshot(errorResponse); err != nil {
		t.Fatal(err)
	}
}