feat(scan): size portions from measured per-answer cost and isolate the kill culprit

The 7.8 trace showed the completion loss was a sizer death spiral, not
expensive families: after two kills the budget fell to 3 sink visits and
never recovered, because portions finishing in 25-30 s sat under an assumed
70 s fixed start and read as "no information", so the budget could only
shrink.

- taint.sc streams each answer's own cost (`ms`); the backend carries it as a
  reserved `__timing__` list through shards, partitions and reuse.
- The sizer budgets predicted seconds from per-family seconds-per-visit rates
  learned from those costs, and measures the fixed start as the wall time the
  answers do not explain (median), instead of assuming it.
- A kill names its culprit: the first unanswered id in request order. The ids
  behind it never started and are re-queued at once; the culprit is asked
  alone after the rest. One that still fails alone after more than a
  portion's target is reported as `query_too_expensive`, apart from
  `query_failed`.
- Portions grow id by id from the queue head, and a single-portion scan now
  goes through the queue too, so a kill there re-asks its remainder.
- Fix: the reuse wrapper wrapped only `run_batch`, so since 7.8 every
  portioned query bypassed the graph-reuse cache and its source-root
  normalization. `run_batch_once` is wrapped too.

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
Claude-Session: https://claude.ai/code/session_01EmM1j46qgKMh3zqD1QMwEo

tests/test_cpg_backend.py

@@ -1271,46 +1271,63 @@ def test_a_killed_batch_retries_only_the_remainder(tmp_path: Path) -> None:
     assert all("a" not in retry for retry in calls[1:]), calls
 
 
-def test_the_portioner_asks_single_attempts_and_halves_on_a_short_answer(tmp_path: Path) -> None:
-    """A killed portion returns what it streamed; the sizer must read that as too big, not as slow.
+def test_the_portioner_asks_single_attempts_and_isolates_the_culprit(tmp_path: Path) -> None:
+    """A killed portion's running request is asked alone, AFTER the rest; the ids behind it go straight back.
 
