From 5b4dbd541e39a0f04f9954712e27f44a95ef16b2 Mon Sep 17 00:00:00 2001
From: Vlad Frangu <me@vladfrangu.dev>
Date: Fri, 12 Dec 2025 15:01:25 +0200
Subject: [PATCH] feat: zstd streaming support (#10758)

* feat(ws): zstd streaming support

* chore: naming of the mode

* fix: remove `@ts-expect-error`

---------

Co-authored-by: Almeida <github@almeidx.dev>
---
 packages/ws/src/utils/constants.ts   |  2 ++
 packages/ws/src/ws/WebSocketShard.ts | 42 ++++++++++++++++++++++++----
 2 files changed, 38 insertions(+), 6 deletions(-)

diff --git a/packages/ws/src/utils/constants.ts b/packages/ws/src/utils/constants.ts
index 413daecf4..35f3da732 100644
--- a/packages/ws/src/utils/constants.ts
+++ b/packages/ws/src/utils/constants.ts
@@ -20,6 +20,7 @@ export enum Encoding {
 export enum CompressionMethod {
 	ZlibNative,
 	ZlibSync,
+	ZstdNative,
 }
 
 export const DefaultDeviceProperty = `@discordjs/ws [VI]{{inject}}[/VI]` as `@discordjs/ws ${string}`;
@@ -29,6 +30,7 @@ const getDefaultSessionStore = lazy(() => new Collection<number, SessionInfo | n
 export const CompressionParameterMap = {
 	[CompressionMethod.ZlibNative]: 'zlib-stream',
 	[CompressionMethod.ZlibSync]: 'zlib-stream',
+	[CompressionMethod.ZstdNative]: 'zstd-stream',
 } as const satisfies Record<CompressionMethod, string>;
 
 /**
diff --git a/packages/ws/src/ws/WebSocketShard.ts b/packages/ws/src/ws/WebSocketShard.ts
index 3d9a61ff3..6b6d0b7f8 100644
--- a/packages/ws/src/ws/WebSocketShard.ts
+++ b/packages/ws/src/ws/WebSocketShard.ts
@@ -217,7 +217,7 @@ export class WebSocketShard extends AsyncEventEmitter<WebSocketShardEventsMap> {
 
 						this.nativeInflate = inflate;
 					} else {
-						console.warn('WebSocketShard: Compression is set to native but node:zlib is not available.');
+						console.warn('WebSocketShard: Compression is set to native zlib but node:zlib is not available.');
 						params.delete('compress');
 					}
 
@@ -238,6 +238,34 @@ export class WebSocketShard extends AsyncEventEmitter<WebSocketShardEventsMap> {
 
 					break;
 				}
+
+				case CompressionMethod.ZstdNative: {
+					const zlib = await getNativeZlib();
+					if (zlib && 'createZstdDecompress' in zlib) {
+						this.inflateBuffer = [];
+
+						const inflate = zlib.createZstdDecompress({
+							chunkSize: 65_535,
+						}) as nativeZlib.Inflate;
+
+						inflate.on('data', (chunk) => {
+							this.inflateBuffer.push(chunk);
+						});
+
+						inflate.on('error', (error) => {
+							this.emit(WebSocketShardEvents.Error, error);
+						});
+
+						this.nativeInflate = inflate;
+					} else {
+						console.warn(
+							'WebSocketShard: Compression is set to native zstd but node:zlib is not available or your node version does not support zstd decompression.',
+						);
+						params.delete('compress');
+					}
+
+					break;
+				}
 			}
 		}
 
@@ -628,12 +656,14 @@ export class WebSocketShard extends AsyncEventEmitter<WebSocketShardEventsMap> {
 
 		// Deal with transport compression
 		if (this.transportCompressionEnabled) {
+			// Each WS message received is a full gateway message for zstd streaming, but for zlib it's chunked
 			const flush =
-				decompressable.length >= 4 &&
-				decompressable.at(-4) === 0x00 &&
-				decompressable.at(-3) === 0x00 &&
-				decompressable.at(-2) === 0xff &&
-				decompressable.at(-1) === 0xff;
+				this.strategy.options.compression === CompressionMethod.ZstdNative ||
+				(decompressable.length >= 4 &&
+					decompressable.at(-4) === 0x00 &&
+					decompressable.at(-3) === 0x00 &&
+					decompressable.at(-2) === 0xff &&
+					decompressable.at(-1) === 0xff);
 
 			if (this.nativeInflate) {
 				const doneWriting = new Promise<void>((resolve) => {