core/tools: persist soft-confirm grants under session.data

This commit is contained in:
2026-07-24 15:41:39 +08:00
parent 3d5380a3d4
commit 90c07061f3
4 changed files with 120 additions and 6 deletions
+3 -3
View File
@@ -329,14 +329,14 @@ class ToolRegistry:
session = self._session_for_grants()
return session is not None and has_grant(session, name)
def _record_grant(self, name: str, *, tags: ToolTag) -> None:
async def _record_grant(self, name: str, *, tags: ToolTag) -> None:
if not (tags & ToolTag.TRUSTABLE):
return
from plyngent.tools.grants import add_grant
session = self._session_for_grants()
if session is not None:
add_grant(session, name)
await add_grant(session, name)
async def _prompt_soft_confirm(
self,
@@ -352,7 +352,7 @@ class ToolRegistry:
if inspect.isawaitable(decision):
decision = await decision
if decision is True:
self._record_grant(name, tags=tags)
await self._record_grant(name, tags=tags)
return None
if isinstance(decision, str) and decision.strip():
return f"error: tool {name!r} denied by user confirm ({reason}); user comment: {decision.strip()}"
+4 -1
View File
@@ -58,7 +58,8 @@ class SessionState:
data: PersistentDataView[Any] = field(default_factory=_default_session_data)
# Live domain object during migration (also under data["todo"] when views bind).
todo: TodoStack | None = None
# Soft-confirm grants: tool_name → True (Phase 1 key is tool name in session).
# Soft-confirm grants live map: tool_name → True (Phase 1 key is tool name).
# Durable copy lives under data["grants"] (see plyngent.tools.grants).
grants: dict[str, bool] = field(default_factory=dict)
extras: dict[str, Any] = field(default_factory=dict)
@@ -66,9 +67,11 @@ class SessionState:
return bool(self.grants.get(tool_name))
def add_grant(self, tool_name: str) -> None:
"""Update the live map only; prefer :func:`plyngent.tools.grants.add_grant` to persist."""
self.grants[tool_name] = True
def clear_grants(self) -> None:
"""Clear the live map; durable view is updated separately when needed."""
self.grants.clear()
+51 -2
View File
@@ -1,4 +1,14 @@
"""Soft-confirm trust grants (session-scoped)."""
"""Soft-confirm trust grants (session-scoped).
Live map: ``SessionState.grants`` (sync reads for the confirm gate).
Durable tree: ``session.data["grants"]`` (tool name → bool) via
:class:`~plyngent.tools.view.PersistentDataView`.
Hosts that resume a session document should call :func:`hydrate_grants`
after binding :class:`~plyngent.tools.context.SessionState`. Soft-confirm
approval calls :func:`add_grant`, which updates the live map and commits
the view.
"""
from __future__ import annotations
@@ -14,12 +24,51 @@ def grant_key(tool_name: str) -> str:
def has_grant(session: SessionState, tool_name: str) -> bool:
"""Return whether *tool_name* is granted in the live session map."""
return session.has_grant(grant_key(tool_name))
def add_grant(session: SessionState, tool_name: str) -> None:
async def add_grant(session: SessionState, tool_name: str) -> None:
"""Grant *tool_name* and commit the grants map under ``session.data``."""
session.add_grant(grant_key(tool_name))
await persist_grants(session)
def clear_grants(session: SessionState) -> None:
"""Clear the live grant map (sync; e.g. CLI YOLO off / ``/new``).
Does not open a view transaction (callers may be outside an event loop).
Use :func:`persist_grants` when the durable tree should also clear.
"""
session.clear_grants()
async def clear_grants_and_persist(session: SessionState) -> None:
"""Clear live grants and write an empty map to ``session.data``."""
session.clear_grants()
await persist_grants(session)
async def persist_grants(session: SessionState) -> None:
"""Write the live grants map to ``session.data["grants"]``."""
async with session.data as data:
data["grants"].store({key: bool(value) for key, value in session.grants.items()})
async def hydrate_grants(session: SessionState) -> None:
"""Load durable grants into the live map (replace)."""
from typing import cast
raw: object | None = None
async with session.data as data:
try:
raw = data["grants"].load()
except RuntimeError:
raw = None
session.grants.clear()
if not isinstance(raw, dict):
return
blob = cast("dict[object, object]", raw)
for key_obj, value in blob.items():
if value:
session.grants[str(key_obj)] = True
+62
View File
@@ -4,6 +4,8 @@ from __future__ import annotations
from plyngent.agent import 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_yolo_only_skips_yolo_tagged_tools() -> None:
@@ -83,3 +85,63 @@ async def test_untagged_trustable_prompts_every_time() -> None:
assert await reg.execute("always_ask", "{}") == "ok"
assert await reg.execute("always_ask", "{}") == "ok"
assert calls == ["always_ask", "always_ask"]
async def test_trustable_grant_persists_to_session_data() -> 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
store = MemoryViewStore({})
session = SessionState(session_id=1, data=session_data_view(store=store))
reg = ToolRegistry(
[trust_me],
danger=danger,
on_confirm=confirm,
yolo=False,
session_state=session,
)
assert await reg.execute("trust_me", "{}") == "ok"
assert has_grant(session, "trust_me")
loaded = await store.load()
assert isinstance(loaded, dict)
grants = loaded.get("grants")
assert isinstance(grants, dict)
assert grants.get("trust_me") is True
# Fresh session + same store: hydrate restores the grant (no second prompt).
session2 = SessionState(session_id=1, data=session_data_view(store=store))
await hydrate_grants(session2)
assert has_grant(session2, "trust_me")
reg2 = ToolRegistry(
[trust_me],
danger=danger,
on_confirm=confirm,
yolo=False,
session_state=session2,
)
assert await reg2.execute("trust_me", "{}") == "ok"
assert calls == ["trust_me"]
async def test_clear_grants_and_persist() -> None:
store = MemoryViewStore({})
session = SessionState(session_id=2, data=session_data_view(store=store))
session.add_grant("tool_a")
await persist_grants(session)
clear_grants(session)
assert not has_grant(session, "tool_a")
# Live clear does not rewrite the store; explicit persist does.
await persist_grants(session)
loaded = await store.load()
assert isinstance(loaded, dict)
assert loaded.get("grants") == {}