Merge branch 'fix-api-key-components' of https://github.com/langflow-ai/langflow into fix-api-key-components

Author: HimavarshaVS

src/backend/base/langflow/api/utils/kb_helpers.py

@@ -2,6 +2,8 @@
 import contextlib
 import gc
 import json
+import shutil
+import time
 import uuid
 from datetime import datetime, timezone
 from functools import lru_cache
@@ -25,8 +27,10 @@
 from langflow.services.deps import get_settings_service
 from langflow.services.jobs.service import JobService
 from langflow.utils.kb_constants import (
+    DELETE_BACKOFF_SECONDS,
     EXPONENTIAL_BACKOFF_MULTIPLIER,
     INGESTION_BATCH_SIZE,
+    MAX_DELETE_RETRIES,
     MAX_RETRY_ATTEMPTS,
 )
 
@@ -81,8 +85,31 @@ def get_fresh_chroma_client(kb_path: Path) -> chromadb.PersistentClient:
         )
 
     @staticmethod
-    def teardown_storage(kb_path: Path, kb_name: str) -> None:
-        """Explicitly flush and invalidate Chroma clients before directory deletion."""
+    def release_chroma_resources(kb_path: Path) -> None:
+        """Release ChromaDB resources by clearing the registry entry and forcing GC."""
+        path_key = str(kb_path)
+        try:
+            if path_key in SharedSystemClient._identifier_to_system:  # noqa: SLF001
+                del SharedSystemClient._identifier_to_system[path_key]  # noqa: SLF001
+        except KeyError:
+            pass
+        gc.collect()
+
+    @staticmethod
+    def delete_storage(kb_path: Path, kb_name: str) -> bool:
+        """Teardown ChromaDB connections and delete KB directory with retry logic.
+
+        Handles ChromaDB SQLite file locks that can prevent deletion, particularly
+        on Windows where mandatory file locks block deletion of open files.
+        Uses retry with exponential backoff and rename-as-fallback strategy.
+
+        Returns:
+            True if deletion succeeded (or path already gone), False otherwise.
+        """
+        if not kb_path.exists():
+            return True
+
+        # Teardown ChromaDB collection to release handles
         try:
             has_data = any((kb_path / m).exists() for m in ["chroma", "chroma.sqlite3", "index"])
             if has_data:
@@ -91,9 +118,70 @@ def teardown_storage(kb_path: Path, kb_name: str) -> None:
                 with contextlib.suppress(Exception):
                     chroma.delete_collection()
                 chroma = None
-                gc.collect()
+                client = None
         except (OSError, ValueError, TypeError, chromadb.errors.ChromaError) as e:
-            logger.debug(f"Storage teardown failed for {kb_path.name} (ignoring): {e}")
+            logger.debug("Collection teardown failed for %s: %s", kb_path.name, e)
+
+        gc.collect()
+
+        for attempt in range(MAX_DELETE_RETRIES):
+            try:
+                if attempt > 0:
+                    time.sleep(DELETE_BACKOFF_SECONDS * (2**attempt))
+
+                _remove_sqlite_lock_files(kb_path)
+                _truncate_sqlite_files(kb_path)
+                gc.collect()
+
+                shutil.rmtree(kb_path, ignore_errors=False)
+
+                if not kb_path.exists():
+                    logger.info("Deleted knowledge base %s on attempt %d", kb_name, attempt + 1)
+                    return True
+
+            except OSError as e:
+                if attempt < MAX_DELETE_RETRIES - 1:
+                    logger.debug("KB deletion attempt %d failed for %s: %s", attempt + 1, kb_name, e)
+                else:
+                    logger.warning(
+                        "KB deletion failed for %s after %d attempts: %s",
+                        kb_name,
+                        MAX_DELETE_RETRIES,
+                        e,
+                    )
+
+        # Last resort: rename for deferred cleanup
+        if kb_path.exists():
+            try:
+                deferred = kb_path.with_name(f".deleted_{kb_name}_{int(time.time())}")
+                kb_path.rename(deferred)
+            except OSError as e:
+                logger.warning("Deferred rename failed for %s: %s", kb_name, e)
+            else:
+                logger.info("Renamed %s for deferred cleanup", kb_name)
+                return True
+
+        return False
+
+
+def _remove_sqlite_lock_files(kb_path: Path) -> None:
+    """Remove SQLite auxiliary files (WAL, SHM, journal) that hold locks."""
+    for pattern in ["*.sqlite3-wal", "*.sqlite3-shm", "*.sqlite3-journal"]:
+        for lock_file in kb_path.glob(pattern):
+            try:
+                lock_file.unlink()
+            except OSError as e:
+                logger.debug("Could not remove lock file %s: %s", lock_file.name, e)
+
+
+def _truncate_sqlite_files(kb_path: Path) -> None:
+    """Truncate SQLite database files to release locks."""
+    for sqlite_file in kb_path.glob("*.sqlite3"):
+        try:
+            with sqlite_file.open("r+b") as f:
+                f.truncate(0)
+        except OSError as e:
+            logger.debug("Could not truncate %s: %s", sqlite_file.name, e)
 
 
 class KBAnalysisHelper:
@@ -154,11 +242,15 @@ def get_metadata(kb_path: Path, *, fast: bool = False) -> dict:
     @staticmethod
     def update_text_metrics(kb_path: Path, metadata: dict, chroma: Chroma | None = None) -> None:
         """Update text metrics (chunks, words, characters) for a knowledge base."""
