From df7651d3855379756e470efbb4cf23b6108b1f90 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A9s?= <andresgutgon@gmail.com>
Date: Thu, 19 Sep 2024 09:57:29 +0200
Subject: [PATCH] Enable connection and add console.logs to debug websockets in
 workers (#219)

---
 apps/infra/src/deployments/gateway.ts     |  2 ++
 apps/infra/src/deployments/web.ts         |  2 ++
 apps/infra/src/deployments/websockets.ts  |  2 ++
 apps/websockets/src/server.ts             |  4 ++++
 apps/workers/src/server.ts                | 11 +++++++++--
 packages/core/src/websockets/constants.ts |  1 +
 6 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/apps/infra/src/deployments/gateway.ts b/apps/infra/src/deployments/gateway.ts
index e22e724de..54e14c476 100644
--- a/apps/infra/src/deployments/gateway.ts
+++ b/apps/infra/src/deployments/gateway.ts
@@ -55,6 +55,7 @@ const taskDefinition = pulumi
         networkMode: 'awsvpc',
         requiresCompatibilities: ['FARGATE'],
         executionRoleArn: ecsTaskExecutionRole,
+        taskRoleArn: ecsTaskExecutionRole,
         containerDefinitions: JSON.stringify([
           {
             name: containerName,
@@ -128,6 +129,7 @@ new aws.ecs.Service('LatitudeLLMGateway', {
   desiredCount: 2,
   launchType: 'FARGATE',
   forceNewDeployment: true,
+  enableExecuteCommand: true,
   networkConfiguration: {
     subnets: privateSubnets.ids,
     assignPublicIp: false,
diff --git a/apps/infra/src/deployments/web.ts b/apps/infra/src/deployments/web.ts
index 5058bb849..8f9a6226d 100644
--- a/apps/infra/src/deployments/web.ts
+++ b/apps/infra/src/deployments/web.ts
@@ -87,6 +87,7 @@ const taskDefinition = pulumi
         networkMode: 'awsvpc',
         requiresCompatibilities: ['FARGATE'],
         executionRoleArn: ecsTaskExecutionRole,
+        taskRoleArn: ecsTaskExecutionRole,
         containerDefinitions: JSON.stringify([
           {
             name: containerName,
@@ -175,6 +176,7 @@ new aws.ecs.Service('LatitudeLLMApp', {
   desiredCount: 2,
   launchType: 'FARGATE',
   forceNewDeployment: true,
+  enableExecuteCommand: true,
   networkConfiguration: {
     subnets: privateSubnets.ids,
     assignPublicIp: false,
diff --git a/apps/infra/src/deployments/websockets.ts b/apps/infra/src/deployments/websockets.ts
index 722a7f923..3f4218027 100644
--- a/apps/infra/src/deployments/websockets.ts
+++ b/apps/infra/src/deployments/websockets.ts
@@ -55,6 +55,7 @@ const taskDefinition = pulumi
         networkMode: 'awsvpc',
         requiresCompatibilities: ['FARGATE'],
         executionRoleArn: ecsTaskExecutionRole,
+        taskRoleArn: ecsTaskExecutionRole,
         containerDefinitions: JSON.stringify([
           {
             name: containerName,
@@ -128,6 +129,7 @@ new aws.ecs.Service('LatitudeLLMWebsockets', {
   desiredCount: 1,
   launchType: 'FARGATE',
   forceNewDeployment: true,
+  enableExecuteCommand: true,
   networkConfiguration: {
     subnets: privateSubnets.ids,
     assignPublicIp: false,
diff --git a/apps/websockets/src/server.ts b/apps/websockets/src/server.ts
index 3d704ddc6..7565accc2 100644
--- a/apps/websockets/src/server.ts
+++ b/apps/websockets/src/server.ts
@@ -130,6 +130,10 @@ workers.use(async (socket, next) => {
 workers.on('connection', (socket) => {
   console.log('DEBUG: Worker connected')
 
+  socket.on('pingFromWorkers', () => {
+    console.log('DEBUG: Ping from workers')
+  })
+
   socket.on('evaluationStatus', (args) => {
     console.log('DEBUG: Evaluation STATUS %s', JSON.stringify(args))
     const { workspaceId, data } = args
diff --git a/apps/workers/src/server.ts b/apps/workers/src/server.ts
index 1cd998dbf..84029dfae 100644
--- a/apps/workers/src/server.ts
+++ b/apps/workers/src/server.ts
@@ -1,5 +1,7 @@
 import http from 'http'
 
+import { WebsocketClient } from '@latitude-data/core/websockets/workers'
+
 import { captureException, captureMessage } from './utils/sentry'
 import startWorkers from './workers'
 
@@ -8,8 +10,13 @@ const workers = startWorkers()
 console.log('Workers started')
 
 const port = process.env.WORKERS_PORT || 3002
-const server = http.createServer((req, res) => {
-  if (req.url === '/health' && req.method === 'GET') {
+const server = http.createServer(async (req, res) => {
+  const websockets = await WebsocketClient.getSocket()
+
+  if (req.url === '/ping' && req.method === 'GET') {
+    websockets.emit('pingFromWorkers')
+    res.end(JSON.stringify({ status: 'OK', message: 'Pong' }))
+  } else if (req.url === '/health' && req.method === 'GET') {
     res.writeHead(200, { 'Content-Type': 'application/json' })
     res.end(JSON.stringify({ status: 'OK', message: 'Workers are healthy' }))
   } else {
diff --git a/packages/core/src/websockets/constants.ts b/packages/core/src/websockets/constants.ts
index e74a6b0cd..9b35f83b0 100644
--- a/packages/core/src/websockets/constants.ts
+++ b/packages/core/src/websockets/constants.ts
@@ -50,4 +50,5 @@ export type WorkersClientToServerEvents = {
     data: EvaluationStatusArgs
     workspaceId: number
   }) => void
+  pingFromWorkers: () => void
 }