claude-code-workflow 6.3.13 → 6.3.15

package/codex-lens/src/codexlens/semantic/vector_store.py

@@ -9,12 +9,13 @@ Optimized for high-performance similarity search using:
 
 from __future__ import annotations
 
-import json
-import logging
-import sqlite3
-import threading
-from pathlib import Path
-from typing import Any, Dict, List, Optional, Tuple
+import json
+import logging
+import sys
+import sqlite3
+import threading
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
 
 from codexlens.entities import SearchResult, SemanticChunk
 from codexlens.errors import StorageError
@@ -39,6 +40,34 @@ logger = logging.getLogger(__name__)
 # Epsilon used to guard against floating point precision edge cases (e.g., near-zero norms).
 EPSILON = 1e-10
 
+# SQLite INTEGER PRIMARY KEY uses signed 64-bit rowids.
+SQLITE_INTEGER_MAX = (1 << 63) - 1
+
+
+def _validate_chunk_id_range(start_id: int, count: int) -> None:
+    """Validate that a batch insert can safely generate sequential chunk IDs."""
+    if count <= 0:
+        return
+
+    last_id = start_id + count - 1
+    if last_id > sys.maxsize or last_id > SQLITE_INTEGER_MAX:
+        raise ValueError(
+            "Chunk ID range overflow: "
+            f"start_id={start_id}, count={count} would allocate up to {last_id}, "
+            f"exceeding limits (sys.maxsize={sys.maxsize}, sqlite_max={SQLITE_INTEGER_MAX}). "
+            "Consider cleaning up the index database or creating a new index database."
+        )
+
+
+def _validate_sql_placeholders(placeholders: str, expected_count: int) -> None:
+    """Validate the placeholder string used for a parameterized SQL IN clause."""
+    expected = ",".join("?" * expected_count)
+    if placeholders != expected:
+        raise ValueError(
+            "Invalid SQL placeholders for IN clause. "
+            f"Expected {expected_count} '?' placeholders."
+        )
+
 
 def _cosine_similarity(a: List[float], b: List[float]) -> float:
     """Compute cosine similarity between two vectors."""
@@ -443,11 +472,11 @@ class VectorStore:
         self._invalidate_cache()
         return ids
 