+        created_locally = chroma is None
+        client = None
         try:
-            if chroma is None:
+            if created_locally:
                 client = KBStorageHelper.get_fresh_chroma_client(kb_path)
                 chroma = Chroma(client=client, collection_name=kb_path.name)
 
+            if chroma is None:
+                return
             collection = chroma._collection  # noqa: SLF001
             metadata["chunks"] = collection.count()
 
@@ -190,6 +282,11 @@ def update_text_metrics(kb_path: Path, metadata: dict, chroma: Chroma | None = N
                 )
         except (OSError, ValueError, TypeError, json.JSONDecodeError, chromadb.errors.ChromaError) as e:
             logger.debug(f"Metrics update failed for {kb_path.name}: {e}")
+        finally:
+            if created_locally:
+                client = None
+                chroma = None
+                KBStorageHelper.release_chroma_resources(kb_path)
 
     @staticmethod
     def _detect_embedding_provider(kb_path: Path) -> str:
@@ -417,8 +514,9 @@ async def perform_ingestion(
             await KBIngestionHelper.cleanup_chroma_chunks_by_job(task_job_id, kb_path, kb_name)
             raise
         finally:
+            client = None
             chroma = None
-            gc.collect()
+            KBStorageHelper.release_chroma_resources(kb_path)
 
     @staticmethod
     async def cleanup_chroma_chunks_by_job(
@@ -438,8 +536,9 @@ async def cleanup_chroma_chunks_by_job(
         except (OSError, ValueError, TypeError, chromadb.errors.ChromaError) as cleanup_error:
             await logger.aerror(f"Failed to clean up chunks for job {job_id}: {cleanup_error}")
         finally:
+            client = None
             chroma = None
-            gc.collect()
+            KBStorageHelper.release_chroma_resources(kb_path)
 
     @staticmethod
     async def _is_job_cancelled(job_service: JobService, job_id: uuid.UUID) -> bool:

src/backend/base/langflow/api/v1/knowledge_bases.py

@@ -1,7 +1,5 @@
 import asyncio
-import gc
 import json
-import shutil
 import uuid
 from datetime import datetime, timezone
 from http import HTTPStatus
@@ -78,11 +76,11 @@ async def create_knowledge_base(
         try:
             client = KBStorageHelper.get_fresh_chroma_client(kb_path)
             client.create_collection(name=kb_name)
-            # Explicitly delete reference to help release handle
-            client = None
-            gc.collect()
         except (OSError, ValueError, chromadb.errors.ChromaError) as e:
             logger.warning("Initial Chroma setup for %s failed: %s", kb_name, e)
+        finally:
+            client = None
+            KBStorageHelper.release_chroma_resources(kb_path)
 
         # Serialize column_config for persistence
         column_config_dicts = None
@@ -130,7 +128,7 @@ async def create_knowledge_base(
     except Exception as e:
         # Clean up if something went wrong
         if kb_path.exists():
-            shutil.rmtree(kb_path)
+            KBStorageHelper.delete_storage(kb_path, kb_name)
         await logger.aerror("Error creating knowledge base: %s", e)
         raise HTTPException(status_code=500, detail="Internal error creating knowledge base") from e
 
@@ -593,8 +591,9 @@ async def get_knowledge_base_chunks(
         await logger.aerror("Error getting chunks for '%s': %s", kb_name, e)
         raise HTTPException(status_code=500, detail="Error getting chunks.") from e
     finally:
+        client = None
         chroma = None
-        gc.collect()
+        KBStorageHelper.release_chroma_resources(kb_path)
 
 
 @router.delete("/{kb_name}", status_code=HTTPStatus.OK)
@@ -603,11 +602,11 @@ async def delete_knowledge_base(kb_name: str, current_user: CurrentActiveUser) -
     try:
         kb_path = _resolve_kb_path(kb_name, current_user)
 
-        # Explicitly teardown KB storage to flush Chroma handles before directory deletion
-        KBStorageHelper.teardown_storage(kb_path, kb_name)
-
-        # Delete the entire knowledge base directory
-        shutil.rmtree(kb_path)
+        if not KBStorageHelper.delete_storage(kb_path, kb_name):
+            raise HTTPException(
+                status_code=500,
+                detail=f"Failed to delete knowledge base '{kb_name}'. The database may be in use.",
+            )
 
     except HTTPException:
         raise
@@ -636,12 +635,8 @@ async def delete_knowledge_bases_bulk(request: BulkDeleteRequest, current_user:
                 continue
 
             try:
-                # Explicitly teardown KB storage to flush Chroma handles before directory deletion
-                KBStorageHelper.teardown_storage(kb_path, kb_name)
-
-                # Delete the entire knowledge base directory
-                shutil.rmtree(kb_path)
-                deleted_count += 1
+                if KBStorageHelper.delete_storage(kb_path, kb_name):
+                    deleted_count += 1
             except (OSError, PermissionError) as e:
                 await logger.aexception("Error deleting knowledge base '%s': %s", kb_name, e)
                 # Continue with other deletions even if one fails

src/backend/base/langflow/utils/kb_constants.py

@@ -3,3 +3,7 @@
 EXPONENTIAL_BACKOFF_MULTIPLIER = 2
 MIN_KB_NAME_LENGTH = 3
 CHUNK_PREVIEW_MULTIPLIER = 3
+
+# KB deletion retry constants
+MAX_DELETE_RETRIES = 5
+DELETE_BACKOFF_SECONDS = 0.5