attocode-code-intel 0.1.17__tar.gz → 0.1.19__tar.gz

{attocode_code_intel-0.1.17 → attocode_code_intel-0.1.19}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: attocode-code-intel
-Version: 0.1.17
+Version: 0.1.19
 Summary: Local code search and navigation for coding agents and people
 Project-URL: Homepage, https://github.com/eren23/attocode
 Project-URL: Repository, https://github.com/eren23/attocode
@@ -139,6 +139,10 @@ To use the latest source, install from a clone of this repository: `uv tool inst
 
 ## What your agent gets
 
+Version 0.1.19 restores source text into spare space after fitting a complete
+compact search page. Existing hits, visible text, source lines and continuations
+remain protected within the same response budget; ranking and retrieval are unchanged.
+
 Version 0.1.7 also provides `changes` and `trace` commands: investigate a diff,
 follow source-backed caller paths, export JSON/Markdown, and check or refresh saved
 evidence. MCP equivalents `change_context` and `trace_flow` are local full-profile
@@ -151,7 +155,9 @@ tests, such as `_test.dart` and Perl `.t` scripts and `t/` modules. Version
 0.1.12 makes `doctor --client claude` read the local-scope entry that
 `claude mcp add` writes to `~/.claude.json`. Version 0.1.13 fills compact search
 pages within their budget, offers bounded deeper searches, and repairs stale
-navigation indexes after successful restart parsing. Version 0.1.17 lets a later explicit inspection retry a failed selected-file parse
+navigation indexes after successful restart parsing. Version 0.1.18 restores the ready phase after a failed snapshot load is followed by a complete fallback rebuild. Incomplete recovery stays guarded.
+
+Version 0.1.17 lets a later explicit inspection retry a failed selected-file parse
 after a one-second cooldown, preserving immediate-repeat and source guards.
 Version 0.1.16 stabilizes tied symbol pages and reduces the Docker image while
 preserving runtime contracts. It adds bounded online maintenance evidence; the

{attocode_code_intel-0.1.17 → attocode_code_intel-0.1.19}/README.md

@@ -26,6 +26,10 @@ To use the latest source, install from a clone of this repository: `uv tool inst
 
 ## What your agent gets
 
+Version 0.1.19 restores source text into spare space after fitting a complete
+compact search page. Existing hits, visible text, source lines and continuations
+remain protected within the same response budget; ranking and retrieval are unchanged.
+
 Version 0.1.7 also provides `changes` and `trace` commands: investigate a diff,
 follow source-backed caller paths, export JSON/Markdown, and check or refresh saved
 evidence. MCP equivalents `change_context` and `trace_flow` are local full-profile
@@ -38,7 +42,9 @@ tests, such as `_test.dart` and Perl `.t` scripts and `t/` modules. Version
 0.1.12 makes `doctor --client claude` read the local-scope entry that
 `claude mcp add` writes to `~/.claude.json`. Version 0.1.13 fills compact search
 pages within their budget, offers bounded deeper searches, and repairs stale
-navigation indexes after successful restart parsing. Version 0.1.17 lets a later explicit inspection retry a failed selected-file parse
+navigation indexes after successful restart parsing. Version 0.1.18 restores the ready phase after a failed snapshot load is followed by a complete fallback rebuild. Incomplete recovery stays guarded.
+
+Version 0.1.17 lets a later explicit inspection retry a failed selected-file parse
 after a one-second cooldown, preserving immediate-repeat and source guards.
 Version 0.1.16 stabilizes tied symbol pages and reduces the Docker image while
 preserving runtime contracts. It adds bounded online maintenance evidence; the

{attocode_code_intel-0.1.17 → attocode_code_intel-0.1.19}/pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "attocode-code-intel"
-version = "0.1.17"
+version = "0.1.19"
 description = "Local code search and navigation for coding agents and people"
 readme = "README.md"
 requires-python = ">=3.12"

{attocode_code_intel-0.1.17 → attocode_code_intel-0.1.19}/src/attocode_intel/__init__.py

@@ -49,4 +49,4 @@ def __getattr__(name: str) -> Any:
     raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
 
 # The other copy is packages/code-intel/pyproject.toml. A test checks that both match.
-__version__ = "0.1.17"
+__version__ = "0.1.19"

