refactor: Move callsites to liteque to be behind a plugin

8 files changed, 121 insertions, 45 deletions

diff --git a/packages/shared-server/package.json b/packages/shared-server/package.json
index 6ba6b6d9..9c1b52a8 100644
--- a/packages/shared-server/package.json
+++ b/packages/shared-server/package.json
@@ -6,6 +6,7 @@
   "type": "module",
   "dependencies": {
     "@karakeep/db": "workspace:^0.1.0",
+    "@karakeep/plugins-queue-liteque": "workspace:^0.1.0",
     "@karakeep/plugins-search-meilisearch": "workspace:^0.1.0",
     "@karakeep/shared": "workspace:^0.1.0"
   },
diff --git a/packages/shared-server/src/index.ts b/packages/shared-server/src/index.ts
index ff3c6abc..d42118c2 100644
--- a/packages/shared-server/src/index.ts
+++ b/packages/shared-server/src/index.ts
@@ -1,2 +1,3 @@
 export { loadAllPlugins } from "./plugins";
 export { QuotaService, StorageQuotaError } from "./services/quotaService";
+export * from "./queues";
diff --git a/packages/shared-server/src/plugins.ts b/packages/shared-server/src/plugins.ts
index 86a0b344..b6a88462 100644
--- a/packages/shared-server/src/plugins.ts
+++ b/packages/shared-server/src/plugins.ts
@@ -6,6 +6,8 @@ export async function loadAllPlugins() {
     return;
   }
   // Load plugins here. Order of plugin loading matter.
+  // Queue provider(s)
+  await import("@karakeep/plugins-queue-liteque");
   await import("@karakeep/plugins-search-meilisearch");
   PluginManager.logAllPlugins();
   pluginsLoaded = true;
diff --git a/packages/shared/queues.ts b/packages/shared-server/src/queues.ts
index cf8920e1..c461c7cb 100644
--- a/packages/shared/queues.ts
+++ b/packages/shared-server/src/queues.ts
@@ -1,18 +1,15 @@
-import path from "node:path";
-import { buildDBClient, EnqueueOptions, migrateDB, SqliteQueue } from "liteque";
 import { z } from "zod";
 
-import serverConfig from "./config";
-import { zRuleEngineEventSchema } from "./types/rules";
+import { EnqueueOptions, getQueueClient } from "@karakeep/shared/queueing";
+import { zRuleEngineEventSchema } from "@karakeep/shared/types/rules";
 
-const QUEUE_DB_PATH = path.join(serverConfig.dataDir, "queue.db");
+import { loadAllPlugins } from ".";
 
-const queueDB = buildDBClient(QUEUE_DB_PATH, {
-  walEnabled: serverConfig.database.walMode,
-});
+await loadAllPlugins();
+const QUEUE_CLIENT = await getQueueClient();
 
 export function runQueueDBMigrations() {
-  migrateDB(queueDB);
+  QUEUE_CLIENT.init();
 }
 
 // Link Crawler
@@ -23,9 +20,8 @@ export const zCrawlLinkRequestSchema = z.object({
 });
 export type ZCrawlLinkRequest = z.input<typeof zCrawlLinkRequestSchema>;
 
