xinyin025/LangBot
1
0
Fork 0
mirror of https://github.com/langbot-app/LangBot.git synced 2026-07-21 20:06:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(tenancy): implement workspace isolation

Browse Source
This commit is contained in:
Junyan Qin
2026-07-19 09:58:59 +08:00
parent 9eb292992d
commit c6f826fe2d
271 changed files with 31162 additions and 6106 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/unit_tests/platform/test_websocket_session_isolation.py
+133 -8
View File
@@ -9,7 +9,25 @@ 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
from langbot.pkg.platform.sources.websocket_manager import (
WebSocketConnectionManager,
WebSocketScope,
is_valid_session_id,
)
SCOPE_A = WebSocketScope('instance-a', 'workspace-a', 1)
SCOPE_B = WebSocketScope('instance-a', 'workspace-b', 1)
def _adapter_logger(scope: WebSocketScope = SCOPE_A):
logger = AsyncMock()
logger.execution_context = Mock(
instance_uuid=scope.instance_uuid,
workspace_uuid=scope.workspace_uuid,
placement_generation=scope.placement_generation,
)
return logger
@pytest.mark.asyncio
@@ -17,18 +35,21 @@ async def test_broadcast_only_reaches_connections_in_same_browser_session():
manager = WebSocketConnectionManager()
first = await manager.add_connection(
websocket=Mock(),
scope=SCOPE_A,
pipeline_uuid='pipeline-1',
session_type='person',
session_id='session-a',
)
second = await manager.add_connection(
websocket=Mock(),
scope=SCOPE_A,
pipeline_uuid='pipeline-1',
session_type='person',
session_id='session-b',
)
dashboard = await manager.add_connection(
websocket=Mock(),
scope=SCOPE_A,
pipeline_uuid='pipeline-1',
session_type='person',
)
@@ -36,6 +57,7 @@ async def test_broadcast_only_reaches_connections_in_same_browser_session():
await manager.broadcast_to_pipeline(
'pipeline-1',
{'type': 'response'},
scope=SCOPE_A,
session_type='person',
session_id='session-a',
)
@@ -47,6 +69,7 @@ async def test_broadcast_only_reaches_connections_in_same_browser_session():
await manager.broadcast_to_pipeline(
'pipeline-1',
{'type': 'dashboard-response'},
scope=SCOPE_A,
session_type='person',
session_id=None,
)
@@ -56,19 +79,49 @@ async def test_broadcast_only_reaches_connections_in_same_browser_session():
assert second.send_queue.empty()
@pytest.mark.asyncio
async def test_pipeline_indexes_and_broadcasts_are_workspace_scoped():
manager = WebSocketConnectionManager()
workspace_a = await manager.add_connection(
websocket=Mock(),
scope=SCOPE_A,
pipeline_uuid='shared-pipeline',
session_type='person',
)
workspace_b = await manager.add_connection(
websocket=Mock(),
scope=SCOPE_B,
pipeline_uuid='shared-pipeline',
session_type='person',
)
await manager.broadcast_to_pipeline(
'shared-pipeline',
{'type': 'workspace-a'},
scope=SCOPE_A,
)
assert await workspace_a.send_queue.get() == {'type': 'workspace-a'}
assert workspace_b.send_queue.empty()
assert await manager.get_connection(workspace_b.connection_id, scope=SCOPE_A) is None
assert await manager.get_connection(workspace_b.connection_id, scope=SCOPE_B) is workspace_b
assert manager.get_stats(scope=SCOPE_A)['total_connections'] == 1
@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(),
scope=SCOPE_A,
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 = WebSocketAdapter.model_construct(ap=Mock(), logger=_adapter_logger())
adapter.websocket_person_session = WebSocketSession(id='person')
adapter.websocket_group_session = WebSocketSession(id='group')
received = []
@@ -92,13 +145,14 @@ async def test_embed_group_event_uses_stable_session_launcher(monkeypatch):
session_id = '31c0f2e9-b115-4ee6-8f15-3e624d6456b1'
connection = await manager.add_connection(
websocket=Mock(),
scope=SCOPE_A,
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 = WebSocketAdapter.model_construct(ap=Mock(), logger=_adapter_logger())
adapter.websocket_person_session = WebSocketSession(id='person')
adapter.websocket_group_session = WebSocketSession(id='group')
received = []
@@ -118,6 +172,7 @@ async def test_embed_group_event_uses_stable_session_launcher(monkeypatch):
dashboard = await manager.add_connection(
websocket=Mock(),
scope=SCOPE_A,
pipeline_uuid='pipeline-1',
session_type='group',
)
@@ -138,30 +193,46 @@ async def test_stable_session_launcher_resolves_to_active_connection(monkeypatch
session_id = '31c0f2e9-b115-4ee6-8f15-3e624d6456b1'
await manager.add_connection(
websocket=Mock(),
scope=SCOPE_A,
pipeline_uuid='pipeline-2',
session_type='person',
session_id=session_id,
)
connection = await manager.add_connection(
websocket=Mock(),
scope=SCOPE_A,
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 = WebSocketAdapter.model_construct(ap=Mock(), logger=_adapter_logger())
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
assert (
await manager.get_connection_by_session_id(
session_id,
scope=SCOPE_A,
pipeline_uuid='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
assert (
await manager.get_connection_by_session_id(
session_id,
scope=SCOPE_A,
pipeline_uuid='pipeline-1',
)
is None
)
def test_session_ids_must_be_canonical_random_uuids():
@@ -171,7 +242,7 @@ def test_session_ids_must_be_canonical_random_uuids():
def test_history_read_does_not_allocate_unknown_session():
adapter = WebSocketAdapter.model_construct(ap=Mock(), logger=AsyncMock())
adapter = WebSocketAdapter.model_construct(ap=Mock(), logger=_adapter_logger())
adapter.websocket_person_session = WebSocketSession(id='person')
adapter.websocket_group_session = WebSocketSession(id='group')
@@ -179,16 +250,70 @@ def test_history_read_does_not_allocate_unknown_session():
assert adapter.websocket_person_session.message_lists == {}
@pytest.mark.asyncio
async def test_attachment_key_must_belong_to_connection_upload_scope():
manager = WebSocketConnectionManager()
connection = await manager.add_connection(
websocket=Mock(),
scope=SCOPE_A,
pipeline_uuid='pipeline-1',
session_type='person',
)
storage_mgr = Mock()
storage_mgr.scoped_prefix.return_value = 'v1/current/upload_image/'
storage_mgr.is_scoped_object_key.return_value = True
storage_mgr.storage_provider.load = AsyncMock(return_value=b'image')
storage_mgr.delete_scoped_object_key = AsyncMock()
adapter = WebSocketAdapter.model_construct(
ap=Mock(storage_mgr=storage_mgr),
logger=_adapter_logger(),
)
message_chain = [{'type': 'Image', 'path': 'v1/current/upload_image/key.png'}]
await adapter._process_image_components(connection, message_chain)
assert message_chain[0]['base64'].startswith('data:image/png;base64,')
assert message_chain[0]['path'] == ''
storage_mgr.scoped_prefix.assert_called_once_with(
connection.execution_context,
owner_type='upload_image',
)
storage_mgr.is_scoped_object_key.assert_called_once_with(
'v1/current/upload_image/key.png',
expected_owner_type='upload_image',
)
storage_mgr.delete_scoped_object_key.assert_awaited_once_with(
connection.execution_context,
'v1/current/upload_image/key.png',
expected_owner_type='upload_image',
)
with pytest.raises(ValueError, match='does not belong'):
await adapter._process_image_components(
connection,
[{'type': 'File', 'path': 'v1/other/upload/key.txt'}],
)
def test_history_and_reset_are_scoped_to_browser_session():
matching_provider_session = Mock(
instance_uuid=SCOPE_A.instance_uuid,
workspace_uuid=SCOPE_A.workspace_uuid,
placement_generation=SCOPE_A.placement_generation,
launcher_type=Mock(value='person'),
launcher_id='websocket_pipeline-1:session-a',
)
matching_group_provider_session = Mock(
instance_uuid=SCOPE_A.instance_uuid,
workspace_uuid=SCOPE_A.workspace_uuid,
placement_generation=SCOPE_A.placement_generation,
launcher_type=Mock(value='group'),
launcher_id='websocketgroup_pipeline-1:session-a',
)
other_session = Mock(
instance_uuid=SCOPE_A.instance_uuid,
workspace_uuid=SCOPE_A.workspace_uuid,
placement_generation=SCOPE_A.placement_generation,
launcher_type=Mock(value='person'),
launcher_id='websocket_pipeline-1:session-b',
)
@@ -200,7 +325,7 @@ def test_history_and_reset_are_scoped_to_browser_session():
]
adapter = WebSocketAdapter.model_construct(
ap=ap,
logger=AsyncMock(),
logger=_adapter_logger(),
)
adapter.websocket_person_session = Mock()
adapter.websocket_group_session = Mock()