feat(WebsocketManager): retroactive token setting (#10418)

Co-authored-by: kodiakhq[bot] <49736102+kodiakhq[bot]@users.noreply.github.com>

Author: didinele

packages/ws/src/strategies/context/IContextFetchingStrategy.ts

@@ -33,7 +33,6 @@ export interface IContextFetchingStrategy {
 }
 
 export async function managerToFetchingStrategyOptions(manager: WebSocketManager): Promise<FetchingStrategyOptions> {
-	/* eslint-disable @typescript-eslint/unbound-method */
 	const {
 		buildIdentifyThrottler,
 		buildStrategy,
@@ -44,10 +43,10 @@ export async function managerToFetchingStrategyOptions(manager: WebSocketManager
 		rest,
 		...managerOptions
 	} = manager.options;
-	/* eslint-enable @typescript-eslint/unbound-method */
 
 	return {
 		...managerOptions,
+		token: manager.token,
 		gatewayInformation: await manager.fetchGatewayInformation(),
 		shardCount: await manager.getShardCount(),
 	};

packages/ws/src/utils/constants.ts

@@ -68,7 +68,7 @@ export const DefaultWebSocketManagerOptions = {
 	handshakeTimeout: 30_000,
 	helloTimeout: 60_000,
 	readyTimeout: 15_000,
-} as const satisfies OptionalWebSocketManagerOptions;
+} as const satisfies Omit<OptionalWebSocketManagerOptions, 'token'>;
 
 export const ImportantGatewayOpcodes = new Set([
 	GatewayOpcodes.Heartbeat,

packages/ws/src/ws/WebSocketManager.ts

@@ -67,10 +67,6 @@ export interface RequiredWebSocketManagerOptions {
 	 * The REST instance to use for fetching gateway information
 	 */
 	rest: REST;
-	/**
-	 * The token to use for identifying with the gateway
-	 */
-	token: string;
 }
 
 /**
@@ -172,6 +168,12 @@ export interface OptionalWebSocketManagerOptions {
 	 * ```
 	 */
 	shardIds: number[] | ShardRange | null;
+	/**
+	 * The token to use for identifying with the gateway
+	 *
+	 * If not provided, the token must be set using {@link WebSocketManager.setToken}
+	 */
+	token: string;
 	/**
 	 * Function used to store session information for a given shard
 	 */
@@ -211,10 +213,12 @@ export interface ManagerShardEventsMap {
 }
 
 export class WebSocketManager extends AsyncEventEmitter<ManagerShardEventsMap> implements AsyncDisposable {
+	#token: string | null = null;
+
 	/**
 	 * The options being used by this manager
 	 */
-	public readonly options: WebSocketManagerOptions;
+	public readonly options: Omit<WebSocketManagerOptions, 'token'>;
 
 	/**
 	 * Internal cache for a GET /gateway/bot result
@@ -236,10 +240,26 @@ export class WebSocketManager extends AsyncEventEmitter<ManagerShardEventsMap> i
 	 */
 	private readonly strategy: IShardingStrategy;
 
+	/**
+	 * Gets the token set for this manager. If no token is set, an error is thrown.
+	 * To set the token, use {@link WebSocketManager.setToken} or pass it in the options.
+	 *
+	 * @remarks
+	 * This getter is mostly used to pass the token to the sharding strategy internally, there's not much reason to use it.
+	 */
+	public get token(): string {
+		if (!this.#token) {
+			throw new Error('Token has not been set');
+		}
+
+		return this.#token;
+	}
+
 	public constructor(options: CreateWebSocketManagerOptions) {
 		super();
 		this.options = { ...DefaultWebSocketManagerOptions, ...options };
 		this.strategy = this.options.buildStrategy(this);
+		this.#token = options.token ?? null;
 	}
 
 	/**
@@ -334,6 +354,14 @@ export class WebSocketManager extends AsyncEventEmitter<ManagerShardEventsMap> i
 		await this.strategy.connect();
 	}
 
+	public setToken(token: string): void {
+		if (this.#token) {
+			throw new Error('Token has already been set');
+		}
+
+		this.#token = token;
+	}
+
 	public destroy(options?: Omit<WebSocketShardDestroyOptions, 'recover'>) {
 		return this.strategy.destroy(options);
 	}