open-swe/agent/middleware/pr_verdict_guard.py

"""Block shell fallbacks that submit PR review verdicts outside publish_review."""

from __future__ import annotations

import json
import re
import shlex
from collections.abc import Awaitable, Callable, Mapping
from typing import Any

from langchain.agents.middleware.types import AgentMiddleware, AgentState
from langchain_core.messages import ToolMessage
from langgraph.prebuilt.tool_node import ToolCallRequest
from langgraph.types import Command

_SHELL_SEPARATORS = {";", "&&", "||", "|", "&"}
_GITHUB_REVIEWS_ENDPOINT = re.compile(
    r"(?:^|/)repos/[^/\s]+/[^/\s]+/pulls/\d+/reviews(?:/\d+/events)?/?$"
)
_GITHUB_REVIEWS_URL = re.compile(
    r"https://api\.github\.com/repos/[^/\s]+/[^/\s]+/pulls/\d+/reviews(?:/\d+/events)?/?"
)
_VERDICT_EVENT_RE = re.compile(r"\bevent\b.{0,4}?(APPROVE|REQUEST_CHANGES)", re.IGNORECASE)
_GH_PR_REVIEW_VERDICT_FLAGS = {"--approve", "-a", "--request-changes", "-r"}
_GH_PR_REVIEW_VERDICT_PREFIXES = ("--approve=", "--request-changes=")
_BLOCK_ERROR = (
    "PR review verdicts (approve / request changes) must go through the "
    "publish_review tool, which enforces verdict authorization, self-review "
    "handling, and reviewer bookkeeping. Do not fall back to gh pr review, "
    "gh api .../reviews, curl, or another direct review-submission path. If a "
    "verdict was requested, call publish_review(verdict=...); otherwise "
    "publish a comment review."
)


def _tool_name(request: ToolCallRequest) -> str | None:
    tool_call = getattr(request, "tool_call", None)
    if isinstance(tool_call, Mapping):
        name = tool_call.get("name")
        return name if isinstance(name, str) else None
    return None


def _tool_args(request: ToolCallRequest) -> dict[str, Any]:
    tool_call = getattr(request, "tool_call", None)
    args = tool_call.get("args") if isinstance(tool_call, Mapping) else None
    return dict(args) if isinstance(args, Mapping) else {}


def _tool_call_id(request: ToolCallRequest) -> str | None:
    tool_call = getattr(request, "tool_call", None)
    if isinstance(tool_call, Mapping):
        value = tool_call.get("id")
        return value if isinstance(value, str) else None
    return None


def _shell_tokens(command: str) -> list[str]:
    try:
        return shlex.split(command, posix=True)
    except ValueError:
        return command.split()


def _subtokens_after(tokens: list[str], index: int) -> list[str]:
    subtokens: list[str] = []
    for token in tokens[index + 1 :]:
        if token in _SHELL_SEPARATORS:
            break
        subtokens.append(token)
    return subtokens


def _contains_gh_pr_review_verdict(tokens: list[str]) -> bool:
    for index, token in enumerate(tokens):
        if token != "gh":
            continue
        subtokens = _subtokens_after(tokens, index)
        is_pr_review = any(
            subtoken == "pr" and subtokens[offset + 1] == "review"
            for offset, subtoken in enumerate(subtokens[:-1])
        )
        if not is_pr_review:
            continue
        for subtoken in subtokens:
            if subtoken in _GH_PR_REVIEW_VERDICT_FLAGS or subtoken.startswith(
                _GH_PR_REVIEW_VERDICT_PREFIXES
            ):
                return True
    return False


def _contains_gh_api_review_verdict(tokens: list[str]) -> bool:
    for index, token in enumerate(tokens):
        if token != "gh":
            continue
        subtokens = _subtokens_after(tokens, index)
        if "api" not in subtokens:
            continue
        targets_reviews = any(
            _GITHUB_REVIEWS_ENDPOINT.search(subtoken.strip("'\""))
            or _GITHUB_REVIEWS_URL.search(subtoken)
            for subtoken in subtokens
        )
        if not targets_reviews:
            continue
        if any(_VERDICT_EVENT_RE.search(subtoken) for subtoken in subtokens):
            return True
    return False


def _contains_direct_review_verdict(tokens: list[str]) -> bool:
    for index, token in enumerate(tokens):
        if token != "curl":
            continue
        subtokens = _subtokens_after(tokens, index)
        if not any(_GITHUB_REVIEWS_URL.search(subtoken) for subtoken in subtokens):
            continue
        if any(_VERDICT_EVENT_RE.search(subtoken) for subtoken in subtokens):
            return True
    return False


def is_pr_verdict_fallback_command(command: str) -> bool:
    """Return True when *command* submits a PR review verdict from the shell.

    Detection is literal-token matching — it catches the primary vectors
    (``gh pr review --approve/-a/--request-changes/-r``, ``gh api`` posting
    ``event=APPROVE|REQUEST_CHANGES`` to a ``/pulls/N/reviews`` endpoint, and
    ``curl`` to the reviews URL with a verdict event in the body) but is
    intentionally fail-open: shell aliases, ``gh`` aliases, ``--input``
    JSON-file bodies, and non-curl HTTP clients will not be blocked. This is
    acceptable because the threat model is an honest agent papering over a
    ``publish_review`` limitation or refusal, not an adversary trying to
    bypass the guardrail. Plain ``gh pr review`` (no verdict flag) and
    ``--comment``/``-c`` are allowed — non-interactive ``gh`` cannot submit a
    verdict without one of the blocked flags.
    """
    tokens = _shell_tokens(command)
    return (
        _contains_gh_pr_review_verdict(tokens)
        or _contains_gh_api_review_verdict(tokens)
        or _contains_direct_review_verdict(tokens)
    )


def _blocked_tool_message(request: ToolCallRequest, command: str) -> ToolMessage:
    content = {
        "status": "error",
        "error_type": "PullRequestVerdictFallbackBlocked",
        "code": "pr_verdict_fallback_blocked",
        "recoverable_by_agent": False,
        "error": _BLOCK_ERROR,
        "blocked_command": command,
    }
    return ToolMessage(
        content=json.dumps(content),
        tool_call_id=_tool_call_id(request),
        status="error",
    )


class PullRequestVerdictGuardMiddleware(AgentMiddleware):
    """Keep APPROVE/REQUEST_CHANGES submissions on the guarded publish_review path."""

    state_schema = AgentState

    def _blocked_message_for_request(self, request: ToolCallRequest) -> ToolMessage | None:
        if _tool_name(request) != "execute":
            return None
        command = _tool_args(request).get("command")
        if not isinstance(command, str) or not is_pr_verdict_fallback_command(command):
            return None
        return _blocked_tool_message(request, command)

    def wrap_tool_call(
        self,
        request: ToolCallRequest,
        handler: Callable[[ToolCallRequest], ToolMessage | Command],
    ) -> ToolMessage | Command:
        blocked = self._blocked_message_for_request(request)
        if blocked is not None:
            return blocked
        return handler(request)

    async def awrap_tool_call(
        self,
        request: ToolCallRequest,
        handler: Callable[[ToolCallRequest], Awaitable[ToolMessage | Command]],
    ) -> ToolMessage | Command:
        blocked = self._blocked_message_for_request(request)
        if blocked is not None:
            return blocked
        return await handler(request)