From 6c6cec05c4eafc29ee603a84be4a0ca83e22c9b3 Mon Sep 17 00:00:00 2001 From: worldmozara Date: Fri, 24 Jul 2026 15:49:57 +0800 Subject: [PATCH] core/cli: bind todo persist via session state, not process global --- src/plyngent/cli/state.py | 65 +++++++++++++++++++-------- tests/test_tools/test_todo_session.py | 26 +++++++++++ 2 files changed, 73 insertions(+), 18 deletions(-) diff --git a/src/plyngent/cli/state.py b/src/plyngent/cli/state.py index 9169b45..6bc57ca 100644 --- a/src/plyngent/cli/state.py +++ b/src/plyngent/cli/state.py @@ -21,9 +21,9 @@ from plyngent.runtime import create_client from plyngent.tools import ( InstanceState, SessionState, - set_todo_stack, set_workspace_root, ) +from plyngent.tools.view import MemoryViewStore, session_data_view if TYPE_CHECKING: from collections.abc import Sequence @@ -79,7 +79,9 @@ class ReplState: self.client = cast("ChatClient", cast("object", create_client(self.provider))) self.workspace = Path(self.workspace).expanduser().resolve() self.instance_state.workspace_root = self.workspace - self.session_state.todo = self.todo_stack + # Process-global workspace remains for tools/tests that run without instance bind. + _ = set_workspace_root(self.workspace) + self.session_state = self._session_data_for_todo() self.agent = self._make_agent() self.sync_display_flags() self._bind_todo_tools() @@ -130,25 +132,52 @@ class ReplState: click.secho("yolo=off (once expired)", fg="bright_black", err=True) - def _bind_todo_tools(self) -> None: - """Point module-level todo tools at this session stack + persist hook.""" + def _todo_on_change(self) -> None: + """Schedule memory persist for the live todo stack (session-scoped hook).""" import asyncio - def on_change() -> None: - if self.session_id is None: - return - try: - loop = asyncio.get_running_loop() - except RuntimeError: - return - task = loop.create_task(self.memory.update_session_todo_stack(self.session_id, self.todo_stack.to_raw())) - # Keep a strong ref until done so the task is not GC'd mid-flight. - self._todo_persist_tasks.add(task) - task.add_done_callback(self._todo_persist_tasks.discard) + if self.session_id is None: + return + try: + loop = asyncio.get_running_loop() + except RuntimeError: + return + task = loop.create_task(self.memory.update_session_todo_stack(self.session_id, self.todo_stack.to_raw())) + # Keep a strong ref until done so the task is not GC'd mid-flight. + self._todo_persist_tasks.add(task) + task.add_done_callback(self._todo_persist_tasks.discard) - self.session_state.todo = self.todo_stack + def _session_data_for_todo(self, *, grants: dict[str, bool] | None = None) -> SessionState: + """Build a SessionState with live todo + MemoryViewStore seed (durable tree). + + *grants* defaults to empty (new session / startup). Pass the previous map + only when intentionally preserving soft-confirm trust across rebinds. + """ + grant_map = dict(grants or {}) + store = MemoryViewStore({"todo": self.todo_stack.to_raw(), "grants": dict(grant_map)}) + return SessionState( + session_id=self.session_id, + data=session_data_view(store=store), + todo=self.todo_stack, + on_todo_change=self._todo_on_change, + grants=grant_map, + ) + + def _bind_todo_tools(self) -> None: + """Bind session/instance state for tools (prefer context over process globals). + + Seeds ``session.data["todo"]`` from the live stack and attaches + :meth:`_todo_on_change` so memory persist no longer depends on + process-global ``set_todo_stack(on_change=...)``. + """ self.session_state.session_id = self.session_id - set_todo_stack(self.todo_stack, on_change=on_change) + self.session_state.todo = self.todo_stack + self.session_state.on_todo_change = self._todo_on_change + # Refresh durable seed so tools that rehydrate from the view see current data. + store = getattr(self.session_state.data, "_store", None) + if isinstance(store, MemoryViewStore): + store.merge_key("todo", self.todo_stack.to_raw()) + self.instance_state.workspace_root = self.workspace if hasattr(self, "agent") and self.agent.tools is not None: self.agent.tools.set_session_state(self.session_state) self.agent.tools.set_instance_state(self.instance_state) @@ -516,7 +545,7 @@ class ReplState: if session.sid not in self._session_id_cache: self._session_id_cache.insert(0, session.sid) self.todo_stack = TodoStack() - self.session_state = SessionState(session_id=session.sid, todo=self.todo_stack) + self.session_state = self._session_data_for_todo() self.agent = self._make_agent() self._bind_todo_tools() diff --git a/tests/test_tools/test_todo_session.py b/tests/test_tools/test_todo_session.py index b50afac..c3e59f5 100644 --- a/tests/test_tools/test_todo_session.py +++ b/tests/test_tools/test_todo_session.py @@ -88,3 +88,29 @@ async def test_with_bound_context_without_registry_session() -> None: loaded = await store.load() assert isinstance(loaded, dict) assert isinstance(loaded.get("todo"), dict) + + +async def test_session_on_todo_change_preferred_over_process() -> None: + """Session.on_todo_change fires instead of process set_todo_stack on_change.""" + set_todo_stack(None) + hits: list[str] = [] + + def process_hook() -> None: + hits.append("process") + + def session_hook() -> None: + hits.append("session") + + stack = TodoStack() + set_todo_stack(stack, on_change=process_hook) + store = MemoryViewStore({}) + session = SessionState( + session_id="hook", + data=session_data_view(store=store), + todo=stack, + on_todo_change=session_hook, + ) + registry = ToolRegistry(list(TODO_TOOLS), session_state=session) + _ = await registry.execute("todo_push", '{"titles": ["H"]}') + assert hits == ["session"] + set_todo_stack(None)