chore: wait for redis connection before returning the connection instance

apps/web/src/actions/evaluations/runBatch.ts

@@ -95,7 +95,7 @@ export const runBatchEvaluationAction = withDataset
     const evaluations = await evaluationsRepo
       .filterById(input.evaluationIds)
       .then((r) => r.unwrap())
-    const queues = setupJobs()
+    const queues = await setupJobs()
     evaluations.forEach((evaluation) => {
       const batchId = `evaluation:${evaluation.id}:${nanoid(5)}`
       queues.defaultQueue.jobs.enqueueRunBatchEvaluationJob({

apps/workers/src/server.ts

@@ -3,7 +3,7 @@ import http from 'http'
 import { captureException, captureMessage } from './utils/sentry'
 import startWorkers from './workers'
 
-const workers = startWorkers()
+const workers = await startWorkers()
 
 console.log('Workers started')
 

apps/workers/src/workers/index.ts

@@ -14,12 +14,13 @@ const WORKER_OPTS = {
 }
 const WORKERS = [defaultWorker]
 
-export default function startWorkers() {
+export default async function startWorkers() {
+  const connection = await queues()
   return WORKERS.flatMap((w) =>
     w.queues.map((q) => {
       const worker = new Worker(q, w.processor, {
         ...WORKER_OPTS,
-        connection: queues(),
+        connection,
       })
 
       worker.on('error', (error: Error) => {

docker-compose.yml

@@ -40,7 +40,20 @@ services:
       dockerfile: apps/gateway/docker/Dockerfile
     profiles: [building]
     environment:
+      - CACHE_HOST=redis
+      - DATABASE_URL=postgresql://latitude:secret@db:5432/latitude
+      - DEFAULT_PROJECT_ID=1
+      - DEFAULT_PROVIDER_API_KEY=test
+      - [email protected]
+      - LATITUDE_DOMAIN=web
+      - LATITUDE_URL=http://web:8080
       - NODE_ENV=production
+      - QUEUE_HOST=redis
+      - WEBSOCKETS_SERVER=websockets
+      - WEBSOCKET_REFRESH_SECRET_TOKEN_KEY=test
+      - WEBSOCKET_SECRET_TOKEN_KEY=test
+      - WORKERS_WEBSOCKET_SECRET_TOKEN=test
+      - HOSTNAME=localhost
     depends_on:
       - db
       - redis
@@ -51,7 +64,19 @@ services:
       dockerfile: apps/workers/docker/Dockerfile
     profiles: [building]
     environment:
+      - CACHE_HOST=redis
+      - DATABASE_URL=postgresql://latitude:secret@db:5432/latitude
+      - DEFAULT_PROJECT_ID=1
+      - DEFAULT_PROVIDER_API_KEY=test
+      - [email protected]
+      - LATITUDE_DOMAIN=web
+      - LATITUDE_URL=http://web:8080
       - NODE_ENV=production
+      - QUEUE_HOST=redis
+      - WEBSOCKETS_SERVER=websockets
+      - WEBSOCKET_REFRESH_SECRET_TOKEN_KEY=test
+      - WEBSOCKET_SECRET_TOKEN_KEY=test
+      - WORKERS_WEBSOCKET_SECRET_TOKEN=test
     depends_on:
       - db
       - redis

packages/core/src/cache/index.ts

@@ -5,10 +5,10 @@ import { buildRedisConnection } from '../redis'
 
 let connection: Redis
 
-export const cache = () => {
+export const cache = async () => {
   if (connection) return connection
 
-  connection = buildRedisConnection({
+  connection = await buildRedisConnection({
     host: env.CACHE_HOST,
     port: env.CACHE_PORT,
   })

packages/core/src/events/publisher.ts

@@ -11,8 +11,8 @@ export const publisher = {
       handlers.map((handler) => handler({ data: event })),
     )
   },
-  publishLater: (event: LatitudeEvent) => {
-    const queues = setupJobs()
+  publishLater: async (event: LatitudeEvent) => {
+    const queues = await setupJobs()
 
     queues.eventsQueue.jobs.enqueueCreateEventJob(event)
     queues.eventsQueue.jobs.enqueuePublishEventJob(event)

packages/core/src/queues/index.ts

@@ -5,10 +5,10 @@ import { buildRedisConnection } from '../redis'
 
 let connection: Redis
 
-export const queues = () => {
+export const queues = async () => {
   if (connection) return connection
 
-  connection = buildRedisConnection({
+  connection = await buildRedisConnection({
     host: env.QUEUE_HOST,
     port: env.QUEUE_PORT,
     password: env.QUEUE_PASSWORD,

packages/core/src/redis/index.ts

@@ -5,5 +5,11 @@ export function buildRedisConnection({
   host,
   ...opts
 }: Omit<RedisOptions, 'port' & 'host'> & { host: string; port: number }) {
-  return new Redis(port, host, opts)
+  return new Promise<Redis>((resolve) => {
+    const instance = new Redis(port, host, opts)
+
+    instance.on('connect', () => {
+      resolve(instance)
+    })
+  })
 }

packages/core/src/services/ai/index.ts

@@ -128,9 +128,8 @@ export async function ai({
       )
       providerLogUuid = providerLog.uuid
     } else {
-      await setupJobs().defaultQueue.jobs.enqueueCreateProviderLogJob(
-        payload.data,
-      )
+      const queues = await setupJobs()
+      queues.defaultQueue.jobs.enqueueCreateProviderLogJob(payload.data)
     }
 
     onFinish?.({ ...event, providerLogUuid })
@@ -175,7 +174,7 @@ const checkDefaultProviderUsage = async ({
   workspace: Workspace
 }) => {
   if (provider.token === env.DEFAULT_PROVIDER_API_KEY) {
-    const c = cache()
+    const c = await cache()
     const value = await c.incr(
       `workspace:${workspace.id}:defaultProviderRunCount`,
     )

packages/jobs/src/index.ts

@@ -2,10 +2,11 @@ import { setupQueues } from './queues'
 
 export { Worker } from 'bullmq'
 
-let queues: ReturnType<typeof setupQueues>
-export function setupJobs() {
+let queues: Awaited<ReturnType<typeof setupQueues>>
+
+export async function setupJobs() {
   if (queues) return queues
-  queues = setupQueues()
+  queues = await setupQueues()
 
   return queues
 }

packages/jobs/src/job-definitions/batchEvaluations/runBatchEvaluationJob.ts

@@ -66,15 +66,15 @@ export const runBatchEvaluationJob = async (
     )
   })
 
-  const progressTracker = new ProgressTracker(queues(), batchId)
+  const progressTracker = new ProgressTracker(await queues(), batchId)
   const firstAttempt = job.attemptsMade === 0
 
   if (firstAttempt) {
     await progressTracker.initializeProgress(parameters.length)
   }
 
   const progress = await progressTracker.getProgress()
-  const jobs = setupJobs()
+  const jobs = await setupJobs()
 
   if (firstAttempt && parameters.length > 0) {
     websockets.emit('evaluationStatus', {

packages/jobs/src/job-definitions/batchEvaluations/runDocumentJob.test.ts

@@ -1,4 +1,4 @@
-import { LogSources } from '@latitude-data/core/browser'
+import { LogSources, Providers } from '@latitude-data/core/browser'
 import * as factories from '@latitude-data/core/factories'
 import { Result } from '@latitude-data/core/lib/Result'
 import { runDocumentAtCommit } from '@latitude-data/core/services/commits/runDocumentAtCommit'
@@ -48,7 +48,10 @@ describe('runDocumentJob', () => {
 
     // Create necessary resources using factories
     const setup = await factories.createProject({
-      documents: { 'test-doc': 'Test content' },
+      providers: [{ type: Providers.OpenAI, name: 'Latitude' }],
+      documents: {
+        'test-doc': factories.helpers.createPrompt({ provider: 'Latitude' }),
+      },
     })
     workspace = setup.workspace
     project = setup.project

packages/jobs/src/job-definitions/batchEvaluations/runDocumentJob.ts

@@ -34,7 +34,7 @@ export const runDocumentJob = async (job: Job<RunDocumentJobData>) => {
     batchId,
   } = job.data
 
-  const progressTracker = new ProgressTracker(queues(), batchId)
+  const progressTracker = new ProgressTracker(await queues(), batchId)
 
   try {
     const workspace = await unsafelyFindWorkspace(workspaceId)
@@ -58,7 +58,7 @@ export const runDocumentJob = async (job: Job<RunDocumentJobData>) => {
 
     await result.response
 
-    const queues = setupJobs()
+    const queues = await setupJobs()
 
     // Enqueue the evaluation job
     await queues.defaultQueue.jobs.enqueueRunEvaluationJob({

packages/jobs/src/job-definitions/batchEvaluations/runEvaluationJob.ts

@@ -29,7 +29,7 @@ export const runEvaluationJob = async (job: Job<RunEvaluationJobData>) => {
     evaluationId,
   } = job.data
 
-  const progressTracker = new ProgressTracker(queues(), batchId)
+  const progressTracker = new ProgressTracker(await queues(), batchId)
   const documentLogsScope = new DocumentLogsRepository(workspaceId)
   const evaluationsScope = new EvaluationsRepository(workspaceId)
   try {

packages/jobs/src/job-definitions/events/publishEventJob.ts

@@ -11,9 +11,8 @@ export const publishEventJob = async (job: Job<LatitudeEvent>) => {
   const handlers = EventHandlers[event.type]
   if (!handlers?.length) return
 
+  const queues = await setupJobs()
   handlers.forEach((handler) => {
-    const queues = setupJobs()
-
     queues.eventsQueue.queue.add(handler.name, event)
   })
 }

packages/jobs/src/queues/index.ts

@@ -1,5 +1,6 @@
 import { queues } from '@latitude-data/core/queues'
 import { Job, JobsOptions, Queue, QueueEvents } from 'bullmq'
+import Redis from 'ioredis'
 
 import { Jobs, Queues, QUEUES } from '../constants'
 import { JobDefinition } from '../job-definitions'
@@ -30,11 +31,12 @@ export const DEFAULT_JOB_OPTIONS: JobsOptions = {
 function setupQueue({
   name,
   jobs,
+  connection,
 }: {
   name: Queues
   jobs: readonly QueueJob[]
+  connection: Redis
 }) {
-  const connection = queues()
   const queue = new Queue(name, {
     connection,
     defaultJobOptions: DEFAULT_JOB_OPTIONS,
@@ -60,14 +62,16 @@ function setupQueue({
 
 type QueueJob = (typeof QUEUES)[keyof typeof QUEUES]['jobs'][number]
 
-export function setupQueues() {
+export async function setupQueues() {
+  const connection = await queues()
+
   return Object.entries(QUEUES).reduce<{
     [K in keyof typeof QUEUES]: ReturnType<typeof setupQueue>
   }>(
     (acc, [name, { jobs }]) => {
       return {
         ...acc,
-        [name]: setupQueue({ name: name as Queues, jobs }),
+        [name]: setupQueue({ name: name as Queues, jobs, connection }),
       }
     },
     {} as { [K in keyof typeof QUEUES]: ReturnType<typeof setupQueue> },

packages/sdks/typescript/src/addMessage.test.ts

@@ -13,6 +13,7 @@ import {
 } from 'vitest'
 
 import { LatitudeSdk } from './index'
+import { parseSSE } from './utils/parseSSE'
 
 const encoder = new TextEncoder()
 let latitudeApiKey = 'fake-api-key'
@@ -60,7 +61,7 @@ describe('addMessage', () => {
             const stream = new ReadableStream({
               start(controller) {
                 CHUNKS.forEach((chunk, index) => {
-                  controller.enqueue(encoder.encode(JSON.stringify(chunk)))
+                  controller.enqueue(encoder.encode(chunk))
                   if (index === CHUNKS.length - 1) {
                     controller.close()
                   }
@@ -86,48 +87,14 @@ describe('addMessage', () => {
       })
 
       CHUNKS.forEach((chunk, index) => {
-        expect(onMessageMock).toHaveBeenNthCalledWith(index + 1, chunk)
+        expect(onMessageMock).toHaveBeenNthCalledWith(
+          index + 1,
+          parseSSE(chunk)!.data,
+        )
       })
     }),
   )
 
-  it(
-    'send on Error callback when chunk is not a valid JSON',
-    server.boundary(async () => {
-      const onMessageMock = vi.fn()
-      const onErrorMock = vi.fn()
-      server.use(
-        http.post(
-          'http://localhost:8787/api/v1/chats/add-message',
-          async () => {
-            const stream = new ReadableStream({
-              start(controller) {
-                controller.enqueue(encoder.encode('invalid json'))
-                controller.close()
-              },
-            })
-
-            return new HttpResponse(stream, {
-              headers: {
-                'Content-Type': 'text/plain',
-              },
-            })
-          },
-        ),
-      )
-      await SDK.addMessages({
-        params: {
-          messages: [],
-          documentLogUuid: 'fake-document-log-uuid',
-          source: LogSources.Playground,
-        },
-        onMessage: onMessageMock,
-        onError: onErrorMock,
-      })
-      expect(onErrorMock).toHaveBeenCalledWith(expect.any(Error))
-    }),
-  )
-
   it(
     'sends all message onFinish callback and final response',
     server.boundary(async () => {
@@ -139,7 +106,7 @@ describe('addMessage', () => {
             const stream = new ReadableStream({
               start(controller) {
                 CHUNKS.forEach((chunk, index) => {
-                  controller.enqueue(encoder.encode(JSON.stringify(chunk)))
+                  controller.enqueue(encoder.encode(chunk))
                   if (index === CHUNKS.length - 1) {
                     controller.close()
                   }

packages/sdks/typescript/src/index.ts

@@ -61,6 +61,7 @@ export class LatitudeSdk {
       params: { projectId, commitUuid },
       body: { documentPath, parameters, source },
     })
+
     return this.handleStreamChainResponse({
       response,
       onMessage,