mirror of
https://github.com/NCBM/plyngent.git
synced 2026-07-25 08:04:57 +08:00
core/tools: catalog, tags, state context, and PersistentDataView
This commit is contained in:
@@ -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
|
||||
|
||||
+187
-20
@@ -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()
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user