{attocode_code_intel-0.1.17 → attocode_code_intel-0.1.19}/src/attocode_intel/_internal/integrations/context/ast_service.py

@@ -471,6 +471,9 @@ class ASTService:
                 for target in targets:
                     self._index.add_file_dependency(src, target)
             state.dep_graph_files = len(dep_graph.forward)
+        if is_full_parse and state.parsed_files == state.reference_indexed_files == source_count:
+            # A failed restore can reset the phase before the full rebuild.
+            state.phase = "ready"
         self._initialized = True
         return state
 

{attocode_code_intel-0.1.17 → attocode_code_intel-0.1.19}/src/attocode_intel/output.py

@@ -140,6 +140,73 @@ def _fill_search_prefix(metadata, data, rows, total, max_tokens, *, upstream_tru
     return compact_result(metadata, data)
 
 
+
+_SEARCH_SNIPPET_PROBES = 12
+_SNIPPET_FIELDS = ("snippet", "snippet_start_line", "snippet_end_line", "snippet_truncated")
+
+
+def _restore_search_snippets(result, originals, max_tokens, *, upstream_truncated=False):
+    """Spend spare space without replacing any evidence on the fitted page.
+
+    BPE counts and anchored windows need not be monotone. These are bounded
+    probes, not a maximality search; every accepted page is checked in full.
+    """
+    page = json.loads(result.content[0].text)
+    remaining = _SEARCH_SNIPPET_PROBES
+    for index, original in enumerate(originals[:len(page["data"]["results"])]):
+        if not remaining:
+            break
+        delivered = page["data"]["results"][index]
+        if (not all(key in original and key in delivered for key in _SNIPPET_FIELDS)
+                or {k: v for k, v in original.items() if k not in _SNIPPET_FIELDS}
+                != {k: v for k, v in delivered.items() if k not in _SNIPPET_FIELDS}):
+            continue
+        body = original["snippet"]
+        if (not isinstance(body, str) or not isinstance(delivered["snippet"], str)
+                or len(body) <= len(delivered["snippet"])):
+            continue
+
+        def probe(window, *, position=index):
+            nonlocal remaining, page, result
+            remaining -= 1
+            current = page["data"]["results"][position]
+            if (len(window["snippet"]) <= len(current["snippet"])
+                    or current["snippet"] not in window["snippet"]):
+                return "skip"
+            old_first, old_last = current["snippet_start_line"], current["snippet_end_line"]
+            first, last = window["snippet_start_line"], window["snippet_end_line"]
+            if old_first is None and old_last is None:
+                if first is not None or last is not None:
+                    return "skip"
+            elif (not all(type(n) is int for n in (old_first, old_last, first, last))
+                  or not 1 <= first <= old_first <= old_last <= last):
+                return "skip"
+            trial = deepcopy(page)
+            trial["data"]["results"][position].update(window)
+            _search_truncation(trial["metadata"], trial["data"], upstream_truncated)
+            candidate = compact_result(trial["metadata"], trial["data"])
+            if response_tokens(candidate) > max_tokens:
+                return "over"
+            page, result = trial, candidate
+            return "fit"
+
+        # The original window is the only source of additional text. Check it
+        # first, then retain the best tested larger anchored window that fits.
+        if probe({key: original[key] for key in _SNIPPET_FIELDS}) == "fit":
+            continue
+        lo, hi = len(delivered["snippet"]) + 1, len(body) - 1
+        while lo <= hi and remaining:
+            cap = (lo + hi) // 2
+            window = snippet_window(body, page["data"].get("query", ""), cap,
+                                    original["snippet_start_line"])
+            window["snippet_truncated"] |= original["snippet_truncated"]
+            if probe(window) == "over":
+                hi = cap - 1
+            else:
+                lo = cap + 1
+    return result
+
+
 def bounded_search(metadata, data, max_tokens, *, all_results, follow_up=None):
     """Fit search delivery, reserving a progressing retry after final omissions."""
     payload = deepcopy(data)
@@ -179,7 +246,8 @@ def bounded_search(metadata, data, max_tokens, *, all_results, follow_up=None):
             and len(delivered["data"]["results"]) >= delivered["data"]["ranking"]["candidate_pool_count"]):
         delivered["data"].pop("follow_up", None)
         result = compact_result(delivered["metadata"], delivered["data"])
