From 1a5a125f15d6ae6731f9608a0e9bfce6ba87b997 Mon Sep 17 00:00:00 2001 From: Fazle Elahee Date: Fri, 3 Jul 2026 09:16:16 +0100 Subject: [PATCH 01/14] docs: token-savings program design spec --- docs/specs/2026-07-03-token-savings-design.md | 160 ++++++++++++++++++ 1 file changed, 160 insertions(+) create mode 100644 docs/specs/2026-07-03-token-savings-design.md diff --git a/docs/specs/2026-07-03-token-savings-design.md b/docs/specs/2026-07-03-token-savings-design.md new file mode 100644 index 0000000..a62c341 --- /dev/null +++ b/docs/specs/2026-07-03-token-savings-design.md @@ -0,0 +1,160 @@ +# Token-Savings Program — Design + +**Date:** 2026-07-03 +**Status:** Approved direction (all four phases), pending spec review +**Goal:** Make CCE measurably better at saving user tokens and reducing model usage — and make that value visible and credible to users. + +## Problem + +The 2026-07-03 full-project review found that CCE's core value proposition leaks in +four places: + +1. **Retrieval wastes tokens.** The retriever packs the token budget full regardless + of chunk quality; overlapping chunks duplicate content; recency scoring is inert + (chunk metadata never persisted); ranking bugs (fixed separately on + `fix/review-findings-2026-07-03`) let keyword-incidental chunks outrank semantic hits. +2. **Savings numbers aren't trustworthy.** The ledger inflates on compression + retries, `saved_pct` can go negative in one code path, and two endpoints compute + savings differently. +3. **Compression is all-or-nothing per level.** Every served chunk gets the same + treatment regardless of how confident retrieval is about it. +4. **Session startup re-explores.** The SessionStart resume exists but is not + token-capped or signal-ranked, so models still spend exploration rounds + re-learning the codebase. + +## Non-goals + +- No new ML dependencies (no cross-encoder reranker — adds latency and its own + energy cost, contrary to the environment goal). +- No remote-backend work (`RemoteBackend` is dead code; separate decision). +- No schema rework of the memory subsystem beyond what Phase 4 needs. + +## Phase 1 — Retrieval precision + +**Files:** `src/context_engine/retrieval/retriever.py`, +`src/context_engine/retrieval/confidence.py`, +`src/context_engine/storage/vector_store.py`, `src/context_engine/config.py`, +`benchmarks/`. + +Depends on the ranking fixes from `fix/review-findings-2026-07-03` (correct +`_distance` handling, per-token FTS queries, cosine metric). + +1. **Confidence cutoff.** New config `retrieval.min_confidence` (default `0.35`). + Chunks scoring below it are dropped before token packing. The top-1 result is + always kept (never return empty for a matching query). Overflow/`expand_chunk` + pointers are still emitted for dropped chunks so nothing becomes unreachable. +2. **Marginal-utility stop.** New config `retrieval.marginal_ratio` (default `0.5`). + After sorting by confidence, stop adding chunks once + `score < marginal_ratio * top_score`. Replaces "fill the budget because it's + there". Fixed `top_k` remains as the hard upper bound. +3. **Overlap dedup.** Two chunks from the same file whose line ranges overlap by + more than 50% collapse to the higher-confidence one before packing. +4. **Persist `modified_ts`.** Add a dedicated `modified_ts REAL` column (not a + generic metadata JSON blob) to the `chunks` table with a schema-version bump + and rebuild-safe migration, populated from file mtime at indexing. This activates the existing 0.1 recency weight in + `ConfidenceScorer` which currently always returns the neutral 0.5. +5. **Benchmark harness.** `benchmarks/retrieval_bench.py`: replay a committed query + set (~30 queries with expected-file labels, drawn from this repo itself) against + the index; report per-query and aggregate: tokens served, hit@k (expected file + present), and chunks served. Run before/after each phase; results table checked + into `benchmarks/results/`. This is the evidence standard for every later claim. + +**Success criteria:** tokens-served per query drops ≥ 25% on the benchmark set with +hit@k unchanged or better. + +## Phase 2 — Honest savings + environment report + +**Files:** `src/context_engine/cli.py`, `src/context_engine/dashboard/server.py`, +`src/context_engine/memory/hooks.py`, `src/context_engine/pricing.py`, docs. + +1. **One baseline definition, documented in code and docs:** + *baseline = tokens the model would have spent Reading the full files containing + the served chunks; saved = baseline − tokens actually served.* Savings recorded + only when a search response is actually served; never on retries, fallbacks, or + re-compressions (ledger-integrity fixes land on the fix branch; this phase adds + the definition and audits every `record_savings` call site against it). +2. **Clamp and unify.** All savings percentages computed by one shared helper, + clamped to `[0, 100]`; `/api/status`, `/api/savings`, the CLI banner, and the + badge all call it. +3. **Per-session summary.** At SessionEnd, write a one-line summary into the + session record; `cce sessions show` and the dashboard display it: + `saved 41,200 tokens this session (63% of baseline)`. +4. **`cce savings --report`.** Cumulative report: total tokens saved, cost saved + (existing pricing module), and an energy/CO₂-equivalent estimate using a + published per-token inference energy figure. Constants live in `pricing.py` + with a source citation and are printed with the report + (`estimates based on , `). Conservative rounding; labeled + "estimate" in every surface. If the number is small, it shows small. + +**Success criteria:** the same savings number appears on every surface; a +from-scratch session's reported savings can be manually reproduced from its +serve log. + +## Phase 3 — Compression tiers by confidence + +**Files:** `src/context_engine/retrieval/retriever.py`, +`src/context_engine/compression/compressor.py`, `src/context_engine/config.py`. + +1. **Tiering rule.** Config `retrieval.tier_full` (default `0.75`) and + `retrieval.tier_compressed` (default `0.5`): confidence ≥ `tier_full` serves + full source; between the two serves the existing compressed form; below + `tier_compressed` (but above Phase 1's `min_confidence`) serves a + signature-skeleton (def/class lines + docstring first line — derivable from the + chunker's existing structure, no LLM needed). Every non-full chunk carries its + `expand_chunk` id — the escape hatch already exists and is documented in + CLAUDE.md templates. +2. **Cache honesty.** Cached compressions record the producing method + (`llm` vs `truncation`); truncation-fallback output is not persisted under the + LLM cache key, so quality recovers automatically when Ollama comes back. +3. Tier thresholds are bypassed when the client explicitly sets + `set_output_compression` to a fixed level (existing behavior wins). + +**Success criteria:** additional ≥ 15% tokens-served reduction on the benchmark +set with hit@k unchanged; `expand_chunk` usage observed in real sessions stays +below ~1 in 5 queries (if models constantly expand, tiers are too aggressive — +thresholds are config, so tuning is cheap). + +## Phase 4 — Smarter session-start briefs + +**Files:** `src/context_engine/memory/hooks.py` (resume builder), +`src/context_engine/memory/db.py` (queries only). + +Depends on the memory-desync fixes (turn-summary triggers, vec bootstrap) from the +fix branch. + +1. **Hard token cap** on the SessionStart resume (config `memory.brief_tokens`, + default `500`), enforced by the same tokenizer used for savings accounting. +2. **Signal ranking within the cap:** binding decisions first, then recent + decisions (deduped), then hot files (code_areas touched in the last N sessions, + ranked by recency × frequency), then an unfinished-work marker (last session's + final turn summary if it ended mid-task). Grammar `expand()` is not applied to + user-authored text (mangle bug fixed on the fix branch). +3. **Measurement:** log brief size and, per session, the count of + `context_search`/Read calls in the first 5 turns — compare across sessions with + briefs on/off to estimate exploration rounds saved. + +**Success criteria:** brief ≤ cap in 100% of sessions; measurable drop in +first-5-turn exploration calls on this repo. + +## Sequencing and delivery + +Each phase is one branch/PR, in order 1 → 2 → 3 → 4, each carrying its +before/after benchmark table in the PR description. Phase 1's harness merges +first and gates the rest. All work builds on top of +`fix/review-findings-2026-07-03`. + +## Error handling + +- Cutoffs/tiers degrade toward current behavior: config values of `0` disable + each mechanism. +- Migration for `modified_ts` must be forward-only and tolerate old rows + (NULL → neutral recency 0.5, exactly today's behavior). +- The environment estimate must never block `cce savings` — pricing fetch + failures fall back to static constants (existing pattern). + +## Testing + +- Unit tests per mechanism (cutoff, marginal stop, dedup, tier selection, clamp + helper, brief cap/ranking). +- The benchmark harness doubles as the integration test for Phases 1 and 3. +- Savings reproducibility test: synthetic serve log → assert reported totals. From 538c8c89114f3fc835679c4c8533b2ff3c58b7cc Mon Sep 17 00:00:00 2001 From: Fazle Elahee Date: Fri, 3 Jul 2026 20:20:20 +0100 Subject: [PATCH 02/14] fix: resolve 2026-07-03 full-project review findings Config writers (data loss): - string-aware JSONC comment stripping; URLs in values no longer corrupt configs - parse failures skip the file with a warning instead of overwriting user MCP configs - non-dict server sections handled; uninstall matches CCE markers only, never bare "cce" - uninstall strips only the CCE block from shared git hooks Memory (index desync): - bootstrap without sqlite-vec no longer stamps CURRENT_VERSION; vec tables created on later connect - turn re-compression uses DELETE+INSERT so FTS/vec delete triggers fire - compression queue dead-letters after 5 attempts; savings recorded only on success - rollup enqueue dedup via -1 sentinel; PII scrub applied to prompt/tool-payload writes - non-numeric started_at no longer 500s Retrieval/storage (ranking + consistency): - FTS-only chunks get worst-case vector distance instead of perfect 0.0 - FTS queries tokenized per-term (OR-joined, stopwords dropped) instead of one phrase - chunks_vec created with distance_metric=cosine; legacy L2 tables detected and rebuilt - vector/FTS/graph stores roll back on mid-batch failure; FTS re-ingest no longer duplicates Indexer (correctness + secrets): - tree-sitter byte offsets slice bytes, not str; non-ASCII files chunk correctly - single-file target path (watcher) now honors is_secret_file and .cceignore - project_dir resolved once; symlinked roots no longer crash targeted reindex - deleted-but-tracked targets prune index + manifest - unquoted credential assignments redacted; *.env suffix files skipped - Stripe/GitHub vendor regexes fixed (capturing groups leaked full key values) Dashboard (XSS): - all API-sourced strings escaped (sessions, decisions, file paths, chart labels) - inline onclick handlers replaced with data attributes + listeners - state.json written atomically 1040 tests pass (~60 new regression tests, all confirmed failing pre-fix). --- src/context_engine/cli.py | 165 +++++++++++++++++---- src/context_engine/dashboard/_page.py | 36 +++-- src/context_engine/dashboard/server.py | 6 +- src/context_engine/editors.py | 111 +++++++++++--- src/context_engine/indexer/chunker.py | 54 ++++--- src/context_engine/indexer/pipeline.py | 70 ++++++++- src/context_engine/indexer/secrets.py | 37 ++++- src/context_engine/memory/compressor.py | 85 +++++++---- src/context_engine/memory/db.py | 7 +- src/context_engine/memory/hooks.py | 35 ++++- src/context_engine/retrieval/retriever.py | 18 ++- src/context_engine/storage/fts_store.py | 80 +++++++--- src/context_engine/storage/graph_store.py | 84 ++++++----- src/context_engine/storage/vector_store.py | 123 +++++++++------ tests/dashboard/test_page_escaping.py | 83 +++++++++++ tests/dashboard/test_server.py | 25 ++++ tests/indexer/test_chunker.py | 39 +++++ tests/indexer/test_pipeline_target_path.py | 159 ++++++++++++++++++++ tests/indexer/test_secrets.py | 83 +++++++++++ tests/memory/test_compressor.py | 151 +++++++++++++++++++ tests/memory/test_db.py | 56 +++++++ tests/memory/test_hooks.py | 113 ++++++++++++++ tests/retrieval/test_retriever.py | 96 ++++++++++++ tests/storage/test_fts_store.py | 81 ++++++++++ tests/storage/test_graph_store.py | 13 ++ tests/storage/test_vector_store.py | 89 +++++++++++ tests/test_cli_mcp_config.py | 52 +++++++ tests/test_cli_uninstall.py | 153 +++++++++++++++++++ tests/test_editors_opencode.py | 130 +++++++++++++++- 29 files changed, 2016 insertions(+), 218 deletions(-) create mode 100644 tests/dashboard/test_page_escaping.py create mode 100644 tests/indexer/test_pipeline_target_path.py diff --git a/src/context_engine/cli.py b/src/context_engine/cli.py index 93a0a6e..b88a5cf 100644 --- a/src/context_engine/cli.py +++ b/src/context_engine/cli.py @@ -139,11 +139,16 @@ def _show_update_notice() -> None: pass -def _configure_mcp(project_dir: Path) -> bool: +def _configure_mcp(project_dir: Path) -> bool | None: """Write MCP server config to .mcp.json in the project directory. - Returns True if the entry was added. Uses an atomic write so a crash or - partial write can't destroy pre-existing MCP server entries in the file. + Returns True if the entry was added/updated, False if already configured, + or None if an existing .mcp.json could not be parsed — in that case the + file is left untouched (overwriting it with `{}` plus our entry would + destroy the user's other MCP servers) and the caller should warn. + + Uses an atomic write so a crash or partial write can't destroy + pre-existing MCP server entries in the file. """ from context_engine.utils import atomic_write_text, resolve_cce_binary @@ -160,20 +165,29 @@ def _configure_mcp(project_dir: Path) -> bool: try: data = json.loads(mcp_path.read_text(encoding="utf-8")) except (json.JSONDecodeError, OSError): - data = {} + # Skip rather than clobber a file we couldn't read/parse. + return None + if not isinstance(data, dict): + return None else: data = {} - servers = data.setdefault("mcpServers", {}) - if "context-engine" in servers: - existing = servers["context-engine"] - if existing.get("command") == command and existing.get("args") == entry["args"]: - return False # already configured and up to date - # Update stale command path or args (e.g. after package rename). - servers["context-engine"] = entry - atomic_write_text(mcp_path, json.dumps(data, indent=2) + "\n") - return True - + # A non-dict value ("mcpServers": null, or a list) can't hold our entry — + # replace it with a fresh dict instead of raising TypeError. + servers = data.get("mcpServers") + if not isinstance(servers, dict): + servers = {} + data["mcpServers"] = servers + + existing = servers.get("context-engine") + if ( + isinstance(existing, dict) + and existing.get("command") == command + and existing.get("args") == entry["args"] + ): + return False # already configured and up to date + + # Add, or update stale command path/args (e.g. after package rename). servers["context-engine"] = entry atomic_write_text(mcp_path, json.dumps(data, indent=2) + "\n") return True @@ -899,7 +913,12 @@ def init(ctx: click.Context, agent: str) -> None: editor_targets = _init_editor_targets(project_dir, agent) if "claude" in editor_targets: configured = _configure_mcp(project_dir) - if configured: + if configured is None: + _warn( + "MCP config skipped — .mcp.json exists but could not be parsed; " + "fix or remove it and re-run `cce init`" + ) + elif configured: _ok("MCP server registered in " + click.style(".mcp.json", fg="cyan")) else: _ok("MCP server already configured in " + click.style(".mcp.json", fg="cyan")) @@ -2296,11 +2315,68 @@ async def _search(): asyncio.run(_search()) +def _strip_cce_git_hook_block(content: str, marker: str) -> str | None: + """Remove the CCE-installed block (the marker line plus the command line + that follows it) from a git hook script. + + Returns the remaining script text, or None when nothing meaningful is + left (only the shebang and blank lines) — meaning the file was created + by CCE and should be deleted outright. + """ + lines = content.splitlines() + kept: list[str] = [] + skip_next = False + for line in lines: + if skip_next: + skip_next = False + continue + if marker in line: + skip_next = True # drop the `cce index ... &` line too + continue + kept.append(line) + meaningful = [ + ln for ln in kept if ln.strip() and not ln.strip().startswith("#!") + ] + if not meaningful: + return None + return "\n".join(kept).rstrip() + "\n" + + +def _is_cce_settings_hook_command(cmd: str) -> bool: + """True only for hook commands CCE itself installs — never for commands + that merely contain "cce" as a substring ("access", "success", ...). + + CCE installs exactly two command shapes: + - memory lifecycle hooks: `/cce_hook.sh ` + (see memory/hook_installer.install_settings) + - SessionStart status hook: ` status --oneline` + (see _ensure_session_hook) + plus legacy invocations of the `code-context-engine` binary. + """ + import re + + from context_engine.memory.hook_installer import HOOK_MARKER as _MEM_HOOK_MARKER + + if _MEM_HOOK_MARKER in cmd: # "cce_hook" + return True + if "context-engine" in cmd: + return True + # `cce ` where cce may be a bare name or an absolute path + # (possibly quoted, possibly cce.exe/cce.cmd on Windows). + return bool( + re.search( + r"""(?:^|[\\/\s"'])cce(?:\.exe|\.cmd)?["']?\s+(?:status|index|serve|sessions)\b""", + cmd, + ) + ) + + @main.command() @click.option("--yes", "-y", is_flag=True, help="Skip confirmation prompt") def uninstall(yes: bool) -> None: """Remove CCE from the current project (hooks, .mcp.json entry, CLAUDE.md block).""" from context_engine.cli_style import section, animate, dim, warn, CROSS, DOT + from context_engine.indexer.git_hooks import HOOK_MARKER as _GIT_HOOK_MARKER project_dir = _safe_cwd() project_name = project_dir.name @@ -2315,7 +2391,11 @@ def uninstall(yes: bool) -> None: lines.append(section(f"Uninstall · {project_name}")) lines.append("") - # Remove git hooks + # Remove git hooks. Match only the exact marker CCE installs (or the + # unambiguous "context-engine" binary name for legacy hooks) — a bare + # "cce" substring also matches user hooks containing words like + # "access" or "success". When CCE appended its block to a pre-existing + # user hook, strip only that block instead of deleting the whole file. hooks_dir = project_dir / ".git" / "hooks" removed_hooks = 0 if hooks_dir.exists(): @@ -2323,7 +2403,16 @@ def uninstall(yes: bool) -> None: hook_file = hooks_dir / hook_name if hook_file.exists(): content = hook_file.read_text(encoding="utf-8") - if "cce" in content.lower() or "context-engine" in content.lower(): + if _GIT_HOOK_MARKER in content: + remaining_hook = _strip_cce_git_hook_block(content, _GIT_HOOK_MARKER) + if remaining_hook is None: + hook_file.unlink() + else: + hook_file.write_text(remaining_hook, encoding="utf-8") + removed_hooks += 1 + elif "context-engine" in content: + # Legacy CCE hooks predate the marker but invoked the + # code-context-engine binary directly. hook_file.unlink() removed_hooks += 1 if removed_hooks: @@ -2409,10 +2498,12 @@ def uninstall(yes: bool) -> None: changed = False for event in list(hooks.keys()): original = hooks[event] + # Match only the command shapes CCE installs — a bare "cce" + # substring also matches user hooks like "log-access ...". filtered = [ h for h in original if not any( - "cce" in cmd.get("command", "") + _is_cce_settings_hook_command(cmd.get("command", "")) for cmd in (h.get("hooks", []) if isinstance(h, dict) else []) ) ] @@ -2438,19 +2529,30 @@ def uninstall(yes: bool) -> None: except (json.JSONDecodeError, OSError): pass - # Remove CCE entries from .gitignore (including comment lines) + # Remove CCE entries from .gitignore (including comment lines). Match + # only the exact entries/comments CCE writes (see + # project_commands._GITIGNORE_ENTRIES / ensure_gitignore) or lines naming + # "context-engine" — never a bare "cce" substring, which also matches + # user lines like "access-logs/". gitignore = project_dir / ".gitignore" if gitignore.exists(): content = gitignore.read_text(encoding="utf-8") - if ".cce" in content or "context-engine" in content.lower() or "cce" in content.lower() or ".claude/settings.local.json" in content: - # These are the exact entries CCE adds (see project_commands._GITIGNORE_ENTRIES) - cce_lines = {".cce/", ".claude/settings.local.json"} - new_lines = [ - line for line in content.splitlines() - if line.strip() not in cce_lines - and "context-engine" not in line.lower() - and not (line.startswith("#") and ("cce" in line.lower() or "claude code local settings" in line.lower())) - ] + cce_entries = {".cce/", ".claude/settings.local.json"} + cce_comments = { + "# CCE (code-context-engine)", + "# CCE local cache (per-machine, not for version control)", + "# Claude Code local settings written by cce init", + "# .mcp.json contains absolute paths regenerated by `cce init`", + } + old_lines = content.splitlines() + new_lines = [ + line for line in old_lines + if line.strip() not in cce_entries + and line.strip() not in cce_comments + and "context-engine" not in line.lower() + and not (line.startswith("#") and "claude code local settings" in line.lower()) + ] + if len(new_lines) != len(old_lines): new_content = "\n".join(new_lines).strip() if new_content: gitignore.write_text(new_content + "\n", encoding="utf-8") @@ -2606,7 +2708,12 @@ def upgrade(ctx: click.Context, check: bool) -> None: click.echo("") click.echo(f" {click.style('Refreshing project config', fg='cyan')}...") configured = _configure_mcp(project_dir) - if configured: + if configured is None: + _warn( + "MCP config skipped — .mcp.json could not be parsed; " + "fix or remove it and re-run `cce init`" + ) + elif configured: _ok("MCP server paths updated in " + click.style(".mcp.json", fg="cyan")) else: _ok("MCP server config is current") diff --git a/src/context_engine/dashboard/_page.py b/src/context_engine/dashboard/_page.py index bd868f1..64e9048 100644 --- a/src/context_engine/dashboard/_page.py +++ b/src/context_engine/dashboard/_page.py @@ -1007,7 +1007,7 @@ var pct = max > 0 ? item.value/max*100 : 0; var name = item.label.split('/').pop() || item.label; return '
' - +'
'+name+'
' + +'
'+_esc(name)+'
' +'
' +'
'+item.value+'
' +'
'; @@ -1032,7 +1032,7 @@ +''; }).join(''); var labels = items.map(function(item) { - return '
'+item.label+'
'; + return '
'+_esc(item.label)+'
'; }).join(''); el.innerHTML = '
'+bars+'
' @@ -1082,7 +1082,7 @@ box.innerHTML = rows.map(function(s) { var rollup = (s.rollup_summary || '').slice(0, 80); return '' - + '
' + + '
' + '
'+_esc(s.id)+'
' + '
'+_esc(s.started_at||'')+'
' + '
'+_esc(s.status||'')+'
' @@ -1090,6 +1090,12 @@ + '
'+_esc(rollup)+(rollup && s.rollup_summary && s.rollup_summary.length>80?'…':'')+'
' + '
'; }).join(''); + // data-sid + listener instead of an inline onclick: _esc() is an HTML + // escaper, and inside onclick="..." its entities decode back to quotes + // before the JS parser runs, so it cannot protect a JS-string context. + box.querySelectorAll('.mem-session-row').forEach(function(row) { + row.addEventListener('click', function() { loadMemoryTimeline(row.dataset.sid); }); + }); } catch(e) { document.getElementById('mem-sessions-rows').innerHTML = '
Memory store unavailable.
'; @@ -1115,7 +1121,7 @@ var turns = data.turns.map(function(t) { return '' + '
' - + '
turn '+t.prompt_number+' · ['+_esc(t.tier)+']
' + + '
turn '+_esc(t.prompt_number)+' · ['+_esc(t.tier)+']
' + '
'+_esc(t.summary)+'
' + '
'; }).join(''); @@ -1347,15 +1353,23 @@ } el.innerHTML = files.map(function(f) { return '
' - +'
'+f.path+'
' - +'
'+f.chunks+'
' - +'
'+f.status+'
' + +'
'+_esc(f.path)+'
' + +'
'+_esc(f.chunks)+'
' + +'
'+_esc(f.status)+'
' +'
' - +'' - +'' + +'' + +'' +'
' +'
'; }).join(''); + // data-path + listener (not inline onclick) so paths containing quotes or + // angle brackets can never escape into HTML-attribute or JS context. + el.querySelectorAll('button[data-path]').forEach(function(btn) { + btn.addEventListener('click', function() { + if (btn.dataset.act === 'reindex') reindexFile(btn.dataset.path); + else deleteFile(btn.dataset.path); + }); + }); } function filterFiles(q) { @@ -1387,7 +1401,7 @@ +'
' +'
'+SVG.chevron+'
' +'
' - +'
'+(s.project||s.id)+'
' + +'
'+_esc(s.project||s.id)+'
' +'
' +''+decs.length+'decisions' +''+areas.length+'code areas' @@ -1399,7 +1413,7 @@ +(decs.length ?'
' +'
Decisions
' - +decs.map(function(d){ return '
'+d.decision+'
'; }).join('') + +decs.map(function(d){ return '
'+_esc(d.decision)+'
'; }).join('') +'
' :'') +'
'; diff --git a/src/context_engine/dashboard/server.py b/src/context_engine/dashboard/server.py index 562dca1..fb39e30 100644 --- a/src/context_engine/dashboard/server.py +++ b/src/context_engine/dashboard/server.py @@ -420,7 +420,11 @@ async def delete_file(file_path: str) -> dict | JSONResponse: async def set_compression(req: CompressionRequest) -> dict: state = _read_state() state["output_level"] = req.level - (storage_base / "state.json").write_text(json.dumps(state), encoding="utf-8") + # state.json is shared with the MCP server, which reads/writes it + # atomically — a plain write_text truncates in place and races it. + from context_engine.utils import atomic_write_text + + atomic_write_text(storage_base / "state.json", json.dumps(state)) return {"level": req.level} @app.get("/api/export") diff --git a/src/context_engine/editors.py b/src/context_engine/editors.py index ce1a425..d827cdd 100644 --- a/src/context_engine/editors.py +++ b/src/context_engine/editors.py @@ -297,26 +297,40 @@ def configure_mcp(project_dir: Path, editor_key: str) -> bool | None: try: data = json.loads(config_path.read_text(encoding="utf-8")) except (json.JSONDecodeError, OSError): - data = {} + # Never rewrite a file we couldn't parse: resetting to {} would + # replace the user's other MCP servers with only our entry. + # Skip and let the caller surface a warning (same contract as + # the TOML path). + return None + if not isinstance(data, dict): + return None else: data = {} - servers = data.setdefault(servers_key, {}) - if "context-engine" in servers: - existing = servers["context-engine"] - if existing.get("command") == command and existing.get("args") == entry["args"]: - return False - servers["context-engine"] = entry - atomic_write_text(config_path, json.dumps(data, indent=2) + "\n") - return True + # A non-dict servers value ("mcpServers": null, or a list) can't hold + # our entry — replace it with a fresh dict instead of raising TypeError. + servers = data.get(servers_key) + if not isinstance(servers, dict): + servers = {} + data[servers_key] = servers + + existing = servers.get("context-engine") + if ( + isinstance(existing, dict) + and existing.get("command") == command + and existing.get("args") == entry["args"] + ): + return False servers["context-engine"] = entry atomic_write_text(config_path, json.dumps(data, indent=2) + "\n") return True -def _configure_opencode(config_path: Path, command: str, project_dir: str) -> bool: - """Add CCE to OpenCode's opencode.json. Returns True if changed. +def _configure_opencode(config_path: Path, command: str, project_dir: str) -> bool | None: + """Add CCE to OpenCode's opencode.json. Returns True if changed, False if + already configured, or None if the existing config could not be parsed + (skipped rather than overwritten). OpenCode uses a different MCP entry format: type "local" with command as an array (not a string + args). @@ -338,18 +352,28 @@ def _configure_opencode(config_path: Path, command: str, project_dir: str) -> bo # Strip JSONC comments for parsing data = json.loads(_strip_jsonc_comments(content)) except (json.JSONDecodeError, OSError): - data = {} + # Never rewrite a file we couldn't parse: resetting to {} would + # replace the user's other MCP servers with only our entry. + return None + if not isinstance(data, dict): + return None else: data = {} - servers = data.setdefault("mcp", {}) - if "context-engine" in servers: - existing = servers["context-engine"] - if existing.get("command") == entry["command"] and existing.get("type") == "local": - return False - servers["context-engine"] = entry - atomic_write_text(config_path, json.dumps(data, indent=2) + "\n") - return True + # A non-dict "mcp" value (null, list) can't hold our entry — replace it + # with a fresh dict instead of raising TypeError. + servers = data.get("mcp") + if not isinstance(servers, dict): + servers = {} + data["mcp"] = servers + + existing = servers.get("context-engine") + if ( + isinstance(existing, dict) + and existing.get("command") == entry["command"] + and existing.get("type") == "local" + ): + return False servers["context-engine"] = entry atomic_write_text(config_path, json.dumps(data, indent=2) + "\n") @@ -357,9 +381,50 @@ def _configure_opencode(config_path: Path, command: str, project_dir: str) -> bo def _strip_jsonc_comments(text: str) -> str: - """Strip single-line // comments from JSONC content for JSON parsing.""" - import re - return re.sub(r'//.*?$', '', text, flags=re.MULTILINE) + """Strip `//` line comments and `/* */` block comments from JSONC content + for JSON parsing — but only outside strings. + + A naive regex (`//.*?$`) truncates any string value containing `//` + (e.g. `"url": "https://..."`), which makes json.loads fail and previously + caused the caller to reset the config to {} and destroy the user's other + MCP servers. This walker tracks in-string state (including escapes) so + string contents are preserved verbatim. + """ + out: list[str] = [] + i = 0 + n = len(text) + in_string = False + while i < n: + ch = text[i] + if in_string: + out.append(ch) + if ch == "\\" and i + 1 < n: + # Escaped char (e.g. \" or \\) — copy it, stay in string. + out.append(text[i + 1]) + i += 2 + continue + if ch == '"': + in_string = False + i += 1 + continue + if ch == '"': + in_string = True + out.append(ch) + i += 1 + continue + if ch == "/" and i + 1 < n and text[i + 1] == "/": + # Line comment: skip to end of line (keep the newline itself). + nl = text.find("\n", i) + i = n if nl == -1 else nl + continue + if ch == "/" and i + 1 < n and text[i + 1] == "*": + # Block comment: skip to the closing */ (or EOF if unterminated). + end = text.find("*/", i + 2) + i = n if end == -1 else end + 2 + continue + out.append(ch) + i += 1 + return "".join(out) _LEGACY_CODEX_SECTION = "mcp_servers.context-engine" diff --git a/src/context_engine/indexer/chunker.py b/src/context_engine/indexer/chunker.py index 5b95de1..58499c7 100644 --- a/src/context_engine/indexer/chunker.py +++ b/src/context_engine/indexer/chunker.py @@ -30,6 +30,16 @@ "use_declaration", # PHP, Rust } +def _node_text(src_bytes: bytes, node) -> str: + """Slice a node's source from the utf-8 BYTES tree-sitter parsed. + + node.start_byte / node.end_byte are byte offsets into the encoded + source, not str indices — slicing the original str garbles content + whenever a multi-byte character precedes the node. + """ + return src_bytes[node.start_byte:node.end_byte].decode("utf-8", errors="replace") + + _LANGUAGES = { "python": Language(tspython.language()), "javascript": Language(tsjavascript.language()), @@ -71,23 +81,28 @@ def chunk(self, source: str, file_path: str, language: str) -> list[Chunk]: parser = self._get_parser(language) if parser is None: return [self._fallback_chunk(source, file_path, language)] - tree = parser.parse(source.encode("utf-8")) + # tree-sitter parses the utf-8 BYTES and reports byte offsets. + # Encode once and slice the bytes — slicing the original str with + # byte offsets silently garbles every chunk after the first + # multi-byte character (emoji, CJK, accents). + src_bytes = source.encode("utf-8") + tree = parser.parse(src_bytes) chunks = [] - self._walk(tree.root_node, source, file_path, language, chunks) + self._walk(tree.root_node, src_bytes, file_path, language, chunks) if not chunks: return [self._fallback_chunk(source, file_path, language)] return chunks - def _walk(self, node, source, file_path, language, chunks): + def _walk(self, node, src_bytes, file_path, language, chunks): if node.type in _FUNCTION_TYPES: - chunks.append(self._node_to_chunk(node, source, file_path, language, ChunkType.FUNCTION)) + chunks.append(self._node_to_chunk(node, src_bytes, file_path, language, ChunkType.FUNCTION)) elif node.type in _CLASS_TYPES: - chunks.append(self._node_to_chunk(node, source, file_path, language, ChunkType.CLASS)) + chunks.append(self._node_to_chunk(node, src_bytes, file_path, language, ChunkType.CLASS)) for child in node.children: - self._walk(child, source, file_path, language, chunks) + self._walk(child, src_bytes, file_path, language, chunks) - def _node_to_chunk(self, node, source, file_path, language, chunk_type): - content = source[node.start_byte:node.end_byte] + def _node_to_chunk(self, node, src_bytes, file_path, language, chunk_type): + content = _node_text(src_bytes, node) start_line = node.start_point.row + 1 end_line = node.end_point.row + 1 chunk_id = hashlib.sha256( @@ -109,38 +124,41 @@ def _extract_imports(self, source: str, language: str) -> list[str]: parser = self._get_parser(language) if parser is None: return [] - tree = parser.parse(source.encode("utf-8")) + # Same byte-offset contract as chunk(): slice the encoded bytes, + # never the str (multi-byte chars shift str indices). + src_bytes = source.encode("utf-8") + tree = parser.parse(src_bytes) imports: list[str] = [] - self._walk_imports(tree.root_node, source, language, imports) + self._walk_imports(tree.root_node, src_bytes, language, imports) return list(dict.fromkeys(imports)) # deduplicate while preserving order - def _walk_imports(self, node, source, language, imports): + def _walk_imports(self, node, src_bytes, language, imports): if node.type in _IMPORT_TYPES: - module = self._parse_import_module(node, source, language) + module = self._parse_import_module(node, src_bytes, language) if module: imports.append(module) for child in node.children: - self._walk_imports(child, source, language, imports) + self._walk_imports(child, src_bytes, language, imports) - def _parse_import_module(self, node, source, language) -> str | None: + def _parse_import_module(self, node, src_bytes, language) -> str | None: if node.type == "import_statement": # Python: "import os" or "import os.path" # Also handles JS/TS: "import React from 'react'" (string child present) for child in node.children: if child.type == "string": # JavaScript/TypeScript import with string module specifier - raw = source[child.start_byte:child.end_byte].strip("'\"") + raw = _node_text(src_bytes, child).strip("'\"") return raw.split("/")[0] if not raw.startswith("@") else "/".join(raw.split("/")[:2]) if child.type in ("dotted_name", "aliased_import"): # Python bare import - name = source[child.start_byte:child.end_byte] + name = _node_text(src_bytes, child) name = name.split(" as ")[0].strip() return name.split(".")[0] elif node.type == "import_from_statement": # Python: "from pathlib import Path" for child in node.children: if child.type in ("dotted_name", "relative_import"): - name = source[child.start_byte:child.end_byte].strip() + name = _node_text(src_bytes, child).strip() name = name.lstrip(".") if name: return name.split(".")[0] @@ -148,7 +166,7 @@ def _parse_import_module(self, node, source, language) -> str | None: # TypeScript (tree-sitter-typescript): "import React from 'react'" for child in node.children: if child.type == "string": - raw = source[child.start_byte:child.end_byte].strip("'\"") + raw = _node_text(src_bytes, child).strip("'\"") return raw.split("/")[0] if not raw.startswith("@") else "/".join(raw.split("/")[:2]) return None diff --git a/src/context_engine/indexer/pipeline.py b/src/context_engine/indexer/pipeline.py index 04d446d..410aa1a 100644 --- a/src/context_engine/indexer/pipeline.py +++ b/src/context_engine/indexer/pipeline.py @@ -253,6 +253,25 @@ def walk(directory: Path) -> Iterable[Path]: _MAX_FILE_BYTES = 2 * 1024 * 1024 +def _cceignore_matches_file(rel_posix: str, patterns: list[str]) -> bool: + """True if a file (given as a /-separated path relative to the project + root) is excluded by `.cceignore`. + + The directory walk gets ancestor exclusion for free — it prunes a + matching directory before ever descending into it. The single-file + target path skips the walk, so directory patterns (`vendor/`) must be + re-checked against every ancestor here. + """ + from context_engine.indexer.ignorefile import matches_any + if matches_any(rel_posix, False, patterns): + return True + parts = rel_posix.split("/") + for i in range(1, len(parts)): + if matches_any("/".join(parts[:i]), True, patterns): + return True + return False + + def _safe_read(file_path: Path) -> str | None: """Read file as UTF-8 text; return None for binary, oversized, or unreadable files.""" try: @@ -289,7 +308,12 @@ async def run_indexing( on large repos) and are complementary: phase_fn is per-phase, embed_ progress_fn is per-batch. """ - project_dir = Path(project_dir) + # Resolve once so every downstream relative_to() call agrees with the + # resolve()d paths produced by _resolve_within / _iter_project_files. + # Without this, a symlinked project root (macOS /tmp → /private/tmp) + # 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.mkdir(parents=True, exist_ok=True) @@ -353,7 +377,33 @@ async def _run_indexing_locked( if target_path: target = _resolve_within(project_dir, target_path) if target.is_file(): - file_iter = [target] if target.suffix not in _SKIP_EXTENSIONS else [] + # The single-file branch is the watcher's only path into the + # pipeline — it must enforce the same filters as the directory + # walk (_iter_project_files), or saving `.env.production` in a + # watched project gets it read and indexed. + from context_engine.indexer.secrets import is_secret_file + rel = str(target.relative_to(project_dir)) + rel_posix = rel.replace("\\", "/") + if target.suffix in _SKIP_EXTENSIONS: + file_iter = [] + elif ( + getattr(config, "indexer_redact_secrets", True) + and is_secret_file(target) + ): + log.info("indexer: skipping secret file %s", target) + if log_fn: + log_fn(f" [skip] {rel} (secret file)") + result.skipped_files.append(rel) + file_iter = [] + elif cceignore_patterns and _cceignore_matches_file( + rel_posix, cceignore_patterns + ): + if log_fn: + log_fn(f" [skip] {rel} (.cceignore)") + result.skipped_files.append(rel) + file_iter = [] + else: + file_iter = [target] elif target.is_dir(): file_iter = list(_iter_project_files( target, ignore_set, _SKIP_EXTENSIONS, @@ -361,6 +411,22 @@ async def _run_indexing_locked( cceignore_patterns=cceignore_patterns, )) else: + # Not on disk. The watcher enqueues deletions through the same + # code path — if the manifest still tracks this file, prune its + # chunks + manifest entry instead of erroring out. + rel = str(target.relative_to(project_dir)) + if manifest.get_hash(rel) is not None: + try: + await backend.delete_by_files([rel]) + except Exception as exc: + result.errors.append(f"Failed to prune deleted file {rel}: {exc}") + return result + manifest.remove(rel) + manifest.save() + result.deleted_files.append(rel) + if log_fn: + log_fn(f" [delete] {rel} (no longer on disk)") + return result result.errors.append(f"Target path not found: {target_path}") return result else: diff --git a/src/context_engine/indexer/secrets.py b/src/context_engine/indexer/secrets.py index e007a68..809bf89 100644 --- a/src/context_engine/indexer/secrets.py +++ b/src/context_engine/indexer/secrets.py @@ -50,6 +50,10 @@ # Filename starts with any of these → skip (handles .env, .env.local, etc.). _SECRET_PREFIXES = (".env",) +# Filename ends with any of these → skip. Catches the `.env` convention +# (prod.env, secrets.env, config.env) that the `.env` prefix rule misses. +_SECRET_SUFFIXES = (".env",) + # File extensions whose presence is a strong signal of a key/cert. Skip outright. _SECRET_EXTENSIONS = frozenset({ ".pem", ".key", ".crt", ".cer", ".der", @@ -75,6 +79,9 @@ def is_secret_file(path: Path) -> bool: for prefix in _SECRET_PREFIXES: if name.startswith(prefix): return True + for suffix in _SECRET_SUFFIXES: + if name.endswith(suffix): + return True return False @@ -98,13 +105,16 @@ def is_secret_file(path: Path) -> bool: # GitHub tokens (classic + fine-grained + app + OAuth). (re.compile(r"\bghp_[A-Za-z0-9]{36}\b"), "GITHUB_PAT"), (re.compile(r"\bgithub_pat_[A-Za-z0-9_]{82}\b"), "GITHUB_FINE_GRAINED_PAT"), - (re.compile(r"\b(ghs|gho|ghu|ghr)_[A-Za-z0-9]{36}\b"), "GITHUB_OAUTH"), + # NOTE: prefix groups must be NON-capturing — `_sub` treats a capture + # group as "the credential value", so `(ghs|gho|...)` would redact + # (or placeholder-skip) just the 3-char prefix and leak the token. + (re.compile(r"\b(?:ghs|gho|ghu|ghr)_[A-Za-z0-9]{36}\b"), "GITHUB_OAUTH"), # Slack tokens. (re.compile(r"\bxox[abprs]-[A-Za-z0-9-]{10,}\b"), "SLACK_TOKEN"), # Stripe live keys (test keys are deliberately not matched — they're # safe to commit and matching them would over-redact every Stripe # quickstart in the wild). - (re.compile(r"\b(sk|rk)_live_[A-Za-z0-9]{24,}\b"), "STRIPE_LIVE_KEY"), + (re.compile(r"\b(?:sk|rk)_live_[A-Za-z0-9]{24,}\b"), "STRIPE_LIVE_KEY"), # OpenAI / Anthropic API keys. (re.compile(r"\bsk-[A-Za-z0-9]{20}T3BlbkFJ[A-Za-z0-9]{20}\b"), "OPENAI_KEY"), (re.compile(r"\bsk-ant-(api03|admin01)-[A-Za-z0-9_\-]{93,}\b"), "ANTHROPIC_KEY"), @@ -130,6 +140,25 @@ def is_secret_file(path: Path) -> bool: r"private[_-]?key|auth[_-]?token|client[_-]?secret)\b" r"['\"\]\s]*[:=]\s*['\"]([^'\"\s]{16,})['\"]" ), "GENERIC_CREDENTIAL"), + # Unquoted variant — dotenv / YAML / TOML / shell assignments don't + # quote their values (`STRIPE_KEY=sk_live_…`, `password: hunter2…`, + # `export AUTH_TOKEN=…`), so the quoted pattern above misses them. + # + # Deliberately stricter than the quoted pattern to limit false + # positives on ordinary code: + # · line-anchored (config formats assign at line start, optionally + # indented / behind `export`) — prose mentions don't fire; + # · the credential keyword must be the TRAILING segment of the + # variable name (`STRIPE_KEY`, `db_password` — but not + # `tokenizer=` or `keyboard_layout=`); + # · the value must be an unbroken run of ≥16 non-quote characters + # (quoted values are the previous pattern's job). + (re.compile( + r"(?im)^[ \t]*(?:export[ \t]+)?" + r"[A-Za-z0-9_.\-]*(?:password|passwd|secret|token|credential|api[_-]?key|key)s?" + r"[ \t]*[:=][ \t]*" + r"([^\s'\"#`]{16,})" + ), "GENERIC_CREDENTIAL"), ] @@ -170,6 +199,10 @@ def _is_placeholder(value: str) -> bool: v = value.strip("'\"<>").lower() if v in _PLACEHOLDER_VALUES: return True + # Already redacted by an earlier (more specific) pattern this run — + # don't re-redact the placeholder into a second label. + if v.startswith("[redacted:"): + return True # Repeated single character ("xxxxxxxxxx", "0000000000") is almost # always a placeholder, never a real key. if len(set(v)) <= 2: diff --git a/src/context_engine/memory/compressor.py b/src/context_engine/memory/compressor.py index f0fac74..05c8a48 100644 --- a/src/context_engine/memory/compressor.py +++ b/src/context_engine/memory/compressor.py @@ -63,6 +63,7 @@ def compress_turn( raw_tokens = _approx_tokens(text) summary, tier = _summarise(text, embedder=embedder, top_k=_DEFAULT_TURN_TOP_K) extractive_tokens = _approx_tokens(summary) + gram_raw = gram_comp = 0 if summary: # Scrub PII before grammar compression — emails / IPs / SSNs that # leaked into a turn (the user pasted a real value into a prompt @@ -71,13 +72,37 @@ def compress_turn( summary, gram_raw, gram_comp = _grammar_compress_counted( summary, level=_GRAMMAR_LEVEL, ) - memory_db.record_savings( - conn, bucket="grammar", baseline=gram_raw, served=gram_comp, - ) + epoch = int(time.time()) + # Explicit DELETE-then-INSERT rather than INSERT OR REPLACE: SQLite only + # fires delete triggers on a REPLACE conflict when recursive_triggers is + # ON, so OR REPLACE silently left dangling turn_summaries_fts entries and + # stale turn_summaries_vec rows every time a turn was re-compressed. The + # explicit DELETE always fires the _ad triggers that clean both up. + conn.execute( + "DELETE FROM turn_summaries WHERE session_id = ? AND prompt_number = ?", + (session_id, prompt_number), + ) + cur = conn.execute( + "INSERT INTO turn_summaries " + "(session_id, prompt_number, summary, tier, created_at_epoch) " + "VALUES (?, ?, ?, ?, ?)", + (session_id, prompt_number, summary, tier, epoch), + ) + if summary: # Scan the clean, PII-free, grammar-compressed summary for decisions. # Running here reuses the pipeline's existing scrub + compress rather # than duplicating that work on the raw turn text. _auto_capture_decisions(conn, summary, session_id=session_id, prompt_number=prompt_number, embedder=embedder) + memory_db.record_turn_summary_vec( + conn, embedder, turn_id=cur.lastrowid, summary=summary, + ) + # Savings are recorded last, after every fallible step above, so a turn + # that fails and gets retried by the worker doesn't inflate the ledger + # once per attempt. + if summary and gram_raw > 0: + memory_db.record_savings( + conn, bucket="grammar", baseline=gram_raw, served=gram_comp, + ) # Turn-summarization savings: raw turn text (prompt + tool inputs/outputs) # vs the extractive summary that ends up in turn_summaries. if raw_tokens > 0 and extractive_tokens > 0: @@ -86,17 +111,6 @@ def compress_turn( baseline=raw_tokens, served=extractive_tokens, meta={"kind": "turn", "tier": tier}, ) - epoch = int(time.time()) - cur = conn.execute( - "INSERT OR REPLACE INTO turn_summaries " - "(session_id, prompt_number, summary, tier, created_at_epoch) " - "VALUES (?, ?, ?, ?, ?)", - (session_id, prompt_number, summary, tier, epoch), - ) - if summary: - memory_db.record_turn_summary_vec( - conn, embedder, turn_id=cur.lastrowid, summary=summary, - ) return summary @@ -119,6 +133,8 @@ def compress_session_rollup( )) text = "\n".join(r["summary"] for r in rows if r["summary"]) raw_tokens = _approx_tokens(text) + extractive_tokens = 0 + gram_raw = gram_comp = 0 if not text: rollup = "" tier = "empty" @@ -137,21 +153,24 @@ def compress_session_rollup( rollup, gram_raw, gram_comp = _grammar_compress_counted( rollup, level=_GRAMMAR_LEVEL, ) - memory_db.record_savings( - conn, bucket="grammar", baseline=gram_raw, served=gram_comp, - ) - if raw_tokens > 0 and extractive_tokens > 0: - memory_db.record_savings( - conn, bucket="turn_summarization", - baseline=raw_tokens, served=extractive_tokens, - meta={"kind": "session_rollup", "tier": tier}, - ) epoch = int(time.time()) conn.execute( "UPDATE sessions SET rollup_summary = ?, rollup_summary_at_epoch = ? " "WHERE id = ?", (rollup, epoch, session_id), ) + # Savings only after the rollup write succeeded — retried failures must + # not re-record into the ledger. + if gram_raw > 0: + memory_db.record_savings( + conn, bucket="grammar", baseline=gram_raw, served=gram_comp, + ) + if raw_tokens > 0 and extractive_tokens > 0: + memory_db.record_savings( + conn, bucket="turn_summarization", + baseline=raw_tokens, served=extractive_tokens, + meta={"kind": "session_rollup", "tier": tier}, + ) log.debug("session rollup tier=%s len=%d", tier, len(rollup)) return rollup @@ -293,14 +312,23 @@ def _auto_capture_decisions( return count +# Rows that have failed this many times are dead-lettered: left in the table +# (visible via `cce sessions status` / last_error) but never picked again, so +# a deterministically-failing row can't starve the queue with retries every +# interval. +_MAX_ATTEMPTS = 5 + + def _drain_one_sync(conn: sqlite3.Connection, embedder) -> bool: - """Pop and process the oldest pending row. Pure-sync; safe for either the - main thread (tests) or a worker thread (production via to_thread). - Returns True iff work was done. + """Pop and process the oldest pending row below the attempt cap. + Pure-sync; safe for either the main thread (tests) or a worker thread + (production via to_thread). Returns True iff work was done. """ row = conn.execute( "SELECT id, kind, session_id, prompt_number, attempts FROM pending_compressions " - "ORDER BY enqueued_at_epoch ASC LIMIT 1" + "WHERE attempts < ? " + "ORDER BY enqueued_at_epoch ASC LIMIT 1", + (_MAX_ATTEMPTS,), ).fetchone() if row is None: return False @@ -323,6 +351,9 @@ def _drain_one_sync(conn: sqlite3.Connection, embedder) -> bool: except Exception as exc: log.exception("Compression failed for %s/%s/%s", row["kind"], row["session_id"], row["prompt_number"]) + # Discard any partial writes from the failed compression before + # persisting the attempt bump. + conn.rollback() conn.execute( "UPDATE pending_compressions SET attempts = attempts + 1, " "last_error = ? WHERE id = ?", diff --git a/src/context_engine/memory/db.py b/src/context_engine/memory/db.py index 09e5858..2c5d514 100644 --- a/src/context_engine/memory/db.py +++ b/src/context_engine/memory/db.py @@ -337,10 +337,15 @@ def _ensure_schema(conn: sqlite3.Connection, *, has_vec: bool) -> None: cur.execute(stmt) for stmt in _SCHEMA_V3: cur.execute(stmt) + # If sqlite-vec was unavailable the v2 vec tables were skipped. + # Stamp only v1 in that case (mirroring the deferred stamp in + # the upgrade branch below) so a future connection with vec + # loaded completes the v1 → v2 step instead of hitting the + # `current >= CURRENT_VERSION` early return forever. cur.execute( "INSERT INTO schema_versions (version, applied_at_epoch) " "VALUES (?, strftime('%s','now'))", - (CURRENT_VERSION,), + (CURRENT_VERSION if has_vec else 1,), ) conn.commit() except Exception: diff --git a/src/context_engine/memory/hooks.py b/src/context_engine/memory/hooks.py index 162953f..a76bf19 100644 --- a/src/context_engine/memory/hooks.py +++ b/src/context_engine/memory/hooks.py @@ -20,8 +20,16 @@ from aiohttp import web +from context_engine.memory import db as memory_db + log = logging.getLogger(__name__) +# Sentinel prompt_number for session_rollup queue rows. SQLite treats NULLs +# as distinct in UNIQUE constraints, so a NULL here would defeat the +# UNIQUE(kind, session_id, prompt_number) dedup and every repeated +# SessionEnd would enqueue (and run) the rollup again. +_ROLLUP_PROMPT_SENTINEL = -1 + def _now_epoch() -> int: return int(time.time()) @@ -215,7 +223,12 @@ async def handle_session_start(request: web.Request) -> web.Response: if not session_id: return web.Response(text="", status=400) project = data.get("project") or request.app.get("project_name", "") - started_epoch = int(data.get("started_at") or _now_epoch()) + # Guarded parse — hooks must never 500, and a non-numeric started_at + # from a misbehaving hook script would raise ValueError here. + try: + started_epoch = int(data.get("started_at") or _now_epoch()) + except (TypeError, ValueError): + started_epoch = _now_epoch() conn = _conn(request) try: @@ -263,7 +276,14 @@ async def handle_user_prompt_submit(request: web.Request) -> web.Response: "INSERT OR IGNORE INTO prompts " "(session_id, prompt_number, prompt_text, created_at_epoch, created_at) " "VALUES (?, ?, ?, ?, ?)", - (session_id, int(prompt_number), str(prompt_text), epoch, _now_iso(epoch)), + ( + session_id, int(prompt_number), + # Scrub before write — prompts_fts indexes this verbatim, so + # PII must never reach the row (same policy as compressor / + # mcp_server writes; scrub_pii no-ops when redaction is off). + memory_db.scrub_pii(str(prompt_text)), + epoch, _now_iso(epoch), + ), ) conn.execute( "UPDATE sessions SET prompt_count = prompt_count + 1 WHERE id = ?", @@ -297,6 +317,11 @@ async def handle_post_tool_use(request: web.Request) -> web.Response: raw_input = tool_input if isinstance(tool_input, str) else json.dumps(tool_input) raw_output = tool_output if isinstance(tool_output, str) else json.dumps(tool_output) + # Scrub raw payloads before they land in tool_event_payloads — they are + # served back verbatim by the session_event MCP tool and feed the turn + # compressor, so PII must be redacted at the write boundary too. + raw_input = memory_db.scrub_pii(raw_input) + raw_output = memory_db.scrub_pii(raw_output) size = len(raw_input) + len(raw_output) conn = _conn(request) @@ -401,7 +426,13 @@ def _enqueue_compression( UNIQUE(kind, session_id, prompt_number) guards against double-enqueue when a prompt fires both Stop *and* the next UserPromptSubmit's "compress prev" trigger in quick succession. + + Rollups have no prompt_number; a NULL would be treated as distinct by the + UNIQUE constraint (defeating INSERT OR IGNORE), so they use the -1 + sentinel instead. The worker ignores prompt_number for rollups. """ + if prompt_number is None: + prompt_number = _ROLLUP_PROMPT_SENTINEL conn.execute( "INSERT OR IGNORE INTO pending_compressions " "(kind, session_id, prompt_number, enqueued_at_epoch) " diff --git a/src/context_engine/retrieval/retriever.py b/src/context_engine/retrieval/retriever.py index 430c915..a758311 100644 --- a/src/context_engine/retrieval/retriever.py +++ b/src/context_engine/retrieval/retriever.py @@ -21,6 +21,9 @@ # When the parsed query looks like a code lookup, give FTS more pull because # exact-identifier hits are usually what the user wants. _FTS_BOOST_CODE_LOOKUP = 1.5 +# Worst possible cosine distance (opposite vectors). Used as the fallback for +# chunks with no vector evidence when no vector results were returned at all. +_WORST_COSINE_DISTANCE = 2.0 class HybridRetriever: @@ -104,6 +107,19 @@ async def retrieve( # almost no signal past the top few. Rank-normalising restores gradient. max_rrf = max(rrf_scores.values()) if rrf_scores else 0.0 + # Chunks hydrated from FTS-only hits carry no _distance. They were + # absent from the vector top-k, so their true distance is at least as + # bad as the worst returned vector hit — never default to 0.0, which + # would grant keyword-incidental matches perfect vector similarity. + observed_distances = [ + c.metadata["_distance"] + for c in vector_results + if "_distance" in c.metadata + ] + no_vector_distance = ( + max(observed_distances) if observed_distances else _WORST_COSINE_DISTANCE + ) + # Score with confidence scorer scored: list[tuple[Chunk, float]] = [] for id_, rrf_score in rrf_scores.items(): @@ -111,7 +127,7 @@ async def retrieve( if chunk is None: continue - distance = chunk.metadata.get("_distance", 0.0) + distance = chunk.metadata.get("_distance", no_vector_distance) normalised_distance = min(max(distance / 2.0, 0.0), 1.0) keyword_distance = self._estimate_keyword_distance(chunk, parsed) conf_score = self._scorer.score( diff --git a/src/context_engine/storage/fts_store.py b/src/context_engine/storage/fts_store.py index 2160353..1e31229 100644 --- a/src/context_engine/storage/fts_store.py +++ b/src/context_engine/storage/fts_store.py @@ -12,9 +12,35 @@ _MAX_CONTENT_CHARS = 5_000 +# Common English words dropped from queries so natural-language phrasing +# ("where is the retry logic") doesn't drown the signal terms. If every token +# is a stopword we fall back to using them all rather than matching nothing. +_STOPWORDS = frozenset({ + "a", "an", "and", "are", "as", "at", "be", "by", "can", "do", "does", + "for", "from", "how", "i", "in", "is", "it", "of", "on", "or", "that", + "the", "this", "to", "was", "we", "what", "when", "where", "which", + "who", "why", "will", "with", "you", "not", +}) + + def _escape_fts5(query: str) -> str: - """Wrap user input as an FTS5 phrase to avoid operator injection.""" - return '"' + query.replace('"', '""') + '"' + """Build a safe FTS5 MATCH expression from free-form user input. + + Each whitespace token is individually quoted (internal quotes doubled), + which neutralises FTS5 operators (NEAR, *, :, AND/OR/NOT, ^, parens). + Tokens are OR-joined: requiring all tokens (AND / a single phrase) makes + multi-word natural-language queries almost never match, while BM25 still + ranks documents matching more tokens higher. Multi-part identifiers like + `delete_by_files` stay as one quoted token, which FTS5 tokenises into the + same consecutive-token phrase unicode61 produced at index time. + + Returns "" when no usable tokens remain; callers must skip the query. + """ + tokens = [t for t in query.split() if any(c.isalnum() for c in t)] + kept = [t for t in tokens if t.lower() not in _STOPWORDS] + if not kept: + kept = tokens + return " OR ".join('"' + t.replace('"', '""') + '"' for t in kept) class FTSStore: @@ -53,12 +79,25 @@ def _ingest_sync(self, chunks: list[Chunk]) -> None: for chunk in chunks ] with self._lock: - self._conn.executemany( - "INSERT OR REPLACE INTO chunks_fts(id, content, file_path, language, chunk_type) " - "VALUES (?, ?, ?, ?, ?)", - rows, - ) - self._conn.commit() + try: + # FTS5 virtual tables have no unique constraints, so + # INSERT OR REPLACE never replaces — delete stale rows for + # these ids first, inside the same transaction. + self._conn.executemany( + "DELETE FROM chunks_fts WHERE id = ?", + [(chunk.id,) for chunk in chunks], + ) + self._conn.executemany( + "INSERT INTO chunks_fts(id, content, file_path, language, chunk_type) " + "VALUES (?, ?, ?, ?, ?)", + rows, + ) + self._conn.commit() + except Exception: + # Roll back so a mid-batch failure doesn't leave pending rows + # that the next unrelated commit() silently flushes. + self._conn.rollback() + raise def _search_sync(self, escaped_query: str, top_k: int) -> list[tuple[str, float]]: with self._lock: @@ -82,14 +121,18 @@ def _delete_files_sync(self, file_paths: list[str]) -> None: from context_engine.utils import batched_params with self._lock: - for batch in batched_params(file_paths): - placeholders = ",".join("?" * len(batch)) - # Safe: placeholders is only "?" chars; values are parameterized. - self._conn.execute( # nosemgrep: sqlalchemy-execute-raw-query - f"DELETE FROM chunks_fts WHERE file_path IN ({placeholders})", - batch, - ) - self._conn.commit() + try: + for batch in batched_params(file_paths): + placeholders = ",".join("?" * len(batch)) + # Safe: placeholders is only "?" chars; values are parameterized. + self._conn.execute( # nosemgrep: sqlalchemy-execute-raw-query + f"DELETE FROM chunks_fts WHERE file_path IN ({placeholders})", + batch, + ) + self._conn.commit() + except Exception: + self._conn.rollback() + raise async def ingest(self, chunks: list[Chunk]) -> None: if not chunks: @@ -99,7 +142,10 @@ async def ingest(self, chunks: list[Chunk]) -> None: async def search(self, query: str, top_k: int = 30) -> list[tuple[str, float]]: if not query.strip(): return [] - return await asyncio.to_thread(self._search_sync, _escape_fts5(query), top_k) + match_expr = _escape_fts5(query) + if not match_expr: + return [] + return await asyncio.to_thread(self._search_sync, match_expr, top_k) def clear(self) -> None: with self._lock: diff --git a/src/context_engine/storage/graph_store.py b/src/context_engine/storage/graph_store.py index 7ba6869..21270f1 100644 --- a/src/context_engine/storage/graph_store.py +++ b/src/context_engine/storage/graph_store.py @@ -63,22 +63,28 @@ def __init__(self, db_path: str) -> None: def _sync_ingest(self, nodes: list[GraphNode], edges: list[GraphEdge]) -> None: with self._lock: - cur = self._conn.cursor() - for node in nodes: - cur.execute( - "INSERT OR REPLACE INTO nodes (id, node_type, name, file_path, properties) " - "VALUES (?, ?, ?, ?, ?)", - (node.id, node.node_type.value, node.name, node.file_path, - json.dumps(node.properties)), - ) - for edge in edges: - cur.execute( - "INSERT OR REPLACE INTO edges (source_id, target_id, edge_type, properties) " - "VALUES (?, ?, ?, ?)", - (edge.source_id, edge.target_id, edge.edge_type.value, - json.dumps(edge.properties)), - ) - self._conn.commit() + try: + cur = self._conn.cursor() + for node in nodes: + cur.execute( + "INSERT OR REPLACE INTO nodes (id, node_type, name, file_path, properties) " + "VALUES (?, ?, ?, ?, ?)", + (node.id, node.node_type.value, node.name, node.file_path, + json.dumps(node.properties)), + ) + for edge in edges: + cur.execute( + "INSERT OR REPLACE INTO edges (source_id, target_id, edge_type, properties) " + "VALUES (?, ?, ?, ?)", + (edge.source_id, edge.target_id, edge.edge_type.value, + json.dumps(edge.properties)), + ) + self._conn.commit() + except Exception: + # Roll back so a mid-batch failure doesn't leave pending rows + # that the next unrelated commit() silently flushes. + self._conn.rollback() + raise def _sync_get_neighbors(self, node_id: str, edge_type: EdgeType | None) -> list[GraphNode]: with self._lock: @@ -160,26 +166,32 @@ def _sync_delete_by_files(self, file_paths: list[str]) -> None: from context_engine.utils import batched_params with self._lock: - cur = self._conn.cursor() - # Collect node IDs in batches to respect SQLite param limits. - # Safe: ph is only "?" chars; values are parameterized. - node_ids: list[str] = [] - for batch in batched_params(file_paths): - ph = ",".join("?" * len(batch)) - cur.execute( # nosemgrep: sqlalchemy-execute-raw-query - f"SELECT id FROM nodes WHERE file_path IN ({ph})", batch - ) - node_ids.extend(row[0] for row in cur.fetchall()) - # Delete edges and nodes in batches. - for batch in batched_params(node_ids): - ph = ",".join("?" * len(batch)) - cur.execute( # nosemgrep: sqlalchemy-execute-raw-query - f"DELETE FROM edges WHERE source_id IN ({ph}) " - f"OR target_id IN ({ph})", - batch + batch, - ) - cur.execute(f"DELETE FROM nodes WHERE id IN ({ph})", batch) # nosemgrep: sqlalchemy-execute-raw-query - self._conn.commit() + try: + cur = self._conn.cursor() + # Collect node IDs in batches to respect SQLite param limits. + # Safe: ph is only "?" chars; values are parameterized. + node_ids: list[str] = [] + for batch in batched_params(file_paths): + ph = ",".join("?" * len(batch)) + cur.execute( # nosemgrep: sqlalchemy-execute-raw-query + f"SELECT id FROM nodes WHERE file_path IN ({ph})", batch + ) + node_ids.extend(row[0] for row in cur.fetchall()) + # Delete edges and nodes in batches. + for batch in batched_params(node_ids): + ph = ",".join("?" * len(batch)) + cur.execute( # nosemgrep: sqlalchemy-execute-raw-query + f"DELETE FROM edges WHERE source_id IN ({ph}) " + f"OR target_id IN ({ph})", + batch + batch, + ) + cur.execute(f"DELETE FROM nodes WHERE id IN ({ph})", batch) # nosemgrep: sqlalchemy-execute-raw-query + self._conn.commit() + except Exception: + # Roll back so a partial multi-batch delete doesn't leave + # pending rows that the next unrelated commit() flushes. + self._conn.rollback() + raise # ------------------------------------------------------------------ # Public async API diff --git a/src/context_engine/storage/vector_store.py b/src/context_engine/storage/vector_store.py index 3385dee..aee590f 100644 --- a/src/context_engine/storage/vector_store.py +++ b/src/context_engine/storage/vector_store.py @@ -98,15 +98,31 @@ def _ensure_tables(self) -> None: """) # Detect vector dimension from existing data row = self._conn.execute( - "SELECT name FROM sqlite_master WHERE type='table' AND name='chunks_vec'" + "SELECT sql FROM sqlite_master WHERE type='table' AND name='chunks_vec'" ).fetchone() if row: - # Table exists — read dim from first row - r = self._conn.execute("SELECT rowid FROM chunks_vec LIMIT 1").fetchone() - if r: - self._dim = self._conn.execute( - "SELECT vec_length(embedding) FROM chunks_vec LIMIT 1" - ).fetchone()[0] + if "distance_metric=cosine" not in (row[0] or ""): + # Legacy table created before distance_metric=cosine was + # specified — sqlite-vec defaulted to L2, which the + # retriever's distance/2.0 normalisation misreads. Rebuild + # empty (embeddings are a cache; reindex repopulates) + # rather than silently mixing metrics. + log.warning( + "Existing vector table uses L2 distance; rebuilding " + "with cosine metric — run a reindex to repopulate." + ) + self._conn.execute("DROP TABLE IF EXISTS chunks_vec") + self._conn.execute("DELETE FROM chunks") + self._conn.execute("DELETE FROM chunk_compressions") + else: + # Table exists — read dim from first row + r = self._conn.execute( + "SELECT rowid FROM chunks_vec LIMIT 1" + ).fetchone() + if r: + self._dim = self._conn.execute( + "SELECT vec_length(embedding) FROM chunks_vec LIMIT 1" + ).fetchone()[0] self._conn.commit() def _ensure_vec_table(self, dim: int) -> None: @@ -127,9 +143,11 @@ def _ensure_vec_table(self, dim: int) -> None: self._conn.execute("DELETE FROM chunks") self._conn.execute("DELETE FROM chunk_compressions") # Safe: dim is a validated integer, never from user input. + # distance_metric=cosine must match the retriever's distance/2.0 + # normalisation (cosine distance lives in [0, 2]). self._conn.execute( # nosemgrep: sqlalchemy-execute-raw-query f"CREATE VIRTUAL TABLE IF NOT EXISTS chunks_vec " - f"USING vec0(embedding float[{dim}])" + f"USING vec0(embedding float[{dim}] distance_metric=cosine)" ) self._dim = dim self._conn.commit() @@ -168,29 +186,36 @@ async def ingest(self, chunks: list[Chunk]) -> None: dim = len(valid[0].embedding) self._ensure_vec_table(dim) with self._lock: - cursor = self._conn.cursor() - for chunk in valid: - row = self._chunk_to_row(chunk) - rowid = cursor.execute( - "INSERT INTO chunks " - "(id, content, chunk_type, file_path, start_line, end_line, language) " - "VALUES (?, ?, ?, ?, ?, ?, ?) " - "ON CONFLICT(id) DO UPDATE SET " - "content = excluded.content, " - "chunk_type = excluded.chunk_type, " - "file_path = excluded.file_path, " - "start_line = excluded.start_line, " - "end_line = excluded.end_line, " - "language = excluded.language " - "RETURNING rowid", - row, - ).fetchone()[0] - cursor.execute("DELETE FROM chunks_vec WHERE rowid = ?", (rowid,)) - cursor.execute( - "INSERT INTO chunks_vec(rowid, embedding) VALUES (?, ?)", - (rowid, _serialize_vec(chunk.embedding)), - ) - self._conn.commit() + try: + cursor = self._conn.cursor() + for chunk in valid: + row = self._chunk_to_row(chunk) + rowid = cursor.execute( + "INSERT INTO chunks " + "(id, content, chunk_type, file_path, start_line, end_line, language) " + "VALUES (?, ?, ?, ?, ?, ?, ?) " + "ON CONFLICT(id) DO UPDATE SET " + "content = excluded.content, " + "chunk_type = excluded.chunk_type, " + "file_path = excluded.file_path, " + "start_line = excluded.start_line, " + "end_line = excluded.end_line, " + "language = excluded.language " + "RETURNING rowid", + row, + ).fetchone()[0] + cursor.execute("DELETE FROM chunks_vec WHERE rowid = ?", (rowid,)) + cursor.execute( + "INSERT INTO chunks_vec(rowid, embedding) VALUES (?, ?)", + (rowid, _serialize_vec(chunk.embedding)), + ) + self._conn.commit() + except Exception: + # Roll back so a mid-batch failure doesn't leave pending rows + # (e.g. chunks with no chunks_vec row — silently unfindable) + # that the next unrelated commit() would flush. + self._conn.rollback() + raise async def search( self, @@ -257,25 +282,29 @@ async def delete_by_files(self, file_paths: list[str]) -> None: from context_engine.utils import batched_params with self._lock: - # Safe: placeholders is only "?" chars; values are parameterized. - for batch in batched_params(file_paths): - placeholders = ",".join("?" * len(batch)) - if self._dim is not None: + try: + # Safe: placeholders is only "?" chars; values are parameterized. + for batch in batched_params(file_paths): + placeholders = ",".join("?" * len(batch)) + if self._dim is not None: + self._conn.execute( # nosemgrep: sqlalchemy-execute-raw-query + f"DELETE FROM chunks_vec " + f"WHERE rowid IN (SELECT rowid FROM chunks WHERE file_path IN ({placeholders}))", + batch, + ) self._conn.execute( # nosemgrep: sqlalchemy-execute-raw-query - f"DELETE FROM chunks_vec " - f"WHERE rowid IN (SELECT rowid FROM chunks WHERE file_path IN ({placeholders}))", + f"DELETE FROM chunk_compressions " + f"WHERE chunk_id IN (SELECT id FROM chunks WHERE file_path IN ({placeholders}))", batch, ) - self._conn.execute( # nosemgrep: sqlalchemy-execute-raw-query - f"DELETE FROM chunk_compressions " - f"WHERE chunk_id IN (SELECT id FROM chunks WHERE file_path IN ({placeholders}))", - batch, - ) - self._conn.execute( # nosemgrep: sqlalchemy-execute-raw-query - f"DELETE FROM chunks WHERE file_path IN ({placeholders})", - batch, - ) - self._conn.commit() + self._conn.execute( # nosemgrep: sqlalchemy-execute-raw-query + f"DELETE FROM chunks WHERE file_path IN ({placeholders})", + batch, + ) + self._conn.commit() + except Exception: + self._conn.rollback() + raise def get_cached_compression(self, chunk_id: str, level: str) -> str | None: """Return the cached compressed text for (chunk_id, level), or None.""" diff --git a/tests/dashboard/test_page_escaping.py b/tests/dashboard/test_page_escaping.py new file mode 100644 index 0000000..5a50b67 --- /dev/null +++ b/tests/dashboard/test_page_escaping.py @@ -0,0 +1,83 @@ +"""Regression tests for XSS escaping in the embedded dashboard page. + +The dashboard JS builds HTML via string concatenation from API-sourced +values (file paths, session projects, decision text, session ids). Every +such interpolation must go through _esc() for HTML context, and event +handlers must not inline API strings into onclick attributes (JS-in-HTML +context) — they must use data-* attributes + addEventListener instead. + +There is no browser test harness, so these tests assert on the generated +page source: the known-bad patterns must be gone and the escaped/safe +forms present. Regression for 2026-07-03 security review. +""" + +from context_engine.dashboard._page import PAGE_HTML + + +# ── Sessions tab (loadSessions) ─────────────────────── + + +def test_sessions_tab_escapes_project_name(): + # Old: '+(s.project||s.id)+' — unescaped stored XSS vector. + assert "'+(s.project||s.id)+'" not in PAGE_HTML + assert "_esc(s.project||s.id)" in PAGE_HTML + + +def test_sessions_tab_escapes_decision_text(): + # Old: '
'+d.decision+'
' + assert "'+d.decision+'" not in PAGE_HTML + assert '
\'+_esc(d.decision)+\'
' in PAGE_HTML + + +# ── Files table (renderFiles) ───────────────────────── + + +def test_files_table_escapes_path_in_title_and_text(): + # Old: title="'+f.path+'">'+f.path+' + assert "'+f.path+'" not in PAGE_HTML + assert 'title="\'+_esc(f.path)+\'"' in PAGE_HTML + + +def test_files_table_escapes_status(): + assert "'+f.status+'" not in PAGE_HTML + assert "_esc(f.status)" in PAGE_HTML + + +def test_files_table_uses_data_attributes_not_inline_onclick(): + # Old: onclick="reindexFile('+JSON.stringify(f.path)+')" — JSON.stringify + # covers JS-string context but not the surrounding HTML attribute context. + assert "JSON.stringify(f.path)" not in PAGE_HTML + assert "onclick=\"reindexFile(" not in PAGE_HTML + assert "onclick=\"deleteFile(" not in PAGE_HTML + assert 'data-path="\'+_esc(f.path)+\'"' in PAGE_HTML + assert "addEventListener" in PAGE_HTML + + +# ── Memory sessions (loadMemorySessions) ────────────── + + +def test_memory_sessions_no_inline_onclick_with_html_escaper(): + # Old: onclick="loadMemoryTimeline(\''+_esc(s.id)+'\')" — HTML escaper + # used in a JS-string-in-attribute context; ' decodes back to a quote + # before the JS parser runs, so it does not neutralize breakouts. + assert "onclick=\"loadMemoryTimeline(" not in PAGE_HTML + assert 'data-sid="\'+_esc(s.id)+\'"' in PAGE_HTML + + +# ── Chart helpers (audit findings) ──────────────────── + + +def test_hbar_chart_escapes_labels(): + # Labels come from f.path via loadOverviewPanels. + assert "title=\"'+item.label+'\"" not in PAGE_HTML + assert "_esc(item.label)" in PAGE_HTML + + +def test_vbar_chart_escapes_labels(): + # Labels come from s.project / s.id via loadOverviewPanels. + assert ">'+item.label+'
" not in PAGE_HTML + + +def test_memory_timeline_escapes_prompt_number(): + assert "'+t.prompt_number+'" not in PAGE_HTML + assert "_esc(t.prompt_number)" in PAGE_HTML diff --git a/tests/dashboard/test_server.py b/tests/dashboard/test_server.py index 74a7d61..5224d2b 100644 --- a/tests/dashboard/test_server.py +++ b/tests/dashboard/test_server.py @@ -267,6 +267,31 @@ def test_set_compression_invalid(tmp_path): assert r.status_code == 422 +def test_set_compression_writes_atomically(tmp_path): + """state.json is shared with the MCP server, which reads/writes it + atomically; the dashboard must use atomic_write_text too, not a plain + write_text that truncates in place. Regression for 2026-07-03 review.""" + from context_engine import utils + + client, storage_base = _make_client(tmp_path) + # Pre-existing state keys must survive the rewrite. + (storage_base / "state.json").write_text( + json.dumps({"output_level": "standard", "other_key": "keep-me"}) + ) + with patch( + "context_engine.utils.atomic_write_text", wraps=utils.atomic_write_text + ) as spy: + r = client.post("/api/compression", json={"level": "max"}) + assert r.status_code == 200 + written_paths = [call.args[0] for call in spy.call_args_list] + assert storage_base / "state.json" in written_paths + state = json.loads((storage_base / "state.json").read_text()) + assert state["output_level"] == "max" + assert state["other_key"] == "keep-me" + # No stray tempfiles left behind. + assert not list(storage_base.glob("*.tmp")) + + def test_status_with_versioned_manifest(tmp_path): """Dashboard must read the real Manifest.save() schema, not treat the top-level dict as {file_path: hash}. Regression for 2026-04-27 review.""" diff --git a/tests/indexer/test_chunker.py b/tests/indexer/test_chunker.py index 3f7c64b..e34f7f5 100644 --- a/tests/indexer/test_chunker.py +++ b/tests/indexer/test_chunker.py @@ -66,6 +66,45 @@ def test_chunk_unsupported_language_falls_back(chunker): assert chunks[0].chunk_type == ChunkType.MODULE +MULTIBYTE_PYTHON = '''"""Módulo 🚀 — docstring with émojis and CJK: 日本語.""" + +def first(): + return "a" + +def second(): + return "b" +''' + + +def test_chunk_multibyte_prefix_does_not_shift_offsets(chunker): + """tree-sitter reports BYTE offsets on the utf-8 encoding; slicing the + original str with them garbles every chunk after a multi-byte char. + Chunk contents must exactly equal the function sources.""" + chunks = chunker.chunk(MULTIBYTE_PYTHON, file_path="emoji.py", language="python") + function_chunks = [c for c in chunks if c.chunk_type == ChunkType.FUNCTION] + assert [c.content for c in function_chunks] == [ + 'def first():\n return "a"', + 'def second():\n return "b"', + ] + + +def test_extract_imports_with_multibyte_prefix(chunker): + """Byte offsets must also be handled in _parse_import_module — a CJK + comment before the imports used to shift the sliced module names.""" + source = "# 日本語のコメント 🚀\nimport os\nfrom pathlib import Path\n\ndef main(): pass\n" + _, imports = chunker.chunk_with_imports(source, file_path="main.py", language="python") + assert imports == ["os", "pathlib"] + + +def test_chunk_multibyte_javascript_import(chunker): + """JS string module specifiers after an emoji comment stay intact.""" + source = "// hello 🎉 world\nimport React from 'react';\nfunction App() { return 1; }\n" + chunks, imports = chunker.chunk_with_imports(source, file_path="app.js", language="javascript") + assert imports == ["react"] + fn = [c for c in chunks if c.chunk_type == ChunkType.FUNCTION] + assert fn and fn[0].content == "function App() { return 1; }" + + def test_extract_imports_python(): source = "import os\nfrom pathlib import Path\n\ndef main(): pass\n" chunker = Chunker() diff --git a/tests/indexer/test_pipeline_target_path.py b/tests/indexer/test_pipeline_target_path.py new file mode 100644 index 0000000..06cb47c --- /dev/null +++ b/tests/indexer/test_pipeline_target_path.py @@ -0,0 +1,159 @@ +"""Single-file `target_path` behaviour of the indexing pipeline. + +The watcher funnels every file event through `run_indexing(target_path=...)`, +so the single-file branch must enforce the same protections as the full walk: + + · secret-named files (.env.production, …) are never read or indexed + · .cceignore patterns apply + · symlinked project roots (macOS /tmp → /private/tmp) don't blow up + · a deleted file that's still in the manifest is pruned, not an error +""" +from __future__ import annotations + +from pathlib import Path + +import pytest + +from context_engine.config import load_config +from context_engine.indexer.manifest import Manifest +from context_engine.indexer.pipeline import run_indexing +from context_engine.storage.local_backend import LocalBackend +from context_engine.utils import project_storage_dir + + +@pytest.fixture +def project(tmp_path): + """Minimal project + isolated storage; returns (project_dir, config).""" + project_dir = tmp_path / "proj" + project_dir.mkdir() + (project_dir / "keep.py").write_text("def keep():\n return 1\n") + storage_base = tmp_path / "storage" + storage_base.mkdir() + config = load_config() + config.storage_path = str(storage_base) + return project_dir, config + + +def _backend(config, project_dir: Path) -> LocalBackend: + storage = project_storage_dir(config, Path(project_dir)) + return LocalBackend(base_path=str(storage)) + + +# ── Bug: single-file target bypassed secret + cceignore filters ───────────── + +@pytest.mark.asyncio +async def test_target_path_secret_file_is_not_indexed(project): + """Saving `.env.production` in a watched project routes through + target_path — is_secret_file must apply there too, not only in the + directory walk.""" + project_dir, config = project + (project_dir / ".env.production").write_text( + "DB_PASSWORD=hunter2longvalue123\nSTRIPE_KEY=" + "sk_live_" + "abcdefghij1234567890abcd\n" + ) + + result = await run_indexing( + config, str(project_dir), target_path=".env.production" + ) + + assert result.indexed_files == [] + assert result.total_chunks == 0 + assert not result.errors + assert _backend(config, project_dir).count_chunks() == 0 + + +@pytest.mark.asyncio +async def test_target_path_respects_cceignore(project): + """A `.cceignore`d file must be skipped on the single-file path too.""" + project_dir, config = project + (project_dir / ".cceignore").write_text("*.generated.py\nvendor/\n") + (project_dir / "schema.generated.py").write_text( + "def generated():\n return 'machine output'\n" + ) + vendor = project_dir / "vendor" + vendor.mkdir() + (vendor / "lib.py").write_text("def vendored():\n return 2\n") + + for target in ("schema.generated.py", "vendor/lib.py"): + result = await run_indexing(config, str(project_dir), target_path=target) + assert result.indexed_files == [], f"{target} should have been ignored" + assert result.total_chunks == 0 + assert not result.errors + + assert _backend(config, project_dir).count_chunks() == 0 + + +@pytest.mark.asyncio +async def test_target_path_normal_file_still_indexed(project): + """Sanity: the filters must not over-block ordinary files.""" + project_dir, config = project + result = await run_indexing(config, str(project_dir), target_path="keep.py") + assert result.indexed_files == ["keep.py"] + assert result.total_chunks > 0 + assert not result.errors + + +# ── Bug: symlinked project root broke relative_to on reindex ──────────────── + +@pytest.mark.asyncio +async def test_target_path_through_symlinked_project_root(tmp_path): + """On macOS /tmp resolves to /private/tmp — target paths resolve() to + the real root while project_dir stayed unresolved, so relative_to() + raised ValueError on every watcher-triggered reindex.""" + real = tmp_path / "real_proj" + real.mkdir() + (real / "main.py").write_text("def main():\n return 42\n") + link = tmp_path / "link_proj" + link.symlink_to(real) + + storage_base = tmp_path / "storage" + storage_base.mkdir() + config = load_config() + config.storage_path = str(storage_base) + + result = await run_indexing(config, str(link), target_path="main.py") + + assert not result.errors + assert result.indexed_files == ["main.py"] + assert result.total_chunks > 0 + + +# ── Bug: watcher deletions never pruned the index ──────────────────────────── + +@pytest.mark.asyncio +async def test_target_path_deleted_file_prunes_index_and_manifest(project): + """The watcher enqueues deleted paths too. When the target no longer + exists on disk but is tracked in the manifest, the pipeline must prune + its chunks and manifest entry instead of erroring.""" + project_dir, config = project + (project_dir / "gone.py").write_text("def gone():\n return 'bye'\n") + + first = await run_indexing(config, str(project_dir), full=True) + assert "gone.py" in first.indexed_files + backend = _backend(config, project_dir) + baseline = backend.count_chunks() + assert baseline > 0 + + (project_dir / "gone.py").unlink() + result = await run_indexing(config, str(project_dir), target_path="gone.py") + + assert not result.errors, result.errors + assert result.deleted_files == ["gone.py"] + assert backend.count_chunks() < baseline + + manifest = Manifest( + manifest_path=project_storage_dir(config, project_dir) / "manifest.json" + ) + assert manifest.get_hash("gone.py") is None + assert manifest.get_hash("keep.py") is not None + + +@pytest.mark.asyncio +async def test_target_path_never_indexed_missing_file_still_errors(project): + """A path that neither exists nor is tracked stays an error — we only + treat manifest-tracked paths as deletions.""" + project_dir, config = project + result = await run_indexing( + config, str(project_dir), target_path="never_existed.py" + ) + assert any("not found" in e.lower() for e in result.errors), result.errors + assert result.deleted_files == [] diff --git a/tests/indexer/test_secrets.py b/tests/indexer/test_secrets.py index 7bc61a0..6c963d6 100644 --- a/tests/indexer/test_secrets.py +++ b/tests/indexer/test_secrets.py @@ -127,6 +127,89 @@ def test_placeholders_are_not_redacted(text): assert out == text +# ── Unquoted credential assignments (dotenv / YAML / TOML / shell) ───────── + +@pytest.mark.parametrize("text,leaked", [ + # dotenv-style, no quotes + ("STRIPE_KEY=abcd1234efgh5678ijkl", "abcd1234efgh5678ijkl"), + ("DB_PASSWORD=hunter2longvalue123", "hunter2longvalue123"), + # YAML mapping, no quotes + ("password: hunter2longvalue123", "hunter2longvalue123"), + ("api_key: yamlvalue1234567890abc", "yamlvalue1234567890abc"), + # TOML-ish bare assignment + ("auth_token = tomlvalue1234567890abc", "tomlvalue1234567890abc"), + # shell export + ("export CLIENT_SECRET=shellvalue1234567890", "shellvalue1234567890"), + # indented (nested YAML) + (" db_password: nestedvalue1234567890", "nestedvalue1234567890"), +]) +def test_unquoted_credential_assignments_redacted(text, leaked): + """Dotenv/YAML/TOML/shell assignments don't quote their values — the + generic credential pattern must still catch them.""" + out, fired = redact_secrets(text) + assert fired, f"nothing fired on {text!r}" + assert leaked not in out, f"secret leaked through: {out!r}" + assert "[REDACTED:" in out + # Key name and assignment syntax survive so structure stays parseable. + key = text.split("=")[0].split(":")[0].replace("export", "").strip() + assert key in out + + +@pytest.mark.parametrize("text", [ + # Keyword is a substring but not the trailing segment of the name. + "tokenizer=bert-base-uncased-whole-word", + "keyboard_layout=english-international-x", + # Short values are never credentials. + "password: abc123", + "API_KEY=short", + # Placeholder values stay untouched even unquoted. + "password=your_secret_value_123456", + "SECRET_TOKEN=changeme", + # Keyword mid-sentence (not line-anchored assignment). + "the password token = somethingsomething1234 is described below", + # Quoted values are the quoted pattern's job; no double handling. + 'password = "short"', +]) +def test_unquoted_pattern_false_positive_guards(text): + out, fired = redact_secrets(text) + assert fired == [], f"unexpected redaction in {text!r}: fired={fired}" + assert out == text + + +def test_unquoted_value_with_trailing_comment_redacted(): + out, fired = redact_secrets("API_KEY=realvalue1234567890 # prod key") + assert "realvalue1234567890" not in out + assert fired == ["GENERIC_CREDENTIAL"] + assert "# prod key" in out + + +def test_vendor_match_not_double_redacted(): + """A vendor pattern fires first; the generic unquoted pattern must not + re-redact the [REDACTED:...] placeholder it left behind.""" + out, fired = redact_secrets("STRIPE_KEY=" + "sk_live_" + "abcdefghij1234567890abcd") + assert fired == ["STRIPE_LIVE_KEY"] + assert out == "STRIPE_KEY=[REDACTED:STRIPE_LIVE_KEY]" + + +# ── *.env filename variants ───────────────────────────────────────────────── + +@pytest.mark.parametrize("name", [ + "prod.env", "secrets.env", "staging.env", "config.env", +]) +def test_env_suffix_files_are_skipped(name): + """`.env`-suffixed files (not just names starting with .env) are + credentials by convention — skip them at the filename layer.""" + assert is_secret_file(Path(name)) + assert is_secret_file(Path("deploy") / name) + + +@pytest.mark.parametrize("name", [ + "environment.py", "env.md", "envoy.yaml", +]) +def test_env_like_names_are_not_skipped(name): + assert not is_secret_file(Path(name)) + + # ── scan_and_redact integration ──────────────────────────────────────────── def test_scan_and_redact_skips_secret_filename(tmp_path): diff --git a/tests/memory/test_compressor.py b/tests/memory/test_compressor.py index ba3d94d..10171ed 100644 --- a/tests/memory/test_compressor.py +++ b/tests/memory/test_compressor.py @@ -133,6 +133,157 @@ def test_session_rollup_with_no_turns_is_empty(conn): assert rollup == "" +class _DimEmbedder: + """Deterministic 384-dim embedder so vec-table writes actually land + (the 2-dim stub trips the float[384] dim check and gets swallowed).""" + + def embed_query(self, text: str) -> list[float]: + import hashlib + digest = hashlib.sha256(text.encode("utf-8")).digest() + vec = [((digest[i % 32] / 255.0) - 0.5) for i in range(384)] + n = sum(x * x for x in vec) ** 0.5 or 1.0 + return [x / n for x in vec] + + +def test_recompress_turn_keeps_fts_and_vec_consistent(conn): + """Re-compressing the same turn (Stop + next UserPromptSubmit both + enqueue it) must not leave dangling FTS entries or stale vec rows. + `INSERT OR REPLACE` silently skipped the delete triggers because + recursive_triggers is OFF by default.""" + _seed_session(conn) + _seed_turn(conn, "s1", 1, "KEY discussion first pass. KEY again matters. Filler text here.") + conn.commit() + + for _ in range(2): + memory_compressor.compress_turn( + conn, session_id="s1", prompt_number=1, embedder=_DimEmbedder(), + ) + conn.commit() + + ids = { + r["id"] for r in conn.execute( + "SELECT id FROM turn_summaries " + "WHERE session_id = 's1' AND prompt_number = 1" + ) + } + assert len(ids) == 1 + + # FTS5 external-content integrity check raises SQLITE_CORRUPT_VTAB on + # dangling index entries. + conn.execute( + "INSERT INTO turn_summaries_fts(turn_summaries_fts) " + "VALUES('integrity-check')" + ) + + # Every FTS hit maps back to a live source row. + fts_ids = { + r["rowid"] for r in conn.execute( + "SELECT rowid FROM turn_summaries_fts " + "WHERE turn_summaries_fts MATCH 'KEY'" + ) + } + assert fts_ids <= ids, f"dangling FTS rowids: {fts_ids - ids}" + + # No stale vec rows pointing at replaced rowids. + vec_ids = { + r["rowid"] for r in conn.execute( + "SELECT rowid FROM turn_summaries_vec" + ) + } + assert vec_ids <= ids, f"stale vec rowids: {vec_ids - ids}" + + +def test_drain_skips_rows_at_attempt_cap(conn): + """A row that has already failed _MAX_ATTEMPTS times is dead-lettered: + left in the table for inspection but never picked again.""" + _seed_session(conn) + _seed_turn(conn, "s1", 1, "some text") + conn.execute( + "INSERT INTO pending_compressions (kind, session_id, prompt_number, " + "enqueued_at_epoch, attempts) VALUES ('turn', 's1', 1, 1700000000, ?)", + (memory_compressor._MAX_ATTEMPTS,), + ) + conn.commit() + + did_work = memory_compressor._drain_one_sync(conn, _StubEmbedder()) + assert did_work is False, "capped row must not be picked" + n = conn.execute( + "SELECT COUNT(*) AS n FROM pending_compressions" + ).fetchone()["n"] + assert n == 1, "dead-letter row stays for inspection" + + +def test_failing_row_does_not_starve_queue(conn, monkeypatch): + """A deterministically-failing older row must stop being retried after + _MAX_ATTEMPTS so younger rows still drain.""" + _seed_session(conn) + _seed_turn(conn, "s1", 1, "poison turn text") + _seed_turn(conn, "s1", 2, "healthy KEY turn text. KEY again. Filler.") + conn.execute( + "INSERT INTO pending_compressions (kind, session_id, prompt_number, " + "enqueued_at_epoch) VALUES ('turn', 's1', 1, 100)" + ) + conn.execute( + "INSERT INTO pending_compressions (kind, session_id, prompt_number, " + "enqueued_at_epoch) VALUES ('turn', 's1', 2, 200)" + ) + conn.commit() + + real_build = memory_compressor._build_turn_text + + def _poisoned(c, *, session_id, prompt_number): + if prompt_number == 1: + raise RuntimeError("boom") + return real_build(c, session_id=session_id, prompt_number=prompt_number) + + monkeypatch.setattr(memory_compressor, "_build_turn_text", _poisoned) + + # Enough drains for the poison row to hit the cap plus one for the + # healthy row. + for _ in range(memory_compressor._MAX_ATTEMPTS + 1): + memory_compressor._drain_one_sync(conn, _StubEmbedder()) + + healthy = conn.execute( + "SELECT COUNT(*) AS n FROM turn_summaries " + "WHERE session_id = 's1' AND prompt_number = 2" + ).fetchone()["n"] + assert healthy == 1, "healthy row must drain despite the poison row" + + poison = conn.execute( + "SELECT attempts FROM pending_compressions " + "WHERE prompt_number = 1" + ).fetchone() + assert poison is not None + assert poison["attempts"] == memory_compressor._MAX_ATTEMPTS + + +def test_failed_compression_records_no_savings(conn, monkeypatch): + """Savings must only be recorded when the summary write succeeds — + otherwise every retry of a failing row inflates the ledger.""" + _seed_session(conn) + _seed_turn(conn, "s1", 1, "KEY text that fails late. KEY again. Filler.") + conn.execute( + "INSERT INTO pending_compressions (kind, session_id, prompt_number, " + "enqueued_at_epoch) VALUES ('turn', 's1', 1, 1700000000)" + ) + conn.commit() + + def _boom(*args, **kwargs): + raise RuntimeError("decision extraction exploded") + + monkeypatch.setattr(memory_compressor, "_auto_capture_decisions", _boom) + + for _ in range(2): + memory_compressor._drain_one_sync(conn, _StubEmbedder()) + + n = conn.execute("SELECT COUNT(*) AS n FROM savings_log").fetchone()["n"] + assert n == 0, "failed compressions must not record savings" + attempts = conn.execute( + "SELECT attempts FROM pending_compressions" + ).fetchone()["attempts"] + assert attempts == 2 + + async def test_drain_one_processes_oldest_pending(conn): _seed_session(conn) _seed_turn(conn, "s1", 1, "Turn one with KEY content here. KEY appears twice. Other text.") diff --git a/tests/memory/test_db.py b/tests/memory/test_db.py index e990bd4..521e3e5 100644 --- a/tests/memory/test_db.py +++ b/tests/memory/test_db.py @@ -444,6 +444,62 @@ async def test_auto_prune_loop_stop_event_short_circuits_initial_delay(tmp_path: await asyncio.wait_for(task, timeout=1.0) +def test_bootstrap_without_vec_completes_on_later_connect( + tmp_path: Path, monkeypatch, +): + """A fresh DB bootstrapped while sqlite-vec is unavailable must not + permanently disable semantic recall. The bootstrap must defer the + version stamp (like the upgrade branch does) so a later connect() + with vec support creates the vec tables, and backfill_vec_tables + then picks up rows written while vec was missing.""" + db_path = tmp_path / "memory.db" + + # Phase 1: bootstrap with sqlite-vec unavailable. + monkeypatch.setattr(memory_db, "_try_load_vec", lambda conn: False) + conn = memory_db.connect(db_path) + try: + assert not memory_db.has_vec_tables(conn) + assert memory_db.schema_version(conn) < memory_db.CURRENT_VERSION, ( + "bootstrap must not stamp CURRENT_VERSION when the vec step " + "was skipped — otherwise the v2 step never runs" + ) + # A row written while vec was unavailable. + conn.execute( + "INSERT INTO decisions (decision, reason, source, " + "created_at_epoch, created_at) " + "VALUES (?, ?, 'manual', 1700000000, '2023-11-14T22:13:20')", + ("Recorded while vec was missing", "should backfill later"), + ) + conn.commit() + finally: + conn.close() + + # Reopening still without vec must be a stable no-op (no crash, no stamp). + conn = memory_db.connect(db_path) + try: + assert not memory_db.has_vec_tables(conn) + assert memory_db.schema_version(conn) < memory_db.CURRENT_VERSION + finally: + conn.close() + + # Phase 2: sqlite-vec is available again. + monkeypatch.undo() + conn = memory_db.connect(db_path) + try: + assert memory_db.has_vec_tables(conn), ( + "vec tables must be created once the extension loads" + ) + assert memory_db.schema_version(conn) == memory_db.CURRENT_VERSION + counts = memory_db.backfill_vec_tables(conn, _FakeEmbedder()) + assert counts["decisions"] == 1 + n = conn.execute( + "SELECT COUNT(*) AS n FROM decisions_vec" + ).fetchone()["n"] + assert n == 1 + finally: + conn.close() + + def test_v1_to_v2_upgrade_in_place(tmp_path: Path): """A db stamped at v1 (no vec tables) gains them on the next connect().""" import sqlite3 diff --git a/tests/memory/test_hooks.py b/tests/memory/test_hooks.py index 4adac76..244b94d 100644 --- a/tests/memory/test_hooks.py +++ b/tests/memory/test_hooks.py @@ -344,6 +344,119 @@ async def test_compression_queue_dedupes(hook_app, aiohttp_client): assert n == 1, "double-enqueue should be deduped by UNIQUE constraint" +async def test_session_end_rollup_enqueue_dedupes(hook_app, aiohttp_client): + """Repeated SessionEnd must not enqueue the rollup more than once. + SQLite treats NULLs as distinct in UNIQUE constraints, so the NULL + prompt_number used by rollups defeated INSERT OR IGNORE.""" + app, conn = hook_app + client = await aiohttp_client(app) + await client.post( + "/hooks/SessionStart", json={"session_id": "abc", "project": "demo"}, + ) + for _ in range(3): + resp = await client.post( + "/hooks/SessionEnd", + json={"session_id": "abc", "exit_reason": "normal"}, + ) + assert resp.status == 200 + n = conn.execute( + "SELECT COUNT(*) AS n FROM pending_compressions " + "WHERE kind = 'session_rollup' AND session_id = 'abc'" + ).fetchone()["n"] + assert n == 1, "repeated SessionEnd must dedupe the rollup enqueue" + + +async def test_session_start_bad_started_at_does_not_500( + hook_app, aiohttp_client, +): + """Hooks promise never to 500 — a garbage started_at must fall back to + the current time instead of raising out of the handler.""" + app, conn = hook_app + client = await aiohttp_client(app) + resp = await client.post( + "/hooks/SessionStart", + json={"session_id": "badts", "started_at": "not-a-number"}, + ) + assert resp.status == 200 + row = conn.execute( + "SELECT started_at_epoch FROM sessions WHERE id = 'badts'" + ).fetchone() + assert row is not None + assert row["started_at_epoch"] > 0 + + +# ── PII scrubbing on hook write paths ───────────────────────────────────── + + +@pytest.fixture +def _pii_on(): + memory_db.set_pii_redaction(True) + yield + memory_db.set_pii_redaction(True) + + +async def test_prompt_text_is_pii_scrubbed(hook_app, aiohttp_client, _pii_on): + app, conn = hook_app + client = await aiohttp_client(app) + await client.post( + "/hooks/UserPromptSubmit", + json={ + "session_id": "pii1", + "prompt_text": "Email alice@example.com when 203.0.113.42 is up", + }, + ) + row = conn.execute( + "SELECT prompt_text FROM prompts WHERE session_id = 'pii1'" + ).fetchone() + assert "alice@example.com" not in row["prompt_text"] + assert "203.0.113.42" not in row["prompt_text"] + assert "[REDACTED:EMAIL]" in row["prompt_text"] + + +async def test_tool_payloads_are_pii_scrubbed(hook_app, aiohttp_client, _pii_on): + app, conn = hook_app + client = await aiohttp_client(app) + await client.post( + "/hooks/PostToolUse", + json={ + "session_id": "pii2", + "tool_name": "Bash", + "tool_input": {"command": "curl -u bob@example.com https://x"}, + "tool_output": "reached 203.0.113.42 as bob@example.com", + }, + ) + row = conn.execute( + "SELECT p.raw_input, p.raw_output FROM tool_event_payloads p " + "JOIN tool_events te ON te.payload_id = p.id " + "WHERE te.session_id = 'pii2'" + ).fetchone() + assert "bob@example.com" not in row["raw_input"] + assert "bob@example.com" not in row["raw_output"] + assert "203.0.113.42" not in row["raw_output"] + + +async def test_hook_pii_scrub_respects_disabled_toggle( + hook_app, aiohttp_client, +): + app, conn = hook_app + client = await aiohttp_client(app) + memory_db.set_pii_redaction(False) + try: + await client.post( + "/hooks/UserPromptSubmit", + json={ + "session_id": "pii3", + "prompt_text": "Email carol@example.com about it", + }, + ) + row = conn.execute( + "SELECT prompt_text FROM prompts WHERE session_id = 'pii3'" + ).fetchone() + assert "carol@example.com" in row["prompt_text"] + finally: + memory_db.set_pii_redaction(True) + + # ── Savings visibility tests ────────────────────────────────────────────── diff --git a/tests/retrieval/test_retriever.py b/tests/retrieval/test_retriever.py index 8bf118c..28a1bc6 100644 --- a/tests/retrieval/test_retriever.py +++ b/tests/retrieval/test_retriever.py @@ -66,3 +66,99 @@ async def test_retrieve_with_max_tokens(seeded_retriever): results = await seeded_retriever.retrieve("function", top_k=10, max_tokens=50) total_tokens = sum(c.token_count for c in results) assert total_tokens <= 50 + + +# --------------------------------------------------------------------------- +# FTS-only chunks must not get free vector-similarity credit +# --------------------------------------------------------------------------- + +def _mk_chunk(chunk_id: str, distance: float | None = None) -> Chunk: + chunk = Chunk(id=chunk_id, content=f"def {chunk_id}(): pass", + chunk_type=ChunkType.FUNCTION, file_path=f"{chunk_id}.py", + start_line=1, end_line=1, language="python") + if distance is not None: + chunk.metadata["_distance"] = distance + return chunk + + +class _StubEmbedder: + def embed_query(self, query): + return (0.1, 0.2, 0.3, 0.4) + + +class _StubBackend: + """Minimal backend: fixed vector results, fixed FTS hits, hydration map. + + Deliberately has no get_related_file_paths so graph expansion is skipped. + """ + + def __init__(self, vector_chunks, fts_results, hydrated): + self._vector_chunks = vector_chunks + self._fts_results = fts_results + self._hydrated = hydrated + + async def vector_search(self, query_embedding, top_k=10, filters=None): + return list(self._vector_chunks) + + async def fts_search(self, query, top_k=30): + return list(self._fts_results) + + async def get_chunks_by_ids(self, chunk_ids): + return [self._hydrated[i] for i in chunk_ids if i in self._hydrated] + + +def _spy_scorer(retriever, seen: dict): + orig = retriever._scorer.score + + def spy(chunk, vector_distance, keyword_distance): + seen[chunk.id] = vector_distance + return orig(chunk, vector_distance=vector_distance, + keyword_distance=keyword_distance) + + retriever._scorer.score = spy + + +@pytest.mark.asyncio +async def test_fts_only_chunk_gets_no_free_vector_credit(): + """A chunk hydrated from an FTS-only hit (no _distance metadata) must not + receive a better (lower) vector distance than a genuine vector hit at + moderate distance. Regression: the 0.0 default meant perfect similarity.""" + vec_moderate = _mk_chunk("vec_moderate", distance=0.4) + vec_far = _mk_chunk("vec_far", distance=1.2) + fts_only = _mk_chunk("fts_only") + + backend = _StubBackend( + vector_chunks=[vec_moderate, vec_far], + fts_results=[("fts_only", -5.0)], + hydrated={"fts_only": fts_only}, + ) + retriever = HybridRetriever(backend=backend, embedder=_StubEmbedder()) + seen: dict[str, float] = {} + _spy_scorer(retriever, seen) + + await retriever.retrieve("some query", top_k=5) + + assert "fts_only" in seen and "vec_moderate" in seen + # Higher normalised distance == lower vector-similarity component. + assert seen["fts_only"] >= seen["vec_moderate"] + # It should also be at least as bad as the worst genuine vector hit. + assert seen["fts_only"] >= seen["vec_far"] + + +@pytest.mark.asyncio +async def test_fts_only_without_vector_results_gets_worst_case_distance(): + """With no vector evidence at all, FTS-only chunks get the worst-case + cosine distance (normalised to 1.0), not a perfect 0.0.""" + fts_only = _mk_chunk("fts_only") + backend = _StubBackend( + vector_chunks=[], + fts_results=[("fts_only", -5.0)], + hydrated={"fts_only": fts_only}, + ) + retriever = HybridRetriever(backend=backend, embedder=_StubEmbedder()) + seen: dict[str, float] = {} + _spy_scorer(retriever, seen) + + await retriever.retrieve("some query", top_k=5) + + assert seen["fts_only"] == pytest.approx(1.0) diff --git a/tests/storage/test_fts_store.py b/tests/storage/test_fts_store.py index d9d2816..eae153f 100644 --- a/tests/storage/test_fts_store.py +++ b/tests/storage/test_fts_store.py @@ -60,3 +60,84 @@ async def test_special_chars_in_query(fts, sample_chunks): for q in ['"quoted"', "a-b", "fn(x)", "col:val", "wild*card"]: results = await fts.search(q, top_k=5) assert isinstance(results, list) + + +@pytest.mark.asyncio +async def test_multiword_query_matches_nonconsecutive_terms(fts, sample_chunks): + """Natural-language queries must not be treated as a single strict phrase. + + 'process' / 'card' / 'payment' all appear in c2's content but never + consecutively in this order — the BM25 leg must still find it.""" + await fts.ingest(sample_chunks) + results = await fts.search("process the card payment", top_k=5) + assert "c2" in [r[0] for r in results] + + +@pytest.mark.asyncio +async def test_multiword_query_with_stopwords(fts, sample_chunks): + await fts.ingest(sample_chunks) + results = await fts.search("where is the shipping costs calculation", top_k=5) + assert "c3" in [r[0] for r in results] + + +@pytest.mark.asyncio +async def test_underscore_identifier_still_matches(fts): + chunk = Chunk(id="c9", + content="async def delete_by_files(self, file_paths): pass", + chunk_type=ChunkType.FUNCTION, file_path="store.py", + start_line=1, end_line=1, language="python") + await fts.ingest([chunk]) + results = await fts.search("delete_by_files", top_k=5) + assert "c9" in [r[0] for r in results] + + +@pytest.mark.asyncio +async def test_fts5_operators_are_neutralized(fts, sample_chunks): + """FTS5 operators in user input must not raise or act as operators.""" + await fts.ingest(sample_chunks) + for q in [ + "payment NEAR card", + "content: payment", + "payment AND card", + "NOT payment", + "(payment)", + 'pay* card"', + "payment OR shipping", + "^payment", + ]: + results = await fts.search(q, top_k=5) + assert isinstance(results, list) + + +@pytest.mark.asyncio +async def test_wildcard_does_not_prefix_match(fts, sample_chunks): + """'pay*' must be treated literally (token 'pay'), not as a prefix query, + so it must not match documents containing the token 'payment'.""" + await fts.ingest(sample_chunks) + results = await fts.search("pay*", top_k=5) + assert results == [] + + +@pytest.mark.asyncio +async def test_reingest_same_id_does_not_duplicate(fts, sample_chunks): + """INSERT OR REPLACE never replaces on FTS5 virtual tables; re-ingesting + the same chunk id must not create duplicate rows.""" + await fts.ingest(sample_chunks) + await fts.ingest(sample_chunks) + results = await fts.search("calculate_tax", top_k=10) + assert [r[0] for r in results].count("c1") == 1 + + +@pytest.mark.asyncio +async def test_ingest_failure_leaves_no_partial_rows(fts, sample_chunks): + """A mid-batch ingest failure must roll back, not leave pending rows that + the next unrelated commit silently flushes.""" + bad = Chunk(id="bad", content=["not", "a", "string"], # type: ignore[arg-type] + chunk_type=ChunkType.FUNCTION, file_path="bad.py", + start_line=1, end_line=1, language="python") + with pytest.raises(Exception): + await fts.ingest(sample_chunks + [bad]) + # Trigger an unrelated commit on the same connection. + await fts.delete_by_file("nonexistent.py") + results = await fts.search("calculate_tax", top_k=5) + assert results == [] diff --git a/tests/storage/test_graph_store.py b/tests/storage/test_graph_store.py index ea99844..a4944a2 100644 --- a/tests/storage/test_graph_store.py +++ b/tests/storage/test_graph_store.py @@ -66,3 +66,16 @@ async def test_ingest_empty(graph): await graph.ingest([], []) neighbors = await graph.get_neighbors("nonexistent") assert neighbors == [] + +@pytest.mark.asyncio +async def test_ingest_rolls_back_on_midbatch_failure(graph): + """A mid-batch ingest failure must roll back already-executed inserts so + the next unrelated commit doesn't flush partial graph state.""" + good = GraphNode(id="n_good", node_type=NodeType.FUNCTION, name="f", file_path="a.py") + bad = GraphNode(id="n_bad", node_type=NodeType.FUNCTION, name="g", file_path="a.py", + properties={"x": object()}) # json.dumps fails + with pytest.raises(TypeError): + await graph.ingest([good, bad], []) + # Trigger an unrelated commit on the same connection. + await graph.ingest([], []) + assert await graph.get_nodes_by_file("a.py") == [] diff --git a/tests/storage/test_vector_store.py b/tests/storage/test_vector_store.py index 455a37b..3c2e9c5 100644 --- a/tests/storage/test_vector_store.py +++ b/tests/storage/test_vector_store.py @@ -149,3 +149,92 @@ def test_file_chunk_counts_after_ingest(tmp_path): counts = vs.file_chunk_counts() assert counts["a.py"] == 2 assert counts["b.py"] == 1 + + +@pytest.mark.asyncio +async def test_search_uses_cosine_distance(store): + """The vec0 table must use cosine distance, matching the retriever's + distance/2.0 normalisation. Under the L2 default, `diff_dir` (small + magnitude, different direction) would wrongly rank above `same_dir` + (large magnitude, same direction as the query).""" + a = Chunk(id="same_dir", content="a", chunk_type=ChunkType.FUNCTION, + file_path="a.py", start_line=1, end_line=1, language="python", + embedding=[10.0, 0.0, 0.0, 0.0]) + b = Chunk(id="diff_dir", content="b", chunk_type=ChunkType.FUNCTION, + file_path="b.py", start_line=1, end_line=1, language="python", + embedding=[0.6, 0.8, 0.0, 0.0]) + await store.ingest([a, b]) + results = await store.search(query_embedding=[1.0, 0.0, 0.0, 0.0], top_k=2) + assert [c.id for c in results] == ["same_dir", "diff_dir"] + assert results[0].metadata["_distance"] == pytest.approx(0.0, abs=1e-5) + assert results[1].metadata["_distance"] == pytest.approx(0.4, abs=1e-3) + + +def test_legacy_l2_table_is_rebuilt_with_cosine(tmp_path): + """Opening a store whose chunks_vec was created without + distance_metric=cosine (the old L2 default) must rebuild the index empty + rather than silently mixing metrics.""" + import asyncio + import sqlite3 + import struct + + import sqlite_vec + + db_dir = tmp_path / "vectors" + db_dir.mkdir() + conn = sqlite3.connect(str(db_dir / "vectors.db")) + conn.enable_load_extension(True) + sqlite_vec.load(conn) + conn.enable_load_extension(False) + conn.execute( + "CREATE TABLE chunks (id TEXT PRIMARY KEY, content TEXT NOT NULL, " + "chunk_type TEXT NOT NULL, file_path TEXT NOT NULL, " + "start_line INTEGER NOT NULL, end_line INTEGER NOT NULL, " + "language TEXT NOT NULL)" + ) + conn.execute("CREATE VIRTUAL TABLE chunks_vec USING vec0(embedding float[4])") + conn.execute( + "INSERT INTO chunks VALUES ('c1', 'x', 'function', 'a.py', 1, 1, 'python')" + ) + conn.execute( + "INSERT INTO chunks_vec(rowid, embedding) VALUES (1, ?)", + (struct.pack("4f", 0.1, 0.2, 0.3, 0.4),), + ) + conn.commit() + conn.close() + + vs = VectorStore(db_path=str(db_dir)) + # Legacy L2 index wiped; repopulated on reindex. + assert vs.count() == 0 + assert asyncio.run(vs.search(query_embedding=[0.1, 0.2, 0.3, 0.4], top_k=5)) == [] + + # New ingest recreates the table with the cosine metric. + asyncio.run(vs.ingest([ + Chunk(id="c2", content="y", chunk_type=ChunkType.FUNCTION, + file_path="b.py", start_line=1, end_line=1, language="python", + embedding=[0.1, 0.2, 0.3, 0.4]), + ])) + ddl = vs._conn.execute( + "SELECT sql FROM sqlite_master WHERE type='table' AND name='chunks_vec'" + ).fetchone()[0] + assert "distance_metric=cosine" in ddl + assert vs.count() == 1 + + +@pytest.mark.asyncio +async def test_ingest_rolls_back_on_midbatch_failure(store): + """A failed chunks_vec insert mid-loop must not leave pending rows that a + later unrelated commit flushes (chunks with no vector rows are silently + unfindable).""" + good = Chunk(id="good", content="a", chunk_type=ChunkType.FUNCTION, + file_path="a.py", start_line=1, end_line=1, language="python", + embedding=[0.1, 0.2, 0.3, 0.4]) + bad = Chunk(id="bad", content="b", chunk_type=ChunkType.FUNCTION, + file_path="b.py", start_line=1, end_line=1, language="python", + embedding=[0.1] * 8) # wrong dim -> chunks_vec insert fails + with pytest.raises(Exception): + await store.ingest([good, bad]) + # Trigger an unrelated commit on the same connection. + store.put_cached_compression("unrelated", "short", "text") + assert store.count() == 0 + assert await store.search(query_embedding=[0.1, 0.2, 0.3, 0.4], top_k=5) == [] diff --git a/tests/test_cli_mcp_config.py b/tests/test_cli_mcp_config.py index b36c27a..a9f01ae 100644 --- a/tests/test_cli_mcp_config.py +++ b/tests/test_cli_mcp_config.py @@ -77,3 +77,55 @@ def test_configure_mcp_preserves_other_mcp_servers(tmp_path): assert "some-other-server" in data["mcpServers"] assert data["mcpServers"]["some-other-server"]["command"] == "/opt/other/bin/x" assert "context-engine" in data["mcpServers"] + + +def test_configure_mcp_skips_unparseable_file(tmp_path): + """An unparseable .mcp.json must be left alone (return None), not reset + to {} and overwritten — that would destroy the user's other MCP servers.""" + mcp_path = tmp_path / ".mcp.json" + original = '{"mcpServers": {"other": ' # truncated JSON + mcp_path.write_text(original) + + with patch("context_engine.utils.resolve_cce_binary", return_value="/bin/cce"): + result = _configure_mcp(tmp_path) + + assert result is None + assert mcp_path.read_text() == original + + +def test_configure_mcp_skips_non_dict_top_level(tmp_path): + mcp_path = tmp_path / ".mcp.json" + original = '["not", "an", "object"]' + mcp_path.write_text(original) + + with patch("context_engine.utils.resolve_cce_binary", return_value="/bin/cce"): + result = _configure_mcp(tmp_path) + + assert result is None + assert mcp_path.read_text() == original + + +def test_configure_mcp_null_servers_key(tmp_path): + """`"mcpServers": null` must not raise TypeError; it is replaced by a dict.""" + mcp_path = tmp_path / ".mcp.json" + mcp_path.write_text('{"mcpServers": null, "other": 1}') + + with patch("context_engine.utils.resolve_cce_binary", return_value="/bin/cce"): + changed = _configure_mcp(tmp_path) + assert changed is True + + data = json.loads(mcp_path.read_text()) + assert data["other"] == 1 + assert "context-engine" in data["mcpServers"] + + +def test_configure_mcp_list_servers_key(tmp_path): + mcp_path = tmp_path / ".mcp.json" + mcp_path.write_text('{"mcpServers": []}') + + with patch("context_engine.utils.resolve_cce_binary", return_value="/bin/cce"): + changed = _configure_mcp(tmp_path) + assert changed is True + + data = json.loads(mcp_path.read_text()) + assert "context-engine" in data["mcpServers"] diff --git a/tests/test_cli_uninstall.py b/tests/test_cli_uninstall.py index f6c6f03..170214f 100644 --- a/tests/test_cli_uninstall.py +++ b/tests/test_cli_uninstall.py @@ -179,6 +179,159 @@ def test_uninstall_removes_gitignore_cce_entries(runner, tmp_path): assert "dist/" in remaining +def test_uninstall_preserves_non_cce_git_hook(runner, tmp_path): + """A user hook whose content merely contains 'cce' as a substring + ('success', 'access') must NOT be deleted.""" + project_dir = tmp_path / "proj" + hooks_dir = project_dir / ".git" / "hooks" + hooks_dir.mkdir(parents=True) + user_hook = '#!/bin/sh\necho "deploy success" >> access.log\n' + (hooks_dir / "post-commit").write_text(user_hook) + + result = _run_uninstall_in(runner, project_dir) + assert result.exit_code == 0, result.output + + assert (hooks_dir / "post-commit").exists() + assert (hooks_dir / "post-commit").read_text() == user_hook + + +def test_uninstall_removes_cce_git_hook_block_preserving_user_lines(runner, tmp_path): + """When CCE appended its block to an existing user hook, uninstall must + strip only the CCE block, not delete the whole hook file.""" + from context_engine.indexer.git_hooks import HOOK_MARKER + + project_dir = tmp_path / "proj" + hooks_dir = project_dir / ".git" / "hooks" + hooks_dir.mkdir(parents=True) + (hooks_dir / "post-commit").write_text( + "#!/bin/sh\necho user-stuff\n\n" + f"{HOOK_MARKER}\n'/usr/local/bin/cce' index >/dev/null 2>&1 &\n" + ) + + result = _run_uninstall_in(runner, project_dir) + assert result.exit_code == 0, result.output + + remaining = (hooks_dir / "post-commit").read_text() + assert "echo user-stuff" in remaining + assert HOOK_MARKER not in remaining + assert "cce' index" not in remaining + + +def test_uninstall_deletes_pure_cce_git_hook(runner, tmp_path): + """A hook file CCE created from scratch is deleted entirely.""" + from context_engine.indexer.git_hooks import HOOK_MARKER + + project_dir = tmp_path / "proj" + hooks_dir = project_dir / ".git" / "hooks" + hooks_dir.mkdir(parents=True) + (hooks_dir / "post-merge").write_text( + f"#!/bin/sh\n\n{HOOK_MARKER}\n'/usr/local/bin/cce' index >/dev/null 2>&1 &\n" + ) + + result = _run_uninstall_in(runner, project_dir) + assert result.exit_code == 0, result.output + assert not (hooks_dir / "post-merge").exists() + + +def test_uninstall_preserves_settings_hook_with_cce_substring(runner, tmp_path): + """A user hook whose command contains 'cce' only as a substring + ('log-access', 'success') must survive; real CCE hooks are removed.""" + import json + project_dir = tmp_path / "proj" + project_dir.mkdir() + settings_dir = project_dir / ".claude" + settings_dir.mkdir() + settings = { + "hooks": { + "SessionStart": [ + {"matcher": "", "hooks": [{"type": "command", "command": "/usr/local/bin/cce status --oneline"}]}, + {"matcher": "", "hooks": [{"type": "command", "command": "log-access --verbose"}]}, + ], + "Stop": [ + {"matcher": "", "hooks": [{"type": "command", "command": "notify-on-success.sh"}]}, + {"matcher": "", "hooks": [{"type": "command", "command": "'/Users/me/.cce/hooks/cce_hook.sh' stop"}]}, + ], + }, + } + (settings_dir / "settings.local.json").write_text(json.dumps(settings)) + + result = _run_uninstall_in(runner, project_dir) + assert result.exit_code == 0, result.output + + remaining = json.loads((settings_dir / "settings.local.json").read_text()) + session_cmds = [ + h["command"] + for entry in remaining["hooks"]["SessionStart"] + for h in entry["hooks"] + ] + stop_cmds = [ + h["command"] + for entry in remaining["hooks"]["Stop"] + for h in entry["hooks"] + ] + assert session_cmds == ["log-access --verbose"] + assert stop_cmds == ["notify-on-success.sh"] + + +def test_uninstall_gitignore_preserves_cce_substring_lines(runner, tmp_path): + """.gitignore lines that merely contain 'cce' as a substring + ('access-logs/', '# my access cache') must NOT be removed.""" + project_dir = tmp_path / "proj" + project_dir.mkdir() + (project_dir / ".gitignore").write_text( + "access-logs/\n" + "# my access cache\n" + ".cce/\n" + "success-reports/\n" + ) + + result = _run_uninstall_in(runner, project_dir) + assert result.exit_code == 0, result.output + + remaining = (project_dir / ".gitignore").read_text() + assert "access-logs/" in remaining + assert "# my access cache" in remaining + assert "success-reports/" in remaining + assert ".cce/" not in remaining + + +def test_uninstall_gitignore_removes_cce_comment_lines(runner, tmp_path): + """The exact comment lines ensure_gitignore writes are removed.""" + project_dir = tmp_path / "proj" + project_dir.mkdir() + (project_dir / ".gitignore").write_text( + "node_modules/\n" + "\n" + "# CCE (code-context-engine)\n" + "# CCE local cache (per-machine, not for version control)\n" + ".cce/\n" + "# Claude Code local settings written by cce init\n" + ".claude/settings.local.json\n" + ) + + result = _run_uninstall_in(runner, project_dir) + assert result.exit_code == 0, result.output + + remaining = (project_dir / ".gitignore").read_text() + assert "node_modules/" in remaining + assert "CCE" not in remaining + assert ".cce/" not in remaining + assert ".claude/settings.local.json" not in remaining + + +def test_uninstall_gitignore_untouched_when_no_cce_entries(runner, tmp_path): + """A .gitignore with only lookalike lines is left byte-identical.""" + project_dir = tmp_path / "proj" + project_dir.mkdir() + original = "access-logs/\nsuccess/\n\n# tools\ndist/\n" + (project_dir / ".gitignore").write_text(original) + + result = _run_uninstall_in(runner, project_dir) + assert result.exit_code == 0, result.output + assert (project_dir / ".gitignore").read_text() == original + assert "Removed CCE entries from .gitignore" not in result.output + + def test_uninstall_removes_index_data(runner, tmp_path): """Index data in ~/.cce/projects/ is deleted.""" from unittest.mock import patch as mock_patch diff --git a/tests/test_editors_opencode.py b/tests/test_editors_opencode.py index aac97e7..823ccae 100644 --- a/tests/test_editors_opencode.py +++ b/tests/test_editors_opencode.py @@ -5,7 +5,7 @@ from unittest.mock import patch from context_engine.editors import ( - configure_mcp, detect_editors, remove_mcp, + _strip_jsonc_comments, configure_mcp, detect_editors, remove_mcp, ) @@ -65,6 +65,134 @@ def test_configure_opencode_uses_jsonc_if_exists(tmp_path): assert data["model"] == "test" +# ── _strip_jsonc_comments string-awareness ──────────────────────────── +# Regression: the old regex `//.*?$` truncated any string value containing +# `//` (e.g. "https://..."), json.loads failed, data was reset to {} and the +# user's opencode.json was overwritten with only the context-engine entry. + +def test_strip_jsonc_preserves_url_in_string(): + text = '{\n "url": "https://example.com/mcp"\n}' + assert json.loads(_strip_jsonc_comments(text)) == { + "url": "https://example.com/mcp" + } + + +def test_strip_jsonc_strips_line_comments(): + text = '{\n // a comment\n "a": 1 // trailing\n}' + assert json.loads(_strip_jsonc_comments(text)) == {"a": 1} + + +def test_strip_jsonc_strips_block_comments(): + text = '{\n /* block\n comment */ "a": 1,\n "b": /* inline */ 2\n}' + assert json.loads(_strip_jsonc_comments(text)) == {"a": 1, "b": 2} + + +def test_strip_jsonc_preserves_comment_markers_inside_strings(): + text = '{\n "a": "not // a comment",\n "b": "not /* a comment */"\n}' + assert json.loads(_strip_jsonc_comments(text)) == { + "a": "not // a comment", + "b": "not /* a comment */", + } + + +def test_strip_jsonc_handles_escaped_quotes_in_strings(): + text = '{\n "a": "quote \\" then //", // real comment\n "b": 1\n}' + assert json.loads(_strip_jsonc_comments(text)) == { + "a": 'quote " then //', + "b": 1, + } + + +def test_configure_opencode_preserves_server_with_url(tmp_path): + """A remote MCP server URL must not be truncated as a // comment.""" + existing = { + "mcp": { + "remote-thing": {"type": "remote", "url": "https://example.com/mcp"}, + } + } + (tmp_path / "opencode.json").write_text(json.dumps(existing)) + + with patch("context_engine.editors.resolve_cce_binary", return_value="/usr/bin/cce"): + configure_mcp(tmp_path, "opencode") + + data = json.loads((tmp_path / "opencode.json").read_text()) + assert data["mcp"]["remote-thing"]["url"] == "https://example.com/mcp" + assert "context-engine" in data["mcp"] + + +def test_configure_opencode_jsonc_with_url_and_comments(tmp_path): + (tmp_path / "opencode.jsonc").write_text( + '{\n' + ' // my servers\n' + ' "mcp": {\n' + ' "remote-thing": {"type": "remote", "url": "https://example.com/mcp"}\n' + ' }\n' + '}\n' + ) + + with patch("context_engine.editors.resolve_cce_binary", return_value="/usr/bin/cce"): + changed = configure_mcp(tmp_path, "opencode") + assert changed is True + + data = json.loads((tmp_path / "opencode.jsonc").read_text()) + assert data["mcp"]["remote-thing"]["url"] == "https://example.com/mcp" + assert "context-engine" in data["mcp"] + + +# ── unparseable configs are skipped, never overwritten ──────────────── + +def test_configure_opencode_skips_unparseable_file(tmp_path): + """Truly invalid JSON must be left alone (return None), not clobbered.""" + original = '{"mcp": {"other": ' # truncated JSON + (tmp_path / "opencode.json").write_text(original) + + with patch("context_engine.editors.resolve_cce_binary", return_value="/usr/bin/cce"): + result = configure_mcp(tmp_path, "opencode") + + assert result is None + assert (tmp_path / "opencode.json").read_text() == original + + +def test_configure_json_editor_skips_unparseable_file(tmp_path): + """Generic json editors (e.g. VS Code) must also skip unparseable files.""" + vscode_dir = tmp_path / ".vscode" + vscode_dir.mkdir() + original = '{\n // VS Code allows comments here\n "servers": {"other": {}}\n}' + (vscode_dir / "mcp.json").write_text(original) + + with patch("context_engine.editors.resolve_cce_binary", return_value="/usr/bin/cce"): + result = configure_mcp(tmp_path, "vscode") + + assert result is None + assert (vscode_dir / "mcp.json").read_text() == original + + +# ── non-dict servers key handled gracefully ─────────────────────────── + +def test_configure_opencode_null_mcp_key(tmp_path): + (tmp_path / "opencode.json").write_text('{"model": "test", "mcp": null}') + + with patch("context_engine.editors.resolve_cce_binary", return_value="/usr/bin/cce"): + changed = configure_mcp(tmp_path, "opencode") + assert changed is True + + data = json.loads((tmp_path / "opencode.json").read_text()) + assert data["model"] == "test" + assert "context-engine" in data["mcp"] + + +def test_configure_json_editor_non_dict_servers_key(tmp_path): + (tmp_path / ".vscode").mkdir() + (tmp_path / ".vscode" / "mcp.json").write_text('{"servers": []}') + + with patch("context_engine.editors.resolve_cce_binary", return_value="/usr/bin/cce"): + changed = configure_mcp(tmp_path, "vscode") + assert changed is True + + data = json.loads((tmp_path / ".vscode" / "mcp.json").read_text()) + assert "context-engine" in data["servers"] + + def test_remove_opencode(tmp_path): config = {"mcp": {"context-engine": {"type": "local", "command": ["/usr/bin/cce"]}}} (tmp_path / "opencode.json").write_text(json.dumps(config)) From 89fe36a8969244c074faf9dff5479d5b48b9a68d Mon Sep 17 00:00:00 2001 From: Fazle Elahee Date: Fri, 3 Jul 2026 20:29:17 +0100 Subject: [PATCH 03/14] docs: Phase 1 retrieval-precision implementation plan --- docs/plans/2026-07-03-retrieval-precision.md | 752 +++++++++++++++++++ 1 file changed, 752 insertions(+) create mode 100644 docs/plans/2026-07-03-retrieval-precision.md diff --git a/docs/plans/2026-07-03-retrieval-precision.md b/docs/plans/2026-07-03-retrieval-precision.md new file mode 100644 index 0000000..f0adb97 --- /dev/null +++ b/docs/plans/2026-07-03-retrieval-precision.md @@ -0,0 +1,752 @@ +# Retrieval Precision (Token-Savings Phase 1) Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Cut tokens-served per `context_search` query ≥25% (benchmark-measured, hit-rate unchanged) via overlap dedup, a marginal-utility stop, an activated recency signal, and evidence-tuned confidence cutoff. + +**Architecture:** All ranking changes live in `HybridRetriever.retrieve()`; the recency signal requires persisting `modified_ts` through the vector store (schema + row mapping) and stamping it in the indexing pipeline. Measurement extends the existing `benchmarks/run_benchmark.py` with an A/B mode that runs each query with old vs new retrieval parameters against the same index. + +**Tech Stack:** Python 3.12+, sqlite-vec, pytest (run via `uv run --no-sync pytest`), existing benchmark harness. + +## Global Constraints + +- Spec: `docs/specs/2026-07-03-token-savings-design.md` (Phase 1 section). +- Branch: work on `fix/review-findings-2026-07-03` (Phase 1 builds on the review fixes). +- Config values of `0` disable each new mechanism (spec "Error handling"). +- `modified_ts` migration is forward-only; old rows read as NULL → recency stays neutral 0.5 (today's behavior). +- Run tests with `uv run --no-sync pytest ...` — plain `uv run` re-syncs a stale wheel in this venv. +- No new dependencies. +- Do NOT add "Co-Authored-By" lines to commits. + +--- + +### Task 1: Persist `modified_ts` through the vector store + +**Files:** +- Modify: `src/context_engine/storage/vector_store.py` (`_ensure_tables` ~line 70, `_chunk_to_row` ~line 155, `_row_to_chunk` ~line 165, `ingest` ~line 179, `search` ~line 220, `get_by_id` ~line 373, `get_chunks_by_ids` ~line 388) +- Test: `tests/storage/test_vector_store.py` + +**Interfaces:** +- Consumes: existing `Chunk.metadata` dict (`context_engine/models.py:60`). +- Produces: chunks returned by `search`/`get_by_id`/`get_chunks_by_ids` carry `chunk.metadata["modified_ts"]` (float epoch seconds) when the stored column is non-NULL. Task 2 writes it; Task 5 measures its effect. `ConfidenceScorer._recency_score` (`retrieval/confidence.py:40-47`) already consumes it — no scorer change needed. + +- [ ] **Step 1: Write the failing tests** + +Append to `tests/storage/test_vector_store.py` (follow the file's existing fixture pattern for constructing a store and chunks — reuse its helper if one exists): + +```python +import sqlite3 +import time + +import pytest + +from context_engine.models import Chunk, ChunkType +from context_engine.storage.vector_store import VectorStore + + +def _mk_chunk(cid: str, mtime: float | None = None) -> Chunk: + c = Chunk( + id=cid, + content=f"def {cid}():\n pass\n", + chunk_type=ChunkType.FUNCTION, + file_path="src/mod.py", + start_line=1, + end_line=2, + language="python", + embedding=[0.1, 0.2, 0.3], + ) + if mtime is not None: + c.metadata["modified_ts"] = mtime + return c + + +@pytest.mark.asyncio +async def test_modified_ts_round_trips(tmp_path): + store = VectorStore(db_path=str(tmp_path)) + now = time.time() + await store.ingest([_mk_chunk("with_ts", mtime=now)]) + + results = await store.search([0.1, 0.2, 0.3], top_k=1) + assert results, "expected one search hit" + assert results[0].metadata["modified_ts"] == pytest.approx(now) + + by_id = await store.get_by_id("with_ts") + assert by_id.metadata["modified_ts"] == pytest.approx(now) + + by_ids = await store.get_chunks_by_ids(["with_ts"]) + assert by_ids[0].metadata["modified_ts"] == pytest.approx(now) + + +@pytest.mark.asyncio +async def test_modified_ts_absent_stays_absent(tmp_path): + store = VectorStore(db_path=str(tmp_path)) + await store.ingest([_mk_chunk("no_ts")]) + results = await store.search([0.1, 0.2, 0.3], top_k=1) + assert "modified_ts" not in results[0].metadata + + +@pytest.mark.asyncio +async def test_legacy_db_without_column_is_migrated(tmp_path): + # Simulate a pre-Phase-1 DB: create the store, then drop the column + # by rebuilding the table without it, then reopen. + store = VectorStore(db_path=str(tmp_path)) + await store.ingest([_mk_chunk("old_row")]) + conn = sqlite3.connect(str(tmp_path / "vectors.db")) + conn.executescript( + """ + CREATE TABLE chunks_old AS + SELECT id, content, chunk_type, file_path, start_line, end_line, language + FROM chunks; + DROP TABLE chunks; + ALTER TABLE chunks_old RENAME TO chunks; + """ + ) + conn.commit() + conn.close() + + reopened = VectorStore(db_path=str(tmp_path)) # must not raise + row = await reopened.get_by_id("old_row") + assert row is not None + assert "modified_ts" not in row.metadata # NULL column → neutral recency +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `uv run --no-sync pytest tests/storage/test_vector_store.py -q -k "modified_ts or legacy_db"` +Expected: FAIL — `KeyError: 'modified_ts'` (round-trip) and/or `sqlite3.OperationalError` (legacy reopen may pass; the two metadata tests must fail). + +- [ ] **Step 3: Implement** + +In `src/context_engine/storage/vector_store.py`: + +1. `_ensure_tables` — add the column to the CREATE and migrate legacy tables. Replace the `CREATE TABLE IF NOT EXISTS chunks (...)` statement with: + +```python + self._conn.execute(""" + CREATE TABLE IF NOT EXISTS chunks ( + id TEXT PRIMARY KEY, + content TEXT NOT NULL, + chunk_type TEXT NOT NULL, + file_path TEXT NOT NULL, + start_line INTEGER NOT NULL, + end_line INTEGER NOT NULL, + language TEXT NOT NULL, + modified_ts REAL + ) + """) + # Forward-only migration: pre-Phase-1 DBs lack modified_ts. + # Old rows stay NULL → ConfidenceScorer keeps neutral recency. + cols = { + r[1] for r in self._conn.execute("PRAGMA table_info(chunks)") + } + if "modified_ts" not in cols: + self._conn.execute( + "ALTER TABLE chunks ADD COLUMN modified_ts REAL" + ) +``` + +2. `_chunk_to_row` — append the value: + +```python + def _chunk_to_row(self, chunk: Chunk) -> tuple: + content = chunk.content + if len(content) > _MAX_CONTENT_CHARS: + content = content[:_MAX_CONTENT_CHARS] + "\n...[truncated]" + return ( + chunk.id, content, chunk.chunk_type.value, + chunk.file_path, chunk.start_line, chunk.end_line, + chunk.language, chunk.metadata.get("modified_ts"), + ) +``` + +3. `_row_to_chunk` — restore it (row layout: 7 base columns, optional 8th `modified_ts`): + +```python + def _row_to_chunk(self, row, distance: float | None = None) -> Chunk: + chunk = Chunk( + id=row[0], + content=row[1], + chunk_type=ChunkType(row[2]), + file_path=row[3], + start_line=row[4], + end_line=row[5], + language=row[6], + ) + if len(row) > 7 and row[7] is not None: + chunk.metadata["modified_ts"] = row[7] + if distance is not None: + chunk.metadata["_distance"] = distance + return chunk +``` + +4. `ingest` — extend the upsert (now 8 placeholders): + +```python + rowid = cursor.execute( + "INSERT INTO chunks " + "(id, content, chunk_type, file_path, start_line, end_line, language, modified_ts) " + "VALUES (?, ?, ?, ?, ?, ?, ?, ?) " + "ON CONFLICT(id) DO UPDATE SET " + "content = excluded.content, " + "chunk_type = excluded.chunk_type, " + "file_path = excluded.file_path, " + "start_line = excluded.start_line, " + "end_line = excluded.end_line, " + "language = excluded.language, " + "modified_ts = excluded.modified_ts " + "RETURNING rowid", + row, + ).fetchone()[0] +``` + +5. `search` — both SELECTs gain `c.modified_ts` before `v.distance`: + +```sql + SELECT c.id, c.content, c.chunk_type, c.file_path, + c.start_line, c.end_line, c.language, + c.modified_ts, v.distance +``` + +and the return line becomes: + +```python + return [self._row_to_chunk(row[:8], distance=row[8]) for row in rows] +``` + +6. `get_by_id` and `get_chunks_by_ids` — SELECT list gains `, modified_ts` after `language` (no other change; `_row_to_chunk` handles the 8th column). + +- [ ] **Step 4: Run the storage suite** + +Run: `uv run --no-sync pytest tests/storage/test_vector_store.py -q` +Expected: PASS (all, including the pre-existing cosine/rollback tests). + +- [ ] **Step 5: Commit** + +```bash +git add src/context_engine/storage/vector_store.py tests/storage/test_vector_store.py +git commit -m "feat(retrieval): persist chunk modified_ts through vector store" +``` + +--- + +### Task 2: Stamp `modified_ts` in the indexing pipeline + +**Files:** +- Modify: `src/context_engine/indexer/pipeline.py` (~line 628, right after `chunks, imported_modules = chunk_outcome` inside the `for (file_path, rel_path, content, content_hash, language), chunk_outcome in zip(...)` loop) +- Test: `tests/indexer/test_pipeline_modified_ts.py` (create) + +**Interfaces:** +- Consumes: Task 1's column (transparent — pipeline only sets `chunk.metadata["modified_ts"]`; the store persists it). +- Produces: every chunk ingested by `run_indexing` carries `metadata["modified_ts"] == source file's st_mtime`. This is what makes `ConfidenceScorer._recency_score` return a real decay value instead of the neutral 0.5. + +- [ ] **Step 1: Write the failing test** + +Create `tests/indexer/test_pipeline_modified_ts.py` (mirror the fixture style of `tests/indexer/test_pipeline_target_path.py`, which already builds a minimal project + storage dir and calls `run_indexing`; reuse its config/fixture helpers rather than inventing new ones): + +```python +import asyncio +import os +from pathlib import Path + +import pytest + +from context_engine.config import Config +from context_engine.indexer.pipeline import run_indexing +from context_engine.storage.local_backend import LocalBackend + + +@pytest.mark.asyncio +async def test_indexed_chunks_carry_file_mtime(tmp_path): + project = tmp_path / "proj" + project.mkdir() + src = project / "app.py" + src.write_text("def handler():\n return 42\n") + known_mtime = 1_700_000_000.0 + os.utime(src, (known_mtime, known_mtime)) + + storage = tmp_path / "storage" + config = Config() + config.storage_path = str(storage) + + result = await run_indexing(config, project, full=True) + assert not result.errors + + backend = LocalBackend(base_path=str(storage / project.name)) + chunks = await backend.get_chunks_by_ids( + [cid for cid in await _all_chunk_ids(backend)] + ) + assert chunks, "expected indexed chunks" + for c in chunks: + assert c.metadata.get("modified_ts") == pytest.approx(known_mtime) + + +async def _all_chunk_ids(backend) -> list[str]: + # LocalBackend exposes the vector store; list ids straight from the table. + store = backend._vector_store + with store._lock: + rows = store._conn.execute("SELECT id FROM chunks").fetchall() + return [r[0] for r in rows] +``` + +(If `LocalBackend`'s attribute is named differently — check `src/context_engine/storage/local_backend.py` for the vector-store attribute name and use that. If the existing pipeline tests construct `Config`/storage differently, follow them.) + +- [ ] **Step 2: Run test to verify it fails** + +Run: `uv run --no-sync pytest tests/indexer/test_pipeline_modified_ts.py -q` +Expected: FAIL — `modified_ts` is None/missing. + +- [ ] **Step 3: Implement** + +In `src/context_engine/indexer/pipeline.py`, immediately after `chunks, imported_modules = chunk_outcome` (~line 628): + +```python + # Stamp source mtime so retrieval's recency weight has + # real signal (spec: Phase 1 item 4). stat() failure is + # non-fatal — chunks just keep neutral recency. + try: + _mtime = file_path.stat().st_mtime + except OSError: + _mtime = None + if _mtime is not None: + for _c in chunks: + _c.metadata["modified_ts"] = _mtime +``` + +- [ ] **Step 4: Run tests** + +Run: `uv run --no-sync pytest tests/indexer/test_pipeline_modified_ts.py tests/indexer -q` +Expected: PASS (new test + full indexer suite). + +- [ ] **Step 5: Commit** + +```bash +git add src/context_engine/indexer/pipeline.py tests/indexer/test_pipeline_modified_ts.py +git commit -m "feat(indexer): stamp chunks with source file mtime for recency scoring" +``` + +--- + +### Task 3: Overlap dedup in the retriever + +**Files:** +- Modify: `src/context_engine/retrieval/retriever.py` (new static method + one call after `scored.sort(...)` at line 150) +- Test: `tests/retrieval/test_retriever.py` + +**Interfaces:** +- Consumes: `scored: list[tuple[Chunk, float]]` sorted descending (existing local in `retrieve`). +- Produces: `HybridRetriever._dedupe_overlaps(scored) -> list[tuple[Chunk, float]]` — same-file chunks whose line ranges overlap >50% of the shorter chunk collapse to the higher-scored one. Task 4 iterates its output. + +- [ ] **Step 1: Write the failing tests** + +Append to `tests/retrieval/test_retriever.py` (the file already has a `_mk_chunk`-style helper and stub-backend fixtures from the scoring-fix work — reuse them; the dedup tests below only need the static method, no backend): + +```python +from context_engine.models import Chunk, ChunkType +from context_engine.retrieval.retriever import HybridRetriever + + +def _chunk_at(cid, fp, start, end): + return Chunk( + id=cid, content="x", chunk_type=ChunkType.FUNCTION, + file_path=fp, start_line=start, end_line=end, language="python", + ) + + +def test_dedupe_overlaps_collapses_majority_overlap(): + scored = [ + (_chunk_at("a", "src/m.py", 10, 30), 0.9), # kept (highest) + (_chunk_at("b", "src/m.py", 12, 28), 0.7), # 17/17 lines inside a → dropped + (_chunk_at("c", "src/m.py", 29, 60), 0.6), # 2/32 overlap → kept + (_chunk_at("d", "src/other.py", 10, 30), 0.5), # other file → kept + ] + kept = HybridRetriever._dedupe_overlaps(scored) + assert [c.id for c, _ in kept] == ["a", "c", "d"] + + +def test_dedupe_overlaps_keeps_disjoint_ranges(): + scored = [ + (_chunk_at("a", "src/m.py", 1, 10), 0.9), + (_chunk_at("b", "src/m.py", 11, 20), 0.8), + ] + kept = HybridRetriever._dedupe_overlaps(scored) + assert len(kept) == 2 +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `uv run --no-sync pytest tests/retrieval/test_retriever.py -q -k dedupe` +Expected: FAIL with `AttributeError: ... has no attribute '_dedupe_overlaps'`. + +- [ ] **Step 3: Implement** + +In `src/context_engine/retrieval/retriever.py`, add below `_apply_path_penalty`: + +```python + @staticmethod + def _dedupe_overlaps( + scored: list[tuple[Chunk, float]], + ) -> list[tuple[Chunk, float]]: + """Collapse same-file chunks whose line ranges overlap by more than + half of the shorter chunk, keeping the higher-scored one. Input must + be sorted by score descending; earlier (better) entries win. + Candidate sets are small (≤ top_k*3 per source), so O(n²) is fine. + """ + kept: list[tuple[Chunk, float]] = [] + for chunk, score in scored: + duplicate = False + for kept_chunk, _ in kept: + if kept_chunk.file_path != chunk.file_path: + continue + overlap = ( + min(chunk.end_line, kept_chunk.end_line) + - max(chunk.start_line, kept_chunk.start_line) + + 1 + ) + if overlap <= 0: + continue + shorter = min( + chunk.end_line - chunk.start_line + 1, + kept_chunk.end_line - kept_chunk.start_line + 1, + ) + if shorter > 0 and overlap / shorter > 0.5: + duplicate = True + break + if not duplicate: + kept.append((chunk, score)) + return kept +``` + +Then wire it in `retrieve()` — directly after `scored.sort(key=lambda x: x[1], reverse=True)`: + +```python + scored = self._dedupe_overlaps(scored) +``` + +- [ ] **Step 4: Run tests** + +Run: `uv run --no-sync pytest tests/retrieval -q` +Expected: PASS (new + existing retriever tests). + +- [ ] **Step 5: Commit** + +```bash +git add src/context_engine/retrieval/retriever.py tests/retrieval/test_retriever.py +git commit -m "feat(retrieval): collapse overlapping same-file chunks before packing" +``` + +--- + +### Task 4: Marginal-utility stop + top-1 guarantee + config plumbing + +**Files:** +- Modify: `src/context_engine/config.py` (dataclass field after `retrieval_top_k` line 70, `_EXPECTED_TYPES` line 131, `_apply_dict_to_config` mapping line 156) +- Modify: `src/context_engine/retrieval/retriever.py` (`retrieve()` signature line 37, threshold filter ~line 147, diversity loop ~line 155) +- Modify: `src/context_engine/integration/mcp_server.py` (the `retrieve(...)` call at line 975) +- Test: `tests/retrieval/test_retriever.py`, `tests/test_config.py` (or wherever config mapping tests live — `grep -rl "retrieval_confidence_threshold" tests/` to find it) + +**Interfaces:** +- Consumes: Task 3's deduped `scored` list. +- Produces: `HybridRetriever.retrieve(query, top_k=10, confidence_threshold=0.0, max_tokens=None, marginal_ratio=0.0)` — new keyword arg, `0.0` = disabled (current behavior). `Config.retrieval_marginal_ratio: float = 0.5`, YAML key `retrieval.marginal_ratio`. Behavior: (a) once at least one chunk is selected, stop selecting when `score < marginal_ratio * top_score`; (b) if the `confidence_threshold` filter empties the candidate set but candidates existed, the single best candidate is returned anyway (top-1 guarantee, spec Phase 1 item 1). + +- [ ] **Step 1: Write the failing tests** + +Append to `tests/retrieval/test_retriever.py`. These need the stub backend/embedder used by the existing FTS-distance tests in this file — reuse that fixture; the sketch below assumes a helper `make_retriever(vector_results=...)` exists or is trivially added from the existing stubs: + +```python +import pytest + + +@pytest.mark.asyncio +async def test_marginal_ratio_stops_low_value_tail(retriever_factory): + # Three vector hits with confidences ~[high, high, low] once scored. + # With marginal_ratio=0.5 the third (score < 0.5 * top) is dropped. + retriever, chunks = retriever_factory( + distances=[0.1, 0.3, 1.8], # → vector scores 0.95, 0.85, 0.1 + ) + results = await retriever.retrieve("query", top_k=10, marginal_ratio=0.5) + assert len(results) == 2 + + all_results = await retriever.retrieve("query", top_k=10, marginal_ratio=0.0) + assert len(all_results) == 3 # 0 disables the stop + + +@pytest.mark.asyncio +async def test_top1_guarantee_when_threshold_filters_everything(retriever_factory): + retriever, chunks = retriever_factory(distances=[1.6, 1.8]) + results = await retriever.retrieve( + "query", top_k=10, confidence_threshold=0.99 + ) + assert len(results) == 1 # best candidate survives an over-tight threshold +``` + +(`retriever_factory` is whatever this test file's existing stub pattern provides — adapt names to the file. The distances → score mapping above is illustrative; assert on relative counts, not absolute scores.) + +Config test (in the file found by `grep -rl "retrieval_confidence_threshold" tests/`): + +```python +def test_marginal_ratio_config_mapping(tmp_path): + cfg_file = tmp_path / "config.yaml" + cfg_file.write_text("retrieval:\n marginal_ratio: 0.7\n") + from context_engine.config import load_config + config = load_config(global_path=cfg_file) + assert config.retrieval_marginal_ratio == 0.7 + + +def test_marginal_ratio_default(): + from context_engine.config import Config + assert Config().retrieval_marginal_ratio == 0.5 +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `uv run --no-sync pytest tests/retrieval/test_retriever.py -q -k "marginal or top1" && uv run --no-sync pytest tests/ -q -k marginal_ratio_config` +Expected: FAIL — `TypeError: retrieve() got an unexpected keyword argument 'marginal_ratio'` / `AttributeError: retrieval_marginal_ratio`. + +- [ ] **Step 3: Implement** + +`src/context_engine/config.py`: + +```python + # Retrieval + retrieval_confidence_threshold: float = 0.2 + retrieval_top_k: int = 20 + # Stop adding result chunks once a chunk's score falls below this + # fraction of the top score. 0 disables (always fill to top_k). + retrieval_marginal_ratio: float = 0.5 + bootstrap_max_tokens: int = 10000 +``` + +`_EXPECTED_TYPES`: add `"retrieval_marginal_ratio": (int, float),` after the `retrieval_confidence_threshold` entry. +`_apply_dict_to_config` mapping: add `("retrieval", "marginal_ratio"): "retrieval_marginal_ratio",` after the `confidence_threshold` line. + +`src/context_engine/retrieval/retriever.py` — signature: + +```python + async def retrieve( + self, + query: str, + top_k: int = 10, + confidence_threshold: float = 0.0, + max_tokens: int | None = None, + marginal_ratio: float = 0.0, + ) -> list[Chunk]: +``` + +Threshold filter + top-1 guarantee — replace the current append-if-above-threshold block (lines 147-148) and sort with: + +```python + scored.append((chunk, final_score)) + + scored.sort(key=lambda x: x[1], reverse=True) + scored = self._dedupe_overlaps(scored) + + # Confidence cutoff with a top-1 guarantee: an over-tight threshold + # must never turn a matching query into an empty result. + filtered = [(c, s) for c, s in scored if s >= confidence_threshold] + if not filtered and scored: + filtered = scored[:1] + scored = filtered +``` + +(The unconditional `scored.append` replaces the old `if final_score >= confidence_threshold:` guard inside the loop.) + +Diversity loop — add the marginal stop (`scored` is sorted, so `top_score` is element 0): + +```python + top_score = scored[0][1] if scored else 0.0 + file_counts: dict[str, int] = {} + diverse: list[Chunk] = [] + for chunk, score in scored: + if diverse and marginal_ratio > 0 and score < marginal_ratio * top_score: + break + count = file_counts.get(chunk.file_path, 0) + if count < _MAX_CHUNKS_PER_FILE: + diverse.append(chunk) + file_counts[chunk.file_path] = count + 1 + if len(diverse) >= top_k: + break + ranked = diverse +``` + +`src/context_engine/integration/mcp_server.py` line 975 — add the parameter to the existing call: + +```python + all_chunks = await self._retriever.retrieve( + ...existing args..., + confidence_threshold=self._config.retrieval_confidence_threshold, + marginal_ratio=self._config.retrieval_marginal_ratio, + ) +``` + +(Keep every existing argument as-is; only add `marginal_ratio`.) + +Also in `mcp_server.py`, directly after that `retrieve` call: the spec requires that cutoff-dropped chunks stay discoverable ("nothing becomes unreachable"). When the cutoff/marginal stop trimmed results below the requested `top_k`, append one compact line to the tool response text (find where the response body is assembled for this handler and add it at the end): + +```python + if len(all_chunks) < self._config.retrieval_top_k: + note = ( + "[note: lower-confidence results omitted — raise top_k or " + "lower retrieval.confidence_threshold to include them]" + ) +``` + +Add a test for this in the MCP-server test file that covers `context_search` (find it via `grep -rl "context_search" tests/`): with a config threshold of 0.99 and stub results, the response contains `lower-confidence results omitted`. + +- [ ] **Step 4: Run tests** + +Run: `uv run --no-sync pytest tests/retrieval tests/integration -q && uv run --no-sync pytest tests/ -q -k "config"` +Expected: PASS. + +- [ ] **Step 5: Commit** + +```bash +git add src/context_engine/config.py src/context_engine/retrieval/retriever.py src/context_engine/integration/mcp_server.py tests/ +git commit -m "feat(retrieval): marginal-utility stop with top-1 guarantee, config-driven" +``` + +--- + +### Task 5: Benchmark A/B mode + evidence-based threshold tuning + +**Files:** +- Modify: `benchmarks/run_benchmark.py` (add `--ab` mode) +- Create: `benchmarks/results/cce-phase1-ab.md` (generated output, committed as evidence) +- Possibly modify: `src/context_engine/config.py` (only the `retrieval_confidence_threshold` default, only if evidence supports it) + +**Interfaces:** +- Consumes: `HybridRetriever.retrieve(..., confidence_threshold=..., marginal_ratio=...)` from Task 4 — same index, two parameterizations. +- Produces: a committed A/B report; the Phase 1 success-gate numbers (spec: ≥25% tokens-served reduction, hit-rate unchanged). + +- [ ] **Step 1: Add the A/B mode** + +In `benchmarks/run_benchmark.py`, add CLI args and a comparison pass. Add to `main()`'s argparse block: + +```python + parser.add_argument("--ab", action="store_true", + help="Run each query twice: baseline retrieval vs tuned " + "(threshold + marginal ratio) and print the delta") + parser.add_argument("--threshold", type=float, default=0.35, + help="Tuned confidence_threshold for --ab (default 0.35)") + parser.add_argument("--marginal-ratio", type=float, default=0.5, + help="Tuned marginal_ratio for --ab (default 0.5)") +``` + +Add this function after `run_benchmark`: + +```python +async def run_ab( + project_dir: Path, + queries: list[dict], + storage_dir: Path, + threshold: float, + marginal_ratio: float, +) -> dict: + """Index once, then run every query with baseline vs tuned retrieval + parameters against the same index. No stacking with compression — + this isolates the retrieval-precision change.""" + config = Config() + config.storage_path = str(storage_dir) + + print("Indexing project once for A/B...") + idx = await run_indexing(config, project_dir, full=True) + print(f" {idx.total_chunks} chunks from {len(idx.indexed_files)} files") + + storage_base = Path(config.storage_path) / project_dir.name + backend = LocalBackend(base_path=str(storage_base)) + embedder = Embedder(model_name=config.embedding_model) + retriever = HybridRetriever(backend=backend, embedder=embedder) + + rows = [] + for q in queries: + base = await retriever.retrieve(q["query"], top_k=10) + tuned = await retriever.retrieve( + q["query"], top_k=10, + confidence_threshold=threshold, + marginal_ratio=marginal_ratio, + ) + expected = set(q.get("expected_files", [])) + + def _measure(chunks): + files = {c.file_path for c in chunks} + return { + "tokens": sum(_count_tokens(c.content) for c in chunks), + "chunks": len(chunks), + "hit": bool(files & expected) if expected else None, + } + + rows.append({"query": q["query"], + "base": _measure(base), "tuned": _measure(tuned)}) + b, t = rows[-1]["base"], rows[-1]["tuned"] + print(f" {q['query'][:45]:<45} tokens {b['tokens']:>6} → {t['tokens']:>6} " + f"chunks {b['chunks']:>2} → {t['chunks']:>2} " + f"hit {b['hit']} → {t['hit']}") + + def _agg(side): + tok = sum(r[side]["tokens"] for r in rows) + hits = sum(1 for r in rows if r[side]["hit"]) + judged = sum(1 for r in rows if r[side]["hit"] is not None) + return tok, hits, judged + + base_tok, base_hits, judged = _agg("base") + tuned_tok, tuned_hits, _ = _agg("tuned") + reduction = (1 - tuned_tok / base_tok) * 100 if base_tok else 0.0 + print(f"\nTokens served: {base_tok:,} → {tuned_tok:,} ({reduction:.1f}% reduction)") + print(f"Hit rate: {base_hits}/{judged} → {tuned_hits}/{judged}") + return { + "threshold": threshold, "marginal_ratio": marginal_ratio, + "base_tokens": base_tok, "tuned_tokens": tuned_tok, + "reduction_pct": round(reduction, 1), + "base_hits": base_hits, "tuned_hits": tuned_hits, + "judged": judged, "rows": rows, + } +``` + +Wire into `main()` before the normal `run_benchmark` call: + +```python + if args.ab: + results = asyncio.run(run_ab( + project_dir, queries, storage_dir, + args.threshold, args.marginal_ratio, + )) + if args.json_output: + Path(args.json_output).write_text(json.dumps(results, indent=2) + "\n") + return +``` + +(Reuse the existing `finally` cleanup — place the `if args.ab:` branch inside the existing `try`.) + +- [ ] **Step 2: Run the A/B benchmark on this repo** + +Run: `uv run --no-sync python benchmarks/run_benchmark.py --ab --json-output benchmarks/results/cce-phase1-ab.json 2>&1 | tail -30` +Expected: per-query table plus aggregate lines. Success gate: `reduction_pct >= 25` AND `tuned_hits >= base_hits`. + +- [ ] **Step 3: Tune if the gate fails** + +- If hit rate DROPPED: retry with `--threshold 0.25`, then `--threshold 0.2` (marginal stop alone often carries the reduction). Use the highest threshold with no hit-rate loss. +- If reduction < 25% but hit rate held: try `--marginal-ratio 0.6`. +- Record every run's numbers; the final report must state which values won and what was rejected. + +- [ ] **Step 4: Apply the winning default + write the report** + +- If the winning threshold ≠ 0.2, update `retrieval_confidence_threshold`'s default in `src/context_engine/config.py` to the winning value (spec proposed 0.35 — evidence decides). +- Write `benchmarks/results/cce-phase1-ab.md` by hand from the JSON: date, repo, chosen defaults, the aggregate table (baseline vs tuned tokens, hit rate), and the rejected parameter values with their numbers. + +- [ ] **Step 5: Full suite + commit** + +Run: `uv run --no-sync pytest -q` +Expected: all pass. + +```bash +git add benchmarks/run_benchmark.py benchmarks/results/cce-phase1-ab.md benchmarks/results/cce-phase1-ab.json src/context_engine/config.py +git commit -m "feat(benchmarks): A/B mode for retrieval precision; tune Phase 1 defaults from evidence" +``` + +--- + +### Task 6: Record the outcome + +**Files:** none (MCP tool call + task bookkeeping) + +- [ ] **Step 1:** Call `record_decision` with the final Phase 1 numbers: chosen `retrieval_confidence_threshold` and `retrieval_marginal_ratio` defaults, measured tokens-served reduction, hit-rate before/after, and a pointer to `benchmarks/results/cce-phase1-ab.md`. +- [ ] **Step 2:** Report the numbers to the user with the success-gate verdict (met / not met, and why). From 35f6a856d8c3012e27835a874e52b477d76cf51e Mon Sep 17 00:00:00 2001 From: Fazle Elahee Date: Fri, 3 Jul 2026 20:37:50 +0100 Subject: [PATCH 04/14] feat(retrieval): persist chunk modified_ts through vector store --- src/context_engine/storage/vector_store.py | 35 +++++++---- tests/storage/test_vector_store.py | 69 ++++++++++++++++++++++ 2 files changed, 94 insertions(+), 10 deletions(-) diff --git a/src/context_engine/storage/vector_store.py b/src/context_engine/storage/vector_store.py index aee590f..e166f3d 100644 --- a/src/context_engine/storage/vector_store.py +++ b/src/context_engine/storage/vector_store.py @@ -77,9 +77,19 @@ def _ensure_tables(self) -> None: file_path TEXT NOT NULL, start_line INTEGER NOT NULL, end_line INTEGER NOT NULL, - language TEXT NOT NULL + language TEXT NOT NULL, + modified_ts REAL ) """) + # Forward-only migration: pre-Phase-1 DBs lack modified_ts. + # Old rows stay NULL → ConfidenceScorer keeps neutral recency. + cols = { + r[1] for r in self._conn.execute("PRAGMA table_info(chunks)") + } + if "modified_ts" not in cols: + self._conn.execute( + "ALTER TABLE chunks ADD COLUMN modified_ts REAL" + ) self._conn.execute(""" CREATE INDEX IF NOT EXISTS idx_chunks_file_path ON chunks(file_path) @@ -159,7 +169,7 @@ def _chunk_to_row(self, chunk: Chunk) -> tuple: return ( chunk.id, content, chunk.chunk_type.value, chunk.file_path, chunk.start_line, chunk.end_line, - chunk.language, + chunk.language, chunk.metadata.get("modified_ts"), ) def _row_to_chunk(self, row, distance: float | None = None) -> Chunk: @@ -172,6 +182,8 @@ def _row_to_chunk(self, row, distance: float | None = None) -> Chunk: end_line=row[5], language=row[6], ) + if len(row) > 7 and row[7] is not None: + chunk.metadata["modified_ts"] = row[7] if distance is not None: chunk.metadata["_distance"] = distance return chunk @@ -192,15 +204,16 @@ async def ingest(self, chunks: list[Chunk]) -> None: row = self._chunk_to_row(chunk) rowid = cursor.execute( "INSERT INTO chunks " - "(id, content, chunk_type, file_path, start_line, end_line, language) " - "VALUES (?, ?, ?, ?, ?, ?, ?) " + "(id, content, chunk_type, file_path, start_line, end_line, language, modified_ts) " + "VALUES (?, ?, ?, ?, ?, ?, ?, ?) " "ON CONFLICT(id) DO UPDATE SET " "content = excluded.content, " "chunk_type = excluded.chunk_type, " "file_path = excluded.file_path, " "start_line = excluded.start_line, " "end_line = excluded.end_line, " - "language = excluded.language " + "language = excluded.language, " + "modified_ts = excluded.modified_ts " "RETURNING rowid", row, ).fetchone()[0] @@ -240,7 +253,8 @@ async def search( rows = self._conn.execute( """ SELECT c.id, c.content, c.chunk_type, c.file_path, - c.start_line, c.end_line, c.language, v.distance + c.start_line, c.end_line, c.language, + c.modified_ts, v.distance FROM chunks_vec v JOIN chunks c ON c.rowid = v.rowid WHERE v.embedding MATCH ? AND k = ? @@ -253,7 +267,8 @@ async def search( rows = self._conn.execute( """ SELECT c.id, c.content, c.chunk_type, c.file_path, - c.start_line, c.end_line, c.language, v.distance + c.start_line, c.end_line, c.language, + c.modified_ts, v.distance FROM chunks_vec v JOIN chunks c ON c.rowid = v.rowid WHERE v.embedding MATCH ? AND k = ? @@ -268,7 +283,7 @@ async def search( exc, ) return [] - return [self._row_to_chunk(row[:7], distance=row[7]) for row in rows] + return [self._row_to_chunk(row[:8], distance=row[8]) for row in rows] async def delete_by_file(self, file_path: str) -> None: await self.delete_by_files([file_path]) @@ -374,7 +389,7 @@ async def get_by_id(self, chunk_id: str) -> Chunk | None: with self._lock: try: row = self._conn.execute( - "SELECT id, content, chunk_type, file_path, start_line, end_line, language " + "SELECT id, content, chunk_type, file_path, start_line, end_line, language, modified_ts " "FROM chunks WHERE id = ?", (chunk_id,), ).fetchone() @@ -392,7 +407,7 @@ async def get_chunks_by_ids(self, chunk_ids: list[str]) -> list[Chunk]: try: placeholders = ",".join("?" for _ in chunk_ids) rows = self._conn.execute( - f"SELECT id, content, chunk_type, file_path, start_line, end_line, language " + f"SELECT id, content, chunk_type, file_path, start_line, end_line, language, modified_ts " f"FROM chunks WHERE id IN ({placeholders})", chunk_ids, ).fetchall() diff --git a/tests/storage/test_vector_store.py b/tests/storage/test_vector_store.py index 3c2e9c5..94eddeb 100644 --- a/tests/storage/test_vector_store.py +++ b/tests/storage/test_vector_store.py @@ -1,3 +1,6 @@ +import sqlite3 +import time + import pytest from context_engine.models import Chunk, ChunkType @@ -238,3 +241,69 @@ async def test_ingest_rolls_back_on_midbatch_failure(store): store.put_cached_compression("unrelated", "short", "text") assert store.count() == 0 assert await store.search(query_embedding=[0.1, 0.2, 0.3, 0.4], top_k=5) == [] + + +def _mk_chunk(cid: str, mtime: float | None = None) -> Chunk: + c = Chunk( + id=cid, + content=f"def {cid}():\n pass\n", + chunk_type=ChunkType.FUNCTION, + file_path="src/mod.py", + start_line=1, + end_line=2, + language="python", + embedding=[0.1, 0.2, 0.3], + ) + if mtime is not None: + c.metadata["modified_ts"] = mtime + return c + + +@pytest.mark.asyncio +async def test_modified_ts_round_trips(tmp_path): + store = VectorStore(db_path=str(tmp_path)) + now = time.time() + await store.ingest([_mk_chunk("with_ts", mtime=now)]) + + results = await store.search([0.1, 0.2, 0.3], top_k=1) + assert results, "expected one search hit" + assert results[0].metadata["modified_ts"] == pytest.approx(now) + + by_id = await store.get_by_id("with_ts") + assert by_id.metadata["modified_ts"] == pytest.approx(now) + + by_ids = await store.get_chunks_by_ids(["with_ts"]) + assert by_ids[0].metadata["modified_ts"] == pytest.approx(now) + + +@pytest.mark.asyncio +async def test_modified_ts_absent_stays_absent(tmp_path): + store = VectorStore(db_path=str(tmp_path)) + await store.ingest([_mk_chunk("no_ts")]) + results = await store.search([0.1, 0.2, 0.3], top_k=1) + assert "modified_ts" not in results[0].metadata + + +@pytest.mark.asyncio +async def test_legacy_db_without_column_is_migrated(tmp_path): + # Simulate a pre-Phase-1 DB: create the store, then drop the column + # by rebuilding the table without it, then reopen. + store = VectorStore(db_path=str(tmp_path)) + await store.ingest([_mk_chunk("old_row")]) + conn = sqlite3.connect(str(tmp_path / "vectors.db")) + conn.executescript( + """ + CREATE TABLE chunks_old AS + SELECT id, content, chunk_type, file_path, start_line, end_line, language + FROM chunks; + DROP TABLE chunks; + ALTER TABLE chunks_old RENAME TO chunks; + """ + ) + conn.commit() + conn.close() + + reopened = VectorStore(db_path=str(tmp_path)) # must not raise + row = await reopened.get_by_id("old_row") + assert row is not None + assert "modified_ts" not in row.metadata # NULL column → neutral recency From 6d152f2f4b6f197564c056298ea9b967a64b8724 Mon Sep 17 00:00:00 2001 From: Fazle Elahee Date: Fri, 3 Jul 2026 20:43:57 +0100 Subject: [PATCH 05/14] feat(indexer): stamp chunks with source file mtime for recency scoring --- src/context_engine/indexer/pipeline.py | 11 ++++ tests/indexer/test_pipeline_modified_ts.py | 65 ++++++++++++++++++++++ 2 files changed, 76 insertions(+) create mode 100644 tests/indexer/test_pipeline_modified_ts.py diff --git a/src/context_engine/indexer/pipeline.py b/src/context_engine/indexer/pipeline.py index 410aa1a..5dbd238 100644 --- a/src/context_engine/indexer/pipeline.py +++ b/src/context_engine/indexer/pipeline.py @@ -627,6 +627,17 @@ async def _embed_and_ingest( continue chunks, imported_modules = chunk_outcome + # Stamp source mtime so retrieval's recency weight has + # real signal (spec: Phase 1 item 4). stat() failure is + # non-fatal — chunks just keep neutral recency. + try: + _mtime = file_path.stat().st_mtime + except OSError: + _mtime = None + if _mtime is not None: + for _c in chunks: + _c.metadata["modified_ts"] = _mtime + batch_files_to_replace.append(rel_path) file_node = GraphNode( diff --git a/tests/indexer/test_pipeline_modified_ts.py b/tests/indexer/test_pipeline_modified_ts.py new file mode 100644 index 0000000..a988124 --- /dev/null +++ b/tests/indexer/test_pipeline_modified_ts.py @@ -0,0 +1,65 @@ +"""Each indexed chunk must carry the source file's st_mtime in metadata. + +run_indexing stamps chunk.metadata["modified_ts"] so that +ConfidenceScorer._recency_score gets a real signal instead of the +neutral 0.5 it returns when the key is absent. +""" +from __future__ import annotations + +import os +from pathlib import Path + +import pytest + +from context_engine.config import load_config +from context_engine.indexer.pipeline import run_indexing +from context_engine.storage.local_backend import LocalBackend +from context_engine.utils import project_storage_dir + + +@pytest.fixture +def project(tmp_path): + """Minimal project + isolated storage; returns (project_dir, config).""" + project_dir = tmp_path / "proj" + project_dir.mkdir() + storage_base = tmp_path / "storage" + storage_base.mkdir() + config = load_config() + config.storage_path = str(storage_base) + return project_dir, config + + +def _backend(config, project_dir: Path) -> LocalBackend: + storage = project_storage_dir(config, Path(project_dir)) + return LocalBackend(base_path=str(storage)) + + +def _all_chunk_ids(backend: LocalBackend) -> list[str]: + """List all chunk ids directly from the SQLite table.""" + store = backend._vector_store + with store._lock: + rows = store._conn.execute("SELECT id FROM chunks").fetchall() + return [r[0] for r in rows] + + +@pytest.mark.asyncio +async def test_indexed_chunks_carry_file_mtime(project): + project_dir, config = project + src = project_dir / "app.py" + src.write_text("def handler():\n return 42\n") + known_mtime = 1_700_000_000.0 + os.utime(src, (known_mtime, known_mtime)) + + result = await run_indexing(config, str(project_dir), full=True) + assert not result.errors, result.errors + + backend = _backend(config, project_dir) + chunk_ids = _all_chunk_ids(backend) + assert chunk_ids, "expected indexed chunks" + + chunks = await backend.get_chunks_by_ids(chunk_ids) + assert chunks, "expected indexed chunks" + for c in chunks: + assert c.metadata.get("modified_ts") == pytest.approx(known_mtime), ( + f"chunk {c.id} missing modified_ts; got metadata={c.metadata}" + ) From 244c4b5e983743281e415700e9d938188c12f681 Mon Sep 17 00:00:00 2001 From: Fazle Elahee Date: Fri, 3 Jul 2026 20:49:09 +0100 Subject: [PATCH 06/14] fix(indexer): stamp git commit chunks with commit-time modified_ts --- src/context_engine/indexer/git_indexer.py | 27 +++++++++++---- tests/indexer/test_git_indexer.py | 42 ++++++++++++++++++++++- 2 files changed, 62 insertions(+), 7 deletions(-) diff --git a/src/context_engine/indexer/git_indexer.py b/src/context_engine/indexer/git_indexer.py index 21bd94a..7d439a4 100644 --- a/src/context_engine/indexer/git_indexer.py +++ b/src/context_engine/indexer/git_indexer.py @@ -1,5 +1,6 @@ """Parse git log into searchable chunks.""" import asyncio +import datetime import logging import re import subprocess @@ -99,6 +100,25 @@ def _parse_meta( content = f"{subject}\n\n{body}".strip() short_hash = commit_hash[:7] + metadata = { + "author": author, + "date": date, + "hash": commit_hash, + "chunk_kind": "commit", + } + # Stamp commit time as epoch so retrieval's recency weight applies + # to git-history chunks too. `date` is git log %ai, e.g. + # "2026-07-03 20:15:25 +0100" — fromisoformat parses the ±HHMM + # offset on Python 3.11+. Parse failure is non-fatal — the chunk + # just keeps neutral recency (matching the stat() convention in + # the file-indexing pipeline). + try: + metadata["modified_ts"] = datetime.datetime.fromisoformat( + date + ).timestamp() + except ValueError: + pass + chunk = Chunk( id=f"commit_{short_hash}", content=content, @@ -107,12 +127,7 @@ def _parse_meta( start_line=0, end_line=0, language="git", - metadata={ - "author": author, - "date": date, - "hash": commit_hash, - "chunk_kind": "commit", - }, + metadata=metadata, ) chunks.append(chunk) diff --git a/tests/indexer/test_git_indexer.py b/tests/indexer/test_git_indexer.py index ae993eb..a2842e1 100644 --- a/tests/indexer/test_git_indexer.py +++ b/tests/indexer/test_git_indexer.py @@ -1,5 +1,8 @@ """Tests for git history indexer.""" +import os import subprocess +from datetime import datetime + import pytest from context_engine.indexer.git_indexer import index_commits from context_engine.models import ChunkType, NodeType, EdgeType @@ -14,7 +17,10 @@ def git_repo(tmp_path): for i in range(3): (tmp_path / f"file{i}.py").write_text(f"def fn{i}(): pass\n") subprocess.run(["git", "add", "."], cwd=tmp_path, capture_output=True, check=True) - subprocess.run(["git", "commit", "-m", f"Add file{i}"], cwd=tmp_path, capture_output=True, check=True) + subprocess.run( + ["git", "-c", "commit.gpgsign=false", "commit", "-m", f"Add file{i}"], + cwd=tmp_path, capture_output=True, check=True, + ) return tmp_path @@ -43,6 +49,40 @@ async def test_commit_nodes_and_edges(git_repo): assert all(e.edge_type == EdgeType.MODIFIES for e in edges) +@pytest.mark.asyncio +async def test_commit_chunks_carry_commit_time_modified_ts(git_repo): + """Commit chunks must stamp modified_ts (epoch float of the author date) + so ConfidenceScorer's recency weight applies to git history, not just + file chunks.""" + fixed_date = "2026-07-03T12:00:00+01:00" + env = { + **os.environ, + "GIT_AUTHOR_DATE": fixed_date, + "GIT_COMMITTER_DATE": fixed_date, + } + (git_repo / "dated.py").write_text("def dated(): pass\n") + subprocess.run(["git", "add", "."], cwd=git_repo, capture_output=True, check=True) + subprocess.run( + ["git", "-c", "commit.gpgsign=false", "commit", "-m", "Dated commit"], + cwd=git_repo, capture_output=True, check=True, env=env, + ) + + chunks, _, _ = await index_commits(git_repo, max_commits=10) + + dated = next(c for c in chunks if c.content.startswith("Dated commit")) + expected = datetime.fromisoformat(fixed_date).timestamp() + assert isinstance(dated.metadata.get("modified_ts"), float), ( + f"expected modified_ts float, got metadata={dated.metadata}" + ) + assert dated.metadata["modified_ts"] == pytest.approx(expected) + + # Every commit chunk carries the stamp, not just the fixed-date one. + for c in chunks: + assert isinstance(c.metadata.get("modified_ts"), float), ( + f"chunk {c.id} missing modified_ts; got metadata={c.metadata}" + ) + + @pytest.mark.asyncio async def test_incremental_since_sha(git_repo): chunks_all, _, _ = await index_commits(git_repo, max_commits=10) From f0f3dede518b4dc5044a07a822dc63118d86f6d5 Mon Sep 17 00:00:00 2001 From: Fazle Elahee Date: Fri, 3 Jul 2026 20:52:22 +0100 Subject: [PATCH 07/14] feat(retrieval): collapse overlapping same-file chunks before packing --- src/context_engine/retrieval/retriever.py | 34 +++++++++++++++++++++++ tests/retrieval/test_retriever.py | 31 +++++++++++++++++++++ 2 files changed, 65 insertions(+) diff --git a/src/context_engine/retrieval/retriever.py b/src/context_engine/retrieval/retriever.py index a758311..8586a34 100644 --- a/src/context_engine/retrieval/retriever.py +++ b/src/context_engine/retrieval/retriever.py @@ -148,6 +148,7 @@ async def retrieve( scored.append((chunk, final_score)) scored.sort(key=lambda x: x[1], reverse=True) + scored = self._dedupe_overlaps(scored) # File diversity: cap chunks per file so one large file doesn't # dominate the result set. This improves precision by letting @@ -220,6 +221,39 @@ def _apply_path_penalty(file_path: str, score: float) -> float: return score * 0.8 return score + @staticmethod + def _dedupe_overlaps( + scored: list[tuple[Chunk, float]], + ) -> list[tuple[Chunk, float]]: + """Collapse same-file chunks whose line ranges overlap by more than + half of the shorter chunk, keeping the higher-scored one. Input must + be sorted by score descending; earlier (better) entries win. + Candidate sets are small (≤ top_k*3 per source), so O(n²) is fine. + """ + kept: list[tuple[Chunk, float]] = [] + for chunk, score in scored: + duplicate = False + for kept_chunk, _ in kept: + if kept_chunk.file_path != chunk.file_path: + continue + overlap = ( + min(chunk.end_line, kept_chunk.end_line) + - max(chunk.start_line, kept_chunk.start_line) + + 1 + ) + if overlap <= 0: + continue + shorter = min( + chunk.end_line - chunk.start_line + 1, + kept_chunk.end_line - kept_chunk.start_line + 1, + ) + if shorter > 0 and overlap / shorter > 0.5: + duplicate = True + break + if not duplicate: + kept.append((chunk, score)) + return kept + def _estimate_keyword_distance(self, chunk, parsed) -> int: if parsed.file_hints: for hint in parsed.file_hints: diff --git a/tests/retrieval/test_retriever.py b/tests/retrieval/test_retriever.py index 28a1bc6..6e0d7fd 100644 --- a/tests/retrieval/test_retriever.py +++ b/tests/retrieval/test_retriever.py @@ -162,3 +162,34 @@ async def test_fts_only_without_vector_results_gets_worst_case_distance(): await retriever.retrieve("some query", top_k=5) assert seen["fts_only"] == pytest.approx(1.0) + + +# --------------------------------------------------------------------------- +# Overlap dedup: collapse same-file chunks with >50% line range overlap +# --------------------------------------------------------------------------- + +def _chunk_at(cid, fp, start, end): + return Chunk( + id=cid, content="x", chunk_type=ChunkType.FUNCTION, + file_path=fp, start_line=start, end_line=end, language="python", + ) + + +def test_dedupe_overlaps_collapses_majority_overlap(): + scored = [ + (_chunk_at("a", "src/m.py", 10, 30), 0.9), # kept (highest) + (_chunk_at("b", "src/m.py", 12, 28), 0.7), # 17/17 lines inside a → dropped + (_chunk_at("c", "src/m.py", 29, 60), 0.6), # 2/32 overlap → kept + (_chunk_at("d", "src/other.py", 10, 30), 0.5), # other file → kept + ] + kept = HybridRetriever._dedupe_overlaps(scored) + assert [c.id for c, _ in kept] == ["a", "c", "d"] + + +def test_dedupe_overlaps_keeps_disjoint_ranges(): + scored = [ + (_chunk_at("a", "src/m.py", 1, 10), 0.9), + (_chunk_at("b", "src/m.py", 11, 20), 0.8), + ] + kept = HybridRetriever._dedupe_overlaps(scored) + assert len(kept) == 2 From 35e060a1225ea31ed23f4dc666517a371d5d3ac8 Mon Sep 17 00:00:00 2001 From: Fazle Elahee Date: Fri, 3 Jul 2026 21:03:09 +0100 Subject: [PATCH 08/14] feat(retrieval): marginal-utility stop with top-1 guarantee, config-driven --- src/context_engine/config.py | 5 ++ src/context_engine/integration/mcp_server.py | 7 +++ src/context_engine/retrieval/retriever.py | 16 +++++-- tests/integration/test_mcp_server.py | 40 ++++++++++++++++ tests/retrieval/test_retriever.py | 49 ++++++++++++++++++++ tests/test_config.py | 13 ++++++ 6 files changed, 127 insertions(+), 3 deletions(-) diff --git a/src/context_engine/config.py b/src/context_engine/config.py index f3b13d0..3d42532 100644 --- a/src/context_engine/config.py +++ b/src/context_engine/config.py @@ -68,6 +68,9 @@ class Config: # Retrieval retrieval_confidence_threshold: float = 0.2 retrieval_top_k: int = 20 + # Stop adding result chunks once a chunk's score falls below this + # fraction of the top score. 0 disables (always fill to top_k). + retrieval_marginal_ratio: float = 0.5 bootstrap_max_tokens: int = 10000 # Indexer @@ -130,6 +133,7 @@ def _deep_merge(base: dict, override: dict) -> dict: "ollama_embed_model": str, "retrieval_confidence_threshold": (int, float), "retrieval_top_k": int, + "retrieval_marginal_ratio": (int, float), "bootstrap_max_tokens": int, "indexer_watch": bool, "indexer_debounce_ms": int, @@ -154,6 +158,7 @@ def _apply_dict_to_config(config: Config, data: dict) -> None: ("embedding", "ollama_model"): "ollama_embed_model", ("retrieval", "confidence_threshold"): "retrieval_confidence_threshold", ("retrieval", "top_k"): "retrieval_top_k", + ("retrieval", "marginal_ratio"): "retrieval_marginal_ratio", ("retrieval", "bootstrap_max_tokens"): "bootstrap_max_tokens", ("indexer", "watch"): "indexer_watch", ("indexer", "debounce_ms"): "indexer_debounce_ms", diff --git a/src/context_engine/integration/mcp_server.py b/src/context_engine/integration/mcp_server.py index a57aa5a..cab8545 100644 --- a/src/context_engine/integration/mcp_server.py +++ b/src/context_engine/integration/mcp_server.py @@ -976,6 +976,7 @@ async def _handle_context_search(self, args): query, top_k=top_k * 2, confidence_threshold=self._config.retrieval_confidence_threshold, + marginal_ratio=self._config.retrieval_marginal_ratio, max_tokens=None, ) all_chunks = await self._compressor.compress(all_chunks, self._config.compression_level) @@ -1026,6 +1027,12 @@ async def _handle_context_search(self, args): body = _format_results_with_overflow(inline_chunks, overflow_chunks) body = self._apply_output_compression(body) + if len(all_chunks) < self._config.retrieval_top_k: + note = ( + "[note: lower-confidence results omitted — raise top_k or " + "lower retrieval.confidence_threshold to include them]" + ) + body = body + "\n" + note self._record(raw_tokens, served_tokens, full_file_tokens) # Compliance audit log — file:line refs of every served chunk + the # score range. Off by default; enable via config.audit_log_enabled. diff --git a/src/context_engine/retrieval/retriever.py b/src/context_engine/retrieval/retriever.py index 8586a34..343385d 100644 --- a/src/context_engine/retrieval/retriever.py +++ b/src/context_engine/retrieval/retriever.py @@ -40,6 +40,7 @@ async def retrieve( top_k: int = 10, confidence_threshold: float = 0.0, max_tokens: int | None = None, + marginal_ratio: float = 0.0, ) -> list[Chunk]: parsed = self._parser.parse(query) query_embedding = self._embedder.embed_query(query) @@ -144,18 +145,27 @@ async def retrieve( final_score = self._apply_path_penalty(chunk.file_path, final_score) chunk.confidence_score = final_score - if final_score >= confidence_threshold: - scored.append((chunk, final_score)) + scored.append((chunk, final_score)) scored.sort(key=lambda x: x[1], reverse=True) scored = self._dedupe_overlaps(scored) + # Confidence cutoff with a top-1 guarantee: an over-tight threshold + # must never turn a matching query into an empty result. + filtered = [(c, s) for c, s in scored if s >= confidence_threshold] + if not filtered and scored: + filtered = scored[:1] + scored = filtered + # File diversity: cap chunks per file so one large file doesn't # dominate the result set. This improves precision by letting # chunks from more files surface into the top-k. + top_score = scored[0][1] if scored else 0.0 file_counts: dict[str, int] = {} diverse: list[Chunk] = [] - for chunk, _ in scored: + for chunk, score in scored: + if diverse and marginal_ratio > 0 and score < marginal_ratio * top_score: + break count = file_counts.get(chunk.file_path, 0) if count < _MAX_CHUNKS_PER_FILE: diverse.append(chunk) diff --git a/tests/integration/test_mcp_server.py b/tests/integration/test_mcp_server.py index 6b9e30e..8072503 100644 --- a/tests/integration/test_mcp_server.py +++ b/tests/integration/test_mcp_server.py @@ -141,3 +141,43 @@ async def test_index_status_with_tracked_stats(tmp_path): assert "400" in text # served assert "600" in text # saved assert "60%" in text + + +@pytest.mark.asyncio +async def test_context_search_appends_omitted_note_when_results_below_top_k(tmp_path): + """When retrieve() returns fewer results than retrieval_top_k, a note is appended.""" + from unittest.mock import AsyncMock + from context_engine.models import Chunk, ChunkType + + server = _make_server(tmp_path) + + stub_chunk = Chunk( + id="c1", content="def foo(): pass", + chunk_type=ChunkType.FUNCTION, file_path="src/foo.py", + start_line=1, end_line=1, language="python", + ) + stub_chunk.confidence_score = 0.8 + + server._retriever = MagicMock() + server._retriever.retrieve = AsyncMock(return_value=[stub_chunk]) + server._compressor = MagicMock() + server._compressor.compress = AsyncMock(return_value=[stub_chunk]) + server._session_capture = MagicMock() + server._session_capture.touch_files = MagicMock() + server._persist_current_session = MagicMock() + server._record = MagicMock() + server._append_audit_log = MagicMock() + server._ensure_indexed = AsyncMock(return_value=True) + + # retrieval_top_k=5, but only 1 chunk returned → note should appear + server._config.retrieval_top_k = 5 + server._config.retrieval_confidence_threshold = 0.99 + server._config.retrieval_marginal_ratio = 0.5 + server._config.output_compression = "off" + server._output_level = "off" + server._project_name = "test-project" + server._session_id = "test-session" + + result = await server._handle_context_search({"query": "find something", "top_k": 5}) + text = result[0].text + assert "lower-confidence results omitted" in text diff --git a/tests/retrieval/test_retriever.py b/tests/retrieval/test_retriever.py index 6e0d7fd..f9ac10e 100644 --- a/tests/retrieval/test_retriever.py +++ b/tests/retrieval/test_retriever.py @@ -193,3 +193,52 @@ def test_dedupe_overlaps_keeps_disjoint_ranges(): ] kept = HybridRetriever._dedupe_overlaps(scored) assert len(kept) == 2 + + +# --------------------------------------------------------------------------- +# Marginal-utility stop + top-1 guarantee +# --------------------------------------------------------------------------- + +@pytest.fixture +def retriever_factory(): + """Stub retriever whose last chunk is vector-only (not in FTS). + + That creates a meaningful RRF gap: the top chunks appear in both vector + and FTS (doubling their RRF score) while the tail chunk only gets half + the RRF credit, driving its final blended score below 0.5 * top_score + when combined with a high cosine distance. + """ + def make(distances): + chunks = [_mk_chunk(f"cm{i}", distance=d) for i, d in enumerate(distances)] + # All but the last chunk appear in FTS — gives them a 2x RRF boost + # relative to the last chunk, which is vector-only. + fts_results = [(c.id, -float(i)) for i, c in enumerate(chunks[:-1])] + backend = _StubBackend( + vector_chunks=chunks, + fts_results=fts_results, + hydrated={}, + ) + return HybridRetriever(backend=backend, embedder=_StubEmbedder()), chunks + return make + + +@pytest.mark.asyncio +async def test_marginal_ratio_stops_low_value_tail(retriever_factory): + # Three vector hits: first two appear in both vector+FTS (high RRF), + # third is vector-only with high distance (low conf + low RRF). + # With marginal_ratio=0.5 the third (score < 0.5 * top) is dropped. + retriever, chunks = retriever_factory(distances=[0.1, 0.3, 1.8]) + results = await retriever.retrieve("query", top_k=10, marginal_ratio=0.5) + assert len(results) == 2 + + all_results = await retriever.retrieve("query", top_k=10, marginal_ratio=0.0) + assert len(all_results) == 3 # 0 disables the stop + + +@pytest.mark.asyncio +async def test_top1_guarantee_when_threshold_filters_everything(retriever_factory): + retriever, chunks = retriever_factory(distances=[1.6, 1.8]) + results = await retriever.retrieve( + "query", top_k=10, confidence_threshold=0.99 + ) + assert len(results) == 1 # best candidate survives an over-tight threshold diff --git a/tests/test_config.py b/tests/test_config.py index 70e65fe..866993a 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -85,3 +85,16 @@ def test_ollama_url_yaml_type_validation(tmp_path): })) with pytest.raises(ValueError, match="ollama_url"): load_config(global_path=config_file) + + +def test_marginal_ratio_config_mapping(tmp_path): + cfg_file = tmp_path / "config.yaml" + cfg_file.write_text("retrieval:\n marginal_ratio: 0.7\n") + from context_engine.config import load_config + config = load_config(global_path=cfg_file) + assert config.retrieval_marginal_ratio == 0.7 + + +def test_marginal_ratio_default(): + from context_engine.config import Config + assert Config().retrieval_marginal_ratio == 0.5 From 58c383417a7fb14e086b61347294ebb700235c15 Mon Sep 17 00:00:00 2001 From: Fazle Elahee Date: Fri, 3 Jul 2026 21:09:02 +0100 Subject: [PATCH 09/14] docs: correct omitted-results note design in Phase 1 plan (truthful dropped-count signal) --- docs/plans/2026-07-03-retrieval-precision.md | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/docs/plans/2026-07-03-retrieval-precision.md b/docs/plans/2026-07-03-retrieval-precision.md index f0adb97..b20f5f9 100644 --- a/docs/plans/2026-07-03-retrieval-precision.md +++ b/docs/plans/2026-07-03-retrieval-precision.md @@ -580,17 +580,18 @@ Diversity loop — add the marginal stop (`scored` is sorted, so `top_score` is (Keep every existing argument as-is; only add `marginal_ratio`.) -Also in `mcp_server.py`, directly after that `retrieve` call: the spec requires that cutoff-dropped chunks stay discoverable ("nothing becomes unreachable"). When the cutoff/marginal stop trimmed results below the requested `top_k`, append one compact line to the tool response text (find where the response body is assembled for this handler and add it at the end): +Also in `mcp_server.py`: the spec requires that cutoff-dropped chunks stay discoverable ("nothing becomes unreachable"). `retrieve()` gains an optional `stats_out: dict | None = None` keyword; when provided it is filled with `candidates` (post-dedup, pre-filter count), `selected` (returned count), and `dropped_low_value` (candidates excluded specifically by the confidence threshold or the marginal stop — NOT by the per-file diversity cap or `top_k`). The `context_search` handler passes a stats dict and appends one compact note line to the response only when `stats["dropped_low_value"] > 0`: ```python - if len(all_chunks) < self._config.retrieval_top_k: note = ( "[note: lower-confidence results omitted — raise top_k or " "lower retrieval.confidence_threshold to include them]" ) ``` -Add a test for this in the MCP-server test file that covers `context_search` (find it via `grep -rl "context_search" tests/`): with a config threshold of 0.99 and stub results, the response contains `lower-confidence results omitted`. +(Rationale: an earlier draft compared the post-compression chunk count against `retrieval_top_k`, which false-positives on nearly every query — compression and unrelated filters shrink the list too.) + +Add tests: retriever `stats_out` accounting (threshold drop, marginal-stop drop, and no-drop → 0), and MCP-server note behavior (present when `dropped_low_value > 0`; absent when retrieval merely returned fewer than `retrieval_top_k` with no drops). - [ ] **Step 4: Run tests** From 78aafce1ae31276e15e8dca5af0f5e9e654eb9d8 Mon Sep 17 00:00:00 2001 From: Fazle Elahee Date: Fri, 3 Jul 2026 21:12:06 +0100 Subject: [PATCH 10/14] fix(retrieval): truthful dropped-results signal for the omitted-results note --- src/context_engine/integration/mcp_server.py | 7 +++- src/context_engine/retrieval/retriever.py | 22 +++++++++- tests/integration/test_mcp_server.py | 34 +++++++++++++--- tests/retrieval/test_retriever.py | 42 ++++++++++++++++++++ 4 files changed, 98 insertions(+), 7 deletions(-) diff --git a/src/context_engine/integration/mcp_server.py b/src/context_engine/integration/mcp_server.py index cab8545..7043f82 100644 --- a/src/context_engine/integration/mcp_server.py +++ b/src/context_engine/integration/mcp_server.py @@ -972,12 +972,14 @@ async def _handle_context_search(self, args): max_tokens = 8000 # Fetch 2x candidates so overflow can offer references + retrieval_stats: dict = {} all_chunks = await self._retriever.retrieve( query, top_k=top_k * 2, confidence_threshold=self._config.retrieval_confidence_threshold, marginal_ratio=self._config.retrieval_marginal_ratio, max_tokens=None, + stats_out=retrieval_stats, ) all_chunks = await self._compressor.compress(all_chunks, self._config.compression_level) @@ -1027,7 +1029,10 @@ async def _handle_context_search(self, args): body = _format_results_with_overflow(inline_chunks, overflow_chunks) body = self._apply_output_compression(body) - if len(all_chunks) < self._config.retrieval_top_k: + # Only note omissions the retriever actually made (threshold or + # marginal stop) — chunk-count heuristics false-positive on every + # query because compression/config filtering also shrink the list. + if retrieval_stats.get("dropped_low_value", 0) > 0: note = ( "[note: lower-confidence results omitted — raise top_k or " "lower retrieval.confidence_threshold to include them]" diff --git a/src/context_engine/retrieval/retriever.py b/src/context_engine/retrieval/retriever.py index 343385d..15e1e21 100644 --- a/src/context_engine/retrieval/retriever.py +++ b/src/context_engine/retrieval/retriever.py @@ -41,6 +41,7 @@ async def retrieve( confidence_threshold: float = 0.0, max_tokens: int | None = None, marginal_ratio: float = 0.0, + stats_out: dict | None = None, ) -> list[Chunk]: parsed = self._parser.parse(query) query_embedding = self._embedder.embed_query(query) @@ -150,11 +151,18 @@ async def retrieve( scored.sort(key=lambda x: x[1], reverse=True) scored = self._dedupe_overlaps(scored) + # Candidate count after dedup, before any confidence filtering — + # the baseline for the dropped_low_value accounting below. + candidates = len(scored) + # Confidence cutoff with a top-1 guarantee: an over-tight threshold # must never turn a matching query into an empty result. filtered = [(c, s) for c, s in scored if s >= confidence_threshold] if not filtered and scored: filtered = scored[:1] + # Candidates excluded specifically by the threshold (a top-1 + # guarantee survivor was not dropped). + dropped_low_value = candidates - len(filtered) scored = filtered # File diversity: cap chunks per file so one large file doesn't @@ -163,8 +171,12 @@ async def retrieve( top_score = scored[0][1] if scored else 0.0 file_counts: dict[str, int] = {} diverse: list[Chunk] = [] - for chunk, score in scored: + for idx, (chunk, score) in enumerate(scored): if diverse and marginal_ratio > 0 and score < marginal_ratio * top_score: + # Everything from here on scores below the marginal cutoff + # (list is sorted) — count them as low-value drops. Chunks + # excluded only by the per-file cap or top_k are NOT counted. + dropped_low_value += len(scored) - idx break count = file_counts.get(chunk.file_path, 0) if count < _MAX_CHUNKS_PER_FILE: @@ -205,6 +217,10 @@ async def retrieve( log.debug("Graph expansion skipped: %s", exc) if max_tokens is None: + if stats_out is not None: + stats_out["candidates"] = candidates + stats_out["selected"] = len(ranked) + stats_out["dropped_low_value"] = dropped_low_value return ranked packed: list[Chunk] = [] @@ -219,6 +235,10 @@ async def retrieve( if compressed_tokens <= budget: packed.append(chunk) budget -= compressed_tokens + if stats_out is not None: + stats_out["candidates"] = candidates + stats_out["selected"] = len(packed) + stats_out["dropped_low_value"] = dropped_low_value return packed @staticmethod diff --git a/tests/integration/test_mcp_server.py b/tests/integration/test_mcp_server.py index 8072503..a31c39f 100644 --- a/tests/integration/test_mcp_server.py +++ b/tests/integration/test_mcp_server.py @@ -143,9 +143,9 @@ async def test_index_status_with_tracked_stats(tmp_path): assert "60%" in text -@pytest.mark.asyncio -async def test_context_search_appends_omitted_note_when_results_below_top_k(tmp_path): - """When retrieve() returns fewer results than retrieval_top_k, a note is appended.""" +def _make_search_server(tmp_path, dropped_low_value): + """Server wired for _handle_context_search with a stub retriever whose + stats_out reports the given dropped_low_value.""" from unittest.mock import AsyncMock from context_engine.models import Chunk, ChunkType @@ -158,8 +158,17 @@ async def test_context_search_appends_omitted_note_when_results_below_top_k(tmp_ ) stub_chunk.confidence_score = 0.8 + async def fake_retrieve(query, top_k=10, confidence_threshold=0.0, + marginal_ratio=0.0, max_tokens=None, + stats_out=None): + if stats_out is not None: + stats_out["candidates"] = 1 + dropped_low_value + stats_out["selected"] = 1 + stats_out["dropped_low_value"] = dropped_low_value + return [stub_chunk] + server._retriever = MagicMock() - server._retriever.retrieve = AsyncMock(return_value=[stub_chunk]) + server._retriever.retrieve = fake_retrieve server._compressor = MagicMock() server._compressor.compress = AsyncMock(return_value=[stub_chunk]) server._session_capture = MagicMock() @@ -169,7 +178,6 @@ async def test_context_search_appends_omitted_note_when_results_below_top_k(tmp_ server._append_audit_log = MagicMock() server._ensure_indexed = AsyncMock(return_value=True) - # retrieval_top_k=5, but only 1 chunk returned → note should appear server._config.retrieval_top_k = 5 server._config.retrieval_confidence_threshold = 0.99 server._config.retrieval_marginal_ratio = 0.5 @@ -177,7 +185,23 @@ async def test_context_search_appends_omitted_note_when_results_below_top_k(tmp_ server._output_level = "off" server._project_name = "test-project" server._session_id = "test-session" + return server + +@pytest.mark.asyncio +async def test_context_search_appends_omitted_note_when_drops_reported(tmp_path): + """Note appears only when the retriever reports threshold/marginal drops.""" + server = _make_search_server(tmp_path, dropped_low_value=2) result = await server._handle_context_search({"query": "find something", "top_k": 5}) text = result[0].text assert "lower-confidence results omitted" in text + + +@pytest.mark.asyncio +async def test_context_search_no_note_when_nothing_dropped(tmp_path): + """Fewer chunks than retrieval_top_k with zero drops must NOT trigger the + note — that was the false-positive the count heuristic produced.""" + server = _make_search_server(tmp_path, dropped_low_value=0) + result = await server._handle_context_search({"query": "find something", "top_k": 5}) + text = result[0].text + assert "lower-confidence results omitted" not in text diff --git a/tests/retrieval/test_retriever.py b/tests/retrieval/test_retriever.py index f9ac10e..4eea35f 100644 --- a/tests/retrieval/test_retriever.py +++ b/tests/retrieval/test_retriever.py @@ -242,3 +242,45 @@ async def test_top1_guarantee_when_threshold_filters_everything(retriever_factor "query", top_k=10, confidence_threshold=0.99 ) assert len(results) == 1 # best candidate survives an over-tight threshold + + +# --------------------------------------------------------------------------- +# stats_out: truthful accounting of what the threshold/marginal stop dropped +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_stats_out_counts_threshold_drops(retriever_factory): + retriever, chunks = retriever_factory(distances=[1.6, 1.8]) + stats: dict = {} + results = await retriever.retrieve( + "query", top_k=10, confidence_threshold=0.99, stats_out=stats, + ) + assert stats["candidates"] == 2 + assert stats["selected"] == len(results) == 1 + # One candidate dropped by threshold; the other survives via top-1 guarantee. + assert stats["dropped_low_value"] == 1 + + +@pytest.mark.asyncio +async def test_stats_out_counts_marginal_stop_drops(retriever_factory): + retriever, chunks = retriever_factory(distances=[0.1, 0.3, 1.8]) + stats: dict = {} + results = await retriever.retrieve( + "query", top_k=10, marginal_ratio=0.5, stats_out=stats, + ) + assert stats["candidates"] == 3 + assert stats["selected"] == len(results) == 2 + assert stats["dropped_low_value"] == 1 # tail chunk cut by marginal stop + + +@pytest.mark.asyncio +async def test_stats_out_zero_drops_when_nothing_fires(retriever_factory): + retriever, chunks = retriever_factory(distances=[0.1, 0.3, 1.8]) + stats: dict = {} + results = await retriever.retrieve( + "query", top_k=10, marginal_ratio=0.0, confidence_threshold=0.0, + stats_out=stats, + ) + assert stats["candidates"] == 3 + assert stats["selected"] == len(results) == 3 + assert stats["dropped_low_value"] == 0 From 657df0bb576d77221dfb56ebb01a59d571490242 Mon Sep 17 00:00:00 2001 From: Fazle Elahee Date: Fri, 3 Jul 2026 22:03:41 +0100 Subject: [PATCH 11/14] feat(benchmarks): A/B mode for retrieval precision; tune Phase 1 defaults from evidence MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds --ab flag to run_benchmark.py: indexes the repo once, then runs each query twice (baseline retrieve() vs tuned with threshold+marginal_ratio) and reports per-query token deltas plus aggregate reduction%. Evidence from 6 A/B runs on this repo shows marginal_ratio=0.75 achieves 26.1% token-served reduction with hit rate unchanged (4/8→4/8), clearing the ≥25% gate. confidence_threshold had no effect (all scores 0.70–0.95). Updates retrieval_marginal_ratio default: 0.5 → 0.75. Results committed as benchmarks/results/cce-phase1-ab.{json,md}. --- benchmarks/results/cce-phase1-ab.json | 116 ++++++++++++++++++++++++++ benchmarks/results/cce-phase1-ab.md | 61 ++++++++++++++ benchmarks/run_benchmark.py | 86 +++++++++++++++++++ src/context_engine/config.py | 4 +- tests/test_config.py | 2 +- 5 files changed, 267 insertions(+), 2 deletions(-) create mode 100644 benchmarks/results/cce-phase1-ab.json create mode 100644 benchmarks/results/cce-phase1-ab.md diff --git a/benchmarks/results/cce-phase1-ab.json b/benchmarks/results/cce-phase1-ab.json new file mode 100644 index 0000000..bc34f6a --- /dev/null +++ b/benchmarks/results/cce-phase1-ab.json @@ -0,0 +1,116 @@ +{ + "threshold": 0.2, + "marginal_ratio": 0.75, + "base_tokens": 32138, + "tuned_tokens": 23741, + "reduction_pct": 26.1, + "base_hits": 4, + "tuned_hits": 4, + "judged": 8, + "rows": [ + { + "query": "How does the chunker split code into chunks?", + "base": { + "tokens": 2533, + "chunks": 10, + "hit": true + }, + "tuned": { + "tokens": 2319, + "chunks": 7, + "hit": true + } + }, + { + "query": "vector search implementation", + "base": { + "tokens": 4600, + "chunks": 10, + "hit": false + }, + "tuned": { + "tokens": 3233, + "chunks": 8, + "hit": false + } + }, + { + "query": "confidence scoring formula", + "base": { + "tokens": 3359, + "chunks": 10, + "hit": false + }, + "tuned": { + "tokens": 3359, + "chunks": 10, + "hit": false + } + }, + { + "query": "MCP server tools", + "base": { + "tokens": 4221, + "chunks": 10, + "hit": true + }, + "tuned": { + "tokens": 4221, + "chunks": 10, + "hit": true + } + }, + { + "query": "how does indexing pipeline work", + "base": { + "tokens": 6331, + "chunks": 10, + "hit": true + }, + "tuned": { + "tokens": 6008, + "chunks": 9, + "hit": true + } + }, + { + "query": "FTS5 full text search", + "base": { + "tokens": 4321, + "chunks": 10, + "hit": false + }, + "tuned": { + "tokens": 736, + "chunks": 2, + "hit": false + } + }, + { + "query": "graph neighbors query", + "base": { + "tokens": 2059, + "chunks": 10, + "hit": false + }, + "tuned": { + "tokens": 1123, + "chunks": 6, + "hit": false + } + }, + { + "query": "git commit history indexing", + "base": { + "tokens": 4714, + "chunks": 10, + "hit": true + }, + "tuned": { + "tokens": 2742, + "chunks": 6, + "hit": true + } + } + ] +} diff --git a/benchmarks/results/cce-phase1-ab.md b/benchmarks/results/cce-phase1-ab.md new file mode 100644 index 0000000..349956d --- /dev/null +++ b/benchmarks/results/cce-phase1-ab.md @@ -0,0 +1,61 @@ +# CCE Phase 1 A/B Benchmark: Retrieval Precision Tuning + +**Date:** 2026-07-03 +**Repo:** code-context-engine (this repo, indexed from `.`) +**Branch:** fix/review-findings-2026-07-03 +**Index:** 4,162 chunks from 285 files +**Queries:** 8 (from `benchmarks/sample_queries.json`) + +## Winning Defaults + +| Parameter | Old Default | New Default | Effect | +|-----------|------------|-------------|--------| +| `retrieval_confidence_threshold` | 0.2 | **0.2** (unchanged) | No effect — all scores ≥ 0.70 on this corpus | +| `retrieval_marginal_ratio` | 0.5 | **0.75** | Achieved 26.1% token reduction | + +## Aggregate Results (Winning Config) + +| Metric | Baseline | Tuned (threshold=0.2, marginal_ratio=0.75) | +|--------|----------|---------------------------------------------| +| Total tokens served | 32,138 | 23,741 | +| Token reduction | — | **26.1%** ✓ (gate: ≥25%) | +| Hit rate | 4/8 | **4/8** ✓ (gate: ≥ baseline) | + +**Gate verdict: PASS** — 26.1% reduction ≥ 25%, hit rate unchanged. + +## Per-Query Detail (Winning Run) + +| Query | Base tokens | Tuned tokens | Base chunks | Tuned chunks | Hit | +|-------|------------|--------------|-------------|--------------|-----| +| How does the chunker split code into chunks? | 2,533 | 2,319 | 10 | 7 | True → True | +| vector search implementation | 4,600 | 3,233 | 10 | 8 | False → False | +| confidence scoring formula | 3,359 | 3,359 | 10 | 10 | False → False | +| MCP server tools | 4,221 | 4,221 | 10 | 10 | True → True | +| how does indexing pipeline work | 6,331 | 6,008 | 10 | 9 | True → True | +| FTS5 full text search | 4,321 | 736 | 10 | 2 | False → False | +| graph neighbors query | 2,059 | 1,123 | 10 | 6 | False → False | +| git commit history indexing | 4,714 | 2,742 | 10 | 6 | True → True | + +## All Runs (Including Rejected) + +| Run | threshold | marginal_ratio | Tokens: base → tuned | Reduction | Hits | Gate | +|-----|-----------|----------------|----------------------|-----------|------|------| +| 1 | 0.35 | 0.50 | 32,026 → 32,026 | 0.0% | 4/8 → 4/8 | FAIL (reduction < 25%) | +| 2 | 0.35 | 0.60 | 32,138 → 32,138 | 0.0% | 4/8 → 4/8 | FAIL (reduction < 25%) | +| 3 | 0.35 | 0.70 | 32,138 → 30,482 | 5.2% | 4/8 → 4/8 | FAIL (reduction < 25%) | +| 4 | 0.35 | 0.75 | 32,138 → 23,741 | 26.1% | 4/8 → 4/8 | PASS | +| 5 | 0.35 | 0.80 | 32,138 → 20,061 | 37.6% | 4/8 → 4/8 | PASS (higher reduction, but 0.75 preferred as minimum-passing) | +| **6 (winner)** | **0.2** | **0.75** | **32,138 → 23,741** | **26.1%** | **4/8 → 4/8** | **PASS** | + +### Notes on Tuning + +- The `confidence_threshold` parameter (0.35, then 0.2) had **no effect** on this corpus: all retrieved chunk confidence scores fell in the 0.70–0.95 range, well above any tested threshold. The ladder's instruction to use `--threshold 0.25` or `--threshold 0.2` when hit rate drops was not needed — hits were stable throughout. +- Reduction came **entirely from `marginal_ratio`**: this parameter stops adding chunks whose score falls below `ratio × top_score`. With scores clustered at the top (spread ≈ 0.25), a ratio of 0.75 is required to create meaningful cuts. +- `marginal_ratio=0.75` is the minimum value that clears the 25% gate. `0.80` provides more reduction (37.6%) with the same hit rate, but `0.75` was chosen as the conservative minimum-passing value. +- The brief's tuning ladder (tries 0.5, 0.6) did not cover the actual operating range of this corpus (needed 0.75+). Extended ladder runs (0.70, 0.75, 0.80) were added to find the gate-passing threshold. + +## Config Changes Applied + +- `src/context_engine/config.py`: `retrieval_marginal_ratio` updated `0.5 → 0.75` +- `tests/test_config.py`: `test_marginal_ratio_default` updated to assert `0.75` +- `retrieval_confidence_threshold` left at `0.2` (unchanged — threshold had no effect) diff --git a/benchmarks/run_benchmark.py b/benchmarks/run_benchmark.py index 1a4a06c..ad612b3 100644 --- a/benchmarks/run_benchmark.py +++ b/benchmarks/run_benchmark.py @@ -42,6 +42,7 @@ from context_engine.memory.grammar import compress_with_counts as grammar_compress from context_engine.retrieval.retriever import HybridRetriever from context_engine.storage.local_backend import LocalBackend +from context_engine.utils import project_storage_dir _CHARS_PER_TOKEN = 4 @@ -420,6 +421,73 @@ def format_markdown(results: dict) -> str: return "\n".join(lines) + "\n" +async def run_ab( + project_dir: Path, + queries: list[dict], + storage_dir: Path, + threshold: float, + marginal_ratio: float, +) -> dict: + """Index once, then run every query with baseline vs tuned retrieval + parameters against the same index. No stacking with compression — + this isolates the retrieval-precision change.""" + config = Config() + config.storage_path = str(storage_dir) + + print("Indexing project once for A/B...") + idx = await run_indexing(config, project_dir, full=True) + print(f" {idx.total_chunks} chunks from {len(idx.indexed_files)} files") + + storage_base = project_storage_dir(config, project_dir) + backend = LocalBackend(base_path=str(storage_base)) + embedder = Embedder(model_name=config.embedding_model) + retriever = HybridRetriever(backend=backend, embedder=embedder) + + rows = [] + for q in queries: + base = await retriever.retrieve(q["query"], top_k=10) + tuned = await retriever.retrieve( + q["query"], top_k=10, + confidence_threshold=threshold, + marginal_ratio=marginal_ratio, + ) + expected = set(q.get("expected_files", [])) + + def _measure(chunks): + files = {c.file_path for c in chunks} + return { + "tokens": sum(_count_tokens(c.content) for c in chunks), + "chunks": len(chunks), + "hit": bool(files & expected) if expected else None, + } + + rows.append({"query": q["query"], + "base": _measure(base), "tuned": _measure(tuned)}) + b, t = rows[-1]["base"], rows[-1]["tuned"] + print(f" {q['query'][:45]:<45} tokens {b['tokens']:>6} → {t['tokens']:>6} " + f"chunks {b['chunks']:>2} → {t['chunks']:>2} " + f"hit {b['hit']} → {t['hit']}") + + def _agg(side): + tok = sum(r[side]["tokens"] for r in rows) + hits = sum(1 for r in rows if r[side]["hit"]) + judged = sum(1 for r in rows if r[side]["hit"] is not None) + return tok, hits, judged + + base_tok, base_hits, judged = _agg("base") + tuned_tok, tuned_hits, _ = _agg("tuned") + reduction = (1 - tuned_tok / base_tok) * 100 if base_tok else 0.0 + print(f"\nTokens served: {base_tok:,} → {tuned_tok:,} ({reduction:.1f}% reduction)") + print(f"Hit rate: {base_hits}/{judged} → {tuned_hits}/{judged}") + return { + "threshold": threshold, "marginal_ratio": marginal_ratio, + "base_tokens": base_tok, "tuned_tokens": tuned_tok, + "reduction_pct": round(reduction, 1), + "base_hits": base_hits, "tuned_hits": tuned_hits, + "judged": judged, "rows": rows, + } + + def main(): parser = argparse.ArgumentParser(description="CCE Benchmark Suite") parser.add_argument("--repo", help="Git repo URL to clone and benchmark") @@ -432,6 +500,13 @@ def main(): parser.add_argument("--queries", help="Path to queries JSON file") parser.add_argument("--output", help="Output path for markdown report") parser.add_argument("--json-output", help="Output path for raw JSON results") + parser.add_argument("--ab", action="store_true", + help="Run each query twice: baseline retrieval vs tuned " + "(threshold + marginal ratio) and print the delta") + parser.add_argument("--threshold", type=float, default=0.35, + help="Tuned confidence_threshold for --ab (default 0.35)") + parser.add_argument("--marginal-ratio", type=float, default=0.5, + help="Tuned marginal_ratio for --ab (default 0.5)") args = parser.parse_args() # Determine project dir and queries @@ -476,6 +551,17 @@ def main(): storage_dir = Path(tempfile.mkdtemp(prefix="cce-bench-storage-")) try: + if args.ab: + results = asyncio.run(run_ab( + project_dir, queries, storage_dir, + args.threshold, args.marginal_ratio, + )) + if args.json_output: + out_path = Path(args.json_output) + out_path.parent.mkdir(parents=True, exist_ok=True) + out_path.write_text(json.dumps(results, indent=2) + "\n") + return + results = asyncio.run(run_benchmark(project_dir, queries, storage_dir)) if args.repo: results["repo_url"] = args.repo diff --git a/src/context_engine/config.py b/src/context_engine/config.py index 3d42532..a4fc52b 100644 --- a/src/context_engine/config.py +++ b/src/context_engine/config.py @@ -70,7 +70,9 @@ class Config: retrieval_top_k: int = 20 # Stop adding result chunks once a chunk's score falls below this # fraction of the top score. 0 disables (always fill to top_k). - retrieval_marginal_ratio: float = 0.5 + # Tuned to 0.75 by Phase 1 A/B benchmark (2026-07-03): achieves ≥25% + # token-served reduction with no hit-rate loss. + retrieval_marginal_ratio: float = 0.75 bootstrap_max_tokens: int = 10000 # Indexer diff --git a/tests/test_config.py b/tests/test_config.py index 866993a..6d808a3 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -97,4 +97,4 @@ def test_marginal_ratio_config_mapping(tmp_path): def test_marginal_ratio_default(): from context_engine.config import Config - assert Config().retrieval_marginal_ratio == 0.5 + assert Config().retrieval_marginal_ratio == 0.75 From 316768ed68313cba28e5a8f8c930e4b246e5d0a7 Mon Sep 17 00:00:00 2001 From: Fazle Elahee Date: Fri, 3 Jul 2026 22:07:53 +0100 Subject: [PATCH 12/14] docs(benchmarks): disclose cliff shape and sample-size limits in Phase 1 A/B evidence --- benchmarks/results/cce-phase1-ab.md | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/benchmarks/results/cce-phase1-ab.md b/benchmarks/results/cce-phase1-ab.md index 349956d..d1dc162 100644 --- a/benchmarks/results/cce-phase1-ab.md +++ b/benchmarks/results/cce-phase1-ab.md @@ -54,6 +54,23 @@ - `marginal_ratio=0.75` is the minimum value that clears the 25% gate. `0.80` provides more reduction (37.6%) with the same hit rate, but `0.75` was chosen as the conservative minimum-passing value. - The brief's tuning ladder (tries 0.5, 0.6) did not cover the actual operating range of this corpus (needed 0.75+). Extended ladder runs (0.70, 0.75, 0.80) were added to find the gate-passing threshold. +### Limitations of this evidence + +- **The 0.75 default sits on a cliff, not a plateau.** The reduction curve is + 0.70 → 5.2%, 0.75 → 26.1%, 0.80 → 37.6% — a 21-point jump across one step. + Because the marginal stop cuts relative to the corpus's score distribution, + a project whose score floor sits slightly higher (e.g. 0.77 instead of + 0.70) would see little or no reduction at 0.75, and one with a wider spread + could see much more. Treat 0.75 as a starting default, not stable guidance; + it is tunable per project via `retrieval.marginal_ratio`. +- **Sample size: 8 queries, one corpus (this repo).** The gate verdict is + real but narrow. Re-validate on at least one external corpus (the existing + `fastapi`/`chi`/`fiber` query sets) before citing these numbers in docs or + README. +- Run 1's baseline total (32,026) differs from later runs (32,138) by 112 + tokens — the index was rebuilt between runs 1 and 2. The winning run's + numbers are internally consistent (per-row sums verified). + ## Config Changes Applied - `src/context_engine/config.py`: `retrieval_marginal_ratio` updated `0.5 → 0.75` From ce918e508f539e5c9b91054e24852f81c14715be Mon Sep 17 00:00:00 2001 From: Fazle Elahee Date: Fri, 3 Jul 2026 22:15:44 +0100 Subject: [PATCH 13/14] fix: point omitted-results note at marginal_ratio; document new retrieval config --- docs-src/src/content/docs/configuration.md | 7 +++++-- src/context_engine/integration/mcp_server.py | 4 ++-- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/docs-src/src/content/docs/configuration.md b/docs-src/src/content/docs/configuration.md index 80b9978..c6e0a81 100644 --- a/docs-src/src/content/docs/configuration.md +++ b/docs-src/src/content/docs/configuration.md @@ -33,7 +33,8 @@ indexer: retrieval: top_k: 20 # Maximum chunks returned per query - confidence_threshold: 0.5 # Minimum score to include a result (0.0 to 1.0) + confidence_threshold: 0.2 # Minimum score to include a result (0.0 to 1.0) + marginal_ratio: 0.75 # Stop adding chunks when score falls below this fraction of top embedding: model: BAAI/bge-small-en-v1.5 # Embedding model (fastembed-compatible) @@ -93,7 +94,9 @@ The default `BAAI/bge-small-en-v1.5` is recommended for most use cases. It balan **`top_k`** controls how many chunks the retriever returns per query. Higher values surface more context but cost more tokens. Default: 20. -**`confidence_threshold`** sets the minimum score to include a result. Range 0.0 to 1.0. Lower values return more results; higher values return only strong matches. Default: 0.5. +**`confidence_threshold`** sets the minimum score to include a result. Range 0.0 to 1.0. Lower values return more results; higher values return only strong matches. Default: 0.2. + +**`marginal_ratio`** stops adding result chunks once a chunk's confidence falls below this fraction of the top result's score. Range 0.0 to 1.0; `0` disables this behavior. The effective reduction depends on your corpus's score distribution, so tune per project. See `benchmarks/results/cce-phase1-ab.md` for tuning evidence. Default: 0.75. At runtime, the agent can pass `top_k` and `max_tokens` directly to `context_search`: diff --git a/src/context_engine/integration/mcp_server.py b/src/context_engine/integration/mcp_server.py index 7043f82..b1a561b 100644 --- a/src/context_engine/integration/mcp_server.py +++ b/src/context_engine/integration/mcp_server.py @@ -1034,8 +1034,8 @@ async def _handle_context_search(self, args): # query because compression/config filtering also shrink the list. if retrieval_stats.get("dropped_low_value", 0) > 0: note = ( - "[note: lower-confidence results omitted — raise top_k or " - "lower retrieval.confidence_threshold to include them]" + "[note: lower-confidence results omitted — lower " + "retrieval.marginal_ratio or retrieval.confidence_threshold to include them]" ) body = body + "\n" + note self._record(raw_tokens, served_tokens, full_file_tokens) From e301dfa737bde9233fd592cb05156f72192b4387 Mon Sep 17 00:00:00 2001 From: Fazle Elahee Date: Thu, 9 Jul 2026 22:11:29 +0100 Subject: [PATCH 14/14] fix(indexer): repopulate index after cosine-metric migration The cosine rebuild in VectorStore._ensure_tables wipes on-disk chunks but left manifest.json claiming every file was still indexed, so an incremental reindex (the default for cce index and the watcher) skipped every unchanged file and left context_search returning nothing until cce index --full. VectorStore now flags metric_rebuilt; run_indexing clears the manifest when set, forcing the scan to re-ingest. Mirrors the existing dim-migration guard. Regression test confirms an incremental reindex repopulates post-migration. --- src/context_engine/indexer/pipeline.py | 11 +++ src/context_engine/storage/vector_store.py | 7 ++ .../indexer/test_pipeline_metric_migration.py | 80 +++++++++++++++++++ 3 files changed, 98 insertions(+) create mode 100644 tests/indexer/test_pipeline_metric_migration.py diff --git a/src/context_engine/indexer/pipeline.py b/src/context_engine/indexer/pipeline.py index 678e0b2..5a78e31 100644 --- a/src/context_engine/indexer/pipeline.py +++ b/src/context_engine/indexer/pipeline.py @@ -364,6 +364,17 @@ async def _run_indexing_locked( backend = LocalBackend(base_path=str(storage_base)) chunker = Chunker() manifest = Manifest(manifest_path=storage_base / "manifest.json") + # If constructing the backend just wiped a legacy L2 vector table for the + # cosine-metric rebuild, the on-disk chunks are gone but the manifest still + # records every file as indexed. Clear it so the scan below re-ingests them + # — otherwise an incremental run (the default for `cce index` and the + # watcher) skips every "unchanged" file and the index stays empty until + # someone runs `cce index --full`. Mirrors the dim-migration guard below. + if getattr(getattr(backend, "_vector_store", None), "metric_rebuilt", False): + manifest.clear_entries() + log.info("Vector table rebuilt for cosine metric; cleared manifest to force full reindex.") + if log_fn: + log_fn(" [migration] vector index rebuilt (cosine metric) — reindexing all files") ignore_set = set(config.indexer_ignore) # Load .cceignore once per indexing run. Patterns are evaluated against # paths relative to project_dir; see indexer/ignorefile.py. diff --git a/src/context_engine/storage/vector_store.py b/src/context_engine/storage/vector_store.py index e166f3d..f46bea7 100644 --- a/src/context_engine/storage/vector_store.py +++ b/src/context_engine/storage/vector_store.py @@ -38,6 +38,12 @@ def __init__(self, db_path: str) -> None: self._db_path = db_path self._lock = RLock() self._dim: int | None = None + # Set True by _ensure_tables when a legacy L2 vector table is wiped for + # the cosine-metric rebuild. The pipeline reads this to clear the + # manifest and force a full reindex — the on-disk chunks are gone but + # the manifest still claims they are indexed, so an incremental run + # would otherwise skip every "unchanged" file and leave the index empty. + self.metric_rebuilt = False os.makedirs(db_path, exist_ok=True) self._db_file = os.path.join(db_path, "vectors.db") self._conn = self._connect() @@ -124,6 +130,7 @@ def _ensure_tables(self) -> None: self._conn.execute("DROP TABLE IF EXISTS chunks_vec") self._conn.execute("DELETE FROM chunks") self._conn.execute("DELETE FROM chunk_compressions") + self.metric_rebuilt = True else: # Table exists — read dim from first row r = self._conn.execute( diff --git a/tests/indexer/test_pipeline_metric_migration.py b/tests/indexer/test_pipeline_metric_migration.py new file mode 100644 index 0000000..7c77862 --- /dev/null +++ b/tests/indexer/test_pipeline_metric_migration.py @@ -0,0 +1,80 @@ +"""Cosine-metric migration must not leave an existing index silently empty. + +Before this fix, opening a pre-cosine (L2) vector store wiped the on-disk +chunks for the cosine rebuild but left manifest.json claiming every file was +still indexed. An incremental reindex (the default for `cce index` and the +watcher) then skipped every "unchanged" file, so `context_search` returned +nothing until the user ran `cce index --full`. + +The pipeline now detects the rebuild (VectorStore.metric_rebuilt) and clears +the manifest so the next incremental run repopulates the index. +""" +from __future__ import annotations + +from pathlib import Path + +import pytest + +from context_engine.config import load_config +from context_engine.indexer.pipeline import run_indexing +from context_engine.storage.local_backend import LocalBackend +from context_engine.utils import project_storage_dir + + +@pytest.fixture +def project(tmp_path): + project_dir = tmp_path / "proj" + project_dir.mkdir() + (project_dir / "app.py").write_text("def handler():\n return 42\n") + storage_base = tmp_path / "storage" + storage_base.mkdir() + config = load_config() + config.storage_path = str(storage_base) + return project_dir, config + + +def _chunk_count(config, project_dir: Path) -> int: + storage = project_storage_dir(config, Path(project_dir)) + return LocalBackend(base_path=str(storage))._vector_store.count() + + +def _downgrade_vec_table_to_l2(config, project_dir: Path) -> None: + """Rewrite chunks_vec without distance_metric=cosine to simulate an index + created by a pre-cosine release.""" + storage = project_storage_dir(config, Path(project_dir)) + vec_db = Path(storage) / "vectors" / "vectors.db" + import sqlite3 + + import sqlite_vec + + conn = sqlite3.connect(str(vec_db)) + conn.enable_load_extension(True) + sqlite_vec.load(conn) + conn.enable_load_extension(False) + dim = conn.execute("SELECT vec_length(embedding) FROM chunks_vec LIMIT 1").fetchone()[0] + conn.execute("DROP TABLE chunks_vec") + conn.execute(f"CREATE VIRTUAL TABLE chunks_vec USING vec0(embedding float[{dim}])") + conn.commit() + conn.close() + + +@pytest.mark.asyncio +async def test_incremental_reindex_repopulates_after_cosine_migration(project): + project_dir, config = project + + # 1. Index normally, then downgrade the vec table to legacy L2 on disk. + result = await run_indexing(config, str(project_dir), full=True) + assert not result.errors, result.errors + assert _chunk_count(config, project_dir) > 0 + _downgrade_vec_table_to_l2(config, project_dir) + + # 2. Incremental reindex (full=False) — the file content is unchanged, so + # without the manifest-clear the wiped index would stay empty. + result = await run_indexing(config, str(project_dir), full=False) + assert not result.errors, result.errors + + # 3. The index must be repopulated, not empty. + assert _chunk_count(config, project_dir) > 0, ( + "cosine migration wiped the index and the incremental reindex did not " + "repopulate it — the manifest-clear guard failed" + )