test/tools: catalog, grants, and async tool helpers

This commit is contained in:
2026-07-24 14:57:03 +08:00
parent 2f6133f685
commit 12c8adda3a
12 changed files with 285 additions and 35 deletions
+6
View File
@@ -0,0 +1,6 @@
"""Shared pytest fixtures.
Ephemeral ``@tool`` tests should pass ``register=False`` so they do not pollute
the process catalog. Do not autouse an empty ``catalog_scope`` here: it would
hide builtins from ``default_tool_definitions`` during tests.
"""
+2 -2
View File
@@ -9,13 +9,13 @@ if TYPE_CHECKING:
from collections.abc import Mapping from collections.abc import Mapping
@tool @tool(register=False)
def delete_path(path: str, *, recursive: bool = False) -> str: def delete_path(path: str, *, recursive: bool = False) -> str:
del recursive del recursive
return f"deleted {path}" return f"deleted {path}"
@tool @tool(register=False)
def read_file(path: str) -> str: def read_file(path: str) -> str:
return f"read {path}" return f"read {path}"
+7 -7
View File
@@ -268,7 +268,7 @@ async def test_chat_agent_accumulates_session_usage() -> None:
async def test_chat_agent_turn_usage_sums_tool_rounds() -> None: async def test_chat_agent_turn_usage_sums_tool_rounds() -> None:
"""Multi-round tool loop: turn usage is billing sum; last_request is final call.""" """Multi-round tool loop: turn usage is billing sum; last_request is final call."""
@tool @tool(register=False)
def ping() -> str: def ping() -> str:
return "pong" return "pong"
@@ -437,7 +437,7 @@ async def test_non_stream_yields_reasoning() -> None:
async def test_run_chat_loop_with_tools() -> None: async def test_run_chat_loop_with_tools() -> None:
@tool @tool(register=False)
def add(a: int, b: int) -> int: def add(a: int, b: int) -> int:
return a + b return a + b
@@ -470,7 +470,7 @@ async def test_run_chat_loop_with_tools() -> None:
async def test_max_rounds() -> None: async def test_max_rounds() -> None:
@tool @tool(register=False)
def ping() -> str: def ping() -> str:
return "pong" return "pong"
@@ -494,7 +494,7 @@ async def test_max_rounds() -> None:
async def test_max_rounds_continue_hook() -> None: async def test_max_rounds_continue_hook() -> None:
@tool @tool(register=False)
def ping() -> str: def ping() -> str:
return "pong" return "pong"
@@ -537,7 +537,7 @@ async def test_max_rounds_continue_hook() -> None:
async def test_max_rounds_async_continue_hook() -> None: async def test_max_rounds_async_continue_hook() -> None:
@tool @tool(register=False)
def ping() -> str: def ping() -> str:
return "pong" return "pong"
@@ -721,7 +721,7 @@ async def test_retry_keeps_committed_tools_after_second_round_fails() -> None:
session = await store.create_session(name="t") session = await store.create_session(name="t")
calls: list[str] = [] calls: list[str] = []
@tool @tool(register=False)
def side_effect() -> str: def side_effect() -> str:
calls.append("ran") calls.append("ran")
return "done-once" return "done-once"
@@ -809,7 +809,7 @@ async def test_chat_agent_system_prompt_prepended() -> None:
async def test_tool_result_char_budget() -> None: async def test_tool_result_char_budget() -> None:
@tool @tool(register=False)
def big() -> str: def big() -> str:
return "x" * 100 return "x" * 100
+2 -2
View File
@@ -126,7 +126,7 @@ async def test_parallel_tool_calls_run_concurrently() -> None:
gate = asyncio.Event() gate = asyncio.Event()
wait_timeout = 2.0 wait_timeout = 2.0
@tool @tool(register=False)
async def slow_a() -> str: async def slow_a() -> str:
nonlocal started nonlocal started
started += 1 started += 1
@@ -135,7 +135,7 @@ async def test_parallel_tool_calls_run_concurrently() -> None:
_ = await asyncio.wait_for(gate.wait(), timeout=wait_timeout) _ = await asyncio.wait_for(gate.wait(), timeout=wait_timeout)
return "a" return "a"
@tool @tool(register=False)
async def slow_b() -> str: async def slow_b() -> str:
nonlocal started nonlocal started
started += 1 started += 1
+7 -7
View File
@@ -4,7 +4,7 @@ from plyngent.agent import ToolRegistry, schema_from_callable, tool
def test_tool_decorator_infers_schema() -> None: def test_tool_decorator_infers_schema() -> None:
@tool @tool(register=False)
def add(a: int, b: int = 0) -> int: def add(a: int, b: int = 0) -> int:
"""Add two numbers.""" """Add two numbers."""
return a + b return a + b
@@ -18,7 +18,7 @@ def test_tool_decorator_infers_schema() -> None:
def test_tool_decorator_overrides() -> None: def test_tool_decorator_overrides() -> None:
@tool(name="ping", description="health") @tool(name="ping", description="health", register=False)
def health() -> str: def health() -> str:
return "ok" return "ok"
@@ -36,11 +36,11 @@ def test_schema_from_callable_optional() -> None:
async def test_registry_execute_sync_and_async() -> None: async def test_registry_execute_sync_and_async() -> None:
@tool @tool(register=False)
def echo(text: str) -> str: def echo(text: str) -> str:
return text return text
@tool @tool(register=False)
async def upper(text: str) -> str: async def upper(text: str) -> str:
return text.upper() return text.upper()
@@ -50,7 +50,7 @@ async def test_registry_execute_sync_and_async() -> None:
async def test_registry_unknown_and_bad_json() -> None: async def test_registry_unknown_and_bad_json() -> None:
@tool @tool(register=False)
def noop() -> str: def noop() -> str:
return "ok" return "ok"
@@ -60,7 +60,7 @@ async def test_registry_unknown_and_bad_json() -> None:
async def test_registry_handler_error() -> None: async def test_registry_handler_error() -> None:
@tool @tool(register=False)
def boom() -> str: def boom() -> str:
msg = "nope" msg = "nope"
raise RuntimeError(msg) raise RuntimeError(msg)
@@ -71,7 +71,7 @@ async def test_registry_handler_error() -> None:
def test_tool_items() -> None: def test_tool_items() -> None:
@tool @tool(register=False)
def f(x: int) -> int: def f(x: int) -> int:
return x return x
+6 -2
View File
@@ -432,16 +432,20 @@ async def test_yolo_rebuilds_tool_registry(tmp_path: Path) -> None:
) )
try: try:
assert st.agent.tools is not None assert st.agent.tools is not None
# Soft-confirm hooks stay installed; YOLO only auto-allows YOLO-tagged tools.
assert st.agent.tools.soft_confirm is True assert st.agent.tools.soft_confirm is True
assert st.agent.tools.yolo is False
st.set_yolo("on") st.set_yolo("on")
assert st.agent.tools is not None assert st.agent.tools is not None
assert st.agent.tools.soft_confirm is False assert st.agent.tools.soft_confirm is True
assert st.agent.tools.yolo is True
st.set_yolo("once") st.set_yolo("once")
assert st.agent.tools is not None assert st.agent.tools is not None
assert st.agent.tools.soft_confirm is False assert st.agent.tools.yolo is True
st.set_yolo("off") st.set_yolo("off")
assert st.agent.tools is not None assert st.agent.tools is not None
assert st.agent.tools.soft_confirm is True assert st.agent.tools.soft_confirm is True
assert st.agent.tools.yolo is False
finally: finally:
await memory.close() await memory.close()
+14 -1
View File
@@ -1,5 +1,6 @@
from __future__ import annotations from __future__ import annotations
import asyncio
import inspect import inspect
from typing import TYPE_CHECKING, Any, cast from typing import TYPE_CHECKING, Any, cast
@@ -8,8 +9,20 @@ if TYPE_CHECKING:
def call_sync(definition: ToolDefinition, *args: object, **kwargs: object) -> str: def call_sync(definition: ToolDefinition, *args: object, **kwargs: object) -> str:
"""Invoke a tool handler from a **sync** test.
Async handlers are run with ``asyncio.run``. Inside an already-running
event loop (async tests), use :func:`call_async` instead.
"""
result: Any = definition.handler(*args, **kwargs) result: Any = definition.handler(*args, **kwargs)
assert not inspect.isawaitable(result) if inspect.isawaitable(result):
try:
asyncio.get_running_loop()
except RuntimeError:
result = asyncio.run(result)
else:
msg = "call_sync cannot await inside a running event loop; use call_async"
raise RuntimeError(msg)
return cast("str", result) return cast("str", result)
+94
View File
@@ -0,0 +1,94 @@
"""Tool catalog, tags, and default select parity."""
from __future__ import annotations
import inspect
from plyngent.agent import ToolTag, tool
from plyngent.tools import DEFAULT_TOOLS, default_tool_definitions, register_builtin_tools
from plyngent.tools.catalog import ToolCatalog, ToolSource, catalog_scope, get_catalog, registration_source
def test_default_tool_names_match_legacy_lists() -> None:
register_builtin_tools()
selected = default_tool_definitions(surface="local")
legacy = [t.name for t in DEFAULT_TOOLS]
assert sorted(t.name for t in selected) == sorted(legacy)
assert len(selected) == len(legacy)
def test_all_builtins_are_async_and_tagged() -> None:
tools = default_tool_definitions()
assert tools
for definition in tools:
assert inspect.iscoroutinefunction(definition.handler), definition.name
assert definition.tags & (ToolTag.LOCAL | ToolTag.PUBLIC)
assert definition.tags & ToolTag.LOCAL or definition.tags & ToolTag.PUBLIC
def test_public_surface_is_subset() -> None:
local = {t.name for t in default_tool_definitions(surface="local")}
public = {t.name for t in default_tool_definitions(surface="public")}
assert public
assert public <= local
# Todo series is the main PUBLIC surface today.
assert "todo_list" in public
assert "read_file" not in public
def test_catalog_scope_empty_isolates_registration() -> None:
with catalog_scope(empty=True) as catalog:
@tool(name="only_in_scope", register=True)
async def only_in_scope() -> str:
return "ok"
_ = only_in_scope
assert catalog.get("only_in_scope") is not None
# default_tool_definitions re-seeds builtins into the override; plugin-like
# names from this scope stay local to the override and leave process catalog.
assert get_catalog().get("only_in_scope") is None
def test_collision_refuses_shadow() -> None:
catalog = ToolCatalog()
@tool(register=False)
async def alpha() -> str:
return "a"
catalog.register(alpha, source=ToolSource(kind="builtin"))
try:
catalog.register(alpha, source=ToolSource(kind="plugin", plugin_id="acme"))
raise AssertionError("expected collision")
except ValueError as exc:
assert "collision" in str(exc)
def test_registration_source_context() -> None:
with catalog_scope(empty=True) as catalog:
with registration_source(ToolSource(kind="plugin", plugin_id="acme")):
@tool(name="plugin_tool")
async def plugin_tool() -> str:
return "p"
_ = plugin_tool
entry = catalog.get("plugin_tool")
assert entry is not None
assert entry.source.kind == "plugin"
assert entry.source.plugin_id == "acme"
def test_tool_tags_reject_neither_surface() -> None:
try:
@tool(tags=ToolTag.YOLO, register=False) # type: ignore[arg-type]
async def bad() -> str:
return "x"
_ = bad
raise AssertionError("expected ValueError")
except ValueError as exc:
assert "LOCAL" in str(exc)
+85
View File
@@ -0,0 +1,85 @@
"""Tag-aware soft confirm: YOLO bit + TRUSTABLE grants."""
from __future__ import annotations
from plyngent.agent import ToolRegistry, ToolTag, tool
from plyngent.tools.context import SessionState, bind_tool_context
async def test_yolo_only_skips_yolo_tagged_tools() -> None:
calls: list[str] = []
@tool(tags=ToolTag.LOCAL | ToolTag.YOLO, register=False)
async def yolo_ok() -> str:
return "y"
@tool(tags=ToolTag.LOCAL, register=False)
async def no_yolo() -> str:
return "n"
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, no_yolo],
danger=danger,
on_confirm=confirm,
yolo=True,
)
assert await reg.execute("yolo_ok", "{}") == "y"
assert calls == []
assert await reg.execute("no_yolo", "{}") == "n"
assert calls == ["no_yolo"]
async def test_trustable_grant_once() -> None:
calls: list[str] = []
@tool(tags=ToolTag.LOCAL | ToolTag.TRUSTABLE, register=False)
async def trust_me() -> str:
return "ok"
def danger(name: str, _args: object) -> str | None:
return "soft"
async def confirm(name: str, _args: object, _reason: str) -> bool:
calls.append(name)
return True
session = SessionState(session_id=1)
reg = ToolRegistry(
[trust_me],
danger=danger,
on_confirm=confirm,
yolo=False,
auto_bind_state=True,
session_state=session,
)
with bind_tool_context(session=session):
assert await reg.execute("trust_me", "{}") == "ok"
assert await reg.execute("trust_me", "{}") == "ok"
assert calls == ["trust_me"]
async def test_untagged_trustable_prompts_every_time() -> None:
calls: list[str] = []
@tool(tags=ToolTag.LOCAL, register=False)
async def always_ask() -> str:
return "ok"
def danger(name: str, _args: object) -> str | None:
return "soft"
async def confirm(name: str, _args: object, _reason: str) -> bool:
calls.append(name)
return True
reg = ToolRegistry([always_ask], danger=danger, on_confirm=confirm, yolo=False)
assert await reg.execute("always_ask", "{}") == "ok"
assert await reg.execute("always_ask", "{}") == "ok"
assert calls == ["always_ask", "always_ask"]
+1 -1
View File
@@ -100,7 +100,7 @@ async def test_confirm_deny_with_comment() -> None:
from plyngent.agent.tools import ToolRegistry, tool from plyngent.agent.tools import ToolRegistry, tool
from plyngent.tools.danger import classify_danger as danger from plyngent.tools.danger import classify_danger as danger
@tool @tool(register=False)
def delete_path(path: str) -> str: def delete_path(path: str) -> str:
return f"deleted {path}" return f"deleted {path}"
+13 -13
View File
@@ -185,7 +185,7 @@ async def test_run_command_env(workspace: object) -> None:
async def test_pty_open_read_close(workspace: object) -> None: async def test_pty_open_read_close(workspace: object) -> None:
del workspace del workspace
try: try:
opened = call_sync(open_pty, _py("import time; time.sleep(30)")) opened = await call_async(open_pty, _py("import time; time.sleep(30)"))
assert "session_id=" in opened assert "session_id=" in opened
session_id = _session_id(opened) session_id = _session_id(opened)
data = await call_async(read_pty, session_id, timeout=0.05) data = await call_async(read_pty, session_id, timeout=0.05)
@@ -207,7 +207,7 @@ def test_pty_denied(workspace: object) -> None:
async def test_pty_echo_output(workspace: object) -> None: async def test_pty_echo_output(workspace: object) -> None:
del workspace del workspace
try: try:
opened = call_sync(open_pty, _py("print('hello-pty')")) opened = await call_async(open_pty, _py("print('hello-pty')"))
session_id = _session_id(opened) session_id = _session_id(opened)
text = await call_async(read_pty, session_id, timeout=2.0, until="hello-pty") text = await call_async(read_pty, session_id, timeout=2.0, until="hello-pty")
assert "hello-pty" in text assert "hello-pty" in text
@@ -222,7 +222,7 @@ async def test_write_pty(workspace: object) -> None:
del workspace del workspace
try: try:
# Line-oriented echo (portable stand-in for ``cat``). # Line-oriented echo (portable stand-in for ``cat``).
opened = call_sync( opened = await call_async(
open_pty, open_pty,
_py( _py(
"import sys\n" "import sys\n"
@@ -235,7 +235,7 @@ async def test_write_pty(workspace: object) -> None:
), ),
) )
session_id = _session_id(opened) session_id = _session_id(opened)
written = call_sync(write_pty, session_id, "pty-input\n") written = await call_async(write_pty, session_id, "pty-input\n")
assert "wrote=" in written assert "wrote=" in written
text = await call_async(read_pty, session_id, timeout=2.0, until="pty-input") text = await call_async(read_pty, session_id, timeout=2.0, until="pty-input")
assert "pty-input" in text assert "pty-input" in text
@@ -258,7 +258,7 @@ async def test_ask_into_pty_writes_without_leaking_secret(workspace: object) ->
secret = "super-secret-password-xyz" secret = "super-secret-password-xyz"
backend = ScriptedBackend([], secrets=[secret]) backend = ScriptedBackend([], secrets=[secret])
try: try:
opened = call_sync( opened = await call_async(
open_pty, open_pty,
_py( _py(
"import sys\n" "import sys\n"
@@ -302,7 +302,7 @@ async def test_ask_into_pty_empty_cancels(workspace: object) -> None:
backend = ScriptedBackend([], secrets=[""]) backend = ScriptedBackend([], secrets=[""])
try: try:
opened = call_sync(open_pty, _py("import time; time.sleep(30)")) opened = await call_async(open_pty, _py("import time; time.sleep(30)"))
session_id = _session_id(opened) session_id = _session_id(opened)
with temporary_backend(backend): with temporary_backend(backend):
result = await call_async(ask_into_pty, session_id, "Pass?", secret=True) result = await call_async(ask_into_pty, session_id, "Pass?", secret=True)
@@ -319,7 +319,7 @@ async def test_ask_into_pty_nonsecret_line(workspace: object) -> None:
backend = ScriptedBackend(["yes"]) backend = ScriptedBackend(["yes"])
try: try:
opened = call_sync( opened = await call_async(
open_pty, open_pty,
_py("import sys\nline = sys.stdin.readline()\nsys.stdout.write(line)\nsys.stdout.flush()\n"), _py("import sys\nline = sys.stdin.readline()\nsys.stdout.write(line)\nsys.stdout.flush()\n"),
) )
@@ -338,7 +338,7 @@ async def test_ask_into_pty_nonsecret_line(workspace: object) -> None:
async def test_pty_exec_failure_surfaces(workspace: object) -> None: async def test_pty_exec_failure_surfaces(workspace: object) -> None:
del workspace del workspace
try: try:
opened = call_sync(open_pty, ["definitely-not-a-real-binary-xyz"]) opened = await call_async(open_pty, ["definitely-not-a-real-binary-xyz"])
# Windows ConPTY may fail at open; POSIX may open then fail on exec. # Windows ConPTY may fail at open; POSIX may open then fail on exec.
if "error:" in opened and "session_id=" not in opened: if "error:" in opened and "session_id=" not in opened:
assert "not found" in opened.lower() or "failed" in opened.lower() assert "not found" in opened.lower() or "failed" in opened.lower()
@@ -394,7 +394,7 @@ async def test_pty_output_budget(workspace: object) -> None:
try: try:
PtyManager.set_limit_continue_hook(None) PtyManager.set_limit_continue_hook(None)
PtyManager.configure(session_output_budget=64) PtyManager.configure(session_output_budget=64)
opened = call_sync(open_pty, _py("print('x' * 1000)")) opened = await call_async(open_pty, _py("print('x' * 1000)"))
session_id = _session_id(opened) session_id = _session_id(opened)
# Drain until budget exhausted or process ends. # Drain until budget exhausted or process ends.
budget_hit = False budget_hit = False
@@ -423,7 +423,7 @@ async def test_pty_output_budget_is_per_session(workspace: object) -> None:
PtyManager.configure(session_output_budget=1024) PtyManager.configure(session_output_budget=1024)
class_budget = PtyManager.session_output_budget class_budget = PtyManager.session_output_budget
PtyManager.set_limit_continue_hook(lambda _reason: True) PtyManager.set_limit_continue_hook(lambda _reason: True)
opened = call_sync(open_pty, _py("print('x' * 200)")) opened = await call_async(open_pty, _py("print('x' * 200)"))
session_id = _session_id(opened) session_id = _session_id(opened)
session = PtyManager.get(session_id) session = PtyManager.get(session_id)
assert session is not None assert session is not None
@@ -495,7 +495,7 @@ def test_close_all_empty_is_noop() -> None:
async def test_read_pty_sanitizes_esc(workspace: object) -> None: async def test_read_pty_sanitizes_esc(workspace: object) -> None:
del workspace del workspace
try: try:
opened = call_sync(open_pty, _py("print(chr(0x1B)+'[31mred'+chr(0x1B)+'[0m')")) opened = await call_async(open_pty, _py("print(chr(0x1B)+'[31mred'+chr(0x1B)+'[0m')"))
session_id = _session_id(opened) session_id = _session_id(opened)
text = await call_async(read_pty, session_id, timeout=2.0, until="red") text = await call_async(read_pty, session_id, timeout=2.0, until="red")
assert "red" in text assert "red" in text
@@ -510,7 +510,7 @@ async def test_read_pty_sanitizes_esc(workspace: object) -> None:
async def test_write_pty_keys_ctrl_escape(workspace: object) -> None: async def test_write_pty_keys_ctrl_escape(workspace: object) -> None:
del workspace del workspace
try: try:
opened = call_sync( opened = await call_async(
open_pty, open_pty,
_py( _py(
"import sys\n" "import sys\n"
@@ -520,7 +520,7 @@ async def test_write_pty_keys_ctrl_escape(workspace: object) -> None:
), ),
) )
session_id = _session_id(opened) session_id = _session_id(opened)
written = call_sync(write_pty_keys, session_id, "ctrl+c") written = await call_async(write_pty_keys, session_id, "ctrl+c")
assert "wrote=1" in written assert "wrote=1" in written
text = await call_async(read_pty, session_id, timeout=2.0) text = await call_async(read_pty, session_id, timeout=2.0)
assert "error" not in text.lower() or "alive=" in text assert "error" not in text.lower() or "alive=" in text
+48
View File
@@ -0,0 +1,48 @@
"""PersistentDataView transaction behavior."""
from __future__ import annotations
import pytest
from plyngent.agent.todo_stack import TodoStack
from plyngent.tools.view import MemoryViewStore, PersistentDataView
async def test_view_txn_commit() -> None:
store = MemoryViewStore({})
root: PersistentDataView[dict[str, object]] = PersistentDataView(store, bound_type=dict)
async with root as data:
data["todo"].store({"groups": [], "next_id": 1})
loaded = await store.load()
assert isinstance(loaded, dict)
assert "todo" in loaded
async def test_view_rollback_on_error() -> None:
store = MemoryViewStore({"keep": 1})
root: PersistentDataView[dict[str, object]] = PersistentDataView(store)
with pytest.raises(RuntimeError, match="boom"):
async with root as data:
data.store({"keep": 2})
msg = "boom"
raise RuntimeError(msg)
assert await store.load() == {"keep": 1}
async def test_typed_todo_stack() -> None:
store = MemoryViewStore({})
root: PersistentDataView[dict[str, object]] = PersistentDataView(store, bound_type=dict)
async with root as data:
todo = data["todo"].typed(TodoStack)
_ = todo.push_group(["A", "B"])
assert todo.depth == 1
loaded = await store.load()
assert isinstance(loaded, dict)
assert isinstance(loaded.get("todo"), TodoStack)
async def test_mutate_outside_txn_errors() -> None:
store = MemoryViewStore({})
root: PersistentDataView[dict[str, object]] = PersistentDataView(store)
with pytest.raises(RuntimeError, match="transaction"):
root.store({"a": 1})