mirror of
https://github.com/NCBM/plyngent.git
synced 2026-11-07 07:17:54 +08:00
core/agent: let YOLO mode skip YOLO-tagged soft confirms
`--yes` / `/yolo on|off|once` clears the prompt of a `YOLO`-tagged tool: interpreter runs (`run_argv`, `run_argv_batch`, `open_pty`), the file mutations (overwrite / copy / move / delete) and `fetch_local`'s cleartext and mutating requests. A tool without the tag still asks, and hard path/command denylists stay untouched. `classify_danger` now answers with a `DangerReason` — the reason text plus `yolo_skip` — so one call can refuse that skip: a command that runs through `sudo` / `su` / `doas` / `pkexec` (as the program or a wrapper) keeps its confirm whatever YOLO mode says. It never adds a reason of its own, so the program the denylist gate already asks about is not asked twice. `prepare_batch_confirm` leaves the covered calls out, so a skipped call neither prompts on its own nor takes a slot in the round's aggregated prompt, and one escalating step keeps the whole batch in front of the human.
This commit is contained in:
@@ -30,6 +30,7 @@ from .notices import skill_tools_notice as skill_tools_notice
|
||||
from .notices import truncated_turn_notice as truncated_turn_notice
|
||||
from .tools import BatchConfirmDecision as BatchConfirmDecision
|
||||
from .tools import DangerClassifier as DangerClassifier
|
||||
from .tools import DangerReason as DangerReason
|
||||
from .tools import ToolBatchConfirmHook as ToolBatchConfirmHook
|
||||
from .tools import ToolConfirmHook as ToolConfirmHook
|
||||
from .tools import ToolConfirmRequest as ToolConfirmRequest
|
||||
|
||||
+39
-11
@@ -16,8 +16,22 @@ from msgspec import Struct
|
||||
from plyngent.lmproto.openai_compatible.model import ChatImagePart, ToolFunction, ToolFunctionItem
|
||||
from plyngent.typedef import JSONSchema # noqa: TC001
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class DangerReason:
|
||||
"""Why a call needs a soft confirm, and whether YOLO mode may drop that prompt.
|
||||
|
||||
``yolo_skip`` is false for a call that reaches rights the session does not
|
||||
have (``sudo`` / ``doas`` / ``pkexec`` and friends): the human answers it
|
||||
however YOLO mode is set.
|
||||
"""
|
||||
|
||||
text: str
|
||||
yolo_skip: bool = True
|
||||
|
||||
|
||||
type ToolHandler = Callable[..., Any | Awaitable[Any]]
|
||||
type DangerClassifier = Callable[[str, Mapping[str, object]], str | None]
|
||||
type DangerClassifier = Callable[[str, Mapping[str, object]], str | DangerReason | None]
|
||||
# Returns True to allow, False to deny, or a non-empty str as denial reason for the model.
|
||||
type ToolConfirmResult = bool | str
|
||||
type ToolConfirmHook = Callable[
|
||||
@@ -313,6 +327,13 @@ def _denied_error(name: str, reason: str, comment: str | None) -> str:
|
||||
return f"error: tool {name!r} denied by user confirm ({reason})"
|
||||
|
||||
|
||||
def _as_danger_reason(raw: str | DangerReason | None) -> DangerReason | None:
|
||||
"""Normalize a classifier answer; a bare string is a reason YOLO mode may drop."""
|
||||
if raw is None:
|
||||
return None
|
||||
return raw if isinstance(raw, DangerReason) else DangerReason(text=raw)
|
||||
|
||||
|
||||
def _write_path_key(value: str) -> str:
|
||||
"""Comparable key for one path argument of a writing tool.
|
||||
|
||||
@@ -559,7 +580,9 @@ class ToolRegistry:
|
||||
|
||||
Mirrors the front half of :meth:`execute`: a call whose arguments fail to
|
||||
decode or miss a required argument is answered with an error there and
|
||||
must not be put in front of the user now.
|
||||
must not be put in front of the user now, and a call whose reason YOLO
|
||||
mode covers runs unprompted, so it must not take a slot in the round's
|
||||
aggregated prompt either.
|
||||
"""
|
||||
definition = self._tools.get(name)
|
||||
if self._danger is None or definition is None:
|
||||
@@ -567,8 +590,10 @@ class ToolRegistry:
|
||||
args = _decode_tool_args(arguments_json)
|
||||
if isinstance(args, str) or _missing_required_error(definition, args) is not None:
|
||||
return None
|
||||
reason = self._danger(name, args)
|
||||
return None if reason is None else ToolConfirmRequest(name=name, args=args, reason=reason)
|
||||
reason = _as_danger_reason(self._danger(name, args))
|
||||
if reason is None or (self._yolo and (definition.tags & ToolTag.YOLO) and reason.yolo_skip):
|
||||
return None
|
||||
return ToolConfirmRequest(name=name, args=args, reason=reason.text)
|
||||
|
||||
def _batch_requests(
|
||||
self,
|
||||
@@ -633,29 +658,32 @@ class ToolRegistry:
|
||||
tags: ToolTag,
|
||||
decision: BatchConfirmDecision | None = None,
|
||||
) -> str | None:
|
||||
"""Soft-confirm gate driven by danger reason + tags + grants.
|
||||
"""Soft-confirm gate driven by danger reason + tags + grants + YOLO.
|
||||
|
||||
Pipeline (soft gate only; hard denylists stay in tool handlers)::
|
||||
|
||||
no soft reason → run
|
||||
YOLO mode, (tags & YOLO) and a reason it may drop → allow
|
||||
batch decision for this call → apply it
|
||||
(tags & TRUSTABLE) and grant exists → allow
|
||||
else → on_confirm (decorator YOLO tag does not override
|
||||
dangerous invocations — internal danger classification
|
||||
always takes precedence)
|
||||
else → on_confirm (a tool without the YOLO tag, or a reason YOLO
|
||||
may not drop — a privilege escalation — still prompts
|
||||
under YOLO mode)
|
||||
"""
|
||||
if self._danger is None:
|
||||
return None
|
||||
reason = self._danger(name, args)
|
||||
reason = _as_danger_reason(self._danger(name, args))
|
||||
if reason is None:
|
||||
return None
|
||||
if self._yolo and (tags & ToolTag.YOLO) and reason.yolo_skip:
|
||||
return None
|
||||
|
||||
async with self._confirm_lock:
|
||||
if decision is not None:
|
||||
return await self._apply_batch_decision(name, reason, decision, tags=tags)
|
||||
return await self._apply_batch_decision(name, reason.text, decision, tags=tags)
|
||||
if self._grant_allows(name, tags=tags):
|
||||
return None
|
||||
return await self._prompt_soft_confirm(name, args, reason, tags=tags)
|
||||
return await self._prompt_soft_confirm(name, args, reason.text, tags=tags)
|
||||
|
||||
async def execute(
|
||||
self,
|
||||
|
||||
@@ -16,7 +16,16 @@ from typing import TYPE_CHECKING
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
__all__ = ["CommandString", "UnwrappedCommand", "basename", "is_interpreter", "unwrap_command"]
|
||||
__all__ = [
|
||||
"ESCALATION_PROGRAMS",
|
||||
"CommandString",
|
||||
"UnwrappedCommand",
|
||||
"basename",
|
||||
"dash_c_code",
|
||||
"escalators_in_chain",
|
||||
"is_interpreter",
|
||||
"unwrap_command",
|
||||
]
|
||||
|
||||
# Wrappers that run the program named on their own command line.
|
||||
_ENV_WRAPPERS: frozenset[str] = frozenset(
|
||||
@@ -99,6 +108,10 @@ _INTERPRETERS: frozenset[str] = frozenset(
|
||||
# Trailing version of an interpreter: ``python3.12`` / ``node20`` / ``php8.2``.
|
||||
_VERSION_SUFFIX = re.compile(r"[0-9][0-9.]*$")
|
||||
|
||||
# Privilege escalation: these run with rights the session itself does not have,
|
||||
# so a call that reaches one is never approved on the caller's behalf
|
||||
# (``danger._command_reason`` keeps the confirm whatever YOLO mode says).
|
||||
ESCALATION_PROGRAMS: frozenset[str] = frozenset({"doas", "pkexec", "run0", "su", "sudo", "systemd-run"})
|
||||
# ``FOO=bar`` in front of a program (env/sudo style).
|
||||
_ASSIGNMENT = re.compile(r"[A-Za-z_][A-Za-z0-9_]*=")
|
||||
|
||||
@@ -245,6 +258,31 @@ def is_interpreter(base: str) -> bool:
|
||||
return base in _INTERPRETERS or _VERSION_SUFFIX.sub("", base) in _INTERPRETERS
|
||||
|
||||
|
||||
def dash_c_code(command: Sequence[str]) -> tuple[int, str] | None:
|
||||
"""Index and text of the ``-c`` code in an interpreter argv, if present.
|
||||
|
||||
Only interpreters hand ``-c`` a block of code; for anything else it is a
|
||||
plain flag (``grep -c``, ``od -c``) with nothing to review.
|
||||
"""
|
||||
for index, part in enumerate(command[1:], start=1):
|
||||
if part == "-c":
|
||||
return index, command[index + 1] if index + 1 < len(command) else ""
|
||||
# Combined short forms are uncommon; only exact -c is supported.
|
||||
return None
|
||||
|
||||
|
||||
def escalators_in_chain(argv: Sequence[str]) -> tuple[str, ...]:
|
||||
"""Escalation programs *argv* runs: the program itself or a wrapper around it.
|
||||
|
||||
``sudo -u root python …`` and ``doas bash`` both count, whatever program they
|
||||
finally run. Code an interpreter only carries as text is out of scope here —
|
||||
this reads the command line, not what the text would do.
|
||||
"""
|
||||
scan = unwrap_command(argv)
|
||||
names = (scan.base, *scan.wrappers)
|
||||
return tuple(dict.fromkeys(name for name in names if name in ESCALATION_PROGRAMS))
|
||||
|
||||
|
||||
def unwrap_command(argv: Sequence[str]) -> UnwrappedCommand:
|
||||
"""Resolve *argv* to the program behind any wrappers and ``run`` launchers.
|
||||
|
||||
|
||||
@@ -4,7 +4,8 @@ import json
|
||||
import shlex
|
||||
from typing import TYPE_CHECKING, cast
|
||||
|
||||
from plyngent.tools.command_scan import is_interpreter, unwrap_command
|
||||
from plyngent.agent.tools import DangerReason
|
||||
from plyngent.tools.command_scan import dash_c_code, escalators_in_chain, is_interpreter, unwrap_command
|
||||
from plyngent.tools.workspace import AccessMode, WorkspaceError, resolve_path
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -23,19 +24,6 @@ def _as_argv(args: Mapping[str, object], *, key: str = "command") -> list[str] |
|
||||
return out
|
||||
|
||||
|
||||
def _dash_c_code(command: Sequence[str]) -> tuple[int, str] | None:
|
||||
"""Index and text of the ``-c`` code in an interpreter argv, if present.
|
||||
|
||||
Only interpreters hand ``-c`` a block of code; for anything else it is a
|
||||
plain flag (``grep -c``, ``od -c``) with nothing to review.
|
||||
"""
|
||||
for index, part in enumerate(command[1:], start=1):
|
||||
if part == "-c":
|
||||
return index, command[index + 1] if index + 1 < len(command) else ""
|
||||
# Combined short forms are uncommon; only exact -c is supported.
|
||||
return None
|
||||
|
||||
|
||||
def _indent_block(text: str, *, prefix: str = " ") -> str:
|
||||
"""Indent every line of *text* (for multi-line -c code in confirm prompts)."""
|
||||
if not text:
|
||||
@@ -51,7 +39,7 @@ def _format_argv_for_confirm(argv: Sequence[str], *, code_index: int | None) ->
|
||||
return shlex.join(parts)
|
||||
|
||||
|
||||
def _command_reason(argv: Sequence[str], *, via: str) -> str | None:
|
||||
def _command_reason(argv: Sequence[str], *, via: str) -> DangerReason | None:
|
||||
"""Confirm a command that runs code or detaches, whatever wraps it.
|
||||
|
||||
Multi-line reason (shown inside the CLI confirm box). ``via`` is a short
|
||||
@@ -65,6 +53,11 @@ def _command_reason(argv: Sequence[str], *, via: str) -> str | None:
|
||||
(``env -S "…"``, ``su -c "…"``) hides code the same way, so it is confirmed
|
||||
too, with the text it carries printed under its own ``<wrapper> <option>:``
|
||||
label.
|
||||
|
||||
A command that escalates privileges (``sudo``/``doas``/``pkexec`` as the
|
||||
program or a wrapper) is confirmed with ``yolo_skip=False``. It never adds a
|
||||
reason of its own: the denylist gate in the handler always asks about the
|
||||
escalation, so a second prompt here would only repeat it.
|
||||
"""
|
||||
scan = unwrap_command(argv)
|
||||
interpreter = is_interpreter(scan.base)
|
||||
@@ -75,7 +68,10 @@ def _command_reason(argv: Sequence[str], *, via: str) -> str | None:
|
||||
findings = list(dict.fromkeys(findings))
|
||||
if not findings:
|
||||
return None
|
||||
code = _dash_c_code(scan.command) if interpreter else None
|
||||
escalators = escalators_in_chain(argv)
|
||||
if escalators:
|
||||
findings.append(f"privilege escalation {', '.join(escalators)}")
|
||||
code = dash_c_code(scan.command) if interpreter else None
|
||||
code_index = scan.offset + code[0] if code is not None else None
|
||||
display = _format_argv_for_confirm(argv, code_index=code_index)
|
||||
reason = f"{via}: {' + '.join(findings)} — review before allow\n argv: {display}"
|
||||
@@ -83,7 +79,7 @@ def _command_reason(argv: Sequence[str], *, via: str) -> str | None:
|
||||
reason += f"\n command:\n{_indent_block(code[1])}"
|
||||
for span in scan.command_strings:
|
||||
reason += f"\n {span.label}:\n{_indent_block(span.text)}"
|
||||
return reason
|
||||
return DangerReason(text=reason, yolo_skip=not escalators)
|
||||
|
||||
|
||||
def _write_file_reason(args: Mapping[str, object]) -> str | None:
|
||||
@@ -139,7 +135,7 @@ def _delete_path_reason(args: Mapping[str, object]) -> str | None:
|
||||
return f"delete path {path!r}{extra}"
|
||||
|
||||
|
||||
def _run_argv_reason(args: Mapping[str, object]) -> str | None:
|
||||
def _run_argv_reason(args: Mapping[str, object]) -> DangerReason | None:
|
||||
argv = _as_argv(args, key="argv")
|
||||
if argv is None:
|
||||
return None
|
||||
@@ -161,8 +157,12 @@ def _batch_step_argv(item: object) -> list[str] | None:
|
||||
return parts or None
|
||||
|
||||
|
||||
def _run_argv_batch_reason(args: Mapping[str, object]) -> str | None:
|
||||
"""One confirm for the whole batch if any step runs code or detaches."""
|
||||
def _run_argv_batch_reason(args: Mapping[str, object]) -> DangerReason | None:
|
||||
"""One confirm for the whole batch if any step runs code or detaches.
|
||||
|
||||
YOLO mode may drop it only when every step's reason allows that, so one
|
||||
escalating step keeps the whole batch in front of the human.
|
||||
"""
|
||||
raw = args.get("steps")
|
||||
if isinstance(raw, str):
|
||||
try:
|
||||
@@ -181,10 +181,13 @@ def _run_argv_batch_reason(args: Mapping[str, object]) -> str | None:
|
||||
if not risky:
|
||||
return None
|
||||
header = f"run_argv_batch: {len(risky)} risky step(s) (review before allow)"
|
||||
return header + "\n" + "\n".join(risky)
|
||||
return DangerReason(
|
||||
text=header + "\n" + "\n".join(reason.text for reason in risky),
|
||||
yolo_skip=all(reason.yolo_skip for reason in risky),
|
||||
)
|
||||
|
||||
|
||||
def _open_pty_reason(args: Mapping[str, object]) -> str | None:
|
||||
def _open_pty_reason(args: Mapping[str, object]) -> DangerReason | None:
|
||||
argv = _as_argv(args)
|
||||
if argv is None:
|
||||
return None
|
||||
@@ -204,14 +207,18 @@ def _fetch_reason(args: Mapping[str, object]) -> str | None:
|
||||
return fetch_soft_reason(method, url, body)
|
||||
|
||||
|
||||
def classify_danger(name: str, args: Mapping[str, object]) -> str | None: # noqa: PLR0911
|
||||
"""Return a short reason if ``name``/``args`` need user confirm, else ``None``.
|
||||
def classify_danger(name: str, args: Mapping[str, object]) -> str | DangerReason | None: # noqa: PLR0911
|
||||
"""Return why ``name``/``args`` need a user confirm, else ``None``.
|
||||
|
||||
Hard denylists (paths/commands) still raise independently. This only covers
|
||||
soft confirms for mutating tools and risky command launches — interpreters
|
||||
(however invoked) and detached runners such as ``nohup``, resolved past
|
||||
environment wrappers and ``<launcher> run``.
|
||||
Private/loopback fetch_local targets use a separate policy grant (not YOLO).
|
||||
|
||||
A reason also says whether YOLO mode may drop its prompt
|
||||
(:attr:`DangerReason.yolo_skip`); a command that escalates privileges
|
||||
refuses.
|
||||
"""
|
||||
if name == "delete_path":
|
||||
return _delete_path_reason(args)
|
||||
|
||||
@@ -93,6 +93,46 @@ async def test_round_asks_once_for_many_mutating_calls() -> None:
|
||||
assert single_prompts == []
|
||||
|
||||
|
||||
async def test_yolo_calls_are_left_out_of_the_prompt() -> None:
|
||||
"""Under YOLO mode a YOLO-tagged call runs unprompted, batch prompt included."""
|
||||
batch_prompts: list[list[str]] = []
|
||||
single_prompts: list[str] = []
|
||||
|
||||
@tool(tags=ToolTag.LOCAL | ToolTag.YOLO, register=False)
|
||||
async def yolo_ok(path: str) -> str:
|
||||
return f"stamped {path}"
|
||||
|
||||
@tool(tags=ToolTag.LOCAL, register=False)
|
||||
async def always_ask(path: str) -> str:
|
||||
return f"stamped {path}"
|
||||
|
||||
def on_batch(requests: Sequence[ToolConfirmRequest]) -> BatchConfirmDecision:
|
||||
batch_prompts.append([request.name for request in requests])
|
||||
return BatchConfirmDecision(allowed=True)
|
||||
|
||||
async def on_single(name: str, _args: Mapping[str, object], _reason: str) -> bool:
|
||||
single_prompts.append(name)
|
||||
return True
|
||||
|
||||
reg = ToolRegistry(
|
||||
[yolo_ok, always_ask],
|
||||
danger=_danger,
|
||||
on_confirm=on_single,
|
||||
on_confirm_batch=on_batch,
|
||||
yolo=True,
|
||||
)
|
||||
calls = [("yolo_ok", '{"path": "a"}'), ("always_ask", '{"path": "b"}'), ("always_ask", '{"path": "c"}')]
|
||||
decisions = await reg.prepare_batch_confirm(calls)
|
||||
|
||||
assert decisions[0] is None
|
||||
assert batch_prompts == [["always_ask", "always_ask"]]
|
||||
results = [
|
||||
await reg.execute(name, arguments, decision=decisions[index]) for index, (name, arguments) in enumerate(calls)
|
||||
]
|
||||
assert results == ["stamped a", "stamped b", "stamped c"]
|
||||
assert single_prompts == []
|
||||
|
||||
|
||||
async def test_single_request_keeps_the_one_call_prompt() -> None:
|
||||
"""One mutating call stays the plain per-call confirm."""
|
||||
batch_prompts: list[int] = []
|
||||
|
||||
@@ -486,8 +486,8 @@ async def test_yolo_rebuilds_tool_registry(tmp_path: Path) -> None:
|
||||
)
|
||||
try:
|
||||
assert st.agent.tools is not None
|
||||
# Soft-confirm hooks stay installed; YOLO mode does not bypass
|
||||
# dangerous invocations — internal danger classification takes precedence.
|
||||
# Soft-confirm hooks stay installed; YOLO mode only auto-approves the
|
||||
# confirms of YOLO-tagged tools, so the others still prompt.
|
||||
assert st.agent.tools.soft_confirm is True
|
||||
assert st.agent.tools.yolo is False
|
||||
st.set_yolo("on")
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from plyngent.tools.command_scan import basename, is_interpreter, unwrap_command
|
||||
from plyngent.tools.command_scan import basename, escalators_in_chain, is_interpreter, unwrap_command
|
||||
|
||||
|
||||
def test_basename() -> None:
|
||||
@@ -9,6 +9,15 @@ def test_basename() -> None:
|
||||
assert basename("pdm") == "pdm"
|
||||
|
||||
|
||||
def test_escalators_in_chain() -> None:
|
||||
assert escalators_in_chain(["sudo", "-u", "root", "python", "-c", "print(1)"]) == ("sudo",)
|
||||
assert escalators_in_chain(["env", "FOO=1", "doas", "id"]) == ("doas",)
|
||||
assert escalators_in_chain(["pdm", "run", "sudo", "id"]) == ("sudo",)
|
||||
# Only the command line counts: the code an interpreter carries is not read.
|
||||
assert escalators_in_chain(["bash", "-c", "sudo id"]) == ()
|
||||
assert escalators_in_chain(["grep", "-c", "sudo", "notes.txt"]) == ()
|
||||
|
||||
|
||||
def test_unwrap_plain_command() -> None:
|
||||
scan = unwrap_command(["grep", "-c", "needle", "f.txt"])
|
||||
assert scan.wrappers == ()
|
||||
|
||||
@@ -2,18 +2,18 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from plyngent.agent import ToolRegistry, ToolTag, tool
|
||||
from plyngent.agent import DangerReason, ToolRegistry, ToolTag, tool
|
||||
from plyngent.tools.context import SessionState, bind_tool_context
|
||||
from plyngent.tools.grants import clear_grants, has_grant, hydrate_grants, persist_grants
|
||||
from plyngent.tools.view import MemoryViewStore, session_data_view
|
||||
|
||||
|
||||
async def test_danger_overrides_yolo_tag() -> None:
|
||||
"""Decorator YOLO tag does not bypass confirm when danger flags the call."""
|
||||
async def test_yolo_only_skips_yolo_tagged_tools() -> None:
|
||||
"""YOLO mode auto-approves the soft confirms of YOLO-tagged tools only."""
|
||||
calls: list[str] = []
|
||||
|
||||
@tool(tags=ToolTag.LOCAL | ToolTag.YOLO, register=False)
|
||||
async def yolo_tool() -> str:
|
||||
async def yolo_ok() -> str:
|
||||
return "y"
|
||||
|
||||
@tool(tags=ToolTag.LOCAL, register=False)
|
||||
@@ -28,15 +28,55 @@ async def test_danger_overrides_yolo_tag() -> None:
|
||||
return True
|
||||
|
||||
reg = ToolRegistry(
|
||||
[yolo_tool, no_yolo],
|
||||
[yolo_ok, no_yolo],
|
||||
danger=danger,
|
||||
on_confirm=confirm,
|
||||
yolo=True,
|
||||
)
|
||||
assert await reg.execute("yolo_tool", "{}") == "y"
|
||||
assert calls == ["yolo_tool"]
|
||||
assert await reg.execute("yolo_ok", "{}") == "y"
|
||||
assert calls == []
|
||||
assert await reg.execute("no_yolo", "{}") == "n"
|
||||
assert calls == ["yolo_tool", "no_yolo"]
|
||||
assert calls == ["no_yolo"]
|
||||
|
||||
|
||||
async def test_yolo_off_keeps_the_yolo_tagged_prompt() -> None:
|
||||
"""The YOLO tag is a policy bit: with YOLO off the call still asks."""
|
||||
calls: list[str] = []
|
||||
|
||||
@tool(tags=ToolTag.LOCAL | ToolTag.YOLO, register=False)
|
||||
async def yolo_ok() -> str:
|
||||
return "y"
|
||||
|
||||
def danger(_name: str, _args: object) -> str | None:
|
||||
return "soft reason"
|
||||
|
||||
async def confirm(name: str, _args: object, _reason: str) -> bool:
|
||||
calls.append(name)
|
||||
return True
|
||||
|
||||
reg = ToolRegistry([yolo_ok], danger=danger, on_confirm=confirm, yolo=False)
|
||||
assert await reg.execute("yolo_ok", "{}") == "y"
|
||||
assert calls == ["yolo_ok"]
|
||||
|
||||
|
||||
async def test_yolo_keeps_a_reason_that_refuses_the_skip() -> None:
|
||||
"""`DangerReason(yolo_skip=False)` — a privilege escalation — still prompts."""
|
||||
seen: list[str] = []
|
||||
|
||||
@tool(tags=ToolTag.LOCAL | ToolTag.YOLO, register=False)
|
||||
async def escalated() -> str:
|
||||
return "x"
|
||||
|
||||
def danger(_name: str, _args: object) -> DangerReason:
|
||||
return DangerReason("interpreter 'python' + privilege escalation sudo", yolo_skip=False)
|
||||
|
||||
async def confirm(_name: str, _args: object, reason: str) -> bool:
|
||||
seen.append(reason)
|
||||
return True
|
||||
|
||||
reg = ToolRegistry([escalated], danger=danger, on_confirm=confirm, yolo=True)
|
||||
assert await reg.execute("escalated", "{}") == "x"
|
||||
assert seen == ["interpreter 'python' + privilege escalation sudo"]
|
||||
|
||||
|
||||
async def test_trustable_grant_once() -> None:
|
||||
|
||||
@@ -2,19 +2,35 @@ from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from plyngent.agent import DangerReason
|
||||
from plyngent.tools.danger import classify_danger
|
||||
|
||||
|
||||
def _reason(name: str, args: dict[str, object]) -> str | None:
|
||||
"""Confirm text of a classifier answer (``None`` when the call needs none)."""
|
||||
found = classify_danger(name, args)
|
||||
if found is None:
|
||||
return None
|
||||
return found if isinstance(found, str) else found.text
|
||||
|
||||
|
||||
def _danger(name: str, args: dict[str, object]) -> DangerReason:
|
||||
"""Classifier answer as a typed reason (the call must need a confirm)."""
|
||||
found = classify_danger(name, args)
|
||||
assert isinstance(found, DangerReason), found
|
||||
return found
|
||||
|
||||
|
||||
def test_classify_delete_and_move() -> None:
|
||||
assert classify_danger("delete_path", {"path": "a.txt"}) == "delete path 'a.txt'"
|
||||
assert "recursively" in (classify_danger("delete_path", {"path": "d", "recursive": True}) or "")
|
||||
assert "move" in (classify_danger("move_path", {"src": "a", "dst": "b"}) or "")
|
||||
assert _reason("delete_path", {"path": "a.txt"}) == "delete path 'a.txt'"
|
||||
assert "recursively" in (_reason("delete_path", {"path": "d", "recursive": True}) or "")
|
||||
assert "move" in (_reason("move_path", {"src": "a", "dst": "b"}) or "")
|
||||
|
||||
|
||||
def test_classify_copy() -> None:
|
||||
# Without overwrite flag, copy is not soft-confirmed.
|
||||
assert classify_danger("copy_path", {"src": "a", "dst": "b"}) is None
|
||||
assert classify_danger("copy_path", {"src": "a", "dst": "b", "overwrite": False}) is None
|
||||
assert _reason("copy_path", {"src": "a", "dst": "b"}) is None
|
||||
assert _reason("copy_path", {"src": "a", "dst": "b", "overwrite": False}) is None
|
||||
|
||||
|
||||
def test_classify_write_overwrite_only(tmp_path: Path) -> None:
|
||||
@@ -25,35 +41,35 @@ def test_classify_write_overwrite_only(tmp_path: Path) -> None:
|
||||
with bind_instance(instance):
|
||||
_ = set_workspace_root(tmp_path)
|
||||
# New file: no soft-confirm, with or without the overwrite flag
|
||||
assert classify_danger("write_file", {"path": "new.txt"}) is None
|
||||
assert classify_danger("write_file", {"path": "new.txt", "overwrite": True}) is None
|
||||
assert _reason("write_file", {"path": "new.txt"}) is None
|
||||
assert _reason("write_file", {"path": "new.txt", "overwrite": True}) is None
|
||||
# Partial edits: never soft-confirm
|
||||
assert classify_danger("edit_replace", {"path": "x.txt"}) is None
|
||||
assert classify_danger("edit_lineno", {"path": "x.txt", "start_line": 1, "end_line": 2}) is None
|
||||
assert _reason("edit_replace", {"path": "x.txt"}) is None
|
||||
assert _reason("edit_lineno", {"path": "x.txt", "start_line": 1, "end_line": 2}) is None
|
||||
# Existing file: only an explicit overwrite asks; an unset/false flag is
|
||||
# refused inside the tool, so no prompt may fire for it.
|
||||
(tmp_path / "x.txt").write_text("old", encoding="utf-8")
|
||||
assert classify_danger("write_file", {"path": "x.txt"}) is None
|
||||
assert classify_danger("write_file", {"path": "x.txt", "overwrite": False}) is None
|
||||
reason = classify_danger("write_file", {"path": "x.txt", "overwrite": True})
|
||||
assert _reason("write_file", {"path": "x.txt"}) is None
|
||||
assert _reason("write_file", {"path": "x.txt", "overwrite": False}) is None
|
||||
reason = _reason("write_file", {"path": "x.txt", "overwrite": True})
|
||||
assert reason is not None and "overwrite" in reason
|
||||
(tmp_path / "dst.txt").write_text("d", encoding="utf-8")
|
||||
(tmp_path / "src.txt").write_text("s", encoding="utf-8")
|
||||
assert classify_danger("copy_path", {"src": "src.txt", "dst": "dst.txt", "overwrite": False}) is None
|
||||
creason = classify_danger("copy_path", {"src": "src.txt", "dst": "dst.txt", "overwrite": True})
|
||||
assert _reason("copy_path", {"src": "src.txt", "dst": "dst.txt", "overwrite": False}) is None
|
||||
creason = _reason("copy_path", {"src": "src.txt", "dst": "dst.txt", "overwrite": True})
|
||||
assert creason is not None and "overwrite" in creason
|
||||
|
||||
|
||||
def test_classify_safe_tools() -> None:
|
||||
assert classify_danger("read_file", {"path": "a"}) is None
|
||||
assert classify_danger("listdir", {"path": "."}) is None
|
||||
assert classify_danger("run_argv", {"argv": ["echo", "hi"]}) is None
|
||||
assert classify_danger("open_pty", {"command": ["true"]}) is None
|
||||
assert classify_danger("run_argv", {"argv": ["ls", "-la"]}) is None
|
||||
assert _reason("read_file", {"path": "a"}) is None
|
||||
assert _reason("listdir", {"path": "."}) is None
|
||||
assert _reason("run_argv", {"argv": ["echo", "hi"]}) is None
|
||||
assert _reason("open_pty", {"command": ["true"]}) is None
|
||||
assert _reason("run_argv", {"argv": ["ls", "-la"]}) is None
|
||||
|
||||
|
||||
def test_classify_run_argv_batch_risky() -> None:
|
||||
reason = classify_danger(
|
||||
reason = _reason(
|
||||
"run_argv_batch",
|
||||
{
|
||||
"steps": [
|
||||
@@ -66,7 +82,7 @@ def test_classify_run_argv_batch_risky() -> None:
|
||||
assert "run_argv_batch" in reason
|
||||
assert "risky" in reason or "bash" in reason
|
||||
assert (
|
||||
classify_danger(
|
||||
_reason(
|
||||
"run_argv_batch",
|
||||
{"steps": [{"argv": ["echo", "ok"]}]},
|
||||
)
|
||||
@@ -75,42 +91,42 @@ def test_classify_run_argv_batch_risky() -> None:
|
||||
|
||||
|
||||
def test_classify_shell_and_dash_c() -> None:
|
||||
r = classify_danger("run_argv", {"argv": ["bash", "-c", "rm -rf /"]})
|
||||
r = _reason("run_argv", {"argv": ["bash", "-c", "rm -rf /"]})
|
||||
assert r is not None
|
||||
assert "bash -c" in r
|
||||
assert "rm -rf" in r
|
||||
|
||||
r2 = classify_danger("run_argv", {"argv": ["python3", "-c", "print(1)"]})
|
||||
r2 = _reason("run_argv", {"argv": ["python3", "-c", "print(1)"]})
|
||||
assert r2 is not None
|
||||
assert "python" in r2
|
||||
assert "-c" in r2
|
||||
assert "print(1)" in r2
|
||||
|
||||
r_py = classify_danger("run_argv", {"argv": ["python", "-c", "import os"]})
|
||||
r_py = _reason("run_argv", {"argv": ["python", "-c", "import os"]})
|
||||
assert r_py is not None and "python -c" in r_py
|
||||
|
||||
r3 = classify_danger("open_pty", {"command": ["bash"]})
|
||||
r3 = _reason("open_pty", {"command": ["bash"]})
|
||||
assert r3 is not None and "bash" in r3
|
||||
assert "interactive" in r3 or "review" in r3
|
||||
|
||||
r4 = classify_danger("open_pty", {"command": ["python3"]})
|
||||
r4 = _reason("open_pty", {"command": ["python3"]})
|
||||
assert r4 is not None and "python" in r4
|
||||
|
||||
r5 = classify_danger("open_pty", {"command": ["python3", "-c", "x=1"]})
|
||||
r5 = _reason("open_pty", {"command": ["python3", "-c", "x=1"]})
|
||||
assert r5 is not None and "-c" in r5
|
||||
|
||||
|
||||
def test_dash_c_means_code_only_for_interpreters() -> None:
|
||||
# ``-c`` is a plain flag elsewhere: no code to review, no confirm.
|
||||
assert classify_danger("run_argv", {"argv": ["grep", "-c", "needle", "f.txt"]}) is None
|
||||
assert classify_danger("run_argv", {"argv": ["od", "-c", "f.bin"]}) is None
|
||||
assert classify_danger("run_argv", {"argv": ["sort", "-c", "f.txt"]}) is None
|
||||
assert _reason("run_argv", {"argv": ["grep", "-c", "needle", "f.txt"]}) is None
|
||||
assert _reason("run_argv", {"argv": ["od", "-c", "f.bin"]}) is None
|
||||
assert _reason("run_argv", {"argv": ["sort", "-c", "f.txt"]}) is None
|
||||
|
||||
|
||||
def test_interpreter_review_does_not_split_on_dash_c() -> None:
|
||||
script = classify_danger("run_argv", {"argv": ["python", "script.py"]})
|
||||
one_liner = classify_danger("run_argv", {"argv": ["python", "-c", "print(1)"]})
|
||||
bare_shell = classify_danger("open_pty", {"command": ["bash"]})
|
||||
script = _reason("run_argv", {"argv": ["python", "script.py"]})
|
||||
one_liner = _reason("run_argv", {"argv": ["python", "-c", "print(1)"]})
|
||||
bare_shell = _reason("open_pty", {"command": ["bash"]})
|
||||
assert script is not None and "interpreter 'python'" in script
|
||||
assert one_liner is not None and "interpreter 'python'" in one_liner
|
||||
assert bare_shell is not None and "interpreter 'bash'" in bare_shell
|
||||
@@ -121,7 +137,7 @@ def test_interpreter_review_does_not_split_on_dash_c() -> None:
|
||||
|
||||
|
||||
def test_versioned_interpreter_is_still_reviewed() -> None:
|
||||
reason = classify_danger("run_argv", {"argv": ["python3.12", "-c", "import os"]})
|
||||
reason = _reason("run_argv", {"argv": ["python3.12", "-c", "import os"]})
|
||||
assert reason is not None and "interpreter 'python3.12'" in reason
|
||||
assert "import os" in reason
|
||||
|
||||
@@ -134,7 +150,7 @@ def test_wrappers_reach_the_interpreter() -> None:
|
||||
["pdm", "run", "python", "-c", "print(1)"],
|
||||
["uv", "run", "--with", "rich", "python", "-c", "print(1)"],
|
||||
):
|
||||
reason = classify_danger("run_argv", {"argv": argv})
|
||||
reason = _reason("run_argv", {"argv": argv})
|
||||
assert reason is not None, argv
|
||||
assert "interpreter 'python'" in reason or "interpreter 'bash'" in reason
|
||||
assert "$(command)" in reason
|
||||
@@ -142,7 +158,7 @@ def test_wrappers_reach_the_interpreter() -> None:
|
||||
|
||||
def test_value_free_flags_do_not_hide_the_interpreter() -> None:
|
||||
"""``exec -c`` takes no value, so the interpreter behind it is still reviewed."""
|
||||
reason = classify_danger("run_argv", {"argv": ["exec", "-c", "python", "-c", "print(1)"]})
|
||||
reason = _reason("run_argv", {"argv": ["exec", "-c", "python", "-c", "print(1)"]})
|
||||
assert reason is not None
|
||||
assert "interpreter 'python'" in reason
|
||||
assert "print(1)" in reason
|
||||
@@ -150,26 +166,26 @@ def test_value_free_flags_do_not_hide_the_interpreter() -> None:
|
||||
|
||||
def test_wrapper_command_line_option_is_reviewed() -> None:
|
||||
"""`env -S` / `su -c` carry a command line, so they are confirmed like `-c`."""
|
||||
reason = classify_danger("run_argv", {"argv": ["env", "-S", "rm -rf /"]})
|
||||
reason = _reason("run_argv", {"argv": ["env", "-S", "rm -rf /"]})
|
||||
assert reason is not None
|
||||
assert "env -S runs a command line" in reason
|
||||
assert " env -S:\n rm -rf /" in reason
|
||||
|
||||
# Plain wrapper options are not a review trigger.
|
||||
assert classify_danger("run_argv", {"argv": ["env", "-u", "FOO=1", "git", "status"]}) is None
|
||||
assert classify_danger("run_argv", {"argv": ["env", "FOO=1", "git", "-c", "x=y", "log"]}) is None
|
||||
assert _reason("run_argv", {"argv": ["env", "-u", "FOO=1", "git", "status"]}) is None
|
||||
assert _reason("run_argv", {"argv": ["env", "FOO=1", "git", "-c", "x=y", "log"]}) is None
|
||||
|
||||
|
||||
def test_wrapper_command_line_value_needs_no_interpreter() -> None:
|
||||
"""The wrapper's value is reviewed even when no interpreter is resolved."""
|
||||
reason = classify_danger("run_argv", {"argv": ["env", "-S", "python -O script.py"]})
|
||||
reason = _reason("run_argv", {"argv": ["env", "-S", "python -O script.py"]})
|
||||
assert reason is not None
|
||||
assert "env -S runs a command line" in reason
|
||||
assert "python -O script.py" in reason
|
||||
|
||||
|
||||
def test_only_the_interpreter_dash_c_becomes_a_placeholder() -> None:
|
||||
reason = classify_danger("run_argv", {"argv": ["ionice", "-c", "2", "python", "-c", "print(1)"]})
|
||||
reason = _reason("run_argv", {"argv": ["ionice", "-c", "2", "python", "-c", "print(1)"]})
|
||||
assert reason is not None
|
||||
argv_line = next(line for line in reason.splitlines() if line.startswith(" argv:"))
|
||||
assert "ionice -c 2" in argv_line
|
||||
@@ -178,15 +194,43 @@ def test_only_the_interpreter_dash_c_becomes_a_placeholder() -> None:
|
||||
|
||||
|
||||
def test_detached_runner_is_always_reviewed() -> None:
|
||||
assert classify_danger("run_argv", {"argv": ["sleep", "5"]}) is None
|
||||
reason = classify_danger("run_argv", {"argv": ["nohup", "sleep", "5"]})
|
||||
assert _reason("run_argv", {"argv": ["sleep", "5"]}) is None
|
||||
reason = _reason("run_argv", {"argv": ["nohup", "sleep", "5"]})
|
||||
assert reason is not None and "nohup (detached run)" in reason
|
||||
nested = classify_danger("run_argv", {"argv": ["nohup", "python", "-c", "print(1)"]})
|
||||
nested = _reason("run_argv", {"argv": ["nohup", "python", "-c", "print(1)"]})
|
||||
assert nested is not None
|
||||
assert "nohup (detached run)" in nested and "interpreter 'python'" in nested
|
||||
assert "print(1)" in nested
|
||||
|
||||
|
||||
def test_privilege_escalation_refuses_the_yolo_skip() -> None:
|
||||
"""A reason that runs through `sudo`/`doas` keeps its prompt under YOLO mode."""
|
||||
reason = _danger("run_argv", {"argv": ["sudo", "python", "-c", "print(1)"]})
|
||||
assert "interpreter 'python'" in reason.text
|
||||
assert "privilege escalation sudo" in reason.text
|
||||
assert reason.yolo_skip is False
|
||||
|
||||
assert _danger("run_argv", {"argv": ["env", "FOO=1", "doas", "bash", "-c", "ls"]}).yolo_skip is False
|
||||
# A plain interpreter run is still YOLO mode's to drop.
|
||||
assert _danger("run_argv", {"argv": ["bash", "-c", "echo hi"]}).yolo_skip is True
|
||||
# Escalation adds no reason of its own: the denylist gate asks about it.
|
||||
assert classify_danger("run_argv", {"argv": ["sudo", "ls"]}) is None
|
||||
|
||||
|
||||
def test_batch_refuses_the_yolo_skip_when_a_step_escalates() -> None:
|
||||
"""One escalating step keeps the batch prompt; all-plain steps may be dropped."""
|
||||
mixed = _danger(
|
||||
"run_argv_batch",
|
||||
{"steps": [{"argv": ["bash", "-c", "ls"]}, {"argv": ["sudo", "python", "-c", "print(1)"]}]},
|
||||
)
|
||||
assert mixed.yolo_skip is False
|
||||
plain = _danger(
|
||||
"run_argv_batch",
|
||||
{"steps": [{"argv": ["bash", "-c", "ls"]}, {"argv": ["python", "-c", "print(1)"]}]},
|
||||
)
|
||||
assert plain.yolo_skip is True
|
||||
|
||||
|
||||
async def test_confirm_deny_with_comment() -> None:
|
||||
from plyngent.agent.tools import ToolRegistry, tool
|
||||
from plyngent.tools.danger import classify_danger as danger
|
||||
@@ -208,7 +252,7 @@ async def test_confirm_deny_with_comment() -> None:
|
||||
|
||||
def test_shell_confirm_formats_command_placeholder() -> None:
|
||||
script = chr(10).join(["line1", "line2", "line3"])
|
||||
reason = classify_danger(
|
||||
reason = _reason(
|
||||
"run_argv",
|
||||
{"argv": ["bash", "-c", script]},
|
||||
)
|
||||
|
||||
@@ -184,9 +184,9 @@ def test_soft_confirm_reason_matrix() -> None:
|
||||
def test_classify_danger_fetch() -> None:
|
||||
assert classify_danger("fetch_local", {"url": "https://example.com/", "method": "GET"}) is None
|
||||
reason = classify_danger("fetch_local", {"url": "http://example.com/", "method": "GET"})
|
||||
assert reason is not None and "cleartext" in reason
|
||||
assert isinstance(reason, str) and "cleartext" in reason
|
||||
reason2 = classify_danger("fetch_local", {"url": "https://example.com/api", "method": "DELETE"})
|
||||
assert reason2 is not None and "DELETE" in reason2
|
||||
assert isinstance(reason2, str) and "DELETE" in reason2
|
||||
|
||||
|
||||
async def test_fetch_get_post_put_delete(workspace: Path, http_server: str) -> None:
|
||||
|
||||
Reference in New Issue
Block a user