diff --git a/src/plyngent/agent/__init__.py b/src/plyngent/agent/__init__.py index f3d9c50..e158b9f 100644 --- a/src/plyngent/agent/__init__.py +++ b/src/plyngent/agent/__init__.py @@ -10,6 +10,8 @@ from .events import ToolCallEvent as ToolCallEvent from .events import ToolResultEvent as ToolResultEvent from .loop import DEFAULT_MAX_ROUNDS as DEFAULT_MAX_ROUNDS from .loop import run_chat_loop as run_chat_loop +from .tools import DangerClassifier as DangerClassifier +from .tools import ToolConfirmHook as ToolConfirmHook from .tools import ToolDefinition as ToolDefinition from .tools import ToolRegistry as ToolRegistry from .tools import schema_from_callable as schema_from_callable diff --git a/src/plyngent/agent/tools.py b/src/plyngent/agent/tools.py index f7070ef..025b6e7 100644 --- a/src/plyngent/agent/tools.py +++ b/src/plyngent/agent/tools.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import inspect import types from collections.abc import Awaitable, Callable, Mapping @@ -11,6 +12,8 @@ from plyngent.lmproto.openai_compatible.model import ToolFunction, ToolFunctionI from plyngent.typedef import JSONSchema # noqa: TC001 type ToolHandler = Callable[..., Any | Awaitable[Any]] +type DangerClassifier = Callable[[str, Mapping[str, object]], str | None] +type ToolConfirmHook = Callable[[str, Mapping[str, object], str], bool] _PRIMITIVE_SCHEMA: dict[type, JSONSchema] = { str: {"type": "string"}, @@ -146,9 +149,21 @@ class ToolRegistry: """Name → tool definition map with execution helpers.""" _tools: dict[str, ToolDefinition] + _danger: DangerClassifier | None + _on_confirm: ToolConfirmHook | None + _confirm_lock: asyncio.Lock - def __init__(self, tools: Mapping[str, ToolDefinition] | list[ToolDefinition] | None = None) -> None: + def __init__( + self, + tools: Mapping[str, ToolDefinition] | list[ToolDefinition] | None = None, + *, + danger: DangerClassifier | None = None, + on_confirm: ToolConfirmHook | None = None, + ) -> None: self._tools = {} + self._danger = danger + self._on_confirm = on_confirm + self._confirm_lock = asyncio.Lock() if tools is None: return if isinstance(tools, list): @@ -186,6 +201,22 @@ class ToolRegistry: return result return msgspec.json.encode(result).decode() + async def _maybe_confirm(self, name: str, args: dict[str, object]) -> str | None: + """Return an error string if the user denies a dangerous tool, else None.""" + if self._danger is None or self._on_confirm is None: + return None + reason = self._danger(name, args) + if reason is None: + return None + async with self._confirm_lock: + # Re-check under the lock so parallel tools do not race the prompt. + reason = self._danger(name, args) + if reason is None: + return None + if self._on_confirm(name, args, reason): + return None + return f"error: user denied tool {name!r} ({reason})" + async def execute(self, name: str, arguments_json: str) -> str: """Run a tool by name; returns a string result (errors become error text).""" definition = self._tools.get(name) @@ -198,4 +229,7 @@ class ToolRegistry: if not isinstance(raw_args, dict): return "error: tool arguments must be a JSON object" args = {str(key): value for key, value in cast("dict[object, object]", raw_args).items()} + denied = await self._maybe_confirm(name, args) + if denied is not None: + return denied return await self._invoke(definition, args) diff --git a/tests/test_agent/test_confirm.py b/tests/test_agent/test_confirm.py new file mode 100644 index 0000000..96d3972 --- /dev/null +++ b/tests/test_agent/test_confirm.py @@ -0,0 +1,59 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +from plyngent.agent import ToolRegistry, tool +from plyngent.tools.danger import classify_danger + +if TYPE_CHECKING: + from collections.abc import Mapping + + +@tool +def delete_path(path: str, *, recursive: bool = False) -> str: + del recursive + return f"deleted {path}" + + +@tool +def read_file(path: str) -> str: + return f"read {path}" + + +async def test_confirm_allow() -> None: + decisions: list[str] = [] + + def on_confirm(name: str, args: Mapping[str, object], reason: str) -> bool: + del args + decisions.append(f"{name}:{reason}") + return True + + registry = ToolRegistry( + [delete_path, read_file], + danger=classify_danger, + on_confirm=on_confirm, + ) + assert await registry.execute("delete_path", '{"path": "a.txt"}') == "deleted a.txt" + assert len(decisions) == 1 + assert await registry.execute("read_file", '{"path": "a.txt"}') == "read a.txt" + assert len(decisions) == 1 + + +async def test_confirm_deny() -> None: + def on_confirm(name: str, args: Mapping[str, object], reason: str) -> bool: + del name, args, reason + return False + + registry = ToolRegistry( + [delete_path], + danger=classify_danger, + on_confirm=on_confirm, + ) + out = await registry.execute("delete_path", '{"path": "a.txt"}') + assert "denied" in out + assert "delete" in out + + +async def test_no_hooks_skips_confirm() -> None: + registry = ToolRegistry([delete_path]) + assert await registry.execute("delete_path", '{"path": "a.txt"}') == "deleted a.txt"