SSE 调整

go.mod

@@ -6,6 +6,7 @@ require (
 	github.com/alexedwards/argon2id v1.0.0
 	github.com/anthropics/anthropic-sdk-go v1.38.0
 	github.com/coder/websocket v1.8.14
+	github.com/felixge/httpsnoop v1.0.4
 	github.com/go-chi/chi/v5 v5.2.5
 	github.com/golang-migrate/migrate/v4 v4.19.1
 	github.com/google/uuid v1.6.0
@@ -43,7 +44,6 @@ require (
 	github.com/docker/go-connections v0.6.0 // indirect
 	github.com/docker/go-units v0.5.0 // indirect
 	github.com/ebitengine/purego v0.10.0 // indirect
-	github.com/felixge/httpsnoop v1.0.4 // indirect
 	github.com/fsnotify/fsnotify v1.9.0 // indirect
 	github.com/go-logr/logr v1.4.3 // indirect
 	github.com/go-logr/stdr v1.2.2 // indirect

internal/chat/sse_test.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"io"
+	"log/slog"
 	"net/http"
 	"net/http/httptest"
 	"strings"
@@ -25,14 +27,22 @@ import (
 type flushRecorder struct {
 	*httptest.ResponseRecorder
 	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.

internal/transport/http/middleware/logger.go

@@ -4,26 +4,10 @@ import (
 	"log/slog"
 	"net/http"
 	"time"
+
+	"github.com/felixge/httpsnoop"
 )
 
-// statusRecorder wraps http.ResponseWriter to capture the response status
-// code so the Logger middleware can include it in the structured log line.
-//
-// The default status is initialised to http.StatusOK because net/http
-// implicitly emits a 200 when a handler writes a body without calling
-// WriteHeader explicitly; mirroring that behaviour keeps logs accurate.
-type statusRecorder struct {
-	http.ResponseWriter
-	status int
-}
-
-// WriteHeader records the status code before delegating to the wrapped
-// ResponseWriter so the original response semantics are preserved.
-func (s *statusRecorder) WriteHeader(code int) {
-	s.status = code
-	s.ResponseWriter.WriteHeader(code)
-}
-
 // Logger returns HTTP middleware that emits a single structured log line
 // per request once the downstream handler has finished. The line includes
 // method, path, response status, elapsed duration in milliseconds, and the
@@ -36,13 +20,12 @@ func Logger(log *slog.Logger) func(http.Handler) http.Handler {
 	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			start := time.Now()
-			rec := &statusRecorder{ResponseWriter: w, status: http.StatusOK}
-			next.ServeHTTP(rec, r)
+			metrics := httpsnoop.CaptureMetrics(next, w, r)
 			ctx := r.Context()
 			log.LogAttrs(ctx, slog.LevelInfo, "http",
 				slog.String("method", r.Method),
 				slog.String("path", r.URL.Path),
-				slog.Int("status", rec.status),
+				slog.Int("status", metrics.Code),
 				slog.Int64("duration_ms", time.Since(start).Milliseconds()),
 				slog.String("request_id", RequestIDFromContext(ctx)),
 			)