From 23c9c2a7105b00270ba440204da343c8823624f7 Mon Sep 17 00:00:00 2001
From: MARCROCK22 <marcos22dev@gmail.com>
Date: Wed, 7 May 2025 18:41:46 -0400
Subject: [PATCH] feat: workerClient.calculateShardId

---
 src/client/workerclient.ts             | 17 +++++++++++++++--
 src/websocket/discord/workermanager.ts |  8 +++++++-
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/src/client/workerclient.ts b/src/client/workerclient.ts
index 2519151..4b54e0a 100644
--- a/src/client/workerclient.ts
+++ b/src/client/workerclient.ts
@@ -1,7 +1,14 @@
 import { type UUID, randomUUID } from 'node:crypto';
 import { ApiHandler, Logger } from '..';
 import { WorkerAdapter } from '../cache';
-import { type DeepPartial, LogLevels, type MakeRequired, type When, lazyLoadPackage } from '../common';
+import {
+	type DeepPartial,
+	LogLevels,
+	type MakeRequired,
+	type When,
+	calculateShardId,
+	lazyLoadPackage,
+} from '../common';
 import { EventHandler } from '../events';
 import type { GatewayDispatchPayload, GatewaySendPayload } from '../types';
 import { Shard, type ShardManagerOptions, ShardSocketCloseCodes, type WorkerData, properties } from '../websocket';
@@ -377,12 +384,18 @@ export class WorkerClient<Ready extends boolean = boolean> extends BaseClient {
 							return this.onPacket(packet, shardId);
 						};
 					}
+					workerData.totalShards = data.totalShards;
+					workerData.shards = [...this.shards.keys()];
 					this.resharding.clear();
 				}
 				break;
 		}
 	}
 
+	calculateShardId(guildId: string) {
+		return calculateShardId(guildId, this.workerData.totalShards);
+	}
+
 	private generateNonce(): UUID {
 		const uuid = randomUUID();
 		if (this.promises.has(uuid)) return this.generateNonce();
@@ -553,7 +566,7 @@ export function generateShardInfo(shard: Shard): WorkerShardInfo {
 	};
 }
 
-interface WorkerClientOptions extends BaseClientOptions {
+export interface WorkerClientOptions extends BaseClientOptions {
 	commands?: NonNullable<Client['options']>['commands'];
 	handlePayload?: ShardManagerOptions['handlePayload'];
 	gateway?: ClientOptions['gateway'];
diff --git a/src/websocket/discord/workermanager.ts b/src/websocket/discord/workermanager.ts
index d2577fd..456589f 100644
--- a/src/websocket/discord/workermanager.ts
+++ b/src/websocket/discord/workermanager.ts
@@ -288,6 +288,7 @@ export class WorkerManager extends Map<
 						for (const [id] of this.entries()) {
 							this.postMessage(id, {
 								type: 'CONNECT_ALL_SHARDS_RESHARDING',
+								totalShards: this.options.totalShards,
 							} satisfies ConnnectAllShardsResharding);
 						}
 						this.forEach(w => {
@@ -656,7 +657,12 @@ export type ManagerSpawnShardsResharding = CreateManagerMessage<
 	Pick<ShardOptions, 'info' | 'properties' | 'compress'>
 >;
 export type DisconnectAllShardsResharding = CreateManagerMessage<'DISCONNECT_ALL_SHARDS_RESHARDING'>;
-export type ConnnectAllShardsResharding = CreateManagerMessage<'CONNECT_ALL_SHARDS_RESHARDING'>;
+export type ConnnectAllShardsResharding = CreateManagerMessage<
+	'CONNECT_ALL_SHARDS_RESHARDING',
+	{
+		totalShards: number;
+	}
+>;
 export type ManagerSendPayload = CreateManagerMessage<
 	'SEND_PAYLOAD',
 	GatewaySendPayload & { shardId: number; nonce: string }