@henrychong-ai/mcp-neo4j-knowledge-graph 2.3.2 → 2.4.0

package/dist/embeddings/EmbeddingJobManager.js

@@ -1,9 +1,6 @@
 import crypto from 'node:crypto';
 import { LRUCache } from 'lru-cache';
-import { v4 as uuidv4 } from 'uuid';
-/**
- * Default logger implementation
- */
+const DEFAULT_STALE_CLAIM_MS = 5 * 60 * 1000;
 const nullLogger = {
     debug: () => { },
     info: () => { },
@@ -11,7 +8,11 @@ const nullLogger = {
     error: () => { },
 };
 /**
- * Manages embedding jobs for semantic search
+ * Manages embedding jobs for semantic search.
+ *
+ * Persistence of the queue lives behind a `JobStore` — production wiring
+ * uses `Neo4jJobStore`, which stores jobs as `:EmbeddingJob` nodes and
+ * makes `claim()` safe under multi-worker contention.
  */
 export class EmbeddingJobManager {
     storageProvider;
@@ -20,20 +21,19 @@ export class EmbeddingJobManager {
     cache;
     cacheOptions = { size: 1000, ttl: 3_600_000 };
     logger;
-    /**
-     * Creates a new embedding job manager
-     *
-     * @param storageProvider - Provider for entity storage
-     * @param embeddingService - Service to generate embeddings
-     * @param rateLimiterOptions - Optional configuration for rate limiting
-     * @param cacheOptions - Optional configuration for caching
-     * @param logger - Optional logger for operation logging
-     */
-    constructor(storageProvider, embeddingService, rateLimiterOptions, cacheOptions, logger) {
+    jobStore;
+    staleClaimMs;
+    /** Stable id for this process — visible in `:EmbeddingJob.claimedBy`. */
+    workerId = crypto.randomUUID();
+    constructor(storageProvider, embeddingService, rateLimiterOptions, cacheOptions, logger, jobStore, staleClaimMs) {
         this.storageProvider = storageProvider;
         this.embeddingService = embeddingService;
         this.logger = logger || nullLogger;
-        // Setup rate limiter with defaults
+        if (!jobStore) {
+            throw new Error('EmbeddingJobManager requires a JobStore (v2.4.0+)');
+        }
+        this.jobStore = jobStore;
+        this.staleClaimMs = staleClaimMs ?? DEFAULT_STALE_CLAIM_MS;
         const defaultRateLimiter = {
             tokensPerInterval: 60,
             interval: 60 * 1000,
@@ -45,9 +45,7 @@ export class EmbeddingJobManager {
             tokensPerInterval: rateOptions.tokensPerInterval,
             interval: rateOptions.interval,
         };
-        // Setup LRU cache
         if (cacheOptions) {
-            // Support both API styles (tests use maxItems/ttlHours)
             this.cacheOptions = {
                 size: cacheOptions.size || cacheOptions.maxItems || 1000,
                 ttl: cacheOptions.ttl ||
@@ -59,191 +57,97 @@ export class EmbeddingJobManager {
             ttl: Math.max(1, Math.round(this.cacheOptions.ttl)),
             updateAgeOnGet: true,
             allowStale: false,
-            // Use a ttlAutopurge option to ensure items are purged when TTL expires
             ttlAutopurge: true,
         });
-        // Initialize database schema
-        this._initializeDatabase();
+        // Schema bootstrap — fire-and-log; safe on every start.
+        this.jobStore.ensureSchema().catch(err => {
+            this.logger.error('Failed to ensure JobStore schema', {
+                error: err instanceof Error ? err.message : String(err),
+            });
+        });
         this.logger.info('EmbeddingJobManager initialized', {
+            workerId: this.workerId,
+            staleClaimMs: this.staleClaimMs,
             cacheSize: this.cacheOptions.size,
             cacheTtl: this.cacheOptions.ttl,
             rateLimit: `${this.rateLimiter.tokensPerInterval} per ${this.rateLimiter.interval}ms`,
         });
     }
-    /**
-     * Initialize the database schema for embedding jobs
-     *
-     * @private
-     */
-    _initializeDatabase() {
-        const createTableSql = `
-      CREATE TABLE IF NOT EXISTS embedding_jobs (
-        id TEXT PRIMARY KEY,
-        entity_name TEXT NOT NULL,
-        status TEXT NOT NULL,
-        priority INTEGER NOT NULL DEFAULT 1,
-        created_at INTEGER NOT NULL,
-        processed_at INTEGER,
-        error TEXT,
-        attempts INTEGER NOT NULL DEFAULT 0,
-        max_attempts INTEGER NOT NULL DEFAULT 3
-      )
-    `;
-        // Create an index for efficient job retrieval
-        const createIndexSql = `
-      CREATE INDEX IF NOT EXISTS idx_embedding_jobs_status_priority
-      ON embedding_jobs (status, priority DESC)
-    `;
-        try {
-            this.storageProvider.db.exec(createTableSql);
-            this.storageProvider.db.exec(createIndexSql);
-            this.logger.debug('Database schema initialized for embedding jobs');
-        }
-        catch (error) {
-            this.logger.error('Failed to initialize database schema', { error });
-            throw error;
-        }
-    }
-    /**
-     * Schedule an entity for embedding generation
-     *
-     * @param entityName - Name of the entity to generate embedding for
-     * @param priority - Optional priority (higher priority jobs are processed first)
-     * @returns Job ID
-     */
     async scheduleEntityEmbedding(entityName, priority = 1) {
-        // Verify entity exists
         const entity = await this.storageProvider.getEntity(entityName);
         if (!entity) {
             const error = `Entity ${entityName} not found`;
             this.logger.error('Failed to schedule embedding', { entityName, error });
             throw new Error(error);
         }
-        // Create a job ID
-        const jobId = uuidv4();
-        // Insert a new job record
-        const stmt = this.storageProvider.db.prepare(`
-      INSERT INTO embedding_jobs (
-        id, entity_name, status, priority, created_at, attempts, max_attempts
-      ) VALUES (?, ?, ?, ?, ?, ?, ?)
-    `);
-        stmt.run(jobId, entityName, 'pending', priority, Date.now(), 0, 3);
-        this.logger.info('Scheduled embedding job', {
-            jobId,
-            entityName,
-            priority,
-        });
+        const jobId = await this.jobStore.enqueue({ entityName, priority });
+        this.logger.info('Scheduled embedding job', { jobId, entityName, priority });
         return jobId;
     }
-    /**
-     * Process a batch of pending embedding jobs
-     *
-     * @param batchSize - Maximum number of jobs to process
-     * @returns Result statistics
-     */
     async processJobs(batchSize = 10) {
-        this.logger.info('Starting job processing', { batchSize });
-        // Get pending jobs, ordered by priority (highest first)
-        const stmt = this.storageProvider.db.prepare(`
-      SELECT * FROM embedding_jobs
-      WHERE status = 'pending'
-      ORDER BY priority DESC, created_at ASC
-      LIMIT ?
-    `);
-        const jobs = stmt.all(batchSize);
-        this.logger.debug('Found pending jobs', { count: jobs.length });
-        // Initialize counters
-        const result = {
-            processed: 0,
-            successful: 0,
-            failed: 0,
-        };
-        // Process each job
+        this.logger.info('Starting job processing', { batchSize, workerId: this.workerId });
+        // Stale-claim sweep — releases jobs orphaned by dead workers before
+        // we attempt to claim. Cheap when nothing's stale; essential for
+        // recovery when a worker died mid-claim.
+        await this.jobStore.releaseStale(this.staleClaimMs);
+        const jobs = await this.jobStore.claim(this.workerId, batchSize);
+        this.logger.debug('Claimed jobs', { count: jobs.length });
+        const result = { processed: 0, successful: 0, failed: 0 };
         for (const job of jobs) {
-            // Check rate limiter before processing
             const rateLimitCheck = this._checkRateLimiter();
             if (!rateLimitCheck.success) {
-                this.logger.warn('Rate limit reached, pausing job processing', {
-                    remaining: jobs.length - result.processed,
+                const remaining = jobs.slice(result.processed).map(j => j.id);
+                this.logger.warn('Rate limit reached, releasing remaining claims', {
+                    remaining: remaining.length,
                 });
-                break; // Stop processing jobs if rate limit is reached
+                // Release without burning retry attempts — these jobs never ran.
+                await this.jobStore.releaseClaims(remaining).catch(() => { });
+                break;
             }
             this.logger.info('Processing embedding job', {
                 jobId: job.id,
-                entityName: job.entity_name,
-                attempt: job.attempts + 1,
-                maxAttempts: job.max_attempts,
+                entityName: job.entityName,
+                attempt: job.attempts,
+                maxAttempts: job.maxAttempts,
             });
-            // Update job status to processing
-            this._updateJobStatus(job.id, 'processing', job.attempts + 1);
             try {
-                // Get the entity
-                const entity = await this.storageProvider.getEntity(job.entity_name);
-                if (!entity) {
-                    throw new Error(`Entity ${job.entity_name} not found`);
-                }
-                // Log entity details for debugging
-                this.logger.debug('Retrieved entity for embedding', {
-                    entityName: job.entity_name,
-                    entityType: entity.entityType,
-                    hasObservations: entity.observations ? 'yes' : 'no',
-                    observationsType: entity.observations ? typeof entity.observations : 'undefined',
-                    observationsLength: entity.observations && Array.isArray(entity.observations)
-                        ? entity.observations.length
-                        : 'n/a',
-                });
-                // Prepare text for embedding
+                const entity = await this.storageProvider.getEntity(job.entityName);
+                if (!entity)
+                    throw new Error(`Entity ${job.entityName} not found`);
                 const text = this._prepareEntityText(entity);
-                // Try to get from cache or generate new embedding
-                this.logger.debug('Generating embedding for entity', { entityName: job.entity_name });
                 const embedding = await this._getCachedEmbeddingOrGenerate(text);
-                // Get model info for embedding metadata
                 const modelInfo = this.embeddingService.getModelInfo();
-                // Store the embedding with the entity
-                this.logger.debug('Storing entity vector', {
-                    entityName: job.entity_name,
-                    vectorLength: embedding.length,
-                    model: modelInfo.name,
-                });
-                await this.storageProvider.storeEntityVector(job.entity_name, {
+                await this.storageProvider.storeEntityVector(job.entityName, {
                     vector: embedding,
                     model: modelInfo.name,
                     lastUpdated: Date.now(),
                 });
-                // Update job status to completed
-                this._updateJobStatus(job.id, 'completed');
+                await this.jobStore.complete(job.id);
                 this.logger.info('Successfully processed embedding job', {
                     jobId: job.id,
-                    entityName: job.entity_name,
+                    entityName: job.entityName,
                     model: modelInfo.name,
                     dimensions: embedding.length,
                 });
                 result.successful++;
             }
             catch (error) {
-                // Handle failures
                 const errorMessage = error instanceof Error ? error.message : String(error);
                 const errorStack = error instanceof Error ? error.stack : undefined;
+                const failResult = await this.jobStore.fail(job.id, errorMessage);
                 this.logger.error('Failed to process embedding job', {
                     jobId: job.id,
-                    entityName: job.entity_name,
+                    entityName: job.entityName,
                     error: errorMessage,
                     errorStack,
-                    attempt: job.attempts + 1,
-                    maxAttempts: job.max_attempts,
+                    attempt: job.attempts,
+                    maxAttempts: job.maxAttempts,
+                    nextStatus: failResult.status,
                 });
-                // Determine if we should mark as failed or keep for retry
-                if (job.attempts + 1 >= job.max_attempts) {
-                    this._updateJobStatus(job.id, 'failed', job.attempts + 1, errorMessage);
-                }
-                else {
-                    this._updateJobStatus(job.id, 'pending', job.attempts + 1, errorMessage);
-                }
                 result.failed++;
             }
             result.processed++;
         }
-        // Log job processing results
         const queueStatus = await this.getQueueStatus();
         this.logger.info('Job processing complete', {
             processed: result.processed,
@@ -253,130 +157,36 @@ export class EmbeddingJobManager {
         });
         return result;
     }
-    /**
-     * Get the current status of the job queue
-     *
-     * @returns Queue statistics
-     */
     async getQueueStatus() {
-        const getCountForStatus = (status) => {
-            let sql = 'SELECT COUNT(*) as count FROM embedding_jobs';
-            const params = [];
-            if (status) {
-                sql += ' WHERE status = ?';
-                params.push(status);
-            }
-            const stmt = this.storageProvider.db.prepare(sql);
-            const result = stmt.get(...params);
-            return result?.count || 0;
-        };
-        const pending = getCountForStatus('pending');
-        const processing = getCountForStatus('processing');
-        const completed = getCountForStatus('completed');
-        const failed = getCountForStatus('failed');
-        const total = getCountForStatus();
-        const result = {
-            pending,
-            processing,
-            completed,
-            failed,
-            totalJobs: total,
-        };
-        this.logger.debug('Retrieved queue status', result);
-        return result;
+        const counts = await this.jobStore.countByStatus();
+        this.logger.debug('Retrieved queue status', { ...counts });
+        return counts;
     }
-    /**
-     * Retry failed embedding jobs
-     *
-     * @returns Number of jobs reset for retry
-     */
     async retryFailedJobs() {
-        const stmt = this.storageProvider.db.prepare(`
-      UPDATE embedding_jobs
-      SET status = 'pending', attempts = 0
-      WHERE status = 'failed'
-    `);
-        const result = stmt.run();
-        const resetCount = result.changes || 0;
-        this.logger.info('Reset failed jobs for retry', { count: resetCount });
-        return resetCount;
+        const reset = await this.jobStore.retryFailed();
+        this.logger.info('Reset failed jobs for retry', { count: reset });
+        return reset;
     }
-    /**
-     * Clean up old completed jobs
-     *
-     * @param threshold - Age in milliseconds after which to delete completed jobs, defaults to 7 days
-     * @returns Number of jobs cleaned up
-     */
     async cleanupJobs(threshold) {
-        const cleanupThreshold = threshold || 7 * 24 * 60 * 60 * 1000; // Default: 7 days
-        const cutoffTime = Date.now() - cleanupThreshold;
-        const stmt = this.storageProvider.db.prepare(`
-      DELETE FROM embedding_jobs
-      WHERE status = 'completed'
-      AND processed_at < ?
-    `);
-        const result = stmt.run(cutoffTime);
-        const deletedCount = result.changes || 0;
+        const cleanupThreshold = threshold || 7 * 24 * 60 * 60 * 1000;
+        const deleted = await this.jobStore.cleanup(cleanupThreshold);
         this.logger.info('Cleaned up old completed jobs', {
-            count: deletedCount,
+            count: deleted,
             threshold: cleanupThreshold,
-            olderThan: new Date(cutoffTime).toISOString(),
+            olderThan: new Date(Date.now() - cleanupThreshold).toISOString(),
         });
-        return deletedCount;
-    }
-    /**
-     * Update a job's status in the database
-     *
-     * @private
-     * @param jobId - ID of the job to update
-     * @param status - New status
-     * @param attempts - Optional attempts count update
-     * @param error - Optional error message
-     * @returns Database result
-     */
-    _updateJobStatus(jobId, status, attempts, error) {
-        let sql = `
-      UPDATE embedding_jobs
-      SET status = ?
-    `;
-        const params = [status];
-        // Add processed_at timestamp for completed/failed statuses
-        if (status === 'completed' || status === 'failed') {
-            sql += ', processed_at = ?';
-            params.push(Date.now());
-        }
-        // Update attempts if provided
-        if (attempts !== undefined) {
-            sql += ', attempts = ?';
-            params.push(attempts);
-        }
-        // Include error message if provided
-        if (error) {
-            sql += ', error = ?';
-            params.push(error);
-        }
-        sql += ' WHERE id = ?';
-        params.push(jobId);
-        const stmt = this.storageProvider.db.prepare(sql);
-        return stmt.run(...params);
+        return deleted;
     }
     /**
-     * Check rate limiter and consume a token if available
-     *
-     * @private
-     * @returns Object with success flag
+     * Token-bucket rate limiter. Public for legacy test compatibility — was
+     * `_checkRateLimiter` historically; kept callable from tests via underscore.
      */
     _checkRateLimiter() {
-        // For testing purposes, make it public by removing 'private'
         const now = Date.now();
         const elapsed = now - this.rateLimiter.lastRefill;
-        // If enough time has passed, refill tokens
         if (elapsed >= this.rateLimiter.interval) {
-            // Calculate how many full intervals have passed
             const intervals = Math.floor(elapsed / this.rateLimiter.interval);
-            // Completely refill tokens (don't accumulate beyond max)
             this.rateLimiter.tokens = this.rateLimiter.tokensPerInterval;
-            // Update last refill time, keeping track of remaining time
             this.rateLimiter.lastRefill = now;
             this.logger.debug('Refilled rate limiter tokens', {
                 current: this.rateLimiter.tokens,
@@ -384,7 +194,6 @@ export class EmbeddingJobManager {
                 intervals,
             });
         }
-        // If we have tokens, consume one and return success
         if (this.rateLimiter.tokens > 0) {
             this.rateLimiter.tokens--;
             this.logger.debug('Consumed rate limiter token', {
@@ -393,7 +202,6 @@ export class EmbeddingJobManager {
             });
             return { success: true };
         }
-        // No tokens available
         this.logger.warn('Rate limit exceeded', {
             availableTokens: 0,
             maxTokens: this.rateLimiter.tokensPerInterval,
@@ -401,15 +209,9 @@ export class EmbeddingJobManager {
         });
         return { success: false };
     }
-    /**
-     * Get the current status of the rate limiter
-     *
-     * @returns Rate limiter status information
-     */
     getRateLimiterStatus() {
         const now = Date.now();
         const elapsed = now - this.rateLimiter.lastRefill;
-        // If enough time has passed for a complete refill
         if (elapsed >= this.rateLimiter.interval) {
             return {
                 availableTokens: this.rateLimiter.tokensPerInterval,
@@ -417,22 +219,14 @@ export class EmbeddingJobManager {
                 resetInMs: this.rateLimiter.interval,
             };
         }
-        // Otherwise return current state
         return {
             availableTokens: this.rateLimiter.tokens,
             maxTokens: this.rateLimiter.tokensPerInterval,
             resetInMs: this.rateLimiter.interval - elapsed,
         };
     }
-    /**
-     * Retrieve a cached embedding or generate a new one
-     *
-     * @param text - Text to generate embedding for
-     * @returns Embedding vector
-     */
     async _getCachedEmbeddingOrGenerate(text) {
         const cacheKey = this._generateCacheKey(text);
-        // Try to get from cache first
         const cachedValue = this.cache.get(cacheKey);
         if (cachedValue) {
             this.logger.debug('Cache hit', {
@@ -443,9 +237,7 @@ export class EmbeddingJobManager {
         }
         this.logger.debug('Cache miss', { textHash: cacheKey.slice(0, 8) });
         try {
-            // Generate new embedding
             const embedding = await this.embeddingService.generateEmbedding(text);
-            // Store in cache
             this._cacheEmbedding(text, embedding);
             return embedding;
         }
@@ -457,13 +249,6 @@ export class EmbeddingJobManager {
             throw error;
         }
     }
-    /**
-     * Store an embedding in the cache
-     *
-     * @private
-     * @param text - Original text
-     * @param embedding - Embedding vector
-     */
     _cacheEmbedding(text, embedding) {
         const cacheKey = this._generateCacheKey(text);
         const modelInfo = this.embeddingService.getModelInfo();
@@ -478,47 +263,27 @@ export class EmbeddingJobManager {
             dimensions: embedding.length,
         });
     }
-    /**
-     * Generate a deterministic cache key for text
-     *
-     * @private
-     * @param text - Text to hash
-     * @returns Cache key
-     */
     _generateCacheKey(text) {
         return crypto.createHash('md5').update(text).digest('hex');
     }
-    /**
-     * Prepare text for embedding from an entity
-     *
-     * @private
-     * @param entity - Entity to prepare text from
-     * @returns Processed text ready for embedding
-     */
     _prepareEntityText(entity) {
-        // Create a descriptive text from entity data
         const lines = [`Name: ${entity.name}`, `Type: ${entity.entityType}`, 'Observations:'];
-        // Add observations, ensuring we handle both string arrays and other formats
         if (entity.observations) {
-            // Handle case where observations might be stored as JSON string in some providers
             let observationsArray = entity.observations;
-            // If observations is a string, try to parse it as JSON
             if (typeof entity.observations === 'string') {
                 try {
                     observationsArray = JSON.parse(entity.observations);
                 }
                 catch {
-                    // If parsing fails, treat it as a single observation
                     observationsArray = [entity.observations];
                 }
             }
-            // Ensure it's an array at this point
             if (!Array.isArray(observationsArray)) {
                 observationsArray = [String(observationsArray)];
             }
-            // Add each observation to the text
-            if (observationsArray.length > 0) {
-                lines.push(...observationsArray.map(obs => `- ${obs}`));
+            const arr = observationsArray;
+            if (arr.length > 0) {
+                lines.push(...arr.map(obs => `- ${obs}`));
             }
             else {
                 lines.push('  (No observations)');
@@ -528,7 +293,6 @@ export class EmbeddingJobManager {
             lines.push('  (No observations)');
         }
         const text = lines.join('\n');
-        // Log the prepared text for debugging
         this.logger.debug('Prepared entity text for embedding', {
             entityName: entity.name,
             entityType: entity.entityType,
@@ -537,37 +301,31 @@ export class EmbeddingJobManager {
         });
         return text;
     }
-    /**
-     * Get a cached embedding entry (used for testing)
-     *
-     * @param key - Cache key
-     * @returns Cached embedding or undefined
-     */
     getCacheEntry(key) {
         return this.cache.get(key);
     }
     /**
-     * Schedule incremental regeneration for entities without embeddings
-     * This method queries all entities and schedules embedding jobs only for those missing embeddings
-     *
-     * @returns Number of entities scheduled for embedding generation
+     * Walk every current entity and enqueue jobs for any that are missing
+     * embeddings. Intended for a server-side cron tick to backfill entities
+     * created by thin clients running with `WRITE_EMBEDDINGS_LOCALLY=false`.
      */
     async scheduleIncrementalRegeneration() {
         this.logger.info('Starting incremental embedding regeneration check');
         try {
-            // Get all entities from storage
             const allEntities = await this._getAllEntitiesFromStorage();
             this.logger.debug('Retrieved entities for embedding check', {
                 totalCount: allEntities.length,
             });
-            // Filter for entities without embeddings
             const entitiesWithoutEmbeddings = allEntities.filter(entity => !entity.embedding);
+            const total = allEntities.length;
+            const coverage = total > 0
+                ? `${Math.round(((total - entitiesWithoutEmbeddings.length) / total) * 100)}%`
+                : '0%';
             this.logger.info('Found entities without embeddings', {
                 count: entitiesWithoutEmbeddings.length,
-                totalEntities: allEntities.length,
-                coverage: `${Math.round(((allEntities.length - entitiesWithoutEmbeddings.length) / allEntities.length) * 100)}%`,
+                totalEntities: total,
+                coverage,
             });
-            // Schedule embedding jobs for entities without embeddings
             let scheduledCount = 0;
             for (const entity of entitiesWithoutEmbeddings) {
                 try {
@@ -595,15 +353,7 @@ export class EmbeddingJobManager {
             throw error;
         }
     }
-    /**
-     * Get all entities from storage provider
-     * This is a helper method to retrieve all entities for incremental regeneration
-     *
-     * @private
-     * @returns Array of all entities
-     */
     async _getAllEntitiesFromStorage() {
-        // Use the storage provider's loadGraph method to get all entities
         // eslint-disable-next-line @typescript-eslint/no-explicit-any
         const storageProviderAny = this.storageProvider;
         if (typeof storageProviderAny.loadGraph === 'function') {