fix(cpg): answer a failed batch in parts instead of losing it whole (M3)

PMPro's 1,274 PHP files produced a response the engine could not serialize: it
failed inside ujson while building its own answer array, and every question in that
batch was lost with it. Three budgets from 900 to 3000 seconds made no difference,
because the size of the answer was the problem rather than the time available.

One invocation stays the first thing tried, so a working scan costs exactly what it
did. Each request in a batch is independent, so a batch that fails as a whole can
still answer in parts: on failure the requests are split and retried, and a part
that fails leaves its own requests unanswered rather than discarding every answer
beside them. Only a batch where no part answered is reported as unavailable, which
keeps "could not ask" distinct from "no rows".

Controls cover all three shapes: a whole batch that fails is answered in parts with
the split recorded, a failing part leaves only its own requests unanswered while the
rest are kept, and a batch no part can answer is still unavailable rather than
empty.

Whether this is enough for PMPro is a separate question the measurement will
answer: a single taint request on that tree already takes 40 seconds and returns
36 KB, so the batch may be slow as well as large.

1,410 tests passed, nine skipped.

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

src/openultrasast/cpg/backend.py

@@ -52,6 +52,9 @@ BUILD_TIMEOUT_SECONDS = 900  # a large tree takes minutes; the go/no-go records
 # The constant was chosen when a "scan" meant an excerpt, and a repository scan is not a slower slice -- it
 # is a different order of question, so the ceiling is configurable rather than a number to keep raising.
 QUERY_TIMEOUT_SECONDS = 300
+# How many parts a failed batch is retried in. Small enough that a response the engine could not
+# serialize becomes several it can, large enough that recovery is not one invocation per request.
+_BATCH_SPLITS = 8
 QUERY_TIMEOUT_ENV = "OPENULTRASAST_CPG_QUERY_TIMEOUT"
 # A JVM with no -Xmx takes a quarter of physical RAM for its heap. On a contributor's laptop that is the
 # difference between a scan running in the background and a scan the machine notices, and it made the
@@ -1035,6 +1038,43 @@ class JoernBackend:
         return _unparsed_files(completed)
 
     def query_batch(self, cpg_path: Path, query: str, requests: Mapping[str, Mapping[str, object]]) -> dict[str, list[object]] | None:
