From d5da6d456220eca647972c1a56729a8782dbc133 Mon Sep 17 00:00:00 2001
From: Nico <47644445+nicotsx@users.noreply.github.com>
Date: Thu, 5 Mar 2026 22:32:44 +0100
Subject: [PATCH] fix: sse connection  (#623)

### TL;DR

Refactored `useServerEvents` hook to use a shared EventSource connection across all components instead of creating individual connections per hook instance.

### What changed?

- Introduced shared state management for EventSource connections with reference counting
- Added comprehensive test coverage with MockEventSource implementation
- Modified the hook to share a single EventSource connection across multiple consumers
- Implemented proper cleanup when the last subscriber unmounts
- Removed automatic query refetching on backup completion, keeping only cache invalidation

### How to test?

Run the new test suite with `bun test use-server-events.test.tsx` to verify:
- Single EventSource instance is shared across multiple hook consumers
- Event listeners work correctly with the shared connection
- Cache invalidation occurs once per event
- Proper cleanup happens when all subscribers unmount

### Why make this change?

This optimization reduces resource usage by preventing multiple EventSource connections when the hook is used in different components. The shared connection approach is more efficient while maintaining the same functionality, and the added tests ensure reliability of the server events system.
---
 .../__tests__/use-server-events.test.tsx      | 132 +++++++++++
 app/client/hooks/use-server-events.ts         | 221 +++++++++++-------
 2 files changed, 270 insertions(+), 83 deletions(-)
 create mode 100644 app/client/hooks/__tests__/use-server-events.test.tsx

diff --git a/app/client/hooks/__tests__/use-server-events.test.tsx b/app/client/hooks/__tests__/use-server-events.test.tsx
new file mode 100644
index 00000000..f5a52a19
--- /dev/null
+++ b/app/client/hooks/__tests__/use-server-events.test.tsx
@@ -0,0 +1,132 @@
+import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test";
+import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
+import { cleanup, render, screen } from "@testing-library/react";
+import { useEffect, useState } from "react";
+import { useServerEvents } from "../use-server-events";
+
+class MockEventSource {
+	static instances: MockEventSource[] = [];
+
+	public onerror: ((event: Event) => void) | null = null;
+	public close = mock(() => {});
+	private listeners = new Map<string, Set<(event: Event) => void>>();
+
+	constructor(public url: string) {
+		MockEventSource.instances.push(this);
+	}
+
+	addEventListener(type: string, listener: EventListenerOrEventListenerObject) {
+		const listeners = this.listeners.get(type) ?? new Set<(event: Event) => void>();
+		const callback = typeof listener === "function" ? listener : (event: Event) => listener.handleEvent(event);
+		listeners.add(callback);
+		this.listeners.set(type, listeners);
+	}
+
+	emit(type: string, data: unknown) {
+		const event = new MessageEvent(type, {
+			data: JSON.stringify(data),
+		});
+		for (const listener of this.listeners.get(type) ?? []) {
+			listener(event);
+		}
+	}
+
+	static reset() {
+		MockEventSource.instances = [];
+	}
+}
+
+const originalEventSource = globalThis.EventSource;
+const originalConsoleInfo = console.info;
+const originalConsoleError = console.error;
+
+const createTestQueryClient = () =>
+	new QueryClient({
+		defaultOptions: {
+			queries: {
+				retry: false,
+				gcTime: Infinity,
+			},
+			mutations: {
+				gcTime: Infinity,
+			},
+		},
+	});
+
+const ConnectionConsumer = () => {
+	useServerEvents();
+	return null;
+};
+
+const BackupCompletedListener = ({ scheduleId }: { scheduleId: string }) => {
+	const { addEventListener } = useServerEvents();
+	const [status, setStatus] = useState("pending");
+
+	useEffect(() => {
+		const abortController = new AbortController();
+
+		addEventListener(
+			"backup:completed",
+			(event) => {
+				if (event.scheduleId === scheduleId) {
+					setStatus(event.status);
+				}
+			},
+			{ signal: abortController.signal },
+		);
+
+		return () => abortController.abort();
+	}, [addEventListener, scheduleId]);
+
+	return <div>{status}</div>;
+};
+
+describe("useServerEvents", () => {
+	beforeEach(() => {
+		MockEventSource.reset();
+		globalThis.EventSource = MockEventSource as unknown as typeof EventSource;
+		console.info = mock(() => {});
+		console.error = mock(() => {});
+	});
+
+	afterEach(() => {
+		cleanup();
+		globalThis.EventSource = originalEventSource;
+		console.info = originalConsoleInfo;
+		console.error = originalConsoleError;
+		MockEventSource.reset();
+	});
+
+	test("shares one EventSource across consumers and invalidates queries once on backup completion", async () => {
+		const queryClient = createTestQueryClient();
+		const invalidateQueries = mock(async () => undefined);
+		const refetchQueries = mock(async () => undefined);
+		queryClient.invalidateQueries = invalidateQueries as typeof queryClient.invalidateQueries;
+		queryClient.refetchQueries = refetchQueries as typeof queryClient.refetchQueries;
+
+		render(
+			<QueryClientProvider client={queryClient}>
+				<ConnectionConsumer />
+				<BackupCompletedListener scheduleId="0b9c940b" />
+			</QueryClientProvider>,
+		);
+
+		expect(MockEventSource.instances).toHaveLength(1);
+
+		MockEventSource.instances[0]?.emit("backup:completed", {
+			organizationId: "default-org",
+			scheduleId: "0b9c940b",
+			volumeName: "synology",
+			repositoryName: "swiss-backup",
+			status: "success",
+		});
+
+		expect(await screen.findByText("success")).toBeTruthy();
+		expect(invalidateQueries).toHaveBeenCalledTimes(1);
+		expect(refetchQueries).not.toHaveBeenCalled();
+
+		cleanup();
+
+		expect(MockEventSource.instances[0]?.close).toHaveBeenCalledTimes(1);
+	});
+});
diff --git a/app/client/hooks/use-server-events.ts b/app/client/hooks/use-server-events.ts
index 88c789a7..62cb4ac8 100644
--- a/app/client/hooks/use-server-events.ts
+++ b/app/client/hooks/use-server-events.ts
@@ -1,4 +1,4 @@
-import { useQueryClient } from "@tanstack/react-query";
+import { type QueryClient, useQueryClient } from "@tanstack/react-query";
 import { useCallback, useEffect, useRef } from "react";
 import { serverEventNames, type ServerEventPayloadMap } from "~/schemas/server-events";
 
@@ -16,6 +16,13 @@ type EventHandlerMap = {
 	[K in ServerEventType]?: EventHandlerSet<K>;
 };
 
+type SharedServerEventsState = {
+	eventSource: EventSource | null;
+	handlers: EventHandlerMap;
+	queryClient: QueryClient | null;
+	subscribers: number;
+};
+
 const invalidatingEvents = new Set<ServerEventType>([
 	"backup:completed",
 	"volume:updated",
@@ -31,103 +38,151 @@ export type RestoreProgressEvent = ServerEventsPayloadMap["restore:progress"];
 export type RestoreCompletedEvent = ServerEventsPayloadMap["restore:completed"];
 export type BackupProgressEvent = ServerEventsPayloadMap["backup:progress"];
 
+const sharedState: SharedServerEventsState = {
+	eventSource: null,
+	handlers: {},
+	queryClient: null,
+	subscribers: 0,
+};
+
 const parseEventData = <T extends ServerEventType>(event: Event): ServerEventsPayloadMap[T] =>
 	JSON.parse((event as MessageEvent<string>).data) as ServerEventsPayloadMap[T];
 
+const isAbortError = (error: unknown): error is Error => error instanceof Error && error.name === "AbortError";
+
+const emit = <T extends ServerEventType>(eventName: T, data: ServerEventsPayloadMap[T]) => {
+	const handlers = sharedState.handlers[eventName] as EventHandlerSet<T> | undefined;
+	handlers?.forEach((handler) => {
+		handler(data);
+	});
+};
+
+const refreshQueriesForEvent = (eventName: ServerEventType) => {
+	if (!invalidatingEvents.has(eventName) || !sharedState.queryClient) {
+		return;
+	}
+
+	void sharedState.queryClient.invalidateQueries().catch((error) => {
+		if (!isAbortError(error)) {
+			console.error(`[SSE] Failed to refresh queries after ${eventName}:`, error);
+		}
+	});
+};
+
+const connectEventSource = (queryClient: QueryClient) => {
+	sharedState.queryClient = queryClient;
+	if (sharedState.eventSource) {
+		return;
+	}
+
+	const eventSource = new EventSource("/api/v1/events");
+	sharedState.eventSource = eventSource;
+
+	eventSource.addEventListener("connected", (event) => {
+		const data = parseEventData<"connected">(event);
+		console.info("[SSE] Connected to server events");
+		emit("connected", data);
+	});
+
+	eventSource.addEventListener("heartbeat", (event) => {
+		emit("heartbeat", parseEventData<"heartbeat">(event));
+	});
+
+	for (const eventName of serverEventNames) {
+		eventSource.addEventListener(eventName, (event) => {
+			const data = parseEventData<typeof eventName>(event);
+			console.info(`[SSE] ${eventName}:`, data);
+
+			refreshQueriesForEvent(eventName);
+
+			if (eventName === "volume:status_changed") {
+				const statusData = data as ServerEventsPayloadMap["volume:status_changed"];
+				emit("volume:status_changed", statusData);
+				emit("volume:updated", statusData);
+				return;
+			}
+
+			emit(eventName, data);
+		});
+	}
+
+	eventSource.onerror = (error) => {
+		console.error("[SSE] Connection error:", error);
+	};
+};
+
+const disconnectEventSource = () => {
+	if (!sharedState.eventSource) {
+		return;
+	}
+
+	console.info("[SSE] Disconnecting from server events");
+	sharedState.eventSource.close();
+	sharedState.eventSource = null;
+	sharedState.queryClient = null;
+	sharedState.handlers = {};
+};
+
+const addSharedEventListener = <T extends ServerEventType>(
+	eventName: T,
+	handler: EventHandler<T>,
+	options?: { signal?: AbortSignal },
+) => {
+	if (options?.signal?.aborted) {
+		return () => {};
+	}
+
+	const existingHandlers = sharedState.handlers[eventName] as EventHandlerSet<T> | undefined;
+	const eventHandlers = existingHandlers ?? new Set<EventHandler<T>>();
+	eventHandlers.add(handler);
+	sharedState.handlers[eventName] = eventHandlers as EventHandlerMap[T];
+
+	const unsubscribe = () => {
+		const handlers = sharedState.handlers[eventName] as EventHandlerSet<T> | undefined;
+		handlers?.delete(handler);
+		if (handlers && handlers.size === 0) {
+			delete sharedState.handlers[eventName];
+		}
+		if (options?.signal) {
+			options.signal.removeEventListener("abort", unsubscribe);
+		}
+	};
+
+	if (options?.signal) {
+		options.signal.addEventListener("abort", unsubscribe, { once: true });
+	}
+
+	return unsubscribe;
+};
+
 /**
  * Hook to listen to Server-Sent Events (SSE) from the backend
  * Automatically handles cache invalidation for backup and volume events
  */
 export function useServerEvents() {
 	const queryClient = useQueryClient();
-	const eventSourceRef = useRef<EventSource | null>(null);
-	const handlersRef = useRef<EventHandlerMap>({});
-	const emit = useCallback(<T extends ServerEventType>(eventName: T, data: ServerEventsPayloadMap[T]) => {
-		const handlers = handlersRef.current[eventName] as EventHandlerSet<T> | undefined;
-		handlers?.forEach((handler) => {
-			handler(data);
-		});
-	}, []);
+	const addEventListener = useCallback(addSharedEventListener, []);
+	const hasMountedRef = useRef(false);
 
 	useEffect(() => {
-		const eventSource = new EventSource("/api/v1/events");
-		eventSourceRef.current = eventSource;
-
-		eventSource.addEventListener("connected", (event) => {
-			const data = parseEventData<"connected">(event);
-			console.info("[SSE] Connected to server events");
-			emit("connected", data);
-		});
-
-		eventSource.addEventListener("heartbeat", (event) => {
-			emit("heartbeat", parseEventData<"heartbeat">(event));
-		});
-
-		for (const eventName of serverEventNames) {
-			eventSource.addEventListener(eventName, (event) => {
-				const data = parseEventData<typeof eventName>(event);
-				console.info(`[SSE] ${eventName}:`, data);
-
-				if (invalidatingEvents.has(eventName)) {
-					void queryClient.invalidateQueries();
-				}
-
-				if (eventName === "backup:completed") {
-					void queryClient.refetchQueries();
-				}
-
-				if (eventName === "volume:status_changed") {
-					const statusData = data as ServerEventsPayloadMap["volume:status_changed"];
-					emit("volume:status_changed", statusData);
-					emit("volume:updated", statusData);
-					return;
-				}
-
-				emit(eventName, data);
-			});
+		connectEventSource(queryClient);
+		if (!hasMountedRef.current) {
+			sharedState.subscribers += 1;
+			hasMountedRef.current = true;
 		}
 
-		eventSource.onerror = (error) => {
-			console.error("[SSE] Connection error:", error);
-		};
-
 		return () => {
-			console.info("[SSE] Disconnecting from server events");
-			eventSource.close();
-			eventSourceRef.current = null;
+			if (!hasMountedRef.current) {
+				return;
+			}
+
+			hasMountedRef.current = false;
+			sharedState.subscribers = Math.max(0, sharedState.subscribers - 1);
+			if (sharedState.subscribers === 0) {
+				disconnectEventSource();
+			}
 		};
-	}, [emit, queryClient]);
-
-	const addEventListener = useCallback(
-		<T extends ServerEventType>(eventName: T, handler: EventHandler<T>, options?: { signal?: AbortSignal }) => {
-			if (options?.signal?.aborted) {
-				return () => {};
-			}
-
-			const existingHandlers = handlersRef.current[eventName] as EventHandlerSet<T> | undefined;
-			const eventHandlers = existingHandlers ?? new Set<EventHandler<T>>();
-			eventHandlers.add(handler);
-			handlersRef.current[eventName] = eventHandlers as EventHandlerMap[T];
-
-			const unsubscribe = () => {
-				const handlers = handlersRef.current[eventName] as EventHandlerSet<T> | undefined;
-				handlers?.delete(handler);
-				if (handlers && handlers.size === 0) {
-					delete handlersRef.current[eventName];
-				}
-				if (options?.signal) {
-					options.signal.removeEventListener("abort", unsubscribe);
-				}
-			};
-
-			if (options?.signal) {
-				options.signal.addEventListener("abort", unsubscribe, { once: true });
-			}
-
-			return unsubscribe;
-		},
-		[],
-	);
+	}, [queryClient]);
 
 	return { addEventListener };
 }