Add hibernation websocket handler for cloudflare DO

Author: wernst

.changeset/nervous-birds-fetch.md

@@ -0,0 +1,5 @@
+---
+'@triplit/server': patch
+---
+
+Add hibernation websocket handler for cloudflare DO

packages/cf-worker-server/src/hibernation.ts

@@ -0,0 +1,96 @@
+import { DurableObject } from 'cloudflare:workers';
+import {
+	upgradeWebSocket,
+	upgradeWebSocketHibernation,
+} from '@triplit/server/cloudflare';
+import { createTriplitHonoServer } from '@triplit/server/hono';
+import { CloudflareDurableObjectKVStore } from '@triplit/db/storage/cf-durable-object';
+
+type ExtendedDurableObjectState = DurableObjectState & {
+	honoWs?: {
+		events?: any;
+		wsContext?: any;
+	};
+};
+
+export class MyDurableObject extends DurableObject {
+	state: ExtendedDurableObjectState;
+	private appPromise: Promise<
+		Awaited<ReturnType<typeof createTriplitHonoServer>>
+	>;
+
+	constructor(ctx: ExtendedDurableObjectState, env: Env) {
+		super(ctx, env);
+		console.log('CALLING CONSTRUCTOR');
+		this.state = ctx;
+		// Create the Triplit server
+		this.appPromise = createTriplitHonoServer(
+			{
+				// add any configuration options here
+				jwtSecret: env.JWT_SECRET,
+				// this is the Triplit storage provider for Durable Objects
+				storage: new CloudflareDurableObjectKVStore(this.state.storage),
+			},
+			// inject the platform-specific WebSocket upgrade function
+			upgradeWebSocketHibernation(ctx)
+		);
+	}
+
+	async fetch(request: Request) {
+		// Await the app initialization before handling the request
+		const app = await this.appPromise;
+		return app.fetch(request);
+	}
+
+	async webSocketMessage(ws: WebSocket, message: string | ArrayBuffer) {
+		await this.appPromise;
+		if (this.state.honoWs?.events?.onMessage) {
+			const wsContext = this.state.honoWs.wsContext;
+			this.state.honoWs.events.onMessage(
+				{
+					data: message,
+				},
+				wsContext
+			);
+		}
+	}
+
+	async webSocketClose(
+		ws: WebSocket,
+		code: number,
+		reason: string,
+		wasClean: boolean
+	) {
+		await this.appPromise;
+		if (this.state.honoWs?.events?.onClose) {
+			const wsContext = this.state.honoWs.wsContext;
+			this.state.honoWs.events.onClose(
+				{
+					code,
+					reason,
+					wasClean,
+				},
+				wsContext
+			);
+		}
+	}
+
+	async webSocketError(ws: WebSocket, error: Error) {
+		await this.appPromise;
+		if (this.state.honoWs?.events?.onError) {
+			const wsContext = this.state.honoWs.wsContext;
+			this.state.honoWs.events.onError(error, wsContext);
+		}
+	}
+}
+
+export default {
+	async fetch(request, env, _ctx): Promise<Response> {
+		// Get the Durable Object ID (this is where you could easily add multi-tenancy)
+		let id: DurableObjectId = env.MY_DURABLE_OBJECT.idFromName('triplitDB');
+		let stub = env.MY_DURABLE_OBJECT.get(id);
+
+		// Forward the request to the Durable Object
+		return await stub.fetch(request);
+	},
+} satisfies ExportedHandler<Env>;

packages/cf-worker-server/src/index.ts packages/cf-worker-server/src/standard.tspackages/cf-worker-server/src/index.ts renamed to packages/cf-worker-server/src/standard.ts

@@ -1,7 +1,9 @@
 // Generated by Wrangler by running `wrangler types --include-runtime=false` (hash: 4c81d17e8de07f650cddf73ce751ba8d)
 declare namespace Cloudflare {
 	interface Env {
-		MY_DURABLE_OBJECT: DurableObjectNamespace<import('./src/index').MyDurableObject>;
+		MY_DURABLE_OBJECT: DurableObjectNamespace<
+			import('./src/standard').MyDurableObject
+		>;
 		JWT_SECRET: string;
 	}
 }

packages/cf-worker-server/worker-configuration.d.ts

@@ -5,7 +5,7 @@
 {
 	"$schema": "node_modules/wrangler/config-schema.json",
 	"name": "cf-worker-server",
-	"main": "src/index.ts",
+	"main": "src/standard.ts",
 	"compatibility_date": "2025-04-30",
 	"migrations": [
 		{

packages/cf-worker-server/wrangler.jsonc

@@ -69,3 +69,51 @@ export const upgradeWebSocket: UpgradeWebSocket<
     webSocket: client,
   });
 });
+
+export function upgradeWebSocketHibernation(
+  ctx: any
+): UpgradeWebSocket<WebSocket, any, Omit<WSEvents<WebSocket>, 'onOpen'>> {
+  return defineWebSocketHelper(async (c, events) => {
+    const upgradeHeader = c.req.header('Upgrade');
+    if (upgradeHeader !== 'websocket') {
+      return;
+    }
+
+    // @ts-expect-error WebSocketPair is not typed
+    const webSocketPair = new WebSocketPair();
+    const client: any = webSocketPair[0];
+    const server: any = webSocketPair[1];
+
+    const wsContext = new WSContext<WebSocket>({
+      close: (code, reason) => server.close(code, reason),
+      get protocol() {
+        return server.protocol;
+      },
+      raw: server,
+      get readyState() {
+        return server.readyState as WSReadyState;
+      },
+      url: server.url ? new URL(server.url) : null,
+      send: (source) => server.send(source),
+    });
+
+    // Add hono content so its accessible by the durable object after initialization
+    ctx.honoWs = {
+      events,
+      wsContext,
+    };
+    // This API allows the durable object to accept the WebSocket connection
+    ctx.acceptWebSocket(server);
+
+    // note: cloudflare workers doesn't support 'open' event
+    if (events.onOpen) {
+      events.onOpen(new Event('open'), wsContext);
+    }
+
+    return new Response(null, {
+      status: 101,
+      // @ts-expect-error - webSocket is not typed
+      webSocket: client,
+    });
+  });
+}