+        """Answer many requests, as one invocation where that works and in parts where it does not.
+
+        One invocation is the whole point of batching, so it stays the first thing tried and the cost
+        of a working scan is unchanged. But a batch is all-or-nothing: PMPro's 1,274 PHP files
+        produced a response the engine could not serialize, failing inside ujson while building its
+        own answer array, and every question in that batch was lost with it. Three budgets from 900
+        to 3000 seconds made no difference, because the size of the answer was the problem.
+
+        Each request in a batch is independent, so a batch that fails as a whole can still answer in
+        parts. On failure the requests are split and retried, and a part that fails leaves its own
+        requests unanswered instead of discarding every answer beside them. Only a batch where no
+        part answered is reported as unavailable.
+        """
+        if not requests:
+            return self._batch_once(cpg_path, query, requests)
+        answered = self._batch_once(cpg_path, query, requests)
+        if answered is not None or len(requests) <= 1:
+            return answered
+        size = max(1, len(requests) // _BATCH_SPLITS)
+        ids = list(requests)
+        merged: dict[str, list[object]] = {}
+        recovered = False
+        for start in range(0, len(ids), size):
+            if not self._time_available():
+                break
+            part = self._batch_once(cpg_path, query, {rid: requests[rid] for rid in ids[start : start + size]})
+            if part is None:
+                continue
+            recovered = True
+            for rid, rows in part.items():
+                merged.setdefault(rid, rows)
+        if recovered:
+            self._note("batch_split_after_failure:" + query)
+            logger.warning("cpg batch %s failed whole; %d of %d requests answered in parts", query, len(merged), len(requests))
+        return merged if recovered else None
+
+    def _batch_once(self, cpg_path: Path, query: str, requests: Mapping[str, Mapping[str, object]]) -> dict[str, list[object]] | None:
         """Run ONE script invocation carrying many requests, keyed back to their ids.
 
         JVM startup, not CPG construction, is what dominates a repository scan: every ``joern --script`` call
@@ -1065,7 +1105,7 @@ class JoernBackend:
         # catch because no verdict was ever produced to label.
         if not self._time_available():
             return None
-        requests_file = cpg_path.parent / f"{query}-requests.json"
+        requests_file = cpg_path.parent / f"{query}-requests-{len(rendered)}-{abs(hash(tuple(sorted(rendered)))) % 10**8}.json"
         try:
             requests_file.write_text(payload)
         except OSError as exc:

tests/test_cpg_backend.py

@@ -1159,3 +1159,61 @@ def test_javascript_partition_uses_declared_frontend_not_legacy_autodetection(tm
     finally:
         if result is not None:
             result.cleanup()
+
+
+def test_a_batch_that_fails_whole_is_answered_in_parts(tmp_path, monkeypatch):
+    """M3: PMPro's batch failed inside the engine's own serialization and lost every answer with it."""
+    from openultrasast.cpg.backend import JoernBackend
+
+    seen = []
+
+    def runner(command, **kwargs):
+        payload = Path(command[command.index("--param") + 3].split("=", 1)[1])
+        requests = json.loads(payload.read_text())
+        seen.append(len(requests))
+        if len(requests) > 4:  # the whole batch is more than this engine can answer
+            return subprocess.CompletedProcess(command, 1, "", "ujson serialization failed")
+        rows = {rid: [{"sink": "q(" + rid + ")"}] for rid in requests}
+        body = "---OUSAST-CPG-BEGIN---\n" + json.dumps(rows) + "\n---OUSAST-CPG-END---"
+        return subprocess.CompletedProcess(command, 0, body, "")
+
+    monkeypatch.setattr("shutil.which", lambda name: f"/opt/joern/{name}")
+    backend = JoernBackend(runner=runner)
+    graph = tmp_path / "cpg.bin"
+    graph.write_bytes(b"graph")
+    answered = backend.query_batch(graph, "taint", {str(i): {"file": f"f{i}.php"} for i in range(16)})
+    assert answered is not None and len(answered) == 16, "every independent request should still be answered"
+    assert seen[0] == 16 and max(seen[1:]) <= 4, "the whole batch is tried first, then parts"
+    assert any("batch_split_after_failure" in note for note in backend._diagnostics)
+
+
+def test_a_part_that_fails_leaves_only_its_own_requests_unanswered(tmp_path, monkeypatch):
+    from openultrasast.cpg.backend import JoernBackend
+
+    def runner(command, **kwargs):
+        payload = Path(command[command.index("--param") + 3].split("=", 1)[1])
+        requests = json.loads(payload.read_text())
+        if len(requests) > 4 or "0" in requests:  # whole batch, and one part, cannot answer
+            return subprocess.CompletedProcess(command, 1, "", "engine failure")
+        rows = {rid: [] for rid in requests}
+        body = "---OUSAST-CPG-BEGIN---\n" + json.dumps(rows) + "\n---OUSAST-CPG-END---"
+        return subprocess.CompletedProcess(command, 0, body, "")
+
+    monkeypatch.setattr("shutil.which", lambda name: f"/opt/joern/{name}")
+    backend = JoernBackend(runner=runner)
+    graph = tmp_path / "cpg.bin"
+    graph.write_bytes(b"graph")
+    answered = backend.query_batch(graph, "taint", {str(i): {"file": f"f{i}.php"} for i in range(16)})
+    assert answered is not None and "0" not in answered, "a failed part stays unanswered, not empty"
+    assert len(answered) == 14, "the parts that answered are kept"
+
+
+def test_a_batch_no_part_can_answer_is_still_unavailable(tmp_path, monkeypatch):
+    from openultrasast.cpg.backend import JoernBackend
+
+    monkeypatch.setattr("shutil.which", lambda name: f"/opt/joern/{name}")
+    backend = JoernBackend(runner=lambda command, **k: subprocess.CompletedProcess(command, 1, "", "dead"))
+    graph = tmp_path / "cpg.bin"
+    graph.write_bytes(b"graph")
+    # None means "could not ask", which must never become an empty answer.
+    assert backend.query_batch(graph, "taint", {str(i): {"file": f"f{i}.php"} for i in range(8)}) is None