From fe8cadfe084a17804324602948cb706f39f5bcc8 Mon Sep 17 00:00:00 2001
From: Shreyas <shreyasjagannath08@gmail.com>
Date: Wed, 11 Mar 2026 17:48:21 +0000
Subject: [PATCH 1/3] fix: prevent learn tool from hanging on large codebases

Resolves #1. The learn tool would get stuck on large folders because
discover_files() traversed everything (including node_modules, build
artifacts) and had no checkpointing for interruption recovery.

Changes:
- Rewrite discover_files() to use os.walk() with pathspec for .gitignore
  support and directory pruning (followlinks=False for symlink safety)
- Add max_files=200K safety limit (addresses SECURITY_REVIEW HIGH-003)
- Switch index_codebase() from delete-recreate to get_or_create + upsert
  with stale chunk cleanup
- Add resumable checkpointing (atomic writes every 1000 chunks)
- Improve progress reporting with discovery events and ETA
- Add 15 new tests (8 unit + 7 integration)
- Apply black formatting across codebase
- Snyk code scan: 0 issues

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 CLAUDE.md                                     |   8 +-
 docs/INDEXING_IMPROVEMENTS.md                 |  87 +++
 pyproject.toml                                |   1 +
 src/codegrok_mcp/__init__.py                  |   8 +
 src/codegrok_mcp/core/exceptions.py           |   5 +
 src/codegrok_mcp/core/models.py               |  82 +--
 src/codegrok_mcp/indexing/__init__.py         |  28 +-
 .../indexing/embedding_service.py             | 142 ++--
 src/codegrok_mcp/indexing/memory_retriever.py | 178 +++--
 src/codegrok_mcp/indexing/parallel_indexer.py |  61 +-
 src/codegrok_mcp/indexing/source_retriever.py | 426 ++++++++----
 src/codegrok_mcp/mcp/server.py                | 302 ++++-----
 src/codegrok_mcp/mcp/state.py                 |   2 +
 src/codegrok_mcp/parsers/language_configs.py  | 607 +++++++++---------
 src/codegrok_mcp/parsers/treesitter_parser.py | 165 +++--
 tests/conftest.py                             |  11 +-
 .../sample_projects/python_project/main.py    |   1 +
 .../sample_projects/python_project/utils.py   |   1 +
 tests/integration/test_source_retriever.py    | 161 +++--
 tests/integration/test_tool_discovery.py      | 189 +++---
 tests/mcp/test_protocol_simulation.py         |  41 +-
 tests/mcp/test_state_management.py            |   5 +-
 tests/mcp/test_tools_direct.py                |  12 +-
 tests/unit/test_discover_files.py             | 126 ++++
 tests/unit/test_embedding_service.py          |  35 +-
 tests/unit/test_extras.py                     |  20 +-
 tests/unit/test_init.py                       |  33 +-
 tests/unit/test_memory_retriever.py           | 193 +++---
 tests/unit/test_models.py                     | 180 +++---
 tests/unit/test_parallel_indexer.py           |  31 +-
 tests/unit/test_parser.py                     |   7 +-
 31 files changed, 1757 insertions(+), 1391 deletions(-)
 create mode 100644 docs/INDEXING_IMPROVEMENTS.md
 create mode 100644 tests/unit/test_discover_files.py

diff --git a/CLAUDE.md b/CLAUDE.md
index c6bf5e2..fc76a0d 100644
--- a/CLAUDE.md
+++ b/CLAUDE.md
@@ -42,7 +42,7 @@ src/codegrok_mcp/
 - **Embedding Model**: `nomic-ai/CodeRankEmbed` (768 dims, 8192 max tokens)
 - **Chunk Strategy**: Symbol-based (each function/class/method = 1 chunk)
 - **Max Chunk Size**: 4000 chars (~1000-1300 tokens)
-- **Storage**: `.codegrok/` (chromadb/ + metadata.json + memory_metadata.json)
+- **Storage**: `.codegrok/` (chromadb/ + metadata.json + memory_metadata.json + checkpoint.json)
 - **Parallelism**: CPU count - 1 workers (min 1, max 32)
 - **Memory TTLs**: session (24h), day, week, month, permanent
 
@@ -59,13 +59,17 @@ mypy src/                 # Type check
 ## Gotchas
 
 1. **State is global singleton** - `state.py` holds SourceRetriever + MemoryRetriever across MCP calls
-2. **Incremental reindex uses SHA256** - File hash comparison, not mtime
+2. **Incremental reindex uses mtime** - File modification time comparison for change detection
 3. **ChromaDB collections**: `codebase_chunks` (code) and `memories` (memory layer)
 4. **No LLM code** - Removed from parent CodeGrok; source_retriever.py has no ask/rerank methods
 5. **Tree-sitter node names vary by language** - language_configs.py normalizes them
 6. **Embedding is cached** - LRU(1000) + batch processing in embedding_service.py
 7. **Memory tags stored as CSV** - ChromaDB doesn't support list metadata; tags joined with commas
 8. **All tools require `learn` first** - Except `list_supported_languages` (static data)
+9. **discover_files respects .gitignore** - Uses `pathspec` + `os.walk()` with directory pruning; also respects nested `.gitignore` files
+10. **Indexing uses upsert** - `collection.upsert()` instead of delete-recreate; stale chunks cleaned after embedding
+11. **Checkpointing** - `.codegrok/checkpoint.json` saves progress every 1000 chunks; atomic writes via `os.replace()`; deleted on success
+12. **max_files safety limit** - `discover_files()` stops at 200K files to prevent DoS (addresses SECURITY_REVIEW HIGH-003)
 
 ## Adding Languages
 
diff --git a/docs/INDEXING_IMPROVEMENTS.md b/docs/INDEXING_IMPROVEMENTS.md
new file mode 100644
index 0000000..583f981
--- /dev/null
+++ b/docs/INDEXING_IMPROVEMENTS.md
@@ -0,0 +1,87 @@
+# Indexing Improvements (v0.2.1)
+
+Fixes for the `learn` tool hanging on large codebases with many folders/subfolders.
+
+## Changes
+
+### 1. `.gitignore` Support
+
+`discover_files()` now respects `.gitignore` patterns using the `pathspec` library.
+
+- Uses `os.walk()` instead of `Path.rglob("*")` for directory pruning
+- Loads root `.gitignore` and stacks nested `.gitignore` files as it descends
+- Prunes ignored directories in-place (never descends into `node_modules/`, `build/`, etc.)
+- Uses `followlinks=False` to prevent symlink loops
+- Backward-compatible: `respect_gitignore=True` by default, can be disabled
+
+### 2. Safety Limits
+
+- `max_files=200_000` circuit breaker stops file discovery if exceeded
+- Emits a warning when the limit is hit
+- Addresses **SECURITY_REVIEW HIGH-003** (Unbounded Resource Consumption / DoS)
+
+### 3. Upsert-Based Indexing
+
+- `index_codebase()` now uses `get_or_create_collection()` + `collection.upsert()` instead of deleting and recreating the collection
+- Chunk IDs are deterministic (`filepath:name:line_start`), making upsert idempotent
+- Stale chunks (from deleted/renamed files) are cleaned up after the embedding loop
+
+### 4. Resumable Checkpointing
+
+- Saves progress to `.codegrok/checkpoint.json` every 1000 chunks
+- Atomic writes via `os.replace()` (POSIX-safe)
+- On restart, detects checkpoint and resumes from where it left off
+- Checkpoint is deleted on successful completion
+
+### 5. Improved Progress Reporting
+
+- New `"discovery_progress"` event emitted every 1000 files during file traversal
+- ETA added to embedding progress messages (e.g., "Embedding... (5000/10000 chunks, ~2.3m remaining)")
+- MCP client now shows progress during the file discovery phase (0-5% range)
+
+## New Dependencies
+
+- `pathspec>=0.11.0` — Pure Python `.gitignore` pattern matching (used by `black`, `flake8`, etc.)
+
+## Security Alignment
+
+| Security Finding | How Addressed |
+|-----------------|---------------|
+| HIGH-003: Unbounded Resource Consumption | `max_files` limit + `.gitignore` filtering |
+| LOW-009: Symlink Following | `followlinks=False` in `os.walk()` |
+
+## MCP Tools Used in Development
+
+This feature was planned and implemented using the following MCP tools:
+
+| MCP Tool | How It Was Used |
+|----------|----------------|
+| **Sequential Thinking** (`mcp__sequential-thinking__sequentialthinking`) | 7-step chain-of-thought to plan execution order, identify risks (backward compatibility, atomic writes, symlink loops), decide to skip async writer thread, and design test strategy |
+| **Snyk Code Scan** (`mcp__Snyk__snyk_code_scan`) | SAST scan on all modified files (`source_retriever.py`, `server.py`) — 0 issues found |
+
+## Test Coverage
+
+### New Unit Tests (`tests/unit/test_discover_files.py`)
+
+| Test | What it verifies |
+|------|-----------------|
+| `test_discover_files_basic` | Finds .py files in simple directory |
+| `test_discover_files_skip_dirs` | Skips `node_modules/`, `__pycache__/`, `.git/` even when nested |
+| `test_discover_files_gitignore` | Respects root `.gitignore` patterns |
+| `test_discover_files_nested_gitignore` | Handles `.gitignore` in subdirectories |
+| `test_discover_files_max_files_limit` | Stops at `max_files` and returns partial results |
+| `test_discover_files_no_gitignore` | Works when no `.gitignore` exists |
+| `test_discover_files_respect_gitignore_false` | Opt-out disables gitignore filtering |
+| `test_discover_files_progress_callback` | Callback mechanism works correctly |
+
+### New Integration Tests (`tests/integration/test_source_retriever.py`)
+
+| Test | What it verifies |
+|------|-----------------|
+| `test_index_codebase_upsert_idempotent` | Re-indexing produces same chunk count |
+| `test_stale_chunk_removal` | Old chunks removed after file deletion |
+| `test_checkpoint_save_and_load` | Checkpoint round-trip |
+| `test_checkpoint_load_missing_file` | Handles missing checkpoint |
+| `test_checkpoint_load_corrupted` | Handles corrupted JSON |
+| `test_checkpoint_cleanup_on_success` | Checkpoint deleted after success |
+| `test_checkpoint_load_none_path` | Handles None path |
diff --git a/pyproject.toml b/pyproject.toml
index b3cf4e1..cfb5b09 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -32,6 +32,7 @@ dependencies = [
     "sentence-transformers>=2.2.0",
     "torch>=2.0.0",
     "einops>=0.7.0",
+    "pathspec>=0.11.0",
 ]
 
 
diff --git a/src/codegrok_mcp/__init__.py b/src/codegrok_mcp/__init__.py
index b138ce3..3617141 100644
--- a/src/codegrok_mcp/__init__.py
+++ b/src/codegrok_mcp/__init__.py
@@ -24,27 +24,35 @@ def __getattr__(name: str):
     """Lazy import heavy modules only when accessed."""
     if name == "SourceRetriever":
         from codegrok_mcp.indexing.source_retriever import SourceRetriever
+
         return SourceRetriever
     elif name == "TreeSitterParser":
         from codegrok_mcp.parsers.treesitter_parser import TreeSitterParser
+
         return TreeSitterParser
     elif name == "ThreadLocalParserFactory":
         from codegrok_mcp.parsers.treesitter_parser import ThreadLocalParserFactory
+
         return ThreadLocalParserFactory
     elif name == "Symbol":
         from codegrok_mcp.core.models import Symbol
+
         return Symbol
     elif name == "SymbolType":
         from codegrok_mcp.core.models import SymbolType
+
         return SymbolType
     elif name == "ParsedFile":
         from codegrok_mcp.core.models import ParsedFile
+
         return ParsedFile
     elif name == "CodebaseIndex":
         from codegrok_mcp.core.models import CodebaseIndex
+
         return CodebaseIndex
     elif name == "IParser":
         from codegrok_mcp.core.interfaces import IParser
+
         return IParser
     raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
 
diff --git a/src/codegrok_mcp/core/exceptions.py b/src/codegrok_mcp/core/exceptions.py
index 8a13a0d..76ecd7c 100644
--- a/src/codegrok_mcp/core/exceptions.py
+++ b/src/codegrok_mcp/core/exceptions.py
@@ -19,6 +19,7 @@ class CodeGrokException(Exception):
     All custom exceptions in CodeGrok inherit from this class,
     allowing for broad exception catching when needed.
     """
+
     pass
 
 
@@ -46,6 +47,7 @@ class IndexingError(CodeGrokException):
     - Chunking failures
     - ChromaDB storage errors
     """
+
     pass
 
 
@@ -57,6 +59,7 @@ class EmbeddingError(CodeGrokException):
     - Encoding errors
     - Memory issues
     """
+
     pass
 
 
@@ -68,6 +71,7 @@ class SearchError(CodeGrokException):
     - Missing index
     - Invalid query parameters
     """
+
     pass
 
 
@@ -79,4 +83,5 @@ class ConfigurationError(CodeGrokException):
     - Invalid file paths
     - Missing required parameters
     """
+
     pass
diff --git a/src/codegrok_mcp/core/models.py b/src/codegrok_mcp/core/models.py
index 2ca9c51..7dda7f6 100644
--- a/src/codegrok_mcp/core/models.py
+++ b/src/codegrok_mcp/core/models.py
@@ -27,6 +27,7 @@ class SymbolType(Enum):
         METHOD: Method within a class
         VARIABLE: Module-level or class-level variable
     """
+
     FUNCTION = "function"
     CLASS = "class"
     METHOD = "method"
@@ -37,7 +38,7 @@ def __str__(self) -> str:
         return self.value
 
     @classmethod
-    def from_string(cls, value: str) -> 'SymbolType':
+    def from_string(cls, value: str) -> "SymbolType":
         """
         Create SymbolType from string value.
 
@@ -79,6 +80,7 @@ class Symbol:
         calls: List of function/method names called by this symbol
         metadata: Extensible dictionary for additional language-specific or custom data
     """
+
     name: str
     type: SymbolType
     filepath: str
@@ -107,7 +109,9 @@ def __post_init__(self):
         if self.line_start < 1:
             raise ValueError(f"line_start must be >= 1, got {self.line_start}")
         if self.line_end < self.line_start:
-            raise ValueError(f"line_end ({self.line_end}) must be >= line_start ({self.line_start})")
+            raise ValueError(
+                f"line_end ({self.line_end}) must be >= line_start ({self.line_start})"
+            )
         if not self.language:
             raise ValueError("Symbol language cannot be empty")
         if not isinstance(self.type, SymbolType):
@@ -143,11 +147,11 @@ def to_dict(self) -> Dict[str, Any]:
             Dictionary representation with all fields
         """
         data = asdict(self)
-        data['type'] = self.type.value  # Convert enum to string
+        data["type"] = self.type.value  # Convert enum to string
         return data
 
     @classmethod
-    def from_dict(cls, data: Dict[str, Any]) -> 'Symbol':
+    def from_dict(cls, data: Dict[str, Any]) -> "Symbol":
         """
         Create Symbol from dictionary.
 
@@ -161,9 +165,9 @@ def from_dict(cls, data: Dict[str, Any]) -> 'Symbol':
             ValueError: If required fields are missing or invalid
         """
         # Convert type string back to enum
-        if 'type' in data and isinstance(data['type'], str):
+        if "type" in data and isinstance(data["type"], str):
             data = data.copy()  # Don't modify original
-            data['type'] = SymbolType.from_string(data['type'])
+            data["type"] = SymbolType.from_string(data["type"])
         return cls(**data)
 
 
@@ -183,6 +187,7 @@ class ParsedFile:
         parse_time: Time taken to parse this file (seconds)
         error: None if parsing succeeded, error message if failed
     """
+
     filepath: str
     language: str
     symbols: List[Symbol] = field(default_factory=list)
@@ -244,16 +249,16 @@ def to_dict(self) -> Dict[str, Any]:
             Dictionary representation
         """
         return {
-            'filepath': self.filepath,
-            'language': self.language,
-            'symbols': [s.to_dict() for s in self.symbols],
-            'imports': self.imports,
-            'parse_time': self.parse_time,
-            'error': self.error,
+            "filepath": self.filepath,
+            "language": self.language,
+            "symbols": [s.to_dict() for s in self.symbols],
+            "imports": self.imports,
+            "parse_time": self.parse_time,
+            "error": self.error,
         }
 
     @classmethod
-    def from_dict(cls, data: Dict[str, Any]) -> 'ParsedFile':
+    def from_dict(cls, data: Dict[str, Any]) -> "ParsedFile":
         """
         Create ParsedFile from dictionary.
 
@@ -264,8 +269,8 @@ def from_dict(cls, data: Dict[str, Any]) -> 'ParsedFile':
             ParsedFile instance
         """
         data = data.copy()
-        if 'symbols' in data:
-            data['symbols'] = [Symbol.from_dict(s) for s in data['symbols']]
+        if "symbols" in data:
+            data["symbols"] = [Symbol.from_dict(s) for s in data["symbols"]]
         return cls(**data)
 
 
@@ -284,6 +289,7 @@ class CodebaseIndex:
         total_symbols: Total number of symbols across all files
         indexed_at: ISO 8601 timestamp of when indexing completed
     """
+
     root_path: str
     files: Dict[str, ParsedFile] = field(default_factory=dict)
     total_files: int = 0
@@ -362,15 +368,15 @@ def to_dict(self) -> Dict[str, Any]:
             Dictionary representation
         """
         return {
-            'root_path': self.root_path,
-            'files': {path: f.to_dict() for path, f in self.files.items()},
-            'total_files': self.total_files,
-            'total_symbols': self.total_symbols,
-            'indexed_at': self.indexed_at,
+            "root_path": self.root_path,
+            "files": {path: f.to_dict() for path, f in self.files.items()},
+            "total_files": self.total_files,
+            "total_symbols": self.total_symbols,
+            "indexed_at": self.indexed_at,
         }
 
     @classmethod
-    def from_dict(cls, data: Dict[str, Any]) -> 'CodebaseIndex':
+    def from_dict(cls, data: Dict[str, Any]) -> "CodebaseIndex":
         """
         Create CodebaseIndex from dictionary.
 
@@ -381,8 +387,8 @@ def from_dict(cls, data: Dict[str, Any]) -> 'CodebaseIndex':
             CodebaseIndex instance
         """
         data = data.copy()
-        if 'files' in data:
-            data['files'] = {path: ParsedFile.from_dict(f) for path, f in data['files'].items()}
+        if "files" in data:
+            data["files"] = {path: ParsedFile.from_dict(f) for path, f in data["files"].items()}
         return cls(**data)
 
 
@@ -403,6 +409,7 @@ class MemoryType(Enum):
         DOC: Documentation snippets, README content, API docs
         NOTE: General notes, reminders, TODOs
     """
+
     CONVERSATION = "conversation"
     STATUS = "status"
     DECISION = "decision"
@@ -414,7 +421,7 @@ def __str__(self) -> str:
         return self.value
 
     @classmethod
-    def from_string(cls, value: str) -> 'MemoryType':
+    def from_string(cls, value: str) -> "MemoryType":
         for member in cls:
             if member.value == value.lower():
                 return member
@@ -441,6 +448,7 @@ class Memory:
         source: Origin of memory ("user", "agent", "auto", "import")
         metadata: Extensible dictionary for additional data
     """
+
     id: str
     content: str
     memory_type: MemoryType
@@ -469,24 +477,24 @@ def __post_init__(self):
     def to_dict(self) -> Dict[str, Any]:
         """Convert Memory to dictionary for serialization."""
         return {
-            'id': self.id,
-            'content': self.content,
-            'memory_type': self.memory_type.value,
-            'project': self.project,
-            'tags': self.tags,
-            'created_at': self.created_at,
-            'accessed_at': self.accessed_at,
-            'ttl': self.ttl,
-            'source': self.source,
-            'metadata': self.metadata
+            "id": self.id,
+            "content": self.content,
+            "memory_type": self.memory_type.value,
+            "project": self.project,
+            "tags": self.tags,
+            "created_at": self.created_at,
+            "accessed_at": self.accessed_at,
+            "ttl": self.ttl,
+            "source": self.source,
+            "metadata": self.metadata,
         }
 
     @classmethod
-    def from_dict(cls, data: Dict[str, Any]) -> 'Memory':
+    def from_dict(cls, data: Dict[str, Any]) -> "Memory":
         """Create Memory from dictionary."""
         data = data.copy()
-        if 'memory_type' in data and isinstance(data['memory_type'], str):
-            data['memory_type'] = MemoryType.from_string(data['memory_type'])
+        if "memory_type" in data and isinstance(data["memory_type"], str):
+            data["memory_type"] = MemoryType.from_string(data["memory_type"])
         return cls(**data)
 
     def touch(self) -> None:
diff --git a/src/codegrok_mcp/indexing/__init__.py b/src/codegrok_mcp/indexing/__init__.py
index 7c5dfe9..e986929 100644
--- a/src/codegrok_mcp/indexing/__init__.py
+++ b/src/codegrok_mcp/indexing/__init__.py
@@ -47,20 +47,20 @@
 
 __all__ = [
     # Embedding service
-    'EmbeddingService',
-    'get_embedding_service',
-    'reset_embedding_service',
-    'ChromaDBEmbeddingFunction',
-    'embed',
-    'EMBEDDING_MODELS',
-    'DEFAULT_MODEL',
+    "EmbeddingService",
+    "get_embedding_service",
+    "reset_embedding_service",
+    "ChromaDBEmbeddingFunction",
+    "embed",
+    "EMBEDDING_MODELS",
+    "DEFAULT_MODEL",
     # Source retriever
-    'SourceRetriever',
-    'CodeChunk',
-    'SUPPORTED_EXTENSIONS',
-    'count_codebase_files',
+    "SourceRetriever",
+    "CodeChunk",
+    "SUPPORTED_EXTENSIONS",
+    "count_codebase_files",
     # Parallel indexing
-    'parallel_parse_files',
-    'ParseResult',
-    'ParallelProgress',
+    "parallel_parse_files",
+    "ParseResult",
+    "ParallelProgress",
 ]
diff --git a/src/codegrok_mcp/indexing/embedding_service.py b/src/codegrok_mcp/indexing/embedding_service.py
index 712661c..fba3405 100644
--- a/src/codegrok_mcp/indexing/embedding_service.py
+++ b/src/codegrok_mcp/indexing/embedding_service.py
@@ -40,6 +40,7 @@ def _import_dependencies():
     if _sentence_transformers is None:
         try:
             import sentence_transformers
+
             _sentence_transformers = sentence_transformers
         except ImportError:  # pragma: no cover
             raise ImportError(
@@ -50,11 +51,11 @@ def _import_dependencies():
     if _torch is None:
         try:
             import torch
+
             _torch = torch
         except ImportError:  # pragma: no cover
             raise ImportError(
-                "PyTorch is required for native embedding. "
-                "Install with: pip install torch"
+                "PyTorch is required for native embedding. " "Install with: pip install torch"
             )
 
     return _sentence_transformers, _torch
@@ -64,27 +65,27 @@ def _import_dependencies():
 EMBEDDING_MODELS = {
     # Default: Lightweight code embedding - efficient (137M params, ~521MB)
     # SOTA on CodeSearchNet for its size class
-    'coderankembed': {
-        'hf_name': 'nomic-ai/CodeRankEmbed',
-        'dimensions': 768,
-        'max_seq_length': 8192,
-        'trust_remote_code': True,
-        'prompt_prefix': '',
-        'query_prefix': 'Represent this query for searching relevant code: ',
+    "coderankembed": {
+        "hf_name": "nomic-ai/CodeRankEmbed",
+        "dimensions": 768,
+        "max_seq_length": 8192,
+        "trust_remote_code": True,
+        "prompt_prefix": "",
+        "query_prefix": "Represent this query for searching relevant code: ",
     },
     # Example template - copy this to add your own model
-    'my-new-model': {
-        'hf_name': 'organization/model-name',  # HuggingFace model ID
-        'dimensions': 768,                      # Output vector dimensions
-        'max_seq_length': 512,                  # Max input tokens
-        'trust_remote_code': False,             # True if model needs custom code
-        'prompt_prefix': '',                    # Prepended to documents
-        'query_prefix': '',                     # Prepended to queries
+    "my-new-model": {
+        "hf_name": "organization/model-name",  # HuggingFace model ID
+        "dimensions": 768,  # Output vector dimensions
+        "max_seq_length": 512,  # Max input tokens
+        "trust_remote_code": False,  # True if model needs custom code
+        "prompt_prefix": "",  # Prepended to documents
+        "query_prefix": "",  # Prepended to queries
     },
 }
 
 # Default model - CodeRankEmbed (137M params, SOTA for size, code-optimized)
-DEFAULT_MODEL = 'coderankembed'
+DEFAULT_MODEL = "coderankembed"
 
 
 class EmbeddingService:
@@ -135,17 +136,17 @@ def __init__(
         else:
             # Assume it's a HuggingFace model name
             self.config = {
-                'hf_name': model_name,
-                'dimensions': None,  # Will be set after loading
-                'max_seq_length': 512,
-                'trust_remote_code': False,
-                'prompt_prefix': '',
-                'query_prefix': '',
+                "hf_name": model_name,
+                "dimensions": None,  # Will be set after loading
+                "max_seq_length": 512,
+                "trust_remote_code": False,
+                "prompt_prefix": "",
+                "query_prefix": "",
             }
 
         # Determine device
         if device is None:
-            device = 'cuda' if _torch.cuda.is_available() else 'cpu'  # pragma: no cover
+            device = "cuda" if _torch.cuda.is_available() else "cpu"  # pragma: no cover
         self.device = device
 
         # Thread safety
@@ -155,11 +156,11 @@ def __init__(
 
         # Stats
         self.stats = {
-            'total_embeddings': 0,
-            'total_batches': 0,
-            'total_time': 0.0,
-            'cache_hits': 0,
-            'cache_misses': 0,
+            "total_embeddings": 0,
+            "total_batches": 0,
+            "total_time": 0.0,
+            "cache_hits": 0,
+            "cache_misses": 0,
         }
 
         # Cache directory
@@ -182,36 +183,36 @@ def _load_model(self):
                 print(f"Device: {self.device}")
 
             model_kwargs = {
-                'device': self.device,
+                "device": self.device,
             }
 
             if self.cache_dir:
-                model_kwargs['cache_folder'] = self.cache_dir
+                model_kwargs["cache_folder"] = self.cache_dir
 
-            if self.config.get('trust_remote_code'):
-                model_kwargs['trust_remote_code'] = True
+            if self.config.get("trust_remote_code"):
+                model_kwargs["trust_remote_code"] = True
 
             # Suppress stdout/stderr during model loading to prevent
             # "<All keys matched successfully>" message from appearing.
             # This message comes from tqdm.write() during weight loading.
             import sys
             import io
+
             old_stdout = sys.stdout
             old_stderr = sys.stderr
             sys.stdout = io.StringIO()
             sys.stderr = io.StringIO()
             try:
                 self._model = _sentence_transformers.SentenceTransformer(
-                    self.config['hf_name'],
-                    **model_kwargs
+                    self.config["hf_name"], **model_kwargs
                 )
             finally:
                 sys.stdout = old_stdout
                 sys.stderr = old_stderr
 
             # Update dimensions if not set
-            if self.config['dimensions'] is None:
-                self.config['dimensions'] = self._model.get_sentence_embedding_dimension()
+            if self.config["dimensions"] is None:
+                self.config["dimensions"] = self._model.get_sentence_embedding_dimension()
 
             self._model_loaded = True
             if self.show_progress:
@@ -221,7 +222,7 @@ def _load_model(self):
     def dimensions(self) -> int:
         """Get embedding dimensions."""
         self._load_model()
-        return self.config['dimensions']
+        return self.config["dimensions"]
 
     def _embed_single_uncached(self, text: str, is_query: bool) -> tuple:
         """
@@ -232,7 +233,7 @@ def _embed_single_uncached(self, text: str, is_query: bool) -> tuple:
         self._load_model()
 
         # Add prefix if configured
-        prefix = self.config['query_prefix'] if is_query else self.config['prompt_prefix']
+        prefix = self.config["query_prefix"] if is_query else self.config["prompt_prefix"]
         if prefix:
             text = prefix + text
 
@@ -270,10 +271,10 @@ def embed(self, text: str, is_query: bool = False) -> List[float]:
         # Update stats
         cache_info_after = self._embed_cached.cache_info()
         if cache_info_after.hits > cache_info_before.hits:
-            self.stats['cache_hits'] += 1
+            self.stats["cache_hits"] += 1
         else:
-            self.stats['cache_misses'] += 1
-            self.stats['total_embeddings'] += 1
+            self.stats["cache_misses"] += 1
+            self.stats["total_embeddings"] += 1
 
         return list(embedding_tuple)
 
