fix(websockets): remove connection trace when it is gone (#953)

packages/websockets/src/Domain/UseCase/SendMessageToClient/SendMessageToClient.spec.ts

@@ -1,4 +1,8 @@
-import { ApiGatewayManagementApiClient } from '@aws-sdk/client-apigatewaymanagementapi'
+import {
+  ApiGatewayManagementApiClient,
+  ApiGatewayManagementApiServiceException,
+  GoneException,
+} from '@aws-sdk/client-apigatewaymanagementapi'
 import { WebSocketsConnectionRepositoryInterface } from '../../WebSockets/WebSocketsConnectionRepositoryInterface'
 import { SendMessageToClient } from './SendMessageToClient'
 import { Logger } from 'winston'
@@ -23,6 +27,7 @@ describe('SendMessageToClient', () => {
 
     webSocketsConnectionRepository = {} as jest.Mocked<WebSocketsConnectionRepositoryInterface>
     webSocketsConnectionRepository.findAllByUserUuid = jest.fn().mockResolvedValue([connection])
+    webSocketsConnectionRepository.removeConnection = jest.fn()
 
     apiGatewayManagementClient = {} as jest.Mocked<ApiGatewayManagementApiClient>
     apiGatewayManagementClient.send = jest.fn().mockResolvedValue({ $metadata: { httpStatusCode: 200 } })
@@ -30,6 +35,7 @@ describe('SendMessageToClient', () => {
     logger = {} as jest.Mocked<Logger>
     logger.debug = jest.fn()
     logger.error = jest.fn()
+    logger.info = jest.fn()
   })
 
   it('sends message to all connections for a user', async () => {
@@ -96,4 +102,26 @@ describe('SendMessageToClient', () => {
 
     expect(result.isFailed()).toBe(true)
   })
+
+  it('removes connection if it is gone', async () => {
+    apiGatewayManagementClient.send = jest.fn().mockRejectedValue(
+      new GoneException(
+        new ApiGatewayManagementApiServiceException({
+          name: 'test',
+          $fault: 'server',
+          $metadata: {},
+        }),
+      ),
+    )
+
+    const useCase = createUseCase()
+
+    const result = await useCase.execute({
+      userUuid: '00000000-0000-0000-0000-000000000000',
+      message: 'message',
+    })
+
+    expect(result.isFailed()).toBe(true)
+    expect(webSocketsConnectionRepository.removeConnection).toHaveBeenCalledTimes(1)
+  })
 })

packages/websockets/src/Domain/UseCase/SendMessageToClient/SendMessageToClient.ts

@@ -1,5 +1,9 @@
 import { Result, UseCaseInterface, Uuid } from '@standardnotes/domain-core'
-import { ApiGatewayManagementApiClient, PostToConnectionCommand } from '@aws-sdk/client-apigatewaymanagementapi'
+import {
+  ApiGatewayManagementApiClient,
+  GoneException,
+  PostToConnectionCommand,
+} from '@aws-sdk/client-apigatewaymanagementapi'
 import { Logger } from 'winston'
 
 import { SendMessageToClientDTO } from './SendMessageToClientDTO'
@@ -44,6 +48,12 @@ export class SendMessageToClient implements UseCaseInterface<void> {
           )
         }
       } catch (error) {
+        if (error instanceof GoneException) {
+          this.logger.info(`Connection ${connection.props.connectionId} for user ${userUuid.value} is gone. Removing.`)
+
+          await this.removeGoneConnection(connection.props.connectionId)
+        }
+
         return Result.fail(
           `Could not send message to connection ${connection.props.connectionId} for user ${
             userUuid.value
@@ -54,4 +64,8 @@ export class SendMessageToClient implements UseCaseInterface<void> {
 
     return Result.ok()
   }
+
+  private async removeGoneConnection(connectionId: string): Promise<void> {
+    await this.webSocketsConnectionRepository.removeConnection(connectionId)
+  }
 }