Agent-System-Backend/internal/server/sse_retry_test.go

package server

import (
	"bufio"
	"context"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"agent-system-backend/internal/bus"
	core "agent-system-core"
)

func TestSSERetryAndLiveDelivery(t *testing.T) {
	srv, _ := newTestServer(t)
	h := srv.Handler()
	_, convID := setupConversation(t, h)
	b := bus.New()
	srv.Subscribe = func(cid string) (<-chan core.StreamEvent, func()) { return b.Subscribe(cid) }

	ts := httptest.NewServer(h)
	t.Cleanup(ts.Close)
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	t.Cleanup(cancel)
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, ts.URL+"/api/events?conversationId="+convID, nil)
	if err != nil {
		t.Fatal(err)
	}
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { resp.Body.Close() })
	reader := bufio.NewReader(resp.Body)
	head := readUntil(t, reader, "retry: 1000")
	if !strings.Contains(head, "retry: 1000") {
		t.Fatalf("missing retry: %s", head)
	}
	b.Publish(core.StreamEvent{
		Seq: 7, ConversationID: core.ID(convID), TS: time.Now().UTC().Format(time.RFC3339), V: 1,
		Type: core.EvMessageDelta, Payload: []byte(`{"index":0,"text":"hi"}`),
	})
	body := readUntil(t, reader, "event: message.delta")
	if !strings.Contains(body, "id: 7") {
		t.Fatalf("live event: %s", body)
	}
}

func TestSSEReconnectFillsOverflowGap(t *testing.T) {
	srv, st := newTestServer(t)
	h := srv.Handler()
	_, convID := setupConversation(t, h)
	evs := []core.StreamEvent{
		{ConversationID: core.ID(convID), TS: time.Now().UTC().Format(time.RFC3339), V: 1, Type: core.EvMessageDelta, Payload: []byte(`{"index":0,"text":"a"}`)},
		{ConversationID: core.ID(convID), TS: time.Now().UTC().Format(time.RFC3339), V: 1, Type: core.EvMessageDelta, Payload: []byte(`{"index":0,"text":"b"}`)},
		{ConversationID: core.ID(convID), TS: time.Now().UTC().Format(time.RFC3339), V: 1, Type: core.EvMessageDelta, Payload: []byte(`{"index":0,"text":"c"}`)},
	}
	if _, err := st.AppendEvents(convID, evs); err != nil {
		t.Fatal(err)
	}
	b := bus.NewWithBuffer(1)
	ch, stop := b.Subscribe(convID)
	b.Publish(evs[1])
	b.Publish(evs[2])
	closed := false
	for i := 0; i < 3; i++ {
		_, ok := <-ch
		if !ok {
			closed = true
			break
		}
	}
	if !closed {
		t.Fatal("overflow must close the subscription")
	}
	stop()

	ts := httptest.NewServer(h)
	t.Cleanup(ts.Close)
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	t.Cleanup(cancel)
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, ts.URL+"/api/events?conversationId="+convID, nil)
	if err != nil {
		t.Fatal(err)
	}
	req.Header.Set("Last-Event-ID", "1")
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { resp.Body.Close() })
	reader := bufio.NewReader(resp.Body)
	body := readUntil(t, reader, "id: 3")
	if !strings.Contains(body, "id: 2") || !strings.Contains(body, "id: 3") {
		t.Fatalf("gap not filled: %s", body)
	}
	if strings.Contains(body, "id: 1\n") {
		t.Fatalf("replayed the event the client already had: %s", body)
	}
}

func readUntil(t *testing.T, r *bufio.Reader, needle string) string {
	t.Helper()
	var b strings.Builder
	for i := 0; i < 40; i++ {
		line, err := r.ReadString('\n')
		b.WriteString(line)
		if strings.Contains(b.String(), needle) {
			return b.String()
		}
		if err != nil {
			t.Fatalf("read until %q: %v\n%s", needle, err, b.String())
		}
	}
	t.Fatalf("gave up waiting for %q\n%s", needle, b.String())
	return ""
}