@lota-sdk/core 0.4.9 → 0.4.11

package/src/queues/queue-factory.ts

@@ -1,26 +1,35 @@
 import { Queue, Worker } from 'bullmq'
 import type { Job, JobsOptions, QueueOptions, WorkerOptions } from 'bullmq'
-import { Effect } from 'effect'
+import { Effect, Schema } from 'effect'
 import type IORedis from 'ioredis'
 
 import type { LotaLogger } from '../config/logger'
 import { serverLogger } from '../config/logger'
-import { getCurrentRuntime } from '../effect/runtime-ref'
-import { RedisServiceTag } from '../effect/services'
+import type { TrackedBullJobLike } from '../services/queue-job.service'
 import {
   attachWorkerEvents,
   createTracedWorkerProcessor,
   createWorkerShutdown,
   DEFAULT_JOB_RETENTION,
   registerShutdownSignals,
-  getQueueJobService,
 } from '../workers/worker-utils'
-import type { WorkerHandle } from '../workers/worker-utils'
+import type { QueueJobService, WorkerHandle } from '../workers/worker-utils'
 
-function getDefaultQueueConnectionProvider(): () => IORedis {
-  const redis = getCurrentRuntime().runSync(Effect.service(RedisServiceTag))
-  return () => redis.getConnectionForBullMQ()
-}
+class QueueFactoryError extends Schema.TaggedErrorClass<QueueFactoryError>()('@lota-sdk/core/QueueFactoryError', {
+  message: Schema.String,
+  cause: Schema.optional(Schema.Defect),
+}) {}
+
+type QueueShape<TJob> = Queue<TJob, unknown, string, TJob, unknown, string>
+type QueueMethod = 'add' | 'close' | 'remove' | 'removeDeduplicationKey' | 'removeJobScheduler'
+
+const queueMethodsThatWaitForClose = new Set<QueueMethod>([
+  'add',
+  'close',
+  'remove',
+  'removeDeduplicationKey',
+  'removeJobScheduler',
+])
 
 interface QueueFactoryConfigBase {
   name: string
@@ -32,7 +41,8 @@ interface QueueFactoryConfigBase {
   stalledInterval?: number
   maxStalledCount?: number
   defaultJobOptions?: JobsOptions
-  connectionProvider?: () => IORedis
+  connectionProvider: () => IORedis
+  queueJobService: QueueJobService
 }
 
 interface QueueFactoryConfigInline<TJob> extends QueueFactoryConfigBase {
@@ -48,32 +58,73 @@ interface QueueFactoryConfigFile extends QueueFactoryConfigBase {
 
 export type QueueFactoryConfig<TJob> = QueueFactoryConfigInline<TJob> | QueueFactoryConfigFile
 
+interface QueueFactoryWithDepsConfig<TJob, TDeps> extends QueueFactoryConfigBase {
+  prepare?: (deps: TDeps, job: Job<TJob>) => Promise<void>
+  processor: (deps: TDeps, job: Job<TJob>) => Promise<unknown>
+}
+
+interface QueueWorkerConfigInline<TJob> {
+  prepare?: (job: Job<TJob>) => Promise<void>
+  processor: (job: Job<TJob>) => Promise<unknown>
+  processorPath?: never
+}
+
+interface QueueWorkerConfigFile {
+  processor?: never
+  processorPath: string
+}
+
+type QueueWorkerConfig<TJob> = QueueWorkerConfigInline<TJob> | QueueWorkerConfigFile
+
+export interface QueueFactoryWorkerOptions {
+  registerSignals?: boolean
+  connectionProvider?: () => IORedis
+}
+
 export interface QueueFactory<TJob> {
-  getQueue: () => Queue<TJob, unknown, string, TJob, unknown, string>
+  getQueue: () => QueueShape<TJob>
   enqueue: (job: TJob, options?: JobsOptions) => Promise<void>
-  startWorker: (options?: { registerSignals?: boolean; connectionProvider?: () => IORedis }) => WorkerHandle
+  startWorker: (options?: QueueFactoryWorkerOptions) => 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
+export interface QueueFactoryWithDeps<TJob, TDeps> {
+  getQueue: () => QueueShape<TJob>
+  enqueue: (job: TJob, options?: JobsOptions) => Promise<void>
+  startWorker: (options: QueueFactoryWorkerOptions & { deps: TDeps }) => WorkerHandle
+}
+
+export function recordEnqueuedJobMetadata(params: {
+  queueName: string
+  job: Omit<TrackedBullJobLike, 'queueName'>
+  queueJobService: QueueJobService
+  logger?: LotaLogger
+}): Effect.Effect<string | undefined> {
+  const logger = params.logger ?? serverLogger
+  return params.queueJobService.recordEnqueued({ queueName: params.queueName, ...params.job }).pipe(
+    Effect.tapError((error) =>
+      Effect.sync(() => {
+        logger.error`Failed to persist queued job metadata (queue=${params.queueName}, job=${params.job.id}): ${error}`
+      }),
+    ),
+    Effect.orElseSucceed(() => undefined),
+  )
+}
+
+function createQueueFactoryRuntime<TJob>(config: QueueFactoryConfigBase): {
+  getQueue: () => QueueShape<TJob>
+  enqueue: (job: TJob, options?: JobsOptions) => Promise<void>
+  startWorker: (workerConfig: QueueWorkerConfig<TJob>, options?: QueueFactoryWorkerOptions) => WorkerHandle
+} {
+  type State = {
+    queue: QueueShape<TJob> | null
+    rawQueue: QueueShape<TJob> | null
     connection: IORedis | null
     pendingClose: Promise<void> | null
-  } = { queue: null, rawQueue: null, connection: null, pendingClose: null }
+  }
+
+  let state: State = { queue: null, rawQueue: null, connection: null, pendingClose: null }
 
-  const resolveConnectionProvider = (): (() => IORedis) =>
-    config.connectionProvider ?? getDefaultQueueConnectionProvider()
+  const resolveConnectionProvider = (): (() => IORedis) => config.connectionProvider
 
   const getConnection = (): IORedis => resolveConnectionProvider()()
 
@@ -90,7 +141,7 @@ export function createQueueFactory<TJob>(config: QueueFactoryConfig<TJob>): Queu
     })
   }
 
-  const wrapQueue = (queue: QueueShape): QueueShape =>
+  const wrapQueue = (queue: QueueShape<TJob>): QueueShape<TJob> =>
     new Proxy(queue, {
       get(target, property, receiver) {
         if (typeof property !== 'string') {
@@ -109,7 +160,7 @@ export function createQueueFactory<TJob>(config: QueueFactoryConfig<TJob>): Queu
       },
     })
 
-  const getQueue = (): QueueShape => {
+  const getQueue = (): QueueShape<TJob> => {
     const connection = getConnection()
     const isStale =
       state.rawQueue === null ||
@@ -146,36 +197,26 @@ export function createQueueFactory<TJob>(config: QueueFactoryConfig<TJob>): Queu
     return wrappedQueue
   }
 
-  const jobName = config.jobName
-  const toData = (job: TJob) => job
-
   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 queuedJob = yield* Effect.tryPromise({
+          try: () => getQueue().add(config.jobName, job, options),
+          catch: (cause) =>
+            new QueueFactoryError({ message: `Failed to enqueue job on queue "${config.name}".`, cause }),
+        })
+        yield* recordEnqueuedJobMetadata({
+          queueName: config.name,
+          job: queuedJob,
+          queueJobService: config.queueJobService,
+          logger: config.logger,
+        })
       }),
     )
 
   const startWorker = (
-    options: { registerSignals?: boolean; connectionProvider?: () => IORedis } = {},
+    workerConfig: QueueWorkerConfig<TJob>,
+    options: QueueFactoryWorkerOptions = {},
   ): WorkerHandle => {
     const { registerSignals = import.meta.main, connectionProvider } = options
     const logger = config.logger ?? serverLogger
@@ -188,23 +229,33 @@ export function createQueueFactory<TJob>(config: QueueFactoryConfig<TJob>): Queu
       ...(config.maxStalledCount !== undefined ? { maxStalledCount: config.maxStalledCount } : {}),
     }
 
