Merge pull request #15543 from kim-sung-jee/feature/websockets-manual-ack

feat(websockets): allow manual acknowledgement handling with @ack() decorator

Author: kamilmysliwiec

integration/websockets/e2e/gateway-ack.spec.ts

@@ -41,5 +41,38 @@ describe('WebSocketGateway (ack)', () => {
     );
   });
 
+  it('should handle manual ack for async operations when @Ack() is used (success case)', async () => {
+    app = await createNestApp(AckGateway);
+    await app.listen(3000);
+
+    ws = io('http://localhost:8080');
+    const payload = { shouldSucceed: true };
+
+    await new Promise<void>(resolve =>
+      ws.emit('manual-ack', payload, response => {
+        expect(response).to.eql({ status: 'success', data: payload });
+        resolve();
+      }),
+    );
+  });
+
+  it('should handle manual ack for async operations when @Ack() is used (error case)', async () => {
+    app = await createNestApp(AckGateway);
+    await app.listen(3000);
+
+    ws = io('http://localhost:8080');
+    const payload = { shouldSucceed: false };
+
+    await new Promise<void>(resolve =>
+      ws.emit('manual-ack', payload, response => {
+        expect(response).to.eql({
+          status: 'error',
+          message: 'Operation failed',
+        });
+        resolve();
+      }),
+    );
+  });
+
   afterEach(() => app.close());
 });

integration/websockets/src/ack.gateway.ts

@@ -1,9 +1,29 @@
-import { SubscribeMessage, WebSocketGateway } from '@nestjs/websockets';
+import {
+  Ack,
+  MessageBody,
+  SubscribeMessage,
+  WebSocketGateway,
+} from '@nestjs/websockets';
 
 @WebSocketGateway(8080)
 export class AckGateway {
   @SubscribeMessage('push')
   onPush() {
     return 'pong';
   }
+
+  @SubscribeMessage('manual-ack')
+  async handleManualAck(
+    @MessageBody() data: any,
+    @Ack() ack: (response: any) => void,
+  ) {
+    await new Promise(resolve => setTimeout(resolve, 20));
+
+    if (data.shouldSucceed) {
+      ack({ status: 'success', data });
+    } else {
+      ack({ status: 'error', message: 'Operation failed' });
+    }
+    return { status: 'ignored' };
+  }
 }

packages/common/enums/route-paramtypes.enum.ts

@@ -12,4 +12,5 @@ export enum RouteParamtypes {
   HOST = 10,
   IP = 11,
   RAW_BODY = 12,
+  ACK = 13,
 }

packages/common/interfaces/websockets/web-socket-adapter.interface.ts

@@ -6,6 +6,7 @@ import { Observable } from 'rxjs';
 export interface WsMessageHandler<T = string> {
   message: T;
   callback: (...args: any[]) => Observable<any> | Promise<any>;
+  isAckHandledManually: boolean;
 }
 
 /**

packages/platform-socket.io/adapters/io-adapter.ts

@@ -44,22 +44,24 @@ export class IoAdapter extends AbstractWsAdapter {
       first(),
     );
 
-    handlers.forEach(({ message, callback }) => {
+    handlers.forEach(({ message, callback, isAckHandledManually }) => {
       const source$ = fromEvent(socket, message).pipe(
         mergeMap((payload: any) => {
           const { data, ack } = this.mapPayload(payload);
           return transform(callback(data, ack)).pipe(
             filter((response: any) => !isNil(response)),
-            map((response: any) => [response, ack]),
+            map((response: any) => [response, ack, isAckHandledManually]),
           );
         }),
         takeUntil(disconnect$),
       );
-      source$.subscribe(([response, ack]) => {
+      source$.subscribe(([response, ack, isAckHandledManually]) => {
         if (response.event) {
           return socket.emit(response.event, response.data);
         }
-        isFunction(ack) && ack(response);
+        if (!isAckHandledManually && isFunction(ack)) {
+          ack(response);
+        }
       });
     });
   }

packages/websockets/context/ws-metadata-constants.ts

@@ -1,6 +1,7 @@
 import { WsParamtype } from '../enums/ws-paramtype.enum';
 
 export const DEFAULT_CALLBACK_METADATA = {
+  [`${WsParamtype.ACK}:2`]: { index: 2, data: undefined, pipes: [] },
   [`${WsParamtype.PAYLOAD}:1`]: { index: 1, data: undefined, pipes: [] },
   [`${WsParamtype.SOCKET}:0`]: { index: 0, data: undefined, pipes: [] },
 };

packages/websockets/decorators/ack.decorator.ts

@@ -0,0 +1,28 @@
+import { WsParamtype } from '../enums/ws-paramtype.enum';
+import { createPipesWsParamDecorator } from '../utils/param.utils';
+
+/**
+ * WebSockets `ack` parameter decorator.
+ * Extracts the `ack` callback function from the arguments of a ws event.
+ *
+ * This decorator signals to the framework that the `ack` callback will be
+ * handled manually within the method, preventing the framework from
+ * automatically sending an acknowledgement based on the return value.
+ *
+ * @example
+ * ```typescript
+ * @SubscribeMessage('events')
+ * onEvent(
+ *   @MessageBody() data: string,
+ *   @Ack() ack: (response: any) => void
+ * ) {
+ *   // Manually call the ack callback
+ *   ack({ status: 'ok' });
+ * }
+ * ```
+ *
+ * @publicApi
+ */
+export function Ack(): ParameterDecorator {
+  return createPipesWsParamDecorator(WsParamtype.ACK)();
+}

packages/websockets/decorators/index.ts

@@ -3,3 +3,4 @@ export * from './gateway-server.decorator';
 export * from './message-body.decorator';
 export * from './socket-gateway.decorator';
 export * from './subscribe-message.decorator';
+export * from './ack.decorator';

packages/websockets/enums/ws-paramtype.enum.ts

@@ -3,4 +3,5 @@ import { RouteParamtypes } from '@nestjs/common/enums/route-paramtypes.enum';
 export enum WsParamtype {
   SOCKET = RouteParamtypes.REQUEST,
   PAYLOAD = RouteParamtypes.BODY,
+  ACK = RouteParamtypes.ACK,
 }

packages/websockets/factories/ws-params-factory.ts

@@ -1,3 +1,4 @@
+import { isFunction } from '@nestjs/common/utils/shared.utils';
 import { WsParamtype } from '../enums/ws-paramtype.enum';
 
 export class WsParamsFactory {
@@ -14,6 +15,9 @@ export class WsParamsFactory {
         return args[0];
       case WsParamtype.PAYLOAD:
         return data ? args[1]?.[data] : args[1];
+      case WsParamtype.ACK: {
+        return args.find(arg => isFunction(arg));
+      }
       default:
         return null;
     }