chore: better types

Author: janthurau

packages/extension-database/src/Database.ts

@@ -1,9 +1,10 @@
 import type {
 	Extension,
+	fetchPayload,
 	onChangePayload,
 	onLoadDocumentPayload,
+	onStoreDocumentPayload,
 	storePayload,
-	fetchPayload,
 } from "@hocuspocus/server";
 import * as Y from "yjs";
 
@@ -52,7 +53,7 @@ export class Database implements Extension {
 	/**
 	 * Store new updates in the database.
 	 */
-	async onStoreDocument(data: onChangePayload) {
+	async onStoreDocument(data: onStoreDocumentPayload) {
 		await this.configuration.store({
 			...data,
 			state: Buffer.from(Y.encodeStateAsUpdate(data.document)),

packages/extension-logger/src/Logger.ts

@@ -3,10 +3,11 @@ import type {
 	onChangePayload,
 	onConfigurePayload,
 	onConnectPayload,
-	onLoadDocumentPayload,
 	onDestroyPayload,
 	onDisconnectPayload,
+	onLoadDocumentPayload,
 	onRequestPayload,
+	onStoreDocumentPayload,
 	onUpgradePayload,
 } from "@hocuspocus/server";
 
@@ -112,7 +113,7 @@ export class Logger implements Extension {
 		}
 	}
 
-	async onStoreDocument(data: onDisconnectPayload) {
+	async onStoreDocument(data: onStoreDocumentPayload) {
 		if (this.configuration.onStoreDocument) {
 			this.log(`Store "${data.documentName}".`);
 		}

packages/extension-redis/src/Redis.ts

@@ -281,7 +281,7 @@ export class Redis implements Extension {
 	 */
 	async afterStoreDocument({
 		documentName,
-		socketId,
+		lastTransactionOrigin,
 	}: afterStoreDocumentPayload) {
 		const lockKey = this.lockKey(documentName);
 		const lock = this.locks.get(lockKey);
@@ -296,9 +296,13 @@ export class Redis implements Extension {
 				this.locks.delete(lockKey);
 			}
 		}
-		// if the change was initiated by a directConnection, we need to delay this hook to make sure sync can finish first.
+
+		// if the change was initiated by a directConnection (or otherwise local source), we need to delay this hook to make sure sync can finish first.
 		// for provider connections, this usually happens in the onDisconnect hook
-		if (socketId === "server") {
+		if (
+			isTransactionOrigin(lastTransactionOrigin) &&
+			lastTransactionOrigin.source === "local"
+		) {
 			const pending = this.pendingAfterStoreDocumentResolves.get(documentName);
 
 			if (pending) {

packages/server/src/DirectConnection.ts

@@ -51,13 +51,14 @@ export class DirectConnection<Context = any>
 				this.document,
 				{
 					clientsCount: this.document.getConnectionsCount(),
-					context: this.context,
+					lastContext: this.context,
+					lastTransactionOrigin: {
+						source: "local",
+						context: this.context,
+					} satisfies LocalTransactionOrigin,
 					document: this.document,
 					documentName: this.document.name,
 					instance: this.instance,
-					requestHeaders: {},
-					requestParameters: new URLSearchParams(),
-					socketId: "server",
 				},
 				true,
 			);

packages/server/src/Hocuspocus.ts

@@ -264,26 +264,36 @@ export class Hocuspocus<Context = any> {
 					: {}
 			: {};
 
-		const hookPayload: onChangePayload | onStoreDocumentPayload = {
+		const changePayload: onChangePayload = {
 			instance: this,
 			clientsCount: document.getConnectionsCount(),
-			context,
 			document,
 			documentName: document.name,
 			requestHeaders: request?.headers ?? {},
 			requestParameters: getParameters(request),
 			socketId: connection?.socketId ?? "",
 			update,
 			transactionOrigin: origin,
+			connection: connection,
+			context,
 		};
 
-		this.hooks("onChange", hookPayload);
+		this.hooks("onChange", changePayload);
 
 		if (shouldSkipStoreHooks(origin)) {
 			return;
 		}
 
-		this.storeDocumentHooks(document, hookPayload);
+		const storePayload: onStoreDocumentPayload = {
+			instance: this,
+			clientsCount: document.getConnectionsCount(),
+			document,
+			lastContext: context,
+			lastTransactionOrigin: origin,
+			documentName: document.name,
+		};
+
+		this.storeDocumentHooks(document, storePayload);
 	}
 
 	/**
@@ -408,14 +418,24 @@ export class Hocuspocus<Context = any> {
 			},
 		);
 
-		document.awareness.on("update", (update: AwarenessUpdate) => {
-			this.hooks("onAwarenessUpdate", {
-				...hookPayload,
-				...update,
-				awareness: document.awareness,
-				states: awarenessStatesToArray(document.awareness.getStates()),
-			});
-		});
+		document.awareness.on(
+			"update",
+			(update: AwarenessUpdate, origin: unknown) => {
+				this.hooks("onAwarenessUpdate", {
+					document,
+					documentName,
+					instance: this,
+					...update,
+					transactionOrigin: origin,
+					connection:
+						isTransactionOrigin(origin) && origin.source === "connection"
+							? origin.connection
+							: undefined,
+					awareness: document.awareness,
+					states: awarenessStatesToArray(document.awareness.getStates()),
+				});
+			},
+		);
 
 		return document;
 	}
@@ -460,6 +480,7 @@ export class Hocuspocus<Context = any> {
 	hooks<T extends HookName>(
 		name: T,
 		payload: HookPayloadByName<Context>[T],
+		// biome-ignore lint/complexity/noBannedTypes: <explanation>
 		callback: Function | null = null,
 	): Promise<any> {
 		const { extensions } = this.configuration;

packages/server/src/types.ts

@@ -301,6 +301,7 @@ export interface onChangePayload<Context = any> {
 	update: Uint8Array;
 	socketId: string;
 	transactionOrigin: unknown;
+	connection?: Connection<Context>;
 }
 
 export interface beforeHandleMessagePayload<Context = any> {
@@ -346,27 +347,23 @@ export interface beforeBroadcastStatelessPayload {
 
 export interface onStoreDocumentPayload<Context = any> {
 	clientsCount: number;
-	context: Context;
 	document: Document;
+	lastContext: Context;
+	lastTransactionOrigin: unknown;
 	documentName: string;
 	instance: Hocuspocus;
-	requestHeaders: IncomingHttpHeaders;
-	requestParameters: URLSearchParams;
-	socketId: string;
-	transactionOrigin?: unknown;
 }
 
 // eslint-disable-next-line @typescript-eslint/no-empty-interface, @typescript-eslint/no-empty-object-type
-export interface afterStoreDocumentPayload<Context = any> extends onStoreDocumentPayload<Context> {}
+export interface afterStoreDocumentPayload<Context = any>
+	extends onStoreDocumentPayload<Context> {}
 
 export interface onAwarenessUpdatePayload<Context = any> {
-	context: Context;
 	document: Document;
 	documentName: string;
 	instance: Hocuspocus;
-	requestHeaders: IncomingHttpHeaders;
-	requestParameters: URLSearchParams;
-	socketId: string;
+	transactionOrigin: unknown;
+	connection?: Connection<Context>;
 	added: number[];
 	updated: number[];
 	removed: number[];
@@ -387,7 +384,8 @@ export interface fetchPayload<Context = any> {
 	connectionConfig: ConnectionConfiguration;
 }
 
-export interface storePayload<Context = any> extends onStoreDocumentPayload<Context> {
+export interface storePayload<Context = any>
+	extends onStoreDocumentPayload<Context> {
 	state: Buffer;
 }