-    const worker = config.processorPath
-      ? new Worker(config.name, config.processorPath, workerOptions)
+    const worker = workerConfig.processorPath
+      ? new Worker(config.name, workerConfig.processorPath, workerOptions)
       : new Worker(
           config.name,
-          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))
-              }),
-            ),
+          createTracedWorkerProcessor(
+            config.name,
+            (job) =>
+              Effect.runPromise(
+                Effect.gen(function* () {
+                  const inlineWorkerConfig = workerConfig as QueueWorkerConfigInline<TJob>
+                  const typedJob = job as Job<TJob>
+                  const prepare = inlineWorkerConfig.prepare
+                  if (prepare) {
+                    yield* Effect.tryPromise({
+                      try: () => prepare(typedJob),
+                      catch: (cause) =>
+                        new QueueFactoryError({ message: `Worker prepare failed for queue "${config.name}".`, cause }),
+                    })
+                  }
+                  return yield* Effect.tryPromise({
+                    try: () => inlineWorkerConfig.processor(typedJob),
+                    catch: (cause) =>
+                      new QueueFactoryError({ message: `Worker processor failed for queue "${config.name}".`, cause }),
+                  })
+                }),
+              ),
+            config.queueJobService,
           ),
           workerOptions,
         )
@@ -222,3 +273,29 @@ export function createQueueFactory<TJob>(config: QueueFactoryConfig<TJob>): Queu
 
   return { getQueue, enqueue, startWorker }
 }
