add WebSocket support for real-time delivery updates with JWT authentication and automatic reconnection

backend/cmd/api/main.go

@@ -10,6 +10,7 @@ import (
 	"github.com/chedius/delivery-tracker/internal/auth"
 	db "github.com/chedius/delivery-tracker/internal/db/sqlc"
 	"github.com/chedius/delivery-tracker/internal/delivery"
+	"github.com/chedius/delivery-tracker/internal/ws"
 	"github.com/gin-contrib/cors"
 	"github.com/gin-gonic/gin"
 	"github.com/jackc/pgx/v5/pgxpool"
@@ -42,7 +43,8 @@ func main() {
 
 	queries := db.New(pool)
 	_, authHandler := initAuth(queries)
-	h := delivery.NewHandler(queries)
+	hub := ws.NewHub()
+	h := delivery.NewHandler(queries, hub)
 
 	r := gin.Default()
 
@@ -61,6 +63,7 @@ func main() {
 
 	r.POST("/api/auth/register", authHandler.Register)
 	r.POST("/api/auth/login", authHandler.Login)
+	r.GET("/api/ws", ws.HandleWS(hub, []byte(os.Getenv("JWT_SECRET"))))
 
 	authorized := r.Group("/api")
 	authorized.Use(auth.AuthMiddleware([]byte(os.Getenv("JWT_SECRET"))))

backend/go.mod

@@ -14,6 +14,8 @@ require (
 
 require github.com/golang-jwt/jwt/v5 v5.3.1
 
+require github.com/gorilla/websocket v1.5.3
+
 require (
 	github.com/bytedance/gopkg v0.1.3 // indirect
 	github.com/bytedance/sonic v1.15.0 // indirect

backend/go.sum

@@ -36,6 +36,8 @@ github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX
 github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
 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/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM=
 github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg=
 github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo=

backend/internal/delivery/handler.go

@@ -2,10 +2,12 @@ package delivery
 
 import (
 	"errors"
+	"log"
 	"net/http"
 	"time"
 
 	sqlc "github.com/chedius/delivery-tracker/internal/db/sqlc"
+	"github.com/chedius/delivery-tracker/internal/ws"
 	"github.com/gin-gonic/gin"
 	"github.com/google/uuid"
 	"github.com/jackc/pgx/v5/pgtype"
@@ -13,6 +15,7 @@ import (
 
 type Handler struct {
 	queries *sqlc.Queries
+	hub     *ws.Hub
 }
 
 // DeliveryRequest represents the request body for creating or updating a delivery
@@ -38,8 +41,8 @@ type DeliveryRequest struct {
 	Comment                 string  `json:"comment"`
 }
 
-func NewHandler(queries *sqlc.Queries) *Handler {
+func NewHandler(queries *sqlc.Queries, hub *ws.Hub) *Handler {
-	return &Handler{queries: queries}
+	return &Handler{queries: queries, hub: hub}
 }
 
 // GET /api/deliveries/:id
@@ -139,6 +142,7 @@ func (h *Handler) CreateDelivery(c *gin.Context) {
 		return
 	}
 
+	h.hub.Broadcast(ws.NewEvent(ws.DeliveryCreated, res))
 	c.JSON(http.StatusOK, gin.H{"message": "Delivery created", "id": res.ID.String()})
 }
 
@@ -199,6 +203,11 @@ func (h *Handler) UpdateDelivery(c *gin.Context) {
 		return
 	}
 
+	if updated, err := h.queries.GetDeliveryByID(c.Request.Context(), pgtype.UUID{Bytes: parsedID, Valid: true}); err == nil {
+		h.hub.Broadcast(ws.NewEvent(ws.DeliveryUpdated, updated))
+	} else {
+		log.Printf("delivery: failed to fetch updated delivery %s for ws broadcast: %v", id, err)
+	}
 	c.JSON(http.StatusOK, gin.H{"message": "Delivery updated"})
 }
 
@@ -240,6 +249,7 @@ func (h *Handler) UpdateDeliveryStatus(c *gin.Context) {
 		return
 	}
 
+	h.hub.Broadcast(ws.NewEvent(ws.DeliveryStatusChanged, ws.StatusPayload{ID: id, Status: status}))
 	c.JSON(http.StatusOK, gin.H{"message": "Delivery status updated"})
 }
 
@@ -263,6 +273,7 @@ func (h *Handler) DeleteDelivery(c *gin.Context) {
 		return
 	}
 
+	h.hub.Broadcast(ws.NewEvent(ws.DeliveryDeleted, ws.DeletePayload{ID: id}))
 	c.JSON(http.StatusOK, gin.H{"message": "Delivery deleted"})
 }
 

backend/internal/ws/client.go

@@ -0,0 +1,74 @@
+package ws
+
+import (
+	"time"
+
+	"github.com/gorilla/websocket"
+)
+
+const (
+	writeWait  = 10 * time.Second
+	pongWait   = 60 * time.Second
+	pingPeriod = (pongWait * 9) / 10
+)
+
+type Client struct {
+	hub  *Hub
+	conn *websocket.Conn
+	send chan []byte
+}
+
+func NewClient(hub *Hub, conn *websocket.Conn) *Client {
+	return &Client{
+		hub:  hub,
+		conn: conn,
+		send: make(chan []byte, 256),
+	}
+}
+
+// ReadPump keeps the connection alive and handles pong frames.
+// We don't expect real messages from the client (read-only WS).
+func (c *Client) ReadPump() {
+	defer func() {
+		c.hub.Unregister(c)
+		c.conn.Close()
+	}()
+	c.conn.SetReadLimit(512)
+	c.conn.SetReadDeadline(time.Now().Add(pongWait))
+	c.conn.SetPongHandler(func(string) error {
+		c.conn.SetReadDeadline(time.Now().Add(pongWait))
+		return nil
+	})
+	for {
+		if _, _, err := c.conn.ReadMessage(); err != nil {
+			break
+		}
+	}
+}
+
+// WritePump sends messages from hub to the client and pings to keep alive.
+func (c *Client) WritePump() {
+	ticker := time.NewTicker(pingPeriod)
+	defer func() {
+		ticker.Stop()
+		c.conn.Close()
+	}()
+	for {
+		select {
+		case msg, ok := <-c.send:
+			c.conn.SetWriteDeadline(time.Now().Add(writeWait))
+			if !ok {
+				c.conn.WriteMessage(websocket.CloseMessage, nil)
+				return
+			}
+			if err := c.conn.WriteMessage(websocket.TextMessage, msg); err != nil {
+				return
+			}
+		case <-ticker.C:
+			c.conn.SetWriteDeadline(time.Now().Add(writeWait))
+			if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
+				return
+			}
+		}
+	}
+}

backend/internal/ws/event.go

@@ -0,0 +1,40 @@
+package ws
+
+import (
+	"encoding/json"
+	"log"
+)
+
+type EventType string
+
+const (
+	DeliveryCreated       EventType = "delivery.created"
+	DeliveryUpdated       EventType = "delivery.updated"
+	DeliveryStatusChanged EventType = "delivery.status_changed"
+	DeliveryDeleted       EventType = "delivery.deleted"
+)
+
+type Event struct {
+	Type    EventType `json:"type"`
+	Payload any       `json:"payload"`
+}
+
+type StatusPayload struct {
+	ID     string `json:"id"`
+	Status string `json:"status"`
+}
+
+type DeletePayload struct {
+	ID string `json:"id"`
+}
+
+// NewEvent serializes an event to JSON. Returns nil on marshal failure;
+// callers should treat nil as "skip broadcast".
+func NewEvent(eventType EventType, payload any) []byte {
+	data, err := json.Marshal(Event{Type: eventType, Payload: payload})
+	if err != nil {
+		log.Printf("ws: failed to marshal event %s: %v", eventType, err)
+		return nil
+	}
+	return data
+}

backend/internal/ws/handler.go

@@ -0,0 +1,41 @@
+package ws
+
+import (
+	"net/http"
+
+	"github.com/chedius/delivery-tracker/internal/auth"
+	"github.com/gin-gonic/gin"
+	"github.com/gorilla/websocket"
+)
+
+var upgrader = websocket.Upgrader{
+	CheckOrigin: func(r *http.Request) bool { return true },
+}
+
+// HandleWS returns a Gin handler that upgrades HTTP to WebSocket
+// after validating the JWT token from the ?token= query param.
+func HandleWS(hub *Hub, jwtSecret []byte) gin.HandlerFunc {
+	return func(c *gin.Context) {
+		token := c.Query("token")
+		if token == "" {
+			c.JSON(http.StatusUnauthorized, gin.H{"error": "missing token"})
+			return
+		}
+
+		if _, err := auth.ParseToken(token, jwtSecret); err != nil {
+			c.JSON(http.StatusUnauthorized, gin.H{"error": "invalid token"})
+			return
+		}
+
+		conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
+		if err != nil {
+			return
+		}
+
+		client := NewClient(hub, conn)
+		hub.Register(client)
+
+		go client.WritePump()
+		go client.ReadPump()
+	}
+}

backend/internal/ws/hub.go

@@ -0,0 +1,48 @@
+package ws
+
+import "sync"
+
+type Hub struct {
+	mu      sync.RWMutex
+	clients map[*Client]struct{}
+}
+
+func NewHub() *Hub {
+	return &Hub{
+		clients: make(map[*Client]struct{}),
+	}
+}
+
+func (h *Hub) Register(c *Client) {
+	h.mu.Lock()
+	h.clients[c] = struct{}{}
+	h.mu.Unlock()
+}
+
+func (h *Hub) Unregister(c *Client) {
+	h.mu.Lock()
+	if _, ok := h.clients[c]; ok {
+		delete(h.clients, c)
+		close(c.send)
+	}
+	h.mu.Unlock()
+}
+
+func (h *Hub) Broadcast(msg []byte) {
+	if msg == nil {
+		return
+	}
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+	for c := range h.clients {
+		select {
+		case c.send <- msg:
+		default:
+			// Client too slow, schedule disconnect
+			go func(c *Client) {
+				h.Unregister(c)
+				c.conn.Close()
+			}(c)
+		}
+	}
+}

frontend/nginx.conf

@@ -49,6 +49,20 @@ server {
         expires 0;
     }
 
+    # WebSocket endpoint — long-lived connections need extended timeouts
+    location = /api/ws {
+        proxy_pass http://backend:8080;
+        proxy_http_version 1.1;
+        proxy_set_header Upgrade $http_upgrade;
+        proxy_set_header Connection "upgrade";
+        proxy_set_header Host $host;
+        proxy_set_header X-Real-IP $remote_addr;
+        proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
+        proxy_set_header X-Forwarded-Proto $scheme;
+        proxy_read_timeout 86400s;
+        proxy_send_timeout 86400s;
+    }
+
     # Proxy API requests to backend
     location /api/ {
         proxy_pass http://backend:8080;

frontend/package-lock.json

@@ -1,12 +1,12 @@
 {
   "name": "delivery-tracker",
-  "version": "0.0.4",
+  "version": "0.0.5",
   "lockfileVersion": 3,
   "requires": true,
   "packages": {
     "": {
       "name": "delivery-tracker",
-      "version": "0.0.4",
+      "version": "0.0.5",
       "dependencies": {
         "@tailwindcss/vite": "^4.2.2",
         "date-fns": "^4.1.0",

frontend/package.json

@@ -1,7 +1,7 @@
 {
   "name": "delivery-tracker",
   "private": true,
-  "version": "0.0.4",
+  "version": "0.0.5",
   "type": "module",
   "scripts": {
     "dev": "vite",

frontend/src/App.tsx

@@ -7,6 +7,7 @@ import { Button } from './components/ui/Button';
 import { UpdatePrompt } from './components/ui/UpdatePrompt';
 import { useDeliveryStore } from './stores/deliveryStore';
 import { useAuthStore } from './stores/authStore';
+import { useWebSocket } from './hooks/useWebSocket';
 
 // Lazy load pages for code splitting
 const Dashboard = lazy(() => import('./pages/Dashboard'));
@@ -30,6 +31,8 @@ function App() {
   const addDelivery = useDeliveryStore(state => state.addDelivery);
   const fetchDeliveryCounts = useDeliveryStore(state => state.fetchDeliveryCounts);
 
+  useWebSocket();
+
   // Restore auth on mount
   useEffect(() => {
     restoreAuth();

frontend/src/api/deliveries.ts

@@ -3,7 +3,7 @@ import { backendDateToFrontend } from '../utils/date';
 import type { Delivery, DeliveryRequestSource, PickupLocation, DeliveryStatus } from '../types';
 
 // Types matching backend responses
-interface BackendDelivery {
+export interface BackendDelivery {
   id: string;
   date: string; // YYYY-MM-DD from pgtype.Date
   pickup_location: PickupLocation;
@@ -57,7 +57,7 @@ interface UpdateDeliveryResponse {
 }
 
 // Map backend delivery to frontend delivery
-function mapBackendToFrontend(backend: BackendDelivery): Delivery {
+export function mapBackendToFrontend(backend: BackendDelivery): Delivery {
   return {
     id: backend.id,
     date: backendDateToFrontend(backend.date),

frontend/src/hooks/useWebSocket.ts

@@ -1,83 +1,133 @@
-import { useEffect, useCallback } from 'react';
+import { useEffect } from 'react';
+import { useAuthStore } from '../stores/authStore';
 import { useDeliveryStore } from '../stores/deliveryStore';
-import type { Delivery } from '../types';
+import { mapBackendToFrontend } from '../api/deliveries';
+import type { BackendDelivery } from '../api/deliveries';
 
-type WebSocketEvent = 
+type WsEventType =
-  | { type: 'delivery.created'; payload: Delivery }
+  | 'delivery.created'
-  | { type: 'delivery.updated'; payload: Delivery }
+  | 'delivery.updated'
-  | { type: 'delivery.deleted'; payload: { id: string } };
+  | 'delivery.status_changed'
+  | 'delivery.deleted';
 
-type EventHandler = (event: WebSocketEvent) => void;
+interface WsEvent {
+  type: WsEventType;
+  payload: unknown;
+}
 
-class MockWebSocket {
+const MAX_RECONNECT_DELAY = 30_000;
-  private handlers: EventHandler[] = [];
+const INITIAL_RECONNECT_DELAY = 1_000;
-  private isConnected = false;
+const COUNTS_DEBOUNCE_MS = 300;
 
-  connect() {
+function getWsUrl(token: string): string {
-    this.isConnected = true;
+  const apiBase = import.meta.env.VITE_API_URL || '';
-    console.log('WebSocket connected (mock)');
+  let wsBase: string;
+
+  if (apiBase) {
+    wsBase = apiBase.replace(/^http/, 'ws');
+  } else {
+    const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
+    wsBase = `${protocol}//${window.location.host}`;
   }
 
-  disconnect() {
+  return `${wsBase}/api/ws?token=${encodeURIComponent(token)}`;
-    this.isConnected = false;
+}
-    console.log('WebSocket disconnected (mock)');
-  }
 
-  subscribe(handler: EventHandler) {
+// Debounced counts refresh — multiple WS events in quick succession
-    this.handlers.push(handler);
+// trigger only a single API call.
-    return () => {
+let countsTimer: ReturnType<typeof setTimeout> | null = null;
-      this.handlers = this.handlers.filter((h) => h !== handler);
+function refreshCountsDebounced() {
-    };
+  if (countsTimer) clearTimeout(countsTimer);
-  }
+  countsTimer = setTimeout(() => {
+    countsTimer = null;
+    useDeliveryStore.getState().fetchDeliveryCounts();
+  }, COUNTS_DEBOUNCE_MS);
+}
 
-  emit(event: WebSocketEvent) {
+function handleEvent(event: WsEvent) {
-    if (!this.isConnected) return;
+  const store = useDeliveryStore.getState();
-    this.handlers.forEach((handler) => handler(event));
-  }
 
-  simulateIncomingEvent(event: WebSocketEvent) {
+  switch (event.type) {
-    this.emit(event);
+    case 'delivery.created': {
+      const delivery = mapBackendToFrontend(event.payload as BackendDelivery);
+      store.handleWsDeliveryCreated(delivery);
+      refreshCountsDebounced();
+      break;
+    }
+    case 'delivery.updated': {
+      const delivery = mapBackendToFrontend(event.payload as BackendDelivery);
+      store.handleWsDeliveryUpdated(delivery);
+      refreshCountsDebounced();
+      break;
+    }
+    case 'delivery.status_changed': {
+      const { id, status } = event.payload as { id: string; status: string };
+      store.handleWsStatusChanged(id, status);
+      break;
+    }
+    case 'delivery.deleted': {
+      const { id } = event.payload as { id: string };
+      store.handleWsDeliveryDeleted(id);
+      refreshCountsDebounced();
+      break;
+    }
   }
 }
 
-const mockWebSocket = new MockWebSocket();
+export function useWebSocket() {
-
+  const token = useAuthStore(state => state.token);
-export const useWebSocket = () => {
-  const { addDelivery, updateDelivery, deleteDelivery } = useDeliveryStore();
 
   useEffect(() => {
-    mockWebSocket.connect();
+    if (!token) return;
 
-    const unsubscribe = mockWebSocket.subscribe((event) => {
+    let cancelled = false;
-      switch (event.type) {
+    let ws: WebSocket | null = null;
-        case 'delivery.created':
+    let reconnectTimer: ReturnType<typeof setTimeout> | undefined;
-          addDelivery(event.payload);
+    let reconnectDelay = INITIAL_RECONNECT_DELAY;
-          break;
+
-        case 'delivery.updated':
+    function connect() {
-          updateDelivery(event.payload.id, event.payload);
+      if (cancelled) return;
-          break;
+
-        case 'delivery.deleted':
+      const socket = new WebSocket(getWsUrl(token!));
-          deleteDelivery(event.payload.id);
+      ws = socket;
-          break;
+
-      }
+      socket.onopen = () => {
-    });
+        reconnectDelay = INITIAL_RECONNECT_DELAY;
+      };
+
+      socket.onmessage = (e) => {
+        try {
+          const event: WsEvent = JSON.parse(e.data);
+          handleEvent(event);
+        } catch {
+          // ignore malformed messages
+        }
+      };
+
+      socket.onclose = () => {
+        // Ignore close events from sockets we no longer track —
+        // protects against stale callbacks after token change.
+        if (cancelled || ws !== socket) return;
+        ws = null;
+
+        reconnectTimer = setTimeout(() => {
+          reconnectDelay = Math.min(reconnectDelay * 2, MAX_RECONNECT_DELAY);
+          connect();
+        }, reconnectDelay);
+      };
+
+      socket.onerror = () => {
+        socket.close();
+      };
+    }
+
+    connect();
 
     return () => {
-      unsubscribe();
+      cancelled = true;
-      mockWebSocket.disconnect();
+      clearTimeout(reconnectTimer);
+      ws?.close();
+      ws = null;
     };
-  }, [addDelivery, updateDelivery, deleteDelivery]);
+  }, [token]);
-
+}
-  const sendEvent = useCallback((event: WebSocketEvent) => {
-    mockWebSocket.emit(event);
-  }, []);
-
-  return { sendEvent, isConnected: true };
-};
-
-export const simulateIncomingDelivery = (delivery: Delivery) => {
-  mockWebSocket.simulateIncomingEvent({
-    type: 'delivery.created',
-    payload: delivery,
-  });
-};

frontend/src/stores/deliveryStore.ts

@@ -7,6 +7,7 @@ interface DeliveryState {
   // Data
   deliveries: Delivery[];
   deliveryCounts: Record<string, number>;
+  currentDate: string | null;
 
   // Loading states
   isLoading: boolean;
@@ -24,19 +25,26 @@ interface DeliveryState {
   getDeliveriesByDateRange: (startDate: string, endDate: string) => Delivery[];
   getDeliveryCountsByDate: () => Record<string, number>;
   clearError: () => void;
+
+  // WebSocket event handlers
+  handleWsDeliveryCreated: (delivery: Delivery) => void;
+  handleWsDeliveryUpdated: (delivery: Delivery) => void;
+  handleWsStatusChanged: (id: string, status: string) => void;
+  handleWsDeliveryDeleted: (id: string) => void;
 }
 
 export const useDeliveryStore = create<DeliveryState>()((set, get) => ({
   // Initial state
   deliveries: [],
   deliveryCounts: {},
+  currentDate: null,
   isLoading: false,
   isLoadingCounts: false,
   error: null,
 
   // Fetch deliveries for a specific date
   fetchDeliveriesByDate: async (date: string) => {
-    set({ isLoading: true, error: null });
+    set({ isLoading: true, error: null, currentDate: date });
     try {
       const deliveries = await deliveriesApi.getByDate(date);
       set({ deliveries, isLoading: false });
@@ -174,4 +182,42 @@ export const useDeliveryStore = create<DeliveryState>()((set, get) => ({
   },
 
   clearError: () => set({ error: null }),
+
+  // WebSocket event handlers (update local state without refetching)
+  handleWsDeliveryCreated: (delivery: Delivery) => {
+    const { currentDate, deliveries } = get();
+    if (currentDate && delivery.date === currentDate) {
+      if (!deliveries.some(d => d.id === delivery.id)) {
+        set({ deliveries: [...deliveries, delivery] });
+      }
+    }
+  },
+
+  handleWsDeliveryUpdated: (delivery: Delivery) => {
+    const { currentDate, deliveries } = get();
+    const exists = deliveries.some(d => d.id === delivery.id);
+    if (exists) {
+      if (delivery.date === currentDate) {
+        set({ deliveries: deliveries.map(d => d.id === delivery.id ? delivery : d) });
+      } else {
+        set({ deliveries: deliveries.filter(d => d.id !== delivery.id) });
+      }
+    } else if (currentDate && delivery.date === currentDate) {
+      set({ deliveries: [...deliveries, delivery] });
+    }
+  },
+
+  handleWsStatusChanged: (id: string, status: string) => {
+    set(state => ({
+      deliveries: state.deliveries.map(d =>
+        d.id === id ? { ...d, status: status as DeliveryStatus, updatedAt: Date.now() } : d
+      ),
+    }));
+  },
+
+  handleWsDeliveryDeleted: (id: string) => {
+    set(state => ({
+      deliveries: state.deliveries.filter(d => d.id !== id),
+    }));
+  },
 }));

frontend/vite.config.ts

@@ -46,6 +46,7 @@ export default defineConfig({
       '/api': {
         target: 'http://localhost:8080',
         changeOrigin: true,
+        ws: true,
       },
     },
   },