1import type * as cf from "@cloudflare/workers-types";
2import * as Effect from "effect/Effect";
3import * as Latch from "effect/Latch";
4import * as Queue from "effect/Queue";
5import * as Schema from "effect/Schema";
6import * as HttpBody from "effect/http/HttpBody";
7import { HttpServerRequest } from "effect/http/HttpServerRequest";
8import * as HttpServerResponse from "effect/http/HttpServerResponse";
9import * as RpcMessage from "effect/rpc/RpcMessage";
10import * as RpcSerialization from "effect/rpc/RpcSerialization";
11import * as RpcServer from "effect/rpc/RpcServer";
12import type { HttpEffect } from "../../Http.ts";
13import { RuntimeContext } from "../../RuntimeContext.ts";
14import { DurableObjectState } from "./DurableObjectState.ts";
15import { fromWebSocket, type WebSocket } from "./WebSocket.ts";
16
17const socketTag = "alchemy:rpc";
18const attachmentKey = "__alchemyRpcWebSocket";
19const restartReason = "Durable Object RPC activation reset";
20const Metadata = Schema.Struct({
21 version: Schema.Literal(1),
22 clientId: Schema.Number,
23 pending: Schema.Boolean,
24 serialization: Schema.String,
25});
26const isAttachment = Schema.is(Schema.Struct({ [attachmentKey]: Metadata }));
27
28interface Connection {
29 readonly id: number;
30 readonly socket: WebSocket;
31 readonly parser: RpcSerialization.Parser;
32 readonly pending: Set<string | number>;
33 readonly idle: Latch.Latch;
34}
35
36export interface Transport {
37 readonly protocol: RpcServer.Protocol["Service"];
38 readonly fetch: HttpEffect;
39 readonly webSocketMessage: (
40 socket: WebSocket,
41 message: string | ArrayBuffer,
42 ) => Effect.Effect<void>;
43 readonly webSocketClose: (socket: WebSocket) => Effect.Effect<void>;
44 readonly webSocketError: (
45 socket: WebSocket,
46 error: unknown,
47 ) => Effect.Effect<void>;
48}
49
50/**
51 * Internal, per-activation transport. Provide `protocol` to `RpcServer.layer`
52 * in an instance-owned scope, not the first fetch or WebSocket event's scope.
53 * This transport exclusively owns the object's WebSocket auto-response pair.
54 * Binary serializers use waking Ping/Pong messages instead of auto-response.
55 * Serializers must emit each response immediately; buffered output fails closed.
56 */
57export const make: Effect.Effect<
58 Transport,
59 never,
60 DurableObjectState | RuntimeContext | RpcSerialization.RpcSerialization
61> = Effect.gen(function* () {
62 const state = yield* DurableObjectState;
63 const runtime = yield* RuntimeContext;
64 const serialization = yield* RpcSerialization.RpcSerialization;
65 const native = globalThis as unknown as {
66 WebSocketPair: typeof cf.WebSocketPair;
67 WebSocketRequestResponsePair: typeof cf.WebSocketRequestResponsePair;
68 Response: typeof cf.Response;
69 };
70 const inRuntime = <A, E>(effect: Effect.Effect<A, E, RuntimeContext>) =>
71 Effect.provideService(effect, RuntimeContext, runtime);
72 const disconnects = yield* Queue.make<number>();
73 const ready = yield* Latch.make();
74 const bySocket = new Map<cf.WebSocket, Connection>();
75 const byId = new Map<number, Connection>();
76 const closed = new WeakSet<cf.WebSocket>();
77 let nextId = 0;
78 let restoring = true;
79 let stopped = false;
80 let receive: Parameters<RpcServer.Protocol["Service"]["run"]>[0] | undefined;
81
82 const heartbeat = yield* Effect.sync(() => {
83 const parser = serialization.makeUnsafe();
84 const ping = parser.encode(RpcMessage.constPing);
85 const pong = parser.encode(RpcMessage.constPong);
86 return typeof ping === "string" && typeof pong === "string"
87 ? new native.WebSocketRequestResponsePair(ping, pong)
88 : undefined;
89 });
90 const existing = yield* inRuntime(state.getWebSocketAutoResponse());
91 if (
92 existing !== null &&
93 (heartbeat === undefined ||
94 existing.request !== heartbeat.request ||
95 existing.response !== heartbeat.response)
96 ) {
97 return yield* Effect.die(
98 new Error("RPC WebSockets require exclusive ownership of auto-response"),
99 );
100 }
101 let heartbeatEnabled = existing !== null;
102
103 const syncHeartbeat = Effect.suspend(() => {
104 const enabled =
105 !restoring &&
106 !stopped &&
107 heartbeat !== undefined &&
108 byId.size > 0 &&
109 !Array.from(byId.values()).some(
110 (connection) => connection.pending.size > 0,
111 );
112 if (enabled === heartbeatEnabled) return Effect.void;
113 return inRuntime(
114 state.setWebSocketAutoResponse(enabled ? heartbeat : undefined),
115 ).pipe(
116 Effect.tap(() =>
117 Effect.sync(() => {
118 heartbeatEnabled = enabled;
119 }),
120 ),
121 Effect.catchCause(() =>
122 inRuntime(state.abort("Unable to update RPC WebSocket heartbeat")),
123 ),
124 );
125 });
126 yield* syncHeartbeat;
127
128 const persist = (connection: Connection) =>
129 Effect.sync(() => {
130 const previous = connection.socket.deserializeAttachment<unknown>();
131 if (
132 previous !== null &&
133 previous !== undefined &&
134 (typeof previous !== "object" || Array.isArray(previous))
135 ) {
136 throw new Error("Invalid RPC WebSocket attachment");
137 }
138 connection.socket.serializeAttachment({
139 ...previous,
140 [attachmentKey]: {
141 version: 1,
142 clientId: connection.id,
143 pending: connection.pending.size > 0,
144 serialization: serialization.contentType,
145 } satisfies typeof Metadata.Type,
146 });
147 });
148
149 const unregister = (socket: WebSocket) =>
150 Effect.gen(function* () {
151 yield* Effect.sync(() => {
152 closed.add(socket.ws);
153 const connection = bySocket.get(socket.ws);
154 if (connection === undefined) return;
155 bySocket.delete(socket.ws);
156 byId.delete(connection.id);
157 connection.pending.clear();
158 connection.idle.openUnsafe();
159 Queue.offerUnsafe(disconnects, connection.id);
160 });
161 yield* syncHeartbeat;
162 });
163
164 const close = (socket: WebSocket, code = 1012, reason = restartReason) =>
165 socket.close(code, reason).pipe(
166 Effect.catchCause(() =>
167 inRuntime(state.abort("Unable to close RPC WebSocket")),
168 ),
169 Effect.ensuring(unregister(socket)),
170 Effect.uninterruptible,
171 );
172
173 const register = (socket: WebSocket, id: number): Connection => {
174 const connection: Connection = {
175 id,
176 socket,
177 parser: serialization.makeUnsafe(),
178 pending: new Set(),
179 idle: Latch.makeUnsafe(true),
180 };
181 bySocket.set(socket.ws, connection);
182 byId.set(id, connection);
183 nextId = Math.max(nextId, id + 1);
184 return connection;
185 };
186
187 for (const socket of yield* inRuntime(state.getWebSockets(socketTag))) {
188 yield* Effect.gen(function* () {
189 const attachment = yield* Effect.sync(() =>
190 socket.deserializeAttachment<unknown>(),
191 );
192 if (!isAttachment(attachment)) return yield* close(socket);
193 const metadata = attachment[attachmentKey];
194 if (
195 metadata.pending ||
196 metadata.serialization !== serialization.contentType ||
197 !Number.isSafeInteger(metadata.clientId) ||
198 metadata.clientId < 0 ||
199 byId.has(metadata.clientId)
200 ) {
201 return yield* close(socket);
202 }
203 const connection = yield* Effect.sync(() =>
204 register(socket, metadata.clientId),
205 );
206 yield* persist(connection);
207 }).pipe(Effect.catchCause(() => close(socket)));
208 }
209 restoring = false;
210 yield* syncHeartbeat;
211
212 const send = (
213 connection: Connection,
214 response: RpcMessage.FromServerEncoded,
215 ) =>
216 Effect.gen(function* () {
217 if (!byId.has(connection.id)) return;
218 const encoded = yield* Effect.sync(() =>
219 connection.parser.encode(response),
220 );
221 // Buffered responses cannot be recovered after hibernation.
222 if (encoded === undefined) return yield* close(connection.socket);
223 yield* connection.socket.send(encoded);
224 if (response._tag === "Exit") {
225 connection.pending.delete(response.requestId);
226 yield* persist(connection);
227 if (connection.pending.size === 0) yield* connection.idle.open;
228 yield* syncHeartbeat;
229 }
230 if (
231 response._tag === "Defect" ||
232 response._tag === "ClientProtocolError"
233 ) {
234 yield* close(connection.socket);
235 }
236 }).pipe(
237 Effect.catchCause(() => close(connection.socket)),
238 Effect.uninterruptible,
239 );
240
241 const protocol = RpcServer.Protocol.of({
242 run: (write) =>
243 Effect.acquireUseRelease(
244 Effect.sync(() => {
245 if (receive !== undefined || stopped) {
246 throw new Error(
247 "RPC WebSocket protocol is already running or stopped",
248 );
249 }
250 receive = write;
251 }),
252 () => ready.open.pipe(Effect.andThen(Effect.never)),
253 () =>
254 Effect.gen(function* () {
255 stopped = true;
256 receive = undefined;
257 yield* ready.open;
258 for (const connection of Array.from(byId.values())) {
259 yield* close(connection.socket);
260 }
261 yield* syncHeartbeat;
262 }),
263 ),
264 disconnects,
265 send: (id, response) =>
266 Effect.suspend(() => {
267 const connection = byId.get(id);
268 return connection === undefined
269 ? Effect.void
270 : send(connection, response);
271 }),
272 end: (id) =>
273 Effect.suspend(() => {
274 const connection = byId.get(id);
275 return connection === undefined
276 ? Effect.void
277 : close(connection.socket, 1000, "");
278 }),
279 clientIds: Effect.sync(() => new Set(byId.keys())),
280 initialMessage: Effect.succeedNone,
281 supportsAck: true,
282 supportsTransferables: false,
283 supportsSpanPropagation: true,
284 supportsNotifications: true,
285 codecFor: serialization.codecFor,
286 });
287
288 const fetch: HttpEffect = Effect.gen(function* () {
289 const request = yield* HttpServerRequest;
290 if (
291 request.method !== "GET" ||
292 request.headers.upgrade?.toLowerCase() !== "websocket"
293 ) {
294 return HttpServerResponse.empty({
295 status: 426,
296 headers: { Upgrade: "websocket" },
297 });
298 }
299 yield* ready.await;
300 if (stopped) return HttpServerResponse.empty({ status: 503 });
301 const pair = yield* Effect.sync(() => new native.WebSocketPair());
302 const socket = fromWebSocket(pair[1]);
303 let accepted = false;
304 return yield* Effect.gen(function* () {
305 yield* inRuntime(state.acceptWebSocket(socket, [socketTag]));
306 accepted = true;
307 const connection = yield* Effect.sync(() => register(socket, nextId));
308 yield* persist(connection);
309 yield* syncHeartbeat;
310 return yield* Effect.sync(() =>
311 HttpServerResponse.setBody(
312 HttpServerResponse.empty({ status: 101 }),
313 HttpBody.raw(
314 new native.Response(null, { status: 101, webSocket: pair[0] }),
315 ),
316 ),
317 );
318 }).pipe(
319 Effect.catchCause(() =>
320 (accepted ? close(socket) : Effect.void).pipe(
321 Effect.as(HttpServerResponse.empty({ status: 503 })),
322 ),
323 ),
324 Effect.uninterruptible,
325 );
326 });
327
328 const webSocketMessage = (socket: WebSocket, message: string | ArrayBuffer) =>
329 Effect.gen(function* () {
330 if (closed.has(socket.ws)) return;
331 yield* ready.await;
332 const connection = bySocket.get(socket.ws);
333 if (connection === undefined || receive === undefined)
334 return yield* close(socket);
335 const messages = yield* Effect.sync(() =>
336 connection.parser.decode(
337 typeof message === "string" ? message : new Uint8Array(message),
338 ),
339 );
340 // The serializer owns the envelopes; RpcServer validates payloads and schemas.
341 const requests = messages as ReadonlyArray<RpcMessage.FromClientEncoded>;
342 let hasRequests = false;
343 for (const request of requests) {
344 if (request._tag !== "Request") continue;
345 if (
346 (typeof request.id !== "string" &&
347 !(typeof request.id === "number" && Number.isFinite(request.id))) ||
348 connection.pending.has(request.id)
349 ) {
350 return yield* close(socket);
351 }
352 connection.pending.add(request.id);
353 hasRequests = true;
354 }
355 if (hasRequests) {
356 yield* connection.idle.close;
357 // Mark the entire batch before dispatching even its first request.
358 yield* persist(connection);
359 yield* syncHeartbeat;
360 yield* inRuntime(state.waitUntil(connection.idle.await));
361 }
362 for (const request of requests) {
363 if (!byId.has(connection.id)) return;
364 // Direct dispatch preserves this event's context; RpcServer creates each RPC scope.
365 yield* receive(connection.id, request);
366 }
367 }).pipe(Effect.catchCause(() => close(socket)));
368
369 return {
370 protocol,
371 fetch,
372 webSocketMessage,
373 webSocketClose: (socket) => close(socket, 1000, ""),
374 webSocketError: (socket) => close(socket),
375 } satisfies Transport;
376});