-    def add_chunks_batch(
-        self,
-        chunks_with_paths: List[Tuple[SemanticChunk, str]],
-        update_ann: bool = True,
-        auto_save_ann: bool = True,
+    def add_chunks_batch(
+        self,
+        chunks_with_paths: List[Tuple[SemanticChunk, str]],
+        update_ann: bool = True,
+        auto_save_ann: bool = True,
     ) -> List[int]:
         """Batch insert chunks from multiple files in a single transaction.
 
@@ -459,16 +488,18 @@ class VectorStore:
             auto_save_ann: If True, save ANN index after update (default: True).
                           Set to False for bulk inserts to reduce I/O overhead.
 
-        Returns:
-            List of inserted chunk IDs
-        """
-        if not chunks_with_paths:
-            return []
-
-        # Prepare batch data
-        batch_data = []
-        embeddings_list = []
-        for chunk, file_path in chunks_with_paths:
+        Returns:
+            List of inserted chunk IDs
+        """
+        if not chunks_with_paths:
+            return []
+
+        batch_size = len(chunks_with_paths)
+
+        # Prepare batch data
+        batch_data = []
+        embeddings_list = []
+        for chunk, file_path in chunks_with_paths:
             if chunk.embedding is None:
                 raise ValueError("All chunks must have embeddings")
             # Optimize: avoid repeated np.array() if already numpy
@@ -481,49 +512,51 @@ class VectorStore:
             batch_data.append((file_path, chunk.content, embedding_blob, metadata_json))
             embeddings_list.append(embedding_arr)
 
-        # Batch insert to SQLite in single transaction
-        with sqlite3.connect(self.db_path) as conn:
-            # Get starting ID before insert
-            row = conn.execute("SELECT MAX(id) FROM semantic_chunks").fetchone()
-            start_id = (row[0] or 0) + 1
-
-            conn.executemany(
-                """
-                INSERT INTO semantic_chunks (file_path, content, embedding, metadata)
-                VALUES (?, ?, ?, ?)
+        # Batch insert to SQLite in single transaction
+        with sqlite3.connect(self.db_path) as conn:
+            # Get starting ID before insert
+            row = conn.execute("SELECT MAX(id) FROM semantic_chunks").fetchone()
+            start_id = (row[0] or 0) + 1
+
+            _validate_chunk_id_range(start_id, batch_size)
+
+            conn.executemany(
+                """
+                INSERT INTO semantic_chunks (file_path, content, embedding, metadata)
+                VALUES (?, ?, ?, ?)
                 """,
                 batch_data
-            )
-            conn.commit()
-            # Calculate inserted IDs based on starting ID
-            ids = list(range(start_id, start_id + len(chunks_with_paths)))
-
-        # Handle ANN index updates
-        if embeddings_list and update_ann and self._ensure_ann_index(len(embeddings_list[0])):
-            # In bulk insert mode, accumulate for later batch update
-            if self._bulk_insert_mode:
-                self._bulk_insert_ids.extend(ids)
-                self._bulk_insert_embeddings.extend(embeddings_list)
-            else:
-                # Normal mode: update immediately
-                with self._ann_write_lock:
-                    try:
-                        embeddings_matrix = np.vstack(embeddings_list)
-                        self._ann_index.add_vectors(ids, embeddings_matrix)
-                        if auto_save_ann:
-                            self._ann_index.save()
-                    except Exception as e:
-                        logger.warning("Failed to add batch to ANN index: %s", e)
+            )
+            conn.commit()
+            # Calculate inserted IDs based on starting ID
+            ids = list(range(start_id, start_id + batch_size))
+
+        # Handle ANN index updates
+        if embeddings_list and update_ann and self._ensure_ann_index(len(embeddings_list[0])):
+            with self._ann_write_lock:
+                # In bulk insert mode, accumulate for later batch update
+                if self._bulk_insert_mode:
+                    self._bulk_insert_ids.extend(ids)
+                    self._bulk_insert_embeddings.extend(embeddings_list)
+                else:
+                    # Normal mode: update immediately
+                    try:
+                        embeddings_matrix = np.vstack(embeddings_list)
+                        self._ann_index.add_vectors(ids, embeddings_matrix)
+                        if auto_save_ann:
+                            self._ann_index.save()
+                    except Exception as e:
+                        logger.warning("Failed to add batch to ANN index: %s", e)
 
         # Invalidate cache after modification
         self._invalidate_cache()
         return ids
 
