feat: collab hocuspocus v4 upgrade (#2351)

* WIP 1 * complete v4 migration * add flushDelay * feat: multiplexing * fix: survive hocuspocus v4 message timeout * fix: searchTerm type error
2026-07-23 23:42:46 +10:00 · 2026-07-20 20:34:09 +01:00
parent ce8fbb86ff
commit a55057db37
13 changed files with 876 additions and 787 deletions
@@ -0,0 +1,46 @@
+import {
+  HocuspocusProviderWebsocket,
+  WebSocketStatus,
+} from "@hocuspocus/provider";
+import { getCollaborationUrl } from "@/lib/config.ts";
+
+const RELEASE_GRACE_MS = 5000;
+
+let socket: HocuspocusProviderWebsocket | null = null;
+let editorCount = 0;
+let releaseTimer: ReturnType<typeof setTimeout> | null = null;
+
+export function getCollabSocket(): HocuspocusProviderWebsocket {
+  if (!socket) {
+    socket = new HocuspocusProviderWebsocket({
+      url: getCollaborationUrl(),
+      autoConnect: false,
+    });
+  }
+  return socket;
+}
+
+export function acquireCollabSocket(): void {
+  editorCount++;
+  if (releaseTimer) {
+    clearTimeout(releaseTimer);
+    releaseTimer = null;
+  }
+  const collabSocket = getCollabSocket();
+  collabSocket.shouldConnect = true;
+  if (collabSocket.status === WebSocketStatus.Disconnected) {
+    collabSocket.connect();
+  }
+}
+
+export function releaseCollabSocket(): void {
+  editorCount--;
+  if (editorCount > 0) return;
+  if (releaseTimer) clearTimeout(releaseTimer);
+  releaseTimer = setTimeout(() => {
+    releaseTimer = null;
+    if (editorCount === 0) {
+      socket?.disconnect();
+    }
+  }, RELEASE_GRACE_MS);
+}
@@ -1,7 +0,0 @@
-import { getCollaborationUrl } from "@/lib/config.ts";
-
-const useCollaborationURL = (): string => {
-  return getCollaborationUrl();
-};
-
-export default useCollaborationURL;
@@ -7,15 +7,16 @@ import React, {
  useState,
 } from "react";
 import { IndexeddbPersistence } from "y-indexeddb";
-import * as Y from "yjs";
 import {
-  HocuspocusProvider,
-  onStatusParameters,
  WebSocketStatus,
-  HocuspocusProviderWebsocket,
-  onSyncedParameters,
  onStatelessParameters,
 } from "@hocuspocus/provider";
+import {
+  HocuspocusProviderWebsocketComponent,
+  HocuspocusRoom,
+  useHocuspocusEvent,
+  useHocuspocusProvider,
+} from "@hocuspocus/provider-react";
 import {
  Editor,
  EditorContent,
@@ -28,7 +29,6 @@ import {
  mainExtensions,
 } from "@/features/editor/extensions/extensions";
 import { useAtom, useAtomValue } from "jotai";
-import useCollaborationUrl from "@/features/editor/hooks/use-collaboration-url";
 import { currentUserAtom } from "@/features/user/atoms/current-user-atom";
 import {
  currentPageEditModeAtom,
@@ -76,6 +76,11 @@ import { EditorLinkMenu } from "@/features/editor/components/link/link-menu";
 import ColumnsMenu from "@/features/editor/components/columns/columns-menu.tsx";
 import { TransclusionLookupProvider } from "@/features/editor/components/transclusion/transclusion-lookup-context";
 import { useTranslation } from "react-i18next";
+import {
+  acquireCollabSocket,
+  getCollabSocket,
+  releaseCollabSocket,
+} from "@/features/editor/collab-socket";

 interface PageEditorProps {
  pageId: string;
@@ -91,7 +96,78 @@ export default function PageEditor({
  canComment,
 }: PageEditorProps) {
  const { t } = useTranslation();
-  const collaborationURL = useCollaborationUrl();
+  const { data: collabQuery, refetch: refetchCollabToken } = useCollabToken();
+  const { pageSlug } = useParams();
+  const slugId = extractPageSlugId(pageSlug);
+  const [socket] = useState(getCollabSocket);
+
+  useEffect(() => {
+    acquireCollabSocket();
+    return () => releaseCollabSocket();
+  }, []);
+
+  const handleStateless = ({ payload }: onStatelessParameters) => {
+    try {
+      const message = JSON.parse(payload);
+      if (message?.type !== "page.updated" || !message.updatedAt) return;
+      const pageData = queryClient.getQueryData<IPage>(["pages", slugId]);
+      if (pageData) {
+        queryClient.setQueryData(["pages", slugId], {
+          ...pageData,
+          updatedAt: message.updatedAt,
+          ...(message.lastUpdatedBy && {
+            lastUpdatedBy: message.lastUpdatedBy,
+          }),
+        });
+      }
+    } catch {
+      // ignore unrelated stateless messages
+    }
+  };
+
+  const handleAuthenticationFailed = () => {
+    const payload = jwtDecode(collabQuery?.token);
+    const now = Date.now().valueOf() / 1000;
+    const isTokenExpired = now >= payload.exp;
+    if (isTokenExpired) {
+      refetchCollabToken();
+    }
+  };
+
+  return (
+    <TransclusionLookupProvider>
+      {collabQuery?.token ? (
+        <HocuspocusProviderWebsocketComponent websocketProvider={socket}>
+          <HocuspocusRoom
+            name={`page.${pageId}`}
+            token={collabQuery.token}
+            flushDelay={500}
+            onStateless={handleStateless}
+            onAuthenticationFailed={handleAuthenticationFailed}
+          >
+            <CollabPageEditor
+              pageId={pageId}
+              editable={editable}
+              content={content}
+              canComment={canComment}
+            />
+          </HocuspocusRoom>
+        </HocuspocusProviderWebsocketComponent>
+      ) : (
+        <StaticPageEditor content={content} ariaLabel={t("Page content")} />
+      )}
+    </TransclusionLookupProvider>
+  );
+}
+
+function CollabPageEditor({
+  pageId,
+  editable,
+  content,
+  canComment,
+}: PageEditorProps) {
+  const { t } = useTranslation();
+  const provider = useHocuspocusProvider();
  const isComponentMounted = useRef(false);
  const editorRef = useRef<Editor | null>(null);

@@ -112,7 +188,6 @@ export default function PageEditor({
  );
  const [, setYjsSynced] = useAtom(yjsSyncedAtom);
  const menuContainerRef = useRef(null);
-  const { data: collabQuery, refetch: refetchCollabToken } = useCollabToken();
  const { isIdle, resetIdle } = useIdle(FIVE_MINUTES, { initialState: false });
  const documentState = useDocumentVisibility();
  const { pageSlug } = useParams();
@@ -123,95 +198,24 @@ export default function PageEditor({
    [isComponentMounted],
  );
  const { handleScrollTo } = useEditorScroll({ canScroll });
-  // Providers only created once per pageId
-  const providersRef = useRef<{
-    local: IndexeddbPersistence;
-    remote: HocuspocusProvider;
-    socket: HocuspocusProviderWebsocket;
-  } | null>(null);
-  const [providersReady, setProvidersReady] = useState(false);

  useEffect(() => {
-    if (!providersRef.current) {
-      const documentName = `page.${pageId}`;
-      const ydoc = new Y.Doc();
-      const local = new IndexeddbPersistence(documentName, ydoc);
-      const socket = new HocuspocusProviderWebsocket({
-        url: collaborationURL,
-      });
-      const onLocalSyncedHandler = () => {
-        setIsLocalSynced(true);
-      };
-      const onStatusHandler = (event: onStatusParameters) => {
-        setYjsConnectionStatus(event.status);
-      };
-      const onSyncedHandler = (event: onSyncedParameters) => {
-        setIsRemoteSynced(event.state);
-      };
-      const onStatelessHandler = ({ payload }: onStatelessParameters) => {
-        try {
-          const message = JSON.parse(payload);
-          if (message?.type !== "page.updated" || !message.updatedAt) return;
-          const pageData = queryClient.getQueryData<IPage>(["pages", slugId]);
-          if (pageData) {
-            queryClient.setQueryData(["pages", slugId], {
-              ...pageData,
-              updatedAt: message.updatedAt,
-              ...(message.lastUpdatedBy && {
-                lastUpdatedBy: message.lastUpdatedBy,
-              }),
-            });
-          }
-        } catch {
-          // ignore unrelated stateless messages
-        }
-      };
-      const onAuthenticationFailedHandler = () => {
-        const payload = jwtDecode(collabQuery?.token);
-        const now = Date.now().valueOf() / 1000;
-        const isTokenExpired = now >= payload.exp;
-        if (isTokenExpired) {
-          refetchCollabToken().then((result) => {
-            if (result.data?.token) {
-              socket.disconnect();
-              setTimeout(() => {
-                remote.configuration.token = result.data.token;
-                socket.connect();
-              }, 100);
-            }
-          });
-        }
-      };
-      const remote = new HocuspocusProvider({
-        websocketProvider: socket,
-        name: documentName,
-        document: ydoc,
-        token: collabQuery?.token,
-        onAuthenticationFailed: onAuthenticationFailedHandler,
-        onStatus: onStatusHandler,
-        onSynced: onSyncedHandler,
-        onStateless: onStatelessHandler,
-      });
-
-      local.on("synced", onLocalSyncedHandler);
-      providersRef.current = { socket, local, remote };
-      setProvidersReady(true);
-    } else {
-      setProvidersReady(true);
-    }
-    // Only destroy on final unmount
+    const local = new IndexeddbPersistence(
+      provider.configuration.name,
+      provider.document,
+    );
+    local.on("synced", () => setIsLocalSynced(true));
    return () => {
-      providersRef.current?.socket.destroy();
-      providersRef.current?.remote.destroy();
-      providersRef.current?.local.destroy();
-      providersRef.current = null;
+      local.destroy();
    };
-  }, [pageId]);
+  }, [provider]);
+
+  useHocuspocusEvent("synced", ({ state }) => setIsRemoteSynced(state));
+  useHocuspocusEvent("status", ({ status }) => setYjsConnectionStatus(status));

  // Only connect/disconnect on tab/idle, not destroy
  useEffect(() => {
-    if (!providersReady || !providersRef.current) return;
-    const socket = providersRef.current.socket;
+    const socket = provider.configuration.websocketProvider;

    if (
      isIdle &&
@@ -228,23 +232,15 @@ export default function PageEditor({
      resetIdle();
      socket.connect();
    }
-  }, [isIdle, documentState, providersReady, resetIdle]);
-
-  // Attach here, to make sure the connection gets properly established
-  providersRef.current?.remote.attach();
+  }, [isIdle, documentState, provider, resetIdle]);

  const extensions = useMemo(() => {
-    if (!providersReady || !providersRef.current || !currentUser?.user) {
+    if (!currentUser?.user) {
      return mainExtensions;
    }

-    const remoteProvider = providersRef.current.remote;
-
-    return [
-      ...mainExtensions,
-      ...collabExtensions(remoteProvider, currentUser?.user),
-    ];
-  }, [providersReady, currentUser?.user]);
+    return [...mainExtensions, ...collabExtensions(provider, currentUser.user)];
+  }, [provider, currentUser?.user]);

  const editor = useEditor(
    {
@@ -416,65 +412,72 @@ export default function PageEditor({
    }
  }, [yjsConnectionStatus, isSynced]);

+  if (showStatic) {
+    return <StaticPageEditor content={content} ariaLabel={t("Page content")} />;
+  }
+
  return (
-    <TransclusionLookupProvider>
-      {showStatic ? (
-        <EditorProvider
-          editable={false}
-          immediatelyRender={true}
-          extensions={mainExtensions}
-          content={content}
-          editorProps={{
-            attributes: {
-              "aria-label": t("Page content"),
-            },
-          }}
-        />
-      ) : (
-        <div className="editor-container" style={{ position: "relative" }}>
-          <div ref={menuContainerRef}>
-            <EditorContent editor={editor} />
+    <div className="editor-container" style={{ position: "relative" }}>
+      <div ref={menuContainerRef}>
+        <EditorContent editor={editor} />

-            {editor && (
-              <SearchAndReplaceDialog editor={editor} editable={editable} />
-            )}
+        {editor && (
+          <SearchAndReplaceDialog editor={editor} editable={editable} />
+        )}

-            {editor && editorIsEditable && (
-              <div>
-                <EditorAiMenu editor={editor} />
-                <EditorLinkMenu editor={editor} />
-                <EditorBubbleMenu editor={editor} />
-                <TableMenu editor={editor} />
-                <TableHandlesLayer editor={editor} />
-                <ImageMenu editor={editor} />
-                <VideoMenu editor={editor} />
-                <PdfMenu editor={editor} />
-                <CalloutMenu editor={editor} />
-                <SubpagesMenu editor={editor} />
-                <ExcalidrawMenu editor={editor} />
-                <DrawioMenu editor={editor} />
-                <ColumnsMenu editor={editor} />
-              </div>
-            )}
-            {editor &&
-              !editorIsEditable &&
-              (editable || canComment) &&
-              providersRef.current && <ReadonlyBubbleMenu editor={editor} />}
-            {showCommentPopup && (
-              <CommentDialog editor={editor} pageId={pageId} />
-            )}
-            {showReadOnlyCommentPopup && (
-              <CommentDialog editor={editor} pageId={pageId} readOnly />
-            )}
+        {editor && editorIsEditable && (
+          <div>
+            <EditorAiMenu editor={editor} />
+            <EditorLinkMenu editor={editor} />
+            <EditorBubbleMenu editor={editor} />
+            <TableMenu editor={editor} />
+            <TableHandlesLayer editor={editor} />
+            <ImageMenu editor={editor} />
+            <VideoMenu editor={editor} />
+            <PdfMenu editor={editor} />
+            <CalloutMenu editor={editor} />
+            <SubpagesMenu editor={editor} />
+            <ExcalidrawMenu editor={editor} />
+            <DrawioMenu editor={editor} />
+            <ColumnsMenu editor={editor} />
          </div>
-          <div
-            onClick={() => {
-              if (editor && !editor.isDestroyed) editor.commands.focus("end");
-            }}
-            style={{ paddingBottom: "20vh" }}
-          ></div>
-        </div>
-      )}
-    </TransclusionLookupProvider>
+        )}
+        {editor && !editorIsEditable && (editable || canComment) && (
+          <ReadonlyBubbleMenu editor={editor} />
+        )}
+        {showCommentPopup && <CommentDialog editor={editor} pageId={pageId} />}
+        {showReadOnlyCommentPopup && (
+          <CommentDialog editor={editor} pageId={pageId} readOnly />
+        )}
+      </div>
+      <div
+        onClick={() => {
+          if (editor && !editor.isDestroyed) editor.commands.focus("end");
+        }}
+        style={{ paddingBottom: "20vh" }}
+      ></div>
+    </div>
+  );
+}
+
+function StaticPageEditor({
+  content,
+  ariaLabel,
+}: {
+  content: any;
+  ariaLabel: string;
+}) {
+  return (
+    <EditorProvider
+      editable={false}
+      immediatelyRender={true}
+      extensions={mainExtensions}
+      content={content}
+      editorProps={{
+        attributes: {
+          "aria-label": ariaLabel,
+        },
+      }}
+    />
  );
 }
@@ -88,7 +88,6 @@
    "kysely-migration-cli": "0.4.2",
    "kysely-postgres-js": "3.0.0",
    "ldapts": "8.1.7",
-    "lib0": "0.2.117",
    "mammoth": "1.12.0",
    "mime-types": "3.0.2",
    "msgpackr": "^1.11.9",
@@ -118,7 +117,6 @@
    "stripe": "^17.7.0",
    "tlds": "1.261.0",
    "tmp-promise": "3.0.3",
-    "tseep": "1.3.1",
    "typesense": "3.0.5",
    "undici": "7.28.0",
    "ws": "8.21.0",
@@ -15,6 +15,7 @@ import {
  RedisSyncExtension,
  SerializedHTTPRequest,
 } from './extensions/redis-sync';
+import { toWebRequest } from './extensions/redis-sync/redis-sync.types';
 import { WsSocketWrapper } from './extensions/redis-sync/ws-socket-wrapper';
 import RedisClient from 'ioredis';
 import { pack, unpack } from 'msgpackr';
@@ -98,34 +99,36 @@ export class CollaborationGateway {
      const serializedHTTPRequest = this.serializeRequest(request);
      const socketId = serializedHTTPRequest.headers['sec-websocket-key'];

-      // Create wrapper socket that only receives events via emit()
-      // This prevents double-handling since Hocuspocus won't listen to raw WebSocket events
      const wrappedSocket = new WsSocketWrapper(client);

      // Route through RedisSync extension (this calls handleConnection internally)
-      this.redisSync.onSocketOpen(wrappedSocket as any, serializedHTTPRequest);
+      this.redisSync.onSocketOpen(wrappedSocket, serializedHTTPRequest);

-      // Forward raw WebSocket messages to the extension
      client.on('message', (data: ArrayBuffer) => {
-        this.redisSync!.onSocketMessage(
-          wrappedSocket as any,
-          serializedHTTPRequest,
-          data,
-        );
+        this.redisSync!.onSocketMessage(serializedHTTPRequest, data);
      });

-      // Forward close events
      client.on('close', (code: number, reason: Buffer) => {
-        this.redisSync!.onSocketClose(socketId, code, reason.buffer as ArrayBuffer);
-      });
-
-      // Forward pong events for keepalive
-      client.on('pong', (data: Buffer) => {
-        wrappedSocket.emit('pong', data);
+        this.redisSync!.onSocketClose(
+          socketId,
+          code,
+          new Uint8Array(reason).buffer,
+        );
      });
    } else {
      // Fallback to direct Hocuspocus connection
-      this.hocuspocus.handleConnection(client, request);
+      const clientConnection = this.hocuspocus.handleConnection(
+        client,
+        toWebRequest(this.serializeRequest(request)),
+      );
+
+      client.on('message', (data: Buffer) => {
+        clientConnection.handleMessage(new Uint8Array(data));
+      });
+
+      client.on('close', (code: number, reason: Buffer) => {
+        clientConnection.handleClose({ code, reason: reason.toString() });
+      });
    }
  }

@@ -178,6 +181,7 @@ export class CollaborationGateway {

        if (this.hocuspocus.getDocumentsCount() === 0) resolve('');
        this.hocuspocus.closeConnections();
+        this.hocuspocus.flushPendingStores();
      } catch (error) {
        console.error(error);
      }
@@ -96,7 +96,7 @@ export class PersistenceExtension implements Extension {
  }

  async onStoreDocument(data: onStoreDocumentPayload) {
-    const { documentName, document, context } = data;
+    const { documentName, document, lastContext } = data;

    const pageId = getPageId(documentName);

@@ -151,7 +151,7 @@ export class PersistenceExtension implements Extension {
            content: tiptapJson,
            textContent: textContent,
            ydoc: ydocState,
-            lastUpdatedById: context.user.id,
+            lastUpdatedById: lastContext.user.id,
            contributorIds: contributorIds,
          },
          pageId,
@@ -169,12 +169,12 @@ export class PersistenceExtension implements Extension {
        JSON.stringify({
          type: 'page.updated',
          updatedAt: new Date().toISOString(),
-          lastUpdatedById: context?.user?.id,
-          lastUpdatedBy: context?.user
+          lastUpdatedById: lastContext?.user?.id,
+          lastUpdatedBy: lastContext?.user
            ? {
-                id: context.user?.id,
-                name: context.user?.name,
-                avatarUrl: context.user?.avatarUrl,
+                id: lastContext.user?.id,
+                name: lastContext.user?.name,
+                avatarUrl: lastContext.user?.avatarUrl,
              }
            : undefined,
        }),
@@ -1,61 +1,37 @@
 import type RedisClient from 'ioredis';
-import { EventEmitter } from 'tseep';
-import type {
-  Pack,
-  RSAMessageClose,
-  RSAMessagePing,
-  RSAMessageSend,
-} from './redis-sync.types';
+import type { WebSocketLike } from '@hocuspocus/server';
+import type { Pack, RSAMessageClose, RSAMessageSend } from './redis-sync.types';

-export class CollabProxySocket extends EventEmitter {
+// Stands in for the client WebSocket on the server that owns the document.
+// Outgoing traffic is relayed over redis to the origin server, which holds the real socket.
+export class CollabProxySocket implements WebSocketLike {
  private readonly replyTo: string;
-  private readonly serverChannel: string;
  private readonly socketId: string;
  private pub: RedisClient;
  private readonly pack: Pack;
  readyState = 1;
+  onClose?: (code?: number, reason?: string) => void;

-  constructor(
-    pub: RedisClient,
-    pack: Pack,
-    replyTo: string,
-    serverChannel: string,
-    socketId: string,
-  ) {
-    super();
+  constructor(pub: RedisClient, pack: Pack, replyTo: string, socketId: string) {
    this.replyTo = replyTo;
    this.socketId = socketId;
-    this.serverChannel = serverChannel;
    this.pub = pub;
    this.pack = pack;
-    this.once('close', () => {
-      this.readyState = 3;
-    });
  }

-  private publish(msg: RSAMessageClose | RSAMessagePing | RSAMessageSend) {
+  private publish(msg: RSAMessageClose | RSAMessageSend) {
    this.pub.publish(this.replyTo, this.pack(msg));
  }

+  // The origin server already closed the real socket; stop relaying without echoing a close back
+  markClosed() {
+    this.readyState = 3;
+  }
+
  close(code?: number, reason?: string) {
    if (this.readyState !== 1) return;
-    const msg: RSAMessageClose = {
-      type: 'close',
-      code,
-      reason,
-      socketId: this.socketId,
-    };
-    this.publish(msg);
-  }
-
-  ping() {
-    if (this.readyState !== 1) return;
-    const msg: RSAMessagePing = {
-      type: 'ping',
-      socketId: this.socketId,
-      replyTo: this.serverChannel,
-    };
-    this.publish(msg);
+    this.readyState = 3;
+    this.onClose?.(code, reason);
  }

  send(message: Uint8Array) {
@@ -3,27 +3,30 @@ import {
  Extension,
  Hocuspocus,
  IncomingMessage,
-  afterUnloadDocumentPayload,
  onConfigurePayload,
  onLoadDocumentPayload,
+  afterUnloadDocumentPayload,
+  WebSocketLike,
 } from '@hocuspocus/server';
+import { ConnectionTimeout, Unauthorized } from '@hocuspocus/common';
 import RedisClient from 'ioredis';
-import { readVarString } from 'lib0/decoding.js';
 import { CollabProxySocket } from './collab-proxy-socket';
 import {
-  BaseWebSocket,
  Configuration,
  CustomEvents,
  Pack,
  RSAMessage,
+  RSAMessageClose,
  RSAMessageCloseProxy,
  RSAMessageCustomEventComplete,
  RSAMessageCustomEventStart,
-  RSAMessagePong,
  RSAMessageProxy,
  RSAMessageUnload,
  SerializedHTTPRequest,
  Unpack,
+  OriginConnection,
+  ProxyConnection,
+  toWebRequest,
 } from './redis-sync.types';

 export type { Pack, SerializedHTTPRequest } from './redis-sync.types';
@@ -38,10 +41,10 @@ export class RedisSyncExtension<TCE extends CustomEvents> implements Extension {
  private sub: RedisClient;
  private readonly pack: Pack;
  private readonly unpack: Unpack;
-  private originSockets: Record<SocketId, BaseWebSocket> = {};
+  private originConnections: Record<SocketId, OriginConnection> = {};
  private locks: Record<DocumentName, NodeJS.Timeout> = {};
  private lockPromises: Record<DocumentName, Promise<ServerId | null>> = {};
-  private proxySockets: Record<SocketId, CollabProxySocket> = {};
+  private proxyConnections: Record<SocketId, ProxyConnection> = {};
  private readonly prefix: string;
  private readonly lockPrefix: string;
  private readonly msgChannel: string;
@@ -54,6 +57,9 @@ export class RedisSyncExtension<TCE extends CustomEvents> implements Extension {
  // @ts-ignore
  private pendingReplies: Record<number, PromiseWithResolvers<any>['resolve']> =
    {};
+  private deriveContext: (
+    serializedHTTPRequest: SerializedHTTPRequest,
+  ) => Record<string, any>;

  constructor(configuration: Configuration<TCE>) {
    const {
@@ -65,6 +71,7 @@ export class RedisSyncExtension<TCE extends CustomEvents> implements Extension {
      prefix,
      customEvents,
      customEventTTL,
+      deriveContext,
    } = configuration;
    this.pub = redis.duplicate();
    this.sub = redis.duplicate();
@@ -77,6 +84,7 @@ export class RedisSyncExtension<TCE extends CustomEvents> implements Extension {
    this.lockPrefix = `${this.prefix}Lock`;
    this.msgChannel = `${this.prefix}Msg`;
    this.customEvents = (customEvents as any) ?? ({} as any as CustomEvents);
+    this.deriveContext = deriveContext ?? (() => ({}));
    this.sub.subscribe(this.msgChannel, `${this.msgChannel}:${this.serverId}`);
    this.sub.on('messageBuffer', this.handleRedisMessage);
    this.pub.on('error', () => {});
@@ -87,44 +95,63 @@ export class RedisSyncExtension<TCE extends CustomEvents> implements Extension {
  }

  private closeProxy(socketId: string) {
-    const proxySocket = this.proxySockets[socketId];
-    if (proxySocket) {
-      proxySocket.emit(
-        'close',
-        1000,
-        Buffer.from('provider_initiated', 'utf-8'),
-      );
-      delete this.proxySockets[socketId];
+    const entry = this.proxyConnections[socketId];
+    if (entry) {
+      delete this.proxyConnections[socketId];
+      const { socket, clientConnection } = entry;
+      // The origin socket is already gone; don't echo a close message back
+      socket.markClosed();
+      clientConnection.handleClose({
+        code: 1000,
+        reason: 'provider_initiated',
+      });
    }
  }

-  private pongProxy(socketId: string) {
-    this.proxySockets[socketId]?.emit('pong');
-  }
-
  private handleProxyMessage(
    msg: Pick<RSAMessageProxy, 'replyTo' | 'message' | 'serializedHTTPRequest'>,
  ) {
    const { replyTo, message, serializedHTTPRequest } = msg;
    const { headers } = serializedHTTPRequest;
-    const socketId = headers['sec-websocket-key']!;
-    let socket = this.proxySockets[socketId];
-    if (!socket) {
-      socket = new CollabProxySocket(
+    const socketId = headers['sec-websocket-key'];
+    let entry = this.proxyConnections[socketId];
+    if (!entry) {
+      const socket = new CollabProxySocket(
        this.pub,
        this.pack,
        replyTo,
-        `${this.msgChannel}:${this.serverId}`,
        socketId,
      );
-      this.proxySockets[socketId] = socket;
-      this.instance.handleConnection(
-        socket as any,
-        serializedHTTPRequest as any,
-        {},
+      // A proxy connection with no live documents (client left the page, auth
+      // failed, or the origin server crashed) is reaped by hocuspocus' message
+      // timeout. Dispose it silently in that case: relaying the timeout close
+      // to the origin would kill the client's real socket, which may be busy
+      // serving other documents. Genuine protocol closes are still relayed.
+      socket.onClose = (code, reason) => {
+        delete this.proxyConnections[socketId];
+        if (code !== ConnectionTimeout.code) {
+          const msg: RSAMessageClose = {
+            type: 'close',
+            code,
+            reason,
+            socketId,
+          };
+          this.pub.publish(replyTo, this.pack(msg));
+        }
+      };
+      const clientConnection = this.instance.handleConnection(
+        socket,
+        toWebRequest(serializedHTTPRequest),
+        this.deriveContext(serializedHTTPRequest),
      );
+      entry = { clientConnection, socket };
+      this.proxyConnections[socketId] = entry;
    }
-    socket.emit('message', message);
+    entry.clientConnection.handleMessage(message);
+  }
+
+  private getLock(documentName: string) {
+    return this.pub.get(this.getKey(documentName));
  }

  private getOrClaimLock(documentName: string) {
@@ -166,10 +193,6 @@ export class RedisSyncExtension<TCE extends CustomEvents> implements Extension {
      this.closeProxy(msg.socketId);
      return;
    }
-    if (type === 'pong') {
-      this.pongProxy(msg.socketId);
-      return;
-    }
    if (type === 'unload') {
      delete this.lockPromises[msg.documentName];
      return;
@@ -198,22 +221,14 @@ export class RedisSyncExtension<TCE extends CustomEvents> implements Extension {
      return;
    }
    const { socketId } = msg;
-    const socket = this.originSockets[socketId];
-    if (!socket) {
+    const entry = this.originConnections[socketId];
+    if (!entry) {
      // origin socket already cleaned up
      return;
    }
+    const { socket } = entry;
    if (type === 'close') {
      socket.close(msg.code, msg.reason);
-    } else if (type === 'ping') {
-      // Reply instantly to the proxy socket, without forwarding to client
-      // The origin socket handles heartbeat for itself
-      const { replyTo, socketId } = msg;
-      const reply: RSAMessagePong = {
-        type: 'pong',
-        socketId,
-      };
-      this.pub.publish(`${replyTo}`, this.pack(reply));
    } else if (type === 'send') {
      socket.send(msg.message);
    }
@@ -251,6 +266,8 @@ export class RedisSyncExtension<TCE extends CustomEvents> implements Extension {
    eventName: TName,
    documentName: string,
    payload: any,
+    // if true, don't claim the lock. Useful for targeting pages that are currently open
+    onlyIfOpen = false,
  ) {
    const isDocLoadedOnInstance = this.instance.documents.has(documentName);

@@ -258,7 +275,14 @@ export class RedisSyncExtension<TCE extends CustomEvents> implements Extension {
      return this.handleEventLocally(eventName, documentName, payload);
    }

-    const proxyTo = await this.getOrClaimLockThrottled(documentName);
+    const proxyTo = await (onlyIfOpen
+      ? this.getLock(documentName)
+      : this.getOrClaimLockThrottled(documentName));
+
+    if (!proxyTo && onlyIfOpen) {
+      return;
+    }
+
    if (proxyTo && proxyTo !== this.serverId) {
      ++this.replyIdCounter; // bug in biome thinks this.replyIdCounter is not used if written on the line below
      const replyId = this.replyIdCounter;
@@ -277,7 +301,8 @@ export class RedisSyncExtension<TCE extends CustomEvents> implements Extension {
      const { promise, resolve, reject } = Promise.withResolvers();
      this.pendingReplies[replyId] = resolve;
      setTimeout(() => {
-        reject('TIMEOUT');
+        delete this.pendingReplies[replyId];
+        reject(new Error('TIMEOUT'));
      }, this.customEventTTL);
      return promise as Promise<ReturnType<TCE[TName]>>;
    }
@@ -296,36 +321,59 @@ export class RedisSyncExtension<TCE extends CustomEvents> implements Extension {

  /* WebSocket Server Hooks */
  onSocketOpen(
-    ws: BaseWebSocket,
+    ws: WebSocketLike,
    serializedHTTPRequest: SerializedHTTPRequest,
-    context = {},
  ) {
-    const socketId = serializedHTTPRequest.headers['sec-websocket-key']!;
-    this.originSockets[socketId] = ws;
-    this.instance.handleConnection(
-      ws as any,
-      serializedHTTPRequest as any,
-      context,
+    const socketId = serializedHTTPRequest.headers['sec-websocket-key'];
+    const clientConnection = this.instance.handleConnection(
+      ws,
+      toWebRequest(serializedHTTPRequest),
+      this.deriveContext(serializedHTTPRequest),
    );
+    this.originConnections[socketId] = { clientConnection, socket: ws };
  }

  async onSocketMessage(
-    ws: BaseWebSocket,
    serializedHTTPRequest: SerializedHTTPRequest,
    detachableMsg: ArrayBuffer,
  ) {
-    const message = new Uint8Array(detachableMsg.slice());
-    const tmpMsg = new IncomingMessage(detachableMsg);
-    const documentName = readVarString(tmpMsg.decoder);
+    const socketId = serializedHTTPRequest.headers['sec-websocket-key'];
+    const entry = this.originConnections[socketId];
+    if (!entry) return;
+    const { clientConnection } = entry;
+
+    let message: Uint8Array;
+    let documentName: string;
+    try {
+      message = new Uint8Array(detachableMsg.slice());
+      const tmpMsg = new IncomingMessage(detachableMsg);
+      const documentNameAndSessionId = tmpMsg.readVarString();
+      // session-aware providers suffix the documentName with \0sessionId
+      const sepIdx = documentNameAndSessionId.indexOf('\0');
+      documentName =
+        sepIdx === -1
+          ? documentNameAndSessionId
+          : documentNameAndSessionId.slice(0, sepIdx);
+    } catch (error) {
+      entry.socket.close(Unauthorized.code, Unauthorized.reason);
+      return;
+    }
    const isDocLoadedOnInstance = this.instance.documents.has(documentName);

    if (isDocLoadedOnInstance) {
-      ws.emit('message', message);
+      clientConnection.handleMessage(message);
      return;
    }

    const proxyTo = await this.getOrClaimLockThrottled(documentName);
    if (proxyTo && proxyTo !== this.serverId) {
+      // Proxied messages bypass handleMessage, so refresh the connection's
+      // liveness fields manually or hocuspocus' message timeout would reap the
+      // real socket every `timeout` ms. connectionEstablishedAt is the
+      // reference while unauthenticated (auth for remote docs is proxied too)
+      // and is private upstream.
+      clientConnection.lastMessageReceivedAt = Date.now();
+      (clientConnection as any).connectionEstablishedAt = Date.now();
      // another server owns the doc
      const proxyMessage: RSAMessageProxy = {
        serializedHTTPRequest: serializedHTTPRequest,
@@ -338,16 +386,17 @@ export class RedisSyncExtension<TCE extends CustomEvents> implements Extension {
      return;
    }
    // This server owns the document, but hocuspocus hasn't loaded it yet
-    ws.emit('message', message);
+    clientConnection.handleMessage(message);
  }

  onSocketClose(socketId: string, code?: number, reason?: ArrayBuffer) {
-    const socket = this.originSockets[socketId];
-    if (!socket) return;
-    // at this point the socket is considered GC'd and we cannot call close
-    // The origin socket did not set up any connections for the proxy, so none of the hooks will work if we just emit
-    socket?.emit('close', code, reason);
-    delete this.originSockets[socketId];
+    const entry = this.originConnections[socketId];
+    if (!entry) return;
+    delete this.originConnections[socketId];
+    entry.clientConnection.handleClose({
+      code: code ?? 1000,
+      reason: reason ? Buffer.from(reason).toString() : '',
+    });
    const msg: RSAMessageCloseProxy = { type: 'closeProxy', socketId };
    this.pub.publish(this.msgChannel, this.pack(msg)).catch(() => {});
  }
@@ -372,6 +421,7 @@ export class RedisSyncExtension<TCE extends CustomEvents> implements Extension {
  }

  async onDestroy() {
+    this.pendingReplies = {};
    this.pub.disconnect(false);
    this.sub.disconnect(false);
  }
@@ -1,12 +1,13 @@
-import EventEmitter from 'node:events';
 import { IncomingHttpHeaders } from 'node:http2';
 import RedisClient from 'ioredis';
+import { CollabProxySocket } from './collab-proxy-socket';
+import { type Hocuspocus, type WebSocketLike } from '@hocuspocus/server';

 export type SecondParam<T> = T extends (
-  arg1: unknown,
+  arg1: any,
  arg2: infer A,
-  ...args: unknown[]
-) => unknown
+  ...args: any[]
+) => any
  ? A
  : never;

@@ -41,17 +42,6 @@ export type RSAMessageClose = {
  socketId: string;
 };

-export type RSAMessagePing = {
-  type: 'ping';
-  socketId: string;
-  replyTo: string;
-};
-
-export type RSAMessagePong = {
-  type: 'pong';
-  socketId: string;
-};
-
 export type RSAMessageSend = {
  type: 'send';
  // @ts-ignore
@@ -59,7 +49,7 @@ export type RSAMessageSend = {
  socketId: string;
 };

-export type RSAMessageCustomEventStart<TName = string, TPayload = unknown> = {
+export type RSAMessageCustomEventStart<TName = string, TPayload = any> = {
  type: 'customEventStart';
  documentName: string;
  eventName: TName;
@@ -71,7 +61,7 @@ export type RSAMessageCustomEventStart<TName = string, TPayload = unknown> = {
 export type RSAMessageCustomEventComplete = {
  type: 'customEventComplete';
  replyId: number;
-  payload: unknown;
+  payload: any;
 };

 export type RSAMessage =
@@ -79,8 +69,6 @@ export type RSAMessage =
  | RSAMessageCloseProxy
  | RSAMessageUnload
  | RSAMessageClose
-  | RSAMessagePing
-  | RSAMessagePong
  | RSAMessageSend
  | RSAMessageCustomEventStart
  | RSAMessageCustomEventComplete;
@@ -99,9 +87,20 @@ type CustomEventName = string;

 export type CustomEvents = Record<
  CustomEventName,
-  (documentName: string, payload: unknown) => Promise<unknown>
+  (documentName: string, payload: any) => Promise<any>
 >;

+// Not exported by @hocuspocus/server
+export type ClientConnection = ReturnType<Hocuspocus['handleConnection']>;
+export type OriginConnection = {
+  clientConnection: ClientConnection;
+  socket: WebSocketLike;
+};
+export type ProxyConnection = {
+  clientConnection: ClientConnection;
+  socket: CollabProxySocket;
+};
+
 export interface Configuration<TCE> {
  redis: RedisClient;
  pack: Pack;
@@ -111,11 +110,29 @@ export interface Configuration<TCE> {
  customEventTTL?: number;
  prefix?: string;
  customEvents?: TCE;
+  // Derive the hocuspocus context once per socket instead of re-deriving it in a
+  // per-document hook like onConnect/onAuthenticate. Runs on the origin server when
+  // the socket opens and on the doc owner when the first proxied message arrives.
+  deriveContext?: (
+    serializedHTTPRequest: SerializedHTTPRequest,
+  ) => Record<string, any>;
 }

-export type BaseWebSocket = EventEmitter & {
-  readyState: number;
-  close(code?: number, reason?: string): void;
-  ping(): void;
-  send(message: Uint8Array): void;
+// Hocuspocus expects a web-standard Request, so rehydrate one from what crossed the wire
+export const toWebRequest = (serializedHTTPRequest: SerializedHTTPRequest) => {
+  const { method, url, headers } = serializedHTTPRequest;
+  const webHeaders = new Headers();
+  Object.entries(headers).forEach(([name, value]) => {
+    if (Array.isArray(value)) {
+      value.forEach((v) => {
+        webHeaders.append(name, v);
+      });
+    } else if (value !== undefined) {
+      webHeaders.set(name, value);
+    }
+  });
+  return new Request(new URL(url, 'http://localhost'), {
+    method,
+    headers: webHeaders,
+  });
 };
@@ -1,20 +1,17 @@
-import { EventEmitter } from 'events';
 import type WebSocket from 'ws';
+import type { WebSocketLike } from '@hocuspocus/server';

 /**
- * Wrapper around ws WebSocket that only receives events via emit().
- * This prevents double-handling when used with RedisSyncExtension.
+ * Wrapper around ws WebSocket that Hocuspocus only writes to.
+ * Incoming socket events are forwarded separately by the gateway,
+ * which prevents double-handling with RedisSyncExtension.
 */
-export class WsSocketWrapper extends EventEmitter {
+export class WsSocketWrapper implements WebSocketLike {
  private ws: WebSocket;
  readyState = 1;

  constructor(ws: WebSocket) {
-    super();
    this.ws = ws;
-    this.once('close', () => {
-      this.readyState = 3;
-    });
  }

  close(code?: number, reason?: string) {
@@ -27,15 +24,6 @@ export class WsSocketWrapper extends EventEmitter {
    }
  }

-  ping() {
-    if (this.readyState !== 1) return;
-    try {
-      this.ws.ping();
-    } catch (e) {
-      // Socket already closed
-    }
-  }
-
  send(message: Uint8Array) {
    if (this.readyState !== 1) return;
    try {