diff --git a/src/openultrasast/learning/detectors.py b/src/openultrasast/learning/detectors.py new file mode 100644 index 0000000..362a303 --- /dev/null +++ b/src/openultrasast/learning/detectors.py @@ -0,0 +1,209 @@ +"""Per-family detector configurations and the runner that uses them (learning-harness, Req 6.1-6.5). + +A detector is a configuration, not code: a system prompt, an evidence checklist, the curated tools it +may call, its step and cost budget, and the hard negatives and counterexamples earlier rounds left it. +Round zero writes one directory per family from the same prompt, so the baseline is controlled, and a +round rewrites exactly one directory and bumps that family's version. Every other family is frozen by +construction, because a round can only touch the directory it targets. + +The runner gives a detector the whole file rather than an excerpt around a hotspot, since context loss +is a documented cause of misjudging patched code, and tags each finding with exactly one family and the +configuration version that produced it, so a finding can always be traced back to what made it. +""" + +from __future__ import annotations + +import json +import tomllib +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path + +from ..findings import StaticFinding +from ..tool_hunter import HUNTER_TOOLS, ChatClient, run_tool_hunter +from .families import FamilyTaxonomy + +MAX_PROMPT_CHARS = 8000 # prompt plus checklist; unbounded prompts are the documented signature of overfitting +GENERALIST = "unknown" # the family every region also runs, so a bug outside the taxonomy is still looked for +CURATED_TOOLS = tuple(str(schema["function"]["name"]) for schema in HUNTER_TOOLS) # type: ignore[index] +_REQUIRED = ("prompt.md", "family.toml") + + +class DetectorConfigError(ValueError): + """Raised when a detector configuration is missing, over its cap, or names something outside the closed sets.""" + + +@dataclass(frozen=True) +class FamilyConfig: + family: str + version: str + prompt: str + checklist: str + tools: tuple[str, ...] + max_steps: int + max_cost_usd: float + max_chars: int + hard_negatives: tuple[Mapping[str, object], ...] = () + counterexamples: tuple[Mapping[str, object], ...] = () + + @property + def detector_id(self) -> str: + return f"{self.family}@{self.version}" + + +@dataclass(frozen=True) +class Region: + """A place a detector is pointed at: one file, optionally one function, and the families admitted for it.""" + + path: str + function: str | None = None + families: tuple[str, ...] = () + + +def load_family_configs(configs_dir: Path, taxonomy: FamilyTaxonomy) -> dict[str, FamilyConfig]: + """Every configuration under ``configs_dir``. A directory that is not a family, or a broken one, fails loud.""" + known = {family.id for family in taxonomy.families} + configs: dict[str, FamilyConfig] = {} + for directory in sorted(path for path in configs_dir.iterdir() if path.is_dir()) if configs_dir.is_dir() else []: + if directory.name not in known: + raise DetectorConfigError(f"{directory}: {directory.name!r} is not a family; expected one of {', '.join(sorted(known))}") + configs[directory.name] = _load_one(directory) + return configs + + +def _load_one(directory: Path) -> FamilyConfig: + for name in _REQUIRED: + if not (directory / name).is_file(): + raise DetectorConfigError(f"{directory}: {name} is required") + try: + settings = tomllib.loads((directory / "family.toml").read_text(encoding="utf-8")) + except tomllib.TOMLDecodeError as exc: + raise DetectorConfigError(f"{directory / 'family.toml'}: not valid TOML: {exc}") from exc + version = str(settings.get("version", "")).strip() + if not version: + raise DetectorConfigError(f"{directory}: version is required") + prompt = (directory / "prompt.md").read_text(encoding="utf-8").strip() + checklist = _read_optional(directory / "checklist.md") + if len(prompt) + len(checklist) > MAX_PROMPT_CHARS: + raise DetectorConfigError( + f"{directory}: prompt and checklist length {len(prompt) + len(checklist)} exceeds the cap of {MAX_PROMPT_CHARS}" + ) + tools = tuple(str(name) for name in settings.get("tools", []) or ()) + unknown_tools = [name for name in tools if name not in CURATED_TOOLS] + if unknown_tools: + raise DetectorConfigError(f"{directory}: unknown tool(s) {', '.join(unknown_tools)}; expected from {', '.join(CURATED_TOOLS)}") + return FamilyConfig( + family=directory.name, + version=version, + prompt=prompt, + checklist=checklist, + tools=tools, + max_steps=int(settings.get("max_steps", 4)), + max_cost_usd=float(settings.get("max_cost_usd", 0.5)), + max_chars=int(settings.get("max_chars", 4000)), + hard_negatives=_read_jsonl(directory / "hard_negatives.jsonl"), + counterexamples=_read_jsonl(directory / "counterexamples.jsonl"), + ) + + +def run_family_detector(root: Path, region: Region, config: FamilyConfig, *, client: ChatClient, model: str) -> list[StaticFinding]: + """One family's detector over one region. Findings carry that family and the configuration that found them.""" + return run_tool_hunter( + root, + [], + client=client, + model=model, + max_steps=config.max_steps, + system_prompt=config.prompt, + user_prompt=_question(region, config), + context_files=(region.path,), + tools=config.tools, + max_chars=config.max_chars, + tags=(f"family:{config.family}", f"detector:{config.detector_id}"), + ) + + +def run_region_detectors( + root: Path, region: Region, configs: Mapping[str, FamilyConfig], *, client: ChatClient, model: str +) -> list[StaticFinding]: + """Every family the classifier admitted for this region, plus the generalist (Req 6.3).""" + wanted = [family for family in region.families if family in configs and family != GENERALIST] + if GENERALIST in configs: + wanted.append(GENERALIST) + findings: list[StaticFinding] = [] + for family in wanted: + findings.extend(run_family_detector(root, region, configs[family], client=client, model=model)) + return findings + + +def write_default_configs( + configs_dir: Path, taxonomy: FamilyTaxonomy, *, prompt: str, version: str, tools: Sequence[str] | None = None +) -> dict[str, Path]: + """Round zero: one directory per family from the same prompt, so the baseline is controlled (Req 8.1).""" + written: dict[str, Path] = {} + for family in taxonomy.families: + directory = configs_dir / family.id + directory.mkdir(parents=True, exist_ok=True) + (directory / "prompt.md").write_text(f"{prompt}\n\nFamily: {family.id}. {family.description}\n", encoding="utf-8") + (directory / "checklist.md").write_text("- name the file, the line and why the code is reachable\n", encoding="utf-8") + selected = list(tools) if tools is not None else list(CURATED_TOOLS) + (directory / "family.toml").write_text( + f'version = "{version}"\ntools = {json.dumps(selected)}\nmax_steps = 4\nmax_cost_usd = 0.5\nmax_chars = 4000\n', + encoding="utf-8", + ) + for name in ("hard_negatives.jsonl", "counterexamples.jsonl"): + if not (directory / name).exists(): + (directory / name).write_text("", encoding="utf-8") + written[family.id] = directory + return written + + +def _question(region: Region, config: FamilyConfig) -> str: + where = f"{region.path}::{region.function}" if region.function else region.path + parts = [f"Look for {config.family} problems in {where}. The whole file follows."] + if config.checklist: + parts.append(f"Evidence checklist:\n{config.checklist}") + if config.hard_negatives: + parts.append("Known non-issues, do not report these:\n" + "\n".join(_line(item) for item in config.hard_negatives)) + if config.counterexamples: + parts.append("Worked examples:\n" + "\n".join(_line(item) for item in config.counterexamples)) + parts.append("Reply with a JSON array of objects with path, line, title and rationale, after using a tool.") + return "\n\n".join(parts) + + +def _line(item: Mapping[str, object]) -> str: + return "- " + "; ".join(f"{key}: {value}" for key, value in sorted(item.items())) + + +def _read_optional(path: Path) -> str: + return path.read_text(encoding="utf-8").strip() if path.is_file() else "" + + +def _read_jsonl(path: Path) -> tuple[Mapping[str, object], ...]: + if not path.is_file(): + return () + rows: list[Mapping[str, object]] = [] + for line in path.read_text(encoding="utf-8").splitlines(): + if not line.strip(): + continue + try: + payload = json.loads(line) + except json.JSONDecodeError as exc: + raise DetectorConfigError(f"{path}: not valid JSON lines: {exc}") from exc + if isinstance(payload, dict): + rows.append(payload) + return tuple(rows) + + +__all__ = [ + "CURATED_TOOLS", + "GENERALIST", + "MAX_PROMPT_CHARS", + "DetectorConfigError", + "FamilyConfig", + "Region", + "load_family_configs", + "run_family_detector", + "run_region_detectors", + "write_default_configs", +] diff --git a/tests/test_learning_detectors.py b/tests/test_learning_detectors.py new file mode 100644 index 0000000..2e97ccf --- /dev/null +++ b/tests/test_learning_detectors.py @@ -0,0 +1,158 @@ +"""learning-harness task 2.5: per-family detector configurations and the family runner (offline).""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from openultrasast.learning.families import load_families +from openultrasast.tool_hunter import ChatResponse, ScriptedChatClient, ToolCall + +APP = ( + "from flask import request\n\n\n" + "@app.route('/books/')\n@login_required\ndef guarded(title):\n" + " return Book.query.filter_by(book_title=title).first()\n\n\n" + "@app.route('/books/any/<title>')\ndef leaky(title):\n" + " owner = request.args.get('owner')\n" + " return Book.query.filter_by(user_id=owner, book_title=title).first()\n" +) +REPLY = '[{"path": "app.py", "line": 12, "title": "unconstrained read", "rationale": "no owner check"}]' + + +@pytest.fixture +def repo(tmp_path: Path) -> Path: + (tmp_path / "app.py").write_text(APP) + return tmp_path + + +def _write_config(root: Path, family: str, **overrides: object) -> Path: + directory = root / family + directory.mkdir(parents=True, exist_ok=True) + (directory / "prompt.md").write_text(str(overrides.get("prompt", f"You hunt {family} only."))) + (directory / "checklist.md").write_text(str(overrides.get("checklist", "- did the handler check ownership?"))) + tools = overrides.get("tools", ["read_definition", "obligations"]) + (directory / "family.toml").write_text( + f'version = "{overrides.get("version", "1")}"\ntools = {json.dumps(tools)}\n' + f"max_steps = {overrides.get('max_steps', 3)}\nmax_cost_usd = {overrides.get('max_cost_usd', 0.5)}\n" + f"max_chars = {overrides.get('max_chars', 4000)}\n" + ) + (directory / "hard_negatives.jsonl").write_text(json.dumps({"note": "a guarded handler is not a finding"}) + "\n") + (directory / "counterexamples.jsonl").write_text("") + return directory + + +def test_a_configuration_loads_every_part_it_needs(tmp_path: Path) -> None: + from openultrasast.learning.detectors import load_family_configs + + _write_config(tmp_path, "access_control") + configs = load_family_configs(tmp_path, load_families()) + config = configs["access_control"] + assert config.family == "access_control" and config.version == "1" + assert config.prompt.startswith("You hunt access_control") and "ownership" in config.checklist + assert config.tools == ("read_definition", "obligations") + assert config.max_steps == 3 and config.max_cost_usd == 0.5 and config.max_chars == 4000 + assert config.hard_negatives[0]["note"].startswith("a guarded handler") + assert config.counterexamples == () + assert load_family_configs(tmp_path, load_families())["access_control"] == config # loading twice is the same config + + +@pytest.mark.parametrize( + ("overrides", "needle"), + [ + ({"prompt": "x" * 9000}, "length"), + ({"tools": ["read_file", "run_shell"]}, "run_shell"), + ({"version": ""}, "version"), + ], + ids=["over_cap_prompt", "unknown_tool", "no_version"], +) +def test_the_loader_refuses_a_broken_configuration_by_name(tmp_path: Path, overrides: dict[str, object], needle: str) -> None: + from openultrasast.learning.detectors import DetectorConfigError, load_family_configs + + _write_config(tmp_path, "access_control", **overrides) + with pytest.raises(DetectorConfigError, match=needle): + load_family_configs(tmp_path, load_families()) + + +def test_a_directory_that_is_not_a_family_is_refused_by_name(tmp_path: Path) -> None: + from openultrasast.learning.detectors import DetectorConfigError, load_family_configs + + _write_config(tmp_path, "access_control") + (tmp_path / "sql_injection").mkdir() + (tmp_path / "sql_injection" / "prompt.md").write_text("x") + with pytest.raises(DetectorConfigError, match="sql_injection"): + load_family_configs(tmp_path, load_families()) + + +def test_the_runner_uses_the_family_prompt_tools_and_tags(repo: Path, tmp_path: Path) -> None: + from openultrasast.learning.detectors import Region, load_family_configs, run_family_detector + + configs_dir = tmp_path / "configs" + _write_config(configs_dir, "access_control") + config = load_family_configs(configs_dir, load_families())["access_control"] + client = ScriptedChatClient( + [ + ChatResponse(tool_calls=(ToolCall(id="1", name="obligations", arguments={"path": "app.py", "function": "leaky"}),)), + ChatResponse(content=REPLY), + ] + ) + findings = run_family_detector(repo, Region(path="app.py", function="leaky"), config, client=client, model="m") + assert [finding.tags for finding in findings] == [["family:access_control", "detector:access_control@1"]] + assert findings[0].evidence_level == "suspicion" + call = client.calls[0] + assert call["messages"][0]["content"] == "You hunt access_control only." + assert [tool["function"]["name"] for tool in call["tools"]] == ["read_definition", "obligations"] + user = str(call["messages"][1]["content"]) + assert "def leaky" in user and "def guarded" in user # the whole file, not an excerpt around a hotspot + assert "ownership" in user and "a guarded handler is not a finding" in user + + +def test_a_region_runs_every_admitted_family_plus_the_generalist(repo: Path, tmp_path: Path) -> None: + from openultrasast.learning.detectors import Region, load_family_configs, run_region_detectors + + configs_dir = tmp_path / "configs" + for family in ("access_control", "injection", "unknown"): + _write_config(configs_dir, family, tools=["read_file"]) + configs = load_family_configs(configs_dir, load_families()) + client = ScriptedChatClient( + [ + ChatResponse(tool_calls=(ToolCall(id="1", name="read_file", arguments={"path": "app.py"}),)), + ChatResponse(content=REPLY), + ] + * 4 + ) + region = Region(path="app.py", function="leaky", families=("access_control",)) + findings = run_region_detectors(repo, region, configs, client=client, model="m") + families = sorted({tag for finding in findings for tag in finding.tags if tag.startswith("family:")}) + assert families == ["family:access_control", "family:unknown"] # the admitted family and the generalist, not injection + assert len(client.calls) == 4 # two turns each, and nothing else was run + + +def test_the_runner_never_offers_a_tool_outside_the_curated_set(repo: Path, tmp_path: Path) -> None: + from openultrasast.learning.detectors import Region, load_family_configs, run_family_detector + from openultrasast.tool_hunter import HUNTER_TOOLS + + configs_dir = tmp_path / "configs" + _write_config( + configs_dir, "injection", tools=["read_file", "grep_repo", "find_refs", "read_definition", "entry_points", "flows", "obligations"] + ) + config = load_family_configs(configs_dir, load_families())["injection"] + client = ScriptedChatClient([ChatResponse(content="")]) + run_family_detector(repo, Region(path="app.py", function="leaky"), config, client=client, model="m") + offered = {tool["function"]["name"] for tool in client.calls[0]["tools"]} + assert offered == {str(tool["function"]["name"]) for tool in HUNTER_TOOLS} + assert "bash" not in offered and "run_shell" not in offered + + +def test_round_zero_can_write_a_configuration_for_every_family(tmp_path: Path) -> None: + from openultrasast.learning.detectors import load_family_configs, write_default_configs + + taxonomy = load_families() + written = write_default_configs(tmp_path, taxonomy, prompt="You are a security hunter.", version="0") + assert set(written) == {family.id for family in taxonomy.families} + configs = load_family_configs(tmp_path, taxonomy) + assert set(configs) == {family.id for family in taxonomy.families} + assert all(config.version == "0" for config in configs.values()) + assert all(config.prompt.startswith("You are a security hunter.") for config in configs.values()) + assert configs["access_control"].prompt != configs["injection"].prompt # each carries its own family's definition