@lota-sdk/core 0.4.8 → 0.4.9

package/src/queues/post-chat-memory.queue.ts

@@ -1,9 +1,14 @@
-type OrganizationOnboardStatus = string
 import type { Job } from 'bullmq'
+import { Effect } from 'effect'
+import type { Context } from 'effect'
+import type IORedis from 'ioredis'
 
-import { databaseService } from '../db/service'
-import { memoryService } from '../services/memory.service'
+import type { MaybeAwaitableService } from '../effect/awaitable-effect'
+import { ConfigurationError } from '../effect/errors'
+import { DatabaseServiceTag, RedisServiceTag } from '../effect/services'
+import { MemoryServiceTag } from '../services/memory/memory.service'
 import { createQueueFactory } from './queue-factory'
+import { runStandaloneQueueWorker } from './standalone-worker'
 
 interface PostChatMemoryMessage {
   role: 'user' | 'agent'
@@ -17,7 +22,7 @@ interface PostChatMemoryExtractionJob {
   sourceId: string
   source?: string
   sourceMetadata?: Record<string, unknown>
-  onboardStatus?: OrganizationOnboardStatus
+  onboardStatus?: string
   userMessage: string
   historyMessages: PostChatMemoryMessage[]
   agentMessages: Array<{ content: string; agentName?: string }>
@@ -25,8 +30,23 @@ interface PostChatMemoryExtractionJob {
   attachmentContext?: string
 }
 
-async function processPostChatMemoryJob(job: Job<PostChatMemoryExtractionJob>): Promise<void> {
-  await databaseService.connect()
+interface PostChatMemoryQueueDeps {
+  databaseService: Context.Service.Shape<typeof DatabaseServiceTag>
+  memoryService: MaybeAwaitableService<Context.Service.Shape<typeof MemoryServiceTag>>
+}
+
+let _deps: PostChatMemoryQueueDeps | null = null
+function getDeps(): PostChatMemoryQueueDeps {
+  if (!_deps)
+    throw new ConfigurationError({
+      message: 'Post-chat memory queue is not configured. Initialize the runtime before starting the worker.',
+      key: 'queue-deps',
+    })
+  return _deps
+}
+
+function processPostChatMemoryJob(job: Job<PostChatMemoryExtractionJob>): Promise<void> {
+  const { memoryService } = getDeps()
 
   const data = job.data
   const userMessage = data.userMessage.trim()
@@ -37,7 +57,9 @@ async function processPostChatMemoryJob(job: Job<PostChatMemoryExtractionJob>):
     }))
     .filter((item) => item.content.length > 0)
 
-  if (!userMessage || agentMessages.length === 0) return
+  if (!userMessage || agentMessages.length === 0) {
+    return Promise.resolve()
+  }
 
   const joinedOutput = agentMessages
     .map((item) => (item.agentName ? `[${item.agentName}] ${item.content}` : item.content))
@@ -51,20 +73,24 @@ async function processPostChatMemoryJob(job: Job<PostChatMemoryExtractionJob>):
     ),
   ]
 