-    def add_chunks_batch_numpy(
-        self,
-        chunks_with_paths: List[Tuple[SemanticChunk, str]],
-        embeddings_matrix: np.ndarray,
-        update_ann: bool = True,
+    def add_chunks_batch_numpy(
+        self,
+        chunks_with_paths: List[Tuple[SemanticChunk, str]],
+        embeddings_matrix: np.ndarray,
+        update_ann: bool = True,
         auto_save_ann: bool = True,
     ) -> List[int]:
         """Batch insert chunks with pre-computed numpy embeddings matrix.
@@ -537,16 +570,18 @@ class VectorStore:
             update_ann: If True, update ANN index with new vectors (default: True)
             auto_save_ann: If True, save ANN index after update (default: True)
 
-        Returns:
-            List of inserted chunk IDs
-        """
-        if not chunks_with_paths:
-            return []
-
-        if len(chunks_with_paths) != embeddings_matrix.shape[0]:
-            raise ValueError(
-                f"Mismatch: {len(chunks_with_paths)} chunks but "
-                f"{embeddings_matrix.shape[0]} embeddings"
+        Returns:
+            List of inserted chunk IDs
+        """
+        if not chunks_with_paths:
+            return []
+
+        batch_size = len(chunks_with_paths)
+
+        if len(chunks_with_paths) != embeddings_matrix.shape[0]:
+            raise ValueError(
+                f"Mismatch: {len(chunks_with_paths)} chunks but "
+                f"{embeddings_matrix.shape[0]} embeddings"
             )
 
         # Ensure float32 format
@@ -560,45 +595,47 @@ class VectorStore:
             metadata_json = json.dumps(chunk.metadata) if chunk.metadata else None
             batch_data.append((file_path, chunk.content, embedding_blob, metadata_json))
 
-        # Batch insert to SQLite in single transaction
-        with sqlite3.connect(self.db_path) as conn:
-            # Get starting ID before insert
-            row = conn.execute("SELECT MAX(id) FROM semantic_chunks").fetchone()
-            start_id = (row[0] or 0) + 1
-
-            conn.executemany(
-                """
-                INSERT INTO semantic_chunks (file_path, content, embedding, metadata)
-                VALUES (?, ?, ?, ?)
+        # Batch insert to SQLite in single transaction
+        with sqlite3.connect(self.db_path) as conn:
+            # Get starting ID before insert
+            row = conn.execute("SELECT MAX(id) FROM semantic_chunks").fetchone()
+            start_id = (row[0] or 0) + 1
+
+            _validate_chunk_id_range(start_id, batch_size)
+
+            conn.executemany(
+                """
+                INSERT INTO semantic_chunks (file_path, content, embedding, metadata)
+                VALUES (?, ?, ?, ?)
                 """,
                 batch_data
-            )
-            conn.commit()
-            # Calculate inserted IDs based on starting ID
-            ids = list(range(start_id, start_id + len(chunks_with_paths)))
-
-        # Handle ANN index updates
-        if update_ann and self._ensure_ann_index(embeddings_matrix.shape[1]):
-            # In bulk insert mode, accumulate for later batch update
-            if self._bulk_insert_mode:
-                self._bulk_insert_ids.extend(ids)
-                # Split matrix into individual arrays for accumulation
-                self._bulk_insert_embeddings.extend([embeddings_matrix[i] for i in range(len(ids))])
-            else:
-                # Normal mode: update immediately
-                with self._ann_write_lock:
-                    try:
-                        self._ann_index.add_vectors(ids, embeddings_matrix)
-                        if auto_save_ann:
-                            self._ann_index.save()
-                    except Exception as e:
-                        logger.warning("Failed to add batch to ANN index: %s", e)
+            )
+            conn.commit()
+            # Calculate inserted IDs based on starting ID
+            ids = list(range(start_id, start_id + batch_size))
+
+        # Handle ANN index updates
+        if update_ann and self._ensure_ann_index(embeddings_matrix.shape[1]):
+            with self._ann_write_lock:
+                # In bulk insert mode, accumulate for later batch update
+                if self._bulk_insert_mode:
+                    self._bulk_insert_ids.extend(ids)
+                    # Split matrix into individual arrays for accumulation
+                    self._bulk_insert_embeddings.extend([embeddings_matrix[i] for i in range(len(ids))])
+                else:
+                    # Normal mode: update immediately
+                    try:
+                        self._ann_index.add_vectors(ids, embeddings_matrix)
+                        if auto_save_ann:
+                            self._ann_index.save()
+                    except Exception as e:
+                        logger.warning("Failed to add batch to ANN index: %s", e)
 
         # Invalidate cache after modification
         self._invalidate_cache()
         return ids
 
-    def begin_bulk_insert(self) -> None:
+    def begin_bulk_insert(self) -> None:
         """Begin bulk insert mode - disable ANN auto-update for better performance.
 
         Usage:
@@ -614,42 +651,45 @@ class VectorStore:
                 for batch in batches:
                     store.add_chunks_batch(batch)
         """
-        self._bulk_insert_mode = True
-        self._bulk_insert_ids.clear()
-        self._bulk_insert_embeddings.clear()
-        logger.debug("Entered bulk insert mode")
-
-    def end_bulk_insert(self) -> None:
+        with self._ann_write_lock:
+            self._bulk_insert_mode = True
+            self._bulk_insert_ids.clear()
+            self._bulk_insert_embeddings.clear()
+        logger.debug("Entered bulk insert mode")
+
+    def end_bulk_insert(self) -> None:
         """End bulk insert mode and rebuild ANN index from accumulated data.
 
         This method should be called after all bulk inserts are complete to
         update the ANN index in a single batch operation.
         """
-        if not self._bulk_insert_mode:
-            logger.warning("end_bulk_insert called but not in bulk insert mode")
-            return
-
-        self._bulk_insert_mode = False
-
-        # Update ANN index with all accumulated data
-        if self._bulk_insert_ids and self._bulk_insert_embeddings:
-            if self._ensure_ann_index(len(self._bulk_insert_embeddings[0])):
-                with self._ann_write_lock:
-                    try:
-                        embeddings_matrix = np.vstack(self._bulk_insert_embeddings)
-                        self._ann_index.add_vectors(self._bulk_insert_ids, embeddings_matrix)
-                        self._ann_index.save()
-                        logger.info(
-                            "Bulk insert complete: added %d vectors to ANN index",
-                            len(self._bulk_insert_ids)
-                        )
-                    except Exception as e:
-                        logger.error("Failed to update ANN index after bulk insert: %s", e)
-
-        # Clear accumulated data
-        self._bulk_insert_ids.clear()
-        self._bulk_insert_embeddings.clear()
-        logger.debug("Exited bulk insert mode")
+        with self._ann_write_lock:
+            if not self._bulk_insert_mode:
+                logger.warning("end_bulk_insert called but not in bulk insert mode")
+                return
+
+            self._bulk_insert_mode = False
+            bulk_ids = list(self._bulk_insert_ids)
+            bulk_embeddings = list(self._bulk_insert_embeddings)
+            self._bulk_insert_ids.clear()
+            self._bulk_insert_embeddings.clear()
+
+        # Update ANN index with accumulated data.
+        if bulk_ids and bulk_embeddings:
+            if self._ensure_ann_index(len(bulk_embeddings[0])):
+                with self._ann_write_lock:
+                    try:
+                        embeddings_matrix = np.vstack(bulk_embeddings)
+                        self._ann_index.add_vectors(bulk_ids, embeddings_matrix)
+                        self._ann_index.save()
+                        logger.info(
+                            "Bulk insert complete: added %d vectors to ANN index",
+                            len(bulk_ids),
+                        )
+                    except Exception as e:
+                        logger.error("Failed to update ANN index after bulk insert: %s", e)
+
+        logger.debug("Exited bulk insert mode")
 
     class BulkInsertContext:
         """Context manager for bulk insert operations."""
@@ -712,34 +752,39 @@ class VectorStore:
             self._invalidate_cache()
         return deleted
 
-    def search_similar(
-        self,
-        query_embedding: List[float],
-        top_k: int = 10,
-        min_score: float = 0.0,
-        return_full_content: bool = True,
-    ) -> List[SearchResult]:
-        """Find chunks most similar to query embedding.
+    def search_similar(
+        self,
+        query_embedding: List[float],
+        top_k: int = 10,
+        min_score: float = 0.0,
+        return_full_content: bool = True,
+    ) -> List[SearchResult]:
+        """Find chunks most similar to query embedding.
 
         Uses HNSW index for O(log N) search when available, falls back to
         brute-force NumPy search otherwise.
 
-        Args:
-            query_embedding: Query vector.
-            top_k: Maximum results to return.
-            min_score: Minimum similarity score (0-1).
-            return_full_content: If True, return full code block content.
-
-        Returns:
-            List of SearchResult ordered by similarity (highest first).
-        """
-        query_vec = np.array(query_embedding, dtype=np.float32)
+        Args:
+            query_embedding: Query vector.
+            top_k: Maximum results to return.
+            min_score: Minimum cosine similarity score in [0.0, 1.0].
+            return_full_content: If True, return full code block content.
 
-        # Try HNSW search first (O(log N))
-        if (
-            HNSWLIB_AVAILABLE
-            and self._ann_index is not None
-            and self._ann_index.is_loaded
+        Returns:
+            List of SearchResult ordered by similarity (highest first).
+        """
+        query_vec = np.array(query_embedding, dtype=np.float32)
+
+        if not 0.0 <= min_score <= 1.0:
+            raise ValueError(
+                f"Invalid min_score: {min_score}. Must be within [0.0, 1.0] for cosine similarity."
+            )
+
+        # Try HNSW search first (O(log N))
+        if (
+            HNSWLIB_AVAILABLE
+            and self._ann_index is not None
+            and self._ann_index.is_loaded
             and self._ann_index.count() > 0
         ):
             try:
@@ -754,20 +799,20 @@ class VectorStore:
             query_vec, top_k, min_score, return_full_content
         )
 