-export const LinkCrawlerQueue = new SqliteQueue<ZCrawlLinkRequest>(
+export const LinkCrawlerQueue = QUEUE_CLIENT.createQueue<ZCrawlLinkRequest>(
   "link_crawler_queue",
-  queueDB,
   {
     defaultJobArgs: {
       numRetries: 5,
@@ -41,9 +37,8 @@ export const zOpenAIRequestSchema = z.object({
 });
 export type ZOpenAIRequest = z.infer<typeof zOpenAIRequestSchema>;
 
-export const OpenAIQueue = new SqliteQueue<ZOpenAIRequest>(
+export const OpenAIQueue = QUEUE_CLIENT.createQueue<ZOpenAIRequest>(
   "openai_queue",
-  queueDB,
   {
     defaultJobArgs: {
       numRetries: 3,
@@ -60,16 +55,13 @@ export const zSearchIndexingRequestSchema = z.object({
 export type ZSearchIndexingRequest = z.infer<
   typeof zSearchIndexingRequestSchema
 >;
-export const SearchIndexingQueue = new SqliteQueue<ZSearchIndexingRequest>(
-  "searching_indexing",
-  queueDB,
-  {
+export const SearchIndexingQueue =
+  QUEUE_CLIENT.createQueue<ZSearchIndexingRequest>("searching_indexing", {
     defaultJobArgs: {
       numRetries: 5,
     },
     keepFailedJobs: false,
-  },
-);
+  });
 
 // Tidy Assets Worker
 export const zTidyAssetsRequestSchema = z.object({
@@ -77,9 +69,8 @@ export const zTidyAssetsRequestSchema = z.object({
   syncAssetMetadata: z.boolean().optional().default(false),
 });
 export type ZTidyAssetsRequest = z.infer<typeof zTidyAssetsRequestSchema>;
-export const TidyAssetsQueue = new SqliteQueue<ZTidyAssetsRequest>(
+export const TidyAssetsQueue = QUEUE_CLIENT.createQueue<ZTidyAssetsRequest>(
   "tidy_assets_queue",
-  queueDB,
   {
     defaultJobArgs: {
       numRetries: 1,
@@ -107,9 +98,8 @@ export const zvideoRequestSchema = z.object({
 });
 export type ZVideoRequest = z.infer<typeof zvideoRequestSchema>;
 
-export const VideoWorkerQueue = new SqliteQueue<ZVideoRequest>(
+export const VideoWorkerQueue = QUEUE_CLIENT.createQueue<ZVideoRequest>(
   "video_queue",
-  queueDB,
   {
     defaultJobArgs: {
       numRetries: 5,
@@ -124,9 +114,8 @@ export const zFeedRequestSchema = z.object({
 });
 export type ZFeedRequestSchema = z.infer<typeof zFeedRequestSchema>;
 
-export const FeedQueue = new SqliteQueue<ZFeedRequestSchema>(
+export const FeedQueue = QUEUE_CLIENT.createQueue<ZFeedRequestSchema>(
   "feed_queue",
-  queueDB,
   {
     defaultJobArgs: {
       // One retry is enough for the feed queue given that it's periodic
@@ -145,9 +134,8 @@ export type AssetPreprocessingRequest = z.infer<
   typeof zAssetPreprocessingRequestSchema
 >;
 export const AssetPreprocessingQueue =
-  new SqliteQueue<AssetPreprocessingRequest>(
+  QUEUE_CLIENT.createQueue<AssetPreprocessingRequest>(
     "asset_preprocessing_queue",
-    queueDB,
     {
       defaultJobArgs: {
         numRetries: 2,
@@ -163,9 +151,8 @@ export const zWebhookRequestSchema = z.object({
   userId: z.string().optional(),
 });
 export type ZWebhookRequest = z.infer<typeof zWebhookRequestSchema>;
-export const WebhookQueue = new SqliteQueue<ZWebhookRequest>(
+export const WebhookQueue = QUEUE_CLIENT.createQueue<ZWebhookRequest>(
   "webhook_queue",
-  queueDB,
   {
     defaultJobArgs: {
       numRetries: 3,
@@ -196,9 +183,8 @@ export const zRuleEngineRequestSchema = z.object({
   events: z.array(zRuleEngineEventSchema),
 });
 export type ZRuleEngineRequest = z.infer<typeof zRuleEngineRequestSchema>;
-export const RuleEngineQueue = new SqliteQueue<ZRuleEngineRequest>(
+export const RuleEngineQueue = QUEUE_CLIENT.createQueue<ZRuleEngineRequest>(
   "rule_engine_queue",
-  queueDB,
   {
     defaultJobArgs: {
       numRetries: 1,
diff --git a/packages/shared/index.ts b/packages/shared/index.ts
index e449443d..e69de29b 100644
--- a/packages/shared/index.ts
+++ b/packages/shared/index.ts
@@ -1 +0,0 @@
-export * as Queues from "./queues";
diff --git a/packages/shared/package.json b/packages/shared/package.json
index 29879868..93739354 100644
--- a/packages/shared/package.json
+++ b/packages/shared/package.json
@@ -9,7 +9,6 @@
     "glob": "^11.0.0",
     "html-to-text": "^9.0.5",
     "js-tiktoken": "^1.0.20",
-    "liteque": "^0.6.0",
     "nodemailer": "^7.0.4",
     "ollama": "^0.5.14",
     "openai": "^4.86.1",
diff --git a/packages/shared/plugins.ts b/packages/shared/plugins.ts
index 2ce5826a..2aa7df4a 100644
--- a/packages/shared/plugins.ts
+++ b/packages/shared/plugins.ts
@@ -1,14 +1,17 @@
 // Implementation inspired from Outline
 
+import type { QueueClient } from "./queueing";
 import logger from "./logger";
 import { SearchIndexClient } from "./search";
 
 export enum PluginType {
   Search = "search",
+  Queue = "queue",
 }
 
 interface PluginTypeMap {
   [PluginType.Search]: SearchIndexClient;
+  [PluginType.Queue]: QueueClient;
 }
 
 export interface TPlugin<T extends PluginType> {
@@ -21,37 +24,38 @@ export interface PluginProvider<T> {
   getClient(): Promise<T | null>;
 }
 
+// Preserve the key-dependent value type: for K, store TPlugin<K>[]
+type ProviderMap = { [K in PluginType]: TPlugin<K>[] };
+
 export class PluginManager {
-  private static providers = new Map<PluginType, TPlugin<PluginType>[]>();
+  private static providers: ProviderMap = {
+    [PluginType.Search]: [],
+    [PluginType.Queue]: [],
+  };
 
   static register<T extends PluginType>(plugin: TPlugin<T>): void {
-    const p = PluginManager.providers.get(plugin.type);
-    if (!p) {
-      PluginManager.providers.set(plugin.type, [plugin]);
-      return;
-    }
-    p.push(plugin);
+    PluginManager.providers[plugin.type].push(plugin);
   }
 
   static async getClient<T extends PluginType>(
     type: T,
   ): Promise<PluginTypeMap[T] | null> {
-    const provider = PluginManager.providers.get(type);
-    if (!provider) {
+    const providers: TPlugin<T>[] = PluginManager.providers[type];
+    if (providers.length === 0) {
       return null;
     }
-    return await provider[provider.length - 1].provider.getClient();
+    return await providers[providers.length - 1]!.provider.getClient();
   }
 
   static isRegistered<T extends PluginType>(type: T): boolean {
-    return !!PluginManager.providers.get(type);
+    return PluginManager.providers[type].length > 0;
   }
 
   static logAllPlugins() {
     logger.info("Plugins (Last one wins):");
     for (const type of Object.values(PluginType)) {
       logger.info(`  ${type}:`);
-      const plugins = PluginManager.providers.get(type);
+      const plugins = PluginManager.providers[type];
       if (!plugins) {
         logger.info("    - None");
         continue;
diff --git a/packages/shared/queueing.ts b/packages/shared/queueing.ts
new file mode 100644
index 00000000..dfe3b31a
--- /dev/null
+++ b/packages/shared/queueing.ts
@@ -0,0 +1,84 @@
+import { ZodType } from "zod";
+
+import { PluginManager, PluginType } from "./plugins";
+
+export interface EnqueueOptions {
+  numRetries?: number;
+  idempotencyKey?: string;
+  priority?: number;
+  delayMs?: number;
+}
+
+export interface QueueOptions {
+  defaultJobArgs: {
+    numRetries: number;
+  };
+  keepFailedJobs: boolean;
+}
+
+export interface DequeuedJob<T> {
+  id: string;
+  data: T;
+  priority: number;
+  runNumber: number;
+  abortSignal: AbortSignal;
+}
+
+export interface DequeuedJobError<T> {
+  id: string;
+  data?: T;
+  priority: number;
+  error: Error;
+  runNumber: number;
+  numRetriesLeft: number;
+}
+
+export interface RunnerFuncs<T> {
+  run: (job: DequeuedJob<T>) => Promise<void>;
+  onComplete?: (job: DequeuedJob<T>) => Promise<void>;
+  onError?: (job: DequeuedJobError<T>) => Promise<void>;
+}
+
+export interface RunnerOptions<T> {
+  pollIntervalMs?: number;
+  timeoutSecs: number;
+  concurrency: number;
+  validator?: ZodType<T>;
+}
+
+export interface Queue<T> {
+  name(): string;
+  enqueue(payload: T, options?: EnqueueOptions): Promise<string | undefined>;
+  stats(): Promise<{
+    pending: number;
+    pending_retry: number;
+    running: number;
+    failed: number;
+  }>;
+  cancelAllNonRunning?(): Promise<number>;
+}
+
+export interface Runner<_T> {
+  run(): Promise<void>;
+  stop(): void;
+  runUntilEmpty?(): Promise<void>;
+}
+
+export interface QueueClient {
+  init(): Promise<void>;
+  createQueue<T>(name: string, options: QueueOptions): Queue<T>;
+  createRunner<T>(
+    queue: Queue<T>,
+    funcs: RunnerFuncs<T>,
+    opts: RunnerOptions<T>,
+  ): Runner<T>;
+  shutdown?(): Promise<void>;
+}
+
+export async function getQueueClient(): Promise<QueueClient> {
+  const client = await PluginManager.getClient(PluginType.Queue);
+  if (!client) {
+    throw new Error("Failed to get queue client");
+  }
+  return client;
+}