-    Timing `run_batch` folds its internal split-retry into the number -- one portion read 953 s that way and
-    the sizer fitted a per-visit rate from it and collapsed. The portioner asks `run_batch_once` instead, and
-    a portion that came back short of the ids it asked halves the budget rather than fitting from a ceiling.
+    The engine answers in order and streams as it goes, so after a kill the first unanswered id is the one
+    that held the ceiling and everything after it never started. Re-asking them together at half the size is
+    what let one expensive request kill portion after portion on the plugin.
     """
     from openultrasast.cpg.backend import CpgResult
     from openultrasast.model import scan
 
-    once_calls: list[list[str]] = []
+    calls: list[list[str]] = []
 
     def once(kind: str, requests: dict[str, dict[str, object]]) -> dict[str, list[object]] | None:
-        once_calls.append(sorted(requests))
-        # Answer everything except a designated heavy id, as a kill would: it streamed the rest and died.
-        return {rid: [] for rid in requests if rid != "heavy"}
+        calls.append(list(requests))
+        answer: dict[str, list[object]] = {"__census__": [{"methods": "3"}]}
+        for rid in requests:
+            if rid == "heavy":
+                return answer  # killed while answering heavy: nothing after it started
+            answer[rid] = []
+        return answer
 
     result = CpgResult(cpg_path=tmp_path / "c.bin", run=lambda k, p: None, run_batch=lambda k, r: None, run_batch_once=once)
+    weights = dict.fromkeys(["a", "heavy", "b", "c"], 1)
+    expensive: dict[str, float] = {}
+    merged = scan._batched(result.run_batch_once, "taint", {rid: {} for rid in weights}, weights=weights, too_expensive=expensive)
+    assert merged is not None and {"a", "b", "c"} <= set(merged) and "heavy" not in merged
+    assert calls[0] == ["a", "heavy", "b", "c"]
+    assert calls[1] == ["b", "c"], "the requests behind the culprit were not re-asked straight away"
+    assert calls[-1] == ["heavy"] and calls.count(["heavy"]) == 1, calls
+    assert not expensive, "a lone failure faster than a portion's target is not 'too expensive'"
 
-    seen: list[bool] = []
-    real_observe = scan._PortionSizer.observe
 
-    def spy(self, chunk, seconds, *, answered):  # type: ignore[no-untyped-def]
-        seen.append(answered)
-        return real_observe(self, chunk, seconds, answered=answered)
+def test_a_single_portion_kill_still_re_asks_its_remainder() -> None:
+    """A scan small enough for one portion used to take one attempt and lose whatever a kill cut off."""
+    from openultrasast.model import scan
 
-    # Two portions over the budget so the portioner actually splits and observes.
-    weights = {"heavy": 500, "a": 400, "b": 400}
-    requests = {rid: {} for rid in weights}
-    monkey = scan._PortionSizer.observe
-    scan._PortionSizer.observe = spy  # type: ignore[method-assign]
-    try:
-        merged = scan._batched(result.run_batch_once, "taint", requests, weights=weights)
-    finally:
-        scan._PortionSizer.observe = monkey  # type: ignore[method-assign]
+    calls: list[list[str]] = []
 
-    assert merged is not None and "a" in merged and "b" in merged and "heavy" not in merged
-    assert once_calls, "the single-attempt path was not used"
-    # The portion that came back short of its ids was reported as not-whole (False) to the sizer.
-    assert False in seen, seen
+    def once(kind: str, requests: dict[str, dict[str, object]]) -> dict[str, list[object]] | None:
+        calls.append(list(requests))
+        if len(calls) == 1:
+            return {"__census__": [], "a": []}  # killed on b
+        return {"__census__": [], **{rid: [] for rid in requests}}
+
+    merged = scan._batched(once, "taint", {"a": {}, "b": {}, "c": {}}, weights={"a": 1, "b": 1, "c": 1})
+    assert merged is not None and {"a", "b", "c"} <= set(merged), (merged, calls)
+
+
+def test_streamed_costs_survive_every_layer() -> None:
+    """The engine's per-answer cost is parsed, carried beside the census, and never read as a request."""
+    from openultrasast.cpg.backend import BEGIN, END, TIMING, extract_payload
+    from openultrasast.model import scan
+
+    body = f'{BEGIN}\n{{"__census__": []}}\n{{"id": "0", "rows": [], "ms": 1500}}\n{{"id": "1", "rows": [1], "ms": 20}}\n{END}\n'
+    parsed = extract_payload(body)
+    assert isinstance(parsed, dict) and parsed[TIMING] == [{"id": "0", "ms": 1500.0}, {"id": "1", "ms": 20.0}]
+    assert scan._answer_timing(parsed) == {"0": 1.5, "1": 0.02}
 
 
 def test_the_queue_re_asks_a_killed_remainder_and_terminates() -> None:

tests/test_cpg_queries_compile.py

@@ -621,3 +621,8 @@ def test_the_taint_query_streams_one_line_per_request(tmp_path: Path) -> None:
     assert json.loads(lines[0]).keys() == {"__census__"}
     assembled = extract_payload(done.stdout)
     assert assembled is not None and set(assembled) >= {"0", "1", "__census__"} and "__partial__" not in assembled
+    # Every answer reports its own cost, which is what the portioner sizes from and how a kill names the
+    # request that was running. Absent, the sizer falls back to guessing from wall time.
+    answers = [json.loads(line) for line in lines[1:]]
+    assert all(isinstance(a.get("ms"), (int, float)) and a["ms"] >= 0 for a in answers), answers
+    assert [t["id"] for t in assembled["__timing__"]] == ["0", "1"]

tests/test_model_scope.py

