Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 31 additions & 5 deletions src/armillary/revive_enhanced.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,10 @@
from pathlib import Path

from armillary.cache import Cache
from armillary.exclude_service import is_excluded
from armillary.models import Status
from armillary.revive_service import revive_show
from armillary.status_override import get_override
from armillary.steal_service import steal


Expand All @@ -36,19 +39,28 @@ def generate_enhanced_brief(
if not query:
return brief

# Overfetch and filter out hits from the project being revived — the
# tool promises quotes from OTHER repos. Without this, a project that
# has indexed itself can crowd out real cross-repo matches.
# Overfetch and filter so we can drop:
# 1. Hits from the project being revived (the tool promises quotes
# from OTHER repos).
# 2. Hits from repos the user has excluded or archived via the panel.
# `armillary_steal` keeps those on purpose (user explicitly mining
# their own dead code), but for revive the panel choice should
# apply, matching every other MCP tool's behaviour.
#
# Multiplier sized for worst case: `steal()` caps at 3 hits per repo
# in its overfetch pool, so three noise repos (own + excluded dupe +
# archived dupe) can consume 9 results. `steal_limit * 6` leaves
# `steal_limit` worth of headroom even in that case.
own_repo = _resolve(project_path)
try:
raw = steal(query, limit=steal_limit * 3)
raw = steal(query, limit=steal_limit * 6)
except Exception: # noqa: BLE001 — graceful: enhanced is bonus over vanilla
# Steal can raise on missing/corrupt code_index.db or a SQLite
# build without FTS5. The vanilla brief is still useful on its
# own, so swallow the failure and fall back to brief-only.
return brief

results = [r for r in raw if _resolve(Path(r.block.repo_path)) != own_repo][
results = [r for r in raw if _is_keepable(r.block.repo_path, own_repo)][
:steal_limit
]
if not results:
Expand Down Expand Up @@ -109,6 +121,20 @@ def _resolve(path: Path) -> str:
return str(path)


def _is_keepable(repo_path: str, own_repo: str) -> bool:
"""True if a steal result should appear in STEAL_HITS.

Drops the project being revived (own-repo filter), repos the user
has excluded via the panel, and repos the user archived via status
override. Other MCP tools honour the same panel choices.
"""
if _resolve(Path(repo_path)) == own_repo:
return False
if is_excluded(repo_path):
return False
return get_override(repo_path) is not Status.ARCHIVED


def _project_name(project_path: Path) -> str:
"""Return the cached project name, or the final path part."""
with Cache() as cache:
Expand Down
70 changes: 67 additions & 3 deletions tests/test_revive_enhanced.py
Original file line number Diff line number Diff line change
Expand Up @@ -110,9 +110,12 @@ def test_generate_enhanced_brief_appends_steal_section(
# signal that yields cross-repo matches without precision-extreme
# AND-collapse on commit subjects.
assert mocks["steal"].call_args.args[0] == "my_project"
# Helper overfetches (limit * 3) so it can drop hits from the same repo
# before slicing back to the requested top-N.
assert mocks["steal"].call_args.kwargs == {"limit": 9}
# Helper overfetches (limit * 6) so it can drop hits from the same
# repo and from panel-excluded / archived repos before slicing back
# to the requested top-N. The 6× multiplier is sized for the worst
# case where three noise repos consume 9 results (steal caps at 3
# hits per repo in its overfetch pool).
assert mocks["steal"].call_args.kwargs == {"limit": 18}


def test_generate_enhanced_brief_skips_section_when_no_hits(
Expand Down Expand Up @@ -236,6 +239,67 @@ def test_generate_enhanced_brief_renders_relative_block_path(
assert "/repos/invoicer/src/price.py" not in out


def test_generate_enhanced_brief_drops_excluded_repos(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Repos the user has excluded via the panel must not appear in
STEAL_HITS, even though `steal()` itself returns them. Mirrors how
`armillary_projects` and `armillary_next` already behave.
"""
excluded_block = _block(
repo_path="/repos/dead-fork", path="/repos/dead-fork/src/x.py"
)
keep_block = _block(
repo_path="/repos/active-sibling", path="/repos/active-sibling/src/y.py"
)
_patch_helper_dependencies(
monkeypatch,
steal_results=[
_result(excluded_block, project_name="dead-fork"),
_result(keep_block, project_name="active-sibling"),
],
cache_project_name="me",
)

def _is_excluded(path: str) -> bool:
return path == "/repos/dead-fork"

monkeypatch.setattr("armillary.revive_enhanced.is_excluded", _is_excluded)

out = generate_enhanced_brief(Path("/repos/me"))

assert "active-sibling/src/y.py" in out
assert "dead-fork" not in out


def test_generate_enhanced_brief_drops_archived_repos(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Repos the user archived via status override must also be dropped."""
from armillary.models import Status

archived_block = _block(repo_path="/repos/old", path="/repos/old/src/x.py")
keep_block = _block(repo_path="/repos/active", path="/repos/active/src/y.py")
_patch_helper_dependencies(
monkeypatch,
steal_results=[
_result(archived_block, project_name="old"),
_result(keep_block, project_name="active"),
],
cache_project_name="me",
)

def _override(path: str) -> Status | None:
return Status.ARCHIVED if path == "/repos/old" else None

monkeypatch.setattr("armillary.revive_enhanced.get_override", _override)

out = generate_enhanced_brief(Path("/repos/me"))

assert "active/src/y.py" in out
assert "old/src/x.py" not in out


def test_generate_enhanced_brief_handles_repo_prefix_collision(
monkeypatch: pytest.MonkeyPatch,
) -> None:
Expand Down
Loading