From 90c07061f33bec4689fe7ebe21a4085c9b731c76 Mon Sep 17 00:00:00 2001 From: worldmozara Date: Fri, 24 Jul 2026 15:41:39 +0800 Subject: [PATCH] core/tools: persist soft-confirm grants under session.data --- src/plyngent/agent/tools.py | 6 +-- src/plyngent/tools/context.py | 5 ++- src/plyngent/tools/grants.py | 53 ++++++++++++++++++++++- tests/test_tools/test_confirm_tags.py | 62 +++++++++++++++++++++++++++ 4 files changed, 120 insertions(+), 6 deletions(-) diff --git a/src/plyngent/agent/tools.py b/src/plyngent/agent/tools.py index 487e1c0..f6f4d59 100644 --- a/src/plyngent/agent/tools.py +++ b/src/plyngent/agent/tools.py @@ -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()}" diff --git a/src/plyngent/tools/context.py b/src/plyngent/tools/context.py index 68b4585..ede73a4 100644 --- a/src/plyngent/tools/context.py +++ b/src/plyngent/tools/context.py @@ -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() diff --git a/src/plyngent/tools/grants.py b/src/plyngent/tools/grants.py index e589d54..73c89ae 100644 --- a/src/plyngent/tools/grants.py +++ b/src/plyngent/tools/grants.py @@ -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 diff --git a/tests/test_tools/test_confirm_tags.py b/tests/test_tools/test_confirm_tags.py index 2a03ddc..290fd1e 100644 --- a/tests/test_tools/test_confirm_tags.py +++ b/tests/test_tools/test_confirm_tags.py @@ -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") == {}