+
+export function createQueueFactory<TJob>(config: QueueFactoryConfig<TJob>): QueueFactory<TJob> {
+  const runtime = createQueueFactoryRuntime<TJob>(config)
+  return {
+    getQueue: runtime.getQueue,
+    enqueue: runtime.enqueue,
+    startWorker: (options) => runtime.startWorker(config, options),
+  }
+}
+
+export function createQueueFactoryWithDeps<TJob, TDeps>(
+  config: QueueFactoryWithDepsConfig<TJob, TDeps>,
+): QueueFactoryWithDeps<TJob, TDeps> {
+  const runtime = createQueueFactoryRuntime<TJob>(config)
+  const prepare = config.prepare
+
+  return {
+    getQueue: runtime.getQueue,
+    enqueue: runtime.enqueue,
+    startWorker: ({ deps, ...options }) =>
+      runtime.startWorker(
+        { prepare: prepare ? (job) => prepare(deps, job) : undefined, processor: (job) => config.processor(deps, job) },
+        options,
+      ),
+  }
+}

package/src/queues/queues.service.ts

@@ -0,0 +1,61 @@
+import { Context, Effect, Layer } from 'effect'
+import type IORedis from 'ioredis'
+
+import { RedisServiceTag } from '../effect/services'
+import { QueueJobServiceTag } from '../services/queue-job.service'
+import { makeAutonomousJobQueueRuntime } from './autonomous-job.queue'
+import type { AutonomousJobQueueRuntime } from './autonomous-job.queue'
+import { makeContextCompactionQueueRuntime } from './context-compaction.queue'
+import type { ContextCompactionQueueRuntime } from './context-compaction.queue'
+import { makeDelayedNodePromotionQueueRuntime } from './delayed-node-promotion.queue'
+import type { DelayedNodePromotionQueueRuntime } from './delayed-node-promotion.queue'
+import { makeMemoryConsolidationQueueRuntime } from './memory-consolidation.queue'
+import type { MemoryConsolidationQueueRuntime } from './memory-consolidation.queue'
+import { makeOrganizationLearningQueueRuntime } from './organization-learning.queue'
+import type { OrganizationLearningQueueRuntime } from './organization-learning.queue'
+import { makePlanAgentHeartbeatQueueRuntime } from './plan-agent-heartbeat.queue'
+import type { PlanAgentHeartbeatQueueRuntime } from './plan-agent-heartbeat.queue'
+import { makePlanSchedulerQueueRuntime } from './plan-scheduler.queue'
+import type { PlanSchedulerQueueRuntime } from './plan-scheduler.queue'
+import { makePostChatMemoryQueueRuntime } from './post-chat-memory.queue'
+import type { PostChatMemoryQueueRuntime } from './post-chat-memory.queue'
+import { makeTitleGenerationQueueRuntime } from './title-generation.queue'
+import type { TitleGenerationQueueRuntime } from './title-generation.queue'
+
+export interface LotaQueuesRuntime {
+  autonomousJob: AutonomousJobQueueRuntime
+  contextCompaction: ContextCompactionQueueRuntime
+  delayedNodePromotion: DelayedNodePromotionQueueRuntime
+  memoryConsolidation: MemoryConsolidationQueueRuntime
+  organizationLearning: OrganizationLearningQueueRuntime
+  planAgentHeartbeat: PlanAgentHeartbeatQueueRuntime
+  planScheduler: PlanSchedulerQueueRuntime
+  postChatMemory: PostChatMemoryQueueRuntime
+  titleGeneration: TitleGenerationQueueRuntime
+}
+
+export class LotaQueuesServiceTag extends Context.Service<LotaQueuesServiceTag, LotaQueuesRuntime>()(
+  '@lota-sdk/core/LotaQueuesService',
+) {}
+
+export const LotaQueuesLive = Layer.effect(
+  LotaQueuesServiceTag,
+  Effect.gen(function* () {
+    const redis = yield* RedisServiceTag
+    const queueJobService = yield* QueueJobServiceTag
+    const connectionProvider: () => IORedis = () => redis.getConnectionForBullMQ()
+    const baseDeps = { connectionProvider, queueJobService }
+
+    return LotaQueuesServiceTag.of({
+      autonomousJob: makeAutonomousJobQueueRuntime(baseDeps),
+      contextCompaction: makeContextCompactionQueueRuntime(baseDeps),
+      delayedNodePromotion: makeDelayedNodePromotionQueueRuntime(baseDeps),
+      memoryConsolidation: makeMemoryConsolidationQueueRuntime(baseDeps),
+      organizationLearning: makeOrganizationLearningQueueRuntime(baseDeps),
+      planAgentHeartbeat: makePlanAgentHeartbeatQueueRuntime(baseDeps),
+      planScheduler: makePlanSchedulerQueueRuntime(baseDeps),
+      postChatMemory: makePostChatMemoryQueueRuntime(baseDeps),
+      titleGeneration: makeTitleGenerationQueueRuntime(baseDeps),
+    })
+  }),
+)

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

