diff --git a/core/harness/classifier.py b/core/harness/classifier.py new file mode 100644 index 00000000..ee3b4213 --- /dev/null +++ b/core/harness/classifier.py @@ -0,0 +1,255 @@ +"""LLM risk classifier for the permission gate (P0-1, Claude Code Auto-mode +lesson). + +The static permission engine (:mod:`core.harness.permissions`) plus hooks can +resolve most tool calls, but an ``ask`` verdict still falls through to a human +approver. In non-interactive runs there is no approver, so every ``ask`` is +denied — the agent stalls on anything the rules didn't anticipate. Claude +Code's Auto mode solves this with an LLM safety classifier that scores each +action and only escalates the genuinely risky ones. This module ports that +idea as an *optional, pluggable* layer: + +* It sits between the PermissionRequest hook and the human approver (in + ``AgentRunSpec``'s approval path), exactly where an ``ask`` would otherwise + block. +* Only ever upgrades an ``ask``: ``low`` risk → auto-allow; ``medium`` / + ``high`` → fall through to the human (or denial in headless runs). It never + turns an ``allow`` into a ``deny`` and never acts when disabled. +* Failure is fail-open to the human path — a classifier crash or timeout must + not escalate risk, only fall back to the existing behavior. +* Off by default: enabled with ``DEEPCODE_RISK_CLASSIFIER=1`` (model via + ``DEEPCODE_RISK_CLASSIFIER_MODEL``); absent provider support degrades to + "no verdict" (human path). + +Design rule (mirrors ``core.harness.__init__``): this module is pure +mechanism — it scores an action and returns a verdict; it never prompts, +executes, or talks to UIs. Enforcement stays in the kernel. +""" + +from __future__ import annotations + +import json +import os +import re +from dataclasses import dataclass +from enum import Enum +from typing import Any + +from loguru import logger + +# --------------------------------------------------------------------------- +# Public surface +# --------------------------------------------------------------------------- + + +class RiskLevel(str, Enum): + """Tri-state verdict the classifier returns.""" + + LOW = "low" # auto-allow the ask + MEDIUM = "medium" # fall through to the human approver + HIGH = "high" # fall through to the human approver (and be loud about it) + + +@dataclass(slots=True) +class RiskVerdict: + """One classification result.""" + + level: RiskLevel + reason: str = "" + model: str | None = None + latency_ms: float | None = None + error: str | None = None + + @property + def decisive(self) -> bool: + """A verdict is decisive only when the classifier actually ran.""" + return self.error is None and self.level is not None + + +# --------------------------------------------------------------------------- +# Prompt +# --------------------------------------------------------------------------- + +_SYSTEM_PROMPT = ( + "You are a conservative safety classifier for a coding agent's tool calls. " + "Given a tool call that the agent wants to perform, decide how risky it is. " + 'Respond with ONLY a JSON object: {"risk": "low"|"medium"|"high", ' + '"reason": ""}.\n' + "Guidance:\n" + "- low: routine, reversible, scoped to the workspace (e.g. editing a file " + "inside the project, running a normal test command).\n" + "- medium: modifies outside the workspace, installs packages, changes " + "system state, or has side effects that are not obviously reversible.\n" + "- high: destructive, exfiltrates data, touches credentials, network " + "writes, removes files, or anything a careful engineer would double-check " + "before approving.\n" + "When unsure, prefer medium over low. Never answer with anything but JSON." +) + +_USER_TEMPLATE = ( + "Tool call to classify:\n" + "tool: {tool_name}\n" + "arguments: {arguments}\n" + "policy note (why this needs confirmation): {reason}" +) + + +# --------------------------------------------------------------------------- +# Verdict parsing +# --------------------------------------------------------------------------- + +_JSON_OBJECT_RE = re.compile(r"\{.*\}", re.DOTALL) + + +def parse_risk_verdict(text: str | None) -> RiskVerdict | None: + """Parse the model's JSON reply into a :class:`RiskVerdict`. + + Tolerates markdown fences and stray prose around the JSON object. + Returns ``None`` when the reply cannot be parsed — callers treat that as + "no verdict" (fail open to the human path). + """ + if not text: + return None + match = _JSON_OBJECT_RE.search(text) + if not match: + return None + try: + payload = json.loads(match.group(0)) + except json.JSONDecodeError: + # Best-effort: some models wrap keys and string values in single + # quotes. Normalize single-quoted strings to double quotes without + # touching escaped quotes inside. + try: + body = match.group(0) + cleaned = re.sub( + r"'((?:[^'\\]|\\.)*)'", + lambda m: '"' + m.group(1).replace('"', '\\"') + '"', + body, + ) + payload = json.loads(cleaned) + except json.JSONDecodeError: + return None + if not isinstance(payload, dict): + return None + raw_risk = str(payload.get("risk", "")).strip().lower() + if raw_risk not in {level.value for level in RiskLevel}: + return None + reason = str(payload.get("reason", "")).strip() + return RiskVerdict(level=RiskLevel(raw_risk), reason=reason[:300]) + + +# --------------------------------------------------------------------------- +# Classifier +# --------------------------------------------------------------------------- + + +def classifier_enabled() -> bool: + """Whether the risk classifier is on (env: ``DEEPCODE_RISK_CLASSIFIER``).""" + value = os.environ.get("DEEPCODE_RISK_CLASSIFIER", "").strip().lower() + return value in {"1", "true", "yes", "on"} + + +def classifier_model() -> str | None: + """Optional explicit model for the classifier (env: + ``DEEPCODE_RISK_CLASSIFIER_MODEL``).""" + value = os.environ.get("DEEPCODE_RISK_CLASSIFIER_MODEL", "").strip() + return value or None + + +class LLMRiskClassifier: + """Score an ``ask``-level tool call with a lightweight LLM. + + Parameters + ---------- + provider: + Any object with ``async chat(messages, model=..., max_tokens=...)`` + returning an ``LLMResponse`` (the ``core.providers`` base interface). + model: + Optional model override; defaults to the provider's own default. + max_tokens: + Tiny budget — a classifier needs a short JSON answer. + timeout_s: + Per-call timeout; on expiry the classifier yields "no verdict". + """ + + def __init__( + self, + provider: Any, + *, + model: str | None = None, + max_tokens: int = 128, + timeout_s: float = 15.0, + ) -> None: + self._provider = provider + self._model = model or classifier_model() + self._max_tokens = max_tokens + self._timeout_s = timeout_s + + async def classify( + self, + tool_name: str, + arguments: dict[str, Any] | None, + reason: str, + ) -> RiskVerdict: + """Score one tool call; never raises, always returns a verdict.""" + import time + + started = time.perf_counter() + try: + user = _USER_TEMPLATE.format( + tool_name=tool_name, + arguments=json.dumps(arguments or {}, ensure_ascii=False)[:2000], + reason=(reason or "")[:500], + ) + response = await asyncio_wait_for( + self._provider.chat( + [ + {"role": "system", "content": _SYSTEM_PROMPT}, + {"role": "user", "content": user}, + ], + model=self._model, + max_tokens=self._max_tokens, + temperature=0.0, + ), + timeout=self._timeout_s, + ) + except Exception as exc: # noqa: BLE001 - provider down, timeout, anything + logger.opt(exception=False).warning( + "risk classifier failed for {}: {}", tool_name, exc + ) + return RiskVerdict( + level=RiskLevel.MEDIUM, + error=str(exc)[:200], + latency_ms=(time.perf_counter() - started) * 1000, + ) + + verdict = parse_risk_verdict(response.content) + latency = (time.perf_counter() - started) * 1000 + if verdict is None: + return RiskVerdict( + level=RiskLevel.MEDIUM, + error="unparseable classifier reply", + model=self._model, + latency_ms=latency, + ) + verdict.model = self._model + verdict.latency_ms = latency + return verdict + + +def asyncio_wait_for(awaitable: Any, *, timeout: float) -> Any: + """Small indirection so the module is importable without asyncio quirks + in sync contexts (the awaitable is only awaited here).""" + import asyncio + + return asyncio.wait_for(awaitable, timeout=timeout) + + +__all__ = [ + "LLMRiskClassifier", + "RiskLevel", + "RiskVerdict", + "classifier_enabled", + "classifier_model", + "parse_risk_verdict", +] diff --git a/tests/test_harness_classifier.py b/tests/test_harness_classifier.py new file mode 100644 index 00000000..e2def567 --- /dev/null +++ b/tests/test_harness_classifier.py @@ -0,0 +1,172 @@ +"""Tests for the P0-1 LLM risk classifier (Claude Code Auto-mode lesson).""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +ROOT = Path(__file__).resolve().parents[1] +if str(ROOT) not in sys.path: + sys.path.insert(0, str(ROOT)) + +from core.harness.classifier import ( + LLMRiskClassifier, + RiskLevel, + classifier_enabled, + classifier_model, + parse_risk_verdict, +) + +LOW = RiskLevel.LOW +MEDIUM = RiskLevel.MEDIUM +HIGH = RiskLevel.HIGH + + +class FakeProvider: + """Minimal provider double matching the ``chat()`` base interface.""" + + def __init__(self, content: str | None = None, error: Exception | None = None): + self._content = content + self._error = error + self.calls: list[dict] = [] + + async def chat(self, messages, model=None, max_tokens=0, temperature=0.0, **kwargs): + self.calls.append( + { + "messages": messages, + "model": model, + "max_tokens": max_tokens, + "temperature": temperature, + } + ) + if self._error is not None: + raise self._error + from core.providers.base import LLMResponse + + return LLMResponse(content=self._content) + + +# ---- verdict parsing --------------------------------------------------------- + + +def test_parse_plain_json(): + v = parse_risk_verdict('{"risk": "low", "reason": "edits workspace file"}') + assert v is not None and v.level is LOW and "workspace" in v.reason + + +def test_parse_markdown_fenced_json(): + v = parse_risk_verdict('```json\n{"risk": "high", "reason": "rm -rf /"}\n```') + assert v is not None and v.level is HIGH + + +def test_parse_stray_prose_around_json(): + v = parse_risk_verdict( + 'Here you go: {"risk": "medium", "reason": "installs a dep"}' + ) + assert v is not None and v.level is MEDIUM + + +def test_parse_single_quoted_keys(): + v = parse_risk_verdict("{'risk': 'low', 'reason': 'fine'}") + assert v is not None and v.level is LOW + + +def test_parse_rejects_garbage(): + assert parse_risk_verdict("I don't know") is None + assert parse_risk_verdict(None) is None + assert parse_risk_verdict('{"risk": "nuclear", "reason": "x"}') is None + assert parse_risk_verdict("[1,2,3]") is None + + +# ---- classifier behavior ---------------------------------------------------- + + +@pytest.mark.asyncio +async def test_classify_low_auto_allows(): + provider = FakeProvider('{"risk": "low", "reason": "in-workspace edit"}') + clf = LLMRiskClassifier(provider) + verdict = await clf.classify("edit", {"file_path": "a.py"}, "needs confirmation") + assert verdict.decisive and verdict.level is LOW + # Prompt shape: system + user with the tool call details. + assert provider.calls[0]["messages"][0]["role"] == "system" + assert "edit" in provider.calls[0]["messages"][1]["content"] + + +@pytest.mark.asyncio +async def test_classify_high_falls_through(): + provider = FakeProvider('{"risk": "high", "reason": "touches credentials"}') + clf = LLMRiskClassifier(provider) + verdict = await clf.classify("bash", {"command": "cat ~/.ssh/id_rsa"}, "ask") + assert verdict.decisive and verdict.level is HIGH + + +@pytest.mark.asyncio +async def test_classify_provider_error_is_fail_open_medium(): + provider = FakeProvider(error=RuntimeError("provider down")) + clf = LLMRiskClassifier(provider, timeout_s=1.0) + verdict = await clf.classify("bash", {"command": "git push"}, "ask") + # Fail-open: not decisive, falls back to the human path, never auto-allows. + assert not verdict.decisive + assert verdict.error is not None + assert verdict.level is MEDIUM + + +@pytest.mark.asyncio +async def test_classify_unparseable_reply_is_fail_open(): + provider = FakeProvider("sorry, cannot classify") + clf = LLMRiskClassifier(provider) + verdict = await clf.classify("write", {"file_path": "x"}, "ask") + assert not verdict.decisive + assert "unparseable" in (verdict.error or "") + + +@pytest.mark.asyncio +async def test_classify_model_and_max_tokens_forwarded(): + provider = FakeProvider('{"risk": "low", "reason": "ok"}') + clf = LLMRiskClassifier(provider, model="deepseek-v4-flash", max_tokens=64) + await clf.classify("read", {"file_path": "a"}, "ask") + call = provider.calls[0] + assert call["model"] == "deepseek-v4-flash" + assert call["max_tokens"] == 64 + assert call["temperature"] == 0.0 + + +@pytest.mark.asyncio +async def test_classify_temperature_is_zero_for_stability(): + provider = FakeProvider('{"risk": "low", "reason": "ok"}') + clf = LLMRiskClassifier(provider) + await clf.classify("read", {"file_path": "a"}, "ask") + assert provider.calls[0]["temperature"] == 0.0 + + +@pytest.mark.asyncio +async def test_classify_arguments_truncated_for_huge_payloads(): + provider = FakeProvider('{"risk": "medium", "reason": "big"}') + clf = LLMRiskClassifier(provider) + huge = {"file_path": "x", "content": "A" * 5000} + await clf.classify("write", huge, "ask") + user_msg = provider.calls[0]["messages"][1]["content"] + assert len(user_msg) < 2600 # 2000-char args cap + template overhead + + +# ---- env switches ----------------------------------------------------------- + + +def test_classifier_env_switches(monkeypatch): + monkeypatch.delenv("DEEPCODE_RISK_CLASSIFIER", raising=False) + assert classifier_enabled() is False + for value in ("1", "true", "yes", "on"): + monkeypatch.setenv("DEEPCODE_RISK_CLASSIFIER", value) + assert classifier_enabled() is True + for value in ("0", "false", "off", "banana"): + monkeypatch.setenv("DEEPCODE_RISK_CLASSIFIER", value) + assert classifier_enabled() is False + + +def test_classifier_model_env(monkeypatch): + monkeypatch.delenv("DEEPCODE_RISK_CLASSIFIER_MODEL", raising=False) + assert classifier_model() is None + monkeypatch.setenv("DEEPCODE_RISK_CLASSIFIER_MODEL", "deepseek-v4-flash") + assert classifier_model() == "deepseek-v4-flash"