Merge pull request #13 from O-Labz/memory-improvements

made more improvements for better performance and reliability for lar…

Author: HopsonO

README.md

@@ -76,6 +76,7 @@ ollama pull nomic-embed-text
 
 docker run -d \
   --name context-simplo \
+  --restart unless-stopped \
   -p 3001:3001 \
   -v "$HOME":/host:ro \
   -v context-simplo-data:/data \
@@ -84,6 +85,7 @@ docker run -d \
   -e LLM_PROVIDER=ollama \
   -e LLM_BASE_URL=http://host.docker.internal:11434 \
   -e LLM_EMBEDDING_MODEL=nomic-embed-text \
+  -e GRAPH_MEMORY_LIMIT_MB=4096 \
   ohopson/context-simplo:latest
 ```
 

src/core/config.ts

@@ -28,14 +28,24 @@ import type {
   ResponseMode,
 } from './types.js';
 
+/**
+ * Default configuration values
+ *
+ * Key boot behavior config:
+ * - autoIndex: When true, resumes incomplete indexing for previously-added repos on boot.
+ *              Does NOT auto-index the mount root. Repos must be explicitly added via API/MCP.
+ * - watchEnabled: Gates per-repo watcher restore on boot. If true, restores watchers for
+ *                 repos where isWatched=true. Does NOT auto-watch the mount root.
+ * - autoWatch: Controls whether newly-added repos are automatically watched (via API/MCP).
+ */
 const DEFAULT_CONFIG = {
   llmProvider: 'none' as LLMProviderType,
   llmApiKey: undefined,
   llmBaseUrl: 'http://host.docker.internal:11434',
   llmEmbeddingModel: undefined,
   dataDir: '/data',
-  autoIndex: true,
-  watchEnabled: true,
+  autoIndex: true, // Resume incomplete indexing on boot for previously-added repos
+  watchEnabled: true, // Gate per-repo watcher restore on boot
   autoWatch: true,
   logLevel: 'info' as const,
   embeddingConcurrency: 5,

src/core/embedding-backfill.ts

@@ -0,0 +1,234 @@
+/**
+ * Embedding Backfiller - Background worker for semantic embeddings
+ *
+ * What it does:
+ * Continuously processes files with embedding_status='pending' from SQLite,
+ * generating embeddings and upserting to LanceDB. Decouples semantic indexing
+ * from the critical path so structural queries work immediately.
+ *
+ * Inputs: SQLite storage (source of truth for pending files)
+ * Outputs: Embeddings written to LanceDB, files marked done/error
+ * Constraints: Bounded concurrency, bounded memory (no unbounded queues)
+ * Assumptions: Files are already parsed and in storage
+ * Failure cases: File deleted, embedding provider down, vectorStore errors
+ *
+ * Design:
+ * - Pull files in batches from SQLite (bounded memory, survives restart)
+ * - Read file content, scrub secrets, chunk by symbols, embed via queue
+ * - Upsert to vectorStore, mark file status (done/error)
+ * - Bounded concurrency (N files in flight at once)
+ * - Graceful error handling (log + mark error + continue, no crash)
+ * - Clean start/stop for graceful shutdown
+ *
+ * Performance: Processes files at the rate the embedding provider allows
+ * Concurrency: Async with configurable concurrency limit
+ * Security: Scrubs secrets before embedding
+ */
+
+import { readFile } from 'node:fs/promises';
+import { resolve } from 'node:path';
+import { EventEmitter } from 'events';
+import type { StorageProvider } from '../store/provider.js';
+import type { EmbeddingQueue } from './embedding-queue.js';
+import type { LanceDBVectorStore } from '../store/lance.js';
+import { scrubSecrets } from '../security/scrubber.js';
+import { chunkCodeForEmbedding } from '../llm/chunker.js';
+import { parseFile } from './parser.js';
+import { StoreError } from './errors.js';
+
+export interface EmbeddingBackfillOptions {
+  concurrency: number;
+  batchSize: number;
+  pollIntervalMs: number;
+  workspaceRoot: string;
+}
+
+export interface EmbeddingBackfillStats {
+  filesProcessed: number;
+  filesErrored: number;
+  chunksCreated: number;
+  isRunning: boolean;
+}
+
+export class EmbeddingBackfiller extends EventEmitter {
+  private storage: StorageProvider;
+  private embeddingQueue: EmbeddingQueue;
+  private vectorStore: LanceDBVectorStore;
+  private options: EmbeddingBackfillOptions;
+  private running = false;
+  private inFlight = 0;
+  private filesProcessed = 0;
+  private filesErrored = 0;
+  private chunksCreated = 0;
+  private pollTimer: NodeJS.Timeout | null = null;
+
+  constructor(
+    storage: StorageProvider,
+    embeddingQueue: EmbeddingQueue,
+    vectorStore: LanceDBVectorStore,
+    options: EmbeddingBackfillOptions
+  ) {
+    super();
+    this.storage = storage;
+    this.embeddingQueue = embeddingQueue;
+    this.vectorStore = vectorStore;
+    this.options = options;
+  }
+
+  start(): void {
+    if (this.running) {
+      console.warn('EmbeddingBackfiller already running');
+      return;
+    }
+
+    this.running = true;
+    console.log(
+      `EmbeddingBackfiller started (concurrency: ${this.options.concurrency}, ` +
+      `batch: ${this.options.batchSize}, poll: ${this.options.pollIntervalMs}ms)`
+    );
+    this.schedulePoll();
+  }
+
+  async stop(): Promise<void> {
+    if (!this.running) {
+      return;
+    }
+
+    this.running = false;
+
+    if (this.pollTimer) {
+      clearTimeout(this.pollTimer);
+      this.pollTimer = null;
+    }
+
+    // Wait for in-flight operations to complete
+    while (this.inFlight > 0) {
+      await this.sleep(100);
+    }
+
+    console.log('EmbeddingBackfiller stopped');
+  }
+
+  getStats(): EmbeddingBackfillStats {
+    return {
+      filesProcessed: this.filesProcessed,
+      filesErrored: this.filesErrored,
+      chunksCreated: this.chunksCreated,
+      isRunning: this.running,
+    };
+  }
+
+  private schedulePoll(): void {
+    if (!this.running) {
+      return;
+    }
+
+    this.pollTimer = setTimeout(() => {
+      this.poll().catch((error) => {
+        console.error('EmbeddingBackfiller poll failed:', error);
+      }).finally(() => {
+        this.schedulePoll();
+      });
+    }, this.options.pollIntervalMs);
+  }
+
+  private async poll(): Promise<void> {
+    if (!this.running || this.inFlight >= this.options.concurrency) {
+      return;
+    }
+
+    const availableSlots = this.options.concurrency - this.inFlight;
+    const files = this.storage.listPendingEmbeddingFiles(availableSlots);
+
+    if (files.length === 0) {
+      return;
+    }
+
+    for (const file of files) {
+      if (!this.running || this.inFlight >= this.options.concurrency) {
+        break;
+      }
+
+      this.inFlight++;
+      this.processFile(file.path, file.repositoryId)
+        .catch((error) => {
+          console.error(`EmbeddingBackfiller failed for ${file.path}:`, error);
+        })
+        .finally(() => {
+          this.inFlight--;
+        });
+    }
+  }
+
+  private async processFile(filePath: string, repositoryId: string): Promise<void> {
+    try {
+      // Read file content
+      const absolutePath = resolve(this.options.workspaceRoot, filePath);
+      let fileContent: string;
+      try {
+        fileContent = await readFile(absolutePath, 'utf-8');
+      } catch (error) {
+        throw new StoreError(
+          'readFile',
+          `File not readable: ${filePath}`,
+          error as Error
+        );
+      }
+
+      // Scrub secrets
+      const { scrubbed } = scrubSecrets(fileContent);
+
+      // Parse file to get AST nodes for chunking
+      let parsed;
+      try {
+        parsed = await parseFile(filePath, repositoryId, this.options.workspaceRoot);
+      } catch (error) {
+        throw new StoreError(
+          'parseFile',
+          `Failed to parse ${filePath}`,
+          error as Error
+        );
+      }
+
+      // Chunk code for embedding
+      const chunks = chunkCodeForEmbedding(parsed, scrubbed);
+
+      if (chunks.length === 0) {
+        // No chunks to embed (e.g., file only has imports)
+        this.storage.updateFileEmbeddingStatus(filePath, 'done');
+        this.filesProcessed++;
+        this.emit('file:complete', filePath, 0);
+        return;
+      }
+
+      // Embed chunks via the queue (with backpressure)
+      const texts = chunks.map((chunk) => chunk.content);
+      const embeddings = await this.embeddingQueue.embed(texts);
+
+      // Attach embeddings to chunks
+      const chunksWithEmbeddings = chunks.map((chunk, i) => ({
+        ...chunk,
+        embedding: embeddings[i],
+      }));
+
+      // Upsert to vectorStore
+      await this.vectorStore.upsertChunks(chunksWithEmbeddings);
+
+      // Mark file as done
+      this.storage.updateFileEmbeddingStatus(filePath, 'done');
+      this.filesProcessed++;
+      this.chunksCreated += chunks.length;
+      this.emit('file:complete', filePath, chunks.length);
+    } catch (error) {
+      // Mark file as error (do not retry in this implementation)
+      this.storage.updateFileEmbeddingStatus(filePath, 'error');
+      this.filesErrored++;
+      this.emit('file:error', filePath, error);
+      console.warn(`EmbeddingBackfiller marked ${filePath} as error:`, error);
+    }
+  }
+
+  private sleep(ms: number): Promise<void> {
+    return new Promise((resolve) => setTimeout(resolve, ms));
+  }
+}

src/core/embedding-queue.ts

@@ -31,6 +31,7 @@ export interface EmbeddingQueueOptions {
   concurrency: number;
   batchSize: number;
   maxRetries: number;
+  maxQueueDepth?: number;
 }
 
 export interface EmbeddingJob {
@@ -59,6 +60,7 @@ export class EmbeddingQueue extends EventEmitter {
   private concurrency: number;
   private maxRetries: number;
   private batchSize: number;
+  private maxQueueDepth: number;
   private draining = false;
 
   constructor(provider: EmbeddingProvider, options: EmbeddingQueueOptions) {
@@ -67,13 +69,19 @@ export class EmbeddingQueue extends EventEmitter {
     this.concurrency = options.concurrency;
     this.maxRetries = options.maxRetries;
     this.batchSize = Math.max(1, options.batchSize || 1);
+    this.maxQueueDepth = options.maxQueueDepth || Number.MAX_SAFE_INTEGER;
   }
 
   async embed(texts: string[]): Promise<number[][]> {
     if (this.draining) {
       throw new LLMError('queue', 'Queue is draining, not accepting new jobs', false);
     }
 
+    // Backpressure: wait for queue to drain below maxQueueDepth
+    while (this.queue.length >= this.maxQueueDepth) {
+      await this.waitForDrain();
+    }
+
     return new Promise((resolve, reject) => {
       const job: EmbeddingJob = {
         id: `job_${Date.now()}_${Math.random().toString(36).substring(2, 9)}`,
@@ -89,6 +97,19 @@ export class EmbeddingQueue extends EventEmitter {
     });
   }
 
+  private waitForDrain(): Promise<void> {
+    return new Promise((resolve) => {
+      const checkDrain = () => {
+        if (this.queue.length < this.maxQueueDepth) {
+          this.removeListener('drained', checkDrain);
+          resolve();
+        }
+      };
+      this.on('drained', checkDrain);
+      checkDrain();
+    });
+  }
+
   private async processQueue(): Promise<void> {
     while (this.inFlight < this.concurrency && this.queue.length > 0) {
       const job = this.queue.shift();
@@ -100,6 +121,7 @@ export class EmbeddingQueue extends EventEmitter {
       this.processJob(job).finally(() => {
         this.inFlight--;
         this.emit('progress', this.getStats());
+        this.emit('drained');
         this.processQueue();
       });
     }