feat(core): add ToolSessionManager for stateful tool execution

Introduce session-level state management for tools that need to maintain
state across conversation turns within the same session (UMO).

- Add ToolSessionManager class for central per-(UMO, tool_name) state
- Add ToolSessionState with set_persistent(key) support
- Add is_stateful flag to FunctionTool base class
- Wire session_manager through run_context in all agent runners
- Update ExecuteShellTool to use framework session manager
- Update CLAUDE.md with stateful tool documentation

BREAKING: Tools that were manually managing session state via _sessions
dict should migrate to use the framework's ToolSessionManager for proper
lifecycle management and persistence support.

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

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]

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,  # 只传递有效的参数
                 )
 

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})"

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]: ...

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
+            }

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:

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,

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:

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,

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,