search-with-proof: family-to-oracle mapping, build memory, tolerant tool calls

Demos declare their oracle within the family allowed set (injection: sql or
command; output_encoding: xss; untrusted_destination: ssrf; path: path);
builds and app starts use the search memory budget instead of the 256 MB tool
limit (npm died with SIGTRAP); malformed tool-call JSON is fed back once.

Co-Authored-By: Claude Fable 5.1 <[email protected]>

tests/test_search_board.py

@@ -65,7 +65,7 @@ def test_snapshot_caps_fact_text_and_preserves_context(tmp_path):
         ],
     )
     view = yaml.safe_load(board.snapshot(budget_left={"tasks": 3}))
-    assert view["candidate"]["file"] == "app.py" and view["family"] == "sql"
+    assert view["candidate"]["file"] == "app.py" and view["family"] == "injection"
     assert view["budget_left"]["tasks"] == 3
     assert {i["status"] for i in view["intents"]} == {"open", "concluded"}
     assert all(len(f["text"]) <= 2000 for f in view["facts"])

tests/test_search_demo.py

@@ -7,6 +7,7 @@ from openultrasast.search.demo import load_demo, validate_demo
 
 def document():
     return {
+        "oracle": "path",
         "build": {"recipe": "none", "arguments": []},
         "start": {"runtime": "python", "path": "app.py", "arguments": [], "mode": "cli"},
         "steps": [{"type": "cli", "arguments": ["hello"]}],
@@ -75,3 +76,38 @@ def test_manifest_recipe_requires_file_path(recipe):
     value["build"] = {"recipe": recipe, "arguments": ["."]}
     with pytest.raises(ValueError):
         validate_demo(value)
+
+
+@pytest.mark.parametrize("oracle", [None, "injection", "unknown", [], 1])
+def test_oracle_required_and_known(oracle):
+    value = document()
+    if oracle is None:
+        value.pop("oracle")
+    else:
+        value["oracle"] = oracle
+    with pytest.raises(ValueError):
+        validate_demo(value)
+
+
+@pytest.mark.parametrize(
+    "family,allowed",
+    [
+        ("injection", ("sql", "command")),
+        ("path", ("path",)),
+        ("output_encoding", ("xss",)),
+        ("untrusted_destination", ("ssrf",)),
+        ("deserialization", ()),
+        ("access_control", ()),
+        ("config_secrets", ()),
+    ],
+)
+def test_family_oracle_contract(family, allowed):
+    from openultrasast.search.demo import FAMILY_ORACLES, validate_oracle
+
+    assert FAMILY_ORACLES[family] == allowed
+    for oracle in ("sql", "command", "path", "xss", "ssrf"):
+        if oracle in allowed:
+            validate_oracle(family, oracle)
+        else:
+            with pytest.raises(ValueError, match="not allowed"):
+                validate_oracle(family, oracle)

tests/test_search_pilot_run.py

@@ -247,7 +247,7 @@ def test_no_oracle_does_not_acquire_or_build(tmp_path, monkeypatch):
     pilot = make_pilot(tmp_path)
     _, demos = materialise(tmp_path / "probe", "path")
     (tmp_path / "demo/demo.json").write_bytes((demos["real"] / "demo.json").read_bytes())
-    pilot.family = "injection"
+    pilot.family = "deserialization"
     monkeypatch.setattr(pilot, "prepare", lambda *args: pytest.fail("no-oracle must not build"))
     assert pilot.verification().outcome == "no_oracle"
     assert pilot.record["tasks_submitted"] == 0
@@ -467,3 +467,16 @@ def test_live_driver_stops_remaining_pairs_on_infrastructure_failure(tmp_path, m
     assert summary["end_reason"] == "instrument_failure"
     assert summary["aborted_reason"] == search["aborted_reason"]
     assert "running=0" in summary["aborted_reason"]
+
+
+@pytest.mark.parametrize("oracle", ["sql", "command"])
+def test_injection_pilot_uses_demo_choice_without_manifest_subtype(tmp_path, oracle):
+    from openultrasast.search.probe import materialise
+
+    pilot = make_pilot(tmp_path, side="vulnerable", dry_run=True)
+    pilot.family = "injection"
+    pilot.row.pop("oracle", None)
+    pilot.sides, demos = materialise(tmp_path / "probe", oracle)
+    (tmp_path / "demo/demo.json").write_bytes((demos["real"] / "demo.json").read_bytes())
+    result = pilot.verification()
+    assert result.outcome == "demonstrated", result

tests/test_search_task_storage.py

@@ -16,6 +16,7 @@ from openultrasast.search.verify import Side, verify_side_task
 def spec():
     return {
         "demo": {
+            "oracle": "path",
             "build": {"recipe": "none", "arguments": []},
             "start": {"runtime": "python", "path": "app.py", "arguments": [], "mode": "cli"},
             "steps": [{"type": "cli", "arguments": []}],
@@ -281,3 +282,38 @@ def test_preparation_with_many_host_processes(workspace, monkeypatch):
             os.kill(pid, signal.SIGKILL)
         for pid in children:
             os.waitpid(pid, 0)
+
+
+@pytest.mark.parametrize("memory_bytes", [None, 256 * 1024**2, 5 * 1024**3])
+def test_prepare_build_virtual_memory(workspace, monkeypatch, memory_bytes):
+    from openultrasast.search.demo import BUILD_RECIPES
+
+    repo = workspace / "checkout"
+    repo.mkdir()
+    source = "import mmap; region=mmap.mmap(-1,512*1024**2); print(len(region))"
+    (repo / "app.py").write_text(source)
+    print("input_bytes=", len((repo / "app.py").read_bytes()))
+    monkeypatch.setitem(BUILD_RECIPES, "npm", (sys.executable, "-c", source))
+    app = executor.InProcessExecutor(repo, task_boundary=True)
+    app.prepared_put_url = "https://files.example/output"
+    blobs = []
+    monkeypatch.setattr(executor.URLTransport, "put", lambda self, url, data, timeout: blobs.append(data))
+    request = spec()
+    request["demo"]["build"] = {"recipe": "npm", "arguments": ["."]}
+    limits = {} if memory_bytes is None else {"memory_bytes": memory_bytes}
+    result = app.submit("prepare_verification", request, limits=limits)
+    if memory_bytes == 256 * 1024**2:
+        assert result["status"] == "could_not_build" and not blobs
+    else:
+        assert result["status"] == "ok" and blobs, result
+
+
+def test_verifier_app_virtual_memory(workspace):
+    repo = workspace / "checkout"
+    repo.mkdir()
+    (repo / "app.py").write_text("import mmap; region=mmap.mmap(-1,512*1024**2); print(len(region))")
+    print("input_bytes=", len((repo / "app.py").read_bytes()))
+    demo = workspace / "demo.json"
+    demo.write_text(json.dumps(spec()["demo"]))
+    result = verify_side_task(Side(repo), demo, "path")
+    assert result.outcome == "observed", result

tests/test_search_verify.py

@@ -33,7 +33,9 @@ def test_real_differential(tmp_path, family):
             pytest.skip("host policy prohibits localhost sockets")
     sides, demos = materialise(tmp_path / family, family)
     assert (sides[0].checkout / "app.py").stat().st_size > 0
-    result = verify(*sides, demos["real"], family)
+    result = verify(
+        *sides, demos["real"], {"sql": "injection", "command": "injection", "ssrf": "untrusted_destination"}.get(family, family)
+    )
     assert result.outcome == "demonstrated", result
     assert [r.observed for r in result.sides[0].runs] == [True] * 3
     assert [r.observed for r in result.sides[1].runs] == [False] * 3
@@ -79,7 +81,7 @@ def test_readiness_failure(pair):
     assert result.sides[1].outcome == "could_not_run"
 
 
-@pytest.mark.parametrize("family", ["xss", "deserialisation", "access_control", "secrets"])
+@pytest.mark.parametrize("family", ["deserialization", "access_control", "config_secrets"])
 def test_no_oracle(pair, family):
     sides, demo = pair
     assert verify(*sides, demo, family).outcome == "no_oracle"
@@ -472,3 +474,43 @@ def test_http_failure_keeps_exception_tail(tmp_path):
             lambda *a, **kw: SandboxResult(1, "", traceback, False),
             lambda: None,
         )
+
+
+@pytest.mark.parametrize(
+    "family,oracle",
+    [
+        ("injection", "sql"),
+        ("injection", "command"),
+        ("path", "path"),
+        ("output_encoding", "xss"),
+        ("untrusted_destination", "ssrf"),
+    ],
+)
+def test_verifier_routes_demo_oracle(pair, monkeypatch, family, oracle):
+    from openultrasast.search import verify as module
+    from openultrasast.search.oracles import BrowserExecutor
+
+    sides, demo = pair
+    edit(demo, lambda value: value.update(oracle=oracle))
+    seen = []
+    monkeypatch.setattr(_sandbox, "isolation_check", lambda **kwargs: None)
+
+    def side(checkout, artefact, kind, *args):
+        seen.append(kind)
+        return SideRecord("could_not_run", (), 0, "stub")
+
+    monkeypatch.setattr(module, "_side", side)
+    result = verify(*sides, demo, family, oracle=oracle, browser=BrowserExecutor() if oracle == "xss" else None)
+    assert seen == [oracle, oracle] and result.outcome != "no_oracle"
+    edit(demo, lambda value: value.update(oracle="path" if oracle != "path" else "sql"))
+    seen.clear()
+    result = verify(*sides, demo, family)
+    assert result.outcome == "inconclusive" and not seen
+    assert "not allowed" in result.reason
+
+
+def test_xss_without_browser_is_unavailable_not_no_oracle(pair):
+    sides, demo = pair
+    edit(demo, lambda value: value.update(oracle="xss"))
+    result = verify(*sides, demo, "output_encoding")
+    assert result.outcome == "inconclusive" and "runtime unavailable" in result.reason

tests/test_search_verify_transport.py

@@ -17,6 +17,7 @@ from openultrasast.search.verify import RunObservation, Side, SideRecord
 
 def demo():
     return {
+        "oracle": "sql",
         "build": {"recipe": "none", "arguments": []},
         "start": {"runtime": "python", "path": "app.py", "arguments": [], "mode": "cli"},
         "steps": [{"type": "cli", "arguments": []}],
@@ -43,7 +44,7 @@ def test_dispatch_uses_fresh_tasks_and_only_data(tmp_path):
     dispatch = transport.AXSideDispatcher(
         SimpleNamespace(), "image@sha256:" + "a" * 64, tmp_path / "out", lane=lane, prepare=lambda side, spec: b"executor-built-archive"
     )
-    result = dispatch(side=Side(root), demo=demo(), family="sql", timeout_seconds=2, fresh_task=True)
+    result = dispatch(side=Side(root), demo=demo(), family="injection", timeout_seconds=2, fresh_task=True)
     assert result.isolation_mode == "task-boundary"
     assert len(result.runs) == 3
     assert len(calls) == 3 and calls[0].id != calls[1].id
@@ -74,7 +75,7 @@ def test_entry_reads_input_and_scrubs_environment(tmp_path, monkeypatch, capsys)
 
     shutil.copytree(root, bundle / "checkout")
     (bundle / "products").mkdir()
-    (bundle / "spec.json").write_text(json.dumps({"demo": demo(), "family": "sql", "timeout_seconds": 2}))
+    (bundle / "spec.json").write_text(json.dumps({"demo": demo(), "family": "injection", "timeout_seconds": 2}))
     blob = entry.pack_checkout(bundle)
     monkeypatch.setattr(entry, "download", lambda url, limit: blob)
     outputs = []
@@ -151,7 +152,7 @@ def test_production_item_targets_packaged_module(tmp_path):
         dispatcher = transport.AXSideDispatcher(
             SimpleNamespace(), "image@sha256:" + "a" * 64, tmp_path / "out", prepare=lambda side, spec: b"built"
         )
-        dispatcher(side=Side(root), demo=demo(), family="sql", timeout_seconds=2, fresh_task=True)
+        dispatcher(side=Side(root), demo=demo(), family="injection", timeout_seconds=2, fresh_task=True)
     assert captured[0].url_env == frozenset({"VERIFY_INPUT_URL", "RESULT_URL"})
 
 
@@ -198,7 +199,7 @@ def test_default_dispatch_builds_in_executor_before_verifier(tmp_path, monkeypat
 
     monkeypatch.setattr(transport, "SearchExecutorTask", Executor)
     dispatch = transport.AXSideDispatcher(SimpleNamespace(), "image@sha256:" + "a" * 64, tmp_path / "out", lane=lane)
-    assert len(dispatch(side=Side(root), demo=demo(), family="sql", timeout_seconds=2, fresh_task=True).runs) == 3
+    assert len(dispatch(side=Side(root), demo=demo(), family="injection", timeout_seconds=2, fresh_task=True).runs) == 3
     assert events[3:] == ["verify"] * 3
 
 
@@ -214,3 +215,23 @@ def test_source_archive_omits_git_metadata(tmp_path):
     unpack_checkout(archive, target)
     assert (target / "app.py").read_text() == "print('source')"
     assert not (target / ".git").exists()
+
+
+@pytest.mark.parametrize(
+    "family,oracle",
+    [
+        ("injection", "sql"),
+        ("injection", "command"),
+        ("path", "path"),
+        ("output_encoding", "xss"),
+        ("untrusted_destination", "ssrf"),
+    ],
+)
+def test_transport_checks_family_and_demo_oracle(family, oracle):
+    schema = demo()
+    schema["oracle"] = oracle
+    spec = {"demo": schema, "family": family, "timeout_seconds": 2}
+    assert entry.validate_spec(spec) == spec
+    schema["oracle"] = "path" if oracle != "path" else "sql"
+    with pytest.raises(ValueError, match="not allowed"):
+        entry.validate_spec(spec)

tests/test_search_worker.py

@@ -85,7 +85,8 @@ def test_coordinator_stub_smoke_real_verify(tmp_path):
     def executor(job):
         if job.step != "verify":
             return worker(job)
-        result = verify(*sides, demo, "path")
+        assert job.payload["oracle"] == "path"
+        result = verify(*sides, demo, "path", oracle=job.payload["oracle"])
         outcomes.append(result.outcome)
         return {"status": "ok", "outcome": result.outcome, "evidence_refs": ["verify:local"], "cost_usd": 0}
 
@@ -198,6 +199,7 @@ def test_brain_refuses_repository_tools_without_executor(tmp_path):
 def test_write_demo_accepts_only_declarative_json(tmp_path):
     worker = Worker(StubModel([]), repo=tmp_path, demo=tmp_path / "demo")
     schema = {
+        "oracle": "path",
         "build": {"recipe": "none", "arguments": []},
         "start": {"runtime": "python", "path": "app.py", "arguments": [], "mode": "cli"},
         "steps": [{"type": "cli", "arguments": ["hello"]}],
@@ -386,3 +388,58 @@ def test_worker_exhaustion_is_structured(tmp_path, monkeypatch, detail):
     result = worker(task("explore"))
     assert result["budget_exhaustion"]["end_detail"] == detail
     assert result["end_detail"] == detail
+
+
+def malformed_tool(arguments="{"):
+    response = tool("read_file")
+    response["tool_calls"][0]["function"]["arguments"] = arguments
+    return response
+
+
+def test_malformed_arguments_recover_with_tool_result(tmp_path):
+    (tmp_path / "input").write_text("readable")
+    model = StubModel(
+        [
+            malformed_tool(),
+            tool("read_file", path="input"),
+            tool("finish", fact={"text": "read input", "evidence_refs": ["tool:1"]}),
+        ]
+    )
+    result = Worker(model, repo=tmp_path, demo=tmp_path / "demo", executor=InProcessExecutor(tmp_path))(task("explore"))
+    assert result["status"] == "ok" and result["fact"]["text"] == "read input"
+    feedback = model.calls[1]["messages"][-1]
+    assert feedback["role"] == "tool" and feedback["tool_call_id"] == "call-1"
+    assert json.loads(feedback["content"])["error"].startswith("invalid JSON arguments: ")
+    assert len(model.calls) == 3
+
+
+def test_two_malformed_arguments_record_failure(tmp_path):
+    model = StubModel([malformed_tool(), malformed_tool()])
+    result = Worker(model, repo=tmp_path, demo=tmp_path / "demo")(task("explore"))
+    assert result["status"] == "execution_failure"
+    assert "two consecutive malformed tool calls" in result["failure"]
+    assert len(model.calls) == 2
+
+
+def test_valid_call_resets_malformed_streak(tmp_path):
+    model = StubModel([malformed_tool(), tool("list_files"), malformed_tool(), tool("finish", failure="done")])
+    result = Worker(model, repo=tmp_path, demo=tmp_path / "demo", executor=InProcessExecutor(tmp_path))(task("explore"))
+    assert result["failure"] == "done" and len(model.calls) == 4
+
+
+def test_malformed_call_consumes_step(tmp_path):
+    model = StubModel([malformed_tool(), tool("finish", failure="unused")])
+    result = Worker(model, repo=tmp_path, demo=tmp_path / "demo", max_steps=1)(task("explore"))
+    assert result["end_detail"] == "worker_steps" and len(model.calls) == 1
+
+
+def test_worker_run_uses_search_memory_budget(tmp_path):
+    class Recorder:
+        def submit(self, name, args, **kwargs):
+            return kwargs["limits"]
+
+    worker = Worker(StubModel([]), repo=tmp_path, demo=tmp_path / "demo", executor=Recorder())
+    assert worker.tool("run", {"command": ["node", "app.js"]})["memory_bytes"] == SearchBudget().memory_bytes
+    assert worker.tool("grep", {"pattern": "needle"})["memory_bytes"] == 256 * 1024**2
+    worker.limits["memory_bytes"] = 128 * 1024**2
+    assert worker.tool("read_file", {"path": "app.py"})["memory_bytes"] == 128 * 1024**2