From d3ef9c941b0ae21f8a7451576bae89e5d703f3b0 Mon Sep 17 00:00:00 2001 From: Jerry Yan <792602257@qq.com> Date: Mon, 4 May 2026 19:28:46 +0800 Subject: [PATCH] feat(chat): SSE writer + multipart upload endpoint with mime whitelist --- internal/chat/sse.go | 83 +++++++++++++ internal/chat/sse_test.go | 206 +++++++++++++++++++++++++++++++ internal/chat/upload.go | 125 +++++++++++++++++++ internal/chat/upload_test.go | 231 +++++++++++++++++++++++++++++++++++ 4 files changed, 645 insertions(+) create mode 100644 internal/chat/sse.go create mode 100644 internal/chat/sse_test.go create mode 100644 internal/chat/upload.go create mode 100644 internal/chat/upload_test.go diff --git a/internal/chat/sse.go b/internal/chat/sse.go new file mode 100644 index 0000000..3ba5ba8 --- /dev/null +++ b/internal/chat/sse.go @@ -0,0 +1,83 @@ +package chat + +import ( + "encoding/json" + "fmt" + "net/http" + "strconv" + + "github.com/go-chi/chi/v5" + + "github.com/yan1h/agent-coding-workflow/internal/infra/errs" +) + +// streamMessages implements GET /api/v1/conversations/{id}/stream. +// It upgrades the connection to a Server-Sent Events stream and forwards +// SSEEvents produced by MessageService.Stream until the channel closes or +// the client disconnects. +func (h *Handler) streamMessages(w http.ResponseWriter, r *http.Request) { + uid, err := h.userID(r) + if err != nil { + writeErr(w, r, err) + return + } + convID, err := parseUUID(chi.URLParam(r, "id")) + if err != nil { + writeErr(w, r, err) + return + } + var since int64 + if s := r.URL.Query().Get("since"); s != "" { + since, _ = strconv.ParseInt(s, 10, 64) + } + + flusher, ok := w.(http.Flusher) + if !ok { + http.Error(w, "streaming unsupported", http.StatusInternalServerError) + return + } + + w.Header().Set("Content-Type", "text/event-stream") + w.Header().Set("Cache-Control", "no-cache") + w.Header().Set("Connection", "keep-alive") + w.Header().Set("X-Accel-Buffering", "no") + w.WriteHeader(http.StatusOK) + + ch, err := h.msgSvc.Stream(r.Context(), uid, convID, since) + if err != nil { + writeSSEError(w, flusher, err) + return + } + for { + select { + case <-r.Context().Done(): + return + case ev, ok := <-ch: + if !ok { + return + } + writeSSE(w, flusher, ev) + } + } +} + +// writeSSE writes a single SSE frame to w and flushes. +func writeSSE(w http.ResponseWriter, f http.Flusher, ev SSEEvent) { + data, _ := json.Marshal(ev.Data) + fmt.Fprintf(w, "id: %d\nevent: %s\ndata: %s\n\n", ev.ID, ev.Event, data) + f.Flush() +} + +// writeSSEError writes an error event to the SSE stream and flushes. +func writeSSEError(w http.ResponseWriter, f http.Flusher, err error) { + code := errs.CodeInternal + msg := err.Error() + if ae, ok := errs.As(err); ok { + code = ae.Code + msg = ae.Message + } + payload := map[string]string{"code": string(code), "message": msg} + data, _ := json.Marshal(payload) + fmt.Fprintf(w, "event: error\ndata: %s\n\n", data) + f.Flush() +} diff --git a/internal/chat/sse_test.go b/internal/chat/sse_test.go new file mode 100644 index 0000000..928d29e --- /dev/null +++ b/internal/chat/sse_test.go @@ -0,0 +1,206 @@ +package chat + +import ( + "context" + "encoding/json" + "fmt" + "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 +} + +func newFlushRecorder() *flushRecorder { + return &flushRecorder{ResponseRecorder: httptest.NewRecorder()} +} + +func (f *flushRecorder) Flush() { + f.flushed++ +} + +// ===== 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, error) { + return 0, 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) +} + +// 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, then cancel. + ch <- SSEEvent{ID: 1, Event: "text_delta", Data: "hi"} + 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") + } + + // The one event before cancel should have been written. + 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) +} diff --git a/internal/chat/upload.go b/internal/chat/upload.go new file mode 100644 index 0000000..6f96569 --- /dev/null +++ b/internal/chat/upload.go @@ -0,0 +1,125 @@ +package chat + +import ( + "context" + "fmt" + "net/http" + "path/filepath" + "strings" + "time" + + "github.com/google/uuid" + + "github.com/yan1h/agent-coding-workflow/internal/infra/errs" + "github.com/yan1h/agent-coding-workflow/internal/infra/storage" + "github.com/yan1h/agent-coding-workflow/internal/transport/http/middleware" +) + +// Uploader handles multipart file uploads for chat attachments. +// It validates the file size and MIME type, stores the file via Storage, +// and records the attachment in the repository. If storage succeeds but DB +// insert fails, the stored blob is deleted to prevent orphaned objects. +type Uploader struct { + repo Repository + storage storage.Storage + mimeWhite map[string]bool + maxFile int64 +} + +// NewUploader constructs an Uploader. +// mimeList is a whitelist of allowed MIME types; "text/*" is a wildcard prefix. +// maxFile is the maximum allowed file size in bytes. +func NewUploader(repo Repository, st storage.Storage, mimeList []string, maxFile int64) *Uploader { + mw := make(map[string]bool, len(mimeList)) + for _, m := range mimeList { + mw[m] = true + } + return &Uploader{repo: repo, storage: st, mimeWhite: mw, maxFile: maxFile} +} + +// Handle implements POST /api/v1/uploads. +func (u *Uploader) Handle(w http.ResponseWriter, r *http.Request) { + uid, ok := middleware.UserIDFromContext(r.Context()) + if !ok { + writeErr(w, r, errs.New(errs.CodeUnauthorized, "unauthorized")) + return + } + + // Parse the multipart form; allow a 1 MiB overhead beyond maxFile for + // form boundary / metadata. + if err := r.ParseMultipartForm(u.maxFile + 1<<20); err != nil { + writeErr(w, r, errs.Wrap(err, errs.CodeInvalidInput, "parse multipart")) + return + } + + file, header, err := r.FormFile("file") + if err != nil { + writeErr(w, r, errs.Wrap(err, errs.CodeInvalidInput, "no file field")) + return + } + defer file.Close() + + if header.Size > u.maxFile { + writeErr(w, r, errs.New(errs.CodeChatAttachmentTooLarge, + fmt.Sprintf("size %d exceeds limit %d", header.Size, u.maxFile))) + return + } + + mt := header.Header.Get("Content-Type") + if !u.mimeAllowed(mt) { + writeErr(w, r, errs.New(errs.CodeChatAttachmentUnsupportedMime, mt)) + return + } + + now := time.Now().UTC() + id, err := uuid.NewV7() + if err != nil { + writeErr(w, r, errs.Wrap(err, errs.CodeInternal, "generate uuid")) + return + } + ext := filepath.Ext(header.Filename) + key := fmt.Sprintf("%04d/%02d/%s%s", now.Year(), int(now.Month()), id.String(), ext) + + res, err := u.storage.Put(r.Context(), key, file, mt) + if err != nil { + writeErr(w, r, errs.Wrap(err, errs.CodeStoragePutFailed, "storage put")) + return + } + + att, err := u.repo.InsertAttachment(r.Context(), InsertAttachmentParams{ + ID: id, + UserID: uid, + Filename: header.Filename, + MimeType: mt, + SizeBytes: res.Size, + SHA256: res.SHA256, + StoragePath: res.StoragePath, + }) + if err != nil { + // Best-effort blob cleanup to prevent orphaned storage objects. + _ = u.storage.Delete(context.Background(), res.StoragePath) + writeErr(w, r, err) + return + } + + writeJSON(w, http.StatusOK, AttachmentDTO{ + ID: att.ID, + Filename: att.Filename, + MimeType: att.MimeType, + SizeBytes: att.SizeBytes, + SHA256: att.SHA256, + CreatedAt: att.CreatedAt, + }) +} + +// mimeAllowed reports whether the given MIME type is in the whitelist. +// "text/*" in the whitelist acts as a wildcard for any text/ subtype. +func (u *Uploader) mimeAllowed(mt string) bool { + if u.mimeWhite[mt] { + return true + } + if u.mimeWhite["text/*"] && strings.HasPrefix(mt, "text/") { + return true + } + return false +} diff --git a/internal/chat/upload_test.go b/internal/chat/upload_test.go new file mode 100644 index 0000000..b8716ea --- /dev/null +++ b/internal/chat/upload_test.go @@ -0,0 +1,231 @@ +package chat + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "io" + "mime/multipart" + "net/http" + "net/http/httptest" + "net/textproto" + "testing" + + "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/infra/storage" + "github.com/yan1h/agent-coding-workflow/internal/transport/http/middleware" +) + +// ===== upload-specific fakes ===== + +// uploadRepo wraps fakeRepo and overrides InsertAttachment to capture or fail. +type uploadRepo struct { + *fakeRepo + inserted []InsertAttachmentParams + insertErr error + storedAtt *Attachment +} + +func (r *uploadRepo) InsertAttachment(_ context.Context, p InsertAttachmentParams) (*Attachment, error) { + r.inserted = append(r.inserted, p) + if r.insertErr != nil { + return nil, r.insertErr + } + att := &Attachment{ + ID: p.ID, UserID: p.UserID, Filename: p.Filename, MimeType: p.MimeType, + SizeBytes: p.SizeBytes, SHA256: p.SHA256, StoragePath: p.StoragePath, + } + r.storedAtt = att + return att, nil +} + +// uploadStorage records put/delete and supports a forced putErr. +type uploadStorage struct { + blobs map[string][]byte + putErr error + deleteCalls []string +} + +func newUploadStorage() *uploadStorage { + return &uploadStorage{blobs: make(map[string][]byte)} +} + +func (s *uploadStorage) Put(_ context.Context, key string, r io.Reader, _ string) (storage.PutResult, error) { + if s.putErr != nil { + return storage.PutResult{}, s.putErr + } + data, err := io.ReadAll(r) + if err != nil { + return storage.PutResult{}, err + } + s.blobs[key] = data + h := sha256.Sum256(data) + return storage.PutResult{StoragePath: key, SHA256: hex.EncodeToString(h[:]), Size: int64(len(data))}, nil +} + +func (s *uploadStorage) Get(_ context.Context, key string) (io.ReadCloser, error) { + data, ok := s.blobs[key] + if !ok { + return nil, errors.New("not found: " + key) + } + return io.NopCloser(bytes.NewReader(data)), nil +} + +func (s *uploadStorage) Delete(_ context.Context, key string) error { + s.deleteCalls = append(s.deleteCalls, key) + delete(s.blobs, key) + return nil +} + +// ===== test helpers ===== + +// uploadHarness builds a chi router with Auth middleware that resolves any +// non-empty bearer token to uid, then mounts the Uploader under POST /uploads. +type uploadHarness struct { + uid uuid.UUID + repo *uploadRepo + stor *uploadStorage + up *Uploader + r chi.Router +} + +func newUploadHarness(t *testing.T, mimes []string, maxFile int64) *uploadHarness { + t.Helper() + uid := uuid.Must(uuid.NewV7()) + repo := &uploadRepo{fakeRepo: newFakeRepo()} + stor := newUploadStorage() + up := NewUploader(repo, stor, mimes, maxFile) + resolver := &fakeSessionResolver{uid: uid} + r := chi.NewRouter() + r.Use(middleware.RequestID) + r.Use(middleware.Auth(resolver, middleware.AuthOptions{})) + r.Post("/uploads", up.Handle) + return &uploadHarness{uid: uid, repo: repo, stor: stor, up: up, r: r} +} + +// buildMultipart constructs a multipart/form-data body with one "file" part. +// declaredMime is what the part's Content-Type header advertises. +func buildMultipart(t *testing.T, filename, declaredMime string, data []byte) (body *bytes.Buffer, contentType string) { + t.Helper() + body = &bytes.Buffer{} + mw := multipart.NewWriter(body) + hdr := textproto.MIMEHeader{} + hdr.Set("Content-Disposition", `form-data; name="file"; filename="`+filename+`"`) + hdr.Set("Content-Type", declaredMime) + part, err := mw.CreatePart(hdr) + require.NoError(t, err) + _, err = part.Write(data) + require.NoError(t, err) + require.NoError(t, mw.Close()) + return body, mw.FormDataContentType() +} + +func newUploadRequest(t *testing.T, body *bytes.Buffer, contentType string, withAuth bool) *http.Request { + t.Helper() + req := httptest.NewRequest(http.MethodPost, "/uploads", body) + req.Header.Set("Content-Type", contentType) + if withAuth { + req.Header.Set("Authorization", "Bearer test-token") + } + return req +} + +// ===== tests ===== + +func TestUploader_HappyPath_ReturnsAttachmentMetadata(t *testing.T) { + h := newUploadHarness(t, []string{"text/plain"}, 1<<20) + data := []byte("hello world") + body, ct := buildMultipart(t, "greeting.txt", "text/plain", data) + w := httptest.NewRecorder() + h.r.ServeHTTP(w, newUploadRequest(t, body, ct, true)) + + require.Equal(t, http.StatusOK, w.Code, "body: %s", w.Body.String()) + var resp AttachmentDTO + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + assert.Equal(t, "greeting.txt", resp.Filename) + assert.Equal(t, "text/plain", resp.MimeType) + assert.Equal(t, int64(len(data)), resp.SizeBytes) + + expectedSHA := sha256.Sum256(data) + assert.Equal(t, hex.EncodeToString(expectedSHA[:]), resp.SHA256) + + require.Len(t, h.repo.inserted, 1) + got := h.repo.inserted[0] + assert.Equal(t, h.uid, got.UserID) + assert.Equal(t, "greeting.txt", got.Filename) +} + +func TestUploader_FileTooLarge_413(t *testing.T) { + h := newUploadHarness(t, []string{"text/plain"}, 5) + body, ct := buildMultipart(t, "big.txt", "text/plain", []byte("more than five bytes")) + w := httptest.NewRecorder() + h.r.ServeHTTP(w, newUploadRequest(t, body, ct, true)) + + require.Equal(t, http.StatusRequestEntityTooLarge, w.Code, "body: %s", w.Body.String()) + assert.Empty(t, h.repo.inserted, "no DB row should be inserted") +} + +func TestUploader_UnsupportedMime_415(t *testing.T) { + h := newUploadHarness(t, []string{"text/plain"}, 1<<20) + body, ct := buildMultipart(t, "evil.exe", "application/octet-stream", []byte("MZ...")) + w := httptest.NewRecorder() + h.r.ServeHTTP(w, newUploadRequest(t, body, ct, true)) + + require.Equal(t, http.StatusUnsupportedMediaType, w.Code) + assert.Empty(t, h.repo.inserted) + assert.Empty(t, h.stor.blobs, "blob must not be stored when mime is rejected") +} + +func TestUploader_StoragePutFails_NoOrphanRow(t *testing.T) { + h := newUploadHarness(t, []string{"text/plain"}, 1<<20) + h.stor.putErr = errors.New("disk full") + body, ct := buildMultipart(t, "x.txt", "text/plain", []byte("hi")) + w := httptest.NewRecorder() + h.r.ServeHTTP(w, newUploadRequest(t, body, ct, true)) + + require.Equal(t, http.StatusBadGateway, w.Code) + assert.Empty(t, h.repo.inserted, "no DB insert when storage fails") + assert.Empty(t, h.stor.deleteCalls, "no delete when nothing was stored") +} + +func TestUploader_DBInsertFails_DeletesBlob(t *testing.T) { + h := newUploadHarness(t, []string{"text/plain"}, 1<<20) + h.repo.insertErr = errs.New(errs.CodeInternal, "db down") + body, ct := buildMultipart(t, "x.txt", "text/plain", []byte("hi")) + w := httptest.NewRecorder() + h.r.ServeHTTP(w, newUploadRequest(t, body, ct, true)) + + require.Equal(t, http.StatusInternalServerError, w.Code) + require.Len(t, h.repo.inserted, 1) + require.Len(t, h.stor.deleteCalls, 1, "blob must be cleaned up after DB failure") + assert.Equal(t, h.repo.inserted[0].StoragePath, h.stor.deleteCalls[0]) +} + +func TestUploader_NoAuth_401(t *testing.T) { + h := newUploadHarness(t, []string{"text/plain"}, 1<<20) + body, ct := buildMultipart(t, "x.txt", "text/plain", []byte("hi")) + w := httptest.NewRecorder() + h.r.ServeHTTP(w, newUploadRequest(t, body, ct, false)) + + require.Equal(t, http.StatusUnauthorized, w.Code) + assert.Empty(t, h.repo.inserted) +} + +func TestUploader_TextWildcard_Allows(t *testing.T) { + h := newUploadHarness(t, []string{"text/*"}, 1<<20) + body, ct := buildMultipart(t, "page.html", "text/html", []byte("

ok

")) + w := httptest.NewRecorder() + h.r.ServeHTTP(w, newUploadRequest(t, body, ct, true)) + + require.Equal(t, http.StatusOK, w.Code, "body: %s", w.Body.String()) + require.Len(t, h.repo.inserted, 1) + assert.Equal(t, "text/html", h.repo.inserted[0].MimeType) +}