package chat import ( "context" "encoding/json" "fmt" "io" "log/slog" "net/http" "net/http/httptest" "strings" "testing" "time" "github.com/go-chi/chi/v5" "github.com/google/uuid" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/yan1h/agent-coding-workflow/internal/infra/errs" "github.com/yan1h/agent-coding-workflow/internal/transport/http/middleware" ) // ===== SSE-capable response writer ===== // flushRecorder is an httptest.ResponseRecorder that also implements http.Flusher. type flushRecorder struct { *httptest.ResponseRecorder flushed int flushedCh chan struct{} } func newFlushRecorder() *flushRecorder { return &flushRecorder{ ResponseRecorder: httptest.NewRecorder(), flushedCh: make(chan struct{}, 1), } } func (f *flushRecorder) Flush() { f.flushed++ select { case f.flushedCh <- struct{}{}: default: } } // ===== Fake MessageService for SSE tests ===== // sseMessageService allows controlling the Stream channel. type sseMessageService struct { streamCh chan SSEEvent streamErr error } func (s *sseMessageService) Send(_ context.Context, _, _ uuid.UUID, _ string, _ []uuid.UUID) (int64, int64, error) { return 0, 0, nil } func (s *sseMessageService) Stream(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ int64) (<-chan SSEEvent, error) { if s.streamErr != nil { return nil, s.streamErr } return s.streamCh, nil } func (s *sseMessageService) Cancel(_ context.Context, _ uuid.UUID, _ int64) error { return nil } func (s *sseMessageService) Retry(_ context.Context, _ uuid.UUID, _ int64) (int64, uuid.UUID, error) { return 0, uuid.Nil, nil } func (s *sseMessageService) ListMessages(_ context.Context, _, _ uuid.UUID, _ *int64, _ int) ([]Message, error) { return nil, nil } // ===== SSE helper unit tests ===== // TestSSE_WritesEventInProperFormat verifies that writeSSE produces the // canonical id/event/data/\n\n format. func TestSSE_WritesEventInProperFormat(t *testing.T) { w := newFlushRecorder() ev := SSEEvent{ ID: 42, Event: "text_delta", Data: map[string]string{"delta": "hello"}, } writeSSE(w, w, ev) body := w.Body.String() assert.True(t, strings.HasPrefix(body, "id: 42\n"), "must start with id line") assert.Contains(t, body, "event: text_delta\n") assert.Contains(t, body, "data: ") assert.True(t, strings.HasSuffix(body, "\n\n"), "must end with double newline") assert.Equal(t, 1, w.flushed) // Verify the data line is valid JSON. lines := strings.Split(strings.TrimRight(body, "\n"), "\n") var dataLine string for _, l := range lines { if strings.HasPrefix(l, "data: ") { dataLine = strings.TrimPrefix(l, "data: ") } } require.NotEmpty(t, dataLine) var payload map[string]string require.NoError(t, json.Unmarshal([]byte(dataLine), &payload)) assert.Equal(t, "hello", payload["delta"]) } // TestSSE_WriteSSEError verifies that writeSSEError produces an error event // with the expected code/message fields from an AppError. func TestSSE_WriteSSEError(t *testing.T) { w := newFlushRecorder() appErr := errs.New(errs.CodeChatConversationNotFound, "conversation gone") writeSSEError(w, w, appErr) body := w.Body.String() assert.Contains(t, body, "event: error\n") assert.True(t, strings.HasSuffix(body, "\n\n")) assert.Equal(t, 1, w.flushed) lines := strings.Split(body, "\n") var dataLine string for _, l := range lines { if strings.HasPrefix(l, "data: ") { dataLine = strings.TrimPrefix(l, "data: ") } } require.NotEmpty(t, dataLine) var payload map[string]string require.NoError(t, json.Unmarshal([]byte(dataLine), &payload)) assert.Equal(t, string(errs.CodeChatConversationNotFound), payload["code"]) assert.Equal(t, "conversation gone", payload["message"]) } // TestStreamMessages_ContentTypeHeader verifies that the SSE handler sets // the required Content-Type and X-Accel-Buffering headers. func TestStreamMessages_ContentTypeHeader(t *testing.T) { uid := uuid.New() convID := uuid.New() ch := make(chan SSEEvent) close(ch) // immediately closed — no events, handler returns quickly. msgSvc := &sseMessageService{streamCh: ch} resolver := &fakeSessionResolver{uid: uid} upload := NewUploader(newFakeRepo(), nil, nil, 1<<20) h := NewHandler(&fakeConvSvc{}, msgSvc, &fakeTplSvc{}, &fakeEpSvc{}, &fakeUsageSvc{}, upload, resolver, newFakeAdminLookup()) r := chi.NewRouter() r.Use(middleware.RequestID) h.Mount(r) req := httptest.NewRequest(http.MethodGet, fmt.Sprintf("/api/v1/conversations/%s/stream", convID), nil) req.Header.Set("Authorization", "Bearer test-token") w := newFlushRecorder() r.ServeHTTP(w, req) assert.Equal(t, "text/event-stream", w.Header().Get("Content-Type")) assert.Equal(t, "no-cache", w.Header().Get("Cache-Control")) assert.Equal(t, "no", w.Header().Get("X-Accel-Buffering")) assert.Equal(t, http.StatusOK, w.Code) } func TestStreamMessages_WithLoggerMiddlewareKeepsStreaming(t *testing.T) { uid := uuid.New() convID := uuid.New() ch := make(chan SSEEvent) close(ch) msgSvc := &sseMessageService{streamCh: ch} resolver := &fakeSessionResolver{uid: uid} upload := NewUploader(newFakeRepo(), nil, nil, 1<<20) h := NewHandler(&fakeConvSvc{}, msgSvc, &fakeTplSvc{}, &fakeEpSvc{}, &fakeUsageSvc{}, upload, resolver, newFakeAdminLookup()) r := chi.NewRouter() r.Use(middleware.RequestID) r.Use(middleware.Logger(slog.New(slog.NewTextHandler(io.Discard, nil)))) h.Mount(r) req := httptest.NewRequest(http.MethodGet, fmt.Sprintf("/api/v1/conversations/%s/stream?since=3", convID), nil) req.Header.Set("Authorization", "Bearer test-token") w := newFlushRecorder() r.ServeHTTP(w, req) assert.Equal(t, http.StatusOK, w.Code) assert.Equal(t, "text/event-stream", w.Header().Get("Content-Type")) assert.NotContains(t, w.Body.String(), "streaming unsupported") } // TestStreamMessages_StopsWhenClientClosesContext verifies that when the // request context is cancelled mid-stream, the handler stops consuming events // and the channel is not blocked. func TestStreamMessages_StopsWhenClientClosesContext(t *testing.T) { uid := uuid.New() convID := uuid.New() // Channel that never closes on its own — cancelling the ctx should stop the handler. ch := make(chan SSEEvent, 1) msgSvc := &sseMessageService{streamCh: ch} resolver := &fakeSessionResolver{uid: uid} upload := NewUploader(newFakeRepo(), nil, nil, 1<<20) h := NewHandler(&fakeConvSvc{}, msgSvc, &fakeTplSvc{}, &fakeEpSvc{}, &fakeUsageSvc{}, upload, resolver, newFakeAdminLookup()) // Build a router identical to newTestHandler but without chi's context copying issues. r := chi.NewRouter() r.Use(middleware.RequestID) h.Mount(r) // Create a cancellable request. ctx, cancel := context.WithCancel(context.Background()) req := httptest.NewRequest(http.MethodGet, fmt.Sprintf("/api/v1/conversations/%s/stream", convID), nil) req.Header.Set("Authorization", "Bearer test-token") req = req.WithContext(ctx) w := newFlushRecorder() done := make(chan struct{}) go func() { defer close(done) r.ServeHTTP(w, req) }() // Send one event and give the handler a moment to drain it before cancelling, // so the body assertion is deterministic. The cancel itself is what we're // verifying — that the handler exits even though the channel never closes. ch <- SSEEvent{ID: 1, Event: "text_delta", Data: "hi"} select { case <-w.flushedCh: case <-waitTimeout(t, time.Second): t.Fatal("event was not flushed before cancel") } cancel() // Handler must return promptly after context cancellation. select { case <-done: // success case <-waitTimeout(t, 2*time.Second): t.Fatal("handler did not stop after context cancellation") } assert.Contains(t, w.Body.String(), "text_delta") } // waitTimeout returns a channel that closes after the given duration. func waitTimeout(t *testing.T, d time.Duration) <-chan time.Time { t.Helper() return time.After(d) }