From d99f589a5909511d3095f4509829c88727abbbc6 Mon Sep 17 00:00:00 2001 From: Xubin Ren <52506698+Re-bin@users.noreply.github.com> Date: Tue, 4 Aug 2026 11:39:39 +0800 Subject: [PATCH] refactor(session): clarify reference boundaries --- nanobot/agent/tools/sessions.py | 57 ++++++++++++------- nanobot/bus/events.py | 2 +- nanobot/channels/websocket/runtime.py | 3 +- .../tests/test_websocket_envelope_media.py | 2 +- nanobot/webui/session_access.py | 29 +++++++--- tests/agent/tools/test_sessions.py | 23 +++++++- tests/webui/test_session_mentions.py | 3 +- .../src/components/thread/ThreadComposer.tsx | 27 ++++++--- webui/src/tests/thread-composer.test.tsx | 5 ++ 9 files changed, 108 insertions(+), 43 deletions(-) diff --git a/nanobot/agent/tools/sessions.py b/nanobot/agent/tools/sessions.py index 1f39b978d..3e678e0de 100644 --- a/nanobot/agent/tools/sessions.py +++ b/nanobot/agent/tools/sessions.py @@ -33,19 +33,19 @@ def session_extra(metadata: Mapping[str, Any] | None) -> dict[str, Any]: def _session_scope() -> SessionAccessScope | None: ctx = current_request_context() - if ctx is None: + if ctx is None or not ctx.session_key: return None - session_key = ctx.session_key + prefix = ctx.metadata.get(INBOUND_META_SESSION_READ_SCOPE) if ( - ctx.channel != "websocket" - or session_key is None - or not session_key.startswith("websocket:") - or ctx.metadata.get(INBOUND_META_SESSION_READ_SCOPE) is not True + not isinstance(prefix, str) + or not prefix.endswith(":") + or not ctx.session_key.startswith(prefix) ): return None workspace = current_workspace_scope() return SessionAccessScope( - current_session_key=session_key, + current_session_key=ctx.session_key, + session_key_prefix=prefix, project_path=workspace.project_path if workspace is not None else ctx.workspace, restrict_to_workspace=workspace.restrict_to_workspace if workspace is not None else False, ) @@ -131,20 +131,30 @@ class SearchSessionsTool(_SessionTool): return ToolResult.error("Error: session search is not available to this client") matches = await asyncio.to_thread(self._access.search, scope, query, _SEARCH_LIMIT) needle = query.casefold() - for match in matches: - match["session_ref"] = _session_ref(match["session_key"]) - match["excerpts"] = [ + result = { + "notice": _UNTRUSTED_NOTICE, + "query": query, + "results": [ { - "message_index": message["message_index"], - "role": message["role"], - "content": _excerpt(message["content"], needle, _SEARCH_EXCERPT_CHARS), + "session_key": match["session_key"], + "session_ref": _session_ref(match["session_key"]), + "title": match["title"], + "updated_at": match["updated_at"], + "excerpts": [ + { + "message_index": message["message_index"], + "role": message["role"], + "content": _excerpt( + message["content"], needle, _SEARCH_EXCERPT_CHARS + ), + } + for message in match["messages"] + ], } - for message in match.pop("messages") - ] - return json.dumps( - {"notice": _UNTRUSTED_NOTICE, "query": query, "results": matches}, - ensure_ascii=False, - ) + for match in matches + ], + } + return json.dumps(result, ensure_ascii=False) @tool_parameters( @@ -205,13 +215,16 @@ class ReadSessionTool(_SessionTool): if match is None: return ToolResult.error(f"Error: session not found: {session_key}") needle = query_text.casefold() - match.update({ + result = { "notice": _UNTRUSTED_NOTICE, + "session_key": match["session_key"], "session_ref": _session_ref(session_key), + "title": match["title"], + "updated_at": match["updated_at"], "query": query_text or None, "messages": [ {**message, "content": _excerpt(message["content"], needle, _READ_MESSAGE_CHARS)} for message in match["messages"] ], - }) - return json.dumps(match, ensure_ascii=False) + } + return json.dumps(result, ensure_ascii=False) diff --git a/nanobot/bus/events.py b/nanobot/bus/events.py index c0fb308a7..93bf3235e 100644 --- a/nanobot/bus/events.py +++ b/nanobot/bus/events.py @@ -15,7 +15,7 @@ OUTBOUND_META_AGENT_UI = "_agent_ui" # Internal-only inbound metadata used by in-process channels to ask the agent # loop to update runtime state without going through a user session. INBOUND_META_RUNTIME_CONTROL = "_runtime_control" -# Trusted WebUI grant for read-only persisted-session tools. +# Trusted namespace grant for read-only persisted-session tools. INBOUND_META_SESSION_READ_SCOPE = "_session_read_scope" RUNTIME_CONTROL_ACK = "_ack" RUNTIME_CONTROL_MCP_RELOAD = "mcp_reload" diff --git a/nanobot/channels/websocket/runtime.py b/nanobot/channels/websocket/runtime.py index 0533489f9..4b20f9d7f 100644 --- a/nanobot/channels/websocket/runtime.py +++ b/nanobot/channels/websocket/runtime.py @@ -814,7 +814,7 @@ class WebSocketChannel(BaseChannel): metadata.update(self._transcripts.client_turn_metadata(envelope.get("turn_id"))) trusted_webui = metadata.get("webui") is True and connection in self._webui_connections if trusted_webui: - metadata[INBOUND_META_SESSION_READ_SCOPE] = True + metadata[INBOUND_META_SESSION_READ_SCOPE] = f"{self.name}:" cli_apps = normalize_cli_app_mentions(envelope.get("cli_apps")) if cli_apps: metadata["cli_apps"] = cli_apps @@ -831,6 +831,7 @@ class WebSocketChannel(BaseChannel): envelope.get("session_mentions"), SessionAccessScope( current_session_key=f"{self.name}:{cid}", + session_key_prefix=f"{self.name}:", project_path=scope.project_path, restrict_to_workspace=scope.restrict_to_workspace, ), diff --git a/nanobot/channels/websocket/tests/test_websocket_envelope_media.py b/nanobot/channels/websocket/tests/test_websocket_envelope_media.py index edce873bc..065ab0209 100644 --- a/nanobot/channels/websocket/tests/test_websocket_envelope_media.py +++ b/nanobot/channels/websocket/tests/test_websocket_envelope_media.py @@ -219,7 +219,7 @@ async def test_webui_message_forwards_verified_session_mentions(tmp_path) -> Non channel._handle_message.assert_awaited_once() metadata = channel._handle_message.call_args.kwargs["metadata"] - assert metadata[INBOUND_META_SESSION_READ_SCOPE] is True + assert metadata[INBOUND_META_SESSION_READ_SCOPE] == "websocket:" assert metadata["session_mentions"] == [{ "name": "pricing", "session_key": "websocket:pricing", diff --git a/nanobot/webui/session_access.py b/nanobot/webui/session_access.py index 3fdc27dda..718bf56ce 100644 --- a/nanobot/webui/session_access.py +++ b/nanobot/webui/session_access.py @@ -7,7 +7,7 @@ from collections.abc import Mapping from dataclasses import dataclass from functools import cache from pathlib import Path -from typing import Any, cast +from typing import Any, TypedDict, cast from nanobot.runtime_context import ( RuntimeContextBlock, @@ -24,24 +24,39 @@ from nanobot.webui.transcript import ( ) _VISIBLE_ROLES = {"user", "assistant"} -_WEBUI_SESSION_PREFIX = "websocket:" -SessionMention = dict[str, str] -SessionMessage = dict[str, Any] -SessionMatch = dict[str, Any] +class SessionMention(TypedDict): + name: str + session_key: str + title: str + + +class SessionMessage(TypedDict): + message_index: int + role: str + timestamp: str | int | None + content: str + + +class SessionMatch(TypedDict): + session_key: str + title: str + updated_at: str | None + messages: list[SessionMessage] @dataclass(frozen=True) class SessionAccessScope: current_session_key: str + session_key_prefix: str project_path: Path | None = None restrict_to_workspace: bool = False def allows(self, session_key: object) -> bool: return ( isinstance(session_key, str) - and session_key.startswith(_WEBUI_SESSION_PREFIX) + and session_key.startswith(self.session_key_prefix) and session_key != self.current_session_key ) @@ -245,7 +260,7 @@ class WebuiSessionAccess: seen_keys: set[str] = set() seen_names: set[str] = set() for raw_mention in normalize_session_mentions_metadata(raw): - mention = raw_mention + mention = cast(SessionMention, raw_mention) key = mention["session_key"] folded_name = mention["name"].lower() payload = self._metadata(key, scope) diff --git a/tests/agent/tools/test_sessions.py b/tests/agent/tools/test_sessions.py index 66db55497..5ebb95677 100644 --- a/tests/agent/tools/test_sessions.py +++ b/tests/agent/tools/test_sessions.py @@ -46,7 +46,7 @@ def _webui_request( channel="websocket", chat_id=session_key.removeprefix("websocket:"), session_key=session_key, - metadata={INBOUND_META_SESSION_READ_SCOPE: True}, + metadata={INBOUND_META_SESSION_READ_SCOPE: "websocket:"}, )) @@ -284,3 +284,24 @@ async def test_session_tools_reject_unscoped_and_out_of_scope_sessions(tmp_path) assert spoofed.is_error assert [row["session_key"] for row in search["results"]] == ["websocket:visible"] assert read.is_error + + +@pytest.mark.asyncio +async def test_session_tools_use_the_scope_granted_by_the_channel(tmp_path): + manager = SessionManager(tmp_path) + _save_session( + manager, + "custom:history", + title="History", + messages=[{"role": "user", "content": "custom needle"}], + ) + + with request_context(RequestContext( + channel="custom", + chat_id="current", + session_key="custom:current", + metadata={INBOUND_META_SESSION_READ_SCOPE: "custom:"}, + )): + result = _decode(await SearchSessionsTool(manager).execute(query="needle")) + + assert [row["session_key"] for row in result["results"]] == ["custom:history"] diff --git a/tests/webui/test_session_mentions.py b/tests/webui/test_session_mentions.py index f6e35d1d9..96c0ac7a6 100644 --- a/tests/webui/test_session_mentions.py +++ b/tests/webui/test_session_mentions.py @@ -51,7 +51,7 @@ def test_normalize_session_mentions_keeps_only_authorized_distinct_targets( {"name": "STRASSE", "session_key": "websocket:upper"}, {"name": "private", "session_key": "telegram:private"}, ], - SessionAccessScope("websocket:current"), + SessionAccessScope("websocket:current", "websocket:"), ) assert mentions == [ @@ -99,6 +99,7 @@ def test_restricted_scope_rejects_sessions_from_other_projects(tmp_path) -> None access = WebuiSessionAccess(manager) scope = SessionAccessScope( "websocket:current", + "websocket:", project_path=project_a, restrict_to_workspace=True, ) diff --git a/webui/src/components/thread/ThreadComposer.tsx b/webui/src/components/thread/ThreadComposer.tsx index ffff9b6ab..1f542a2ef 100644 --- a/webui/src/components/thread/ThreadComposer.tsx +++ b/webui/src/components/thread/ThreadComposer.tsx @@ -1325,14 +1325,23 @@ export function ThreadComposer({ logoUrl: preset.logo_url ?? null, initials: mcpPresetInitials(preset), })); - return [ - ...sessionCandidates.slice(0, 4), - ...cliCandidates.slice(0, 2), - ...mcpCandidates.slice(0, 2), - ...sessionCandidates.slice(4), - ...cliCandidates.slice(2), - ...mcpCandidates.slice(2), - ].slice(0, 8); + const groups = [ + { candidates: cliCandidates, reserved: 2 }, + { candidates: mcpCandidates, reserved: 2 }, + { candidates: sessionCandidates, reserved: 4 }, + ]; + let remaining = 8; + const counts = groups.map(({ candidates, reserved }) => { + const count = Math.min(candidates.length, reserved); + remaining -= count; + return count; + }); + for (const index of [2, 0, 1]) { + const extra = Math.min(remaining, groups[index].candidates.length - counts[index]); + counts[index] += extra; + remaining -= extra; + } + return groups.flatMap(({ candidates }, index) => candidates.slice(0, counts[index])); }, [activeSessionMentions, availableSessionMentions, cliAppMention, cliApps, mcpPresets]); const showCliAppMenu = filteredMentionCandidates.length > 0; @@ -2660,7 +2669,7 @@ function CliAppMentionPalette({ layout.maxHeight - SLASH_PALETTE_CHROME_PX, ); const listRef = useSelectedOptionScroll(selectedIndex); - const groupedCandidates = (["session", "cli", "mcp"] as const) + const groupedCandidates = (["cli", "mcp", "session"] as const) .map((kind) => ({ kind, label: kind === "session" diff --git a/webui/src/tests/thread-composer.test.tsx b/webui/src/tests/thread-composer.test.tsx index d029c9265..af99b1a20 100644 --- a/webui/src/tests/thread-composer.test.tsx +++ b/webui/src/tests/thread-composer.test.tsx @@ -1596,6 +1596,7 @@ describe("ThreadComposer", () => { onSend={vi.fn()} placeholder="Type your message..." cliApps={CLI_APPS} + mcpPresets={MCP_PRESETS} sessions={[ ...["a", "b"].map((chatId) => session(chatId, "Plan")), session("blender-chat", "Blender", "3D notes"), @@ -1606,6 +1607,10 @@ describe("ThreadComposer", () => { const input = screen.getByLabelText("Message input"); fireEvent.change(input, { target: { value: "@", selectionStart: 1 } }); + const palette = screen.getByRole("listbox", { name: "Mentions" }); + expect(within(palette).getAllByRole("group").map((group) => ( + group.getAttribute("aria-label") + ))).toEqual(["CLI apps", "MCP services", "Nanobot conversations"]); const options = screen.getAllByRole("option", { name: /Plan @Plan/i }); expect(options.map((option) => option.textContent)).toEqual([ expect.stringContaining("@Plan"),