engine: one resident Joern session per pin (frontend + one query JVM instead of four), ActiveProcessorCount=2 on the cluster

Implemented with Codex.

Co-Authored-By: Claude Opus 5.5 <[email protected]>

tests/test_cpg_session.py

@@ -343,3 +343,42 @@ def test_a_startup_wait_is_bounded_by_its_allowance_not_the_deadline(tmp_path, f
     # Bounded by the allowance, nowhere near the 600-second deadline.
     assert time.monotonic() - started < 10.0
     live.close()
+
+
+@pytest.mark.parametrize("cancel", [False, True])
+def test_live_receipt_observation_without_socket_or_jvm(tmp_path, monkeypatch, cancel):
+    import threading
+    from types import SimpleNamespace
+
+    live = EngineSession(tmp_path, ExecutionBudget(time.monotonic() + 10, 1), {})
+    live._process = SimpleNamespace(poll=lambda: None)
+    stopped = threading.Event()
+    observed = []
+
+    def close():
+        live._process = None
+        stopped.set()
+
+    def post(code, **kwargs):
+        target = code.split('java.nio.file.Path.of("', 1)[1].split('"', 1)[0]
+        nonce = code.split('__ousastPath, "', 1)[1].split("\\n", 1)[0]
+        Path(target).write_text(nonce + '\n{"__question__":"a"}\n')
+        if cancel:
+            assert stopped.wait(2)
+        return {"success": True}
+
+    def observe(output):
+        observed.append(output)
+        return not cancel
+
+    monkeypatch.setattr(live, "_post", post)
+    monkeypatch.setattr(live, "close", close)
+    live.output_observer = observe
+    answer = live.evaluate("work()", timeout=5)
+    assert any("__question__" in output for output in observed)
+    if cancel:
+        assert answer is None and live.failure == "session_observer_cancelled"
+        assert stopped.is_set()
+    else:
+        assert answer.stdout == '{"__question__":"a"}\n'
+    assert not list(tmp_path.glob("answer-*.txt"))

tests/test_engine_task_entry.py

@@ -30,7 +30,7 @@ def test_engine_cluster_heap_flags(monkeypatch):
     backend = worker.MeasuredBackend(30, 10, lambda: None, heap_profile="cluster")
     flags = (
         "-Xms1700m -Xmx1700m -XX:MaxMetaspaceSize=256m -XX:ReservedCodeCacheSize=128m "
-        "-XX:MaxDirectMemorySize=256m -Xss512k -XX:ActiveProcessorCount=1 -XX:+UseG1GC"
+        "-XX:MaxDirectMemorySize=256m -Xss512k -XX:ActiveProcessorCount=2 -XX:+UseG1GC"
     )
     assert backend._jvm_env()["JAVA_TOOL_OPTIONS"] == flags
     assert backend._heap_flag() == "-J-Xmx1700m"

tests/test_engine_trace.py

@@ -1598,3 +1598,124 @@ def test_ax_pending_fallback_resumes_in_docker_and_retains_attempt(tmp_path, mon
     progress = json.loads((tmp_path / "progress.json").read_text())
     assert progress["pins_done"] == 1 and progress["pins_in_flight"] == 0
     assert progress["fallbacks_pending"] == 0 and progress["statuses"] == {"asked-nothing": 1}
+
+
+@pytest.mark.parametrize("failure", ["", "start", "taint"])
+def test_worker_resident_session_and_fallback(tmp_path, monkeypatch, failure):
+    from types import SimpleNamespace
+
+    calls = []
+    graph = tmp_path / "cpg.bin"
+    graph.write_bytes(b"g" * 2048)
+    payloads = {
+        "overlay": {"files": "1", "methods": "1"},
+        "census": {"files": "1", "file_names": ["app.py"]},
+        "taint": {"__census__": [{"files": 1, "file_names": ["app.py"]}], "a": []},
+    }
+
+    def output(query):
+        return BEGIN + "\n" + json.dumps(payloads[query]) + "\n" + END
+
+    class Session:
+        loaded_graph = ""
+        failure = "fake failure"
+        output_observer = None
+
+        def __init__(self, scratch, budget, env, **kwargs):
+            self.scratch = scratch
+            assert "ActiveProcessorCount=2" in env["JAVA_TOOL_OPTIONS"]
+
+        def start(self):
+            calls.append("start")
+            return failure != "start"
+
+        def load(self, path, **kwargs):
+            calls.append("load")
+            assert path.read_bytes() == b"g" * 2048
+            self.loaded_graph = str(path)
+            return True
+
+        def define(self, source, **kwargs):
+            assert "  importCpg(cpgFile)" not in source
+            return ""
+
+        def evaluate(self, code, **kwargs):
+            query = code.split("(")[0].removeprefix("ousast_")
+            calls.append(query)
+            if query == "overlay":
+                saved = self.scratch / "workspace" / graph.name / "cpg.bin"
+                saved.parent.mkdir(parents=True)
+                saved.write_bytes(b"s" * 2048)
+            if query == "taint":
+                if failure == "taint":
+                    return None
+                self.output_observer('{"__question__":"a"}\n')
+                self.output_observer('{"__question__":"a"}\n{"id":"a","rows":[]}\n')
+            return SimpleNamespace(stdout=output(query))
+
+        def close(self):
+            calls.append("close")
+
+        def log_tail(self):
+            return "fake log"
+
+    monkeypatch.setattr(worker, "EngineSession", Session)
+    backend = worker.MeasuredBackend(900, 120, lambda: None, heap_profile="cluster")
+    disposable = []
+
+    def run(command, **kwargs):
+        query = Path(command[command.index("--script") + 1]).stem
+        disposable.append(query)
+        if query == "overlay":
+            saved = tmp_path / "workspace" / graph.name / "cpg.bin"
+            saved.parent.mkdir(parents=True, exist_ok=True)
+            saved.write_bytes(b"s" * 2048)
+        return subprocess.CompletedProcess(command, 0, output(query), "")
+
+    monkeypatch.setattr(backend, "_run", run)
+    monkeypatch.setattr("openultrasast.cpg.backend.shutil.which", lambda name: name)
+    assert backend._apply_overlays(graph, tmp_path)
+    assert backend.query(graph, "census", {})["files"] == "1"
+    assert backend._batch_once(graph, "taint", {"a": {}})["a"] == []
+    backend.close_session()
+    assert calls.count("start") == 1
+    if failure == "start":
+        assert disposable == ["overlay", "census", "taint"]
+        assert backend.session_fallback == "fake failure"
+    else:
+        assert calls.count("load") == 1
+        assert [stage["step"] for stage in backend.stages] == ["start", "load", "overlay", "census", "taint"]
+        assert all(stage["seconds"] >= 0 for stage in backend.stages)
+        assert backend.stages[3]["files"] == 1
+        assert disposable == (["taint"] if failure else [])
+        if not failure:
+            assert backend.asked == ["a"]
+            assert backend.answers == {"a": []}
+            assert backend.completion_counts["a"] == 1
+    assert calls.count("close") == 1
+
+
+def test_worker_session_question_deadline_keeps_markers_and_does_not_replay(tmp_path):
+    from types import SimpleNamespace
+
+    graph = tmp_path / "cpg.bin"
+    graph.write_bytes(b"g" * 2048)
+    backend = worker.MeasuredBackend(900, 0, lambda: None)
+    closed = []
+
+    def evaluate(code, **kwargs):
+        assert not fake.output_observer('{"__question__":"a"}\n')
+        return None
+
+    fake = SimpleNamespace(
+        loaded_graph=str(graph),
+        define=lambda *a, **kw: "",
+        evaluate=evaluate,
+        close=lambda: closed.append(True),
+    )
+    backend._session = fake
+    body = backend._session_payload(graph, "taint", {"cpgFile": str(graph)})
+    assert body == '{"__question__":"a"}\n'
+    assert backend.timed_out and backend.asked == ["a"]
+    assert not backend.answers and not backend.session_fallback
+    assert closed == [True]