diff --git a/packages/websockets/src/Domain/UseCase/SendMessageToClient/SendMessageToClient.spec.ts b/packages/websockets/src/Domain/UseCase/SendMessageToClient/SendMessageToClient.spec.ts index 0296044c5..db9b4d84e 100644 --- a/packages/websockets/src/Domain/UseCase/SendMessageToClient/SendMessageToClient.spec.ts +++ b/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 webSocketsConnectionRepository.findAllByUserUuid = jest.fn().mockResolvedValue([connection]) + webSocketsConnectionRepository.removeConnection = jest.fn() apiGatewayManagementClient = {} as jest.Mocked apiGatewayManagementClient.send = jest.fn().mockResolvedValue({ $metadata: { httpStatusCode: 200 } }) @@ -30,6 +35,7 @@ describe('SendMessageToClient', () => { logger = {} as jest.Mocked 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) + }) }) diff --git a/packages/websockets/src/Domain/UseCase/SendMessageToClient/SendMessageToClient.ts b/packages/websockets/src/Domain/UseCase/SendMessageToClient/SendMessageToClient.ts index d8ce37d70..440a5b909 100644 --- a/packages/websockets/src/Domain/UseCase/SendMessageToClient/SendMessageToClient.ts +++ b/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 { ) } } 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 { return Result.ok() } + + private async removeGoneConnection(connectionId: string): Promise { + await this.webSocketsConnectionRepository.removeConnection(connectionId) + } }