Milestone 5/4: Add real-time email updates via WebSockets

ROADMAP.md

@@ -500,58 +500,60 @@ This is the most complex but most rewarding "quality of life" feature.
 
 #### **Backend**
 
-* [ ] **Add WebSocket library:**
+* [x] **Add WebSocket library:**
     * `go get github.com/gorilla/websocket`
-* [ ] **Create WebSocket Hub:**
-    * Create `/backend/internal/websocket/hub.go`.
-    * The `Hub` struct will manage active connections: `clients map[string]*websocket.Conn` (mapping `userID` to their connection).
-    * It needs methods: `Register(userID, conn)`, `Unregister(userID)`, and `Send(userID, message []byte)`.
-* [ ] **Create <code>GET /api/v1/ws</code> endpoint:**
-    * Add this route in `routes.go`.
-    * The handler (`wsHandler`) upgrades the HTTP connection to a WebSocket.
-    * It gets the `userID` from the auth context.
-    * It calls `hub.Register(userID, conn)`.
-    * It must handle disconnects by calling `hub.Unregister(userID)`.
-* [ ] **Create IMAP IDLE listener:**
-    * In `/backend/internal/imap/idle.go`, create `func (s *Service) StartIdleListener(ctx context.Context, userID string, hub *websocket.Hub)`.
-    * **Launch it:** When a user successfully connects to the WebSocket (`hub.Register`), launch this function in a **new goroutine** for that `userID`.
+* [x] **Create WebSocket Hub:**
+    * Created `/backend/internal/websocket/hub.go`.
+    * The `Hub` struct manages active connections: `userID -> set of *websocket.Conn` (via a `Client` wrapper).
+    * It exposes methods: `Register(userID, conn)`, `Unregister(userID, client)`, and `Send(userID, message []byte)`.
+    * It supports multiple connections per user with a per-user limit (currently 10).
+* [x] **Create <code>GET /api/v1/ws</code> endpoint:**
+    * Added the route in `cmd/server/main.go`.
+    * The handler (`WebSocketHandler`) upgrades the HTTP connection to a WebSocket.
+    * It gets the `userID` from the auth context using `GetUserIDFromContext`.
+    * It calls `hub.Register(userID, conn)` and starts a read loop to detect disconnects.
+    * On disconnect, it calls `hub.Unregister(userID, client)` and stops the IMAP IDLE listener when there are no more active connections for that user.
+* [x] **Create IMAP IDLE listener:**
+    * Implemented in `/backend/internal/imap/idle.go` as `func (s *Service) StartIdleListener(ctx context.Context, userID string, hub *websocket.Hub)`.
+    * **Launch:** When a user successfully connects to the WebSocket (`hub.Register`), `WebSocketHandler` starts this function in a **new goroutine** for that `userID` (if not already running).
     * **Logic:**
-        1. Get a *dedicated* IMAP connection (do not use the pool).
+        1. Get a dedicated IMAP listener connection from the pool.
         2. Run `SELECT INBOX`.
-        3. Start a `for` loop (to handle disconnects).
-        4. Inside the loop, run `client.Idle()`.
-        5. Listen for updates. When an update arrives (e.g., `* 1 EXISTS`), call `hub.Send(userID, []byte('{"type": "new_email", "folder": "INBOX"}'))`.
-        6. If `client.Idle()` returns an error (e.g., timeout), `log.Println` and `time.Sleep(10 * time.Second)` before the loop retries.
+        3. Start an IDLE loop (with fallback) using `go-imap-idle`.
+        4. Listen for updates via the client's `Updates` channel. When an update indicates new messages in `INBOX`, call `SyncThreadsForFolder` for `INBOX` immediately.
+        5. After syncing, call `hub.Send(userID, []byte('{"type":"new_email","folder":"INBOX"}'))`.
+        6. On errors (e.g., timeout), log, remove the listener connection, and retry after a short sleep.
 
 #### **Frontend**
 
-* [ ] **Create <code>useWebSocket</code> hook:**
-    * Create `hooks/useWebSocket.ts`.
-    * It should be called *once* from your main `Layout.tsx`.
+* [x] **Create <code>useWebSocket</code> hook:**
+    * Created `hooks/useWebSocket.ts`.
+    * It is called *once* from the main `Layout.tsx`.
     * `useEffect` on mount:
