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

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: 'delivery.created'; payload: Delivery }
-  | { type: 'delivery.updated'; payload: Delivery }
-  | { type: 'delivery.deleted'; payload: { id: string } };
+type WsEventType =
+  | 'delivery.created'
+  | 'delivery.updated'
+  | 'delivery.status_changed'
+  | 'delivery.deleted';
 
-type EventHandler = (event: WebSocketEvent) => void;
+interface WsEvent {
+  type: WsEventType;
+  payload: unknown;
+}
 
-class MockWebSocket {
-  private handlers: EventHandler[] = [];
-  private isConnected = false;
+const MAX_RECONNECT_DELAY = 30_000;
+const INITIAL_RECONNECT_DELAY = 1_000;
+const COUNTS_DEBOUNCE_MS = 300;
 
-  connect() {
-    this.isConnected = true;
-    console.log('WebSocket connected (mock)');
+function getWsUrl(token: string): string {
+  const apiBase = import.meta.env.VITE_API_URL || '';
+  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() {
-    this.isConnected = false;
-    console.log('WebSocket disconnected (mock)');
-  }
+  return `${wsBase}/api/ws?token=${encodeURIComponent(token)}`;
+}
 
-  subscribe(handler: EventHandler) {
-    this.handlers.push(handler);
-    return () => {
-      this.handlers = this.handlers.filter((h) => h !== handler);
-    };
-  }
+// Debounced counts refresh — multiple WS events in quick succession
+// trigger only a single API call.
+let countsTimer: ReturnType<typeof setTimeout> | null = null;
+function refreshCountsDebounced() {
+  if (countsTimer) clearTimeout(countsTimer);
+  countsTimer = setTimeout(() => {
+    countsTimer = null;
+    useDeliveryStore.getState().fetchDeliveryCounts();
+  }, COUNTS_DEBOUNCE_MS);
+}
 
-  emit(event: WebSocketEvent) {
-    if (!this.isConnected) return;
-    this.handlers.forEach((handler) => handler(event));
-  }
+function handleEvent(event: WsEvent) {
+  const store = useDeliveryStore.getState();
 
-  simulateIncomingEvent(event: WebSocketEvent) {
-    this.emit(event);
+  switch (event.type) {
+    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 const useWebSocket = () => {
-  const { addDelivery, updateDelivery, deleteDelivery } = useDeliveryStore();
+export function useWebSocket() {
+  const token = useAuthStore(state => state.token);
 
   useEffect(() => {
-    mockWebSocket.connect();
+    if (!token) return;
 
-    const unsubscribe = mockWebSocket.subscribe((event) => {
-      switch (event.type) {
-        case 'delivery.created':
-          addDelivery(event.payload);
-          break;
-        case 'delivery.updated':
-          updateDelivery(event.payload.id, event.payload);
-          break;
-        case 'delivery.deleted':
-          deleteDelivery(event.payload.id);
-          break;
-      }
-    });
+    let cancelled = false;
+    let ws: WebSocket | null = null;
+    let reconnectTimer: ReturnType<typeof setTimeout> | undefined;
+    let reconnectDelay = INITIAL_RECONNECT_DELAY;
+
+    function connect() {
+      if (cancelled) return;
+
+      const socket = new WebSocket(getWsUrl(token!));
+      ws = socket;
+
+      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();
-      mockWebSocket.disconnect();
+      cancelled = true;
+      clearTimeout(reconnectTimer);
+      ws?.close();
+      ws = null;
     };
-  }, [addDelivery, updateDelivery, deleteDelivery]);
-
-  const sendEvent = useCallback((event: WebSocketEvent) => {
-    mockWebSocket.emit(event);
-  }, []);
-
-  return { sendEvent, isConnected: true };
-};
-
-export const simulateIncomingDelivery = (delivery: Delivery) => {
-  mockWebSocket.simulateIncomingEvent({
-    type: 'delivery.created',
-    payload: delivery,
-  });
-};
+  }, [token]);
+}

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),
+    }));
+  },
 }));