feat(server): add beforeHandleAwareness hook (#1096)

Fires after decoding an inbound awareness message and before
applyAwarenessUpdate, exposing per-client states as a mutable
Map<clientId, state>. Mutations rewrite the broadcast; throwing rejects
the update. Extensions chain naturally and run before the
configuration-level hook.

Payload shape matches onAwarenessUpdatePayload: transactionOrigin plus
connection? convenience shortcut.

Author: dschmidt

packages/server/src/Document.ts

@@ -21,6 +21,11 @@ export class Document extends Doc {
 		// eslint-disable-next-line @typescript-eslint/no-empty-function
 		onUpdate: (document: Document, origin: unknown, update: Uint8Array) => {},
 		beforeBroadcastStateless: (document: Document, stateless: string) => {},
+		beforeHandleAwareness: (
+			document: Document,
+			states: Map<number, Record<string, any>>,
+			transactionOrigin: unknown,
+		) => Promise.resolve(),
 	};
 
 	connections: Map<
@@ -100,6 +105,27 @@ export class Document extends Doc {
 		return this;
 	}
 
+	/**
+	 * Set a callback that will be triggered before an inbound awareness update
+	 * is applied to this document's awareness state. The callback receives the
+	 * document, the decoded per-client states as a mutable `Map`, and the
+	 * `TransactionOrigin` that will be forwarded to `applyAwarenessUpdate`.
+	 * Use `isTransactionOrigin(origin)` to discriminate sources. Mutate the
+	 * map in place (set/delete/field changes) to rewrite the update, or throw
+	 * to reject it entirely.
+	 */
+	beforeHandleAwareness(
+		callback: (
+			document: Document,
+			states: Map<number, Record<string, any>>,
+			transactionOrigin: unknown,
+		) => Promise<any>,
+	): Document {
+		this.callbacks.beforeHandleAwareness = callback;
+
+		return this;
+	}
+
 	/**
 	 * Register a connection and a set of clients on this document keyed by the
 	 * underlying websocket connection

packages/server/src/Hocuspocus.ts

@@ -46,6 +46,7 @@ export class Hocuspocus<Context = any> {
 		onConnect: () => new Promise((r) => r(null)),
 		connected: () => new Promise((r) => r(null)),
 		beforeHandleMessage: () => new Promise((r) => r(null)),
+		beforeHandleAwareness: () => new Promise<void>((r) => r()),
 		beforeSync: () => new Promise((r) => r(null)),
 		beforeBroadcastStateless: () => new Promise((r) => r(null)),
 		onStateless: () => new Promise((r) => r(null)),
@@ -112,6 +113,7 @@ export class Hocuspocus<Context = any> {
 			onLoadDocument: this.configuration.onLoadDocument,
 			afterLoadDocument: this.configuration.afterLoadDocument,
 			beforeHandleMessage: this.configuration.beforeHandleMessage,
+			beforeHandleAwareness: this.configuration.beforeHandleAwareness,
 			beforeBroadcastStateless: this.configuration.beforeBroadcastStateless,
 			beforeSync: this.configuration.beforeSync,
 			onStateless: this.configuration.onStateless,
@@ -436,6 +438,31 @@ export class Hocuspocus<Context = any> {
 			},
 		);
 
+		document.beforeHandleAwareness((document, states, transactionOrigin) => {
+			const connection =
+				isTransactionOrigin(transactionOrigin) &&
+				transactionOrigin.source === "connection"
+					? transactionOrigin.connection
+					: undefined;
+			const request = connection?.request;
+			return this.hooks("beforeHandleAwareness", {
+				awareness: document.awareness,
+				clientsCount: document.getConnectionsCount(),
+				context: connection?.context,
+				document,
+				documentName: document.name,
+				instance: this,
+				requestHeaders: request?.headers ?? new Headers(),
+				requestParameters: request
+					? getParameters(request)
+					: new URLSearchParams(),
+				socketId: connection?.socketId ?? "",
+				transactionOrigin,
+				connection,
+				states,
+			});
+		});
+
 		document.awareness.on(
 			"update",
 			(update: AwarenessUpdate, origin: unknown) => {

packages/server/src/MessageReceiver.ts

@@ -1,7 +1,11 @@
 import { AuthMessageType } from "@hocuspocus/common";
 import * as decoding from "lib0/decoding";
 import { readVarString } from "lib0/decoding";
-import { applyAwarenessUpdate } from "y-protocols/awareness";
+import {
+	Awareness,
+	applyAwarenessUpdate,
+	encodeAwarenessUpdate,
+} from "y-protocols/awareness";
 import {
 	messageYjsSyncStep1,
 	messageYjsSyncStep2,
@@ -66,19 +70,36 @@ export class MessageReceiver {
 				break;
 			}
 			case MessageType.Awareness: {
+				let update = message.readVarUint8Array();
+
 				const origin: TransactionOrigin = connection
 					? ({
 							source: "connection",
 							connection,
 						} satisfies ConnectionTransactionOrigin)
 					: (this.defaultTransactionOrigin ?? { source: "local" });
 
-				applyAwarenessUpdate(
-					document.awareness,
-					message.readVarUint8Array(),
+				// Decode the inbound update into a scratch Awareness so the hook
+				// chain sees a high-level Map<clientId, state>. Mutations to that
+				// map (including `set`, `delete`, and field changes on each state
+				// object) are picked up by the re-encode below and forwarded as
+				// the broadcast payload. Hooks may also throw to reject the
+				// update entirely.
+				const scratch = new Awareness(new Y.Doc());
+				applyAwarenessUpdate(scratch, update, null);
+
+				await document.callbacks.beforeHandleAwareness(
+					document,
+					scratch.getStates(),
 					origin,
 				);
 
+				update = encodeAwarenessUpdate(scratch, [
+					...scratch.getStates().keys(),
+				]);
+
+				applyAwarenessUpdate(document.awareness, update, origin);
+
 				break;
 			}
 			case MessageType.QueryAwareness: {

packages/server/src/types.ts

@@ -98,6 +98,20 @@ export interface Extension<Context = any> {
 	onLoadDocument?(data: onLoadDocumentPayload<Context>): Promise<any>;
 	afterLoadDocument?(data: afterLoadDocumentPayload<Context>): Promise<any>;
 	beforeHandleMessage?(data: beforeHandleMessagePayload<Context>): Promise<any>;
+	/**
+	 * Fired before an inbound awareness update is applied to the document's
+	 * awareness state. The hook receives the decoded per-client `states` as a
+	 * mutable `Map` keyed by Yjs clientId. Mutate the map and the contained
+	 * state objects in place to rewrite fields, drop peers (`states.delete`),
+	 * or add synthetic ones (`states.set`); mutations are reflected in the
+	 * broadcast. Throw to reject the update without applying anything.
+	 *
+	 * Multiple extensions chain naturally: each extension sees the map as
+	 * mutated by previous extensions and can mutate it further.
+	 */
+	beforeHandleAwareness?(
+		data: beforeHandleAwarenessPayload<Context>,
+	): Promise<any>;
 	beforeSync?(data: beforeSyncPayload<Context>): Promise<any>;
 	beforeBroadcastStateless?(
 		data: beforeBroadcastStatelessPayload,
@@ -126,6 +140,7 @@ export type HookName =
 	| "onLoadDocument"
 	| "afterLoadDocument"
 	| "beforeHandleMessage"
+	| "beforeHandleAwareness"
 	| "beforeBroadcastStateless"
 	| "beforeSync"
 	| "onStateless"
@@ -151,6 +166,7 @@ export type HookPayloadByName<Context = any> = {
 	onLoadDocument: onLoadDocumentPayload<Context>;
 	afterLoadDocument: afterLoadDocumentPayload<Context>;
 	beforeHandleMessage: beforeHandleMessagePayload<Context>;
+	beforeHandleAwareness: beforeHandleAwarenessPayload<Context>;
 	beforeBroadcastStateless: beforeBroadcastStatelessPayload;
 	beforeSync: beforeSyncPayload<Context>;
 	onStateless: onStatelessPayload;
@@ -326,6 +342,44 @@ export interface beforeHandleMessagePayload<Context = any> {
 	connection: Connection<Context>;
 }
 
+export interface beforeHandleAwarenessPayload<Context = any> {
+	awareness: Awareness;
+	clientsCount: number;
+	/**
+	 * Connection context populated by `onAuthenticate`. `undefined` when the
+	 * update did not originate from a client connection (e.g. server-internal
+	 * writes via `DirectConnection`).
+	 */
+	context: Context | undefined;
+	document: Document;
+	documentName: string;
+	instance: Hocuspocus;
+	requestHeaders: Headers;
+	requestParameters: URLSearchParams;
+	/**
+	 * Per-client awareness states decoded from the inbound update, keyed by
+	 * Yjs clientId. Mutate this map in place to rewrite the update: change
+	 * fields on a state object, `states.delete(clientId)` to drop a peer, or
+	 * `states.set(clientId, ...)` to add or replace one. The encoded update
+	 * sent to peers reflects whatever the map looks like after every hook in
+	 * the chain has run.
+	 */
+	states: Map<number, Record<string, any>>;
+	socketId: string;
+	/**
+	 * The `TransactionOrigin` that will be passed to `applyAwarenessUpdate`.
+	 * Use `isTransactionOrigin(origin)` to discriminate sources. Matches the
+	 * `transactionOrigin` shape of `onAwarenessUpdatePayload`.
+	 */
+	transactionOrigin: unknown;
+	/**
+	 * Convenience shortcut: `origin.connection` when `transactionOrigin` is a
+	 * `ConnectionTransactionOrigin`, otherwise `undefined`. Matches the
+	 * `connection?` shape of `onAwarenessUpdatePayload`.
+	 */
+	connection?: Connection<Context>;
+}
+
 export interface beforeSyncPayload<Context = any> {
 	clientsCount: number;
 	context: Context;