working improved sockets

api/src/db/schema.ts

@@ -52,7 +52,10 @@ export const partyMember = pgTable(
 		joinedAt: timestamp().defaultNow().notNull(),
 		lastSeen: timestamp().defaultNow().notNull(),
 	},
-	(partyMember) => [uniqueIndex().on(partyMember.partyId, partyMember.userId)],
+	(partyMember) => [
+		uniqueIndex().on(partyMember.partyId, partyMember.userId),
+		index().on(partyMember.userId, partyMember.joinedAt),
+	],
 );
 
 export const platform = pgEnum("enum_platform", ["spotify", "apple"]);

api/src/index.ts

@@ -7,7 +7,7 @@ import "./workflows/sync";
 import "./workflows/party-analysis";
 import "./dbos.ts";
 import { partyApp } from "./routes/party";
-import { partySocketApp } from "./routes/party-socket";
+import { partySocketApp, pubsub } from "./routes/party-socket";
 import { statsApp } from "./routes/stats.ts";
 
 const app = new Elysia()
@@ -23,6 +23,8 @@ const app = new Elysia()
 	)
 	.listen(4000);
 
+pubsub.setServer(app.server);
+
 export type App = typeof app;
 
 await DBOS.launch({

api/src/party-data.ts

@@ -1,6 +1,7 @@
 import { eq } from "drizzle-orm";
 import { db } from "./db";
 import { party, partyMember } from "./db/schema";
+import type { PartySnapshot } from "./party-types";
 
 type DbClient = typeof db;
 type DbTransaction = Parameters<typeof db.transaction>[0] extends (
@@ -29,11 +30,16 @@ export async function getMemberRecord(dbClient: DbLike, userId: string) {
 			where: {
 				userId,
 			},
+			orderBy: {
+				joinedAt: "desc",
+			},
 		})) ?? null
 	);
 }
 
