From bd35b793e19177c9e49769c99933d0a8646cf2b7 Mon Sep 17 00:00:00 2001 From: RockChinQ Date: Mon, 13 Jul 2026 15:08:45 +0800 Subject: [PATCH] fix(embed): isolate browser chat sessions (#2335) * fix(embed): isolate browser chat sessions * fix(embed): harden reset and reconnect recovery --- .../http/controller/groups/pipelines/embed.py | 20 +- .../pkg/platform/sources/websocket_adapter.py | 186 ++++++++---- .../pkg/platform/sources/websocket_manager.py | 60 +++- src/langbot/templates/embed/widget.js | 92 +++++- tests/integration/api/test_embed.py | 37 ++- .../test_websocket_session_isolation.py | 267 ++++++++++++++++++ 6 files changed, 588 insertions(+), 74 deletions(-) create mode 100644 tests/unit_tests/platform/test_websocket_session_isolation.py diff --git a/src/langbot/pkg/api/http/controller/groups/pipelines/embed.py b/src/langbot/pkg/api/http/controller/groups/pipelines/embed.py index 99c00944f..50e9112b7 100644 --- a/src/langbot/pkg/api/http/controller/groups/pipelines/embed.py +++ b/src/langbot/pkg/api/http/controller/groups/pipelines/embed.py @@ -21,7 +21,7 @@ import quart from ... import group from ......utils import paths -from ......platform.sources.websocket_manager import ws_connection_manager +from ......platform.sources.websocket_manager import is_valid_session_id, ws_connection_manager logger = logging.getLogger(__name__) @@ -203,11 +203,15 @@ class EmbedRouterGroup(group.RouterGroup): if session_type not in ['person', 'group']: return self.http_status(400, -1, 'session_type must be person or group') + session_id = quart.request.args.get('session_id', '') + if not is_valid_session_id(session_id): + return self.http_status(400, -1, 'Valid session_id is required') + websocket_adapter = self.ap.platform_mgr.websocket_proxy_bot.adapter if not websocket_adapter: return self.http_status(404, -1, 'WebSocket adapter not found') - messages = websocket_adapter.get_websocket_messages(pipeline_uuid, session_type) + messages = websocket_adapter.get_websocket_messages(pipeline_uuid, session_type, session_id) return self.success(data={'messages': messages}) except Exception as e: @@ -227,11 +231,15 @@ class EmbedRouterGroup(group.RouterGroup): if session_type not in ['person', 'group']: return self.http_status(400, -1, 'session_type must be person or group') + session_id = quart.request.args.get('session_id', '') + if not is_valid_session_id(session_id): + return self.http_status(400, -1, 'Valid session_id is required') + websocket_adapter = self.ap.platform_mgr.websocket_proxy_bot.adapter if not websocket_adapter: return self.http_status(404, -1, 'WebSocket adapter not found') - websocket_adapter.reset_session(pipeline_uuid, session_type) + websocket_adapter.reset_session(pipeline_uuid, session_type, session_id) return self.success(data={'message': 'Session reset successfully'}) except Exception as e: @@ -294,6 +302,11 @@ class EmbedRouterGroup(group.RouterGroup): ) return + session_id = quart.websocket.args.get('session_id', '') + if not is_valid_session_id(session_id): + await quart.websocket.send(json.dumps({'type': 'error', 'message': 'Valid session_id is required'})) + return + websocket_adapter = self.ap.platform_mgr.websocket_proxy_bot.adapter if not websocket_adapter: await quart.websocket.send(json.dumps({'type': 'error', 'message': 'WebSocket adapter not found'})) @@ -304,6 +317,7 @@ class EmbedRouterGroup(group.RouterGroup): websocket=quart.websocket._get_current_object(), pipeline_uuid=pipeline_uuid, session_type=session_type, + session_id=session_id, metadata={'user_agent': quart.websocket.headers.get('User-Agent', '')}, ) diff --git a/src/langbot/pkg/platform/sources/websocket_adapter.py b/src/langbot/pkg/platform/sources/websocket_adapter.py index 0574292f3..a66d4ac9a 100644 --- a/src/langbot/pkg/platform/sources/websocket_adapter.py +++ b/src/langbot/pkg/platform/sources/websocket_adapter.py @@ -13,7 +13,7 @@ import langbot_plugin.api.entities.builtin.platform.events as platform_events import langbot_plugin.api.entities.builtin.platform.entities as platform_entities import langbot_plugin.api.definition.abstract.platform.event_logger as abstract_platform_logger from ...core import app -from .websocket_manager import ws_connection_manager, WebSocketConnection +from .websocket_manager import WebSocketConnection, is_valid_session_id, ws_connection_manager logger = logging.getLogger(__name__) @@ -91,6 +91,59 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter) self.outbound_message_queue = asyncio.Queue() self.stream_enabled = True + @staticmethod + def _conversation_key(pipeline_uuid: str, session_id: str | None = None) -> str: + """Return the history key for a pipeline/client conversation.""" + return f'{pipeline_uuid}:{session_id}' if session_id else pipeline_uuid + + @staticmethod + def _parse_embed_target(target_id: str) -> tuple[str, str] | None: + """Extract pipeline and session identifiers from a stable embed launcher.""" + target_value = str(target_id) + for prefix in ('websocket_', 'websocketgroup_'): + if target_value.startswith(prefix): + target = target_value[len(prefix) :] + break + else: + return None + if ':' not in target: + return None + pipeline_uuid, session_id = target.rsplit(':', 1) + if not pipeline_uuid or not is_valid_session_id(session_id): + return None + return pipeline_uuid, session_id + + @classmethod + async def _get_connection_from_target(cls, target_id: str): + """Resolve a person or group WebSocket launcher to its connection.""" + target_value = str(target_id) + for prefix in ('websocket_', 'websocketgroup_'): + if target_value.startswith(prefix): + target = target_value[len(prefix) :] + break + else: + return None + connection = await ws_connection_manager.get_connection(target) + if connection is not None: + return connection + embed_target = cls._parse_embed_target(target_id) + if embed_target is not None: + pipeline_uuid, session_id = embed_target + return await ws_connection_manager.get_connection_by_session_id(session_id, pipeline_uuid) + return await ws_connection_manager.get_connection_by_session_id(target) + + async def _get_message_context(self, message_source) -> tuple[str, str | None]: + """Resolve the originating pipeline and browser session for a reply.""" + sender = getattr(message_source, 'sender', None) + sender_id = getattr(sender, 'id', '') + connection = await self._get_connection_from_target(sender_id) + if connection is not None: + return connection.pipeline_uuid, connection.session_id + embed_target = self._parse_embed_target(sender_id) + if embed_target is not None: + return embed_target + return typing.cast(str, self.ap.platform_mgr.websocket_proxy_bot.bot_entity.use_pipeline_uuid), None + async def send_message( self, target_type: str, @@ -103,15 +156,26 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter) target_id 可能是 launcher_id（如 websocket_xxx）或 pipeline_uuid。 我们需要尝试两种方式来确保消息能够送达。 """ - # 获取当前的 pipeline_uuid - pipeline_uuid = self.ap.platform_mgr.websocket_proxy_bot.bot_entity.use_pipeline_uuid + connection = await self._get_connection_from_target(target_id) + if connection is not None: + pipeline_uuid = connection.pipeline_uuid + session_id = connection.session_id + else: + embed_target = self._parse_embed_target(target_id) + if embed_target is not None: + pipeline_uuid, session_id = embed_target + else: + pipeline_uuid = typing.cast( + str, + self.ap.platform_mgr.websocket_proxy_bot.bot_entity.use_pipeline_uuid, + ) + session_id = None session_type = 'group' if target_type == 'group' else 'person' + conversation_key = self._conversation_key(pipeline_uuid, session_id) - # 选择会话 session = self.websocket_group_session if session_type == 'group' else self.websocket_person_session - # 生成唯一消息ID - msg_id = len(session.get_message_list(pipeline_uuid)) + 1 + msg_id = len(session.get_message_list(conversation_key)) + 1 message_data = WebSocketMessage( id=msg_id, @@ -122,10 +186,8 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter) is_final=True, ) - # 保存到历史记录 - session.get_message_list(pipeline_uuid).append(message_data) + session.get_message_list(conversation_key).append(message_data) - # 直接广播到当前pipeline的连接 await ws_connection_manager.broadcast_to_pipeline( pipeline_uuid, { @@ -134,6 +196,7 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter) 'data': message_data.model_dump(), }, session_type=session_type, + session_id=session_id, ) return message_data.model_dump() @@ -152,12 +215,11 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter) else self.websocket_person_session ) - # 从message_source获取pipeline_uuid和connection_id - pipeline_uuid = self.ap.platform_mgr.websocket_proxy_bot.bot_entity.use_pipeline_uuid + pipeline_uuid, session_id = await self._get_message_context(message_source) session_type = 'group' if isinstance(message_source, platform_events.GroupMessage) else 'person' + conversation_key = self._conversation_key(pipeline_uuid, session_id) - # 生成新的消息ID - msg_id = len(session.get_message_list(pipeline_uuid)) + 1 + msg_id = len(session.get_message_list(conversation_key)) + 1 message_data = WebSocketMessage( id=msg_id, @@ -168,10 +230,8 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter) is_final=True, ) - # 保存到历史记录 - session.get_message_list(pipeline_uuid).append(message_data) + session.get_message_list(conversation_key).append(message_data) - # 直接广播到所有该pipeline的连接，包含session_type信息 await ws_connection_manager.broadcast_to_pipeline( pipeline_uuid, { @@ -180,6 +240,7 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter) 'data': message_data.model_dump(), }, session_type=session_type, + session_id=session_id, ) return message_data.model_dump() @@ -200,10 +261,11 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter) else self.websocket_person_session ) - pipeline_uuid = self.ap.platform_mgr.websocket_proxy_bot.bot_entity.use_pipeline_uuid + pipeline_uuid, session_id = await self._get_message_context(message_source) session_type = 'group' if isinstance(message_source, platform_events.GroupMessage) else 'person' - message_list = session.get_message_list(pipeline_uuid) - stream_message_indexes = session.get_stream_message_indexes(pipeline_uuid) + conversation_key = self._conversation_key(pipeline_uuid, session_id) + message_list = session.get_message_list(conversation_key) + stream_message_indexes = session.get_stream_message_indexes(conversation_key) # Streaming messages in LangBot have a stable resp_message_id during the same assistant reply. # Use it as the primary key to avoid overwriting an old card from a previous reply. @@ -247,7 +309,6 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter) if message_is_final and resp_message_id: stream_message_indexes.pop(resp_message_id, None) - # 直接广播到所有该pipeline的连接，包含session_type信息 await ws_connection_manager.broadcast_to_pipeline( pipeline_uuid, { @@ -256,6 +317,7 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter) 'data': message_data.model_dump(), }, session_type=session_type, + session_id=session_id, ) return message_data.model_dump() @@ -381,23 +443,19 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter) """ pipeline_uuid = connection.pipeline_uuid session_type = connection.session_type + conversation_key = self._conversation_key(pipeline_uuid, connection.session_id) - # 获取stream参数，默认为True self.stream_enabled = message_data.get('stream', True) - # 选择会话 use_session = self.websocket_group_session if session_type == 'group' else self.websocket_person_session - # 解析消息链 message_chain_obj = message_data.get('message', []) - # 处理图片组件：将path转换为base64 await self._process_image_components(message_chain_obj) message_chain = platform_message.MessageChain.model_validate(message_chain_obj) - # 生成消息ID - message_id = len(use_session.get_message_list(pipeline_uuid)) + 1 + message_id = len(use_session.get_message_list(conversation_key)) + 1 # 保存用户消息 user_message = WebSocketMessage( @@ -409,9 +467,8 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter) connection_id=connection.connection_id, is_final=True, # 用户消息始终是完整的，非流式 ) - use_session.get_message_list(pipeline_uuid).append(user_message) + use_session.get_message_list(conversation_key).append(user_message) - # 广播用户消息到所有连接（包括发送者），包含session_type信息 await ws_connection_manager.broadcast_to_pipeline( pipeline_uuid, { @@ -420,25 +477,27 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter) 'data': user_message.model_dump(), }, session_type=session_type, + session_id=connection.session_id, ) # 添加消息源 message_chain.insert(0, platform_message.Source(id=message_id, time=datetime.now().timestamp())) # 创建事件 + launcher_id = f'{pipeline_uuid}:{connection.session_id}' if connection.session_id else connection.connection_id if session_type == 'person': - sender = platform_entities.Friend( - id=f'websocket_{connection.connection_id}', nickname='User', remark='User' - ) + sender = platform_entities.Friend(id=f'websocket_{launcher_id}', nickname='User', remark='User') event = platform_events.FriendMessage( sender=sender, message_chain=message_chain, time=datetime.now().timestamp() ) else: group = platform_entities.Group( - id='websocketgroup', name='Group', permission=platform_entities.Permission.Member + id=f'websocketgroup_{launcher_id}' if connection.session_id else 'websocketgroup', + name='Group', + permission=platform_entities.Permission.Member, ) sender = platform_entities.GroupMember( - id=f'websocket_{connection.connection_id}', + id=f'websocket_{launcher_id}', member_name='User', group=group, permission=platform_entities.Permission.Member, @@ -468,22 +527,47 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter) if event.__class__ in listeners: asyncio.create_task(listeners[event.__class__](event, callback_adapter)) - def get_websocket_messages(self, pipeline_uuid: str, session_type: str) -> list[dict]: - """获取消息历史""" - if session_type == 'person': - return [message.model_dump() for message in self.websocket_person_session.get_message_list(pipeline_uuid)] - else: - return [message.model_dump() for message in self.websocket_group_session.get_message_list(pipeline_uuid)] + def get_websocket_messages( + self, + pipeline_uuid: str, + session_type: str, + session_id: str | None = None, + ) -> list[dict]: + """Return history for one pipeline/client conversation.""" + conversation_key = self._conversation_key(pipeline_uuid, session_id) + session = self.websocket_person_session if session_type == 'person' else self.websocket_group_session + return [message.model_dump() for message in session.message_lists.get(conversation_key, [])] - def reset_session(self, pipeline_uuid: str, session_type: str): - """重置会话""" - if session_type == 'person': - if pipeline_uuid in self.websocket_person_session.message_lists: - self.websocket_person_session.message_lists[pipeline_uuid] = [] - if pipeline_uuid in self.websocket_person_session.stream_message_indexes: - self.websocket_person_session.stream_message_indexes[pipeline_uuid] = {} - else: - if pipeline_uuid in self.websocket_group_session.message_lists: - self.websocket_group_session.message_lists[pipeline_uuid] = [] - if pipeline_uuid in self.websocket_group_session.stream_message_indexes: - self.websocket_group_session.stream_message_indexes[pipeline_uuid] = {} + def reset_session( + self, + pipeline_uuid: str, + session_type: str, + session_id: str | None = None, + ): + """Reset one pipeline/client conversation.""" + conversation_key = self._conversation_key(pipeline_uuid, session_id) + session = self.websocket_person_session if session_type == 'person' else self.websocket_group_session + if conversation_key in session.message_lists: + session.message_lists[conversation_key] = [] + if conversation_key in session.stream_message_indexes: + session.stream_message_indexes[conversation_key] = {} + + if session_id: + launcher_id = ( + f'websocketgroup_{pipeline_uuid}:{session_id}' + if session_type == 'group' + else f'websocket_{pipeline_uuid}:{session_id}' + ) + self.ap.sess_mgr.session_list = [ + candidate_session + for candidate_session in self.ap.sess_mgr.session_list + if not ( + str( + candidate_session.launcher_type.value + if hasattr(candidate_session.launcher_type, 'value') + else candidate_session.launcher_type + ) + == session_type + and str(candidate_session.launcher_id) == launcher_id + ) + ] diff --git a/src/langbot/pkg/platform/sources/websocket_manager.py b/src/langbot/pkg/platform/sources/websocket_manager.py index 425dc9e09..c0c4d1807 100644 --- a/src/langbot/pkg/platform/sources/websocket_manager.py +++ b/src/langbot/pkg/platform/sources/websocket_manager.py @@ -9,6 +9,16 @@ from datetime import datetime import pydantic logger = logging.getLogger(__name__) +_SESSION_FILTER_UNSET = object() + + +def is_valid_session_id(value: str) -> bool: + """Accept only canonical random UUIDs for client conversation identifiers.""" + try: + parsed = uuid.UUID(value) + except (ValueError, TypeError, AttributeError): + return False + return parsed.version == 4 and str(parsed) == value class WebSocketConnection(pydantic.BaseModel): @@ -25,6 +35,9 @@ class WebSocketConnection(pydantic.BaseModel): session_type: str # 'person' or 'group' """会话类型""" + session_id: str | None = None + """Optional client conversation identifier used by embed widgets.""" + websocket: typing.Any = pydantic.Field(exclude=True) """WebSocket连接对象 (quart.websocket)""" @@ -65,13 +78,15 @@ class WebSocketConnectionManager: websocket: typing.Any, pipeline_uuid: str, session_type: str, - metadata: dict = None, + metadata: dict | None = None, + session_id: str | None = None, ) -> WebSocketConnection: - """添加新的WebSocket连接""" + """Register a WebSocket connection and its optional embed session.""" async with self._lock: connection = WebSocketConnection( pipeline_uuid=pipeline_uuid, session_type=session_type, + session_id=session_id, websocket=websocket, metadata=metadata or {}, ) @@ -120,10 +135,25 @@ class WebSocketConnectionManager: logger.debug(f'WebSocket connection disconnected: {connection_id}') - async def get_connection(self, connection_id: str) -> typing.Optional[WebSocketConnection]: - """获取指定连接""" + async def get_connection(self, connection_id: str) -> WebSocketConnection | None: + """Get a connection by its transport identifier.""" return self.connections.get(connection_id) + async def get_connection_by_session_id( + self, + session_id: str, + pipeline_uuid: str | None = None, + ) -> WebSocketConnection | None: + """Get an active embed connection by its stable browser session identifier.""" + for connection in self.connections.values(): + if ( + connection.session_id == session_id + and connection.is_active + and (pipeline_uuid is None or connection.pipeline_uuid == pipeline_uuid) + ): + return connection + return None + async def get_connections_by_pipeline(self, pipeline_uuid: str) -> list[WebSocketConnection]: """获取指定流水线的所有连接""" connection_ids = self.pipeline_connections.get(pipeline_uuid, set()) @@ -134,20 +164,30 @@ class WebSocketConnectionManager: connection_ids = self.session_connections.get(session_type, set()) return [self.connections[cid] for cid in connection_ids if cid in self.connections] - async def broadcast_to_pipeline(self, pipeline_uuid: str, message: dict, session_type: str = None): - """向指定流水线的所有连接广播消息 + async def broadcast_to_pipeline( + self, + pipeline_uuid: str, + message: dict, + session_type: str | None = None, + session_id: typing.Any = _SESSION_FILTER_UNSET, + ): + """Broadcast a message to matching connections for one pipeline. Args: - pipeline_uuid: 流水线UUID - message: 要广播的消息 - session_type: 可选的会话类型过滤器，如果提供则只向匹配的session_type连接广播 + pipeline_uuid: Pipeline identifier. + message: Serialized message to enqueue. + session_type: Optional session-type filter. + session_id: Embed conversation filter. Omit it to broadcast across + conversations; pass ``None`` to target non-embed connections. """ connections = await self.get_connections_by_pipeline(pipeline_uuid) - # 如果指定了session_type，只向匹配的连接广播 if session_type is not None: connections = [conn for conn in connections if conn.session_type == session_type] + if session_id is not _SESSION_FILTER_UNSET: + connections = [conn for conn in connections if conn.session_id == session_id] + tasks = [] for conn in connections: tasks.append(self.send_to_connection(conn.connection_id, message)) diff --git a/src/langbot/templates/embed/widget.js b/src/langbot/templates/embed/widget.js index 72be5eeba..e08ffe180 100644 --- a/src/langbot/templates/embed/widget.js +++ b/src/langbot/templates/embed/widget.js @@ -303,11 +303,60 @@ ''; // ========== State ========== + function createSessionId() { + if (window.crypto && typeof window.crypto.randomUUID === "function") { + return window.crypto.randomUUID(); + } + + var bytes = new Uint8Array(16); + if (!window.crypto || typeof window.crypto.getRandomValues !== "function") { + throw new Error("Secure random number generation is unavailable"); + } + window.crypto.getRandomValues(bytes); + bytes[6] = (bytes[6] & 15) | 64; + bytes[8] = (bytes[8] & 63) | 128; + var hex = Array.prototype.map + .call(bytes, function (value) { + return value.toString(16).padStart(2, "0"); + }) + .join(""); + return ( + hex.slice(0, 8) + + "-" + + hex.slice(8, 12) + + "-" + + hex.slice(12, 16) + + "-" + + hex.slice(16, 20) + + "-" + + hex.slice(20) + ); + } + + function getOrCreateSessionId() { + var storageKey = "langbot_embed_session_" + CONFIG.botUuid; + try { + var stored = window.sessionStorage.getItem(storageKey); + if ( + /^[a-f0-9]{8}-[a-f0-9]{4}-4[a-f0-9]{3}-[89ab][a-f0-9]{3}-[a-f0-9]{12}$/.test( + stored || "", + ) + ) + return stored; + var created = createSessionId(); + window.sessionStorage.setItem(storageKey, created); + return created; + } catch (e) { + return createSessionId(); + } + } + var state = { isOpen: false, isConnected: false, ws: null, connectionId: null, + sessionId: getOrCreateSessionId(), reconnectAttempts: 0, heartbeatTimer: null, messages: [], @@ -315,6 +364,9 @@ isStreaming: false, streamingMsgId: null, historyLoaded: false, + hasConnected: false, + messageVersion: 0, + historyReloadTimer: null, pendingImage: null, feedbackState: {}, }; @@ -473,7 +525,9 @@ "/api/v1/embed/" + CONFIG.botUuid + "/ws/connect?session_type=" + - CONFIG.sessionType; + CONFIG.sessionType + + "&session_id=" + + encodeURIComponent(state.sessionId); try { state.ws = new WebSocket(url); @@ -520,6 +574,8 @@ case "connected": state.isConnected = true; state.connectionId = data.connection_id; + if (state.hasConnected) loadHistory(true); + state.hasConnected = true; updateStatusDot(); updateSendBtn(); break; @@ -587,6 +643,7 @@ if (existingIdx >= 0) { state.messages[existingIdx] = msg; + state.messageVersion++; updateMessageEl(existingIdx, msg); } else { addMessage(msg); @@ -656,16 +713,27 @@ } // ========== Message History ========== - function loadHistory() { - if (state.historyLoaded) return; + function scheduleHistoryReload() { + if (state.historyReloadTimer) clearTimeout(state.historyReloadTimer); + state.historyReloadTimer = setTimeout(function () { + state.historyReloadTimer = null; + loadHistory(true); + }, 100); + } + + function loadHistory(force) { + if (state.historyLoaded && !force) return; state.historyLoaded = true; + var messageVersion = state.messageVersion; var url = CONFIG.baseUrl + "/api/v1/embed/" + CONFIG.botUuid + "/messages/" + - CONFIG.sessionType; + CONFIG.sessionType + + "?session_id=" + + encodeURIComponent(state.sessionId); var headers = {}; if (state.sessionToken) headers["Authorization"] = "Bearer " + state.sessionToken; @@ -675,6 +743,16 @@ }) .then(function (json) { if (json.code === 0 && json.data && json.data.messages) { + if (force && messageVersion !== state.messageVersion) { + scheduleHistoryReload(); + return; + } + if (force) { + state.messages = []; + state.isStreaming = false; + state.streamingMsgId = null; + renderMessages(); + } var msgs = json.data.messages; for (var i = 0; i < msgs.length; i++) { addMessage(msgs[i], true); @@ -693,13 +771,16 @@ "/api/v1/embed/" + CONFIG.botUuid + "/reset/" + - CONFIG.sessionType; + CONFIG.sessionType + + "?session_id=" + + encodeURIComponent(state.sessionId); var headers = {}; if (state.sessionToken) headers["Authorization"] = "Bearer " + state.sessionToken; fetch(url, { method: "POST", headers: headers }) .then(function () { state.messages = []; + state.messageVersion++; state.isStreaming = false; state.streamingMsgId = null; state.historyLoaded = true; @@ -713,6 +794,7 @@ // ========== UI Rendering ========== function addMessage(msg, silent) { state.messages.push(msg); + if (!silent) state.messageVersion++; var el = createMessageEl(msg); if (els.welcome) { els.welcome.style.display = "none"; diff --git a/tests/integration/api/test_embed.py b/tests/integration/api/test_embed.py index 5d034d133..2862337f5 100644 --- a/tests/integration/api/test_embed.py +++ b/tests/integration/api/test_embed.py @@ -15,6 +15,7 @@ from tests.factories import FakeApp pytestmark = pytest.mark.integration +SESSION_ID = '31c0f2e9-b115-4ee6-8f15-3e624d6456b1' @pytest.fixture(scope='module') @@ -191,10 +192,10 @@ class TestEmbedMessagesEndpoint: """Tests for messages endpoint.""" @pytest.mark.asyncio - async def test_get_messages_person_success(self, quart_test_client): + async def test_get_messages_person_success(self, quart_test_client, fake_embed_app): """GET messages/person returns messages.""" response = await quart_test_client.get( - '/api/v1/embed/a1b2c3d4-5678-90ab-cdef-123456789abc/messages/person', + f'/api/v1/embed/a1b2c3d4-5678-90ab-cdef-123456789abc/messages/person?session_id={SESSION_ID}', headers={'Authorization': 'Bearer 1234567890.dummy'}, ) @@ -202,17 +203,30 @@ class TestEmbedMessagesEndpoint: data = await response.get_json() assert data['code'] == 0 assert 'messages' in data['data'] + fake_embed_app.platform_mgr.websocket_proxy_bot.adapter.get_websocket_messages.assert_called_with( + 'test-pipeline-uuid', 'person', SESSION_ID + ) @pytest.mark.asyncio async def test_get_messages_group_success(self, quart_test_client): """GET messages/group returns messages.""" response = await quart_test_client.get( - '/api/v1/embed/a1b2c3d4-5678-90ab-cdef-123456789abc/messages/group', + f'/api/v1/embed/a1b2c3d4-5678-90ab-cdef-123456789abc/messages/group?session_id={SESSION_ID}', headers={'Authorization': 'Bearer 1234567890.dummy'}, ) assert response.status_code == 200 + @pytest.mark.asyncio + async def test_get_messages_requires_session_id(self, quart_test_client): + """GET messages without a client session identifier returns 400.""" + response = await quart_test_client.get( + '/api/v1/embed/a1b2c3d4-5678-90ab-cdef-123456789abc/messages/person', + headers={'Authorization': 'Bearer 1234567890.dummy'}, + ) + + assert response.status_code == 400 + @pytest.mark.asyncio async def test_get_messages_invalid_session_type(self, quart_test_client): """GET messages with invalid session_type returns 400.""" @@ -229,16 +243,29 @@ class TestEmbedResetEndpoint: """Tests for session reset endpoint.""" @pytest.mark.asyncio - async def test_reset_session_person_success(self, quart_test_client): + async def test_reset_session_person_success(self, quart_test_client, fake_embed_app): """POST reset/person resets session.""" response = await quart_test_client.post( - '/api/v1/embed/a1b2c3d4-5678-90ab-cdef-123456789abc/reset/person', + f'/api/v1/embed/a1b2c3d4-5678-90ab-cdef-123456789abc/reset/person?session_id={SESSION_ID}', headers={'Authorization': 'Bearer 1234567890.dummy'}, ) assert response.status_code == 200 data = await response.get_json() assert data['code'] == 0 + fake_embed_app.platform_mgr.websocket_proxy_bot.adapter.reset_session.assert_called_with( + 'test-pipeline-uuid', 'person', SESSION_ID + ) + + @pytest.mark.asyncio + async def test_reset_session_requires_session_id(self, quart_test_client): + """POST reset without a client session identifier returns 400.""" + response = await quart_test_client.post( + '/api/v1/embed/a1b2c3d4-5678-90ab-cdef-123456789abc/reset/person', + headers={'Authorization': 'Bearer 1234567890.dummy'}, + ) + + assert response.status_code == 400 @pytest.mark.asyncio async def test_reset_session_invalid_uuid(self, quart_test_client): diff --git a/tests/unit_tests/platform/test_websocket_session_isolation.py b/tests/unit_tests/platform/test_websocket_session_isolation.py new file mode 100644 index 000000000..d86580000 --- /dev/null +++ b/tests/unit_tests/platform/test_websocket_session_isolation.py @@ -0,0 +1,267 @@ +"""Regression tests for isolated embed-widget conversations.""" + +import asyncio +from pathlib import Path +from unittest.mock import AsyncMock, Mock + +import pytest + +import langbot_plugin.api.entities.builtin.platform.events as platform_events +from langbot.pkg.platform.sources import websocket_adapter as websocket_adapter_module +from langbot.pkg.platform.sources.websocket_adapter import WebSocketAdapter, WebSocketMessage, WebSocketSession +from langbot.pkg.platform.sources.websocket_manager import WebSocketConnectionManager, is_valid_session_id + + +@pytest.mark.asyncio +async def test_broadcast_only_reaches_connections_in_same_browser_session(): + manager = WebSocketConnectionManager() + first = await manager.add_connection( + websocket=Mock(), + pipeline_uuid='pipeline-1', + session_type='person', + session_id='session-a', + ) + second = await manager.add_connection( + websocket=Mock(), + pipeline_uuid='pipeline-1', + session_type='person', + session_id='session-b', + ) + dashboard = await manager.add_connection( + websocket=Mock(), + pipeline_uuid='pipeline-1', + session_type='person', + ) + + await manager.broadcast_to_pipeline( + 'pipeline-1', + {'type': 'response'}, + session_type='person', + session_id='session-a', + ) + + assert await first.send_queue.get() == {'type': 'response'} + assert second.send_queue.empty() + assert dashboard.send_queue.empty() + + await manager.broadcast_to_pipeline( + 'pipeline-1', + {'type': 'dashboard-response'}, + session_type='person', + session_id=None, + ) + + assert await dashboard.send_queue.get() == {'type': 'dashboard-response'} + assert first.send_queue.empty() + assert second.send_queue.empty() + + +@pytest.mark.asyncio +async def test_embed_event_uses_stable_session_launcher(monkeypatch): + manager = WebSocketConnectionManager() + session_id = '31c0f2e9-b115-4ee6-8f15-3e624d6456b1' + connection = await manager.add_connection( + websocket=Mock(), + pipeline_uuid='pipeline-1', + session_type='person', + session_id=session_id, + ) + monkeypatch.setattr(websocket_adapter_module, 'ws_connection_manager', manager) + + adapter = WebSocketAdapter.model_construct(ap=Mock(), logger=AsyncMock()) + adapter.websocket_person_session = WebSocketSession(id='person') + adapter.websocket_group_session = WebSocketSession(id='group') + received = [] + + async def listener(event, _callback_adapter): + received.append(event) + + adapter.listeners = {platform_events.FriendMessage: listener} + await adapter.handle_websocket_message( + connection, + {'message': [{'type': 'Plain', 'text': 'hello'}], 'stream': False}, + ) + await asyncio.sleep(0) + + assert received[0].sender.id == f'websocket_pipeline-1:{session_id}' + + +@pytest.mark.asyncio +async def test_embed_group_event_uses_stable_session_launcher(monkeypatch): + manager = WebSocketConnectionManager() + session_id = '31c0f2e9-b115-4ee6-8f15-3e624d6456b1' + connection = await manager.add_connection( + websocket=Mock(), + pipeline_uuid='pipeline-1', + session_type='group', + session_id=session_id, + ) + monkeypatch.setattr(websocket_adapter_module, 'ws_connection_manager', manager) + + adapter = WebSocketAdapter.model_construct(ap=Mock(), logger=AsyncMock()) + adapter.websocket_person_session = WebSocketSession(id='person') + adapter.websocket_group_session = WebSocketSession(id='group') + received = [] + + async def listener(event, _callback_adapter): + received.append(event) + + adapter.listeners = {platform_events.GroupMessage: listener} + await adapter.handle_websocket_message( + connection, + {'message': [{'type': 'Plain', 'text': 'hello'}], 'stream': False}, + ) + await asyncio.sleep(0) + + assert received[0].sender.id == f'websocket_pipeline-1:{session_id}' + assert received[0].sender.group.id == f'websocketgroup_pipeline-1:{session_id}' + + dashboard = await manager.add_connection( + websocket=Mock(), + pipeline_uuid='pipeline-1', + session_type='group', + ) + await adapter.handle_websocket_message( + dashboard, + {'stream': False, 'message': [{'type': 'Plain', 'text': 'dashboard'}]}, + None, + ) + await asyncio.sleep(0) + + assert received[1].sender.id == f'websocket_{dashboard.connection_id}' + assert received[1].sender.group.id == 'websocketgroup' + + +@pytest.mark.asyncio +async def test_stable_session_launcher_resolves_to_active_connection(monkeypatch): + manager = WebSocketConnectionManager() + session_id = '31c0f2e9-b115-4ee6-8f15-3e624d6456b1' + await manager.add_connection( + websocket=Mock(), + pipeline_uuid='pipeline-2', + session_type='person', + session_id=session_id, + ) + connection = await manager.add_connection( + websocket=Mock(), + pipeline_uuid='pipeline-1', + session_type='person', + session_id=session_id, + ) + monkeypatch.setattr(websocket_adapter_module, 'ws_connection_manager', manager) + + adapter = WebSocketAdapter.model_construct(ap=Mock(), logger=AsyncMock()) + message_source = Mock() + message_source.sender.id = f'websocket_pipeline-1:{session_id}' + + assert await adapter._get_message_context(message_source) == ('pipeline-1', session_id) + assert await adapter._get_connection_from_target(f'websocketgroup_pipeline-1:{session_id}') is connection + assert await manager.get_connection_by_session_id(session_id, 'pipeline-1') is connection + + await manager.remove_connection(connection.connection_id) + + assert await adapter._get_message_context(message_source) == ('pipeline-1', session_id) + assert await manager.get_connection_by_session_id(session_id, 'pipeline-1') is None + + +def test_session_ids_must_be_canonical_random_uuids(): + assert is_valid_session_id('31c0f2e9-b115-4ee6-8f15-3e624d6456b1') + assert not is_valid_session_id('session-a') + assert not is_valid_session_id('00000000-0000-0000-0000-000000000000') + + +def test_history_read_does_not_allocate_unknown_session(): + adapter = WebSocketAdapter.model_construct(ap=Mock(), logger=AsyncMock()) + adapter.websocket_person_session = WebSocketSession(id='person') + adapter.websocket_group_session = WebSocketSession(id='group') + + assert adapter.get_websocket_messages('pipeline-1', 'person', 'missing-session') == [] + assert adapter.websocket_person_session.message_lists == {} + + +def test_history_and_reset_are_scoped_to_browser_session(): + matching_provider_session = Mock( + launcher_type=Mock(value='person'), + launcher_id='websocket_pipeline-1:session-a', + ) + matching_group_provider_session = Mock( + launcher_type=Mock(value='group'), + launcher_id='websocketgroup_pipeline-1:session-a', + ) + other_session = Mock( + launcher_type=Mock(value='person'), + launcher_id='websocket_pipeline-1:session-b', + ) + ap = Mock() + ap.sess_mgr.session_list = [ + matching_provider_session, + matching_group_provider_session, + other_session, + ] + adapter = WebSocketAdapter.model_construct( + ap=ap, + logger=AsyncMock(), + ) + adapter.websocket_person_session = Mock() + adapter.websocket_group_session = Mock() + + session_a = [ + WebSocketMessage( + id=1, + role='user', + content='private-a', + message_chain=[], + timestamp='2026-07-13T00:00:00', + ) + ] + session_b = [ + WebSocketMessage( + id=1, + role='user', + content='private-b', + message_chain=[], + timestamp='2026-07-13T00:00:00', + ) + ] + histories = { + 'pipeline-1:session-a': session_a, + 'pipeline-1:session-b': session_b, + } + stream_indexes = { + 'pipeline-1:session-a': {'response-a': 0}, + 'pipeline-1:session-b': {'response-b': 0}, + } + adapter.websocket_person_session.get_message_list.side_effect = histories.__getitem__ + adapter.websocket_person_session.message_lists = histories + adapter.websocket_person_session.stream_message_indexes = stream_indexes + adapter.websocket_group_session.message_lists = {} + adapter.websocket_group_session.stream_message_indexes = {} + + assert adapter.get_websocket_messages('pipeline-1', 'person', 'session-a')[0]['content'] == 'private-a' + assert adapter.get_websocket_messages('pipeline-1', 'person', 'session-b')[0]['content'] == 'private-b' + + adapter.reset_session('pipeline-1', 'person', 'session-a') + + assert histories['pipeline-1:session-a'] == [] + assert stream_indexes['pipeline-1:session-a'] == {} + assert histories['pipeline-1:session-b'] == session_b + assert stream_indexes['pipeline-1:session-b'] == {'response-b': 0} + assert ap.sess_mgr.session_list == [matching_group_provider_session, other_session] + + adapter.reset_session('pipeline-1', 'group', 'session-a') + + assert ap.sess_mgr.session_list == [other_session] + + +def test_widget_sends_stable_session_id_to_all_conversation_endpoints(): + widget_path = Path(__file__).parents[3] / 'src/langbot/templates/embed/widget.js' + widget = widget_path.read_text(encoding='utf-8') + + assert 'langbot_embed_session_' in widget + assert 'window.sessionStorage' in widget + assert 'window.localStorage' not in widget + assert 'session_id=' in widget + assert widget.count('encodeURIComponent(state.sessionId)') >= 3 + assert 'loadHistory(true)' in widget + assert 'messageVersion !== state.messageVersion' in widget + assert 'scheduleHistoryReload();' in widget