diff --git a/src/openultrasast/hunt/artifact.py b/src/openultrasast/hunt/artifact.py index 96ae253..bf6cb90 100644 --- a/src/openultrasast/hunt/artifact.py +++ b/src/openultrasast/hunt/artifact.py @@ -77,8 +77,12 @@ class UsageLedger: def add(self, call_row: Mapping[str, Any]) -> None: replay_hits = call_row.get("replay_hits", int(bool(call_row.get("replay_hit", False)))) - values = {"usd": call_row.get("usd", 0), "calls": call_row.get("calls", 0 if replay_hits else 1), - "replay_hits": replay_hits, "wall_seconds": call_row.get("wall_seconds", 0)} + values = { + "usd": call_row.get("usd", 0), + "calls": call_row.get("calls", 0 if replay_hits else 1), + "replay_hits": replay_hits, + "wall_seconds": call_row.get("wall_seconds", 0), + } if any(not isinstance(v, int | float) or not math.isfinite(v) or v < 0 for v in values.values()): raise ValueError("usage values must be finite and non-negative") if any(int(values[k]) != values[k] for k in ("calls", "replay_hits")): diff --git a/src/openultrasast/hunt/labels.py b/src/openultrasast/hunt/labels.py index 371f036..a944c6c 100644 --- a/src/openultrasast/hunt/labels.py +++ b/src/openultrasast/hunt/labels.py @@ -21,13 +21,17 @@ def label_rows(rows: Sequence[Mapping[str, Any]], change: Mapping[str, Any]) -> data = deepcopy(dict(row)) span = change.get("function_lines") function = row.get("function", row.get("feature.source.function")) - site = kind == "vulnerability" and row.get("path") == change["path"] and ( - function == change["function"] or (span[0] <= row.get("line", 0) <= span[1]) + site = ( + kind == "vulnerability" + and row.get("path") == change["path"] + and (function == change["function"] or (span[0] <= row.get("line", 0) <= span[1])) + ) + data.update( + label=1 if site else 0 if kind == "ordinary" else None, + family=change["family"] if site else None, + fix_span=list(span) if site else None, + label_ref=change["case_id"] if site or kind == "ordinary" else None, + verification=None, ) - data.update(label=1 if site else 0 if kind == "ordinary" else None, - family=change["family"] if site else None, - fix_span=list(span) if site else None, - label_ref=change["case_id"] if site or kind == "ordinary" else None, - verification=None) result.append(data) return result diff --git a/src/openultrasast/hunt/metric.py b/src/openultrasast/hunt/metric.py index 399cc20..80e2c57 100644 --- a/src/openultrasast/hunt/metric.py +++ b/src/openultrasast/hunt/metric.py @@ -34,9 +34,11 @@ def _location_match(change: Mapping[str, Any], finding: Mapping[str, Any], toler def _on_diff(change: Mapping[str, Any], finding: Mapping[str, Any], tolerance: int) -> bool: - return any(finding.get("path") == file["path"] and - any(start - tolerance <= (finding.get("line", 0) or 0) <= end + tolerance - for start, end in file.get("head_lines", [])) for file in change.get("files", [])) + return any( + finding.get("path") == file["path"] + and any(start - tolerance <= (finding.get("line", 0) or 0) <= end + tolerance for start, end in file.get("head_lines", [])) + for file in change.get("files", []) + ) def search_outcome(change: Mapping[str, Any], findings: Sequence[Mapping[str, Any]], *, tolerance: int = 10) -> Outcome: @@ -51,11 +53,11 @@ def search_outcome(change: Mapping[str, Any], findings: Sequence[Mapping[str, An on_diff = sum(_on_diff(change, finding, tolerance) for finding in findings) feedback = [] if change["kind"] == "vulnerability" and not hit: - span = change.get("function_lines") or [next((start for f in change.get("files", []) - for start, _ in f.get("head_lines", [])), 0)] + span = change.get("function_lines") or [next((start for f in change.get("files", []) for start, _ in f.get("head_lines", [])), 0)] feedback.append(f"missed {change.get('path', '')}:{span[0]} {change.get('function', '')} family {change.get('family', '')}") - off_diff = [f"{f.get('path', '')}:{f.get('line', 0)} family {f.get('family', '')}" - for f in findings if not _on_diff(change, f, tolerance)] + off_diff = [ + f"{f.get('path', '')}:{f.get('line', 0)} family {f.get('family', '')}" for f in findings if not _on_diff(change, f, tolerance) + ] if off_diff: feedback.append("off-diff suspicions: " + ", ".join(off_diff)) return Outcome(hit, len(findings), on_diff, "; ".join(feedback)) diff --git a/src/openultrasast/hunt/producers.py b/src/openultrasast/hunt/producers.py index 289a078..4b9a8af 100644 --- a/src/openultrasast/hunt/producers.py +++ b/src/openultrasast/hunt/producers.py @@ -98,9 +98,18 @@ class Context: def hunks(self) -> dict[str, list[list[int]]]: self.verify_head() # Exclude private files before git opens the patch; only hunk coordinates become features. - patch = _git(self.root, "diff", "--no-ext-diff", "--no-textconv", "-U0", - f"{self.change['base']}..{self.change['head']}", "--", ".", - ":(glob,exclude)**/.env*", ":(exclude).env*") + patch = _git( + self.root, + "diff", + "--no-ext-diff", + "--no-textconv", + "-U0", + f"{self.change['base']}..{self.change['head']}", + "--", + ".", + ":(glob,exclude)**/.env*", + ":(exclude).env*", + ) self.diff_bytes = len(patch.encode("utf-8")) result: dict[str, list[list[int]]] = {} path = None @@ -154,8 +163,14 @@ def _producer(name: str, compute: Compute) -> Producer: bytes_read = ctx.diff_bytes elif name == "taxonomy" and reason is None: bytes_read = len(DEFAULT_FAMILIES_PATH.read_bytes()) - result.update({f"provenance.{name}.ok": reason is None, f"provenance.{name}.reason": reason, - f"provenance.{name}.version": "1", f"provenance.{name}.bytes_read": bytes_read}) + result.update( + { + f"provenance.{name}.ok": reason is None, + f"provenance.{name}.reason": reason, + f"provenance.{name}.version": "1", + f"provenance.{name}.bytes_read": bytes_read, + } + ) return result return Producer(name, "1", columns + provenance, guarded) @@ -175,10 +190,13 @@ def _source(root: Path, change: Change, path: str, line: int, ctx: Context) -> d # Select the occurrence at this line when several methods share a name. spans = [r for r in ctx.ranges(path) if r.name == facts.function and r.start <= line <= r.end] span = min(spans, key=lambda r: r.end - r.start) if spans else None - return {"function": facts.function, "span": [span.start, span.end] if span else None, - "callers": [list(caller) for caller in facts.callers], - "entry_points": [{"path": entry.path, "line": entry.line, "kind": entry.kind} for entry in facts.entry_points], - "language": facts.language} + return { + "function": facts.function, + "span": [span.start, span.end] if span else None, + "callers": [list(caller) for caller in facts.callers], + "entry_points": [{"path": entry.path, "line": entry.line, "kind": entry.kind} for entry in facts.entry_points], + "language": facts.language, + } def _roles(function: Any, facts: Any) -> dict[str, str]: @@ -192,8 +210,11 @@ def _roles(function: Any, facts: Any) -> dict[str, str]: for arguments in call.arg_names: names.update(arguments) roles = {} - for role, entries, attribute in (("source", facts.sources, "patterns"), ("sink", facts.sinks, "calls"), - ("sanitizer", facts.sanitizers, "calls")): + for role, entries, attribute in ( + ("source", facts.sources, "patterns"), + ("sink", facts.sinks, "calls"), + ("sanitizer", facts.sanitizers, "calls"), + ): for fact in entries: for pattern in getattr(fact, attribute): token = pattern.rstrip("(") @@ -211,13 +232,29 @@ def _obligations(path: str, ctx: Context) -> list[dict[str, object]]: if not scoped.operations and not scoped.dischargers: raise _Absent("obligations_language_unsupported") dominance = OrderDominance({path: text}) - result = check_obligations(irs={path: (ir, text)}, entries=ctx.entry_points, facts=facts, flow_facts=flow, - policy=None, paths=(), dominance=dominance, store_shapes=(), min_siblings=3) + result = check_obligations( + irs={path: (ir, text)}, + entries=ctx.entry_points, + facts=facts, + flow_facts=flow, + policy=None, + paths=(), + dominance=dominance, + store_shapes=(), + min_siblings=3, + ) # No finding does not imply discharge: the checker also gates on entry-point reachability. - return [{"operation": operation.kind, "line": operation.line, - "discharged": any(valid_witness(witness) and covers(witness, operation) and - dominance.dominates(witness, operation, None) for witness in result.discharges)} - for operation in find_operations(ir, scoped, text=text)] + return [ + { + "operation": operation.kind, + "line": operation.line, + "discharged": any( + valid_witness(witness) and covers(witness, operation) and dominance.dominates(witness, operation, None) + for witness in result.discharges + ), + } + for operation in find_operations(ir, scoped, text=text) + ] def _semantic(root: Path, change: Change, path: str, line: int, ctx: Context) -> dict[str, object]: @@ -228,20 +265,21 @@ def _semantic(root: Path, change: Change, path: str, line: int, ctx: Context) -> if not functions: raise _Absent("function_unavailable") function = min(functions, key=lambda fn: fn.end_line - fn.start_line) - return {"roles": _roles(function, load_facts().for_language(ir.language)), - "def_use": {"binds": [{"line": b.line, "name": b.name, "names": list(b.names), "call_name": b.call_name} - for b in function.binds], - "calls": [{"line": c.line, "name": c.name, "names": [list(names) for names in c.arg_names]} - for c in function.calls]}, - "obligations": _obligations(path, ctx)} + return { + "roles": _roles(function, load_facts().for_language(ir.language)), + "def_use": { + "binds": [{"line": b.line, "name": b.name, "names": list(b.names), "call_name": b.call_name} for b in function.binds], + "calls": [{"line": c.line, "name": c.name, "names": [list(names) for names in c.arg_names]} for c in function.calls], + }, + "obligations": _obligations(path, ctx), + } def _taxonomy(root: Path, change: Change, path: str, line: int, ctx: Context) -> dict[str, object]: return {"families": [{"id": f.id, "definition": " ".join(f.description.split())} for f in load_families().families]} -PRODUCERS = (_producer("diff", _diff), _producer("source", _source), - _producer("semantic", _semantic), _producer("taxonomy", _taxonomy)) +PRODUCERS = (_producer("diff", _diff), _producer("source", _source), _producer("semantic", _semantic), _producer("taxonomy", _taxonomy)) def _candidates(root: Path, change: Change, ctx: Context) -> list[dict[str, Any]]: @@ -259,20 +297,35 @@ def _candidates(root: Path, change: Change, ctx: Context) -> list[dict[str, Any] line = next((start for start, _ in file.get("head_lines", [])), 1) selected.add((path, max(1, line))) continue - touched = [r for r in ranges if r.kind == "function" and - any(r.start <= end and start <= r.end for start, end in file.get("head_lines", []))] + touched = [ + r + for r in ranges + if r.kind == "function" and any(r.start <= end and start <= r.end for start, end in file.get("head_lines", [])) + ] selected.update((path, r.start) for r in touched) if not touched and file.get("head_lines") and ranges: selected.add((path, max(1, file["head_lines"][0][0]))) if change.get("kind") == "vulnerability" and path == change["path"]: - labelled = [r for r in ranges if r.name == change["function"] and - r.start <= change["function_lines"][1] and change["function_lines"][0] <= r.end] + labelled = [ + r + for r in ranges + if r.name == change["function"] and r.start <= change["function_lines"][1] and change["function_lines"][0] <= r.end + ] selected.update((path, r.start) for r in labelled) if not labelled: selected.add((path, change["function_lines"][0])) - return [{"change_id": change.get("id", change.get("case_id")), "repo": change.get("repo", change.get("repository")), - "head": change["head"], "path": path, "line": line, "kind": change["kind"], "fold": change.get("fold")} - for path, line in sorted(selected)] + return [ + { + "change_id": change.get("id", change.get("case_id")), + "repo": change.get("repo", change.get("repository")), + "head": change["head"], + "path": path, + "line": line, + "kind": change["kind"], + "fold": change.get("fold"), + } + for path, line in sorted(selected) + ] def candidate_rows(root: Path, change: Change) -> list[dict[str, Any]]: diff --git a/src/openultrasast/hunt/program.py b/src/openultrasast/hunt/program.py index 14d20d7..2bc24e5 100644 --- a/src/openultrasast/hunt/program.py +++ b/src/openultrasast/hunt/program.py @@ -117,10 +117,14 @@ def shown_lines(inputs: Mapping[str, str]) -> dict[str, set[int]]: return shown -def check(findings: Iterable[Finding | Mapping[str, Any]], inputs: Mapping[str, str]) -> list[dict[str, Any]]: +def check( + findings: Iterable[Finding | Mapping[str, Any]], inputs: Mapping[str, str], seen: Mapping[str, set[int]] | None = None +) -> list[dict[str, Any]]: """Apply the answer rules: a family outside the closed list, or a cited line the model was not shown, scores against the program by becoming `unknown` with a reason. Nothing is dropped, so the metric sees it.""" shown = shown_lines(inputs) + for path, lines in (seen or {}).items(): + shown.setdefault(path, set()).update(lines) allowed = {family for family in FAMILY_IDS if family != "unknown"} out = [] for item in findings: @@ -139,22 +143,35 @@ def check(findings: Iterable[Finding | Mapping[str, Any]], inputs: Mapping[str, return out -def repository_tools(root: Path, *, max_chars: int = 4000, max_matches: int = 50) -> list[dspy.Tool]: - """The three tools the shipped hunter has, bound to one checkout and clamped to it.""" +def repository_tools( + root: Path, *, max_chars: int = 4000, max_matches: int = 50, seen: dict[str, set[int]] | None = None +) -> list[dspy.Tool]: + """The three tools the shipped hunter has, bound to one checkout and clamped to it. Every line a tool returns + is recorded in `seen` (path -> line numbers), because a line the model read through a tool was shown to it.""" root = Path(root) + seen = seen if seen is not None else {} + + def note(path: str, lines: Iterable[int]) -> None: + seen.setdefault(path, set()).update(int(line) for line in lines) def read_file(path: str) -> str: """Read a file of the repository by its relative path, numbered lines.""" - return hunter_tools.read_file(root, path, max_chars=max_chars) + text = hunter_tools.read_file(root, path, max_chars=max_chars) + note(path, range(1, text.count("\n") + 2)) + return text def grep_repo(pattern: str) -> str: """Search the repository for a regular expression; returns path:line: text, one per line.""" rows = hunter_tools.grep_repo(root, pattern, max_matches=max_matches) + for row in rows: + note(str(row.get("path")), [int(row.get("line") or 0)]) return "\n".join(f"{row.get('path')}:{row.get('line')}: {row.get('text', '')}" for row in rows) or "(no matches)" def find_refs(symbol: str) -> str: """Find where a symbol is referenced in the repository; returns path:line per reference.""" rows = hunter_tools.find_refs(root, symbol) + for row in rows: + note(str(row.get("path")), [int(row.get("line") or 0)]) return "\n".join(f"{row.get('path')}:{row.get('line')}" for row in rows) or "(no references)" return [dspy.Tool(read_file), dspy.Tool(grep_repo), dspy.Tool(find_refs)] @@ -166,13 +183,25 @@ class HuntProgram(dspy.Module): def __init__(self, *, mask: Iterable[str] = (), tools: Sequence[dspy.Tool] = (), max_iters: int = 8) -> None: super().__init__() self.mask = tuple(sorted(set(mask))) + self.seen: dict[str, set[int]] = {} signature = signature_for(self.mask) self.hunt = dspy.ReAct(signature, tools=list(tools), max_iters=max_iters) if tools else dspy.Predict(signature) + def bind(self, root: Path) -> None: + """Point the tools at one checkout and start a fresh record of what they show.""" + self.seen = {} + if isinstance(self.hunt, dspy.ReAct): + for tool in repository_tools(root, seen=self.seen): + self.hunt.tools[tool.name] = tool + def forward(self, **inputs: str) -> dspy.Prediction: - prediction = self.hunt(**inputs) - findings = check(getattr(prediction, "findings", []) or [], inputs) - return dspy.Prediction(findings=findings, raw=getattr(prediction, "findings", [])) + try: + prediction = self.hunt(**inputs) + except dspy.utils.exceptions.AdapterParseError: + # An answer the program could not shape is an empty answer that scores against it, not a broken run. + return dspy.Prediction(findings=[], raw=[], parse_failed=True) + findings = check(getattr(prediction, "findings", []) or [], inputs, self.seen) + return dspy.Prediction(findings=findings, raw=getattr(prediction, "findings", []), parse_failed=False) def program_id(program: dspy.Module, *, mask: Iterable[str] = ()) -> str: @@ -193,7 +222,9 @@ def build_lm( if not key: raise RuntimeError("DEEPSEEK_API_KEY is not set") if cache_dir is not None: - os.environ["DSPY_CACHEDIR"] = str(cache_dir) + # A run's own response cache, kept with its records so the run can be replayed at no cost. + Path(cache_dir).mkdir(parents=True, exist_ok=True) + dspy.configure_cache(enable_disk_cache=True, enable_memory_cache=True, disk_cache_dir=str(cache_dir)) try: import litellm diff --git a/src/openultrasast/hunt/runner.py b/src/openultrasast/hunt/runner.py index c67618b..034cf2e 100644 --- a/src/openultrasast/hunt/runner.py +++ b/src/openultrasast/hunt/runner.py @@ -138,11 +138,10 @@ def hunt_change( record["rows"] = len(rows) inputs = render(change, rows, texts, mask=mask) record["inputs_present"] = sorted(inputs) - if tools_for is not None: - program.hunt.tools = ( - {tool.name: tool for tool in tools_for(root)} if hasattr(program.hunt, "tools") else program.hunt.tools - ) + if tools_for is not None and hasattr(program, "bind"): + program.bind(root) prediction = program(**inputs) + record["parse_failed"] = bool(getattr(prediction, "parse_failed", False)) for finding in prediction.findings: if finding.get("path") in texts and not finding.get("function") and (span := _function_at(root, finding)): finding["function"] = span diff --git a/src/openultrasast/hunt/table.py b/src/openultrasast/hunt/table.py index 92f85f7..5ba4480 100644 --- a/src/openultrasast/hunt/table.py +++ b/src/openultrasast/hunt/table.py @@ -87,8 +87,7 @@ class FeatureTable: if unknown: raise ValueError("mask names an unknown producer or feature column") columns = {c for c in FEATURE_COLUMNS if c in names or c.split(".")[1] in names} - return FeatureTable(({k: v for k, v in row.items() if k not in columns} for row in self.rows), - masked=(*self.masked, *names)) + return FeatureTable(({k: v for k, v in row.items() if k not in columns} for row in self.rows), masked=(*self.masked, *names)) def feature_columns(self) -> list[str]: return sorted({column for row in self.rows for column in row if column.startswith("feature.")}) diff --git a/tests/test_hunt_program.py b/tests/test_hunt_program.py index d4a4f51..22352b1 100644 --- a/tests/test_hunt_program.py +++ b/tests/test_hunt_program.py @@ -95,3 +95,23 @@ def test_build_lm_needs_a_key_from_the_environment(monkeypatch): monkeypatch.delenv("DEEPSEEK_API_KEY", raising=False) with pytest.raises(RuntimeError, match="DEEPSEEK_API_KEY"): hp.build_lm() + + +def test_lines_read_through_a_tool_count_as_shown(tmp_path): + (tmp_path / "app.py").write_text("a = 1\nb = 2\nc = eval(input())\n") + seen: dict[str, set[int]] = {} + tools = {tool.name: tool for tool in hp.repository_tools(tmp_path, seen=seen)} + tools["read_file"](path="app.py") + assert seen["app.py"] >= {1, 2, 3} + inputs = hp.render_rows(CHANGE, ROWS, TEXT) + finding = {"path": "app.py", "line": 3, "family": "injection", "confidence": 0.9, "rationale": "eval", "cited_lines": [3]} + assert hp.check([finding], inputs)[0]["family"] == "unknown" # line 3 is outside the rendered span 5-7 + assert hp.check([finding], inputs, seen)[0]["family"] == "injection" # but the model read it through a tool + + +def test_an_unparseable_answer_is_an_empty_flagged_answer(): + lm = DummyLM([{"findings": "not json at all"}] * 4) + dspy.configure(lm=lm, adapter=dspy.ChatAdapter()) + program = hp.HuntProgram() + prediction = program(**hp.render_rows(CHANGE, ROWS, TEXT)) + assert prediction.findings == [] and prediction.parse_failed is True