@@ -4,11 +4,12 @@ import type { Context } from 'effect'
 import type IORedis from 'ioredis'
 
 import { ensureRecordId } from '../db/record-id'
-import { ConfigurationError } from '../effect/errors'
 import { DatabaseServiceTag, RedisServiceTag } from '../effect/services'
+import { QueueJobServiceTag } from '../services/queue-job.service'
 import { RecentActivityTitleServiceTag } from '../services/recent-activity-title.service'
 import { ThreadTitleServiceTag } from '../services/thread/thread-title.service'
-import { createQueueFactory } from './queue-factory'
+import type { QueueJobService, WorkerHandle } from '../workers/worker-utils'
+import { createQueueFactoryWithDeps } from './queue-factory'
 import { runStandaloneQueueWorker } from './standalone-worker'
 
 export const TITLE_GENERATION_QUEUE = 'title-generation'
@@ -30,24 +31,14 @@ interface RecentActivityTitleRefinementJob {
 
 type TitleGenerationJob = ThreadTitleGenerationJob | RecentActivityTitleRefinementJob
 
-interface TitleGenerationQueueDeps {
+export interface TitleGenerationWorkerDeps {
   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()
+function processTitleGenerationJob(deps: TitleGenerationWorkerDeps, job: Job<TitleGenerationJob>): Promise<void> {
+  const { threadTitleService, recentActivityTitleService } = deps
   if (job.data.kind === 'thread-title') {
     return Effect.runPromise(
       Effect.asVoid(threadTitleService.generateAndPersistTitle(ensureRecordId(job.data.threadId), job.data.sourceText)),
@@ -57,44 +48,53 @@ function processTitleGenerationJob(job: Job<TitleGenerationJob>): Promise<void>
   return Effect.runPromise(Effect.asVoid(recentActivityTitleService.refineRecentActivityTitle(job.data.activityId)))
 }
 
-const titleGeneration = createQueueFactory<TitleGenerationJob>({
-  name: TITLE_GENERATION_QUEUE,
-  displayName: 'Title generation',
-  jobName: 'title-generation',
-  concurrency: 10,
-  lockDuration: 300_000,
-  defaultJobOptions: { attempts: 3, backoff: { type: 'exponential', delay: 2_000 } },
-  prepare: () => getDeps().databaseService.connect(),
-  processor: processTitleGenerationJob,
-})
-
-export function enqueueThreadTitleGeneration(job: Omit<ThreadTitleGenerationJob, 'kind'>) {
-  return titleGeneration.enqueue({ kind: 'thread-title', ...job }, { jobId: `thread-title:${job.threadId}` })
+export interface TitleGenerationQueueRuntime {
+  enqueueThreadTitleGeneration(job: Omit<ThreadTitleGenerationJob, 'kind'>): Promise<void>
+  enqueueRecentActivityTitleRefinement(job: Omit<RecentActivityTitleRefinementJob, 'kind'>): Promise<void>
+  startWorker(options: { registerSignals?: boolean; deps: TitleGenerationWorkerDeps }): WorkerHandle
 }
 
-export function enqueueRecentActivityTitleRefinement(job: Omit<RecentActivityTitleRefinementJob, 'kind'>) {
-  return titleGeneration.enqueue(
-    { kind: 'recent-activity-title', ...job },
-    { jobId: `recent-activity-title:${job.activityId}` },
-  )
+interface MakeTitleGenerationQueueRuntimeParams {
+  connectionProvider: () => IORedis
+  queueJobService: QueueJobService
 }
 
-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,
+export function makeTitleGenerationQueueRuntime(
+  params: MakeTitleGenerationQueueRuntimeParams,
+): TitleGenerationQueueRuntime {
+  const { connectionProvider, queueJobService } = params
+
+  const queue = createQueueFactoryWithDeps<TitleGenerationJob, TitleGenerationWorkerDeps>({
+    name: TITLE_GENERATION_QUEUE,
+    displayName: 'Title generation',
+    jobName: 'title-generation',
+    concurrency: 10,
+    lockDuration: 300_000,
+    defaultJobOptions: { attempts: 3, backoff: { type: 'exponential', delay: 2_000 } },
+    connectionProvider,
+    queueJobService,
+    prepare: ({ databaseService }) => Effect.runPromise(databaseService.connect()),
+    processor: processTitleGenerationJob,
   })
+
+  return {
+    enqueueThreadTitleGeneration: (job) =>
+      queue.enqueue({ kind: 'thread-title', ...job }, { jobId: `thread-title:${job.threadId}` }),
+    enqueueRecentActivityTitleRefinement: (job) =>
+      queue.enqueue({ kind: 'recent-activity-title', ...job }, { jobId: `recent-activity-title:${job.activityId}` }),
+    startWorker: (options) =>
+      queue.startWorker({ deps: options.deps, registerSignals: options.registerSignals, connectionProvider }),
+  }
 }
 
 runStandaloneQueueWorker((runtime) => {
   const resolve = <I, T>(tag: Context.Key<I, T>): T => runtime.runSync(Effect.service(tag))
-  startTitleGenerationWorker({
-    connectionProvider: () => resolve(RedisServiceTag).getConnectionForBullMQ(),
+  const redis = resolve(RedisServiceTag)
+  const titleGenerationQueue = makeTitleGenerationQueueRuntime({
+    connectionProvider: () => redis.getConnectionForBullMQ(),
+    queueJobService: resolve(QueueJobServiceTag),
+  })
+  titleGenerationQueue.startWorker({
     deps: {
       databaseService: resolve(DatabaseServiceTag),
       threadTitleService: resolve(ThreadTitleServiceTag),