From d0dd88fa6ab4242bf2ab956da7d54941a29cb203 Mon Sep 17 00:00:00 2001
From: didinele <didinele.dev@gmail.com>
Date: Wed, 3 Jul 2024 00:18:52 +0300
Subject: [PATCH] fix: worker event forwarding

---
 .../strategy/WorkerShardingStrategy.test.ts         | 13 ++++++++-----
 .../strategies/sharding/WorkerShardingStrategy.ts   |  2 +-
 packages/ws/src/utils/WorkerBootstrapper.ts         |  5 ++---
 3 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/packages/ws/__tests__/strategy/WorkerShardingStrategy.test.ts b/packages/ws/__tests__/strategy/WorkerShardingStrategy.test.ts
index 3277dcb3d..4a96408d3 100644
--- a/packages/ws/__tests__/strategy/WorkerShardingStrategy.test.ts
+++ b/packages/ws/__tests__/strategy/WorkerShardingStrategy.test.ts
@@ -91,7 +91,7 @@ vi.mock('node:worker_threads', async () => {
 							op: WorkerReceivePayloadOp.Event,
 							shardId: message.shardId,
 							event: WebSocketShardEvents.Dispatch,
-							data: memberChunkData,
+							data: [memberChunkData],
 						};
 						this.emit('message', response);
 
@@ -204,10 +204,13 @@ test('spawn, connect, send a message, session info, and destroy', async () => {
 	};
 	await manager.send(0, payload);
 	expect(mockSend).toHaveBeenCalledWith(0, payload);
-	expect(managerEmitSpy).toHaveBeenCalledWith(WebSocketShardEvents.Dispatch, {
-		...memberChunkData,
-		shardId: 0,
-	});
+	expect(managerEmitSpy).toHaveBeenCalledWith(
+		WebSocketShardEvents.Dispatch,
+		{
+			...memberChunkData,
+		},
+		0,
+	);
 	expect(mockRetrieveSessionInfo).toHaveBeenCalledWith(0);
 	expect(mockUpdateSessionInfo).toHaveBeenCalledWith(0, { ...sessionInfo, sequence: sessionInfo.sequence + 1 });
 
diff --git a/packages/ws/src/strategies/sharding/WorkerShardingStrategy.ts b/packages/ws/src/strategies/sharding/WorkerShardingStrategy.ts
index 78473916b..a3ed5450e 100644
--- a/packages/ws/src/strategies/sharding/WorkerShardingStrategy.ts
+++ b/packages/ws/src/strategies/sharding/WorkerShardingStrategy.ts
@@ -48,7 +48,7 @@ export enum WorkerReceivePayloadOp {
 
 export type WorkerReceivePayload =
 	// Can't seem to get a type-safe union based off of the event, so I'm sadly leaving data as any for now
-	| { data: any; event: WebSocketShardEvents; op: WorkerReceivePayloadOp.Event; shardId: number }
+	| { data: any[]; event: WebSocketShardEvents; op: WorkerReceivePayloadOp.Event; shardId: number }
 	| { nonce: number; op: WorkerReceivePayloadOp.CancelIdentify }
 	| { nonce: number; op: WorkerReceivePayloadOp.FetchStatusResponse; status: WebSocketShardStatus }
 	| { nonce: number; op: WorkerReceivePayloadOp.RetrieveSessionInfo; shardId: number }
diff --git a/packages/ws/src/utils/WorkerBootstrapper.ts b/packages/ws/src/utils/WorkerBootstrapper.ts
index dd101293e..86913f4f6 100644
--- a/packages/ws/src/utils/WorkerBootstrapper.ts
+++ b/packages/ws/src/utils/WorkerBootstrapper.ts
@@ -148,12 +148,11 @@ export class WorkerBootstrapper {
 		for (const shardId of this.data.shardIds) {
 			const shard = new WebSocketShard(new WorkerContextFetchingStrategy(this.data), shardId);
 			for (const event of options.forwardEvents ?? Object.values(WebSocketShardEvents)) {
-				// @ts-expect-error: Event types incompatible
-				shard.on(event, (data) => {
+				shard.on(event, (...args) => {
 					const payload: WorkerReceivePayload = {
 						op: WorkerReceivePayloadOp.Event,
 						event,
-						data,
+						data: args,
 						shardId,
 					};
 					parentPort!.postMessage(payload);