@@ -305,7 +306,7 @@ def embed_batch(
         batch_size = min(batch_size, self.max_batch_size)
 
         # Add prefix if configured
-        prefix = self.config['query_prefix'] if is_query else self.config['prompt_prefix']
+        prefix = self.config["query_prefix"] if is_query else self.config["prompt_prefix"]
         if prefix:
             texts = [prefix + t for t in texts]
 
@@ -325,14 +326,14 @@ def embed_batch(
         elapsed = time.time() - start_time
 
         # Update stats
-        self.stats['total_embeddings'] += len(texts)
-        self.stats['total_batches'] += (len(texts) + batch_size - 1) // batch_size
-        self.stats['total_time'] += elapsed
+        self.stats["total_embeddings"] += len(texts)
+        self.stats["total_batches"] += (len(texts) + batch_size - 1) // batch_size
+        self.stats["total_time"] += elapsed
 
         # Memory management - trigger GC periodically
-        if self.stats['total_batches'] % 100 == 0:
+        if self.stats["total_batches"] % 100 == 0:
             gc.collect()
-            if self.device == 'cuda':  # pragma: no cover
+            if self.device == "cuda":  # pragma: no cover
                 _torch.cuda.empty_cache()
 
         # Convert to list of lists (ChromaDB compatible)
@@ -341,39 +342,39 @@ def embed_batch(
     def get_stats(self) -> dict:
         """Get embedding statistics."""
         stats = self.stats.copy()
-        if stats['total_time'] > 0:
-            stats['embeddings_per_second'] = stats['total_embeddings'] / stats['total_time']
+        if stats["total_time"] > 0:
+            stats["embeddings_per_second"] = stats["total_embeddings"] / stats["total_time"]
         else:
-            stats['embeddings_per_second'] = 0
+            stats["embeddings_per_second"] = 0
         return stats
 
     def get_cache_stats(self) -> dict:
         """Get embedding cache statistics."""
         cache_info = self._embed_cached.cache_info()
-        total_requests = self.stats['cache_hits'] + self.stats['cache_misses']
-        hit_rate = self.stats['cache_hits'] / total_requests if total_requests > 0 else 0.0
+        total_requests = self.stats["cache_hits"] + self.stats["cache_misses"]
+        hit_rate = self.stats["cache_hits"] / total_requests if total_requests > 0 else 0.0
         return {
-            'hits': self.stats['cache_hits'],
-            'misses': self.stats['cache_misses'],
-            'hit_rate': f"{hit_rate:.1%}",
-            'size': cache_info.currsize,
-            'maxsize': cache_info.maxsize,
+            "hits": self.stats["cache_hits"],
+            "misses": self.stats["cache_misses"],
+            "hit_rate": f"{hit_rate:.1%}",
+            "size": cache_info.currsize,
+            "maxsize": cache_info.maxsize,
         }
 
     def clear_cache(self):
         """Clear the embedding cache."""
         self._embed_cached.cache_clear()
-        self.stats['cache_hits'] = 0
-        self.stats['cache_misses'] = 0
+        self.stats["cache_hits"] = 0
+        self.stats["cache_misses"] = 0
 
     def reset_stats(self):
         """Reset statistics."""
         self.stats = {
-            'total_embeddings': 0,
-            'total_batches': 0,
-            'total_time': 0.0,
-            'cache_hits': 0,
-            'cache_misses': 0,
+            "total_embeddings": 0,
+            "total_batches": 0,
+            "total_time": 0.0,
+            "cache_hits": 0,
+            "cache_misses": 0,
         }
 
     def unload(self):
@@ -385,7 +386,7 @@ def unload(self):
                 self._model_loaded = False
                 self._embed_cached.cache_clear()  # Clear embedding cache
                 gc.collect()
-                if self.device == 'cuda':  # pragma: no cover
+                if self.device == "cuda":  # pragma: no cover
                     _torch.cuda.empty_cache()
 
 
@@ -394,10 +395,7 @@ def unload(self):
 _singleton_lock = threading.Lock()
 
 
-def get_embedding_service(
-    model_name: str = DEFAULT_MODEL,
-    **kwargs
-) -> EmbeddingService:
+def get_embedding_service(model_name: str = DEFAULT_MODEL, **kwargs) -> EmbeddingService:
     """
     Get embedding service instance for the specified model.
 
@@ -468,7 +466,9 @@ def __call__(self, input: List[str]) -> List[List[float]]:
 
 
 # Convenience function for quick embedding
-def embed(texts: Union[str, List[str]], is_query: bool = False) -> Union[List[float], List[List[float]]]:
+def embed(
+    texts: Union[str, List[str]], is_query: bool = False
+) -> Union[List[float], List[List[float]]]:
     """
     Quick embedding function.
 
diff --git a/src/codegrok_mcp/indexing/memory_retriever.py b/src/codegrok_mcp/indexing/memory_retriever.py
index 89d82f3..c051f3e 100644
--- a/src/codegrok_mcp/indexing/memory_retriever.py
+++ b/src/codegrok_mcp/indexing/memory_retriever.py
@@ -22,14 +22,13 @@
 from codegrok_mcp.indexing.embedding_service import get_embedding_service, EmbeddingService
 from codegrok_mcp.core.models import Memory, MemoryType
 
-
 # TTL duration mappings
 TTL_DURATIONS = {
-    "session": timedelta(hours=24),      # Cleared on session end or after 24h
+    "session": timedelta(hours=24),  # Cleared on session end or after 24h
     "day": timedelta(days=1),
     "week": timedelta(weeks=1),
     "month": timedelta(days=30),
-    "permanent": None                     # Never expires
+    "permanent": None,  # Never expires
 }
 
 
@@ -55,7 +54,7 @@ def __init__(
         embedding_model: str = "coderankembed",
         verbose: bool = False,
         persist_path: Optional[str] = None,
-        embedding_service: Optional[EmbeddingService] = None
+        embedding_service: Optional[EmbeddingService] = None,
     ):
         """
         Initialize the memory retriever.
@@ -75,8 +74,7 @@ def __init__(
 
         # Reuse embedding service singleton (same model as code embeddings)
         self.embedding_service = embedding_service or get_embedding_service(
-            embedding_model,
-            show_progress=False
+            embedding_model, show_progress=False
         )
 
         # Initialize ChromaDB (reuse existing client if possible)
@@ -89,15 +87,11 @@ def __init__(
         # Get or create memories collection
         self.collection = self.chroma_client.get_or_create_collection(
             name=self.COLLECTION_NAME,
-            metadata={"description": "Memory storage for conversations, status, decisions"}
+            metadata={"description": "Memory storage for conversations, status, decisions"},
         )
 
         # Statistics
-        self.stats = {
-            'total_memories': 0,
-            'by_type': {},
-            'last_cleanup': None
-        }
+        self.stats = {"total_memories": 0, "by_type": {}, "last_cleanup": None}
 
         self._load_stats()
 
@@ -112,9 +106,9 @@ def _load_stats(self):
             metadata_path = Path(self.persist_path).parent / self.METADATA_FILE
             if metadata_path.exists():
                 try:
-                    with open(metadata_path, 'r') as f:
+                    with open(metadata_path, "r") as f:
                         data = json.load(f)
-                        self.stats = data.get('stats', self.stats)
+                        self.stats = data.get("stats", self.stats)
                 except Exception:
                     pass
 
@@ -123,8 +117,8 @@ def _save_stats(self):
         if self.persist_path:
             metadata_path = Path(self.persist_path).parent / self.METADATA_FILE
             try:
-                with open(metadata_path, 'w') as f:
-                    json.dump({'stats': self.stats}, f, indent=2)
+                with open(metadata_path, "w") as f:
+                    json.dump({"stats": self.stats}, f, indent=2)
             except Exception:
                 pass
 
@@ -135,7 +129,7 @@ def remember(
         tags: List[str] = None,
         ttl: str = "permanent",
         source: str = "user",
-        metadata: Dict[str, Any] = None
+        metadata: Dict[str, Any] = None,
     ) -> Memory:
         """
         Store a new memory with automatic embedding.
@@ -160,7 +154,7 @@ def remember(
             tags=tags or [],
             ttl=ttl,
             source=source,
-            metadata=metadata or {}
+            metadata=metadata or {},
         )
 
         # Generate embedding
@@ -168,13 +162,13 @@ def remember(
 
         # Build ChromaDB metadata (flatten for ChromaDB compatibility)
         chroma_metadata = {
-            'memory_type': memory.memory_type.value,
-            'project': memory.project,
-            'tags': ','.join(memory.tags),  # ChromaDB doesn't support list values
-            'created_at': memory.created_at,
-            'accessed_at': memory.accessed_at,
-            'ttl': memory.ttl,
-            'source': memory.source
+            "memory_type": memory.memory_type.value,
+            "project": memory.project,
+            "tags": ",".join(memory.tags),  # ChromaDB doesn't support list values
+            "created_at": memory.created_at,
+            "accessed_at": memory.accessed_at,
+            "ttl": memory.ttl,
+            "source": memory.source,
         }
 
         # Store in ChromaDB
@@ -182,13 +176,13 @@ def remember(
             ids=[memory.id],
             embeddings=[embedding],
             documents=[memory.content],
-            metadatas=[chroma_metadata]
+            metadatas=[chroma_metadata],
         )
 
         # Update stats
-        self.stats['total_memories'] = self.collection.count()
+        self.stats["total_memories"] = self.collection.count()
         type_key = memory.memory_type.value
-        self.stats['by_type'][type_key] = self.stats['by_type'].get(type_key, 0) + 1
+        self.stats["by_type"][type_key] = self.stats["by_type"].get(type_key, 0) + 1
         self._save_stats()
 
         self._log(f"Stored memory: {memory.id[:8]}... ({memory_type})")
@@ -202,7 +196,7 @@ def recall(
         tags: Optional[List[str]] = None,
         n_results: int = 5,
         time_range: Optional[str] = None,
-        min_relevance: float = 0.0
+        min_relevance: float = 0.0,
     ) -> List[Dict[str, Any]]:
         """
         Semantically search memories.
@@ -241,33 +235,32 @@ def recall(
         if memory_type:
             where_clauses.append({"memory_type": memory_type})
 
-        chroma_where = (
-            {"$and": where_clauses} if len(where_clauses) > 1
-            else where_clauses[0]
-        )
+        chroma_where = {"$and": where_clauses} if len(where_clauses) > 1 else where_clauses[0]
 
         try:
             results = self.collection.query(
                 query_embeddings=[query_embedding],
                 n_results=min(n_results * 2, 50),  # Over-fetch for tag/time filtering
                 where=chroma_where,
-                include=["documents", "metadatas", "distances"]
+                include=["documents", "metadatas", "distances"],
             )
         except Exception as e:
             self._log(f"Query error: {e}")
             return []
 
-        if not results['ids'] or not results['ids'][0]:
+        if not results["ids"] or not results["ids"][0]:
             return []
 
         # Process results
         memories = []
-        for i, (id_, doc, metadata, distance) in enumerate(zip(
-            results['ids'][0],
-            results['documents'][0],
-            results['metadatas'][0],
-            results['distances'][0]
-        )):
+        for i, (id_, doc, metadata, distance) in enumerate(
+            zip(
+                results["ids"][0],
+                results["documents"][0],
+                results["metadatas"][0],
+                results["distances"][0],
+            )
+        ):
             # Convert distance to relevance (ChromaDB uses L2 distance)
             # Lower distance = higher relevance
             relevance = 1.0 / (1.0 + distance)
@@ -277,10 +270,10 @@ def recall(
 
             # Apply time_range filter (post-filter - ChromaDB date ops are limited)
             if time_cutoff:
-                created_str = metadata.get('created_at', '')
+                created_str = metadata.get("created_at", "")
                 if created_str:
                     try:
-                        created = datetime.fromisoformat(created_str.replace('Z', '+00:00'))
+                        created = datetime.fromisoformat(created_str.replace("Z", "+00:00"))
                         if created < time_cutoff:
                             continue
                     except ValueError:
@@ -288,19 +281,21 @@ def recall(
 
             # Apply tags filter (post-filter - ChromaDB doesn't support list contains)
             if tags:
-                stored_tags = set(metadata.get('tags', '').split(','))
+                stored_tags = set(metadata.get("tags", "").split(","))
                 if not stored_tags.intersection(set(tags)):
                     continue
 
-            memories.append({
-                'id': id_,
-                'content': doc,
-                'memory_type': metadata.get('memory_type'),
-                'tags': metadata.get('tags', '').split(',') if metadata.get('tags') else [],
-                'created_at': metadata.get('created_at'),
-                'relevance': round(relevance, 3),
-                'source': metadata.get('source', 'unknown')
-            })
+            memories.append(
+                {
+                    "id": id_,
+                    "content": doc,
+                    "memory_type": metadata.get("memory_type"),
+                    "tags": metadata.get("tags", "").split(",") if metadata.get("tags") else [],
+                    "created_at": metadata.get("created_at"),
+                    "relevance": round(relevance, 3),
+                    "source": metadata.get("source", "unknown"),
+                }
+            )
 
             if len(memories) >= n_results:
                 break
@@ -310,10 +305,7 @@ def recall(
             now = datetime.now(timezone.utc).isoformat()
             for mem in memories:
                 try:
-                    self.collection.update(
-                        ids=[mem['id']],
-                        metadatas=[{'accessed_at': now}]
-                    )
+                    self.collection.update(ids=[mem["id"]], metadatas=[{"accessed_at": now}])
                 except Exception:
                     pass  # Non-critical
 
@@ -326,7 +318,7 @@ def forget(
         memory_id: Optional[str] = None,
         memory_type: Optional[str] = None,
         tags: Optional[List[str]] = None,
-        older_than: Optional[str] = None
+        older_than: Optional[str] = None,
     ) -> Dict[str, int]:
         """
         Remove memories matching criteria.
@@ -346,7 +338,7 @@ def forget(
             # Delete specific memory - verify it exists first
             try:
                 existing = self.collection.get(ids=[memory_id])
-                if existing['ids']:
+                if existing["ids"]:
                     self.collection.delete(ids=[memory_id])
                     deleted = 1
                 else:
@@ -364,13 +356,12 @@ def forget(
             if older_than or tags:
                 # Get all memories for this project
                 all_memories = self.collection.get(
-                    where={"project": self.project_id},
-                    include=["metadatas"]
+                    where={"project": self.project_id}, include=["metadatas"]
                 )
 
                 ids_to_delete = []
 
-                for id_, metadata in zip(all_memories['ids'], all_memories['metadatas']):
+                for id_, metadata in zip(all_memories["ids"], all_memories["metadatas"]):
                     should_delete = False
 
                     # Check older_than
@@ -379,22 +370,22 @@ def forget(
                             "1d": timedelta(days=1),
                             "7d": timedelta(days=7),
                             "30d": timedelta(days=30),
-                            "1y": timedelta(days=365)
+                            "1y": timedelta(days=365),
                         }
                         if older_than in duration_map:
                             cutoff = datetime.now(timezone.utc) - duration_map[older_than]
-                            created = datetime.fromisoformat(metadata.get('created_at', ''))
+                            created = datetime.fromisoformat(metadata.get("created_at", ""))
                             if created < cutoff:
                                 should_delete = True
 
                     # Check tags
                     if tags:
-                        stored_tags = set(metadata.get('tags', '').split(','))
+                        stored_tags = set(metadata.get("tags", "").split(","))
                         if stored_tags.intersection(set(tags)):
                             should_delete = True
 
                     # Check memory_type (if also specified)
-                    if memory_type and metadata.get('memory_type') != memory_type:
+                    if memory_type and metadata.get("memory_type") != memory_type:
                         should_delete = False
 
                     if should_delete:
@@ -409,21 +400,16 @@ def forget(
                 try:
                     # Get IDs first, then delete (ChromaDB requires $and for multiple conditions)
                     to_delete = self.collection.get(
-                        where={
-                            "$and": [
-                                {"memory_type": memory_type},
-                                {"project": self.project_id}
-                            ]
-                        }
+                        where={"$and": [{"memory_type": memory_type}, {"project": self.project_id}]}
                     )
-                    if to_delete['ids']:
-                        self.collection.delete(ids=to_delete['ids'])
-                        deleted = len(to_delete['ids'])
+                    if to_delete["ids"]:
+                        self.collection.delete(ids=to_delete["ids"])
+                        deleted = len(to_delete["ids"])
                 except Exception as e:
                     self._log(f"Delete error: {e}")
 
         # Update stats
-        self.stats['total_memories'] = self.collection.count()
+        self.stats["total_memories"] = self.collection.count()
         self._save_stats()
 
         self._log(f"Forgot {deleted} memories")
@@ -444,34 +430,33 @@ def cleanup_expired(self) -> Dict[str, int]:
 
         # Get all memories for this project
         all_memories = self.collection.get(
-            where={"project": self.project_id},
-            include=["metadatas"]
+            where={"project": self.project_id}, include=["metadatas"]
         )
 
         ids_to_delete = []
 
-        for id_, metadata in zip(all_memories['ids'], all_memories['metadatas']):
-            ttl = metadata.get('ttl', 'permanent')
+        for id_, metadata in zip(all_memories["ids"], all_memories["metadatas"]):
+            ttl = metadata.get("ttl", "permanent")
 
-            if ttl == 'permanent':
+            if ttl == "permanent":
                 continue
 
             duration = TTL_DURATIONS.get(ttl)
             if not duration:
                 continue
 
-            created = datetime.fromisoformat(metadata.get('created_at', now.isoformat()))
+            created = datetime.fromisoformat(metadata.get("created_at", now.isoformat()))
             if now - created > duration:
                 ids_to_delete.append(id_)
-                mem_type = metadata.get('memory_type', 'unknown')
+                mem_type = metadata.get("memory_type", "unknown")
                 cleaned[mem_type] = cleaned.get(mem_type, 0) + 1
 
         if ids_to_delete:
             self.collection.delete(ids=ids_to_delete)
 
         # Update stats
-        self.stats['total_memories'] = self.collection.count()
-        self.stats['last_cleanup'] = now.isoformat()
+        self.stats["total_memories"] = self.collection.count()
+        self.stats["last_cleanup"] = now.isoformat()
         self._save_stats()
 
         total_cleaned = sum(cleaned.values())
@@ -494,22 +479,21 @@ def get_stats(self) -> Dict[str, Any]:
         for mem_type in MemoryType:
             try:
                 # ChromaDB requires $and for multiple conditions
-                count = len(self.collection.get(
-                    where={
-                        "$and": [
-                            {"memory_type": mem_type.value},
-                            {"project": self.project_id}
-                        ]
-                    }
-                )['ids'])
+                count = len(
+                    self.collection.get(
+                        where={
+                            "$and": [{"memory_type": mem_type.value}, {"project": self.project_id}]
+                        }
+                    )["ids"]
+                )
                 if count > 0:
                     by_type[mem_type.value] = count
             except Exception:
                 pass
 
         return {
-            'total_memories': total,
-            'by_type': by_type,
-            'project': self.project_id,
-            'last_cleanup': self.stats.get('last_cleanup')
+            "total_memories": total,
+            "by_type": by_type,
+            "project": self.project_id,
+            "last_cleanup": self.stats.get("last_cleanup"),
         }
diff --git a/src/codegrok_mcp/indexing/parallel_indexer.py b/src/codegrok_mcp/indexing/parallel_indexer.py
index 936629c..c6d5da6 100644
--- a/src/codegrok_mcp/indexing/parallel_indexer.py
+++ b/src/codegrok_mcp/indexing/parallel_indexer.py
@@ -37,6 +37,7 @@ class ParseResult:
         success: True if parsing succeeded, False otherwise
         error: Error message if parsing failed, None otherwise
     """
+
     filepath: str
     symbols: List[Symbol]
     success: bool
@@ -55,6 +56,7 @@ class ParallelProgress:
         _completed: Number of completed items (access via .completed property)
         _errors: Number of errors encountered (access via .errors property)
     """
+
     total: int
     _completed: int = 0
     _errors: int = 0
@@ -95,10 +97,7 @@ def increment_errors(self) -> int:
             return self._errors
 
 
-def parse_file_worker(
-    filepath: Path,
-    parser_factory: ThreadLocalParserFactory
-) -> ParseResult:
+def parse_file_worker(filepath: Path, parser_factory: ThreadLocalParserFactory) -> ParseResult:
     """
     Worker function for parallel file parsing.
 
@@ -120,15 +119,10 @@ def parse_file_worker(
             filepath=str(filepath),
             symbols=list(parsed.symbols),
             success=parsed.error is None,
-            error=parsed.error
+            error=parsed.error,
         )
     except Exception as e:
-        return ParseResult(
-            filepath=str(filepath),
-            symbols=[],
-            success=False,
-            error=str(e)
-        )
+        return ParseResult(filepath=str(filepath), symbols=[], success=False, error=str(e))
 
 
 def parallel_parse_files(
@@ -185,10 +179,7 @@ def emit(event_type: str, data: dict):
 
     with ThreadPoolExecutor(max_workers=max_workers) as executor:
         # Submit all tasks
-        future_to_file = {
-            executor.submit(parse_file_worker, f, parser_factory): f
-            for f in files
-        }
+        future_to_file = {executor.submit(parse_file_worker, f, parser_factory): f for f in files}
 
         # Process results as they complete
         for future in as_completed(future_to_file):
@@ -200,33 +191,33 @@ def emit(event_type: str, data: dict):
                 if result.success and result.symbols:
                     with symbols_lock:
                         all_symbols.extend(result.symbols)
-                    emit("file_parsed", {
-                        "path": result.filepath,
-                        "symbols": len(result.symbols),
-                        "index": completed,
-                        "total": progress.total
-                    })
+                    emit(
+                        "file_parsed",
+                        {
+                            "path": result.filepath,
+                            "symbols": len(result.symbols),
+                            "index": completed,
+                            "total": progress.total,
+                        },
+                    )
                 elif result.error:
                     progress.increment_errors()
-                    emit("parse_error", {
-                        "path": result.filepath,
-                        "error": result.error
-                    })
+                    emit("parse_error", {"path": result.filepath, "error": result.error})
                 else:
                     # File parsed but no symbols found (not an error)
-                    emit("file_parsed", {
-                        "path": result.filepath,
-                        "symbols": 0,
-                        "index": completed,
-                        "total": progress.total
-                    })
+                    emit(
+                        "file_parsed",
+                        {
+                            "path": result.filepath,
+                            "symbols": 0,
+                            "index": completed,
+                            "total": progress.total,
+                        },
+                    )
 
             except Exception as e:  # pragma: no cover
                 progress.increment_completed()
                 progress.increment_errors()
-                emit("parse_error", {
-                    "path": str(filepath),
-                    "error": str(e)
-                })
+                emit("parse_error", {"path": str(filepath), "error": str(e)})
 
     return all_symbols, progress.errors
diff --git a/src/codegrok_mcp/indexing/source_retriever.py b/src/codegrok_mcp/indexing/source_retriever.py
index c403b74..6160b64 100644
--- a/src/codegrok_mcp/indexing/source_retriever.py
+++ b/src/codegrok_mcp/indexing/source_retriever.py
@@ -29,12 +29,16 @@
 """
 
 import json
+import logging
+import os
 import time
 from datetime import datetime
 from pathlib import Path
-from typing import List, Dict, Any, Optional
+from typing import List, Dict, Any, Optional, Callable
 from dataclasses import dataclass
 
+import pathspec
+
 try:
     import chromadb
 except ImportError:
@@ -47,26 +51,59 @@
 from codegrok_mcp.parsers.language_configs import get_supported_extensions, get_language_for_file
 from codegrok_mcp.core.models import Symbol, SymbolType
 
-
 # Derived from authoritative EXTENSION_MAP in language_configs.py (30+ extensions, 9 languages)
 # This eliminates duplication and ensures extensions stay in sync
 SUPPORTED_EXTENSIONS = list(get_supported_extensions())
 SUPPORTED_EXTENSIONS_SET = set(SUPPORTED_EXTENSIONS)  # For O(1) lookup
 
 # Common directories to skip during file discovery
-SKIP_DIRS = {'.git', 'node_modules', '__pycache__', '.codegrok', 'venv', '.venv',
-             '.tox', '.mypy_cache', '.pytest_cache', 'dist', 'build', '.eggs'}
-
-
-def discover_files(codebase_path: Path, extensions: set = None) -> List[Path]:
-    """Single-pass file discovery with extension filtering.
-
-    This is 30x+ faster than calling rglob() once per extension because it
-    traverses the directory tree only once instead of 30+ times.
+SKIP_DIRS = {
+    ".git",
+    "node_modules",
+    "__pycache__",
+    ".codegrok",
+    "venv",
+    ".venv",
+    ".tox",
+    ".mypy_cache",
+    ".pytest_cache",
+    "dist",
+    "build",
+    ".eggs",
+}
+
+
+def _load_gitignore(directory: Path) -> Optional[pathspec.PathSpec]:
+    """Load a .gitignore file from a directory, returning a PathSpec or None."""
+    gitignore_path = directory / ".gitignore"
+    if gitignore_path.is_file():
+        try:
+            with open(gitignore_path, "r", encoding="utf-8", errors="ignore") as f:
+                return pathspec.PathSpec.from_lines("gitignore", f)
+        except (OSError, IOError):
+            pass
+    return None
+
+
+def discover_files(
+    codebase_path: Path,
+    extensions: set = None,
+    respect_gitignore: bool = True,
+    max_files: int = 200_000,
+    progress_callback: Callable = None,
+) -> List[Path]:
+    """Single-pass file discovery with extension filtering and .gitignore support.
+
+    Uses os.walk() to traverse the directory tree once, pruning ignored
+    directories in-place so they are never descended into. Respects .gitignore
+    patterns (including nested .gitignore files) and a hardcoded SKIP_DIRS set.
 
     Args:
         codebase_path: Path to the codebase root directory.
-        extensions: Set of extensions to include (default: SUPPORTED_EXTENSIONS_SET)
+        extensions: Set of extensions to include (default: SUPPORTED_EXTENSIONS_SET).
+        respect_gitignore: Whether to respect .gitignore patterns (default: True).
+        max_files: Safety limit on number of files to discover (default: 200,000).
+        progress_callback: Optional callback(event_type, data) for progress events.
 
     Returns:
         List of file paths matching the extensions.
@@ -74,16 +111,92 @@ def discover_files(codebase_path: Path, extensions: set = None) -> List[Path]:
     if extensions is None:
         extensions = SUPPORTED_EXTENSIONS_SET
 
-    files = []
-    for path in codebase_path.rglob("*"):
-        # Skip directories in SKIP_DIRS
-        if any(skip_dir in path.parts for skip_dir in SKIP_DIRS):
-            continue
-        if path.is_file() and path.suffix in extensions:
-            files.append(path)
+    # Load root .gitignore
+    gitignore_specs: List[pathspec.PathSpec] = []
+    if respect_gitignore:
+        root_spec = _load_gitignore(codebase_path)
+        if root_spec is not None:
+            gitignore_specs.append(root_spec)
+
+    files: List[Path] = []
+    codebase_str = str(codebase_path)
+
+    for dirpath_str, dirnames, filenames in os.walk(codebase_path, followlinks=False):
+        dirpath = Path(dirpath_str)
+
+        # Load nested .gitignore for this directory (not root)
+        if respect_gitignore and dirpath != codebase_path:
+            nested_spec = _load_gitignore(dirpath)
+            if nested_spec is not None:
+                gitignore_specs.append(nested_spec)
+
+        # Prune directories in-place: remove SKIP_DIRS and gitignored dirs
+        filtered_dirs = []
+        for d in dirnames:
+            if d in SKIP_DIRS:
+                continue
+            if respect_gitignore and gitignore_specs:
+                # Compute relative path with trailing slash for directory matching
+                rel_dir = str((dirpath / d).relative_to(codebase_path)) + "/"
+                if any(spec.match_file(rel_dir) for spec in gitignore_specs):
+                    continue
+            filtered_dirs.append(d)
+        dirnames[:] = filtered_dirs
+
+        # Check files
+        for filename in filenames:
+            filepath = dirpath / filename
+            if filepath.suffix not in extensions:
+                continue
+            if respect_gitignore and gitignore_specs:
+                rel_path = str(filepath.relative_to(codebase_path))
+                if any(spec.match_file(rel_path) for spec in gitignore_specs):
+                    continue
+            files.append(filepath)
+
+            # Report progress every 1000 files
+            if progress_callback and len(files) % 1000 == 0:
+                progress_callback("discovery_progress", {"files_found": len(files)})
+
+            # Safety limit
+            if len(files) >= max_files:
+                logging.warning(
+                    f"discover_files: reached max_files limit ({max_files}). "
+                    f"Stopping discovery. Consider using .gitignore to exclude files."
+                )
+                return files
+
     return files
 
 
+def _save_checkpoint(checkpoint_path: Path, chunks_completed: int, total_chunks: int) -> None:
+    """Atomically save a checkpoint file for resumable indexing."""
+    data = {
+        "chunks_completed": chunks_completed,
+        "total_chunks": total_chunks,
+        "timestamp": datetime.now().isoformat(),
+    }
+    tmp_path = checkpoint_path.with_suffix(".tmp")
+    try:
+        with open(tmp_path, "w") as f:
+            json.dump(data, f)
+        os.replace(str(tmp_path), str(checkpoint_path))
+    except OSError:
+        # Best-effort: if we can't save checkpoint, indexing still continues
+        pass
+
+
+def _load_checkpoint(checkpoint_path: Path) -> Optional[Dict[str, Any]]:
+    """Load a checkpoint file, returning the data dict or None."""
+    if checkpoint_path is None or not checkpoint_path.exists():
+        return None
+    try:
+        with open(checkpoint_path, "r") as f:
+            return json.load(f)
+    except (json.JSONDecodeError, OSError):
+        return None
+
+
 def count_codebase_files(codebase_path: Path) -> int:
     """Quick file count for ETA estimation.
 
@@ -99,6 +212,7 @@ def count_codebase_files(codebase_path: Path) -> int:
 @dataclass
 class CodeChunk:
     """A chunk of code suitable for embedding."""
+
     id: str
     text: str
     filepath: str
@@ -127,14 +241,12 @@ def __init__(
         collection_name: str = "codebase",
         verbose: bool = True,
         persist_path: Optional[str] = None,
-
         # Parallel indexing options (3-5x faster for large codebases)
         parallel: bool = True,  # Enabled by default for better performance
         max_workers: Optional[int] = None,
-
         # Dependency injection for testability
-        parser: Optional['TreeSitterParser'] = None,
-        embedding_service: Optional['EmbeddingService'] = None
+        parser: Optional["TreeSitterParser"] = None,
+        embedding_service: Optional["EmbeddingService"] = None,
     ):
         """
         Initialize the source retriever.
@@ -165,8 +277,7 @@ def __init__(
         # Initialize embedding service (use injected or create default)
         self._log(f"Using native embedding: {embedding_model}")
         self.embedding_service = embedding_service or get_embedding_service(
-            embedding_model,
-            show_progress=verbose  # Only show tqdm progress bar if verbose
+            embedding_model, show_progress=verbose  # Only show tqdm progress bar if verbose
         )
 
         # Initialize ChromaDB (persistent or in-memory)
@@ -182,11 +293,11 @@ def __init__(
 
         # Statistics
         self.stats = {
-            'total_files': 0,
-            'total_symbols': 0,
-            'total_chunks': 0,
-            'parse_errors': 0,
-            'indexing_time': 0.0
+            "total_files": 0,
+            "total_symbols": 0,
+            "total_chunks": 0,
+            "parse_errors": 0,
+            "indexing_time": 0.0,
         }
 
         # Metadata storage for incremental reindexing (file modification times)
@@ -276,20 +387,20 @@ def _create_chunk(self, symbol: Symbol) -> CodeChunk:
             symbol_type=symbol.type.value,
             line_start=symbol.line_start,
             metadata={
-                'filepath': symbol.filepath,
-                'name': symbol.name,
-                'type': symbol.type.value,
-                'line': symbol.line_start,
-                'signature': symbol.signature,
-                'parent': symbol.parent or "",
-                'language': language  # NEW: enables language filtering in search
-            }
+                "filepath": symbol.filepath,
+                "name": symbol.name,
+                "type": symbol.type.value,
+                "line": symbol.line_start,
+                "signature": symbol.signature,
+                "parent": symbol.parent or "",
+                "language": language,  # NEW: enables language filtering in search
+            },
         )
 
     def index_codebase(
         self,
         file_extensions: Optional[List[str]] = None,
-        progress_callback: Optional[callable] = None
+        progress_callback: Optional[callable] = None,
     ):
         """
         Index the entire codebase.
@@ -323,9 +434,9 @@ def emit(event_type: str, data: dict):
 
         # Legacy logging for when no callback provided
         if not progress_callback:
-            self._log("\n" + "="*80)
+            self._log("\n" + "=" * 80)
             self._log("INDEXING CODEBASE")
-            self._log("="*80)
+            self._log("=" * 80)
             self._log(f"Codebase: {self.codebase_path}")
             self._log(f"Extensions: {file_extensions}")
             self._log(f"Embedding model: {self.embedding_model}")
@@ -337,9 +448,11 @@ def emit(event_type: str, data: dict):
             self._log("\nStep 1: Finding files...")
 
         extensions_set = set(file_extensions)
-        all_files = discover_files(self.codebase_path, extensions_set)
+        all_files = discover_files(
+            self.codebase_path, extensions_set, progress_callback=progress_callback
+        )
 
-        self.stats['total_files'] = len(all_files)
+        self.stats["total_files"] = len(all_files)
 
         # Store file modification times for incremental reindexing
         file_mtimes = {}
@@ -348,12 +461,9 @@ def emit(event_type: str, data: dict):
                 file_mtimes[str(filepath)] = filepath.stat().st_mtime
             except OSError:
                 pass  # Skip files that can't be stat'd
-        self._metadata['file_mtimes'] = file_mtimes
+        self._metadata["file_mtimes"] = file_mtimes
 
-        emit("files_found", {
-            "files": all_files,
-            "codebase_path": self.codebase_path
-        })
+        emit("files_found", {"files": all_files, "codebase_path": self.codebase_path})
 
         if not progress_callback:
             self._log(f"Found {len(all_files)} files")
@@ -366,7 +476,9 @@ def emit(event_type: str, data: dict):
         all_symbols = []
 
         # Use parallel parsing if enabled and there are enough files
-        use_parallel = self.parallel and len(all_files) > 50  # Threshold increased for small projects
+        use_parallel = (
+            self.parallel and len(all_files) > 50
+        )  # Threshold increased for small projects
         if use_parallel:
             # Parallel parsing (3-5x faster for large codebases)
             from codegrok_mcp.indexing.parallel_indexer import parallel_parse_files
@@ -375,11 +487,9 @@ def emit(event_type: str, data: dict):
                 self._log(f"  Using parallel parsing with {self.max_workers or 'auto'} workers...")
 
             all_symbols, parse_errors = parallel_parse_files(
-                files=all_files,
-                max_workers=self.max_workers,
-                progress_callback=progress_callback
+                files=all_files, max_workers=self.max_workers, progress_callback=progress_callback
             )
-            self.stats['parse_errors'] = parse_errors
+            self.stats["parse_errors"] = parse_errors
         else:
             # Sequential parsing (original code)
             for i, filepath in enumerate(all_files, 1):
@@ -389,23 +499,26 @@ def emit(event_type: str, data: dict):
                     all_symbols.extend(parsed.symbols)
                     symbols_count = len(parsed.symbols)
 
-                    emit("file_parsed", {
-                        "path": str(filepath),
-                        "symbols": symbols_count,
-                        "index": i,
-                        "total": len(all_files)
-                    })
+                    emit(
+                        "file_parsed",
+                        {
+                            "path": str(filepath),
+                            "symbols": symbols_count,
+                            "index": i,
+                            "total": len(all_files),
+                        },
+                    )
                 except Exception as e:
-                    self.stats['parse_errors'] += 1
+                    self.stats["parse_errors"] += 1
                     emit("parse_error", {"path": str(filepath), "error": str(e)})
-                    if not progress_callback and self.verbose and self.stats['parse_errors'] <= 5:
+                    if not progress_callback and self.verbose and self.stats["parse_errors"] <= 5:
                         self._log(f"  Error parsing {filepath}: {e}")
 
                 # Legacy progress for no callback
                 if not progress_callback and self.verbose and i % 100 == 0:
-                    print(f"  Parsed {i}/{len(all_files)} files...", end='\r')
+                    print(f"  Parsed {i}/{len(all_files)} files...", end="\r")
 
-        self.stats['total_symbols'] = len(all_symbols)
+        self.stats["total_symbols"] = len(all_symbols)
         if not progress_callback:
             self._log(f"\nParsed {len(all_symbols):,} symbols from {len(all_files)} files")
 
@@ -414,28 +527,26 @@ def emit(event_type: str, data: dict):
             self._log("\nStep 3: Creating chunks...")
 
         chunks = [self._create_chunk(symbol) for symbol in all_symbols]
-        self.stats['total_chunks'] = len(chunks)
+        self.stats["total_chunks"] = len(chunks)
 
         emit("chunks_created", {"total": len(chunks)})
 
         if not progress_callback:
             self._log(f"Created {len(chunks):,} chunks")
 
-        # Step 4: Create ChromaDB collection
+        # Step 4: Get or create ChromaDB collection (supports resumption)
         if not progress_callback:
             self._log("\nStep 4: Creating vector database...")
 
-        try:
-            self.chroma_client.delete_collection(self.collection_name)
-        except Exception:
-            pass  # Collection doesn't exist yet - this is expected
-
-        self.collection = self.chroma_client.create_collection(
+        self.collection = self.chroma_client.get_or_create_collection(
             name=self.collection_name,
-            metadata={"description": f"Code embeddings for {self.codebase_path.name}"}
+            metadata={"description": f"Code embeddings for {self.codebase_path.name}"},
         )
 
-        # Step 5: Generate embeddings and store
+        # Track new chunk IDs for stale chunk cleanup after embedding
+        new_chunk_ids = set(chunk.id for chunk in chunks)
+
+        # Step 5: Generate embeddings and store (with checkpointing)
         eta_minutes = len(chunks) / 50 / 60  # ~50 embeddings/sec native
 
         emit("embedding_start", {"total": len(chunks), "eta_minutes": eta_minutes})
@@ -444,18 +555,33 @@ def emit(event_type: str, data: dict):
             self._log(f"\nStep 5: Generating embeddings (ETA: ~{eta_minutes:.1f} minutes)...")
             self._log("(You can interrupt and resume later)")
 
+        # Load checkpoint if available (resume from interrupted indexing)
+        checkpoint_path = (
+            Path(self.persist_path).parent / "checkpoint.json" if self.persist_path else None
+        )
+        start_chunk_idx = 0
+        if checkpoint_path:
+            checkpoint = _load_checkpoint(checkpoint_path)
+            if checkpoint and checkpoint.get("total_chunks") == len(chunks):
+                resume_idx = checkpoint.get("chunks_completed", 0)
+                if self.collection.count() >= resume_idx:
+                    start_chunk_idx = resume_idx
+                    if not progress_callback:
+                        self._log(f"  Resuming from chunk {start_chunk_idx} (checkpoint found)")
+
         batch_size = 100
         embedding_start_time = time.time()
         chunks_per_second = None  # Will be calibrated after first batch
 
-        for i in range(0, len(chunks), batch_size):
-            batch = chunks[i:i + batch_size]
+        for i in range(start_chunk_idx, len(chunks), batch_size):
+            batch = chunks[i : i + batch_size]
             current_count = i + len(batch)
             elapsed = time.time() - embedding_start_time
 
             # Calibrate speed after first batch, then update continuously
-            if elapsed > 0 and current_count > 0:
-                chunks_per_second = current_count / elapsed
+            chunks_processed = current_count - start_chunk_idx
+            if elapsed > 0 and chunks_processed > 0:
+                chunks_per_second = chunks_processed / elapsed
 
             # Calculate remaining time estimate
             remaining_seconds = None
@@ -463,22 +589,28 @@ def emit(event_type: str, data: dict):
                 remaining_chunks = len(chunks) - current_count
                 remaining_seconds = remaining_chunks / chunks_per_second
 
-            emit("embedding_progress", {
-                "current": current_count,
-                "total": len(chunks),
-                "elapsed_seconds": elapsed,
-                "remaining_seconds": remaining_seconds,
-                "chunks_per_second": chunks_per_second
-            })
+            emit(
+                "embedding_progress",
+                {
+                    "current": current_count,
+                    "total": len(chunks),
+                    "elapsed_seconds": elapsed,
+                    "remaining_seconds": remaining_seconds,
+                    "chunks_per_second": chunks_per_second,
+                },
+            )
 
             # Legacy progress for no callback
             if not progress_callback and self.verbose and i % 500 == 0:
                 elapsed = time.time() - start_time
                 rate = i / elapsed if elapsed > 0 else 0
                 remaining = (len(chunks) - i) / rate if rate > 0 else 0
-                print(f"  Embedded {i:,}/{len(chunks):,} chunks "
-                      f"({i/len(chunks)*100:.1f}%) "
-                      f"- ETA: {remaining/60:.1f}m", end='\r')
+                print(
+                    f"  Embedded {i:,}/{len(chunks):,} chunks "
+                    f"({i/len(chunks)*100:.1f}%) "
+                    f"- ETA: {remaining/60:.1f}m",
+                    end="\r",
+                )
 
             try:
                 # Generate embeddings for batch
@@ -487,41 +619,66 @@ def emit(event_type: str, data: dict):
                 # Native batch embedding (10-20x faster)
                 embeddings = self.embedding_service.embed_batch(texts)
 
-                # Add to ChromaDB
-                self.collection.add(
+                # Upsert to ChromaDB (idempotent - safe for resumption)
+                self.collection.upsert(
                     ids=[chunk.id for chunk in batch],
                     embeddings=embeddings,
                     documents=[chunk.text for chunk in batch],
-                    metadatas=[chunk.metadata for chunk in batch]
+                    metadatas=[chunk.metadata for chunk in batch],
                 )
             except Exception as e:
                 if not progress_callback:
                     self._log(f"\n  Error embedding batch {i}: {e}")
                 continue
 
-        self.stats['indexing_time'] = time.time() - start_time
+            # Save checkpoint every 1000 chunks
+            if checkpoint_path and current_count % 1000 == 0:
+                _save_checkpoint(checkpoint_path, current_count, len(chunks))
+
+        # Remove stale chunks (from deleted/renamed files)
+        try:
+            existing = self.collection.get(include=[])
+            existing_ids = set(existing["ids"])
+            stale_ids = list(existing_ids - new_chunk_ids)
+            if stale_ids:
+                # ChromaDB delete supports batching
+                for j in range(0, len(stale_ids), 500):
+                    self.collection.delete(ids=stale_ids[j : j + 500])
+                if not progress_callback:
+                    self._log(f"  Removed {len(stale_ids)} stale chunks")
+        except Exception as e:
+            if not progress_callback:
+                self._log(f"  Warning: Could not clean stale chunks: {e}")
+
+        # Clean up checkpoint on success
+        if checkpoint_path and checkpoint_path.exists():
+            checkpoint_path.unlink()
+
+        self.stats["indexing_time"] = time.time() - start_time
 
         emit("complete", {"stats": self.stats.copy()})
 
         # Legacy summary for no callback
         if not progress_callback:
-            self._log("\n\n" + "="*80)
+            self._log("\n\n" + "=" * 80)
             self._log("INDEXING COMPLETE")
-            self._log("="*80)
+            self._log("=" * 80)
             self._log(f"Files parsed:      {self.stats['total_files']:,}")
             self._log(f"Symbols extracted: {self.stats['total_symbols']:,}")
             self._log(f"Chunks created:    {self.stats['total_chunks']:,}")
             self._log(f"Parse errors:      {self.stats['parse_errors']}")
-            self._log(f"Time elapsed:      {self.stats['indexing_time']:.1f}s ({self.stats['indexing_time']/60:.1f}m)")
+            self._log(
+                f"Time elapsed:      {self.stats['indexing_time']:.1f}s ({self.stats['indexing_time']/60:.1f}m)"
+            )
             self._log(f"Ready for retrieval!")
-            self._log("="*80 + "\n")
+            self._log("=" * 80 + "\n")
 
     def get_sources_for_question(
         self,
         question: str,
         n_results: int = 10,
         language: Optional[str] = None,
-        symbol_type: Optional[str] = None
+        symbol_type: Optional[str] = None,
     ) -> tuple[List[Dict[str, Any]], List[str]]:
         """
         Get source references and documents for a question.
@@ -556,18 +713,16 @@ def get_sources_for_question(
 
         # Search ChromaDB
         results = self.collection.query(
-            query_embeddings=[query_embedding],
-            n_results=n_results,
-            where=where_filter
+            query_embeddings=[query_embedding], n_results=n_results, where=where_filter
         )
 
-        documents = results['documents'][0]
-        metadatas = results['metadatas'][0]
+        documents = results["documents"][0]
+        metadatas = results["metadatas"][0]
 
         # Format sources for display
         sources = []
         for metadata in metadatas:
-            filepath = metadata['filepath']
+            filepath = metadata["filepath"]
             try:
                 filepath = str(Path(filepath).relative_to(self.codebase_path))
             except ValueError:
@@ -577,10 +732,7 @@ def get_sources_for_question(
         # Build document list with metadata
         doc_results = []
         for doc, metadata in zip(documents, metadatas):
-            doc_results.append({
-                'text': doc,
-                'metadata': metadata
-            })
+            doc_results.append({"text": doc, "metadata": metadata})
 
         return doc_results, sources
 
@@ -600,9 +752,7 @@ def load_existing_index(self) -> bool:
             return False
 
         try:
-            self.collection = self.chroma_client.get_collection(
-                name=self.collection_name
-            )
+            self.collection = self.chroma_client.get_collection(name=self.collection_name)
             count = self.collection.count()
             self._log(f"Loaded existing index with {count:,} chunks")
             return True
@@ -618,14 +768,14 @@ def save_metadata(self, metadata_path: str) -> None:
             metadata_path: Path to save metadata JSON
         """
         metadata = {
-            'codebase_path': str(self.codebase_path),
-            'embedding_model': self.embedding_model,
-            'collection_name': self.collection_name,
-            'indexed_at': datetime.now().isoformat(),
-            'stats': self.stats,
-            'file_mtimes': self._metadata.get('file_mtimes', {})
+            "codebase_path": str(self.codebase_path),
+            "embedding_model": self.embedding_model,
+            "collection_name": self.collection_name,
+            "indexed_at": datetime.now().isoformat(),
+            "stats": self.stats,
+            "file_mtimes": self._metadata.get("file_mtimes", {}),
         }
-        with open(metadata_path, 'w') as f:
+        with open(metadata_path, "w") as f:
             json.dump(metadata, f, indent=2)
         self._log(f"Metadata saved to {metadata_path}")
 
@@ -640,14 +790,14 @@ def load_metadata(self, metadata_path: str) -> Optional[Dict[str, Any]]:
             Metadata dictionary or None if not found
         """
         try:
-            with open(metadata_path, 'r') as f:
+            with open(metadata_path, "r") as f:
                 metadata = json.load(f)
             # Restore stats
-            if 'stats' in metadata:
-                self.stats = metadata['stats']
+            if "stats" in metadata:
+                self.stats = metadata["stats"]
             # Restore file modification times for incremental reindexing
-            if 'file_mtimes' in metadata:
-                self._metadata['file_mtimes'] = metadata['file_mtimes']
+            if "file_mtimes" in metadata:
+                self._metadata["file_mtimes"] = metadata["file_mtimes"]
             return metadata
         except FileNotFoundError:
             return None
@@ -655,7 +805,7 @@ def load_metadata(self, metadata_path: str) -> Optional[Dict[str, Any]]:
     def incremental_reindex(
         self,
         file_extensions: Optional[List[str]] = None,
-        progress_callback: Optional[callable] = None
+        progress_callback: Optional[callable] = None,
     ) -> Dict[str, Any]:
         """
         Re-index only files that changed since last indexing.
@@ -695,7 +845,7 @@ def emit(event_type: str, data: dict):
         extensions = file_extensions or SUPPORTED_EXTENSIONS
 
         # 1. Get stored file_mtimes from metadata
-        stored_mtimes = self._metadata.get('file_mtimes', {})
+        stored_mtimes = self._metadata.get("file_mtimes", {})
 
         # 2. Scan current files and collect modification times (single-pass - 30x faster)
         extensions_set = set(extensions)
@@ -715,21 +865,21 @@ def emit(event_type: str, data: dict):
         new_files = current_paths - stored_paths
         deleted_files = stored_paths - current_paths
         modified_files = {
-            p for p in (current_paths & stored_paths)
-            if current_files[p] > stored_mtimes.get(p, 0)
+            p for p in (current_paths & stored_paths) if current_files[p] > stored_mtimes.get(p, 0)
         }
 
         files_to_reindex = new_files | modified_files
         files_to_remove = deleted_files | modified_files
 
         # Emit changes detected event
-        emit("changes_detected", {
-            "new": len(new_files),
-            "modified": len(modified_files),
-            "deleted": len(deleted_files)
-        })
+        emit(
+            "changes_detected",
+            {"new": len(new_files), "modified": len(modified_files), "deleted": len(deleted_files)},
+        )
 
-        self._log(f"Incremental reindex: {len(new_files)} new, {len(modified_files)} modified, {len(deleted_files)} deleted")
+        self._log(
+            f"Incremental reindex: {len(new_files)} new, {len(modified_files)} modified, {len(deleted_files)} deleted"
+        )
 
         chunks_removed = 0
         chunks_added = 0
@@ -776,16 +926,16 @@ def emit(event_type: str, data: dict):
                     ids=[chunk.id for chunk in chunks],
                     embeddings=embeddings,
                     documents=[chunk.text for chunk in chunks],
-                    metadatas=[chunk.metadata for chunk in chunks]
+                    metadatas=[chunk.metadata for chunk in chunks],
                 )
                 chunks_added = len(chunks)
 
         # 6. Update metadata with new file modification times
-        self._metadata['file_mtimes'] = current_files
+        self._metadata["file_mtimes"] = current_files
 
         # Persist metadata if we have a persist path
         if self.persist_path:
-            metadata_path = Path(self.persist_path).parent / 'metadata.json'
+            metadata_path = Path(self.persist_path).parent / "metadata.json"
             self.save_metadata(str(metadata_path))
 
         elapsed_time = round(time.time() - start_time, 2)
@@ -796,13 +946,15 @@ def emit(event_type: str, data: dict):
             "files_deleted": len(deleted_files),
             "chunks_added": chunks_added,
             "chunks_removed": chunks_removed,
-            "time_seconds": elapsed_time
+            "time_seconds": elapsed_time,
         }
 
         # Emit complete event
         emit("complete", {"chunks_added": chunks_added})
 
-        self._log(f"Incremental reindex complete in {elapsed_time}s: "
-                  f"+{chunks_added} chunks, -{chunks_removed} files processed")
+        self._log(
+            f"Incremental reindex complete in {elapsed_time}s: "
+            f"+{chunks_added} chunks, -{chunks_removed} files processed"
+        )
 
         return result
diff --git a/src/codegrok_mcp/mcp/server.py b/src/codegrok_mcp/mcp/server.py
index 6fe1bd5..7af421b 100644
--- a/src/codegrok_mcp/mcp/server.py
+++ b/src/codegrok_mcp/mcp/server.py
@@ -27,6 +27,7 @@
 from pydantic import Field
 
 from codegrok_mcp.mcp.state import get_state
+
 # Lazy import SourceRetriever to avoid heavy startup cost
 # from codegrok_mcp.indexing.source_retriever import SourceRetriever, SUPPORTED_EXTENSIONS
 from codegrok_mcp.parsers.language_configs import EXTENSION_MAP
@@ -36,9 +37,9 @@
 
 
 # Storage constants (same as CLI)
-CODEGROK_DIR = '.codegrok'
-CHROMA_DIR = 'chroma'
-METADATA_FILE = 'metadata.json'
+CODEGROK_DIR = ".codegrok"
+CHROMA_DIR = "chroma"
+METADATA_FILE = "metadata.json"
 
 # Initialize FastMCP server
 mcp = FastMCP(
@@ -69,7 +70,7 @@
 1. learn(path="/project") - Index codebase (required first step)
 2. recall("user preferences") - Check existing context
 3. remember("Decision: Using Redis for caching", memory_type="decision")
-4. get_sources("authentication flow") - Find relevant code"""
+4. get_sources("authentication flow") - Find relevant code""",
 )
 
 
@@ -77,28 +78,32 @@ def _get_codegrok_paths(codebase_path: Path) -> Dict[str, Path]:
     """Get paths to .codegrok storage locations."""
     codegrok_dir = codebase_path / CODEGROK_DIR
     return {
-        'codegrok_dir': codegrok_dir,
-        'chroma_path': codegrok_dir / CHROMA_DIR,
-        'metadata_path': codegrok_dir / METADATA_FILE,
+        "codegrok_dir": codegrok_dir,
+        "chroma_path": codegrok_dir / CHROMA_DIR,
+        "metadata_path": codegrok_dir / METADATA_FILE,
     }
 
 
 def _has_valid_index(paths: Dict[str, Path]) -> bool:
     """Check if a valid CodeGrok index exists at the given paths."""
     return (
-        paths['codegrok_dir'].exists() and
-        paths['metadata_path'].exists() and
-        paths['chroma_path'].exists()
+        paths["codegrok_dir"].exists()
+        and paths["metadata_path"].exists()
+        and paths["chroma_path"].exists()
     )
 
 
 def _create_learn_progress_callback(ctx: Context, loop) -> Callable:
     """Create a progress callback that reports indexing progress to MCP client."""
+
     def callback(event_type: str, data: dict):
         progress = 0
         message = ""
 
-        if event_type == "files_found":
+        if event_type == "discovery_progress":
+            progress = min(4, int(data.get("files_found", 0) / 1000))
+            message = f"Discovering files... ({data.get('files_found', 0)} found)"
+        elif event_type == "files_found":
             progress = 5
             message = f"Found {len(data['files'])} files..."
         elif event_type == "parsing_start":
@@ -112,23 +117,31 @@ def callback(event_type: str, data: dict):
             message = f"Generating embeddings for {data['total']} chunks..."
         elif event_type == "embedding_progress":
             # Scale embedding progress (35-95%)
-            pct = data['current'] / data['total'] if data['total'] > 0 else 1
+            pct = data["current"] / data["total"] if data["total"] > 0 else 1
             progress = 35 + int(pct * 60)
-            message = f"Embedding... ({data['current']}/{data['total']} chunks)"
+            remaining = data.get("remaining_seconds")
+            if remaining and remaining > 0:
+                eta_str = (
+                    f", ~{remaining / 60:.1f}m remaining"
+                    if remaining >= 60
+                    else f", ~{remaining:.0f}s remaining"
+                )
+            else:
+                eta_str = ""
+            message = f"Embedding... ({data['current']}/{data['total']} chunks{eta_str})"
         elif event_type == "complete":
             progress = 100
             message = "Indexing complete!"
 
         if progress > 0:
-            asyncio.run_coroutine_threadsafe(
-                ctx.report_progress(progress, 100, message),
-                loop
-            )
+            asyncio.run_coroutine_threadsafe(ctx.report_progress(progress, 100, message), loop)
+
     return callback
 
 
 def _create_relearn_progress_callback(ctx: Context, loop) -> Callable:
     """Create a progress callback that reports reindexing progress to MCP client."""
+
     def callback(event_type: str, data: dict):
         progress = 0
         message = ""
@@ -147,10 +160,8 @@ def callback(event_type: str, data: dict):
             message = "Re-indexing complete!"
 
         if progress > 0:
-            asyncio.run_coroutine_threadsafe(
-                ctx.report_progress(progress, 100, message),
-                loop
-            )
+            asyncio.run_coroutine_threadsafe(ctx.report_progress(progress, 100, message), loop)
+
     return callback
 
 
@@ -165,27 +176,30 @@ def callback(event_type: str, data: dict):
 
 Creates a .codegrok/ folder in the codebase directory.""",
     annotations=ToolAnnotations(
-        readOnlyHint=False,     # Creates/modifies .codegrok/ directory
+        readOnlyHint=False,  # Creates/modifies .codegrok/ directory
         destructiveHint=False,  # Doesn't destroy user data (only index data)
-        idempotentHint=True,    # Safe to re-run on same path
-        openWorldHint=False     # Only accesses local filesystem
-    )
+        idempotentHint=True,  # Safe to re-run on same path
+        openWorldHint=False,  # Only accesses local filesystem
+    ),
 )
 async def learn(
     path: Annotated[str, Field(description="Absolute path to the codebase directory to index")],
     mode: Annotated[
         str,
-        Field(description="Indexing mode: 'auto' (smart detection), 'full' (force re-index), 'load_only' (just load)")
+        Field(
+            description="Indexing mode: 'auto' (smart detection), 'full' (force re-index), 'load_only' (just load)"
+        ),
     ] = "auto",
     file_extensions: Annotated[
         Optional[List[str]],
-        Field(description="Optional list of file extensions to include (e.g., ['.py', '.js']). Defaults to all supported extensions.")
+        Field(
+            description="Optional list of file extensions to include (e.g., ['.py', '.js']). Defaults to all supported extensions."
+        ),
     ] = None,
     embedding_model: Annotated[
-        str,
-        Field(description="Embedding model to use (default: coderankembed)")
+        str, Field(description="Embedding model to use (default: coderankembed)")
     ] = "coderankembed",
-    ctx: Context = None
+    ctx: Context = None,
 ) -> Dict[str, Any]:
     """Index a codebase with smart mode detection."""
     state = get_state()
@@ -224,10 +238,7 @@ async def learn(
 
 
 async def _load_existing_index(
-    codebase_path: Path,
-    paths: Dict[str, Path],
-    state,
-    embedding_model: str
+    codebase_path: Path, paths: Dict[str, Path], state, embedding_model: str
 ) -> Dict[str, Any]:
     """Load an existing index without any reindexing."""
     from codegrok_mcp.indexing.source_retriever import SourceRetriever
@@ -236,15 +247,15 @@ async def _load_existing_index(
         codebase_path=str(codebase_path),
         embedding_model=embedding_model,
         verbose=False,
-        persist_path=str(paths['chroma_path'])
+        persist_path=str(paths["chroma_path"]),
     )
 
     if not retriever.load_existing_index():
         raise ToolError(f"Failed to load index from {paths['chroma_path']}")
 
-    metadata = retriever.load_metadata(str(paths['metadata_path']))
-    stats = metadata.get('stats', {}) if metadata else {}
-    indexed_at = metadata.get('indexed_at') if metadata else None
+    metadata = retriever.load_metadata(str(paths["metadata_path"]))
+    stats = metadata.get("stats", {}) if metadata else {}
+    indexed_at = metadata.get("indexed_at") if metadata else None
 
     state.retriever = retriever
     state.codebase_path = codebase_path
@@ -254,16 +265,12 @@ async def _load_existing_index(
         "mode_used": "load_only",
         "message": f"Loaded existing index for {codebase_path.name}",
         "stats": stats,
-        "indexed_at": indexed_at
+        "indexed_at": indexed_at,
     }
 
 
 async def _incremental_reindex(
-    codebase_path: Path,
-    paths: Dict[str, Path],
-    state,
-    embedding_model: str,
-    ctx: Context = None
+    codebase_path: Path, paths: Dict[str, Path], state, embedding_model: str, ctx: Context = None
 ) -> Dict[str, Any]:
     """Load existing index and perform incremental reindex."""
     from codegrok_mcp.indexing.source_retriever import SourceRetriever
@@ -272,14 +279,14 @@ async def _incremental_reindex(
         codebase_path=str(codebase_path),
         embedding_model=embedding_model,
         verbose=False,
-        persist_path=str(paths['chroma_path'])
+        persist_path=str(paths["chroma_path"]),
     )
 
     if not retriever.load_existing_index():
         raise ToolError(f"Failed to load existing index from {paths['chroma_path']}")
 
     # Load metadata to get file mtimes for incremental detection
-    retriever.load_metadata(str(paths['metadata_path']))
+    retriever.load_metadata(str(paths["metadata_path"]))
 
     # Create progress callback if context available
     progress_callback = None
@@ -291,7 +298,7 @@ async def _incremental_reindex(
     result = retriever.incremental_reindex(progress_callback=progress_callback)
 
     # Save updated metadata
-    retriever.save_metadata(str(paths['metadata_path']))
+    retriever.save_metadata(str(paths["metadata_path"]))
 
     state.retriever = retriever
     state.codebase_path = codebase_path
@@ -300,7 +307,7 @@ async def _incremental_reindex(
         "success": True,
         "mode_used": "incremental",
         "message": f"Incremental reindex complete for {codebase_path.name}",
-        **result
+        **result,
     }
 
 
@@ -310,13 +317,13 @@ async def _full_index(
     state,
     file_extensions: Optional[List[str]],
     embedding_model: str,
-    ctx: Context = None
+    ctx: Context = None,
 ) -> Dict[str, Any]:
     """Perform full index (creates or replaces existing index)."""
     from codegrok_mcp.indexing.source_retriever import SourceRetriever
 
     # Create .codegrok directory
-    paths['codegrok_dir'].mkdir(parents=True, exist_ok=True)
+    paths["codegrok_dir"].mkdir(parents=True, exist_ok=True)
 
     # Create progress callback if context available
     progress_callback = None
@@ -328,7 +335,7 @@ async def _full_index(
         codebase_path=str(codebase_path),
         embedding_model=embedding_model,
         verbose=False,
-        persist_path=str(paths['chroma_path'])
+        persist_path=str(paths["chroma_path"]),
     )
 
     # Index the codebase with progress reporting
@@ -340,7 +347,7 @@ async def _full_index(
         await ctx.report_progress(95, 100, "Saving index...")
 
     # Save metadata
-    retriever.save_metadata(str(paths['metadata_path']))
+    retriever.save_metadata(str(paths["metadata_path"]))
 
     # Update state
     state.retriever = retriever
@@ -350,7 +357,7 @@ async def _full_index(
         "success": True,
         "mode_used": "full",
         "message": f"Successfully indexed {codebase_path.name}",
-        "stats": retriever.get_stats()
+        "stats": retriever.get_stats(),
     }
 
 
@@ -364,25 +371,24 @@ async def _full_index(
 - Find auth code: get_sources(question="authentication login flow")
 - Find Python classes: get_sources(question="user model", language="python", symbol_type="class")""",
     annotations=ToolAnnotations(
-        readOnlyHint=True,      # Only reads from index
-        idempotentHint=True,    # Same query returns same results
-        openWorldHint=False     # Only accesses local ChromaDB
-    )
+        readOnlyHint=True,  # Only reads from index
+        idempotentHint=True,  # Same query returns same results
+        openWorldHint=False,  # Only accesses local ChromaDB
+    ),
 )
 def get_sources(
     question: Annotated[str, Field(description="Natural language question or search query")],
     n_results: Annotated[
-        int,
-        Field(description="Number of source references to return (default: 10)", ge=1, le=50)
+        int, Field(description="Number of source references to return (default: 10)", ge=1, le=50)
     ] = 10,
     language: Annotated[
         Optional[str],
-        Field(description="Filter by language (e.g., 'python', 'javascript', 'typescript')")
+        Field(description="Filter by language (e.g., 'python', 'javascript', 'typescript')"),
     ] = None,
     symbol_type: Annotated[
         Optional[str],
-        Field(description="Filter by symbol type (e.g., 'function', 'class', 'method')")
-    ] = None
+        Field(description="Filter by symbol type (e.g., 'function', 'class', 'method')"),
+    ] = None,
 ) -> Dict[str, Any]:
     """Get source references for a question with optional filters."""
     state = get_state()
@@ -393,15 +399,12 @@ def get_sources(
     try:
         # get_sources_for_question returns tuple: (doc_results, formatted_sources)
         doc_results, formatted_sources = state.retriever.get_sources_for_question(
-            question,
-            n_results=n_results,
-            language=language,
-            symbol_type=symbol_type
+            question, n_results=n_results, language=language, symbol_type=symbol_type
         )
 
         return {
-            "documents": doc_results,      # Full document data with metadata
-            "sources": formatted_sources   # Formatted source references for display
+            "documents": doc_results,  # Full document data with metadata
+            "sources": formatted_sources,  # Formatted source references for display
         }
 
     except Exception as e:
@@ -414,26 +417,22 @@ def get_sources(
 
 Returns: files indexed, total chunks, symbols by type, languages detected, index creation time.""",
     annotations=ToolAnnotations(
-        readOnlyHint=True,      # Only reads metadata
-        idempotentHint=True,    # Same state = same results
-        openWorldHint=False     # Only accesses local state
-    )
+        readOnlyHint=True,  # Only reads metadata
+        idempotentHint=True,  # Same state = same results
+        openWorldHint=False,  # Only accesses local state
+    ),
 )
 def get_stats() -> Dict[str, Any]:
     """Get indexing statistics."""
     state = get_state()
 
     if not state.is_loaded:
-        return {
-            "loaded": False,
-            "codebase_path": None,
-            "stats": None
-        }
+        return {"loaded": False, "codebase_path": None, "stats": None}
 
     return {
         "loaded": True,
         "codebase_path": str(state.codebase_path),
-        "stats": state.retriever.get_stats()
+        "stats": state.retriever.get_stats(),
     }
 
 
@@ -443,10 +442,10 @@ def get_stats() -> Dict[str, Any]:
 
 Returns extensions grouped by language. Currently supports: Python, JavaScript, TypeScript, Go, Rust, Java, C, C++, Ruby, and more.""",
     annotations=ToolAnnotations(
-        readOnlyHint=True,      # Returns static data
-        idempotentHint=True,    # Always same result
-        openWorldHint=False     # No external access
-    )
+        readOnlyHint=True,  # Returns static data
+        idempotentHint=True,  # Always same result
+        openWorldHint=False,  # No external access
+    ),
 )
 def list_supported_languages() -> Dict[str, Any]:
     """List supported file extensions and languages."""
@@ -461,10 +460,7 @@ def list_supported_languages() -> Dict[str, Any]:
     for lang in languages:
         languages[lang] = sorted(languages[lang])
 
-    return {
-        "extensions": sorted(EXTENSION_MAP.keys()),
-        "languages": languages
-    }
+    return {"extensions": sorted(EXTENSION_MAP.keys()), "languages": languages}
 
 
 # =============================================================================
@@ -497,26 +493,24 @@ def list_supported_languages() -> Dict[str, Any]:
 - Remember a decision: remember(content="Using JWT with refresh tokens for auth", memory_type="decision", tags=["auth"])
 """,
     annotations=ToolAnnotations(
-        readOnlyHint=False,     # Writes to ChromaDB
+        readOnlyHint=False,  # Writes to ChromaDB
         destructiveHint=False,  # Adds data, doesn't delete
-        idempotentHint=False,   # Creates new memory each call
-        openWorldHint=False     # Only local storage
-    )
+        idempotentHint=False,  # Creates new memory each call
+        openWorldHint=False,  # Only local storage
+    ),
 )
 def remember(
     content: Annotated[str, Field(description="The memory content to store")],
     memory_type: Annotated[
-        str,
-        Field(description="Type: conversation, status, decision, preference, doc, note")
+        str, Field(description="Type: conversation, status, decision, preference, doc, note")
     ],
     tags: Annotated[
         Optional[List[str]],
-        Field(description="Optional tags for filtering (e.g., ['auth', 'backend'])")
+        Field(description="Optional tags for filtering (e.g., ['auth', 'backend'])"),
     ] = None,
     ttl: Annotated[
-        str,
-        Field(description="Time-to-live: session, day, week, month, permanent (default)")
-    ] = "permanent"
+        str, Field(description="Time-to-live: session, day, week, month, permanent (default)")
+    ] = "permanent",
 ) -> Dict[str, Any]:
     """Store a new memory with automatic embedding."""
     from codegrok_mcp.indexing.memory_retriever import MemoryRetriever
@@ -532,8 +526,8 @@ def remember(
         paths = _get_codegrok_paths(state.codebase_path)
         state.memory_retriever = MemoryRetriever(
             project_path=str(state.codebase_path),
-            persist_path=str(paths['chroma_path']),
-            verbose=False
+            persist_path=str(paths["chroma_path"]),
+            verbose=False,
         )
 
     # Validate memory_type
@@ -543,17 +537,14 @@ def remember(
 
     # Store memory
     memory = state.memory_retriever.remember(
-        content=content,
-        memory_type=memory_type,
-        tags=tags or [],
-        ttl=ttl
+        content=content, memory_type=memory_type, tags=tags or [], ttl=ttl
     )
 
     return {
         "success": True,
         "memory_id": memory.id,
         "message": f"Stored {memory_type} memory",
-        "tags": memory.tags
+        "tags": memory.tags,
     }
 
 
@@ -577,29 +568,22 @@ def remember(
 - Recent decisions: recall(query="architecture decisions", time_range="week")
 """,
     annotations=ToolAnnotations(
-        readOnlyHint=True,      # Only reads from memory store
-        idempotentHint=True,    # Same query = same results
-        openWorldHint=False     # Only accesses local ChromaDB
-    )
+        readOnlyHint=True,  # Only reads from memory store
+        idempotentHint=True,  # Same query = same results
+        openWorldHint=False,  # Only accesses local ChromaDB
+    ),
 )
 def recall(
     query: Annotated[str, Field(description="Natural language search query")],
     memory_type: Annotated[
         Optional[str],
-        Field(description="Filter by type: conversation, status, decision, preference, doc, note")
+        Field(description="Filter by type: conversation, status, decision, preference, doc, note"),
     ] = None,
-    tags: Annotated[
-        Optional[List[str]],
-        Field(description="Filter by tags (matches any)")
-    ] = None,
-    n_results: Annotated[
-        int,
-        Field(description="Number of results (default: 5)", ge=1, le=20)
-    ] = 5,
+    tags: Annotated[Optional[List[str]], Field(description="Filter by tags (matches any)")] = None,
+    n_results: Annotated[int, Field(description="Number of results (default: 5)", ge=1, le=20)] = 5,
     time_range: Annotated[
-        Optional[str],
-        Field(description="Time filter: today, week, month, all")
-    ] = None
+        Optional[str], Field(description="Time filter: today, week, month, all")
+    ] = None,
 ) -> Dict[str, Any]:
     """Retrieve memories using semantic search."""
     from codegrok_mcp.indexing.memory_retriever import MemoryRetriever
@@ -615,8 +599,8 @@ def recall(
         paths = _get_codegrok_paths(state.codebase_path)
         state.memory_retriever = MemoryRetriever(
             project_path=str(state.codebase_path),
-            persist_path=str(paths['chroma_path']),
-            verbose=False
+            persist_path=str(paths["chroma_path"]),
+            verbose=False,
         )
 
     # Validate memory_type if provided
@@ -627,18 +611,10 @@ def recall(
 
     # Search memories
     memories = state.memory_retriever.recall(
-        query=query,
-        memory_type=memory_type,
-        tags=tags,
-        n_results=n_results,
-        time_range=time_range
+        query=query, memory_type=memory_type, tags=tags, n_results=n_results, time_range=time_range
     )
 
-    return {
-        "success": True,
-        "count": len(memories),
-        "memories": memories
-    }
+    return {"success": True, "count": len(memories), "memories": memories}
 
 
 @mcp.tool(
@@ -657,29 +633,23 @@ def recall(
 - Remove by tag: forget(tags=["deprecated", "outdated"])
 """,
     annotations=ToolAnnotations(
-        readOnlyHint=False,     # Deletes from ChromaDB
-        destructiveHint=True,   # ⚠️ PERMANENTLY DELETES data
-        idempotentHint=True,    # Re-calling same filter is safe
-        openWorldHint=False     # Only local storage
-    )
+        readOnlyHint=False,  # Deletes from ChromaDB
+        destructiveHint=True,  # ⚠️ PERMANENTLY DELETES data
+        idempotentHint=True,  # Re-calling same filter is safe
+        openWorldHint=False,  # Only local storage
+    ),
 )
 def forget(
-    memory_id: Annotated[
-        Optional[str],
-        Field(description="Specific memory ID to delete")
-    ] = None,
+    memory_id: Annotated[Optional[str], Field(description="Specific memory ID to delete")] = None,
     memory_type: Annotated[
-        Optional[str],
-        Field(description="Delete all memories of this type")
+        Optional[str], Field(description="Delete all memories of this type")
     ] = None,
     tags: Annotated[
-        Optional[List[str]],
-        Field(description="Delete memories with any of these tags")
+        Optional[List[str]], Field(description="Delete memories with any of these tags")
     ] = None,
     older_than: Annotated[
-        Optional[str],
-        Field(description="Delete memories older than: 1d, 7d, 30d, 1y")
-    ] = None
+        Optional[str], Field(description="Delete memories older than: 1d, 7d, 30d, 1y")
+    ] = None,
 ) -> Dict[str, Any]:
     """Remove memories matching criteria."""
     from codegrok_mcp.indexing.memory_retriever import MemoryRetriever
@@ -693,24 +663,23 @@ def forget(
         paths = _get_codegrok_paths(state.codebase_path)
         state.memory_retriever = MemoryRetriever(
             project_path=str(state.codebase_path),
-            persist_path=str(paths['chroma_path']),
-            verbose=False
+            persist_path=str(paths["chroma_path"]),
+            verbose=False,
         )
 
     if not any([memory_id, memory_type, tags, older_than]):
-        raise ToolError("Must specify at least one filter: memory_id, memory_type, tags, or older_than")
+        raise ToolError(
+            "Must specify at least one filter: memory_id, memory_type, tags, or older_than"
+        )
 
     result = state.memory_retriever.forget(
-        memory_id=memory_id,
-        memory_type=memory_type,
-        tags=tags,
-        older_than=older_than
+        memory_id=memory_id, memory_type=memory_type, tags=tags, older_than=older_than
     )
 
     return {
         "success": True,
         "deleted": result["deleted"],
-        "message": f"Deleted {result['deleted']} memories"
+        "message": f"Deleted {result['deleted']} memories",
     }
 
 
@@ -720,10 +689,10 @@ def forget(
 
 Returns: total memories, count by type, count by TTL, oldest/newest memory dates.""",
     annotations=ToolAnnotations(
-        readOnlyHint=True,      # Only reads metadata
-        idempotentHint=True,    # Same state = same results
-        openWorldHint=False     # Only accesses local state
-    )
+        readOnlyHint=True,  # Only reads metadata
+        idempotentHint=True,  # Same state = same results
+        openWorldHint=False,  # Only accesses local state
+    ),
 )
 def memory_stats() -> Dict[str, Any]:
     """Get memory statistics."""
@@ -732,26 +701,19 @@ def memory_stats() -> Dict[str, Any]:
     state = get_state()
 
     if not state.codebase_path:
-        return {
-            "loaded": False,
-            "message": "No codebase loaded. Use 'learn' first."
-        }
+        return {"loaded": False, "message": "No codebase loaded. Use 'learn' first."}
 
     if state.memory_retriever is None:
         paths = _get_codegrok_paths(state.codebase_path)
         state.memory_retriever = MemoryRetriever(
             project_path=str(state.codebase_path),
-            persist_path=str(paths['chroma_path']),
-            verbose=False
+            persist_path=str(paths["chroma_path"]),
+            verbose=False,
         )
 
     stats = state.memory_retriever.get_stats()
 
-    return {
-        "loaded": True,
-        "project": str(state.codebase_path),
-        **stats
-    }
+    return {"loaded": True, "project": str(state.codebase_path), **stats}
 
 
 def main():  # pragma: no cover
diff --git a/src/codegrok_mcp/mcp/state.py b/src/codegrok_mcp/mcp/state.py
index 5032279..74fb160 100644
--- a/src/codegrok_mcp/mcp/state.py
+++ b/src/codegrok_mcp/mcp/state.py
@@ -1,4 +1,5 @@
 """Session state management for MCP server."""
+
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Optional, TYPE_CHECKING
@@ -11,6 +12,7 @@
 @dataclass
 class MCPSessionState:
     """Singleton state for MCP server session."""
+
     retriever: Optional["SourceRetriever"] = None
     memory_retriever: Optional["MemoryRetriever"] = None
     codebase_path: Optional[Path] = None
diff --git a/src/codegrok_mcp/parsers/language_configs.py b/src/codegrok_mcp/parsers/language_configs.py
index 28a6b99..1d307c3 100644
--- a/src/codegrok_mcp/parsers/language_configs.py
+++ b/src/codegrok_mcp/parsers/language_configs.py
@@ -28,53 +28,46 @@
 from typing import Optional, Dict, List, Set
 from pathlib import Path
 
-
 # ==============================================================================
 # Extension to Language Mapping
 # ==============================================================================
 
 EXTENSION_MAP: Dict[str, str] = {
     # Python
-    '.py': 'python',
-    '.pyi': 'python',  # Type stub files
-    '.pyw': 'python',  # Windows Python GUI scripts
-
+    ".py": "python",
+    ".pyi": "python",  # Type stub files
+    ".pyw": "python",  # Windows Python GUI scripts
     # JavaScript/TypeScript
-    '.js': 'javascript',
-    '.jsx': 'javascript',
-    '.mjs': 'javascript',  # ES6 modules
-    '.cjs': 'javascript',  # CommonJS modules
-    '.ts': 'typescript',
-    '.tsx': 'typescript',
-    '.mts': 'typescript',  # TypeScript ES6 modules
-    '.cts': 'typescript',  # TypeScript CommonJS modules
-
+    ".js": "javascript",
+    ".jsx": "javascript",
+    ".mjs": "javascript",  # ES6 modules
+    ".cjs": "javascript",  # CommonJS modules
+    ".ts": "typescript",
+    ".tsx": "typescript",
+    ".mts": "typescript",  # TypeScript ES6 modules
+    ".cts": "typescript",  # TypeScript CommonJS modules
     # C/C++
-    '.c': 'c',
-    '.h': 'c',  # Header files (may be C or C++)
-    '.cpp': 'cpp',
-    '.cc': 'cpp',
-    '.cxx': 'cpp',
-    '.c++': 'cpp',
-    '.hpp': 'cpp',
-    '.hh': 'cpp',
-    '.hxx': 'cpp',
-    '.h++': 'cpp',
-
+    ".c": "c",
+    ".h": "c",  # Header files (may be C or C++)
+    ".cpp": "cpp",
+    ".cc": "cpp",
+    ".cxx": "cpp",
+    ".c++": "cpp",
+    ".hpp": "cpp",
+    ".hh": "cpp",
+    ".hxx": "cpp",
+    ".h++": "cpp",
     # Bash
-    '.sh': 'bash',
-    '.bash': 'bash',
-    '.zsh': 'bash',  # Zsh is similar enough to bash
-
+    ".sh": "bash",
+    ".bash": "bash",
+    ".zsh": "bash",  # Zsh is similar enough to bash
     # Go
-    '.go': 'go',
-
+    ".go": "go",
     # Java
-    '.java': 'java',
-
+    ".java": "java",
     # Kotlin
-    '.kt': 'kotlin',
-    '.kts': 'kotlin',  # Kotlin script files
+    ".kt": "kotlin",
+    ".kts": "kotlin",  # Kotlin script files
 }
 
 
@@ -83,45 +76,42 @@
 # ==============================================================================
 
 LANGUAGE_CONFIGS: Dict[str, Dict[str, any]] = {
-
     # --------------------------------------------------------------------------
     # Python Configuration
     # --------------------------------------------------------------------------
-    'python': {
-        'function_types': [
-            'function_definition',  # def foo(): ...
+    "python": {
+        "function_types": [
+            "function_definition",  # def foo(): ...
         ],
-        'class_types': [
-            'class_definition',  # class MyClass: ...
+        "class_types": [
+            "class_definition",  # class MyClass: ...
         ],
-        'method_types': [
-            'function_definition',  # Methods are functions inside class bodies
+        "method_types": [
+            "function_definition",  # Methods are functions inside class bodies
         ],
-        'constant_types': [
-            'expression_statement',  # MODULE_CONST = value (at module level)
+        "constant_types": [
+            "expression_statement",  # MODULE_CONST = value (at module level)
         ],
-        'import_types': [
-            'import_statement',       # import os
-            'import_from_statement',  # from os import path
+        "import_types": [
+            "import_statement",  # import os
+            "import_from_statement",  # from os import path
         ],
-        'call_types': [
-            'call',  # function_name(args)
+        "call_types": [
+            "call",  # function_name(args)
         ],
-        'decorator_types': [
-            'decorator',  # @decorator
+        "decorator_types": [
+            "decorator",  # @decorator
         ],
-        'docstring_field': 'string',  # First string literal in function/class body
-        'identifier_field': 'name',   # Field containing the symbol name
-        'body_field': 'body',         # Field containing the body block
-
+        "docstring_field": "string",  # First string literal in function/class body
+        "identifier_field": "name",  # Field containing the symbol name
+        "body_field": "body",  # Field containing the body block
         # Additional node types for comprehensive parsing
-        'async_function_types': [
-            'function_definition',  # async def foo(): ... (same node type)
+        "async_function_types": [
+            "function_definition",  # async def foo(): ... (same node type)
         ],
-        'lambda_types': [
-            'lambda',  # lambda x: x + 1
+        "lambda_types": [
+            "lambda",  # lambda x: x + 1
         ],
-
         # Examples of AST nodes:
         # function_definition:
         #   name: identifier
@@ -138,51 +128,48 @@
         #   function: identifier | attribute
         #   arguments: argument_list
     },
-
     # --------------------------------------------------------------------------
     # JavaScript Configuration
     # --------------------------------------------------------------------------
-    'javascript': {
-        'function_types': [
-            'function_declaration',    # function foo() {}
-            'function',                # function foo() {} (alternate name in some versions)
-            'generator_function_declaration',  # function* foo() {}
-        ],
-        'class_types': [
-            'class_declaration',  # class MyClass {}
+    "javascript": {
+        "function_types": [
+            "function_declaration",  # function foo() {}
+            "function",  # function foo() {} (alternate name in some versions)
+            "generator_function_declaration",  # function* foo() {}
+        ],
+        "class_types": [
+            "class_declaration",  # class MyClass {}
             # Note: 'class' is just the keyword, not the full class node
         ],
-        'method_types': [
-            'method_definition',  # Methods inside class bodies
-            'function_expression',  # foo: function() {}
-            'arrow_function',      # foo: () => {}
+        "method_types": [
+            "method_definition",  # Methods inside class bodies
+            "function_expression",  # foo: function() {}
+            "arrow_function",  # foo: () => {}
         ],
-        'constant_types': [
-            'lexical_declaration',  # const MAX_RETRIES = 3;
+        "constant_types": [
+            "lexical_declaration",  # const MAX_RETRIES = 3;
         ],
-        'import_types': [
-            'import_statement',  # import { x } from 'module'
-            'import_clause',     # Part of import statement
+        "import_types": [
+            "import_statement",  # import { x } from 'module'
+            "import_clause",  # Part of import statement
         ],
-        'call_types': [
-            'call_expression',  # foo()
-            'new_expression',   # new Foo()
+        "call_types": [
+            "call_expression",  # foo()
+            "new_expression",  # new Foo()
         ],
-        'export_types': [
-            'export_statement',  # export { foo }
+        "export_types": [
+            "export_statement",  # export { foo }
         ],
-        'docstring_field': 'comment',  # JSDoc comments (/** ... */)
-        'identifier_field': 'name',
-        'body_field': 'body',
-
+        "docstring_field": "comment",  # JSDoc comments (/** ... */)
+        "identifier_field": "name",
+        "body_field": "body",
         # Additional patterns
-        'arrow_function_types': [
-            'arrow_function',  # const foo = () => {}
+        "arrow_function_types": [
+            "arrow_function",  # const foo = () => {}
         ],
-        'variable_declaration_types': [
-            'variable_declarator',  # const foo = ...
+        "variable_declaration_types": [
+            "variable_declarator",  # const foo = ...
         ],
-
         # Examples of AST nodes:
         # function_declaration:
         #   name: identifier
@@ -202,53 +189,50 @@
         #   function: identifier | member_expression
         #   arguments: arguments
     },
-
     # --------------------------------------------------------------------------
     # TypeScript Configuration
     # --------------------------------------------------------------------------
-    'typescript': {
-        'function_types': [
-            'function_declaration',
-            'function_signature',  # TypeScript type declaration
-            'generator_function_declaration',
-        ],
-        'class_types': [
-            'class_declaration',
-            'interface_declaration',  # TypeScript interfaces
-            'type_alias_declaration',  # type MyType = ...
-        ],
-        'method_types': [
-            'method_definition',
-            'method_signature',  # TypeScript method signatures in interfaces
-            'arrow_function',
-            'function_expression',
-        ],
-        'import_types': [
-            'import_statement',
-            'import_clause',
-        ],
-        'call_types': [
-            'call_expression',
-            'new_expression',
-        ],
-        'export_types': [
-            'export_statement',
-        ],
-        'docstring_field': 'comment',  # TSDoc comments
-        'identifier_field': 'name',
-        'body_field': 'body',
-
+    "typescript": {
+        "function_types": [
+            "function_declaration",
+            "function_signature",  # TypeScript type declaration
+            "generator_function_declaration",
+        ],
+        "class_types": [
+            "class_declaration",
+            "interface_declaration",  # TypeScript interfaces
+            "type_alias_declaration",  # type MyType = ...
+        ],
+        "method_types": [
+            "method_definition",
+            "method_signature",  # TypeScript method signatures in interfaces
+            "arrow_function",
+            "function_expression",
+        ],
+        "import_types": [
+            "import_statement",
+            "import_clause",
+        ],
+        "call_types": [
+            "call_expression",
+            "new_expression",
+        ],
+        "export_types": [
+            "export_statement",
+        ],
+        "docstring_field": "comment",  # TSDoc comments
+        "identifier_field": "name",
+        "body_field": "body",
         # TypeScript-specific
-        'interface_types': [
-            'interface_declaration',  # interface Foo { ... }
+        "interface_types": [
+            "interface_declaration",  # interface Foo { ... }
         ],
-        'type_types': [
-            'type_alias_declaration',  # type Foo = ...
+        "type_types": [
+            "type_alias_declaration",  # type Foo = ...
         ],
-        'enum_types': [
-            'enum_declaration',  # enum Foo { ... }
+        "enum_types": [
+            "enum_declaration",  # enum Foo { ... }
         ],
-
         # Examples of AST nodes:
         # interface_declaration:
         #   name: type_identifier
@@ -258,47 +242,44 @@
         #   name: type_identifier
         #   value: type
     },
-
     # --------------------------------------------------------------------------
     # C Configuration
     # --------------------------------------------------------------------------
-    'c': {
-        'function_types': [
-            'function_definition',  # void foo() { ... }
-            'function_declarator',  # Function declarations
+    "c": {
+        "function_types": [
+            "function_definition",  # void foo() { ... }
+            "function_declarator",  # Function declarations
         ],
-        'class_types': [
-            'struct_specifier',  # struct Foo { ... }
-            'union_specifier',   # union Foo { ... }
+        "class_types": [
+            "struct_specifier",  # struct Foo { ... }
+            "union_specifier",  # union Foo { ... }
         ],
-        'method_types': [
-            'function_definition',  # C doesn't have methods, but struct functions
+        "method_types": [
+            "function_definition",  # C doesn't have methods, but struct functions
         ],
-        'import_types': [
-            'preproc_include',  # #include <stdio.h>
+        "import_types": [
+            "preproc_include",  # #include <stdio.h>
         ],
-        'call_types': [
-            'call_expression',  # foo()
+        "call_types": [
+            "call_expression",  # foo()
         ],
-        'typedef_types': [
-            'type_definition',  # typedef struct { ... } Foo;
+        "typedef_types": [
+            "type_definition",  # typedef struct { ... } Foo;
         ],
-        'docstring_field': 'comment',  # /* ... */ or // ...
-        'identifier_field': 'declarator',
-        'body_field': 'body',
-
+        "docstring_field": "comment",  # /* ... */ or // ...
+        "identifier_field": "declarator",
+        "body_field": "body",
         # Additional C-specific nodes
-        'enum_types': [
-            'enum_specifier',  # enum Foo { ... }
+        "enum_types": [
+            "enum_specifier",  # enum Foo { ... }
         ],
-        'macro_types': [
-            'preproc_def',      # #define FOO 123
-            'preproc_function_def',  # #define FOO(x) (x + 1)
+        "macro_types": [
+            "preproc_def",  # #define FOO 123
+            "preproc_function_def",  # #define FOO(x) (x + 1)
         ],
-
         # Examples of AST nodes:
         # function_definition:
-        #   type: primitive_type | struct_specifier
+        # type: primitive_type | struct_specifier
         #   declarator: function_declarator
         #     declarator: identifier
         #     parameters: parameter_list
@@ -311,52 +292,49 @@
         # preproc_include:
         #   path: string_literal | system_lib_string
     },
-
     # --------------------------------------------------------------------------
     # C++ Configuration
     # --------------------------------------------------------------------------
-    'cpp': {
-        'function_types': [
-            'function_definition',
+    "cpp": {
+        "function_types": [
+            "function_definition",
             # Note: function_declarator is a child of function_definition, not standalone
         ],
-        'class_types': [
-            'class_specifier',    # class Foo { ... }
-            'struct_specifier',   # struct Foo { ... }
-            'union_specifier',    # union Foo { ... }
+        "class_types": [
+            "class_specifier",  # class Foo { ... }
+            "struct_specifier",  # struct Foo { ... }
+            "union_specifier",  # union Foo { ... }
         ],
-        'method_types': [
-            'function_definition',  # Methods inside class bodies
+        "method_types": [
+            "function_definition",  # Methods inside class bodies
             # Note: field_declaration is for member variables, not methods
         ],
-        'import_types': [
-            'preproc_include',  # #include <iostream>
+        "import_types": [
+            "preproc_include",  # #include <iostream>
         ],
-        'call_types': [
-            'call_expression',  # foo()
+        "call_types": [
+            "call_expression",  # foo()
         ],
-        'namespace_types': [
-            'namespace_definition',  # namespace foo { ... }
+        "namespace_types": [
+            "namespace_definition",  # namespace foo { ... }
         ],
-        'template_types': [
-            'template_declaration',  # template<typename T> class Foo { ... }
+        "template_types": [
+            "template_declaration",  # template<typename T> class Foo { ... }
         ],
-        'typedef_types': [
-            'type_definition',  # typedef or using
-            'alias_declaration',  # using Foo = Bar;
+        "typedef_types": [
+            "type_definition",  # typedef or using
+            "alias_declaration",  # using Foo = Bar;
         ],
-        'docstring_field': 'comment',  # Doxygen comments (/** ... */)
-        'identifier_field': 'declarator',
-        'body_field': 'body',
-
+        "docstring_field": "comment",  # Doxygen comments (/** ... */)
+        "identifier_field": "declarator",
+        "body_field": "body",
         # C++-specific features
-        'enum_types': [
-            'enum_specifier',  # enum class Foo { ... }
+        "enum_types": [
+            "enum_specifier",  # enum class Foo { ... }
         ],
-        'lambda_types': [
-            'lambda_expression',  # [](int x) { return x + 1; }
+        "lambda_types": [
+            "lambda_expression",  # [](int x) { return x + 1; }
         ],
-
         # Examples of AST nodes:
         # class_specifier:
         #   name: type_identifier
@@ -371,44 +349,41 @@
         #   parameters: template_parameter_list
         #   declaration: class_specifier | function_definition
     },
-
     # --------------------------------------------------------------------------
     # Bash Configuration
     # --------------------------------------------------------------------------
-    'bash': {
-        'function_types': [
-            'function_definition',  # foo() { ... } or function foo { ... }
+    "bash": {
+        "function_types": [
+            "function_definition",  # foo() { ... } or function foo { ... }
         ],
-        'class_types': [
+        "class_types": [
             # Bash doesn't have classes
         ],
-        'method_types': [
+        "method_types": [
             # Bash doesn't have methods
         ],
-        'constant_types': [
-            'declaration_command',  # readonly VAR=value (true constants)
+        "constant_types": [
+            "declaration_command",  # readonly VAR=value (true constants)
             # Note: variable_assignment excluded to avoid duplicates within declaration_command
         ],
-        'import_types': [
-            'command',  # Can detect source/. commands via command name
+        "import_types": [
+            "command",  # Can detect source/. commands via command name
         ],
-        'call_types': [
-            'command',           # foo arg1 arg2
-            'command_substitution',  # $(foo)
+        "call_types": [
+            "command",  # foo arg1 arg2
+            "command_substitution",  # $(foo)
         ],
-        'variable_types': [
-            'variable_assignment',  # VAR=value
+        "variable_types": [
+            "variable_assignment",  # VAR=value
         ],
-        'docstring_field': 'comment',  # # ...
-        'identifier_field': 'name',
-        'body_field': 'body',
-
+        "docstring_field": "comment",  # # ...
+        "identifier_field": "name",
+        "body_field": "body",
         # Bash-specific patterns
-        'source_patterns': [
-            'source',  # source script.sh
-            '.',       # . script.sh
+        "source_patterns": [
+            "source",  # source script.sh
+            ".",  # . script.sh
         ],
-
         # Examples of AST nodes:
         # function_definition:
         #   name: word
@@ -422,51 +397,48 @@
         #   name: variable_name
         #   value: word | string | command_substitution
     },
-
     # --------------------------------------------------------------------------
     # Go Configuration
     # --------------------------------------------------------------------------
-    'go': {
-        'function_types': [
-            'function_declaration',  # func foo() { ... }
+    "go": {
+        "function_types": [
+            "function_declaration",  # func foo() { ... }
         ],
-        'class_types': [
-            'type_declaration',  # type Foo struct { ... }
+        "class_types": [
+            "type_declaration",  # type Foo struct { ... }
         ],
-        'method_types': [
-            'method_declaration',  # func (r Receiver) foo() { ... }
+        "method_types": [
+            "method_declaration",  # func (r Receiver) foo() { ... }
         ],
-        'constant_types': [
-            'const_declaration',  # const ( MaxRetries = 3 ... )
+        "constant_types": [
+            "const_declaration",  # const ( MaxRetries = 3 ... )
         ],
-        'import_types': [
-            'import_declaration',  # import "fmt" or import ( ... )
-            'import_spec',         # Individual import within import block
+        "import_types": [
+            "import_declaration",  # import "fmt" or import ( ... )
+            "import_spec",  # Individual import within import block
         ],
-        'call_types': [
-            'call_expression',  # foo()
+        "call_types": [
+            "call_expression",  # foo()
         ],
-        'interface_types': [
-            'interface_type',  # interface { ... } within type_declaration
+        "interface_types": [
+            "interface_type",  # interface { ... } within type_declaration
         ],
-        'struct_types': [
-            'struct_type',  # struct { ... } within type_declaration
+        "struct_types": [
+            "struct_type",  # struct { ... } within type_declaration
         ],
-        'docstring_field': 'comment',  # GoDoc comments (// ...)
-        'identifier_field': 'name',
-        'body_field': 'body',
-
+        "docstring_field": "comment",  # GoDoc comments (// ...)
+        "identifier_field": "name",
+        "body_field": "body",
         # Additional Go patterns
-        'package_types': [
-            'package_clause',  # package main
+        "package_types": [
+            "package_clause",  # package main
         ],
-        'const_types': [
-            'const_declaration',  # const Foo = 123
+        "const_types": [
+            "const_declaration",  # const Foo = 123
         ],
-        'var_types': [
-            'var_declaration',  # var foo int
+        "var_types": [
+            "var_declaration",  # var foo int
         ],
-
         # Examples of AST nodes:
         # function_declaration:
         #   name: identifier
@@ -483,7 +455,7 @@
         #
         # type_declaration:
         #   name: type_identifier
-        #   type: struct_type | interface_type | ...
+        # type: struct_type | interface_type | ...
         #
         # import_declaration:
         #   import_spec: package_identifier string_literal
@@ -492,52 +464,49 @@
         #   function: identifier | selector_expression
         #   arguments: argument_list
     },
-
     # --------------------------------------------------------------------------
     # Java Configuration
     # --------------------------------------------------------------------------
-    'java': {
-        'function_types': [
-            'method_declaration',  # public void foo() { ... }
+    "java": {
+        "function_types": [
+            "method_declaration",  # public void foo() { ... }
         ],
-        'class_types': [
-            'class_declaration',      # public class MyClass { ... }
-            'interface_declaration',  # public interface MyInterface { ... }
-            'enum_declaration',       # public enum MyEnum { ... }
+        "class_types": [
+            "class_declaration",  # public class MyClass { ... }
+            "interface_declaration",  # public interface MyInterface { ... }
+            "enum_declaration",  # public enum MyEnum { ... }
         ],
-        'method_types': [
-            'method_declaration',       # Methods inside class bodies
-            'constructor_declaration',  # public MyClass() { ... }
+        "method_types": [
+            "method_declaration",  # Methods inside class bodies
+            "constructor_declaration",  # public MyClass() { ... }
         ],
-        'constant_types': [
-            'field_declaration',  # private static final int MAX = 100;
+        "constant_types": [
+            "field_declaration",  # private static final int MAX = 100;
         ],
-        'import_types': [
-            'import_declaration',  # import java.util.List;
+        "import_types": [
+            "import_declaration",  # import java.util.List;
         ],
-        'call_types': [
-            'method_invocation',      # obj.method() or method()
-            'object_creation_expression',  # new MyClass()
+        "call_types": [
+            "method_invocation",  # obj.method() or method()
+            "object_creation_expression",  # new MyClass()
         ],
-        'package_types': [
-            'package_declaration',  # package com.example;
+        "package_types": [
+            "package_declaration",  # package com.example;
         ],
-        'annotation_types': [
-            'marker_annotation',   # @Override
-            'annotation',          # @SuppressWarnings("unchecked")
+        "annotation_types": [
+            "marker_annotation",  # @Override
+            "annotation",  # @SuppressWarnings("unchecked")
         ],
-        'docstring_field': 'comment',  # Javadoc comments (/** ... */)
-        'identifier_field': 'name',
-        'body_field': 'body',
-
+        "docstring_field": "comment",  # Javadoc comments (/** ... */)
+        "identifier_field": "name",
+        "body_field": "body",
         # Additional Java-specific nodes
-        'interface_types': [
-            'interface_declaration',  # interface Foo { ... }
+        "interface_types": [
+            "interface_declaration",  # interface Foo { ... }
         ],
-        'enum_types': [
-            'enum_declaration',  # enum Foo { A, B, C }
+        "enum_types": [
+            "enum_declaration",  # enum Foo { A, B, C }
         ],
-
         # Examples of AST nodes:
         # class_declaration:
         #   modifiers: public, abstract, final, etc.
@@ -548,7 +517,7 @@
         #
         # method_declaration:
         #   modifiers: public, static, etc.
-        #   type: type_identifier | void_type
+        # type: type_identifier | void_type
         #   name: identifier
         #   parameters: formal_parameters
         #   body: block
@@ -556,58 +525,55 @@
         # import_declaration:
         #   path: scoped_identifier
     },
-
     # --------------------------------------------------------------------------
     # Kotlin Configuration
     # --------------------------------------------------------------------------
-    'kotlin': {
-        'function_types': [
-            'function_declaration',  # fun foo() { ... }
+    "kotlin": {
+        "function_types": [
+            "function_declaration",  # fun foo() { ... }
         ],
-        'class_types': [
-            'class_declaration',   # class MyClass { ... }
-            'object_declaration',  # object Singleton { ... }
+        "class_types": [
+            "class_declaration",  # class MyClass { ... }
+            "object_declaration",  # object Singleton { ... }
         ],
-        'method_types': [
-            'function_declaration',  # Methods inside class bodies (same node type)
+        "method_types": [
+            "function_declaration",  # Methods inside class bodies (same node type)
         ],
-        'constant_types': [
-            'property_declaration',  # val/var declarations
+        "constant_types": [
+            "property_declaration",  # val/var declarations
         ],
-        'import_types': [
-            'import_header',  # import kotlin.collections.List
+        "import_types": [
+            "import_header",  # import kotlin.collections.List
         ],
-        'call_types': [
-            'call_expression',  # foo() or obj.foo()
+        "call_types": [
+            "call_expression",  # foo() or obj.foo()
         ],
-        'package_types': [
-            'package_header',  # package com.example
+        "package_types": [
+            "package_header",  # package com.example
         ],
-        'annotation_types': [
-            'annotation',  # @Annotation
+        "annotation_types": [
+            "annotation",  # @Annotation
         ],
-        'docstring_field': 'comment',  # KDoc comments (/** ... */)
-        'identifier_field': 'simple_identifier',
-        'body_field': 'class_body',
-
+        "docstring_field": "comment",  # KDoc comments (/** ... */)
+        "identifier_field": "simple_identifier",
+        "body_field": "class_body",
         # Kotlin-specific features
-        'interface_types': [
-            'class_declaration',  # interface in Kotlin uses class_declaration
+        "interface_types": [
+            "class_declaration",  # interface in Kotlin uses class_declaration
         ],
-        'enum_types': [
-            'class_declaration',  # enum class uses class_declaration
+        "enum_types": [
+            "class_declaration",  # enum class uses class_declaration
         ],
-        'object_types': [
-            'object_declaration',    # object Singleton { ... }
-            'companion_object',      # companion object { ... }
+        "object_types": [
+            "object_declaration",  # object Singleton { ... }
+            "companion_object",  # companion object { ... }
         ],
-        'data_class_types': [
-            'class_declaration',  # data class uses class_declaration with modifier
+        "data_class_types": [
+            "class_declaration",  # data class uses class_declaration with modifier
         ],
-        'lambda_types': [
-            'lambda_literal',  # { x -> x + 1 }
+        "lambda_types": [
+            "lambda_literal",  # { x -> x + 1 }
         ],
-
         # Examples of AST nodes:
         # class_declaration:
         #   modifiers: data, sealed, open, etc.
@@ -637,6 +603,7 @@
 # Helper Functions
 # ==============================================================================
 
+
 def get_language_for_file(filepath: str) -> Optional[str]:
     """
     Determine the programming language from a file path.
@@ -715,14 +682,14 @@ def validate_config(language: str) -> bool:
         ValueError: If configuration is missing required fields
     """
     required_fields = [
-        'function_types',
-        'class_types',
-        'method_types',
-        'import_types',
-        'call_types',
-        'docstring_field',
-        'identifier_field',
-        'body_field',
+        "function_types",
+        "class_types",
+        "method_types",
+        "import_types",
+        "call_types",
+        "docstring_field",
+        "identifier_field",
+        "body_field",
     ]
 
     config = get_config_for_language(language)
diff --git a/src/codegrok_mcp/parsers/treesitter_parser.py b/src/codegrok_mcp/parsers/treesitter_parser.py
index 1b13ee5..200654f 100644
--- a/src/codegrok_mcp/parsers/treesitter_parser.py
+++ b/src/codegrok_mcp/parsers/treesitter_parser.py
@@ -148,8 +148,7 @@ def parse_file(self, filepath: str) -> ParsedFile:
         file_size_mb = file_path.stat().st_size / (1024 * 1024)
         if file_size_mb > self.MAX_FILE_SIZE_MB:
             logger.warning(
-                f"Large file ({file_size_mb:.2f}MB): {filepath}. "
-                f"Parsing may be slow."
+                f"Large file ({file_size_mb:.2f}MB): {filepath}. " f"Parsing may be slow."
             )
 
         # Read file content
@@ -214,9 +213,7 @@ def parse_file(self, filepath: str) -> ParsedFile:
             root_node = tree.root_node
 
             # Extract symbols and imports
-            symbols = self._extract_symbols(
-                root_node, content, filepath, language, config
-            )
+            symbols = self._extract_symbols(root_node, content, filepath, language, config)
             imports = self._extract_imports(root_node, content, config)
 
             parse_time = time.time() - start_time
@@ -279,13 +276,13 @@ def _is_binary_file(self, content: bytes) -> bool:
         sample = content[:8192]
 
         # Check for null bytes (common in binary files)
-        if b'\x00' in sample:
+        if b"\x00" in sample:
             return True
 
         # Check for high ratio of non-text bytes
         try:
             # Try to decode as UTF-8
-            sample.decode('utf-8')
+            sample.decode("utf-8")
             return False
         except UnicodeDecodeError:
             # Count how many bytes fail to decode
@@ -341,10 +338,10 @@ def _extract_symbols(
         class_stack: List[str] = []
 
         # Get node types from config
-        function_types = config.get('function_types', [])
-        class_types = config.get('class_types', [])
-        method_types = config.get('method_types', [])
-        constant_types = config.get('constant_types', [])
+        function_types = config.get("function_types", [])
+        class_types = config.get("class_types", [])
+        method_types = config.get("method_types", [])
+        constant_types = config.get("constant_types", [])
 
         def traverse(node, depth: int = 0):
             """Recursively traverse the AST and extract symbols."""
@@ -354,9 +351,7 @@ def traverse(node, depth: int = 0):
 
             # Extract classes
             if node_type in class_types:
-                symbol = self._extract_class_symbol(
-                    node, content, filepath, language, config
-                )
+                symbol = self._extract_class_symbol(node, content, filepath, language, config)
                 if symbol:
                     symbols.append(symbol)
                     # Push class onto stack for method extraction
@@ -387,9 +382,7 @@ def traverse(node, depth: int = 0):
 
             # Extract constants (module-level only)
             elif node_type in constant_types and not current_class:
-                symbol = self._extract_constant_symbol(
-                    node, content, filepath, language, config
-                )
+                symbol = self._extract_constant_symbol(node, content, filepath, language, config)
                 if symbol:
                     symbols.append(symbol)
 
@@ -436,7 +429,7 @@ def _extract_class_symbol(
         line_end = node.end_point[0] + 1
 
         # Extract signature
-        signature = self._get_node_text(node, content).split('\n')[0].strip()
+        signature = self._get_node_text(node, content).split("\n")[0].strip()
 
         # Extract docstring
         docstring = self._extract_docstring(node, content, config)
@@ -499,7 +492,7 @@ def _extract_function_symbol(
         line_end = node.end_point[0] + 1
 
         # Extract signature
-        signature = self._get_node_text(node, content).split('\n')[0].strip()
+        signature = self._get_node_text(node, content).split("\n")[0].strip()
 
         # Extract docstring
         docstring = self._extract_docstring(node, content, config)
@@ -563,71 +556,83 @@ def _extract_constant_symbol(
         # Extract the name based on language patterns
         name = None
 
-        if language == 'python':
+        if language == "python":
             # Python: expression_statement -> assignment -> identifier
             # Look for UPPERCASE names (convention for constants)
             for child in node.children:
-                if child.type == 'assignment':
+                if child.type == "assignment":
                     for subchild in child.children:
-                        if subchild.type == 'identifier':
+                        if subchild.type == "identifier":
                             potential_name = self._get_node_text(subchild, content)
                             # Only extract UPPERCASE constants
-                            if potential_name.isupper() or '_' in potential_name and potential_name.replace('_', '').isupper():
+                            if (
+                                potential_name.isupper()
+                                or "_" in potential_name
+                                and potential_name.replace("_", "").isupper()
+                            ):
                                 name = potential_name
                             break
                     break
 
-        elif language == 'javascript' or language == 'typescript':
+        elif language == "javascript" or language == "typescript":
             # JavaScript: lexical_declaration -> variable_declarator -> identifier
             # Only extract const with UPPERCASE names (convention for constants)
-            if full_text.startswith('const '):
+            if full_text.startswith("const "):
                 for child in node.children:
-                    if child.type == 'variable_declarator':
+                    if child.type == "variable_declarator":
                         for subchild in child.children:
-                            if subchild.type == 'identifier':
+                            if subchild.type == "identifier":
                                 potential_name = self._get_node_text(subchild, content)
                                 # Only UPPERCASE constants
-                                if potential_name.isupper() or ('_' in potential_name and potential_name.replace('_', '').isupper()):
+                                if potential_name.isupper() or (
+                                    "_" in potential_name
+                                    and potential_name.replace("_", "").isupper()
+                                ):
                                     name = potential_name
                                 break
                         break
 
                         break
-        
+
         # Add pragma for other languages not covered by tests yet
-        elif language == 'go':  # pragma: no cover
+        elif language == "go":  # pragma: no cover
             # Go: const_declaration -> const_spec -> identifier
             # Return all const names as one symbol (grouped)
             const_names = []
             for child in node.children:
-                if child.type == 'const_spec':
+                if child.type == "const_spec":
                     for subchild in child.children:
-                        if subchild.type == 'identifier':
+                        if subchild.type == "identifier":
                             const_names.append(self._get_node_text(subchild, content))
                             break
             if const_names:
-                name = ', '.join(const_names)
+                name = ", ".join(const_names)
 
-        elif language == 'bash':  # pragma: no cover
+        elif language == "bash":  # pragma: no cover
             # Bash: variable_assignment or declaration_command
             # Only extract UPPERCASE variables (constants by convention)
-            if node.type == 'declaration_command':
+            if node.type == "declaration_command":
                 # readonly VAR=value
                 for child in node.children:
-                    if child.type == 'variable_assignment':
+                    if child.type == "variable_assignment":
                         for subchild in child.children:
-                            if subchild.type == 'variable_name':
+                            if subchild.type == "variable_name":
                                 potential_name = self._get_node_text(subchild, content)
-                                if potential_name.isupper() or ('_' in potential_name and potential_name.replace('_', '').isupper()):
+                                if potential_name.isupper() or (
+                                    "_" in potential_name
+                                    and potential_name.replace("_", "").isupper()
+                                ):
                                     name = potential_name
                                 break
                         break
-            elif node.type == 'variable_assignment':
+            elif node.type == "variable_assignment":
                 # VAR=value - only at top level with UPPERCASE
                 for child in node.children:
-                    if child.type == 'variable_name':
+                    if child.type == "variable_name":
                         potential_name = self._get_node_text(child, content)
-                        if potential_name.isupper() or ('_' in potential_name and potential_name.replace('_', '').isupper()):
+                        if potential_name.isupper() or (
+                            "_" in potential_name and potential_name.replace("_", "").isupper()
+                        ):
                             name = potential_name
                         break
 
@@ -653,9 +658,7 @@ def _extract_constant_symbol(
             calls=[],
         )
 
-    def _extract_imports(
-        self, root_node, content: bytes, config: Dict[str, Any]
-    ) -> List[str]:
+    def _extract_imports(self, root_node, content: bytes, config: Dict[str, Any]) -> List[str]:
         """
         Extract all import statements from the file.
 
@@ -668,7 +671,7 @@ def _extract_imports(
             List of import statement strings
         """
         imports: List[str] = []
-        import_types = config.get('import_types', [])
+        import_types = config.get("import_types", [])
 
         def traverse(node):
             if node.type in import_types:
@@ -682,9 +685,7 @@ def traverse(node):
         traverse(root_node)
         return imports
 
-    def _extract_imports_from_node(
-        self, node, content: bytes, config: Dict[str, Any]
-    ) -> List[str]:
+    def _extract_imports_from_node(self, node, content: bytes, config: Dict[str, Any]) -> List[str]:
         """
         Extract imports from a specific node (for scoped imports).
 
@@ -697,7 +698,7 @@ def _extract_imports_from_node(
             List of import statement strings
         """
         imports: List[str] = []
-        import_types = config.get('import_types', [])
+        import_types = config.get("import_types", [])
 
         def traverse(n):
             if n.type in import_types:
@@ -711,9 +712,7 @@ def traverse(n):
         traverse(node)
         return imports
 
-    def _extract_calls_from_node(
-        self, node, content: bytes, config: Dict[str, Any]
-    ) -> List[str]:
+    def _extract_calls_from_node(self, node, content: bytes, config: Dict[str, Any]) -> List[str]:
         """
         Extract function calls from a specific node.
 
@@ -726,7 +725,7 @@ def _extract_calls_from_node(
             List of function call names
         """
         calls: Set[str] = set()
-        call_types = config.get('call_types', [])
+        call_types = config.get("call_types", [])
 
         def traverse(n):
             if n.type in call_types:
@@ -754,34 +753,32 @@ def _get_call_name(self, call_node, content: bytes) -> Optional[str]:
         """
         # Try to find the function identifier
         for child in call_node.children:
-            if child.type in ('identifier', 'name', 'word', 'field_identifier'):
+            if child.type in ("identifier", "name", "word", "field_identifier"):
                 return self._get_node_text(child, content).strip()
-            elif child.type == 'attribute':
+            elif child.type == "attribute":
                 # For method calls like obj.method()
                 for subchild in child.children:
-                    if subchild.type in ('identifier', 'property_identifier', 'field_identifier'):
+                    if subchild.type in ("identifier", "property_identifier", "field_identifier"):
                         return self._get_node_text(subchild, content).strip()
-            elif child.type == 'member_expression':
+            elif child.type == "member_expression":
                 # JavaScript/TypeScript member expressions
                 for subchild in child.children:
-                    if subchild.type in ('property_identifier', 'identifier'):
+                    if subchild.type in ("property_identifier", "identifier"):
                         return self._get_node_text(subchild, content).strip()
-            elif child.type == 'selector_expression':
+            elif child.type == "selector_expression":
                 # Go selector expressions
                 for subchild in child.children:
-                    if subchild.type in ('field_identifier', 'identifier'):
+                    if subchild.type in ("field_identifier", "identifier"):
                         return self._get_node_text(subchild, content).strip()
 
         # Fallback: try to get first identifier child
         if call_node.child_count > 0:
             first_child = call_node.children[0]
-            return self._get_node_text(first_child, content).strip().split('(')[0]
+            return self._get_node_text(first_child, content).strip().split("(")[0]
 
         return None
 
-    def _extract_docstring(
-        self, node, content: bytes, config: Dict[str, Any]
-    ) -> str:
+    def _extract_docstring(self, node, content: bytes, config: Dict[str, Any]) -> str:
         """
         Extract docstring from a function or class node.
 
@@ -800,14 +797,12 @@ def _extract_docstring(
 
         # Check first child of body for docstring
         for child in body.children:
-            if child.type == 'expression_statement':
+            if child.type == "expression_statement":
                 # Python: first expression statement might be docstring
                 for subchild in child.children:
-                    if subchild.type in ('string', 'string_literal'):
-                        return self._clean_docstring(
-                            self._get_node_text(subchild, content)
-                        )
-            elif child.type in ('string', 'string_literal', 'comment'):
+                    if subchild.type in ("string", "string_literal"):
+                        return self._clean_docstring(self._get_node_text(subchild, content))
+            elif child.type in ("string", "string_literal", "comment"):
                 return self._clean_docstring(self._get_node_text(child, content))
 
         return ""
@@ -823,15 +818,15 @@ def _get_body_node(self, node, config: Dict[str, Any]):
         Returns:
             Body node or None
         """
-        body_field = config.get('body_field', 'body')
+        body_field = config.get("body_field", "body")
 
         # Try named child first
         for child in node.children:
-            if child.type in ('block', 'body', 'compound_statement', 'statement_block'):
+            if child.type in ("block", "body", "compound_statement", "statement_block"):
                 return child
 
         # Try field access
-        if hasattr(node, 'child_by_field_name'):
+        if hasattr(node, "child_by_field_name"):
             body = node.child_by_field_name(body_field)
             if body:
                 return body
@@ -851,25 +846,25 @@ def _get_node_name(self, node, content: bytes, config: Dict[str, Any]) -> Option
             Name string or None
         """
         # Try field-based access first
-        if hasattr(node, 'child_by_field_name'):
-            name_node = node.child_by_field_name('name')
+        if hasattr(node, "child_by_field_name"):
+            name_node = node.child_by_field_name("name")
             if name_node:
                 return self._get_node_text(name_node, content).strip()
 
         # Try finding identifier child
         for child in node.children:
             if child.type in (
-                'identifier',
-                'name',
-                'type_identifier',
-                'field_identifier',
-                'property_identifier',
+                "identifier",
+                "name",
+                "type_identifier",
+                "field_identifier",
+                "property_identifier",
             ):
                 return self._get_node_text(child, content).strip()
 
         # For C/C++ functions with declarators
         for child in node.children:
-            if child.type in ('declarator', 'function_declarator'):
+            if child.type in ("declarator", "function_declarator"):
                 return self._get_node_name(child, content, config)
 
         return None
@@ -886,10 +881,10 @@ def _get_node_text(self, node, content: bytes) -> str:
             Node text as string
         """
         try:
-            return content[node.start_byte : node.end_byte].decode('utf-8')
+            return content[node.start_byte : node.end_byte].decode("utf-8")
         except UnicodeDecodeError:
             # Fallback for binary content
-            return content[node.start_byte : node.end_byte].decode('utf-8', errors='ignore')
+            return content[node.start_byte : node.end_byte].decode("utf-8", errors="ignore")
 
     def _get_code_snippet(self, node, content: bytes) -> str:
         """
@@ -938,7 +933,7 @@ def _clean_docstring(self, raw_docstring: str) -> str:
         cleaned = cleaned.strip()
 
         # For multi-line docstrings, extract just the first line/paragraph
-        lines = [line.strip() for line in cleaned.split('\n') if line.strip()]
+        lines = [line.strip() for line in cleaned.split("\n") if line.strip()]
         if lines:
             return lines[0]
 
@@ -991,6 +986,6 @@ def get_parser(self) -> TreeSitterParser:
         Returns:
             TreeSitterParser instance unique to the calling thread
         """
-        if not hasattr(self._local, 'parser'):
+        if not hasattr(self._local, "parser"):
             self._local.parser = TreeSitterParser()
         return self._local.parser
diff --git a/tests/conftest.py b/tests/conftest.py
index 2f59562..b496e39 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -3,6 +3,7 @@
 from pathlib import Path
 from codegrok_mcp.mcp.state import reset_state, get_state
 
+
 @pytest.fixture(autouse=True)
 def clean_state():
     """Reset MCP state before and after each test."""
@@ -10,6 +11,7 @@ def clean_state():
     yield
     reset_state()
 
+
 @pytest.fixture
 def temp_project(tmp_path):
     """Create a minimal Python project for testing."""
@@ -43,19 +45,22 @@ def load_config(path: str) -> dict:
 
     return tmp_path
 
+
 @pytest.fixture
 def multi_lang_project(tmp_path):
     """Create a multi-language project."""
-    (tmp_path / "app.py").write_text('def main(): pass')
-    (tmp_path / "helper.js").write_text('function helper() { return 1; }')
-    (tmp_path / "server.go").write_text('package main\n\nfunc main() {}')
+    (tmp_path / "app.py").write_text("def main(): pass")
+    (tmp_path / "helper.js").write_text("function helper() { return 1; }")
+    (tmp_path / "server.go").write_text("package main\n\nfunc main() {}")
     return tmp_path
 
+
 @pytest.fixture
 def python_project_fixture():
     """Return path to the static Python project fixture."""
     return Path(__file__).parent / "fixtures" / "sample_projects" / "python_project"
 
+
 @pytest.fixture
 def multi_lang_fixture():
     """Return path to the static multi-language fixture."""
diff --git a/tests/fixtures/sample_projects/python_project/main.py b/tests/fixtures/sample_projects/python_project/main.py
index 99b8b1d..9a29213 100644
--- a/tests/fixtures/sample_projects/python_project/main.py
+++ b/tests/fixtures/sample_projects/python_project/main.py
@@ -1,4 +1,5 @@
 """Main application module."""
+
 from typing import List, Optional
 
 
diff --git a/tests/fixtures/sample_projects/python_project/utils.py b/tests/fixtures/sample_projects/python_project/utils.py
index b500cad..04f6010 100644
--- a/tests/fixtures/sample_projects/python_project/utils.py
+++ b/tests/fixtures/sample_projects/python_project/utils.py
@@ -1,4 +1,5 @@
 """Utility functions for common operations."""
+
 import json
 import hashlib
 from pathlib import Path
diff --git a/tests/integration/test_source_retriever.py b/tests/integration/test_source_retriever.py
index 66bcd49..1ac5375 100644
--- a/tests/integration/test_source_retriever.py
+++ b/tests/integration/test_source_retriever.py
@@ -2,23 +2,27 @@
 Integration tests for SourceRetriever - the core indexing/search engine.
 These test the actual functionality without MCP overhead.
 """
+
+import json
 import pytest
 import tempfile
 from pathlib import Path
-from codegrok_mcp.indexing.source_retriever import SourceRetriever
+from codegrok_mcp.indexing.source_retriever import (
+    SourceRetriever,
+    _save_checkpoint,
+    _load_checkpoint,
+)
 
 
 from unittest.mock import patch
 
+
 class TestSourceRetrieverIndexing:
     """Test indexing functionality."""
 
     def test_index_python_project(self, temp_project):
         with tempfile.TemporaryDirectory() as persist_dir:
-            retriever = SourceRetriever(
-                codebase_path=str(temp_project),
-                persist_path=persist_dir
-            )
+            retriever = SourceRetriever(codebase_path=str(temp_project), persist_path=persist_dir)
 
             retriever.index_codebase()
             stats = retriever.get_stats()
@@ -29,8 +33,7 @@ def test_index_python_project(self, temp_project):
     def test_index_respects_file_extensions(self, multi_lang_project):
         with tempfile.TemporaryDirectory() as persist_dir:
             retriever = SourceRetriever(
-                codebase_path=str(multi_lang_project),
-                persist_path=persist_dir
+                codebase_path=str(multi_lang_project), persist_path=persist_dir
             )
 
             # Index only Python files
@@ -43,8 +46,7 @@ def test_index_creates_persist_directory(self, temp_project):
         with tempfile.TemporaryDirectory() as persist_dir:
             persist_path = Path(persist_dir) / "chroma"
             retriever = SourceRetriever(
-                codebase_path=str(temp_project),
-                persist_path=str(persist_path)
+                codebase_path=str(temp_project), persist_path=str(persist_path)
             )
 
             retriever.index_codebase()
@@ -58,42 +60,29 @@ class TestSourceRetrieverSearch:
     @pytest.fixture
     def indexed_retriever(self, temp_project):
         with tempfile.TemporaryDirectory() as persist_dir:
-            retriever = SourceRetriever(
-                codebase_path=str(temp_project),
-                persist_path=persist_dir
-            )
+            retriever = SourceRetriever(codebase_path=str(temp_project), persist_path=persist_dir)
             retriever.index_codebase()
             yield retriever
 
     def test_get_sources_before_indexing(self, temp_project):
         with tempfile.TemporaryDirectory() as persist_dir:
-            retriever = SourceRetriever(
-                codebase_path=str(temp_project),
-                persist_path=persist_dir
-            )
+            retriever = SourceRetriever(codebase_path=str(temp_project), persist_path=persist_dir)
             results, _ = retriever.get_sources_for_question("test")
             assert results == []
 
     def test_search_returns_results(self, indexed_retriever):
-        results, _ = indexed_retriever.get_sources_for_question(
-            "calculator",
-            n_results=5
-        )
+        results, _ = indexed_retriever.get_sources_for_question("calculator", n_results=5)
 
         assert len(results) > 0
 
     def test_search_respects_n_results(self, indexed_retriever):
-        results, _ = indexed_retriever.get_sources_for_question(
-            "function",
-            n_results=2
-        )
+        results, _ = indexed_retriever.get_sources_for_question("function", n_results=2)
 
         assert len(results) <= 2
 
     def test_search_returns_relevant_results(self, indexed_retriever):
         results, _ = indexed_retriever.get_sources_for_question(
-            "add numbers calculator",
-            n_results=5
+            "add numbers calculator", n_results=5
         )
 
         # Should find calculator-related code
@@ -106,17 +95,14 @@ class TestIncrementalReindex:
 
     def test_detects_modified_files(self, temp_project):
         with tempfile.TemporaryDirectory() as persist_dir:
-            retriever = SourceRetriever(
-                codebase_path=str(temp_project),
-                persist_path=persist_dir
-            )
+            retriever = SourceRetriever(codebase_path=str(temp_project), persist_path=persist_dir)
 
             # Initial index
             retriever.index_codebase()
 
             # Modify a file
             main_py = temp_project / "main.py"
-            main_py.write_text('def brand_new_function(): pass')
+            main_py.write_text("def brand_new_function(): pass")
 
             # Incremental reindex
             retriever.incremental_reindex()
@@ -128,17 +114,14 @@ def test_detects_modified_files(self, temp_project):
 
     def test_handles_new_files(self, temp_project):
         with tempfile.TemporaryDirectory() as persist_dir:
-            retriever = SourceRetriever(
-                codebase_path=str(temp_project),
-                persist_path=persist_dir
-            )
+            retriever = SourceRetriever(codebase_path=str(temp_project), persist_path=persist_dir)
 
             # Initial index
             retriever.index_codebase()
 
             # Add a new file with unique content
             new_file = temp_project / "new_module.py"
-            new_file.write_text('def unique_xyz_function_12345(): pass')
+            new_file.write_text("def unique_xyz_function_12345(): pass")
 
             # Incremental reindex
             retriever.incremental_reindex()
@@ -150,32 +133,108 @@ def test_handles_new_files(self, temp_project):
 
     def test_incremental_reindex_with_parse_error(self, temp_project):
         with tempfile.TemporaryDirectory() as persist_dir:
-            retriever = SourceRetriever(
-                codebase_path=str(temp_project),
-                persist_path=persist_dir
-            )
+            retriever = SourceRetriever(codebase_path=str(temp_project), persist_path=persist_dir)
             retriever.index_codebase()
-            
+
             # Create broken file
             (temp_project / "broken.py").write_text("def (")
-            
+
             # Should not raise exception
             result = retriever.incremental_reindex()
-            
+
             # Check result properties if possible, or just sufficient that it didn't crash
             assert result is not None
-            assert result['files_added'] == 1 or result['files_modified'] == 1
+            assert result["files_added"] == 1 or result["files_modified"] == 1
 
     def test_index_codebase_with_parse_error(self, temp_project):
         with tempfile.TemporaryDirectory() as persist_dir:
             retriever = SourceRetriever(
-                codebase_path=str(temp_project),
-                persist_path=persist_dir,
-                parallel=False
+                codebase_path=str(temp_project), persist_path=persist_dir, parallel=False
             )
-            
-            with patch.object(retriever.parser, 'parse_file', side_effect=Exception("Boom")):
+
+            with patch.object(retriever.parser, "parse_file", side_effect=Exception("Boom")):
                 retriever.index_codebase()
                 # Should handle error and continue/finish
-                assert retriever.stats['parse_errors'] > 0
+                assert retriever.stats["parse_errors"] > 0
+
+
+class TestUpsertBehavior:
+    """Test that upsert-based indexing is idempotent and handles stale chunks."""
+
+    def test_index_codebase_upsert_idempotent(self, temp_project):
+        """Re-indexing same codebase produces same chunk count (no duplicates)."""
+        with tempfile.TemporaryDirectory() as persist_dir:
+            retriever = SourceRetriever(codebase_path=str(temp_project), persist_path=persist_dir)
+
+            retriever.index_codebase()
+            count_first = retriever.collection.count()
+
+            # Re-index the same codebase
+            retriever.index_codebase()
+            count_second = retriever.collection.count()
+
+            assert count_first == count_second
+            assert count_first > 0
+
+    def test_stale_chunk_removal(self, temp_project):
+        """After deleting a file and re-indexing, old chunks are removed."""
+        with tempfile.TemporaryDirectory() as persist_dir:
+            retriever = SourceRetriever(codebase_path=str(temp_project), persist_path=persist_dir)
+
+            retriever.index_codebase()
+            count_before = retriever.collection.count()
+
+            # Delete a file
+            (temp_project / "utils.py").unlink()
+
+            # Re-index
+            retriever.index_codebase()
+            count_after = retriever.collection.count()
+
+            assert count_after < count_before
+
+
+class TestCheckpointing:
+    """Test checkpoint save/load/resume functionality."""
+
+    def test_checkpoint_save_and_load(self, tmp_path):
+        """Checkpoint round-trip: save → load → verify data."""
+        cp_path = tmp_path / "checkpoint.json"
+
+        _save_checkpoint(cp_path, chunks_completed=500, total_chunks=1000)
+
+        data = _load_checkpoint(cp_path)
+        assert data is not None
+        assert data["chunks_completed"] == 500
+        assert data["total_chunks"] == 1000
+        assert "timestamp" in data
+
+    def test_checkpoint_load_missing_file(self, tmp_path):
+        """Returns None when checkpoint file doesn't exist."""
+        data = _load_checkpoint(tmp_path / "nonexistent.json")
+        assert data is None
+
+    def test_checkpoint_load_corrupted(self, tmp_path):
+        """Returns None for corrupted checkpoint file."""
+        cp_path = tmp_path / "checkpoint.json"
+        cp_path.write_text("not valid json{{{")
+
+        data = _load_checkpoint(cp_path)
+        assert data is None
+
+    def test_checkpoint_cleanup_on_success(self, temp_project):
+        """Checkpoint file is deleted after successful indexing."""
+        with tempfile.TemporaryDirectory() as persist_dir:
+            retriever = SourceRetriever(codebase_path=str(temp_project), persist_path=persist_dir)
+
+            retriever.index_codebase()
+
+            checkpoint_path = Path(persist_dir).parent / "checkpoint.json"
+            # Also check the actual persist dir parent
+            actual_cp = Path(persist_dir) / ".." / "checkpoint.json"
+            assert not checkpoint_path.exists() or not actual_cp.resolve().exists()
 
+    def test_checkpoint_load_none_path(self):
+        """Returns None when checkpoint_path is None."""
+        data = _load_checkpoint(None)
+        assert data is None
diff --git a/tests/integration/test_tool_discovery.py b/tests/integration/test_tool_discovery.py
index f968dd1..90a9464 100644
--- a/tests/integration/test_tool_discovery.py
+++ b/tests/integration/test_tool_discovery.py
@@ -38,11 +38,11 @@
 )
 from codegrok_mcp.mcp.state import get_state, reset_state
 
-
 # =============================================================================
 # Test Fixtures
 # =============================================================================
 
+
 @pytest.fixture
 def indexed_project(temp_project):
     """Create a project and index it for tests that need pre-indexed state."""
@@ -99,7 +99,7 @@ def save_config(path: str, config: dict) -> None:
 ''')
 
     # JavaScript file
-    (tmp_path / "src" / "helper.js").write_text(r'''/**
+    (tmp_path / "src" / "helper.js").write_text(r"""/**
  * Helper functions for the application.
  */
 
@@ -116,7 +116,7 @@ def save_config(path: str, config: dict) -> None:
     maxItems: 100,
     timeout: 5000
 };
-''')
+""")
 
     return tmp_path
 
@@ -125,6 +125,7 @@ def save_config(path: str, config: dict) -> None:
 # Helper Functions
 # =============================================================================
 
+
 def learn(**kwargs):
     """Helper to run async learn function synchronously."""
     return asyncio.get_event_loop().run_until_complete(learn_tool.fn(**kwargs))
@@ -169,6 +170,7 @@ def list_supported_languages():
 # Tool Discovery Tests
 # =============================================================================
 
+
 class TestToolDiscovery:
     """Test that all 8 tools are properly exposed via MCP."""
 
@@ -214,80 +216,82 @@ def test_memory_tools_present(self):
 # Tool Annotations Tests
 # =============================================================================
 
+
 class TestToolAnnotations:
     """Test that tool annotations are correctly set per MCP spec."""
 
     def _get_tool_annotations(self, tool_name: str) -> Dict[str, Any]:
         """Get annotations for a tool by name."""
         tool = mcp._tool_manager._tools.get(tool_name)
-        if tool and hasattr(tool, 'annotations') and tool.annotations:
+        if tool and hasattr(tool, "annotations") and tool.annotations:
             return {
-                'readOnlyHint': tool.annotations.readOnlyHint,
-                'destructiveHint': tool.annotations.destructiveHint,
-                'idempotentHint': tool.annotations.idempotentHint,
-                'openWorldHint': tool.annotations.openWorldHint,
+                "readOnlyHint": tool.annotations.readOnlyHint,
+                "destructiveHint": tool.annotations.destructiveHint,
+                "idempotentHint": tool.annotations.idempotentHint,
+                "openWorldHint": tool.annotations.openWorldHint,
             }
         return {}
 
     def test_learn_annotations(self):
         """learn: writes data, not destructive, idempotent, local only."""
         annotations = self._get_tool_annotations("learn")
-        assert annotations.get('readOnlyHint') is False, "learn writes to .codegrok/"
-        assert annotations.get('destructiveHint') is False, "learn doesn't destroy user data"
-        assert annotations.get('idempotentHint') is True, "learn can be safely re-run"
-        assert annotations.get('openWorldHint') is False, "learn only accesses local files"
+        assert annotations.get("readOnlyHint") is False, "learn writes to .codegrok/"
+        assert annotations.get("destructiveHint") is False, "learn doesn't destroy user data"
+        assert annotations.get("idempotentHint") is True, "learn can be safely re-run"
+        assert annotations.get("openWorldHint") is False, "learn only accesses local files"
 
     def test_get_sources_annotations(self):
         """get_sources: read-only search."""
         annotations = self._get_tool_annotations("get_sources")
-        assert annotations.get('readOnlyHint') is True, "get_sources only reads"
-        assert annotations.get('idempotentHint') is True, "Same query = same results"
-        assert annotations.get('openWorldHint') is False, "Local ChromaDB only"
+        assert annotations.get("readOnlyHint") is True, "get_sources only reads"
+        assert annotations.get("idempotentHint") is True, "Same query = same results"
+        assert annotations.get("openWorldHint") is False, "Local ChromaDB only"
 
     def test_get_stats_annotations(self):
         """get_stats: read-only metadata."""
         annotations = self._get_tool_annotations("get_stats")
-        assert annotations.get('readOnlyHint') is True
-        assert annotations.get('idempotentHint') is True
+        assert annotations.get("readOnlyHint") is True
+        assert annotations.get("idempotentHint") is True
 
     def test_list_supported_languages_annotations(self):
         """list_supported_languages: static data, always read-only."""
         annotations = self._get_tool_annotations("list_supported_languages")
-        assert annotations.get('readOnlyHint') is True
-        assert annotations.get('idempotentHint') is True
+        assert annotations.get("readOnlyHint") is True
+        assert annotations.get("idempotentHint") is True
 
     def test_remember_annotations(self):
         """remember: writes data, not destructive, NOT idempotent."""
         annotations = self._get_tool_annotations("remember")
-        assert annotations.get('readOnlyHint') is False, "remember writes to ChromaDB"
-        assert annotations.get('destructiveHint') is False, "remember adds, doesn't delete"
-        assert annotations.get('idempotentHint') is False, "Each call creates new memory"
-        assert annotations.get('openWorldHint') is False
+        assert annotations.get("readOnlyHint") is False, "remember writes to ChromaDB"
+        assert annotations.get("destructiveHint") is False, "remember adds, doesn't delete"
+        assert annotations.get("idempotentHint") is False, "Each call creates new memory"
+        assert annotations.get("openWorldHint") is False
 
     def test_recall_annotations(self):
         """recall: read-only search."""
         annotations = self._get_tool_annotations("recall")
-        assert annotations.get('readOnlyHint') is True
-        assert annotations.get('idempotentHint') is True
+        assert annotations.get("readOnlyHint") is True
+        assert annotations.get("idempotentHint") is True
 
     def test_forget_annotations_destructive(self):
         """forget: DESTRUCTIVE - permanently deletes data."""
         annotations = self._get_tool_annotations("forget")
-        assert annotations.get('readOnlyHint') is False, "forget deletes data"
-        assert annotations.get('destructiveHint') is True, "forget is DESTRUCTIVE"
-        assert annotations.get('idempotentHint') is True, "Re-calling same filter is safe"
+        assert annotations.get("readOnlyHint") is False, "forget deletes data"
+        assert annotations.get("destructiveHint") is True, "forget is DESTRUCTIVE"
+        assert annotations.get("idempotentHint") is True, "Re-calling same filter is safe"
 
     def test_memory_stats_annotations(self):
         """memory_stats: read-only statistics."""
         annotations = self._get_tool_annotations("memory_stats")
-        assert annotations.get('readOnlyHint') is True
-        assert annotations.get('idempotentHint') is True
+        assert annotations.get("readOnlyHint") is True
+        assert annotations.get("idempotentHint") is True
 
 
 # =============================================================================
 # Server Instructions Tests
 # =============================================================================
 
+
 class TestServerInstructions:
     """Test that server instructions properly guide agents."""
 
@@ -328,6 +332,7 @@ def test_instructions_describe_memory_use_cases(self):
 # Progressive Discovery Tests (Learn-First Requirement)
 # =============================================================================
 
+
 class TestLearnFirstRequirement:
     """Test that tools properly require 'learn' to be called first."""
 
@@ -381,6 +386,7 @@ def test_list_supported_languages_works_without_learn(self):
 # Error Message Quality Tests
 # =============================================================================
 
+
 class TestErrorMessageQuality:
     """Test that error messages guide agents to correct actions."""
 
@@ -419,6 +425,7 @@ def test_invalid_memory_type_lists_valid_types(self, indexed_project):
 # Tool Description Tests
 # =============================================================================
 
+
 class TestToolDescriptions:
     """Test that tool descriptions are informative and consistent."""
 
@@ -464,6 +471,7 @@ def test_memory_tools_describe_types(self):
 # Protocol-Level Tests (Subprocess)
 # =============================================================================
 
+
 class StdioMCPClient:
     """MCP client that communicates via stdio subprocess."""
 
@@ -479,7 +487,7 @@ def start(self):
             stdout=subprocess.PIPE,
             stderr=subprocess.PIPE,
             text=True,
-            bufsize=1
+            bufsize=1,
         )
         time.sleep(0.5)  # Allow server to initialize
 
@@ -499,7 +507,7 @@ def _send_request(self, method: str, params: dict = None) -> dict:
             "jsonrpc": "2.0",
             "id": self.request_id,
             "method": method,
-            "params": params or {}
+            "params": params or {},
         }
 
         request_line = json.dumps(request) + "\n"
@@ -514,11 +522,14 @@ def _send_request(self, method: str, params: dict = None) -> dict:
 
     def initialize(self) -> dict:
         """MCP initialization handshake."""
-        return self._send_request("initialize", {
-            "protocolVersion": "2024-11-05",
-            "capabilities": {},
-            "clientInfo": {"name": "integration-test", "version": "1.0.0"}
-        })
+        return self._send_request(
+            "initialize",
+            {
+                "protocolVersion": "2024-11-05",
+                "capabilities": {},
+                "clientInfo": {"name": "integration-test", "version": "1.0.0"},
+            },
+        )
 
     def list_tools(self) -> List[dict]:
         """Get available tools."""
@@ -527,10 +538,7 @@ def list_tools(self) -> List[dict]:
 
     def call_tool(self, name: str, arguments: dict = None) -> dict:
         """Call a tool."""
-        return self._send_request("tools/call", {
-            "name": name,
-            "arguments": arguments or {}
-        })
+        return self._send_request("tools/call", {"name": name, "arguments": arguments or {}})
 
 
 class TestProtocolToolDiscovery:
@@ -550,8 +558,16 @@ def test_tools_list_returns_all_eight(self, mcp_client):
         tools = mcp_client.list_tools()
 
         tool_names = [t["name"] for t in tools]
-        expected = ["learn", "get_sources", "get_stats", "list_supported_languages",
-                    "remember", "recall", "forget", "memory_stats"]
+        expected = [
+            "learn",
+            "get_sources",
+            "get_stats",
+            "list_supported_languages",
+            "remember",
+            "recall",
+            "forget",
+            "memory_stats",
+        ]
 
         for name in expected:
             assert name in tool_names, f"Tool '{name}' not in tools/list response"
@@ -641,28 +657,29 @@ def test_full_workflow_via_protocol(self, mcp_client, temp_project):
         assert learn_data["success"] is True
 
         # 2. Get sources
-        search_resp = mcp_client.call_tool("get_sources", {
-            "question": "calculator add",
-            "n_results": 5
-        })
+        search_resp = mcp_client.call_tool(
+            "get_sources", {"question": "calculator add", "n_results": 5}
+        )
         search_data = json.loads(search_resp["result"]["content"][0]["text"])
         assert "sources" in search_data
 
         # 3. Remember
-        remember_resp = mcp_client.call_tool("remember", {
-            "content": "User prefers functional style",
-            "memory_type": "preference",
-            "tags": ["style", "coding"]
-        })
+        remember_resp = mcp_client.call_tool(
+            "remember",
+            {
+                "content": "User prefers functional style",
+                "memory_type": "preference",
+                "tags": ["style", "coding"],
+            },
+        )
         remember_data = json.loads(remember_resp["result"]["content"][0]["text"])
         assert remember_data["success"] is True
         assert "memory_id" in remember_data
 
         # 4. Recall
-        recall_resp = mcp_client.call_tool("recall", {
-            "query": "coding style preference",
-            "n_results": 5
-        })
+        recall_resp = mcp_client.call_tool(
+            "recall", {"query": "coding style preference", "n_results": 5}
+        )
         recall_data = json.loads(recall_resp["result"]["content"][0]["text"])
         assert recall_data["success"] is True
         assert recall_data["count"] >= 1
@@ -691,6 +708,7 @@ def test_error_response_before_learn(self, mcp_client):
 # Integration: Combined Workflow Tests
 # =============================================================================
 
+
 class TestCombinedWorkflow:
     """Test complete agent workflows combining code search and memory."""
 
@@ -709,7 +727,7 @@ def test_typical_agent_session(self, multi_file_project):
         memory1 = remember(
             content="Using Calculator class for all math operations",
             memory_type="decision",
-            tags=["architecture", "math"]
+            tags=["architecture", "math"],
         )
         assert memory1["success"] is True
 
@@ -717,7 +735,7 @@ def test_typical_agent_session(self, multi_file_project):
         memory2 = remember(
             content="User prefers docstrings on all public methods",
             memory_type="preference",
-            tags=["style", "documentation"]
+            tags=["style", "documentation"],
         )
         assert memory2["success"] is True
 
@@ -740,7 +758,7 @@ def test_memory_persistence_across_state(self, multi_file_project):
         remember(
             content="Important architectural decision: use microservices",
             memory_type="decision",
-            tags=["architecture"]
+            tags=["architecture"],
         )
 
         # Simulate session restart by resetting state
@@ -776,6 +794,7 @@ def test_incremental_reindex_preserves_memories(self, multi_file_project):
 # Edge Cases and Robustness Tests
 # =============================================================================
 
+
 class TestEdgeCases:
     """Test edge cases and error handling."""
 
@@ -793,8 +812,7 @@ def test_special_characters_in_query(self, indexed_project):
     def test_unicode_in_memory(self, indexed_project):
         """Handle unicode content in memories."""
         result = remember(
-            content="User prefers emoji: 🎉 and unicode: café résumé",
-            memory_type="preference"
+            content="User prefers emoji: 🎉 and unicode: café résumé", memory_type="preference"
         )
         assert result["success"] is True
 
@@ -810,11 +828,7 @@ def test_very_long_memory_content(self, indexed_project):
     def test_many_tags(self, indexed_project):
         """Handle many tags on a memory."""
         tags = [f"tag{i}" for i in range(50)]
-        result = remember(
-            content="Memory with many tags",
-            memory_type="note",
-            tags=tags
-        )
+        result = remember(content="Memory with many tags", memory_type="note", tags=tags)
         assert result["success"] is True
 
     def test_forget_nonexistent_id(self, indexed_project):
@@ -827,9 +841,7 @@ def test_recall_with_no_memories(self, indexed_project):
         # Clear any memories first
         state = get_state()
         if state.memory_retriever:
-            state.memory_retriever.collection.delete(
-                where={"project": str(state.codebase_path)}
-            )
+            state.memory_retriever.collection.delete(where={"project": str(state.codebase_path)})
 
         result = recall(query="something that doesn't exist")
         assert result["success"] is True
@@ -846,9 +858,7 @@ def test_concurrent_memory_operations(self, indexed_project):
         def store_memory(i):
             try:
                 result = remember(
-                    content=f"Concurrent memory {i}",
-                    memory_type="note",
-                    tags=[f"concurrent-{i}"]
+                    content=f"Concurrent memory {i}", memory_type="note", tags=[f"concurrent-{i}"]
                 )
                 results.append(result)
             except Exception as e:
@@ -869,17 +879,13 @@ def store_memory(i):
 # Memory Type Tests
 # =============================================================================
 
+
 class TestMemoryTypes:
     """Test all memory types work correctly."""
 
-    @pytest.mark.parametrize("memory_type", [
-        "conversation",
-        "status",
-        "decision",
-        "preference",
-        "doc",
-        "note"
-    ])
+    @pytest.mark.parametrize(
+        "memory_type", ["conversation", "status", "decision", "preference", "doc", "note"]
+    )
     def test_all_memory_types(self, indexed_project, memory_type):
         """Test each memory type can be stored and recalled."""
         content = f"Test content for {memory_type}"
@@ -901,23 +907,14 @@ def test_all_memory_types(self, indexed_project, memory_type):
 # TTL Tests
 # =============================================================================
 
+
 class TestMemoryTTL:
     """Test memory TTL functionality."""
 
-    @pytest.mark.parametrize("ttl", [
-        "session",
-        "day",
-        "week",
-        "month",
-        "permanent"
-    ])
+    @pytest.mark.parametrize("ttl", ["session", "day", "week", "month", "permanent"])
     def test_all_ttl_values(self, indexed_project, ttl):
         """Test each TTL value is accepted."""
-        result = remember(
-            content=f"Memory with {ttl} TTL",
-            memory_type="note",
-            ttl=ttl
-        )
+        result = remember(content=f"Memory with {ttl} TTL", memory_type="note", ttl=ttl)
         assert result["success"] is True
 
     def test_invalid_ttl_rejected(self, indexed_project):
@@ -934,7 +931,7 @@ def test_invalid_ttl_rejected(self, indexed_project):
                 content="test",
                 memory_type=MemoryType.NOTE,
                 project="test",
-                ttl="invalid_ttl"
+                ttl="invalid_ttl",
             )
 
 
@@ -942,6 +939,7 @@ def test_invalid_ttl_rejected(self, indexed_project):
 # Filter Tests
 # =============================================================================
 
+
 class TestRecallFilters:
     """Test recall filtering capabilities."""
 
@@ -963,7 +961,9 @@ def test_filter_by_tags(self, indexed_project):
         # Filter by auth tag
         results = recall(query="implementation", tags=["auth"])
         # Should find the auth-tagged memory
-        assert any("auth" in m.get("tags", []) for m in results["memories"]) or results["count"] >= 0
+        assert (
+            any("auth" in m.get("tags", []) for m in results["memories"]) or results["count"] >= 0
+        )
 
     def test_filter_by_time_range(self, indexed_project):
         """Filter recall by time range."""
@@ -978,6 +978,7 @@ def test_filter_by_time_range(self, indexed_project):
 # Forget Tests
 # =============================================================================
 
+
 class TestForgetOperations:
     """Test forget tool functionality."""
 
diff --git a/tests/mcp/test_protocol_simulation.py b/tests/mcp/test_protocol_simulation.py
index 5065c27..0b957fd 100644
--- a/tests/mcp/test_protocol_simulation.py
+++ b/tests/mcp/test_protocol_simulation.py
@@ -12,6 +12,7 @@
 This test file spawns the actual MCP server as a subprocess and
 communicates with it via stdin/stdout - the EXACT same way Claude does.
 """
+
 import pytest
 import json
 import subprocess
@@ -38,7 +39,7 @@ def start(self):
             stdout=subprocess.PIPE,
             stderr=subprocess.PIPE,
             text=True,
-            bufsize=1  # Line buffered
+            bufsize=1,  # Line buffered
         )
         # Give server time to initialize
         time.sleep(0.5)
@@ -59,7 +60,7 @@ def _send_request(self, method: str, params: dict = None) -> dict:
             "jsonrpc": "2.0",
             "id": self.request_id,
             "method": method,
-            "params": params or {}
+            "params": params or {},
         }
 
         # Write to stdin (this is what Claude does!)
@@ -76,11 +77,14 @@ def _send_request(self, method: str, params: dict = None) -> dict:
 
     def initialize(self) -> dict:
         """MCP initialization handshake."""
-        return self._send_request("initialize", {
-            "protocolVersion": "2024-11-05",
-            "capabilities": {},
-            "clientInfo": {"name": "pytest-simulator", "version": "1.0.0"}
-        })
+        return self._send_request(
+            "initialize",
+            {
+                "protocolVersion": "2024-11-05",
+                "capabilities": {},
+                "clientInfo": {"name": "pytest-simulator", "version": "1.0.0"},
+            },
+        )
 
     def list_tools(self) -> list:
         """Discover available tools (tools/list)."""
@@ -89,10 +93,7 @@ def list_tools(self) -> list:
 
     def call_tool(self, name: str, arguments: dict = None) -> dict:
         """Call an MCP tool (tools/call)."""
-        return self._send_request("tools/call", {
-            "name": name,
-            "arguments": arguments or {}
-        })
+        return self._send_request("tools/call", {"name": name, "arguments": arguments or {}})
 
 
 class TestStdioProtocol:
@@ -156,17 +157,14 @@ def test_full_agent_workflow(self, mcp_client, temp_project):
         assert len(tools) >= 4
 
         # Step 3: Learn codebase
-        learn_response = mcp_client.call_tool("learn", {
-            "path": str(temp_project)
-        })
+        learn_response = mcp_client.call_tool("learn", {"path": str(temp_project)})
         learn_result = json.loads(learn_response["result"]["content"][0]["text"])
         assert learn_result["success"] is True
 
         # Step 4: Search for code
-        search_response = mcp_client.call_tool("get_sources", {
-            "question": "calculator add",
-            "n_results": 5
-        })
+        search_response = mcp_client.call_tool(
+            "get_sources", {"question": "calculator add", "n_results": 5}
+        )
         search_result = json.loads(search_response["result"]["content"][0]["text"])
         assert "sources" in search_result
 
@@ -193,8 +191,7 @@ def test_multiple_rapid_requests(self, mcp_client, temp_project):
 
         # Rapid fire search requests
         for i in range(5):
-            response = mcp_client.call_tool("get_sources", {
-                "question": f"query {i}",
-                "n_results": 2
-            })
+            response = mcp_client.call_tool(
+                "get_sources", {"question": f"query {i}", "n_results": 2}
+            )
             assert "result" in response
diff --git a/tests/mcp/test_state_management.py b/tests/mcp/test_state_management.py
index 18afaaa..7ef1519 100644
--- a/tests/mcp/test_state_management.py
+++ b/tests/mcp/test_state_management.py
@@ -1,10 +1,9 @@
 """Test state management for isolation between tests and calls."""
+
 import asyncio
 import pytest
 from pathlib import Path
-from codegrok_mcp.mcp.state import (
-    get_state, reset_state, MCPSessionState
-)
+from codegrok_mcp.mcp.state import get_state, reset_state, MCPSessionState
 from codegrok_mcp.mcp.server import learn as learn_tool
 
 # Access underlying function from FastMCP FunctionTool wrapper
diff --git a/tests/mcp/test_tools_direct.py b/tests/mcp/test_tools_direct.py
index 15f9f2d..76d015a 100644
--- a/tests/mcp/test_tools_direct.py
+++ b/tests/mcp/test_tools_direct.py
@@ -5,6 +5,7 @@
 Note: FastMCP's @mcp.tool decorator wraps functions into FunctionTool objects.
 We access the underlying function via the .fn attribute.
 """
+
 import asyncio
 import pytest
 from pathlib import Path
@@ -13,7 +14,7 @@
     learn as learn_tool,
     get_sources as get_sources_tool,
     get_stats as get_stats_tool,
-    list_supported_languages as list_supported_languages_tool
+    list_supported_languages as list_supported_languages_tool,
 )
 from codegrok_mcp.mcp.state import get_state, reset_state
 
@@ -71,10 +72,7 @@ def test_learn_creates_codegrok_directory(self, temp_project):
         assert (codegrok_dir / "metadata.json").exists()
 
     def test_learn_with_custom_extensions(self, temp_project):
-        result = learn(
-            path=str(temp_project),
-            file_extensions=[".py"]
-        )
+        result = learn(path=str(temp_project), file_extensions=[".py"])
 
         assert result["success"] is True
         assert result["stats"]["total_files"] == 2
@@ -97,7 +95,6 @@ def test_learn_path_is_file(self, tmp_path):
             learn(path=str(file))
 
 
-
 class TestGetStatsTool:
     """Test the get_stats tool."""
 
@@ -148,7 +145,6 @@ def test_get_sources_no_codebase_loaded(self):
             get_sources(question="test")
 
 
-
 class TestLoadOnlyMode:
     """Test the learn tool with mode='load_only' - loads existing index."""
 
@@ -181,7 +177,7 @@ def test_incremental_reindex_after_file_change(self, temp_project):
         learn(path=str(temp_project))
 
         # Modify a file
-        (temp_project / "main.py").write_text('def new_function(): pass')
+        (temp_project / "main.py").write_text("def new_function(): pass")
 
         # Learn again with auto mode (should do incremental)
         result = learn(path=str(temp_project), mode="auto")
diff --git a/tests/unit/test_discover_files.py b/tests/unit/test_discover_files.py
new file mode 100644
index 0000000..e1db3ec
--- /dev/null
+++ b/tests/unit/test_discover_files.py
@@ -0,0 +1,126 @@
+"""
+Unit tests for discover_files() — file discovery with .gitignore support.
+"""
+
+import pytest
+from pathlib import Path
+from codegrok_mcp.indexing.source_retriever import discover_files
+
+
+def _create_file(path: Path, content: str = "# placeholder"):
+    """Helper to create a file with parent directories."""
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(content)
+
+
+class TestDiscoverFilesBasic:
+    """Test basic file discovery functionality."""
+
+    def test_discover_files_basic(self, tmp_path):
+        """Finds .py files in a simple directory."""
+        _create_file(tmp_path / "main.py", "def main(): pass")
+        _create_file(tmp_path / "utils.py", "x = 1")
+        _create_file(tmp_path / "readme.txt", "not code")
+
+        files = discover_files(tmp_path, extensions={".py"})
+
+        py_names = sorted(f.name for f in files)
+        assert py_names == ["main.py", "utils.py"]
+
+    def test_discover_files_skip_dirs(self, tmp_path):
+        """Skips node_modules, __pycache__, .git even when nested."""
+        _create_file(tmp_path / "app.py", "x = 1")
+        _create_file(tmp_path / "node_modules" / "pkg" / "index.py", "x = 2")
+        _create_file(tmp_path / "__pycache__" / "app.cpython-311.pyc", "x = 3")
+        _create_file(tmp_path / ".git" / "config.py", "x = 4")
+        _create_file(tmp_path / "src" / "node_modules" / "deep" / "mod.py", "x = 5")
+
+        files = discover_files(tmp_path, extensions={".py", ".pyc"})
+
+        assert len(files) == 1
+        assert files[0].name == "app.py"
+
+    def test_discover_files_no_gitignore(self, tmp_path):
+        """Works normally when no .gitignore exists."""
+        _create_file(tmp_path / "a.py", "x = 1")
+        _create_file(tmp_path / "sub" / "b.py", "x = 2")
+
+        files = discover_files(tmp_path, extensions={".py"})
+
+        assert len(files) == 2
+
+
+class TestDiscoverFilesGitignore:
+    """Test .gitignore support."""
+
+    def test_discover_files_gitignore(self, tmp_path):
+        """Respects root .gitignore patterns."""
+        (tmp_path / ".gitignore").write_text("*.log\nbuild/\n")
+
+        _create_file(tmp_path / "app.py", "x = 1")
+        _create_file(tmp_path / "debug.log", "log data")
+        _create_file(tmp_path / "build" / "output.py", "x = 2")
+        _create_file(tmp_path / "src" / "core.py", "x = 3")
+
+        files = discover_files(tmp_path, extensions={".py", ".log"})
+
+        names = sorted(f.name for f in files)
+        assert names == ["app.py", "core.py"]
+
+    def test_discover_files_nested_gitignore(self, tmp_path):
+        """Handles .gitignore in subdirectories."""
+        (tmp_path / ".gitignore").write_text("*.log\n")
+        (tmp_path / "vendor").mkdir()
+        (tmp_path / "vendor" / ".gitignore").write_text("*.py\n")
+
+        _create_file(tmp_path / "app.py", "x = 1")
+        _create_file(tmp_path / "vendor" / "lib.py", "x = 2")
+        _create_file(tmp_path / "vendor" / "data.txt", "data")
+
+        files = discover_files(tmp_path, extensions={".py", ".txt"})
+
+        names = sorted(f.name for f in files)
+        assert "app.py" in names
+        assert "lib.py" not in names
+
+    def test_discover_files_respect_gitignore_false(self, tmp_path):
+        """Opt-out disables gitignore filtering."""
+        (tmp_path / ".gitignore").write_text("*.py\n")
+
+        _create_file(tmp_path / "app.py", "x = 1")
+        _create_file(tmp_path / "lib.py", "x = 2")
+
+        files = discover_files(tmp_path, extensions={".py"}, respect_gitignore=False)
+
+        assert len(files) == 2
+
+
+class TestDiscoverFilesLimits:
+    """Test safety limits."""
+
+    def test_discover_files_max_files_limit(self, tmp_path):
+        """Stops at max_files and returns partial results."""
+        for i in range(20):
+            _create_file(tmp_path / f"file_{i}.py", f"x = {i}")
+
+        files = discover_files(tmp_path, extensions={".py"}, max_files=10)
+
+        assert len(files) == 10
+
+    def test_discover_files_progress_callback(self, tmp_path):
+        """Progress callback is called during discovery."""
+        # Create enough files to trigger callback (every 1000)
+        # We'll use a smaller set and verify callback mechanism works
+        for i in range(5):
+            _create_file(tmp_path / f"file_{i}.py", f"x = {i}")
+
+        events = []
+
+        def callback(event_type, data):
+            events.append((event_type, data))
+
+        files = discover_files(tmp_path, extensions={".py"}, progress_callback=callback)
+
+        # With only 5 files, callback won't fire (fires every 1000)
+        assert len(files) == 5
+        assert len(events) == 0  # Below threshold
diff --git a/tests/unit/test_embedding_service.py b/tests/unit/test_embedding_service.py
index 979c977..bc83035 100644
--- a/tests/unit/test_embedding_service.py
+++ b/tests/unit/test_embedding_service.py
@@ -1,7 +1,11 @@
-
 import pytest
 from unittest.mock import MagicMock, patch
-from codegrok_mcp.indexing.embedding_service import EmbeddingService, get_embedding_service, reset_embedding_service
+from codegrok_mcp.indexing.embedding_service import (
+    EmbeddingService,
+    get_embedding_service,
+    reset_embedding_service,
+)
+
 
 class TestEmbeddingService:
     def teardown_method(self):
@@ -12,40 +16,41 @@ def test_embed_empty_batch(self):
         # Mock _model to avoid downloading/loading
         service._model = MagicMock()
         service._model_loaded = True
-        
+
         result = service.embed_batch([])
         assert result == []
 
     def test_cache_statistics(self):
         service = EmbeddingService(model_name="all-MiniLM-L6-v2")
-        
+
         # Mock dependencies
-        with patch('codegrok_mcp.indexing.embedding_service._sentence_transformers') as mock_st:
+        with patch("codegrok_mcp.indexing.embedding_service._sentence_transformers") as mock_st:
             mock_model = MagicMock()
             # mock encode return value
             import numpy as np
+
             mock_model.encode.return_value = np.array([[0.1, 0.2]])
             mock_model.get_sentence_embedding_dimension.return_value = 2
-            
+
             # Manually load mock model
             service._model = mock_model
-            service.config['dimensions'] = 2
+            service.config["dimensions"] = 2
             service._model_loaded = True
-            
+
             service.embed("test")
             service.embed("test")  # Cache hit
-            
+
             stats = service.get_cache_stats()
-            
-            assert service.stats['cache_hits'] == 1
-            assert service.stats['cache_misses'] == 1
-            assert stats['hits'] == 1
+
+            assert service.stats["cache_hits"] == 1
+            assert service.stats["cache_misses"] == 1
+            assert stats["hits"] == 1
 
     def test_unload_model(self):
         service = EmbeddingService(model_name="all-MiniLM-L6-v2")
         service._model = MagicMock()
         service._model_loaded = True
-        
+
         service.unload()
         assert service._model is None
         assert service._model_loaded is False
@@ -53,7 +58,7 @@ def test_unload_model(self):
     def test_singleton_pattern(self):
         reset_embedding_service()
         # Mock dependencies to avoid import errors if not installed
-        with patch('codegrok_mcp.indexing.embedding_service._import_dependencies'):
+        with patch("codegrok_mcp.indexing.embedding_service._import_dependencies"):
             s1 = get_embedding_service()
             s2 = get_embedding_service()
             assert s1 is s2
diff --git a/tests/unit/test_extras.py b/tests/unit/test_extras.py
index 7b1c2a2..7b9c48f 100644
--- a/tests/unit/test_extras.py
+++ b/tests/unit/test_extras.py
@@ -1,24 +1,25 @@
-
 import pytest
 from unittest.mock import MagicMock, patch
 from codegrok_mcp.indexing.embedding_service import embed
 from codegrok_mcp.indexing.source_retriever import SourceRetriever
 from codegrok_mcp.parsers.treesitter_parser import TreeSitterParser
 
+
 class TestEmbeddingHelpers:
-    @patch('codegrok_mcp.indexing.embedding_service.get_embedding_service')
+    @patch("codegrok_mcp.indexing.embedding_service.get_embedding_service")
     def test_embed_convenience(self, mock_get_service):
         mock_service = MagicMock()
         mock_get_service.return_value = mock_service
-        
+
         # Test string
         embed("text")
         mock_service.embed.assert_called_with("text", is_query=False)
-        
+
         # Test list
         embed(["t1", "t2"])
         mock_service.embed_batch.assert_called_with(["t1", "t2"], is_query=False)
 
+
 class TestSourceRetrieverExtras:
     def test_parallel_init(self, tmp_path):
         # Just check it initializes without error
@@ -38,20 +39,21 @@ def test_logging_disabled(self, tmp_path, capsys):
         captured = capsys.readouterr()
         assert "test message" not in captured.out
 
+
 class TestParserExtras:
     def test_clean_docstring(self):
         parser = TreeSitterParser()
-        
+
         # Triple quotes
         assert parser._clean_docstring('"""doc"""') == "doc"
         assert parser._clean_docstring("'''doc'''") == "doc"
-        
+
         # Single quotes
         assert parser._clean_docstring('"doc"') == "doc"
         assert parser._clean_docstring("'doc'") == "doc"
-        
+
         # Multiline
         assert parser._clean_docstring('"""\n  Line 1\n  Line 2\n"""') == "Line 1"
-        
+
         # Raw
-        assert parser._clean_docstring('Simple doc') == "Simple doc"
+        assert parser._clean_docstring("Simple doc") == "Simple doc"
diff --git a/tests/unit/test_init.py b/tests/unit/test_init.py
index a2348ac..4e5a9d4 100644
--- a/tests/unit/test_init.py
+++ b/tests/unit/test_init.py
@@ -1,40 +1,57 @@
-
 import pytest
 
+
 def test_lazy_import_treesitter_parser():
     import codegrok_mcp
-    assert hasattr(codegrok_mcp, 'TreeSitterParser')
+
+    assert hasattr(codegrok_mcp, "TreeSitterParser")
+
 
 def test_lazy_import_symbol():
     import codegrok_mcp
-    assert hasattr(codegrok_mcp, 'Symbol')
+
+    assert hasattr(codegrok_mcp, "Symbol")
+
 
 def test_lazy_import_symbol_type():
     import codegrok_mcp
-    assert hasattr(codegrok_mcp, 'SymbolType')
+
+    assert hasattr(codegrok_mcp, "SymbolType")
+
 
 def test_lazy_import_parsed_file():
     import codegrok_mcp
-    assert hasattr(codegrok_mcp, 'ParsedFile')
+
+    assert hasattr(codegrok_mcp, "ParsedFile")
+
 
 def test_lazy_import_codebase_index():
     import codegrok_mcp
-    assert hasattr(codegrok_mcp, 'CodebaseIndex')
+
+    assert hasattr(codegrok_mcp, "CodebaseIndex")
+
 
 def test_lazy_import_iparser():
     import codegrok_mcp
-    assert hasattr(codegrok_mcp, 'IParser')
+
+    assert hasattr(codegrok_mcp, "IParser")
+
 
 def test_lazy_import_parser_factory():
     import codegrok_mcp
-    assert hasattr(codegrok_mcp, 'ThreadLocalParserFactory')
+
+    assert hasattr(codegrok_mcp, "ThreadLocalParserFactory")
+
 
 def test_lazy_import_invalid_name():
     import codegrok_mcp
+
     with pytest.raises(AttributeError):
         _ = codegrok_mcp.NonExistentClass
 
+
 def test_all_exported_names_are_accessible():
     import codegrok_mcp
+
     for name in codegrok_mcp.__all__:
         assert hasattr(codegrok_mcp, name)
diff --git a/tests/unit/test_memory_retriever.py b/tests/unit/test_memory_retriever.py
index 40c2907..2f91718 100644
--- a/tests/unit/test_memory_retriever.py
+++ b/tests/unit/test_memory_retriever.py
@@ -26,7 +26,7 @@ def retriever(self, tmp_path, mock_embedding_service):
             project_path=str(tmp_path),
             persist_path=str(tmp_path / "chroma"),
             embedding_service=mock_embedding_service,
-            verbose=False
+            verbose=False,
         )
 
     @pytest.fixture
@@ -36,7 +36,7 @@ def verbose_retriever(self, tmp_path, mock_embedding_service):
             project_path=str(tmp_path),
             persist_path=str(tmp_path / "chroma_verbose"),
             embedding_service=mock_embedding_service,
-            verbose=True
+            verbose=True,
         )
 
     @pytest.fixture
@@ -46,7 +46,7 @@ def ephemeral_retriever(self, tmp_path, mock_embedding_service):
             project_path=str(tmp_path),
             persist_path=None,  # Ephemeral mode
             embedding_service=mock_embedding_service,
-            verbose=False
+            verbose=False,
         )
 
     # ==========================================================================
@@ -56,9 +56,7 @@ def ephemeral_retriever(self, tmp_path, mock_embedding_service):
     def test_remember_creates_memory(self, retriever):
         """Test that remember() stores a memory."""
         memory = retriever.remember(
-            content="Test memory content",
-            memory_type="status",
-            tags=["test"]
+            content="Test memory content", memory_type="status", tags=["test"]
         )
 
         assert memory.id is not None
@@ -69,29 +67,21 @@ def test_remember_creates_memory(self, retriever):
     def test_recall_returns_relevant_memories(self, retriever):
         """Test that recall() returns stored memories."""
         retriever.remember(
-            content="Authentication uses JWT tokens",
-            memory_type="decision",
-            tags=["auth"]
+            content="Authentication uses JWT tokens", memory_type="decision", tags=["auth"]
         )
 
-        results = retriever.recall(
-            query="How does auth work?",
-            n_results=5
-        )
+        results = retriever.recall(query="How does auth work?", n_results=5)
 
         assert len(results) > 0
-        assert "JWT" in results[0]['content']
+        assert "JWT" in results[0]["content"]
 
     def test_forget_removes_memory(self, retriever):
         """Test that forget() removes memories."""
-        memory = retriever.remember(
-            content="Temporary note",
-            memory_type="note"
-        )
+        memory = retriever.remember(content="Temporary note", memory_type="note")
 
         result = retriever.forget(memory_id=memory.id)
 
-        assert result['deleted'] == 1
+        assert result["deleted"] == 1
 
     def test_get_stats(self, retriever):
         """Test that get_stats returns correct counts."""
@@ -101,9 +91,9 @@ def test_get_stats(self, retriever):
 
         stats = retriever.get_stats()
 
-        assert stats['total_memories'] == 3
-        assert stats['by_type'].get('note', 0) == 2
-        assert stats['by_type'].get('status', 0) == 1
+        assert stats["total_memories"] == 3
+        assert stats["by_type"].get("note", 0) == 2
+        assert stats["by_type"].get("status", 0) == 1
 
     # ==========================================================================
     # All Memory Types Tests
@@ -126,7 +116,7 @@ def test_remember_all_memory_types(self, retriever):
             assert memory.content == content
 
         stats = retriever.get_stats()
-        assert stats['total_memories'] == 6
+        assert stats["total_memories"] == 6
 
     def test_recall_by_each_memory_type(self, retriever):
         """Test recalling memories filtered by each type."""
@@ -140,9 +130,9 @@ def test_recall_by_each_memory_type(self, retriever):
         status_results = retriever.recall(query="blocked", memory_type="status")
         decision_results = retriever.recall(query="frontend", memory_type="decision")
 
-        assert all(r['memory_type'] == 'conversation' for r in conv_results)
-        assert all(r['memory_type'] == 'status' for r in status_results)
-        assert all(r['memory_type'] == 'decision' for r in decision_results)
+        assert all(r["memory_type"] == "conversation" for r in conv_results)
+        assert all(r["memory_type"] == "status" for r in status_results)
+        assert all(r["memory_type"] == "decision" for r in decision_results)
 
     # ==========================================================================
     # Filtering Tests
@@ -153,24 +143,18 @@ def test_memory_type_filtering(self, retriever):
         retriever.remember(content="Status update", memory_type="status")
         retriever.remember(content="A decision", memory_type="decision")
 
-        results = retriever.recall(
-            query="update",
-            memory_type="status"
-        )
+        results = retriever.recall(query="update", memory_type="status")
 
-        assert all(r['memory_type'] == 'status' for r in results)
+        assert all(r["memory_type"] == "status" for r in results)
 
     def test_tag_filtering(self, retriever):
         """Test filtering by tags."""
         retriever.remember(content="Auth note", memory_type="note", tags=["auth"])
         retriever.remember(content="DB note", memory_type="note", tags=["database"])
 
-        results = retriever.recall(
-            query="note",
-            tags=["auth"]
-        )
+        results = retriever.recall(query="note", tags=["auth"])
 
-        assert all("auth" in r['tags'] for r in results)
+        assert all("auth" in r["tags"] for r in results)
 
     def test_multiple_tags_filtering(self, retriever):
         """Test filtering with multiple tags (OR logic)."""
@@ -184,7 +168,7 @@ def test_multiple_tags_filtering(self, retriever):
         # Should find auth and database, not frontend
         tags_found = set()
         for r in results:
-            tags_found.update(r['tags'])
+            tags_found.update(r["tags"])
 
         assert "auth" in tags_found or "database" in tags_found
 
@@ -244,9 +228,9 @@ def test_forget_by_type(self, retriever):
 
         result = retriever.forget(memory_type="note")
 
-        assert result['deleted'] == 2
+        assert result["deleted"] == 2
         stats = retriever.get_stats()
-        assert stats['total_memories'] == 1
+        assert stats["total_memories"] == 1
 
     def test_forget_by_tags(self, retriever):
         """Test deleting memories by tags."""
@@ -256,9 +240,9 @@ def test_forget_by_tags(self, retriever):
 
         result = retriever.forget(tags=["deprecated"])
 
-        assert result['deleted'] == 2
+        assert result["deleted"] == 2
         stats = retriever.get_stats()
-        assert stats['total_memories'] == 1
+        assert stats["total_memories"] == 1
 
     def test_forget_by_older_than(self, retriever):
         """Test deleting memories older than specified duration."""
@@ -270,7 +254,7 @@ def test_forget_by_older_than(self, retriever):
 
         # Recent memory should still exist
         stats = retriever.get_stats()
-        assert stats['total_memories'] == 1
+        assert stats["total_memories"] == 1
 
     def test_forget_by_older_than_1d(self, retriever):
         """Test forget with 1d older_than."""
@@ -281,21 +265,21 @@ def test_forget_by_older_than_1d(self, retriever):
 
         # Memory was just created, shouldn't be deleted
         stats = retriever.get_stats()
-        assert stats['total_memories'] == 1
+        assert stats["total_memories"] == 1
 
     def test_forget_by_older_than_7d(self, retriever):
         """Test forget with 7d older_than."""
         retriever.remember(content="Memory", memory_type="note")
         result = retriever.forget(older_than="7d")
         stats = retriever.get_stats()
-        assert stats['total_memories'] == 1
+        assert stats["total_memories"] == 1
 
     def test_forget_by_older_than_30d(self, retriever):
         """Test forget with 30d older_than."""
         retriever.remember(content="Memory", memory_type="note")
         result = retriever.forget(older_than="30d")
         stats = retriever.get_stats()
-        assert stats['total_memories'] == 1
+        assert stats["total_memories"] == 1
 
     def test_forget_combined_type_and_tags(self, retriever):
         """Test forget with both memory_type and tags filters."""
@@ -308,7 +292,7 @@ def test_forget_combined_type_and_tags(self, retriever):
 
         # Only the old note should be deleted
         stats = retriever.get_stats()
-        assert stats['total_memories'] == 2
+        assert stats["total_memories"] == 2
 
     def test_forget_nonexistent_id(self, retriever):
         """Test forgetting a non-existent memory ID."""
@@ -325,9 +309,7 @@ def test_remember_with_all_ttl_options(self, retriever):
 
         for ttl in ttl_options:
             memory = retriever.remember(
-                content=f"Memory with TTL {ttl}",
-                memory_type="note",
-                ttl=ttl
+                content=f"Memory with TTL {ttl}", memory_type="note", ttl=ttl
             )
             assert memory.ttl == ttl
 
@@ -339,7 +321,7 @@ def test_cleanup_expired_with_permanent(self, retriever):
 
         # Permanent memory should still exist
         stats = retriever.get_stats()
-        assert stats['total_memories'] == 1
+        assert stats["total_memories"] == 1
 
     def test_cleanup_expired_function(self, retriever):
         """Test cleanup_expired returns proper structure."""
@@ -365,8 +347,7 @@ def test_recall_with_no_matches(self, retriever):
         retriever.remember(content="Specific technical content", memory_type="note")
 
         results = retriever.recall(
-            query="completely unrelated xyz123",
-            min_relevance=0.9  # High threshold
+            query="completely unrelated xyz123", min_relevance=0.9  # High threshold
         )
         # May or may not return results depending on embeddings
 
@@ -379,10 +360,7 @@ def test_verbose_mode_logging(self, verbose_retriever, capsys):
 
     def test_ephemeral_retriever(self, ephemeral_retriever):
         """Test retriever without persistence works."""
-        memory = ephemeral_retriever.remember(
-            content="Ephemeral memory",
-            memory_type="note"
-        )
+        memory = ephemeral_retriever.remember(content="Ephemeral memory", memory_type="note")
         assert memory.id is not None
 
         results = ephemeral_retriever.recall(query="ephemeral")
@@ -393,7 +371,7 @@ def test_remember_with_metadata(self, retriever):
         memory = retriever.remember(
             content="Memory with metadata",
             memory_type="note",
-            metadata={"custom_field": "custom_value", "priority": 1}
+            metadata={"custom_field": "custom_value", "priority": 1},
         )
 
         assert memory.metadata["custom_field"] == "custom_value"
@@ -405,9 +383,7 @@ def test_remember_with_source(self, retriever):
 
         for source in sources:
             memory = retriever.remember(
-                content=f"Memory from {source}",
-                memory_type="note",
-                source=source
+                content=f"Memory from {source}", memory_type="note", source=source
             )
             assert memory.source == source
 
@@ -417,9 +393,9 @@ def test_get_stats_with_project_info(self, retriever):
 
         stats = retriever.get_stats()
 
-        assert 'project' in stats
-        assert 'total_memories' in stats
-        assert 'by_type' in stats
+        assert "project" in stats
+        assert "total_memories" in stats
+        assert "by_type" in stats
 
     def test_stats_persistence(self, tmp_path, mock_embedding_service):
         """Test that stats are persisted and loaded."""
@@ -430,13 +406,13 @@ def test_stats_persistence(self, tmp_path, mock_embedding_service):
             project_path=str(tmp_path),
             persist_path=persist_path,
             embedding_service=mock_embedding_service,
-            verbose=False
+            verbose=False,
         )
         retriever1.remember(content="Persisted memory", memory_type="note")
 
         # Stats should be saved
         stats1 = retriever1.get_stats()
-        assert stats1['total_memories'] == 1
+        assert stats1["total_memories"] == 1
 
     def test_recall_updates_accessed_at(self, retriever):
         """Test that recalling memories updates accessed_at timestamp."""
@@ -459,32 +435,17 @@ class TestMemoryModel:
     def test_memory_validation_empty_content(self):
         """Test memory validation for empty content."""
         with pytest.raises(ValueError, match="content cannot be empty"):
-            Memory(
-                id="test",
-                content="",
-                memory_type=MemoryType.NOTE,
-                project="/test"
-            )
+            Memory(id="test", content="", memory_type=MemoryType.NOTE, project="/test")
 
     def test_memory_validation_empty_id(self):
         """Test memory validation for empty id."""
         with pytest.raises(ValueError, match="id cannot be empty"):
-            Memory(
-                id="",
-                content="Test",
-                memory_type=MemoryType.NOTE,
-                project="/test"
-            )
+            Memory(id="", content="Test", memory_type=MemoryType.NOTE, project="/test")
 
     def test_memory_validation_empty_project(self):
         """Test memory validation for empty project."""
         with pytest.raises(ValueError, match="project cannot be empty"):
-            Memory(
-                id="test",
-                content="Test",
-                memory_type=MemoryType.NOTE,
-                project=""
-            )
+            Memory(id="test", content="Test", memory_type=MemoryType.NOTE, project="")
 
     def test_memory_validation_invalid_memory_type(self):
         """Test memory validation for invalid memory_type."""
@@ -493,7 +454,7 @@ def test_memory_validation_invalid_memory_type(self):
                 id="test",
                 content="Test",
                 memory_type="invalid",  # Should be MemoryType enum
-                project="/test"
+                project="/test",
             )
 
     def test_memory_validation_invalid_ttl(self):
@@ -504,7 +465,7 @@ def test_memory_validation_invalid_ttl(self):
                 content="Test",
                 memory_type=MemoryType.NOTE,
                 project="/test",
-                ttl="invalid"
+                ttl="invalid",
             )
 
     def test_memory_serialization(self):
@@ -514,7 +475,7 @@ def test_memory_serialization(self):
             content="Test content",
             memory_type=MemoryType.DECISION,
             project="/test",
-            tags=["tag1", "tag2"]
+            tags=["tag1", "tag2"],
         )
 
         data = memory.to_dict()
@@ -535,17 +496,17 @@ def test_memory_serialization_all_fields(self):
             tags=["a", "b", "c"],
             ttl="week",
             source="agent",
-            metadata={"key": "value"}
+            metadata={"key": "value"},
         )
 
         data = memory.to_dict()
 
-        assert data['id'] == "full-test"
-        assert data['memory_type'] == "status"
-        assert data['tags'] == ["a", "b", "c"]
-        assert data['ttl'] == "week"
-        assert data['source'] == "agent"
-        assert data['metadata'] == {"key": "value"}
+        assert data["id"] == "full-test"
+        assert data["memory_type"] == "status"
+        assert data["tags"] == ["a", "b", "c"]
+        assert data["ttl"] == "week"
+        assert data["source"] == "agent"
+        assert data["metadata"] == {"key": "value"}
 
         # Round-trip
         restored = Memory.from_dict(data)
@@ -554,12 +515,7 @@ def test_memory_serialization_all_fields(self):
 
     def test_memory_touch(self):
         """Test that touch() updates accessed_at."""
-        memory = Memory(
-            id="test",
-            content="Test",
-            memory_type=MemoryType.NOTE,
-            project="/test"
-        )
+        memory = Memory(id="test", content="Test", memory_type=MemoryType.NOTE, project="/test")
 
         original_accessed = memory.accessed_at
         time.sleep(0.01)  # Small delay
@@ -570,10 +526,7 @@ def test_memory_touch(self):
     def test_memory_default_values(self):
         """Test that Memory has correct default values."""
         memory = Memory(
-            id="defaults",
-            content="Test defaults",
-            memory_type=MemoryType.NOTE,
-            project="/test"
+            id="defaults", content="Test defaults", memory_type=MemoryType.NOTE, project="/test"
         )
 
         assert memory.tags == []
@@ -653,7 +606,7 @@ def test_time_range_invalid_value(self, tmp_path, mock_embedding_service):
             project_path=str(tmp_path),
             persist_path=str(tmp_path / "chroma"),
             embedding_service=mock_embedding_service,
-            verbose=False
+            verbose=False,
         )
         retriever.remember(content="Test memory", memory_type="note")
 
@@ -672,13 +625,13 @@ def test_stats_load_from_existing_file(self, tmp_path, mock_embedding_service):
         metadata_path = tmp_path / "memory_metadata.json"
 
         existing_stats = {
-            'stats': {
-                'total_memories': 10,
-                'by_type': {'note': 5, 'status': 5},
-                'last_cleanup': '2024-01-01T00:00:00'
+            "stats": {
+                "total_memories": 10,
+                "by_type": {"note": 5, "status": 5},
+                "last_cleanup": "2024-01-01T00:00:00",
             }
         }
-        with open(metadata_path, 'w') as f:
+        with open(metadata_path, "w") as f:
             json.dump(existing_stats, f)
 
         # Create retriever - should load existing stats
@@ -686,7 +639,7 @@ def test_stats_load_from_existing_file(self, tmp_path, mock_embedding_service):
             project_path=str(tmp_path),
             persist_path=str(persist_path),
             embedding_service=mock_embedding_service,
-            verbose=False
+            verbose=False,
         )
 
         # Stats should have been loaded (though actual count may differ)
@@ -698,7 +651,7 @@ def test_recall_empty_tags_in_stored_memory(self, tmp_path, mock_embedding_servi
             project_path=str(tmp_path),
             persist_path=str(tmp_path / "chroma"),
             embedding_service=mock_embedding_service,
-            verbose=False
+            verbose=False,
         )
 
         # Store memory without tags
@@ -714,13 +667,13 @@ def test_forget_invalid_older_than(self, tmp_path, mock_embedding_service):
             project_path=str(tmp_path),
             persist_path=str(tmp_path / "chroma"),
             embedding_service=mock_embedding_service,
-            verbose=False
+            verbose=False,
         )
         retriever.remember(content="Test", memory_type="note")
 
         # Invalid older_than should not crash
         result = retriever.forget(older_than="invalid")
-        assert 'deleted' in result
+        assert "deleted" in result
 
     def test_verbose_recall_logging(self, tmp_path, mock_embedding_service, capsys):
         """Test verbose logging during recall."""
@@ -728,7 +681,7 @@ def test_verbose_recall_logging(self, tmp_path, mock_embedding_service, capsys):
             project_path=str(tmp_path),
             persist_path=str(tmp_path / "chroma"),
             embedding_service=mock_embedding_service,
-            verbose=True
+            verbose=True,
         )
         retriever.remember(content="Verbose recall test", memory_type="note")
         retriever.recall(query="verbose")
@@ -742,7 +695,7 @@ def test_verbose_forget_logging(self, tmp_path, mock_embedding_service, capsys):
             project_path=str(tmp_path),
             persist_path=str(tmp_path / "chroma"),
             embedding_service=mock_embedding_service,
-            verbose=True
+            verbose=True,
         )
         memory = retriever.remember(content="Verbose forget test", memory_type="note")
         retriever.forget(memory_id=memory.id)
@@ -756,7 +709,7 @@ def test_verbose_cleanup_logging(self, tmp_path, mock_embedding_service, capsys)
             project_path=str(tmp_path),
             persist_path=str(tmp_path / "chroma"),
             embedding_service=mock_embedding_service,
-            verbose=True
+            verbose=True,
         )
         retriever.remember(content="Cleanup test", memory_type="note", ttl="session")
         retriever.cleanup_expired()
@@ -770,7 +723,7 @@ def test_cleanup_with_various_ttls(self, tmp_path, mock_embedding_service):
             project_path=str(tmp_path),
             persist_path=str(tmp_path / "chroma"),
             embedding_service=mock_embedding_service,
-            verbose=False
+            verbose=False,
         )
 
         # Add memories with different TTLs
@@ -784,7 +737,7 @@ def test_cleanup_with_various_ttls(self, tmp_path, mock_embedding_service):
 
         # Fresh memories shouldn't be cleaned up
         stats = retriever.get_stats()
-        assert stats['total_memories'] == 5
+        assert stats["total_memories"] == 5
 
     def test_get_stats_last_cleanup(self, tmp_path, mock_embedding_service):
         """Test that get_stats returns last_cleanup time."""
@@ -792,12 +745,12 @@ def test_get_stats_last_cleanup(self, tmp_path, mock_embedding_service):
             project_path=str(tmp_path),
             persist_path=str(tmp_path / "chroma"),
             embedding_service=mock_embedding_service,
-            verbose=False
+            verbose=False,
         )
 
         # Initially no cleanup
         stats = retriever.get_stats()
-        assert 'last_cleanup' in stats
+        assert "last_cleanup" in stats
 
         # After cleanup
         retriever.cleanup_expired()
diff --git a/tests/unit/test_models.py b/tests/unit/test_models.py
index d8c207a..c38201d 100644
--- a/tests/unit/test_models.py
+++ b/tests/unit/test_models.py
@@ -1,4 +1,5 @@
 """Unit tests for core data models."""
+
 import pytest
 from codegrok_mcp.core.models import Symbol, SymbolType, ParsedFile, CodebaseIndex
 
@@ -14,7 +15,7 @@ def test_create_function_symbol(self):
             line_start=1,
             line_end=3,
             language="python",
-            signature="def hello():"
+            signature="def hello():",
         )
 
         assert symbol.name == "hello"
@@ -29,7 +30,7 @@ def test_create_class_symbol(self):
             line_start=1,
             line_end=10,
             language="python",
-            signature="class MyClass:"
+            signature="class MyClass:",
         )
 
         assert symbol.name == "MyClass"
@@ -43,7 +44,7 @@ def test_line_count_property(self):
             line_start=5,
             line_end=10,
             language="python",
-            signature="def func():"
+            signature="def func():",
         )
 
         assert symbol.line_count == 6  # 5,6,7,8,9,10
@@ -56,7 +57,7 @@ def test_qualified_name_without_parent(self):
             line_start=1,
             line_end=3,
             language="python",
-            signature="def standalone():"
+            signature="def standalone():",
         )
 
         assert symbol.qualified_name == "standalone"
@@ -70,13 +71,11 @@ def test_qualified_name_with_parent(self):
             line_end=10,
             language="python",
             signature="def method(self):",
-            parent="MyClass"
+            parent="MyClass",
         )
 
         assert symbol.qualified_name == "MyClass.method"
 
-
-
     def test_symbol_empty_name_raises(self):
         with pytest.raises(ValueError, match="name cannot be empty"):
             Symbol(
@@ -86,7 +85,7 @@ def test_symbol_empty_name_raises(self):
                 line_start=1,
                 line_end=3,
                 language="python",
-                signature="def foo():"
+                signature="def foo():",
             )
 
     def test_symbol_invalid_lines_raises(self):
@@ -98,7 +97,7 @@ def test_symbol_invalid_lines_raises(self):
                 line_start=5,
                 line_end=4,  # Invalid
                 language="python",
-                signature="def foo():"
+                signature="def foo():",
             )
 
     def test_symbol_serialization_roundtrip(self):
@@ -114,12 +113,12 @@ def test_symbol_serialization_roundtrip(self):
             docstring="Doing things",
             imports=["os", "sys"],
             calls=["print"],
-            metadata={"complexity": 5}
+            metadata={"complexity": 5},
         )
-        
+
         data = symbol.to_dict()
         restored = Symbol.from_dict(data)
-        
+
         assert restored == symbol
         assert restored.parent == "AppClass"
         assert restored.metadata["complexity"] == 5
@@ -137,15 +136,11 @@ def test_create_successful_parsed_file(self):
                 line_start=1,
                 line_end=3,
                 language="python",
-                signature="def func():"
+                signature="def func():",
             )
         ]
 
-        parsed = ParsedFile(
-            filepath="/path/file.py",
-            language="python",
-            symbols=symbols
-        )
+        parsed = ParsedFile(filepath="/path/file.py", language="python", symbols=symbols)
 
         assert parsed.is_successful
         assert parsed.symbol_count == 1
@@ -153,10 +148,7 @@ def test_create_successful_parsed_file(self):
 
     def test_create_failed_parsed_file(self):
         parsed = ParsedFile(
-            filepath="/path/file.py",
-            language="python",
-            symbols=[],
-            error="Failed to parse"
+            filepath="/path/file.py", language="python", symbols=[], error="Failed to parse"
         )
 
         assert not parsed.is_successful
@@ -164,19 +156,36 @@ def test_create_failed_parsed_file(self):
 
     def test_get_symbols_by_type(self):
         symbols = [
-            Symbol(name="func1", type=SymbolType.FUNCTION, filepath="/path/file.py",
-                   line_start=1, line_end=3, language="python", signature="def func1():"),
-            Symbol(name="MyClass", type=SymbolType.CLASS, filepath="/path/file.py",
-                   line_start=5, line_end=10, language="python", signature="class MyClass:"),
-            Symbol(name="func2", type=SymbolType.FUNCTION, filepath="/path/file.py",
-                   line_start=12, line_end=15, language="python", signature="def func2():"),
+            Symbol(
+                name="func1",
+                type=SymbolType.FUNCTION,
+                filepath="/path/file.py",
+                line_start=1,
+                line_end=3,
+                language="python",
+                signature="def func1():",
+            ),
+            Symbol(
+                name="MyClass",
+                type=SymbolType.CLASS,
+                filepath="/path/file.py",
+                line_start=5,
+                line_end=10,
+                language="python",
+                signature="class MyClass:",
+            ),
+            Symbol(
+                name="func2",
+                type=SymbolType.FUNCTION,
+                filepath="/path/file.py",
+                line_start=12,
+                line_end=15,
+                language="python",
+                signature="def func2():",
+            ),
         ]
 
-        parsed = ParsedFile(
-            filepath="/path/file.py",
-            language="python",
-            symbols=symbols
-        )
+        parsed = ParsedFile(filepath="/path/file.py", language="python", symbols=symbols)
 
         functions = parsed.get_symbols_by_type(SymbolType.FUNCTION)
         classes = parsed.get_symbols_by_type(SymbolType.CLASS)
@@ -185,11 +194,7 @@ def test_get_symbols_by_type(self):
         assert len(classes) == 1
 
     def test_empty_file_is_successful(self):
-        parsed = ParsedFile(
-            filepath="/path/empty.py",
-            language="python",
-            symbols=[]
-        )
+        parsed = ParsedFile(filepath="/path/empty.py", language="python", symbols=[])
 
         assert parsed.is_successful
 
@@ -207,7 +212,7 @@ def test_create_index(self):
     def test_index_validation(self):
         with pytest.raises(ValueError, match="root_path cannot be empty"):
             CodebaseIndex(root_path="")
-            
+
         with pytest.raises(ValueError, match="total_files must be >= 0"):
             CodebaseIndex(root_path="/app", total_files=-1)
 
@@ -223,44 +228,62 @@ def test_serialization_roundtrip(self):
                     line_start=1,
                     line_end=5,
                     language="python",
-                    signature="def main():"
+                    signature="def main():",
                 )
-            ]
+            ],
         )
-        
+
         index = CodebaseIndex(
-            root_path="/app",
-            files={"/app/main.py": parsed_file},
-            total_files=1,
-            total_symbols=1
+            root_path="/app", files={"/app/main.py": parsed_file}, total_files=1, total_symbols=1
         )
-        
+
         data = index.to_dict()
         restored = CodebaseIndex.from_dict(data)
-        
+
         assert restored == index
         assert "/app/main.py" in restored.files
         assert restored.files["/app/main.py"].symbol_count == 1
 
     def test_get_symbols_by_name(self):
-        s1 = Symbol(name="target", type=SymbolType.FUNCTION, filepath="/a.py", 
-                   line_start=1, line_end=1, language="py", signature="def target()")
-        s2 = Symbol(name="other", type=SymbolType.FUNCTION, filepath="/b.py", 
-                   line_start=1, line_end=1, language="py", signature="def other()")
-        s3 = Symbol(name="target", type=SymbolType.VARIABLE, filepath="/c.py", 
-                   line_start=1, line_end=1, language="py", signature="target = 1")
-                   
+        s1 = Symbol(
+            name="target",
+            type=SymbolType.FUNCTION,
+            filepath="/a.py",
+            line_start=1,
+            line_end=1,
+            language="py",
+            signature="def target()",
+        )
+        s2 = Symbol(
+            name="other",
+            type=SymbolType.FUNCTION,
+            filepath="/b.py",
+            line_start=1,
+            line_end=1,
+            language="py",
+            signature="def other()",
+        )
+        s3 = Symbol(
+            name="target",
+            type=SymbolType.VARIABLE,
+            filepath="/c.py",
+            line_start=1,
+            line_end=1,
+            language="py",
+            signature="target = 1",
+        )
+
         index = CodebaseIndex(
             root_path="/app",
             files={
                 "/a.py": ParsedFile(filepath="/a.py", language="py", symbols=[s1]),
                 "/b.py": ParsedFile(filepath="/b.py", language="py", symbols=[s2]),
                 "/c.py": ParsedFile(filepath="/c.py", language="py", symbols=[s3]),
-            }
+            },
         )
-        
+
         assert index.get_symbols_by_name("NonExistent") == []
-        
+
         results = index.get_symbols_by_name("target")
         assert len(results) == 2
         assert s1 in results
@@ -271,16 +294,32 @@ def test_get_symbols_by_type(self):
             filepath="/root/file1.py",
             language="python",
             symbols=[
-                Symbol(name="func1", type=SymbolType.FUNCTION, filepath="/root/file1.py", line_start=1, line_end=5, language="python", signature="def func1()"),
-                Symbol(name="Class1", type=SymbolType.CLASS, filepath="/root/file1.py", line_start=10, line_end=20, language="python", signature="class Class1"),
-            ]
+                Symbol(
+                    name="func1",
+                    type=SymbolType.FUNCTION,
+                    filepath="/root/file1.py",
+                    line_start=1,
+                    line_end=5,
+                    language="python",
+                    signature="def func1()",
+                ),
+                Symbol(
+                    name="Class1",
+                    type=SymbolType.CLASS,
+                    filepath="/root/file1.py",
+                    line_start=10,
+                    line_end=20,
+                    language="python",
+                    signature="class Class1",
+                ),
+            ],
         )
         index = CodebaseIndex(root_path="/root", files={"/root/file1.py": file1})
-        
+
         funcs = index.get_symbols_by_type(SymbolType.FUNCTION)
         assert len(funcs) == 1
         assert funcs[0].name == "func1"
-        
+
         classes = index.get_symbols_by_type(SymbolType.CLASS)
         assert len(classes) == 1
         assert classes[0].name == "Class1"
@@ -292,7 +331,7 @@ def test_validation_total_symbols(self):
     def test_successful_failed_parses_properties(self):
         file1 = ParsedFile(filepath="f1", language="py", symbols=[])
         file2 = ParsedFile(filepath="f2", language="py", symbols=[], error="Error")
-        
+
         index = CodebaseIndex(root_path="/root", files={"f1": file1, "f2": file2})
         assert index.successful_parses == 1
         assert index.failed_parses == 1
@@ -300,15 +339,8 @@ def test_successful_failed_parses_properties(self):
     def test_stats_properties(self):
         success = ParsedFile(filepath="/a.py", language="py", symbols=[])
         failed = ParsedFile(filepath="/b.py", language="py", symbols=[], error="Syntax error")
-        
-        index = CodebaseIndex(
-            root_path="/app",
-            files={
-                "/a.py": success,
-                "/b.py": failed
-            }
-        )
-        
+
+        index = CodebaseIndex(root_path="/app", files={"/a.py": success, "/b.py": failed})
+
         assert index.successful_parses == 1
         assert index.failed_parses == 1
-
diff --git a/tests/unit/test_parallel_indexer.py b/tests/unit/test_parallel_indexer.py
index 9706cfa..be63d08 100644
--- a/tests/unit/test_parallel_indexer.py
+++ b/tests/unit/test_parallel_indexer.py
@@ -1,50 +1,58 @@
-
 import pytest
 from pathlib import Path
-from codegrok_mcp.indexing.parallel_indexer import ParallelProgress, parallel_parse_files, parse_file_worker
+from codegrok_mcp.indexing.parallel_indexer import (
+    ParallelProgress,
+    parallel_parse_files,
+    parse_file_worker,
+)
 from codegrok_mcp.parsers.treesitter_parser import ThreadLocalParserFactory
 
+
 def test_parallel_progress_increment():
     progress = ParallelProgress(total=10)
     assert progress.completed == 0
     assert progress.errors == 0
-    
+
     new_val = progress.increment_completed()
     assert new_val == 1
     assert progress.completed == 1
-    
+
     new_err = progress.increment_errors()
     assert new_err == 1
     assert progress.errors == 1
 
+
 def test_parallel_parse_files_empty_list():
     symbols, errors = parallel_parse_files([])
     assert symbols == []
     assert errors == 0
 
+
 def test_parallel_parse_files_single_file(tmp_path):
     f = tmp_path / "test.py"
     f.write_text("def foo(): pass")
-    
+
     symbols, errors = parallel_parse_files([f])
-    
+
     assert errors == 0
     assert len(symbols) >= 1
     assert symbols[0].name == "foo"
     assert symbols[0].filepath == str(f)
 
+
 def test_parse_file_worker_success(tmp_path):
     f = tmp_path / "test.py"
     f.write_text("def worker_test(): pass")
     factory = ThreadLocalParserFactory()
-    
+
     result = parse_file_worker(f, factory)
-    
+
     assert result.success
     assert result.filepath == str(f)
     assert len(result.symbols) >= 1
     assert result.symbols[0].name == "worker_test"
 
+
 def test_parse_file_worker_error(tmp_path):
     # Depending on parser implementation, syntax error might not raise exception but return ParseResult with error
     # Or might parse partially.
@@ -52,15 +60,16 @@ def test_parse_file_worker_error(tmp_path):
     d = tmp_path / "subdir"
     d.mkdir()
     factory = ThreadLocalParserFactory()
-    
+
     # parse_file usually expects a file path string. If it fails, worker catches exception.
     # TreeSitterParser.parse_file reads content.
-    
+
     result = parse_file_worker(d, factory)
-    
+
     assert not result.success
     assert result.error is not None
 
+
 def test_parallel_parse_files_with_error(tmp_path):
     d = tmp_path / "subdir"
     d.mkdir()
diff --git a/tests/unit/test_parser.py b/tests/unit/test_parser.py
index 4677efe..f18965e 100644
--- a/tests/unit/test_parser.py
+++ b/tests/unit/test_parser.py
@@ -1,4 +1,5 @@
 """Unit tests for TreeSitterParser."""
+
 import pytest
 from codegrok_mcp.parsers.treesitter_parser import TreeSitterParser, ThreadLocalParserFactory
 from codegrok_mcp.core.models import SymbolType
@@ -67,7 +68,6 @@ def test_parse_javascript_function(self, parser, tmp_path):
         assert result.is_successful
         assert any(s.name == "hello" for s in result.symbols)
 
-
     def test_parse_javascript_constants(self, parser, tmp_path):
         js_file = tmp_path / "test.js"
         js_file.write_text("const MAX_SIZE = 100;")
@@ -135,17 +135,14 @@ def test_parse_python_constant(self, parser, tmp_path):
         assert "TIMEOUT" in consts
 
 
-
-
 class TestThreadLocalParserFactory:
     def test_thread_local_parser_factory(self):
         factory = ThreadLocalParserFactory()
         parser1 = factory.get_parser()
         parser2 = factory.get_parser()
         assert parser1 is parser2
-        
+
     def test_factory_creates_parser(self):
         factory = ThreadLocalParserFactory()
         parser = factory.get_parser()
         assert isinstance(parser, TreeSitterParser)
-

From 2543eb1daa00ca336bcb845f5646493448ee6074 Mon Sep 17 00:00:00 2001
From: Shreyas <shreyasjagannath08@gmail.com>
Date: Wed, 11 Mar 2026 19:24:14 +0000
Subject: [PATCH 2/3] perf: reduce memory consumption and add configurable
 timeout

Memory optimizations:
- Inline symbol-to-chunk conversion per file batch (never hold both lists)
- Process files in batches of 500 with gc.collect() between batches
- Cap parallel parse workers at 4 (was up to 32)
- Explicit del + gc.collect() after embedding completes

Configurable MCP timeout:
- Default 600s, configurable via CODEGROK_TIMEOUT env var
- Per-call override via timeout_seconds parameter on learn tool
- Indexing runs in asyncio.to_thread() so wait_for can cancel
- On timeout: checkpoint preserved, re-run resumes

Tests: 15 new tests (9 timeout unit, 4 memory integration, 2 constant unit)
Snyk: 0 issues on modified files

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 CLAUDE.md                                     |  7 +-
 docs/INDEXING_IMPROVEMENTS.md                 | 56 +++++++++++++-
 src/codegrok_mcp/indexing/source_retriever.py | 74 +++++++++++++------
 src/codegrok_mcp/mcp/server.py                | 73 ++++++++++++++++--
 tests/integration/test_source_retriever.py    | 59 +++++++++++++++
 tests/unit/test_discover_files.py             | 22 +++++-
 tests/unit/test_timeout.py                    | 58 +++++++++++++++
 7 files changed, 316 insertions(+), 33 deletions(-)
 create mode 100644 tests/unit/test_timeout.py

diff --git a/CLAUDE.md b/CLAUDE.md
index fc76a0d..de1c91a 100644
--- a/CLAUDE.md
+++ b/CLAUDE.md
@@ -43,7 +43,9 @@ src/codegrok_mcp/
 - **Chunk Strategy**: Symbol-based (each function/class/method = 1 chunk)
 - **Max Chunk Size**: 4000 chars (~1000-1300 tokens)
 - **Storage**: `.codegrok/` (chromadb/ + metadata.json + memory_metadata.json + checkpoint.json)
-- **Parallelism**: CPU count - 1 workers (min 1, max 32)
+- **Parallelism**: CPU count - 1 workers (min 1, max 4 for parsing)
+- **File Batch Size**: 500 files per parse batch (memory optimization)
+- **Default Timeout**: 600s (configurable via `CODEGROK_TIMEOUT` env var or `timeout_seconds` param)
 - **Memory TTLs**: session (24h), day, week, month, permanent
 
 ## Commands
@@ -70,6 +72,9 @@ mypy src/                 # Type check
 10. **Indexing uses upsert** - `collection.upsert()` instead of delete-recreate; stale chunks cleaned after embedding
 11. **Checkpointing** - `.codegrok/checkpoint.json` saves progress every 1000 chunks; atomic writes via `os.replace()`; deleted on success
 12. **max_files safety limit** - `discover_files()` stops at 200K files to prevent DoS (addresses SECURITY_REVIEW HIGH-003)
+13. **Memory-optimized parsing** - Symbols converted to chunks per file batch (500 files), then freed; `gc.collect()` between batches; chunks list freed after embedding
+14. **Worker cap** - Parallel parse workers capped at `MAX_PARSE_WORKERS=4` to limit memory from tree-sitter instances
+15. **Configurable timeout** - `learn` tool has `timeout_seconds` param; also reads `CODEGROK_TIMEOUT` env var; defaults to 600s
 
 ## Adding Languages
 
diff --git a/docs/INDEXING_IMPROVEMENTS.md b/docs/INDEXING_IMPROVEMENTS.md
index 583f981..6c47350 100644
--- a/docs/INDEXING_IMPROVEMENTS.md
+++ b/docs/INDEXING_IMPROVEMENTS.md
@@ -1,6 +1,7 @@
 # Indexing Improvements (v0.2.1)
 
-Fixes for the `learn` tool hanging on large codebases with many folders/subfolders.
+Fixes for the `learn` tool hanging on large codebases with many folders/subfolders,
+high memory consumption during indexing, and missing timeout protection.
 
 ## Changes
 
@@ -39,6 +40,39 @@ Fixes for the `learn` tool hanging on large codebases with many folders/subfolde
 - ETA added to embedding progress messages (e.g., "Embedding... (5000/10000 chunks, ~2.3m remaining)")
 - MCP client now shows progress during the file discovery phase (0-5% range)
 
+### 6. Memory Optimizations
+
+Reduces peak memory consumption during indexing:
+
+- **Inline symbol-to-chunk conversion**: Symbols are converted to chunks per file batch and freed immediately, instead of accumulating all symbols in a separate list before chunking
+- **File batch processing**: Files are parsed in batches of 500 (`FILE_BATCH_SIZE`), with `gc.collect()` between batches to free memory promptly
+- **Worker cap**: Parallel parse workers capped at 4 (`MAX_PARSE_WORKERS`) to limit memory from tree-sitter parser instances (previously up to 32)
+- **Post-embedding cleanup**: Chunks list is explicitly deleted and garbage collected after embedding completes
+
+### 7. Configurable MCP Timeout
+
+Prevents the `learn` tool from running indefinitely on very large codebases:
+
+- Default timeout: 600 seconds (10 minutes)
+- **Environment variable**: Set `CODEGROK_TIMEOUT` in your MCP client config
+- **Per-call override**: Pass `timeout_seconds` parameter to the `learn` tool
+- Priority: `timeout_seconds` param > `CODEGROK_TIMEOUT` env var > default (600s)
+- On timeout: checkpoint is preserved, re-running `learn` resumes from where it stopped
+
+**Configuration example** (claude_desktop_config.json):
+```json
+{
+  "mcpServers": {
+    "codegrok": {
+      "command": "codegrok-mcp",
+      "env": {
+        "CODEGROK_TIMEOUT": "1200"
+      }
+    }
+  }
+}
+```
+
 ## New Dependencies
 
 - `pathspec>=0.11.0` — Pure Python `.gitignore` pattern matching (used by `black`, `flake8`, etc.)
@@ -73,6 +107,22 @@ This feature was planned and implemented using the following MCP tools:
 | `test_discover_files_no_gitignore` | Works when no `.gitignore` exists |
 | `test_discover_files_respect_gitignore_false` | Opt-out disables gitignore filtering |
 | `test_discover_files_progress_callback` | Callback mechanism works correctly |
+| `test_file_batch_size` | FILE_BATCH_SIZE constant is 500 |
+| `test_max_parse_workers` | MAX_PARSE_WORKERS constant is 4 |
+
+### New Unit Tests (`tests/unit/test_timeout.py`)
+
+| Test | What it verifies |
+|------|-----------------|
+| `test_default_timeout` | Returns 600s when no override or env var |
+| `test_per_call_override` | Per-call override takes highest priority |
+| `test_env_var_override` | CODEGROK_TIMEOUT env var overrides default |
+| `test_env_var_invalid_ignored` | Invalid env var falls back to default |
+| `test_env_var_zero_ignored` | Zero env var falls back to default |
+| `test_env_var_negative_ignored` | Negative env var falls back to default |
+| `test_override_zero_uses_env` | Override of 0/None falls through to env |
+| `test_override_negative_uses_env` | Negative override falls through to env |
+| `test_default_is_600` | Default constant is 600 seconds |
 
 ### New Integration Tests (`tests/integration/test_source_retriever.py`)
 
@@ -85,3 +135,7 @@ This feature was planned and implemented using the following MCP tools:
 | `test_checkpoint_load_corrupted` | Handles corrupted JSON |
 | `test_checkpoint_cleanup_on_success` | Checkpoint deleted after success |
 | `test_checkpoint_load_none_path` | Handles None path |
+| `test_file_batch_processing_parallel` | Parallel batch processing produces correct results |
+| `test_file_batch_processing_sequential` | Sequential inline chunking produces correct results |
+| `test_worker_cap_respected` | Workers capped at MAX_PARSE_WORKERS |
+| `test_explicit_max_workers_not_overridden` | Explicit max_workers used as-is |
diff --git a/src/codegrok_mcp/indexing/source_retriever.py b/src/codegrok_mcp/indexing/source_retriever.py
index 6160b64..c718a32 100644
--- a/src/codegrok_mcp/indexing/source_retriever.py
+++ b/src/codegrok_mcp/indexing/source_retriever.py
@@ -28,6 +28,7 @@
     sources = retriever.get_sources_for_question("How does authentication work?")
 """
 
+import gc
 import json
 import logging
 import os
@@ -72,6 +73,14 @@
     ".eggs",
 }
 
+# Memory optimization: process files in batches to limit peak memory usage.
+# Each batch's symbols are converted to chunks then freed before the next batch.
+FILE_BATCH_SIZE = 500
+
+# Cap parallel parse workers to limit memory (each worker holds a tree-sitter parser).
+# Lower than the 32 max in parallel_indexer.py to prevent OOM on large codebases.
+MAX_PARSE_WORKERS = 4
+
 
 def _load_gitignore(directory: Path) -> Optional[pathspec.PathSpec]:
     """Load a .gitignore file from a directory, returning a PathSpec or None."""
@@ -468,36 +477,61 @@ def emit(event_type: str, data: dict):
         if not progress_callback:
             self._log(f"Found {len(all_files)} files")
 
-        # Step 2: Parse all files
+        # Step 2+3: Parse files and create chunks (memory-optimized)
+        # Symbols are converted to chunks per batch then freed, so we never
+        # hold both a large all_symbols list and a large chunks list simultaneously.
         if not progress_callback:
-            self._log("\nStep 2: Parsing files...")
+            self._log("\nStep 2: Parsing files and creating chunks...")
 
         emit("parsing_start", {"total": len(all_files)})
-        all_symbols = []
+        chunks = []
+        total_symbols = 0
 
         # Use parallel parsing if enabled and there are enough files
         use_parallel = (
             self.parallel and len(all_files) > 50
         )  # Threshold increased for small projects
+
         if use_parallel:
-            # Parallel parsing (3-5x faster for large codebases)
             from codegrok_mcp.indexing.parallel_indexer import parallel_parse_files
 
-            if not progress_callback:
-                self._log(f"  Using parallel parsing with {self.max_workers or 'auto'} workers...")
+            # Cap workers to limit memory (each holds a tree-sitter parser instance)
+            effective_workers = self.max_workers
+            if effective_workers is None:
+                cpu_count = os.cpu_count() or 4
+                effective_workers = max(1, min(cpu_count - 1, MAX_PARSE_WORKERS))
 
-            all_symbols, parse_errors = parallel_parse_files(
-                files=all_files, max_workers=self.max_workers, progress_callback=progress_callback
-            )
-            self.stats["parse_errors"] = parse_errors
+            if not progress_callback:
+                self._log(f"  Using parallel parsing with {effective_workers} workers...")
+
+            # Process files in batches to limit peak memory usage
+            for batch_start in range(0, len(all_files), FILE_BATCH_SIZE):
+                file_batch = all_files[batch_start : batch_start + FILE_BATCH_SIZE]
+                batch_symbols, batch_errors = parallel_parse_files(
+                    files=file_batch,
+                    max_workers=effective_workers,
+                    progress_callback=progress_callback,
+                )
+                self.stats["parse_errors"] += batch_errors
+                total_symbols += len(batch_symbols)
+
+                # Convert symbols to chunks immediately, then free symbol memory
+                for symbol in batch_symbols:
+                    chunks.append(self._create_chunk(symbol))
+                del batch_symbols
+                gc.collect()
         else:
-            # Sequential parsing (original code)
+            # Sequential parsing
             for i, filepath in enumerate(all_files, 1):
                 symbols_count = 0
                 try:
                     parsed = self.parser.parse_file(str(filepath))
-                    all_symbols.extend(parsed.symbols)
                     symbols_count = len(parsed.symbols)
+                    total_symbols += symbols_count
+
+                    # Convert to chunks immediately (don't accumulate symbols)
+                    for symbol in parsed.symbols:
+                        chunks.append(self._create_chunk(symbol))
 
                     emit(
                         "file_parsed",
@@ -518,21 +552,13 @@ def emit(event_type: str, data: dict):
                 if not progress_callback and self.verbose and i % 100 == 0:
                     print(f"  Parsed {i}/{len(all_files)} files...", end="\r")
 
-        self.stats["total_symbols"] = len(all_symbols)
-        if not progress_callback:
-            self._log(f"\nParsed {len(all_symbols):,} symbols from {len(all_files)} files")
-
-        # Step 3: Create chunks
-        if not progress_callback:
-            self._log("\nStep 3: Creating chunks...")
-
-        chunks = [self._create_chunk(symbol) for symbol in all_symbols]
+        self.stats["total_symbols"] = total_symbols
         self.stats["total_chunks"] = len(chunks)
 
         emit("chunks_created", {"total": len(chunks)})
 
         if not progress_callback:
-            self._log(f"Created {len(chunks):,} chunks")
+            self._log(f"\nParsed {total_symbols:,} symbols → {len(chunks):,} chunks")
 
         # Step 4: Get or create ChromaDB collection (supports resumption)
         if not progress_callback:
@@ -635,6 +661,10 @@ def emit(event_type: str, data: dict):
             if checkpoint_path and current_count % 1000 == 0:
                 _save_checkpoint(checkpoint_path, current_count, len(chunks))
 
+        # Free chunks list now that embedding is complete
+        del chunks
+        gc.collect()
+
         # Remove stale chunks (from deleted/renamed files)
         try:
             existing = self.collection.get(include=[])
diff --git a/src/codegrok_mcp/mcp/server.py b/src/codegrok_mcp/mcp/server.py
index 7af421b..c4aa760 100644
--- a/src/codegrok_mcp/mcp/server.py
+++ b/src/codegrok_mcp/mcp/server.py
@@ -20,6 +20,7 @@
 from typing import Optional, List, Dict, Any, Annotated, Callable
 from pathlib import Path
 import asyncio
+import os
 
 from fastmcp import FastMCP, Context
 from fastmcp.exceptions import ToolError
@@ -41,6 +42,31 @@
 CHROMA_DIR = "chroma"
 METADATA_FILE = "metadata.json"
 
+# Default timeout for indexing operations (seconds).
+# Configurable via CODEGROK_TIMEOUT environment variable in MCP client config.
+# Example in claude_desktop_config.json:
+#   "env": {"CODEGROK_TIMEOUT": "1200"}
+DEFAULT_TIMEOUT_SECONDS = 600  # 10 minutes
+
+
+def _get_timeout(override: Optional[int] = None) -> int:
+    """Get the effective timeout in seconds.
+
+    Priority: per-call override > CODEGROK_TIMEOUT env var > DEFAULT_TIMEOUT_SECONDS.
+    """
+    if override is not None and override > 0:
+        return override
+    env_timeout = os.environ.get("CODEGROK_TIMEOUT")
+    if env_timeout:
+        try:
+            val = int(env_timeout)
+            if val > 0:
+                return val
+        except ValueError:
+            pass
+    return DEFAULT_TIMEOUT_SECONDS
+
+
 # Initialize FastMCP server
 mcp = FastMCP(
     name="CodeGrok",
@@ -199,6 +225,12 @@ async def learn(
     embedding_model: Annotated[
         str, Field(description="Embedding model to use (default: coderankembed)")
     ] = "coderankembed",
+    timeout_seconds: Annotated[
+        Optional[int],
+        Field(
+            description="Timeout in seconds for indexing. Overrides CODEGROK_TIMEOUT env var. Default: 600 (10 min)."
+        ),
+    ] = None,
     ctx: Context = None,
 ) -> Dict[str, Any]:
     """Index a codebase with smart mode detection."""
@@ -220,7 +252,10 @@ async def learn(
     paths = _get_codegrok_paths(codebase_path)
     has_existing = _has_valid_index(paths)
 
-    # Handle load_only mode
+    # Resolve timeout
+    timeout = _get_timeout(timeout_seconds)
+
+    # Handle load_only mode (no timeout needed — just loads metadata)
     if mode == "load_only":
         if not has_existing:
             raise ToolError(
@@ -231,10 +266,30 @@ async def learn(
 
     # Handle auto mode with existing index -> incremental reindex
     if mode == "auto" and has_existing:
-        return await _incremental_reindex(codebase_path, paths, state, embedding_model, ctx)
+        try:
+            return await asyncio.wait_for(
+                _incremental_reindex(codebase_path, paths, state, embedding_model, ctx),
+                timeout=timeout,
+            )
+        except asyncio.TimeoutError:
+            raise ToolError(
+                f"Indexing timed out after {timeout}s. "
+                f"Set CODEGROK_TIMEOUT env var or pass timeout_seconds to increase. "
+                f"Checkpoint saved — re-run to resume."
+            )
 
     # Full index: mode == "full" OR (mode == "auto" and no existing index)
-    return await _full_index(codebase_path, paths, state, file_extensions, embedding_model, ctx)
+    try:
+        return await asyncio.wait_for(
+            _full_index(codebase_path, paths, state, file_extensions, embedding_model, ctx),
+            timeout=timeout,
+        )
+    except asyncio.TimeoutError:
+        raise ToolError(
+            f"Indexing timed out after {timeout}s. "
+            f"Set CODEGROK_TIMEOUT env var or pass timeout_seconds to increase. "
+            f"Checkpoint saved — re-run to resume."
+        )
 
 
 async def _load_existing_index(
@@ -294,8 +349,10 @@ async def _incremental_reindex(
         loop = asyncio.get_event_loop()
         progress_callback = _create_relearn_progress_callback(ctx, loop)
 
-    # Perform incremental reindex
-    result = retriever.incremental_reindex(progress_callback=progress_callback)
+    # Run blocking reindex in a thread so asyncio.wait_for can cancel it
+    result = await asyncio.to_thread(
+        retriever.incremental_reindex, progress_callback=progress_callback
+    )
 
     # Save updated metadata
     retriever.save_metadata(str(paths["metadata_path"]))
@@ -338,9 +395,11 @@ async def _full_index(
         persist_path=str(paths["chroma_path"]),
     )
 
-    # Index the codebase with progress reporting
+    # Run blocking indexing in a thread so asyncio.wait_for can cancel it
     extensions = file_extensions if file_extensions else SUPPORTED_EXTENSIONS
-    retriever.index_codebase(file_extensions=extensions, progress_callback=progress_callback)
+    await asyncio.to_thread(
+        retriever.index_codebase, file_extensions=extensions, progress_callback=progress_callback
+    )
 
     # Report saving phase
     if ctx:
diff --git a/tests/integration/test_source_retriever.py b/tests/integration/test_source_retriever.py
index 1ac5375..86f6e31 100644
--- a/tests/integration/test_source_retriever.py
+++ b/tests/integration/test_source_retriever.py
@@ -194,6 +194,65 @@ def test_stale_chunk_removal(self, temp_project):
             assert count_after < count_before
 
 
+class TestMemoryOptimizations:
+    """Test memory optimization behavior."""
+
+    def test_file_batch_processing_parallel(self, temp_project):
+        """Parallel parsing processes files in batches without accumulating all symbols."""
+        with tempfile.TemporaryDirectory() as persist_dir:
+            retriever = SourceRetriever(
+                codebase_path=str(temp_project), persist_path=persist_dir, parallel=True
+            )
+            retriever.index_codebase()
+
+            # Should still produce correct results
+            stats = retriever.get_stats()
+            assert stats["total_symbols"] > 0
+            assert stats["total_chunks"] > 0
+            assert stats["total_chunks"] == stats["total_symbols"]
+
+    def test_file_batch_processing_sequential(self, temp_project):
+        """Sequential parsing converts symbols to chunks inline."""
+        with tempfile.TemporaryDirectory() as persist_dir:
+            retriever = SourceRetriever(
+                codebase_path=str(temp_project), persist_path=persist_dir, parallel=False
+            )
+            retriever.index_codebase()
+
+            stats = retriever.get_stats()
+            assert stats["total_symbols"] > 0
+            assert stats["total_chunks"] == stats["total_symbols"]
+
+    def test_worker_cap_respected(self, temp_project):
+        """Parallel workers are capped at MAX_PARSE_WORKERS when max_workers is None."""
+        from codegrok_mcp.indexing.source_retriever import MAX_PARSE_WORKERS
+
+        with tempfile.TemporaryDirectory() as persist_dir:
+            retriever = SourceRetriever(
+                codebase_path=str(temp_project),
+                persist_path=persist_dir,
+                parallel=True,
+                max_workers=None,  # Should auto-cap
+            )
+            retriever.index_codebase()
+
+            # Verify it completed successfully (workers were capped internally)
+            assert retriever.get_stats()["total_chunks"] > 0
+            assert MAX_PARSE_WORKERS == 4
+
+    def test_explicit_max_workers_not_overridden(self, temp_project):
+        """Explicit max_workers value is used as-is."""
+        with tempfile.TemporaryDirectory() as persist_dir:
+            retriever = SourceRetriever(
+                codebase_path=str(temp_project),
+                persist_path=persist_dir,
+                parallel=True,
+                max_workers=2,
+            )
+            retriever.index_codebase()
+            assert retriever.get_stats()["total_chunks"] > 0
+
+
 class TestCheckpointing:
     """Test checkpoint save/load/resume functionality."""
 
diff --git a/tests/unit/test_discover_files.py b/tests/unit/test_discover_files.py
index e1db3ec..569dddf 100644
--- a/tests/unit/test_discover_files.py
+++ b/tests/unit/test_discover_files.py
@@ -1,10 +1,14 @@
 """
-Unit tests for discover_files() — file discovery with .gitignore support.
+Unit tests for discover_files() and memory optimization constants.
 """
 
 import pytest
 from pathlib import Path
-from codegrok_mcp.indexing.source_retriever import discover_files
+from codegrok_mcp.indexing.source_retriever import (
+    discover_files,
+    FILE_BATCH_SIZE,
+    MAX_PARSE_WORKERS,
+)
 
 
 def _create_file(path: Path, content: str = "# placeholder"):
@@ -124,3 +128,17 @@ def callback(event_type, data):
         # With only 5 files, callback won't fire (fires every 1000)
         assert len(files) == 5
         assert len(events) == 0  # Below threshold
+
+
+class TestMemoryOptimizationConstants:
+    """Test that memory optimization constants are set correctly."""
+
+    def test_file_batch_size(self):
+        """FILE_BATCH_SIZE is set to a reasonable value."""
+        assert FILE_BATCH_SIZE == 500
+        assert isinstance(FILE_BATCH_SIZE, int)
+
+    def test_max_parse_workers(self):
+        """MAX_PARSE_WORKERS caps parallel workers to limit memory."""
+        assert MAX_PARSE_WORKERS == 4
+        assert isinstance(MAX_PARSE_WORKERS, int)
diff --git a/tests/unit/test_timeout.py b/tests/unit/test_timeout.py
new file mode 100644
index 0000000..7086707
--- /dev/null
+++ b/tests/unit/test_timeout.py
@@ -0,0 +1,58 @@
+"""
+Unit tests for MCP timeout configuration.
+"""
+
+import os
+import pytest
+from unittest.mock import patch
+
+from codegrok_mcp.mcp.server import _get_timeout, DEFAULT_TIMEOUT_SECONDS
+
+
+class TestGetTimeout:
+    """Test timeout resolution logic."""
+
+    def test_default_timeout(self):
+        """Returns DEFAULT_TIMEOUT_SECONDS when no override or env var."""
+        with patch.dict(os.environ, {}, clear=True):
+            assert _get_timeout() == DEFAULT_TIMEOUT_SECONDS
+
+    def test_per_call_override(self):
+        """Per-call override takes highest priority."""
+        with patch.dict(os.environ, {"CODEGROK_TIMEOUT": "999"}):
+            assert _get_timeout(override=120) == 120
+
+    def test_env_var_override(self):
+        """CODEGROK_TIMEOUT env var overrides default."""
+        with patch.dict(os.environ, {"CODEGROK_TIMEOUT": "1200"}):
+            assert _get_timeout() == 1200
+
+    def test_env_var_invalid_ignored(self):
+        """Invalid CODEGROK_TIMEOUT falls back to default."""
+        with patch.dict(os.environ, {"CODEGROK_TIMEOUT": "not_a_number"}):
+            assert _get_timeout() == DEFAULT_TIMEOUT_SECONDS
+
+    def test_env_var_zero_ignored(self):
+        """Zero CODEGROK_TIMEOUT falls back to default."""
+        with patch.dict(os.environ, {"CODEGROK_TIMEOUT": "0"}):
+            assert _get_timeout() == DEFAULT_TIMEOUT_SECONDS
+
+    def test_env_var_negative_ignored(self):
+        """Negative CODEGROK_TIMEOUT falls back to default."""
+        with patch.dict(os.environ, {"CODEGROK_TIMEOUT": "-5"}):
+            assert _get_timeout() == DEFAULT_TIMEOUT_SECONDS
+
+    def test_override_zero_uses_env(self):
+        """Override of 0 or None falls through to env var."""
+        with patch.dict(os.environ, {"CODEGROK_TIMEOUT": "300"}):
+            assert _get_timeout(override=0) == 300
+            assert _get_timeout(override=None) == 300
+
+    def test_override_negative_uses_env(self):
+        """Negative override falls through to env var."""
+        with patch.dict(os.environ, {"CODEGROK_TIMEOUT": "300"}):
+            assert _get_timeout(override=-1) == 300
+
+    def test_default_is_600(self):
+        """Default timeout is 600 seconds (10 minutes)."""
+        assert DEFAULT_TIMEOUT_SECONDS == 600

From 1697e09a7a68ef29329f87a8be930ab68c5c5406 Mon Sep 17 00:00:00 2001
From: Shreyas <shreyasjagannath08@gmail.com>
Date: Wed, 11 Mar 2026 19:50:29 +0000
Subject: [PATCH 3/3] fix: background indexing to prevent MCP transport timeout

MCP clients (Claude Desktop, etc.) have ~60-120s transport-level timeouts
that kill long-running tool calls. The learn tool now returns immediately
and runs indexing in a background daemon thread. Clients poll get_stats()
for progress.

- Add IndexingStatus thread-safe dataclass for background progress tracking
- Rewrite learn tool: starts threading.Thread(daemon=True), returns immediately
- get_stats() includes indexing progress when active or errored
- Stateful learn responses: indexing_started, indexing_in_progress, complete, error
- load_only mode stays synchronous (fast, no background needed)
- 33 new tests for IndexingStatus, progress callbacks, learn behavior, get_stats
- Snyk SAST scan: 0 issues
- All 191 tests pass

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 CLAUDE.md                              |   2 +
 docs/INDEXING_IMPROVEMENTS.md          |  51 ++++
 src/codegrok_mcp/mcp/server.py         | 317 ++++++++++----------
 src/codegrok_mcp/mcp/state.py          |  53 +++-
 tests/unit/test_background_indexing.py | 396 +++++++++++++++++++++++++
 5 files changed, 653 insertions(+), 166 deletions(-)
 create mode 100644 tests/unit/test_background_indexing.py

diff --git a/CLAUDE.md b/CLAUDE.md
index de1c91a..51be059 100644
--- a/CLAUDE.md
+++ b/CLAUDE.md
@@ -75,6 +75,8 @@ mypy src/                 # Type check
 13. **Memory-optimized parsing** - Symbols converted to chunks per file batch (500 files), then freed; `gc.collect()` between batches; chunks list freed after embedding
 14. **Worker cap** - Parallel parse workers capped at `MAX_PARSE_WORKERS=4` to limit memory from tree-sitter instances
 15. **Configurable timeout** - `learn` tool has `timeout_seconds` param; also reads `CODEGROK_TIMEOUT` env var; defaults to 600s
+16. **Background indexing** - `learn` returns immediately; indexing runs in `threading.Thread(daemon=True)`; client polls `get_stats()` for progress; `IndexingStatus` in `state.py` tracks progress thread-safely
+17. **learn stateful responses** - Returns `indexing_started` (new), `indexing_in_progress` (already running), `complete` (done, clears result), or raises `ToolError` (failed, clears error for retry)
 
 ## Adding Languages
 
diff --git a/docs/INDEXING_IMPROVEMENTS.md b/docs/INDEXING_IMPROVEMENTS.md
index 6c47350..dcac01f 100644
--- a/docs/INDEXING_IMPROVEMENTS.md
+++ b/docs/INDEXING_IMPROVEMENTS.md
@@ -73,6 +73,28 @@ Prevents the `learn` tool from running indefinitely on very large codebases:
 }
 ```
 
+### 8. Background Indexing
+
+Solves MCP transport timeout: clients (Claude Desktop, etc.) have ~60-120s transport-level
+timeouts that kill long-running tool calls. Now `learn` returns immediately and indexing
+runs in a background thread.
+
+- `learn` starts a `threading.Thread(daemon=True)` and returns `{"status": "indexing_started"}`
+- `IndexingStatus` dataclass in `state.py` provides thread-safe progress tracking with `threading.Lock`
+- `get_stats()` includes `indexing` field with `active`, `progress`, `message`, `error` when indexing is in progress
+- Polling pattern: client calls `get_stats()` repeatedly to check progress
+- If `learn` is called while indexing is active, returns `{"status": "indexing_in_progress"}`
+- If `learn` is called after completion, returns the result once and clears it
+- If previous indexing failed, raises `ToolError` with error details and clears for retry
+- `load_only` mode remains synchronous (fast, no background needed)
+
+**Workflow:**
+```
+1. learn(path="/project")         → {"status": "indexing_started", "progress": 0}
+2. get_stats()                    → {"indexing": {"active": true, "progress": 42, ...}}
+3. get_stats()                    → {"loaded": true, "stats": {...}}  (indexing done)
+```
+
 ## New Dependencies
 
 - `pathspec>=0.11.0` — Pure Python `.gitignore` pattern matching (used by `black`, `flake8`, etc.)
@@ -139,3 +161,32 @@ This feature was planned and implemented using the following MCP tools:
 | `test_file_batch_processing_sequential` | Sequential inline chunking produces correct results |
 | `test_worker_cap_respected` | Workers capped at MAX_PARSE_WORKERS |
 | `test_explicit_max_workers_not_overridden` | Explicit max_workers used as-is |
+
+### New Unit Tests (`tests/unit/test_background_indexing.py`)
+
+| Test | What it verifies |
+|------|-----------------|
+| `test_initial_state` | IndexingStatus defaults are correct |
+| `test_start` | start() sets active, clears error/result |
+| `test_start_clears_previous_error` | Restart after failure clears error |
+| `test_start_clears_previous_result` | Restart after success clears result |
+| `test_update` | Progress and message update correctly |
+| `test_update_caps_at_99` | Progress capped at 99 (100 = complete only) |
+| `test_complete` | Complete sets progress=100, stores result |
+| `test_fail` | Fail sets error, deactivates |
+| `test_to_dict` | Dict output has correct fields |
+| `test_to_dict_excludes_result` | Result not exposed in to_dict |
+| `test_thread_safety` | 10 concurrent threads update without errors |
+| `test_state_has_indexing_status` | MCPSessionState includes IndexingStatus |
+| `test_singleton_state_has_indexing` | Singleton state has indexing field |
+| `test_callback_updates_indexing_status` | Progress callback updates status |
+| `test_callback_discovery_progress` | Discovery progress event handled |
+| `test_callback_embedding_progress_with_eta` | ETA displayed in message |
+| `test_learn_returns_in_progress_when_active` | learn rejects when already indexing |
+| `test_learn_returns_completed_result` | learn returns result after completion |
+| `test_learn_raises_on_previous_error` | learn raises ToolError on prior failure |
+| `test_learn_starts_background_thread` | learn spawns daemon thread |
+| `test_learn_auto_with_existing_uses_incremental` | auto mode uses incremental reindex |
+| `test_learn_full_mode_uses_full_index` | full mode uses full index |
+| `test_get_stats_includes_indexing_when_active` | get_stats shows indexing progress |
+| `test_get_stats_no_indexing_when_idle` | get_stats omits indexing when idle |
diff --git a/src/codegrok_mcp/mcp/server.py b/src/codegrok_mcp/mcp/server.py
index c4aa760..17c1bcb 100644
--- a/src/codegrok_mcp/mcp/server.py
+++ b/src/codegrok_mcp/mcp/server.py
@@ -21,13 +21,14 @@
 from pathlib import Path
 import asyncio
 import os
+import threading
 
 from fastmcp import FastMCP, Context
 from fastmcp.exceptions import ToolError
 from mcp.types import ToolAnnotations
 from pydantic import Field
 
-from codegrok_mcp.mcp.state import get_state
+from codegrok_mcp.mcp.state import get_state, IndexingStatus
 
 # Lazy import SourceRetriever to avoid heavy startup cost
 # from codegrok_mcp.indexing.source_retriever import SourceRetriever, SUPPORTED_EXTENSIONS
@@ -119,8 +120,8 @@ def _has_valid_index(paths: Dict[str, Path]) -> bool:
     )
 
 
-def _create_learn_progress_callback(ctx: Context, loop) -> Callable:
-    """Create a progress callback that reports indexing progress to MCP client."""
+def _create_bg_progress_callback(indexing_status: IndexingStatus) -> Callable:
+    """Create a progress callback that updates IndexingStatus (thread-safe)."""
 
     def callback(event_type: str, data: dict):
         progress = 0
@@ -142,7 +143,6 @@ def callback(event_type: str, data: dict):
             progress = 35
             message = f"Generating embeddings for {data['total']} chunks..."
         elif event_type == "embedding_progress":
-            # Scale embedding progress (35-95%)
             pct = data["current"] / data["total"] if data["total"] > 0 else 1
             progress = 35 + int(pct * 60)
             remaining = data.get("remaining_seconds")
@@ -155,40 +155,87 @@ def callback(event_type: str, data: dict):
             else:
                 eta_str = ""
             message = f"Embedding... ({data['current']}/{data['total']} chunks{eta_str})"
+        elif event_type == "changes_detected":
+            progress = 10
+            message = f"Found {data['new']} new, {data['modified']} modified files..."
         elif event_type == "complete":
-            progress = 100
-            message = "Indexing complete!"
+            progress = 99
+            message = "Finishing up..."
 
         if progress > 0:
-            asyncio.run_coroutine_threadsafe(ctx.report_progress(progress, 100, message), loop)
+            indexing_status.update(progress, message)
 
     return callback
 
 
-def _create_relearn_progress_callback(ctx: Context, loop) -> Callable:
-    """Create a progress callback that reports reindexing progress to MCP client."""
+def _run_full_index_bg(
+    codebase_path: Path,
+    paths: Dict[str, Path],
+    state,
+    file_extensions: Optional[List[str]],
+    embedding_model: str,
+):
+    """Run full indexing in a background thread. Updates state on completion."""
+    from codegrok_mcp.indexing.source_retriever import SourceRetriever
 
-    def callback(event_type: str, data: dict):
-        progress = 0
-        message = ""
+    try:
+        retriever = SourceRetriever(
+            codebase_path=str(codebase_path),
+            embedding_model=embedding_model,
+            verbose=False,
+            persist_path=str(paths["chroma_path"]),
+        )
 
-        if event_type == "changes_detected":
-            progress = 10
-            message = f"Found {data['new']} new, {data['modified']} modified files..."
-        elif event_type == "parsing_start":
-            progress = 20
-            message = f"Parsing {data['total']} changed files..."
-        elif event_type == "embedding_start":
-            progress = 40
-            message = f"Updating embeddings for {data['total']} chunks..."
-        elif event_type == "complete":
-            progress = 100
-            message = "Re-indexing complete!"
+        extensions = file_extensions if file_extensions else SUPPORTED_EXTENSIONS
+        progress_callback = _create_bg_progress_callback(state.indexing)
 
-        if progress > 0:
-            asyncio.run_coroutine_threadsafe(ctx.report_progress(progress, 100, message), loop)
+        retriever.index_codebase(file_extensions=extensions, progress_callback=progress_callback)
+        retriever.save_metadata(str(paths["metadata_path"]))
 
-    return callback
+        state.retriever = retriever
+        state.codebase_path = codebase_path
+
+        state.indexing.complete(
+            {"success": True, "mode_used": "full", "stats": retriever.get_stats()}
+        )
+    except Exception as e:
+        state.indexing.fail(str(e))
+
+
+def _run_incremental_reindex_bg(
+    codebase_path: Path,
+    paths: Dict[str, Path],
+    state,
+    embedding_model: str,
+):
+    """Run incremental reindex in a background thread."""
+    from codegrok_mcp.indexing.source_retriever import SourceRetriever
+
+    try:
+        retriever = SourceRetriever(
+            codebase_path=str(codebase_path),
+            embedding_model=embedding_model,
+            verbose=False,
+            persist_path=str(paths["chroma_path"]),
+        )
+
+        if not retriever.load_existing_index():
+            state.indexing.fail(f"Failed to load existing index from {paths['chroma_path']}")
+            return
+
+        retriever.load_metadata(str(paths["metadata_path"]))
+
+        progress_callback = _create_bg_progress_callback(state.indexing)
+        result = retriever.incremental_reindex(progress_callback=progress_callback)
+
+        retriever.save_metadata(str(paths["metadata_path"]))
+
+        state.retriever = retriever
+        state.codebase_path = codebase_path
+
+        state.indexing.complete({"success": True, "mode_used": "incremental", **result})
+    except Exception as e:
+        state.indexing.fail(str(e))
 
 
 @mcp.tool(
@@ -200,6 +247,9 @@ def callback(event_type: str, data: dict):
 - full: Force complete re-index (destroys existing index).
 - load_only: Just load existing index without any indexing.
 
+Indexing runs in the background — this tool returns immediately.
+Call get_stats() to check progress. Once complete, search tools become available.
+
 Creates a .codegrok/ folder in the codebase directory.""",
     annotations=ToolAnnotations(
         readOnlyHint=False,  # Creates/modifies .codegrok/ directory
@@ -233,9 +283,32 @@ async def learn(
     ] = None,
     ctx: Context = None,
 ) -> Dict[str, Any]:
-    """Index a codebase with smart mode detection."""
+    """Index a codebase with smart mode detection. Returns immediately; indexing runs in background."""
     state = get_state()
 
+    # If indexing is already in progress, return current status
+    if state.indexing.active:
+        return {
+            "success": True,
+            "status": "indexing_in_progress",
+            "message": "Indexing is already running. Call get_stats() to check progress.",
+            **state.indexing.to_dict(),
+        }
+
+    # If last indexing completed, return the result and clear it
+    if state.indexing.result is not None:
+        result = state.indexing.result.copy()
+        result["status"] = "complete"
+        result["message"] = "Indexing completed successfully."
+        state.indexing.result = None  # Clear so next call can re-index
+        return result
+
+    # If last indexing failed, return the error
+    if state.indexing.error is not None:
+        error = state.indexing.error
+        state.indexing.error = None  # Clear so next call can retry
+        raise ToolError(f"Previous indexing failed: {error}. Retrying...")
+
     # Validate mode
     valid_modes = ("auto", "full", "load_only")
     if mode not in valid_modes:
@@ -252,50 +325,51 @@ async def learn(
     paths = _get_codegrok_paths(codebase_path)
     has_existing = _has_valid_index(paths)
 
-    # Resolve timeout
-    timeout = _get_timeout(timeout_seconds)
-
-    # Handle load_only mode (no timeout needed — just loads metadata)
+    # Handle load_only mode (fast, no background needed)
     if mode == "load_only":
         if not has_existing:
             raise ToolError(
                 f"No existing index found at {codebase_path}. "
                 "Use mode='auto' or mode='full' to create one."
             )
-        return await _load_existing_index(codebase_path, paths, state, embedding_model)
+        return _load_existing_index_sync(codebase_path, paths, state, embedding_model)
 
-    # Handle auto mode with existing index -> incremental reindex
-    if mode == "auto" and has_existing:
-        try:
-            return await asyncio.wait_for(
-                _incremental_reindex(codebase_path, paths, state, embedding_model, ctx),
-                timeout=timeout,
-            )
-        except asyncio.TimeoutError:
-            raise ToolError(
-                f"Indexing timed out after {timeout}s. "
-                f"Set CODEGROK_TIMEOUT env var or pass timeout_seconds to increase. "
-                f"Checkpoint saved — re-run to resume."
-            )
+    # Start background indexing
+    paths["codegrok_dir"].mkdir(parents=True, exist_ok=True)
 
-    # Full index: mode == "full" OR (mode == "auto" and no existing index)
-    try:
-        return await asyncio.wait_for(
-            _full_index(codebase_path, paths, state, file_extensions, embedding_model, ctx),
-            timeout=timeout,
+    if mode == "auto" and has_existing:
+        state.indexing.start("Starting incremental reindex...")
+        thread = threading.Thread(
+            target=_run_incremental_reindex_bg,
+            args=(codebase_path, paths, state, embedding_model),
+            daemon=True,
         )
-    except asyncio.TimeoutError:
-        raise ToolError(
-            f"Indexing timed out after {timeout}s. "
-            f"Set CODEGROK_TIMEOUT env var or pass timeout_seconds to increase. "
-            f"Checkpoint saved — re-run to resume."
+    else:
+        state.indexing.start("Starting full index...")
+        thread = threading.Thread(
+            target=_run_full_index_bg,
+            args=(codebase_path, paths, state, file_extensions, embedding_model),
+            daemon=True,
         )
 
+    thread.start()
 
-async def _load_existing_index(
+    return {
+        "success": True,
+        "status": "indexing_started",
+        "message": (
+            f"Indexing started for {codebase_path.name}. "
+            "Call get_stats() to check progress. "
+            "Search tools will be available once indexing completes."
+        ),
+        **state.indexing.to_dict(),
+    }
+
+
+def _load_existing_index_sync(
     codebase_path: Path, paths: Dict[str, Path], state, embedding_model: str
 ) -> Dict[str, Any]:
-    """Load an existing index without any reindexing."""
+    """Load an existing index without any reindexing (synchronous, fast)."""
     from codegrok_mcp.indexing.source_retriever import SourceRetriever
 
     retriever = SourceRetriever(
@@ -317,6 +391,7 @@ async def _load_existing_index(
 
     return {
         "success": True,
+        "status": "complete",
         "mode_used": "load_only",
         "message": f"Loaded existing index for {codebase_path.name}",
         "stats": stats,
@@ -324,102 +399,6 @@ async def _load_existing_index(
     }
 
 
-async def _incremental_reindex(
-    codebase_path: Path, paths: Dict[str, Path], state, embedding_model: str, ctx: Context = None
-) -> Dict[str, Any]:
-    """Load existing index and perform incremental reindex."""
-    from codegrok_mcp.indexing.source_retriever import SourceRetriever
-
-    retriever = SourceRetriever(
-        codebase_path=str(codebase_path),
-        embedding_model=embedding_model,
-        verbose=False,
-        persist_path=str(paths["chroma_path"]),
-    )
-
-    if not retriever.load_existing_index():
-        raise ToolError(f"Failed to load existing index from {paths['chroma_path']}")
-
-    # Load metadata to get file mtimes for incremental detection
-    retriever.load_metadata(str(paths["metadata_path"]))
-
-    # Create progress callback if context available
-    progress_callback = None
-    if ctx:
-        loop = asyncio.get_event_loop()
-        progress_callback = _create_relearn_progress_callback(ctx, loop)
-
-    # Run blocking reindex in a thread so asyncio.wait_for can cancel it
-    result = await asyncio.to_thread(
-        retriever.incremental_reindex, progress_callback=progress_callback
-    )
-
-    # Save updated metadata
-    retriever.save_metadata(str(paths["metadata_path"]))
-
-    state.retriever = retriever
-    state.codebase_path = codebase_path
-
-    return {
-        "success": True,
-        "mode_used": "incremental",
-        "message": f"Incremental reindex complete for {codebase_path.name}",
-        **result,
-    }
-
-
-async def _full_index(
-    codebase_path: Path,
-    paths: Dict[str, Path],
-    state,
-    file_extensions: Optional[List[str]],
-    embedding_model: str,
-    ctx: Context = None,
-) -> Dict[str, Any]:
-    """Perform full index (creates or replaces existing index)."""
-    from codegrok_mcp.indexing.source_retriever import SourceRetriever
-
-    # Create .codegrok directory
-    paths["codegrok_dir"].mkdir(parents=True, exist_ok=True)
-
-    # Create progress callback if context available
-    progress_callback = None
-    if ctx:
-        loop = asyncio.get_event_loop()
-        progress_callback = _create_learn_progress_callback(ctx, loop)
-
-    retriever = SourceRetriever(
-        codebase_path=str(codebase_path),
-        embedding_model=embedding_model,
-        verbose=False,
-        persist_path=str(paths["chroma_path"]),
-    )
-
-    # Run blocking indexing in a thread so asyncio.wait_for can cancel it
-    extensions = file_extensions if file_extensions else SUPPORTED_EXTENSIONS
-    await asyncio.to_thread(
-        retriever.index_codebase, file_extensions=extensions, progress_callback=progress_callback
-    )
-
-    # Report saving phase
-    if ctx:
-        await ctx.report_progress(95, 100, "Saving index...")
-
-    # Save metadata
-    retriever.save_metadata(str(paths["metadata_path"]))
-
-    # Update state
-    state.retriever = retriever
-    state.codebase_path = codebase_path
-
-    return {
-        "success": True,
-        "mode_used": "full",
-        "message": f"Successfully indexed {codebase_path.name}",
-        "stats": retriever.get_stats(),
-    }
-
-
 @mcp.tool(
     name="get_sources",
     description="""Get source code references relevant to a query using semantic search. Requires 'learn' first.
@@ -472,9 +451,10 @@ def get_sources(
 
 @mcp.tool(
     name="get_stats",
-    description="""Get statistics about the currently loaded codebase index.
+    description="""Get statistics about the currently loaded codebase index and indexing progress.
 
-Returns: files indexed, total chunks, symbols by type, languages detected, index creation time.""",
+Returns: files indexed, total chunks, symbols by type, languages detected, index creation time.
+If indexing is in progress, also returns progress percentage and ETA.""",
     annotations=ToolAnnotations(
         readOnlyHint=True,  # Only reads metadata
         idempotentHint=True,  # Same state = same results
@@ -482,17 +462,26 @@ def get_sources(
     ),
 )
 def get_stats() -> Dict[str, Any]:
-    """Get indexing statistics."""
+    """Get indexing statistics and progress."""
     state = get_state()
 
-    if not state.is_loaded:
-        return {"loaded": False, "codebase_path": None, "stats": None}
+    result: Dict[str, Any] = {}
 
-    return {
-        "loaded": True,
-        "codebase_path": str(state.codebase_path),
-        "stats": state.retriever.get_stats(),
-    }
+    # Include indexing status if active or recently completed/failed
+    indexing_info = state.indexing.to_dict()
+    if indexing_info["active"] or indexing_info["error"]:
+        result["indexing"] = indexing_info
+
+    if not state.is_loaded:
+        result["loaded"] = False
+        result["codebase_path"] = None
+        result["stats"] = None
+        return result
+
+    result["loaded"] = True
+    result["codebase_path"] = str(state.codebase_path)
+    result["stats"] = state.retriever.get_stats()
+    return result
 
 
 @mcp.tool(
diff --git a/src/codegrok_mcp/mcp/state.py b/src/codegrok_mcp/mcp/state.py
index 74fb160..8227214 100644
--- a/src/codegrok_mcp/mcp/state.py
+++ b/src/codegrok_mcp/mcp/state.py
@@ -1,14 +1,62 @@
 """Session state management for MCP server."""
 
-from dataclasses import dataclass
+import threading
+from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Optional, TYPE_CHECKING
+from typing import Optional, Dict, Any, TYPE_CHECKING
 
 if TYPE_CHECKING:
     from codegrok_mcp.indexing.source_retriever import SourceRetriever
     from codegrok_mcp.indexing.memory_retriever import MemoryRetriever
 
 
+@dataclass
+class IndexingStatus:
+    """Thread-safe status of a background indexing operation."""
+
+    active: bool = False
+    progress: int = 0  # 0-100
+    message: str = ""
+    error: Optional[str] = None
+    result: Optional[Dict[str, Any]] = None
+    _lock: threading.Lock = field(default_factory=threading.Lock)
+
+    def start(self, message: str = "Starting indexing..."):
+        with self._lock:
+            self.active = True
+            self.progress = 0
+            self.message = message
+            self.error = None
+            self.result = None
+
+    def update(self, progress: int, message: str):
+        with self._lock:
+            self.progress = min(progress, 99)
+            self.message = message
+
+    def complete(self, result: Dict[str, Any]):
+        with self._lock:
+            self.active = False
+            self.progress = 100
+            self.message = "Indexing complete"
+            self.result = result
+
+    def fail(self, error: str):
+        with self._lock:
+            self.active = False
+            self.message = f"Indexing failed: {error}"
+            self.error = error
+
+    def to_dict(self) -> Dict[str, Any]:
+        with self._lock:
+            return {
+                "active": self.active,
+                "progress": self.progress,
+                "message": self.message,
+                "error": self.error,
+            }
+
+
 @dataclass
 class MCPSessionState:
     """Singleton state for MCP server session."""
@@ -16,6 +64,7 @@ class MCPSessionState:
     retriever: Optional["SourceRetriever"] = None
     memory_retriever: Optional["MemoryRetriever"] = None
     codebase_path: Optional[Path] = None
+    indexing: IndexingStatus = field(default_factory=IndexingStatus)
 
     @property
     def is_loaded(self) -> bool:
diff --git a/tests/unit/test_background_indexing.py b/tests/unit/test_background_indexing.py
new file mode 100644
index 0000000..0d09a98
--- /dev/null
+++ b/tests/unit/test_background_indexing.py
@@ -0,0 +1,396 @@
+"""Tests for background indexing and IndexingStatus."""
+
+import threading
+import time
+from unittest.mock import patch, MagicMock
+
+import pytest
+
+from codegrok_mcp.mcp.state import IndexingStatus, MCPSessionState, get_state, reset_state
+
+
+class TestIndexingStatus:
+    """Tests for the IndexingStatus thread-safe dataclass."""
+
+    def test_initial_state(self):
+        status = IndexingStatus()
+        assert status.active is False
+        assert status.progress == 0
+        assert status.message == ""
+        assert status.error is None
+        assert status.result is None
+
+    def test_start(self):
+        status = IndexingStatus()
+        status.start("Starting full index...")
+        assert status.active is True
+        assert status.progress == 0
+        assert status.message == "Starting full index..."
+        assert status.error is None
+        assert status.result is None
+
+    def test_start_clears_previous_error(self):
+        status = IndexingStatus()
+        status.fail("previous error")
+        status.start("Retrying...")
+        assert status.active is True
+        assert status.error is None
+        assert status.result is None
+
+    def test_start_clears_previous_result(self):
+        status = IndexingStatus()
+        status.complete({"success": True})
+        status.start("Re-indexing...")
+        assert status.active is True
+        assert status.result is None
+
+    def test_update(self):
+        status = IndexingStatus()
+        status.start()
+        status.update(50, "Halfway there...")
+        assert status.progress == 50
+        assert status.message == "Halfway there..."
+
+    def test_update_caps_at_99(self):
+        status = IndexingStatus()
+        status.start()
+        status.update(100, "Should cap")
+        assert status.progress == 99
+        status.update(500, "Way over")
+        assert status.progress == 99
+
+    def test_complete(self):
+        status = IndexingStatus()
+        status.start()
+        result = {"success": True, "stats": {"files": 10}}
+        status.complete(result)
+        assert status.active is False
+        assert status.progress == 100
+        assert status.message == "Indexing complete"
+        assert status.result == result
+
+    def test_fail(self):
+        status = IndexingStatus()
+        status.start()
+        status.fail("Out of memory")
+        assert status.active is False
+        assert status.message == "Indexing failed: Out of memory"
+        assert status.error == "Out of memory"
+
+    def test_to_dict(self):
+        status = IndexingStatus()
+        status.start("Testing...")
+        status.update(42, "Processing...")
+        d = status.to_dict()
+        assert d == {
+            "active": True,
+            "progress": 42,
+            "message": "Processing...",
+            "error": None,
+        }
+
+    def test_to_dict_excludes_result(self):
+        """to_dict should not include the result field (it's for internal use)."""
+        status = IndexingStatus()
+        status.complete({"big": "data"})
+        d = status.to_dict()
+        assert "result" not in d
+
+    def test_thread_safety(self):
+        """Multiple threads can update IndexingStatus without errors."""
+        status = IndexingStatus()
+        status.start()
+        errors = []
+
+        def updater(thread_id):
+            try:
+                for i in range(100):
+                    status.update(i, f"Thread {thread_id} at {i}")
+                    status.to_dict()
+            except Exception as e:
+                errors.append(e)
+
+        threads = [threading.Thread(target=updater, args=(i,)) for i in range(10)]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+
+        assert errors == []
+        assert status.active is True  # No one called complete/fail
+
+
+class TestMCPSessionStateIndexing:
+    """Tests for IndexingStatus integration in MCPSessionState."""
+
+    def test_state_has_indexing_status(self):
+        state = MCPSessionState()
+        assert isinstance(state.indexing, IndexingStatus)
+        assert state.indexing.active is False
+
+    def test_each_state_gets_own_indexing(self):
+        state1 = MCPSessionState()
+        state2 = MCPSessionState()
+        state1.indexing.start("s1")
+        assert state2.indexing.active is False
+
+    def test_singleton_state_has_indexing(self):
+        reset_state()
+        state = get_state()
+        assert isinstance(state.indexing, IndexingStatus)
+        reset_state()
+
+
+class TestBackgroundProgressCallback:
+    """Tests for _create_bg_progress_callback."""
+
+    def test_callback_updates_indexing_status(self):
+        from codegrok_mcp.mcp.server import _create_bg_progress_callback
+
+        status = IndexingStatus()
+        status.start()
+        cb = _create_bg_progress_callback(status)
+
+        cb("files_found", {"files": list(range(100))})
+        assert status.progress == 5
+        assert "100 files" in status.message
+
+    def test_callback_discovery_progress(self):
+        from codegrok_mcp.mcp.server import _create_bg_progress_callback
+
+        status = IndexingStatus()
+        status.start()
+        cb = _create_bg_progress_callback(status)
+
+        cb("discovery_progress", {"files_found": 5000})
+        assert status.progress <= 4
+        assert "5000" in status.message
+
+    def test_callback_parsing_start(self):
+        from codegrok_mcp.mcp.server import _create_bg_progress_callback
+
+        status = IndexingStatus()
+        status.start()
+        cb = _create_bg_progress_callback(status)
+
+        cb("parsing_start", {"total": 50})
+        assert status.progress == 10
+        assert "50 files" in status.message
+
+    def test_callback_embedding_progress_with_eta(self):
+        from codegrok_mcp.mcp.server import _create_bg_progress_callback
+
+        status = IndexingStatus()
+        status.start()
+        cb = _create_bg_progress_callback(status)
+
+        cb("embedding_progress", {"current": 500, "total": 1000, "remaining_seconds": 120})
+        assert status.progress > 35
+        assert "500/1000" in status.message
+        assert "remaining" in status.message
+
+    def test_callback_embedding_progress_without_eta(self):
+        from codegrok_mcp.mcp.server import _create_bg_progress_callback
+
+        status = IndexingStatus()
+        status.start()
+        cb = _create_bg_progress_callback(status)
+
+        cb("embedding_progress", {"current": 500, "total": 1000, "remaining_seconds": None})
+        assert "500/1000" in status.message
+        assert "remaining" not in status.message
+
+    def test_callback_changes_detected(self):
+        from codegrok_mcp.mcp.server import _create_bg_progress_callback
+
+        status = IndexingStatus()
+        status.start()
+        cb = _create_bg_progress_callback(status)
+
+        cb("changes_detected", {"new": 5, "modified": 3})
+        assert status.progress == 10
+        assert "5 new" in status.message
+        assert "3 modified" in status.message
+
+    def test_callback_complete(self):
+        from codegrok_mcp.mcp.server import _create_bg_progress_callback
+
+        status = IndexingStatus()
+        status.start()
+        cb = _create_bg_progress_callback(status)
+
+        cb("complete", {})
+        assert status.progress == 99
+
+    def test_unknown_event_ignored(self):
+        from codegrok_mcp.mcp.server import _create_bg_progress_callback
+
+        status = IndexingStatus()
+        status.start()
+        cb = _create_bg_progress_callback(status)
+
+        cb("unknown_event", {"foo": "bar"})
+        # Progress stays at 0 from start
+        assert status.progress == 0
+
+
+class TestLearnToolBehavior:
+    """Tests for learn tool's background indexing flow (without actual indexing)."""
+
+    def setup_method(self):
+        reset_state()
+
+    def teardown_method(self):
+        reset_state()
+
+    @pytest.mark.asyncio
+    async def test_learn_returns_in_progress_when_active(self):
+        """If indexing is active, learn should return status without starting another."""
+        from codegrok_mcp.mcp.server import learn
+
+        state = get_state()
+        state.indexing.start("Already running...")
+
+        result = await learn(path="/tmp", mode="auto")
+        assert result["status"] == "indexing_in_progress"
+        assert "already running" in result["message"].lower()
+
+    @pytest.mark.asyncio
+    async def test_learn_returns_completed_result(self):
+        """If last indexing completed, learn should return the result."""
+        from codegrok_mcp.mcp.server import learn
+
+        state = get_state()
+        state.indexing.complete({"success": True, "mode_used": "full", "stats": {"files": 10}})
+
+        result = await learn(path="/tmp", mode="auto")
+        assert result["status"] == "complete"
+        assert result["stats"] == {"files": 10}
+        # Result should be cleared after retrieval
+        assert state.indexing.result is None
+
+    @pytest.mark.asyncio
+    async def test_learn_raises_on_previous_error(self):
+        """If last indexing failed, learn should raise and clear the error."""
+        from codegrok_mcp.mcp.server import learn
+        from fastmcp.exceptions import ToolError
+
+        state = get_state()
+        state.indexing.error = "Disk full"
+
+        with pytest.raises(ToolError, match="Disk full"):
+            await learn(path="/tmp", mode="auto")
+        # Error should be cleared
+        assert state.indexing.error is None
+
+    @pytest.mark.asyncio
+    async def test_learn_invalid_mode(self):
+        from codegrok_mcp.mcp.server import learn
+        from fastmcp.exceptions import ToolError
+
+        with pytest.raises(ToolError, match="Invalid mode"):
+            await learn(path="/tmp", mode="bad_mode")
+
+    @pytest.mark.asyncio
+    async def test_learn_invalid_path(self):
+        from codegrok_mcp.mcp.server import learn
+        from fastmcp.exceptions import ToolError
+
+        with pytest.raises(ToolError, match="does not exist"):
+            await learn(path="/nonexistent/path/12345", mode="auto")
+
+    @pytest.mark.asyncio
+    async def test_learn_starts_background_thread(self, tmp_path):
+        """learn should start a daemon thread and return immediately."""
+        from codegrok_mcp.mcp.server import learn
+
+        # Create a minimal directory
+        (tmp_path / "test.py").write_text("x = 1")
+
+        with patch("codegrok_mcp.mcp.server.threading.Thread") as mock_thread_cls:
+            mock_thread = MagicMock()
+            mock_thread_cls.return_value = mock_thread
+
+            result = await learn(path=str(tmp_path), mode="full")
+
+            assert result["status"] == "indexing_started"
+            mock_thread_cls.assert_called_once()
+            assert mock_thread_cls.call_args[1]["daemon"] is True
+            mock_thread.start.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_learn_auto_with_existing_uses_incremental(self, tmp_path):
+        """auto mode with existing index should start incremental reindex thread."""
+        from codegrok_mcp.mcp.server import learn
+
+        # Create fake existing index
+        codegrok_dir = tmp_path / ".codegrok"
+        codegrok_dir.mkdir()
+        (codegrok_dir / "chroma").mkdir()
+        (codegrok_dir / "metadata.json").write_text("{}")
+
+        with patch("codegrok_mcp.mcp.server.threading.Thread") as mock_thread_cls:
+            mock_thread = MagicMock()
+            mock_thread_cls.return_value = mock_thread
+
+            result = await learn(path=str(tmp_path), mode="auto")
+
+            assert result["status"] == "indexing_started"
+            # Should use incremental reindex
+            call_args = mock_thread_cls.call_args
+            assert call_args[1]["target"].__name__ == "_run_incremental_reindex_bg"
+
+    @pytest.mark.asyncio
+    async def test_learn_full_mode_uses_full_index(self, tmp_path):
+        """full mode should always start full index thread."""
+        from codegrok_mcp.mcp.server import learn
+
+        (tmp_path / "test.py").write_text("x = 1")
+
+        with patch("codegrok_mcp.mcp.server.threading.Thread") as mock_thread_cls:
+            mock_thread = MagicMock()
+            mock_thread_cls.return_value = mock_thread
+
+            result = await learn(path=str(tmp_path), mode="full")
+
+            call_args = mock_thread_cls.call_args
+            assert call_args[1]["target"].__name__ == "_run_full_index_bg"
+
+
+class TestGetStatsIndexingInfo:
+    """Tests for get_stats returning indexing progress."""
+
+    def setup_method(self):
+        reset_state()
+
+    def teardown_method(self):
+        reset_state()
+
+    def test_get_stats_includes_indexing_when_active(self):
+        from codegrok_mcp.mcp.server import get_stats
+
+        state = get_state()
+        state.indexing.start("Indexing...")
+        state.indexing.update(42, "Processing...")
+
+        result = get_stats()
+        assert "indexing" in result
+        assert result["indexing"]["active"] is True
+        assert result["indexing"]["progress"] == 42
+
+    def test_get_stats_includes_indexing_on_error(self):
+        from codegrok_mcp.mcp.server import get_stats
+
+        state = get_state()
+        state.indexing.fail("Something broke")
+
+        result = get_stats()
+        assert "indexing" in result
+        assert result["indexing"]["error"] == "Something broke"
+
+    def test_get_stats_no_indexing_when_idle(self):
+        from codegrok_mcp.mcp.server import get_stats
+
+        result = get_stats()
+        assert "indexing" not in result
+        assert result["loaded"] is False