1import { RpcSession, type RpcCompatible, type RpcTransport } from "capnweb";
2
3/**
4 * A platform-agnostic interface for a server-side WebSocket.
5 */
6export interface ServerWebSocketLike {
7 send: (message: string) => any | Promise<any>;
8 close: (code?: number, reason?: string) => void;
9}
10
11/**
12 * Represents a websocket RPC session with capnweb.
13 * Use the `dispatch` object to handle messages that are received from the client.
14 */
15export type ServerRpcSession<T extends RpcCompatible<T>> = ReturnType<
16 typeof makeServerRpcSession<T>
17>;
18
19/**
20 * Constructs a ServerRpcSession using capnweb.
21 * @param ws - The WebSocket to use for the session.
22 * @param main - The main object to use for the session.
23 * @returns A ServerRpcSession.
24 */
25export function makeServerRpcSession<T extends RpcCompatible<T>>(
26 ws: ServerWebSocketLike,
27 main: T,
28) {
29 const { transport, dispatch } = makeWebSocketRpcTransport(ws);
30 const session = new RpcSession(transport, main);
31 return { session, dispatch };
32}
33
34function makeWebSocketRpcTransport(ws: ServerWebSocketLike) {
35 let receiveQueue: Array<string> = [];
36 let receiveResolver: ((value: string) => void) | undefined;
37 let receiveRejecter: ((reason: unknown) => void) | undefined;
38 let error: unknown | undefined;
39 return {
40 transport: {
41 send: async (message: string) => await ws.send(message),
42 receive: async () => {
43 const next = receiveQueue.shift();
44 if (next) {
45 return next;
46 } else if (error) {
47 throw error;
48 }
49 return new Promise<string>((resolve, reject) => {
50 receiveResolver = resolve;
51 receiveRejecter = reject;
52 });
53 },
54 abort: (reason: unknown) => {
55 const message =
56 reason instanceof Error ? reason.message : String(reason);
57 ws.close(3000, message);
58 error ??= reason;
59 },
60 } satisfies RpcTransport,
61 dispatch: {
62 message: (data: string | Buffer<ArrayBuffer>) => {
63 if (error) {
64 return;
65 }
66 data = typeof data === "string" ? data : data.toString();
67 if (receiveResolver) {
68 receiveResolver(data);
69 receiveResolver = undefined;
70 receiveRejecter = undefined;
71 } else {
72 receiveQueue.push(data);
73 }
74 },
75 close: (code: number, reason: string) => {
76 if (!error) {
77 error = new Error(`WebSocket closed with code ${code}: ${reason}`);
78 if (receiveRejecter) {
79 receiveRejecter(error);
80 receiveRejecter = undefined;
81 receiveResolver = undefined;
82 }
83 }
84 },
85 },
86 };
87}