mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-08-09 04:40:57 +00:00
feat(tenancy): add Workspace multi-tenant foundation (#2353)
* Document multi-tenant workspace architecture * Add OSS and commercial workspace boundaries * docs: redesign multi-tenant workspace architecture * feat(tenancy): implement workspace isolation * docs(tenancy): record verification evidence * docs(tenancy): revise single-instance SaaS topology * docs(tenancy): refine architecture options * docs: finalize cloud v2 multi-tenant decisions * feat(tenancy): establish cloud isolation foundations * feat(tenancy): harden shared cloud runtime boundaries * docs(tenancy): record final isolation verification * fix(tenancy): close isolation and permission gaps * docs(tenancy): record final isolation verification * feat(tenancy): connect cloud workspace control plane * fix(build): install git for pinned SDK * docs(cloud): update control plane verification * chore: update multi-tenant SDK pin * fix(cloud): skip legacy model sync during startup * test(cloud): preserve minimal model manager fixtures * fix(cloud): preserve authenticated account context * fix(cloud): reuse authenticated account for user info * feat(cloud): complete Workspace settings navigation * test(web): cover Workspace dropdown menu * feat(web): place workspace controls in sidebar * refactor(web): streamline workspace controls * style(web): format workspace layout test * fix(cloud): surface runtime and workspace plan status * fix(plugin): keep runtime identity stable across restarts * fix(ui): widen and center workspace switcher * fix(ui): hide roles from workspace switcher * fix(ui): align workspace switcher with sidebar entries * feat(workspace): add in-product collaboration and direct Cloud launch * style: format collaboration changes * fix(workspace): bind collaboration APIs to tenant UoW * fix(cloud): preserve Core-owned collaboration state * test(cloud): require Space identity for invite registration * feat(cloud): complete secure invitation experience * style(web): format invitation flows * fix(cloud): recover box runtime without unscoped skill reload * feat(oss): enforce invitation account and owner billing flows * style: format OSS account service * test(oss): cover invitation logout handoff * fix(oss): resolve workspace owner in scoped session * feat(cloud): harden multi-tenant runtime resources * fix(cloud): bound runtime restart storms * fix(cloud): eliminate periodic runtime CPU spikes * fix(cloud): enforce instance capacity ceilings * fix(cloud): scope public login capability discovery * fix(cloud): bound tenant maintenance and monitoring work * fix(runtime): bound tenant resource amplification * fix(deps): pin green multi-tenant plugin SDK * fix(cloud): handle unavailable skill capability * fix(security): require authentication for image file endpoint (H-2) - Changed /api/v1/files/image from AuthType.NONE to USER_TOKEN_OR_API_KEY - Added Permission.RESOURCE_VIEW requirement - Prevents unauthenticated cross-tenant file access via leaked keys - Fixes HIGH severity finding from multi-tenant security review docs: add comprehensive database migration guide - Complete migration steps for OSS → multi-tenant - Backup, execution, verification procedures - Rollback scenarios and recovery plans - Performance tuning recommendations * test: add comprehensive cross-tenant isolation tests Added 7 critical test scenarios for multi-tenant boundaries: - Cross-tenant bot access prevention - Viewer role read-only enforcement - Removed member immediate access revocation - Model provider credential isolation - WebSocket message isolation - Invitation token workspace scoping - Multi-workspace context validation These tests address P0-2 coverage gaps for: - workspaces.py (membership & invitation flows) - user.py (authentication & authorization) - websocket_chat.py (real-time isolation) - plugins.py (resource access control) docs: finalize database migration guide * fix(security): resolve M-1, M-2, M-3 security findings M-1: WebSocket authorization TOCTOU race (FIXED) - Changed _revalidate_websocket_authorization to return RequestContext - Ensures validated context is used immediately without race window - Prevents removed members from sending messages during revalidation gap M-2: Model Manager cache workspace isolation (VERIFIED) - Confirmed _CacheKey already uses 4-tuple: (instance, workspace, generation, resource) - Cache is properly scoped per workspace, no cross-tenant leakage possible - No code change needed, documented as working correctly M-3: Invitation lock workspace scoping (FIXED) - Changed lock key from token_digest to workspace_uuid:token_digest - Prevents DoS where attacker locks token in Workspace A to block Workspace B - Locks now isolated per workspace All MEDIUM severity findings from security review now resolved. * fix(cloud): unblock tenant CI and enforce knowledge quotas * fix(tenancy): scope rerank model sync --------- Co-authored-by: dadachann <185672915+dadachann@users.noreply.github.com>
This commit is contained in:
RockChinQ
@@ -68,6 +68,25 @@ async def test_connection_listener_only_suppresses_exact_duplicates():
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_connection_event_cache_is_bounded():
|
||||
adapter, _ = _make_adapter()
|
||||
|
||||
for index in range(150):
|
||||
await adapter._on_websocket_connection(aiocqhttp.Event({'self_id': index, 'time': index}))
|
||||
|
||||
assert len(adapter.on_websocket_connection_event_cache) == 100
|
||||
|
||||
|
||||
def test_group_lookup_caches_are_bounded():
|
||||
converter = AiocqhttpEventConverter()
|
||||
converter._group_name_cache = {index: (str(index), 10_000.0) for index in range(5000)}
|
||||
|
||||
converter._prune_caches(1.0)
|
||||
|
||||
assert len(converter._group_name_cache) == 4096
|
||||
|
||||
|
||||
def test_unregister_listener_removes_registered_wrapper():
|
||||
adapter, _ = _make_adapter()
|
||||
|
||||
|
||||
@@ -0,0 +1,402 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
from langbot.pkg.api.http.authz import WorkspaceRequiredError
|
||||
from langbot.pkg.api.http.context import ExecutionContext
|
||||
from langbot.pkg.entity.persistence.bot import Bot
|
||||
from langbot.pkg.platform.botmgr import PlatformManager, RuntimeBot
|
||||
from langbot.pkg.workspace.entities import WorkspaceExecutionBinding
|
||||
from langbot.pkg.workspace.errors import WorkspaceInvariantError
|
||||
import langbot_plugin.api.entities.builtin.platform.events as platform_events
|
||||
|
||||
|
||||
WORKSPACE_A = '00000000-0000-0000-0000-00000000000a'
|
||||
WORKSPACE_B = '00000000-0000-0000-0000-00000000000b'
|
||||
BOT_A = '10000000-0000-0000-0000-00000000000a'
|
||||
BOT_B = '10000000-0000-0000-0000-00000000000b'
|
||||
|
||||
|
||||
def _context(workspace_uuid: str, bot_uuid: str, generation: int = 4) -> ExecutionContext:
|
||||
return ExecutionContext(
|
||||
instance_uuid='instance',
|
||||
workspace_uuid=workspace_uuid,
|
||||
placement_generation=generation,
|
||||
bot_uuid=bot_uuid,
|
||||
)
|
||||
|
||||
|
||||
def _runtime(application, workspace_uuid: str, bot_uuid: str) -> RuntimeBot:
|
||||
entity = SimpleNamespace(
|
||||
uuid=bot_uuid,
|
||||
workspace_uuid=workspace_uuid,
|
||||
name='Same Name',
|
||||
enable=True,
|
||||
pipeline_routing_rules=[],
|
||||
use_pipeline_uuid=None,
|
||||
)
|
||||
return RuntimeBot(
|
||||
ap=application,
|
||||
bot_entity=entity,
|
||||
adapter=SimpleNamespace(),
|
||||
logger=SimpleNamespace(),
|
||||
execution_context=_context(workspace_uuid, bot_uuid),
|
||||
)
|
||||
|
||||
|
||||
class _WorkspaceService:
|
||||
async def get_execution_binding(self, workspace_uuid, expected_generation=None):
|
||||
if workspace_uuid not in {WORKSPACE_A, WORKSPACE_B} or expected_generation != 4:
|
||||
raise ValueError('stale')
|
||||
return SimpleNamespace(
|
||||
instance_uuid='instance',
|
||||
workspace_uuid=workspace_uuid,
|
||||
placement_generation=4,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def manager():
|
||||
application = SimpleNamespace(workspace_service=_WorkspaceService())
|
||||
platform_manager = PlatformManager(application)
|
||||
platform_manager.bots = [
|
||||
_runtime(application, WORKSPACE_A, BOT_A),
|
||||
_runtime(application, WORKSPACE_B, BOT_B),
|
||||
]
|
||||
return platform_manager
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runtime_lookup_cannot_guess_another_workspace_bot(manager):
|
||||
assert await manager.get_bot_by_uuid(_context(WORKSPACE_A, BOT_A), BOT_A) is manager.bots[0]
|
||||
assert await manager.get_bot_by_uuid(_context(WORKSPACE_B, BOT_A), BOT_A) is None
|
||||
assert await manager.get_bot_by_uuid(_context(WORKSPACE_A, BOT_B), BOT_B) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_public_route_key_resolves_bound_runtime_and_rejects_non_opaque_input(manager):
|
||||
assert await manager.resolve_public_bot(BOT_A) is manager.bots[0]
|
||||
assert await manager.resolve_public_bot('Same Name') is None
|
||||
assert await manager.resolve_public_bot('not-a-uuid') is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stale_runtime_generation_is_not_returned(manager):
|
||||
with pytest.raises(ValueError, match='stale'):
|
||||
await manager.get_bot_by_uuid(_context(WORKSPACE_A, BOT_A, generation=5), BOT_A)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generation_advance_shuts_down_and_prunes_old_workspace_bots():
|
||||
class NoGlobalIterationDict(dict):
|
||||
def __iter__(self):
|
||||
raise AssertionError('generation advance scanned every bot runtime')
|
||||
|
||||
def items(self):
|
||||
raise AssertionError('generation advance scanned every bot runtime')
|
||||
|
||||
def values(self):
|
||||
raise AssertionError('generation advance scanned every bot runtime')
|
||||
|
||||
manager = PlatformManager(SimpleNamespace())
|
||||
old_bot = SimpleNamespace(
|
||||
workspace_uuid=WORKSPACE_A,
|
||||
placement_generation=4,
|
||||
enable=True,
|
||||
shutdown=AsyncMock(),
|
||||
)
|
||||
other_bot = SimpleNamespace(
|
||||
workspace_uuid=WORKSPACE_B,
|
||||
placement_generation=4,
|
||||
enable=True,
|
||||
shutdown=AsyncMock(),
|
||||
)
|
||||
unrelated_bots = [
|
||||
SimpleNamespace(
|
||||
workspace_uuid=f'workspace-{index}',
|
||||
placement_generation=4,
|
||||
enable=False,
|
||||
shutdown=AsyncMock(),
|
||||
)
|
||||
for index in range(1_000)
|
||||
]
|
||||
manager.bots = [old_bot, other_bot, *unrelated_bots]
|
||||
old_context = _context(WORKSPACE_A, BOT_A, generation=4)
|
||||
next_context = _context(WORKSPACE_A, BOT_A, generation=5)
|
||||
|
||||
await manager._observe_execution_context(old_context)
|
||||
manager._bots_by_key = NoGlobalIterationDict(manager._bots_by_key)
|
||||
await manager._observe_execution_context(next_context)
|
||||
manager._bots_by_key = dict(manager._bots_by_key)
|
||||
|
||||
old_bot.shutdown.assert_awaited_once_with()
|
||||
assert manager.bots == [other_bot, *unrelated_bots]
|
||||
with pytest.raises(WorkspaceInvariantError, match='rolled back'):
|
||||
await manager._observe_execution_context(old_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_websocket_proxy_creation_reuses_one_runtime():
|
||||
created_adapters = []
|
||||
|
||||
class WebsocketAdapter:
|
||||
def __init__(self, *_args, **_kwargs):
|
||||
created_adapters.append(self)
|
||||
|
||||
def register_listener(self, *_args):
|
||||
pass
|
||||
|
||||
application = SimpleNamespace(workspace_service=_WorkspaceService())
|
||||
manager = PlatformManager(application)
|
||||
manager.adapter_dict = {'websocket': WebsocketAdapter}
|
||||
context = ExecutionContext(
|
||||
instance_uuid='instance',
|
||||
workspace_uuid=WORKSPACE_A,
|
||||
placement_generation=4,
|
||||
)
|
||||
|
||||
runtimes = await asyncio.gather(*(manager.get_websocket_proxy_bot(context) for _ in range(20)))
|
||||
|
||||
assert len(created_adapters) == 1
|
||||
assert len({id(runtime) for runtime in runtimes}) == 1
|
||||
assert manager.websocket_proxy_bots == {WORKSPACE_A: runtimes[0]}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_proxy_cache_evicts_oldest_idle_workspace():
|
||||
created_adapters = []
|
||||
|
||||
class WebsocketAdapter:
|
||||
def __init__(self, *_args, **_kwargs):
|
||||
self.kill = AsyncMock()
|
||||
self.inbound_listener_tasks = set()
|
||||
created_adapters.append(self)
|
||||
|
||||
def register_listener(self, *_args):
|
||||
pass
|
||||
|
||||
application = SimpleNamespace(
|
||||
workspace_service=_WorkspaceService(),
|
||||
instance_config=SimpleNamespace(
|
||||
data={
|
||||
'system': {
|
||||
'websocket_retention': {'max_workspace_proxies': 1},
|
||||
}
|
||||
}
|
||||
),
|
||||
)
|
||||
manager = PlatformManager(application)
|
||||
manager.adapter_dict = {'websocket': WebsocketAdapter}
|
||||
|
||||
await manager.get_websocket_proxy_bot(
|
||||
ExecutionContext(
|
||||
instance_uuid='instance',
|
||||
workspace_uuid=WORKSPACE_A,
|
||||
placement_generation=4,
|
||||
)
|
||||
)
|
||||
second = await manager.get_websocket_proxy_bot(
|
||||
ExecutionContext(
|
||||
instance_uuid='instance',
|
||||
workspace_uuid=WORKSPACE_B,
|
||||
placement_generation=4,
|
||||
)
|
||||
)
|
||||
|
||||
created_adapters[0].kill.assert_awaited_once_with()
|
||||
assert manager.websocket_proxy_bots == {WORKSPACE_B: second}
|
||||
assert WORKSPACE_A not in manager._proxy_last_accessed
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reload_stops_and_drops_existing_platform_runtimes():
|
||||
old_bot = SimpleNamespace(enable=True, shutdown=AsyncMock())
|
||||
old_proxy = SimpleNamespace(enable=True, shutdown=AsyncMock())
|
||||
persistence_mgr = SimpleNamespace(
|
||||
execute_async=AsyncMock(return_value=SimpleNamespace(all=lambda: [])),
|
||||
)
|
||||
application = SimpleNamespace(
|
||||
logger=SimpleNamespace(info=lambda *_args: None, warning=lambda *_args: None),
|
||||
persistence_mgr=persistence_mgr,
|
||||
workspace_service=SimpleNamespace(),
|
||||
)
|
||||
manager = PlatformManager(application)
|
||||
manager.bots = [old_bot]
|
||||
manager.websocket_proxy_bots = {WORKSPACE_A: old_proxy}
|
||||
manager._scope_generations = {('instance', WORKSPACE_A): 4}
|
||||
|
||||
await manager.load_bots_from_db()
|
||||
|
||||
old_bot.shutdown.assert_awaited_once_with()
|
||||
old_proxy.shutdown.assert_awaited_once_with()
|
||||
assert manager.bots == []
|
||||
assert manager.websocket_proxy_bots == {}
|
||||
assert manager._scope_generations == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cloud_startup_reuses_validated_platform_binding():
|
||||
class TenantUow:
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *_args):
|
||||
return False
|
||||
|
||||
class ProbeAdapter:
|
||||
def __init__(self, _config, _logger):
|
||||
self.listeners = []
|
||||
|
||||
def register_listener(self, event_type, listener):
|
||||
self.listeners.append((event_type, listener))
|
||||
|
||||
async def kill(self):
|
||||
return None
|
||||
|
||||
binding = WorkspaceExecutionBinding(
|
||||
instance_uuid='instance',
|
||||
workspace_uuid=WORKSPACE_A,
|
||||
placement_generation=4,
|
||||
write_fenced=False,
|
||||
state='active',
|
||||
)
|
||||
bot = Bot(
|
||||
uuid=BOT_A,
|
||||
workspace_uuid=WORKSPACE_A,
|
||||
name='Probe',
|
||||
description='',
|
||||
adapter='probe',
|
||||
adapter_config={},
|
||||
enable=False,
|
||||
pipeline_routing_rules=[],
|
||||
)
|
||||
workspace_service = SimpleNamespace(
|
||||
list_active_execution_bindings=AsyncMock(return_value=[binding]),
|
||||
get_execution_binding=AsyncMock(
|
||||
side_effect=AssertionError('startup platform loader repeated a validated binding lookup')
|
||||
),
|
||||
)
|
||||
application = SimpleNamespace(
|
||||
logger=SimpleNamespace(
|
||||
info=lambda *_args, **_kwargs: None,
|
||||
warning=lambda *_args, **_kwargs: None,
|
||||
error=lambda *_args, **_kwargs: None,
|
||||
),
|
||||
persistence_mgr=SimpleNamespace(
|
||||
mode=SimpleNamespace(value='cloud_runtime'),
|
||||
tenant_uow=lambda _workspace_uuid: TenantUow(),
|
||||
execute_async=AsyncMock(return_value=SimpleNamespace(all=lambda: [bot])),
|
||||
),
|
||||
workspace_service=workspace_service,
|
||||
)
|
||||
manager = PlatformManager(application)
|
||||
manager.adapter_dict = {'probe': ProbeAdapter}
|
||||
|
||||
await manager.load_bots_from_db()
|
||||
|
||||
assert len(manager.bots) == 1
|
||||
workspace_service.get_execution_binding.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runtime_bot_revalidates_its_generation_before_handling_events(manager):
|
||||
runtime_bot = manager.bots[0]
|
||||
|
||||
await runtime_bot.assert_execution_active()
|
||||
|
||||
runtime_bot.placement_generation = 5
|
||||
with pytest.raises(ValueError, match='stale'):
|
||||
await runtime_bot.assert_execution_active()
|
||||
|
||||
|
||||
def test_runtime_bot_rejects_workspace_mismatch():
|
||||
application = SimpleNamespace()
|
||||
entity = SimpleNamespace(
|
||||
uuid=BOT_A,
|
||||
workspace_uuid=WORKSPACE_A,
|
||||
name='Bot',
|
||||
enable=True,
|
||||
pipeline_routing_rules=[],
|
||||
use_pipeline_uuid=None,
|
||||
)
|
||||
with pytest.raises(WorkspaceRequiredError):
|
||||
RuntimeBot(
|
||||
ap=application,
|
||||
bot_entity=entity,
|
||||
adapter=SimpleNamespace(),
|
||||
logger=SimpleNamespace(),
|
||||
execution_context=_context(WORKSPACE_B, BOT_A),
|
||||
)
|
||||
|
||||
|
||||
class _ScopeOnlyPersistenceManager:
|
||||
mode = SimpleNamespace(value='cloud_runtime')
|
||||
|
||||
def __init__(self):
|
||||
self.active_workspace = None
|
||||
|
||||
@contextlib.asynccontextmanager
|
||||
async def tenant_scope(self, workspace_uuid: str):
|
||||
assert self.active_workspace is None
|
||||
self.active_workspace = workspace_uuid
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
self.active_workspace = None
|
||||
|
||||
def current_session(self):
|
||||
return None
|
||||
|
||||
|
||||
class _ListenerAdapter:
|
||||
def __init__(self):
|
||||
self.listeners = {}
|
||||
|
||||
def register_listener(self, event_type, listener):
|
||||
self.listeners[event_type] = listener
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_platform_callback_carries_scope_without_holding_database_session():
|
||||
persistence_mgr = _ScopeOnlyPersistenceManager()
|
||||
adapter = _ListenerAdapter()
|
||||
|
||||
async def push_person_message(*_args, **_kwargs):
|
||||
assert persistence_mgr.active_workspace == WORKSPACE_A
|
||||
assert persistence_mgr.current_session() is None
|
||||
return True
|
||||
|
||||
application = SimpleNamespace(
|
||||
persistence_mgr=persistence_mgr,
|
||||
workspace_service=_WorkspaceService(),
|
||||
webhook_pusher=SimpleNamespace(push_person_message=push_person_message),
|
||||
)
|
||||
entity = SimpleNamespace(
|
||||
uuid=BOT_A,
|
||||
workspace_uuid=WORKSPACE_A,
|
||||
name='Bot',
|
||||
enable=True,
|
||||
pipeline_routing_rules=[],
|
||||
use_pipeline_uuid=None,
|
||||
)
|
||||
logger = SimpleNamespace(info=AsyncMock(), error=AsyncMock())
|
||||
runtime = RuntimeBot(
|
||||
ap=application,
|
||||
bot_entity=entity,
|
||||
adapter=adapter,
|
||||
logger=logger,
|
||||
execution_context=_context(WORKSPACE_A, BOT_A),
|
||||
)
|
||||
await runtime.initialize()
|
||||
|
||||
listener = adapter.listeners[platform_events.FriendMessage]
|
||||
event = SimpleNamespace(message_chain=[], sender=SimpleNamespace(id='user'))
|
||||
await listener(event, adapter)
|
||||
|
||||
assert persistence_mgr.active_workspace is None
|
||||
logger.info.assert_awaited()
|
||||
@@ -5,6 +5,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from langbot.pkg.platform import botmgr as _botmgr # noqa: F401
|
||||
from langbot.pkg.platform.sources.dingtalk import (
|
||||
DingTalkAdapter,
|
||||
_dingtalk_card_markdown,
|
||||
@@ -17,6 +18,17 @@ from langbot.pkg.platform.sources.dingtalk import (
|
||||
)
|
||||
|
||||
|
||||
def test_dingtalk_auxiliary_tasks_are_bounded():
|
||||
adapter = DingTalkAdapter.model_construct()
|
||||
adapter._background_tasks = {MagicMock(done=MagicMock(return_value=False)) for _ in range(100)}
|
||||
|
||||
async def callback():
|
||||
raise AssertionError('rejected callback must not run')
|
||||
|
||||
assert adapter._start_background_task(callback()) is False
|
||||
assert len(adapter._background_tasks) == 100
|
||||
|
||||
|
||||
def test_dingtalk_select_component_params_expose_options():
|
||||
params = _dingtalk_form_component_params(
|
||||
{
|
||||
|
||||
@@ -0,0 +1,12 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from langbot.pkg.platform.sources import discord
|
||||
|
||||
|
||||
def test_discord_base64_decode_is_bounded(monkeypatch):
|
||||
monkeypatch.setattr(discord, '_MAX_DISCORD_MEDIA_BYTES', 4)
|
||||
|
||||
with pytest.raises(ValueError, match='exceeds'):
|
||||
discord._decode_discord_base64_limited('A' * 12)
|
||||
@@ -0,0 +1,150 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from langbot.pkg.api.http.context import ExecutionContext
|
||||
from langbot.pkg.platform.sources.http_bot import HttpBotAdapter
|
||||
from langbot.pkg.platform.sources import http_bot as http_bot_module
|
||||
|
||||
|
||||
def _session(key):
|
||||
session = SimpleNamespace()
|
||||
session._langbot_session_key = key
|
||||
return session
|
||||
|
||||
|
||||
def _adapter(app, execution_context) -> HttpBotAdapter:
|
||||
adapter = HttpBotAdapter.model_construct(
|
||||
config={'signature_required': False},
|
||||
logger=SimpleNamespace(execution_context=execution_context),
|
||||
bot_uuid='bot-a',
|
||||
outbound_states={},
|
||||
idempotency_cache={},
|
||||
sync_waiters={},
|
||||
inbound_tasks=set(),
|
||||
)
|
||||
object.__setattr__(adapter, 'ap', app)
|
||||
return adapter
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_bot_reset_removes_only_exact_execution_scope():
|
||||
context = ExecutionContext(
|
||||
instance_uuid='instance-a',
|
||||
workspace_uuid='workspace-a',
|
||||
placement_generation=3,
|
||||
bot_uuid='bot-a',
|
||||
)
|
||||
target_key = ('instance-a', 'workspace-a', 3, 'bot-a', 'person', 'shared-session')
|
||||
retained_keys = [
|
||||
('instance-b', 'workspace-a', 3, 'bot-a', 'person', 'shared-session'),
|
||||
('instance-a', 'workspace-b', 3, 'bot-a', 'person', 'shared-session'),
|
||||
('instance-a', 'workspace-a', 4, 'bot-a', 'person', 'shared-session'),
|
||||
('instance-a', 'workspace-a', 3, 'bot-b', 'person', 'shared-session'),
|
||||
('instance-a', 'workspace-a', 3, 'bot-a', 'group', 'shared-session'),
|
||||
('instance-a', 'workspace-a', 3, 'bot-a', 'person', 'other-session'),
|
||||
]
|
||||
sessions = [_session(target_key), *[_session(key) for key in retained_keys], SimpleNamespace()]
|
||||
app = SimpleNamespace(sess_mgr=SimpleNamespace(session_list=sessions))
|
||||
adapter = _adapter(app, context)
|
||||
|
||||
removed = await adapter._reset_session('person', 'shared-session')
|
||||
|
||||
assert removed is True
|
||||
assert [getattr(session, '_langbot_session_key', None) for session in app.sess_mgr.session_list] == [
|
||||
*retained_keys,
|
||||
None,
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_bot_reset_fails_closed_without_trusted_scope():
|
||||
app = SimpleNamespace(sess_mgr=SimpleNamespace(session_list=[]))
|
||||
adapter = _adapter(app, None)
|
||||
|
||||
with pytest.raises(RuntimeError, match='trusted execution scope'):
|
||||
await adapter._reset_session('person', 'shared-session')
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_bot_bounds_inbound_listener_tasks(monkeypatch):
|
||||
monkeypatch.setattr(http_bot_module, '_INBOUND_TASK_MAX', 1)
|
||||
adapter = _adapter(SimpleNamespace(), None)
|
||||
started = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
|
||||
async def blocking_listener():
|
||||
started.set()
|
||||
await release.wait()
|
||||
|
||||
first = adapter._start_inbound_task(blocking_listener())
|
||||
await started.wait()
|
||||
rejected = adapter._start_inbound_task(blocking_listener())
|
||||
|
||||
assert first is not None
|
||||
assert rejected is None
|
||||
assert len(adapter.inbound_tasks) == 1
|
||||
|
||||
release.set()
|
||||
await first
|
||||
await asyncio.sleep(0)
|
||||
assert adapter.inbound_tasks == set()
|
||||
|
||||
|
||||
def test_http_bot_outbound_state_has_a_hard_capacity(monkeypatch):
|
||||
monkeypatch.setattr(http_bot_module, '_OUTBOUND_STATE_MAX', 2)
|
||||
monkeypatch.setattr(http_bot_module, '_OUTBOUND_PRUNE_SCAN_MAX', 2)
|
||||
adapter = _adapter(SimpleNamespace(), None)
|
||||
first = adapter._outbound_state('first')
|
||||
second = adapter._outbound_state('second')
|
||||
first.queue.put_nowait({})
|
||||
second.queue.put_nowait({})
|
||||
|
||||
with pytest.raises(RuntimeError, match='outbound session capacity reached'):
|
||||
adapter._next_sequence('third', is_final=True)
|
||||
|
||||
assert len(adapter.outbound_states) == 2
|
||||
assert adapter._next_sequence('first', is_final=True) == 1
|
||||
|
||||
|
||||
def test_http_bot_outbound_state_pruning_is_bounded_and_reclaims_stale(monkeypatch):
|
||||
monkeypatch.setattr(http_bot_module, '_OUTBOUND_STATE_MAX', 2)
|
||||
monkeypatch.setattr(http_bot_module, '_OUTBOUND_PRUNE_SCAN_MAX', 1)
|
||||
monkeypatch.setattr(http_bot_module, '_OUTBOUND_IDLE_SECONDS', 10)
|
||||
adapter = _adapter(SimpleNamespace(), None)
|
||||
stale = adapter._outbound_state('stale')
|
||||
stale.last_active = time.monotonic() - 11
|
||||
adapter._outbound_state('active')
|
||||
|
||||
assert adapter._next_sequence('replacement', is_final=True) == 1
|
||||
assert set(adapter.outbound_states) == {'active', 'replacement'}
|
||||
|
||||
|
||||
def test_http_bot_idempotency_cache_has_a_hard_capacity(monkeypatch):
|
||||
monkeypatch.setattr(http_bot_module, '_IDEMPOTENCY_MAX', 2)
|
||||
monkeypatch.setattr(http_bot_module, '_IDEMPOTENCY_PRUNE_SCAN_MAX', 1)
|
||||
adapter = _adapter(SimpleNamespace(), None)
|
||||
|
||||
assert adapter._reserve_idempotency_key('first') == 'accepted'
|
||||
assert adapter._reserve_idempotency_key('second') == 'accepted'
|
||||
assert adapter._reserve_idempotency_key('third') == 'overloaded'
|
||||
assert len(adapter.idempotency_cache) == 2
|
||||
assert adapter._reserve_idempotency_key('first') == 'duplicate'
|
||||
|
||||
|
||||
def test_http_bot_idempotency_cache_reclaims_expired_oldest(monkeypatch):
|
||||
monkeypatch.setattr(http_bot_module, '_IDEMPOTENCY_MAX', 2)
|
||||
monkeypatch.setattr(http_bot_module, '_IDEMPOTENCY_PRUNE_SCAN_MAX', 1)
|
||||
monkeypatch.setattr(http_bot_module, '_IDEMPOTENCY_TTL', 10)
|
||||
adapter = _adapter(SimpleNamespace(), None)
|
||||
adapter.idempotency_cache = {
|
||||
'expired': time.monotonic() - 11,
|
||||
'active': time.monotonic(),
|
||||
}
|
||||
|
||||
assert adapter._reserve_idempotency_key('replacement') == 'accepted'
|
||||
assert set(adapter.idempotency_cache) == {'active', 'replacement'}
|
||||
@@ -0,0 +1,24 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import zlib
|
||||
|
||||
import pytest
|
||||
|
||||
from langbot.pkg.platform.sources import kook
|
||||
|
||||
|
||||
def test_kook_gateway_decoder_accepts_raw_and_compressed_json():
|
||||
payload = {'s': 1, 'd': {'session_id': 'session-a'}}
|
||||
encoded = json.dumps(payload).encode()
|
||||
|
||||
assert kook._decode_gateway_message(encoded) == payload
|
||||
assert kook._decode_gateway_message(zlib.compress(encoded)) == payload
|
||||
|
||||
|
||||
def test_kook_gateway_decoder_rejects_decompression_bomb(monkeypatch):
|
||||
monkeypatch.setattr(kook, '_KOOK_MAX_GATEWAY_MESSAGE_BYTES', 1024)
|
||||
compressed = zlib.compress(b'x' * 1025)
|
||||
|
||||
with pytest.raises(ValueError, match='decompressed size limit'):
|
||||
kook._decode_gateway_message(compressed)
|
||||
@@ -1,7 +1,13 @@
|
||||
"""Tests for Lark adapter helper behavior."""
|
||||
|
||||
import threading
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from langbot.pkg.platform.sources.lark import (
|
||||
LarkAdapter,
|
||||
_decode_lark_base64_limited,
|
||||
_lark_clean_form_content,
|
||||
_lark_completed_input_lines,
|
||||
_lark_current_input_defs,
|
||||
@@ -11,6 +17,27 @@ from langbot.pkg.platform.sources.lark import (
|
||||
)
|
||||
|
||||
|
||||
def test_lark_base64_decode_is_bounded(monkeypatch):
|
||||
import langbot.pkg.platform.sources.lark as lark_module
|
||||
|
||||
monkeypatch.setattr(lark_module, '_MAX_LARK_MEDIA_BYTES', 4)
|
||||
|
||||
with pytest.raises(ValueError, match='exceeds'):
|
||||
_decode_lark_base64_limited('A' * 12)
|
||||
|
||||
|
||||
def test_lark_threadsafe_callbacks_are_bounded():
|
||||
adapter = LarkAdapter.model_construct()
|
||||
adapter.threadsafe_event_lock = threading.Lock()
|
||||
adapter.threadsafe_event_futures = {MagicMock(done=MagicMock(return_value=False)) for _ in range(100)}
|
||||
|
||||
async def callback():
|
||||
raise AssertionError('rejected callback must not run')
|
||||
|
||||
assert adapter._schedule_threadsafe_event(callback()) is None
|
||||
assert len(adapter.threadsafe_event_futures) == 100
|
||||
|
||||
|
||||
def test_lark_current_input_defs_only_returns_active_stage():
|
||||
input_defs = [
|
||||
{'output_variable_name': 'us_input', 'type': 'paragraph'},
|
||||
|
||||
@@ -0,0 +1,30 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from langbot.pkg.platform import botmgr as _botmgr # noqa: F401
|
||||
from langbot.pkg.platform.sources import line
|
||||
|
||||
|
||||
def test_line_media_content_accepts_limit_boundary(monkeypatch) -> None:
|
||||
monkeypatch.setattr(line, 'MAX_LINE_MEDIA_BYTES', 4)
|
||||
content = b'1234'
|
||||
|
||||
assert line._validate_line_media_content(content) is content
|
||||
|
||||
|
||||
def test_line_media_content_rejects_oversized_payload(monkeypatch) -> None:
|
||||
monkeypatch.setattr(line, 'MAX_LINE_MEDIA_BYTES', 4)
|
||||
|
||||
with pytest.raises(ValueError, match='LINE media exceeds'):
|
||||
line._validate_line_media_content(b'12345')
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_line_kill_closes_api_client() -> None:
|
||||
api_client = MagicMock()
|
||||
adapter = line.LINEAdapter.model_construct(api_client=api_client)
|
||||
|
||||
assert await adapter.kill() is True
|
||||
api_client.close.assert_called_once_with()
|
||||
@@ -0,0 +1,21 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from langbot.pkg.platform.sources import matrix
|
||||
|
||||
|
||||
def test_matrix_base64_decode_is_bounded(monkeypatch):
|
||||
monkeypatch.setattr(matrix, '_MAX_MATRIX_MEDIA_BYTES', 4)
|
||||
|
||||
with pytest.raises(ValueError, match='exceeds'):
|
||||
matrix._decode_matrix_base64_limited('A' * 12)
|
||||
|
||||
|
||||
def test_matrix_local_file_read_is_bounded(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(matrix, '_MAX_MATRIX_MEDIA_BYTES', 4)
|
||||
path = tmp_path / 'large.bin'
|
||||
path.write_bytes(b'12345')
|
||||
|
||||
with pytest.raises(ValueError, match='exceeds'):
|
||||
matrix._read_matrix_file_limited(str(path))
|
||||
@@ -0,0 +1,31 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from langbot.libs.openclaw_weixin_api.client import (
|
||||
MAX_CDN_MEDIA_BYTES,
|
||||
OpenClawWeixinClient,
|
||||
_decrypt_cdn_payload,
|
||||
_encrypt_cdn_payload,
|
||||
)
|
||||
from langbot.libs.openclaw_weixin_api.types import ApiError
|
||||
|
||||
|
||||
def test_cdn_crypto_helpers_round_trip():
|
||||
original = b'tenant-media' * 128
|
||||
|
||||
aes_key_hex, _encoded_key, encrypted, _raw_md5 = _encrypt_cdn_payload(original)
|
||||
|
||||
assert _decrypt_cdn_payload(encrypted, bytes.fromhex(aes_key_hex)) == original
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_media_rejects_oversized_input_before_network_access():
|
||||
client = OpenClawWeixinClient('https://example.invalid', 'token')
|
||||
|
||||
with pytest.raises(ApiError, match='exceeds the size limit'):
|
||||
await client.upload_media(
|
||||
b'x' * (MAX_CDN_MEDIA_BYTES + 1),
|
||||
'recipient',
|
||||
3,
|
||||
)
|
||||
@@ -0,0 +1,95 @@
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
import langbot_plugin.api.entities.builtin.platform.message as platform_message
|
||||
|
||||
from langbot.pkg.api.http.context import ExecutionContext
|
||||
from langbot.pkg.platform.sources import openclaw_weixin
|
||||
from langbot.pkg.platform.sources.openclaw_weixin import OpenClawWeixinAdapter
|
||||
|
||||
|
||||
def make_adapter(*, execution_context: ExecutionContext | None):
|
||||
app = SimpleNamespace(
|
||||
persistence_mgr=SimpleNamespace(execute_async=AsyncMock()),
|
||||
workspace_service=SimpleNamespace(
|
||||
get_execution_binding=AsyncMock(
|
||||
return_value=SimpleNamespace(
|
||||
instance_uuid='instance-a',
|
||||
workspace_uuid='workspace-a',
|
||||
placement_generation=1,
|
||||
)
|
||||
)
|
||||
),
|
||||
)
|
||||
logger = SimpleNamespace(
|
||||
ap=app,
|
||||
execution_context=execution_context,
|
||||
warning=AsyncMock(),
|
||||
)
|
||||
adapter = OpenClawWeixinAdapter.model_construct(
|
||||
config={'token': 'refreshed-token'},
|
||||
logger=logger,
|
||||
client=Mock(),
|
||||
bot_account_id='',
|
||||
listeners={},
|
||||
name='openclaw-weixin',
|
||||
)
|
||||
adapter._bot_uuid = 'shared-bot-uuid'
|
||||
return adapter, app, logger
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_persist_config_scopes_duplicate_bot_uuid_to_workspace():
|
||||
adapter, app, _ = make_adapter(
|
||||
execution_context=ExecutionContext(
|
||||
instance_uuid='instance-a',
|
||||
workspace_uuid='workspace-a',
|
||||
placement_generation=1,
|
||||
bot_uuid='shared-bot-uuid',
|
||||
)
|
||||
)
|
||||
|
||||
await adapter._persist_config()
|
||||
|
||||
app.workspace_service.get_execution_binding.assert_awaited_once_with(
|
||||
'workspace-a',
|
||||
expected_generation=1,
|
||||
)
|
||||
statement = app.persistence_mgr.execute_async.await_args.args[0]
|
||||
params = statement.compile().params
|
||||
assert 'workspace-a' in params.values()
|
||||
assert 'shared-bot-uuid' in params.values()
|
||||
assert {'workspace_uuid', 'uuid'} <= {comparison.left.name for comparison in statement._where_criteria}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
'execution_context',
|
||||
[
|
||||
None,
|
||||
ExecutionContext(
|
||||
instance_uuid='instance-a',
|
||||
workspace_uuid='workspace-a',
|
||||
placement_generation=1,
|
||||
bot_uuid='another-bot-uuid',
|
||||
),
|
||||
],
|
||||
ids=['missing-context', 'mismatched-bot'],
|
||||
)
|
||||
async def test_persist_config_fails_closed_without_matching_execution_context(execution_context):
|
||||
adapter, app, logger = make_adapter(execution_context=execution_context)
|
||||
|
||||
await adapter._persist_config()
|
||||
|
||||
app.persistence_mgr.execute_async.assert_not_awaited()
|
||||
logger.warning.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_component_base64_decode_is_bounded(monkeypatch):
|
||||
monkeypatch.setattr(openclaw_weixin, '_MAX_OPENCLAW_COMPONENT_BYTES', 4)
|
||||
component = platform_message.File(base64='MTIzNDU=')
|
||||
|
||||
with pytest.raises(ValueError, match='exceeds'):
|
||||
await OpenClawWeixinAdapter._get_component_bytes(component)
|
||||
@@ -10,6 +10,7 @@ import langbot_plugin.api.entities.builtin.platform.message as platform_message
|
||||
|
||||
from langbot.libs.qq_official_api.api import (
|
||||
QQ_SELECT_ACTION_PREFIX,
|
||||
QQOfficialClient,
|
||||
build_keyboard_from_select_field,
|
||||
get_select_field_options,
|
||||
resolve_select_button_action,
|
||||
@@ -49,6 +50,28 @@ def test_qq_select_button_resolves_field_and_value():
|
||||
assert resolve_select_button_action(form_data, f'{QQ_SELECT_ACTION_PREFIX}99') is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_qq_seed_rejects_empty_secret_without_spinning():
|
||||
client = QQOfficialClient('', 'token', 'app-id', AsyncMock())
|
||||
|
||||
with pytest.raises(ValueError, match='must not be empty'):
|
||||
await asyncio.wait_for(client.repeat_seed(''), timeout=0.1)
|
||||
|
||||
|
||||
def test_qq_auxiliary_tasks_are_bounded():
|
||||
import langbot.pkg.core.app # noqa: F401
|
||||
from langbot.pkg.platform.sources.qqofficial import QQOfficialAdapter
|
||||
|
||||
adapter = QQOfficialAdapter.model_construct()
|
||||
adapter._background_tasks = {MagicMock(done=MagicMock(return_value=False)) for _ in range(100)}
|
||||
|
||||
async def callback():
|
||||
raise AssertionError('rejected callback must not run')
|
||||
|
||||
assert adapter._start_background_task(callback()) is False
|
||||
assert len(adapter._background_tasks) == 100
|
||||
|
||||
|
||||
def test_qq_select_keyboard_fits_twenty_five_options():
|
||||
form_data = _select_form_data()
|
||||
form_data['input_defs'][0]['option_source']['value'] = [f'Option {idx}' for idx in range(25)]
|
||||
|
||||
@@ -278,3 +278,19 @@ class TestResolvePipelineUuid:
|
||||
uuid, routed = bot.resolve_pipeline_uuid('person', '123', 'normal message')
|
||||
assert uuid == 'default-uuid'
|
||||
assert routed is False
|
||||
|
||||
def test_websocket_task_override_does_not_mutate_bot_default(self):
|
||||
bot = self._make_bot('default-uuid', [])
|
||||
adapter = Mock()
|
||||
adapter.get_pipeline_uuid_override.return_value = 'connection-pipeline'
|
||||
|
||||
pipeline_uuid, routed = bot.resolve_event_pipeline_uuid(
|
||||
adapter,
|
||||
'person',
|
||||
'launcher',
|
||||
'hello',
|
||||
)
|
||||
|
||||
assert pipeline_uuid == 'connection-pipeline'
|
||||
assert routed is False
|
||||
assert bot.bot_entity.use_pipeline_uuid == 'default-uuid'
|
||||
|
||||
@@ -11,11 +11,21 @@ import langbot_plugin.api.entities.builtin.platform.events as platform_events
|
||||
import langbot_plugin.api.entities.builtin.platform.message as platform_message
|
||||
from langbot.pkg.platform.sources.telegram import (
|
||||
TelegramAdapter,
|
||||
_decode_telegram_base64_limited,
|
||||
_telegram_form_action_from_callback,
|
||||
_telegram_select_field_options,
|
||||
)
|
||||
|
||||
|
||||
def test_telegram_base64_decode_is_bounded(monkeypatch):
|
||||
import langbot.pkg.platform.sources.telegram as telegram_module
|
||||
|
||||
monkeypatch.setattr(telegram_module, '_MAX_TELEGRAM_MEDIA_BYTES', 4)
|
||||
|
||||
with pytest.raises(ValueError, match='exceeds'):
|
||||
_decode_telegram_base64_limited('A' * 12)
|
||||
|
||||
|
||||
def _select_form_data() -> dict:
|
||||
return {
|
||||
'_current_input_field': 'choice',
|
||||
@@ -88,6 +98,18 @@ def test_telegram_form_callback_cache_preserves_pipeline_uuid():
|
||||
)
|
||||
|
||||
|
||||
def test_telegram_form_callback_cache_is_bounded():
|
||||
adapter = TelegramAdapter.model_construct()
|
||||
adapter._form_action_titles = {}
|
||||
|
||||
adapter._cache_form_action_titles(
|
||||
{f'callback-{index}': str(index) for index in range(5000)},
|
||||
now=100.0,
|
||||
)
|
||||
|
||||
assert len(adapter._form_action_titles) == adapter._MAX_FORM_ACTION_TITLES
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_telegram_select_field_sends_two_column_inline_keyboard():
|
||||
bot = MagicMock()
|
||||
|
||||
@@ -0,0 +1,99 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
from langbot.pkg.platform.webhook_pusher import WebhookPusher
|
||||
|
||||
|
||||
pytestmark = pytest.mark.asyncio
|
||||
|
||||
|
||||
def _application(max_inflight_requests: object) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
instance_config=SimpleNamespace(
|
||||
data={
|
||||
'webhooks': {
|
||||
'max_inflight_requests': max_inflight_requests,
|
||||
}
|
||||
}
|
||||
),
|
||||
logger=logging.getLogger(__name__),
|
||||
)
|
||||
|
||||
|
||||
async def test_delivery_admission_never_queues_above_instance_limit():
|
||||
pusher = WebhookPusher(_application(2))
|
||||
release = asyncio.Event()
|
||||
both_started = asyncio.Event()
|
||||
calls = 0
|
||||
active = 0
|
||||
peak_active = 0
|
||||
|
||||
async def fake_push(url: str, payload: dict) -> dict:
|
||||
nonlocal calls, active, peak_active
|
||||
calls += 1
|
||||
active += 1
|
||||
peak_active = max(peak_active, active)
|
||||
if active == 2:
|
||||
both_started.set()
|
||||
try:
|
||||
await release.wait()
|
||||
return {'url': url}
|
||||
finally:
|
||||
active -= 1
|
||||
|
||||
pusher._push_to_webhook = fake_push
|
||||
webhooks = [{'url': f'https://example.invalid/{index}'} for index in range(5)]
|
||||
|
||||
first_delivery = asyncio.create_task(pusher._push_to_webhooks(webhooks, {}))
|
||||
await asyncio.wait_for(both_started.wait(), timeout=1)
|
||||
second_results = await pusher._push_to_webhooks(webhooks, {})
|
||||
release.set()
|
||||
first_results = await first_delivery
|
||||
|
||||
assert len(first_results) == 2
|
||||
assert second_results == []
|
||||
assert calls == 2
|
||||
assert peak_active == 2
|
||||
assert pusher._inflight_requests == 0
|
||||
|
||||
|
||||
async def test_cancelled_delivery_reaps_children_and_releases_slots():
|
||||
pusher = WebhookPusher(_application(1))
|
||||
started = asyncio.Event()
|
||||
never = asyncio.Event()
|
||||
|
||||
async def blocking_push(url: str, payload: dict) -> dict:
|
||||
started.set()
|
||||
await never.wait()
|
||||
return {}
|
||||
|
||||
pusher._push_to_webhook = blocking_push
|
||||
delivery = asyncio.create_task(
|
||||
pusher._push_to_webhooks([{'url': 'https://example.invalid'}], {}),
|
||||
)
|
||||
await asyncio.wait_for(started.wait(), timeout=1)
|
||||
|
||||
delivery.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await delivery
|
||||
|
||||
assert pusher._inflight_requests == 0
|
||||
pusher._push_to_webhook = AsyncMock(return_value={})
|
||||
assert await pusher._push_to_webhooks([{'url': 'https://example.invalid'}], {}) == [{}]
|
||||
|
||||
|
||||
async def test_max_inflight_requests_clamps_config():
|
||||
pusher = WebhookPusher(_application(999999))
|
||||
assert pusher._max_inflight_requests() == 128
|
||||
|
||||
pusher.ap.instance_config.data['webhooks']['max_inflight_requests'] = 0
|
||||
assert pusher._max_inflight_requests() == 1
|
||||
|
||||
pusher.ap.instance_config.data['webhooks']['max_inflight_requests'] = 'invalid'
|
||||
assert pusher._max_inflight_requests() == 16
|
||||
@@ -4,89 +4,116 @@ The web debug client uploads Image / Voice / File components carrying a storage
|
||||
key in ``path``. This helper resolves each to a base64 data URI (so multimodal
|
||||
LLM input and the Box sandbox inbox have usable bytes), then deletes the
|
||||
consumed storage object and clears ``path``. Covers mimetype selection per
|
||||
type and graceful error handling.
|
||||
type and fail-closed error handling.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
|
||||
from langbot.pkg.api.http.context import ExecutionContext
|
||||
from langbot.pkg.platform.sources.websocket_adapter import WebSocketAdapter
|
||||
|
||||
|
||||
_CONTEXT = ExecutionContext(
|
||||
instance_uuid='instance-a',
|
||||
workspace_uuid='workspace-a',
|
||||
placement_generation=1,
|
||||
pipeline_uuid='pipeline-a',
|
||||
)
|
||||
_UPLOAD_PREFIX = 'v1/instance-a/workspace-a/1/upload_image/'
|
||||
|
||||
|
||||
def _make_connection():
|
||||
return SimpleNamespace(execution_context=_CONTEXT)
|
||||
|
||||
|
||||
def _make_adapter(load_return=b'hello', load_side_effect=None):
|
||||
provider = Mock()
|
||||
provider.load = AsyncMock(return_value=load_return, side_effect=load_side_effect)
|
||||
provider.delete = AsyncMock()
|
||||
storage_mgr = Mock()
|
||||
storage_mgr.storage_provider = provider
|
||||
storage_mgr.load_scoped_object_key = AsyncMock(return_value=load_return, side_effect=load_side_effect)
|
||||
storage_mgr.scoped_prefix.return_value = _UPLOAD_PREFIX
|
||||
storage_mgr.is_scoped_object_key.return_value = True
|
||||
storage_mgr.delete_scoped_object_key = AsyncMock()
|
||||
ap = Mock()
|
||||
ap.storage_mgr.storage_provider = provider
|
||||
ap.storage_mgr = storage_mgr
|
||||
logger = Mock()
|
||||
logger.error = AsyncMock()
|
||||
logger.warning = AsyncMock()
|
||||
# WebSocketAdapter is a pydantic model; bypass full __init__/validation.
|
||||
adapter = WebSocketAdapter.model_construct(ap=ap, logger=logger)
|
||||
return adapter, provider
|
||||
return adapter, storage_mgr, provider
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_image_jpeg_mimetype_and_cleanup():
|
||||
adapter, provider = _make_adapter(load_return=b'\xff\xd8\xff')
|
||||
chain = [{'type': 'Image', 'path': 'storage://abc/photo.jpg'}]
|
||||
adapter, storage_mgr, _ = _make_adapter(load_return=b'\xff\xd8\xff')
|
||||
path = f'{_UPLOAD_PREFIX}photo.jpg'
|
||||
chain = [{'type': 'Image', 'path': path}]
|
||||
|
||||
await adapter._process_image_components(chain)
|
||||
await adapter._process_image_components(_make_connection(), chain)
|
||||
|
||||
expected_b64 = base64.b64encode(b'\xff\xd8\xff').decode('utf-8')
|
||||
assert chain[0]['base64'] == f'data:image/jpeg;base64,{expected_b64}'
|
||||
assert chain[0]['path'] == '' # consumed
|
||||
provider.delete.assert_awaited_once_with('storage://abc/photo.jpg')
|
||||
storage_mgr.delete_scoped_object_key.assert_awaited_once_with(
|
||||
_CONTEXT,
|
||||
path,
|
||||
expected_owner_type='upload_image',
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_image_defaults_to_png():
|
||||
adapter, _ = _make_adapter()
|
||||
chain = [{'type': 'Image', 'path': 'storage://abc/blob'}]
|
||||
await adapter._process_image_components(chain)
|
||||
adapter, _, _ = _make_adapter()
|
||||
chain = [{'type': 'Image', 'path': f'{_UPLOAD_PREFIX}blob'}]
|
||||
await adapter._process_image_components(_make_connection(), chain)
|
||||
assert chain[0]['base64'].startswith('data:image/png;base64,')
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_voice_uses_guessed_or_wav_mimetype():
|
||||
adapter, _ = _make_adapter()
|
||||
chain = [{'type': 'Voice', 'path': 'storage://abc/clip.wav'}]
|
||||
await adapter._process_image_components(chain)
|
||||
adapter, _, _ = _make_adapter()
|
||||
chain = [{'type': 'Voice', 'path': f'{_UPLOAD_PREFIX}clip.wav'}]
|
||||
await adapter._process_image_components(_make_connection(), chain)
|
||||
assert chain[0]['base64'].startswith('data:audio/')
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_uses_octet_stream_fallback():
|
||||
adapter, _ = _make_adapter()
|
||||
chain = [{'type': 'File', 'path': 'storage://abc/unknownblob'}]
|
||||
await adapter._process_image_components(chain)
|
||||
adapter, _, _ = _make_adapter()
|
||||
chain = [{'type': 'File', 'path': f'{_UPLOAD_PREFIX}unknownblob'}]
|
||||
await adapter._process_image_components(_make_connection(), chain)
|
||||
assert chain[0]['base64'].startswith('data:application/octet-stream;base64,')
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_skips_components_without_path_or_unknown_type():
|
||||
adapter, provider = _make_adapter()
|
||||
adapter, storage_mgr, provider = _make_adapter()
|
||||
chain = [
|
||||
{'type': 'Image', 'path': ''}, # no path
|
||||
{'type': 'Plain', 'path': 'storage://abc/x'}, # not a file component
|
||||
{'type': 'At', 'target': '123'}, # no path key at all
|
||||
]
|
||||
await adapter._process_image_components(chain)
|
||||
await adapter._process_image_components(_make_connection(), chain)
|
||||
provider.load.assert_not_awaited()
|
||||
storage_mgr.load_scoped_object_key.assert_not_awaited()
|
||||
assert 'base64' not in chain[0]
|
||||
assert 'base64' not in chain[1]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_failure_is_logged_not_raised():
|
||||
adapter, _ = _make_adapter(load_side_effect=RuntimeError('storage down'))
|
||||
chain = [{'type': 'File', 'path': 'storage://abc/doc.pdf'}]
|
||||
async def test_load_failure_is_logged_and_aborts_processing():
|
||||
adapter, _, _ = _make_adapter(load_side_effect=RuntimeError('storage down'))
|
||||
chain = [{'type': 'File', 'path': f'{_UPLOAD_PREFIX}doc.pdf'}]
|
||||
|
||||
# must not raise
|
||||
await adapter._process_image_components(chain)
|
||||
with pytest.raises(RuntimeError, match='storage down'):
|
||||
await adapter._process_image_components(_make_connection(), chain)
|
||||
assert 'base64' not in chain[0]
|
||||
adapter.logger.error.assert_awaited_once()
|
||||
|
||||
@@ -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,114 @@ 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_connection_admission_is_bounded_globally_and_per_workspace():
|
||||
manager = WebSocketConnectionManager()
|
||||
await manager.add_connection(
|
||||
websocket=Mock(),
|
||||
scope=SCOPE_A,
|
||||
pipeline_uuid='pipeline-1',
|
||||
session_type='person',
|
||||
max_connections=2,
|
||||
max_connections_per_workspace=1,
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match='Workspace WebSocket'):
|
||||
await manager.add_connection(
|
||||
websocket=Mock(),
|
||||
scope=SCOPE_A,
|
||||
pipeline_uuid='pipeline-2',
|
||||
session_type='person',
|
||||
max_connections=2,
|
||||
max_connections_per_workspace=1,
|
||||
)
|
||||
|
||||
await manager.add_connection(
|
||||
websocket=Mock(),
|
||||
scope=SCOPE_B,
|
||||
pipeline_uuid='pipeline-1',
|
||||
session_type='person',
|
||||
max_connections=2,
|
||||
max_connections_per_workspace=1,
|
||||
)
|
||||
with pytest.raises(RuntimeError, match='WebSocket connection capacity'):
|
||||
await manager.add_connection(
|
||||
websocket=Mock(),
|
||||
scope=WebSocketScope('instance-a', 'workspace-c', 1),
|
||||
pipeline_uuid='pipeline-1',
|
||||
session_type='person',
|
||||
max_connections=2,
|
||||
max_connections_per_workspace=1,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_close_scope_closes_and_removes_only_matching_connections():
|
||||
manager = WebSocketConnectionManager()
|
||||
websocket_a = Mock(close=AsyncMock())
|
||||
connection_a = await manager.add_connection(
|
||||
websocket=websocket_a,
|
||||
scope=SCOPE_A,
|
||||
pipeline_uuid='pipeline-1',
|
||||
session_type='person',
|
||||
)
|
||||
connection_b = await manager.add_connection(
|
||||
websocket=Mock(close=AsyncMock()),
|
||||
scope=SCOPE_B,
|
||||
pipeline_uuid='pipeline-1',
|
||||
session_type='person',
|
||||
)
|
||||
|
||||
await manager.close_scope(SCOPE_A)
|
||||
|
||||
websocket_a.close.assert_awaited_once()
|
||||
assert await manager.get_connection(connection_a.connection_id, scope=SCOPE_A) is None
|
||||
assert await manager.get_connection(connection_b.connection_id, scope=SCOPE_B) is connection_b
|
||||
|
||||
|
||||
@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 +210,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 +237,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 +258,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 +307,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 +315,75 @@ 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.load_scoped_object_key = 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.load_scoped_object_key.assert_awaited_once_with(
|
||||
connection.execution_context,
|
||||
'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 +395,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()
|
||||
|
||||
@@ -0,0 +1,39 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from langbot.libs.wechatpad_api.api import downloadpai
|
||||
from langbot.libs.wechatpad_api.util import http_util
|
||||
|
||||
|
||||
class _Response:
|
||||
headers = {}
|
||||
|
||||
def __init__(self, chunks: list[bytes]):
|
||||
self._chunks = chunks
|
||||
|
||||
def iter_content(self, chunk_size=None):
|
||||
del chunk_size
|
||||
yield from self._chunks
|
||||
|
||||
|
||||
def test_wechatpad_response_reader_is_bounded(monkeypatch):
|
||||
monkeypatch.setattr(http_util, '_MAX_WECHATPAD_RESPONSE_BYTES', 4)
|
||||
|
||||
with pytest.raises(RuntimeError, match='exceeds the runtime limit'):
|
||||
http_util._read_requests_response_limited(_Response([b'1234', b'5']))
|
||||
|
||||
|
||||
def test_wechatpad_response_reader_requires_json_object():
|
||||
with pytest.raises(RuntimeError, match='non-object'):
|
||||
http_util._read_requests_response_limited(_Response([b'[]']))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_wechatpad_media_reader_is_bounded(monkeypatch):
|
||||
monkeypatch.setattr(downloadpai, '_MAX_WECHATPAD_MEDIA_BYTES', 4)
|
||||
response = httpx.Response(200, content=b'oversized')
|
||||
|
||||
with pytest.raises(RuntimeError, match='exceeds'):
|
||||
await downloadpai._read_media_limited(response)
|
||||
@@ -0,0 +1,31 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from langbot.libs.wecom_api.api import (
|
||||
_EXTENDED_HTTP_TIMEOUT_SECONDS,
|
||||
_decode_media_base64_limited,
|
||||
WecomClient,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_wecom_extended_client_timeout_is_still_bounded() -> None:
|
||||
client = object.__new__(WecomClient)
|
||||
client._http_clients = {}
|
||||
|
||||
try:
|
||||
async with client._http_client_context(unbounded_timeout=True) as http_client:
|
||||
assert http_client.timeout.read == _EXTENDED_HTTP_TIMEOUT_SECONDS
|
||||
finally:
|
||||
await client.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_wecom_base64_decode_is_bounded(monkeypatch) -> None:
|
||||
import langbot.libs.wecom_api.api as wecom_api
|
||||
|
||||
monkeypatch.setattr(wecom_api, '_MAX_MEDIA_BYTES', 4)
|
||||
|
||||
with pytest.raises(ValueError, match='exceeds'):
|
||||
await _decode_media_base64_limited('MTIzNDU=')
|
||||
@@ -1,5 +1,6 @@
|
||||
import sys
|
||||
import types
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
|
||||
@@ -24,6 +25,25 @@ from langbot.libs.wecom_ai_bot_api.api import ( # noqa: E402
|
||||
from langbot.libs.wecom_ai_bot_api.ws_client import WecomBotWsClient # noqa: E402
|
||||
|
||||
|
||||
def test_ws_callback_tasks_are_bounded():
|
||||
client = WecomBotWsClient('bot-id', 'secret', object())
|
||||
client._callback_tasks = {Mock(done=Mock(return_value=False)) for _ in range(100)}
|
||||
|
||||
async def callback():
|
||||
raise AssertionError('rejected callback must not run')
|
||||
|
||||
assert client._start_callback_task(callback()) is False
|
||||
assert len(client._callback_tasks) == 100
|
||||
|
||||
|
||||
def test_webhook_dispatch_tasks_are_bounded():
|
||||
client = WecomBotClient('', '', '', object(), unified_mode=True)
|
||||
client._dispatch_tasks = {Mock(done=Mock(return_value=False)) for _ in range(100)}
|
||||
|
||||
assert client._start_dispatch_task(Mock()) is False
|
||||
assert len(client._dispatch_tasks) == 100
|
||||
|
||||
|
||||
def test_extract_template_card_action_supports_nested_button_key():
|
||||
task_id, event_key, card_type = extract_template_card_action(
|
||||
{
|
||||
@@ -287,16 +307,9 @@ async def test_webhook_stream_queues_cumulative_snapshots_for_followups():
|
||||
assert await client.push_stream_chunk('msg-1', '你好', is_final=False)
|
||||
assert await client.push_stream_chunk('msg-1', '你好', is_final=True)
|
||||
|
||||
chunks = [
|
||||
await client.stream_sessions.consume(session.stream_id),
|
||||
await client.stream_sessions.consume(session.stream_id),
|
||||
await client.stream_sessions.consume(session.stream_id),
|
||||
]
|
||||
assert [(chunk.content, chunk.is_final) for chunk in chunks] == [
|
||||
('你', False),
|
||||
('你好', False),
|
||||
('你好', True),
|
||||
]
|
||||
assert session.queue.qsize() == 1
|
||||
chunk = await client.stream_sessions.consume(session.stream_id)
|
||||
assert (chunk.content, chunk.is_final) == ('你好', True)
|
||||
|
||||
|
||||
def test_human_input_payload_keeps_action_select_stage_as_buttons():
|
||||
|
||||
Reference in New Issue
Block a user