From 8cf190e516f6a4c6e1ee4ea295803e5384aa17c0 Mon Sep 17 00:00:00 2001
From: Mohamed Boudra <boudra.moha@gmail.com>
Date: Fri, 17 Jul 2026 20:01:06 +0000
Subject: [PATCH] fix(server): tighten relay backpressure accounting

Count pending encryption and transport backlog together, and release socket leases before forced termination.
---
 .../server/src/server/websocket-server.ts     |  6 ++---
 .../websocket/encrypted-relay-socket.test.ts  | 23 +++++++++++++++++++
 .../websocket/encrypted-relay-socket.ts       |  4 ++--
 3 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/packages/server/src/server/websocket-server.ts b/packages/server/src/server/websocket-server.ts
index 3a152908c..af945a744 100644
--- a/packages/server/src/server/websocket-server.ts
+++ b/packages/server/src/server/websocket-server.ts
@@ -1008,7 +1008,7 @@ export class VoiceAssistantWebSocketServer {
   }
 
   private sendMessageToSockets(sockets: Iterable<WebSocketLike>, message: WSOutboundMessage): void {
-    const writableSockets = [...sockets].filter((ws) => this.canBufferOutbound(ws, 0));
+    const writableSockets = [...sockets].filter((ws) => this.ensureOutboundCapacity(ws, 0));
     if (writableSockets.length === 0) {
       return;
     }
@@ -1054,7 +1054,7 @@ export class VoiceAssistantWebSocketServer {
     }
   }
 
-  private canBufferOutbound(ws: WebSocketLike, frameBytes: number): boolean {
+  private ensureOutboundCapacity(ws: WebSocketLike, frameBytes: number): boolean {
     if (ws.readyState !== 1) return false;
     if (physicalSocketHasCapacity(ws, frameBytes)) return true;
 
@@ -1077,8 +1077,8 @@ export class VoiceAssistantWebSocketServer {
 
   private closePhysicalSocket(params: ClosePhysicalSocketParams): void {
     const { ws, code, reason, logMessage, logFields } = params;
+    this.applicationSocketLease.release(ws);
     if (ws.readyState !== 1) {
-      this.applicationSocketLease.release(ws);
       return;
     }
     const identity = this.socketIdentities.get(ws);
diff --git a/packages/server/src/server/websocket/encrypted-relay-socket.test.ts b/packages/server/src/server/websocket/encrypted-relay-socket.test.ts
index 91a4c006f..3add059fc 100644
--- a/packages/server/src/server/websocket/encrypted-relay-socket.test.ts
+++ b/packages/server/src/server/websocket/encrypted-relay-socket.test.ts
@@ -77,6 +77,29 @@ test("underlying relay backpressure rejects binary before encryption and termina
   expect(terminations).toBe(1);
 });
 
+test("pending encryption and underlying relay backpressure share one hard bound", () => {
+  const channel = new BlockingChannel();
+  let transportBufferedAmount = 3 * 1024 * 1024;
+  let terminations = 0;
+  const socket = createEncryptedRelaySocket({
+    channel,
+    emitter: new EventEmitter(),
+    getTransportBufferedAmount: () => transportBufferedAmount,
+    terminateTransport: () => {
+      terminations += 1;
+    },
+  });
+
+  socket.send(new Uint8Array(3 * 1024 * 1024));
+  expect(channel.sent).toHaveLength(1);
+
+  transportBufferedAmount = 4 * 1024 * 1024;
+  socket.send(new Uint8Array(1));
+
+  expect(channel.sent).toHaveLength(1);
+  expect(terminations).toBe(1);
+});
+
 test("explicit encrypted-socket termination forcibly terminates the relay transport", () => {
   const channel = new BlockingChannel();
   let terminations = 0;
diff --git a/packages/server/src/server/websocket/encrypted-relay-socket.ts b/packages/server/src/server/websocket/encrypted-relay-socket.ts
index d5c4c48fc..4b00f8200 100644
--- a/packages/server/src/server/websocket/encrypted-relay-socket.ts
+++ b/packages/server/src/server/websocket/encrypted-relay-socket.ts
@@ -53,13 +53,13 @@ export function createEncryptedRelaySocket(params: {
       return readyState;
     },
     get bufferedAmount() {
-      return Math.max(pendingEncryptedBytes, getTransportBufferedAmount() ?? 0);
+      return pendingEncryptedBytes + (getTransportBufferedAmount() ?? 0);
     },
     send: (data) => {
       if (readyState !== 1) return;
       const outbound = normalizeRelaySendPayload(data);
       const outboundBytes = encryptedRelayFrameByteLength(outbound);
-      const queuedBytes = Math.max(pendingEncryptedBytes, getTransportBufferedAmount() ?? 0);
+      const queuedBytes = pendingEncryptedBytes + (getTransportBufferedAmount() ?? 0);
       if (queuedBytes + outboundBytes > MAX_PHYSICAL_SOCKET_BUFFERED_BYTES) {
         terminate();
         return;