-    def _search_with_ann(
-        self,
-        query_vec: np.ndarray,
-        top_k: int,
-        min_score: float,
-        return_full_content: bool,
-    ) -> List[SearchResult]:
-        """Search using HNSW index (O(log N)).
-
-        Args:
-            query_vec: Query vector as numpy array
-            top_k: Maximum results to return
-            min_score: Minimum similarity score (0-1)
-            return_full_content: If True, return full code block content
+    def _search_with_ann(
+        self,
+        query_vec: np.ndarray,
+        top_k: int,
+        min_score: float,
+        return_full_content: bool,
+    ) -> List[SearchResult]:
+        """Search using HNSW index (O(log N)).
+
+        Args:
+            query_vec: Query vector as numpy array
+            top_k: Maximum results to return
+            min_score: Minimum cosine similarity score in [0.0, 1.0]
+            return_full_content: If True, return full code block content
 
         Returns:
             List of SearchResult ordered by similarity (highest first)
@@ -779,15 +824,36 @@ class VectorStore:
         if effective_top_k == 0:
             return []
 
-        # HNSW search returns (ids, distances)
-        # For cosine space: distance = 1 - similarity
-        ids, distances = self._ann_index.search(query_vec, effective_top_k)
-
-        if not ids:
-            return []
-
-        # Convert distances to similarity scores
-        scores = [1.0 - d for d in distances]
+        # HNSW search returns (ids, distances)
+        # For cosine space: distance = 1 - similarity
+        ids, distances = self._ann_index.search(query_vec, effective_top_k)
+
+        if ids is None or distances is None:
+            logger.debug(
+                "ANN search returned null results (ids=%s, distances=%s)",
+                ids,
+                distances,
+            )
+            return []
+
+        if len(ids) == 0 or len(distances) == 0:
+            logger.debug(
+                "ANN search returned empty results (ids=%s, distances=%s)",
+                ids,
+                distances,
+            )
+            return []
+
+        if len(ids) != len(distances):
+            logger.warning(
+                "ANN search returned mismatched result lengths (%d ids, %d distances)",
+                len(ids),
+                len(distances),
+            )
+            return []
+
+        # Convert distances to similarity scores
+        scores = [1.0 - d for d in distances]
 
         # Filter by min_score
         filtered = [
@@ -805,20 +871,20 @@ class VectorStore:
         # Fetch content from SQLite
         return self._fetch_results_by_ids(top_ids, top_scores, return_full_content)
 
-    def _search_brute_force(
-        self,
-        query_vec: np.ndarray,
-        top_k: int,
-        min_score: float,
-        return_full_content: bool,
-    ) -> List[SearchResult]:
-        """Brute-force search using NumPy (O(N) fallback).
-
-        Args:
-            query_vec: Query vector as numpy array
-            top_k: Maximum results to return
-            min_score: Minimum similarity score (0-1)
-            return_full_content: If True, return full code block content
+    def _search_brute_force(
+        self,
+        query_vec: np.ndarray,
+        top_k: int,
+        min_score: float,
+        return_full_content: bool,
+    ) -> List[SearchResult]:
+        """Brute-force search using NumPy (O(N) fallback).
+
+        Args:
+            query_vec: Query vector as numpy array
+            top_k: Maximum results to return
+            min_score: Minimum cosine similarity score in [0.0, 1.0]
+            return_full_content: If True, return full code block content
 
         Returns:
             List of SearchResult ordered by similarity (highest first)
@@ -885,16 +951,21 @@ class VectorStore:
         Returns:
             List of SearchResult objects.
         """
-        if not chunk_ids:
-            return []
-
-        # Build parameterized query for IN clause
-        placeholders = ",".join("?" * len(chunk_ids))
-        query = f"""
-            SELECT id, file_path, content, metadata
-            FROM semantic_chunks
-            WHERE id IN ({placeholders})
-        """
+        if not chunk_ids:
+            return []
+
+        # Build parameterized query for IN clause
+        placeholders = ",".join("?" * len(chunk_ids))
+        _validate_sql_placeholders(placeholders, len(chunk_ids))
+
+        # SQL injection prevention:
+        # - Only a validated placeholders string (commas + '?') is interpolated into the query.
+        # - User-provided values are passed separately via sqlite3 parameters.
+        query = """
+            SELECT id, file_path, content, metadata
+            FROM semantic_chunks
+            WHERE id IN ({placeholders})
+        """.format(placeholders=placeholders)
 
         with sqlite3.connect(self.db_path) as conn:
             conn.execute("PRAGMA mmap_size = 30000000000")