diff --git a/CLAUDE.md b/CLAUDE.md index 42d36bf24..0de35236e 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -107,6 +107,38 @@ Stars are plugins in `astrbot/builtin_stars/`: - Use decorators for command handlers: `@star.on_command`, `@star.on_message`, etc. - Access via `context` object +### Stateful Tool Execution (Session Lifecycle) + +Tools can maintain state across conversation turns within a session via `ToolSessionManager`. + +**Key classes:** +- `ToolSessionManager` (`astrbot/core/agent/tool_session_manager.py`) — central manager, keyed by `(umo, tool_name)` +- `ToolSessionState` — dict-like per-tool session state with `set_persistent(key)` support +- `FunctionTool.is_stateful` — opt-in flag for stateful tools +- `FunctionTool.get_session_state(umo)` — get/create session state dict + +**Usage in a tool:** +```python +@dataclass +class MyTool(FunctionTool): + is_stateful = True # declare stateful + + async def call(self, context, **kwargs): + umo = context.context.event.unified_msg_origin + state = self.get_session_state(umo) + state["counter"] = state.get("counter", 0) + 1 + # Mark to survive session clear: + state.set_persistent("persistent_data") +``` + +**Architecture flow:** +``` +AgentContextWrapper(session_manager=ToolSessionManager()) + → ToolLoopAgentRunner.run_context.session_manager + → executor.execute(..., session_manager=run_context.session_manager) + → tool.call(context) # context.session_manager available +``` + ## Testing 1. Tests go in `tests/` directory diff --git a/astrbot/core/agent/run_context.py b/astrbot/core/agent/run_context.py index 3c500b2d6..dfdf256c1 100644 --- a/astrbot/core/agent/run_context.py +++ b/astrbot/core/agent/run_context.py @@ -17,6 +17,12 @@ class ContextWrapper(Generic[TContext]): messages: list[Message] = Field(default_factory=list) """This field stores the llm message context for the agent run, agent runners will maintain this field automatically.""" tool_call_timeout: int = 120 # Default tool call timeout in seconds + session_manager: Any = None + """ + Optional session manager (ToolSessionManager) for stateful tool execution. + When provided, stateful tools can maintain state across + conversation turns within the same session (UMO). + """ NoContext = ContextWrapper[None] diff --git a/astrbot/core/agent/runners/tool_loop_agent_runner.py b/astrbot/core/agent/runners/tool_loop_agent_runner.py index ff937ed37..e57507980 100644 --- a/astrbot/core/agent/runners/tool_loop_agent_runner.py +++ b/astrbot/core/agent/runners/tool_loop_agent_runner.py @@ -767,6 +767,7 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): executor = self.tool_executor.execute( tool=func_tool, run_context=self.run_context, + session_manager=self.run_context.session_manager, **valid_params, # 只传递有效的参数 ) diff --git a/astrbot/core/agent/tool.py b/astrbot/core/agent/tool.py index 11f0aa741..7aa3e143a 100644 --- a/astrbot/core/agent/tool.py +++ b/astrbot/core/agent/tool.py @@ -1,5 +1,6 @@ import copy from collections.abc import AsyncGenerator, Awaitable, Callable +from dataclasses import field from typing import Any, Generic import jsonschema @@ -68,6 +69,27 @@ class FunctionTool(ToolSchema, Generic[TContext]): Origin of this tool: 'plugin' (from star plugins), 'internal' (AstrBot built-in), or 'mcp' (from MCP servers). Used by WebUI for display grouping. """ + is_stateful: bool = False + """ + Declare this tool as stateful. Stateful tools maintain state + across conversation turns within the same session (UMO). + When True, the tool can use get_session_state(umo) to access + per-session state that persists across tool calls. + """ + _session_state: dict[str, dict[str, Any]] = field(default_factory=dict, repr=False) + """ + Internal: per-UMO session state storage for stateful tools. + Managed by ToolSessionManager; use get_session_state(umo) instead. + """ + + def get_session_state(self, umo: str) -> dict[str, Any]: + """Get or create session state for the given UMO. + + Only valid when is_stateful=True. Otherwise returns empty dict. + """ + if umo not in self._session_state: + self._session_state[umo] = {} + return self._session_state[umo] def __repr__(self) -> str: return f"FuncTool(name={self.name}, parameters={self.parameters}, description={self.description})" diff --git a/astrbot/core/agent/tool_executor.py b/astrbot/core/agent/tool_executor.py index 14fe4beee..0ba2c9522 100644 --- a/astrbot/core/agent/tool_executor.py +++ b/astrbot/core/agent/tool_executor.py @@ -15,5 +15,6 @@ class BaseFunctionToolExecutor(abc.ABC, Generic[TContext]): cls, tool: FunctionTool, run_context: ContextWrapper[TContext], + session_manager: Any = None, **tool_args, ) -> AsyncGenerator[Any | mcp.types.CallToolResult, None]: ... diff --git a/astrbot/core/agent/tool_session_manager.py b/astrbot/core/agent/tool_session_manager.py new file mode 100644 index 000000000..e185bc734 --- /dev/null +++ b/astrbot/core/agent/tool_session_manager.py @@ -0,0 +1,118 @@ +""" +ToolSessionManager - Session-level state management for stateful tools. + +Provides per-(UMO, tool_name) session state that persists across conversation +turns within the same session, with optional persistence via SharedPreferences. +""" + +from collections.abc import MutableMapping +from dataclasses import dataclass, field +from typing import Any + +from astrbot.core.utils.shared_preferences import SharedPreferences + + +@dataclass +class ToolSessionState(MutableMapping[str, Any]): + """ + Represents the session state for a single tool within a session. + Acts like a dict but supports persistence markers. + + Use `set_persistent(key)` to mark keys that survive session clear. + """ + + umo: str + tool_name: str + _data: dict[str, Any] = field(default_factory=dict) + _persistent_keys: set[str] = field(default_factory=set) + + def __getitem__(self, key: str) -> Any: + return self._data[key] + + def __setitem__(self, key: str, value: Any) -> None: + self._data[key] = value + + def __delitem__(self, key: str) -> None: + del self._data[key] + + def __iter__(self): + return iter(self._data) + + def __len__(self) -> int: + return len(self._data) + + def set_persistent(self, key: str) -> None: + """Mark a key as persistent (survives session clear).""" + self._persistent_keys.add(key) + + def is_persistent(self, key: str) -> bool: + """Check if a key is marked as persistent.""" + return key in self._persistent_keys + + +class ToolSessionManager: + """ + Central manager for all tool session states. + + Maintains in-memory state per (umo, tool_name) combination. + Optional SharedPreferences integration for persistence across sessions. + + Example: + mgr = ToolSessionManager() + state = mgr.get_state(umo, "shell") + state["cwd"] = "/tmp" + state.set_persistent("env") # env survives session clear + """ + + def __init__(self, sp: SharedPreferences | None = None) -> None: + self._states: dict[tuple[str, str], ToolSessionState] = {} + self._sp = sp + + def get_state(self, umo: str, tool_name: str) -> ToolSessionState: + """Get or create session state for a tool in a session.""" + key = (umo, tool_name) + if key not in self._states: + self._states[key] = ToolSessionState(umo=umo, tool_name=tool_name) + return self._states[key] + + async def persist_state(self, umo: str, tool_name: str) -> None: + """Persist marked keys to SharedPreferences.""" + if not self._sp: + return + state = self.get_state(umo, tool_name) + for key, value in state._data.items(): + if key in state._persistent_keys: + storage_key = f"tool_state:{tool_name}:{key}" + await self._sp.session_put(umo, storage_key, value) + + async def load_persistent_state(self, umo: str, tool_name: str) -> None: + """Load persistent state from SharedPreferences into the session state.""" + if not self._sp: + return + state = self.get_state(umo, tool_name) + storage_prefix = f"tool_state:{tool_name}:" + # session_get(umo, None) returns list[Preference] for all prefs in this UMO + prefs: list = await self._sp.session_get(umo, None) + for pref in prefs: + key = getattr(pref, "key", None) or "" + if key.startswith(storage_prefix): + actual_key = key[len(storage_prefix) :] + val = getattr(pref, "value", None) + state._data[actual_key] = ( + val.get("val") if isinstance(val, dict) else val + ) + state.set_persistent(actual_key) + + def clear_session(self, umo: str) -> None: + """ + Clear non-persistent state for all tools in a session. + + Persistent keys (marked via `set_persistent`) are preserved. + """ + keys_to_clear = [k for k in self._states if k[0] == umo] + for key in keys_to_clear: + state = self._states[key] + # Keep only persistent keys + state._data = { + k: v for k, v in state._data.items() if k in state._persistent_keys + } diff --git a/astrbot/core/astr_agent_tool_exec.py b/astrbot/core/astr_agent_tool_exec.py index e43b9447b..2468d5c3c 100644 --- a/astrbot/core/astr_agent_tool_exec.py +++ b/astrbot/core/astr_agent_tool_exec.py @@ -114,11 +114,14 @@ class FunctionToolExecutor(BaseFunctionToolExecutor[AstrAgentContext]): return sanitized @classmethod - async def execute(cls, tool, run_context, **tool_args): + async def execute(cls, tool, run_context, session_manager=None, **tool_args): """执行函数调用｡ Args: - event (AstrMessageEvent): 事件对象, 当 origin 为 local 时必须提供｡ + tool: The tool to execute. + run_context: The run context. + session_manager: Optional ToolSessionManager for stateful tool execution. + **tool_args: Tool-specific arguments. **kwargs: 函数调用的参数｡ Returns: diff --git a/astrbot/core/astr_main_agent.py b/astrbot/core/astr_main_agent.py index f671fb74c..f34e3a19a 100644 --- a/astrbot/core/astr_main_agent.py +++ b/astrbot/core/astr_main_agent.py @@ -15,6 +15,7 @@ from astrbot.core.agent.handoff import HandoffTool from astrbot.core.agent.mcp_client import MCPTool from astrbot.core.agent.message import TextPart from astrbot.core.agent.tool import ToolSet +from astrbot.core.agent.tool_session_manager import ToolSessionManager from astrbot.core.astr_agent_context import AgentContextWrapper, AstrAgentContext from astrbot.core.astr_agent_hooks import MAIN_AGENT_HOOKS from astrbot.core.astr_agent_run_util import AgentRunner @@ -1337,6 +1338,7 @@ async def build_main_agent( run_context=AgentContextWrapper( context=astr_agent_ctx, tool_call_timeout=config.tool_call_timeout, + session_manager=ToolSessionManager(), ), tool_executor=FunctionToolExecutor(), agent_hooks=MAIN_AGENT_HOOKS, diff --git a/astrbot/core/computer/tools/shell.py b/astrbot/core/computer/tools/shell.py index f9187f5a6..8e2146e9e 100644 --- a/astrbot/core/computer/tools/shell.py +++ b/astrbot/core/computer/tools/shell.py @@ -22,7 +22,7 @@ import os import shlex import subprocess from dataclasses import dataclass, field -from typing import Any, cast +from typing import Any from astrbot.api import FunctionTool from astrbot.core.agent.run_context import ContextWrapper @@ -67,6 +67,7 @@ class ExecuteShellTool(FunctionTool): ) is_local: bool = False + is_stateful: bool = True # session_id -> {"cwd": str, "env": dict} _sessions: dict[str, dict[str, Any]] = field( default_factory=dict, init=False, repr=False @@ -87,6 +88,18 @@ class ExecuteShellTool(FunctionTool): } return self._sessions[session_id] + def _get_framework_session_state( + self, context: ContextWrapper[AstrAgentContext], session_id: str + ) -> dict[str, Any]: + """ + Get session state via the framework's ToolSessionManager if available. + Falls back to the tool's own _sessions if session_manager is not set. + """ + session_mgr = getattr(context, "session_manager", None) + if session_mgr is None: + return self._get_session_state(session_id) + return session_mgr.get_state(session_id, self.name) + async def call( # type: ignore[override] self, context: ContextWrapper[AstrAgentContext], **kwargs: Any ) -> ToolExecResult: @@ -98,9 +111,8 @@ class ExecuteShellTool(FunctionTool): - background (bool): whether to run in background - env (dict): environment variables to merge for this execution """ - # Cast the generic ContextWrapper to the concrete AstrAgentContext wrapper so - # subsequent permission checks and attribute access use the expected type. - astr_ctx = cast(ContextWrapper[AstrAgentContext], context) + # Use the context directly - already typed as ContextWrapper[AstrAgentContext] + astr_ctx = context # Permission check (use the cast wrapper) if permission_error := check_admin_permission(astr_ctx, "Shell execution"): @@ -113,9 +125,13 @@ class ExecuteShellTool(FunctionTool): # Resolve session id and session state (use the cast wrapper) session_id = astr_ctx.context.event.unified_msg_origin - state = self._get_session_state(session_id) - session_cwd = state["cwd"] - session_env = state["env"].copy() + # Use framework ToolSessionManager if available, otherwise fall back + if astr_ctx.session_manager is not None: + state = self._get_framework_session_state(astr_ctx, session_id) + else: + state = self._get_session_state(session_id) + session_cwd = state.get("cwd", os.getcwd()) + session_env = state.get("env", dict(os.environ)).copy() # Merge provided env into execution env (do not mutate saved session env) if env: diff --git a/astrbot/core/pipeline/process_stage/method/agent_sub_stages/third_party.py b/astrbot/core/pipeline/process_stage/method/agent_sub_stages/third_party.py index 15d57e66c..687d15334 100644 --- a/astrbot/core/pipeline/process_stage/method/agent_sub_stages/third_party.py +++ b/astrbot/core/pipeline/process_stage/method/agent_sub_stages/third_party.py @@ -24,6 +24,7 @@ from astrbot.core.persona_error_reply import ( if TYPE_CHECKING: from astrbot.core.agent.runners.base import BaseAgentRunner from astrbot.core.provider.entities import LLMResponse +from astrbot.core.agent.tool_session_manager import ToolSessionManager from astrbot.core.astr_agent_context import AgentContextWrapper, AstrAgentContext from astrbot.core.pipeline.context import PipelineContext, call_event_hook from astrbot.core.pipeline.stage import Stage @@ -406,6 +407,7 @@ class ThirdPartyAgentSubStage(Stage): run_context=AgentContextWrapper( context=astr_agent_ctx, tool_call_timeout=120, + session_manager=ToolSessionManager(), ), agent_hooks=MAIN_AGENT_HOOKS, provider_config=self.prov_cfg, diff --git a/astrbot/core/star/context.py b/astrbot/core/star/context.py index 757d689b0..3f346c1d2 100644 --- a/astrbot/core/star/context.py +++ b/astrbot/core/star/context.py @@ -205,6 +205,7 @@ class Context: Exception: For other errors during LLM generation """ # Import here to avoid circular imports + from astrbot.core.agent.tool_session_manager import ToolSessionManager from astrbot.core.astr_agent_context import ( AgentContextWrapper, AstrAgentContext, @@ -245,6 +246,7 @@ class Context: run_context=AgentContextWrapper( context=agent_context, tool_call_timeout=tool_call_timeout, + session_manager=ToolSessionManager(), ), tool_executor=tool_executor, agent_hooks=agent_hooks,