core/agent: optional confirm hook on ToolRegistry execute

Serialise dangerous-tool prompts under a lock so parallel tool rounds
cannot race the TTY confirm.
This commit is contained in:
2026-07-14 23:30:10 +08:00
parent 510dd67a54
commit 35e89c4dcc
3 changed files with 96 additions and 1 deletions
+2
View File
@@ -10,6 +10,8 @@ from .events import ToolCallEvent as ToolCallEvent
from .events import ToolResultEvent as ToolResultEvent from .events import ToolResultEvent as ToolResultEvent
from .loop import DEFAULT_MAX_ROUNDS as DEFAULT_MAX_ROUNDS from .loop import DEFAULT_MAX_ROUNDS as DEFAULT_MAX_ROUNDS
from .loop import run_chat_loop as run_chat_loop 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 ToolDefinition as ToolDefinition
from .tools import ToolRegistry as ToolRegistry from .tools import ToolRegistry as ToolRegistry
from .tools import schema_from_callable as schema_from_callable from .tools import schema_from_callable as schema_from_callable
+35 -1
View File
@@ -1,5 +1,6 @@
from __future__ import annotations from __future__ import annotations
import asyncio
import inspect import inspect
import types import types
from collections.abc import Awaitable, Callable, Mapping 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 from plyngent.typedef import JSONSchema # noqa: TC001
type ToolHandler = Callable[..., Any | Awaitable[Any]] 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] = { _PRIMITIVE_SCHEMA: dict[type, JSONSchema] = {
str: {"type": "string"}, str: {"type": "string"},
@@ -146,9 +149,21 @@ class ToolRegistry:
"""Name → tool definition map with execution helpers.""" """Name → tool definition map with execution helpers."""
_tools: dict[str, ToolDefinition] _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._tools = {}
self._danger = danger
self._on_confirm = on_confirm
self._confirm_lock = asyncio.Lock()
if tools is None: if tools is None:
return return
if isinstance(tools, list): if isinstance(tools, list):
@@ -186,6 +201,22 @@ class ToolRegistry:
return result return result
return msgspec.json.encode(result).decode() 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: async def execute(self, name: str, arguments_json: str) -> str:
"""Run a tool by name; returns a string result (errors become error text).""" """Run a tool by name; returns a string result (errors become error text)."""
definition = self._tools.get(name) definition = self._tools.get(name)
@@ -198,4 +229,7 @@ class ToolRegistry:
if not isinstance(raw_args, dict): if not isinstance(raw_args, dict):
return "error: tool arguments must be a JSON object" return "error: tool arguments must be a JSON object"
args = {str(key): value for key, value in cast("dict[object, object]", raw_args).items()} 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) return await self._invoke(definition, args)
+59
View File
@@ -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"