@@ -322,26 +322,56 @@ def test_expired_planning_keeps_census_without_more_fact_reads(tmp_path, monkeyp
     assert "deadline_exhausted" in result.scope.unresolved_boundaries
 
 
-def test_a_portion_sizer_shrinks_to_budget_and_learns_from_what_answered() -> None:
-    """Portions follow the engine's observed cost; a portion that did not answer halves the budget."""
+def test_a_portion_sizer_learns_measured_costs_per_family() -> None:
+    """Portions are sized in predicted seconds from what each family's answers reported costing."""
     from openultrasast.model import scan
 
-    weights = {"w": 500, "x": 400, "y": 100, "z": 100}
-    sizer = scan._PortionSizer(weights)
-    sizer.budget = 600.0
-    pending = [{"y": {}}, {"z": {}}]
-    fitted = sizer.fit({"w": {}, "x": {}}, pending=pending)
-    assert sorted(fitted) == ["w"] and pending[0] == {"x": {}}, (fitted, pending)  # x pushed to the front, not lost
-    grown = sizer.fit({"y": {}}, pending=[{"z": {}}])
-    assert sorted(grown) == ["y", "z"]
-    sizer.observe({"w": {}}, 300.0, answered=False)
-    assert sizer.budget == 300.0
-    sizer.observe({"y": {}, "z": {}}, scan.PORTION_FIXED_SECONDS + 20.0, answered=True)  # 0.1 s per visit
-    assert 450 < sizer.budget < 550, sizer.budget
-    # A portion that finished inside the fixed start says nothing: the budget must not move on it.
-    before = sizer.budget
-    sizer.observe({"y": {}}, scan.PORTION_FIXED_SECONDS + 1.0, answered=True)
-    assert sizer.budget == before
-    # And a genuinely cheap portion may double the budget, never more: 0.01 s per visit would fit 5,000.
-    sizer.observe({"w": {}}, scan.PORTION_FIXED_SECONDS + 5.0 + 5.0, answered=True)
-    assert sizer.budget == before * scan.PORTION_MAX_GROWTH, sizer.budget
+    weights = {"i1": 10, "i2": 10, "o1": 10, **{f"x{i}": 10 for i in range(8)}}
+    families = {"i1": "injection", "i2": "injection", "o1": "output_encoding", **{f"x{i}": "output_encoding" for i in range(8)}}
+    sizer = scan._PortionSizer(weights, families)
+    # A whole portion: answers explain 2 s, the other 23 s is the fixed start -- not the 70 s prior.
+    sizer.observe({"i1": {}, "i2": {}}, 25.0, answered={"i1", "i2"}, timing={"i1": 1.0, "i2": 1.0})
+    assert sizer.fixed == 23.0
+    assert abs(sizer.rate("injection") - 0.1) < 1e-9
+    # A family never seen is costed at the dearest rate known, never below the prior.
+    assert sizer.rate("output_encoding") >= scan.PORTION_PRIOR_SECONDS_PER_VISIT
+    sizer.observe({"o1": {}}, 43.0, answered={"o1"}, timing={"o1": 20.0})
+    assert abs(sizer.rate("output_encoding") - 2.0) < 1e-9
+    assert abs(sizer.rate("injection") - 0.1) < 1e-9, "one family's cost leaked into another's"
+    # A portion is cut where the predicted seconds run out: 97 s of allowance holds four 20 s requests.
+    pending: list[dict[str, dict[str, object]]] = []
+    fitted = sizer.fit({f"x{i}": {} for i in range(8)}, pending=pending)
+    assert len(fitted) == 4 and len(pending[0]) == 4, (fitted, pending)
+
+
+def test_a_fast_portion_grows_the_next_one_after_a_collapse() -> None:
+    """The trace that motivated this: after two kills the budget fell to 3 visits and never came back.
+
+    Portions of 3 visits finished in 27 s. Under an assumed 70 s start that was 'no information', so the
+    budget could only shrink. Measured, 27 s is almost all start, the work is cheap, and the next portion
+    must carry far more.
+    """
+    from openultrasast.model import scan
+
+    weights = {f"r{i}": 3 for i in range(200)}
+    families = dict.fromkeys(weights, "injection")
+    sizer = scan._PortionSizer(weights, families)
+    sizer.observe({"r0": {}}, 27.0, answered={"r0"}, timing={"r0": 0.6})
+    pending = [{f"r{i}": {} for i in range(2, 200)}]
+    fitted = sizer.fit({"r1": {}}, pending=pending)
+    assert len(fitted) >= 100, len(fitted)
+
+
+def test_a_killed_culprit_makes_its_family_dearer_and_nothing_else() -> None:
+    """The request running at the kill ran at least as long as nothing else explains."""
+    from openultrasast.model import scan
+
+    weights = {"a": 5, "b": 5, "c": 5}
+    families = {"a": "injection", "b": "output_encoding", "c": "output_encoding"}
+    sizer = scan._PortionSizer(weights, families)
+    sizer.fixed = 25.0
+    sizer.observe({"a": {}, "b": {}, "c": {}}, 300.0, answered={"a"}, timing={"a": 5.0}, culprit="b")
+    # b ran at least 300 - 25 - 5 = 270 s over 5 visits.
+    assert sizer.rate("output_encoding") >= 54.0
+    assert abs(sizer.rate("injection") - 1.0) < 1e-9
+    assert sizer.fixed == 25.0, "a killed portion must not be read as a fixed-start sample"

tests/test_push_reuse.py

@@ -87,3 +87,37 @@ def test_fresh_missing_census_cannot_borrow_a_cached_query_census(tmp_path):
     assert "q" in result.run_batch("taint", {"q": {"file": "a.js"}})
     answer = result.run_batch("taint", {"q": {"file": "a.js"}, "fresh": {"file": "a.js", "callDepth": 3}})
     assert "q" in answer and "fresh" not in answer
+
+
+def test_the_single_attempt_path_is_reused_and_normalized_too(tmp_path):
+    """The portioner asks `run_batch_once`; wrapping only `run_batch` let every portioned query skip reuse."""
+    from openultrasast.cpg.artifact import GraphArtifact, GraphCensus, GraphIdentity
+    from openultrasast.cpg.backend import TIMING
+
+    identity = GraphIdentity(
+        "source", "declarations", "exclude", "repository", "javascript", "engine", "frontend", "version", ("dataflowOss",), ()
+    )
+    artifact = GraphArtifact(identity, "graph-bytes", 10, GraphCensus(1, 1, 1), "complete", (), "/old", ("a.js",))
+    cache = ArtifactCache(tmp_path / "cache", max_bytes=100000)
+    semantics = SemanticKeys("facts", "queries", "config", "ranking", "admission", "disabled", "empty", "advisory")
+    wrapper = ReusingBackend(JoernBackend(), cache, semantics, declarations="d", exclusions="e")
+    census = [{"files": 1, "methods": 1, "file_names": ["a.js"]}]
+    once_calls = []
+
+    def once(kind, requests):
+        once_calls.append(dict(requests))
+        return {
+            **{rid: [{"file": "/old/a.js"}] for rid in requests},
+            "__census__": census,
+            TIMING: [{"id": rid, "ms": 5.0} for rid in requests],
+        }
+
+    graph = CpgResult(tmp_path / "graph", lambda *_: None, lambda k, r: None, run_batch_once=once)
+    result = wrapper._answers(graph, artifact, ExecutionBudget(time.monotonic() + 5, 0.2))
+    assert result.run_batch_once is not None
+    first = result.run_batch_once("taint", {"q": {"file": "a.js"}})
+    assert first["q"] == [{"file": "a.js"}], "the source root was not stripped on the single-attempt path"
+    assert first[TIMING] == [{"id": "q", "ms": 5.0}]
+    second = result.run_batch_once("taint", {"q": {"file": "a.js"}})
+    assert second["q"] == [{"file": "a.js"}] and len(once_calls) == 1, "the single-attempt path bypassed the cache"
+    assert TIMING not in second, "a cached answer cost the engine nothing and must report no time"