mirror of
https://github.com/NCBM/plyngent.git
synced 2026-07-25 08:04:57 +08:00
core/tools: persist soft-confirm grants under session.data
This commit is contained in:
@@ -329,14 +329,14 @@ class ToolRegistry:
|
|||||||
session = self._session_for_grants()
|
session = self._session_for_grants()
|
||||||
return session is not None and has_grant(session, name)
|
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):
|
if not (tags & ToolTag.TRUSTABLE):
|
||||||
return
|
return
|
||||||
from plyngent.tools.grants import add_grant
|
from plyngent.tools.grants import add_grant
|
||||||
|
|
||||||
session = self._session_for_grants()
|
session = self._session_for_grants()
|
||||||
if session is not None:
|
if session is not None:
|
||||||
add_grant(session, name)
|
await add_grant(session, name)
|
||||||
|
|
||||||
async def _prompt_soft_confirm(
|
async def _prompt_soft_confirm(
|
||||||
self,
|
self,
|
||||||
@@ -352,7 +352,7 @@ class ToolRegistry:
|
|||||||
if inspect.isawaitable(decision):
|
if inspect.isawaitable(decision):
|
||||||
decision = await decision
|
decision = await decision
|
||||||
if decision is True:
|
if decision is True:
|
||||||
self._record_grant(name, tags=tags)
|
await self._record_grant(name, tags=tags)
|
||||||
return None
|
return None
|
||||||
if isinstance(decision, str) and decision.strip():
|
if isinstance(decision, str) and decision.strip():
|
||||||
return f"error: tool {name!r} denied by user confirm ({reason}); user comment: {decision.strip()}"
|
return f"error: tool {name!r} denied by user confirm ({reason}); user comment: {decision.strip()}"
|
||||||
|
|||||||
@@ -58,7 +58,8 @@ class SessionState:
|
|||||||
data: PersistentDataView[Any] = field(default_factory=_default_session_data)
|
data: PersistentDataView[Any] = field(default_factory=_default_session_data)
|
||||||
# Live domain object during migration (also under data["todo"] when views bind).
|
# Live domain object during migration (also under data["todo"] when views bind).
|
||||||
todo: TodoStack | None = None
|
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)
|
grants: dict[str, bool] = field(default_factory=dict)
|
||||||
extras: dict[str, Any] = 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))
|
return bool(self.grants.get(tool_name))
|
||||||
|
|
||||||
def add_grant(self, tool_name: str) -> None:
|
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
|
self.grants[tool_name] = True
|
||||||
|
|
||||||
def clear_grants(self) -> None:
|
def clear_grants(self) -> None:
|
||||||
|
"""Clear the live map; durable view is updated separately when needed."""
|
||||||
self.grants.clear()
|
self.grants.clear()
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -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
|
from __future__ import annotations
|
||||||
|
|
||||||
@@ -14,12 +24,51 @@ def grant_key(tool_name: str) -> str:
|
|||||||
|
|
||||||
|
|
||||||
def has_grant(session: SessionState, tool_name: str) -> bool:
|
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))
|
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))
|
session.add_grant(grant_key(tool_name))
|
||||||
|
await persist_grants(session)
|
||||||
|
|
||||||
|
|
||||||
def clear_grants(session: SessionState) -> None:
|
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()
|
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
|
||||||
|
|||||||
@@ -4,6 +4,8 @@ from __future__ import annotations
|
|||||||
|
|
||||||
from plyngent.agent import ToolRegistry, ToolTag, tool
|
from plyngent.agent import ToolRegistry, ToolTag, tool
|
||||||
from plyngent.tools.context import SessionState, bind_tool_context
|
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:
|
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 await reg.execute("always_ask", "{}") == "ok"
|
assert await reg.execute("always_ask", "{}") == "ok"
|
||||||
assert calls == ["always_ask", "always_ask"]
|
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") == {}
|
||||||
|
|||||||
Reference in New Issue
Block a user