@lota-sdk/core 0.4.28 → 0.4.29

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@lota-sdk/core",
-  "version": "0.4.28",
+  "version": "0.4.29",
   "files": [
     "src",
     "infrastructure/schema"
@@ -32,7 +32,7 @@
     "@ai-sdk/provider": "^3.0.9",
     "@chat-adapter/slack": "^4.26.0",
     "@chat-adapter/state-ioredis": "^4.26.0",
-    "@lota-sdk/shared": "0.4.28",
+    "@lota-sdk/shared": "0.4.29",
     "@mendable/firecrawl-js": "^4.20.0",
     "@surrealdb/node": "^3.0.3",
     "ai": "^6.0.170",

package/src/ai-gateway/ai-gateway.ts

@@ -38,8 +38,8 @@ class AiGatewayStreamAttemptTag extends Context.Service<
 
 const EXPECTED_GATEWAY_KEY_PREFIX = 'sk-bf-'
 const AI_GATEWAY_VIRTUAL_KEY_HEADER = 'x-bf-vk'
-const AI_GATEWAY_TIMEOUT_MS = 180_000
-const AI_GATEWAY_STREAM_IDLE_TIMEOUT_MS = 30_000
+const AI_GATEWAY_TIMEOUT_MS = 360_000
+const AI_GATEWAY_STREAM_IDLE_TIMEOUT_MS = 180_000
 const AI_GATEWAY_MAX_RETRIES = 4
 const AI_GATEWAY_MAX_RETRY_DELAY_MS = 15_000
 const OPENAI_RESPONSES_PROVIDER_ID = 'openai.responses'
@@ -748,13 +748,60 @@ function isOpenRouterModel(modelId: string): boolean {
   return modelId.trim().toLowerCase().startsWith('openrouter/')
 }
 
+function mergeAbortSignals(signals: Array<AbortSignal | undefined>): { signal?: AbortSignal; cleanup: () => void } {
+  const activeSignals = signals.filter((signal): signal is AbortSignal => Boolean(signal))
+  if (activeSignals.length === 0) return { cleanup: () => undefined }
+  if (activeSignals.length === 1) return { signal: activeSignals[0], cleanup: () => undefined }
+
+  const controller = new AbortController()
+  const listeners: Array<() => void> = []
+  const abortFrom = (signal: AbortSignal) => {
+    if (!controller.signal.aborted) controller.abort(signal.reason)
+  }
+
+  for (const signal of activeSignals) {
+    if (signal.aborted) {
+      abortFrom(signal)
+      continue
+    }
+
+    const listener = () => abortFrom(signal)
+    signal.addEventListener('abort', listener, { once: true })
+    listeners.push(() => signal.removeEventListener('abort', listener))
+  }
+
+  return {
+    signal: controller.signal,
+    cleanup: () => {
+      for (const cleanup of listeners) cleanup()
+    },
+  }
+}
+
+async function callAbortableProvider<T>(
+  params: AiGatewayCallOptions,
+  effectSignal: AbortSignal,
+  evaluate: (params: AiGatewayCallOptions) => PromiseLike<T>,
+): Promise<T> {
+  const { signal, cleanup } = mergeAbortSignals([effectSignal, params.abortSignal])
+  try {
+    return await evaluate(signal ? { ...params, abortSignal: signal } : params)
+  } finally {
+    cleanup()
+  }
+}
+
 function attemptAiGatewayGenerate(
   source: string,
-  evaluate: () => PromiseLike<AiGatewayGenerateResult>,
+  params: AiGatewayCallOptions,
+  evaluate: (params: AiGatewayCallOptions) => PromiseLike<AiGatewayGenerateResult>,
 ): Effect.Effect<AiGatewayAttemptResult<AiGatewayGenerateResult>, AiGenerationError> {
   return withAiGatewayResilience(
     source,
-    Effect.tryPromise({ try: evaluate, catch: (cause) => classifyAiGatewayError(source, cause) }),
+    Effect.tryPromise({
+      try: (signal) => callAbortableProvider(params, signal, evaluate),
+      catch: (cause) => classifyAiGatewayError(source, cause),
+    }),
   ).pipe(
     Effect.map((result) => ({ source, result })),
     Effect.withSpan('AiGateway.generateAttempt'),
@@ -764,11 +811,15 @@ function attemptAiGatewayGenerate(
 
 function attemptAiGatewayStream(
   source: string,
-  evaluate: () => PromiseLike<AiGatewayStreamResult>,
+  params: AiGatewayCallOptions,
+  evaluate: (params: AiGatewayCallOptions) => PromiseLike<AiGatewayStreamResult>,
 ): Effect.Effect<AiGatewayAttemptResult<AiGatewayStreamResult>, AiGenerationError> {
   return withAiGatewayResilience(
     source,
-    Effect.tryPromise({ try: evaluate, catch: (cause) => classifyAiGatewayError(source, cause) }),
+    Effect.tryPromise({
+      try: (signal) => callAbortableProvider(params, signal, evaluate),
+      catch: (cause) => classifyAiGatewayError(source, cause),
+    }),
   ).pipe(
     Effect.map((result) => ({ source, result })),
     Effect.withSpan('AiGateway.streamAttempt'),
@@ -778,10 +829,11 @@ function attemptAiGatewayStream(
 
 function executeGenerateAttemptPlan(
   modelId: string,
-  doGenerate: () => PromiseLike<AiGatewayGenerateResult>,
+  params: AiGatewayCallOptions,
+  doGenerate: (params: AiGatewayCallOptions) => PromiseLike<AiGatewayGenerateResult>,
 ): Effect.Effect<AiGatewayAttemptResult<AiGatewayGenerateResult>, AiGenerationError> {
   const primary = Layer.succeed(AiGatewayGenerateAttemptTag, {
-    execute: attemptAiGatewayGenerate('ai-gateway.generate', doGenerate),
+    execute: attemptAiGatewayGenerate('ai-gateway.generate', params, doGenerate),
   })
   const effect = Effect.gen(function* () {
     const attempt = yield* AiGatewayGenerateAttemptTag
@@ -797,10 +849,11 @@ function executeGenerateAttemptPlan(
 
 function executeStreamAttemptPlan(
   modelId: string,
-  doStream: () => PromiseLike<AiGatewayStreamResult>,
+  params: AiGatewayCallOptions,
+  doStream: (params: AiGatewayCallOptions) => PromiseLike<AiGatewayStreamResult>,
 ): Effect.Effect<AiGatewayAttemptResult<AiGatewayStreamResult>, AiGenerationError> {
   const primary = Layer.succeed(AiGatewayStreamAttemptTag, {
-    execute: attemptAiGatewayStream('ai-gateway.stream', doStream),
+    execute: attemptAiGatewayStream('ai-gateway.stream', params, doStream),
   })
   const effect = Effect.gen(function* () {
     const attempt = yield* AiGatewayStreamAttemptTag
@@ -1046,7 +1099,7 @@ function createAiGatewayLanguageModelMiddleware(
       const model = resolveProviderModel(resolvedDeps.gateway.provider, modelId, providerId)
       return resolvedDeps.runPromise(
         withAiGatewayConcurrency(
-          executeGenerateAttemptPlan(modelId, () => model.doGenerate(params)).pipe(
+          executeGenerateAttemptPlan(modelId, params, (attemptParams) => model.doGenerate(attemptParams)).pipe(
             Effect.map(({ result }) => ({
               ...result,
               content: injectAiGatewayChatReasoningContent(
@@ -1063,7 +1116,7 @@ function createAiGatewayLanguageModelMiddleware(
       const model = resolveProviderModel(resolvedDeps.gateway.provider, modelId, providerId)
       return resolvedDeps.runPromise(
         withAiGatewayStreamConcurrency(
-          executeStreamAttemptPlan(modelId, () => model.doStream(params)).pipe(
+          executeStreamAttemptPlan(modelId, params, (attemptParams) => model.doStream(attemptParams)).pipe(
             Effect.map((attempt) => ({
               ...attempt,
               result: isReasoningEnabled(params)

package/src/queues/document-processor.queue.ts

@@ -66,11 +66,13 @@ export function createDocumentProcessorQueueRuntime<TJob extends DocumentProcess
   queueName?: string
   workerName?: string
   concurrency?: number
+  attempts?: number
   lockDuration?: number
 }): { enqueue: (job: TJob) => Promise<void>; startWorker: (options?: { registerSignals?: boolean }) => WorkerHandle } {
   const queueName = params.queueName ?? DEFAULT_DOCUMENT_PROCESSOR_QUEUE
   const workerName = params.workerName ?? DEFAULT_WORKER_NAME
   const concurrency = params.concurrency ?? 10
+  const attempts = params.attempts ?? 3
   const lockDuration = params.lockDuration ?? 300_000
   const queueRuntime = createQueueFactory<TJob>({
     name: queueName,
@@ -80,7 +82,7 @@ export function createDocumentProcessorQueueRuntime<TJob extends DocumentProcess
     lockDuration,
     logger: params.logger,
     connectionProvider: params.getConnectionForBullMQ,
-    defaultJobOptions: { ...DEFAULT_JOB_RETENTION, attempts: 3, backoff: { type: 'exponential', delay: 1000 } },
+    defaultJobOptions: { ...DEFAULT_JOB_RETENTION, attempts, backoff: { type: 'exponential', delay: 1000 } },
     processorPath: params.getWorkerPath(),
     queueJobService: params.queueJobService,
   })

package/src/workers/worker-utils.ts

@@ -1,3 +1,4 @@
+import { setTimeout as delay } from 'node:timers/promises'
 import { fileURLToPath } from 'node:url'
 
 import type { Job, Worker } from 'bullmq'
@@ -30,6 +31,32 @@ export interface WorkerHandle {
   shutdown: () => Promise<void>
 }
 
+async function closeWorkerWithTimeout(
+  worker: Worker,
+  name: string,
+  logger: typeof chatLogger,
+  timeoutMs = DEFAULT_SHUTDOWN_TIMEOUT_MS,
+): Promise<void> {
+  const gracefulClose = worker.close(false)
+  const gracefulOutcome = gracefulClose.then(
+    () => ({ status: 'closed' as const }),
+    (error: unknown) => ({ status: 'failed' as const, error }),
+  )
+
+  const outcome = await Promise.race([gracefulOutcome, delay(timeoutMs).then(() => ({ status: 'timed_out' as const }))])
+
+  if (outcome.status === 'closed') return
+
+  if (outcome.status === 'timed_out') {
+    logger.warn`${name} worker did not close within ${timeoutMs}ms; force-closing`
+  } else {
+    logger.warn`${name} worker graceful close failed; force-closing: ${outcome.error}`
+  }
+
+  await worker.close(true)
+  if (outcome.status === 'failed') throw outcome.error
+}
+
 interface TracedWorkerJobLike {
   id?: unknown
   name: string
@@ -87,7 +114,7 @@ export const createWorkerShutdown = (worker: Worker, name: string, logger: typeo
     return Effect.runPromise(
       Effect.asVoid(
         Effect.tryPromise({
-          try: () => worker.close(true),
+          try: () => closeWorkerWithTimeout(worker, name, logger),
           catch: (cause) => new QueueWorkerError({ phase: 'close', cause }),
         }),
       ),