Enable connection and add console.logs to debug websockets in workers (…

…#219)
latitude-dev · Sep 19, 2024 · df7651d · df7651d
1 parent df38a01
commit df7651d
Show file tree

Hide file tree

Showing 6 changed files with 20 additions and 2 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -50,4 +50,5 @@ export type WorkersClientToServerEvents = {
     data: EvaluationStatusArgs
     workspaceId: number
   }) => void
+  pingFromWorkers: () => void
 }