-export async function getPartyStatus(partyId: string) {
+export async function getPartyStatus(
+	partyId: string,
+): Promise<PartySnapshot | null> {
 	const partyRecord = await db.query.party.findFirst({
 		where: {
 			id: partyId,

api/src/party-sockets.ts

@@ -1,159 +0,0 @@
-type PartySocketEvent = {
-	type: string;
-	[key: string]: unknown;
-};
-
-type WebSocketLike = {
-	send: (data: string) => void;
-	close?: (code?: number, reason?: string) => void;
-};
-
-const partySockets = new Map<string, Map<string, Set<WebSocketLike>>>();
-const userSockets = new Map<string, Set<WebSocketLike>>();
-
-function getPartyUserSockets(partyId: string, userId: string) {
-	const partyMap = partySockets.get(partyId);
-	if (!partyMap) return null;
-	return partyMap.get(userId) ?? null;
-}
-
-export function registerPartySocket(
-	partyId: string,
-	userId: string,
-	ws: WebSocketLike,
-) {
-	let partyMap = partySockets.get(partyId);
-	if (!partyMap) {
-		partyMap = new Map();
-		partySockets.set(partyId, partyMap);
-	}
-
-	let userSockets = partyMap.get(userId);
-	if (!userSockets) {
-		userSockets = new Set();
-		partyMap.set(userId, userSockets);
-	}
-
-	userSockets.add(ws);
-}
-
-export function unregisterPartySocket(
-	partyId: string,
-	userId: string,
-	ws: WebSocketLike,
-) {
-	const partyMap = partySockets.get(partyId);
-	if (!partyMap) return;
-
-	const userSockets = partyMap.get(userId);
-	if (!userSockets) return;
-
-	userSockets.delete(ws);
-
-	if (userSockets.size === 0) {
-		partyMap.delete(userId);
-	}
-
-	if (partyMap.size === 0) {
-		partySockets.delete(partyId);
-	}
-}
-
-export function registerUserSocket(userId: string, ws: WebSocketLike) {
-	let sockets = userSockets.get(userId);
-	if (!sockets) {
-		sockets = new Set();
-		userSockets.set(userId, sockets);
-	}
-
-	sockets.add(ws);
-}
-
-export function unregisterUserSocket(userId: string, ws: WebSocketLike) {
-	const sockets = userSockets.get(userId);
-	if (!sockets) return;
-
-	sockets.delete(ws);
-
-	if (sockets.size === 0) {
-		userSockets.delete(userId);
-	}
-}
-
-export function unregisterUserSocketFromAllParties(
-	userId: string,
-	ws: WebSocketLike,
-) {
-	for (const [partyId, partyMap] of partySockets) {
-		const userSockets = partyMap.get(userId);
-		if (!userSockets) continue;
-		userSockets.delete(ws);
-		if (userSockets.size === 0) {
-			partyMap.delete(userId);
-		}
-		if (partyMap.size === 0) {
-			partySockets.delete(partyId);
-		}
-	}
-}
-
-export function broadcastPartyEvent(partyId: string, event: PartySocketEvent) {
-	const partyMap = partySockets.get(partyId);
-	if (!partyMap) return;
-
-	const payload = JSON.stringify(event);
-	for (const userSockets of partyMap.values()) {
-		for (const ws of userSockets) {
-			ws.send(payload);
-		}
-	}
-}
-
-export function sendPartyEventToUser(
-	partyId: string,
-	userId: string,
-	event: PartySocketEvent,
-) {
-	const userSockets = getPartyUserSockets(partyId, userId);
-	if (!userSockets) return;
-
-	const payload = JSON.stringify(event);
-	for (const ws of userSockets) {
-		ws.send(payload);
-	}
-}
-
-export function sendDirectEventToUser(userId: string, event: PartySocketEvent) {
-	const sockets = userSockets.get(userId);
-	if (!sockets) return;
-
-	const payload = JSON.stringify(event);
-	for (const ws of sockets) {
-		ws.send(payload);
-	}
-}
-
-export function reassignUserSocketsToParty(
-	userId: string,
-	partyId: string | null,
-) {
-	for (const [existingPartyId, partyMap] of partySockets) {
-		if (!partyMap.has(userId)) continue;
-		partyMap.delete(userId);
-		if (partyMap.size === 0) {
-			partySockets.delete(existingPartyId);
-		}
-	}
-
-	if (!partyId) return;
-	const sockets = userSockets.get(userId);
-	if (!sockets) return;
-
-	let partyMap = partySockets.get(partyId);
-	if (!partyMap) {
-		partyMap = new Map();
-		partySockets.set(partyId, partyMap);
-	}
-
-	partyMap.set(userId, new Set(sockets));
-}

api/src/party-types.ts

@@ -0,0 +1,32 @@
+import type { InferSelectModel } from "drizzle-orm";
+import type { party, partyMember, user } from "./db/schema";
+
+export type Party = InferSelectModel<typeof party>;
+export type PartyMember = InferSelectModel<typeof partyMember>;
+export type User = InferSelectModel<typeof user>;
+
+export type PartyMemberWithUser = PartyMember & { user: User | null };
+
+export const PARTY_STATUS = ["created", "started", "ended"] as const;
+export type PartyStatus = (typeof PARTY_STATUS)[number];
+
+export type PartySnapshot = {
+	party: Party;
+	members: PartyMemberWithUser[];
+};
+
+export type PartyState = {
+	party: Party | null;
+	members: PartyMemberWithUser[];
+};
+
+export type PartySocketOutgoing =
+	| { type: "ping" }
+	| { type: "member_payload"; payload: unknown };
+
+export type PartySocketEvent =
+	| { type: "snapshot"; party: Party | null; members: PartyMemberWithUser[] }
+	| { type: "party_status"; party: Party; members: PartyMemberWithUser[] }
+	| { type: "member_payload"; fromUserId: string; payload: unknown }
+	| { type: "error"; message: string }
+	| { type: "pong" };

api/src/routes/party-socket.ts

@@ -1,137 +1,133 @@
-import Elysia, { t } from "elysia";
-import { auth, betterAuthElysia } from "../auth";
+import { Elysia } from "elysia";
+
+import { betterAuthElysia } from "../auth";
+
 import { db } from "../db";
 import { getMemberRecord, getPartyStatus } from "../party-data";
-import {
-	registerPartySocket,
-	registerUserSocket,
-	sendPartyEventToUser,
-	unregisterPartySocket,
-	unregisterUserSocket,
-	unregisterUserSocketFromAllParties,
-} from "../party-sockets";
 
-type PartySocketMessage =
-	| {
-			type: "member_payload";
-			payload: unknown;
-	  }
-	| {
-			type: "ping";
-	  };
+function userTopic(userId: string) {
+	return `user:${userId}`;
+}
 
-const MAX_MEMBER_PAYLOAD_SIZE = 8_000;
+function partyTopic(partyId: string) {
+	return `party:${partyId}`;
+}
 
-type PartyWsData = {
-	user?: { id: string };
-	partyId?: string | null;
+const socketPartyId = new WeakMap<object, string>();
+
+export const pubsub = {
+	_server: null as ReturnType<typeof Bun.serve> | null,
+	setServer(server: ReturnType<typeof Bun.serve> | null) {
+		this._server = server;
+	},
+	publish(topic: string, data: string) {
+		this._server?.publish(topic, data);
+	},
 };
 
-function getPayloadSize(payload: unknown) {
-	try {
-		return JSON.stringify(payload).length;
-	} catch {
-		return Infinity;
-	}
-}
+export const topic = {
+	user: userTopic,
+	party: partyTopic,
+};
 
 export const partySocketApp = new Elysia()
 	.use(betterAuthElysia)
 	.group("/party-socket", (app) =>
-		app.ws("/ws", {
-			beforeHandle: async ({ request, set }) => {
-				const session = await auth.api.getSession({
-					headers: request.headers,
-				});
-				if (!session) {
-					set.status = 401;
-					return;
-				}
-				return {
-					user: session.user,
-					session: session.session,
-				};
-			},
-			open: async (ws) => {
-				const data = ws.data as unknown as PartyWsData;
-				const user = data.user;
-				if (!user) return;
-				registerUserSocket(user.id, ws);
-				const membership = await getMemberRecord(db, user.id);
-				if (!membership) {
-					ws.send(
+		app
+			.get("/test", () => ({ ok: 1 }))
+			.ws("/ws", {
+				auth: true,
+				publishToSelf: true,
+				open: async (ws) => {
+					const user = ws.data.user;
+					if (!user) return;
+
+					ws.subscribe(userTopic(user.id));
+
+					const membership = await getMemberRecord(db, user.id);
+					if (!membership) {
+						ws.send(
+							JSON.stringify({
+								type: "snapshot",
+								party: null,
+								members: [],
+							}),
+						);
+						return;
+					}
+
+					socketPartyId.set(ws, membership.partyId);
+					ws.subscribe(partyTopic(membership.partyId));
+
+					const snapshot = await getPartyStatus(membership.partyId);
+					if (snapshot) {
+						ws.send(
+							JSON.stringify({
+								type: "snapshot",
+								party: snapshot.party,
+								members: snapshot.members,
+							}),
+						);
+					}
+				},
+				message: async (ws, message) => {
+					const data = ws.data;
+					const user = data.user;
+					if (!user) return;
+
+					if (typeof message !== "string") return;
+
+					let parsed: { type: string; payload?: unknown };
+					try {
+						parsed = JSON.parse(message);
+					} catch {
+						ws.send(JSON.stringify({ type: "error", message: "Invalid JSON" }));
+						return;
+					}
+
+					if (parsed.type === "ping") {
+						ws.send(JSON.stringify({ type: "pong" }));
+						return;
+					}
+
+					if (parsed.type !== "member_payload") return;
+
+					const MAX_MEMBER_PAYLOAD_SIZE = 8_000;
+					const payloadString = JSON.stringify(parsed.payload);
+					if (payloadString.length > MAX_MEMBER_PAYLOAD_SIZE) {
+						ws.send(
+							JSON.stringify({ type: "error", message: "Payload too large." }),
+						);
+						return;
+					}
+
+					const membership = await getMemberRecord(db, user.id);
+					if (!membership) return;
+
+					const currentParty = await db.query.party.findFirst({
+						where: { id: membership.partyId },
+					});
+					if (!currentParty) return;
+
+					ws.publish(
+						partyTopic(membership.partyId),
 						JSON.stringify({
-							type: "snapshot",
-							party: null,
-							members: [],
+							type: "member_payload",
+							fromUserId: user.id,
+							payload: parsed.payload,
 						}),
 					);
-					return;
-				}
+				},
+				close: async (ws) => {
+					const user = ws.data.user;
+					if (!user) return;
 
-				const snapshot = await getPartyStatus(membership.partyId);
-				data.partyId = membership.partyId;
-				registerPartySocket(membership.partyId, user.id, ws);
-				if (snapshot) {
-					ws.send(
-						JSON.stringify({
-							type: "snapshot",
-							party: snapshot.party,
-							members: snapshot.members,
-						}),
-					);
-				}
-			},
-			message: async (ws, message: PartySocketMessage) => {
-				const data = ws.data as unknown as PartyWsData;
-				const user = data.user;
-				if (!user) return;
-				if (message.type === "ping") {
-					ws.send(JSON.stringify({ type: "pong" }));
-					return;
-				}
+					ws.unsubscribe(userTopic(user.id));
 
-				if (message.type !== "member_payload") return;
-				const membership = await getMemberRecord(db, user.id);
-				if (!membership) return;
+					const partyId = socketPartyId.get(ws);
+					if (!partyId) return;
 
-				if (getPayloadSize(message.payload) > MAX_MEMBER_PAYLOAD_SIZE) {
-					ws.send(
-						JSON.stringify({
-							type: "error",
-							message: "Payload too large.",
-						}),
-					);
-					return;
-				}
-
-				const currentParty = await db.query.party.findFirst({
-					where: { id: membership.partyId },
-				});
-				if (!currentParty) return;
-
-				sendPartyEventToUser(membership.partyId, currentParty.hostId, {
-					type: "member_payload",
-					fromUserId: user.id,
-					payload: message.payload,
-				});
-			},
-			close: async (ws) => {
-				const data = ws.data as unknown as PartyWsData;
-				const user = data.user;
-				const { partyId } = data;
-				if (!user) return;
-				if (!partyId) {
-					unregisterUserSocketFromAllParties(user.id, ws);
-					unregisterUserSocket(user.id, ws);
-					return;
-				}
-				unregisterPartySocket(partyId, user.id, ws);
-				unregisterUserSocket(user.id, ws);
-			},
-			body: t.Union([
-				t.Object({ type: t.Literal("ping") }),
-				t.Object({ type: t.Literal("member_payload"), payload: t.Any() }),
-			]),
-		}),
+					ws.unsubscribe(partyTopic(partyId));
+				},
+			}),
 	);

api/src/routes/party.ts

@@ -10,29 +10,29 @@ import {
 	getPartyStatus,
 	leaveParty,
 } from "../party-data";
-import {
-	broadcastPartyEvent,
-	reassignUserSocketsToParty,
-	sendDirectEventToUser,
-} from "../party-sockets";
-
-const PARTY_STATUS = ["created", "started", "ended"] as const;
-
-type PartyStatus = (typeof PARTY_STATUS)[number];
-
-type PartySnapshot = NonNullable<Awaited<ReturnType<typeof getPartyStatus>>>;
+import type { PartySnapshot } from "../party-types";
+import { pubsub, topic } from "./party-socket";
 
 function broadcastSnapshot(partyId: string, snapshot: PartySnapshot | null) {
 	if (!snapshot) return;
-	broadcastPartyEvent(partyId, {
-		type: "party_status",
-		party: snapshot.party,
-		members: snapshot.members,
-	});
+	pubsub.publish(
+		topic.party(partyId),
+		JSON.stringify({
+			type: "party_status",
+			party: snapshot.party,
+			members: snapshot.members,
+		}),
+	);
 }
 
-function isValidStatus(status: string): status is PartyStatus {
-	return PARTY_STATUS.includes(status as PartyStatus);
+function broadcastToUser(userId: string, event: Record<string, unknown>) {
+	pubsub.publish(topic.user(userId), JSON.stringify(event));
+}
+
+function isValidStatus(
+	status: string,
+): status is import("../party-types").PartyStatus {
+	return ["created", "started", "ended"].includes(status);
 }
 
 export const partyApp = new Elysia()
@@ -127,31 +127,19 @@ export const partyApp = new Elysia()
 					if (!partyId) return { party: null, members: [] };
 					const status = await getPartyStatus(partyId);
 					if (leaveResult?.newHostId) {
-						broadcastPartyEvent(leaveResult.partyId, {
-							type: "host_changed",
-							hostId: leaveResult.newHostId,
-						});
+						broadcastSnapshot(leaveResult.partyId, status);
 					}
 					if (hostChanged) {
-						broadcastPartyEvent(partyId, {
-							type: "host_changed",
-							hostId: targetUserId,
-						});
+						broadcastSnapshot(partyId, status);
 					}
-					broadcastPartyEvent(partyId, {
-						type: "member_joined",
-						userId: user.id,
-					});
 					broadcastSnapshot(partyId, status);
-					reassignUserSocketsToParty(user.id, partyId);
-					reassignUserSocketsToParty(targetUserId, partyId);
 					if (status) {
-						sendDirectEventToUser(targetUserId, {
+						broadcastToUser(targetUserId, {
 							type: "party_status",
 							party: status.party,
 							members: status.members,
 						});
-						sendDirectEventToUser(user.id, {
+						broadcastToUser(user.id, {
 							type: "party_status",
 							party: status.party,
 							members: status.members,
@@ -174,18 +162,7 @@ export const partyApp = new Elysia()
 					});
 					if (!result) return { party: null, members: [] };
 					const status = await getPartyStatus(result.partyId);
-					broadcastPartyEvent(result.partyId, {
-						type: "member_left",
-						userId: user.id,
-					});
-					if (result.newHostId) {
-						broadcastPartyEvent(result.partyId, {
-							type: "host_changed",
-							hostId: result.newHostId,
-						});
-					}
 					broadcastSnapshot(result.partyId, status);
-					reassignUserSocketsToParty(user.id, null);
 					return status ?? { party: null, members: [] };
 				},
 				{ auth: true },
@@ -226,13 +203,7 @@ export const partyApp = new Elysia()
 						await cleanupPartyIfEmpty(tx, currentMembership.partyId);
 					});
 					const status = await getPartyStatus(currentMembership.partyId);
-					broadcastPartyEvent(currentMembership.partyId, {
-						type: "member_left",
-						userId: body.memberUserId,
-						kickedBy: user.id,
-					});
 					broadcastSnapshot(currentMembership.partyId, status);
-					reassignUserSocketsToParty(body.memberUserId, null);
 					return status ?? { party: null, members: [] };
 				},
 				{

web/src/components/user-info.tsx

@@ -1,29 +1,39 @@
 import { useRouteContext } from "@tanstack/react-router";
-import { Avatar, AvatarFallback, AvatarImage } from "./ui/avatar";
-import {
-  Item,
-  ItemContent,
-  ItemDescription,
-  ItemMedia,
-  ItemTitle,
-} from "./ui/item";
+import { useParty } from "#/hooks/use-party";
 import { useUser } from "#/hooks/user";
 import { initials } from "#/lib/utils";
+import { Avatar, AvatarFallback, AvatarImage } from "./ui/avatar";
+import {
+	Item,
+	ItemContent,
+	ItemDescription,
+	ItemMedia,
+	ItemTitle,
+} from "./ui/item";
 
 export function UserInfo() {
-  const { user } = useUser();
-  return (
-    <Item>
-      <ItemMedia>
-        <Avatar>
-          <AvatarImage src={user?.image || undefined} />
-          <AvatarFallback>{initials(user?.name || "")}</AvatarFallback>
-        </Avatar>
-      </ItemMedia>
-      <ItemContent>
-        <ItemTitle>{user?.name}</ItemTitle>
-        <ItemDescription>No party yet</ItemDescription>
-      </ItemContent>
-    </Item>
-  );
+	const { user } = useUser();
+	const { party, members, isConnecting, isReconnecting } = useParty();
+	return (
+		<Item>
+			<ItemMedia>
+				<Avatar>
+					<AvatarImage src={user?.image || undefined} />
+					<AvatarFallback>{initials(user?.name || "")}</AvatarFallback>
+				</Avatar>
+			</ItemMedia>
+			<ItemContent>
+				<ItemTitle>{user?.name}</ItemTitle>
+				<ItemDescription>
+					{isConnecting
+						? "Connecting..."
+						: isReconnecting
+							? "Reconnecting..."
+							: party
+								? `${members.length} in party`
+								: "No party yet"}
+				</ItemDescription>
+			</ItemContent>
+		</Item>
+	);
 }

web/src/hooks/use-party-socket.ts

@@ -0,0 +1,127 @@
+import { useCallback, useEffect, useMemo, useRef, useState } from "react";
+import type { PartySocketEvent } from "../../../api/src/party-types";
+
+type Handler = (event: PartySocketEvent) => void;
+
+const PING_INTERVAL_MS = 30_000;
+const RECONNECT_BASE_MS = 1_000;
+const RECONNECT_MAX_MS = 30_000;
+
+export function usePartySocket({
+	apiUrl,
+	onMessage,
+}: {
+	apiUrl: string | null;
+	onMessage: Handler | null;
+}) {
+	const [connectionState, setConnectionState] = useState<
+		"disconnected" | "connecting" | "connected" | "reconnecting"
+	>("disconnected");
+
+	const wsRef = useRef<WebSocket | null>(null);
+	const pingTimerRef = useRef<ReturnType<typeof setInterval> | null>(null);
+	const reconnectTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null);
+	const reconnectAttemptRef = useRef(0);
+	const handlerRef = useRef(onMessage);
+
+	useEffect(() => {
+		handlerRef.current = onMessage;
+	}, [onMessage]);
+
+	const setupWs = useCallback(
+		(ws: WebSocket) => {
+			ws.onopen = () => {
+				reconnectAttemptRef.current = 0;
+				setConnectionState("connected");
+				pingTimerRef.current = setInterval(() => {
+					if (ws.readyState === WebSocket.OPEN) {
+						ws.send(JSON.stringify({ type: "ping" }));
+					}
+				}, PING_INTERVAL_MS);
+			};
+
+			ws.onmessage = (event) => {
+				const parsed = JSON.parse(event.data) as PartySocketEvent;
+				handlerRef.current?.(parsed);
+			};
+
+			ws.onclose = () => {
+				if (pingTimerRef.current) {
+					clearInterval(pingTimerRef.current);
+					pingTimerRef.current = null;
+				}
+				wsRef.current = null;
+				setConnectionState("reconnecting");
+
+				const delay = Math.min(
+					RECONNECT_BASE_MS * 2 ** reconnectAttemptRef.current,
+					RECONNECT_MAX_MS,
+				);
+				reconnectAttemptRef.current++;
+				reconnectTimerRef.current = setTimeout(() => {
+					if (!apiUrl) return;
+					const protocol = apiUrl.startsWith("https") ? "wss" : "ws";
+					const newWs = new WebSocket(
+						`${protocol}://${apiUrl.replace(/https?:\/\//, "")}/api/party-socket/ws`,
+					);
+					wsRef.current = newWs;
+					setupWs(newWs);
+				}, delay);
+			};
+		},
+		[apiUrl],
+	);
+
+	useEffect(() => {
+		if (!apiUrl) {
+			if (wsRef.current) {
+				wsRef.current.close();
+				wsRef.current = null;
+			}
+			if (pingTimerRef.current) {
+				clearInterval(pingTimerRef.current);
+				pingTimerRef.current = null;
+			}
+			if (reconnectTimerRef.current) {
+				clearTimeout(reconnectTimerRef.current);
+				reconnectTimerRef.current = null;
+			}
+			setConnectionState("disconnected");
+			reconnectAttemptRef.current = 0;
+			return;
+		}
+
+		setConnectionState("connecting");
+		const protocol = apiUrl.startsWith("https") ? "wss" : "ws";
+		const ws = new WebSocket(
+			`${protocol}://${apiUrl.replace(/https?:\/\//, "")}/api/party-socket/ws`,
+		);
+		wsRef.current = ws;
+		setupWs(ws);
+
+		return () => {
+			ws.close();
+			wsRef.current = null;
+			if (pingTimerRef.current) {
+				clearInterval(pingTimerRef.current);
+				pingTimerRef.current = null;
+			}
+			if (reconnectTimerRef.current) {
+				clearTimeout(reconnectTimerRef.current);
+				reconnectTimerRef.current = null;
+			}
+		};
+	}, [apiUrl, setupWs]);
+
+	const state = useMemo(
+		() => ({
+			connectionState,
+			isConnected: connectionState === "connected",
+			isConnecting: connectionState === "connecting",
+			isReconnecting: connectionState === "reconnecting",
+		}),
+		[connectionState],
+	);
+
+	return state;
+}

web/src/hooks/use-party.ts

@@ -0,0 +1,58 @@
+import { useCallback, useMemo, useState } from "react";
+import type {
+	PartyMember,
+	PartySocketEvent,
+	PartyState,
+} from "../../../api/src/party-types";
+import { usePartySocket } from "./use-party-socket";
+import { useUser } from "./user";
+
+function reducePartyState(
+	state: PartyState,
+	event: PartySocketEvent,
+): PartyState {
+	switch (event.type) {
+		case "snapshot":
+		case "party_status":
+			return { party: event.party, members: event.members };
+		case "member_payload":
+		case "pong":
+		case "error":
+			return state;
+	}
+}
+
+function getApiUrl(): string | null {
+	if (typeof window === "undefined") return null;
+	const envUrl = import.meta.env.VITE_BETTER_AUTH_URL;
+	if (envUrl) return envUrl;
+	return `${window.location.protocol}//${window.location.host}`;
+}
+
+export function useParty() {
+	const { session } = useUser();
+	const [state, setState] = useState<PartyState>({
+		party: null,
+		members: [],
+	});
+
+	const handleMessage = useCallback((event: PartySocketEvent) => {
+		setState((prev: PartyState) => reducePartyState(prev, event));
+	}, []);
+
+	const apiUrl = useMemo(() => {
+		const url = getApiUrl();
+		if (!url) return null;
+		return url;
+	}, []);
+
+	const wsState = usePartySocket({
+		apiUrl,
+		onMessage: session ? handleMessage : null,
+	});
+
+	return {
+		...state,
+		...wsState,
+	};
+}

web/vite.config.ts

@@ -20,7 +20,17 @@ const config = defineConfig({
 	],
 	server: {
 		proxy: {
-			"/api": "http://localhost:4000",
+			"/api": {
+				target: "http://localhost:4000",
+				changeOrigin: true,
+				rewrite: (path) =>
+					path.replace(/^\/api/, "/api"),
+			},
+			"/api/party-socket/ws": {
+				target: "ws://localhost:4000",
+				ws: true,
+				rewriteWsOrigin: true,
+			},
 		},
 	},
 });