/** * Badge WebSocket hook with multi-project isolation. * * ## Scoped Key Format * All badge state is keyed by `${projectId}:${taskId}` where `projectId` defaults to * `"default"` when not in a project-scoped context. This prevents overlapping task IDs * across projects from sharing badge state (e.g., `proj-A/FN-063` vs `proj-B/FN-063`). * * ## Project-Switch Reset Behavior (FN-1657) * When `projectId` changes, the store: * 1. Increments `projectContextVersionRef` to invalidate any pending callbacks from the old context * 2. Clears badge snapshots immediately to prevent stale data visibility * 3. Re-subscribes all active tasks over the new project-scoped WebSocket * 4. Rejects incoming messages from old context via context version guard * * ## Backward Compatibility * - WebSocket URL: `/api/ws` (default) or `/api/ws?projectId=` when projectId is present * - Subscribe payload: `{ type: "subscribe", taskId, projectId? }` * - Unsubscribe payload: `{ type: "unsubscribe", taskId, projectId? }` * - Incoming badge message: `{ type: "badge:updated", taskId, timestamp, prInfo?, issueInfo?, projectId? }` * * If incoming messages omit `projectId`, they are treated as current-connection scope only. * If `projectId` is present and does not match active context, the message is ignored. */ import { useCallback, useEffect, useRef, useSyncExternalStore } from "react"; import type { IssueInfo, PrInfo } from "@fusion/core"; import { appendTokenQuery } from "../auth"; interface BadgeUpdatedMessage { type: "badge:updated"; taskId: string; prInfo?: PrInfo | null; issueInfo?: IssueInfo | null; timestamp: string; /** Present when message is from a project-scoped channel (FN-1745+ server) */ projectId?: string; } interface BadgeSnapshot { prInfo?: PrInfo | null; issueInfo?: IssueInfo | null; timestamp: string; } interface StoreSnapshot { badgeUpdates: Map; isConnected: boolean; } /** * Scoped key helper for multi-project isolation. * Keys are formatted as `${projectId}:${taskId}` to prevent * overlapping task IDs across projects from sharing badge state. */ function toScopedKey(projectId: string | null, taskId: string): string { return `${projectId ?? "default"}:${taskId}`; } class BadgeWebSocketStore { private ws: WebSocket | null = null; private listeners = new Set<() => void>(); private badgeUpdates = new Map(); private retainedBadgeUpdates = new Map(); private subscriptionsByTask = new Map>(); private reconnectTimeout: ReturnType | null = null; private reconnectDelayMs = 1_000; private shouldReconnect = false; private isConnected = false; private projectId: string | null = null; private snapshot: StoreSnapshot = { badgeUpdates: new Map(), isConnected: false, }; // FN-1657 context-version guard: prevents stale callbacks from old project contexts // from corrupting badge state after project switches. private previousProjectIdRef: string | null = null; private projectContextVersionRef = 0; private contextVersionAtStart = 0; private static readonly RETAINED_BADGE_CAP = 200; subscribe(listener: () => void): () => void { this.listeners.add(listener); return () => { this.listeners.delete(listener); }; } getSnapshot(): StoreSnapshot { return this.snapshot; } setProjectId(projectId: string | null): void { if (this.projectId === projectId) return; // FN-1657: Increment context version to invalidate any pending callbacks from old context this.previousProjectIdRef = this.projectId; this.projectContextVersionRef++; this.contextVersionAtStart = this.projectContextVersionRef; const hadSubscriptions = this.subscriptionsByTask.size > 0; // Collect all (hookId, taskId) pairs that were subscribed const previousSubscriptions: Array<{ hookId: string; taskId: string }> = []; for (const [scopedTaskId, subscribers] of this.subscriptionsByTask) { // Extract taskId from scoped key (format: "projectId:taskId") const taskId = scopedTaskId.split(":").slice(1).join(":"); for (const hookId of subscribers) { previousSubscriptions.push({ hookId, taskId }); } } this.projectId = projectId; this.reset(); // Re-subscribe to all previous subscriptions after project change // This ensures badge subscriptions survive project switches if (hadSubscriptions) { for (const { hookId, taskId } of previousSubscriptions) { const scopedKey = toScopedKey(this.projectId, taskId); this.subscriptionsByTask.set(scopedKey, new Set([hookId])); } // Set shouldReconnect BEFORE calling connect() to ensure the connection is made this.shouldReconnect = this.subscriptionsByTask.size > 0; // Note: we restore subscriptions BEFORE calling connect() so that // onopen will send the subscribe messages over the new socket this.connect(); } } subscribeTask(hookId: string, taskId: string): void { const scopedKey = toScopedKey(this.projectId, taskId); const retainedSnapshot = this.retainedBadgeUpdates.get(scopedKey); if (retainedSnapshot && !this.badgeUpdates.has(scopedKey)) { this.badgeUpdates.set(scopedKey, retainedSnapshot); this.retainedBadgeUpdates.delete(scopedKey); this.emit(); } const subscribers = this.subscriptionsByTask.get(scopedKey) ?? new Set(); const isNewSubscription = !subscribers.has(hookId); subscribers.add(hookId); this.subscriptionsByTask.set(scopedKey, subscribers); this.shouldReconnect = this.subscriptionsByTask.size > 0; this.connect(); // Only send subscribe message if this is a genuinely new subscription // (not a re-subscription from project switch or unmount/remount) if (isNewSubscription) { this.send({ type: "subscribe", taskId, projectId: this.projectId }); } } unsubscribeTask(hookId: string, taskId: string): void { const scopedKey = toScopedKey(this.projectId, taskId); const subscribers = this.subscriptionsByTask.get(scopedKey); if (!subscribers) return; subscribers.delete(hookId); if (subscribers.size === 0) { this.subscriptionsByTask.delete(scopedKey); const snapshot = this.badgeUpdates.get(scopedKey); if (snapshot) { // Keep a bounded cache of last-known snapshots so cards can restore badge state // immediately after transient viewport unsubscriptions. Cap prevents unbounded growth. this.retainedBadgeUpdates.delete(scopedKey); this.retainedBadgeUpdates.set(scopedKey, snapshot); while (this.retainedBadgeUpdates.size > BadgeWebSocketStore.RETAINED_BADGE_CAP) { const oldestKey = this.retainedBadgeUpdates.keys().next().value; if (!oldestKey) break; this.retainedBadgeUpdates.delete(oldestKey); } } this.badgeUpdates.delete(scopedKey); this.send({ type: "unsubscribe", taskId, projectId: this.projectId }); this.emit(); } this.shouldReconnect = this.subscriptionsByTask.size > 0; if (!this.shouldReconnect) { this.disconnect(); } } cleanupHook(hookId: string): void { for (const scopedTaskId of [...this.subscriptionsByTask.keys()]) { // Extract taskId from scoped key const taskId = scopedTaskId.split(":").slice(1).join(":"); this.unsubscribeTask(hookId, taskId); } } reset(): void { this.disconnect(); this.badgeUpdates.clear(); this.retainedBadgeUpdates.clear(); this.subscriptionsByTask.clear(); this.shouldReconnect = false; this.emit(); } private connect(): void { if (!this.shouldReconnect || typeof window === "undefined") { return; } if (this.ws && (this.ws.readyState === WebSocket.OPEN || this.ws.readyState === WebSocket.CONNECTING)) { return; } // Capture context version at socket start to detect stale callbacks later // This is captured in the closure so each socket remembers its creation version const contextVersionAtConnect = this.projectContextVersionRef; const protocol = window.location.protocol === "https:" ? "wss:" : "ws:"; let url = `${protocol}//${window.location.host}/api/ws`; if (this.projectId) { url += `?projectId=${encodeURIComponent(this.projectId)}`; } // Bearer token must be on the URL — WebSocket construction doesn't // support custom headers. No-op when auth is disabled. const ws = new WebSocket(appendTokenQuery(url)); this.ws = ws; ws.onopen = () => { // FN-1657: Reject callback if context changed since socket creation if (this.projectContextVersionRef !== contextVersionAtConnect) { ws.close(); return; } this.isConnected = true; this.reconnectDelayMs = 1_000; this.emit(); // Extract raw taskIds from scoped keys (format: "projectId:taskId") const uniqueTaskIds = new Set(); for (const scopedKey of this.subscriptionsByTask.keys()) { const taskId = scopedKey.split(":").slice(1).join(":"); uniqueTaskIds.add(taskId); } for (const taskId of uniqueTaskIds) { this.send({ type: "subscribe", taskId, projectId: this.projectId }); } }; ws.onmessage = (event) => { // FN-1657: Reject callbacks from stale context using captured version if (this.projectContextVersionRef !== contextVersionAtConnect) { return; } try { const message = JSON.parse(event.data as string) as BadgeUpdatedMessage; if (message.type !== "badge:updated") { return; } // Backward compatibility: if incoming message has projectId and it doesn't match // our active context, ignore the message (FN-1745+ server behavior) if (message.projectId !== undefined && message.projectId !== this.projectId) { return; } const scopedKey = toScopedKey(this.projectId, message.taskId); const previous = this.badgeUpdates.get(scopedKey); this.badgeUpdates.set(scopedKey, { prInfo: hasMessageField(message, "prInfo") ? message.prInfo ?? null : previous?.prInfo, issueInfo: hasMessageField(message, "issueInfo") ? message.issueInfo ?? null : previous?.issueInfo, timestamp: message.timestamp, }); this.emit(); } catch { // Ignore malformed messages. } }; ws.onclose = () => { // FN-1657: Reject callbacks from stale context using captured version if (this.projectContextVersionRef !== contextVersionAtConnect) { return; } this.ws = null; if (this.isConnected) { this.isConnected = false; this.emit(); } if (!this.shouldReconnect) { return; } this.scheduleReconnect(contextVersionAtConnect); }; ws.onerror = () => { // Closed socket is handled by onclose. }; } private disconnect(): void { if (this.reconnectTimeout) { clearTimeout(this.reconnectTimeout); this.reconnectTimeout = null; } if (this.ws) { const socket = this.ws; this.ws = null; socket.close(); } if (this.isConnected) { this.isConnected = false; this.emit(); } } private scheduleReconnect(contextVersionAtSchedule?: number): void { if (this.reconnectTimeout) { return; } // Capture version at schedule time if not provided const versionAtSchedule = contextVersionAtSchedule ?? this.projectContextVersionRef; const delay = Math.min(this.reconnectDelayMs, 5_000); this.reconnectTimeout = setTimeout(() => { this.reconnectTimeout = null; // FN-1657: Skip reconnect if context changed since scheduling if (this.projectContextVersionRef !== versionAtSchedule) { return; } if (!this.shouldReconnect) { return; } this.connect(); }, delay); this.reconnectDelayMs = Math.min(this.reconnectDelayMs * 2, 5_000); } private send(message: { type: "subscribe" | "unsubscribe"; taskId: string; projectId?: string | null }): void { if (!this.ws || this.ws.readyState !== WebSocket.OPEN) { return; } this.ws.send(JSON.stringify(message)); } private emit(): void { this.snapshot = { badgeUpdates: new Map(this.badgeUpdates), isConnected: this.isConnected, }; for (const listener of this.listeners) { listener(); } } } const badgeWebSocketStore = new BadgeWebSocketStore(); let nextHookId = 0; function hasMessageField(message: BadgeUpdatedMessage, field: "prInfo" | "issueInfo"): boolean { return Object.prototype.hasOwnProperty.call(message, field); } export function useBadgeWebSocket(projectId?: string): { badgeUpdates: Map; isConnected: boolean; subscribeToBadge: (taskId: string) => void; unsubscribeFromBadge: (taskId: string) => void; } { const hookIdRef = useRef(null); if (hookIdRef.current === null) { hookIdRef.current = `badge-hook-${nextHookId++}`; } const snapshot = useSyncExternalStore( (listener) => badgeWebSocketStore.subscribe(listener), () => badgeWebSocketStore.getSnapshot(), () => badgeWebSocketStore.getSnapshot(), ); const subscribeToBadge = useCallback((taskId: string) => { badgeWebSocketStore.subscribeTask(hookIdRef.current!, taskId); }, []); const unsubscribeFromBadge = useCallback((taskId: string) => { badgeWebSocketStore.unsubscribeTask(hookIdRef.current!, taskId); }, []); // Update project context when projectId changes useEffect(() => { badgeWebSocketStore.setProjectId(projectId ?? null); }, [projectId]); useEffect(() => { return () => { badgeWebSocketStore.cleanupHook(hookIdRef.current!); }; }, []); return { badgeUpdates: snapshot.badgeUpdates, isConnected: snapshot.isConnected, subscribeToBadge, unsubscribeFromBadge, }; } export function __resetBadgeWebSocketStoreForTests(): void { badgeWebSocketStore.reset(); nextHookId = 0; }