-    return result
+    return _restore_search_snippets(result, all_results, max_tokens,
+                                    upstream_truncated=metadata.get("truncated", False))
 
 
 def bounded_compact(metadata: dict, data, max_tokens: int, *, search_results=None) -> types.CallToolResult:

attocode_code_intel-0.1.19/tests/test_recovery_readiness.py

@@ -0,0 +1,157 @@
+"""A completed snapshot fallback must publish its reconstructed readiness."""
+
+import json
+import sqlite3
+from contextlib import nullcontext
+from dataclasses import asdict
+from unittest.mock import patch
+
+import pytest
+from attocode_intel._internal.integrations.context import ast_service
+from attocode_intel._internal.integrations.context.ast_service import ASTService
+from attocode_intel._internal.integrations.context.index_store import SNAPSHOT_KEY, IndexStore
+from attocode_intel.service import CodeIntelService
+
+from eval.quality.readiness import coverage_ready
+
+
+def close(service):
+    service.stop_hydration()
+    service._store.close()
+
+
+def seed(root):
+    (root / 'helper.py').write_text('def helper():\n    return 1\n')
+    (root / 'caller.py').write_text('from helper import helper\ndef caller():\n    return helper()\n')
+    (root / 'empty.py').write_text('# intentionally no definitions\n')
+    service = ASTService(str(root))
+    assert service.initialize_skeleton().phase == 'ready'
+    return service
+
+
+def navigation(service):
+    index = service._index
+    return {
+        'asts': {p: asdict(ast) for p, ast in service._ast_cache.items()},
+        'definitions': {p: [asdict(v) for v in values] for p, values in index.definitions.items()},
+        'references': {p: [asdict(v) for v in values] for p, values in index.references.items()},
+        'dependencies': index.file_dependencies,
+        'calls': index.call_edges,
+        'callers': index.callers_of,
+        'stats': service._store.stats(),
+    }
+
+
+@pytest.mark.parametrize('fault', ['sqlite', 'malformed', 'shape'])
+def test_recovered_snapshot_publishes_ready(tmp_path, fault):
+    first = seed(tmp_path)
+    expected = navigation(first)
+    if fault != 'sqlite':
+        first._store.set_meta(SNAPSHOT_KEY, '{bad json' if fault == 'malformed' else 'null')
+    close(first)
+    second = ASTService(str(tmp_path))
+    failure = patch.object(IndexStore, 'load_symbols', side_effect=sqlite3.DatabaseError('read failed')) if fault == 'sqlite' else nullcontext()
+    try:
+        with failure, patch.object(second, '_index_references', wraps=second._index_references) as extracted:
+            state = second.initialize_skeleton()
+        assert extracted.call_count == state.total_files == state.parsed_files == state.reference_indexed_files == 3
+        assert navigation(second) == expected
+        assert json.loads(second._store.get_meta(SNAPSHOT_KEY))['parsed'] == sorted(expected['asts'])
+        coverage = second.hydration_snapshot()
+        assert coverage['parse_coverage'] == coverage['reference_coverage'] == 1.0
+        assert state.phase == coverage['phase'] == 'ready', 'completed snapshot fallback still reports skeleton'
+        assert coverage_ready(coverage)
+    finally:
+        close(second)
+    third = ASTService(str(tmp_path))
+    try:
+        with patch.object(third, '_index_references', side_effect=AssertionError('reparsed healthy recovery')):
+            assert third.initialize_skeleton().phase == 'ready'
+        assert navigation(third) == expected
+    finally:
+        close(third)
+
+
+@pytest.mark.parametrize('partial', ['budget', 'parse'])
+def test_partial_snapshot_fallback_stays_skeleton(tmp_path, monkeypatch, partial):
+    first = seed(tmp_path)
+    close(first)
+    if partial == 'budget':
+        monkeypatch.setattr(ast_service, 'skeleton_budget', lambda *_: 1)
+    second = ASTService(str(tmp_path))
+    original = ast_service.parse_file
+
+    def parse(path):
+        if partial == 'parse' and str(path).endswith('/caller.py'):
+            raise OSError('source read failed')
+        return original(path)
+
+    try:
+        with (patch.object(IndexStore, 'load_symbols', side_effect=sqlite3.DatabaseError('read failed')),
+              patch.object(ast_service, 'parse_file', parse)):
+            state = second.initialize_skeleton()
+        assert state.parsed_files < state.total_files
+        assert state.phase == second.hydration_snapshot()['phase'] == 'skeleton'
+        assert not coverage_ready(second.hydration_snapshot())
+        assert not second._store.get_meta(SNAPSHOT_KEY)
+    finally:
+        close(second)
+
+
+def test_minimal_initialization_keeps_its_skeleton(tmp_path):
+    first = seed(tmp_path)
+    close(first)
+    second = ASTService(str(tmp_path))
+    try:
+        with patch.object(IndexStore, 'load_symbols', side_effect=AssertionError('minimal loaded snapshot')):
+            state = second.initialize_skeleton(indexing_depth='minimal')
+        assert state.phase == 'skeleton' and state.parsed_files == state.reference_indexed_files == 0
+        assert not coverage_ready(second.hydration_snapshot())
+    finally:
+        close(second)
+
+
+def test_empty_fallback_publishes_completed_empty_state(tmp_path):
+    service = ASTService(str(tmp_path))
+    service._store.set_meta(SNAPSHOT_KEY, 'null')
+    try:
+        state = service.initialize_skeleton()
+        assert state.total_files == state.parsed_files == state.reference_indexed_files == 0
+        assert state.phase == 'ready'
+        assert service._store.get_meta(SNAPSHOT_KEY)
+        assert not coverage_ready(service.hydration_snapshot())  # experiment barrier requires a nonempty fixture
+    finally:
+        close(service)
+
+
+def test_interrupted_fallback_does_not_publish_ready(tmp_path):
+    first = seed(tmp_path)
+    close(first)
+    second = ASTService(str(tmp_path))
+    try:
+        with (patch.object(IndexStore, 'load_symbols', side_effect=sqlite3.DatabaseError('read failed')),
+              patch.object(IndexStore, 'save_files_batch', side_effect=OSError('write interrupted')),
+              pytest.raises(OSError, match='write interrupted')):
+            second.initialize_skeleton()
+        assert second._hydration_state.phase == 'skeleton'
+        assert not second.initialized and not second._store.get_meta(SNAPSHOT_KEY)
+    finally:
+        close(second)
+
+
+def test_service_uses_completed_recovery_without_background_hydration(tmp_path):
+    first = seed(tmp_path)
+    close(first)
+    ast = ASTService(str(tmp_path))
+    service = CodeIntelService(str(tmp_path))
+    try:
+        # Suppress the real background thread so it cannot hide the returned phase.
+        with (patch.object(ASTService, 'get_instance', return_value=ast),
+              patch.object(IndexStore, 'load_symbols', side_effect=sqlite3.DatabaseError('read failed')),
+              patch.object(ASTService, 'start_hydration') as background):
+            restored = service._get_ast_service()
+        assert restored is ast and coverage_ready(restored.hydration_snapshot())
+        assert background.call_count == 0
+        assert ast._index.get_references('helper')
+    finally:
+        close(ast)