-        1. `const socket = new WebSocket('ws://localhost:8080/api/v1/ws')` (use wss in prod).
-        2. `socket.onmessage = (event) => { ... }`
-        3. `socket.onclose = () => { ... }`
-    * The `onmessage` handler parses the `event.data`.
-    * `if (message.type === 'new_email') { ... }`
-* [ ] **Invalidate cache on message:**
+        1. Opens `new WebSocket(VITE_WS_URL || '<origin>/api/v1/ws')`.
+        2. Sets status in a `connection.store.ts` (Zustand) to `connecting`/`connected`/`disconnected`.
+        3. Handles `onmessage` and `onclose` to update connection state.
+    * The `onmessage` handler parses `event.data` and, when `message.type === 'new_email'`, invalidates queries for that folder.
+* [x] **Invalidate cache on message:**
     * Inside the `onmessage` handler:
-    * Get the `queryClient` using `useQueryClient()`.
-    * Call `queryClient.invalidateQueries({ queryKey: ['threads', message.folder] })`.
-    * This will automatically make `TanStack Query` refetch the thread list, and the new email will appear.
+    * Gets the `queryClient` using `useQueryClient()`.
+    * Calls `queryClient.invalidateQueries({ queryKey: ['threads', message.folder] })`.
+    * This automatically makes `TanStack Query` refetch the thread list, and the new email appears.
+* [x] **Connection status banner and manual reconnect:**
+    * Added a `ConnectionStatusBanner` component, shown when the WebSocket status is `disconnected`.
+    * The banner displays a Gmail-style "Connection lost. New emails may be delayed." message with a "Try now" link that triggers a reconnect of the WebSocket.
 
 #### **Testing**
 
-* [ ] **Frontend Integration (RTL + Mock WebSocket):**
-    * You'll need a library like `mock-socket`.
-    * Render the `Inbox.page.tsx` (which is inside `Layout.tsx`, so the hook runs).
-    * Simulate a message from the mock socket: `mockSocket.send('{"type": "new_email", "folder": "INBOX"}')`.
-    * **Assert** that `queryClient.invalidateQueries` was called with `['threads', 'INBOX']`.
-* [ ] **E2E:**
-    * This is the only true test.
-    * Log in to V-Mail. Have the Inbox page open.
-    * Use a *different* email client (or your `spike` script!) to send a new email to your test account.
-    * **Assert** the new email appears in the V-Mail inbox *without* a page reload.
+* [x] **Frontend Integration (RTL + WebSocket mocking via MSW):**
+    * Uses `msw`'s WebSocket support instead of `mock-socket`.
+    * Renders a component that uses `useWebSocket` under a `QueryClientProvider`.
+    * Simulates a message from the mock socket: `server.send('{"type": "new_email", "folder": "INBOX"}')`.
+    * **Asserts** that `queryClient.invalidateQueries` was called with `{ queryKey: ['threads', 'INBOX'] }`.
+* [x] **E2E:**
+    * Adds a new E2E test in `e2e/tests/inbox.spec.ts`.
+    * With the Inbox page open, the test calls `/test/add-imap-message` (a test-only backend endpoint) to append a message to `INBOX` on the IMAP server.
+    * **Asserts** the new email appears in the V-Mail inbox *without* a page reload.
 
 ## Milestone 6: Offline
 

backend/cmd/server/main.go

