From 4feb0e0d262dca4aef110b85f12063b84546a140 Mon Sep 17 00:00:00 2001 From: worldmozara Date: Fri, 24 Jul 2026 14:56:19 +0800 Subject: [PATCH] core/tools: catalog, tags, state context, and PersistentDataView --- src/plyngent/agent/__init__.py | 1 + src/plyngent/agent/tools.py | 207 ++++++++++++++++++++++--- src/plyngent/tools/__init__.py | 16 ++ src/plyngent/tools/catalog.py | 264 ++++++++++++++++++++++++++++++++ src/plyngent/tools/context.py | 114 ++++++++++++++ src/plyngent/tools/grants.py | 25 +++ src/plyngent/tools/view.py | 270 +++++++++++++++++++++++++++++++++ 7 files changed, 877 insertions(+), 20 deletions(-) create mode 100644 src/plyngent/tools/catalog.py create mode 100644 src/plyngent/tools/context.py create mode 100644 src/plyngent/tools/grants.py create mode 100644 src/plyngent/tools/view.py diff --git a/src/plyngent/agent/__init__.py b/src/plyngent/agent/__init__.py index 58eebe6..00f1349 100644 --- a/src/plyngent/agent/__init__.py +++ b/src/plyngent/agent/__init__.py @@ -20,6 +20,7 @@ from .tools import DangerClassifier as DangerClassifier from .tools import ToolConfirmHook as ToolConfirmHook from .tools import ToolDefinition as ToolDefinition from .tools import ToolRegistry as ToolRegistry +from .tools import ToolTag as ToolTag from .tools import schema_from_callable as schema_from_callable from .tools import tool as tool from .usage import TokenUsage as TokenUsage diff --git a/src/plyngent/agent/tools.py b/src/plyngent/agent/tools.py index cb19c88..87362a6 100644 --- a/src/plyngent/agent/tools.py +++ b/src/plyngent/agent/tools.py @@ -4,6 +4,7 @@ import asyncio import inspect import types from collections.abc import Awaitable, Callable, Mapping +from enum import Flag, auto from typing import Any, cast, get_args, get_origin, get_type_hints, overload import msgspec @@ -28,6 +29,21 @@ _PRIMITIVE_SCHEMA: dict[type, JSONSchema] = { } +class ToolTag(Flag): + """Host policy / affinity bits for a tool (not I/O taxonomy). + + Default when omitted on ``@tool`` is :attr:`LOCAL`. A tool should set at + least one of :attr:`LOCAL` or :attr:`PUBLIC` (register rejects neither). + """ + + LOCAL = auto() # local agent frontend (CLI) + PUBLIC = auto() # may be exposed on shared / multi-tenant frontends + TRUSTABLE = auto() # soft-confirm: grant once, then reuse + YOLO = auto() # soft-confirm: eligible for YOLO auto-approve + INSTANCE_STATE = auto() # needs instance-scoped state + SESSION_STATE = auto() # needs session-scoped state + + class ToolDefinition: """A registered tool: schema for the model plus a callable handler.""" @@ -35,6 +51,7 @@ class ToolDefinition: description: str parameters: JSONSchema handler: ToolHandler + tags: ToolTag def __init__( self, @@ -42,11 +59,17 @@ class ToolDefinition: description: str, parameters: JSONSchema, handler: ToolHandler, + *, + tags: ToolTag = ToolTag.LOCAL, ) -> None: + if not (tags & (ToolTag.LOCAL | ToolTag.PUBLIC)): + msg = f"tool {name!r} tags must include LOCAL and/or PUBLIC" + raise ValueError(msg) self.name = name self.description = description self.parameters = parameters self.handler = handler + self.tags = tags def to_tool_item(self) -> ToolFunctionItem: return ToolFunctionItem( @@ -105,6 +128,7 @@ def _build_definition( *, name: str | None, description: str | None, + tags: ToolTag, ) -> ToolDefinition: tool_name = name or func.__name__ tool_description = description if description is not None else (inspect.getdoc(func) or "") @@ -113,9 +137,18 @@ def _build_definition( description=tool_description, parameters=schema_from_callable(func), handler=func, + tags=tags, ) +def _register_definition(definition: ToolDefinition) -> None: + """Push *definition* into the process tool catalog (lazy import).""" + # Imported lazily so agent.tools does not hard-depend on plyngent.tools. + from plyngent.tools.catalog import register_tool + + register_tool(definition) + + @overload def tool[**PS, R](func: Callable[PS, R], /) -> ToolDefinition: ... @@ -127,6 +160,8 @@ def tool[**PS, R]( *, name: str | None = None, description: str | None = None, + tags: ToolTag = ToolTag.LOCAL, + register: bool = True, ) -> Callable[[Callable[PS, R]], ToolDefinition]: ... @@ -136,14 +171,29 @@ def tool[**PS, R]( *, name: str | None = None, description: str | None = None, + tags: ToolTag = ToolTag.LOCAL, + register: bool = True, ) -> ToolDefinition | Callable[[Callable[PS, R]], ToolDefinition]: - """Register a function as an agent tool (decorator). + """Define (and by default catalog-register) a function as an agent tool. Schema is inferred from type hints; description defaults to the docstring. + Default ``tags`` is :attr:`ToolTag.LOCAL`. When ``register`` is true, the + definition is added to the process :class:`~plyngent.tools.catalog.ToolCatalog` + with the current registration source (builtin unless a plugin context is set). + Catalog registration does **not** alone make the tool model-visible — hosts + still **select** tools into a :class:`ToolRegistry`. """ def decorator(fn: Callable[PS, R]) -> ToolDefinition: - return _build_definition(fn, name=name, description=description) + definition = _build_definition( + fn, + name=name, + description=description, + tags=tags, + ) + if register: + _register_definition(definition) + return definition if func is not None: return decorator(func) @@ -157,6 +207,10 @@ class ToolRegistry: _danger: DangerClassifier | None _on_confirm: ToolConfirmHook | None _confirm_lock: asyncio.Lock + _yolo: bool + _auto_bind_state: bool + _instance: object | None + _session: object | None def __init__( self, @@ -164,11 +218,19 @@ class ToolRegistry: *, danger: DangerClassifier | None = None, on_confirm: ToolConfirmHook | None = None, + yolo: bool = False, + auto_bind_state: bool = False, + instance_state: object | None = None, + session_state: object | None = None, ) -> None: self._tools = {} self._danger = danger self._on_confirm = on_confirm self._confirm_lock = asyncio.Lock() + self._yolo = yolo + self._auto_bind_state = auto_bind_state + self._instance = instance_state + self._session = session_state if tools is None: return if isinstance(tools, list): @@ -193,11 +255,49 @@ class ToolRegistry: def __len__(self) -> int: return len(self._tools) + def set_yolo(self, *, enabled: bool) -> None: + self._yolo = enabled + + def set_session_state(self, session_state: object | None) -> None: + self._session = session_state + + def set_instance_state(self, instance_state: object | None) -> None: + self._instance = instance_state + + @property + def yolo(self) -> bool: + """Whether YOLO mode may auto-approve YOLO-tagged soft confirms.""" + return self._yolo + @property def soft_confirm(self) -> bool: - """True when dangerous tools are gated by ``on_confirm``.""" + """True when a soft-confirm path is configured (danger + on_confirm). + + YOLO mode does not clear this: non-YOLO-tagged tools still prompt. + """ return self._danger is not None and self._on_confirm is not None + def _check_state_tags(self, definition: ToolDefinition) -> str | None: + """Return an error string if required state context is missing. + + Only enforced when the host opted into ``auto_bind_state`` (CLI). Hand + registries in tests may still rely on process globals during migration. + """ + if not self._auto_bind_state: + return None + tags = definition.tags + if tags & ToolTag.INSTANCE_STATE: + from plyngent.tools.context import get_instance + + if get_instance() is None and self._instance is None: + return f"error: tool {definition.name!r} requires instance state (INSTANCE_STATE) but none is bound" + if tags & ToolTag.SESSION_STATE: + from plyngent.tools.context import get_session + + if get_session() is None and self._session is None: + return f"error: tool {definition.name!r} requires session state (SESSION_STATE) but none is bound" + return None + async def _invoke(self, definition: ToolDefinition, args: dict[str, object]) -> str: try: result = definition.handler(**args) @@ -213,30 +313,82 @@ class ToolRegistry: return result return msgspec.json.encode(result).decode() - async def _maybe_confirm(self, name: str, args: dict[str, object]) -> str | None: - """If confirm is required and denied, return an error string for the model. + def _session_for_grants(self) -> Any | None: + from plyngent.tools.context import get_session - The confirm hook may return ``True`` (allow), ``False`` (deny), or a - non-empty string (deny with user comment for the model). + session = get_session() + if session is None and self._session is not None: + return cast("Any", self._session) + return session + + def _grant_allows(self, name: str, *, tags: ToolTag) -> bool: + if not (tags & ToolTag.TRUSTABLE): + return False + from plyngent.tools.grants import has_grant + + session = self._session_for_grants() + return session is not None and has_grant(session, name) + + 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) + + async def _prompt_soft_confirm( + self, + name: str, + args: dict[str, object], + reason: str, + *, + tags: ToolTag, + ) -> str | None: + if self._on_confirm is None: + return f"error: tool {name!r} denied by policy ({reason}; no confirm hook)" + decision = self._on_confirm(name, args, reason) + if inspect.isawaitable(decision): + decision = await decision + if decision is True: + 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()}" + return f"error: tool {name!r} denied by user confirm ({reason})" + + async def _maybe_confirm( + self, + name: str, + args: dict[str, object], + *, + tags: ToolTag, + ) -> str | None: + """Soft-confirm gate driven by danger reason + tags + grants + YOLO. + + Pipeline (soft gate only; hard denylists stay in tool handlers):: + + no soft reason → run + YOLO mode and (tags & YOLO) → allow + (tags & TRUSTABLE) and grant exists → allow + else on_confirm; on approve + TRUSTABLE → store grant """ - if self._danger is None or self._on_confirm is None: + if self._danger 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 - decision = self._on_confirm(name, args, reason) - if inspect.isawaitable(decision): - decision = await decision - if decision is True: + if self._yolo and (tags & ToolTag.YOLO): return None - 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})" + if self._grant_allows(name, tags=tags): + return None + return await self._prompt_soft_confirm(name, args, reason, tags=tags) async def execute(self, name: str, arguments_json: str) -> str: """Run a tool by name; returns a string result (errors become error text).""" @@ -250,7 +402,22 @@ class ToolRegistry: if not isinstance(raw_args, dict): return "error: tool arguments must be a JSON object" 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) + + async def _run() -> str: + missing = self._check_state_tags(definition) + if missing is not None: + return missing + denied = await self._maybe_confirm(name, args, tags=definition.tags) + if denied is not None: + return denied + return await self._invoke(definition, args) + + if self._auto_bind_state and (self._instance is not None or self._session is not None): + from plyngent.tools.context import bind_tool_context + + with bind_tool_context( + instance=cast("Any", self._instance), + session=cast("Any", self._session), + ): + return await _run() + return await _run() diff --git a/src/plyngent/tools/__init__.py b/src/plyngent/tools/__init__.py index b31ea1b..b4f0e12 100644 --- a/src/plyngent/tools/__init__.py +++ b/src/plyngent/tools/__init__.py @@ -1,7 +1,20 @@ +from .catalog import ToolCatalog as ToolCatalog +from .catalog import ToolSource as ToolSource +from .catalog import catalog_scope as catalog_scope +from .catalog import default_tool_definitions as default_tool_definitions +from .catalog import get_catalog as get_catalog +from .catalog import register_builtin_tools as register_builtin_tools from .chat import CHAT_TOOLS as CHAT_TOOLS from .chat import ask_user as ask_user from .chat import choose_user as choose_user from .chat import form_user as form_user +from .context import InstanceState as InstanceState +from .context import SessionState as SessionState +from .context import bind_instance as bind_instance +from .context import bind_session as bind_session +from .context import bind_tool_context as bind_tool_context +from .context import require_instance as require_instance +from .context import require_session as require_session from .danger import classify_danger as classify_danger from .file import FILE_TOOLS as FILE_TOOLS from .file import copy_path as copy_path @@ -40,6 +53,8 @@ from .vcs import vcs_diff as vcs_diff from .vcs import vcs_kind as vcs_kind from .vcs import vcs_log as vcs_log from .vcs import vcs_status as vcs_status +from .view import MemoryViewStore as MemoryViewStore +from .view import PersistentDataView as PersistentDataView from .workspace import ( DEFAULT_COMMAND_DENYLIST as DEFAULT_COMMAND_DENYLIST, ) @@ -65,4 +80,5 @@ from .workspace import set_policy_confirm_hook as set_policy_confirm_hook from .workspace import set_policy_confirm_timeout as set_policy_confirm_timeout from .workspace import set_workspace_root as set_workspace_root +# Deprecated alias: prefer default_tool_definitions() (catalog select). DEFAULT_TOOLS = [*FILE_TOOLS, *PROCESS_TOOLS, *VCS_TOOLS, *CHAT_TOOLS, *TODO_TOOLS] diff --git a/src/plyngent/tools/catalog.py b/src/plyngent/tools/catalog.py new file mode 100644 index 0000000..19ed9e2 --- /dev/null +++ b/src/plyngent/tools/catalog.py @@ -0,0 +1,264 @@ +"""Process-wide tool catalog: define+register vs select for model visibility.""" + +from __future__ import annotations + +from contextlib import contextmanager +from contextvars import ContextVar, Token +from dataclasses import dataclass +from typing import TYPE_CHECKING, Literal, override + +from plyngent.agent.tools import ToolDefinition, ToolTag + +if TYPE_CHECKING: + from collections.abc import Container, Generator + +type ToolSourceKind = Literal["builtin", "plugin"] +type ToolSurface = Literal["local", "public"] + + +@dataclass(frozen=True, slots=True) +class ToolSource: + """Where a catalog entry came from (not a :class:`~plyngent.agent.tools.ToolTag`).""" + + kind: ToolSourceKind + plugin_id: str | None = None + package: str | None = None + module: str | None = None + + @override + def __str__(self) -> str: + if self.kind == "builtin": + return "builtin" + plugin = self.plugin_id or "?" + return f"plugin:{plugin}" + + +@dataclass(frozen=True, slots=True) +class RegisteredTool: + definition: ToolDefinition + source: ToolSource + + +_BUILTIN_SOURCE = ToolSource(kind="builtin") +_registration_source: ContextVar[ToolSource] = ContextVar( + "plyngent_tool_registration_source", + default=_BUILTIN_SOURCE, +) + + +def get_registration_source() -> ToolSource: + return _registration_source.get() + + +@contextmanager +def registration_source(source: ToolSource) -> Generator[None]: + """Set :func:`get_registration_source` for the duration of a load block.""" + token = _registration_source.set(source) + try: + yield + finally: + _registration_source.reset(token) + + +class ToolCatalog: + """Name → registered tool map for one process (or test-scoped snapshot).""" + + _by_name: dict[str, RegisteredTool] + + def __init__(self) -> None: + self._by_name = {} + + def register(self, definition: ToolDefinition, *, source: ToolSource | None = None) -> None: + """Add *definition*; never shadow an existing name (builtin or plugin).""" + resolved = source if source is not None else get_registration_source() + existing = self._by_name.get(definition.name) + if existing is not None: + msg = ( + f"tool name collision: {definition.name!r} already registered " + f"from {existing.source} (refusing {resolved})" + ) + raise ValueError(msg) + self._by_name[definition.name] = RegisteredTool(definition=definition, source=resolved) + + def get(self, name: str) -> RegisteredTool | None: + return self._by_name.get(name) + + def names(self) -> list[str]: + return sorted(self._by_name) + + @staticmethod + def _matches( # noqa: PLR0911 — filter predicates are clearer as early exits + entry: RegisteredTool, + *, + surface: ToolSurface, + sources: Container[ToolSourceKind] | None, + plugin_ids: Container[str] | None, + require_tags: ToolTag | None, + exclude_tags: ToolTag | None, + include_names: set[str] | None, + exclude_names: set[str] | None, + ) -> bool: + name = entry.definition.name + if include_names is not None and name not in include_names: + return False + if exclude_names is not None and name in exclude_names: + return False + if sources is not None and entry.source.kind not in sources: + return False + if plugin_ids is not None: + plugin_id = entry.source.plugin_id + if entry.source.kind != "plugin" or plugin_id is None or plugin_id not in plugin_ids: + return False + tags = entry.definition.tags + if surface == "public" and not (tags & ToolTag.PUBLIC): + return False + if surface == "local" and not (tags & (ToolTag.LOCAL | ToolTag.PUBLIC)): + return False + if require_tags is not None and (tags & require_tags) != require_tags: + return False + return not (exclude_tags is not None and tags & exclude_tags) + + def select( + self, + *, + surface: ToolSurface = "local", + sources: Container[ToolSourceKind] | None = None, + plugin_ids: Container[str] | None = None, + require_tags: ToolTag | None = None, + exclude_tags: ToolTag | None = None, + include_names: set[str] | None = None, + exclude_names: set[str] | None = None, + ) -> list[ToolDefinition]: + """Return definitions allowed for this host surface / filters. + + *surface* ``local`` keeps tools with LOCAL and/or PUBLIC. + *surface* ``public`` keeps only tools with PUBLIC. + """ + return [ + self._by_name[name].definition + for name in sorted(self._by_name) + if self._matches( + self._by_name[name], + surface=surface, + sources=sources, + plugin_ids=plugin_ids, + require_tags=require_tags, + exclude_tags=exclude_tags, + include_names=include_names, + exclude_names=exclude_names, + ) + ] + + def clear(self) -> None: + self._by_name.clear() + + def snapshot(self) -> dict[str, RegisteredTool]: + return dict(self._by_name) + + def restore(self, snapshot: dict[str, RegisteredTool]) -> None: + self._by_name = dict(snapshot) + + +_catalog = ToolCatalog() +_catalog_override: ContextVar[ToolCatalog | None] = ContextVar( + "plyngent_tool_catalog_override", + default=None, +) + + +def get_catalog() -> ToolCatalog: + override = _catalog_override.get() + if override is not None: + return override + return _catalog + + +def register_tool(definition: ToolDefinition, *, source: ToolSource | None = None) -> None: + get_catalog().register(definition, source=source) + + +@contextmanager +def catalog_scope(*, empty: bool = True) -> Generator[ToolCatalog]: + """Snapshot the process catalog; optionally start empty for unit tests. + + Restores the previous catalog contents on exit. Nested scopes stack via + contextvars when *empty* installs a fresh catalog override. + """ + if empty: + scoped = ToolCatalog() + token: Token[ToolCatalog | None] = _catalog_override.set(scoped) + try: + yield scoped + finally: + _catalog_override.reset(token) + return + + catalog = get_catalog() + snap = catalog.snapshot() + try: + yield catalog + finally: + catalog.restore(snap) + + +_builtins_loaded = False + + +def _ensure_builtin_definitions_registered(catalog: ToolCatalog) -> None: + """Register builtin ToolDefinitions if the active catalog is missing them. + + Importing tool modules only registers into the catalog that was active at + import time (usually the process catalog). Test ``catalog_scope(empty=True)`` + overrides must still see builtins for ``default_tool_definitions``. + """ + from plyngent.tools import DEFAULT_TOOLS + + for definition in DEFAULT_TOOLS: + if catalog.get(definition.name) is None: + catalog.register(definition, source=_BUILTIN_SOURCE) + + +def register_builtin_tools(*, force: bool = False) -> ToolCatalog: + """Import builtin tool modules so ``@tool`` registrations run. + + Idempotent unless *force* is true (re-import path still no-ops if names + already exist — callers should use :func:`catalog_scope` for isolation). + Always ensures the **active** catalog (including overrides) has builtins. + """ + global _builtins_loaded # noqa: PLW0603 — process load flag + if not _builtins_loaded or force: + # Import side effects: each module's @tool(...) registers into the catalog. + # Group packages pull their leaf modules (FILE_TOOLS etc. still exported). + import plyngent.tools.chat as _chat_tools + import plyngent.tools.file as _file_tools + import plyngent.tools.process as _process_tools + import plyngent.tools.temp_workspace as _temp_workspace_tools + import plyngent.tools.todo as _todo_tools + import plyngent.tools.vcs as _vcs_tools + + _ = ( + _chat_tools, + _file_tools, + _process_tools, + _temp_workspace_tools, + _todo_tools, + _vcs_tools, + ) + _builtins_loaded = True + + catalog = get_catalog() + _ensure_builtin_definitions_registered(catalog) + return catalog + + +def default_tool_definitions( + *, + surface: ToolSurface = "local", + sources: Container[ToolSourceKind] | None = None, +) -> list[ToolDefinition]: + """Load builtins (if needed) and select for the given surface.""" + catalog = register_builtin_tools() + kind_filter: Container[ToolSourceKind] | None = sources + if kind_filter is None: + kind_filter = ("builtin",) + return catalog.select(surface=surface, sources=kind_filter) diff --git a/src/plyngent/tools/context.py b/src/plyngent/tools/context.py new file mode 100644 index 0000000..95fbe65 --- /dev/null +++ b/src/plyngent/tools/context.py @@ -0,0 +1,114 @@ +"""Instance / session tool context (contextvars) for state affinity tags.""" + +from __future__ import annotations + +from contextlib import contextmanager +from contextvars import ContextVar, Token +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from collections.abc import Generator + from pathlib import Path + + from plyngent.agent.todo_stack import TodoStack + from plyngent.tools.view import PersistentDataView + + +@dataclass +class InstanceState: + """Process / agent-host scoped state for INSTANCE_STATE tools.""" + + # Optional fixed facets (workspace path still also lives in workspace module + # during migration; hosts should set both until globals retire). + workspace_root: Path | None = None + data: PersistentDataView[Any] | None = None + # Ephemeral process maps (PTY etc.) may hang here later. + extras: dict[str, Any] = field(default_factory=dict) + + async def shutdown(self) -> None: + """Best-effort cleanup hooks (PTY close, temp workspaces).""" + from plyngent.tools.process.pty_session import PtyManager + from plyngent.tools.temp_workspace import cleanup_temporary_workspaces + + PtyManager.close_all() + _ = cleanup_temporary_workspaces() + + +@dataclass +class SessionState: + """One chat session for SESSION_STATE tools.""" + + session_id: int | str | None = None + data: PersistentDataView[Any] | None = None + # 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). + grants: dict[str, bool] = field(default_factory=dict) + extras: dict[str, Any] = field(default_factory=dict) + + def has_grant(self, tool_name: str) -> bool: + return bool(self.grants.get(tool_name)) + + def add_grant(self, tool_name: str) -> None: + self.grants[tool_name] = True + + def clear_grants(self) -> None: + self.grants.clear() + + +_instance: ContextVar[InstanceState | None] = ContextVar("plyngent_instance_state", default=None) +_session: ContextVar[SessionState | None] = ContextVar("plyngent_session_state", default=None) + + +def get_instance() -> InstanceState | None: + return _instance.get() + + +def get_session() -> SessionState | None: + return _session.get() + + +def require_instance() -> InstanceState: + state = _instance.get() + if state is None: + msg = "instance state is not bound; host must set InstanceState around tool execution" + raise RuntimeError(msg) + return state + + +def require_session() -> SessionState: + state = _session.get() + if state is None: + msg = "session state is not bound; host must set SessionState around tool execution" + raise RuntimeError(msg) + return state + + +@contextmanager +def bind_instance(state: InstanceState | None) -> Generator[InstanceState | None]: + token: Token[InstanceState | None] = _instance.set(state) + try: + yield state + finally: + _instance.reset(token) + + +@contextmanager +def bind_session(state: SessionState | None) -> Generator[SessionState | None]: + token: Token[SessionState | None] = _session.set(state) + try: + yield state + finally: + _session.reset(token) + + +@contextmanager +def bind_tool_context( + *, + instance: InstanceState | None = None, + session: SessionState | None = None, +) -> Generator[tuple[InstanceState | None, SessionState | None]]: + """Bind instance and session contextvars for a tool batch / test.""" + with bind_instance(instance), bind_session(session): + yield instance, session diff --git a/src/plyngent/tools/grants.py b/src/plyngent/tools/grants.py new file mode 100644 index 0000000..e589d54 --- /dev/null +++ b/src/plyngent/tools/grants.py @@ -0,0 +1,25 @@ +"""Soft-confirm trust grants (session-scoped).""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from plyngent.tools.context import SessionState + + +def grant_key(tool_name: str) -> str: + """Phase 1 grant key is the tool name (session isolation via SessionState).""" + return tool_name + + +def has_grant(session: SessionState, tool_name: str) -> bool: + return session.has_grant(grant_key(tool_name)) + + +def add_grant(session: SessionState, tool_name: str) -> None: + session.add_grant(grant_key(tool_name)) + + +def clear_grants(session: SessionState) -> None: + session.clear_grants() diff --git a/src/plyngent/tools/view.py b/src/plyngent/tools/view.py new file mode 100644 index 0000000..a8e814d --- /dev/null +++ b/src/plyngent/tools/view.py @@ -0,0 +1,270 @@ +"""Transactional path-scoped views over durable / publishable trees.""" + +from __future__ import annotations + +from contextlib import AbstractAsyncContextManager +from dataclasses import dataclass +from typing import Any, Protocol, Self, TypeVar, cast, overload, override + +U = TypeVar("U") + + +class ViewStore(Protocol): + """Backend for a root PersistentDataView tree.""" + + async def load(self) -> object: ... + + async def store(self, root: object) -> None: ... + + +class MemoryViewStore: + """In-memory root store (tests / process-only state).""" + + def __init__(self, initial: object | None = None) -> None: + self._root: object = {} if initial is None else initial + + async def load(self) -> object: + return self._root + + async def store(self, root: object) -> None: + self._root = root + + +@dataclass +class _TxnState: + """Private buffer for an open root transaction.""" + + root: object + dirty: bool = False + depth: int = 0 + + +class PersistentDataView[T](AbstractAsyncContextManager["PersistentDataView[T]"]): + """Path-scoped view over a durable/publishable tree. + + ``T`` is the type of data referenced at this path (for load/store/typed()). + The view **is** the transaction context manager: ``async with view``. + Mutations outside an open txn raise. Nested ``async with`` joins the same + root txn (savepoint-lite: depth counter; full rollback on any exception). + """ + + _store: ViewStore + _path: tuple[str | int, ...] + _bound_type: type[T] | None + _txn: _TxnState | None + _domain: dict[tuple[str | int, ...], object] + _child_cache: dict[tuple[str | int, ...], PersistentDataView[Any]] + + def __init__( + self, + store: ViewStore, + *, + path: tuple[str | int, ...] = (), + bound_type: type[T] | None = None, + _txn: _TxnState | None = None, + _domain: dict[tuple[str | int, ...], object] | None = None, + _child_cache: dict[tuple[str | int, ...], PersistentDataView[Any]] | None = None, + ) -> None: + self._store = store + self._path = path + self._bound_type = bound_type + self._txn = _txn + self._domain = _domain if _domain is not None else {} + self._child_cache = _child_cache if _child_cache is not None else {} + + @property + def path(self) -> tuple[str | int, ...]: + return self._path + + def _require_txn(self) -> _TxnState: + if self._txn is None or self._txn.depth < 1: + msg = "PersistentDataView mutation/read of live buffer requires an open transaction (async with view)" + raise RuntimeError(msg) + return self._txn + + def _navigate(self, root: object) -> object: + current: object = root + for key in self._path: + if isinstance(current, dict): + current = cast("dict[object, object]", current).get(key) + elif isinstance(current, list) and isinstance(key, int): + lst = cast("list[object]", current) + current = lst[key] if 0 <= key < len(lst) else None + else: + return None + return current + + def _ensure_parent(self, root: object) -> object: + """Ensure dict parents along path exist; return parent container for last key.""" + if not self._path: + return root + current: object = root + for key in self._path[:-1]: + if not isinstance(current, dict): + msg = f"cannot navigate non-dict parent at {key!r}" + raise TypeError(msg) + mapping = cast("dict[object, object]", current) + child: object | None = mapping.get(key) + if child is None: + child = cast("object", {}) + mapping[key] = child + current = child + return current + + def __getitem__(self, key: str | int) -> PersistentDataView[Any]: + child_path = (*self._path, key) + cached = self._child_cache.get(child_path) + if cached is not None: + return cached + child: PersistentDataView[Any] = PersistentDataView( + self._store, + path=child_path, + bound_type=None, + _txn=self._txn, + _domain=self._domain, + _child_cache=self._child_cache, + ) + self._child_cache[child_path] = child + return child + + def load(self) -> T: + """Materialize the value at this path (from txn buffer or by reading store root snapshot).""" + if self._txn is not None and self._txn.depth >= 1: + if self._path in self._domain: + return cast("T", self._domain[self._path]) + value = self._navigate(self._txn.root) + return cast("T", value) + # Outside txn: only allowed as a frozen read via store is async — sync load + # outside txn is not supported for remote stores; require txn. + msg = "load() requires an open transaction" + raise RuntimeError(msg) + + def store(self, value: T) -> None: + txn = self._require_txn() + self._domain[self._path] = value + if not self._path: + txn.root = value + else: + parent = self._ensure_parent(txn.root) + last = self._path[-1] + if isinstance(parent, dict): + cast("dict[object, object]", parent)[last] = value + elif isinstance(parent, list) and isinstance(last, int): + lst = cast("list[object]", parent) + while len(lst) <= last: + lst.append(None) + lst[last] = value + else: + msg = f"cannot store at path {self._path!r}" + raise TypeError(msg) + txn.dirty = True + + def _materialize_domain(self, typ: type[Any], current: object) -> object: + ctor = cast("Any", typ) + if current is None: + try: + created: object = ctor() + except TypeError: + created = None + self._domain[self._path] = created + if created is not None: + self.store(cast("T", created)) + return created + if not isinstance(current, typ): + try: + converted: object = ctor(current) + except TypeError: + converted = current + self._domain[self._path] = converted + self.store(cast("T", converted)) + return converted + return current + + @overload + def typed(self, typ: None = None) -> T: ... + + @overload + def typed(self, typ: type[U]) -> U: ... + + def typed(self, typ: type[U] | None = None) -> T | U: + """Return the live bound data at this path (not a view). + + - ``typed()`` → ``T`` when bound + - ``typed(U)`` → rebind / convert as ``U`` + """ + txn = self._require_txn() + if self._path in self._domain: + current = self._domain[self._path] + else: + current = self._navigate(txn.root) + if current is not None: + self._domain[self._path] = current + + if typ is not None: + return cast("U", self._materialize_domain(typ, current)) + + if self._bound_type is not None and current is None: + return cast("T", self._materialize_domain(self._bound_type, current)) + return cast("T", current) + + async def save(self) -> None: + """Flush the root buffer to the store without closing the txn.""" + txn = self._require_txn() + if txn.dirty: + await self._store.store(txn.root) + txn.dirty = False + + @override + async def __aenter__(self) -> Self: + if self._txn is None: + # Root view owns the txn state. + root = await self._store.load() + self._txn = _TxnState(root=root, dirty=False, depth=1) + # Propagate txn to cached children created before enter (rare). + for child in self._child_cache.values(): + child._txn = self._txn + return self + self._txn.depth += 1 + return self + + @override + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: object, + ) -> bool: + txn = self._txn + if txn is None: + return False + txn.depth -= 1 + if txn.depth > 0: + # Nested exit: on exception, mark for full rollback at root. + if exc_type is not None: + txn.dirty = False + # Reload root discarded on root exit via exception path. + txn.root = await self._store.load() + self._domain.clear() + return False + # Root exit + try: + if exc_type is None and txn.dirty: + await self._store.store(txn.root) + finally: + self._txn = None + self._domain.clear() + # Drop child txn links + for child in self._child_cache.values(): + child._txn = None + self._child_cache.clear() + return False + + +def session_data_view( + initial: dict[str, object] | None = None, + *, + store: ViewStore | None = None, +) -> PersistentDataView[dict[str, object]]: + """Convenience root view for session documents (todo, grants, …).""" + backend: ViewStore = store if store is not None else MemoryViewStore(initial if initial is not None else {}) + return PersistentDataView(backend, path=(), bound_type=dict)