SSE 调整

internal/chat/sse_test.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"io"
+	"log/slog"
 	"net/http"
 	"net/http/httptest"
 	"strings"
@@ -24,15 +26,23 @@ import (
 // flushRecorder is an httptest.ResponseRecorder that also implements http.Flusher.
 type flushRecorder struct {
 	*httptest.ResponseRecorder
-	flushed int
+	flushed   int
+	flushedCh chan struct{}
 }
 
 func newFlushRecorder() *flushRecorder {
-	return &flushRecorder{ResponseRecorder: httptest.NewRecorder()}
+	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 =====
@@ -149,6 +159,33 @@ func TestStreamMessages_ContentTypeHeader(t *testing.T) {
 	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.
@@ -187,9 +224,11 @@ func TestStreamMessages_StopsWhenClientClosesContext(t *testing.T) {
 	// 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"}
-	require.Eventually(t, func() bool {
-		return strings.Contains(w.Body.String(), "text_delta")
-	}, time.Second, 10*time.Millisecond, "event was not flushed before cancel")
+	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.