@@ -14,6 +14,7 @@ import (
 	"github.com/vdavid/vmail/backend/internal/crypto"
 	"github.com/vdavid/vmail/backend/internal/db"
 	"github.com/vdavid/vmail/backend/internal/imap"
+	ws "github.com/vdavid/vmail/backend/internal/websocket"
 )
 
 func main() {
@@ -50,13 +51,16 @@ func NewServer(cfg *config.Config, dbPool *pgxpool.Pool) http.Handler {
 
 	imapPool := imap.NewPoolWithMaxWorkers(cfg.IMAPMaxWorkers)
 	imapService := imap.NewService(dbPool, imapPool, encryptor)
+	wsHub := ws.NewHub(10)
 
 	authHandler := api.NewAuthHandler(dbPool)
 	settingsHandler := api.NewSettingsHandler(dbPool, encryptor)
 	foldersHandler := api.NewFoldersHandler(dbPool, encryptor, imapPool)
 	threadsHandler := api.NewThreadsHandler(dbPool, encryptor, imapService)
 	threadHandler := api.NewThreadHandler(dbPool, encryptor, imapService)
 	searchHandler := api.NewSearchHandler(dbPool, encryptor, imapService)
+	wsHandler := api.NewWebSocketHandler(dbPool, imapService, wsHub)
+	testHandler := api.NewTestHandler(dbPool, encryptor, imapService, wsHub)
 
 	mux := http.NewServeMux()
 
@@ -76,6 +80,13 @@ func NewServer(cfg *config.Config, dbPool *pgxpool.Pool) http.Handler {
 	mux.Handle("/api/v1/folders", auth.RequireAuth(http.HandlerFunc(foldersHandler.GetFolders)))
 	mux.Handle("/api/v1/threads", auth.RequireAuth(http.HandlerFunc(threadsHandler.GetThreads)))
 	mux.Handle("/api/v1/search", auth.RequireAuth(http.HandlerFunc(searchHandler.Search)))
+	// WebSocket handler handles its own authentication via query parameter
+	// (since browsers can't set headers on WebSocket connections).
+	mux.Handle("/api/v1/ws", http.HandlerFunc(wsHandler.Handle))
+	// Add test endpoints
+	if cfg.Environment == "test" {
+		mux.Handle("/test/add-imap-message", auth.RequireAuth(http.HandlerFunc(testHandler.AddIMAPMessage)))
+	}
 
 	// Handle /api/v1/thread/{thread_id} pattern
 	mux.Handle("/api/v1/thread/", auth.RequireAuth(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {

backend/cmd/test-server/main.go

@@ -23,6 +23,7 @@ import (
 	"github.com/vdavid/vmail/backend/internal/imap"
 	"github.com/vdavid/vmail/backend/internal/models"
 	"github.com/vdavid/vmail/backend/internal/testutil"
+	ws "github.com/vdavid/vmail/backend/internal/websocket"
 )
 
 func main() {
@@ -246,12 +247,15 @@ func NewServer(cfg *config.Config, dbPool *pgxpool.Pool) http.Handler {
 	imapPool := imap.NewPoolWithMaxWorkers(cfg.IMAPMaxWorkers)
 	imapService := imap.NewService(dbPool, imapPool, encryptor)
 
+	tsHub := ws.NewHub(10)
 	authHandler := api.NewAuthHandler(dbPool)
 	settingsHandler := api.NewSettingsHandler(dbPool, encryptor)
 	foldersHandler := api.NewFoldersHandler(dbPool, encryptor, imapPool)
 	threadsHandler := api.NewThreadsHandler(dbPool, encryptor, imapService)
 	threadHandler := api.NewThreadHandler(dbPool, encryptor, imapService)
 	searchHandler := api.NewSearchHandler(dbPool, encryptor, imapService)
+	wsHandler := api.NewWebSocketHandler(dbPool, imapService, tsHub)
+	testHandler := api.NewTestHandler(dbPool, encryptor, imapService, tsHub)
 
 	mux := http.NewServeMux()
 
@@ -271,6 +275,11 @@ func NewServer(cfg *config.Config, dbPool *pgxpool.Pool) http.Handler {
 	mux.Handle("/api/v1/folders", auth.RequireAuth(http.HandlerFunc(foldersHandler.GetFolders)))
 	mux.Handle("/api/v1/threads", auth.RequireAuth(http.HandlerFunc(threadsHandler.GetThreads)))
 	mux.Handle("/api/v1/search", auth.RequireAuth(http.HandlerFunc(searchHandler.Search)))
+	// WebSocket handler handles its own authentication via query parameter
+	// (since browsers can't set headers on WebSocket connections).
+	mux.Handle("/api/v1/ws", http.HandlerFunc(wsHandler.Handle))
+	// Test endpoints are only available in test environment
+	mux.Handle("/test/add-imap-message", auth.RequireAuth(http.HandlerFunc(testHandler.AddIMAPMessage)))
 
 	// Handle /api/v1/thread/{thread_id} pattern
 	mux.Handle("/api/v1/thread/", auth.RequireAuth(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {

backend/go.mod

@@ -4,8 +4,10 @@ go 1.25.3
 
 require (
 	github.com/emersion/go-imap v1.2.1
+	github.com/emersion/go-imap-idle v0.0.0-20210907174914-db2568431445
 	github.com/emersion/go-imap-sortthread v1.2.0
 	github.com/emersion/go-smtp v0.24.0
+	github.com/gorilla/websocket v1.5.3
 	github.com/jackc/pgx/v5 v5.7.6
 	github.com/jhillyerd/enmime v1.3.0
 	github.com/joho/godotenv v1.5.1

backend/go.sum

@@ -40,8 +40,11 @@ github.com/docker/go-units v0.5.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDD
 github.com/ebitengine/purego v0.8.4 h1:CF7LEKg5FFOsASUj0+QwaXf8Ht6TlFxg09+S9wz0omw=
 github.com/ebitengine/purego v0.8.4/go.mod h1:iIjxzd6CiRiOG0UyXP+V1+jWqUXVjPKLAI0mRfJZTmQ=
 github.com/emersion/go-imap v1.0.5/go.mod h1:yKASt+C3ZiDAiCSssxg9caIckWF/JG7ZQTO7GAmvicU=
+github.com/emersion/go-imap v1.0.6/go.mod h1:yKASt+C3ZiDAiCSssxg9caIckWF/JG7ZQTO7GAmvicU=
 github.com/emersion/go-imap v1.2.1 h1:+s9ZjMEjOB8NzZMVTM3cCenz2JrQIGGo5j1df19WjTA=
 github.com/emersion/go-imap v1.2.1/go.mod h1:Qlx1FSx2FTxjnjWpIlVNEuX+ylerZQNFE5NsmKFSejY=
+github.com/emersion/go-imap-idle v0.0.0-20210907174914-db2568431445 h1:dAGbaaU4LLupO7dnYZaELOoI3RoVDNi5DCGejLe8a7c=
+github.com/emersion/go-imap-idle v0.0.0-20210907174914-db2568431445/go.mod h1:N/6S3dRTVt8xT867m+476C16+v/Fq4WZYvh2Chg0nmg=
 github.com/emersion/go-imap-sortthread v1.2.0 h1:EMVEJXPWAhXMWECjR82Rn/tza6MddcvTwGAdTu1vJKU=
 github.com/emersion/go-imap-sortthread v1.2.0/go.mod h1:UhenCBupR+vSYRnqJkpjSq84INUCsyAK1MLpogv14pE=
 github.com/emersion/go-message v0.11.1/go.mod h1:C4jnca5HOTo4bGN9YdqNQM9sITuT3Y0K6bSUw9RklvY=
@@ -74,6 +77,8 @@ github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
 github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
 github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
 github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
+github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
+github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
 github.com/grpc-ecosystem/grpc-gateway/v2 v2.27.3 h1:NmZ1PKzSTQbuGHw9DGPFomqkkLWMC+vZCkfs+FHv1Vg=
 github.com/grpc-ecosystem/grpc-gateway/v2 v2.27.3/go.mod h1:zQrxl1YP88HQlA6i9c63DSVPFklWpGX4OWAc9bFuaH4=
 github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM=

backend/internal/api/auth_handler_test.go

@@ -102,10 +102,10 @@ func TestAuthHandler_GetAuthStatus(t *testing.T) {
 	t.Run("returns 500 when GetOrCreateUser returns an error", func(t *testing.T) {
 		req := httptest.NewRequest("GET", "/api/v1/auth/status", nil)
 
-		// Use a cancelled context to simulate database connection failure
-		cancelledCtx, cancel := context.WithCancel(context.Background())
+		// Use a canceled context to simulate database connection failure
+		canceledCtx, cancel := context.WithCancel(context.Background())
 		cancel()
-		reqCtx := context.WithValue(cancelledCtx, auth.UserEmailKey, "test@example.com")
+		reqCtx := context.WithValue(canceledCtx, auth.UserEmailKey, "test@example.com")
 		req = req.WithContext(reqCtx)
 
 		rr := httptest.NewRecorder()

backend/internal/api/folders_handler_test.go

@@ -60,11 +60,11 @@ func TestFoldersHandler_GetFolders(t *testing.T) {
 	t.Run("returns 500 when GetOrCreateUser returns an error", func(t *testing.T) {
 		email := "dberror@example.com"
 
-		// Use a cancelled context to simulate database connection failure
-		cancelledCtx, cancel := context.WithCancel(context.Background())
+		// Use a canceled context to simulate database connection failure
+		canceledCtx, cancel := context.WithCancel(context.Background())
 		cancel()
 		req := httptest.NewRequest("GET", "/api/v1/folders", nil)
-		reqCtx := context.WithValue(cancelledCtx, auth.UserEmailKey, email)
+		reqCtx := context.WithValue(canceledCtx, auth.UserEmailKey, email)
 		req = req.WithContext(reqCtx)
 
 		rr := httptest.NewRecorder()
@@ -98,8 +98,11 @@ type mockIMAPPool struct {
 	getClientPass      string
 	removeClientCalled map[string]bool
 	// For retry scenarios: the first call returns one client, the second call returns another
-	retryClient    imap.IMAPClient
-	retryClientErr error
+	retryClient          imap.IMAPClient
+	retryClientErr       error
+	listenerClient       imap.ListenerClient
+	listenerClientErr    error
+	removeListenerCalled map[string]bool
 }
 
 func (m *mockIMAPPool) WithClient(userID, server, username, password string, fn func(imap.IMAPClient) error) error {
@@ -138,6 +141,20 @@ func (m *mockIMAPPool) RemoveClient(userID string) {
 
 func (m *mockIMAPPool) Close() {}
 
+func (m *mockIMAPPool) GetListenerConnection(string, string, string, string) (imap.ListenerClient, error) {
+	if m.listenerClientErr != nil {
+		return nil, m.listenerClientErr
+	}
+	return m.listenerClient, nil
+}
+
+func (m *mockIMAPPool) RemoveListenerConnection(userID string) {
+	if m.removeListenerCalled == nil {
+		m.removeListenerCalled = make(map[string]bool)
+	}
+	m.removeListenerCalled[userID] = true
+}
+
 // callGetFolders is a helper function that sets up and calls GetFolders handler.
 // It returns the response recorder for assertions.
 func callGetFolders(t *testing.T, handler *FoldersHandler, email string) *httptest.ResponseRecorder {

backend/internal/api/search_handler_test.go

@@ -12,6 +12,7 @@ import (
 	"github.com/vdavid/vmail/backend/internal/imap"
 	"github.com/vdavid/vmail/backend/internal/models"
 	"github.com/vdavid/vmail/backend/internal/testutil"
+	ws "github.com/vdavid/vmail/backend/internal/websocket"
 )
 
 func TestSearchHandler_Search(t *testing.T) {
@@ -321,6 +322,10 @@ func (m *mockIMAPServiceForSearch) Search(_ context.Context, _ string, query str
 
 func (m *mockIMAPServiceForSearch) Close() {}
 
+// StartIdleListener is part of the IMAPService interface but is not used in search tests.
+func (m *mockIMAPServiceForSearch) StartIdleListener(context.Context, string, *ws.Hub) {
+}
+
 type imapError struct {
 	message string
 }

backend/internal/api/settings_handler_test.go

@@ -334,8 +334,8 @@ func TestSettingsHandler_PostSettings(t *testing.T) {
 	t.Run("returns 500 when GetUserSettings returns non-NotFound error in PostSettings", func(t *testing.T) {
 		email := "dberror-post@example.com"
 
-		// Use a cancelled context to simulate database connection failure
-		cancelledCtx, cancel := context.WithCancel(context.Background())
+		// Use a canceled context to simulate database connection failure
+		canceledCtx, cancel := context.WithCancel(context.Background())
 		cancel()
 
 		reqBody := models.UserSettingsRequest{
@@ -351,7 +351,7 @@ func TestSettingsHandler_PostSettings(t *testing.T) {
 
 		body, _ := json.Marshal(reqBody)
 		req := httptest.NewRequest("POST", "/api/v1/settings", bytes.NewReader(body))
-		reqCtx := context.WithValue(cancelledCtx, auth.UserEmailKey, email)
+		reqCtx := context.WithValue(canceledCtx, auth.UserEmailKey, email)
 		req = req.WithContext(reqCtx)
 
 		rr := httptest.NewRecorder()
@@ -371,12 +371,12 @@ func TestSettingsHandler_PostSettings(t *testing.T) {
 	t.Run("returns 500 when GetUserSettings returns non-NotFound error in GetSettings", func(t *testing.T) {
 		email := "dberror-get@example.com"
 
-		// Use a cancelled context to simulate database connection failure
-		cancelledCtx, cancel := context.WithCancel(context.Background())
+		// Use a canceled context to simulate database connection failure
+		canceledCtx, cancel := context.WithCancel(context.Background())
 		cancel()
 
 		req := httptest.NewRequest("GET", "/api/v1/settings", nil)
-		reqCtx := context.WithValue(cancelledCtx, auth.UserEmailKey, email)
+		reqCtx := context.WithValue(canceledCtx, auth.UserEmailKey, email)
 		req = req.WithContext(reqCtx)
 
 		rr := httptest.NewRecorder()