diff --git a/README.md b/README.md index 0488dc8..c087676 100644 --- a/README.md +++ b/README.md @@ -519,6 +519,7 @@ All other text files are chunked by line range. Binary files are skipped. | [How It Works](https://github.com/elara-labs/code-context-engine/blob/main/docs/wiki/How-It-Works.md) | Full 9-stage pipeline | | [CLI Reference](https://github.com/elara-labs/code-context-engine/blob/main/docs/wiki/CLI-Reference.md) | Every command with output | | [Configuration](https://github.com/elara-labs/code-context-engine/blob/main/docs/wiki/Configuration.md) | All config options | +| [CodeGraph Worktree Integration Discovery](docs/codegraph-worktree-integration-discovery.md) | Phase 0 discovery for shared base plus worktree overlay integration | --- diff --git a/docs/codegraph-worktree-integration-discovery.md b/docs/codegraph-worktree-integration-discovery.md new file mode 100644 index 0000000..93c24bb --- /dev/null +++ b/docs/codegraph-worktree-integration-discovery.md @@ -0,0 +1,190 @@ +[← Back to README](../README.md) + +# CodeGraph Worktree Integration Discovery + +## Overview + +This document records Phase 0 discovery for integrating Code Context Engine +(CCE) with CodeGraph using shared repository context plus per-worktree overlays. +It is evidence for the staged implementation plan, not production behavior. + +## Current CCE Storage Identity + +CCE currently keys project storage by the absolute project checkout path. + +File references: + +- `src/context_engine/utils.py` +- `src/context_engine/integration/mcp_server.py` +- `tests/test_project_storage_dir.py` + +Observed behavior: + +- `project_storage_dir(config, project_dir)` resolves `project_dir`, hashes that + absolute path, and stores data under `/-<6hex>`. +- `ContextEngineMCP.__init__` calls `project_storage_dir(config, Path.cwd())`. +- Existing migration only renames legacy basename-only storage to the + path-hashed slug. + +Consequence: + +- Two linked Git worktrees of the same repository get different complete CCE + storage directories. +- CCE does not yet model a shared repository identity plus a separate worktree + identity. + +## CodeGraph Worktree Mismatch Behavior + +CodeGraph already detects a dangerous shared-index case. + +File references: + +- `/home/joe/Webstorm_Projects/codegraph/src/sync/worktree.ts` +- `/home/joe/Webstorm_Projects/codegraph/__tests__/worktree-detection.test.ts` +- `/home/joe/Webstorm_Projects/codegraph/src/directory.ts` + +Observed behavior: + +- CodeGraph resolves an index by walking upward to the nearest `.codegraph/` + directory. +- `gitWorktreeRoot(dir)` runs `git rev-parse --show-toplevel`. +- `gitCommonDir(dir)` runs `git rev-parse --git-common-dir`. +- `detectWorktreeIndexMismatch(startPath, indexRoot)` warns when a command runs + inside one Git worktree but uses another worktree's `.codegraph/` index from + the same Git common directory. +- The warning says results may reflect another branch and suggests + `codegraph init -i` for a worktree-local index. + +Implication: + +- CCE must not symlink or blindly reuse another worktree's mutable + `.codegraph/` directory. +- Shared base reuse is only safe when CCE explicitly records the base revision + and overlays worktree changes. + +## CodeGraph Machine-Readable Interface + +Best stable interface order for the MVP: + +1. CodeGraph package API when available to the CCE process through a Node + subprocess wrapper. +2. Structured CodeGraph CLI JSON output. +3. Human CLI output only for `explore`, as a temporary read-only fallback. + +Evidence: + +- CodeGraph README documents package API usage: + `CodeGraph.init`, `CodeGraph.open`, `searchNodes`, `getCallers`, + `buildContext`, and `getImpactRadius`. +- `src/index.ts` exports `CodeGraph`, `getDatabasePath`, `DatabaseConnection`, + `QueryBuilder`, `findNearestCodeGraphRoot`, and related types. +- `codegraph query`, `codegraph files`, `codegraph callers`, + `codegraph callees`, `codegraph impact`, and `codegraph affected` expose + `--json`. +- `codegraph explore` is the primary MCP-equivalent high-level tool, but the + CLI path is Markdown text rather than JSON in the inspected source. + +Constraints: + +- CCE must invoke CodeGraph without `shell=True`. +- CCE must use bounded timeouts and stdout limits. +- Any Node wrapper must be treated as a boundary process, not imported into the + Python runtime directly. + +## CodeGraph Index Identity And Freshness Signals + +File references: + +- `/home/joe/Webstorm_Projects/codegraph/__tests__/status-json.test.ts` +- `/home/joe/Webstorm_Projects/codegraph/src/directory.ts` + +Observed behavior: + +- `codegraph status --json` exposes `initialized`, `version`, `indexPath`, + `lastIndexed`, and an `index.state` value. +- Tests cover `index.state == "complete"` after clean full index and + `"indexing"` for interrupted index work. +- CodeGraph uses `.codegraph/codegraph.db` as the SQLite index. +- `CODEGRAPH_DIR` can point one checkout to a different local index directory, + but it remains a per-project-root directory name, not a shared base overlay + model. + +MVP freshness use: + +- Treat `version`, `indexPath`, `lastIndexed`, and `index.state` as status + evidence. +- Treat non-complete or missing status as degraded provider state. +- Keep CCE overlay freshness independent from CodeGraph base freshness. + +## Stable Symbol Enumeration + +CodeGraph exposes enough read APIs for shared-base exploration: + +- `searchNodes(query, options)` for symbol search. +- `getCallers(node_id)` and `getCallees(node_id)` for relationships. +- `getImpactRadius(node_id, depth)` for blast radius. +- `files --json` for indexed file inventory. + +Limitations: + +- CodeGraph node IDs are database-local. They must not be used as stable + identities across a shared base and a worktree overlay. +- CCE overlay merge should use logical symbol identity: qualified name, kind, + path, and signature where available. + +## Reusable Base Index Safety + +Safe reuse: + +- Query a CodeGraph index only as the shared base for the base revision it + represents. +- Store CCE metadata that ties that base to repository identity, base SHA, + CodeGraph index path, CodeGraph version, and freshness status. + +Unsafe reuse: + +- Do not point every worktree at another checkout's mutable `.codegraph/`. +- Do not return CodeGraph base source for a file known modified or deleted in + the worktree overlay. + +## Selected Git Base-Ref Strategy + +Use this deterministic hierarchy: + +1. Explicit configured base SHA/ref. +2. Upstream branch merge-base. +3. Unambiguous `origin/main`, `origin/master`, `main`, or `master` merge-base. +4. Current `HEAD` for dirty-only worktrees. +5. No base SHA when no safe base can be established. + +Rules: + +- Never silently compare against an unrelated branch. +- Include staged, unstaged, committed branch delta, untracked, deleted, and + renamed files in the diff model. +- Treat renames as old-path tombstone plus new-path addition for the MVP. + +## Limitations Requiring CCE Overlay Analysis + +CCE must own overlay semantics because CodeGraph has no inspected stable API for +querying a shared base plus a separate worktree overlay. + +Required CCE responsibilities: + +- Worktree identity detection. +- Worktree diff detection. +- Overlay semantic chunks for changed and added files. +- Tombstones for deleted and renamed-away files. +- Modified-path shadowing so stale base chunks do not leak into answers. +- Logical-symbol merge rules for structural overlay results. + +## Stage 1 Implementation Notes + +Recommended first production branch: + +- Add a CCE Git repository context module. +- Add repository/worktree IDs based on realpath-normalized Git common directory + and worktree root. +- Add base SHA resolution using the selected hierarchy. +- Add a diff model using Git plumbing. +- Add unit tests with real temporary Git repositories and linked worktrees. diff --git a/docs/wiki/Configuration.md b/docs/wiki/Configuration.md index 1c9d4e7..42d2fe2 100644 --- a/docs/wiki/Configuration.md +++ b/docs/wiki/Configuration.md @@ -35,6 +35,10 @@ retrieval: marginal_ratio: 0.75 # Stop adding results once score drops below this fraction of the # top score. 0 disables (always fill to top_k). Default 0.75. +structural: + provider: off # off | codegraph + codegraph_executable: codegraph + embedding: model: BAAI/bge-small-en-v1.5 # Embedding model (fastembed-compatible) @@ -137,6 +141,22 @@ At runtime, Claude can pass `top_k` and `max_tokens` directly to `context_search context_search(query="payment processing", top_k=5, max_tokens=3000) ``` +## Structural Context + +Enable CodeGraph structural context in global or project config: + +```yaml +structural: + provider: codegraph + codegraph_executable: codegraph +``` + +When enabled, `context_search` adds CodeGraph-derived structural sources, +relationships, and impact items to the normal semantic chunks. All structural +items include file/line provenance when CodeGraph provides it. If CodeGraph is +missing, uninitialized, indexing, or degraded, `context_search` falls back to +semantic-only retrieval. + --- ## Ignoring Files diff --git a/src/context_engine/cli.py b/src/context_engine/cli.py index abd61a4..ef69094 100644 --- a/src/context_engine/cli.py +++ b/src/context_engine/cli.py @@ -1282,6 +1282,79 @@ def status(ctx: click.Context, output_json: bool, oneline: bool) -> None: animate(lines) +@main.command() +@click.option("--json", "output_json", is_flag=True, help="Output as JSON") +@click.pass_context +def doctor(ctx: click.Context, output_json: bool) -> None: + """Show worktree-aware repository diagnostics.""" + from context_engine.git.doctor import worktree_doctor_report + + config = ctx.obj["config"] + report = worktree_doctor_report(config, _safe_cwd()) + + if output_json: + click.echo(json.dumps(report, indent=2)) + return + + if not report.get("git", {}).get("available"): + click.echo("Git: unavailable") + return + + repo = report["repository"] + worktree = report["worktree"] + storage = report["storage"] + overlay = report["overlay"] + click.echo( + "\n".join( + [ + "Repository", + f" Common git dir: {repo['common_dir']}", + f" Repository ID: {repo['id']}", + "Worktree", + f" Root: {worktree['root']}", + f" Worktree ID: {worktree['id']}", + f" HEAD: {worktree['head_sha']}", + f" Base: {worktree['base_sha']}", + "Storage", + f" Shared base: {storage['base_dir']}", + f" Overlay: {storage['worktree_dir']}", + "Overlay", + f" Modified: {overlay['modified_count']}", + f" Added: {overlay['added_count']}", + f" Deleted: {overlay['deleted_count']}", + ] + ) + ) + + +@main.command(name="worktree-benchmark") +@click.option("--json", "output_json", is_flag=True, help="Output as JSON") +def worktree_benchmark(output_json: bool) -> None: + """Measure worktree overlay discovery cost.""" + from context_engine.git.benchmark import benchmark_worktree_overlay + + result = benchmark_worktree_overlay(_safe_cwd()) + if result is None: + raise click.ClickException("Not inside a Git worktree") + data = result.to_dict() + if output_json: + click.echo(json.dumps(data, indent=2)) + return + click.echo( + "\n".join( + [ + "Worktree Overlay Benchmark", + f" Changed files: {data['changed_file_count']}", + f" Modified: {data['modified_count']}", + f" Added: {data['added_count']}", + f" Deleted: {data['deleted_count']}", + f" Renamed: {data['renamed_count']}", + f" Elapsed: {data['elapsed_ms']:.2f} ms", + ] + ) + ) + + @main.command("list") def list_commands() -> None: """Show all available CCE commands with usage examples.""" @@ -3379,6 +3452,7 @@ async def _run_index( """Run indexing pipeline (thin wrapper over `indexer.pipeline.run_indexing`).""" from context_engine.indexer.pipeline import run_indexing + storage_base = _runtime_storage_base(config, Path(project_dir)) log_fn = (lambda msg: click.echo(msg)) if verbose else None from context_engine.cli_style import warn, dim, CHECK, CROSS @@ -3431,6 +3505,7 @@ def phase_fn(msg: str) -> None: result = await run_indexing( config, project_dir, full=full, target_path=target_path, + storage_base_override=storage_base, log_fn=log_fn, progress_fn=progress_fn, embed_progress_fn=embed_progress_fn, phase_fn=phase_fn, ) @@ -3461,7 +3536,7 @@ def phase_fn(msg: str) -> None: ) # Update full_file_tokens baseline so cce savings shows codebase size - _storage_dir = project_storage_dir(config, Path(project_dir)) + _storage_dir = storage_base stats_path = _storage_dir / "stats.json" try: stats = json.loads(stats_path.read_text(encoding="utf-8")) if stats_path.exists() else {} @@ -3483,6 +3558,58 @@ def phase_fn(msg: str) -> None: stats_path.write_text(json.dumps(stats), encoding="utf-8") +def _runtime_storage_base(config, project_dir: Path) -> Path: + from context_engine.git import ( + repository_storage_layout, + resolve_git_repository_context, + ) + + context = resolve_git_repository_context(project_dir) + if context is None: + return project_storage_dir(config, project_dir) + + layout = repository_storage_layout(config, project_dir, context) + return layout.base_dir + + +def _runtime_storage_backend(config, project_dir: Path): + from context_engine.git import ( + get_worktree_diff, + repository_storage_layout, + resolve_git_repository_context, + ) + from context_engine.storage.local_backend import LocalBackend + from context_engine.storage.worktree_overlay import WorktreeOverlayBackend + + context = resolve_git_repository_context(project_dir) + if context is None: + storage_base = project_storage_dir(config, project_dir) + return storage_base, LocalBackend(base_path=str(storage_base)), None + + layout = repository_storage_layout(config, project_dir, context) + diff = get_worktree_diff( + context.worktree_root, + base_sha=context.base_sha, + head_sha=context.head_sha, + ) + backend = WorktreeOverlayBackend( + base=LocalBackend(base_path=str(layout.base_dir)), + overlay=LocalBackend(base_path=str(layout.worktree_dir)), + diff=diff, + ) + + def refresh_diff() -> None: + backend.update_diff( + get_worktree_diff( + context.worktree_root, + base_sha=context.base_sha, + head_sha=context.head_sha, + ) + ) + + return layout.worktree_dir, backend, refresh_diff + + async def _run_serve(config) -> None: """Start MCP server with live file watcher.""" import logging @@ -3504,7 +3631,6 @@ async def _run_serve(config) -> None: ) cap_ort_threads(max_threads=getattr(config, "serve_max_ort_threads", None)) - from context_engine.storage.local_backend import LocalBackend from context_engine.indexer.embedder import Embedder from context_engine.retrieval.retriever import HybridRetriever from context_engine.compression.compressor import Compressor @@ -3516,8 +3642,7 @@ async def _run_serve(config) -> None: project_dir = str(_safe_cwd()) project_name = _safe_cwd().name - storage_base = project_storage_dir(config, _safe_cwd()) - backend = LocalBackend(base_path=str(storage_base)) + storage_base, backend, refresh_diff = _runtime_storage_backend(config, _safe_cwd()) embedder = Embedder(model_name=config.embedding_model) retriever = HybridRetriever(backend=backend, embedder=embedder) compressor = Compressor( @@ -3527,7 +3652,7 @@ async def _run_serve(config) -> None: ) mcp = ContextEngineMCP( retriever=retriever, backend=backend, compressor=compressor, - embedder=embedder, config=config, + embedder=embedder, config=config, index_storage_base=storage_base, ) # Idle tracker — shuts down the server after prolonged inactivity (#139). @@ -3539,9 +3664,28 @@ async def _run_serve(config) -> None: # all indexing the same project simultaneously (#139). index_lock = ProjectIndexLock(storage_base) - chunk_count = backend._vector_store.count() + count_chunks = getattr(backend, "count_chunks", None) + chunk_count = ( + count_chunks() + if count_chunks is not None + else backend._vector_store.count() + ) import sys + needs_overlay_index = getattr(backend, "needs_overlay_index", None) + if needs_overlay_index is not None and needs_overlay_index(): + try: + await run_indexing( + config, + project_dir, + full=False, + storage_base_override=storage_base, + ) + if refresh_diff is not None: + refresh_diff() + except Exception as exc: + _log.warning("Startup worktree overlay indexing failed: %s", exc) + watcher = None worker_task = None @@ -3599,7 +3743,14 @@ async def _reindex_worker(): await _reindex_queue.put(rel) continue try: - await run_indexing(config, project_dir, target_path=rel) + await run_indexing( + config, + project_dir, + target_path=rel, + storage_base_override=storage_base, + ) + if refresh_diff is not None: + refresh_diff() _log.debug("Re-indexed: %s", rel) except Exception as exc: _log.warning("Watch re-index failed for %s: %s", rel, exc) diff --git a/src/context_engine/config.py b/src/context_engine/config.py index 63bd92e..a82ea34 100644 --- a/src/context_engine/config.py +++ b/src/context_engine/config.py @@ -74,6 +74,8 @@ class Config: # token-served reduction with no hit-rate loss. retrieval_marginal_ratio: float = 0.75 bootstrap_max_tokens: int = 10000 + structural_provider: str = "off" # off | codegraph + structural_codegraph_executable: str = "codegraph" # Indexer indexer_watch: bool = True @@ -143,6 +145,8 @@ def _deep_merge(base: dict, override: dict) -> dict: "retrieval_top_k": int, "retrieval_marginal_ratio": (int, float), "bootstrap_max_tokens": int, + "structural_provider": str, + "structural_codegraph_executable": str, "indexer_watch": bool, "indexer_debounce_ms": int, "indexer_ignore": list, @@ -170,6 +174,8 @@ def _apply_dict_to_config(config: Config, data: dict) -> None: ("retrieval", "top_k"): "retrieval_top_k", ("retrieval", "marginal_ratio"): "retrieval_marginal_ratio", ("retrieval", "bootstrap_max_tokens"): "bootstrap_max_tokens", + ("structural", "provider"): "structural_provider", + ("structural", "codegraph_executable"): "structural_codegraph_executable", ("serve", "idle_timeout_minutes"): "serve_idle_timeout_minutes", ("serve", "max_ort_threads"): "serve_max_ort_threads", ("indexer", "watch"): "indexer_watch", diff --git a/src/context_engine/data/instructions.md b/src/context_engine/data/instructions.md index 90d226b..5375506 100644 --- a/src/context_engine/data/instructions.md +++ b/src/context_engine/data/instructions.md @@ -10,6 +10,16 @@ the codebase, answering questions about code, or understanding how things work. `context_search` returns the most relevant code chunks with confidence scores instead of whole files. +When configured with CodeGraph, `context_search` may combine a shared +repository base with the current Git worktree overlay. Prefer its returned +worktree-aware sources; read files directly only when exact source is needed +before editing. + +If `structural.provider: codegraph` is enabled, `context_search` may include +structural sources, relationships, and impact items after the semantic chunks. +Treat file/line references in that section as provenance, and prefer worktree +overlay entries over base CodeGraph entries. + When to use `context_search`: - Answering questions about the codebase ("how does X work?", "where is Y?") - Exploring structure or architecture diff --git a/src/context_engine/data/tools_reference.md b/src/context_engine/data/tools_reference.md index 1ee6c1a..5a2af99 100644 --- a/src/context_engine/data/tools_reference.md +++ b/src/context_engine/data/tools_reference.md @@ -13,8 +13,11 @@ Search the codebase using hybrid vector + BM25 retrieval. | top_k | integer | no | 10 | Maximum results to return | | max_tokens | integer | no | 8000 | Token budget for results | -Returns ranked code chunks with confidence scores. Use this instead of -Read, Grep, or Glob when exploring code. +Returns ranked code chunks with confidence scores. When +`structural.provider: codegraph` is enabled, the same response may also include +structural sources, relationships, and impact items with exact file/line +provenance. If CodeGraph is unavailable, the tool returns semantic results +only. Use this instead of Read, Grep, or Glob when exploring code. ## expand_chunk diff --git a/src/context_engine/git/__init__.py b/src/context_engine/git/__init__.py new file mode 100644 index 0000000..9bfc197 --- /dev/null +++ b/src/context_engine/git/__init__.py @@ -0,0 +1,19 @@ +"""Git repository and worktree context helpers.""" + +from context_engine.git.benchmark import WorktreeBenchmarkResult, benchmark_worktree_overlay +from context_engine.git.diff import GitWorktreeDiff, get_worktree_diff +from context_engine.git.doctor import worktree_doctor_report +from context_engine.git.repository import GitRepositoryContext, resolve_git_repository_context +from context_engine.git.storage import RepositoryStorageLayout, repository_storage_layout + +__all__ = [ + "GitRepositoryContext", + "GitWorktreeDiff", + "RepositoryStorageLayout", + "WorktreeBenchmarkResult", + "benchmark_worktree_overlay", + "get_worktree_diff", + "repository_storage_layout", + "resolve_git_repository_context", + "worktree_doctor_report", +] diff --git a/src/context_engine/git/benchmark.py b/src/context_engine/git/benchmark.py new file mode 100644 index 0000000..296aa2f --- /dev/null +++ b/src/context_engine/git/benchmark.py @@ -0,0 +1,56 @@ +"""Worktree overlay benchmark helpers.""" + +from __future__ import annotations + +from dataclasses import dataclass, asdict +from pathlib import Path +import time +from typing import Any + +from context_engine.git.diff import get_worktree_diff +from context_engine.git.repository import resolve_git_repository_context + + +@dataclass(frozen=True) +class WorktreeBenchmarkResult: + repository_id: str + worktree_id: str + base_sha: str | None + head_sha: str + changed_file_count: int + modified_count: int + added_count: int + deleted_count: int + renamed_count: int + elapsed_ms: float + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +def benchmark_worktree_overlay(project_dir: Path) -> WorktreeBenchmarkResult | None: + """Measure O(diff) overlay discovery work for one worktree.""" + context = resolve_git_repository_context(project_dir) + if context is None: + return None + + start = time.perf_counter() + diff = get_worktree_diff( + context.worktree_root, + base_sha=context.base_sha, + head_sha=context.head_sha, + ) + elapsed_ms = (time.perf_counter() - start) * 1000 + changed = diff.modified | diff.added | diff.deleted | set(diff.renamed) + return WorktreeBenchmarkResult( + repository_id=context.repository_id, + worktree_id=context.worktree_id, + base_sha=context.base_sha, + head_sha=context.head_sha, + changed_file_count=len(changed), + modified_count=len(diff.modified), + added_count=len(diff.added), + deleted_count=len(diff.deleted), + renamed_count=len(diff.renamed), + elapsed_ms=elapsed_ms, + ) diff --git a/src/context_engine/git/diff.py b/src/context_engine/git/diff.py new file mode 100644 index 0000000..5904717 --- /dev/null +++ b/src/context_engine/git/diff.py @@ -0,0 +1,100 @@ +"""Git worktree diff model for overlay indexing.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path + +from context_engine.git.repository import _git + + +@dataclass(frozen=True) +class GitWorktreeDiff: + modified: set[str] = field(default_factory=set) + added: set[str] = field(default_factory=set) + deleted: set[str] = field(default_factory=set) + renamed: dict[str, str] = field(default_factory=dict) + base_sha: str | None = None + head_sha: str | None = None + + +def get_worktree_diff( + worktree_root: Path, + *, + base_sha: str | None, + head_sha: str | None, +) -> GitWorktreeDiff: + """Return committed, staged, unstaged, and untracked worktree changes.""" + builder = _DiffBuilder(base_sha=base_sha, head_sha=head_sha) + + if base_sha and head_sha and base_sha != head_sha: + builder.apply_name_status(_git(worktree_root, ["diff", "--name-status", f"{base_sha}...HEAD"])) + + builder.apply_name_status(_git(worktree_root, ["diff", "--cached", "--name-status"])) + builder.apply_name_status(_git(worktree_root, ["diff", "--name-status"])) + builder.apply_untracked(_git(worktree_root, ["ls-files", "--others", "--exclude-standard"])) + + return builder.build() + + +@dataclass +class _DiffBuilder: + modified: set[str] = field(default_factory=set) + added: set[str] = field(default_factory=set) + deleted: set[str] = field(default_factory=set) + renamed: dict[str, str] = field(default_factory=dict) + base_sha: str | None = None + head_sha: str | None = None + + def apply_name_status(self, output: str | None) -> None: + if not output: + return + for line in output.splitlines(): + parts = line.split("\t") + if not parts: + continue + status = parts[0] + if status.startswith("R") and len(parts) >= 3: + self._rename(parts[1], parts[2]) + elif status.startswith("A") and len(parts) >= 2: + self._add(parts[1]) + elif status.startswith("D") and len(parts) >= 2: + self._delete(parts[1]) + elif len(parts) >= 2: + self._modify(parts[1]) + + def apply_untracked(self, output: str | None) -> None: + if not output: + return + for path in output.splitlines(): + if path: + self._add(path) + + def build(self) -> GitWorktreeDiff: + return GitWorktreeDiff( + modified=set(self.modified), + added=set(self.added), + deleted=set(self.deleted), + renamed=dict(self.renamed), + base_sha=self.base_sha, + head_sha=self.head_sha, + ) + + def _add(self, path: str) -> None: + self.deleted.discard(path) + self.modified.discard(path) + self.added.add(path) + + def _delete(self, path: str) -> None: + self.added.discard(path) + self.modified.discard(path) + self.deleted.add(path) + + def _modify(self, path: str) -> None: + if path not in self.added and path not in self.deleted: + self.modified.add(path) + + def _rename(self, old_path: str, new_path: str) -> None: + self.renamed[old_path] = new_path + self._delete(old_path) + self._add(new_path) diff --git a/src/context_engine/git/doctor.py b/src/context_engine/git/doctor.py new file mode 100644 index 0000000..3186f12 --- /dev/null +++ b/src/context_engine/git/doctor.py @@ -0,0 +1,54 @@ +"""Worktree-aware repository diagnostics.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from context_engine.git.diff import get_worktree_diff +from context_engine.git.repository import resolve_git_repository_context +from context_engine.git.storage import repository_storage_layout + + +def worktree_doctor_report(config: object, project_dir: Path) -> dict[str, Any]: + """Return repository/worktree/storage diagnostics for current project.""" + context = resolve_git_repository_context(project_dir) + if context is None: + return {"git": {"available": False, "project_dir": str(project_dir.resolve())}} + + layout = repository_storage_layout(config, project_dir, context, migrate_legacy=False) + diff = get_worktree_diff( + context.worktree_root, + base_sha=context.base_sha, + head_sha=context.head_sha, + ) + return { + "git": {"available": True}, + "repository": { + "id": context.repository_id, + "common_dir": str(context.git_common_dir), + }, + "worktree": { + "id": context.worktree_id, + "root": str(context.worktree_root), + "head_sha": context.head_sha, + "base_sha": context.base_sha, + }, + "storage": { + "repository_root": str(layout.repository_root), + "base_dir": str(layout.base_dir), + "worktree_dir": str(layout.worktree_dir), + "legacy_project_dir": str(layout.legacy_project_dir), + }, + "overlay": { + "modified": sorted(diff.modified), + "added": sorted(diff.added), + "deleted": sorted(diff.deleted), + "renamed": dict(sorted(diff.renamed.items())), + "base_sha": diff.base_sha, + "head_sha": diff.head_sha, + "modified_count": len(diff.modified), + "added_count": len(diff.added), + "deleted_count": len(diff.deleted), + }, + } diff --git a/src/context_engine/git/repository.py b/src/context_engine/git/repository.py new file mode 100644 index 0000000..8dfcb3e --- /dev/null +++ b/src/context_engine/git/repository.py @@ -0,0 +1,135 @@ +"""Git repository identity and base-revision resolution.""" + +from __future__ import annotations + +from dataclasses import dataclass +import hashlib +import subprocess +from pathlib import Path + + +GIT_TIMEOUT_SECONDS = 5 +DEFAULT_BASE_REFS = ("origin/main", "origin/master", "main", "master") + + +@dataclass(frozen=True) +class GitRepositoryContext: + repository_id: str + git_common_dir: Path + worktree_id: str + worktree_root: Path + head_sha: str + base_sha: str | None + + +def resolve_git_repository_context( + project_dir: Path, + *, + base_ref: str | None = None, + fallback_base_refs: tuple[str, ...] = DEFAULT_BASE_REFS, +) -> GitRepositoryContext | None: + """Return Git repository/worktree identity, or None outside a usable Git repo.""" + worktree_root = _git_path(project_dir, ["rev-parse", "--show-toplevel"]) + git_common_dir = _git_path(project_dir, ["rev-parse", "--git-common-dir"]) + head_sha = _git(project_dir, ["rev-parse", "--verify", "HEAD"]) + if worktree_root is None or git_common_dir is None or head_sha is None: + return None + + return GitRepositoryContext( + repository_id=_path_id(git_common_dir), + git_common_dir=git_common_dir, + worktree_id=_path_id(worktree_root), + worktree_root=worktree_root, + head_sha=head_sha, + base_sha=resolve_base_sha( + worktree_root, + head_sha=head_sha, + base_ref=base_ref, + fallback_base_refs=fallback_base_refs, + ), + ) + + +def resolve_base_sha( + worktree_root: Path, + *, + head_sha: str | None = None, + base_ref: str | None = None, + fallback_base_refs: tuple[str, ...] = DEFAULT_BASE_REFS, +) -> str | None: + """Resolve the safest base SHA for a worktree diff.""" + head_sha = head_sha or _git(worktree_root, ["rev-parse", "--verify", "HEAD"]) + if head_sha is None: + return None + + if base_ref: + return _merge_base(worktree_root, base_ref) + + remote_default_ref = _remote_default_ref(worktree_root) + if remote_default_ref: + remote_default_base = _merge_base(worktree_root, remote_default_ref) + if remote_default_base: + return remote_default_base + + fallback_ref = _unambiguous_base_ref(worktree_root, fallback_base_refs) + if fallback_ref: + fallback_base = _merge_base(worktree_root, fallback_ref) + if fallback_base: + return fallback_base + + return head_sha + + +def _merge_base(worktree_root: Path, ref: str) -> str | None: + if _git(worktree_root, ["rev-parse", "--verify", f"{ref}^{{commit}}"]) is None: + return None + return _git(worktree_root, ["merge-base", "HEAD", ref]) + + +def _unambiguous_base_ref(worktree_root: Path, refs: tuple[str, ...]) -> str | None: + found: dict[str, str] = {} + for ref in refs: + sha = _git(worktree_root, ["rev-parse", "--verify", f"{ref}^{{commit}}"]) + if sha: + found[ref] = sha + if not found: + return None + if len(set(found.values())) != 1: + return None + return next(iter(found)) + + +def _remote_default_ref(worktree_root: Path) -> str | None: + return _git(worktree_root, ["symbolic-ref", "--quiet", "--short", "refs/remotes/origin/HEAD"]) + + +def _git_path(cwd: Path, args: list[str]) -> Path | None: + value = _git(cwd, args) + if not value: + return None + path = Path(value) + if not path.is_absolute(): + path = cwd / path + return path.resolve() + + +def _git(cwd: Path, args: list[str]) -> str | None: + try: + result = subprocess.run( + ["git", *args], + cwd=cwd, + capture_output=True, + text=True, + timeout=GIT_TIMEOUT_SECONDS, + check=False, + ) + except (FileNotFoundError, subprocess.TimeoutExpired, OSError): + return None + if result.returncode != 0: + return None + value = result.stdout.strip() + return value or None + + +def _path_id(path: Path) -> str: + return hashlib.sha256(str(path.resolve()).encode("utf-8")).hexdigest() diff --git a/src/context_engine/git/storage.py b/src/context_engine/git/storage.py new file mode 100644 index 0000000..d9fe2c7 --- /dev/null +++ b/src/context_engine/git/storage.py @@ -0,0 +1,44 @@ +"""Repository-scoped storage layout for shared base plus worktree overlays.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +from context_engine.git.repository import GitRepositoryContext +from context_engine.utils import project_storage_dir, resolve_project_storage_dir + + +@dataclass(frozen=True) +class RepositoryStorageLayout: + repository_root: Path + base_dir: Path + worktree_dir: Path + legacy_project_dir: Path + + +def repository_storage_layout( + config: object, + project_dir: Path, + git_context: GitRepositoryContext, + *, + migrate_legacy: bool = True, +) -> RepositoryStorageLayout: + """Return shared repository storage plus this worktree's overlay namespace.""" + repository_root = _repository_storage_root(config) / git_context.repository_id + legacy_project_dir = ( + project_storage_dir(config, project_dir) + if migrate_legacy + else resolve_project_storage_dir(config, project_dir) + ) + return RepositoryStorageLayout( + repository_root=repository_root, + base_dir=repository_root / "base", + worktree_dir=repository_root / "worktrees" / git_context.worktree_id, + legacy_project_dir=legacy_project_dir, + ) + + +def _repository_storage_root(config: object) -> Path: + project_storage_root = Path(config.storage_path) # type: ignore[union-attr] + return project_storage_root.parent / "repos" diff --git a/src/context_engine/indexer/pipeline.py b/src/context_engine/indexer/pipeline.py index 5a78e31..5de16c0 100644 --- a/src/context_engine/indexer/pipeline.py +++ b/src/context_engine/indexer/pipeline.py @@ -288,6 +288,7 @@ async def run_indexing( *, full: bool = False, target_path: str | None = None, + storage_base_override: str | Path | None = None, log_fn=None, progress_fn=None, embed_progress_fn=None, @@ -314,7 +315,11 @@ async def run_indexing( # makes `file_path.relative_to(project_dir)` raise ValueError on every # watcher-triggered reindex. project_dir = Path(project_dir).resolve() - storage_base = project_storage_dir(config, project_dir) + storage_base = ( + Path(storage_base_override) + if storage_base_override is not None + else project_storage_dir(config, project_dir) + ) storage_base.mkdir(parents=True, exist_ok=True) async with _pipeline_lock(str(storage_base)): diff --git a/src/context_engine/integration/mcp_server.py b/src/context_engine/integration/mcp_server.py index 7bf8dc1..4af2e94 100644 --- a/src/context_engine/integration/mcp_server.py +++ b/src/context_engine/integration/mcp_server.py @@ -12,6 +12,7 @@ from mcp.server import Server from mcp.types import Tool, TextContent +from context_engine.git.diff import GitWorktreeDiff from context_engine.compression.output_rules import ( ADVERTISED_PCT, ESTIMATED_AVG_REPLY_TOKENS, @@ -33,6 +34,14 @@ expand as _grammar_expand, DEFAULT_LEVEL as _GRAMMAR_LEVEL, ) +from context_engine.retrieval import RetrievalFusionInput, fuse_retrieval_context +from context_engine.structural import ( + CodeGraphBaseProvider, + CodeGraphClient, + SourceRange, + StructuralContext, + merge_structural_contexts, +) log = logging.getLogger(__name__) @@ -368,6 +377,71 @@ def _format_results_with_overflow(inline_chunks: list, overflow_chunks: list) -> return "\n\n---\n\n".join(parts) if parts else "No results found." +def _format_structural_context(structural: StructuralContext) -> str: + lines: list[str] = [] + if structural.sources: + lines.append("Structural sources:") + for source in structural.sources: + lines.append( + f"- {source.path}:{source.start_line}-{source.end_line}" + + (f"\n{source.content}" if source.content else "") + ) + if structural.relationships: + lines.append("Structural relationships:") + for relationship in structural.relationships: + lines.append( + f"- {_symbol_ref(relationship.source)} " + f"{relationship.kind} {_symbol_ref(relationship.target)}" + ) + if structural.impact: + lines.append("Structural impact:") + for symbol in structural.impact: + lines.append(f"- {_symbol_ref(symbol)}") + return "\n".join(lines) + + +def _symbol_ref(symbol) -> str: + location = symbol.path or "" + if getattr(symbol, "line", None): + location = f"{location}:{symbol.line}" + return f"{symbol.qualified_name} ({symbol.kind or 'symbol'} at {location})" + + +def _worktree_structural_overlay( + *, + query: str, + worktree_root: Path, + diff, +) -> StructuralContext: + sources = [] + for rel_path in sorted(diff.added | diff.modified): + source = _source_for_worktree_path(query, worktree_root, rel_path) + if source is not None: + sources.append(source) + return StructuralContext(sources=sources, provider="worktree") + + +def _source_for_worktree_path(query: str, worktree_root: Path, rel_path: str) -> SourceRange | None: + path = worktree_root / rel_path + try: + content = path.read_text(encoding="utf-8", errors="ignore") + except OSError: + return None + lines = content.splitlines() + if not lines: + return None + query_terms = [term.lower() for term in re.findall(r"\w+", query) if len(term) > 2] + start = 1 + for index, line in enumerate(lines, start=1): + haystack = line.lower() + if any(term in haystack for term in query_terms): + start = max(1, index - 3) + break + end = min(len(lines), start + 39) + excerpt = "\n".join(lines[start - 1:end]) + return SourceRange(path=rel_path, start_line=start, end_line=end, content=excerpt) + + class ContextEngineMCP: TOOL_NAMES = [ "context_search", @@ -383,12 +457,22 @@ class ContextEngineMCP: "set_output_compression", ] - def __init__(self, retriever, backend, compressor, embedder, config) -> None: + def __init__( + self, + retriever, + backend, + compressor, + embedder, + config, + *, + index_storage_base=None, + ) -> None: self._retriever = retriever self._backend = backend self._compressor = compressor self._embedder = embedder self._config = config + self._index_storage_base = index_storage_base # Set by _run_serve after construction; reset on every tool call # so the idle-shutdown watchdog knows the server is in use (#139). self._idle_tracker = None @@ -973,6 +1057,56 @@ async def call_tool(name: str, arguments: dict): # ── tool handlers ─────────────────────────────────────────────────────── + async def _build_structural_context(self, query: str): + if getattr(self._config, "structural_provider", "off") != "codegraph": + return StructuralContext(), None + try: + from context_engine.git import get_worktree_diff, resolve_git_repository_context + + git_context = resolve_git_repository_context(Path(self._project_dir)) + if git_context is None: + return StructuralContext(), None + diff = get_worktree_diff( + git_context.worktree_root, + base_sha=git_context.base_sha, + head_sha=git_context.head_sha, + ) + provider = getattr(self, "_structural_provider", None) + if provider is None: + provider = CodeGraphBaseProvider( + CodeGraphClient( + executable=getattr( + self._config, + "structural_codegraph_executable", + "codegraph", + ) + ) + ) + status = await provider.status(git_context.worktree_root) + if not status.available: + return StructuralContext(), diff + base_explore = await provider.explore(query, git_context.worktree_root) + base_impact = await provider.impact(query, git_context.worktree_root) + base = StructuralContext( + sources=base_explore.sources, + relationships=[ + *base_explore.relationships, + *base_impact.relationships, + ], + impact=[*base_explore.impact, *base_impact.impact], + provider=base_explore.provider, + metadata={"status": status.metadata}, + ) + overlay = _worktree_structural_overlay( + query=query, + worktree_root=git_context.worktree_root, + diff=diff, + ) + return merge_structural_contexts(base=base, overlay=overlay, diff=diff), diff + except Exception as exc: + log.debug("Structural context skipped: %s", exc) + return StructuralContext(), None + async def _ensure_indexed(self) -> bool: """Lazy indexing on empty index. @@ -985,7 +1119,12 @@ async def _ensure_indexed(self) -> bool: client's side it looked like the call had hung (#67). """ try: - count = self._backend._vector_store.count() + count_chunks = getattr(self._backend, "count_chunks", None) + count = ( + count_chunks() + if count_chunks is not None + else self._backend._vector_store.count() + ) if count > 0: return True except Exception: @@ -1007,7 +1146,12 @@ async def _ensure_indexed(self) -> bool: async def _bg_index(): try: from context_engine.indexer.pipeline import run_indexing - await run_indexing(self._config, self._project_dir, full=False) + await run_indexing( + self._config, + self._project_dir, + full=False, + storage_base_override=getattr(self, "_index_storage_base", None), + ) log.info("Background indexing complete for %s", self._project_name) except Exception as exc: log.warning("Background indexing failed: %s", exc) @@ -1019,7 +1163,12 @@ async def _bg_index(): # to blocking. Better than swallowing the request. try: from context_engine.indexer.pipeline import run_indexing - await run_indexing(self._config, self._project_dir, full=False) + await run_indexing( + self._config, + self._project_dir, + full=False, + storage_base_override=getattr(self, "_index_storage_base", None), + ) return True except Exception as exc: log.warning("Lazy indexing failed: %s", exc) @@ -1087,7 +1236,18 @@ async def _handle_context_search(self, args): # small projects and aren't what users are searching for. all_chunks = [c for c in all_chunks if not _is_cce_config(c.file_path)] - inline_chunks, overflow_chunks = _split_inline_overflow(all_chunks, max_tokens) + structural_context, structural_diff = await self._build_structural_context(query) + fused_context = fuse_retrieval_context( + RetrievalFusionInput( + overlay_chunks=all_chunks, + structural=structural_context, + diff=structural_diff or GitWorktreeDiff(), + ), + max_tokens=max_tokens, + ) + inline_chunks = fused_context.chunks + inline_ids = {chunk.id for chunk in inline_chunks} + overflow_chunks = [chunk for chunk in all_chunks if chunk.id not in inline_ids] # Accounting — track per-file to cap overlapping chunks per_file_raw: dict[str, int] = {} @@ -1128,6 +1288,9 @@ async def _handle_context_search(self, args): self._persist_current_session() body = _format_results_with_overflow(inline_chunks, overflow_chunks) + structural_text = _format_structural_context(fused_context.structural) + if structural_text: + body = "Relevant source chunks:\n" + body + "\n\n---\n\n" + structural_text # Memory nudge — appended before output compression so the agent # sees it as part of the tool result. Fires only when thresholds # are met (see _build_nudge). Increment BEFORE building the nudge diff --git a/src/context_engine/retrieval/__init__.py b/src/context_engine/retrieval/__init__.py index e69de29..07b33f8 100644 --- a/src/context_engine/retrieval/__init__.py +++ b/src/context_engine/retrieval/__init__.py @@ -0,0 +1,19 @@ +from context_engine.retrieval.budgeting import ( + DEFAULT_DISCLOSURE_BUDGETS, + DisclosureLevel, + resolve_context_budget, +) +from context_engine.retrieval.fusion import ( + FusedRetrievalContext, + RetrievalFusionInput, + fuse_retrieval_context, +) + +__all__ = [ + "FusedRetrievalContext", + "RetrievalFusionInput", + "DEFAULT_DISCLOSURE_BUDGETS", + "DisclosureLevel", + "fuse_retrieval_context", + "resolve_context_budget", +] diff --git a/src/context_engine/retrieval/budgeting.py b/src/context_engine/retrieval/budgeting.py new file mode 100644 index 0000000..128ac60 --- /dev/null +++ b/src/context_engine/retrieval/budgeting.py @@ -0,0 +1,24 @@ +"""Progressive disclosure token budgets.""" + +from __future__ import annotations + +from typing import Literal + + +DisclosureLevel = Literal["summary", "standard", "full"] + +DEFAULT_DISCLOSURE_BUDGETS: dict[DisclosureLevel, int] = { + "summary": 1500, + "standard": 6000, + "full": 50000, +} + + +def resolve_context_budget( + *, + level: DisclosureLevel = "standard", + max_tokens: int | None = None, +) -> int: + """Return hard token budget for a disclosure level, honoring explicit caps.""" + budget = DEFAULT_DISCLOSURE_BUDGETS[level] + return min(budget, max_tokens) if max_tokens is not None else budget diff --git a/src/context_engine/retrieval/fusion.py b/src/context_engine/retrieval/fusion.py new file mode 100644 index 0000000..47262e4 --- /dev/null +++ b/src/context_engine/retrieval/fusion.py @@ -0,0 +1,92 @@ +"""Query-time fusion for semantic and structural worktree context.""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +from context_engine.git.diff import GitWorktreeDiff +from context_engine.models import Chunk +from context_engine.retrieval.budgeting import DisclosureLevel, resolve_context_budget +from context_engine.structural.models import StructuralContext + + +@dataclass(frozen=True) +class RetrievalFusionInput: + base_chunks: list[Chunk] = field(default_factory=list) + overlay_chunks: list[Chunk] = field(default_factory=list) + structural: StructuralContext = field(default_factory=StructuralContext) + diff: GitWorktreeDiff = field(default_factory=GitWorktreeDiff) + + +@dataclass(frozen=True) +class FusedRetrievalContext: + chunks: list[Chunk] + structural: StructuralContext + omitted: int = 0 + + +def fuse_retrieval_context( + data: RetrievalFusionInput, + *, + level: DisclosureLevel = "standard", + max_tokens: int | None = None, +) -> FusedRetrievalContext: + """Merge semantic overlay/base chunks and keep hard token budget.""" + chunks = _merged_chunks(data.base_chunks, data.overlay_chunks, data.diff) + budget = resolve_context_budget(level=level, max_tokens=max_tokens) + packed, omitted = _pack_chunks( + chunks, + max(0, budget - _structural_token_count(data.structural)), + ) + return FusedRetrievalContext(chunks=packed, structural=data.structural, omitted=omitted) + + +def _merged_chunks( + base_chunks: list[Chunk], + overlay_chunks: list[Chunk], + diff: GitWorktreeDiff, +) -> list[Chunk]: + shadowed_paths = diff.added | diff.modified | diff.deleted + merged: dict[tuple[str, int, int], Chunk] = {} + for chunk in base_chunks: + if chunk.file_path in shadowed_paths: + continue + merged[_chunk_key(chunk)] = chunk + for chunk in overlay_chunks: + merged[_chunk_key(chunk)] = chunk + return sorted(merged.values(), key=_chunk_rank_key) + + +def _pack_chunks(chunks: list[Chunk], max_tokens: int | None) -> tuple[list[Chunk], int]: + if max_tokens is None: + return chunks, 0 + budget = max_tokens + packed: list[Chunk] = [] + for chunk in chunks: + if chunk.token_count > budget: + continue + packed.append(chunk) + budget -= chunk.token_count + return packed, len(chunks) - len(packed) + + +def _chunk_key(chunk: Chunk) -> tuple[str, int, int]: + return (chunk.file_path, chunk.start_line, chunk.end_line) + + +def _chunk_rank_key(chunk: Chunk) -> tuple[float, float, str]: + distance = chunk.metadata.get("_distance") + if isinstance(distance, int | float): + return (0.0, float(distance), chunk.id) + return (1.0, -chunk.confidence_score, chunk.id) + + +def _structural_token_count(structural: StructuralContext) -> int: + text_parts: list[str] = [] + text_parts.extend(source.content or source.path for source in structural.sources) + text_parts.extend( + f"{relationship.source.qualified_name} {relationship.kind} {relationship.target.qualified_name}" + for relationship in structural.relationships + ) + text_parts.extend(symbol.qualified_name for symbol in structural.impact) + return max(0, int(sum(len(part) for part in text_parts) / 3.3)) diff --git a/src/context_engine/storage/worktree_overlay.py b/src/context_engine/storage/worktree_overlay.py new file mode 100644 index 0000000..0fca189 --- /dev/null +++ b/src/context_engine/storage/worktree_overlay.py @@ -0,0 +1,251 @@ +"""Semantic base+worktree overlay storage backend.""" + +from __future__ import annotations + +import asyncio + +from context_engine.git.diff import GitWorktreeDiff +from context_engine.models import Chunk, EdgeType, GraphEdge, GraphNode +from context_engine.storage.backend import StorageBackend + + +class WorktreeOverlayBackend: + """Search shared base plus overlay while hiding stale base paths.""" + + def __init__( + self, + *, + base: StorageBackend, + overlay: StorageBackend, + diff: GitWorktreeDiff, + ) -> None: + self._base = base + self._overlay = overlay + self._diff = diff + + def update_diff(self, diff: GitWorktreeDiff) -> None: + self._diff = diff + + async def ingest( + self, + chunks: list[Chunk], + nodes: list[GraphNode], + edges: list[GraphEdge], + ) -> None: + await self._overlay.ingest(chunks, nodes, edges) + + async def vector_search( + self, + query_embedding: list[float], + top_k: int = 10, + filters: dict | None = None, + ) -> list[Chunk]: + base_results, overlay_results = await asyncio.gather( + self._base.vector_search(query_embedding, top_k=top_k, filters=filters), + self._overlay.vector_search(query_embedding, top_k=top_k, filters=filters), + ) + return self._merge_chunks(base_results, overlay_results, top_k) + + async def fts_search(self, query: str, top_k: int = 30) -> list[tuple[str, float]]: + base_hits, overlay_hits = await asyncio.gather( + self._base.fts_search(query, top_k=top_k), + self._overlay.fts_search(query, top_k=top_k), + ) + base_chunks = await self._base.get_chunks_by_ids( + [chunk_id for chunk_id, _ in base_hits] + ) + visible_base_ids = {chunk.id for chunk in self._visible_base_chunks(base_chunks)} + return _merge_ranked_hits( + [ + (rank, hit) + for rank, hit in enumerate(base_hits, start=1) + if hit[0] in visible_base_ids + ], + list(enumerate(overlay_hits, start=1)), + top_k, + ) + + async def graph_neighbors( + self, + node_id: str, + edge_type: EdgeType | None = None, + ) -> list[GraphNode]: + if _node_file_path(node_id) in self._diff.added | self._diff.modified: + overlay_nodes = await self._overlay.graph_neighbors(node_id, edge_type) + return _dedupe_nodes( + [node for node in overlay_nodes if not self._is_deleted(node.file_path)] + ) + + base_nodes, overlay_nodes = await asyncio.gather( + self._base.graph_neighbors(node_id, edge_type), + self._overlay.graph_neighbors(node_id, edge_type), + ) + return _dedupe_nodes( + [ + *[node for node in overlay_nodes if not self._is_deleted(node.file_path)], + *[node for node in base_nodes if not self._is_deleted(node.file_path)], + ] + ) + + async def get_chunk_by_id(self, chunk_id: str) -> Chunk | None: + chunk = await self._overlay.get_chunk_by_id(chunk_id) + if chunk is not None: + return chunk + chunk = await self._base.get_chunk_by_id(chunk_id) + if chunk is not None and self._is_shadowed(chunk.file_path): + return None + return chunk + + async def get_chunks_by_ids(self, chunk_ids: list[str]) -> list[Chunk]: + overlay_chunks, base_chunks = await asyncio.gather( + self._overlay.get_chunks_by_ids(chunk_ids), + self._base.get_chunks_by_ids(chunk_ids), + ) + return [*overlay_chunks, *self._visible_base_chunks(base_chunks)] + + def count_chunks(self) -> int: + overlay_count = _count_chunks(self._overlay) + changed_paths = self._diff.added | self._diff.modified + if changed_paths and overlay_count == 0: + return 0 + return _count_chunks(self._base) + overlay_count + + def needs_overlay_index(self) -> bool: + return ( + bool(self._diff.added | self._diff.modified) + and _count_chunks(self._overlay) == 0 + ) + + async def delete_by_file(self, file_path: str) -> None: + await self._overlay.delete_by_file(file_path) + + async def delete_by_files(self, file_paths: list[str]) -> None: + delete_many = getattr(self._overlay, "delete_by_files", None) + if delete_many is not None: + await delete_many(file_paths) + return + for file_path in file_paths: + await self._overlay.delete_by_file(file_path) + + async def get_related_file_paths(self, file_paths: list[str]) -> list[str]: + base_related = getattr(self._base, "get_related_file_paths", None) + overlay_related = getattr(self._overlay, "get_related_file_paths", None) + active_base_files = [ + file_path + for file_path in file_paths + if file_path not in self._diff.added | self._diff.modified + ] + base_paths, overlay_paths = await asyncio.gather( + base_related(active_base_files) + if base_related is not None and active_base_files + else _empty_paths(), + overlay_related(file_paths) if overlay_related is not None else _empty_paths(), + ) + return _dedupe_paths([*overlay_paths, *base_paths], deleted=self._diff.deleted) + + def get_cached_compression(self, chunk_id: str, level: str) -> str | None: + overlay_cache = getattr(self._overlay, "get_cached_compression", None) + if overlay_cache is not None: + cached = overlay_cache(chunk_id, level) + if cached is not None: + return cached + base_cache = getattr(self._base, "get_cached_compression", None) + return base_cache(chunk_id, level) if base_cache is not None else None + + def put_cached_compression(self, chunk_id: str, level: str, content: str) -> None: + overlay_cache = getattr(self._overlay, "put_cached_compression", None) + if overlay_cache is not None: + overlay_cache(chunk_id, level, content) + + def _merge_chunks( + self, + base_chunks: list[Chunk], + overlay_chunks: list[Chunk], + top_k: int, + ) -> list[Chunk]: + ranked: dict[tuple[str, int, int], Chunk] = {} + for chunk in self._visible_base_chunks(base_chunks): + ranked[_source_key(chunk)] = chunk + for chunk in overlay_chunks: + ranked[_source_key(chunk)] = chunk + return sorted(ranked.values(), key=_chunk_rank_key)[:top_k] + + def _visible_base_chunks(self, chunks: list[Chunk]) -> list[Chunk]: + return [chunk for chunk in chunks if not self._is_shadowed(chunk.file_path)] + + def _is_shadowed(self, file_path: str) -> bool: + return ( + file_path in self._diff.added + or file_path in self._diff.modified + or file_path in self._diff.deleted + ) + + def _is_deleted(self, file_path: str) -> bool: + return file_path in self._diff.deleted + + +async def _empty_paths() -> list[str]: + return [] + + +def _dedupe_paths(paths: list[str], *, deleted: set[str]) -> list[str]: + deduped = [] + seen = set() + for path in paths: + if path in deleted or path in seen: + continue + deduped.append(path) + seen.add(path) + return deduped + + +def _dedupe_nodes(nodes: list[GraphNode]) -> list[GraphNode]: + deduped = [] + seen = set() + for node in nodes: + key = (node.id, node.file_path, node.name) + if key in seen: + continue + deduped.append(node) + seen.add(key) + return deduped + + +def _merge_ranked_hits( + base_hits: list[tuple[int, tuple[str, float]]], + overlay_hits: list[tuple[int, tuple[str, float]]], + top_k: int, +) -> list[tuple[str, float]]: + hits_by_id: dict[str, tuple[str, float]] = {} + rank_scores: dict[str, float] = {} + for hits in (base_hits, overlay_hits): + for rank, hit in hits: + chunk_id, _ = hit + hits_by_id[chunk_id] = hit + rank_scores[chunk_id] = rank_scores.get(chunk_id, 0.0) + 1.0 / (60 + rank) + return sorted( + hits_by_id.values(), + key=lambda hit: (-rank_scores[hit[0]], hit[0]), + )[:top_k] + + +def _chunk_rank_key(chunk: Chunk) -> tuple[float, float, str]: + distance = chunk.metadata.get("_distance") + if isinstance(distance, int | float): + return (0.0, float(distance), chunk.id) + return (1.0, -chunk.confidence_score, chunk.id) + + +def _source_key(chunk: Chunk) -> tuple[str, int, int]: + return (chunk.file_path, chunk.start_line, chunk.end_line) + + +def _node_file_path(node_id: str) -> str: + return node_id.split(":", 1)[0] + + +def _count_chunks(backend: StorageBackend) -> int: + count_chunks = getattr(backend, "count_chunks", None) + if count_chunks is not None: + return count_chunks() + return backend._vector_store.count() diff --git a/src/context_engine/structural/__init__.py b/src/context_engine/structural/__init__.py new file mode 100644 index 0000000..18dd09e --- /dev/null +++ b/src/context_engine/structural/__init__.py @@ -0,0 +1,18 @@ +"""Provider-neutral structural context models and providers.""" + +from context_engine.structural.base import BaseStructuralProvider, ProviderStatus +from context_engine.structural.codegraph import CodeGraphBaseProvider, CodeGraphClient +from context_engine.structural.merge import merge_structural_contexts +from context_engine.structural.models import Relationship, SourceRange, StructuralContext, SymbolKey + +__all__ = [ + "BaseStructuralProvider", + "CodeGraphBaseProvider", + "CodeGraphClient", + "ProviderStatus", + "Relationship", + "SourceRange", + "StructuralContext", + "SymbolKey", + "merge_structural_contexts", +] diff --git a/src/context_engine/structural/base.py b/src/context_engine/structural/base.py new file mode 100644 index 0000000..46c87a6 --- /dev/null +++ b/src/context_engine/structural/base.py @@ -0,0 +1,25 @@ +"""Structural provider protocols.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Protocol + +from context_engine.structural.models import StructuralContext + + +@dataclass(frozen=True) +class ProviderStatus: + provider: str + available: bool + detail: str = "" + metadata: dict = field(default_factory=dict) + + +class BaseStructuralProvider(Protocol): + async def status(self, project_root: Path) -> ProviderStatus: ... + + async def explore(self, query: str, project_root: Path) -> StructuralContext: ... + + async def impact(self, symbol: str, project_root: Path) -> StructuralContext: ... diff --git a/src/context_engine/structural/codegraph.py b/src/context_engine/structural/codegraph.py new file mode 100644 index 0000000..74ddf84 --- /dev/null +++ b/src/context_engine/structural/codegraph.py @@ -0,0 +1,140 @@ +"""CodeGraph shared-base structural provider.""" + +from __future__ import annotations + +import json +import subprocess +import tempfile +from pathlib import Path +from typing import Any + +from context_engine.structural.base import ProviderStatus +from context_engine.structural.models import SourceRange, StructuralContext, SymbolKey + + +DEFAULT_CODEGRAPH_TIMEOUT_SECONDS = 8 +DEFAULT_CODEGRAPH_STDOUT_LIMIT = 1_000_000 + + +class CodeGraphClient: + def __init__( + self, + *, + executable: str = "codegraph", + timeout_seconds: int = DEFAULT_CODEGRAPH_TIMEOUT_SECONDS, + stdout_limit: int = DEFAULT_CODEGRAPH_STDOUT_LIMIT, + ) -> None: + self._executable = executable + self._timeout_seconds = timeout_seconds + self._stdout_limit = stdout_limit + + def status(self, project_root: Path) -> dict[str, Any]: + return self._run_json(project_root, ["status", "--json"]) + + def query(self, project_root: Path, query: str, *, limit: int = 20) -> list[dict[str, Any]]: + data = self._run_json(project_root, ["query", query, "--limit", str(limit), "--json"]) + if isinstance(data, list): + return data + if isinstance(data, dict): + results = data.get("results", []) + return results if isinstance(results, list) else [] + return [] + + def impact(self, project_root: Path, symbol: str, *, depth: int = 2) -> dict[str, Any]: + return self._run_json(project_root, ["impact", symbol, "--depth", str(depth), "--json"]) + + def _run_json(self, project_root: Path, args: list[str]) -> Any: + with tempfile.TemporaryFile(mode="w+t", encoding="utf-8") as stdout_file: + process = subprocess.Popen( + [self._executable, *args], + cwd=project_root, + stdout=stdout_file, + stderr=subprocess.PIPE, + text=True, + ) + try: + _, stderr = process.communicate(timeout=self._timeout_seconds) + except subprocess.TimeoutExpired: + process.kill() + process.communicate() + raise + + stdout_file.seek(0, 2) + stdout_size = stdout_file.tell() + stdout_file.seek(0) + stdout = stdout_file.read(self._stdout_limit + 1) + + if stdout_size > self._stdout_limit: + raise CodeGraphError("CodeGraph JSON response too large") + if process.returncode != 0: + detail = (stderr or "").strip() or stdout.strip() or f"exit {process.returncode}" + raise CodeGraphError(detail[:1000]) + try: + return json.loads(stdout) + except json.JSONDecodeError as exc: + raise CodeGraphError(f"invalid JSON from CodeGraph: {exc.msg}") from exc + + +class CodeGraphBaseProvider: + provider_name = "codegraph" + + def __init__(self, client: CodeGraphClient | None = None) -> None: + self._client = client or CodeGraphClient() + + async def status(self, project_root: Path) -> ProviderStatus: + try: + status = self._client.status(project_root) + except (CodeGraphError, FileNotFoundError, subprocess.TimeoutExpired) as exc: + return ProviderStatus(self.provider_name, False, str(exc)) + available = ( + bool(status.get("initialized")) + and status.get("index", {}).get("state") == "complete" + ) + return ProviderStatus( + provider=self.provider_name, + available=available, + detail="" if available else "CodeGraph index unavailable or incomplete", + metadata=status, + ) + + async def explore(self, query: str, project_root: Path) -> StructuralContext: + rows = self._client.query(project_root, query) + return StructuralContext( + sources=[_source_from_query_row(row) for row in rows], + provider=self.provider_name, + metadata={"query": query}, + ) + + async def impact(self, symbol: str, project_root: Path) -> StructuralContext: + data = self._client.impact(project_root, symbol) + affected = data.get("affected", []) if isinstance(data, dict) else [] + return StructuralContext( + impact=[_symbol_from_impact_row(row) for row in affected if isinstance(row, dict)], + provider=self.provider_name, + metadata={"symbol": symbol, "raw": data}, + ) + + +class CodeGraphError(RuntimeError): + pass + + +def _source_from_query_row(row: dict[str, Any]) -> SourceRange: + node = row.get("node", row) + return SourceRange( + path=str(node.get("filePath") or node.get("file_path") or ""), + start_line=int(node.get("startLine") or node.get("start_line") or 0), + end_line=int(node.get("endLine") or node.get("end_line") or node.get("startLine") or 0), + content=node.get("code") or node.get("content"), + ) + + +def _symbol_from_impact_row(row: dict[str, Any]) -> SymbolKey: + line = row.get("startLine") or row.get("start_line") + return SymbolKey( + qualified_name=str(row.get("name") or ""), + kind=row.get("kind"), + path=row.get("filePath") or row.get("file_path"), + signature=row.get("signature"), + line=int(line) if line is not None else None, + ) diff --git a/src/context_engine/structural/merge.py b/src/context_engine/structural/merge.py new file mode 100644 index 0000000..ec21036 --- /dev/null +++ b/src/context_engine/structural/merge.py @@ -0,0 +1,86 @@ +"""Merge shared structural base with worktree overlay context.""" + +from __future__ import annotations + +from context_engine.git.diff import GitWorktreeDiff +from context_engine.structural.models import Relationship, SourceRange, StructuralContext, SymbolKey + + +def merge_structural_contexts( + *, + base: StructuralContext, + overlay: StructuralContext, + diff: GitWorktreeDiff, +) -> StructuralContext: + """Return effective structural context for a worktree.""" + shadowed_paths = diff.added | diff.modified | diff.deleted + overlay_symbols = set(overlay.impact) + overlay_relationship_sources = {relationship.source for relationship in overlay.relationships} + + return StructuralContext( + sources=_merge_sources(base.sources, overlay.sources, shadowed_paths), + relationships=_merge_relationships( + base.relationships, + overlay.relationships, + shadowed_paths, + diff.deleted, + overlay_relationship_sources, + ), + impact=[ + *overlay.impact, + *[ + symbol + for symbol in base.impact + if symbol not in overlay_symbols and not _symbol_is_shadowed(symbol, shadowed_paths) + ], + ], + provider=f"{base.provider}+{overlay.provider}", + metadata={"base": base.metadata, "overlay": overlay.metadata}, + ) + + +def _merge_sources( + base_sources: list[SourceRange], + overlay_sources: list[SourceRange], + shadowed_paths: set[str], +) -> list[SourceRange]: + merged = list(overlay_sources) + seen = {_source_key(source) for source in overlay_sources} + for source in base_sources: + if source.path in shadowed_paths: + continue + key = _source_key(source) + if key not in seen: + merged.append(source) + seen.add(key) + return merged + + +def _merge_relationships( + base_relationships: list[Relationship], + overlay_relationships: list[Relationship], + shadowed_paths: set[str], + deleted_paths: set[str], + overlay_relationship_sources: set[SymbolKey], +) -> list[Relationship]: + merged = list(overlay_relationships) + seen = set(overlay_relationships) + for relationship in base_relationships: + if _symbol_is_shadowed(relationship.source, shadowed_paths): + continue + if relationship.target.path in deleted_paths: + continue + if relationship.source in overlay_relationship_sources: + continue + if relationship not in seen: + merged.append(relationship) + seen.add(relationship) + return merged + + +def _source_key(source: SourceRange) -> tuple[str, int, int]: + return (source.path, source.start_line, source.end_line) + + +def _symbol_is_shadowed(symbol: SymbolKey, shadowed_paths: set[str]) -> bool: + return symbol.path in shadowed_paths diff --git a/src/context_engine/structural/models.py b/src/context_engine/structural/models.py new file mode 100644 index 0000000..143a3e1 --- /dev/null +++ b/src/context_engine/structural/models.py @@ -0,0 +1,39 @@ +"""Provider-independent structural context models.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + + +@dataclass(frozen=True) +class SymbolKey: + qualified_name: str + kind: str | None + path: str | None = None + signature: str | None = None + line: int | None = None + + +@dataclass(frozen=True) +class SourceRange: + path: str + start_line: int + end_line: int + content: str | None = None + + +@dataclass(frozen=True) +class Relationship: + source: SymbolKey + target: SymbolKey + kind: str + + +@dataclass(frozen=True) +class StructuralContext: + sources: list[SourceRange] = field(default_factory=list) + relationships: list[Relationship] = field(default_factory=list) + impact: list[SymbolKey] = field(default_factory=list) + provider: str = "unknown" + metadata: dict[str, Any] = field(default_factory=dict) diff --git a/src/context_engine/utils.py b/src/context_engine/utils.py index 2c7e3bc..801404d 100644 --- a/src/context_engine/utils.py +++ b/src/context_engine/utils.py @@ -70,6 +70,13 @@ def _project_slug(project_dir: Path) -> str: return f"{safe or 'project'}-{h}" +def resolve_project_storage_dir(config: object, project_dir: Path) -> Path: + """Return the slugged per-project storage directory without filesystem changes.""" + slug = _project_slug(project_dir) + storage_root = Path(config.storage_path) # type: ignore[union-attr] + return storage_root / slug + + def project_storage_dir(config: object, project_dir: Path) -> Path: """Return the per-project storage directory under ``config.storage_path``. @@ -81,9 +88,8 @@ def project_storage_dir(config: object, project_dir: Path) -> Path: exists but the new slug directory does not, the legacy directory is renamed in place to preserve existing users' data. """ - slug = _project_slug(project_dir) + slug_path = resolve_project_storage_dir(config, project_dir) storage_root = Path(config.storage_path) # type: ignore[union-attr] - slug_path = storage_root / slug legacy_path = storage_root / project_dir.resolve().name if not slug_path.exists() and legacy_path.exists(): diff --git a/tests/integration/test_mcp_empty_index.py b/tests/integration/test_mcp_empty_index.py index 1d96a8e..95da351 100644 --- a/tests/integration/test_mcp_empty_index.py +++ b/tests/integration/test_mcp_empty_index.py @@ -26,6 +26,7 @@ def _make_mcp(tmp_path, monkeypatch, *, chunk_count: int): embedding_model="BAAI/bge-small-en-v1.5", ) backend = MagicMock() + backend.count_chunks.return_value = chunk_count backend._vector_store.count.return_value = chunk_count compressor = MagicMock() embedder = MagicMock() @@ -103,8 +104,11 @@ async def test_ensure_indexed_returns_false_for_empty_and_true_for_populated( tmp_path, monkeypatch ): mcp_empty = _make_mcp(tmp_path / "a", monkeypatch, chunk_count=0) + mcp_empty._index_storage_base = tmp_path / "overlay-index" + indexing_kwargs: list[dict] = [] async def _fake_run_indexing(*a, **kw): + indexing_kwargs.append(kw) return None monkeypatch.setattr( @@ -112,6 +116,10 @@ async def _fake_run_indexing(*a, **kw): ) assert await mcp_empty._ensure_indexed() is False + await asyncio.sleep(0) + assert indexing_kwargs == [ + {"full": False, "storage_base_override": tmp_path / "overlay-index"} + ] mcp_full = _make_mcp(tmp_path / "b", monkeypatch, chunk_count=10) assert await mcp_full._ensure_indexed() is True diff --git a/tests/integration/test_mcp_server.py b/tests/integration/test_mcp_server.py index 93234d0..7ed6901 100644 --- a/tests/integration/test_mcp_server.py +++ b/tests/integration/test_mcp_server.py @@ -1,6 +1,17 @@ import pytest -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock +from context_engine.cli import _run_index, _runtime_storage_backend, _runtime_storage_base from context_engine.integration.mcp_server import ContextEngineMCP +from context_engine.models import Chunk, ChunkType +from context_engine.retrieval.retriever import HybridRetriever +from context_engine.structural import ( + ProviderStatus, + Relationship, + SourceRange, + StructuralContext, + SymbolKey, +) +from tests.test_git_repository_context import init_repo def test_mcp_server_has_required_tools(): @@ -36,6 +47,287 @@ def _make_server(tmp_path): return server +def _embedded_chunk(chunk_id, file_path, content): + return Chunk( + id=chunk_id, + content=content, + chunk_type=ChunkType.FUNCTION, + file_path=file_path, + start_line=1, + end_line=2, + language="python", + embedding=[0.1, 0.2, 0.3, 0.4], + ) + + +class _StubEmbedder: + def embed_query(self, query): + return [0.1, 0.2, 0.3, 0.4] + + +def _symbol(name, path, line=1): + return SymbolKey(qualified_name=name, kind="function", path=path, line=line) + + +class _StubStructuralProvider: + def __init__(self, context, *, available=True): + self._context = context + self._available = available + + async def status(self, project_root): + return ProviderStatus("codegraph", self._available, metadata={"index": {"state": "complete"}}) + + async def explore(self, query, project_root): + return self._context + + async def impact(self, symbol, project_root): + return StructuralContext(impact=list(self._context.impact), provider="codegraph") + + +@pytest.mark.asyncio +async def test_context_search_uses_worktree_overlay_version(tmp_path, monkeypatch): + repo = init_repo(tmp_path / "repo") + (repo / "base.py").write_text("def base():\n return 'worktree value'\n") + cfg = MagicMock() + cfg.storage_path = str(tmp_path / "storage") + cfg.output_compression = "off" + cfg.compression_level = "off" + cfg.retrieval_confidence_threshold = 0.0 + cfg.retrieval_marginal_ratio = 0.0 + monkeypatch.chdir(repo) + + _, backend, _ = _runtime_storage_backend(cfg, repo) + await backend._base.ingest( + [_embedded_chunk("base-old", "base.py", "def base():\n return 'base value'\n")], + [], + [], + ) + await backend.ingest( + [_embedded_chunk("overlay-new", "base.py", "def base():\n return 'worktree value'\n")], + [], + [], + ) + + server = _make_server(tmp_path) + server._config = cfg + server._backend = backend + server._retriever = HybridRetriever(backend=backend, embedder=_StubEmbedder()) + server._compressor = MagicMock() + server._compressor.compress = AsyncMock(side_effect=lambda chunks, _level: chunks) + server._session_capture = MagicMock() + server._session_capture.touch_files = MagicMock() + server._session_capture.get_session_snapshot = MagicMock(return_value={ + "decisions": [], + "code_areas": [], + "touched_files": {}, + }) + server._persist_current_session = MagicMock() + server._record = MagicMock() + server._append_audit_log = MagicMock() + server._session_id = "test-session" + + result = await server._handle_context_search({"query": "base value", "top_k": 5}) + text = result[0].text + + assert "worktree value" in text + assert "base value" not in text + + +@pytest.mark.asyncio +async def test_context_search_returns_codegraph_structural_context(tmp_path, monkeypatch): + repo = init_repo(tmp_path / "repo") + cfg = MagicMock() + cfg.storage_path = str(tmp_path / "storage") + cfg.output_compression = "off" + cfg.compression_level = "off" + cfg.retrieval_confidence_threshold = 0.0 + cfg.retrieval_marginal_ratio = 0.0 + cfg.structural_provider = "codegraph" + monkeypatch.chdir(repo) + + _, backend, _ = _runtime_storage_backend(cfg, repo) + await backend._base.ingest( + [_embedded_chunk("semantic", "base.py", "def base():\n return 'semantic'\n")], + [], + [], + ) + server = _make_server(tmp_path) + server._config = cfg + server._project_dir = str(repo) + server._backend = backend + server._retriever = HybridRetriever(backend=backend, embedder=_StubEmbedder()) + server._compressor = MagicMock() + server._compressor.compress = AsyncMock(side_effect=lambda chunks, _level: chunks) + server._session_capture = MagicMock() + server._session_capture.touch_files = MagicMock() + server._session_capture.get_session_snapshot = MagicMock(return_value={ + "decisions": [], + "code_areas": [], + "touched_files": {}, + }) + server._persist_current_session = MagicMock() + server._record = MagicMock() + server._append_audit_log = MagicMock() + server._session_id = "test-session" + auth = _symbol("auth.login", "auth.py", line=12) + token = _symbol("token.validate", "token.py", line=3) + server._structural_provider = _StubStructuralProvider(StructuralContext( + sources=[SourceRange("auth.py", 12, 18, "def login(): pass")], + relationships=[Relationship(auth, token, "calls")], + impact=[token], + provider="codegraph", + )) + + result = await server._handle_context_search({"query": "login", "top_k": 5}) + text = result[0].text + + assert "Relevant source chunks:" in text + assert "def base()" in text + assert "Structural sources:" in text + assert "auth.py:12-18" in text + assert "auth.login (function at auth.py:12) calls token.validate (function at token.py:3)" in text + assert "Structural impact:" in text + + +@pytest.mark.asyncio +async def test_context_search_shadows_modified_codegraph_base_source(tmp_path, monkeypatch): + repo = init_repo(tmp_path / "repo") + (repo / "base.py").write_text("def base():\n return 'worktree auth'\n") + cfg = MagicMock() + cfg.storage_path = str(tmp_path / "storage") + cfg.output_compression = "off" + cfg.compression_level = "off" + cfg.retrieval_confidence_threshold = 0.0 + cfg.retrieval_marginal_ratio = 0.0 + cfg.structural_provider = "codegraph" + monkeypatch.chdir(repo) + + _, backend, _ = _runtime_storage_backend(cfg, repo) + await backend.ingest( + [_embedded_chunk("overlay", "base.py", "def base():\n return 'worktree auth'\n")], + [], + [], + ) + server = _make_server(tmp_path) + server._config = cfg + server._project_dir = str(repo) + server._backend = backend + server._retriever = HybridRetriever(backend=backend, embedder=_StubEmbedder()) + server._compressor = MagicMock() + server._compressor.compress = AsyncMock(side_effect=lambda chunks, _level: chunks) + server._session_capture = MagicMock() + server._session_capture.touch_files = MagicMock() + server._session_capture.get_session_snapshot = MagicMock(return_value={ + "decisions": [], + "code_areas": [], + "touched_files": {}, + }) + server._persist_current_session = MagicMock() + server._record = MagicMock() + server._append_audit_log = MagicMock() + server._session_id = "test-session" + server._structural_provider = _StubStructuralProvider(StructuralContext( + sources=[SourceRange("base.py", 1, 2, "def base():\n return 'base auth'")], + provider="codegraph", + )) + + result = await server._handle_context_search({"query": "auth", "top_k": 5}) + text = result[0].text + + assert "worktree auth" in text + assert "base auth" not in text + + +@pytest.mark.asyncio +async def test_context_search_degrades_when_codegraph_unavailable(tmp_path, monkeypatch): + repo = init_repo(tmp_path / "repo") + cfg = MagicMock() + cfg.storage_path = str(tmp_path / "storage") + cfg.output_compression = "off" + cfg.compression_level = "off" + cfg.retrieval_confidence_threshold = 0.0 + cfg.retrieval_marginal_ratio = 0.0 + cfg.structural_provider = "codegraph" + monkeypatch.chdir(repo) + + _, backend, _ = _runtime_storage_backend(cfg, repo) + await backend._base.ingest( + [_embedded_chunk("semantic", "base.py", "def base():\n return 'semantic'\n")], + [], + [], + ) + server = _make_server(tmp_path) + server._config = cfg + server._project_dir = str(repo) + server._backend = backend + server._retriever = HybridRetriever(backend=backend, embedder=_StubEmbedder()) + server._compressor = MagicMock() + server._compressor.compress = AsyncMock(side_effect=lambda chunks, _level: chunks) + server._session_capture = MagicMock() + server._session_capture.touch_files = MagicMock() + server._session_capture.get_session_snapshot = MagicMock(return_value={ + "decisions": [], + "code_areas": [], + "touched_files": {}, + }) + server._persist_current_session = MagicMock() + server._record = MagicMock() + server._append_audit_log = MagicMock() + server._session_id = "test-session" + server._structural_provider = _StubStructuralProvider( + StructuralContext(), + available=False, + ) + + result = await server._handle_context_search({"query": "semantic", "top_k": 5}) + text = result[0].text + + assert "def base()" in text + assert "Structural sources:" not in text + + +@pytest.mark.asyncio +async def test_runtime_backend_refreshes_diff_after_worktree_change(tmp_path): + repo = init_repo(tmp_path / "repo") + cfg = MagicMock() + cfg.storage_path = str(tmp_path / "storage") + _, backend, refresh_diff = _runtime_storage_backend(cfg, repo) + await backend._base.ingest( + [_embedded_chunk("base-old", "base.py", "def base():\n return 'base value'\n")], + [], + [], + ) + + assert await backend.get_chunk_by_id("base-old") is not None + + (repo / "base.py").write_text("def base():\n return 'worktree value'\n") + assert refresh_diff is not None + refresh_diff() + + assert await backend.get_chunk_by_id("base-old") is None + + +@pytest.mark.asyncio +async def test_run_index_uses_runtime_storage_base_for_git_repositories(tmp_path, monkeypatch): + repo = init_repo(tmp_path / "repo") + cfg = MagicMock() + cfg.storage_path = str(tmp_path / "storage") + expected_storage = _runtime_storage_base(cfg, repo) + indexing_kwargs: list[dict] = [] + + async def fake_run_indexing(*args, **kwargs): + from context_engine.indexer.pipeline import IndexResult + + indexing_kwargs.append(kwargs) + return IndexResult() + + monkeypatch.setattr("context_engine.indexer.pipeline.run_indexing", fake_run_indexing) + + await _run_index(cfg, str(repo)) + + assert indexing_kwargs[0]["storage_base_override"] == expected_storage + + def test_apply_output_compression_appends_directive(tmp_path): """When level != off, the helper appends the directive and bumps the bucket.""" server = _make_server(tmp_path) diff --git a/tests/retrieval/test_budgeting.py b/tests/retrieval/test_budgeting.py new file mode 100644 index 0000000..e2d9f7b --- /dev/null +++ b/tests/retrieval/test_budgeting.py @@ -0,0 +1,16 @@ +"""Tests for progressive disclosure budgets.""" + +from context_engine.retrieval import resolve_context_budget + + +def test_resolve_context_budget_uses_level_default(): + assert resolve_context_budget(level="summary") == 1500 + assert resolve_context_budget(level="standard") == 6000 + + +def test_resolve_context_budget_honors_explicit_lower_cap(): + assert resolve_context_budget(level="full", max_tokens=2000) == 2000 + + +def test_resolve_context_budget_keeps_level_when_explicit_cap_is_higher(): + assert resolve_context_budget(level="summary", max_tokens=9000) == 1500 diff --git a/tests/retrieval/test_fusion.py b/tests/retrieval/test_fusion.py new file mode 100644 index 0000000..4ec98f7 --- /dev/null +++ b/tests/retrieval/test_fusion.py @@ -0,0 +1,127 @@ +"""Tests for semantic/structural retrieval fusion.""" + +from context_engine.git.diff import GitWorktreeDiff +from context_engine.models import Chunk, ChunkType +from context_engine.retrieval import RetrievalFusionInput, fuse_retrieval_context +from context_engine.structural import SourceRange, StructuralContext + + +def chunk(chunk_id: str, file_path: str, content: str = "x") -> Chunk: + return Chunk( + id=chunk_id, + content=content, + chunk_type=ChunkType.FUNCTION, + file_path=file_path, + start_line=1, + end_line=2, + language="python", + ) + + +def ranked_chunk(chunk_id: str, file_path: str, *, distance: float) -> Chunk: + item = chunk(chunk_id, file_path) + item.metadata["_distance"] = distance + return item + + +def test_fusion_suppresses_base_chunks_from_modified_and_deleted_paths(): + result = fuse_retrieval_context( + RetrievalFusionInput( + base_chunks=[ + chunk("base_changed", "changed.py"), + chunk("base_deleted", "deleted.py"), + chunk("base_stable", "stable.py"), + ], + overlay_chunks=[chunk("overlay_changed", "changed.py")], + diff=GitWorktreeDiff(modified={"changed.py"}, deleted={"deleted.py"}), + ) + ) + + assert {item.id for item in result.chunks} == {"overlay_changed", "base_stable"} + + +def test_fusion_preserves_relevance_across_overlay_and_base_results(): + result = fuse_retrieval_context( + RetrievalFusionInput( + base_chunks=[ranked_chunk("base_strong", "stable.py", distance=0.1)], + overlay_chunks=[ranked_chunk("overlay_weak", "new.py", distance=1.2)], + diff=GitWorktreeDiff(added={"new.py"}), + ) + ) + + assert [item.id for item in result.chunks] == ["base_strong", "overlay_weak"] + + +def test_fusion_suppresses_base_chunks_from_recreated_paths(): + result = fuse_retrieval_context( + RetrievalFusionInput( + base_chunks=[chunk("base_recreated", "recreated.py")], + overlay_chunks=[chunk("overlay_recreated", "recreated.py")], + diff=GitWorktreeDiff(added={"recreated.py"}), + ) + ) + + assert [item.id for item in result.chunks] == ["overlay_recreated"] + + +def test_fusion_dedupes_overlay_and_base_source_ranges(): + result = fuse_retrieval_context( + RetrievalFusionInput( + base_chunks=[chunk("base", "same.py")], + overlay_chunks=[chunk("overlay", "same.py")], + ) + ) + + assert [item.id for item in result.chunks] == ["overlay"] + + +def test_fusion_respects_hard_token_budget_and_counts_omitted_chunks(): + result = fuse_retrieval_context( + RetrievalFusionInput( + overlay_chunks=[ + chunk("small", "a.py", "short"), + chunk("large", "b.py", "x" * 200), + ], + ), + max_tokens=10, + ) + + assert [item.id for item in result.chunks] == ["small"] + assert result.omitted == 1 + + +def test_fusion_summary_level_applies_default_budget(): + result = fuse_retrieval_context( + RetrievalFusionInput(overlay_chunks=[chunk("large", "a.py", "x" * 6000)]), + level="summary", + ) + + assert result.chunks == [] + assert result.omitted == 1 + + +def test_fusion_preserves_structural_context(): + structural = StructuralContext(sources=[SourceRange("a.py", 1, 2)], provider="test") + + result = fuse_retrieval_context(RetrievalFusionInput(structural=structural)) + + assert result.structural is structural + + +def test_fusion_counts_structural_context_against_token_budget(): + structural = StructuralContext( + sources=[SourceRange("a.py", 1, 2, content="x" * 40)], + provider="test", + ) + + result = fuse_retrieval_context( + RetrievalFusionInput( + overlay_chunks=[chunk("semantic", "semantic.py", "short")], + structural=structural, + ), + max_tokens=10, + ) + + assert result.chunks == [] + assert result.omitted == 1 + assert result.structural is structural diff --git a/tests/storage/test_worktree_overlay.py b/tests/storage/test_worktree_overlay.py new file mode 100644 index 0000000..d0358a2 --- /dev/null +++ b/tests/storage/test_worktree_overlay.py @@ -0,0 +1,267 @@ +"""Tests for semantic base+worktree overlay backend.""" + +import pytest + +from context_engine.git.diff import GitWorktreeDiff +from context_engine.models import Chunk, ChunkType, EdgeType, GraphNode, NodeType +from context_engine.storage.worktree_overlay import WorktreeOverlayBackend + + +def chunk(chunk_id: str, file_path: str, score: float = 0.5) -> Chunk: + return Chunk( + id=chunk_id, + content=chunk_id, + chunk_type=ChunkType.FUNCTION, + file_path=file_path, + start_line=1, + end_line=2, + language="python", + confidence_score=score, + ) + + +def node(node_id: str, file_path: str) -> GraphNode: + return GraphNode( + id=node_id, + node_type=NodeType.FUNCTION, + name=node_id, + file_path=file_path, + ) + + +class StubBackend: + def __init__( + self, + chunks: list[Chunk] | None = None, + *, + neighbors: list[GraphNode] | None = None, + related_paths: list[str] | None = None, + ) -> None: + self.chunks = chunks or [] + self.neighbors = neighbors or [] + self.related_paths = related_paths or ["changed.py", "stable.py"] + self.deleted: list[str] = [] + self.cache: dict[tuple[str, str], str] = {} + + async def ingest(self, chunks, nodes, edges): + self.chunks.extend(chunks) + + async def vector_search(self, query_embedding, top_k=10, filters=None): + return self.chunks[:top_k] + + async def fts_search(self, query, top_k=30): + return [(item.id, -rank) for rank, item in enumerate(self.chunks[:top_k])] + + async def graph_neighbors(self, node_id, edge_type=None): + return self.neighbors + + async def get_chunk_by_id(self, chunk_id): + return next((item for item in self.chunks if item.id == chunk_id), None) + + async def get_chunks_by_ids(self, chunk_ids): + wanted = set(chunk_ids) + return [item for item in self.chunks if item.id in wanted] + + async def delete_by_file(self, file_path): + self.deleted.append(file_path) + + async def get_related_file_paths(self, file_paths): + return self.related_paths + + def count_chunks(self): + return len(self.chunks) + + def get_cached_compression(self, chunk_id, level): + return self.cache.get((chunk_id, level)) + + def put_cached_compression(self, chunk_id, level, content): + self.cache[(chunk_id, level)] = content + + +@pytest.mark.asyncio +async def test_vector_search_suppresses_base_chunks_from_modified_and_deleted_paths(): + base = StubBackend([ + chunk("base_changed", "changed.py", 0.99), + chunk("base_deleted", "deleted.py", 0.98), + chunk("base_stable", "stable.py", 0.2), + ]) + overlay = StubBackend([chunk("overlay_changed", "changed.py", 0.5)]) + backend = WorktreeOverlayBackend( + base=base, + overlay=overlay, + diff=GitWorktreeDiff(modified={"changed.py"}, deleted={"deleted.py"}), + ) + + results = await backend.vector_search([0.1], top_k=10) + + assert [item.id for item in results] == ["overlay_changed", "base_stable"] + + +@pytest.mark.asyncio +async def test_vector_search_preserves_relevance_across_overlay_and_base_results(): + base_chunk = chunk("base_strong", "stable.py", 0.99) + base_chunk.metadata["_distance"] = 0.1 + overlay_chunk = chunk("overlay_weak", "new.py", 0.2) + overlay_chunk.metadata["_distance"] = 1.2 + backend = WorktreeOverlayBackend( + base=StubBackend([base_chunk]), + overlay=StubBackend([overlay_chunk]), + diff=GitWorktreeDiff(added={"new.py"}), + ) + + results = await backend.vector_search([0.1], top_k=10) + + assert [item.id for item in results] == ["base_strong", "overlay_weak"] + + +@pytest.mark.asyncio +async def test_vector_search_suppresses_base_chunks_from_recreated_paths(): + backend = WorktreeOverlayBackend( + base=StubBackend([chunk("base_recreated", "recreated.py")]), + overlay=StubBackend([chunk("overlay_recreated", "recreated.py")]), + diff=GitWorktreeDiff(added={"recreated.py"}), + ) + + results = await backend.vector_search([0.1], top_k=10) + + assert [item.id for item in results] == ["overlay_recreated"] + + +@pytest.mark.asyncio +async def test_fts_search_filters_stale_base_ids_before_hydration(): + base = StubBackend([chunk("base_changed", "changed.py"), chunk("base_stable", "stable.py")]) + overlay = StubBackend([chunk("overlay_changed", "changed.py")]) + backend = WorktreeOverlayBackend( + base=base, + overlay=overlay, + diff=GitWorktreeDiff(modified={"changed.py"}), + ) + + results = await backend.fts_search("changed", top_k=10) + + assert [chunk_id for chunk_id, _ in results] == ["overlay_changed", "base_stable"] + + +@pytest.mark.asyncio +async def test_fts_search_preserves_ranked_base_hits_ahead_of_weak_overlay_hits(): + base = StubBackend([ + chunk("base_strong", "stable.py"), + chunk("base_second", "stable2.py"), + ]) + overlay = StubBackend([chunk("overlay_weak", "new.py")]) + backend = WorktreeOverlayBackend( + base=base, + overlay=overlay, + diff=GitWorktreeDiff(added={"new.py"}), + ) + + results = await backend.fts_search("query", top_k=10) + + assert [chunk_id for chunk_id, _ in results][:2] == ["base_strong", "overlay_weak"] + + +@pytest.mark.asyncio +async def test_get_chunk_by_id_never_returns_shadowed_base_chunk(): + backend = WorktreeOverlayBackend( + base=StubBackend([chunk("base_changed", "changed.py")]), + overlay=StubBackend([]), + diff=GitWorktreeDiff(modified={"changed.py"}), + ) + + assert await backend.get_chunk_by_id("base_changed") is None + + +@pytest.mark.asyncio +async def test_update_diff_refreshes_shadowed_paths(): + backend = WorktreeOverlayBackend( + base=StubBackend([chunk("base_changed", "changed.py")]), + overlay=StubBackend([]), + diff=GitWorktreeDiff(), + ) + + assert await backend.get_chunk_by_id("base_changed") is not None + + backend.update_diff(GitWorktreeDiff(modified={"changed.py"})) + + assert await backend.get_chunk_by_id("base_changed") is None + + +def test_count_chunks_requires_overlay_when_worktree_has_changed_paths(): + backend = WorktreeOverlayBackend( + base=StubBackend([chunk("base_changed", "changed.py")]), + overlay=StubBackend([]), + diff=GitWorktreeDiff(modified={"changed.py"}), + ) + + assert backend.count_chunks() == 0 + + +@pytest.mark.asyncio +async def test_ingest_and_delete_target_overlay_only(): + base = StubBackend([]) + overlay = StubBackend([]) + backend = WorktreeOverlayBackend(base=base, overlay=overlay, diff=GitWorktreeDiff()) + + await backend.ingest([chunk("overlay", "new.py")], [], []) + await backend.delete_by_file("new.py") + + assert [item.id for item in overlay.chunks] == ["overlay"] + assert overlay.deleted == ["new.py"] + assert base.chunks == [] + + +@pytest.mark.asyncio +async def test_related_paths_suppress_deleted_base_paths(): + backend = WorktreeOverlayBackend( + base=StubBackend([]), + overlay=StubBackend([]), + diff=GitWorktreeDiff(deleted={"changed.py"}), + ) + + assert await backend.get_related_file_paths(["seed.py"]) == ["stable.py"] + + +@pytest.mark.asyncio +async def test_graph_neighbors_include_overlay_only_relationships(): + backend = WorktreeOverlayBackend( + base=StubBackend([chunk("base", "stable.py")], neighbors=[node("stable", "stable.py")]), + overlay=StubBackend( + [chunk("overlay", "new.py")], + neighbors=[node("changed", "changed.py")], + related_paths=["changed.py"], + ), + diff=GitWorktreeDiff(modified={"changed.py"}, added={"new.py"}), + ) + + neighbors = await backend.graph_neighbors("new.py:caller", EdgeType.CALLS) + related_paths = await backend.get_related_file_paths(["new.py"]) + + assert [item.id for item in neighbors] == ["changed"] + assert related_paths == ["changed.py"] + + +@pytest.mark.asyncio +async def test_graph_expansion_drops_base_outgoing_edges_for_modified_sources(): + backend = WorktreeOverlayBackend( + base=StubBackend(related_paths=["old_auth.py"]), + overlay=StubBackend(related_paths=["new_auth.py"]), + diff=GitWorktreeDiff(modified={"auth.py"}), + ) + + related_paths = await backend.get_related_file_paths(["auth.py"]) + + assert related_paths == ["new_auth.py"] + + +def test_compression_cache_reads_overlay_then_base_and_writes_overlay(): + base = StubBackend([]) + overlay = StubBackend([]) + base.cache[("chunk", "standard")] = "base cached" + backend = WorktreeOverlayBackend(base=base, overlay=overlay, diff=GitWorktreeDiff()) + + assert backend.get_cached_compression("chunk", "standard") == "base cached" + + backend.put_cached_compression("chunk", "standard", "overlay cached") + + assert backend.get_cached_compression("chunk", "standard") == "overlay cached" + assert overlay.cache[("chunk", "standard")] == "overlay cached" diff --git a/tests/structural/test_codegraph_provider.py b/tests/structural/test_codegraph_provider.py new file mode 100644 index 0000000..1b7f642 --- /dev/null +++ b/tests/structural/test_codegraph_provider.py @@ -0,0 +1,129 @@ +"""Tests for CodeGraph structural provider adapter.""" + +import json +from pathlib import Path +from unittest.mock import patch + +import pytest + +from context_engine.structural import CodeGraphBaseProvider, CodeGraphClient + + +class CompletedProcess: + def __init__(self, stdout: object, returncode: int = 0, stderr: str = "") -> None: + self.stdout = json.dumps(stdout) + self.returncode = returncode + self.stderr = stderr + + def communicate(self, timeout=None): + self.stdout_file.write(self.stdout) + self.stdout_file.flush() + return None, self.stderr + + +def test_client_runs_codegraph_json_without_shell(): + process = CompletedProcess({"initialized": True}) + with patch("subprocess.Popen", side_effect=bind_process(process)) as popen: + result = CodeGraphClient(executable="cg", timeout_seconds=3).status(Path("/repo")) + + assert result == {"initialized": True} + args, kwargs = popen.call_args + assert args[0] == ["cg", "status", "--json"] + assert kwargs["cwd"] == Path("/repo") + assert kwargs["text"] is True + assert "shell" not in kwargs + + +def test_client_raises_short_error_on_nonzero_exit(): + process = CompletedProcess({}, returncode=1, stderr="missing index") + with patch("subprocess.Popen", side_effect=bind_process(process)): + with pytest.raises(RuntimeError, match="missing index"): + CodeGraphClient().status(Path("/repo")) + + +def test_client_reports_oversized_json_without_parsing_truncated_stdout(): + process = CompletedProcess({"items": ["x" * 200]}) + with patch("subprocess.Popen", side_effect=bind_process(process)): + with pytest.raises(RuntimeError, match="response too large"): + CodeGraphClient(stdout_limit=10).status(Path("/repo")) + + +def bind_process(process): + def _popen(*args, **kwargs): + process.stdout_file = kwargs["stdout"] + return process + + return _popen + + +@pytest.mark.asyncio +async def test_provider_status_reports_complete_index_available(): + provider = CodeGraphBaseProvider(client=StubClient(status={"initialized": True, "index": {"state": "complete"}})) + + status = await provider.status(Path("/repo")) + + assert status.available is True + assert status.provider == "codegraph" + + +@pytest.mark.asyncio +async def test_provider_status_reports_indexing_unavailable(): + provider = CodeGraphBaseProvider(client=StubClient(status={"initialized": True, "index": {"state": "indexing"}})) + + status = await provider.status(Path("/repo")) + + assert status.available is False + + +@pytest.mark.asyncio +async def test_provider_status_reports_degraded_state_unavailable(): + provider = CodeGraphBaseProvider(client=StubClient(status={"initialized": True, "index": {"state": "degraded"}})) + + status = await provider.status(Path("/repo")) + + assert status.available is False + + +@pytest.mark.asyncio +async def test_provider_maps_query_rows_to_source_ranges(): + provider = CodeGraphBaseProvider( + client=StubClient(query=[ + {"node": {"filePath": "src/auth.py", "startLine": 4, "endLine": 9, "code": "def login(): ..."}} + ]) + ) + + context = await provider.explore("login", Path("/repo")) + + assert context.provider == "codegraph" + assert context.sources[0].path == "src/auth.py" + assert context.sources[0].start_line == 4 + assert context.sources[0].content == "def login(): ..." + + +@pytest.mark.asyncio +async def test_provider_maps_impact_rows_to_symbol_keys(): + provider = CodeGraphBaseProvider( + client=StubClient(impact={"affected": [{"name": "login", "kind": "function", "filePath": "src/auth.py"}]}) + ) + + context = await provider.impact("login", Path("/repo")) + + assert context.impact[0].qualified_name == "login" + assert context.impact[0].kind == "function" + assert context.impact[0].path == "src/auth.py" + + +class StubClient: + def __init__(self, *, status=None, query=None, impact=None): + self._status = status or {} + self._query = query or [] + self._impact = impact or {} + + def status(self, project_root): + return self._status + + def query(self, project_root, query, *, limit=20): + return self._query + + def impact(self, project_root, symbol, *, depth=2): + return self._impact diff --git a/tests/structural/test_merge.py b/tests/structural/test_merge.py new file mode 100644 index 0000000..b4df234 --- /dev/null +++ b/tests/structural/test_merge.py @@ -0,0 +1,115 @@ +"""Tests for structural base+overlay merge semantics.""" + +from context_engine.git.diff import GitWorktreeDiff +from context_engine.structural import ( + Relationship, + SourceRange, + StructuralContext, + SymbolKey, + merge_structural_contexts, +) + + +def symbol(name: str, path: str) -> SymbolKey: + return SymbolKey(qualified_name=name, kind="function", path=path) + + +def source(path: str) -> SourceRange: + return SourceRange(path=path, start_line=1, end_line=3, content=path) + + +def relationship(source_symbol: SymbolKey, target_symbol: SymbolKey) -> Relationship: + return Relationship(source=source_symbol, target=target_symbol, kind="calls") + + +def test_modified_file_sources_shadow_base_sources(): + merged = merge_structural_contexts( + base=StructuralContext( + sources=[source("auth.py"), source("stable.py")], + provider="base", + ), + overlay=StructuralContext( + sources=[SourceRange("auth.py", 1, 4, "new auth")], + provider="overlay", + ), + diff=GitWorktreeDiff(modified={"auth.py"}), + ) + + assert [item.content for item in merged.sources] == ["new auth", "stable.py"] + + +def test_deleted_file_sources_and_symbols_are_tombstoned(): + deleted = symbol("deleted", "deleted.py") + stable = symbol("stable", "stable.py") + + merged = merge_structural_contexts( + base=StructuralContext( + sources=[source("deleted.py"), source("stable.py")], + impact=[deleted, stable], + provider="base", + ), + overlay=StructuralContext(provider="overlay"), + diff=GitWorktreeDiff(deleted={"deleted.py"}), + ) + + assert [item.path for item in merged.sources] == ["stable.py"] + assert merged.impact == [stable] + + +def test_overlay_outgoing_edges_replace_base_outgoing_edges_for_same_symbol(): + base_b = symbol("B", "b.py") + c = symbol("C", "c.py") + d = symbol("D", "d.py") + e = symbol("E", "e.py") + + merged = merge_structural_contexts( + base=StructuralContext( + relationships=[relationship(base_b, c), relationship(base_b, d)], + provider="base", + ), + overlay=StructuralContext( + relationships=[relationship(base_b, c), relationship(base_b, e)], + provider="overlay", + ), + diff=GitWorktreeDiff(modified={"b.py"}), + ) + + assert merged.relationships == [relationship(base_b, c), relationship(base_b, e)] + + +def test_unmodified_incoming_base_edge_survives_when_target_replaced(): + a = symbol("A", "a.py") + b = symbol("B", "b.py") + c = symbol("C", "c.py") + + merged = merge_structural_contexts( + base=StructuralContext( + relationships=[relationship(a, b), relationship(b, c)], + provider="base", + ), + overlay=StructuralContext( + relationships=[relationship(b, c)], + provider="overlay", + ), + diff=GitWorktreeDiff(modified={"b.py"}), + ) + + assert relationship(a, b) in merged.relationships + assert relationship(b, c) in merged.relationships + assert len(merged.relationships) == 2 + + +def test_unmodified_incoming_base_edge_drops_when_target_deleted(): + stable = symbol("A", "stable.py") + deleted = symbol("B", "deleted.py") + + merged = merge_structural_contexts( + base=StructuralContext( + relationships=[relationship(stable, deleted)], + provider="base", + ), + overlay=StructuralContext(provider="overlay"), + diff=GitWorktreeDiff(deleted={"deleted.py"}), + ) + + assert merged.relationships == [] diff --git a/tests/test_config.py b/tests/test_config.py index 86c6eae..910ba68 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -100,6 +100,20 @@ def test_marginal_ratio_default(): assert Config().retrieval_marginal_ratio == 0.75 +def test_structural_config_defaults(): + config = Config() + assert config.structural_provider == "off" + assert config.structural_codegraph_executable == "codegraph" + + +def test_structural_config_yaml_mapping(tmp_path): + cfg_file = tmp_path / "config.yaml" + cfg_file.write_text("structural:\n provider: codegraph\n codegraph_executable: cg\n") + config = load_config(global_path=cfg_file) + assert config.structural_provider == "codegraph" + assert config.structural_codegraph_executable == "cg" + + def test_serve_config_defaults(): config = Config() assert config.serve_idle_timeout_minutes == 30 diff --git a/tests/test_git_repository_context.py b/tests/test_git_repository_context.py new file mode 100644 index 0000000..ce7a1b2 --- /dev/null +++ b/tests/test_git_repository_context.py @@ -0,0 +1,153 @@ +"""Tests for Git repository/worktree identity and diff modeling.""" + +from pathlib import Path +import subprocess + +from context_engine.git import get_worktree_diff, resolve_git_repository_context +from context_engine.git.repository import resolve_base_sha + + +def git(cwd: Path, *args: str) -> str: + result = subprocess.run( + ["git", *args], + cwd=cwd, + capture_output=True, + text=True, + check=True, + ) + return result.stdout.strip() + + +def init_repo(path: Path) -> Path: + path.mkdir() + git(path, "init") + git(path, "config", "user.email", "test@example.com") + git(path, "config", "user.name", "Test") + git(path, "config", "commit.gpgsign", "false") + (path / "base.py").write_text("def base():\n return 'base'\n") + git(path, "add", ".") + git(path, "commit", "-m", "init") + return path + + +def test_repository_and_worktree_identity_are_separate_for_linked_worktree(tmp_path): + repo = init_repo(tmp_path / "repo") + worktree = tmp_path / "repo-linked" + git(repo, "worktree", "add", "-b", "feature", str(worktree)) + + repo_context = resolve_git_repository_context(repo) + worktree_context = resolve_git_repository_context(worktree) + + assert repo_context is not None + assert worktree_context is not None + assert repo_context.repository_id == worktree_context.repository_id + assert repo_context.git_common_dir == worktree_context.git_common_dir + assert repo_context.worktree_id != worktree_context.worktree_id + assert repo_context.worktree_root == repo.resolve() + assert worktree_context.worktree_root == worktree.resolve() + + +def test_repository_context_normalizes_symlinked_checkout(tmp_path): + repo = init_repo(tmp_path / "repo") + symlink = tmp_path / "repo-link" + symlink.symlink_to(repo, target_is_directory=True) + + repo_context = resolve_git_repository_context(repo) + symlink_context = resolve_git_repository_context(symlink) + + assert repo_context is not None + assert symlink_context is not None + assert repo_context == symlink_context + + +def test_repository_context_returns_none_outside_git(tmp_path): + assert resolve_git_repository_context(tmp_path) is None + + +def test_base_sha_uses_unambiguous_main_merge_base(tmp_path): + repo = init_repo(tmp_path / "repo") + base_sha = git(repo, "rev-parse", "HEAD") + git(repo, "checkout", "-b", "feature") + (repo / "feature.py").write_text("def feature():\n return 'feature'\n") + git(repo, "add", ".") + git(repo, "commit", "-m", "feature") + + actual = resolve_base_sha(repo, fallback_base_refs=("master",)) + + assert actual == base_sha + + +def test_base_sha_uses_remote_default_not_tracking_feature_branch(tmp_path): + remote = tmp_path / "remote.git" + remote.mkdir() + git(remote, "init", "--bare") + + repo = init_repo(tmp_path / "repo") + git(repo, "branch", "-M", "main") + main_sha = git(repo, "rev-parse", "HEAD") + git(repo, "remote", "add", "origin", str(remote)) + git(repo, "push", "-u", "origin", "main") + git(repo, "remote", "set-head", "origin", "main") + + git(repo, "checkout", "-b", "feature") + (repo / "feature.py").write_text("def feature():\n return 'pushed'\n") + git(repo, "add", ".") + git(repo, "commit", "-m", "pushed feature") + git(repo, "push", "-u", "origin", "feature") + (repo / "local.py").write_text("def local():\n return 'unpushed'\n") + git(repo, "add", ".") + git(repo, "commit", "-m", "unpushed feature") + + actual = resolve_base_sha(repo) + + assert actual == main_sha + + +def test_base_sha_falls_back_to_head_when_no_safe_base_exists(tmp_path): + repo = init_repo(tmp_path / "repo") + head_sha = git(repo, "rev-parse", "HEAD") + + actual = resolve_base_sha(repo, fallback_base_refs=()) + + assert actual == head_sha + + +def test_base_sha_falls_back_to_head_when_fallback_refs_disagree(tmp_path): + repo = init_repo(tmp_path / "repo") + git(repo, "checkout", "-b", "main") + (repo / "main.py").write_text("def main():\n return True\n") + git(repo, "add", ".") + git(repo, "commit", "-m", "main") + git(repo, "checkout", "-b", "feature") + head_sha = git(repo, "rev-parse", "HEAD") + + actual = resolve_base_sha(repo, fallback_base_refs=("main", "master")) + + assert actual == head_sha + + +def test_worktree_diff_collects_committed_dirty_untracked_deleted_and_renamed(tmp_path): + repo = init_repo(tmp_path / "repo") + base_sha = git(repo, "rev-parse", "HEAD") + + (repo / "committed.py").write_text("def committed():\n return True\n") + git(repo, "add", ".") + git(repo, "commit", "-m", "committed") + head_sha = git(repo, "rev-parse", "HEAD") + + (repo / "base.py").write_text("def base():\n return 'dirty'\n") + (repo / "staged.py").write_text("def staged():\n return True\n") + git(repo, "add", "staged.py") + (repo / "untracked.py").write_text("def untracked():\n return True\n") + git(repo, "mv", "committed.py", "renamed.py") + + diff = get_worktree_diff(repo, base_sha=base_sha, head_sha=head_sha) + + assert diff.base_sha == base_sha + assert diff.head_sha == head_sha + assert "base.py" in diff.modified + assert "staged.py" in diff.added + assert "untracked.py" in diff.added + assert "committed.py" in diff.deleted + assert "renamed.py" in diff.added + assert diff.renamed == {"committed.py": "renamed.py"} diff --git a/tests/test_git_storage_layout.py b/tests/test_git_storage_layout.py new file mode 100644 index 0000000..ad0d03b --- /dev/null +++ b/tests/test_git_storage_layout.py @@ -0,0 +1,66 @@ +"""Tests for repository-scoped CCE storage layout.""" + +from pathlib import Path +from types import SimpleNamespace + +from context_engine.git import repository_storage_layout, resolve_git_repository_context + +from tests.test_git_repository_context import git, init_repo + + +def config(storage_root: Path) -> SimpleNamespace: + return SimpleNamespace(storage_path=str(storage_root)) + + +def test_repository_storage_layout_splits_shared_base_from_worktree_overlay(tmp_path): + repo = init_repo(tmp_path / "repo") + worktree = tmp_path / "repo-linked" + git(repo, "worktree", "add", "-b", "feature", str(worktree)) + + repo_context = resolve_git_repository_context(repo) + worktree_context = resolve_git_repository_context(worktree) + assert repo_context is not None + assert worktree_context is not None + + cfg = config(tmp_path / ".cce" / "projects") + repo_layout = repository_storage_layout(cfg, repo, repo_context) + worktree_layout = repository_storage_layout(cfg, worktree, worktree_context) + + assert repo_layout.repository_root == worktree_layout.repository_root + assert repo_layout.repository_root == tmp_path / ".cce" / "repos" / repo_context.repository_id + assert repo_layout.base_dir == worktree_layout.base_dir + assert repo_layout.worktree_dir != worktree_layout.worktree_dir + assert repo_layout.worktree_dir.parent == worktree_layout.worktree_dir.parent + + +def test_repository_storage_layout_preserves_legacy_project_storage_reference(tmp_path): + repo = init_repo(tmp_path / "repo") + repo_context = resolve_git_repository_context(repo) + assert repo_context is not None + + cfg = config(tmp_path / ".cce" / "projects") + legacy = tmp_path / ".cce" / "projects" / "repo" + legacy.mkdir(parents=True) + (legacy / "marker.txt").write_text("legacy") + + layout = repository_storage_layout(cfg, repo, repo_context) + + assert layout.legacy_project_dir.exists() + assert (layout.legacy_project_dir / "marker.txt").read_text() == "legacy" + assert layout.base_dir == tmp_path / ".cce" / "repos" / repo_context.repository_id / "base" + + +def test_repository_storage_layout_can_resolve_without_migrating_legacy_storage(tmp_path): + repo = init_repo(tmp_path / "repo") + repo_context = resolve_git_repository_context(repo) + assert repo_context is not None + + cfg = config(tmp_path / ".cce" / "projects") + legacy = tmp_path / ".cce" / "projects" / "repo" + legacy.mkdir(parents=True) + + layout = repository_storage_layout(cfg, repo, repo_context, migrate_legacy=False) + + assert legacy.exists() + assert layout.legacy_project_dir.name.startswith("repo-") + assert not layout.legacy_project_dir.exists() diff --git a/tests/test_worktree_benchmark.py b/tests/test_worktree_benchmark.py new file mode 100644 index 0000000..f7476cb --- /dev/null +++ b/tests/test_worktree_benchmark.py @@ -0,0 +1,39 @@ +"""Tests for worktree overlay benchmark helper.""" + +from click.testing import CliRunner + +from context_engine.cli import main +from context_engine.git import benchmark_worktree_overlay +from tests.test_git_repository_context import git, init_repo + + +def test_worktree_benchmark_counts_only_diff_files(tmp_path): + repo = init_repo(tmp_path / "repo") + for index in range(20): + (repo / f"stable_{index}.py").write_text(f"VALUE = {index}\n") + git(repo, "add", ".") + git(repo, "commit", "-m", "stable files") + + (repo / "base.py").write_text("def base():\n return 'changed'\n") + (repo / "new.py").write_text("def new():\n return True\n") + + result = benchmark_worktree_overlay(repo) + + assert result is not None + assert result.changed_file_count == 2 + assert result.modified_count == 1 + assert result.added_count == 1 + + +def test_worktree_benchmark_returns_none_outside_git(tmp_path): + assert benchmark_worktree_overlay(tmp_path) is None + + +def test_worktree_benchmark_cli_outputs_json(tmp_path, monkeypatch): + repo = init_repo(tmp_path / "repo") + monkeypatch.chdir(repo) + + result = CliRunner().invoke(main, ["worktree-benchmark", "--json"]) + + assert result.exit_code == 0 + assert '"changed_file_count"' in result.output diff --git a/tests/test_worktree_doctor.py b/tests/test_worktree_doctor.py new file mode 100644 index 0000000..e76c612 --- /dev/null +++ b/tests/test_worktree_doctor.py @@ -0,0 +1,61 @@ +"""Tests for worktree-aware doctor diagnostics.""" + +from types import SimpleNamespace + +from click.testing import CliRunner + +from context_engine.cli import main +from context_engine.git import worktree_doctor_report +from tests.test_git_repository_context import git, init_repo + + +def config(tmp_path): + return SimpleNamespace(storage_path=str(tmp_path / ".cce" / "projects")) + + +def test_worktree_doctor_report_includes_repo_storage_and_overlay_counts(tmp_path): + repo = init_repo(tmp_path / "repo") + (repo / "base.py").write_text("def base():\n return 'changed'\n") + (repo / "new.py").write_text("def new():\n return True\n") + + report = worktree_doctor_report(config(tmp_path), repo) + + assert report["git"]["available"] is True + assert report["repository"]["common_dir"].endswith(".git") + assert report["worktree"]["root"] == str(repo.resolve()) + assert report["storage"]["base_dir"].endswith("/base") + assert report["storage"]["worktree_dir"].find("/worktrees/") >= 0 + assert report["overlay"]["modified_count"] == 1 + assert report["overlay"]["added_count"] == 1 + + +def test_worktree_doctor_report_does_not_migrate_legacy_project_storage(tmp_path): + repo = init_repo(tmp_path / "repo") + legacy = tmp_path / ".cce" / "projects" / "repo" + legacy.mkdir(parents=True) + (legacy / "marker.txt").write_text("legacy") + + report = worktree_doctor_report(config(tmp_path), repo) + + assert legacy.exists() + assert (legacy / "marker.txt").read_text() == "legacy" + assert report["storage"]["legacy_project_dir"] != str(legacy) + + +def test_worktree_doctor_report_handles_non_git_directory(tmp_path): + report = worktree_doctor_report(config(tmp_path), tmp_path) + + assert report == {"git": {"available": False, "project_dir": str(tmp_path.resolve())}} + + +def test_doctor_json_command_outputs_report(tmp_path, monkeypatch): + repo = init_repo(tmp_path / "repo") + git(repo, "status") + monkeypatch.chdir(repo) + + runner = CliRunner() + result = runner.invoke(main, ["doctor", "--json"]) + + assert result.exit_code == 0 + assert '"repository"' in result.output + assert '"worktree"' in result.output