-  await memoryService.addConversationMemories({
-    orgId: data.orgId,
-    input: userMessage,
-    output: joinedOutput,
-    sourceId: data.sourceId,
-    source: data.source,
-    sourceMetadata: data.sourceMetadata,
-    onboardStatus: data.onboardStatus,
-    ...(uniqueAgentNames.length > 0 ? { agentName: uniqueAgentNames[0] } : {}),
-    historyMessages: data.historyMessages,
-    memoryBlock: data.memoryBlock,
-    attachmentContext: data.attachmentContext,
-    agentNames: uniqueAgentNames,
-  })
+  return Effect.runPromise(
+    Effect.asVoid(
+      memoryService.addConversationMemories({
+        orgId: data.orgId,
+        input: userMessage,
+        output: joinedOutput,
+        sourceId: data.sourceId,
+        source: data.source,
+        sourceMetadata: data.sourceMetadata,
+        onboardStatus: data.onboardStatus,
+        ...(uniqueAgentNames.length > 0 ? { agentName: uniqueAgentNames[0] } : {}),
+        historyMessages: data.historyMessages,
+        memoryBlock: data.memoryBlock,
+        attachmentContext: data.attachmentContext,
+        agentNames: uniqueAgentNames,
+      }),
+    ),
+  )
 }
 
 const postChatMemory = createQueueFactory<PostChatMemoryExtractionJob>({
@@ -76,14 +102,30 @@ const postChatMemory = createQueueFactory<PostChatMemoryExtractionJob>({
   maxStalledCount: 10,
   stalledInterval: 120_000,
   defaultJobOptions: { attempts: 3, backoff: { type: 'exponential', delay: 2_000 } },
+  prepare: () => getDeps().databaseService.connect(),
   processor: processPostChatMemoryJob,
 })
 
 export function enqueuePostChatMemory(job: PostChatMemoryExtractionJob, options?: { dedupeKey?: string }) {
   return postChatMemory.enqueue(job, options?.dedupeKey ? { jobId: options.dedupeKey } : undefined)
 }
-export const startPostChatMemoryWorker = postChatMemory.startWorker
 
-if (import.meta.main) {
-  startPostChatMemoryWorker()
+export function startPostChatMemoryWorker(options: {
+  registerSignals?: boolean
+  connectionProvider: () => IORedis
+  deps: PostChatMemoryQueueDeps
+}) {
+  _deps = options.deps
+  return postChatMemory.startWorker({
+    registerSignals: options.registerSignals,
+    connectionProvider: options.connectionProvider,
+  })
 }
+
+runStandaloneQueueWorker((runtime) => {
+  const resolve = <I, T>(tag: Context.Key<I, T>): T => runtime.runSync(Effect.service(tag))
+  startPostChatMemoryWorker({
+    connectionProvider: () => resolve(RedisServiceTag).getConnectionForBullMQ(),
+    deps: { databaseService: resolve(DatabaseServiceTag), memoryService: resolve(MemoryServiceTag) },
+  })
+})

package/src/queues/queue-factory.ts

@@ -1,24 +1,33 @@
 import { Queue, Worker } from 'bullmq'
 import type { Job, JobsOptions, QueueOptions, WorkerOptions } from 'bullmq'
+import { Effect } from 'effect'
 import type IORedis from 'ioredis'
 
+import type { LotaLogger } from '../config/logger'
 import { serverLogger } from '../config/logger'
-import { getRedisConnectionForBullMQ } from '../redis'
-import { queueJobService } from '../services/queue-job.service'
+import { getCurrentRuntime } from '../effect/runtime-ref'
+import { RedisServiceTag } from '../effect/services'
 import {
   attachWorkerEvents,
   createTracedWorkerProcessor,
   createWorkerShutdown,
   DEFAULT_JOB_RETENTION,
   registerShutdownSignals,
+  getQueueJobService,
 } from '../workers/worker-utils'
 import type { WorkerHandle } from '../workers/worker-utils'
 
+function getDefaultQueueConnectionProvider(): () => IORedis {
+  const redis = getCurrentRuntime().runSync(Effect.service(RedisServiceTag))
+  return () => redis.getConnectionForBullMQ()
+}
+
 interface QueueFactoryConfigBase {
   name: string
   displayName: string
   jobName: string
   concurrency: number
+  logger?: LotaLogger
   lockDuration?: number
   stalledInterval?: number
   maxStalledCount?: number
@@ -27,6 +36,7 @@ interface QueueFactoryConfigBase {
 }
 
 interface QueueFactoryConfigInline<TJob> extends QueueFactoryConfigBase {
+  prepare?: (job: Job<TJob>) => Promise<void>
   processor: (job: Job<TJob>) => Promise<unknown>
   processorPath?: never
 }
@@ -41,70 +51,137 @@ export type QueueFactoryConfig<TJob> = QueueFactoryConfigInline<TJob> | QueueFac
 export interface QueueFactory<TJob> {
   getQueue: () => Queue<TJob, unknown, string, TJob, unknown, string>
   enqueue: (job: TJob, options?: JobsOptions) => Promise<void>
-  startWorker: (options?: { registerSignals?: boolean }) => WorkerHandle
+  startWorker: (options?: { registerSignals?: boolean; connectionProvider?: () => IORedis }) => WorkerHandle
 }
 
 export function createQueueFactory<TJob>(config: QueueFactoryConfig<TJob>): QueueFactory<TJob> {
   type QueueShape = Queue<TJob, unknown, string, TJob, unknown, string>
+  type QueueMethod = 'add' | 'close' | 'remove' | 'removeDeduplicationKey' | 'removeJobScheduler'
+  const queueMethodsThatWaitForClose = new Set<QueueMethod>([
+    'add',
+    'close',
+    'remove',
+    'removeDeduplicationKey',
+    'removeJobScheduler',
+  ])
+
+  let state: {
+    queue: QueueShape | null
+    rawQueue: QueueShape | null
+    connection: IORedis | null
+    pendingClose: Promise<void> | null
+  } = { queue: null, rawQueue: null, connection: null, pendingClose: null }
+
+  const resolveConnectionProvider = (): (() => IORedis) =>
+    config.connectionProvider ?? getDefaultQueueConnectionProvider()
+
+  const getConnection = (): IORedis => resolveConnectionProvider()()
+
+  const waitForPendingClose = (): Promise<void> => {
+    const pendingClose = state.pendingClose
+    if (!pendingClose) {
+      return Promise.resolve()
+    }
 
-  let _queue: QueueShape | null = null
-  let _queueConnection: IORedis | null = null
+    return pendingClose.finally(() => {
+      if (state.pendingClose === pendingClose) {
+        state.pendingClose = null
+      }
+    })
+  }
 
-  const getConnection = (): IORedis => config.connectionProvider?.() ?? getRedisConnectionForBullMQ()
+  const wrapQueue = (queue: QueueShape): QueueShape =>
+    new Proxy(queue, {
+      get(target, property, receiver) {
+        if (typeof property !== 'string') {
+          return Reflect.get(target, property, receiver) as unknown
+        }
+
+        const value = (target as unknown as Record<string, unknown>)[property]
+        if (typeof value !== 'function' || !queueMethodsThatWaitForClose.has(property as QueueMethod)) {
+          return value
+        }
+
+        return (...args: Array<unknown>) =>
+          waitForPendingClose().then(() =>
+            Reflect.apply(value as (...methodArgs: Array<unknown>) => unknown, target, args),
+          )
+      },
+    })
 
   const getQueue = (): QueueShape => {
     const connection = getConnection()
-    const shouldRecreateQueue =
-      _queue === null ||
-      _queueConnection === null ||
-      _queueConnection !== connection ||
-      _queueConnection.status === 'close' ||
-      _queueConnection.status === 'end'
-
-    if (shouldRecreateQueue) {
-      if (_queue) {
-        void _queue.close().catch((error: unknown) => {
-          serverLogger.warn`Failed to close stale ${config.displayName} queue: ${error}`
-        })
-      }
-
-      _queue = new Queue<TJob, unknown, string, TJob, unknown, string>(config.name, {
-        connection: connection as QueueOptions['connection'],
-        defaultJobOptions: { ...DEFAULT_JOB_RETENTION, ...config.defaultJobOptions },
-      })
-      _queueConnection = connection
+    const isStale =
+      state.rawQueue === null ||
+      state.queue === null ||
+      state.connection === null ||
+      state.connection !== connection ||
+      state.connection.status === 'close' ||
+      state.connection.status === 'end'
+
+    if (!isStale && state.queue) {
+      return state.queue
     }
-    if (_queue === null) {
-      throw new Error(`Failed to initialize queue: ${config.name}`)
+
+    if (state.rawQueue) {
+      const staleQueue = state.rawQueue
+      const previousPendingClose = state.pendingClose ?? Promise.resolve()
+      state.pendingClose = previousPendingClose
+        .catch(() => undefined)
+        .then(() =>
+          staleQueue.close().catch((error: unknown) => {
+            serverLogger.warn`Failed to close stale ${config.displayName} queue: ${error}`
+          }),
+        )
+        .then(() => undefined)
     }
-    return _queue
+
+    const queue = new Queue<TJob, unknown, string, TJob, unknown, string>(config.name, {
+      connection: connection as QueueOptions['connection'],
+      defaultJobOptions: { ...DEFAULT_JOB_RETENTION, ...config.defaultJobOptions },
+    })
+    const wrappedQueue = wrapQueue(queue)
+
+    state = { queue: wrappedQueue, rawQueue: queue, connection, pendingClose: state.pendingClose }
+    return wrappedQueue
   }
 
   const jobName = config.jobName
   const toData = (job: TJob) => job
 
-  const enqueue = async (job: TJob, options?: JobsOptions): Promise<void> => {
-    const queuedJob = await getQueue().add(jobName, toData(job), options)
-    try {
-      await queueJobService.recordEnqueued({
-        queueName: config.name,
-        id: queuedJob.id,
-        name: queuedJob.name,
-        data: queuedJob.data,
-        opts: queuedJob.opts,
-        attemptsMade: queuedJob.attemptsMade,
-        timestamp: queuedJob.timestamp,
-      })
-    } catch (error) {
-      serverLogger.error`Failed to persist queued job metadata (queue=${config.name}, job=${queuedJob.id}): ${error}`
-    }
-  }
-
-  const startWorker = (options: { registerSignals?: boolean } = {}): WorkerHandle => {
-    const { registerSignals = import.meta.main } = options
+  const enqueue = (job: TJob, options?: JobsOptions): Promise<void> =>
+    Effect.runPromise(
+      Effect.gen(function* () {
+        const queuedJob = yield* Effect.tryPromise(() => getQueue().add(jobName, toData(job), options))
+        yield* getQueueJobService()
+          .recordEnqueued({
+            queueName: config.name,
+            id: queuedJob.id,
+            name: queuedJob.name,
+            data: queuedJob.data,
+            opts: queuedJob.opts,
+            attemptsMade: queuedJob.attemptsMade,
+            timestamp: queuedJob.timestamp,
+          })
+          .pipe(
+            Effect.tapError((error) =>
+              Effect.sync(() => {
+                serverLogger.error`Failed to persist queued job metadata (queue=${config.name}, job=${queuedJob.id}): ${error}`
+              }),
+            ),
+            Effect.orElseSucceed(() => undefined),
+          )
+      }),
+    )
+
+  const startWorker = (
+    options: { registerSignals?: boolean; connectionProvider?: () => IORedis } = {},
+  ): WorkerHandle => {
+    const { registerSignals = import.meta.main, connectionProvider } = options
+    const logger = config.logger ?? serverLogger
 
     const workerOptions: WorkerOptions = {
-      connection: getConnection() as QueueOptions['connection'],
+      connection: (connectionProvider ?? resolveConnectionProvider())() as QueueOptions['connection'],
       concurrency: config.concurrency,
       ...(config.lockDuration !== undefined ? { lockDuration: config.lockDuration } : {}),
       ...(config.stalledInterval !== undefined ? { stalledInterval: config.stalledInterval } : {}),
@@ -115,16 +192,29 @@ export function createQueueFactory<TJob>(config: QueueFactoryConfig<TJob>): Queu
       ? new Worker(config.name, config.processorPath, workerOptions)
       : new Worker(
           config.name,
-          createTracedWorkerProcessor(config.name, (config as QueueFactoryConfigInline<TJob>).processor),
+          createTracedWorkerProcessor(config.name, (job) =>
+            Effect.runPromise(
+              Effect.gen(function* () {
+                const inlineConfig = config as QueueFactoryConfigInline<TJob>
+                const typedJob = job as Job<TJob>
+                const prepare = inlineConfig.prepare
+                if (prepare) {
+                  yield* Effect.tryPromise(() => prepare(typedJob))
+                }
+                const processor = inlineConfig.processor
+                return yield* Effect.tryPromise(() => processor(typedJob))
+              }),
+            ),
+          ),
           workerOptions,
         )
 
-    attachWorkerEvents(worker, config.displayName, serverLogger)
+    attachWorkerEvents(worker, config.displayName, logger)
 
-    const shutdown = createWorkerShutdown(worker, config.displayName, serverLogger)
+    const shutdown = createWorkerShutdown(worker, config.displayName, logger)
 
     if (registerSignals) {
-      registerShutdownSignals({ name: config.displayName, shutdown, logger: serverLogger })
+      registerShutdownSignals({ name: config.displayName, shutdown, logger })
     }
 
     return { worker, shutdown }

package/src/queues/standalone-worker.ts

@@ -0,0 +1,39 @@
+import type { ManagedRuntime } from 'effect'
+import { Schema, Effect } from 'effect'
+
+import { serverLogger } from '../config/logger'
+import { initializeSandboxedWorkerRuntime } from '../workers/bootstrap'
+
+class StandaloneQueueWorkerError extends Schema.TaggedErrorClass<StandaloneQueueWorkerError>()(
+  'StandaloneQueueWorkerError',
+  { message: Schema.String, cause: Schema.Defect },
+) {}
+
+function toStandaloneQueueWorkerError(cause: unknown): StandaloneQueueWorkerError {
+  return new StandaloneQueueWorkerError({ message: cause instanceof Error ? cause.message : String(cause), cause })
+}
+
+// eslint-disable-next-line typescript-eslint/no-explicit-any -- wildcard for host-provided ManagedRuntime
+export function runStandaloneQueueWorker(start: (runtime: ManagedRuntime.ManagedRuntime<any, any>) => void): void {
+  if (!import.meta.main) {
+    return
+  }
+
+  void Effect.runFork(
+    Effect.gen(function* () {
+      const runtime = yield* Effect.tryPromise({
+        try: () => initializeSandboxedWorkerRuntime(),
+        catch: (cause) => toStandaloneQueueWorkerError(cause),
+      })
+
+      yield* Effect.sync(() => start(runtime))
+    }).pipe(
+      Effect.catchTag('StandaloneQueueWorkerError', (error) =>
+        Effect.sync(() => {
+          serverLogger.error`Standalone queue worker failed: ${error.message}`
+          process.exit(1)
+        }),
+      ),
+    ),
+  )
+}

package/src/queues/title-generation.queue.ts

@@ -1,10 +1,15 @@
 import type { Job } from 'bullmq'
+import { Effect } from 'effect'
+import type { Context } from 'effect'
+import type IORedis from 'ioredis'
 
 import { ensureRecordId } from '../db/record-id'
-import { databaseService } from '../db/service'
-import { recentActivityTitleService } from '../services/recent-activity-title.service'
-import { threadTitleService } from '../services/thread-title.service'
+import { ConfigurationError } from '../effect/errors'
+import { DatabaseServiceTag, RedisServiceTag } from '../effect/services'
+import { RecentActivityTitleServiceTag } from '../services/recent-activity-title.service'
+import { ThreadTitleServiceTag } from '../services/thread/thread-title.service'
 import { createQueueFactory } from './queue-factory'
+import { runStandaloneQueueWorker } from './standalone-worker'
 
 export const TITLE_GENERATION_QUEUE = 'title-generation'
 
@@ -25,14 +30,31 @@ interface RecentActivityTitleRefinementJob {
 
 type TitleGenerationJob = ThreadTitleGenerationJob | RecentActivityTitleRefinementJob
 
-async function processTitleGenerationJob(job: Job<TitleGenerationJob>): Promise<void> {
-  await databaseService.connect()
+interface TitleGenerationQueueDeps {
+  databaseService: Context.Service.Shape<typeof DatabaseServiceTag>
+  threadTitleService: Context.Service.Shape<typeof ThreadTitleServiceTag>
+  recentActivityTitleService: Context.Service.Shape<typeof RecentActivityTitleServiceTag>
+}
+
+let _deps: TitleGenerationQueueDeps | null = null
+function getDeps(): TitleGenerationQueueDeps {
+  if (!_deps)
+    throw new ConfigurationError({
+      message: 'Title generation queue is not configured. Initialize the runtime before starting the worker.',
+      key: 'queue-deps',
+    })
+  return _deps
+}
+
+function processTitleGenerationJob(job: Job<TitleGenerationJob>): Promise<void> {
+  const { threadTitleService, recentActivityTitleService } = getDeps()
   if (job.data.kind === 'thread-title') {
-    await threadTitleService.generateAndPersistTitle(ensureRecordId(job.data.threadId), job.data.sourceText)
-    return
+    return Effect.runPromise(
+      Effect.asVoid(threadTitleService.generateAndPersistTitle(ensureRecordId(job.data.threadId), job.data.sourceText)),
+    )
   }
 
-  await recentActivityTitleService.refineRecentActivityTitle(job.data.activityId)
+  return Effect.runPromise(Effect.asVoid(recentActivityTitleService.refineRecentActivityTitle(job.data.activityId)))
 }
 
 const titleGeneration = createQueueFactory<TitleGenerationJob>({
@@ -42,6 +64,7 @@ const titleGeneration = createQueueFactory<TitleGenerationJob>({
   concurrency: 10,
   lockDuration: 300_000,
   defaultJobOptions: { attempts: 3, backoff: { type: 'exponential', delay: 2_000 } },
+  prepare: () => getDeps().databaseService.connect(),
   processor: processTitleGenerationJob,
 })
 
@@ -56,4 +79,26 @@ export function enqueueRecentActivityTitleRefinement(job: Omit<RecentActivityTit
   )
 }
 
-export const startTitleGenerationWorker = titleGeneration.startWorker
+export function startTitleGenerationWorker(options: {
+  registerSignals?: boolean
+  connectionProvider: () => IORedis
+  deps: TitleGenerationQueueDeps
+}): ReturnType<typeof titleGeneration.startWorker> {
+  _deps = options.deps
+  return titleGeneration.startWorker({
+    registerSignals: options.registerSignals,
+    connectionProvider: options.connectionProvider,
+  })
+}
+
+runStandaloneQueueWorker((runtime) => {
+  const resolve = <I, T>(tag: Context.Key<I, T>): T => runtime.runSync(Effect.service(tag))
+  startTitleGenerationWorker({
+    connectionProvider: () => resolve(RedisServiceTag).getConnectionForBullMQ(),
+    deps: {
+      databaseService: resolve(DatabaseServiceTag),
+      threadTitleService: resolve(ThreadTitleServiceTag),
+      recentActivityTitleService: resolve(RecentActivityTitleServiceTag),
+    },
+  })
+})