From 5f015a5c75fe4b49ae943a54ae17003b1d7260be Mon Sep 17 00:00:00 2001 From: Jerry Yan <792602257@qq.com> Date: Wed, 10 Jun 2026 14:48:28 +0800 Subject: [PATCH] =?UTF-8?q?SSE=20=E8=B0=83=E6=95=B4?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- go.mod | 2 +- internal/chat/sse_test.go | 49 ++++++++++++++++++-- internal/transport/http/middleware/logger.go | 25 ++-------- 3 files changed, 49 insertions(+), 27 deletions(-) diff --git a/go.mod b/go.mod index ff8a148..5e397d1 100644 --- a/go.mod +++ b/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 diff --git a/internal/chat/sse_test.go b/internal/chat/sse_test.go index 20315c7..a684756 100644 --- a/internal/chat/sse_test.go +++ b/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. diff --git a/internal/transport/http/middleware/logger.go b/internal/transport/http/middleware/logger.go index dac3362..efbfb8f 100644 --- a/internal/transport/http/middleware/logger.go +++ b/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)), )