{attocode_code_intel-0.1.17 → attocode_code_intel-0.1.19}/tests/test_search_budget_repair.py

@@ -115,7 +115,17 @@ def test_new_follow_up_never_crowds_out_previously_delivered_evidence(path_size,
     result = bounded_search(metadata, data, budget, all_results=data["results"], follow_up=follow)
     assert response_tokens(result) <= budget
     actual = json.loads(result.content[0].text)["data"]["results"]
-    assert actual[:len(protected)] == protected
+    assert len(actual) >= len(protected)
+    preview_fields = {"snippet", "snippet_start_line", "snippet_end_line", "snippet_truncated"}
+    for before, after in zip(protected, actual, strict=False):
+        assert {k: v for k, v in before.items() if k not in preview_fields} == {
+            k: v for k, v in after.items() if k not in preview_fields}
+        assert before["snippet"] in after["snippet"]
+        if before["snippet_start_line"] is not None:
+            assert after["snippet_start_line"] <= before["snippet_start_line"]
+            assert after["snippet_end_line"] >= before["snippet_end_line"]
+        else:
+            assert after["snippet_start_line"] is after["snippet_end_line"] is None
 
 
 @pytest.mark.asyncio

attocode_code_intel-0.1.19/tests/test_snippet_budget.py

@@ -0,0 +1,300 @@
+"""Source text discarded by coarse fitting can be restored without losing the page."""
+
+from __future__ import annotations
+
+import json
+from copy import deepcopy
+
+import pytest
+from attocode_intel.output import bounded_compact, bounded_search, compact_result, response_tokens
+
+
+def source_page(kind="unicode", workspace="/repo"):
+    body = (
+        ('    needle = "' + "雪境風雷光雨山海" * 40 + '"\n')
+        if kind == "unicode"
+        else ("    # needle " + '"\\\t雪"' * 64 + "\r\n")
+    )
+    row = {
+        "file_path": "src/worker.py",
+        "name": "worker",
+        "chunk_type": "function",
+        "score": 0.5,
+        "line": 9,
+        "end_line": 12,
+        "snippet": body,
+        "snippet_start_line": 10,
+        "snippet_end_line": 10,
+        "snippet_truncated": True,
+        "matched_terms": ["needle"],
+        "match_fields": ["source"],
+    }
+    return {"workspace": workspace, "truncated": True}, {
+        "query": "needle",
+        "results": [row],
+        "ranking": {"candidate_pool_count": 1, "returned_count": 1, "delivered_count": 1},
+    }
+
+
+@pytest.mark.parametrize("kind", ["unicode", "escaped"])
+@pytest.mark.parametrize(
+    "workspace", ["/repo", "/" + "long_checkout_name/" * 30], ids=["short-root", "long-root"]
+)
+def test_small_overflow_preserves_more_than_three_quarters_of_source(kind, workspace):
+    metadata, data = source_page(kind, workspace)
+    budget = response_tokens(compact_result(metadata, data)) - 1
+    original = data["results"][0]["snippet"]
+    legacy = bounded_compact(metadata, data, budget, search_results=[])
+    old = json.loads(legacy.content[0].text)["data"]["results"][0]["snippet"]
+    assert len(old) <= len(original) // 2
+    actual = bounded_search(metadata, data, budget, all_results=data["results"])
+    delivered = json.loads(actual.content[0].text)["data"]["results"][0]
+    assert response_tokens(actual) <= budget
+    assert len(delivered["snippet"]) > len(original) * 3 // 4, (
+        "snippet restoration boundary: excessive source loss"
+    )
+    assert old in delivered["snippet"] and delivered["snippet"] in original
+    assert delivered["snippet_start_line"] == delivered["snippet_end_line"] == 10
+
+
+def test_fitting_search_payload_is_byte_identical():
+    metadata, data = source_page()
+    budget = response_tokens(compact_result(metadata, data)) + 1000
+    legacy = bounded_compact(metadata, data, budget, search_results=[])
+    actual = bounded_search(metadata, data, budget, all_results=data["results"])
+    assert actual.model_dump() == legacy.model_dump()
+
+
+def test_restoration_does_not_mutate_inputs():
+    metadata, data = source_page()
+    original = deepcopy((metadata, data))
+    budget = response_tokens(compact_result(metadata, data)) - 1
+    bounded_search(metadata, data, budget, all_results=data["results"])
+    assert (metadata, data) == original
+
+
+def test_delivered_floor_and_source_coordinates_remain_valid():
+    metadata, data = source_page()
+    budget = response_tokens(compact_result(metadata, data)) - 1
+    legacy = bounded_compact(metadata, data, budget, search_results=[])
+    old = json.loads(legacy.content[0].text)["data"]["results"]
+    response = bounded_search(metadata, data, budget, all_results=data["results"])
+    actual = json.loads(response.content[0].text)["data"]["results"]
+    previews = {"snippet", "snippet_start_line", "snippet_end_line", "snippet_truncated"}
+    assert len(actual) == len(old) == 1
+    for before, after in zip(old, actual, strict=True):
+        assert {k: v for k, v in before.items() if k not in previews} == {
+            k: v for k, v in after.items() if k not in previews
+        }
+        assert before["snippet"] in after["snippet"] in data["results"][0]["snippet"]
+        assert after["snippet_start_line"] == after["snippet_end_line"] == 10
+    assert response_tokens(response) <= budget
+
+
+def legacy_search(monkeypatch, metadata, data, budget, follow_up=None):
+    from attocode_intel import output
+
+    with monkeypatch.context() as patch:
+        patch.setattr(output, "_restore_search_snippets", lambda result, *args, **kwargs: result)
+        return output.bounded_search(
+            metadata, data, budget, all_results=data["results"], follow_up=follow_up
+        )
+
+
+def assert_protected_page(before, after):
+    old = json.loads(before.content[0].text)
+    new = json.loads(after.content[0].text)
+    assert len(old["data"]["results"]) == len(new["data"]["results"])
+    previews = {"snippet", "snippet_start_line", "snippet_end_line", "snippet_truncated"}
+    for previous, current in zip(old["data"]["results"], new["data"]["results"], strict=True):
+        assert {k: v for k, v in previous.items() if k not in previews} == {
+            k: v for k, v in current.items() if k not in previews
+        }
+        assert previous["snippet"] in current["snippet"]
+        if previous["snippet_start_line"] is None:
+            assert current["snippet_start_line"] is current["snippet_end_line"] is None
+        else:
+            assert current["snippet_start_line"] <= previous["snippet_start_line"]
+            assert current["snippet_end_line"] >= previous["snippet_end_line"]
+    assert {k: v for k, v in old["data"].items() if k != "results"} == {
+        k: v for k, v in new["data"].items() if k != "results"
+    }
+    assert {k: v for k, v in old["metadata"].items() if k not in {"truncated", "omitted"}} == {
+        k: v for k, v in new["metadata"].items() if k not in {"truncated", "omitted"}
+    }
+
+
+def many_source_rows(count=5):
+    metadata, data = source_page()
+    row = data["results"][0]
+    data["results"] = [
+        {**deepcopy(row), "file_path": f"src/worker_{i}.py", "score": 1 - i / 100}
+        for i in range(count)
+    ]
+    data["ranking"].update(candidate_pool_count=count, returned_count=count, delivered_count=count)
+    return metadata, data
+
+
+@pytest.mark.parametrize("budget", [600, 1200, 2000, 8000])
+def test_complete_fitted_prefix_and_continuation_are_protected(monkeypatch, budget):
+    metadata, data = many_source_rows()
+    follow = {
+        "tool": "semantic_search",
+        "arguments": {
+            "query": "needle",
+            "top_k": 10,
+            "workspace": metadata["workspace"],
+            "max_tokens": 16000,
+        },
+        "reruns_current_index": True,
+        "reason": "search_deeper",
+    }
+    before = legacy_search(monkeypatch, metadata, data, budget, follow)
+    after = bounded_search(metadata, data, budget, all_results=data["results"], follow_up=follow)
+    assert_protected_page(before, after)
+    assert response_tokens(after) <= budget
+
+
+def test_restoration_has_a_global_serialization_probe_cap(monkeypatch):
+    from attocode_intel import output
+
+    metadata, data = many_source_rows(24)
+    budget = 8000
+    before = legacy_search(monkeypatch, metadata, data, budget)
+    real_count = output.response_tokens
+    calls = []
+
+    def count(result):
+        calls.append(result)
+        return real_count(result)
+
+    monkeypatch.setattr(output, "response_tokens", count)
+    after = output._restore_search_snippets(
+        before, data["results"], budget, upstream_truncated=True
+    )
+    assert 0 < len(calls) <= 12
+    assert real_count(after) <= budget
+    assert_protected_page(before, after)
+
+
+def test_nonmonotone_probes_cannot_admit_overbudget_or_lose_text(monkeypatch):
+    from attocode_intel import output
+
+    metadata, data = source_page()
+    budget = response_tokens(compact_result(metadata, data)) - 1
+    before = legacy_search(monkeypatch, metadata, data, budget)
+    real_count = output.response_tokens
+    calls = []
+
+    def conservative_nonmonotone(result):
+        actual = real_count(result)
+        measured = max(actual, budget + 1) if len(calls) % 2 else actual
+        calls.append((actual, measured))
+        return measured
+
+    monkeypatch.setattr(output, "response_tokens", conservative_nonmonotone)
+    after = output._restore_search_snippets(
+        before, data["results"], budget, upstream_truncated=True
+    )
+    assert any(measured > actual for actual, measured in calls)
+    assert len(calls) <= 12 and real_count(after) <= budget
+    assert_protected_page(before, after)
+
+
+@pytest.mark.parametrize("upstream", [False, True])
+@pytest.mark.parametrize("unrelated", [False, True])
+def test_full_restoration_cleans_only_its_own_omission(upstream, unrelated):
+    from attocode_intel.output import _restore_search_snippets
+
+    metadata, data = source_page()
+    original = data["results"][0]
+    original.update(line=10, end_line=10, snippet_truncated=False)
+    rows = deepcopy(data["results"])
+    data["results"][0].update(snippet=original["snippet"][:40], snippet_truncated=True)
+    metadata["truncated"] = True
+    metadata["omitted"] = {"data.results.snippet": "shortened"}
+    if unrelated:
+        metadata["omitted"]["data.other"] = "shortened"
+    before = compact_result(metadata, data)
+    after = _restore_search_snippets(before, rows, 8000, upstream_truncated=upstream)
+    page = json.loads(after.content[0].text)
+    assert page["data"]["results"] == rows
+    assert not page["metadata"].get("omitted", {}).get("data.results.snippet")
+    assert page["metadata"]["truncated"] is (upstream or unrelated)
+    assert bool(page["metadata"].get("omitted", {}).get("data.other")) is unrelated
+    assert_protected_page(before, after)
+    assert response_tokens(after) <= 8000
+
+
+def test_repeated_text_keeps_corresponding_source_rows():
+    from attocode_intel.output import _restore_search_snippets
+
+    metadata, data = many_source_rows(2)
+    for i, row in enumerate(data["results"]):
+        row.update(
+            file_path="src/worker.py",
+            line=10 + i * 10,
+            end_line=10 + i * 10,
+            snippet_start_line=10 + i * 10,
+            snippet_end_line=10 + i * 10,
+            snippet=row["snippet"] + f"# source marker {i}",
+        )
+        row["snippet_end_line"] += 1
+        row["end_line"] += 1
+    originals = deepcopy(data["results"])
+    for row in data["results"]:
+        row.update(snippet=row["snippet"][:32], snippet_end_line=row["snippet_start_line"])
+    before = compact_result(metadata, data)
+    after = _restore_search_snippets(before, originals, 8000, upstream_truncated=True)
+    assert json.loads(after.content[0].text)["data"]["results"] == originals
+    assert_protected_page(before, after)
+
+
+async def test_actual_gateway_preserves_source_and_budget_after_restoration(tmp_path, monkeypatch):
+    import asyncio
+    import time
+
+    from attocode_intel.gateway import OperationGateway
+    from attocode_intel.service import CodeIntelService
+
+    monkeypatch.setenv("ATTOCODE_INTEL_RANKING_PROVIDER", "off")
+    monkeypatch.setenv("ATTOCODE_INTEL_PRECISION", "off")
+    body = 'def worker():\n    needle = "' + "雪境風雷光雨山海" * 80 + '"\n    return needle\n'
+    (tmp_path / "worker.py").write_text(body)
+    service = CodeIntelService(str(tmp_path))
+    manager = service._get_semantic_search()
+    manager.warm()
+    assert manager.wait_for_body_index(timeout=10)
+    gateway = OperationGateway(str(tmp_path), "daily", watch=False)
+    arguments = {"query": "needle", "top_k": 24, "max_tokens": 8000}
+    try:
+        deadline = time.monotonic() + 20
+        while True:
+            full = await gateway.execute_mcp("semantic_search", arguments)
+            assert not full.isError
+            page = json.loads(full.content[0].text)
+            if page["data"]["ranking"]["index"]["status"] == "ready":
+                break
+            assert time.monotonic() < deadline
+            await asyncio.sleep(0.05)
+        budget = response_tokens(full) - 1
+        from attocode_intel import output
+
+        with monkeypatch.context() as patch:
+            patch.setattr(output, "_restore_search_snippets", lambda result, *a, **k: result)
+            before = await gateway.execute_mcp(
+                "semantic_search", {**arguments, "max_tokens": budget}
+            )
+        after = await gateway.execute_mcp("semantic_search", {**arguments, "max_tokens": budget})
+        assert not before.isError and not after.isError
+        assert_protected_page(before, after)
+        assert response_tokens(after) <= budget
+        for hit in json.loads(after.content[0].text)["data"]["results"]:
+            assert hit["file_path"] == "worker.py"
+            lines = body.splitlines(keepends=True)
+            assert hit["snippet"] in "".join(
+                lines[hit["snippet_start_line"] - 1 : hit["snippet_end_line"]]
+            )
+    finally:
+        await gateway.close()
+        manager.close()