xinyin025/LangBot
1
0
Fork 0
mirror of https://github.com/langbot-app/LangBot.git synced 2026-08-09 04:40:57 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(tenancy): add Workspace multi-tenant foundation (#2353)

Browse Source 
* 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
2026-07-30 21:43:35 +08:00
committed by GitHub
parent 463b120923
commit e1ac5e0fc8
468 changed files with 78320 additions and 13137 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/unit_tests/api/http/service/test_bot_service.py
+7 -2
View File
@@ -6,6 +6,9 @@ from sqlalchemy.sql.dml import Update
from langbot.pkg.api.http.service.bot import BotService
WORKSPACE_UUID = 'workspace-a'
class _FakeResult:
def __init__(self, value):
self.value = value
@@ -21,7 +24,9 @@ class _PersistenceManager:
async def execute_async(self, statement):
if isinstance(statement, Update):
self.update_values = {
key: value for key, value in statement.compile().params.items() if not key.startswith('uuid_')
key: value
for key, value in statement.compile().params.items()
if not key.startswith(('uuid_', 'workspace_uuid_'))
}
return None
@@ -48,7 +53,7 @@ async def test_update_bot_copies_input_before_filtering_and_setting_pipeline_nam
'use_pipeline_uuid': 'pipeline-1',
}
await service.update_bot('bot-1', payload)
await service.update_bot(WORKSPACE_UUID, 'bot-1', payload)
assert payload == {
'uuid': 'caller-owned-uuid',
tests/unit_tests/api/http/service/test_tenant.py
+34
View File
@@ -0,0 +1,34 @@
import pytest
import sqlalchemy
from langbot.pkg.api.http.authz import WorkspaceRequiredError
from langbot.pkg.api.http.context import ExecutionContext, PrincipalContext, PrincipalType
from langbot.pkg.api.http.service.tenant import require_workspace_uuid, scope_statement
class _TenantRow:
workspace_uuid = sqlalchemy.column('workspace_uuid')
def test_require_workspace_uuid_accepts_execution_context():
context = ExecutionContext(
instance_uuid='instance-test',
workspace_uuid='workspace-test',
placement_generation=1,
trigger_principal=PrincipalContext(PrincipalType.SYSTEM),
)
assert require_workspace_uuid(context) == 'workspace-test'
@pytest.mark.parametrize('context', [None, '', ' '])
def test_require_workspace_uuid_rejects_missing_context(context):
with pytest.raises(WorkspaceRequiredError):
require_workspace_uuid(context)
def test_scope_statement_adds_workspace_predicate():
statement = scope_statement(sqlalchemy.select(_TenantRow.workspace_uuid), _TenantRow, 'workspace-test')
assert 'workspace_uuid = :workspace_uuid_1' in str(statement)
assert statement.compile().params == {'workspace_uuid_1': 'workspace-test'}
tests/unit_tests/api/http/service/test_tenant_resource_isolation.py
+373
View File
@@ -0,0 +1,373 @@
from __future__ import annotations
import datetime
from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
import sqlalchemy
from sqlalchemy.ext.asyncio import create_async_engine
from langbot.pkg.api.http.service.bot import BotService
from langbot.pkg.api.http.service.model import LLMModelsService
from langbot.pkg.api.http.service.pipeline import PipelineService
from langbot.pkg.api.http.service.provider import ModelProviderService
from langbot.pkg.api.http.service.tenant import require_workspace_uuid
from langbot.pkg.api.http.authz import WorkspaceRequiredError
from langbot.pkg.entity.persistence.base import Base
from langbot.pkg.entity.persistence.bot import Bot
from langbot.pkg.entity.persistence.model import LLMModel, ModelProvider
from langbot.pkg.entity.persistence.pipeline import LegacyPipeline
from langbot.pkg.entity.persistence.workspace import Workspace
from langbot.pkg.workspace.errors import WorkspaceNotFoundError
pytestmark = pytest.mark.asyncio
WORKSPACE_A = '00000000-0000-0000-0000-00000000000a'
WORKSPACE_B = '00000000-0000-0000-0000-00000000000b'
class _PersistenceManager:
def __init__(self, engine):
self.engine = engine
async def execute_async(self, *args, **kwargs):
async with self.engine.connect() as connection:
result = await connection.execute(*args, **kwargs)
await connection.commit()
return result
@staticmethod
def serialize_model(model, data, masked_columns=None):
masked_columns = masked_columns or []
return {
column.name: (
getattr(data, column.name).isoformat()
if isinstance(getattr(data, column.name), datetime.datetime)
else getattr(data, column.name)
)
for column in model.__table__.columns
if column.name not in masked_columns
}
@pytest.fixture
async def tenant_services(tmp_path):
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "tenant-resources.db"}')
async with engine.begin() as connection:
await connection.run_sync(Base.metadata.create_all)
await connection.execute(
sqlalchemy.insert(Workspace),
[
{
'uuid': WORKSPACE_A,
'instance_uuid': 'instance-a',
'name': 'Workspace A',
'slug': 'workspace-a',
'source': 'cloud_projection',
},
{
'uuid': WORKSPACE_B,
'instance_uuid': 'instance-b',
'name': 'Workspace B',
'slug': 'workspace-b',
'source': 'cloud_projection',
},
],
)
await connection.execute(
sqlalchemy.insert(ModelProvider),
[
{
'uuid': 'provider-a',
'workspace_uuid': WORKSPACE_A,
'name': 'Same Provider',
'requester': 'chatcmpl',
'base_url': 'https://a.invalid',
'api_keys': ['secret-a'],
},
{
'uuid': 'provider-b',
'workspace_uuid': WORKSPACE_B,
'name': 'Same Provider',
'requester': 'chatcmpl',
'base_url': 'https://b.invalid',
'api_keys': ['secret-b'],
},
],
)
await connection.execute(
sqlalchemy.insert(LLMModel),
[
{
'uuid': 'model-a',
'workspace_uuid': WORKSPACE_A,
'name': 'Same Model',
'provider_uuid': 'provider-a',
'abilities': [],
'extra_args': {},
'prefered_ranking': 0,
},
{
'uuid': 'model-b',
'workspace_uuid': WORKSPACE_B,
'name': 'Same Model',
'provider_uuid': 'provider-b',
'abilities': [],
'extra_args': {},
'prefered_ranking': 0,
},
],
)
await connection.execute(
sqlalchemy.insert(LegacyPipeline),
[
{
'uuid': 'pipeline-a',
'workspace_uuid': WORKSPACE_A,
'name': 'Same Pipeline',
'description': 'A',
'for_version': 'test',
'is_default': False,
'stages': [],
'config': {},
'extensions_preferences': {},
},
{
'uuid': 'pipeline-b',
'workspace_uuid': WORKSPACE_B,
'name': 'Same Pipeline',
'description': 'B',
'for_version': 'test',
'is_default': False,
'stages': [],
'config': {},
'extensions_preferences': {},
},
],
)
await connection.execute(
sqlalchemy.insert(Bot),
[
{
'uuid': 'bot-a',
'workspace_uuid': WORKSPACE_A,
'name': 'Same Bot',
'description': 'A',
'adapter': 'test',
'adapter_config': {},
'enable': False,
'use_pipeline_uuid': 'pipeline-a',
'use_pipeline_name': 'Same Pipeline',
'pipeline_routing_rules': [],
},
{
'uuid': 'bot-b',
'workspace_uuid': WORKSPACE_B,
'name': 'Same Bot',
'description': 'B',
'adapter': 'test',
'adapter_config': {},
'enable': False,
'use_pipeline_uuid': 'pipeline-b',
'use_pipeline_name': 'Same Pipeline',
'pipeline_routing_rules': [],
},
],
)
runtime_provider_a = SimpleNamespace(provider_entity=SimpleNamespace(uuid='provider-a'))
runtime_provider_b = SimpleNamespace(provider_entity=SimpleNamespace(uuid='provider-b'))
application = SimpleNamespace(
persistence_mgr=_PersistenceManager(engine),
instance_config=SimpleNamespace(data={'system': {'limitation': {}}, 'api': {}}),
ver_mgr=SimpleNamespace(get_current_version=lambda: 'test'),
platform_mgr=SimpleNamespace(
load_bot=AsyncMock(return_value=SimpleNamespace(enable=False)),
remove_bot=AsyncMock(),
get_bot_by_uuid=AsyncMock(return_value=None),
),
pipeline_mgr=SimpleNamespace(
load_pipeline=AsyncMock(),
remove_pipeline=AsyncMock(),
),
model_mgr=SimpleNamespace(
provider_dict={'provider-a': runtime_provider_a, 'provider-b': runtime_provider_b},
llm_models=[],
embedding_models=[],
rerank_models=[],
load_provider=AsyncMock(),
cache_provider=AsyncMock(),
get_provider_by_uuid=AsyncMock(return_value=runtime_provider_a),
reload_provider=AsyncMock(),
remove_provider=AsyncMock(),
load_llm_model_with_provider=AsyncMock(return_value=SimpleNamespace()),
cache_llm_model=AsyncMock(),
remove_llm_model=AsyncMock(),
),
sess_mgr=SimpleNamespace(session_list=[]),
)
application.provider_service = ModelProviderService(application)
application.llm_model_service = LLMModelsService(application)
application.pipeline_service = PipelineService(application)
application.bot_service = BotService(application)
yield application, engine
await engine.dispose()
async def test_context_is_mandatory_and_fails_closed(tenant_services):
application, _engine = tenant_services
with pytest.raises(WorkspaceRequiredError):
require_workspace_uuid(None)
with pytest.raises(WorkspaceRequiredError):
await application.bot_service.get_bots(None)
with pytest.raises(WorkspaceRequiredError):
await application.provider_service.get_providers(None)
with pytest.raises(WorkspaceRequiredError):
await application.pipeline_service.get_pipelines(None)
with pytest.raises(WorkspaceRequiredError):
await application.llm_model_service.get_llm_models(None)
async def test_lists_and_same_names_are_isolated(tenant_services):
application, _engine = tenant_services
assert [item['uuid'] for item in await application.bot_service.get_bots(WORKSPACE_A)] == ['bot-a']
assert [item['uuid'] for item in await application.pipeline_service.get_pipelines(WORKSPACE_A)] == ['pipeline-a']
assert [item['uuid'] for item in await application.provider_service.get_providers(WORKSPACE_A)] == ['provider-a']
assert [item['uuid'] for item in await application.llm_model_service.get_llm_models(WORKSPACE_A)] == ['model-a']
async def test_cross_workspace_uuid_guessing_cannot_read_update_or_delete(tenant_services):
application, engine = tenant_services
assert await application.bot_service.get_bot(WORKSPACE_A, 'bot-b') is None
assert await application.pipeline_service.get_pipeline(WORKSPACE_A, 'pipeline-b') is None
assert await application.provider_service.get_provider(WORKSPACE_A, 'provider-b') is None
assert await application.llm_model_service.get_llm_model(WORKSPACE_A, 'model-b') is None
with pytest.raises(WorkspaceNotFoundError):
await application.bot_service.update_bot(WORKSPACE_A, 'bot-b', {'name': 'stolen'})
with pytest.raises(WorkspaceNotFoundError):
await application.pipeline_service.update_pipeline(
WORKSPACE_A,
'pipeline-b',
{'description': 'stolen'},
)
with pytest.raises(WorkspaceNotFoundError):
await application.provider_service.update_provider(WORKSPACE_A, 'provider-b', {'name': 'stolen'})
with pytest.raises(WorkspaceNotFoundError):
await application.llm_model_service.update_llm_model(
WORKSPACE_A,
'model-b',
{'name': 'stolen'},
)
with pytest.raises(WorkspaceNotFoundError):
await application.bot_service.delete_bot(WORKSPACE_A, 'bot-b')
with pytest.raises(WorkspaceNotFoundError):
await application.pipeline_service.delete_pipeline(WORKSPACE_A, 'pipeline-b')
with pytest.raises(WorkspaceNotFoundError):
await application.provider_service.delete_provider(WORKSPACE_A, 'provider-b')
with pytest.raises(WorkspaceNotFoundError):
await application.llm_model_service.delete_llm_model(WORKSPACE_A, 'model-b')
async with engine.connect() as connection:
assert await connection.scalar(sqlalchemy.select(Bot.name).where(Bot.uuid == 'bot-b')) == 'Same Bot'
assert (
await connection.scalar(sqlalchemy.select(LegacyPipeline.uuid).where(LegacyPipeline.uuid == 'pipeline-b'))
== 'pipeline-b'
)
assert (
await connection.scalar(sqlalchemy.select(ModelProvider.name).where(ModelProvider.uuid == 'provider-b'))
== 'Same Provider'
)
assert await connection.scalar(sqlalchemy.select(LLMModel.uuid).where(LLMModel.uuid == 'model-b')) == 'model-b'
async def test_cross_workspace_parent_references_are_rejected(tenant_services):
application, _engine = tenant_services
with pytest.raises(WorkspaceNotFoundError):
await application.bot_service.update_bot(
WORKSPACE_A,
'bot-a',
{'use_pipeline_uuid': 'pipeline-b'},
)
with pytest.raises(WorkspaceNotFoundError):
await application.llm_model_service.create_llm_model(
WORKSPACE_A,
{
'name': 'Cross reference',
'provider_uuid': 'provider-b',
'abilities': [],
'extra_args': {},
'prefered_ranking': 0,
},
auto_set_to_default_pipeline=False,
)
async def test_created_resources_are_bound_to_callers_workspace(tenant_services):
application, engine = tenant_services
runtime_provider = SimpleNamespace(provider_entity=SimpleNamespace(uuid='provider-created'))
application.model_mgr.load_provider.return_value = runtime_provider
provider_uuid = await application.provider_service.create_provider(
WORKSPACE_A,
{
'name': 'Created Provider',
'requester': 'chatcmpl',
'base_url': 'https://created.invalid',
'api_keys': [],
},
)
pipeline_uuid = await application.pipeline_service.create_pipeline(
WORKSPACE_A,
{'name': 'Created Pipeline', 'description': 'created'},
)
bot_uuid = await application.bot_service.create_bot(
WORKSPACE_A,
{
'name': 'Created Bot',
'description': 'created',
'adapter': 'test',
'adapter_config': {},
'enable': False,
'pipeline_routing_rules': [],
},
)
model_uuid = await application.llm_model_service.create_llm_model(
WORKSPACE_A,
{
'name': 'Created Model',
'provider_uuid': 'provider-a',
'abilities': [],
'extra_args': {},
'prefered_ranking': 0,
},
auto_set_to_default_pipeline=False,
)
async with engine.connect() as connection:
assert (
await connection.scalar(
sqlalchemy.select(ModelProvider.workspace_uuid).where(ModelProvider.uuid == provider_uuid)
)
== WORKSPACE_A
)
assert (
await connection.scalar(
sqlalchemy.select(LegacyPipeline.workspace_uuid).where(LegacyPipeline.uuid == pipeline_uuid)
)
== WORKSPACE_A
)
assert await connection.scalar(sqlalchemy.select(Bot.workspace_uuid).where(Bot.uuid == bot_uuid)) == WORKSPACE_A
assert (
await connection.scalar(sqlalchemy.select(LLMModel.workspace_uuid).where(LLMModel.uuid == model_uuid))
== WORKSPACE_A
)
tests/unit_tests/api/http/test_authz.py
+74
View File
@@ -0,0 +1,74 @@
from langbot.pkg.api.http import authz
from langbot.pkg.api.http.context import PrincipalContext, PrincipalType, RequestContext, WorkspaceContext
def _context(role: authz.WorkspaceRole) -> RequestContext:
return RequestContext(
instance_uuid='instance-test',
placement_generation=1,
request_id='request-test',
auth_type='user-token',
principal=PrincipalContext(
principal_type=PrincipalType.ACCOUNT,
account_uuid='account-test',
),
workspace=WorkspaceContext(
workspace_uuid='workspace-test',
membership_uuid='membership-test',
role=role.value,
permissions=authz.permissions_for_role(role),
),
)
def test_owner_has_every_fixed_permission():
ctx = _context(authz.WorkspaceRole.OWNER)
assert ctx.workspace.permissions == frozenset(permission.value for permission in authz.Permission)
def test_admin_cannot_transfer_owner_delete_workspace_or_link_billing():
ctx = _context(authz.WorkspaceRole.ADMIN)
assert not authz.has_permission(ctx, authz.Permission.OWNER_TRANSFER)
assert not authz.has_permission(ctx, authz.Permission.WORKSPACE_DELETE)
assert not authz.has_permission(ctx, authz.Permission.BILLING_LINK_MANAGE)
assert authz.has_permission(ctx, authz.Permission.MEMBER_INVITE)
def test_operator_can_run_but_cannot_manage_resources_or_secrets():
ctx = _context(authz.WorkspaceRole.OPERATOR)
assert authz.has_permission(ctx, authz.Permission.RUNTIME_OPERATE)
assert not authz.has_permission(ctx, authz.Permission.RESOURCE_MANAGE)
assert not authz.has_permission(ctx, authz.Permission.PROVIDER_SECRET_MANAGE)
def test_unknown_role_has_no_permissions():
assert authz.permissions_for_role('unknown') == frozenset()
def test_require_permission_reports_stable_permission():
ctx = _context(authz.WorkspaceRole.VIEWER)
try:
authz.require_permission(ctx, authz.Permission.RESOURCE_MANAGE)
except authz.PermissionDeniedError as exc:
assert exc.permission == authz.Permission.RESOURCE_MANAGE.value
assert exc.error_code == 'permission_denied'
else:
raise AssertionError('PermissionDeniedError was not raised')
def test_execution_context_preserves_workspace_and_generation():
from langbot.pkg.api.http.context import ExecutionContext
ctx = _context(authz.WorkspaceRole.DEVELOPER)
execution = ExecutionContext.from_request(ctx, bot_uuid='bot-test', pipeline_uuid='pipeline-test')
assert execution.instance_uuid == 'instance-test'
assert execution.workspace_uuid == 'workspace-test'
assert execution.placement_generation == 1
assert execution.bot_uuid == 'bot-test'
assert execution.pipeline_uuid == 'pipeline-test'
assert execution.trigger_principal == ctx.principal
tests/unit_tests/api/http/test_bounded_json_request.py
+39
View File
@@ -0,0 +1,39 @@
from __future__ import annotations
import quart
from langbot.pkg.api.http.controller import main as controller_main
from langbot.pkg.utils import bounded_executor
async def test_bounded_json_request_decodes_off_loop_in_workspace_scope(
monkeypatch,
):
app = quart.Quart(__name__)
app.request_class = controller_main.BoundedJSONRequest
observed_scopes: list[str | None] = []
async def fake_to_thread(fn, *args, **kwargs):
observed_scopes.append(bounded_executor.current_blocking_work_scope())
return fn(*args, **kwargs)
monkeypatch.setattr(
controller_main.asyncio,
'to_thread',
fake_to_thread,
)
@app.post('/json')
async def parse_json():
with bounded_executor.blocking_work_scope('workspace-a'):
payload = await quart.request.get_json()
return quart.jsonify(payload)
response = await app.test_client().post(
'/json',
json={'nested': {'value': 1}},
)
assert response.status_code == 200
assert await response.get_json() == {'nested': {'value': 1}}
assert observed_scopes == ['workspace-a']
tests/unit_tests/api/http/test_entitlement_context.py
+78
View File
@@ -0,0 +1,78 @@
from __future__ import annotations
from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
import quart
from langbot.pkg.api.http.context import (
ExecutionContext,
PrincipalContext,
PrincipalType,
RequestContext,
WorkspaceContext,
)
from langbot.pkg.api.http.controller.group import RouterGroup
from langbot.pkg.cloud.entitlements import EntitlementSnapshot, EntitlementUnavailableError
from langbot.pkg.cloud.entitlements import EntitlementResolver
class _Group(RouterGroup):
async def initialize(self) -> None:
return None
def _router(deployment) -> _Group:
provider = getattr(deployment, 'entitlement_provider', None)
resolver = EntitlementResolver('instance-a', provider) if provider is not None else None
ap = SimpleNamespace(deployment=deployment, entitlement_resolver=resolver)
return _Group(ap, quart.Quart(__name__))
@pytest.mark.asyncio
async def test_cloud_request_resolves_verified_entitlement_revision():
snapshot = EntitlementSnapshot(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
entitlement_revision=9,
status='active',
not_before=1,
expires_at=4_000_000_000,
features={},
limits={},
)
provider = SimpleNamespace(get_workspace_entitlement=AsyncMock(return_value=snapshot))
router = _router(SimpleNamespace(multi_workspace_enabled=True, entitlement_provider=provider))
revision = await router._resolve_entitlement_revision('instance-a', 'workspace-a')
assert revision == 9
provider.get_workspace_entitlement.assert_awaited_once_with('workspace-a')
@pytest.mark.asyncio
async def test_cloud_request_fails_closed_without_entitlement_provider():
router = _router(SimpleNamespace(multi_workspace_enabled=True, entitlement_provider=None))
with pytest.raises(EntitlementUnavailableError):
await router._resolve_entitlement_revision('instance-a', 'workspace-a')
def test_execution_context_preserves_entitlement_revision():
request = RequestContext(
instance_uuid='instance-a',
placement_generation=1,
request_id='request-a',
auth_type='user-token',
principal=PrincipalContext(PrincipalType.ACCOUNT, account_uuid='account-a'),
workspace=WorkspaceContext(
workspace_uuid='workspace-a',
membership_uuid='membership-a',
role='owner',
permissions=frozenset(),
),
entitlement_revision=11,
)
assert ExecutionContext.from_request(request).entitlement_revision == 11
tests/unit_tests/api/http/test_internal_error_responses.py
+257
View File
@@ -0,0 +1,257 @@
from __future__ import annotations
import contextlib
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
import pytest
import quart
from langbot.pkg.api.http.controller import group
from langbot.pkg.api.http.controller.groups.webhooks import WebhookRouterGroup
from langbot.pkg.utils.bounded_executor import (
BlockingWorkCapacityError,
current_blocking_work_scope,
)
pytestmark = pytest.mark.asyncio
class _FailingRouterGroup(group.RouterGroup):
name = 'failing-test'
path = '/failing-test'
async def initialize(self) -> None:
@self.route('', methods=['GET'], auth_type=group.AuthType.NONE)
async def _():
raise RuntimeError('database password=do-not-return')
class _AuthenticatedRouterGroup(group.RouterGroup):
name = 'authenticated-test'
path = '/authenticated-test'
async def initialize(self) -> None:
@self.route('', methods=['GET'], auth_type=group.AuthType.USER_TOKEN)
async def _():
return self.success()
class _BlockingCapacityRouterGroup(group.RouterGroup):
name = 'blocking-capacity-test'
path = '/blocking-capacity-test'
async def initialize(self) -> None:
@self.route('', methods=['GET'], auth_type=group.AuthType.NONE)
async def _():
raise BlockingWorkCapacityError('Workspace blocking executor capacity reached')
class _InvalidAccountRouterGroup(group.RouterGroup):
name = 'invalid-account-test'
path = '/invalid-account-test'
async def initialize(self) -> None:
@self.route(
'',
methods=['GET'],
auth_type=group.AuthType.ACCOUNT_TOKEN,
permission='workspace.view',
)
async def _():
return self.success()
async def test_unhandled_http_error_returns_generic_body_and_correlated_request_id():
logger = Mock()
application = SimpleNamespace(logger=logger)
quart_app = quart.Quart(__name__)
await _FailingRouterGroup(application, quart_app).initialize()
response = await quart_app.test_client().get(
'/failing-test',
headers={'X-Request-Id': 'request-http-test'},
)
assert response.status_code == 500
assert await response.get_json() == {
'code': 'internal_error',
'msg': 'Internal server error',
'request_id': 'request-http-test',
}
assert response.headers['X-Request-Id'] == 'request-http-test'
log_message = logger.error.call_args.args[0]
assert 'request_id=request-http-test' in log_message
assert 'database password=do-not-return' in log_message
assert 'do-not-return' not in (await response.get_data(as_text=True))
async def test_public_webhook_error_uses_same_generic_error_contract():
logger = Mock()
application = SimpleNamespace(
logger=logger,
platform_mgr=SimpleNamespace(
resolve_public_bot=AsyncMock(side_effect=RuntimeError('adapter credential=do-not-return'))
),
)
quart_app = quart.Quart(__name__)
await WebhookRouterGroup(application, quart_app).initialize()
response = await quart_app.test_client().post(
'/bots/11111111-1111-4111-8111-111111111111',
headers={'X-Request-Id': 'request-webhook-test'},
)
assert response.status_code == 500
assert await response.get_json() == {
'code': 'internal_error',
'msg': 'Internal server error',
'request_id': 'request-webhook-test',
}
assert response.headers['X-Request-Id'] == 'request-webhook-test'
log_message = logger.error.call_args.args[0]
assert 'request_id=request-webhook-test' in log_message
assert 'adapter credential=do-not-return' in log_message
assert 'do-not-return' not in (await response.get_data(as_text=True))
async def test_blocking_work_capacity_maps_to_retryable_http_response():
application = SimpleNamespace(logger=Mock())
quart_app = quart.Quart(__name__)
await _BlockingCapacityRouterGroup(application, quart_app).initialize()
response = await quart_app.test_client().get('/blocking-capacity-test')
assert response.status_code == 429
assert await response.get_json() == {
'code': 'blocking_work_capacity_exceeded',
'msg': 'Workspace blocking executor capacity reached',
}
async def test_public_webhook_carries_scope_without_holding_database_session():
class ScopeOnlyPersistenceManager:
mode = SimpleNamespace(value='cloud_runtime')
def __init__(self):
self.active_workspace = None
@contextlib.asynccontextmanager
async def tenant_scope(self, workspace_uuid):
self.active_workspace = workspace_uuid
try:
yield
finally:
self.active_workspace = None
def current_session(self):
return None
persistence_mgr = ScopeOnlyPersistenceManager()
workspace_uuid = '00000000-0000-0000-0000-00000000000a'
bot_uuid = '11111111-1111-4111-8111-111111111111'
class Adapter:
async def handle_unified_webhook(self, **_kwargs):
assert persistence_mgr.active_workspace == workspace_uuid
assert persistence_mgr.current_session() is None
assert current_blocking_work_scope() == workspace_uuid
return {'ok': True}
async def get_execution_binding(resolved_workspace_uuid, expected_generation=None):
assert resolved_workspace_uuid == workspace_uuid
assert expected_generation == 4
runtime_bot = SimpleNamespace(
workspace_uuid=workspace_uuid,
placement_generation=4,
enable=True,
adapter=Adapter(),
)
application = SimpleNamespace(
logger=Mock(),
persistence_mgr=persistence_mgr,
platform_mgr=SimpleNamespace(resolve_public_bot=AsyncMock(return_value=runtime_bot)),
workspace_service=SimpleNamespace(get_execution_binding=get_execution_binding),
)
quart_app = quart.Quart(__name__)
await WebhookRouterGroup(application, quart_app).initialize()
response = await quart_app.test_client().post(f'/bots/{bot_uuid}')
assert response.status_code == 200
assert await response.get_json() == {'ok': True}
assert persistence_mgr.active_workspace is None
async def test_public_webhook_blocking_capacity_is_retryable():
workspace_uuid = '00000000-0000-0000-0000-00000000000a'
bot_uuid = '11111111-1111-4111-8111-111111111111'
class Adapter:
async def handle_unified_webhook(self, **_kwargs):
raise BlockingWorkCapacityError(
'Workspace blocking executor capacity reached',
scope=workspace_uuid,
)
runtime_bot = SimpleNamespace(
workspace_uuid=workspace_uuid,
placement_generation=4,
enable=True,
adapter=Adapter(),
)
application = SimpleNamespace(
logger=Mock(),
persistence_mgr=SimpleNamespace(mode=SimpleNamespace(value='oss')),
platform_mgr=SimpleNamespace(resolve_public_bot=AsyncMock(return_value=runtime_bot)),
workspace_service=SimpleNamespace(get_execution_binding=AsyncMock(return_value=None)),
)
quart_app = quart.Quart(__name__)
await WebhookRouterGroup(application, quart_app).initialize()
response = await quart_app.test_client().post(f'/bots/{bot_uuid}')
assert response.status_code == 429
assert await response.get_json() == {
'code': 'blocking_work_capacity_exceeded',
'msg': 'Workspace blocking executor capacity reached',
}
async def test_authentication_failure_does_not_return_internal_exception_text():
logger = Mock()
application = SimpleNamespace(
logger=logger,
user_service=SimpleNamespace(
get_authenticated_account=AsyncMock(side_effect=RuntimeError('database password=do-not-return'))
),
)
quart_app = quart.Quart(__name__)
await _AuthenticatedRouterGroup(application, quart_app).initialize()
response = await quart_app.test_client().get(
'/authenticated-test',
headers={
'Authorization': 'Bearer invalid',
'X-Request-Id': 'request-auth-test',
},
)
assert response.status_code == 401
assert await response.get_json() == {
'code': 'invalid_authentication',
'msg': 'Invalid authentication credentials',
}
assert 'do-not-return' not in (await response.get_data(as_text=True))
assert 'request_id=request-auth-test' in logger.warning.call_args.args[0]
assert 'database password=do-not-return' in logger.warning.call_args.args[0]
async def test_account_token_route_cannot_declare_workspace_permission():
application = SimpleNamespace(logger=Mock())
quart_app = quart.Quart(__name__)
with pytest.raises(ValueError, match='cannot declare Workspace permissions'):
await _InvalidAccountRouterGroup(application, quart_app).initialize()
tests/unit_tests/api/http/test_route_tenant_scope.py
+73
View File
@@ -0,0 +1,73 @@
from __future__ import annotations
import asyncio
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
import pytest
import quart
from sqlalchemy.ext.asyncio import create_async_engine
from langbot.pkg.api.http.controller import group
from langbot.pkg.persistence.mgr import PersistenceManager, PersistenceMode
pytestmark = pytest.mark.asyncio
async def test_authenticated_route_does_not_hold_database_session_during_external_wait():
entered = asyncio.Event()
release = asyncio.Event()
observations: list[bool] = []
engine = create_async_engine('sqlite+aiosqlite:///:memory:')
persistence = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
persistence.db = SimpleNamespace(get_engine=lambda: engine)
class BlockingRouter(group.RouterGroup):
name = 'blocking-route-test'
path = '/blocking-route-test'
async def initialize(self) -> None:
@self.route('', methods=['GET'], auth_type=group.AuthType.USER_TOKEN)
async def _():
observations.append(persistence.current_session() is None)
entered.set()
await release.wait()
observations.append(persistence.current_session() is None)
return self.success(data={})
account = SimpleNamespace(uuid='account-a', user='owner@example.com')
access = SimpleNamespace(
execution=SimpleNamespace(instance_uuid='instance-a', placement_generation=1),
workspace=SimpleNamespace(uuid='workspace-a'),
membership=SimpleNamespace(uuid='membership-a', role='owner', projection_revision=1),
)
application = SimpleNamespace(
persistence_mgr=persistence,
deployment=SimpleNamespace(multi_workspace_enabled=False),
user_service=SimpleNamespace(get_authenticated_account=AsyncMock(return_value=account)),
workspace_collaboration_service=SimpleNamespace(resolve_account_workspace=AsyncMock(return_value=access)),
logger=Mock(),
)
quart_app = quart.Quart(__name__)
await BlockingRouter(application, quart_app).initialize()
client = quart_app.test_client()
request = asyncio.create_task(
client.get(
'/blocking-route-test',
headers={'Authorization': 'Bearer token', 'X-Workspace-Id': 'workspace-a'},
)
)
try:
await entered.wait()
assert observations == [True]
release.set()
response = await request
assert response.status_code == 200
assert observations == [True, True]
finally:
release.set()
if not request.done():
await request
await engine.dispose()
tests/unit_tests/api/service/test_apikey_service.py
+382 -398
View File
@@ -1,482 +1,466 @@
"""
Unit tests for ApiKeyService.
Tests API key CRUD operations with mocked persistence layer.
Source: src/langbot/pkg/api/http/service/apikey.py
"""
from __future__ import annotations
import pytest
from unittest.mock import AsyncMock, Mock, patch
import datetime
import hashlib
import logging
import uuid
from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
import sqlalchemy
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
from langbot.pkg.api.http.authz import Permission, PermissionDeniedError
from langbot.pkg.api.http.context import PrincipalContext, PrincipalType, RequestContext, WorkspaceContext
from langbot.pkg.api.http.service.apikey import ApiKeyService
from langbot.pkg.entity.persistence.apikey import ApiKey
from langbot.pkg.entity.persistence.base import Base
from langbot.pkg.entity.persistence.user import User
from langbot.pkg.entity.persistence.workspace import (
Workspace,
WorkspaceExecutionSource,
WorkspaceExecutionState,
WorkspaceSource,
)
from langbot.pkg.workspace.policy import SingleWorkspacePolicy
from langbot.pkg.workspace.errors import WorkspaceNotFoundError
from langbot.pkg.workspace.service import WorkspaceService
pytestmark = pytest.mark.asyncio
class _PersistenceManager:
def __init__(self, engine):
self.engine = engine
def get_db_engine(self):
return self.engine
async def execute_async(self, *args, **kwargs):
async with self.engine.connect() as connection:
result = await connection.execute(*args, **kwargs)
await connection.commit()
return result
@staticmethod
def serialize_model(model, row, masked_columns=()):
return {
column.name: (
getattr(row, column.name).isoformat()
if isinstance(getattr(row, column.name), datetime.datetime)
else getattr(row, column.name)
)
for column in model.__table__.columns
if column.name not in masked_columns
}
def _context(workspace_uuid: str, account_uuid: str, permissions: set[Permission]) -> RequestContext:
return RequestContext(
instance_uuid='api-key-instance',
placement_generation=1,
request_id=str(uuid.uuid4()),
auth_type='user-token',
principal=PrincipalContext(PrincipalType.ACCOUNT, account_uuid=account_uuid),
workspace=WorkspaceContext(
workspace_uuid=workspace_uuid,
membership_uuid=str(uuid.uuid4()),
role='owner',
permissions=frozenset(permission.value for permission in permissions),
),
)
@pytest.fixture
async def api_key_context(tmp_path):
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "api-keys.db"}')
async with engine.begin() as connection:
await connection.run_sync(Base.metadata.create_all)
application = SimpleNamespace(
persistence_mgr=_PersistenceManager(engine),
instance_config=SimpleNamespace(data={'api': {'global_api_key': ''}}),
logger=logging.getLogger('api-key-test'),
)
application.workspace_service = WorkspaceService(application, instance_uuid='api-key-instance')
workspace = await application.workspace_service.ensure_singleton_workspace()
account_uuid = str(uuid.uuid4())
session_factory = async_sessionmaker(engine, expire_on_commit=False)
async with session_factory.begin() as session:
session.add(
User(
uuid=account_uuid,
user='owner@example.com',
normalized_email='owner@example.com',
password='hash',
account_type='local',
)
)
service = ApiKeyService(application)
context = _context(workspace.uuid, account_uuid, set(Permission))
yield application, service, context, engine
await engine.dispose()
async def test_secret_is_returned_once_and_only_hash_is_persisted(api_key_context):
_application, service, context, engine = api_key_context
created = await service.create_api_key(context, 'Automation', 'CI key')
secret = created['key']
assert secret.startswith('lbk_')
assert created['secret_available'] is True
assert 'key_hash' not in created
listed = await service.get_api_keys(context)
assert len(listed) == 1
assert 'key' not in listed[0]
assert 'key_hash' not in listed[0]
assert listed[0]['secret_available'] is False
async with engine.connect() as connection:
stored = await connection.scalar(sqlalchemy.select(ApiKey.key_hash))
assert stored == hashlib.sha256(secret.encode()).hexdigest()
assert secret not in stored
async def test_authentication_derives_workspace_scopes_and_updates_usage(api_key_context):
_application, service, context, engine = api_key_context
created = await service.create_api_key(
context,
'Read only',
scopes=[Permission.RESOURCE_VIEW.value],
)
identity = await service.authenticate_api_key(created['key'])
assert identity is not None
assert identity.workspace_uuid == context.workspace_uuid
assert identity.permissions == frozenset({Permission.RESOURCE_VIEW.value})
async with engine.connect() as connection:
last_used_at = await connection.scalar(sqlalchemy.select(ApiKey.last_used_at))
assert last_used_at is not None
async def test_revoked_expired_and_unknown_keys_fail_closed(api_key_context):
_application, service, context, _engine = api_key_context
created = await service.create_api_key(context, 'Revocable')
await service.delete_api_key(context, created['id'])
assert await service.authenticate_api_key(created['key']) is None
assert await service.verify_api_key('') is False
assert await service.verify_api_key('plain-secret') is False
assert await service.verify_api_key('lbk_unknown') is False
expired_secret = 'lbk_expired'
await service.ap.persistence_mgr.execute_async(
sqlalchemy.insert(ApiKey).values(
workspace_uuid=context.workspace_uuid,
name='Expired',
key_hash=hashlib.sha256(expired_secret.encode()).hexdigest(),
scopes=[Permission.RESOURCE_VIEW.value],
status='active',
expires_at=datetime.datetime.now(datetime.UTC).replace(tzinfo=None) - datetime.timedelta(seconds=1),
)
)
assert await service.authenticate_api_key(expired_secret) is None
async def test_revoke_winning_last_used_update_race_fails_authentication(api_key_context):
application, service, context, _engine = api_key_context
created = await service.create_api_key(context, 'Racing revoke')
original_execute = application.persistence_mgr.execute_async
injected_revoke = False
async def execute_with_revoke(statement, *args, **kwargs):
nonlocal injected_revoke
if (
not injected_revoke
and isinstance(statement, sqlalchemy.sql.dml.Update)
and statement.table.name == ApiKey.__tablename__
):
injected_revoke = True
await original_execute(sqlalchemy.update(ApiKey).where(ApiKey.id == created['id']).values(status='revoked'))
return await original_execute(statement, *args, **kwargs)
application.persistence_mgr.execute_async = execute_with_revoke
assert await service.authenticate_api_key(created['key']) is None
assert injected_revoke is True
async def test_cross_workspace_crud_and_secret_guessing_are_isolated(api_key_context):
application, service, first_context, engine = api_key_context
second_workspace_uuid = str(uuid.uuid4())
async with async_sessionmaker(engine, expire_on_commit=False).begin() as session:
session.add(
Workspace(
uuid=second_workspace_uuid,
instance_uuid='api-key-instance',
name='Second',
slug='second',
source=WorkspaceSource.CLOUD_PROJECTION.value,
)
)
session.add(
WorkspaceExecutionState(
workspace_uuid=second_workspace_uuid,
instance_uuid='api-key-instance',
active_generation=3,
state='active',
write_fenced=False,
source=WorkspaceExecutionSource.CLOUD.value,
)
)
second_context = _context(second_workspace_uuid, first_context.account_uuid or '', set(Permission))
created = await service.create_api_key(first_context, 'First only')
assert await service.get_api_key(second_context, created['id']) is None
assert await service.get_api_keys(second_context) == []
identity = await service.authenticate_api_key(created['key'])
assert identity is not None
assert identity.workspace_uuid == first_context.workspace_uuid
assert identity.workspace_uuid != second_workspace_uuid
# Prove the explicit multi-Workspace policy does not change key-derived routing.
application.workspace_service.policy = SingleWorkspacePolicy(workspace_limit=10, multi_workspace_enabled=True)
identity = await service.authenticate_api_key(created['key'])
assert identity is not None
assert identity.workspace_uuid == first_context.workspace_uuid
async def test_global_config_key_is_oss_singleton_only(api_key_context):
application, service, _context_value, _engine = api_key_context
application.instance_config.data['api']['global_api_key'] = 'configured-secret'
identity = await service.authenticate_api_key('configured-secret')
assert identity is not None
assert identity.api_key_uuid == 'global-oss-api-key'
application.workspace_service.policy = SingleWorkspacePolicy(workspace_limit=10, multi_workspace_enabled=True)
assert await service.authenticate_api_key('configured-secret') is None
async def test_explicit_scopes_cannot_exceed_callers_workspace_permissions(api_key_context):
_application, service, context, _engine = api_key_context
limited_context = _context(
context.workspace_uuid,
context.account_uuid or '',
{Permission.API_KEY_MANAGE, Permission.RESOURCE_VIEW},
)
created = await service.create_api_key(
limited_context,
'Read only',
scopes=[Permission.RESOURCE_VIEW.value],
)
identity = await service.authenticate_api_key(created['key'])
assert identity is not None
assert identity.permissions == frozenset({Permission.RESOURCE_VIEW.value})
with pytest.raises(PermissionDeniedError) as exc_info:
await service.create_api_key(
limited_context,
'Escalated',
scopes=[Permission.WORKSPACE_DELETE.value],
)
assert exc_info.value.permission == Permission.WORKSPACE_DELETE.value
# Preserve the pre-tenancy CRUD and verification regression matrix while
# exercising it through the new Workspace-bound API. The assertions reflect
# intentional security changes: secrets are returned once, deletion revokes,
# and missing Workspace resources are reported as not found.
class TestApiKeyServiceGetApiKeys:
"""Tests for get_api_keys method."""
async def test_get_api_keys_empty_list(self, api_key_context):
_application, service, context, _engine = api_key_context
async def test_get_api_keys_empty_list(self):
"""Returns empty list when no API keys exist."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_result = Mock()
mock_result.all = Mock(return_value=[])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(
side_effect=lambda model_cls, entity: {
'id': entity.id,
'name': entity.name,
'key': entity.key,
'description': entity.description,
}
if entity
else {}
)
assert await service.get_api_keys(context) == []
service = ApiKeyService(ap)
async def test_get_api_keys_returns_serialized_list(self, api_key_context):
_application, service, context, _engine = api_key_context
await service.create_api_key(context, 'Test Key 1', 'First test key')
await service.create_api_key(context, 'Test Key 2', 'Second test key')
# Execute
result = await service.get_api_keys()
result = await service.get_api_keys(context)
# Verify
assert result == []
ap.persistence_mgr.execute_async.assert_called_once()
async def test_get_api_keys_returns_serialized_list(self):
"""Returns serialized list of API keys."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
# Create mock API key entities
key1 = Mock(spec=ApiKey)
key1.id = 1
key1.name = 'Test Key 1'
key1.key = 'lbk_test_key_1'
key1.description = 'First test key'
key2 = Mock(spec=ApiKey)
key2.id = 2
key2.name = 'Test Key 2'
key2.key = 'lbk_test_key_2'
key2.description = 'Second test key'
mock_result = Mock()
mock_result.all = Mock(return_value=[key1, key2])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(
side_effect=lambda model_cls, entity: {
'id': entity.id,
'name': entity.name,
'key': entity.key,
'description': entity.description,
}
)
service = ApiKeyService(ap)
# Execute
result = await service.get_api_keys()
# Verify
assert len(result) == 2
assert result[0]['name'] == 'Test Key 1'
assert result[1]['name'] == 'Test Key 2'
assert [item['name'] for item in result] == ['Test Key 1', 'Test Key 2']
assert [item['description'] for item in result] == ['First test key', 'Second test key']
assert all('key' not in item and 'key_hash' not in item for item in result)
class TestApiKeyServiceCreateApiKey:
"""Tests for create_api_key method."""
async def test_create_api_key_generates_key_with_prefix(self, api_key_context):
_application, service, context, _engine = api_key_context
async def test_create_api_key_generates_key_with_prefix(self):
"""Creates API key with 'lbk_' prefix."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
with pytest.MonkeyPatch.context() as monkeypatch:
monkeypatch.setattr(
'langbot.pkg.api.http.service.apikey.secrets.token_urlsafe', lambda _size: 'fixed-token'
)
result = await service.create_api_key(context, 'New Key', 'Test description')
created_key = Mock(spec=ApiKey)
created_key.id = 1
created_key.name = 'New Key'
created_key.key = 'lbk_fixed-token'
created_key.description = 'Test description'
select_result = Mock()
select_result.first = Mock(return_value=created_key)
insert_params = []
async def mock_execute(query):
params = query.compile().params
if {'name', 'key', 'description'}.issubset(params):
insert_params.append(params)
return Mock()
return select_result
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
ap.persistence_mgr.serialize_model = Mock(
side_effect=lambda model_cls, entity: {
'id': 1,
'name': entity.name,
'key': entity.key,
'description': entity.description,
}
)
service = ApiKeyService(ap)
with patch('langbot.pkg.api.http.service.apikey.secrets.token_urlsafe', return_value='fixed-token'):
result = await service.create_api_key('New Key', 'Test description')
assert insert_params == [{'name': 'New Key', 'key': 'lbk_fixed-token', 'description': 'Test description'}]
assert result['key'].startswith('lbk_')
assert result['key'] == 'lbk_fixed-token'
assert result['name'] == 'New Key'
assert result['description'] == 'Test description'
assert result['secret_available'] is True
async def test_create_api_key_without_description(self):
"""Creates API key with empty description when not provided."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
async def test_create_api_key_without_description(self, api_key_context):
_application, service, context, _engine = api_key_context
created_key = Mock(spec=ApiKey)
created_key.id = 1
created_key.name = 'No Desc Key'
created_key.key = 'lbk_no_desc_key'
created_key.description = ''
result = await service.create_api_key(context, 'No Desc Key')
select_result = Mock()
select_result.first = Mock(return_value=created_key)
insert_result = Mock()
async def mock_execute(query):
if hasattr(query, 'values'):
return insert_result
return select_result
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
ap.persistence_mgr.serialize_model = Mock(
return_value={
'id': 1,
'name': 'No Desc Key',
'key': 'lbk_no_desc_key',
'description': '',
}
)
service = ApiKeyService(ap)
# Execute
result = await service.create_api_key('No Desc Key')
# Verify
assert result['description'] == ''
class TestApiKeyServiceGetApiKey:
"""Tests for get_api_key method."""
async def test_get_api_key_by_id_found(self, api_key_context):
_application, service, context, _engine = api_key_context
created = await service.create_api_key(context, 'Found Key', 'Found')
async def test_get_api_key_by_id_found(self):
"""Returns API key when found by ID."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
result = await service.get_api_key(context, created['id'])
key = Mock(spec=ApiKey)
key.id = 1
key.name = 'Found Key'
key.key = 'lbk_found_key'
key.description = 'Found'
mock_result = Mock()
mock_result.first = Mock(return_value=key)
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(
return_value={
'id': 1,
'name': 'Found Key',
'key': 'lbk_found_key',
'description': 'Found',
}
)
service = ApiKeyService(ap)
# Execute
result = await service.get_api_key(1)
# Verify
assert result is not None
assert result['id'] == 1
assert result['id'] == created['id']
assert result['name'] == 'Found Key'
assert 'key' not in result and 'key_hash' not in result
async def test_get_api_key_by_id_not_found(self):
"""Returns None when API key not found."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
async def test_get_api_key_by_id_not_found(self, api_key_context):
_application, service, context, _engine = api_key_context
mock_result = Mock()
mock_result.first = Mock(return_value=None)
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
assert await service.get_api_key(context, 999) is None
service = ApiKeyService(ap)
async def test_get_api_key_by_id_zero(self, api_key_context):
_application, service, context, _engine = api_key_context
# Execute
result = await service.get_api_key(999)
# Verify
assert result is None
async def test_get_api_key_by_id_zero(self):
"""Handles ID=0 (edge case) correctly."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_result = Mock()
mock_result.first = Mock(return_value=None)
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = ApiKeyService(ap)
# Execute
result = await service.get_api_key(0)
# Verify - should return None (no key with ID 0)
assert result is None
assert await service.get_api_key(context, 0) is None
class TestApiKeyServiceVerifyApiKey:
"""Tests for verify_api_key method."""
async def test_verify_api_key_valid(self, api_key_context):
_application, service, context, _engine = api_key_context
created = await service.create_api_key(context, 'Valid')
@staticmethod
def _make_ap(db_key=None, global_api_key=''):
"""Build a mock Application with persistence + instance_config."""
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_result = Mock()
mock_result.first = Mock(return_value=db_key)
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.instance_config = SimpleNamespace(data={'api': {'global_api_key': global_api_key}})
return ap
assert await service.verify_api_key(created['key']) is True
async def test_verify_api_key_valid(self):
"""Returns True for valid API key."""
# Setup
key = Mock(spec=ApiKey)
ap = self._make_ap(db_key=key)
async def test_verify_api_key_invalid(self, api_key_context):
_application, service, _context, _engine = api_key_context
service = ApiKeyService(ap)
assert await service.verify_api_key('lbk_invalid_key') is False
# Execute
result = await service.verify_api_key('lbk_valid_key')
async def test_verify_api_key_empty_string(self, api_key_context):
_application, service, _context, _engine = api_key_context
# Verify
assert result is True
assert await service.verify_api_key('') is False
async def test_verify_api_key_invalid(self):
"""Returns False for invalid API key."""
# Setup
ap = self._make_ap(db_key=None)
async def test_verify_api_key_unknown_key(self, api_key_context):
_application, service, _context, _engine = api_key_context
service = ApiKeyService(ap)
assert await service.verify_api_key('unknown_key') is False
# Execute
result = await service.verify_api_key('lbk_invalid_key')
async def test_verify_global_api_key_match(self, api_key_context):
application, service, context, _engine = api_key_context
application.instance_config.data['api']['global_api_key'] = 'my-global-secret'
# Verify
assert result is False
identity = await service.authenticate_api_key('my-global-secret')
async def test_verify_api_key_empty_string(self):
"""Returns False for empty key string."""
# Setup
ap = self._make_ap(db_key=None)
assert identity is not None
assert identity.workspace_uuid == context.workspace_uuid
assert identity.api_key_uuid == 'global-oss-api-key'
service = ApiKeyService(ap)
async def test_verify_global_api_key_no_prefix_required(self, api_key_context):
application, service, _context, _engine = api_key_context
application.instance_config.data['api']['global_api_key'] = 'plainsecret123'
# Execute
result = await service.verify_api_key('')
assert await service.verify_api_key('plainsecret123') is True
# Verify
assert result is False
async def test_verify_global_api_key_mismatch_falls_back_to_db(self, api_key_context):
application, service, context, _engine = api_key_context
application.instance_config.data['api']['global_api_key'] = 'my-global-secret'
created = await service.create_api_key(context, 'DB key')
async def test_verify_api_key_unknown_key(self):
"""Returns False when the key is not present in persistence."""
# Setup
ap = self._make_ap(db_key=None)
identity = await service.authenticate_api_key(created['key'])
service = ApiKeyService(ap)
assert identity is not None
assert identity.api_key_uuid == created['uuid']
# Execute
result = await service.verify_api_key('unknown_key')
async def test_verify_empty_global_api_key_disabled(self, api_key_context):
application, service, _context, _engine = api_key_context
application.instance_config.data['api']['global_api_key'] = ''
# Verify
assert result is False
async def test_verify_global_api_key_match(self):
"""Returns True when key matches the config.yaml global API key (no DB lookup)."""
# Setup: no DB record, but a global key is configured
ap = self._make_ap(db_key=None, global_api_key='my-global-secret')
service = ApiKeyService(ap)
# Execute
result = await service.verify_api_key('my-global-secret')
# Verify: accepted purely on config match
assert result is True
# DB should not have been consulted for the global-key path
ap.persistence_mgr.execute_async.assert_not_called()
async def test_verify_global_api_key_no_prefix_required(self):
"""Global API key is accepted even without the lbk_ prefix."""
ap = self._make_ap(db_key=None, global_api_key='plainsecret123')
service = ApiKeyService(ap)
result = await service.verify_api_key('plainsecret123')
assert result is True
async def test_verify_global_api_key_mismatch_falls_back_to_db(self):
"""A non-matching key still falls through to the DB lookup."""
# Global key set, but request uses a different lbk_ key that IS in DB
key = Mock(spec=ApiKey)
ap = self._make_ap(db_key=key, global_api_key='my-global-secret')
service = ApiKeyService(ap)
result = await service.verify_api_key('lbk_db_key')
assert result is True
ap.persistence_mgr.execute_async.assert_called_once()
async def test_verify_empty_global_api_key_disabled(self):
"""An empty global_api_key must never authenticate an empty/blank request."""
ap = self._make_ap(db_key=None, global_api_key='')
service = ApiKeyService(ap)
# Empty request key is rejected, and a blank global key never matches
assert await service.verify_api_key('') is False
assert await service.verify_api_key(' ') is False
async def test_verify_api_key_missing_global_config_key(self):
"""Works even when api.global_api_key is absent (existing installs)."""
# instance_config without the global_api_key field at all
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
mock_result = Mock()
mock_result.first = Mock(return_value=None)
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.instance_config = SimpleNamespace(data={'api': {}})
async def test_verify_api_key_missing_global_config_key(self, api_key_context):
application, service, _context, _engine = api_key_context
application.instance_config.data = {'api': {}}
service = ApiKeyService(ap)
result = await service.verify_api_key('lbk_some_key')
assert result is False
assert await service.verify_api_key('lbk_some_key') is False
class TestApiKeyServiceDeleteApiKey:
"""Tests for delete_api_key method."""
async def test_delete_api_key_by_id(self, api_key_context):
_application, service, context, _engine = api_key_context
created = await service.create_api_key(context, 'Delete me')
async def test_delete_api_key_by_id(self):
"""Deletes API key by ID."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.persistence_mgr.execute_async = AsyncMock()
await service.delete_api_key(context, created['id'])
service = ApiKeyService(ap)
stored = await service.get_api_key(context, created['id'])
assert stored is not None
assert stored['status'] == 'revoked'
assert await service.verify_api_key(created['key']) is False
# Execute
await service.delete_api_key(1)
async def test_delete_api_key_nonexistent_id(self, api_key_context):
_application, service, context, _engine = api_key_context
# Verify - execute_async was called (delete operation)
ap.persistence_mgr.execute_async.assert_called_once()
async def test_delete_api_key_nonexistent_id(self):
"""Delete operation completes even for nonexistent ID (no error raised)."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.persistence_mgr.execute_async = AsyncMock()
service = ApiKeyService(ap)
# Execute - should not raise error
await service.delete_api_key(999)
# Verify - execute_async was called regardless
ap.persistence_mgr.execute_async.assert_called_once()
with pytest.raises(WorkspaceNotFoundError, match='API key not found'):
await service.delete_api_key(context, 999)
class TestApiKeyServiceUpdateApiKey:
"""Tests for update_api_key method."""
async def test_update_api_key_name_only(self, api_key_context):
_application, service, context, _engine = api_key_context
created = await service.create_api_key(context, 'Original', 'Description')
async def test_update_api_key_name_only(self):
"""Updates only the name field."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.persistence_mgr.execute_async = AsyncMock()
await service.update_api_key(context, created['id'], name='Updated Name')
service = ApiKeyService(ap)
stored = await service.get_api_key(context, created['id'])
assert stored is not None
assert stored['name'] == 'Updated Name'
assert stored['description'] == 'Description'
# Execute
await service.update_api_key(1, name='Updated Name')
async def test_update_api_key_description_only(self, api_key_context):
_application, service, context, _engine = api_key_context
created = await service.create_api_key(context, 'Original', 'Description')
# Verify - execute_async was called with update
ap.persistence_mgr.execute_async.assert_called_once()
await service.update_api_key(context, created['id'], description='Updated description')
async def test_update_api_key_description_only(self):
"""Updates only the description field."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.persistence_mgr.execute_async = AsyncMock()
stored = await service.get_api_key(context, created['id'])
assert stored is not None
assert stored['name'] == 'Original'
assert stored['description'] == 'Updated description'
service = ApiKeyService(ap)
async def test_update_api_key_both_fields(self, api_key_context):
_application, service, context, _engine = api_key_context
created = await service.create_api_key(context, 'Original', 'Description')
# Execute
await service.update_api_key(1, description='Updated description')
await service.update_api_key(
context,
created['id'],
name='New Name',
description='New description',
)
# Verify
ap.persistence_mgr.execute_async.assert_called_once()
stored = await service.get_api_key(context, created['id'])
assert stored is not None
assert stored['name'] == 'New Name'
assert stored['description'] == 'New description'
async def test_update_api_key_both_fields(self):
"""Updates both name and description."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.persistence_mgr.execute_async = AsyncMock()
async def test_update_api_key_no_fields(self, api_key_context):
application, service, context, _engine = api_key_context
created = await service.create_api_key(context, 'Original')
original_execute = application.persistence_mgr.execute_async
application.persistence_mgr.execute_async = AsyncMock(wraps=original_execute)
service = ApiKeyService(ap)
await service.update_api_key(context, created['id'])
# Execute
await service.update_api_key(1, name='New Name', description='New description')
# Verify
ap.persistence_mgr.execute_async.assert_called_once()
async def test_update_api_key_no_fields(self):
"""Does nothing when no fields provided."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.persistence_mgr.execute_async = AsyncMock()
service = ApiKeyService(ap)
# Execute
await service.update_api_key(1)
# Verify - no execute call since no update_data
ap.persistence_mgr.execute_async.assert_not_called()
application.persistence_mgr.execute_async.assert_not_awaited()
tests/unit_tests/api/service/test_bot_service.py
+43 -26
View File
@@ -19,6 +19,8 @@ from langbot.pkg.entity.persistence.bot import Bot
pytestmark = pytest.mark.asyncio
WORKSPACE_UUID = 'workspace-a'
def _create_mock_bot(
bot_uuid: str = None,
@@ -73,7 +75,9 @@ class TestBotServiceGetBots:
service = BotService(ap)
# Execute
result = await service.get_bots()
result = await service.get_bots(
WORKSPACE_UUID,
)
# Verify
assert result == []
@@ -101,7 +105,7 @@ class TestBotServiceGetBots:
service = BotService(ap)
# Execute
result = await service.get_bots(include_secret=True)
result = await service.get_bots(WORKSPACE_UUID, include_secret=True)
# Verify
assert len(result) == 2
@@ -130,7 +134,7 @@ class TestBotServiceGetBots:
service = BotService(ap)
# Execute
result = await service.get_bots(include_secret=False)
result = await service.get_bots(WORKSPACE_UUID, include_secret=False)
# Verify - adapter_config should be masked
assert result[0]['adapter_config'] is None
@@ -159,7 +163,7 @@ class TestBotServiceGetBot:
service = BotService(ap)
# Execute
result = await service.get_bot('test-uuid')
result = await service.get_bot(WORKSPACE_UUID, 'test-uuid')
# Verify
assert result is not None
@@ -178,7 +182,7 @@ class TestBotServiceGetBot:
service = BotService(ap)
# Execute
result = await service.get_bot('nonexistent-uuid')
result = await service.get_bot(WORKSPACE_UUID, 'nonexistent-uuid')
# Verify
assert result is None
@@ -203,7 +207,7 @@ class TestBotServiceGetRuntimeBotInfo:
# Execute & Verify
with pytest.raises(Exception, match='Bot not found'):
await service.get_runtime_bot_info('nonexistent-uuid')
await service.get_runtime_bot_info(WORKSPACE_UUID, 'nonexistent-uuid')
async def test_get_runtime_bot_info_returns_webhook_for_wecom(self):
"""Returns webhook URL for wecom adapter."""
@@ -231,7 +235,7 @@ class TestBotServiceGetRuntimeBotInfo:
service.get_bot = AsyncMock(return_value=bot_data)
# Execute
result = await service.get_runtime_bot_info('wecom-uuid')
result = await service.get_runtime_bot_info(WORKSPACE_UUID, 'wecom-uuid')
# Verify
assert result['adapter_runtime_values']['webhook_url'] == '/bots/wecom-uuid'
@@ -257,7 +261,7 @@ class TestBotServiceGetRuntimeBotInfo:
service.get_bot = AsyncMock(return_value=bot_data)
# Execute
result = await service.get_runtime_bot_info('telegram-uuid')
result = await service.get_runtime_bot_info(WORKSPACE_UUID, 'telegram-uuid')
# Verify - no webhook for telegram
assert result['adapter_runtime_values']['webhook_url'] is None
@@ -288,7 +292,7 @@ class TestBotServiceGetRuntimeBotInfo:
service.get_bot = AsyncMock(return_value=bot_data)
# Execute
result = await service.get_runtime_bot_info('runtime-uuid')
result = await service.get_runtime_bot_info(WORKSPACE_UUID, 'runtime-uuid')
# Verify
assert result['adapter_runtime_values']['bot_account_id'] == 'runtime-account-123'
@@ -318,7 +322,7 @@ class TestBotServiceCreateBot:
# Execute & Verify
with pytest.raises(ValueError, match='Maximum number of bots'):
await service.create_bot({'name': 'New Bot'})
await service.create_bot(WORKSPACE_UUID, {'name': 'New Bot'})
async def test_create_bot_no_limit(self):
"""Creates bot without limit check when max_bots=-1."""
@@ -360,7 +364,9 @@ class TestBotServiceCreateBot:
service = BotService(ap)
# Execute
bot_uuid = await service.create_bot({'name': 'New Bot', 'adapter': 'telegram', 'adapter_config': {}})
bot_uuid = await service.create_bot(
WORKSPACE_UUID, {'name': 'New Bot', 'adapter': 'telegram', 'adapter_config': {}}
)
# Verify
assert bot_uuid is not None
@@ -412,11 +418,15 @@ class TestBotServiceCreateBot:
# Execute
bot_data = {'name': 'New Bot', 'adapter': 'telegram', 'adapter_config': {}}
bot_uuid = await service.create_bot(bot_data)
bot_uuid = await service.create_bot(WORKSPACE_UUID, bot_data)
# Verify - pipeline uuid and name were set
assert 'use_pipeline_uuid' in bot_data
assert 'use_pipeline_name' in bot_data
# The service owns a copy and cannot mutate caller input while adding tenant data.
assert bot_data == {'name': 'New Bot', 'adapter': 'telegram', 'adapter_config': {}}
insert_statement = ap.persistence_mgr.execute_async.await_args_list[1].args[0]
insert_values = insert_statement.compile().params
assert insert_values['workspace_uuid'] == WORKSPACE_UUID
assert insert_values['use_pipeline_uuid'] == 'default-pipeline-uuid'
assert insert_values['use_pipeline_name'] == 'Default Pipeline'
assert bot_uuid is not None # Verify UUID was returned
@@ -446,7 +456,7 @@ class TestBotServiceUpdateBot:
# Execute
update_data = {'uuid': 'should-be-removed', 'name': 'Updated Name'}
await service.update_bot('test-uuid', update_data)
await service.update_bot(WORKSPACE_UUID, 'test-uuid', update_data)
update_params = ap.persistence_mgr.execute_async.await_args_list[0].args[0].compile().params
assert update_params['name'] == 'Updated Name'
@@ -467,7 +477,7 @@ class TestBotServiceUpdateBot:
# Execute & Verify
with pytest.raises(Exception, match='Pipeline not found'):
await service.update_bot('test-uuid', {'use_pipeline_uuid': 'nonexistent-pipeline'})
await service.update_bot(WORKSPACE_UUID, 'test-uuid', {'use_pipeline_uuid': 'nonexistent-pipeline'})
async def test_update_bot_sets_pipeline_name(self):
"""Sets use_pipeline_name when updating use_pipeline_uuid."""
@@ -504,7 +514,7 @@ class TestBotServiceUpdateBot:
ap.platform_mgr.load_bot = AsyncMock(return_value=runtime_bot)
# Execute
await service.update_bot('test-uuid', {'use_pipeline_uuid': 'pipeline-uuid'})
await service.update_bot(WORKSPACE_UUID, 'test-uuid', {'use_pipeline_uuid': 'pipeline-uuid'})
update_params = ap.persistence_mgr.execute_async.await_args_list[1].args[0].compile().params
assert update_params['use_pipeline_uuid'] == 'pipeline-uuid'
@@ -524,12 +534,13 @@ class TestBotServiceDeleteBot:
ap.platform_mgr.remove_bot = AsyncMock()
service = BotService(ap)
service.get_bot = AsyncMock(return_value={'uuid': 'bot-uuid'})
# Execute
await service.delete_bot('test-uuid')
await service.delete_bot(WORKSPACE_UUID, 'test-uuid')
# Verify
ap.platform_mgr.remove_bot.assert_called_once_with('test-uuid')
ap.platform_mgr.remove_bot.assert_called_once_with(WORKSPACE_UUID, 'test-uuid')
ap.persistence_mgr.execute_async.assert_called_once()
async def test_delete_bot_nonexistent_uuid(self):
@@ -542,9 +553,10 @@ class TestBotServiceDeleteBot:
ap.platform_mgr.remove_bot = AsyncMock()
service = BotService(ap)
service.get_bot = AsyncMock(return_value={'uuid': 'bot-uuid'})
# Execute - should not raise
await service.delete_bot('nonexistent-uuid')
await service.delete_bot(WORKSPACE_UUID, 'nonexistent-uuid')
# Verify - both called regardless
ap.platform_mgr.remove_bot.assert_called_once()
@@ -561,10 +573,11 @@ class TestBotServiceListEventLogs:
ap.platform_mgr.get_bot_by_uuid = AsyncMock(return_value=None)
service = BotService(ap)
service.get_bot = AsyncMock(return_value={'uuid': 'nonexistent-uuid'})
# Execute & Verify
with pytest.raises(Exception, match='Bot not found'):
await service.list_event_logs('nonexistent-uuid', 0, 10)
await service.list_event_logs(WORKSPACE_UUID, 'nonexistent-uuid', 0, 10)
async def test_list_event_logs_returns_logs(self):
"""Returns logs from runtime bot logger."""
@@ -581,9 +594,10 @@ class TestBotServiceListEventLogs:
ap.platform_mgr.get_bot_by_uuid = AsyncMock(return_value=runtime_bot)
service = BotService(ap)
service.get_bot = AsyncMock(return_value={'uuid': 'bot-uuid'})
# Execute
logs, total = await service.list_event_logs('bot-uuid', 0, 10)
logs, total = await service.list_event_logs(WORKSPACE_UUID, 'bot-uuid', 0, 10)
# Verify
assert len(logs) == 1
@@ -602,10 +616,11 @@ class TestBotServiceSendMessage:
ap.platform_mgr.get_bot_by_uuid = AsyncMock(return_value=None)
service = BotService(ap)
service.get_bot = AsyncMock(return_value={'uuid': 'nonexistent-uuid'})
# Execute & Verify
with pytest.raises(Exception, match='Bot not found'):
await service.send_message('nonexistent-uuid', 'group', '123', {'test': 'data'})
await service.send_message(WORKSPACE_UUID, 'nonexistent-uuid', 'group', '123', {'test': 'data'})
async def test_send_message_invalid_message_chain_raises(self):
"""Raises Exception when message_chain_data is invalid."""
@@ -619,10 +634,11 @@ class TestBotServiceSendMessage:
ap.platform_mgr.get_bot_by_uuid = AsyncMock(return_value=runtime_bot)
service = BotService(ap)
service.get_bot = AsyncMock(return_value={'uuid': 'bot-uuid'})
# Execute & Verify - invalid format should raise
with pytest.raises(Exception, match='Invalid message_chain format'):
await service.send_message('bot-uuid', 'group', '123', {'invalid': 'format'})
await service.send_message(WORKSPACE_UUID, 'bot-uuid', 'group', '123', {'invalid': 'format'})
async def test_send_message_valid_call(self):
"""Sends message through adapter when all valid."""
@@ -636,6 +652,7 @@ class TestBotServiceSendMessage:
ap.platform_mgr.get_bot_by_uuid = AsyncMock(return_value=runtime_bot)
service = BotService(ap)
service.get_bot = AsyncMock(return_value={'uuid': 'bot-uuid'})
# Execute with valid message chain format
message_chain_data = {'messages': [{'type': 'text', 'data': {'text': 'Hello'}}]}
@@ -644,7 +661,7 @@ class TestBotServiceSendMessage:
with patch('langbot_plugin.api.entities.builtin.platform.message.MessageChain') as MockMessageChain:
mock_chain = Mock()
MockMessageChain.model_validate = Mock(return_value=mock_chain)
await service.send_message('bot-uuid', 'group', '123', message_chain_data)
await service.send_message(WORKSPACE_UUID, 'bot-uuid', 'group', '123', message_chain_data)
# Verify adapter.send_message was called
runtime_bot.adapter.send_message.assert_called_once_with('group', '123', mock_chain)
tests/unit_tests/api/service/test_knowledge_service.py
+437 -245
View File
@@ -1,389 +1,581 @@
"""Unit tests for API knowledge service.
Tests cover:
- Knowledge base CRUD operations
- Capability checking
- Knowledge engine discovery
- File operations
"""
"""Tests for the tenant-aware knowledge service facade."""
from __future__ import annotations
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
import pytest
from unittest.mock import Mock, AsyncMock
from importlib import import_module
from langbot.pkg.api.http.authz import WorkspaceRequiredError
from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.api.http.service.knowledge import KnowledgeService
from langbot.pkg.workspace.errors import WorkspaceNotFoundError
def get_knowledge_service_module():
"""Lazy import to avoid circular import issues."""
return import_module('langbot.pkg.api.http.service.knowledge')
CONTEXT = ExecutionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=2,
)
def create_mock_app():
"""Create mock Application for testing."""
mock_app = Mock()
mock_app.logger = Mock()
mock_app.rag_mgr = AsyncMock()
mock_app.persistence_mgr = AsyncMock()
mock_app.persistence_mgr.execute_async = AsyncMock()
mock_app.persistence_mgr.serialize_model = Mock(return_value={})
mock_app.plugin_connector = AsyncMock()
mock_app.plugin_connector.is_enable_plugin = True
return mock_app
class _Rows:
def __init__(self, rows=()):
self.rows = list(rows)
def all(self):
return self.rows
def __iter__(self):
return iter(self.rows)
def _app():
return SimpleNamespace(
logger=Mock(),
instance_config=SimpleNamespace(data={}),
rag_mgr=SimpleNamespace(
get_all_knowledge_base_details=AsyncMock(return_value=[]),
get_knowledge_base_details=AsyncMock(return_value=None),
create_knowledge_base=AsyncMock(),
remove_knowledge_base_from_runtime=AsyncMock(),
load_knowledge_base=AsyncMock(),
get_knowledge_base_by_uuid=AsyncMock(return_value=None),
delete_knowledge_base=AsyncMock(),
),
persistence_mgr=SimpleNamespace(
execute_async=AsyncMock(return_value=_Rows()),
serialize_model=Mock(return_value={}),
),
plugin_connector=SimpleNamespace(
is_enable_plugin=True,
require_workspace_context=AsyncMock(side_effect=lambda context: context),
get_rag_creation_schema=AsyncMock(return_value={}),
get_rag_retrieval_schema=AsyncMock(return_value={}),
list_knowledge_engines=AsyncMock(return_value=[]),
list_parsers=AsyncMock(return_value=[]),
),
)
@pytest.mark.asyncio
async def test_list_and_get_forward_explicit_context():
app = _app()
app.rag_mgr.get_all_knowledge_base_details.return_value = [{'uuid': 'kb-a'}]
app.rag_mgr.get_knowledge_base_details.return_value = {'uuid': 'kb-a'}
service = KnowledgeService(app)
assert await service.get_knowledge_bases(CONTEXT) == [{'uuid': 'kb-a'}]
assert await service.get_knowledge_base(CONTEXT, 'kb-a') == {'uuid': 'kb-a'}
app.rag_mgr.get_all_knowledge_base_details.assert_awaited_once_with(CONTEXT)
app.rag_mgr.get_knowledge_base_details.assert_awaited_once_with(CONTEXT, 'kb-a')
@pytest.mark.asyncio
async def test_none_context_fails_closed_before_plugin_or_manager_access():
app = _app()
service = KnowledgeService(app)
with pytest.raises(WorkspaceRequiredError):
await service.get_knowledge_bases(None)
with pytest.raises(WorkspaceRequiredError):
await service.create_knowledge_base(None, {'knowledge_engine_plugin_id': 'author/engine'})
app.plugin_connector.get_rag_creation_schema.assert_not_awaited()
app.rag_mgr.create_knowledge_base.assert_not_awaited()
@pytest.mark.asyncio
async def test_create_validates_schema_and_binds_context():
app = _app()
app.plugin_connector.get_rag_creation_schema.return_value = {
'schema': [{'name': 'endpoint', 'label': {'en_US': 'Endpoint'}, 'required': True}]
}
app.rag_mgr.create_knowledge_base.return_value = SimpleNamespace(uuid='kb-created')
service = KnowledgeService(app)
with pytest.raises(ValueError, match='Endpoint is required'):
await service.create_knowledge_base(
CONTEXT,
{'knowledge_engine_plugin_id': 'author/engine'},
)
result = await service.create_knowledge_base(
CONTEXT,
{
'name': 'KB',
'description': 'desc',
'knowledge_engine_plugin_id': 'author/engine',
'creation_settings': {'endpoint': 'https://example.invalid'},
},
)
assert result == 'kb-created'
app.rag_mgr.create_knowledge_base.assert_awaited_once_with(
CONTEXT,
name='KB',
knowledge_engine_plugin_id='author/engine',
creation_settings={'endpoint': 'https://example.invalid'},
retrieval_settings={},
description='desc',
)
@pytest.mark.asyncio
async def test_create_enforces_workspace_knowledge_base_limit():
app = _app()
app.instance_config.data = {'system': {'limitation': {'max_knowledge_bases': 2}}}
app.rag_mgr.get_all_knowledge_base_details.return_value = [{'uuid': 'kb-a'}, {'uuid': 'kb-b'}]
service = KnowledgeService(app)
with pytest.raises(ValueError, match=r'Maximum number of knowledge bases \(2\) reached'):
await service.create_knowledge_base(
CONTEXT,
{'knowledge_engine_plugin_id': 'author/engine'},
)
app.rag_mgr.create_knowledge_base.assert_not_awaited()
@pytest.mark.asyncio
async def test_update_rejects_guessed_uuid_and_scopes_reload():
app = _app()
service = KnowledgeService(app)
with pytest.raises(WorkspaceNotFoundError):
await service.update_knowledge_base(CONTEXT, 'kb-other', {'name': 'stolen'})
app.persistence_mgr.execute_async.assert_not_awaited()
app.rag_mgr.get_knowledge_base_details.return_value = {'uuid': 'kb-a', 'workspace_uuid': 'workspace-a'}
await service.update_knowledge_base(CONTEXT, 'kb-a', {'name': 'updated', 'uuid': 'ignored'})
app.rag_mgr.remove_knowledge_base_from_runtime.assert_awaited_once_with(CONTEXT, 'kb-a')
app.rag_mgr.load_knowledge_base.assert_awaited_once_with(
CONTEXT,
{'uuid': 'kb-a', 'workspace_uuid': 'workspace-a'},
)
@pytest.mark.asyncio
async def test_runtime_retrieve_uses_execution_context():
app = _app()
entry = SimpleNamespace(model_dump=Mock(return_value={'id': 'entry-a'}))
runtime_kb = SimpleNamespace(retrieve=AsyncMock(return_value=[entry]))
app.rag_mgr.get_knowledge_base_by_uuid.return_value = runtime_kb
service = KnowledgeService(app)
assert await service.retrieve_knowledge_base(CONTEXT, 'kb-a', 'query', {'top_k': 3}) == [{'id': 'entry-a'}]
runtime_kb.retrieve.assert_awaited_once_with(CONTEXT, 'query', settings={'top_k': 3})
@pytest.mark.asyncio
async def test_runtime_retrieve_cross_workspace_uuid_is_not_found():
app = _app()
service = KnowledgeService(app)
with pytest.raises(WorkspaceNotFoundError):
await service.retrieve_knowledge_base(CONTEXT, 'kb-other', 'query')
@pytest.mark.asyncio
async def test_file_listing_checks_parent_knowledge_base_first():
app = _app()
service = KnowledgeService(app)
with pytest.raises(WorkspaceNotFoundError):
await service.get_files_by_knowledge_base(CONTEXT, 'kb-other')
app.persistence_mgr.execute_async.assert_not_awaited()
app.rag_mgr.get_knowledge_base_details.return_value = {'uuid': 'kb-a'}
row = SimpleNamespace(uuid='file-a')
app.persistence_mgr.execute_async.return_value = _Rows([row])
app.persistence_mgr.serialize_model.return_value = {'uuid': 'file-a'}
assert await service.get_files_by_knowledge_base(CONTEXT, 'kb-a') == [{'uuid': 'file-a'}]
@pytest.mark.asyncio
async def test_store_and_delete_file_require_runtime_parent_and_capability():
app = _app()
runtime_kb = SimpleNamespace(
store_file=AsyncMock(return_value='task-a'),
delete_file=AsyncMock(),
)
app.rag_mgr.get_knowledge_base_by_uuid.return_value = runtime_kb
app.rag_mgr.get_knowledge_base_details.return_value = {'knowledge_engine': {'capabilities': ['doc_ingestion']}}
service = KnowledgeService(app)
assert await service.store_file(CONTEXT, 'kb-a', 'upload.pdf', 'author/parser') == 'task-a'
runtime_kb.store_file.assert_awaited_once_with(CONTEXT, 'upload.pdf', parser_plugin_id='author/parser')
await service.delete_file(CONTEXT, 'kb-a', 'file-a')
runtime_kb.delete_file.assert_awaited_once_with(CONTEXT, 'file-a')
@pytest.mark.asyncio
async def test_delete_knowledge_base_rejects_cross_workspace_uuid():
app = _app()
service = KnowledgeService(app)
with pytest.raises(WorkspaceNotFoundError):
await service.delete_knowledge_base(CONTEXT, 'kb-other')
app.rag_mgr.delete_knowledge_base.assert_not_awaited()
@pytest.mark.asyncio
async def test_engine_and_parser_discovery_require_context_and_filter_results():
app = _app()
app.plugin_connector.list_knowledge_engines.return_value = [{'plugin_id': 'author/engine'}]
app.plugin_connector.list_parsers.return_value = [
{'id': 'text', 'supported_mime_types': ['text/plain']},
{'id': 'pdf', 'supported_mime_types': ['application/pdf']},
]
service = KnowledgeService(app)
assert await service.list_knowledge_engines(CONTEXT) == [{'plugin_id': 'author/engine'}]
assert await service.list_parsers(CONTEXT, 'application/pdf') == [
{'id': 'pdf', 'supported_mime_types': ['application/pdf']}
]
with pytest.raises(WorkspaceRequiredError):
await service.list_parsers(None)
@pytest.mark.asyncio
async def test_engine_discovery_rejects_connector_workspace_or_generation_mismatch():
app = _app()
app.plugin_connector.require_workspace_context.side_effect = WorkspaceNotFoundError('Plugin resource not found')
service = KnowledgeService(app)
with pytest.raises(WorkspaceNotFoundError, match='Plugin resource not found'):
await service.list_knowledge_engines(CONTEXT)
app.plugin_connector.list_knowledge_engines.assert_not_awaited()
@pytest.mark.asyncio
async def test_schema_validation_refences_before_second_runtime_call():
app = _app()
app.plugin_connector.require_workspace_context.side_effect = [
CONTEXT,
WorkspaceNotFoundError('Plugin resource not found'),
]
service = KnowledgeService(app)
with pytest.raises(WorkspaceNotFoundError, match='Plugin resource not found'):
await service.create_knowledge_base(
CONTEXT,
{'knowledge_engine_plugin_id': 'author/engine'},
)
app.plugin_connector.get_rag_creation_schema.assert_awaited_once()
app.plugin_connector.get_rag_retrieval_schema.assert_not_awaited()
app.rag_mgr.create_knowledge_base.assert_not_awaited()
@pytest.mark.asyncio
async def test_engine_schemas_are_context_gated_and_fail_soft_on_connector_error():
app = _app()
app.plugin_connector.get_rag_creation_schema.return_value = {'schema': ['creation']}
app.plugin_connector.get_rag_retrieval_schema.side_effect = RuntimeError('offline')
service = KnowledgeService(app)
assert await service.get_engine_creation_schema(CONTEXT, 'author/engine') == {'schema': ['creation']}
assert await service.get_engine_retrieval_schema(CONTEXT, 'author/engine') == {}
with pytest.raises(WorkspaceRequiredError):
await service.get_engine_creation_schema(None, 'author/engine')
# Preserve the original service regression matrix with the new explicit
# Workspace context. These intentionally overlap a few isolation-focused
# tests above so legacy business behavior cannot disappear behind new guards.
class TestKnowledgeServiceInit:
"""Tests for KnowledgeService initialization."""
def test_init_stores_app_reference(self):
"""Test that __init__ stores Application reference."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
app = _app()
service = knowledge_module.KnowledgeService(mock_app)
service = KnowledgeService(app)
assert service.ap is mock_app
assert service.ap is app
class TestGetKnowledgeBases:
"""Tests for get_knowledge_bases method."""
@pytest.mark.asyncio
async def test_returns_all_kb_details(self):
"""Test that it returns all knowledge base details."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_app.rag_mgr.get_all_knowledge_base_details = AsyncMock(return_value=[{'uuid': 'kb1', 'name': 'KB1'}])
app = _app()
app.rag_mgr.get_all_knowledge_base_details.return_value = [{'uuid': 'kb1', 'name': 'KB1'}]
service = knowledge_module.KnowledgeService(mock_app)
result = await service.get_knowledge_bases()
result = await KnowledgeService(app).get_knowledge_bases(CONTEXT)
assert len(result) == 1
assert result[0]['uuid'] == 'kb1'
assert result == [{'uuid': 'kb1', 'name': 'KB1'}]
app.rag_mgr.get_all_knowledge_base_details.assert_awaited_once_with(CONTEXT)
@pytest.mark.asyncio
async def test_returns_empty_list_when_no_kbs(self):
"""Test that it returns empty list when no knowledge bases."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_app.rag_mgr.get_all_knowledge_base_details = AsyncMock(return_value=[])
app = _app()
service = knowledge_module.KnowledgeService(mock_app)
result = await service.get_knowledge_bases()
assert result == []
assert await KnowledgeService(app).get_knowledge_bases(CONTEXT) == []
class TestGetKnowledgeBase:
"""Tests for get_knowledge_base method."""
@pytest.mark.asyncio
async def test_returns_kb_details_by_uuid(self):
"""Test that it returns specific KB details."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_app.rag_mgr.get_knowledge_base_details = AsyncMock(return_value={'uuid': 'kb1', 'name': 'KB1'})
app = _app()
app.rag_mgr.get_knowledge_base_details.return_value = {'uuid': 'kb1', 'name': 'KB1'}
service = knowledge_module.KnowledgeService(mock_app)
result = await service.get_knowledge_base('kb1')
result = await KnowledgeService(app).get_knowledge_base(CONTEXT, 'kb1')
assert result['uuid'] == 'kb1'
assert result == {'uuid': 'kb1', 'name': 'KB1'}
@pytest.mark.asyncio
async def test_returns_none_when_not_found(self):
"""Test that it returns None when KB not found."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_app.rag_mgr.get_knowledge_base_details = AsyncMock(return_value=None)
app = _app()
service = knowledge_module.KnowledgeService(mock_app)
result = await service.get_knowledge_base('nonexistent')
assert result is None
assert await KnowledgeService(app).get_knowledge_base(CONTEXT, 'nonexistent') is None
class TestCreateKnowledgeBase:
"""Tests for create_knowledge_base method."""
@pytest.mark.asyncio
async def test_creates_kb_with_required_fields(self):
"""Test creating KB with required plugin ID."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_kb = Mock()
mock_kb.uuid = 'new_kb_uuid'
mock_app.rag_mgr.create_knowledge_base = AsyncMock(return_value=mock_kb)
service = knowledge_module.KnowledgeService(mock_app)
app = _app()
app.rag_mgr.create_knowledge_base.return_value = SimpleNamespace(uuid='new_kb_uuid')
service = KnowledgeService(app)
kb_data = {
'name': 'Test KB',
'knowledge_engine_plugin_id': 'author/engine',
'description': 'Test description',
}
result = await service.create_knowledge_base(kb_data)
result = await service.create_knowledge_base(CONTEXT, kb_data)
assert result == 'new_kb_uuid'
mock_app.rag_mgr.create_knowledge_base.assert_called_once()
app.rag_mgr.create_knowledge_base.assert_awaited_once_with(
CONTEXT,
name='Test KB',
knowledge_engine_plugin_id='author/engine',
creation_settings={},
retrieval_settings={},
description='Test description',
)
@pytest.mark.asyncio
async def test_raises_when_missing_plugin_id(self):
"""Test that ValueError is raised when plugin ID missing."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
app = _app()
service = knowledge_module.KnowledgeService(mock_app)
with pytest.raises(ValueError, match='knowledge_engine_plugin_id is required'):
await KnowledgeService(app).create_knowledge_base(CONTEXT, {'name': 'Test'})
with pytest.raises(ValueError) as exc_info:
await service.create_knowledge_base({'name': 'Test'})
assert 'knowledge_engine_plugin_id is required' in str(exc_info.value)
app.rag_mgr.create_knowledge_base.assert_not_awaited()
@pytest.mark.asyncio
async def test_creates_with_default_name(self):
"""Test that KB is created with default name if not provided."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_kb = Mock()
mock_kb.uuid = 'new_kb_uuid'
mock_app.rag_mgr.create_knowledge_base = AsyncMock(return_value=mock_kb)
app = _app()
app.rag_mgr.create_knowledge_base.return_value = SimpleNamespace(uuid='new_kb_uuid')
service = knowledge_module.KnowledgeService(mock_app)
await KnowledgeService(app).create_knowledge_base(
CONTEXT,
{'knowledge_engine_plugin_id': 'author/engine'},
)
await service.create_knowledge_base({'knowledge_engine_plugin_id': 'author/engine'})
# Check that default name 'Untitled' was used
call_args = mock_app.rag_mgr.create_knowledge_base.call_args
assert call_args.kwargs['name'] == 'Untitled'
assert app.rag_mgr.create_knowledge_base.await_args.kwargs['name'] == 'Untitled'
class TestUpdateKnowledgeBase:
"""Tests for update_knowledge_base method."""
@pytest.mark.asyncio
async def test_updates_mutable_fields_only(self):
"""Test that only mutable fields are updated."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_app.rag_mgr.get_knowledge_base_details = AsyncMock(return_value={'uuid': 'kb1', 'name': 'Updated'})
mock_app.rag_mgr.remove_knowledge_base_from_runtime = AsyncMock()
mock_app.rag_mgr.load_knowledge_base = AsyncMock()
app = _app()
app.rag_mgr.get_knowledge_base_details.return_value = {'uuid': 'kb1', 'name': 'Updated'}
service = KnowledgeService(app)
service = knowledge_module.KnowledgeService(mock_app)
# Pass both mutable and immutable fields
await service.update_knowledge_base(
CONTEXT,
'kb1',
{
'name': 'New Name',
'description': 'New desc',
'uuid': 'should_be_filtered', # immutable
'uuid': 'should_be_filtered',
},
)
# Check that only mutable fields were passed to update
call_args = mock_app.persistence_mgr.execute_async.call_args
assert call_args is not None
update_statement = app.persistence_mgr.execute_async.await_args_list[0].args[0]
params = update_statement.compile().params
assert params['name'] == 'New Name'
assert params['description'] == 'New desc'
assert 'uuid' not in params
app.rag_mgr.remove_knowledge_base_from_runtime.assert_awaited_once_with(CONTEXT, 'kb1')
@pytest.mark.asyncio
async def test_returns_early_when_no_mutable_fields(self):
"""Test that update returns early when no mutable fields provided."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
app = _app()
app.rag_mgr.get_knowledge_base_details.return_value = {'uuid': 'kb1'}
service = knowledge_module.KnowledgeService(mock_app)
await KnowledgeService(app).update_knowledge_base(
CONTEXT,
'kb1',
{'uuid': 'should_be_filtered'},
)
# Pass only immutable fields
await service.update_knowledge_base('kb1', {'uuid': 'should_be_filtered'})
# No DB update should be called
mock_app.persistence_mgr.execute_async.assert_not_called()
app.persistence_mgr.execute_async.assert_not_awaited()
app.rag_mgr.remove_knowledge_base_from_runtime.assert_not_awaited()
class TestCheckDocCapability:
"""Tests for _check_doc_capability method."""
@pytest.mark.asyncio
async def test_passes_when_capability_supported(self):
"""Test that check passes when doc_ingestion capability exists."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_app.rag_mgr.get_knowledge_base_details = AsyncMock(
return_value={'knowledge_engine': {'capabilities': ['doc_ingestion']}}
)
app = _app()
app.rag_mgr.get_knowledge_base_details.return_value = {'knowledge_engine': {'capabilities': ['doc_ingestion']}}
service = knowledge_module.KnowledgeService(mock_app)
await service._check_doc_capability('kb1', 'document upload')
# No exception raised means success
await KnowledgeService(app)._check_doc_capability(CONTEXT, 'kb1', 'document upload')
@pytest.mark.asyncio
async def test_raises_when_kb_not_found(self):
"""Test that Exception is raised when KB not found."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_app.rag_mgr.get_knowledge_base_details = AsyncMock(return_value=None)
app = _app()
service = knowledge_module.KnowledgeService(mock_app)
with pytest.raises(Exception) as exc_info:
await service._check_doc_capability('nonexistent', 'test operation')
assert 'Knowledge base not found' in str(exc_info.value)
with pytest.raises(WorkspaceNotFoundError, match='Knowledge base not found'):
await KnowledgeService(app)._check_doc_capability(
CONTEXT,
'nonexistent',
'test operation',
)
@pytest.mark.asyncio
async def test_raises_when_capability_not_supported(self):
"""Test that Exception is raised when doc_ingestion not in capabilities."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_app.rag_mgr.get_knowledge_base_details = AsyncMock(
return_value={'knowledge_engine': {'capabilities': ['other_capability']}}
)
app = _app()
app.rag_mgr.get_knowledge_base_details.return_value = {
'knowledge_engine': {'capabilities': ['other_capability']}
}
service = knowledge_module.KnowledgeService(mock_app)
with pytest.raises(Exception) as exc_info:
await service._check_doc_capability('kb1', 'document upload')
assert 'does not support document upload' in str(exc_info.value)
with pytest.raises(Exception, match='does not support document upload'):
await KnowledgeService(app)._check_doc_capability(
CONTEXT,
'kb1',
'document upload',
)
class TestListKnowledgeEngines:
"""Tests for list_knowledge_engines method."""
@pytest.mark.asyncio
async def test_returns_engines_from_plugin_connector(self):
"""Test that it returns knowledge engines from plugin connector."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_app.plugin_connector.list_knowledge_engines = AsyncMock(
return_value=[{'id': 'engine1', 'name': 'Engine 1'}]
)
app = _app()
app.plugin_connector.list_knowledge_engines.return_value = [{'id': 'engine1', 'name': 'Engine 1'}]
service = knowledge_module.KnowledgeService(mock_app)
result = await service.list_knowledge_engines()
result = await KnowledgeService(app).list_knowledge_engines(CONTEXT)
assert len(result) == 1
assert result[0]['id'] == 'engine1'
assert result == [{'id': 'engine1', 'name': 'Engine 1'}]
@pytest.mark.asyncio
async def test_returns_empty_when_plugin_disabled(self):
"""Test that it returns empty list when plugin disabled."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_app.plugin_connector.is_enable_plugin = False
app = _app()
app.plugin_connector.is_enable_plugin = False
service = knowledge_module.KnowledgeService(mock_app)
result = await service.list_knowledge_engines()
assert result == []
assert await KnowledgeService(app).list_knowledge_engines(CONTEXT) == []
app.plugin_connector.list_knowledge_engines.assert_not_awaited()
@pytest.mark.asyncio
async def test_returns_empty_on_exception(self):
"""Test that it returns empty list and logs warning on exception."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_app.plugin_connector.list_knowledge_engines = AsyncMock(side_effect=Exception('Connection error'))
app = _app()
app.plugin_connector.list_knowledge_engines.side_effect = RuntimeError('Connection error')
service = knowledge_module.KnowledgeService(mock_app)
result = await service.list_knowledge_engines()
assert result == []
mock_app.logger.warning.assert_called_once()
assert await KnowledgeService(app).list_knowledge_engines(CONTEXT) == []
app.logger.warning.assert_called_once()
class TestListParsers:
"""Tests for list_parsers method."""
@pytest.mark.asyncio
async def test_returns_all_parsers(self):
"""Test that it returns all parsers when no MIME type filter."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_app.plugin_connector.list_parsers = AsyncMock(
return_value=[
{'id': 'parser1', 'supported_mime_types': ['text/plain']},
{'id': 'parser2', 'supported_mime_types': ['application/pdf']},
]
)
app = _app()
app.plugin_connector.list_parsers.return_value = [
{'id': 'parser1', 'supported_mime_types': ['text/plain']},
{'id': 'parser2', 'supported_mime_types': ['application/pdf']},
]
service = knowledge_module.KnowledgeService(mock_app)
result = await service.list_parsers()
result = await KnowledgeService(app).list_parsers(CONTEXT)
assert len(result) == 2
@pytest.mark.asyncio
async def test_filters_by_mime_type(self):
"""Test that it filters parsers by MIME type."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_app.plugin_connector.list_parsers = AsyncMock(
return_value=[
{'id': 'parser1', 'supported_mime_types': ['text/plain']},
{'id': 'parser2', 'supported_mime_types': ['application/pdf']},
]
)
app = _app()
app.plugin_connector.list_parsers.return_value = [
{'id': 'parser1', 'supported_mime_types': ['text/plain']},
{'id': 'parser2', 'supported_mime_types': ['application/pdf']},
]
service = knowledge_module.KnowledgeService(mock_app)
result = await service.list_parsers(mime_type='application/pdf')
result = await KnowledgeService(app).list_parsers(CONTEXT, 'application/pdf')
assert len(result) == 1
assert result[0]['id'] == 'parser2'
assert result == [{'id': 'parser2', 'supported_mime_types': ['application/pdf']}]
@pytest.mark.asyncio
async def test_returns_empty_when_plugin_disabled(self):
"""Test that it returns empty list when plugin disabled."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_app.plugin_connector.is_enable_plugin = False
app = _app()
app.plugin_connector.is_enable_plugin = False
service = knowledge_module.KnowledgeService(mock_app)
result = await service.list_parsers()
assert result == []
assert await KnowledgeService(app).list_parsers(CONTEXT) == []
app.plugin_connector.list_parsers.assert_not_awaited()
class TestGetEngineSchemas:
"""Tests for get_engine_creation_schema and get_engine_retrieval_schema."""
@pytest.mark.asyncio
async def test_returns_creation_schema(self):
"""Test that it returns creation schema for engine."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_app.plugin_connector.get_rag_creation_schema = AsyncMock(
return_value={'properties': {'name': {'type': 'string'}}}
)
app = _app()
app.plugin_connector.get_rag_creation_schema.return_value = {'properties': {'name': {'type': 'string'}}}
service = knowledge_module.KnowledgeService(mock_app)
result = await service.get_engine_creation_schema('author/engine')
result = await KnowledgeService(app).get_engine_creation_schema(
CONTEXT,
'author/engine',
)
assert 'properties' in result
@pytest.mark.asyncio
async def test_returns_retrieval_schema(self):
"""Test that it returns retrieval schema for engine."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_app.plugin_connector.get_rag_retrieval_schema = AsyncMock(
return_value={'properties': {'top_k': {'type': 'integer'}}}
)
app = _app()
app.plugin_connector.get_rag_retrieval_schema.return_value = {'properties': {'top_k': {'type': 'integer'}}}
service = knowledge_module.KnowledgeService(mock_app)
result = await service.get_engine_retrieval_schema('author/engine')
result = await KnowledgeService(app).get_engine_retrieval_schema(
CONTEXT,
'author/engine',
)
assert 'properties' in result
@pytest.mark.asyncio
async def test_returns_empty_dict_on_exception(self):
"""Test that it returns empty dict and logs warning on exception."""
knowledge_module = get_knowledge_service_module()
mock_app = create_mock_app()
mock_app.plugin_connector.get_rag_creation_schema = AsyncMock(side_effect=Exception('Plugin error'))
app = _app()
app.plugin_connector.get_rag_creation_schema.side_effect = RuntimeError('Plugin error')
service = knowledge_module.KnowledgeService(mock_app)
result = await service.get_engine_creation_schema('author/engine')
result = await KnowledgeService(app).get_engine_creation_schema(
CONTEXT,
'author/engine',
)
assert result == {}
mock_app.logger.warning.assert_called_once()
app.logger.warning.assert_called_once()
class TestKnowledgeBaseSecretViews:
@pytest.mark.asyncio
async def test_creation_settings_are_redacted_for_resource_view_only(self):
app = _app()
raw = {
'uuid': 'kb-secret',
'creation_settings': {
'dify_apikey': 'dify-secret',
'headers': {'Authorization': 'Bearer secret'},
},
}
app.rag_mgr.get_all_knowledge_base_details.return_value = [raw]
service = KnowledgeService(app)
redacted = await service.get_knowledge_bases(CONTEXT)
manager_view = await service.get_knowledge_bases(CONTEXT, include_secret=True)
assert redacted[0]['creation_settings']['dify_apikey'] == '***'
assert redacted[0]['creation_settings']['headers']['Authorization'] == '***'
assert manager_view[0]['creation_settings']['dify_apikey'] == 'dify-secret'
assert raw['creation_settings']['dify_apikey'] == 'dify-secret'
@pytest.mark.asyncio
async def test_new_masked_creation_secret_is_rejected(self):
app = _app()
with pytest.raises(ValueError, match='no existing value'):
await KnowledgeService(app).create_knowledge_base(
CONTEXT,
{
'knowledge_engine_plugin_id': 'author/engine',
'creation_settings': {'dify_apikey': '***'},
},
)
app.rag_mgr.create_knowledge_base.assert_not_awaited()
tests/unit_tests/api/service/test_maintenance_service.py
+278 -30
View File
@@ -13,16 +13,39 @@ Source: src/langbot/pkg/api/http/service/maintenance.py
from __future__ import annotations
import contextlib
import pytest
from unittest.mock import AsyncMock, Mock, patch, MagicMock
from types import SimpleNamespace
import datetime
from pathlib import Path
import sqlalchemy
from sqlalchemy.ext.asyncio import create_async_engine
from langbot.pkg.api.http.authz import WorkspaceRequiredError
from langbot.pkg.api.http.service.maintenance import MaintenanceService
from langbot.pkg.api.http.context import ExecutionContext, PrincipalContext, PrincipalType
from langbot.pkg.entity.persistence.base import Base
from langbot.pkg.entity.persistence.bstorage import BinaryStorage
from langbot.pkg.entity.persistence.monitoring import MonitoringMessage
from langbot.pkg.entity.persistence.workspace import Workspace
pytestmark = pytest.mark.asyncio
TEST_CONTEXT = ExecutionContext(
instance_uuid='test-instance',
workspace_uuid='test-workspace',
placement_generation=1,
)
@pytest.fixture(autouse=True)
def assume_oss_singleton(monkeypatch):
async def is_oss_singleton(_self, _context):
return True
monkeypatch.setattr(MaintenanceService, '_is_oss_singleton', is_oss_singleton)
def _create_mock_result(scalar_value=None):
@@ -32,6 +55,14 @@ def _create_mock_result(scalar_value=None):
return result
def _scoped_storage_manager():
prefix = 'instances/i/workspaces/w/generations/1/owners/upload/o/'
return SimpleNamespace(
scoped_prefix=Mock(return_value=prefix),
is_scoped_object_key=Mock(side_effect=lambda key, **_: key == f'{prefix}uploaded_file.txt'),
)
class TestMaintenanceServiceCleanupExpiredFiles:
"""Tests for cleanup_expired_files method."""
@@ -39,6 +70,7 @@ class TestMaintenanceServiceCleanupExpiredFiles:
"""Uses default retention days when config not set."""
# Setup
ap = SimpleNamespace()
ap.storage_mgr = _scoped_storage_manager()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {}
ap.storage_mgr = SimpleNamespace()
@@ -58,7 +90,7 @@ class TestMaintenanceServiceCleanupExpiredFiles:
service._cleanup_expired_log_files = Mock(return_value=0) # NOT async!
# Execute
result = await service.cleanup_expired_files()
result = await service.cleanup_expired_files(TEST_CONTEXT)
# Verify - returns counts
assert 'uploaded_files' in result
@@ -95,7 +127,7 @@ class TestMaintenanceServiceCleanupExpiredFiles:
service._cleanup_expired_log_files = Mock(return_value=3) # NOT async
# Execute
result = await service.cleanup_expired_files()
result = await service.cleanup_expired_files(TEST_CONTEXT)
# Verify
assert result['uploaded_files'] == 2
@@ -124,7 +156,7 @@ class TestMaintenanceServiceCleanupExpiredFiles:
service._cleanup_expired_log_files = Mock(return_value=0) # NOT async
# Execute
result = await service.cleanup_expired_files()
result = await service.cleanup_expired_files(TEST_CONTEXT)
# Verify
assert result['uploaded_files'] == 1
@@ -159,12 +191,57 @@ class TestMaintenanceServiceCleanupExpiredFiles:
service._cleanup_expired_log_files = Mock(return_value=0) # NOT async
# Execute
result = await service.cleanup_expired_files()
result = await service.cleanup_expired_files(TEST_CONTEXT)
# Verify - warning logged, defaults used
assert ap.logger.warning.called
assert 'uploaded_files' in result
async def test_cloud_cleanup_carries_scope_without_holding_database_session(self):
class ScopeOnlyPersistenceManager:
mode = SimpleNamespace(value='cloud_runtime')
def __init__(self):
self.active_workspace = None
@contextlib.asynccontextmanager
async def tenant_scope(self, workspace_uuid):
self.active_workspace = workspace_uuid
try:
yield
finally:
self.active_workspace = None
def current_session(self):
return None
persistence_mgr = ScopeOnlyPersistenceManager()
application = SimpleNamespace(
persistence_mgr=persistence_mgr,
instance_config=SimpleNamespace(data={}),
logger=SimpleNamespace(warning=Mock()),
)
service = MaintenanceService(application)
async def cleanup_uploads(_context, _retention_days):
assert persistence_mgr.active_workspace == TEST_CONTEXT.workspace_uuid
assert persistence_mgr.current_session() is None
return 2
def cleanup_logs(_retention_days):
assert persistence_mgr.active_workspace == TEST_CONTEXT.workspace_uuid
assert persistence_mgr.current_session() is None
return 1
service._cleanup_expired_uploaded_files = cleanup_uploads
service._cleanup_expired_log_files = cleanup_logs
assert await service.cleanup_expired_files(TEST_CONTEXT) == {
'uploaded_files': 2,
'log_files': 1,
}
assert persistence_mgr.active_workspace is None
class TestMaintenanceServiceGetStorageAnalysis:
"""Tests for get_storage_analysis method."""
@@ -196,7 +273,7 @@ class TestMaintenanceServiceGetStorageAnalysis:
service._expired_log_candidates = Mock(return_value=[])
# Execute
result = await service.get_storage_analysis()
result = await service.get_storage_analysis(TEST_CONTEXT)
# Verify
assert 'generated_at' in result
@@ -229,7 +306,7 @@ class TestMaintenanceServiceGetStorageAnalysis:
service._expired_log_candidates = Mock(return_value=[])
# Execute
result = await service.get_storage_analysis()
result = await service.get_storage_analysis(TEST_CONTEXT)
# Verify - all sections present
sections = {s['key'] for s in result['sections']}
@@ -265,7 +342,7 @@ class TestMaintenanceServiceGetStorageAnalysis:
service._expired_log_candidates = Mock(return_value=[])
# Execute
result = await service.get_storage_analysis()
result = await service.get_storage_analysis(TEST_CONTEXT)
# Verify
assert result['database']['type'] == 'postgresql'
@@ -294,7 +371,7 @@ class TestMaintenanceServiceGetStorageAnalysis:
service._expired_log_candidates = Mock(return_value=[{'name': 'old_log', 'size_bytes': 50}])
# Execute
result = await service.get_storage_analysis()
result = await service.get_storage_analysis(TEST_CONTEXT)
# Verify
assert len(result['cleanup_candidates']['uploaded_files']) == 1
@@ -316,7 +393,7 @@ class TestMaintenanceServiceMonitoringCounts:
service = MaintenanceService(ap)
# Execute
result = await service._monitoring_counts()
result = await service._monitoring_counts(TEST_CONTEXT)
# Verify - all table keys present
assert 'messages' in result
@@ -338,7 +415,7 @@ class TestMaintenanceServiceMonitoringCounts:
service = MaintenanceService(ap)
# Execute
result = await service._monitoring_counts()
result = await service._monitoring_counts(TEST_CONTEXT)
# Verify - all zero
assert all(v == 0 for v in result.values())
@@ -374,7 +451,7 @@ class TestMaintenanceServiceBinaryStorageStats:
service = MaintenanceService(ap)
# Execute
result = await service._binary_storage_stats()
result = await service._binary_storage_stats(TEST_CONTEXT)
# Verify
assert result['count'] == 10
@@ -404,7 +481,7 @@ class TestMaintenanceServiceBinaryStorageStats:
service = MaintenanceService(ap)
# Execute
result = await service._binary_storage_stats()
result = await service._binary_storage_stats(TEST_CONTEXT)
# Verify - warning logged, size_bytes None or 0
assert ap.logger.warning.called
@@ -618,11 +695,13 @@ class TestMaintenanceServiceIsUploadedFileKey:
"""Returns True for valid upload file key."""
# Setup
ap = SimpleNamespace()
ap.storage_mgr = _scoped_storage_manager()
service = MaintenanceService(ap)
# Execute - simple filename without path
result = service._is_uploaded_file_key('uploaded_file.txt')
key = f'{ap.storage_mgr.scoped_prefix(TEST_CONTEXT, owner_type="upload")}uploaded_file.txt'
result = service._is_uploaded_file_key(TEST_CONTEXT, key)
# Verify
assert result is True
@@ -631,11 +710,12 @@ class TestMaintenanceServiceIsUploadedFileKey:
"""Returns False for key with path separator."""
# Setup
ap = SimpleNamespace()
ap.storage_mgr = _scoped_storage_manager()
service = MaintenanceService(ap)
# Execute - key with path
result = service._is_uploaded_file_key('path/to/file.txt')
result = service._is_uploaded_file_key(TEST_CONTEXT, 'path/to/file.txt')
# Verify
assert result is False
@@ -644,11 +724,12 @@ class TestMaintenanceServiceIsUploadedFileKey:
"""Returns False for plugin config prefix."""
# Setup
ap = SimpleNamespace()
ap.storage_mgr = _scoped_storage_manager()
service = MaintenanceService(ap)
# Execute - plugin config file
result = service._is_uploaded_file_key('plugin_config_some_plugin.json')
result = service._is_uploaded_file_key(TEST_CONTEXT, 'plugin_config_some_plugin.json')
# Verify
assert result is False
@@ -662,6 +743,7 @@ class TestMaintenanceServiceExpiredLogCandidates:
# Setup
ap = SimpleNamespace()
ap.logger = SimpleNamespace()
ap.storage_mgr = _scoped_storage_manager()
service = MaintenanceService(ap)
@@ -748,11 +830,12 @@ class TestMaintenanceServiceExpiredLocalUploadCandidates:
# Setup
ap = SimpleNamespace()
ap.logger = SimpleNamespace()
ap.storage_mgr = _scoped_storage_manager()
service = MaintenanceService(ap)
with patch.object(Path, 'exists', return_value=False):
result = service._expired_local_upload_candidates(7)
result = service._expired_local_upload_candidates(TEST_CONTEXT, 7)
# Verify
assert result == []
@@ -762,12 +845,10 @@ class TestMaintenanceServiceExpiredLocalUploadCandidates:
# Setup
ap = SimpleNamespace()
ap.logger = SimpleNamespace()
ap.storage_mgr = _scoped_storage_manager()
service = MaintenanceService(ap)
# Mock _is_uploaded_file_key
service._is_uploaded_file_key = Mock(side_effect=lambda key: 'plugin_config_' not in key and '/' not in key)
# Create mock files - one valid, one plugin config
# Create one file and one non-file entry under the scoped upload root.
mock_entry_valid = Mock(spec=Path)
mock_entry_valid.is_file = Mock(return_value=True)
mock_entry_valid.name = 'valid_upload.txt'
@@ -775,9 +856,10 @@ class TestMaintenanceServiceExpiredLocalUploadCandidates:
mock_stat.st_size = 100
mock_stat.st_mtime = 0 # Very old
mock_entry_valid.stat = Mock(return_value=mock_stat)
mock_entry_valid.relative_to = Mock(return_value=Path('scoped/valid_upload.txt'))
mock_entry_plugin = Mock(spec=Path)
mock_entry_plugin.is_file = Mock(return_value=True)
mock_entry_plugin.is_file = Mock(return_value=False)
mock_entry_plugin.name = 'plugin_config_test.json'
mock_stat2 = Mock()
mock_stat2.st_size = 200
@@ -785,23 +867,22 @@ class TestMaintenanceServiceExpiredLocalUploadCandidates:
mock_entry_plugin.stat = Mock(return_value=mock_stat2)
with patch.object(Path, 'exists', return_value=True):
with patch.object(Path, 'iterdir') as mock_iterdir:
mock_iterdir.return_value = [mock_entry_valid, mock_entry_plugin]
result = service._expired_local_upload_candidates(7)
with patch.object(Path, 'rglob') as mock_rglob:
mock_rglob.return_value = [mock_entry_valid, mock_entry_plugin]
result = service._expired_local_upload_candidates(TEST_CONTEXT, 7)
# Verify - only valid upload included
assert len(result) == 1
assert result[0]['key'] == 'valid_upload.txt'
assert result[0]['key'] == 'scoped/valid_upload.txt'
def test_expired_local_upload_candidates_includes_path(self):
"""Includes path when include_paths=True."""
# Setup
ap = SimpleNamespace()
ap.logger = SimpleNamespace()
ap.storage_mgr = _scoped_storage_manager()
service = MaintenanceService(ap)
service._is_uploaded_file_key = Mock(return_value=True)
mock_entry = Mock(spec=Path)
mock_entry.is_file = Mock(return_value=True)
mock_entry.name = 'old_file.txt'
@@ -810,11 +891,178 @@ class TestMaintenanceServiceExpiredLocalUploadCandidates:
mock_stat.st_size = 100
mock_stat.st_mtime = 0
mock_entry.stat = Mock(return_value=mock_stat)
mock_entry.relative_to = Mock(return_value=Path('scoped/old_file.txt'))
with patch.object(Path, 'exists', return_value=True):
with patch.object(Path, 'iterdir') as mock_iterdir:
mock_iterdir.return_value = [mock_entry]
result = service._expired_local_upload_candidates(7, include_paths=True)
with patch.object(Path, 'rglob') as mock_rglob:
mock_rglob.return_value = [mock_entry]
result = service._expired_local_upload_candidates(
TEST_CONTEXT,
7,
include_paths=True,
)
# Verify - path included
assert 'path' in result[0]
def test_expired_local_upload_candidates_respects_run_limit(self):
ap = SimpleNamespace(
logger=SimpleNamespace(warning=Mock()),
storage_mgr=_scoped_storage_manager(),
instance_config=SimpleNamespace(data={'storage': {'cleanup': {'max_files_per_run': 2}}}),
)
service = MaintenanceService(ap)
entries = []
for index in range(3):
entry = Mock(spec=Path)
entry.is_file = Mock(return_value=True)
entry.stat = Mock(return_value=SimpleNamespace(st_size=100, st_mtime=0))
entry.relative_to = Mock(return_value=Path(f'scoped/old-{index}.txt'))
entries.append(entry)
with patch.object(Path, 'exists', return_value=True):
with patch.object(Path, 'rglob', return_value=entries):
result = service._expired_local_upload_candidates(TEST_CONTEXT, 7)
assert [item['key'] for item in result] == [
'scoped/old-0.txt',
'scoped/old-1.txt',
]
ap.instance_config.data['storage']['cleanup']['max_files_per_run'] = 999999
assert service._max_files_per_run() == 10000
ISOLATION_WORKSPACE_A = '00000000-0000-0000-0000-00000000000a'
ISOLATION_WORKSPACE_B = '00000000-0000-0000-0000-00000000000b'
def _tenant_context(workspace_uuid: str) -> ExecutionContext:
return ExecutionContext(
instance_uuid='instance',
workspace_uuid=workspace_uuid,
placement_generation=1,
trigger_principal=PrincipalContext(PrincipalType.SYSTEM),
)
class _RealPersistenceManager:
def __init__(self, engine):
self.engine = engine
async def execute_async(self, *args, **kwargs):
async with self.engine.connect() as connection:
result = await connection.execute(*args, **kwargs)
await connection.commit()
return result
@pytest.fixture
async def tenant_maintenance_service(tmp_path):
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "maintenance.db"}')
async with engine.begin() as connection:
await connection.run_sync(Base.metadata.create_all)
await connection.execute(
sqlalchemy.insert(Workspace),
[
{
'uuid': ISOLATION_WORKSPACE_A,
'instance_uuid': 'instance',
'name': 'A',
'slug': 'a',
'source': 'cloud_projection',
},
{
'uuid': ISOLATION_WORKSPACE_B,
'instance_uuid': 'instance',
'name': 'B',
'slug': 'b',
'source': 'cloud_projection',
},
],
)
now = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
await connection.execute(
sqlalchemy.insert(MonitoringMessage),
[
{
'id': 'message-a',
'workspace_uuid': ISOLATION_WORKSPACE_A,
'timestamp': now,
'bot_id': 'bot',
'bot_name': 'Bot',
'pipeline_id': 'pipeline',
'pipeline_name': 'Pipeline',
'message_content': 'A',
'session_id': 'same-session',
'status': 'success',
'level': 'info',
},
{
'id': 'message-b',
'workspace_uuid': ISOLATION_WORKSPACE_B,
'timestamp': now,
'bot_id': 'bot',
'bot_name': 'Bot',
'pipeline_id': 'pipeline',
'pipeline_name': 'Pipeline',
'message_content': 'B',
'session_id': 'same-session',
'status': 'success',
'level': 'info',
},
],
)
await connection.execute(
sqlalchemy.insert(BinaryStorage),
[
{
'workspace_uuid': ISOLATION_WORKSPACE_A,
'unique_key': 'a',
'key': 'same',
'owner_type': 'plugin',
'owner': 'same',
'value': b'aaa',
},
{
'workspace_uuid': ISOLATION_WORKSPACE_B,
'unique_key': 'b',
'key': 'same',
'owner_type': 'plugin',
'owner': 'same',
'value': b'bbbbb',
},
],
)
application = SimpleNamespace(
persistence_mgr=_RealPersistenceManager(engine),
instance_config=SimpleNamespace(data={}),
logger=SimpleNamespace(warning=lambda *_: None),
)
yield MaintenanceService(application)
await engine.dispose()
async def test_cleanup_requires_execution_context(tenant_maintenance_service):
with pytest.raises(WorkspaceRequiredError):
await tenant_maintenance_service.cleanup_expired_files(None)
async def test_monitoring_counts_are_workspace_scoped(tenant_maintenance_service):
counts_a = await tenant_maintenance_service._monitoring_counts(_tenant_context(ISOLATION_WORKSPACE_A))
counts_b = await tenant_maintenance_service._monitoring_counts(_tenant_context(ISOLATION_WORKSPACE_B))
assert counts_a['messages'] == 1
assert counts_b['messages'] == 1
async def test_binary_storage_stats_are_workspace_scoped(tenant_maintenance_service):
stats_a = await tenant_maintenance_service._binary_storage_stats(_tenant_context(ISOLATION_WORKSPACE_A))
stats_b = await tenant_maintenance_service._binary_storage_stats(_tenant_context(ISOLATION_WORKSPACE_B))
assert stats_a == {'count': 1, 'size_bytes': 3}
assert stats_b == {'count': 1, 'size_bytes': 5}
async def test_path_helpers_handle_missing_paths(tenant_maintenance_service, tmp_path):
missing = tmp_path / 'missing'
assert tenant_maintenance_service._path_size(missing) == 0
assert tenant_maintenance_service._file_count(missing) == 0
tests/unit_tests/api/service/test_mcp_service.py
+427 -69
View File
@@ -13,17 +13,62 @@ Source: src/langbot/pkg/api/http/service/mcp.py
from __future__ import annotations
import asyncio
import copy
import pytest
from unittest.mock import AsyncMock, Mock, MagicMock
from types import SimpleNamespace
import uuid
from langbot.pkg.api.http.service.mcp import MCPService
from langbot.pkg.api.http.authz import Permission
from langbot.pkg.api.http.context import (
ExecutionContext,
PrincipalContext,
PrincipalType,
RequestContext,
WorkspaceContext,
)
from langbot.pkg.api.http.service.mcp import MCPService, redact_mcp_secrets, restore_mcp_secret_placeholders
from langbot.pkg.core.taskmgr import TaskCapacityError
from langbot.pkg.entity.persistence.mcp import MCPServer
from langbot.pkg.provider.tools.loaders.mcp_policy import MCPStdioDisabledError
from langbot.pkg.workspace.errors import WorkspaceNotFoundError
pytestmark = pytest.mark.asyncio
_CONTEXT = ExecutionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=1,
)
_VIEWER_CONTEXT = RequestContext(
instance_uuid='instance-a',
placement_generation=1,
request_id='request-a',
auth_type='user_token',
principal=PrincipalContext(
principal_type=PrincipalType.ACCOUNT,
account_uuid='account-a',
),
workspace=WorkspaceContext(
workspace_uuid='workspace-a',
membership_uuid='membership-a',
role='viewer',
permissions=frozenset({Permission.RESOURCE_VIEW.value}),
),
)
def _service(ap: SimpleNamespace) -> MCPService:
ap.workspace_service = SimpleNamespace(
get_execution_binding=AsyncMock(return_value=SimpleNamespace(instance_uuid=_CONTEXT.instance_uuid))
)
if not hasattr(ap, 'logger'):
ap.logger = Mock()
return MCPService(ap)
def _create_mock_mcp_server(
server_uuid: str = None,
@@ -42,11 +87,13 @@ def _create_mock_mcp_server(
return server
def _create_mock_result(items: list = None, first_item=None):
def _create_mock_result(items: list = None, first_item=None, *, scalar_value=0, rowcount=1):
"""Create mock result object for persistence queries."""
result = Mock()
result.all = Mock(return_value=items or [])
result.first = Mock(return_value=first_item)
result.scalar = Mock(return_value=scalar_value)
result.rowcount = rowcount
return result
@@ -64,10 +111,10 @@ class TestMCPServiceGetRuntimeInfo:
mock_session.get_runtime_info_dict = Mock(return_value={'status': 'running', 'tools': 5})
ap.tool_mgr.mcp_tool_loader.get_session = Mock(return_value=mock_session)
service = MCPService(ap)
service = _service(ap)
# Execute
result = await service.get_runtime_info('test-server')
result = await service.get_runtime_info(_CONTEXT, 'test-server')
# Verify
assert result is not None
@@ -81,10 +128,10 @@ class TestMCPServiceGetRuntimeInfo:
ap.tool_mgr.mcp_tool_loader = SimpleNamespace()
ap.tool_mgr.mcp_tool_loader.get_session = Mock(return_value=None)
service = MCPService(ap)
service = _service(ap)
# Execute
result = await service.get_runtime_info('nonexistent-server')
result = await service.get_runtime_info(_CONTEXT, 'nonexistent-server')
# Verify
assert result is None
@@ -101,12 +148,13 @@ class TestMCPServiceResources:
return_value=[{'uri_template': 'file:///{path}', 'name': 'files'}]
)
service = MCPService(ap)
service = _service(ap)
service._require_server = AsyncMock(return_value=(_CONTEXT, {'name': 'docs'}))
result = await service.get_mcp_server_resource_templates('docs')
result = await service.get_mcp_server_resource_templates(_CONTEXT, 'docs')
assert result == [{'uri_template': 'file:///{path}', 'name': 'files'}]
ap.tool_mgr.mcp_tool_loader.get_resource_templates.assert_awaited_once_with('docs')
ap.tool_mgr.mcp_tool_loader.get_resource_templates.assert_awaited_once_with(_CONTEXT, 'docs')
async def test_read_resource_envelope_uses_ui_preview_source(self):
ap = SimpleNamespace()
@@ -121,9 +169,11 @@ class TestMCPServiceResources:
}
)
service = MCPService(ap)
service = _service(ap)
service._require_server = AsyncMock(return_value=(_CONTEXT, {'name': 'docs'}))
result = await service.read_mcp_server_resource_envelope(
_CONTEXT,
'docs',
'file:///README.md',
max_bytes=4096,
@@ -132,6 +182,7 @@ class TestMCPServiceResources:
assert result['source'] == 'ui_preview'
ap.tool_mgr.mcp_tool_loader.read_resource_envelope.assert_awaited_once_with(
_CONTEXT,
'docs',
'file:///README.md',
include_blob=True,
@@ -156,12 +207,12 @@ class TestMCPServiceGetMCPServers:
'name': entity.name,
}
)
ap.tool_mgr = None
ap.tool_mgr = SimpleNamespace(mcp_tool_loader=SimpleNamespace(get_session=Mock(return_value=None)))
service = MCPService(ap)
service = _service(ap)
# Execute
result = await service.get_mcp_servers()
result = await service.get_mcp_servers(_CONTEXT)
# Verify
assert result == []
@@ -185,12 +236,12 @@ class TestMCPServiceGetMCPServers:
'mode': entity.mode,
}
)
ap.tool_mgr = None
ap.tool_mgr = SimpleNamespace(mcp_tool_loader=SimpleNamespace(get_session=Mock(return_value=None)))
service = MCPService(ap)
service = _service(ap)
# Execute
result = await service.get_mcp_servers()
result = await service.get_mcp_servers(_CONTEXT)
# Verify
assert len(result) == 2
@@ -215,21 +266,115 @@ class TestMCPServiceGetMCPServers:
)
ap.tool_mgr = SimpleNamespace()
ap.tool_mgr.mcp_tool_loader = SimpleNamespace()
ap.tool_mgr.mcp_tool_loader.get_session = Mock(return_value=None)
runtime_session = SimpleNamespace(get_runtime_info_dict=Mock(return_value={'status': 'connected'}))
ap.tool_mgr.mcp_tool_loader.get_session = Mock(return_value=runtime_session)
service = MCPService(ap)
service.get_runtime_info = AsyncMock(return_value={'status': 'connected'})
service = _service(ap)
# Execute
result = await service.get_mcp_servers(contain_runtime_info=True)
result = await service.get_mcp_servers(_CONTEXT, contain_runtime_info=True)
# Verify - runtime info included
assert result[0]['runtime_info'] == {'status': 'connected'}
async def test_resource_view_list_and_detail_redact_secrets_without_mutating_raw_data(self):
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
server = _create_mock_mcp_server(name='Secret Server')
serialized = {
'uuid': 'secret-uuid',
'name': 'Secret Server',
'enable': True,
'extra_args': {
'url': (
'https://mcp-user:mcp-password@mcp.invalid/connect'
'?token=url-secret&transport=streamable&sig=signed-secret'
),
'headers': {
'Authorization': 'Bearer top-secret',
'X-API-Key': 'api-secret',
'Accept': 'application/json',
},
'env': {
'ACCESS_TOKEN': 'access-secret',
'TOKENIZER': 'public-model-name',
},
'credentials': {
'username': 'service-user',
'password': 'password-secret',
},
'public_key': 'public-value',
},
}
original = copy.deepcopy(serialized)
ap.persistence_mgr.execute_async = AsyncMock(
side_effect=[
_create_mock_result([server]),
_create_mock_result(first_item=server),
]
)
ap.persistence_mgr.serialize_model = Mock(return_value=serialized)
ap.tool_mgr = SimpleNamespace(mcp_tool_loader=SimpleNamespace(get_session=Mock(return_value=None)))
service = _service(ap)
listed = await service.get_mcp_servers(_VIEWER_CONTEXT)
detail = await service.get_mcp_server_by_name(_VIEWER_CONTEXT, 'Secret Server')
for response in (listed[0], detail):
assert response['extra_args']['url'] == (
'https://***@mcp.invalid/connect?token=***&transport=streamable&sig=***'
)
assert response['extra_args']['headers'] == {
'Authorization': '***',
'X-API-Key': '***',
'Accept': 'application/json',
}
assert response['extra_args']['env'] == {
'ACCESS_TOKEN': '***',
'TOKENIZER': 'public-model-name',
}
assert response['extra_args']['credentials'] == {
'username': '***',
'password': '***',
}
assert response['extra_args']['public_key'] == 'public-value'
assert serialized == original
async def test_redacted_url_roundtrip_restores_persisted_credentials(self):
persisted = {
'extra_args': {'url': 'https://mcp-user:mcp-password@mcp.invalid/connect?token=url-secret&transport=http'}
}
submitted = redact_mcp_secrets(persisted)
assert submitted['extra_args']['url'] == 'https://***@mcp.invalid/connect?token=***&transport=http'
assert restore_mcp_secret_placeholders(submitted, persisted) == persisted
class TestMCPServiceCreateMCPServer:
"""Tests for create_mcp_server method."""
async def test_create_stdio_rejected_by_independent_instance_gate(self):
ap = SimpleNamespace(
instance_config=SimpleNamespace(
data={
'mcp': {'stdio': {'enabled': False}},
'system': {'limitation': {'max_extensions': -1}},
}
),
persistence_mgr=SimpleNamespace(execute_async=AsyncMock()),
tool_mgr=None,
)
service = _service(ap)
with pytest.raises(MCPStdioDisabledError, match='disabled by instance policy'):
await service.create_mcp_server(
_CONTEXT,
{'name': 'local', 'mode': 'stdio', 'enable': True, 'extra_args': {}},
)
ap.persistence_mgr.execute_async.assert_not_awaited()
async def test_create_mcp_server_max_extensions_reached_raises(self):
"""Raises ValueError when max_extensions limit reached."""
# Setup
@@ -241,16 +386,20 @@ class TestMCPServiceCreateMCPServer:
ap.plugin_connector.list_plugins = AsyncMock(return_value=[Mock(), Mock()]) # 2 plugins
# Mock get_mcp_servers to return 0 servers (2 plugins already)
mock_result = _create_mock_result([])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.execute_async = AsyncMock(
side_effect=[
_create_mock_result(scalar_value=0),
_create_mock_result(scalar_value=2),
]
)
ap.persistence_mgr.serialize_model = Mock(return_value={})
ap.tool_mgr = None
ap.tool_mgr = SimpleNamespace(mcp_tool_loader=SimpleNamespace(get_session=Mock(return_value=None)))
service = MCPService(ap)
service = _service(ap)
# Execute & Verify - 2 plugins + new server would exceed limit
with pytest.raises(ValueError, match='Maximum number of extensions'):
await service.create_mcp_server({'name': 'New Server'})
await service.create_mcp_server(_CONTEXT, {'name': 'New Server'})
async def test_create_mcp_server_no_limit(self):
"""Creates MCP server without limit when max_extensions=-1."""
@@ -271,10 +420,10 @@ class TestMCPServiceCreateMCPServer:
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(return_value={'uuid': 'new-uuid'})
service = MCPService(ap)
service = _service(ap)
# Execute
server_uuid = await service.create_mcp_server({'name': 'New Server'})
server_uuid = await service.create_mcp_server(_CONTEXT, {'name': 'New Server'})
# Verify
assert server_uuid is not None
@@ -293,11 +442,11 @@ class TestMCPServiceCreateMCPServer:
ap.persistence_mgr.execute_async = AsyncMock(return_value=_create_mock_result(first_item=existing_server))
ap.persistence_mgr.serialize_model = Mock(return_value={})
service = MCPService(ap)
service = _service(ap)
# Execute & Verify
with pytest.raises(ValueError, match='MCP server already exists: Existing Server'):
await service.create_mcp_server({'name': 'Existing Server'})
await service.create_mcp_server(_CONTEXT, {'name': 'Existing Server'})
async def test_create_mcp_server_loads_server(self):
"""Loads server into tool_mgr when enabled."""
@@ -330,14 +479,62 @@ class TestMCPServiceCreateMCPServer:
return_value={'uuid': 'new-uuid', 'name': 'New Server', 'enable': True}
)
service = MCPService(ap)
service = _service(ap)
# Execute
await service.create_mcp_server({'name': 'New Server', 'enable': True})
await service.create_mcp_server(_CONTEXT, {'name': 'New Server', 'enable': True})
# Verify - host_mcp_server was called
ap.tool_mgr.mcp_tool_loader.host_mcp_server.assert_called_once()
async def test_create_mcp_server_does_not_start_host_until_transaction_commits(self):
"""The Runtime must not observe a server row that can still roll back."""
gate = asyncio.get_running_loop().create_future()
class PersistenceManagerStub:
def create_after_commit_gate(self):
return gate
ap = SimpleNamespace()
ap.persistence_mgr = PersistenceManagerStub()
ap.instance_config = SimpleNamespace(data={'system': {'limitation': {'max_extensions': -1}}})
observed = []
async def host_mcp_server(context, config):
observed.append((context, config))
ap.tool_mgr = SimpleNamespace(
mcp_tool_loader=SimpleNamespace(
host_mcp_server=host_mcp_server,
_hosted_mcp_tasks=[],
)
)
server_entity = _create_mock_mcp_server(server_uuid='new-uuid', enable=True)
results = [
_create_mock_result([]),
Mock(),
_create_mock_result(first_item=server_entity),
]
ap.persistence_mgr.execute_async = AsyncMock(side_effect=results)
ap.persistence_mgr.serialize_model = Mock(
return_value={'uuid': 'new-uuid', 'name': 'New Server', 'enable': True}
)
service = _service(ap)
await service.create_mcp_server(_CONTEXT, {'name': 'New Server', 'enable': True})
await asyncio.sleep(0)
assert observed == []
gate.set_result(None)
await ap.tool_mgr.mcp_tool_loader._hosted_mcp_tasks[0]
assert observed == [
(
_CONTEXT,
{'uuid': 'new-uuid', 'name': 'New Server', 'enable': True},
)
]
async def test_create_mcp_server_disabled_no_load(self):
"""Does not load server when disabled."""
# Setup
@@ -351,10 +548,10 @@ class TestMCPServiceCreateMCPServer:
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
ap.persistence_mgr.serialize_model = Mock(return_value={'uuid': 'new-uuid'})
service = MCPService(ap)
service = _service(ap)
# Execute with enable=False
server_uuid = await service.create_mcp_server({'name': 'New Server', 'enable': False})
server_uuid = await service.create_mcp_server(_CONTEXT, {'name': 'New Server', 'enable': False})
# Verify - no tool_mgr load attempt
assert server_uuid is not None
@@ -379,13 +576,11 @@ class TestMCPServiceGetMCPServerByName:
'runtime_info': None,
}
)
ap.tool_mgr = None
service = MCPService(ap)
service.get_runtime_info = AsyncMock(return_value=None)
ap.tool_mgr = SimpleNamespace(mcp_tool_loader=SimpleNamespace(get_session=Mock(return_value=None)))
service = _service(ap)
# Execute
result = await service.get_mcp_server_by_name('Found Server')
result = await service.get_mcp_server_by_name(_CONTEXT, 'Found Server')
# Verify
assert result is not None
@@ -400,10 +595,10 @@ class TestMCPServiceGetMCPServerByName:
mock_result = _create_mock_result(first_item=None)
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
service = MCPService(ap)
service = _service(ap)
# Execute
result = await service.get_mcp_server_by_name('Nonexistent Server')
result = await service.get_mcp_server_by_name(_CONTEXT, 'Nonexistent Server')
# Verify
assert result is None
@@ -421,8 +616,10 @@ class TestMCPServiceUpdateMCPServer:
ap.tool_mgr.mcp_tool_loader = SimpleNamespace()
ap.tool_mgr.mcp_tool_loader.sessions = {'Old Server': Mock()}
ap.tool_mgr.mcp_tool_loader.remove_mcp_server = AsyncMock()
ap.tool_mgr.mcp_tool_loader.has_session = Mock(return_value=True)
old_server = _create_mock_mcp_server(name='Old Server', enable=True)
updated_server = _create_mock_mcp_server(name='Old Server', enable=False)
call_count = 0
@@ -431,14 +628,23 @@ class TestMCPServiceUpdateMCPServer:
call_count += 1
if call_count == 1:
return _create_mock_result(first_item=old_server)
return Mock() # Update
if call_count == 2:
return _create_mock_result()
return _create_mock_result(first_item=updated_server)
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
ap.persistence_mgr.serialize_model = Mock(
side_effect=lambda _model, entity: {
'uuid': 'test-uuid',
'name': entity.name,
'enable': entity.enable,
}
)
service = MCPService(ap)
service = _service(ap)
# Execute - disable server
await service.update_mcp_server('test-uuid', {'enable': False})
await service.update_mcp_server(_CONTEXT, 'test-uuid', {'enable': False})
# Verify - server was removed
ap.tool_mgr.mcp_tool_loader.remove_mcp_server.assert_called_once()
@@ -453,6 +659,7 @@ class TestMCPServiceUpdateMCPServer:
ap.tool_mgr.mcp_tool_loader.sessions = {}
ap.tool_mgr.mcp_tool_loader.host_mcp_server = AsyncMock()
ap.tool_mgr.mcp_tool_loader._hosted_mcp_tasks = []
ap.tool_mgr.mcp_tool_loader.has_session = Mock(return_value=False)
old_server = _create_mock_mcp_server(name='Old Server', enable=False)
@@ -474,10 +681,10 @@ class TestMCPServiceUpdateMCPServer:
return_value={'uuid': 'test-uuid', 'name': 'Old Server', 'enable': True}
)
service = MCPService(ap)
service = _service(ap)
# Execute - enable server
await service.update_mcp_server('test-uuid', {'enable': True})
await service.update_mcp_server(_CONTEXT, 'test-uuid', {'enable': True})
# Verify - server was loaded
ap.tool_mgr.mcp_tool_loader.host_mcp_server.assert_called_once()
@@ -493,6 +700,7 @@ class TestMCPServiceUpdateMCPServer:
ap.tool_mgr.mcp_tool_loader.remove_mcp_server = AsyncMock()
ap.tool_mgr.mcp_tool_loader.host_mcp_server = AsyncMock()
ap.tool_mgr.mcp_tool_loader._hosted_mcp_tasks = []
ap.tool_mgr.mcp_tool_loader.has_session = Mock(return_value=True)
old_server = _create_mock_mcp_server(name='Old Server', enable=True)
@@ -510,13 +718,13 @@ class TestMCPServiceUpdateMCPServer:
return_value={'uuid': 'test-uuid', 'name': 'Old Server', 'enable': True}
)
service = MCPService(ap)
service = _service(ap)
# Execute - update enabled server (keep enabled, update extra_args)
await service.update_mcp_server('test-uuid', {'enable': True, 'extra_args': {'new': 'args'}})
await service.update_mcp_server(_CONTEXT, 'test-uuid', {'enable': True, 'extra_args': {'new': 'args'}})
# Verify - remove and reload
ap.tool_mgr.mcp_tool_loader.remove_mcp_server.assert_called_once_with('Old Server')
ap.tool_mgr.mcp_tool_loader.remove_mcp_server.assert_called_once_with(_CONTEXT, 'Old Server')
ap.tool_mgr.mcp_tool_loader.host_mcp_server.assert_called_once()
async def test_update_mcp_server_no_tool_mgr(self):
@@ -541,15 +749,99 @@ class TestMCPServiceUpdateMCPServer:
return Mock() # Update
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
ap.persistence_mgr.serialize_model = Mock(
return_value={
'uuid': 'test-uuid',
'name': 'Server',
'enable': True,
}
)
service = MCPService(ap)
service = _service(ap)
# Execute - should not raise
await service.update_mcp_server('test-uuid', {'name': 'New Name'})
await service.update_mcp_server(_CONTEXT, 'test-uuid', {'enable': False})
# Verify - persistence was called
assert ap.persistence_mgr.execute_async.call_count >= 2
async def test_update_restores_existing_masked_secrets_and_preserves_explicit_changes(self):
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.tool_mgr = SimpleNamespace(mcp_tool_loader=None)
old_server = _create_mock_mcp_server(name='Server', enable=True)
old_data = {
'uuid': 'test-uuid',
'name': 'Server',
'enable': True,
'mode': 'streamable_http',
'extra_args': {
'headers': {
'Authorization': 'Bearer original-secret',
'X-API-Key': 'original-api-key',
'Cookie': 'original-cookie',
}
},
}
captured_updates = []
async def mock_execute(statement):
if not captured_updates:
captured_updates.append(None)
return _create_mock_result(first_item=old_server)
captured_updates[0] = statement
return _create_mock_result()
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
ap.persistence_mgr.serialize_model = Mock(return_value=old_data)
service = _service(ap)
await service.update_mcp_server(
_CONTEXT,
'test-uuid',
{
'extra_args': {
'headers': {
'Authorization': '***',
'X-API-Key': 'replacement-api-key',
'Cookie': '',
}
}
},
)
persisted = captured_updates[0].compile().params['extra_args']
assert persisted['headers'] == {
'Authorization': 'Bearer original-secret',
'X-API-Key': 'replacement-api-key',
'Cookie': '',
}
async def test_update_rejects_masked_secret_without_existing_value(self):
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.tool_mgr = SimpleNamespace(mcp_tool_loader=None)
old_server = _create_mock_mcp_server(name='Server', enable=True)
ap.persistence_mgr.execute_async = AsyncMock(return_value=_create_mock_result(first_item=old_server))
ap.persistence_mgr.serialize_model = Mock(
return_value={
'uuid': 'test-uuid',
'name': 'Server',
'enable': True,
'extra_args': {'headers': {'Accept': 'application/json'}},
}
)
service = _service(ap)
with pytest.raises(ValueError, match='Masked MCP secret has no existing value'):
await service.update_mcp_server(
_CONTEXT,
'test-uuid',
{'extra_args': {'headers': {'Authorization': '***'}}},
)
assert ap.persistence_mgr.execute_async.await_count == 1
class TestMCPServiceDeleteMCPServer:
"""Tests for delete_mcp_server method."""
@@ -563,6 +855,7 @@ class TestMCPServiceDeleteMCPServer:
ap.tool_mgr.mcp_tool_loader = SimpleNamespace()
ap.tool_mgr.mcp_tool_loader.sessions = {'Server to Delete': Mock()}
ap.tool_mgr.mcp_tool_loader.remove_mcp_server = AsyncMock()
ap.tool_mgr.mcp_tool_loader.has_session = Mock(return_value=True)
server = _create_mock_mcp_server(name='Server to Delete')
@@ -576,14 +869,21 @@ class TestMCPServiceDeleteMCPServer:
return Mock() # Delete
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
ap.persistence_mgr.serialize_model = Mock(
return_value={
'uuid': 'test-uuid',
'name': 'Server to Delete',
'enable': True,
}
)
service = MCPService(ap)
service = _service(ap)
# Execute
await service.delete_mcp_server('test-uuid')
await service.delete_mcp_server(_CONTEXT, 'test-uuid')
# Verify
ap.tool_mgr.mcp_tool_loader.remove_mcp_server.assert_called_once_with('Server to Delete')
ap.tool_mgr.mcp_tool_loader.remove_mcp_server.assert_called_once_with(_CONTEXT, 'Server to Delete')
ap.persistence_mgr.execute_async.assert_called()
async def test_delete_mcp_server_not_in_sessions(self):
@@ -595,6 +895,7 @@ class TestMCPServiceDeleteMCPServer:
ap.tool_mgr.mcp_tool_loader = SimpleNamespace()
ap.tool_mgr.mcp_tool_loader.sessions = {} # Server not in sessions
ap.tool_mgr.mcp_tool_loader.remove_mcp_server = AsyncMock()
ap.tool_mgr.mcp_tool_loader.has_session = Mock(return_value=False)
server = _create_mock_mcp_server(name='Not in Sessions')
@@ -608,11 +909,18 @@ class TestMCPServiceDeleteMCPServer:
return Mock()
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
ap.persistence_mgr.serialize_model = Mock(
return_value={
'uuid': 'test-uuid',
'name': 'Not in Sessions',
'enable': True,
}
)
service = MCPService(ap)
service = _service(ap)
# Execute
await service.delete_mcp_server('test-uuid')
await service.delete_mcp_server(_CONTEXT, 'test-uuid')
# Verify - remove not called (server not in sessions)
ap.tool_mgr.mcp_tool_loader.remove_mcp_server.assert_not_called()
@@ -626,6 +934,7 @@ class TestMCPServiceDeleteMCPServer:
ap.tool_mgr.mcp_tool_loader = SimpleNamespace()
ap.tool_mgr.mcp_tool_loader.sessions = {}
ap.tool_mgr.mcp_tool_loader.remove_mcp_server = AsyncMock()
ap.tool_mgr.mcp_tool_loader.has_session = Mock(return_value=False)
# No server found
call_count = 0
@@ -639,18 +948,35 @@ class TestMCPServiceDeleteMCPServer:
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
service = MCPService(ap)
service = _service(ap)
# Execute - should not raise
await service.delete_mcp_server('nonexistent-uuid')
with pytest.raises(WorkspaceNotFoundError, match='MCP server not found'):
await service.delete_mcp_server(_CONTEXT, 'nonexistent-uuid')
# Verify - delete was called regardless
ap.persistence_mgr.execute_async.assert_called()
assert ap.persistence_mgr.execute_async.await_count == 1
class TestMCPServiceTestMCPServer:
"""Tests for test_mcp_server method."""
async def test_transient_stdio_test_rejected_by_instance_gate(self):
ap = SimpleNamespace(
instance_config=SimpleNamespace(data={'mcp': {'stdio': {'enabled': False}}}),
tool_mgr=SimpleNamespace(mcp_tool_loader=SimpleNamespace(load_mcp_server=AsyncMock())),
task_mgr=SimpleNamespace(create_user_task=Mock()),
)
service = _service(ap)
with pytest.raises(MCPStdioDisabledError, match='disabled by instance policy'):
await service.test_mcp_server(
_CONTEXT,
'_',
{'name': 'local', 'mode': 'stdio', 'enable': True, 'extra_args': {}},
)
ap.tool_mgr.mcp_tool_loader.load_mcp_server.assert_not_awaited()
ap.task_mgr.create_user_task.assert_not_called()
async def test_test_mcp_server_existing_server(self):
"""Tests existing MCP server connection."""
# Setup
@@ -667,12 +993,18 @@ class TestMCPServiceTestMCPServer:
ap.tool_mgr.mcp_tool_loader.get_session = Mock(return_value=mock_session)
ap.task_mgr = SimpleNamespace()
ap.task_mgr.create_user_task = Mock(return_value=SimpleNamespace(id=123))
service = MCPService(ap)
service = _service(ap)
service._require_server = AsyncMock(return_value=(_CONTEXT, {'name': 'existing-server'}))
def create_user_task(coroutine, **_kwargs):
coroutine.close()
return SimpleNamespace(id=123)
ap.task_mgr.create_user_task = Mock(side_effect=create_user_task)
# Execute
task_id = await service.test_mcp_server('existing-server', {})
task_id = await service.test_mcp_server(_CONTEXT, 'existing-server', {})
# Verify - returns task ID
assert task_id == 123
@@ -685,11 +1017,12 @@ class TestMCPServiceTestMCPServer:
ap.tool_mgr.mcp_tool_loader = SimpleNamespace()
ap.tool_mgr.mcp_tool_loader.get_session = Mock(return_value=None)
service = MCPService(ap)
service = _service(ap)
service._require_server = AsyncMock(side_effect=WorkspaceNotFoundError('MCP server not found'))
# Execute & Verify
with pytest.raises(ValueError, match='Server not found'):
await service.test_mcp_server('nonexistent-server', {})
with pytest.raises(WorkspaceNotFoundError, match='MCP server not found'):
await service.test_mcp_server(_CONTEXT, 'nonexistent-server', {})
async def test_test_mcp_server_new_server(self):
"""Tests new MCP server with underscore name."""
@@ -703,13 +1036,38 @@ class TestMCPServiceTestMCPServer:
ap.tool_mgr.mcp_tool_loader.load_mcp_server = AsyncMock(return_value=mock_session)
ap.task_mgr = SimpleNamespace()
ap.task_mgr.create_user_task = Mock(return_value=SimpleNamespace(id=456))
service = MCPService(ap)
service = _service(ap)
def create_user_task(coroutine, **_kwargs):
coroutine.close()
return SimpleNamespace(id=456)
ap.task_mgr.create_user_task = Mock(side_effect=create_user_task)
# Execute with '_' name (new server)
task_id = await service.test_mcp_server('_', {'name': 'New Server'})
task_id = await service.test_mcp_server(_CONTEXT, '_', {'name': 'New Server'})
# Verify - load_mcp_server called
ap.tool_mgr.mcp_tool_loader.load_mcp_server.assert_called_once()
assert task_id == 456
async def test_rejected_transient_test_session_is_shut_down(self):
ap = SimpleNamespace()
mock_session = MagicMock()
mock_session.shutdown = AsyncMock()
ap.tool_mgr = SimpleNamespace(
mcp_tool_loader=SimpleNamespace(load_mcp_server=AsyncMock(return_value=mock_session))
)
def reject(coroutine, **_kwargs):
coroutine.close()
raise TaskCapacityError('capacity')
ap.task_mgr = SimpleNamespace(create_user_task=Mock(side_effect=reject))
service = _service(ap)
with pytest.raises(TaskCapacityError, match='capacity'):
await service.test_mcp_server(_CONTEXT, '_', {'name': 'New Server'})
mock_session.shutdown.assert_awaited_once_with()
tests/unit_tests/api/service/test_model_service.py
+203 -58
View File
@@ -25,11 +25,37 @@ from langbot.pkg.api.http.service.model import (
_runtime_model_data,
_validate_provider_supports,
)
from langbot.pkg.api.http.service import model as model_service_module
from langbot.pkg.entity.persistence.model import LLMModel, EmbeddingModel, RerankModel, ModelProvider
pytestmark = pytest.mark.asyncio
WORKSPACE_UUID = 'workspace-a'
@pytest.fixture(autouse=True)
def assume_test_provider_belongs_to_workspace(monkeypatch):
"""Keep legacy runtime-focused tests isolated from the new ownership lookup."""
async def _allow_provider(_ap, _context, provider_uuid):
return {'uuid': provider_uuid}
monkeypatch.setattr(model_service_module, '_require_workspace_provider', _allow_provider)
def _existing_llm_data(provider_uuid: str = 'provider-uuid') -> dict:
return {
'uuid': 'existing-uuid',
'workspace_uuid': WORKSPACE_UUID,
'name': 'Existing Model',
'provider_uuid': provider_uuid,
'abilities': [],
'context_length': None,
'extra_args': {},
'prefered_ranking': 0,
}
def _create_mock_llm_model(
model_uuid: str = 'llm-uuid',
@@ -101,6 +127,35 @@ def _create_mock_result(items: list = None, first_item=None):
return result
def _create_runtime_model_mgr() -> SimpleNamespace:
"""Build a context-aware runtime-manager double for service tests."""
manager = SimpleNamespace(
provider_dict={},
llm_models=[],
embedding_models=[],
rerank_models=[],
load_llm_model_with_provider=AsyncMock(return_value=Mock()),
load_embedding_model_with_provider=AsyncMock(return_value=Mock()),
load_rerank_model_with_provider=AsyncMock(return_value=Mock()),
cache_llm_model=AsyncMock(),
cache_embedding_model=AsyncMock(),
cache_rerank_model=AsyncMock(),
remove_llm_model=AsyncMock(),
remove_embedding_model=AsyncMock(),
remove_rerank_model=AsyncMock(),
)
async def get_provider(_context, provider_uuid):
provider = manager.provider_dict.get(provider_uuid)
if provider is None:
raise ValueError(f'Model provider {provider_uuid} not found')
return provider
manager.get_provider_by_uuid = AsyncMock(side_effect=get_provider)
return manager
class TestParseProviderApiKeys:
"""Tests for _parse_provider_api_keys helper function."""
@@ -183,7 +238,9 @@ class TestLLMModelsServiceGetLLMModels:
service = LLMModelsService(ap)
# Execute
result = await service.get_llm_models()
result = await service.get_llm_models(
WORKSPACE_UUID,
)
# Verify
assert result == []
@@ -221,7 +278,9 @@ class TestLLMModelsServiceGetLLMModels:
service = LLMModelsService(ap)
# Execute
result = await service.get_llm_models()
result = await service.get_llm_models(
WORKSPACE_UUID,
)
# Verify
assert len(result) == 1
@@ -260,7 +319,7 @@ class TestLLMModelsServiceGetLLMModels:
service = LLMModelsService(ap)
# Execute
result = await service.get_llm_models(include_secret=False)
result = await service.get_llm_models(WORKSPACE_UUID, include_secret=False)
# Verify - keys should be masked
assert result[0]['provider']['api_keys'] == ['***', '***']
@@ -302,7 +361,7 @@ class TestLLMModelsServiceGetLLMModel:
service = LLMModelsService(ap)
# Execute
result = await service.get_llm_model('found-uuid')
result = await service.get_llm_model(WORKSPACE_UUID, 'found-uuid')
# Verify
assert result is not None
@@ -321,7 +380,7 @@ class TestLLMModelsServiceGetLLMModel:
service = LLMModelsService(ap)
# Execute
result = await service.get_llm_model('nonexistent-uuid')
result = await service.get_llm_model(WORKSPACE_UUID, 'nonexistent-uuid')
# Verify
assert result is None
@@ -346,7 +405,7 @@ class TestLLMModelsServiceGetLLMModelsByProvider:
service = LLMModelsService(ap)
# Execute
result = await service.get_llm_models_by_provider('target-provider')
result = await service.get_llm_models_by_provider(WORKSPACE_UUID, 'target-provider')
# Verify
assert len(result) == 2
@@ -360,7 +419,7 @@ class TestLLMModelsServiceCreateLLMModel:
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.model_mgr = SimpleNamespace()
ap.model_mgr = _create_runtime_model_mgr()
ap.model_mgr.provider_dict = {'provider-uuid': Mock()}
ap.model_mgr.llm_models = []
ap.model_mgr.load_llm_model_with_provider = AsyncMock(return_value=Mock())
@@ -374,12 +433,13 @@ class TestLLMModelsServiceCreateLLMModel:
# Execute
model_uuid = await service.create_llm_model(
WORKSPACE_UUID,
{
'name': 'New LLM',
'provider_uuid': 'provider-uuid',
'abilities': [],
'extra_args': {},
}
},
)
# Verify
@@ -391,7 +451,7 @@ class TestLLMModelsServiceCreateLLMModel:
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.model_mgr = SimpleNamespace()
ap.model_mgr = _create_runtime_model_mgr()
ap.model_mgr.provider_dict = {'provider-uuid': Mock()}
ap.model_mgr.llm_models = []
ap.model_mgr.load_llm_model_with_provider = AsyncMock(return_value=Mock())
@@ -405,6 +465,7 @@ class TestLLMModelsServiceCreateLLMModel:
# Execute
model_uuid = await service.create_llm_model(
WORKSPACE_UUID,
{
'uuid': 'preserved-uuid',
'name': 'Preserved UUID Model',
@@ -422,7 +483,7 @@ class TestLLMModelsServiceCreateLLMModel:
"""Creates LLM model with context_length outside extra_args."""
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.model_mgr = SimpleNamespace()
ap.model_mgr = _create_runtime_model_mgr()
ap.model_mgr.provider_dict = {'provider-uuid': Mock()}
ap.model_mgr.llm_models = []
ap.model_mgr.load_llm_model_with_provider = AsyncMock(return_value=Mock())
@@ -434,6 +495,7 @@ class TestLLMModelsServiceCreateLLMModel:
service = LLMModelsService(ap)
await service.create_llm_model(
WORKSPACE_UUID,
{
'uuid': 'model-with-context',
'name': 'Context Model',
@@ -446,7 +508,7 @@ class TestLLMModelsServiceCreateLLMModel:
auto_set_to_default_pipeline=False,
)
runtime_entity = ap.model_mgr.load_llm_model_with_provider.await_args.args[0]
runtime_entity = ap.model_mgr.load_llm_model_with_provider.await_args.args[1]
assert runtime_entity.context_length == 128000
assert runtime_entity.extra_args == {'temperature': 0.2}
assert 'context_length' not in runtime_entity.extra_args
@@ -456,7 +518,7 @@ class TestLLMModelsServiceCreateLLMModel:
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.model_mgr = SimpleNamespace()
ap.model_mgr = _create_runtime_model_mgr()
ap.model_mgr.provider_dict = {} # Empty - no provider
mock_result = _create_mock_result([])
@@ -467,12 +529,13 @@ class TestLLMModelsServiceCreateLLMModel:
# Execute & Verify
with pytest.raises(Exception, match='provider not found'):
await service.create_llm_model(
WORKSPACE_UUID,
{
'name': 'No Provider Model',
'provider_uuid': 'nonexistent-provider',
'abilities': [],
'extra_args': {},
}
},
)
async def test_create_llm_model_with_provider_data(self):
@@ -480,7 +543,7 @@ class TestLLMModelsServiceCreateLLMModel:
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.model_mgr = SimpleNamespace()
ap.model_mgr = _create_runtime_model_mgr()
ap.model_mgr.provider_dict = {}
ap.model_mgr.llm_models = []
ap.model_mgr.load_llm_model_with_provider = AsyncMock(return_value=Mock())
@@ -500,6 +563,7 @@ class TestLLMModelsServiceCreateLLMModel:
# Execute - with provider data (no UUID)
result_uuid = await service.create_llm_model(
WORKSPACE_UUID,
{
'name': 'Model with New Provider',
'provider': {
@@ -509,7 +573,7 @@ class TestLLMModelsServiceCreateLLMModel:
},
'abilities': [],
'extra_args': {},
}
},
)
# Verify - provider_service was called and UUID generated
@@ -525,7 +589,7 @@ class TestLLMModelsServiceUpdateLLMModel:
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.model_mgr = SimpleNamespace()
ap.model_mgr = _create_runtime_model_mgr()
ap.model_mgr.provider_dict = {'provider-uuid': Mock()}
ap.model_mgr.llm_models = []
ap.model_mgr.remove_llm_model = AsyncMock()
@@ -534,9 +598,11 @@ class TestLLMModelsServiceUpdateLLMModel:
ap.persistence_mgr.execute_async = AsyncMock()
service = LLMModelsService(ap)
service.get_llm_model = AsyncMock(return_value=_existing_llm_data())
# Execute
await service.update_llm_model(
WORKSPACE_UUID,
'existing-uuid',
{
'uuid': 'should-be-removed',
@@ -546,24 +612,26 @@ class TestLLMModelsServiceUpdateLLMModel:
)
# Verify - remove and load called
ap.model_mgr.remove_llm_model.assert_called_once_with('existing-uuid')
ap.model_mgr.remove_llm_model.assert_called_once_with(WORKSPACE_UUID, 'existing-uuid')
async def test_update_llm_model_provider_not_found_raises_error(self):
"""Raises Exception when provider not found after update."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.model_mgr = SimpleNamespace()
ap.model_mgr = _create_runtime_model_mgr()
ap.model_mgr.provider_dict = {} # Empty
ap.model_mgr.remove_llm_model = AsyncMock()
ap.persistence_mgr.execute_async = AsyncMock()
service = LLMModelsService(ap)
service.get_llm_model = AsyncMock(return_value=_existing_llm_data('nonexistent-provider'))
# Execute & Verify
with pytest.raises(Exception, match='provider not found'):
await service.update_llm_model(
WORKSPACE_UUID,
'model-uuid',
{
'name': 'Update',
@@ -575,15 +643,17 @@ class TestLLMModelsServiceUpdateLLMModel:
"""Updates runtime model with context_length outside extra_args."""
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace(execute_async=AsyncMock())
ap.model_mgr = SimpleNamespace()
ap.model_mgr = _create_runtime_model_mgr()
ap.model_mgr.provider_dict = {'provider-uuid': Mock()}
ap.model_mgr.llm_models = []
ap.model_mgr.remove_llm_model = AsyncMock()
ap.model_mgr.load_llm_model_with_provider = AsyncMock(return_value=Mock())
service = LLMModelsService(ap)
service.get_llm_model = AsyncMock(return_value=_existing_llm_data())
await service.update_llm_model(
WORKSPACE_UUID,
'existing-uuid',
{
'name': 'Updated Name',
@@ -594,7 +664,7 @@ class TestLLMModelsServiceUpdateLLMModel:
},
)
runtime_entity = ap.model_mgr.load_llm_model_with_provider.await_args.args[0]
runtime_entity = ap.model_mgr.load_llm_model_with_provider.await_args.args[1]
assert runtime_entity.uuid == 'existing-uuid'
assert runtime_entity.context_length == 64000
assert runtime_entity.extra_args == {'temperature': 0.4}
@@ -609,7 +679,7 @@ class TestLLMModelsServiceDeleteLLMModel:
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.model_mgr = SimpleNamespace()
ap.model_mgr = _create_runtime_model_mgr()
ap.model_mgr.remove_llm_model = AsyncMock()
ap.persistence_mgr.execute_async = AsyncMock()
@@ -617,11 +687,11 @@ class TestLLMModelsServiceDeleteLLMModel:
service = LLMModelsService(ap)
# Execute
await service.delete_llm_model('delete-uuid')
await service.delete_llm_model(WORKSPACE_UUID, 'delete-uuid')
# Verify
ap.persistence_mgr.execute_async.assert_called_once()
ap.model_mgr.remove_llm_model.assert_called_once_with('delete-uuid')
ap.model_mgr.remove_llm_model.assert_called_once_with(WORKSPACE_UUID, 'delete-uuid')
class TestEmbeddingModelsServiceGetEmbeddingModels:
@@ -640,7 +710,9 @@ class TestEmbeddingModelsServiceGetEmbeddingModels:
service = EmbeddingModelsService(ap)
# Execute
result = await service.get_embedding_models()
result = await service.get_embedding_models(
WORKSPACE_UUID,
)
# Verify
assert result == []
@@ -677,7 +749,9 @@ class TestEmbeddingModelsServiceGetEmbeddingModels:
service = EmbeddingModelsService(ap)
# Execute
result = await service.get_embedding_models()
result = await service.get_embedding_models(
WORKSPACE_UUID,
)
# Verify
assert len(result) == 1
@@ -717,7 +791,7 @@ class TestEmbeddingModelsServiceGetEmbeddingModel:
service = EmbeddingModelsService(ap)
# Execute
result = await service.get_embedding_model('found-embedding')
result = await service.get_embedding_model(WORKSPACE_UUID, 'found-embedding')
# Verify
assert result is not None
@@ -734,7 +808,7 @@ class TestEmbeddingModelsServiceGetEmbeddingModel:
service = EmbeddingModelsService(ap)
# Execute
result = await service.get_embedding_model('nonexistent-embedding')
result = await service.get_embedding_model(WORKSPACE_UUID, 'nonexistent-embedding')
# Verify
assert result is None
@@ -748,7 +822,7 @@ class TestEmbeddingModelsServiceCreateEmbeddingModel:
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.model_mgr = SimpleNamespace()
ap.model_mgr = _create_runtime_model_mgr()
ap.model_mgr.provider_dict = {'provider-uuid': Mock()}
ap.model_mgr.embedding_models = []
ap.model_mgr.load_embedding_model_with_provider = AsyncMock(return_value=Mock())
@@ -760,11 +834,12 @@ class TestEmbeddingModelsServiceCreateEmbeddingModel:
# Execute
model_uuid = await service.create_embedding_model(
WORKSPACE_UUID,
{
'name': 'New Embedding',
'provider_uuid': 'provider-uuid',
'extra_args': {},
}
},
)
# Verify
@@ -776,7 +851,7 @@ class TestEmbeddingModelsServiceCreateEmbeddingModel:
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.model_mgr = SimpleNamespace()
ap.model_mgr = _create_runtime_model_mgr()
ap.model_mgr.provider_dict = {} # Empty
mock_result = _create_mock_result([])
@@ -787,11 +862,12 @@ class TestEmbeddingModelsServiceCreateEmbeddingModel:
# Execute & Verify
with pytest.raises(Exception, match='provider not found'):
await service.create_embedding_model(
WORKSPACE_UUID,
{
'name': 'No Provider Embedding',
'provider_uuid': 'nonexistent',
'extra_args': {},
}
},
)
@@ -803,7 +879,7 @@ class TestEmbeddingModelsServiceDeleteEmbeddingModel:
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.model_mgr = SimpleNamespace()
ap.model_mgr = _create_runtime_model_mgr()
ap.model_mgr.remove_embedding_model = AsyncMock()
ap.persistence_mgr.execute_async = AsyncMock()
@@ -811,7 +887,7 @@ class TestEmbeddingModelsServiceDeleteEmbeddingModel:
service = EmbeddingModelsService(ap)
# Execute
await service.delete_embedding_model('delete-embedding-uuid')
await service.delete_embedding_model(WORKSPACE_UUID, 'delete-embedding-uuid')
# Verify
ap.model_mgr.remove_embedding_model.assert_called_once()
@@ -832,7 +908,9 @@ class TestRerankModelsServiceGetRerankModels:
service = RerankModelsService(ap)
# Execute
result = await service.get_rerank_models()
result = await service.get_rerank_models(
WORKSPACE_UUID,
)
# Verify
assert result == []
@@ -869,7 +947,9 @@ class TestRerankModelsServiceGetRerankModels:
service = RerankModelsService(ap)
# Execute
result = await service.get_rerank_models()
result = await service.get_rerank_models(
WORKSPACE_UUID,
)
# Verify
assert len(result) == 1
@@ -909,7 +989,7 @@ class TestRerankModelsServiceGetRerankModel:
service = RerankModelsService(ap)
# Execute
result = await service.get_rerank_model('found-rerank')
result = await service.get_rerank_model(WORKSPACE_UUID, 'found-rerank')
# Verify
assert result is not None
@@ -926,7 +1006,7 @@ class TestRerankModelsServiceGetRerankModel:
service = RerankModelsService(ap)
# Execute
result = await service.get_rerank_model('nonexistent-rerank')
result = await service.get_rerank_model(WORKSPACE_UUID, 'nonexistent-rerank')
# Verify
assert result is None
@@ -940,7 +1020,7 @@ class TestRerankModelsServiceCreateRerankModel:
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.model_mgr = SimpleNamespace()
ap.model_mgr = _create_runtime_model_mgr()
ap.model_mgr.provider_dict = {'provider-uuid': Mock()}
ap.model_mgr.rerank_models = []
ap.model_mgr.load_rerank_model_with_provider = AsyncMock(return_value=Mock())
@@ -952,11 +1032,12 @@ class TestRerankModelsServiceCreateRerankModel:
# Execute
model_uuid = await service.create_rerank_model(
WORKSPACE_UUID,
{
'name': 'New Rerank',
'provider_uuid': 'provider-uuid',
'extra_args': {},
}
},
)
# Verify
@@ -967,7 +1048,7 @@ class TestRerankModelsServiceCreateRerankModel:
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.model_mgr = SimpleNamespace()
ap.model_mgr = _create_runtime_model_mgr()
ap.model_mgr.provider_dict = {}
mock_result = _create_mock_result([])
@@ -978,11 +1059,12 @@ class TestRerankModelsServiceCreateRerankModel:
# Execute & Verify
with pytest.raises(Exception, match='provider not found'):
await service.create_rerank_model(
WORKSPACE_UUID,
{
'name': 'No Provider Rerank',
'provider_uuid': 'nonexistent',
'extra_args': {},
}
},
)
@@ -994,7 +1076,7 @@ class TestRerankModelsServiceDeleteRerankModel:
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.model_mgr = SimpleNamespace()
ap.model_mgr = _create_runtime_model_mgr()
ap.model_mgr.remove_rerank_model = AsyncMock()
ap.persistence_mgr.execute_async = AsyncMock()
@@ -1002,7 +1084,7 @@ class TestRerankModelsServiceDeleteRerankModel:
service = RerankModelsService(ap)
# Execute
await service.delete_rerank_model('delete-rerank-uuid')
await service.delete_rerank_model(WORKSPACE_UUID, 'delete-rerank-uuid')
# Verify
ap.model_mgr.remove_rerank_model.assert_called_once()
@@ -1027,7 +1109,7 @@ class TestEmbeddingModelsServiceGetEmbeddingModelsByProvider:
service = EmbeddingModelsService(ap)
# Execute
result = await service.get_embedding_models_by_provider('provider-uuid')
result = await service.get_embedding_models_by_provider(WORKSPACE_UUID, 'provider-uuid')
# Verify
assert len(result) == 2
@@ -1052,7 +1134,7 @@ class TestRerankModelsServiceGetRerankModelsByProvider:
service = RerankModelsService(ap)
# Execute
result = await service.get_rerank_models_by_provider('provider-uuid')
result = await service.get_rerank_models_by_provider(WORKSPACE_UUID, 'provider-uuid')
# Verify
assert len(result) == 2
@@ -1066,39 +1148,102 @@ class TestValidateProviderSupports:
"""Build a fake ap whose model_mgr resolves a manifest with support_type."""
manifest = SimpleNamespace(spec={'support_type': support_type})
runtime_provider = SimpleNamespace(provider_entity=SimpleNamespace(requester=requester_name))
model_mgr = SimpleNamespace(
provider_dict={'p1': runtime_provider},
get_available_requester_manifest_by_name=lambda name: manifest if name == requester_name else None,
)
model_mgr = _create_runtime_model_mgr()
model_mgr.provider_dict = {'p1': runtime_provider}
model_mgr.get_available_requester_manifest_by_name = lambda name: manifest if name == requester_name else None
return SimpleNamespace(model_mgr=model_mgr)
async def test_allows_supported_type(self):
ap = self._make_ap('cohere-rerank', ['rerank'])
# Should not raise
await _validate_provider_supports(ap, 'p1', 'rerank')
await _validate_provider_supports(ap, WORKSPACE_UUID, 'p1', 'rerank')
async def test_rejects_unsupported_type(self):
ap = self._make_ap('cohere-rerank', ['rerank'])
with pytest.raises(ValueError, match='does not support llm'):
await _validate_provider_supports(ap, 'p1', 'llm')
await _validate_provider_supports(ap, WORKSPACE_UUID, 'p1', 'llm')
async def test_allows_when_support_type_missing(self):
# Manifest without support_type must not block (backward compatible)
manifest = SimpleNamespace(spec={})
runtime_provider = SimpleNamespace(provider_entity=SimpleNamespace(requester='legacy'))
model_mgr = SimpleNamespace(
provider_dict={'p1': runtime_provider},
get_available_requester_manifest_by_name=lambda name: manifest,
)
model_mgr = _create_runtime_model_mgr()
model_mgr.provider_dict = {'p1': runtime_provider}
model_mgr.get_available_requester_manifest_by_name = lambda name: manifest
ap = SimpleNamespace(model_mgr=model_mgr)
await _validate_provider_supports(ap, 'p1', 'rerank')
await _validate_provider_supports(ap, WORKSPACE_UUID, 'p1', 'rerank')
async def test_allows_when_provider_unknown(self):
ap = self._make_ap('cohere-rerank', ['rerank'])
# Unknown provider uuid -> no entry -> no block
await _validate_provider_supports(ap, 'missing', 'llm')
await _validate_provider_supports(ap, WORKSPACE_UUID, 'missing', 'llm')
async def test_degrades_when_model_mgr_incomplete(self):
# A bare ap without a usable model_mgr must not raise (defensive)
ap = SimpleNamespace(model_mgr=SimpleNamespace())
await _validate_provider_supports(ap, 'p1', 'llm')
await _validate_provider_supports(ap, WORKSPACE_UUID, 'p1', 'llm')
class TestModelSecretRoundtrip:
async def test_provider_filtered_list_redacts_extra_args_without_mutating_source(self):
model = _create_mock_llm_model(extra_args={'headers': {'Authorization': 'Bearer secret'}})
raw = {
'uuid': model.uuid,
'provider_uuid': model.provider_uuid,
'extra_args': {'headers': {'Authorization': 'Bearer secret'}},
}
ap = SimpleNamespace(
persistence_mgr=SimpleNamespace(
execute_async=AsyncMock(return_value=_create_mock_result([model])),
serialize_model=Mock(return_value=raw),
)
)
service = LLMModelsService(ap)
redacted = await service.get_llm_models_by_provider(WORKSPACE_UUID, model.provider_uuid)
unredacted = await service.get_llm_models_by_provider(
WORKSPACE_UUID,
model.provider_uuid,
include_secret=True,
)
assert redacted[0]['extra_args']['headers']['Authorization'] == '***'
assert unredacted[0]['extra_args']['headers']['Authorization'] == 'Bearer secret'
assert raw['extra_args']['headers']['Authorization'] == 'Bearer secret'
async def test_masked_extra_args_update_restores_existing_header(self):
existing = _existing_llm_data()
existing['extra_args'] = {
'headers': {'Authorization': 'Bearer secret', 'X-API-Key': 'key-secret'},
'timeout': 30,
}
runtime_provider = SimpleNamespace(provider_entity=SimpleNamespace(requester=None))
write_result = Mock(rowcount=1)
model_mgr = _create_runtime_model_mgr()
model_mgr.provider_dict = {'provider-uuid': runtime_provider}
ap = SimpleNamespace(
persistence_mgr=SimpleNamespace(execute_async=AsyncMock(return_value=write_result)),
model_mgr=model_mgr,
)
service = LLMModelsService(ap)
service.get_llm_model = AsyncMock(return_value=existing)
await service.update_llm_model(
WORKSPACE_UUID,
'existing-uuid',
{
'extra_args': {
'headers': {'Authorization': '***', 'X-API-Key': ''},
'timeout': 60,
}
},
)
statement = ap.persistence_mgr.execute_async.await_args.args[0]
stored_extra_args = next(
value.value for column, value in statement._values.items() if column.key == 'extra_args'
)
assert stored_extra_args == {
'headers': {'Authorization': 'Bearer secret', 'X-API-Key': ''},
'timeout': 60,
}
tests/unit_tests/api/service/test_monitoring_tenancy.py
+392
View File
@@ -0,0 +1,392 @@
from __future__ import annotations
import datetime
from types import SimpleNamespace
import pytest
import sqlalchemy
from sqlalchemy.ext.asyncio import create_async_engine
from langbot.pkg.api.http.authz import WorkspaceRequiredError
from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.api.http.service.monitoring import MonitoringService
from langbot.pkg.entity.persistence.base import Base
from langbot.pkg.entity.persistence.monitoring import MonitoringLLMCall, MonitoringMessage
from langbot.pkg.entity.persistence.workspace import Workspace
from langbot.pkg.persistence.mgr import PersistenceManager
pytestmark = pytest.mark.asyncio
WORKSPACE_A = '00000000-0000-0000-0000-00000000000a'
WORKSPACE_B = '00000000-0000-0000-0000-00000000000b'
def _context(workspace_uuid: str) -> ExecutionContext:
return ExecutionContext(
instance_uuid='instance',
workspace_uuid=workspace_uuid,
placement_generation=3,
bot_uuid='same-bot',
pipeline_uuid='same-pipeline',
)
class _PersistenceManager:
def __init__(self, engine):
self.engine = engine
async def execute_async(self, *args, **kwargs):
async with self.engine.connect() as connection:
result = await connection.execute(*args, **kwargs)
await connection.commit()
return result
def get_db_engine(self):
return self.engine
@staticmethod
def serialize_model(model, data, masked_columns=None):
return {
column.name: (
getattr(data, column.name).isoformat()
if isinstance(getattr(data, column.name), datetime.datetime)
else getattr(data, column.name)
)
for column in model.__table__.columns
if column.name not in (masked_columns or [])
}
@pytest.fixture
async def service(tmp_path):
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "monitoring.db"}')
async with engine.begin() as connection:
await connection.run_sync(Base.metadata.create_all)
await connection.execute(
sqlalchemy.insert(Workspace),
[
{
'uuid': WORKSPACE_A,
'instance_uuid': 'instance',
'name': 'A',
'slug': 'a',
'source': 'cloud_projection',
},
{
'uuid': WORKSPACE_B,
'instance_uuid': 'instance',
'name': 'B',
'slug': 'b',
'source': 'cloud_projection',
},
],
)
application = SimpleNamespace(
persistence_mgr=_PersistenceManager(engine),
instance_config=SimpleNamespace(data={'database': {'use': 'sqlite'}}),
)
yield MonitoringService(application)
await engine.dispose()
async def _record_message(service, context, content):
return await service.record_message(
context,
bot_id='same-bot',
bot_name='Same Bot',
pipeline_id='same-pipeline',
pipeline_name='Same Pipeline',
message_content=content,
session_id='same-session',
)
async def test_monitoring_write_without_execution_context_fails_closed(service):
with pytest.raises(WorkspaceRequiredError):
await _record_message(service, None, 'unscoped')
async def test_same_session_and_resource_ids_do_not_collide(service):
context_a = _context(WORKSPACE_A)
context_b = _context(WORKSPACE_B)
message_a = await _record_message(service, context_a, 'tenant-a')
message_b = await _record_message(service, context_b, 'tenant-b')
await service.record_session_start(
context_a,
session_id='same-session',
bot_id='same-bot',
bot_name='Same Bot',
pipeline_id='same-pipeline',
pipeline_name='Same Pipeline',
)
await service.record_session_start(
context_b,
session_id='same-session',
bot_id='same-bot',
bot_name='Same Bot',
pipeline_id='same-pipeline',
pipeline_name='Same Pipeline',
)
messages_a, total_a = await service.get_messages(context_a)
messages_b, total_b = await service.get_messages(context_b)
assert total_a == total_b == 1
assert messages_a[0]['message_content'] == 'tenant-a'
assert messages_b[0]['message_content'] == 'tenant-b'
assert (await service.get_message_details(context_b, message_a))['found'] is False
assert (await service.get_message_details(context_a, message_b))['found'] is False
async def test_tool_call_inherits_context_from_connection_message_row(service):
context = _context(WORKSPACE_A)
message_id = await _record_message(service, context, 'tool context')
await service.record_tool_call(
context,
tool_name='search',
tool_source='native',
duration=12,
message_id=message_id,
)
tool_calls, total = await service.get_tool_calls(context)
assert total == 1
assert tool_calls[0]['bot_id'] == 'same-bot'
assert tool_calls[0]['pipeline_id'] == 'same-pipeline'
assert tool_calls[0]['session_id'] == 'same-session'
assert tool_calls[0]['message_id'] == message_id
async def test_feedback_upsert_and_cancel_are_workspace_scoped(service):
context_a = _context(WORKSPACE_A)
context_b = _context(WORKSPACE_B)
await service.record_feedback(context_a, feedback_id='same-feedback', feedback_type=1)
await service.record_feedback(context_b, feedback_id='same-feedback', feedback_type=2)
stats_a = await service.get_feedback_stats(context_a)
stats_b = await service.get_feedback_stats(context_b)
assert stats_a['total_likes'] == 1
assert stats_a['total_dislikes'] == 0
assert stats_b['total_likes'] == 0
assert stats_b['total_dislikes'] == 1
await service.record_feedback(context_a, feedback_id='same-feedback', feedback_type=3)
assert (await service.get_feedback_stats(context_a))['total_feedback'] == 0
assert (await service.get_feedback_stats(context_b))['total_feedback'] == 1
async def test_monitoring_queries_and_detail_views_are_strictly_bounded(service):
context = _context(WORKSPACE_A)
service.ap.instance_config.data['monitoring'] = {
'query_limits': {
'page_rows': 2,
'export_rows': 2,
'detail_rows': 2,
'timeseries_buckets': 2,
'max_offset': 10,
}
}
await service.record_session_start(
context,
session_id='same-session',
bot_id='same-bot',
bot_name='Same Bot',
pipeline_id='same-pipeline',
pipeline_name='Same Pipeline',
)
message_ids = [await _record_message(service, context, f'message-{index}') for index in range(4)]
for index in range(3):
await service.record_llm_call(
context,
bot_id='same-bot',
bot_name='Same Bot',
pipeline_id='same-pipeline',
pipeline_name='Same Pipeline',
session_id='same-session',
model_name='model',
input_tokens=1,
output_tokens=2,
duration=10,
message_id=message_ids[0],
)
await service.record_tool_call(
context,
tool_name=f'tool-{index}',
tool_source='native',
duration=5,
session_id='same-session',
message_id=message_ids[0],
)
await service.record_error(
context,
bot_id='same-bot',
bot_name='Same Bot',
pipeline_id='same-pipeline',
pipeline_name='Same Pipeline',
error_type='Failure',
error_message=f'error-{index}',
session_id='same-session',
message_id=message_ids[0],
)
page, total = await service.get_messages(context, limit=100000, offset=-5)
exported = await service.export_messages(context, limit=100000)
session_detail = await service.get_session_analysis(context, 'same-session')
message_detail = await service.get_message_details(context, message_ids[0])
assert total == 4
assert len(page) == 2
assert len(exported) == 2
assert session_detail['message_stats']['total'] == 4
assert session_detail['llm_stats']['total_calls'] == 3
assert session_detail['tool_stats']['total_calls'] == 3
assert len(session_detail['tool_calls']) == 2
assert len(session_detail['errors']) == 2
assert session_detail['detail_truncated'] == {
'tool_calls': True,
'errors': True,
}
assert message_detail['llm_stats']['total_calls'] == 3
assert len(message_detail['llm_calls']) == 2
assert len(message_detail['errors']) == 2
assert message_detail['detail_truncated'] == {
'llm_calls': True,
'errors': True,
}
service.ap.instance_config.data['monitoring']['query_limits'] = {
'page_rows': 999999,
'export_rows': 999999,
'detail_rows': 999999,
'timeseries_buckets': 999999,
'max_offset': 99999999,
}
assert service.normalize_page_window(999999, 99999999) == (5000, 10000000)
assert service.normalize_export_limit(999999) == 50000
assert service._detail_limit() == 10000
assert service._timeseries_bucket_limit() == 10000
async def test_token_statistics_aggregate_and_limit_groups_in_database(service):
context = _context(WORKSPACE_A)
service.ap.instance_config.data['monitoring'] = {
'query_limits': {
'page_rows': 1,
'timeseries_buckets': 2,
}
}
first_hour = datetime.datetime(2026, 7, 28, 10, 0)
rows = [
{
'id': f'llm-{index}',
'workspace_uuid': WORKSPACE_A,
'timestamp': first_hour + datetime.timedelta(hours=hour, minutes=index),
'model_name': model,
'input_tokens': input_tokens,
'output_tokens': output_tokens,
'total_tokens': input_tokens + output_tokens,
'duration': 100,
'cost': 0.01,
'status': 'success',
'bot_id': 'same-bot',
'bot_name': 'Same Bot',
'pipeline_id': 'same-pipeline',
'pipeline_name': 'Same Pipeline',
'session_id': 'same-session',
}
for index, (hour, model, input_tokens, output_tokens) in enumerate(
[
(0, 'small-model', 1, 2),
(1, 'large-model', 3, 4),
(2, 'large-model', 5, 6),
(2, 'large-model', 7, 8),
]
)
]
await service.ap.persistence_mgr.execute_async(sqlalchemy.insert(MonitoringLLMCall), rows)
stats = await service.get_token_statistics(context, bucket='hour')
assert stats['summary']['total_calls'] == 4
assert stats['summary']['total_tokens'] == 36
assert stats['by_model_truncated'] is True
assert [model['model_name'] for model in stats['by_model']] == ['large-model']
assert stats['timeseries_truncated'] is True
assert stats['timeseries'] == [
{
'bucket': '2026-07-28 11:00',
'input_tokens': 3,
'output_tokens': 4,
'total_tokens': 7,
'calls': 1,
},
{
'bucket': '2026-07-28 12:00',
'input_tokens': 12,
'output_tokens': 14,
'total_tokens': 26,
'calls': 2,
},
]
async def test_cleanup_commits_sqlite_delete_before_vacuum(tmp_path):
engine = create_async_engine(
f'sqlite+aiosqlite:///{tmp_path / "monitoring-cleanup.db"}',
connect_args={'timeout': 0.1},
)
application = SimpleNamespace(
instance_config=SimpleNamespace(data={'database': {'use': 'sqlite'}}),
)
manager = PersistenceManager(application)
manager.db = SimpleNamespace(get_engine=lambda: engine)
application.persistence_mgr = manager
try:
async with engine.begin() as connection:
await connection.run_sync(Base.metadata.create_all)
await connection.execute(
sqlalchemy.insert(Workspace).values(
uuid=WORKSPACE_A,
instance_uuid='instance',
name='A',
slug='a',
source='cloud_projection',
)
)
await connection.execute(
sqlalchemy.insert(MonitoringMessage),
[
{
'id': f'expired-message-{index}',
'workspace_uuid': WORKSPACE_A,
'timestamp': datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)
- datetime.timedelta(days=30),
'bot_id': 'bot',
'bot_name': 'Bot',
'pipeline_id': 'pipeline',
'pipeline_name': 'Pipeline',
'message_content': 'expired',
'session_id': 'session',
'status': 'success',
'level': 'info',
}
for index in range(5)
],
)
deleted = await MonitoringService(application).cleanup_expired_records(
_context(WORKSPACE_A),
retention_days=1,
batch_size=2,
max_batches_per_table=1,
)
assert deleted['monitoring_messages'] == 2
async with engine.connect() as connection:
remaining = await connection.scalar(
sqlalchemy.select(sqlalchemy.func.count()).select_from(MonitoringMessage)
)
assert remaining == 3
finally:
await engine.dispose()
tests/unit_tests/api/service/test_pipeline_service.py
+112 -25
View File
@@ -21,10 +21,13 @@ import json
from langbot.pkg.api.http.service.pipeline import PipelineService, default_stage_order
from langbot.pkg.entity.persistence.pipeline import LegacyPipeline
from langbot.pkg.workspace.errors import WorkspaceNotFoundError
pytestmark = pytest.mark.asyncio
WORKSPACE_UUID = 'workspace-a'
def _create_mock_pipeline(
pipeline_uuid: str = None,
@@ -77,7 +80,9 @@ class TestPipelineServiceGetPipelineMetadata:
service = PipelineService(ap)
# Execute
result = await service.get_pipeline_metadata()
result = await service.get_pipeline_metadata(
WORKSPACE_UUID,
)
# Verify
assert len(result) == 4
@@ -107,7 +112,9 @@ class TestPipelineServiceGetPipelines:
service = PipelineService(ap)
# Execute
result = await service.get_pipelines()
result = await service.get_pipelines(
WORKSPACE_UUID,
)
# Verify
assert result == []
@@ -133,7 +140,9 @@ class TestPipelineServiceGetPipelines:
service = PipelineService(ap)
# Execute
result = await service.get_pipelines()
result = await service.get_pipelines(
WORKSPACE_UUID,
)
# Verify
assert len(result) == 2
@@ -152,7 +161,7 @@ class TestPipelineServiceGetPipelines:
service = PipelineService(ap)
# Execute
await service.get_pipelines(sort_by='updated_at', sort_order='ASC')
await service.get_pipelines(WORKSPACE_UUID, sort_by='updated_at', sort_order='ASC')
# Verify - execute was called with sort parameters
ap.persistence_mgr.execute_async.assert_called_once()
@@ -181,7 +190,7 @@ class TestPipelineServiceGetPipeline:
service = PipelineService(ap)
# Execute
result = await service.get_pipeline('test-uuid')
result = await service.get_pipeline(WORKSPACE_UUID, 'test-uuid')
# Verify
assert result is not None
@@ -200,7 +209,7 @@ class TestPipelineServiceGetPipeline:
service = PipelineService(ap)
# Execute
result = await service.get_pipeline('nonexistent-uuid')
result = await service.get_pipeline(WORKSPACE_UUID, 'nonexistent-uuid')
# Verify
assert result is None
@@ -229,7 +238,7 @@ class TestPipelineServiceCreatePipeline:
# Execute & Verify
with pytest.raises(ValueError, match='Maximum number of pipelines'):
await service.create_pipeline({'name': 'New Pipeline'})
await service.create_pipeline(WORKSPACE_UUID, {'name': 'New Pipeline'})
async def test_create_pipeline_no_limit(self):
"""Creates pipeline without limit when max_pipelines=-1."""
@@ -258,7 +267,7 @@ class TestPipelineServiceCreatePipeline:
with patch(
'langbot.pkg.utils.paths.get_resource_path', return_value='templates/default-pipeline-config.json'
):
bot_uuid = await service.create_pipeline({'name': 'New Pipeline'})
bot_uuid = await service.create_pipeline(WORKSPACE_UUID, {'name': 'New Pipeline'})
# Verify
assert bot_uuid is not None
@@ -293,7 +302,7 @@ class TestPipelineServiceCreatePipeline:
with patch(
'langbot.pkg.utils.paths.get_resource_path', return_value='templates/default-pipeline-config.json'
):
await service.create_pipeline({'name': 'Default Pipeline'}, default=True)
await service.create_pipeline(WORKSPACE_UUID, {'name': 'Default Pipeline'}, default=True)
# Verify - execute was called
ap.persistence_mgr.execute_async.assert_called()
@@ -340,7 +349,7 @@ class TestPipelineServiceCreatePipeline:
with patch(
'langbot.pkg.utils.paths.get_resource_path', return_value='templates/default-pipeline-config.json'
):
await service.create_pipeline({'name': 'New Pipeline'})
await service.create_pipeline(WORKSPACE_UUID, {'name': 'New Pipeline'})
assert len(insert_params) == 1
assert insert_params[0]['extensions_preferences'] == {
@@ -394,7 +403,7 @@ class TestPipelineServiceUpdatePipeline:
'is_default': True,
'description': 'New description', # Not name change, so no bot_service needed
}
await service.update_pipeline('test-uuid', pipeline_data)
await service.update_pipeline(WORKSPACE_UUID, 'test-uuid', pipeline_data)
update_params = ap.persistence_mgr.execute_async.await_args_list[0].args[0].compile().params
assert update_params['description'] == 'New description'
@@ -450,7 +459,7 @@ class TestPipelineServiceUpdatePipeline:
service.get_pipeline = AsyncMock(return_value={'uuid': 'test-uuid', 'name': 'New Name'})
# Execute with name change
await service.update_pipeline('test-uuid', {'name': 'New Name'})
await service.update_pipeline(WORKSPACE_UUID, 'test-uuid', {'name': 'New Name'})
# Verify - bot_service.update_bot was called for each bot
assert ap.bot_service.update_bot.call_count == 2
@@ -478,7 +487,7 @@ class TestPipelineServiceUpdatePipeline:
service.get_pipeline = AsyncMock(return_value={'uuid': 'test-uuid'})
# Execute
await service.update_pipeline('test-uuid', {'description': 'Updated'})
await service.update_pipeline(WORKSPACE_UUID, 'test-uuid', {'description': 'Updated'})
# Verify - conversation was cleared
assert session.using_conversation is None
@@ -499,10 +508,10 @@ class TestPipelineServiceDeletePipeline:
service = PipelineService(ap)
# Execute
await service.delete_pipeline('test-uuid')
await service.delete_pipeline(WORKSPACE_UUID, 'test-uuid')
# Verify
ap.pipeline_mgr.remove_pipeline.assert_called_once_with('test-uuid')
ap.pipeline_mgr.remove_pipeline.assert_called_once_with(WORKSPACE_UUID, 'test-uuid')
ap.persistence_mgr.execute_async.assert_called_once()
async def test_delete_pipeline_nonexistent_uuid(self):
@@ -517,7 +526,7 @@ class TestPipelineServiceDeletePipeline:
service = PipelineService(ap)
# Execute - should not raise
await service.delete_pipeline('nonexistent-uuid')
await service.delete_pipeline(WORKSPACE_UUID, 'nonexistent-uuid')
# Verify
ap.pipeline_mgr.remove_pipeline.assert_called_once()
@@ -549,7 +558,7 @@ class TestPipelineServiceCopyPipeline:
# Execute & Verify
with pytest.raises(ValueError, match='Maximum number of pipelines'):
await service.copy_pipeline('original-uuid')
await service.copy_pipeline(WORKSPACE_UUID, 'original-uuid')
async def test_copy_pipeline_not_found_raises(self):
"""Raises ValueError when original pipeline not found."""
@@ -570,8 +579,8 @@ class TestPipelineServiceCopyPipeline:
ap.persistence_mgr.serialize_model = Mock(return_value={})
# Execute & Verify
with pytest.raises(ValueError, match='Pipeline original-uuid not found'):
await service.copy_pipeline('original-uuid')
with pytest.raises(WorkspaceNotFoundError, match='Pipeline original-uuid not found'):
await service.copy_pipeline(WORKSPACE_UUID, 'original-uuid')
async def test_copy_pipeline_creates_copy(self):
"""Creates a copy with (Copy) suffix."""
@@ -614,7 +623,7 @@ class TestPipelineServiceCopyPipeline:
)
# Execute
new_uuid = await service.copy_pipeline('original-uuid')
new_uuid = await service.copy_pipeline(WORKSPACE_UUID, 'original-uuid')
# Verify
assert new_uuid is not None
@@ -647,7 +656,7 @@ class TestPipelineServiceCopyPipeline:
service.get_pipeline = AsyncMock(return_value={'uuid': 'copy-uuid', 'is_default': False})
# Execute
await service.copy_pipeline('original-uuid')
await service.copy_pipeline(WORKSPACE_UUID, 'original-uuid')
# Verify - pipeline_mgr.load_pipeline called (copy created)
ap.pipeline_mgr.load_pipeline.assert_called_once()
@@ -667,8 +676,8 @@ class TestPipelineServiceUpdatePipelineExtensions:
service = PipelineService(ap)
# Execute & Verify
with pytest.raises(ValueError, match='Pipeline nonexistent-uuid not found'):
await service.update_pipeline_extensions('nonexistent-uuid', [])
with pytest.raises(WorkspaceNotFoundError, match='Pipeline nonexistent-uuid not found'):
await service.update_pipeline_extensions(WORKSPACE_UUID, 'nonexistent-uuid', [])
async def test_update_extensions_sets_plugins(self):
"""Updates plugins in extensions_preferences."""
@@ -715,6 +724,7 @@ class TestPipelineServiceUpdatePipelineExtensions:
# Execute
bound_plugins = [{'plugin_uuid': 'plugin-1'}]
await service.update_pipeline_extensions(
WORKSPACE_UUID,
'test-uuid',
bound_plugins=bound_plugins,
enable_all_plugins=False,
@@ -764,6 +774,7 @@ class TestPipelineServiceUpdatePipelineExtensions:
# Execute
await service.update_pipeline_extensions(
WORKSPACE_UUID,
'test-uuid',
bound_plugins=[],
bound_mcp_servers=['mcp-server-1'],
@@ -811,7 +822,7 @@ class TestPipelineServiceUpdatePipelineExtensions:
)
# Execute - bound_mcp_servers is None (not provided)
await service.update_pipeline_extensions('test-uuid', bound_plugins=[])
await service.update_pipeline_extensions(WORKSPACE_UUID, 'test-uuid', bound_plugins=[])
# Verify - persistence was called
ap.persistence_mgr.execute_async.assert_called()
@@ -850,7 +861,7 @@ class TestPipelineServiceUpdatePipelineExtensions:
service = PipelineService(ap)
service.get_pipeline = AsyncMock(return_value={'uuid': 'test-uuid'})
await service.update_pipeline_extensions('test-uuid', bound_plugins=[])
await service.update_pipeline_extensions(WORKSPACE_UUID, 'test-uuid', bound_plugins=[])
assert original_pipeline.extensions_preferences['mcp_resource_agent_read_enabled'] is False
assert original_pipeline.extensions_preferences['mcp_resources'] == [
@@ -858,6 +869,82 @@ class TestPipelineServiceUpdatePipelineExtensions:
]
class TestPipelineSecretRoundtrip:
async def test_resource_view_redacts_runner_secrets_without_mutating_serialized_data(self):
raw = {
'uuid': 'pipeline-secret',
'config': {
'ai': {
'n8n': {
'webhook-url': 'https://hook.invalid/bearer-secret',
'headers': {'Authorization': 'Bearer secret'},
}
}
},
}
pipeline = _create_mock_pipeline(pipeline_uuid='pipeline-secret')
ap = SimpleNamespace(
persistence_mgr=SimpleNamespace(
execute_async=AsyncMock(return_value=_create_mock_result([pipeline])),
serialize_model=Mock(return_value=raw),
)
)
redacted = await PipelineService(ap).get_pipelines(WORKSPACE_UUID)
assert redacted[0]['config']['ai']['n8n']['webhook-url'] == '***'
assert redacted[0]['config']['ai']['n8n']['headers']['Authorization'] == '***'
assert raw['config']['ai']['n8n']['webhook-url'] == 'https://hook.invalid/bearer-secret'
async def test_masked_runner_config_update_restores_existing_secret(self):
raw_config = {
'ai': {
'n8n': {
'webhook-url': 'https://hook.invalid/bearer-secret',
'headers': {'Authorization': 'Bearer secret'},
'timeout': 30,
}
}
}
current_pipeline = {'uuid': 'pipeline-secret', 'config': raw_config}
write_result = Mock(rowcount=1)
ap = SimpleNamespace(
persistence_mgr=SimpleNamespace(execute_async=AsyncMock(return_value=write_result)),
pipeline_mgr=SimpleNamespace(remove_pipeline=AsyncMock(), load_pipeline=AsyncMock()),
sess_mgr=SimpleNamespace(session_list=[]),
)
service = PipelineService(ap)
service.get_pipeline = AsyncMock(side_effect=[current_pipeline, current_pipeline])
await service.update_pipeline(
WORKSPACE_UUID,
'pipeline-secret',
{
'config': {
'ai': {
'n8n': {
'webhook-url': '***',
'headers': {'Authorization': '***'},
'timeout': 60,
}
}
}
},
)
statement = ap.persistence_mgr.execute_async.await_args.args[0]
stored_config = next(value.value for column, value in statement._values.items() if column.key == 'config')
assert stored_config == {
'ai': {
'n8n': {
'webhook-url': 'https://hook.invalid/bearer-secret',
'headers': {'Authorization': 'Bearer secret'},
'timeout': 60,
}
}
}
class TestDefaultStageOrder:
"""Tests for default_stage_order constant."""
tests/unit_tests/api/service/test_provider_service.py
+129 -30
View File
@@ -19,10 +19,13 @@ from types import SimpleNamespace
from langbot.pkg.api.http.service.provider import ModelProviderService
from langbot.pkg.entity.persistence.model import ModelProvider, LLMModel, EmbeddingModel, RerankModel
from langbot.pkg.workspace.errors import WorkspaceNotFoundError
pytestmark = pytest.mark.asyncio
WORKSPACE_UUID = 'workspace-a'
def _create_mock_provider(
provider_uuid: str = 'test-provider-uuid',
@@ -86,7 +89,9 @@ class TestModelProviderServiceGetProviders:
service = ModelProviderService(ap)
# Execute
result = await service.get_providers()
result = await service.get_providers(
WORKSPACE_UUID,
)
# Verify
assert result == []
@@ -115,7 +120,9 @@ class TestModelProviderServiceGetProviders:
service = ModelProviderService(ap)
# Execute
result = await service.get_providers()
result = await service.get_providers(
WORKSPACE_UUID,
)
# Verify
assert len(result) == 2
@@ -143,7 +150,10 @@ class TestModelProviderServiceGetProviders:
service = ModelProviderService(ap)
# Execute
result = await service.get_providers()
result = await service.get_providers(
WORKSPACE_UUID,
include_secret=True,
)
# Verify - api_keys should be parsed from string
assert result[0]['api_keys'] == ['key1', 'key2']
@@ -169,11 +179,41 @@ class TestModelProviderServiceGetProviders:
service = ModelProviderService(ap)
# Execute
result = await service.get_providers()
result = await service.get_providers(
WORKSPACE_UUID,
)
# Verify - invalid JSON returns empty list
assert result[0]['api_keys'] == []
async def test_get_providers_masks_api_keys_for_resource_view(self):
ap = SimpleNamespace()
provider = _create_mock_provider(
api_keys=['first', 'second'],
base_url=(
'https://provider-user:provider-password@api.provider.invalid/v1?access_token=url-secret&region=sg'
),
)
ap.persistence_mgr = SimpleNamespace(
execute_async=AsyncMock(return_value=_create_mock_result([provider])),
serialize_model=Mock(
return_value={
'uuid': provider.uuid,
'name': provider.name,
'base_url': provider.base_url,
'api_keys': provider.api_keys,
}
),
)
result = await ModelProviderService(ap).get_providers(
WORKSPACE_UUID,
include_secret=False,
)
assert result[0]['api_keys'] == ['***', '***']
assert result[0]['base_url'] == ('https://***@api.provider.invalid/v1?access_token=***&region=sg')
class TestModelProviderServiceGetProvider:
"""Tests for get_provider method."""
@@ -199,7 +239,7 @@ class TestModelProviderServiceGetProvider:
service = ModelProviderService(ap)
# Execute
result = await service.get_provider('found-uuid')
result = await service.get_provider(WORKSPACE_UUID, 'found-uuid')
# Verify
assert result is not None
@@ -217,7 +257,7 @@ class TestModelProviderServiceGetProvider:
service = ModelProviderService(ap)
# Execute
result = await service.get_provider('nonexistent-uuid')
result = await service.get_provider(WORKSPACE_UUID, 'nonexistent-uuid')
# Verify
assert result is None
@@ -239,6 +279,7 @@ class TestModelProviderServiceCreateProvider:
runtime_provider.provider_entity = Mock()
runtime_provider.provider_entity.uuid = 'generated-uuid'
ap.model_mgr.load_provider = AsyncMock(return_value=runtime_provider)
ap.model_mgr.cache_provider = AsyncMock()
ap.persistence_mgr.execute_async = AsyncMock()
@@ -246,12 +287,13 @@ class TestModelProviderServiceCreateProvider:
# Execute
provider_uuid = await service.create_provider(
WORKSPACE_UUID,
{
'name': 'New Provider',
'requester': 'openai',
'base_url': 'https://api.openai.com',
'api_keys': ['key'],
}
},
)
# Verify - UUID is generated
@@ -270,6 +312,7 @@ class TestModelProviderServiceCreateProvider:
runtime_provider.provider_entity = Mock()
runtime_provider.provider_entity.uuid = 'runtime-uuid'
ap.model_mgr.load_provider = AsyncMock(return_value=runtime_provider)
ap.model_mgr.cache_provider = AsyncMock()
ap.persistence_mgr.execute_async = AsyncMock()
@@ -277,12 +320,13 @@ class TestModelProviderServiceCreateProvider:
# Execute
result_uuid = await service.create_provider(
WORKSPACE_UUID,
{
'name': 'Runtime Provider',
'requester': 'openai',
'base_url': 'https://api.openai.com',
'api_keys': ['key'],
}
},
)
# Verify - provider added to runtime dict and UUID generated
@@ -307,6 +351,7 @@ class TestModelProviderServiceUpdateProvider:
# Execute
await service.update_provider(
WORKSPACE_UUID,
'existing-uuid',
{
'uuid': 'should-be-removed', # Will be removed
@@ -315,7 +360,7 @@ class TestModelProviderServiceUpdateProvider:
)
# Verify - reload called
ap.model_mgr.reload_provider.assert_called_once_with('existing-uuid')
ap.model_mgr.reload_provider.assert_called_once_with(WORKSPACE_UUID, 'existing-uuid')
async def test_update_provider_reloads_runtime(self):
"""Reloads provider in runtime after update."""
@@ -330,7 +375,7 @@ class TestModelProviderServiceUpdateProvider:
service = ModelProviderService(ap)
# Execute
await service.update_provider('update-uuid', {'name': 'New Name'})
await service.update_provider(WORKSPACE_UUID, 'update-uuid', {'name': 'New Name'})
# Verify
ap.model_mgr.reload_provider.assert_called_once()
@@ -354,7 +399,7 @@ class TestModelProviderServiceDeleteProvider:
# Execute & Verify
with pytest.raises(ValueError, match='Cannot delete provider: LLM models'):
await service.delete_provider('provider-with-llm')
await service.delete_provider(WORKSPACE_UUID, 'provider-with-llm')
async def test_delete_provider_with_embedding_models_raises_error(self):
"""Raises ValueError when Embedding models reference provider."""
@@ -387,7 +432,7 @@ class TestModelProviderServiceDeleteProvider:
# Execute & Verify - should raise embedding error (LLM check passes, embedding check fails)
with pytest.raises(ValueError, match='Cannot delete provider: Embedding models'):
await service.delete_provider('provider-with-embedding')
await service.delete_provider(WORKSPACE_UUID, 'provider-with-embedding')
async def test_delete_provider_with_rerank_models_raises_error(self):
"""Raises ValueError when Rerank models reference provider."""
@@ -420,7 +465,7 @@ class TestModelProviderServiceDeleteProvider:
# Execute & Verify - should raise rerank error (LLM and embedding checks pass, rerank check fails)
with pytest.raises(ValueError, match='Cannot delete provider: Rerank models'):
await service.delete_provider('provider-with-rerank')
await service.delete_provider(WORKSPACE_UUID, 'provider-with-rerank')
async def test_delete_provider_no_models_success(self):
"""Deletes provider when no models reference it."""
@@ -439,10 +484,10 @@ class TestModelProviderServiceDeleteProvider:
service = ModelProviderService(ap)
# Execute
await service.delete_provider('provider-no-models')
await service.delete_provider(WORKSPACE_UUID, 'provider-no-models')
# Verify - delete and remove called
ap.model_mgr.remove_provider.assert_called_once_with('provider-no-models')
ap.model_mgr.remove_provider.assert_called_once_with(WORKSPACE_UUID, 'provider-no-models')
class TestModelProviderServiceGetProviderModelCounts:
@@ -476,9 +521,10 @@ class TestModelProviderServiceGetProviderModelCounts:
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
service = ModelProviderService(ap)
service.get_provider = AsyncMock(return_value={'uuid': 'provider-uuid'})
# Execute
result = await service.get_provider_model_counts('provider-uuid')
result = await service.get_provider_model_counts(WORKSPACE_UUID, 'provider-uuid')
# Verify
assert result['llm_count'] == 3
@@ -497,9 +543,10 @@ class TestModelProviderServiceGetProviderModelCounts:
ap.persistence_mgr.execute_async = AsyncMock(return_value=zero_result)
service = ModelProviderService(ap)
service.get_provider = AsyncMock(return_value={'uuid': 'empty-provider'})
# Execute
result = await service.get_provider_model_counts('empty-provider')
result = await service.get_provider_model_counts(WORKSPACE_UUID, 'empty-provider')
# Verify
assert result['llm_count'] == 0
@@ -530,6 +577,7 @@ class TestModelProviderServiceFindOrCreateProvider:
# Execute
result = await service.find_or_create_provider(
WORKSPACE_UUID,
requester='openai',
base_url='https://api.openai.com',
api_keys=['key1', 'key2'], # Same keys (sorted)
@@ -558,6 +606,7 @@ class TestModelProviderServiceFindOrCreateProvider:
# Execute with reversed key order
result = await service.find_or_create_provider(
WORKSPACE_UUID,
requester='openai',
base_url='https://api.openai.com',
api_keys=['key2', 'key1'], # Different order, should still match
@@ -578,6 +627,7 @@ class TestModelProviderServiceFindOrCreateProvider:
runtime_provider.provider_entity = Mock()
runtime_provider.provider_entity.uuid = None # Will be set by uuid.uuid4()
ap.model_mgr.load_provider = AsyncMock(return_value=runtime_provider)
ap.model_mgr.cache_provider = AsyncMock()
# Mock no existing providers
mock_result = _create_mock_result([])
@@ -587,6 +637,7 @@ class TestModelProviderServiceFindOrCreateProvider:
# Execute
result = await service.find_or_create_provider(
WORKSPACE_UUID,
requester='new-requester',
base_url='https://new.api.com',
api_keys=['new-key'],
@@ -610,6 +661,7 @@ class TestModelProviderServiceFindOrCreateProvider:
runtime_provider.provider_entity = Mock()
runtime_provider.provider_entity.uuid = 'parsed-url-uuid'
ap.model_mgr.load_provider = AsyncMock(return_value=runtime_provider)
ap.model_mgr.cache_provider = AsyncMock()
mock_result = _create_mock_result([])
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
@@ -618,6 +670,7 @@ class TestModelProviderServiceFindOrCreateProvider:
# Execute
result_uuid = await service.find_or_create_provider(
WORKSPACE_UUID,
requester='custom',
base_url='https://api.example.com/v1',
api_keys=['key'],
@@ -644,17 +697,20 @@ class TestModelProviderServiceUpdateSpaceModelProviderApiKeys:
service = ModelProviderService(ap)
# Execute
await service.update_space_model_provider_api_keys('space-api-key')
await service.update_space_model_provider_api_keys(WORKSPACE_UUID, 'space-api-key')
# Verify - update and reload called for Space provider UUID
ap.model_mgr.reload_provider.assert_called_once_with('00000000-0000-0000-0000-000000000000')
ap.model_mgr.reload_provider.assert_called_once_with(
WORKSPACE_UUID,
'00000000-0000-0000-0000-000000000000',
)
class TestModelProviderServiceScanProviderModels:
"""Tests for scan_provider_models method."""
async def test_scan_provider_not_found_raises_error(self):
"""Raises ValueError when provider not found."""
"""Raises a non-enumerating not-found error when provider is outside the Workspace."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
@@ -665,8 +721,8 @@ class TestModelProviderServiceScanProviderModels:
service = ModelProviderService(ap)
# Execute & Verify
with pytest.raises(ValueError, match='provider not found'):
await service.scan_provider_models('nonexistent-uuid')
with pytest.raises(WorkspaceNotFoundError, match='Provider not found'):
await service.scan_provider_models(WORKSPACE_UUID, 'nonexistent-uuid')
async def test_scan_provider_returns_models_list(self):
"""Returns scanned models list."""
@@ -718,7 +774,7 @@ class TestModelProviderServiceScanProviderModels:
service = ModelProviderService(ap)
# Execute
result = await service.scan_provider_models('scan-uuid')
result = await service.scan_provider_models(WORKSPACE_UUID, 'scan-uuid')
# Verify
assert 'models' in result
@@ -771,7 +827,7 @@ class TestModelProviderServiceScanProviderModels:
service = ModelProviderService(ap)
# Execute - filter for LLM only
result = await service.scan_provider_models('filter-uuid', model_type='llm')
result = await service.scan_provider_models(WORKSPACE_UUID, 'filter-uuid', model_type='llm')
# Verify - only LLM models returned
assert len(result['models']) == 1
@@ -806,11 +862,11 @@ class TestModelProviderServiceScanProviderModels:
ap.model_mgr.load_provider = AsyncMock(return_value=runtime_provider)
ap.llm_model_service.get_llm_models_by_provider = AsyncMock(return_value=[])
ap.embedding_models_service.get_embedding_models_by_provider = AsyncMock(return_value=[])
ap.rerank_models_service.get_rerank_models_by_provider = AsyncMock(
return_value=[{'name': 'Qwen3-Reranker-8B'}]
)
ap.rerank_models_service.get_rerank_models_by_provider = AsyncMock(return_value=[{'name': 'Qwen3-Reranker-8B'}])
result = await ModelProviderService(ap).scan_provider_models('rerank-scan-uuid', model_type='rerank')
result = await ModelProviderService(ap).scan_provider_models(
WORKSPACE_UUID, 'rerank-scan-uuid', model_type='rerank'
)
assert result['models'][0]['type'] == 'rerank'
assert result['models'][0]['already_added'] is True
@@ -848,7 +904,7 @@ class TestModelProviderServiceScanProviderModels:
# Execute & Verify
with pytest.raises(ValueError, match='current provider does not support model scanning'):
await service.scan_provider_models('no-scan-uuid')
await service.scan_provider_models(WORKSPACE_UUID, 'no-scan-uuid')
async def test_scan_provider_marks_already_added_models(self):
"""Marks models that are already added."""
@@ -898,7 +954,7 @@ class TestModelProviderServiceScanProviderModels:
service = ModelProviderService(ap)
# Execute
result = await service.scan_provider_models('already-added-uuid')
result = await service.scan_provider_models(WORKSPACE_UUID, 'already-added-uuid')
# Verify - existing model marked as already_added
existing_model = next(m for m in result['models'] if m['name'] == 'Existing Model')
@@ -906,3 +962,46 @@ class TestModelProviderServiceScanProviderModels:
new_model = next(m for m in result['models'] if m['name'] == 'New Model')
assert new_model['already_added'] is False
class TestProviderSecretRoundtrip:
async def test_masked_api_keys_update_preserves_existing_values(self):
write_result = Mock(rowcount=1)
ap = SimpleNamespace(
persistence_mgr=SimpleNamespace(execute_async=AsyncMock(return_value=write_result)),
model_mgr=SimpleNamespace(reload_provider=AsyncMock()),
)
service = ModelProviderService(ap)
service.get_provider = AsyncMock(
return_value={
'uuid': 'provider-secret',
'api_keys': ['first-secret', 'second-secret'],
}
)
await service.update_provider(
WORKSPACE_UUID,
'provider-secret',
{'name': 'Updated', 'api_keys': ['***', 'replacement-secret']},
)
statement = ap.persistence_mgr.execute_async.await_args.args[0]
stored_api_keys = next(value.value for column, value in statement._values.items() if column.key == 'api_keys')
assert stored_api_keys == ['first-secret', 'replacement-secret']
async def test_extra_masked_api_key_is_rejected(self):
ap = SimpleNamespace(
persistence_mgr=SimpleNamespace(execute_async=AsyncMock()),
model_mgr=SimpleNamespace(reload_provider=AsyncMock()),
)
service = ModelProviderService(ap)
service.get_provider = AsyncMock(return_value={'uuid': 'provider-secret', 'api_keys': ['only-secret']})
with pytest.raises(ValueError, match='no existing value'):
await service.update_provider(
WORKSPACE_UUID,
'provider-secret',
{'api_keys': ['***', '***']},
)
ap.persistence_mgr.execute_async.assert_not_awaited()
tests/unit_tests/api/service/test_secret_redaction.py
+88
View File
@@ -0,0 +1,88 @@
from __future__ import annotations
import copy
import pytest
from langbot.pkg.api.http.service.secrets import (
contains_secret_placeholder,
redact_secrets,
restore_secret_placeholders,
)
RAW_CONFIG = {
'apiKey': 'api-secret',
'dify_apikey': 'dify-secret',
'base_url': (
'https://service-user:service-password@api.invalid/v1'
'?api_key=query-secret&region=sg&X-Amz-Signature=signed-secret'
),
'nested': {
'headers': {
'Authorization': 'Bearer nested-secret',
'X-API-Key': 'header-secret',
'Accept': 'application/json',
},
'webhook-url': 'https://hooks.invalid/path?token=secret',
'public_key': 'public-material',
'tokenizer': 'not-a-secret',
},
'credentials': {'username': 'service-user', 'password': 'service-password'},
'secret_list': ['first-secret', {'value': 'second-secret'}],
'empty_secret': '',
'enabled': True,
}
def test_recursive_redaction_is_shape_preserving_and_does_not_mutate_source():
source = copy.deepcopy(RAW_CONFIG)
redacted = redact_secrets(source)
assert redacted['apiKey'] == '***'
assert redacted['dify_apikey'] == '***'
assert redacted['base_url'] == ('https://***@api.invalid/v1?api_key=***&region=sg&X-Amz-Signature=***')
assert redacted['nested']['headers'] == {
'Authorization': '***',
'X-API-Key': '***',
'Accept': 'application/json',
}
assert redacted['nested']['webhook-url'] == '***'
assert redacted['nested']['public_key'] == 'public-material'
assert redacted['nested']['tokenizer'] == 'not-a-secret'
assert redacted['credentials'] == {'username': '***', 'password': '***'}
assert redacted['secret_list'] == ['***', {'value': '***'}]
assert redacted['empty_secret'] == ''
assert redacted['enabled'] is True
assert source == RAW_CONFIG
def test_masked_roundtrip_preserves_existing_secrets_and_accepts_replace_and_clear():
submitted = redact_secrets(RAW_CONFIG)
submitted['enabled'] = False
submitted['apiKey'] = 'replacement-secret'
submitted['nested']['headers']['X-API-Key'] = ''
restored = restore_secret_placeholders(submitted, RAW_CONFIG)
assert restored['apiKey'] == 'replacement-secret'
assert restored['dify_apikey'] == 'dify-secret'
assert restored['nested']['headers']['Authorization'] == 'Bearer nested-secret'
assert restored['nested']['headers']['X-API-Key'] == ''
assert restored['nested']['webhook-url'] == RAW_CONFIG['nested']['webhook-url']
assert restored['base_url'] == RAW_CONFIG['base_url']
assert restored['enabled'] is False
assert RAW_CONFIG['apiKey'] == 'api-secret'
def test_new_or_extra_masked_secret_fails_closed():
assert contains_secret_placeholder({'headers': {'Authorization': '***'}})
assert contains_secret_placeholder({'base_url': 'https://***@api.invalid?token=***'})
with pytest.raises(ValueError, match='no existing value'):
restore_secret_placeholders({'api_key': '***'})
with pytest.raises(ValueError, match='no existing value'):
restore_secret_placeholders(
{'api_keys': ['***', '***']},
{'api_keys': ['existing']},
)
tests/unit_tests/api/service/test_space_service.py
+68 -3
View File
@@ -13,6 +13,10 @@ Source: src/langbot/pkg/api/http/service/space.py
from __future__ import annotations
from collections import OrderedDict
import json
from urllib.parse import parse_qs, urlsplit
import pytest
from unittest.mock import AsyncMock, Mock, patch, MagicMock
from types import SimpleNamespace
@@ -26,6 +30,23 @@ from langbot.pkg.entity.persistence.user import User
pytestmark = pytest.mark.asyncio
def _set_response_body(response: MagicMock, body: dict | str) -> None:
"""Configure an aiohttp-like streaming body on an HTTP response mock."""
raw_body = body.encode() if isinstance(body, str) else json.dumps(body).encode()
class Content:
async def iter_chunked(self, _chunk_size: int):
midpoint = max(len(raw_body) // 2, 1)
yield raw_body[:midpoint]
if midpoint < len(raw_body):
yield raw_body[midpoint:]
response.headers = {}
response.content = Content()
response.charset = 'utf-8'
def _create_mock_user(
email: str = 'test@example.com',
account_type: str = 'space',
@@ -73,7 +94,7 @@ class TestSpaceServiceGetOAuthAuthorizeUrl:
result = service.get_oauth_authorize_url('http://localhost/callback')
# Verify
assert 'redirect_uri=http://localhost/callback' in result
assert parse_qs(urlsplit(result).query)['redirect_uri'] == ['http://localhost/callback']
assert 'https://space.langbot.app/auth/authorize' in result
def test_get_oauth_authorize_url_with_state(self):
@@ -93,8 +114,9 @@ class TestSpaceServiceGetOAuthAuthorizeUrl:
result = service.get_oauth_authorize_url('http://localhost/callback', state='random_state')
# Verify
assert 'redirect_uri=http://localhost/callback' in result
assert 'state=random_state' in result
params = parse_qs(urlsplit(result).query)
assert params['redirect_uri'] == ['http://localhost/callback']
assert params['state'] == ['random_state']
def test_get_oauth_authorize_url_default_config(self):
"""Uses default OAuth URL when config not set."""
@@ -289,6 +311,40 @@ class TestSpaceServiceGetCredits:
# Verify - returns cached value without API call
assert result == 100
async def test_cached_credit_lookup_does_not_scan_all_users(self):
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace(data={})
ap.persistence_mgr = SimpleNamespace()
service = SpaceService(ap)
class AtMostOneStepOrderedDict(OrderedDict):
def __iter__(self):
iterator = super().__iter__()
yielded = False
def next_entry():
nonlocal yielded
if yielded:
raise AssertionError('credits cache scanned all users')
yielded = True
return next(iterator)
class AtMostOneStepIterator:
def __iter__(self):
return self
def __next__(self):
return next_entry()
return AtMostOneStepIterator()
now = time.time()
service._credits_cache = AtMostOneStepOrderedDict(
(f'user-{index}@example.com', (index, now)) for index in range(512)
)
assert await service.get_credits('user-511@example.com') == 511
async def test_get_credits_cache_expired_refreshes(self):
"""Refreshes expired cache."""
# Setup
@@ -403,6 +459,7 @@ class TestSpaceServiceRefreshToken:
},
}
)
_set_response_body(mock_response, mock_response.json.return_value)
with patch('langbot.pkg.api.http.service.space.httpclient.get_session') as mock_session:
mock_session_obj = MagicMock()
@@ -438,6 +495,7 @@ class TestSpaceServiceRefreshToken:
}
)
mock_response.text = AsyncMock(return_value='{"code":1,"msg":"Invalid refresh token"}')
_set_response_body(mock_response, mock_response.json.return_value)
with patch('langbot.pkg.api.http.service.space.httpclient.get_session') as mock_session:
mock_session_obj = MagicMock()
@@ -464,6 +522,7 @@ class TestSpaceServiceRefreshToken:
mock_response = MagicMock()
mock_response.status = 500
mock_response.text = AsyncMock(return_value='Internal Server Error')
_set_response_body(mock_response, mock_response.text.return_value)
with patch('langbot.pkg.api.http.service.space.httpclient.get_session') as mock_session:
mock_session_obj = MagicMock()
@@ -503,6 +562,7 @@ class TestSpaceServiceExchangeOAuthCode:
},
}
)
_set_response_body(mock_response, mock_response.json.return_value)
with patch('langbot.pkg.api.http.service.space.httpclient.get_session') as mock_session:
mock_session_obj = MagicMock()
@@ -532,6 +592,7 @@ class TestSpaceServiceExchangeOAuthCode:
mock_response.status = 200
mock_response.json = AsyncMock(return_value={'code': 1, 'msg': 'Invalid code'})
mock_response.text = AsyncMock(return_value='{"code":1,"msg":"Invalid code"}')
_set_response_body(mock_response, mock_response.json.return_value)
with patch('langbot.pkg.api.http.service.space.httpclient.get_session') as mock_session:
mock_session_obj = MagicMock()
@@ -570,6 +631,7 @@ class TestSpaceServiceGetUserInfoRaw:
},
}
)
_set_response_body(mock_response, mock_response.json.return_value)
with patch('langbot.pkg.api.http.service.space.httpclient.get_session') as mock_session:
mock_session_obj = MagicMock()
@@ -600,6 +662,7 @@ class TestSpaceServiceGetUserInfoRaw:
mock_response.status = 200
mock_response.json = AsyncMock(return_value={'code': 1, 'msg': 'Unauthorized'})
mock_response.text = AsyncMock(return_value='{"code":1,"msg":"Unauthorized"}')
_set_response_body(mock_response, mock_response.json.return_value)
with patch('langbot.pkg.api.http.service.space.httpclient.get_session') as mock_session:
mock_session_obj = MagicMock()
@@ -700,6 +763,7 @@ class TestSpaceServiceGetModels:
},
}
)
_set_response_body(mock_response, mock_response.json.return_value)
with patch('langbot.pkg.api.http.service.space.httpclient.get_session') as mock_session:
mock_session_obj = MagicMock()
@@ -730,6 +794,7 @@ class TestSpaceServiceGetModels:
mock_response.status = 200
mock_response.json = AsyncMock(return_value={'code': 1, 'msg': 'Unauthorized'})
mock_response.text = AsyncMock(return_value='{"code":1,"msg":"Unauthorized"}')
_set_response_body(mock_response, mock_response.json.return_value)
with patch('langbot.pkg.api.http.service.space.httpclient.get_session') as mock_session:
mock_session_obj = MagicMock()
tests/unit_tests/api/service/test_user_service.py
+382 -6
View File
@@ -14,17 +14,124 @@ Source: src/langbot/pkg/api/http/service/user.py
from __future__ import annotations
import pytest
import jwt
import datetime
from unittest.mock import AsyncMock, Mock
from types import SimpleNamespace
from langbot.pkg.api.http.service.user import UserService
from langbot.pkg.entity.persistence.user import User
from langbot.pkg.entity.errors.account import AccountEmailMismatchError
from langbot.pkg.api.http.service.user import (
ControlPlaneDirectoryRequiredError,
UserService,
)
from langbot.pkg.entity.persistence.user import AccountSource, AccountStatus, User
from langbot.pkg.entity.errors.account import (
AccountEmailMismatchError,
SpaceAccountBindingRequiredError,
SpaceAccountNotRegisteredError,
)
from langbot.pkg.utils.bounded_executor import BlockingWorkCapacityError
pytestmark = pytest.mark.asyncio
async def test_password_hashing_rejects_concurrent_waiters() -> None:
service = UserService(SimpleNamespace())
await service._password_hash_lock.acquire()
try:
with pytest.raises(
BlockingWorkCapacityError,
match='Password hashing capacity reached',
):
await service._hash_password('secret')
finally:
service._password_hash_lock.release()
class TestSpaceOAuthState:
async def test_login_state_is_opaque_single_use(self):
service = UserService(SimpleNamespace())
state = await service.issue_space_oauth_state('login')
assert state.count('.') == 0
assert await service.consume_space_oauth_state(state, 'login') is None
with pytest.raises(ValueError, match='Invalid or expired OAuth state'):
await service.consume_space_oauth_state(state, 'login')
async def test_bind_state_resolves_only_bound_active_account(self):
service = UserService(SimpleNamespace())
account = SimpleNamespace(uuid='account-a', status=AccountStatus.ACTIVE.value)
service.get_user_by_uuid = AsyncMock(return_value=account)
state = await service.issue_space_oauth_state('bind', account_uuid='account-a')
assert await service.consume_space_oauth_state(state, 'bind') is account
service.get_user_by_uuid.assert_awaited_once_with('account-a')
async def test_state_purpose_mismatch_is_rejected_and_consumed(self):
service = UserService(SimpleNamespace())
state = await service.issue_space_oauth_state('login')
with pytest.raises(ValueError, match='Invalid or expired OAuth state'):
await service.consume_space_oauth_state(state, 'bind')
with pytest.raises(ValueError, match='Invalid or expired OAuth state'):
await service.consume_space_oauth_state(state, 'login')
async def test_expired_state_is_rejected(self):
service = UserService(SimpleNamespace())
state = await service.issue_space_oauth_state('login')
digest = service._space_oauth_state_digest(state)
purpose, account_uuid, _, launch_workspace_uuid = service._space_oauth_states[digest]
service._space_oauth_states[digest] = (purpose, account_uuid, 0, launch_workspace_uuid)
with pytest.raises(ValueError, match='Invalid or expired OAuth state'):
await service.consume_space_oauth_state(state, 'login')
async def test_login_state_can_carry_launch_workspace_without_changing_normal_return(self):
service = UserService(SimpleNamespace())
state = await service.issue_space_oauth_state(
'login',
launch_workspace_uuid='workspace-a',
)
assert await service.consume_space_oauth_state(state, 'login') is None
state = await service.issue_space_oauth_state(
'login',
launch_workspace_uuid='workspace-a',
)
consumed = await service.consume_space_oauth_state_details(state, 'login')
assert consumed.account is None
assert consumed.launch_workspace_uuid == 'workspace-a'
async def test_issue_state_does_not_scan_all_live_states(self, monkeypatch):
service = UserService(SimpleNamespace())
for _ in range(512):
await service.issue_space_oauth_state('login')
class NoGlobalIterationDict(dict):
def __iter__(self):
raise AssertionError('OAuth state issuance scanned all live states')
def keys(self):
raise AssertionError('OAuth state issuance scanned all live states')
def items(self):
raise AssertionError('OAuth state issuance scanned all live states')
def values(self):
raise AssertionError('OAuth state issuance scanned all live states')
guarded_states = NoGlobalIterationDict(service._space_oauth_states)
monkeypatch.setattr(service, '_space_oauth_states', guarded_states)
state = await service.issue_space_oauth_state('login')
assert await service.consume_space_oauth_state(state, 'login') is None
assert len(guarded_states) == 512
def _create_mock_user(
email: str = 'test@example.com',
password: str = 'hashed_password',
@@ -34,6 +141,7 @@ def _create_mock_user(
"""Helper to create mock User entity."""
user = Mock(spec=User)
user.user = email
user.uuid = f'account-{email}'
user.password = password
user.account_type = account_type
user.space_account_uuid = space_account_uuid
@@ -102,6 +210,41 @@ class TestUserServiceIsInitialized:
assert result is False
class TestUserServiceGetLoginCapabilities:
"""Tests for public login capability discovery."""
async def test_uses_explicit_identity_discovery_scope(self):
discovery_result = Mock()
discovery_result.one = Mock(return_value=(1, 2))
discovery_session = SimpleNamespace(execute=AsyncMock(return_value=discovery_result))
class DiscoveryContext:
async def __aenter__(self):
return SimpleNamespace(session=discovery_session)
async def __aexit__(self, exc_type, exc, tb):
return False
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace(
current_session=Mock(return_value=None),
identity_discovery_uow=Mock(return_value=DiscoveryContext()),
execute_async=AsyncMock(side_effect=AssertionError('unscoped persistence access')),
)
ap.workspace_service = SimpleNamespace(instance_uuid='instance-a')
service = UserService(ap)
result = await service.get_login_capabilities()
assert result == {
'password_login_enabled': True,
'space_login_enabled': True,
}
ap.persistence_mgr.identity_discovery_uow.assert_called_once()
discovery_session.execute.assert_awaited_once()
ap.persistence_mgr.execute_async.assert_not_awaited()
class TestUserServiceGetUserByEmail:
"""Tests for get_user_by_email method."""
@@ -309,6 +452,50 @@ class TestUserServiceVerifyJwtToken:
with pytest.raises(Exception): # jwt.DecodeError or similar
await service.verify_jwt_token('invalid.token.here')
async def test_verify_jwt_token_rejects_foreign_audience(self):
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {'system': {'jwt': {'secret': 'test_secret', 'expire': 3600}}}
service = UserService(ap)
token = jwt.encode(
{
'user': 'verify@example.com',
'iss': 'langbot-core',
'aud': 'langbot-instance:another-instance',
'exp': datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(hours=1),
},
'test_secret',
algorithm='HS256',
)
with pytest.raises(jwt.InvalidAudienceError):
await service.verify_jwt_token(token)
async def test_verify_jwt_token_accepts_legacy_community_token_only_in_oss(self):
ap = SimpleNamespace()
ap.instance_config = SimpleNamespace()
ap.instance_config.data = {'system': {'jwt': {'secret': 'test_secret', 'expire': 3600}}}
ap.workspace_service = SimpleNamespace(
instance_uuid='instance-a',
policy=SimpleNamespace(multi_workspace_enabled=False),
)
service = UserService(ap)
legacy_token = jwt.encode(
{
'user': 'legacy@example.com',
'iss': 'LangBot-community',
'exp': datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(hours=1),
},
'test_secret',
algorithm='HS256',
)
assert await service.verify_jwt_token(legacy_token) == 'legacy@example.com'
ap.workspace_service.policy.multi_workspace_enabled = True
with pytest.raises(jwt.MissingRequiredClaimError):
await service.verify_jwt_token(legacy_token)
class TestUserServiceResetPassword:
"""Tests for reset_password method."""
@@ -476,6 +663,71 @@ class TestUserServiceCreateOrUpdateSpaceUser:
ap.persistence_mgr.execute_async.assert_called()
assert updated_user.space_account_uuid == 'existing-space-uuid'
async def test_cloud_login_updates_only_the_projected_space_account(self):
projected = SimpleNamespace(
uuid='projected-space-uuid',
user='Cloud Owner',
normalized_email='owner@example.com',
password='',
account_type='space',
status=AccountStatus.ACTIVE.value,
source=AccountSource.CLOUD_PROJECTION.value,
projection_revision=7,
space_account_uuid='projected-space-uuid',
)
persistence = SimpleNamespace(execute_async=AsyncMock())
ap = SimpleNamespace(
persistence_mgr=persistence,
workspace_service=SimpleNamespace(policy=SimpleNamespace(multi_workspace_enabled=True)),
)
service = UserService(ap)
service.get_user_by_space_account_uuid = AsyncMock(side_effect=[projected, projected])
result = await service.create_or_update_space_user(
space_account_uuid='projected-space-uuid',
email='OWNER@example.com',
access_token='access-token',
refresh_token='refresh-token',
api_key='api-key',
expires_in=3600,
)
assert result is projected
persistence.execute_async.assert_awaited_once()
async def test_cloud_login_never_creates_an_unprojected_account(self):
persistence = SimpleNamespace(execute_async=AsyncMock())
ap = SimpleNamespace(
persistence_mgr=persistence,
workspace_service=SimpleNamespace(policy=SimpleNamespace(multi_workspace_enabled=True)),
)
service = UserService(ap)
service.get_user_by_space_account_uuid = AsyncMock(return_value=None)
with pytest.raises(
ControlPlaneDirectoryRequiredError,
match='verified Cloud directory',
):
await service.create_or_update_space_user(
space_account_uuid='unknown-space-uuid',
email='unknown@example.com',
access_token='access-token',
refresh_token='refresh-token',
api_key='api-key',
expires_in=3600,
)
persistence.execute_async.assert_not_awaited()
async def test_cloud_invitation_registration_requires_space_identity(self):
ap = SimpleNamespace(
workspace_service=SimpleNamespace(policy=SimpleNamespace(multi_workspace_enabled=True)),
)
service = UserService(ap)
with pytest.raises(ControlPlaneDirectoryRequiredError, match='Space account'):
await service.register_invited_account('invite-token', 'member@example.com', 'password')
async def test_create_or_update_new_space_user_first_init(self):
"""Creates new Space user on first initialization."""
# Setup
@@ -522,8 +774,8 @@ class TestUserServiceCreateOrUpdateSpaceUser:
# Verify
assert result.space_account_uuid == 'new-space-uuid'
async def test_create_or_update_space_user_already_initialized_raises_error(self):
"""Raises AccountEmailMismatchError when system already initialized and user not found."""
async def test_create_or_update_space_user_already_initialized_reports_unknown_space_email(self):
"""Unknown Space email is distinct from an existing local Account collision."""
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
@@ -538,7 +790,7 @@ class TestUserServiceCreateOrUpdateSpaceUser:
service.is_initialized = AsyncMock(return_value=True) # Already initialized
# Execute & Verify
with pytest.raises(AccountEmailMismatchError):
with pytest.raises(SpaceAccountNotRegisteredError):
await service.create_or_update_space_user(
space_account_uuid='unknown-space-uuid',
email='unknown@example.com',
@@ -548,6 +800,78 @@ class TestUserServiceCreateOrUpdateSpaceUser:
expires_in=3600,
)
async def test_unknown_space_subject_cannot_claim_existing_account_by_email(self):
"""An OAuth login collision requires the explicit account-bound bind flow."""
existing_user = _create_mock_user(
email='owner@example.com',
account_type='local',
space_account_uuid=None,
)
ap = SimpleNamespace(
persistence_mgr=SimpleNamespace(execute_async=AsyncMock()),
provider_service=SimpleNamespace(update_space_model_provider_api_keys=AsyncMock()),
space_service=SimpleNamespace(
get_user_info_raw=AsyncMock(
return_value={
'account': {
'uuid': 'attacker-space-subject',
'email': 'owner@example.com',
},
'api_key': 'attacker-api-key',
}
)
),
)
service = UserService(ap)
service.get_user_by_space_account_uuid = AsyncMock(return_value=None)
service.get_user_by_email = AsyncMock(return_value=existing_user)
service.generate_jwt_token = AsyncMock(return_value='must-not-be-issued')
with pytest.raises(SpaceAccountBindingRequiredError):
await service.authenticate_space_user(
'attacker-access-token',
'attacker-refresh-token',
3600,
)
ap.persistence_mgr.execute_async.assert_not_awaited()
ap.provider_service.update_space_model_provider_api_keys.assert_not_awaited()
service.generate_jwt_token.assert_not_awaited()
async def test_oss_space_provider_refresh_requires_workspace_owner(self):
member_account = _create_mock_user(email='member@example.com', space_account_uuid='space-member')
access = SimpleNamespace(
workspace=SimpleNamespace(uuid='workspace-a'),
membership=SimpleNamespace(role='admin'),
)
provider_service = SimpleNamespace(update_space_model_provider_api_keys=AsyncMock())
ap = SimpleNamespace(
workspace_service=SimpleNamespace(policy=SimpleNamespace(multi_workspace_enabled=False)),
workspace_collaboration_service=SimpleNamespace(list_account_workspaces=AsyncMock(return_value=[access])),
provider_service=provider_service,
)
await UserService(ap)._update_space_provider_for_account(member_account, 'member-api-key')
provider_service.update_space_model_provider_api_keys.assert_not_awaited()
async def test_oss_space_provider_refresh_uses_workspace_owner_credentials(self):
owner_account = _create_mock_user(email='owner@example.com', space_account_uuid='space-owner')
access = SimpleNamespace(
workspace=SimpleNamespace(uuid='workspace-a'),
membership=SimpleNamespace(role='owner'),
)
provider_service = SimpleNamespace(update_space_model_provider_api_keys=AsyncMock())
ap = SimpleNamespace(
workspace_service=SimpleNamespace(policy=SimpleNamespace(multi_workspace_enabled=False)),
workspace_collaboration_service=SimpleNamespace(list_account_workspaces=AsyncMock(return_value=[access])),
provider_service=provider_service,
)
await UserService(ap)._update_space_provider_for_account(owner_account, 'owner-api-key')
provider_service.update_space_model_provider_api_keys.assert_awaited_once_with('workspace-a', 'owner-api-key')
async def test_create_or_update_space_user_no_expiry(self):
"""Creates Space user without token expiry."""
# Setup
@@ -594,6 +918,58 @@ class TestUserServiceCreateOrUpdateSpaceUser:
assert result is not None
assert result.space_account_uuid == 'noexpiry-uuid'
async def test_bind_space_account_rejects_different_email(self):
service = UserService(SimpleNamespace())
service.get_user_by_email = AsyncMock(return_value=_create_mock_user(email='invited@example.com'))
service.ap.space_service = SimpleNamespace(
exchange_oauth_code=AsyncMock(
return_value={'access_token': 'access', 'refresh_token': 'refresh', 'expires_in': 3600}
),
get_user_info_raw=AsyncMock(
return_value={
'account': {'uuid': 'space-other', 'email': 'other@example.com'},
'api_key': 'key',
}
),
)
service.get_user_by_space_account_uuid = AsyncMock(return_value=None)
service._identity_execute = AsyncMock()
with pytest.raises(AccountEmailMismatchError):
await service.bind_space_account('invited@example.com', 'code')
service._identity_execute.assert_not_awaited()
@pytest.mark.asyncio
async def test_get_workspace_owner_returns_user_object_from_core_connection_result(self):
service = UserService(SimpleNamespace())
owner = _create_mock_user('owner@example.com', password='pw')
service.ap.persistence_mgr = SimpleNamespace(
current_session=lambda: SimpleNamespace(scalar=AsyncMock(return_value=owner)),
)
resolved = await service.get_workspace_owner('workspace-1')
assert resolved is owner
class TestUserServiceLoginCapabilities:
async def test_capabilities_are_derived_from_all_accounts(self):
result = SimpleNamespace(one=lambda: (2, 1))
ap = SimpleNamespace(persistence_mgr=SimpleNamespace(execute_async=AsyncMock(return_value=result)))
capabilities = await UserService(ap).get_login_capabilities()
assert capabilities == {'password_login_enabled': True, 'space_login_enabled': True}
async def test_capabilities_disable_absent_login_methods(self):
result = SimpleNamespace(one=lambda: (0, 0))
ap = SimpleNamespace(persistence_mgr=SimpleNamespace(execute_async=AsyncMock(return_value=result)))
capabilities = await UserService(ap).get_login_capabilities()
assert capabilities == {'password_login_enabled': False, 'space_login_enabled': False}
class TestUserServiceCreateUserLock:
"""Tests for create_user_lock attribute."""
tests/unit_tests/api/service/test_webhook_service.py
+266 -30
View File
@@ -14,15 +14,23 @@ Source: src/langbot/pkg/api/http/service/webhook.py
from __future__ import annotations
import datetime
import pytest
import sqlalchemy
from sqlalchemy.ext.asyncio import create_async_engine
from unittest.mock import AsyncMock, Mock
from types import SimpleNamespace
from langbot.pkg.api.http.authz import WorkspaceRequiredError
from langbot.pkg.api.http.service.webhook import WebhookService
from langbot.pkg.entity.persistence.base import Base
from langbot.pkg.entity.persistence.webhook import Webhook
from langbot.pkg.entity.persistence.workspace import Workspace
pytestmark = pytest.mark.asyncio
WORKSPACE_UUID = 'workspace-a'
def _create_mock_webhook(
@@ -42,11 +50,20 @@ def _create_mock_webhook(
return webhook
def _create_mock_result(items: list = None, first_item=None):
def _create_mock_result(items: list = None, first_item=None, scalar_value=None):
"""Create mock result object for persistence queries."""
result = Mock()
result.all = Mock(return_value=items or [])
result.first = Mock(return_value=first_item)
result.scalar = Mock(return_value=scalar_value)
result.rowcount = 1
return result
def _create_write_result(rowcount: int = 1, inserted_id: int = 1):
result = Mock()
result.rowcount = rowcount
result.inserted_primary_key = [inserted_id]
return result
@@ -71,7 +88,7 @@ class TestWebhookServiceGetWebhooks:
service = WebhookService(ap)
# Execute
result = await service.get_webhooks()
result = await service.get_webhooks(WORKSPACE_UUID)
# Verify
assert result == []
@@ -100,7 +117,7 @@ class TestWebhookServiceGetWebhooks:
service = WebhookService(ap)
# Execute
result = await service.get_webhooks()
result = await service.get_webhooks(WORKSPACE_UUID)
# Verify
assert len(result) == 2
@@ -119,6 +136,7 @@ class TestWebhookServiceCreateWebhook:
# Mock insert result
insert_result = Mock()
insert_result.inserted_primary_key = [1]
# Mock select result for retrieving created webhook
created_webhook = _create_mock_webhook(
@@ -137,6 +155,8 @@ class TestWebhookServiceCreateWebhook:
nonlocal call_count
call_count += 1
if call_count == 1:
return _create_mock_result(scalar_value=0) # Count
if call_count == 2:
return insert_result # Insert
return select_result # Select
@@ -155,6 +175,7 @@ class TestWebhookServiceCreateWebhook:
# Execute
result = await service.create_webhook(
WORKSPACE_UUID,
name='New Webhook',
url='http://new.example.com/webhook',
description='New Description',
@@ -187,7 +208,9 @@ class TestWebhookServiceCreateWebhook:
nonlocal call_count
call_count += 1
if call_count == 1:
return Mock() # Insert
return _create_mock_result(scalar_value=0)
if call_count == 2:
return _create_write_result() # Insert
return _create_mock_result(first_item=created_webhook)
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
@@ -204,7 +227,11 @@ class TestWebhookServiceCreateWebhook:
service = WebhookService(ap)
# Execute - only name and url required
result = await service.create_webhook(name='Minimal Webhook', url='http://minimal.example.com')
result = await service.create_webhook(
WORKSPACE_UUID,
name='Minimal Webhook',
url='http://minimal.example.com',
)
# Verify defaults
assert result['description'] == ''
@@ -224,7 +251,9 @@ class TestWebhookServiceCreateWebhook:
nonlocal call_count
call_count += 1
if call_count == 1:
return Mock()
return _create_mock_result(scalar_value=0)
if call_count == 2:
return _create_write_result()
return _create_mock_result(first_item=created_webhook)
ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
@@ -233,11 +262,52 @@ class TestWebhookServiceCreateWebhook:
service = WebhookService(ap)
# Execute
result = await service.create_webhook(name='Disabled', url='http://disabled.com', enabled=False)
result = await service.create_webhook(
WORKSPACE_UUID,
name='Disabled',
url='http://disabled.com',
enabled=False,
)
# Verify
assert result['enabled'] is False
async def test_create_webhook_rejects_workspace_at_capacity(self):
ap = SimpleNamespace(
instance_config=SimpleNamespace(
data={'webhooks': {'max_per_workspace': 2}},
),
persistence_mgr=SimpleNamespace(
execute_async=AsyncMock(return_value=_create_mock_result(scalar_value=2)),
),
)
service = WebhookService(ap)
with pytest.raises(ValueError, match=r'Maximum number of webhooks \(2\) reached'):
await service.create_webhook(
WORKSPACE_UUID,
name='Too many',
url='https://example.invalid',
)
ap.persistence_mgr.execute_async.assert_awaited_once()
async def test_max_per_workspace_clamps_invalid_and_oversized_values(self):
ap = SimpleNamespace(
instance_config=SimpleNamespace(
data={'webhooks': {'max_per_workspace': 999999}},
)
)
service = WebhookService(ap)
assert service.max_per_workspace() == 64
ap.instance_config.data['webhooks']['max_per_workspace'] = 0
assert service.max_per_workspace() == 1
ap.instance_config.data['webhooks']['max_per_workspace'] = 'invalid'
assert service.max_per_workspace() == 16
class TestWebhookServiceGetWebhook:
"""Tests for get_webhook method."""
@@ -262,7 +332,7 @@ class TestWebhookServiceGetWebhook:
service = WebhookService(ap)
# Execute
result = await service.get_webhook(1)
result = await service.get_webhook(WORKSPACE_UUID, 1)
# Verify
assert result is not None
@@ -281,7 +351,7 @@ class TestWebhookServiceGetWebhook:
service = WebhookService(ap)
# Execute
result = await service.get_webhook(999)
result = await service.get_webhook(WORKSPACE_UUID, 999)
# Verify
assert result is None
@@ -298,7 +368,7 @@ class TestWebhookServiceGetWebhook:
service = WebhookService(ap)
# Execute
result = await service.get_webhook(0)
result = await service.get_webhook(WORKSPACE_UUID, 0)
# Verify - should return None (no webhook with ID 0)
assert result is None
@@ -312,12 +382,12 @@ class TestWebhookServiceUpdateWebhook:
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.persistence_mgr.execute_async = AsyncMock()
ap.persistence_mgr.execute_async = AsyncMock(return_value=_create_write_result())
service = WebhookService(ap)
# Execute
await service.update_webhook(1, name='Updated Name')
await service.update_webhook(WORKSPACE_UUID, 1, name='Updated Name')
# Verify
ap.persistence_mgr.execute_async.assert_called_once()
@@ -327,12 +397,12 @@ class TestWebhookServiceUpdateWebhook:
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.persistence_mgr.execute_async = AsyncMock()
ap.persistence_mgr.execute_async = AsyncMock(return_value=_create_write_result())
service = WebhookService(ap)
# Execute
await service.update_webhook(1, url='http://updated.example.com')
await service.update_webhook(WORKSPACE_UUID, 1, url='http://updated.example.com')
# Verify
ap.persistence_mgr.execute_async.assert_called_once()
@@ -342,12 +412,12 @@ class TestWebhookServiceUpdateWebhook:
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.persistence_mgr.execute_async = AsyncMock()
ap.persistence_mgr.execute_async = AsyncMock(return_value=_create_write_result())
service = WebhookService(ap)
# Execute
await service.update_webhook(1, description='Updated description')
await service.update_webhook(WORKSPACE_UUID, 1, description='Updated description')
# Verify
ap.persistence_mgr.execute_async.assert_called_once()
@@ -357,12 +427,12 @@ class TestWebhookServiceUpdateWebhook:
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.persistence_mgr.execute_async = AsyncMock()
ap.persistence_mgr.execute_async = AsyncMock(return_value=_create_write_result())
service = WebhookService(ap)
# Execute
await service.update_webhook(1, enabled=False)
await service.update_webhook(WORKSPACE_UUID, 1, enabled=False)
# Verify
ap.persistence_mgr.execute_async.assert_called_once()
@@ -372,12 +442,13 @@ class TestWebhookServiceUpdateWebhook:
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.persistence_mgr.execute_async = AsyncMock()
ap.persistence_mgr.execute_async = AsyncMock(return_value=_create_write_result())
service = WebhookService(ap)
# Execute
await service.update_webhook(
WORKSPACE_UUID,
1,
name='All Updated',
url='http://all.updated.com',
@@ -393,15 +464,17 @@ class TestWebhookServiceUpdateWebhook:
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.persistence_mgr.execute_async = AsyncMock()
existing = _create_mock_webhook(webhook_id=1)
ap.persistence_mgr.execute_async = AsyncMock(return_value=_create_mock_result(first_item=existing))
ap.persistence_mgr.serialize_model = Mock(return_value={'id': 1})
service = WebhookService(ap)
# Execute - no update parameters
await service.update_webhook(1)
await service.update_webhook(WORKSPACE_UUID, 1)
# Verify - no execute call since no update_data
ap.persistence_mgr.execute_async.assert_not_called()
# No write is issued; one scoped existence lookup is performed.
ap.persistence_mgr.execute_async.assert_called_once()
class TestWebhookServiceDeleteWebhook:
@@ -412,12 +485,12 @@ class TestWebhookServiceDeleteWebhook:
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.persistence_mgr.execute_async = AsyncMock()
ap.persistence_mgr.execute_async = AsyncMock(return_value=_create_write_result())
service = WebhookService(ap)
# Execute
await service.delete_webhook(1)
await service.delete_webhook(WORKSPACE_UUID, 1)
# Verify
ap.persistence_mgr.execute_async.assert_called_once()
@@ -427,12 +500,12 @@ class TestWebhookServiceDeleteWebhook:
# Setup
ap = SimpleNamespace()
ap.persistence_mgr = SimpleNamespace()
ap.persistence_mgr.execute_async = AsyncMock()
ap.persistence_mgr.execute_async = AsyncMock(return_value=_create_write_result(rowcount=0))
service = WebhookService(ap)
# Execute - should not raise
await service.delete_webhook(999)
await service.delete_webhook(WORKSPACE_UUID, 999)
# Verify - still called
ap.persistence_mgr.execute_async.assert_called_once()
@@ -453,7 +526,7 @@ class TestWebhookServiceGetEnabledWebhooks:
service = WebhookService(ap)
# Execute
result = await service.get_enabled_webhooks()
result = await service.get_enabled_webhooks(WORKSPACE_UUID)
# Verify
assert result == []
@@ -481,7 +554,7 @@ class TestWebhookServiceGetEnabledWebhooks:
service = WebhookService(ap)
# Execute
result = await service.get_enabled_webhooks()
result = await service.get_enabled_webhooks(WORKSPACE_UUID)
# Verify
assert len(result) == 2
@@ -501,7 +574,170 @@ class TestWebhookServiceGetEnabledWebhooks:
service = WebhookService(ap)
# Execute
result = await service.get_enabled_webhooks()
result = await service.get_enabled_webhooks(WORKSPACE_UUID)
# Verify - should be empty (SQL would filter disabled)
assert result == []
ISOLATION_WORKSPACE_A = '00000000-0000-0000-0000-00000000000a'
ISOLATION_WORKSPACE_B = '00000000-0000-0000-0000-00000000000b'
class _RealPersistenceManager:
def __init__(self, engine):
self.engine = engine
async def execute_async(self, *args, **kwargs):
async with self.engine.connect() as connection:
result = await connection.execute(*args, **kwargs)
await connection.commit()
return result
@staticmethod
def serialize_model(model, data, masked_columns=None):
return {
column.name: (
getattr(data, column.name).isoformat()
if isinstance(getattr(data, column.name), datetime.datetime)
else getattr(data, column.name)
)
for column in model.__table__.columns
if column.name not in (masked_columns or [])
}
@pytest.fixture
async def tenant_webhook_service(tmp_path):
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "webhooks.db"}')
async with engine.begin() as connection:
await connection.run_sync(Base.metadata.create_all)
await connection.execute(
sqlalchemy.insert(Workspace),
[
{
'uuid': ISOLATION_WORKSPACE_A,
'instance_uuid': 'instance',
'name': 'A',
'slug': 'a',
'source': 'cloud_projection',
},
{
'uuid': ISOLATION_WORKSPACE_B,
'instance_uuid': 'instance',
'name': 'B',
'slug': 'b',
'source': 'cloud_projection',
},
],
)
service = WebhookService(SimpleNamespace(persistence_mgr=_RealPersistenceManager(engine)))
yield service
await engine.dispose()
async def test_webhook_service_requires_workspace(tenant_webhook_service):
with pytest.raises(WorkspaceRequiredError):
await tenant_webhook_service.get_webhooks(None)
async def test_same_name_webhooks_are_isolated(tenant_webhook_service):
created_a = await tenant_webhook_service.create_webhook(
ISOLATION_WORKSPACE_A,
'deploy',
'https://a.invalid',
)
created_b = await tenant_webhook_service.create_webhook(
ISOLATION_WORKSPACE_B,
'deploy',
'https://b.invalid',
)
assert created_a['workspace_uuid'] == ISOLATION_WORKSPACE_A
assert created_b['workspace_uuid'] == ISOLATION_WORKSPACE_B
assert [item['url'] for item in await tenant_webhook_service.get_webhooks(ISOLATION_WORKSPACE_A)] == ['***']
assert [item['url'] for item in await tenant_webhook_service.get_webhooks(ISOLATION_WORKSPACE_B)] == ['***']
assert [
item['url'] for item in await tenant_webhook_service.get_webhooks(ISOLATION_WORKSPACE_A, include_secret=True)
] == ['https://a.invalid']
assert [
item['url'] for item in await tenant_webhook_service.get_webhooks(ISOLATION_WORKSPACE_B, include_secret=True)
] == ['https://b.invalid']
async def test_cross_workspace_id_guessing_is_not_found(tenant_webhook_service):
created = await tenant_webhook_service.create_webhook(
ISOLATION_WORKSPACE_A,
'secret',
'https://a.invalid/hook',
)
webhook_id = created['id']
assert await tenant_webhook_service.get_webhook(ISOLATION_WORKSPACE_B, webhook_id) is None
assert not await tenant_webhook_service.update_webhook(
ISOLATION_WORKSPACE_B,
webhook_id,
name='stolen',
)
assert not await tenant_webhook_service.delete_webhook(ISOLATION_WORKSPACE_B, webhook_id)
assert (await tenant_webhook_service.get_webhook(ISOLATION_WORKSPACE_A, webhook_id))['name'] == 'secret'
async def test_update_and_delete_are_scoped(tenant_webhook_service):
created = await tenant_webhook_service.create_webhook(
ISOLATION_WORKSPACE_A,
'old',
'https://a.invalid/old',
)
assert await tenant_webhook_service.update_webhook(
ISOLATION_WORKSPACE_A,
created['id'],
name='new',
enabled=False,
)
assert await tenant_webhook_service.get_enabled_webhooks(ISOLATION_WORKSPACE_A) == []
assert await tenant_webhook_service.delete_webhook(ISOLATION_WORKSPACE_A, created['id'])
assert await tenant_webhook_service.get_webhook(ISOLATION_WORKSPACE_A, created['id']) is None
async def test_masked_webhook_url_roundtrip_preserves_replace_and_clear(tenant_webhook_service):
created = await tenant_webhook_service.create_webhook(
ISOLATION_WORKSPACE_A,
'roundtrip',
'https://a.invalid/bearer-secret',
)
masked = await tenant_webhook_service.get_webhook(ISOLATION_WORKSPACE_A, created['id'])
assert masked['url'] == '***'
assert await tenant_webhook_service.update_webhook(
ISOLATION_WORKSPACE_A,
created['id'],
name='preserved',
url=masked['url'],
)
preserved = await tenant_webhook_service.get_webhook(
ISOLATION_WORKSPACE_A,
created['id'],
include_secret=True,
)
assert preserved['url'] == 'https://a.invalid/bearer-secret'
assert await tenant_webhook_service.update_webhook(
ISOLATION_WORKSPACE_A,
created['id'],
url='https://a.invalid/replacement',
)
replaced = await tenant_webhook_service.get_webhook(
ISOLATION_WORKSPACE_A,
created['id'],
include_secret=True,
)
assert replaced['url'] == 'https://a.invalid/replacement'
assert await tenant_webhook_service.update_webhook(ISOLATION_WORKSPACE_A, created['id'], url='')
cleared = await tenant_webhook_service.get_webhook(
ISOLATION_WORKSPACE_A,
created['id'],
include_secret=True,
)
assert cleared['url'] == ''
tests/unit_tests/api/test_adapter_session_scoping.py
+254
View File
@@ -0,0 +1,254 @@
from __future__ import annotations
import asyncio
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
import lark_oapi
import pytest
import quart
from langbot.pkg.api.http.context import (
PrincipalContext,
PrincipalType,
RequestContext,
WorkspaceContext,
)
from langbot.pkg.api.http.controller.groups.platform.adapters import (
AdaptersRouterGroup,
_AdapterSessionScope,
_bind_session_scope,
_get_owned_session,
_make_room_for_session,
_pop_owned_session,
_start_adapter_session_task,
)
pytestmark = pytest.mark.asyncio
SENSITIVE_ADAPTER_ROUTES = (
('post', '/api/v1/platform/adapters/lark/create-app'),
('get', '/api/v1/platform/adapters/lark/create-app/status/missing'),
('delete', '/api/v1/platform/adapters/lark/create-app/missing'),
('post', '/api/v1/platform/adapters/weixin/login'),
('get', '/api/v1/platform/adapters/weixin/login/status/missing'),
('delete', '/api/v1/platform/adapters/weixin/login/missing'),
('post', '/api/v1/platform/adapters/dingtalk/create-app'),
('get', '/api/v1/platform/adapters/dingtalk/create-app/status/missing'),
('delete', '/api/v1/platform/adapters/dingtalk/create-app/missing'),
('post', '/api/v1/platform/adapters/wecombot/create-bot'),
('get', '/api/v1/platform/adapters/wecombot/create-bot/status/missing'),
('delete', '/api/v1/platform/adapters/wecombot/create-bot/missing'),
('post', '/api/v1/platform/adapters/qqofficial/bind'),
('get', '/api/v1/platform/adapters/qqofficial/bind/status/missing'),
('delete', '/api/v1/platform/adapters/qqofficial/bind/missing'),
)
def _request_context(
*,
account_uuid: str = 'account-a',
workspace_uuid: str = 'workspace-a',
placement_generation: int = 1,
) -> RequestContext:
return RequestContext(
instance_uuid='instance-test',
placement_generation=placement_generation,
request_id='request-test',
auth_type='user-token',
principal=PrincipalContext(
principal_type=PrincipalType.ACCOUNT,
account_uuid=account_uuid,
),
workspace=WorkspaceContext(
workspace_uuid=workspace_uuid,
membership_uuid='membership-test',
role='developer',
permissions=frozenset({'resource.manage'}),
),
)
async def _create_client(*, role: str = 'developer'):
quart_app = quart.Quart(__name__)
accounts = {
'owner-token': SimpleNamespace(uuid='account-a', user='owner@example.com'),
'other-token': SimpleNamespace(uuid='account-b', user='other@example.com'),
}
async def get_authenticated_account(token: str):
return accounts[token]
async def resolve_account_workspace(account_uuid: str, requested_workspace_uuid: str | None):
workspace_uuid = requested_workspace_uuid or 'workspace-a'
return SimpleNamespace(
execution=SimpleNamespace(
instance_uuid='instance-test',
placement_generation=1,
),
workspace=SimpleNamespace(uuid=workspace_uuid),
membership=SimpleNamespace(
uuid=f'membership-{account_uuid}-{workspace_uuid}',
role=role,
projection_revision=1,
),
)
class TestTaskManager:
def create_user_task(self, coro, **_kwargs):
return SimpleNamespace(task=asyncio.create_task(coro))
application = SimpleNamespace(
user_service=SimpleNamespace(
get_authenticated_account=AsyncMock(side_effect=get_authenticated_account),
),
workspace_collaboration_service=SimpleNamespace(
resolve_account_workspace=AsyncMock(side_effect=resolve_account_workspace),
),
platform_mgr=SimpleNamespace(),
task_mgr=TestTaskManager(),
)
router = AdaptersRouterGroup(application, quart_app)
await router.initialize()
return quart_app.test_client()
@pytest.mark.parametrize(('method', 'path'), SENSITIVE_ADAPTER_ROUTES)
async def test_sensitive_adapter_flows_require_resource_manage(method: str, path: str):
client = await _create_client(role='viewer')
response = await getattr(client, method)(
path,
headers={'Authorization': 'Bearer owner-token'},
)
assert response.status_code == 403
assert (await response.get_json())['code'] == 'permission_denied'
async def test_session_scope_matches_exact_tenant_placement_and_principal():
owner_context = _request_context()
sessions: dict[str, dict] = {'session-test': {'status': 'waiting'}}
_bind_session_scope(sessions['session-test'], owner_context)
assert sessions['session-test']['scope'] == _AdapterSessionScope.from_request_context(owner_context)
assert _get_owned_session(sessions, 'session-test', owner_context) is sessions['session-test']
for other_context in (
_request_context(account_uuid='account-b'),
_request_context(workspace_uuid='workspace-b'),
_request_context(placement_generation=2),
):
assert _get_owned_session(sessions, 'session-test', other_context) is None
assert _pop_owned_session(sessions, 'session-test', other_context) is None
assert 'session-test' in sessions
assert _pop_owned_session(sessions, 'session-test', owner_context) is not None
assert sessions == {}
async def test_session_capacity_evicts_oldest_session_in_same_workspace():
owner_context = _request_context()
sessions: dict[str, dict] = {}
tasks = []
for index in range(10):
task = SimpleNamespace(done=Mock(return_value=False), cancel=Mock())
tasks.append(task)
session = {'created_at': float(index), 'task': task}
_bind_session_scope(session, owner_context)
sessions[f'session-{index}'] = session
_make_room_for_session(sessions, owner_context)
assert 'session-0' not in sessions
assert len(sessions) == 9
tasks[0].cancel.assert_called_once_with()
async def test_adapter_session_task_uses_tenant_task_admission():
blocker = asyncio.Event()
async def credential_exchange():
await blocker.wait()
task_manager = SimpleNamespace(create_user_task=Mock())
def create_user_task(coro, **_kwargs):
return SimpleNamespace(task=asyncio.create_task(coro))
task_manager.create_user_task.side_effect = create_user_task
application = SimpleNamespace(task_mgr=task_manager)
request_context = _request_context()
returned = _start_adapter_session_task(
application,
credential_exchange(),
adapter='lark',
session_id='session-test',
request_context=request_context,
)
assert returned is not None
task_manager.create_user_task.assert_called_once()
kwargs = task_manager.create_user_task.call_args.kwargs
assert kwargs['kind'] == 'platform-adapter-credential-exchange'
assert kwargs['instance_uuid'] == request_context.instance_uuid
assert kwargs['workspace_uuid'] == request_context.workspace_uuid
assert kwargs['placement_generation'] == request_context.placement_generation
blocker.set()
await returned
async def test_lark_session_status_and_delete_hide_cross_scope_sessions(monkeypatch):
registration_blocker = asyncio.Event()
async def fake_register_app(*, on_qr_code, source: str):
assert source == 'langbot'
on_qr_code({'url': 'https://example.test/lark-qr'})
await registration_blocker.wait()
raise AssertionError('registration should have been cancelled')
monkeypatch.setattr(lark_oapi, 'aregister_app', fake_register_app)
client = await _create_client()
owner_headers = {
'Authorization': 'Bearer owner-token',
'X-Workspace-Id': 'workspace-a',
}
create_response = await client.post(
'/api/v1/platform/adapters/lark/create-app',
headers=owner_headers,
)
assert create_response.status_code == 200
session_id = (await create_response.get_json())['data']['session_id']
status_path = f'/api/v1/platform/adapters/lark/create-app/status/{session_id}'
delete_path = f'/api/v1/platform/adapters/lark/create-app/{session_id}'
for headers in (
{
'Authorization': 'Bearer other-token',
'X-Workspace-Id': 'workspace-a',
},
{
'Authorization': 'Bearer owner-token',
'X-Workspace-Id': 'workspace-b',
},
):
status_response = await client.get(status_path, headers=headers)
delete_response = await client.delete(delete_path, headers=headers)
assert status_response.status_code == 404
assert delete_response.status_code == 404
assert (await status_response.get_json())['msg'] == 'Session not found'
assert (await delete_response.get_json())['msg'] == 'Session not found'
owner_status_response = await client.get(status_path, headers=owner_headers)
assert owner_status_response.status_code == 200
assert (await owner_status_response.get_json())['data']['status'] == 'waiting'
owner_delete_response = await client.delete(delete_path, headers=owner_headers)
assert owner_delete_response.status_code == 200
missing_delete_response = await client.delete(delete_path, headers=owner_headers)
assert missing_delete_response.status_code == 404
await asyncio.sleep(0)
tests/unit_tests/api/test_apikey_service.py
+56 -15
View File
@@ -6,6 +6,7 @@ from unittest.mock import AsyncMock, Mock
import pytest
from langbot.pkg.api.http.service.apikey import ApiKeyService
from langbot.pkg.entity.persistence.apikey import ApiKeyStatus
@pytest.mark.asyncio
@@ -13,30 +14,70 @@ from langbot.pkg.api.http.service.apikey import ApiKeyService
async def test_verify_api_key_rejects_non_lbk_keys_without_db_query(api_key):
persistence_mgr = SimpleNamespace(execute_async=AsyncMock())
instance_config = SimpleNamespace(data={'api': {'global_api_key': ''}})
service = ApiKeyService(SimpleNamespace(persistence_mgr=persistence_mgr, instance_config=instance_config))
workspace_service = SimpleNamespace(get_execution_binding=AsyncMock())
service = ApiKeyService(
SimpleNamespace(
persistence_mgr=persistence_mgr,
instance_config=instance_config,
workspace_service=workspace_service,
)
)
result = await service.verify_api_key(api_key)
assert result is False
persistence_mgr.execute_async.assert_not_awaited()
workspace_service.get_execution_binding.assert_not_awaited()
@pytest.mark.asyncio
@pytest.mark.parametrize(
('db_row', 'expected'),
[
(object(), True),
(None, False),
],
)
async def test_verify_api_key_keeps_db_validation_for_lbk_keys(db_row, expected):
query_result = Mock()
query_result.first.return_value = db_row
persistence_mgr = SimpleNamespace(execute_async=AsyncMock(return_value=query_result))
@pytest.mark.parametrize('key_exists', [True, False])
async def test_verify_api_key_keeps_db_validation_for_lbk_keys(key_exists):
key = (
SimpleNamespace(
id=1,
uuid='key-uuid',
workspace_uuid='workspace-a',
status=ApiKeyStatus.ACTIVE.value,
expires_at=None,
scopes=[],
)
if key_exists
else None
)
discovery_result = Mock()
discovery_result.first.return_value = key
query_results = [discovery_result]
if key_exists:
scoped_result = Mock()
scoped_result.first.return_value = key
update_result = Mock()
update_result.scalar_one_or_none.return_value = key.id
query_results.extend([scoped_result, update_result])
persistence_mgr = SimpleNamespace(execute_async=AsyncMock(side_effect=query_results))
instance_config = SimpleNamespace(data={'api': {'global_api_key': ''}})
service = ApiKeyService(SimpleNamespace(persistence_mgr=persistence_mgr, instance_config=instance_config))
workspace_service = SimpleNamespace(
get_execution_binding=AsyncMock(
return_value=SimpleNamespace(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=1,
)
)
)
service = ApiKeyService(
SimpleNamespace(
persistence_mgr=persistence_mgr,
instance_config=instance_config,
workspace_service=workspace_service,
)
)
result = await service.verify_api_key('lbk_valid_format')
assert result is expected
persistence_mgr.execute_async.assert_awaited_once()
assert result is key_exists
assert persistence_mgr.execute_async.await_count == (3 if key_exists else 1)
if key_exists:
workspace_service.get_execution_binding.assert_awaited_once_with('workspace-a')
else:
workspace_service.get_execution_binding.assert_not_awaited()
tests/unit_tests/api/test_bot_controller_secrets.py
+113
View File
@@ -0,0 +1,113 @@
from __future__ import annotations
from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
import quart
from langbot.pkg.api.http.controller.groups.platform.bots import BotsRouterGroup
pytestmark = pytest.mark.asyncio
SECRET_CONFIG = {'token': 'tenant-secret', 'app_secret': 'also-secret'}
async def create_client(*, role: str):
quart_app = quart.Quart(__name__)
account = SimpleNamespace(uuid='account-test', user='test@example.com')
user_service = SimpleNamespace(
get_authenticated_account=AsyncMock(return_value=account),
)
access = SimpleNamespace(
execution=SimpleNamespace(
instance_uuid='instance-test',
placement_generation=1,
),
workspace=SimpleNamespace(uuid='workspace-test'),
membership=SimpleNamespace(
uuid='membership-test',
role=role,
projection_revision=1,
),
)
async def get_bots(_context, *, include_secret=False):
bot = {'uuid': 'bot-test', 'name': 'Test Bot'}
if include_secret:
bot['adapter_config'] = SECRET_CONFIG
return [bot]
async def get_runtime_bot_info(_context, _bot_uuid, *, include_secret=False):
bot = {'uuid': 'bot-test', 'name': 'Test Bot'}
if include_secret:
bot['adapter_config'] = SECRET_CONFIG
return bot
bot_service = SimpleNamespace(
get_bots=AsyncMock(side_effect=get_bots),
get_runtime_bot_info=AsyncMock(side_effect=get_runtime_bot_info),
update_bot=AsyncMock(),
)
application = SimpleNamespace(
user_service=user_service,
apikey_service=SimpleNamespace(
authenticate_api_key=AsyncMock(
return_value=SimpleNamespace(
instance_uuid='instance-test',
placement_generation=1,
api_key_uuid='api-key-test',
workspace_uuid='workspace-test',
permissions=frozenset({'resource.view'}),
)
)
),
workspace_collaboration_service=SimpleNamespace(resolve_account_workspace=AsyncMock(return_value=access)),
bot_service=bot_service,
)
router = BotsRouterGroup(application, quart_app)
await router.initialize()
return quart_app.test_client(), bot_service
async def test_viewer_list_and_detail_never_receive_adapter_credentials():
client, bot_service = await create_client(role='viewer')
headers = {'Authorization': 'Bearer test-token'}
list_response = await client.get('/api/v1/platform/bots', headers=headers)
detail_response = await client.get('/api/v1/platform/bots/bot-test', headers=headers)
assert list_response.status_code == 200
assert detail_response.status_code == 200
assert 'adapter_config' not in (await list_response.get_json())['data']['bots'][0]
assert 'adapter_config' not in (await detail_response.get_json())['data']['bot']
assert bot_service.get_bots.await_args.kwargs['include_secret'] is False
assert bot_service.get_runtime_bot_info.await_args.kwargs['include_secret'] is False
async def test_resource_manager_can_read_adapter_credentials():
client, bot_service = await create_client(role='developer')
headers = {'Authorization': 'Bearer test-token'}
list_response = await client.get('/api/v1/platform/bots', headers=headers)
detail_response = await client.get('/api/v1/platform/bots/bot-test', headers=headers)
assert (await list_response.get_json())['data']['bots'][0]['adapter_config'] == SECRET_CONFIG
assert (await detail_response.get_json())['data']['bot']['adapter_config'] == SECRET_CONFIG
assert bot_service.get_bots.await_args.kwargs['include_secret'] is True
assert bot_service.get_runtime_bot_info.await_args.kwargs['include_secret'] is True
async def test_viewer_cannot_write_adapter_credentials():
client, bot_service = await create_client(role='viewer')
response = await client.put(
'/api/v1/platform/bots/bot-test',
headers={'Authorization': 'Bearer test-token'},
json={'adapter_config': SECRET_CONFIG},
)
assert response.status_code == 403
assert (await response.get_json())['code'] == 'permission_denied'
bot_service.update_bot.assert_not_awaited()
tests/unit_tests/api/test_extensions_runtime_fence.py
+142
View File
@@ -0,0 +1,142 @@
from __future__ import annotations
from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
import quart
import sqlalchemy
from sqlalchemy.ext.asyncio import create_async_engine
from langbot.pkg.api.http.controller.groups.extensions import ExtensionsRouterGroup
from langbot.pkg.persistence.mgr import PersistenceManager
from langbot.pkg.workspace.errors import WorkspaceNotFoundError
@pytest.mark.asyncio
async def test_extensions_route_hides_runtime_bound_to_another_workspace():
account = SimpleNamespace(uuid='account-a', user='owner@example.com')
connector = SimpleNamespace(
is_enable_plugin=True,
require_workspace_context=AsyncMock(side_effect=WorkspaceNotFoundError('Plugin resource not found')),
list_plugins=AsyncMock(return_value=[]),
)
ap = SimpleNamespace(
user_service=SimpleNamespace(
get_authenticated_account=AsyncMock(return_value=account),
),
workspace_collaboration_service=SimpleNamespace(
resolve_account_workspace=AsyncMock(
return_value=SimpleNamespace(
workspace=SimpleNamespace(uuid='workspace-a'),
membership=SimpleNamespace(uuid='membership-a', role='owner', projection_revision=0),
execution=SimpleNamespace(instance_uuid='instance-a', placement_generation=2),
)
)
),
plugin_connector=connector,
mcp_service=SimpleNamespace(get_mcp_servers=AsyncMock(return_value=[])),
skill_service=SimpleNamespace(list_skills=AsyncMock(return_value=[])),
)
quart_app = quart.Quart(__name__)
router = ExtensionsRouterGroup(ap, quart_app)
await router.initialize()
response = await quart_app.test_client().get(
'/api/v1/extensions',
headers={'Authorization': 'Bearer token'},
)
assert response.status_code == 404
connector.list_plugins.assert_not_awaited()
ap.mcp_service.get_mcp_servers.assert_not_awaited()
ap.skill_service.list_skills.assert_not_awaited()
@pytest.mark.asyncio
async def test_extensions_route_redacts_plugin_secrets_without_mutating_runtime_data():
account = SimpleNamespace(uuid='account-a', user='viewer@example.com')
raw_plugin = {
'plugin_config': {'apiKey': 'plugin-secret', 'nested': {'token': 'nested-secret'}},
'debug': {'plugin_debug_key': 'debug-secret'},
}
connector = SimpleNamespace(
is_enable_plugin=True,
require_workspace_context=AsyncMock(),
list_plugins=AsyncMock(return_value=[raw_plugin]),
)
ap = SimpleNamespace(
user_service=SimpleNamespace(get_authenticated_account=AsyncMock(return_value=account)),
workspace_collaboration_service=SimpleNamespace(
resolve_account_workspace=AsyncMock(
return_value=SimpleNamespace(
workspace=SimpleNamespace(uuid='workspace-a'),
membership=SimpleNamespace(uuid='membership-a', role='viewer', projection_revision=0),
execution=SimpleNamespace(instance_uuid='instance-a', placement_generation=2),
)
)
),
plugin_connector=connector,
mcp_service=SimpleNamespace(get_mcp_servers=AsyncMock(return_value=[])),
skill_service=SimpleNamespace(list_skills=AsyncMock(return_value=[])),
)
quart_app = quart.Quart(__name__)
router = ExtensionsRouterGroup(ap, quart_app)
await router.initialize()
response = await quart_app.test_client().get(
'/api/v1/extensions',
headers={'Authorization': 'Bearer token', 'X-Workspace-Id': 'workspace-a'},
)
assert response.status_code == 200
plugin = (await response.get_json())['data']['extensions'][0]['plugin']
assert plugin['plugin_config']['apiKey'] == '***'
assert plugin['plugin_config']['nested']['token'] == '***'
assert plugin['debug']['plugin_debug_key'] == '***'
assert raw_plugin['plugin_config']['apiKey'] == 'plugin-secret'
@pytest.mark.asyncio
async def test_extensions_parallel_reads_open_explicit_child_task_scopes():
engine = create_async_engine('sqlite+aiosqlite:///:memory:')
account = SimpleNamespace(uuid='account-a', user='owner@example.com')
ap = SimpleNamespace()
ap.persistence_mgr = PersistenceManager(ap)
ap.persistence_mgr.db = SimpleNamespace(get_engine=lambda: engine)
ap.user_service = SimpleNamespace(get_authenticated_account=AsyncMock(return_value=account))
ap.workspace_collaboration_service = SimpleNamespace(
resolve_account_workspace=AsyncMock(
return_value=SimpleNamespace(
workspace=SimpleNamespace(uuid='workspace-a'),
membership=SimpleNamespace(uuid='membership-a', role='owner', projection_revision=0),
execution=SimpleNamespace(instance_uuid='instance-a', placement_generation=2),
)
)
)
ap.plugin_connector = SimpleNamespace(
is_enable_plugin=False,
list_plugins=AsyncMock(return_value=[]),
)
async def list_mcp_servers(_context, *, contain_runtime_info):
await ap.persistence_mgr.execute_async(sqlalchemy.select(sqlalchemy.literal(1)))
assert contain_runtime_info is True
return [{'name': 'Scoped MCP'}]
ap.mcp_service = SimpleNamespace(get_mcp_servers=AsyncMock(side_effect=list_mcp_servers))
ap.skill_service = SimpleNamespace(list_skills=AsyncMock(return_value=[]))
quart_app = quart.Quart(__name__)
router = ExtensionsRouterGroup(ap, quart_app)
await router.initialize()
try:
response = await quart_app.test_client().get(
'/api/v1/extensions',
headers={'Authorization': 'Bearer token', 'X-Workspace-Id': 'workspace-a'},
)
finally:
await engine.dispose()
assert response.status_code == 200
assert (await response.get_json())['data']['extensions'] == [{'type': 'mcp', 'server': {'name': 'Scoped MCP'}}]
tests/unit_tests/api/test_file_upload_scoping.py
+59
View File
@@ -0,0 +1,59 @@
from __future__ import annotations
import io
from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
import quart
from quart.datastructures import FileStorage
from langbot.pkg.api.http.controller.groups.files import FilesRouterGroup
pytestmark = pytest.mark.asyncio
async def test_document_upload_uses_dedicated_scoped_owner_type():
quart_app = quart.Quart(__name__)
account = SimpleNamespace(uuid='account-test', user='test@example.com')
access = SimpleNamespace(
execution=SimpleNamespace(
instance_uuid='instance-test',
placement_generation=3,
),
workspace=SimpleNamespace(uuid='00000000-0000-0000-0000-00000000000a'),
membership=SimpleNamespace(
uuid='membership-test',
role='developer',
projection_revision=1,
),
)
storage_mgr = SimpleNamespace(save_scoped=AsyncMock(return_value='scoped-document-key'))
application = SimpleNamespace(
user_service=SimpleNamespace(get_authenticated_account=AsyncMock(return_value=account)),
workspace_collaboration_service=SimpleNamespace(resolve_account_workspace=AsyncMock(return_value=access)),
storage_mgr=storage_mgr,
)
router = FilesRouterGroup(application, quart_app)
await router.initialize()
client = quart_app.test_client()
response = await client.post(
'/api/v1/files/documents',
headers={'Authorization': 'Bearer test-token'},
files={
'file': FileStorage(
stream=io.BytesIO(b'document bytes'),
filename='report.pdf',
)
},
)
assert response.status_code == 200
assert (await response.get_json())['data']['file_id'] == 'scoped-document-key'
kwargs = storage_mgr.save_scoped.await_args.kwargs
assert kwargs['owner_type'] == 'upload_document'
assert kwargs['owner'] == 'account:account-test'
assert kwargs['key'].endswith('.pdf')
assert kwargs['value'] == b'document bytes'
tests/unit_tests/api/test_knowledge_migration_runtime_fence.py
+293
View File
@@ -0,0 +1,293 @@
from __future__ import annotations
from datetime import datetime
import json
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
import pytest
import sqlalchemy
from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.api.http.controller.groups.knowledge.migration import KnowledgeMigrationRouterGroup
from langbot.pkg.persistence.tenant_uow import _validate_scoped_statement_call
from langbot.pkg.workspace.errors import WorkspaceInvariantError, WorkspaceNotFoundError
CONTEXT = ExecutionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=3,
)
@pytest.mark.asyncio
async def test_background_migration_propagates_generation_change_before_runtime_call():
connector = SimpleNamespace(
require_workspace_context=AsyncMock(side_effect=[CONTEXT, WorkspaceNotFoundError('Plugin resource not found')]),
list_knowledge_engines=AsyncMock(return_value=[]),
)
router = object.__new__(KnowledgeMigrationRouterGroup)
router.ap = SimpleNamespace(
plugin_connector=connector,
workspace_service=SimpleNamespace(
get_local_execution_binding=AsyncMock(return_value=CONTEXT),
),
logger=Mock(),
)
router._table_exists = AsyncMock(return_value=False)
router._set_migration_flag = AsyncMock()
task_context = SimpleNamespace(trace=Mock())
with pytest.raises(WorkspaceNotFoundError, match='Plugin resource not found'):
await router._execute_rag_migration(
CONTEXT,
task_context,
install_plugin=False,
)
assert connector.require_workspace_context.await_count == 2
connector.list_knowledge_engines.assert_not_awaited()
router._set_migration_flag.assert_not_awaited()
@pytest.mark.asyncio
async def test_cloud_migration_is_rejected_before_legacy_table_access():
router = object.__new__(KnowledgeMigrationRouterGroup)
router.ap = SimpleNamespace(
workspace_service=SimpleNamespace(
get_local_execution_binding=AsyncMock(side_effect=WorkspaceInvariantError('not an OSS local workspace')),
),
plugin_connector=SimpleNamespace(require_workspace_context=AsyncMock()),
logger=Mock(),
)
router._table_exists = AsyncMock()
router._set_migration_flag = AsyncMock()
task_context = SimpleNamespace(trace=Mock())
with pytest.raises(WorkspaceNotFoundError, match='migration is unavailable'):
await router._execute_rag_migration(CONTEXT, task_context, install_plugin=False)
router._table_exists.assert_not_awaited()
router.ap.plugin_connector.require_workspace_context.assert_not_awaited()
router._set_migration_flag.assert_not_awaited()
@pytest.mark.asyncio
@pytest.mark.parametrize('database_name, exists', [('postgresql', True), ('sqlite', False)])
async def test_legacy_table_discovery_uses_scoped_structured_queries(database_name: str, exists: bool):
result = Mock()
result.first.return_value = ('knowledge_bases_backup',) if exists else None
execute_async = AsyncMock(return_value=result)
router = object.__new__(KnowledgeMigrationRouterGroup)
router.ap = SimpleNamespace(
persistence_mgr=SimpleNamespace(
db=SimpleNamespace(name=database_name),
execute_async=execute_async,
)
)
assert await router._table_exists('knowledge_bases_backup') is exists
statement = execute_async.await_args.args[0]
assert isinstance(statement, sqlalchemy.sql.selectable.SelectBase)
_validate_scoped_statement_call((statement,), {})
@pytest.mark.asyncio
async def test_legacy_restore_emits_only_scoped_structured_statements():
missing_table_result = Mock()
missing_table_result.first.return_value = None
existing_table_result = Mock()
existing_table_result.first.return_value = ('knowledge_bases_backup',)
backup_result = Mock()
backup_result.keys.return_value = [
'uuid',
'name',
'description',
'emoji',
'embedding_model_uuid',
'top_k',
'created_at',
'updated_at',
]
now = datetime.now()
backup_result.fetchall.return_value = [
('kb-legacy', 'Legacy KB', 'Description', 'U0001f4da', 'embedding-model', 7, now, now)
]
execute_async = AsyncMock(
side_effect=[
missing_table_result,
existing_table_result,
backup_result,
Mock(),
Mock(),
]
)
connector = SimpleNamespace(
require_workspace_context=AsyncMock(return_value=CONTEXT),
list_knowledge_engines=AsyncMock(return_value=[]),
rag_on_kb_create=AsyncMock(),
)
router = object.__new__(KnowledgeMigrationRouterGroup)
router.ap = SimpleNamespace(
workspace_service=SimpleNamespace(
get_local_execution_binding=AsyncMock(return_value=CONTEXT),
),
plugin_connector=connector,
persistence_mgr=SimpleNamespace(
db=SimpleNamespace(name='sqlite'),
execute_async=execute_async,
),
rag_mgr=SimpleNamespace(load_knowledge_bases_from_db=AsyncMock()),
logger=Mock(),
)
task_context = SimpleNamespace(trace=Mock())
await router._execute_rag_migration(CONTEXT, task_context, install_plugin=False)
statements = [call.args[0] for call in execute_async.await_args_list]
assert any(isinstance(statement, sqlalchemy.sql.dml.Insert) for statement in statements)
assert any(isinstance(statement, sqlalchemy.sql.dml.Update) for statement in statements)
for statement in statements:
assert not isinstance(statement, sqlalchemy.sql.elements.TextClause)
_validate_scoped_statement_call((statement,), {})
connector.rag_on_kb_create.assert_awaited_once_with(
'langbot-team/LangRAG',
'kb-legacy',
{'embedding_model_uuid': 'embedding-model'},
)
@pytest.mark.asyncio
async def test_legacy_restore_accepts_sqlite_string_dates_and_text_json_columns():
engine = sqlalchemy.ext.asyncio.create_async_engine('sqlite+aiosqlite:///:memory:')
try:
async with engine.begin() as connection:
await connection.exec_driver_sql(
"""
CREATE TABLE knowledge_bases_backup (
uuid TEXT PRIMARY KEY,
name TEXT,
description TEXT,
emoji TEXT,
embedding_model_uuid TEXT,
top_k INTEGER,
created_at DATETIME,
updated_at DATETIME
)
"""
)
await connection.exec_driver_sql(
"""
CREATE TABLE external_knowledge_bases (
uuid TEXT PRIMARY KEY,
name TEXT,
description TEXT,
emoji TEXT,
plugin_author TEXT,
plugin_name TEXT,
retriever_config TEXT,
created_at DATETIME
)
"""
)
await connection.exec_driver_sql(
"""
CREATE TABLE knowledge_bases (
uuid TEXT PRIMARY KEY,
workspace_uuid TEXT NOT NULL,
name TEXT,
description TEXT,
emoji TEXT,
created_at DATETIME,
updated_at DATETIME,
knowledge_engine_plugin_id TEXT,
collection_id TEXT,
creation_settings TEXT,
retrieval_settings TEXT
)
"""
)
await connection.exec_driver_sql(
'CREATE TABLE workspace_metadata (workspace_uuid TEXT, key TEXT, value TEXT)'
)
legacy_timestamp = '2026-07-20 03:00:00.123456'
await connection.exec_driver_sql(
"""
INSERT INTO knowledge_bases_backup
(uuid, name, description, emoji, embedding_model_uuid, top_k, created_at, updated_at)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
""",
(
'kb-internal',
'Internal',
'Internal legacy KB',
'U0001f4da',
'embedding-model',
5,
legacy_timestamp,
legacy_timestamp,
),
)
await connection.exec_driver_sql(
"""
INSERT INTO external_knowledge_bases
(uuid, name, description, emoji, plugin_author, plugin_name, retriever_config, created_at)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
""",
(
'kb-external',
'External',
'External legacy KB',
'U0001f517',
'langbot-team',
'DifyDatasetsRetriever',
json.dumps({'api_base_url': 'https://example.invalid', 'top_k': 8}),
legacy_timestamp,
),
)
connector = SimpleNamespace(
require_workspace_context=AsyncMock(return_value=CONTEXT),
list_knowledge_engines=AsyncMock(return_value=[]),
rag_on_kb_create=AsyncMock(),
)
router = object.__new__(KnowledgeMigrationRouterGroup)
router.ap = SimpleNamespace(
workspace_service=SimpleNamespace(
get_local_execution_binding=AsyncMock(return_value=CONTEXT),
),
plugin_connector=connector,
persistence_mgr=SimpleNamespace(
db=SimpleNamespace(name='sqlite'),
execute_async=connection.execute,
),
rag_mgr=SimpleNamespace(load_knowledge_bases_from_db=AsyncMock()),
logger=Mock(),
)
await router._execute_rag_migration(
CONTEXT,
SimpleNamespace(trace=Mock()),
install_plugin=False,
)
restored = (
await connection.exec_driver_sql(
"""
SELECT uuid, created_at, updated_at, creation_settings, retrieval_settings
FROM knowledge_bases
ORDER BY uuid
"""
)
).all()
assert [row.uuid for row in restored] == ['kb-external', 'kb-internal']
assert all(row.created_at == legacy_timestamp for row in restored)
assert all(row.updated_at == legacy_timestamp for row in restored)
assert json.loads(restored[0].creation_settings)['api_base_url'] == 'https://example.invalid'
assert json.loads(restored[0].retrieval_settings) == {'top_k': 8}
assert json.loads(restored[1].creation_settings) == {'embedding_model_uuid': 'embedding-model'}
assert json.loads(restored[1].retrieval_settings) == {'top_k': 5}
finally:
await engine.dispose()
tests/unit_tests/api/test_mcp_controller.py
+63 -7
View File
@@ -17,19 +17,63 @@ sys.modules.setdefault('langbot.pkg.core.app', core_app_module)
pytestmark = pytest.mark.asyncio
async def _create_test_client(mcp_service: SimpleNamespace):
async def _create_test_client(mcp_service: SimpleNamespace, *, role: str = 'owner'):
app = quart.Quart(__name__)
user_service = SimpleNamespace(
verify_jwt_token=AsyncMock(return_value='test@example.com'),
get_user_by_email=AsyncMock(return_value=SimpleNamespace(user='test@example.com')),
get_user_by_email=AsyncMock(
return_value=SimpleNamespace(
user='test@example.com',
uuid='account-a',
)
),
)
workspace_collaboration_service = SimpleNamespace(
resolve_account_workspace=AsyncMock(
return_value=SimpleNamespace(
execution=SimpleNamespace(
instance_uuid='instance-a',
placement_generation=1,
),
workspace=SimpleNamespace(uuid='workspace-a'),
membership=SimpleNamespace(
uuid='membership-a',
role=role,
projection_revision=1,
),
)
)
)
ap = SimpleNamespace(
mcp_service=mcp_service,
user_service=user_service,
workspace_collaboration_service=workspace_collaboration_service,
)
ap = SimpleNamespace(mcp_service=mcp_service, user_service=user_service)
MCPRouterGroup = import_module('langbot.pkg.api.http.controller.groups.resources.mcp').MCPRouterGroup
group = MCPRouterGroup(ap, app)
await group.initialize()
return app.test_client()
async def test_viewer_cannot_read_mcp_runtime_logs():
mcp_service = SimpleNamespace(
get_mcp_server_logs=AsyncMock(return_value=['private runtime line']),
)
client = await _create_test_client(mcp_service, role='viewer')
response = await client.get(
'/api/v1/mcp/servers/example/logs',
headers={
'Authorization': 'Bearer test-token',
'X-Workspace-Id': 'workspace-a',
},
)
assert response.status_code == 403
assert (await response.get_json())['code'] == 'permission_denied'
mcp_service.get_mcp_server_logs.assert_not_awaited()
async def test_mcp_server_route_accepts_encoded_slash_name():
mcp_service = SimpleNamespace(
get_mcp_server_by_name=AsyncMock(
@@ -46,11 +90,17 @@ async def test_mcp_server_route_accepts_encoded_slash_name():
response = await client.get(
'/api/v1/mcp/servers/pab1it0%2Fprometheus',
headers={'Authorization': 'Bearer test-token'},
headers={
'Authorization': 'Bearer test-token',
'X-Workspace-Id': 'workspace-a',
},
)
assert response.status_code == 200
mcp_service.get_mcp_server_by_name.assert_awaited_once_with('pab1it0/prometheus')
mcp_service.get_mcp_server_by_name.assert_awaited_once()
context, server_name = mcp_service.get_mcp_server_by_name.await_args.args
assert context.workspace_uuid == 'workspace-a'
assert server_name == 'pab1it0/prometheus'
payload = await response.get_json()
assert payload['data']['server']['name'] == 'pab1it0/prometheus'
@@ -66,11 +116,17 @@ async def test_mcp_resource_route_accepts_encoded_slash_name():
response = await client.get(
'/api/v1/mcp/servers/pab1it0%2Fprometheus/resources',
headers={'Authorization': 'Bearer test-token'},
headers={
'Authorization': 'Bearer test-token',
'X-Workspace-Id': 'workspace-a',
},
)
assert response.status_code == 200
mcp_service.get_mcp_server_by_name.assert_not_awaited()
mcp_service.get_mcp_server_resources.assert_awaited_once_with('pab1it0/prometheus')
mcp_service.get_mcp_server_resources.assert_awaited_once()
context, server_name = mcp_service.get_mcp_server_resources.await_args.args
assert context.workspace_uuid == 'workspace-a'
assert server_name == 'pab1it0/prometheus'
payload = await response.get_json()
assert payload['data']['resource_capabilities'] == {'subscribe': False}
tests/unit_tests/api/test_mcp_mount_tenant_scope.py
+121
View File
@@ -0,0 +1,121 @@
from __future__ import annotations
import asyncio
from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
import sqlalchemy as sa
from sqlalchemy.ext.asyncio import create_async_engine
from langbot.pkg.api.http.service.apikey import ApiKeyIdentity
from langbot.pkg.api.mcp.context import get_request_context
from langbot.pkg.api.mcp.mount import MCPMount
from langbot.pkg.persistence.mgr import PersistenceManager, PersistenceMode
from langbot.pkg.persistence.tenant_uow import PersistenceScopeKind
@pytest.mark.asyncio
async def test_mcp_mount_keeps_request_context_but_no_session_during_stream_wait(tmp_path) -> None:
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "mcp-short-scope.db"}')
table = sa.Table('mcp_scope_probe', sa.MetaData(), sa.Column('id', sa.Integer, primary_key=True))
manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
manager.db = SimpleNamespace(get_engine=lambda: engine)
checked_out = 0
def on_checkout(*_args):
nonlocal checked_out
checked_out += 1
def on_checkin(*_args):
nonlocal checked_out
checked_out -= 1
sa.event.listen(engine.sync_engine, 'checkout', on_checkout)
sa.event.listen(engine.sync_engine, 'checkin', on_checkin)
try:
async with engine.begin() as conn:
await conn.run_sync(table.metadata.create_all)
identity = ApiKeyIdentity(
instance_uuid='instance-1',
workspace_uuid='workspace-1',
placement_generation=7,
api_key_uuid='key-1',
permissions=frozenset({'pipelines:read'}),
)
app = SimpleNamespace(
apikey_service=SimpleNamespace(authenticate_api_key=AsyncMock(return_value=identity)),
persistence_mgr=manager,
deployment_admission=None,
deployment=None,
)
stream_waiting = asyncio.Event()
release_stream = asyncio.Event()
observations: list[tuple[str, str, bool]] = []
async def fake_mcp_asgi(scope, receive, send):
del scope, receive
context = get_request_context()
assert manager.current_scope().kind is PersistenceScopeKind.WORKSPACE
assert manager.current_session() is None
await manager.execute_async(sa.select(table.c.id))
assert manager.current_session() is None
observations.append((context.request_id, context.workspace_uuid, manager.current_session() is None))
stream_waiting.set()
await release_stream.wait()
preserved_context = get_request_context()
observations.append(
(
preserved_context.request_id,
preserved_context.workspace_uuid,
manager.current_session() is None,
)
)
await manager.execute_async(sa.select(table.c.id))
assert manager.current_session() is None
await send({'type': 'http.response.start', 'status': 200, 'headers': []})
await send({'type': 'http.response.body', 'body': b'{}'})
async def unused_quart_asgi(scope, receive, send):
del scope, receive, send
raise AssertionError('MCP request was routed to Quart')
mount = MCPMount.__new__(MCPMount)
mount.ap = app
mount._mcp_asgi = fake_mcp_asgi
sent_messages: list[dict] = []
async def receive():
return {'type': 'http.request', 'body': b'', 'more_body': False}
async def send(message):
sent_messages.append(message)
async def release_after_observation() -> None:
await asyncio.wait_for(stream_waiting.wait(), timeout=2)
assert checked_out == 0
release_stream.set()
release_task = asyncio.create_task(release_after_observation())
await mount.wrap(unused_quart_asgi)(
{
'type': 'http',
'path': '/mcp',
'headers': [(b'x-api-key', b'secret')],
},
receive,
send,
)
await release_task
assert sent_messages[0]['status'] == 200
assert len(observations) == 2
assert observations[0] == observations[1]
assert observations[0][1:] == ('workspace-1', True)
assert checked_out == 0
assert manager.current_scope() is None
with pytest.raises(RuntimeError, match='context is unavailable'):
get_request_context()
finally:
await engine.dispose()
tests/unit_tests/api/test_plugin_runtime_route_fence.py
+139
View File
@@ -0,0 +1,139 @@
from __future__ import annotations
from contextlib import asynccontextmanager
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
import pytest
import quart
from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.workspace.errors import WorkspaceNotFoundError
CONTEXT = ExecutionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=4,
)
@pytest.fixture(scope='module')
def plugin_router_cls():
from tests.utils.import_isolation import MockLifecycleControlScope, isolated_sys_modules
class FakeMinimalApplication:
pass
mock_app = Mock(Application=FakeMinimalApplication)
mock_entities = Mock(LifecycleControlScope=MockLifecycleControlScope)
clear = [
'langbot.pkg.core.taskmgr',
'langbot.pkg.api.http.controller.group',
'langbot.pkg.api.http.controller.groups',
'langbot.pkg.api.http.controller.groups.plugins',
'langbot.pkg.api.http.controller.main',
]
with isolated_sys_modules(
mocks={
'langbot.pkg.core.app': mock_app,
'langbot.pkg.core.entities': mock_entities,
},
clear=clear,
):
from langbot.pkg.api.http.controller.groups.plugins import PluginsRouterGroup
yield PluginsRouterGroup
@pytest.mark.asyncio
async def test_public_plugin_asset_route_is_disabled_for_multi_workspace_policy(plugin_router_cls):
connector = SimpleNamespace(
get_plugin_icon=AsyncMock(),
require_workspace_context=AsyncMock(),
)
ap = SimpleNamespace(
plugin_connector=connector,
workspace_service=SimpleNamespace(
policy=SimpleNamespace(multi_workspace_enabled=True),
),
)
quart_app = quart.Quart(__name__)
router = plugin_router_cls(ap, quart_app)
await router.initialize()
response = await quart_app.test_client().get('/api/v1/plugins/author/plugin/icon')
assert response.status_code == 404
connector.require_workspace_context.assert_not_awaited()
connector.get_plugin_icon.assert_not_awaited()
@pytest.mark.asyncio
async def test_public_plugin_asset_uses_trusted_oss_singleton_binding(plugin_router_cls):
connector = SimpleNamespace(
require_workspace_context=AsyncMock(side_effect=lambda context: context),
)
binding = SimpleNamespace(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=4,
)
router = object.__new__(plugin_router_cls)
router.ap = SimpleNamespace(
plugin_connector=connector,
workspace_service=SimpleNamespace(
policy=SimpleNamespace(multi_workspace_enabled=False),
get_local_execution_binding=AsyncMock(return_value=binding),
),
)
result = await router._require_public_plugin_runtime_context()
assert result == CONTEXT
connector.require_workspace_context.assert_awaited_once_with(CONTEXT)
@pytest.mark.asyncio
async def test_background_plugin_operation_refences_captured_generation(plugin_router_cls):
operation = AsyncMock()
connector = SimpleNamespace(
require_workspace_context=AsyncMock(side_effect=WorkspaceNotFoundError('Plugin resource not found')),
)
router = object.__new__(plugin_router_cls)
router.ap = SimpleNamespace(plugin_connector=connector)
with pytest.raises(WorkspaceNotFoundError, match='Plugin resource not found'):
await router._run_fenced_plugin_operation(CONTEXT, operation)
operation.assert_not_awaited()
@pytest.mark.asyncio
async def test_background_plugin_operation_revalidates_inside_short_tenant_uow(plugin_router_cls):
scopes = []
@asynccontextmanager
async def tenant_uow(workspace_uuid):
scopes.append(workspace_uuid)
yield
connector = SimpleNamespace(
require_workspace_context=AsyncMock(side_effect=lambda context: context),
)
operation = AsyncMock(return_value='done')
router = object.__new__(plugin_router_cls)
router.ap = SimpleNamespace(
plugin_connector=connector,
persistence_mgr=SimpleNamespace(
mode=SimpleNamespace(value='cloud_runtime'),
tenant_uow=tenant_uow,
),
)
result = await router._run_fenced_plugin_operation(CONTEXT, operation)
assert result == 'done'
assert scopes == [CONTEXT.workspace_uuid]
connector.require_workspace_context.assert_awaited_once_with(CONTEXT)
operation.assert_awaited_once()
tests/unit_tests/api/test_resource_secret_permissions.py
+196
View File
@@ -0,0 +1,196 @@
from __future__ import annotations
import copy
from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
import quart
from langbot.pkg.api.http.controller.groups.knowledge.base import KnowledgeBaseRouterGroup
from langbot.pkg.api.http.controller.groups.pipelines.pipelines import PipelinesRouterGroup
from langbot.pkg.api.http.controller.groups.provider.models import LLMModelsRouterGroup
from langbot.pkg.api.http.controller.groups.provider.providers import ModelProvidersRouterGroup
from langbot.pkg.api.http.controller.groups.resources.mcp import MCPRouterGroup
from langbot.pkg.api.http.controller.groups.webhook_mgmt import WebhookManagementRouterGroup
from langbot.pkg.api.http.service.secrets import mask_secret_value, redact_secrets
pytestmark = pytest.mark.asyncio
WORKSPACE_UUID = '11111111-1111-4111-8111-111111111111'
RAW_PIPELINE = {
'uuid': 'pipeline-test',
'config': {'ai': {'n8n': {'webhook-url': 'https://hook.invalid/bearer-secret'}}},
}
RAW_MODEL = {
'uuid': 'model-test',
'provider_uuid': 'provider-test',
'extra_args': {'headers': {'Authorization': 'Bearer model-secret'}},
}
RAW_PROVIDER = {
'uuid': 'provider-test',
'base_url': 'https://provider-user:provider-password@provider.invalid/v1?token=url-secret&region=sg',
'api_keys': ['provider-secret'],
}
RAW_MCP_SERVER = {
'uuid': 'mcp-test',
'name': 'MCP Test',
'extra_args': {'url': 'https://mcp-user:mcp-password@mcp.invalid/connect?api_key=url-secret&transport=http'},
}
RAW_KNOWLEDGE_BASE = {
'uuid': 'kb-test',
'creation_settings': {'dify_apikey': 'knowledge-secret'},
}
RAW_WEBHOOK = {'id': 1, 'url': 'https://hook.invalid/path?token=webhook-secret'}
def _access(role: str):
return SimpleNamespace(
workspace=SimpleNamespace(uuid=WORKSPACE_UUID),
membership=SimpleNamespace(uuid='membership-test', role=role, projection_revision=1),
execution=SimpleNamespace(instance_uuid='instance-test', placement_generation=1),
)
async def _create_client(role: str):
application = SimpleNamespace()
account = SimpleNamespace(uuid='account-test', user='test@example.com')
application.user_service = SimpleNamespace(get_authenticated_account=AsyncMock(return_value=account))
application.apikey_service = SimpleNamespace(authenticate_api_key=AsyncMock(return_value=None))
application.workspace_collaboration_service = SimpleNamespace(
resolve_account_workspace=AsyncMock(return_value=_access(role))
)
async def get_pipelines(_context, *_args, include_secret=False):
value = copy.deepcopy(RAW_PIPELINE)
return [value] if include_secret else [redact_secrets(value)]
async def get_pipeline(_context, _uuid, *, include_secret=False):
value = copy.deepcopy(RAW_PIPELINE)
return value if include_secret else redact_secrets(value)
application.pipeline_service = SimpleNamespace(
get_pipelines=AsyncMock(side_effect=get_pipelines),
get_pipeline=AsyncMock(side_effect=get_pipeline),
)
application.plugin_connector = SimpleNamespace(list_plugins=AsyncMock(return_value=[]))
application.mcp_service = SimpleNamespace(
get_mcp_servers=AsyncMock(return_value=[redact_secrets(copy.deepcopy(RAW_MCP_SERVER))])
)
application.skill_service = SimpleNamespace(list_skills=AsyncMock(return_value=[]))
async def get_models_by_provider(_context, _provider_uuid, *, include_secret=False):
value = copy.deepcopy(RAW_MODEL)
return [value] if include_secret else [redact_secrets(value)]
application.llm_model_service = SimpleNamespace(
get_llm_models_by_provider=AsyncMock(side_effect=get_models_by_provider)
)
async def get_providers(_context, *, include_secret=False):
value = copy.deepcopy(RAW_PROVIDER)
return [value] if include_secret else [redact_secrets(value)]
application.provider_service = SimpleNamespace(
get_providers=AsyncMock(side_effect=get_providers),
get_provider_model_counts=AsyncMock(return_value={'llm_count': 0, 'embedding_count': 0, 'rerank_count': 0}),
)
async def get_knowledge_bases(_context, *, include_secret=False):
value = copy.deepcopy(RAW_KNOWLEDGE_BASE)
return [value] if include_secret else [redact_secrets(value)]
application.knowledge_service = SimpleNamespace(get_knowledge_bases=AsyncMock(side_effect=get_knowledge_bases))
async def get_webhooks(_context, *, include_secret=False):
value = copy.deepcopy(RAW_WEBHOOK)
if not include_secret:
value['url'] = mask_secret_value(value['url'])
return [value]
application.webhook_service = SimpleNamespace(get_webhooks=AsyncMock(side_effect=get_webhooks))
quart_app = quart.Quart(__name__)
for router_type in (
PipelinesRouterGroup,
LLMModelsRouterGroup,
ModelProvidersRouterGroup,
MCPRouterGroup,
KnowledgeBaseRouterGroup,
WebhookManagementRouterGroup,
):
await router_type(application, quart_app).initialize()
return application, quart_app.test_client()
def _headers() -> dict[str, str]:
return {'Authorization': 'Bearer test-token', 'X-Workspace-Id': WORKSPACE_UUID}
@pytest.mark.parametrize('role', ['viewer', 'operator'])
async def test_viewer_and_operator_resource_reads_are_redacted(role: str):
application, client = await _create_client(role)
pipeline = (await (await client.get('/api/v1/pipelines', headers=_headers())).get_json())['data']['pipelines'][0]
model = (
await (
await client.get(
'/api/v1/provider/models/llm?provider_uuid=provider-test',
headers=_headers(),
)
).get_json()
)['data']['models'][0]
provider = (await (await client.get('/api/v1/provider/providers', headers=_headers())).get_json())['data'][
'providers'
][0]
mcp_server = (await (await client.get('/api/v1/mcp/servers', headers=_headers())).get_json())['data']['servers'][0]
knowledge_base = (await (await client.get('/api/v1/knowledge/bases', headers=_headers())).get_json())['data'][
'bases'
][0]
webhook = (await (await client.get('/api/v1/webhooks', headers=_headers())).get_json())['data']['webhooks'][0]
assert pipeline['config']['ai']['n8n']['webhook-url'] == '***'
assert model['extra_args']['headers']['Authorization'] == '***'
assert provider['api_keys'] == ['***']
assert provider['base_url'] == 'https://***@provider.invalid/v1?token=***&region=sg'
assert mcp_server['extra_args']['url'] == 'https://***@mcp.invalid/connect?api_key=***&transport=http'
assert knowledge_base['creation_settings']['dify_apikey'] == '***'
assert webhook['url'] == '***'
assert application.pipeline_service.get_pipelines.await_args.kwargs['include_secret'] is False
assert application.llm_model_service.get_llm_models_by_provider.await_args.kwargs['include_secret'] is False
assert application.provider_service.get_providers.await_args.kwargs['include_secret'] is False
assert application.knowledge_service.get_knowledge_bases.await_args.kwargs['include_secret'] is False
assert application.webhook_service.get_webhooks.await_args.kwargs['include_secret'] is False
async def test_resource_manager_receives_credentials_needed_for_management():
application, client = await _create_client('developer')
pipeline = (await (await client.get('/api/v1/pipelines', headers=_headers())).get_json())['data']['pipelines'][0]
model = (
await (
await client.get(
'/api/v1/provider/models/llm?provider_uuid=provider-test',
headers=_headers(),
)
).get_json()
)['data']['models'][0]
provider = (await (await client.get('/api/v1/provider/providers', headers=_headers())).get_json())['data'][
'providers'
][0]
knowledge_base = (await (await client.get('/api/v1/knowledge/bases', headers=_headers())).get_json())['data'][
'bases'
][0]
webhook = (await (await client.get('/api/v1/webhooks', headers=_headers())).get_json())['data']['webhooks'][0]
assert pipeline == RAW_PIPELINE
assert model == RAW_MODEL
assert provider['api_keys'] == ['provider-secret']
assert knowledge_base == RAW_KNOWLEDGE_BASE
assert webhook == RAW_WEBHOOK
assert application.pipeline_service.get_pipelines.await_args.kwargs['include_secret'] is True
assert application.llm_model_service.get_llm_models_by_provider.await_args.kwargs['include_secret'] is True
assert application.provider_service.get_providers.await_args.kwargs['include_secret'] is True
assert application.knowledge_service.get_knowledge_bases.await_args.kwargs['include_secret'] is True
assert application.webhook_service.get_webhooks.await_args.kwargs['include_secret'] is True
tests/unit_tests/api/test_stats_controller.py
+110
View File
@@ -0,0 +1,110 @@
from __future__ import annotations
from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
import quart
from langbot.pkg.api.http.controller.groups.stats import StatsRouterGroup
pytestmark = pytest.mark.asyncio
def session(
workspace_uuid: str,
*,
placement_generation: int = 1,
conversation_count: int = 0,
):
return SimpleNamespace(
instance_uuid='instance-test',
workspace_uuid=workspace_uuid,
placement_generation=placement_generation,
conversations=[object() for _ in range(conversation_count)],
)
async def create_client(*, role='viewer'):
quart_app = quart.Quart(__name__)
account = SimpleNamespace(uuid='account-test', user='test@example.com')
user_service = SimpleNamespace(
get_authenticated_account=AsyncMock(return_value=account),
)
access = SimpleNamespace(
execution=SimpleNamespace(
instance_uuid='instance-test',
placement_generation=1,
),
workspace=SimpleNamespace(uuid='workspace-a'),
membership=SimpleNamespace(
uuid='membership-test',
role=role,
projection_revision=1,
),
)
collaboration_service = SimpleNamespace(
resolve_account_workspace=AsyncMock(return_value=access),
)
def get_query_count(context):
assert context.instance_uuid == 'instance-test'
assert context.workspace_uuid == 'workspace-a'
assert context.placement_generation == 1
return 7
ap = SimpleNamespace(
user_service=user_service,
workspace_collaboration_service=collaboration_service,
sess_mgr=SimpleNamespace(
session_list=[
session('workspace-a', conversation_count=2),
session('workspace-b', conversation_count=5),
session(
'workspace-a',
placement_generation=2,
conversation_count=3,
),
SimpleNamespace(conversations=[object()] * 11),
]
),
query_pool=SimpleNamespace(get_query_count=get_query_count),
)
router = StatsRouterGroup(ap, quart_app)
await router.initialize()
return quart_app.test_client(), collaboration_service
async def test_basic_stats_are_scoped_to_selected_workspace_placement():
client, collaboration_service = await create_client()
response = await client.get(
'/api/v1/stats/basic',
headers={
'Authorization': 'Bearer test-token',
'X-Workspace-Id': 'workspace-a',
},
)
assert response.status_code == 200
payload = await response.get_json()
assert payload['data'] == {
'active_session_count': 1,
'conversation_count': 2,
'query_count': 7,
}
collaboration_service.resolve_account_workspace.assert_awaited_once_with('account-test', 'workspace-a')
async def test_basic_stats_requires_resource_view_permission():
client, _ = await create_client(role='unknown-role')
response = await client.get(
'/api/v1/stats/basic',
headers={'Authorization': 'Bearer test-token'},
)
assert response.status_code == 403
payload = await response.get_json()
assert payload['code'] == 'permission_denied'
tests/unit_tests/api/test_websocket_chat_uow.py
+144
View File
@@ -0,0 +1,144 @@
from __future__ import annotations
import asyncio
from contextlib import asynccontextmanager
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
import pytest
from langbot.pkg.api.http.context import (
PrincipalContext,
PrincipalType,
RequestContext,
WorkspaceContext,
)
from langbot.pkg.api.http.controller.groups.pipelines.websocket_chat import WebSocketChatRouterGroup
from langbot.pkg.api.http.controller.groups.pipelines.websocket_chat import (
create_scoped_duplex_tasks,
)
from langbot.pkg.api.http.controller.groups.pipelines.websocket_chat import wait_for_duplex_tasks
from langbot.pkg.utils.bounded_executor import current_blocking_work_scope
@pytest.mark.asyncio
async def test_websocket_pipeline_lookup_opens_workspace_uow_after_auth_scope_closed() -> None:
workspace_uuid = 'workspace-a'
scopes: list[str] = []
in_scope = False
@asynccontextmanager
async def tenant_uow(selected_workspace_uuid: str):
nonlocal in_scope
assert not in_scope
in_scope = True
scopes.append(selected_workspace_uuid)
try:
yield
finally:
in_scope = False
async def get_pipeline(_context, _pipeline_uuid):
assert in_scope
return {'uuid': 'pipeline-a'}
adapter = Mock()
router = object.__new__(WebSocketChatRouterGroup)
router.ap = SimpleNamespace(
persistence_mgr=SimpleNamespace(
mode=SimpleNamespace(value='cloud_runtime'),
tenant_uow=tenant_uow,
),
pipeline_service=SimpleNamespace(get_pipeline=AsyncMock(side_effect=get_pipeline)),
platform_mgr=SimpleNamespace(get_websocket_proxy_bot=AsyncMock(return_value=SimpleNamespace(adapter=adapter))),
)
request_context = RequestContext(
instance_uuid='instance-a',
placement_generation=1,
request_id='request-a',
auth_type='user_token',
principal=PrincipalContext(
principal_type=PrincipalType.ACCOUNT,
account_uuid='account-a',
),
workspace=WorkspaceContext(
workspace_uuid=workspace_uuid,
membership_uuid='membership-a',
role='owner',
permissions=frozenset(),
),
)
result = await router._get_scoped_adapter(request_context, 'pipeline-a')
assert result is adapter
assert scopes == [workspace_uuid]
@pytest.mark.asyncio
async def test_duplex_websocket_tasks_cancel_blocked_peer_when_one_direction_ends() -> None:
blocked = asyncio.Event()
async def receive_forever() -> None:
blocked.set()
await asyncio.Future()
async def send_finishes() -> None:
await blocked.wait()
receive_task = asyncio.create_task(receive_forever())
send_task = asyncio.create_task(send_finishes())
await asyncio.wait_for(
wait_for_duplex_tasks(receive_task, send_task),
timeout=1,
)
assert receive_task.cancelled()
assert send_task.done()
@pytest.mark.asyncio
async def test_duplex_websocket_tasks_allow_terminal_send_to_drain() -> None:
receive_finished = asyncio.Event()
send_drained = asyncio.Event()
async def receive_finishes() -> None:
receive_finished.set()
async def send_terminal_frame() -> None:
await receive_finished.wait()
await asyncio.sleep(0)
send_drained.set()
receive_task = asyncio.create_task(receive_finishes())
send_task = asyncio.create_task(send_terminal_frame())
await wait_for_duplex_tasks(receive_task, send_task)
assert send_drained.is_set()
assert send_task.done()
assert not send_task.cancelled()
@pytest.mark.asyncio
async def test_duplex_websocket_tasks_share_trusted_workspace_budget() -> None:
observed: list[tuple[str, str | None]] = []
async def observe(direction: str) -> None:
await asyncio.sleep(0)
observed.append((direction, current_blocking_work_scope()))
receive_task, send_task = create_scoped_duplex_tasks(
observe('receive'),
observe('send'),
'workspace-a',
)
await asyncio.gather(receive_task, send_task)
assert sorted(observed) == [
('receive', 'workspace-a'),
('send', 'workspace-a'),
]
assert current_blocking_work_scope() is None
tests/unit_tests/box/test_box_admission.py
+230
View File
@@ -0,0 +1,230 @@
from __future__ import annotations
import datetime as dt
import hashlib
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
import pytest
from langbot_plugin.box.models import SandboxAdmissionPolicy
from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.box.admission import (
BoxAdmissionError,
SandboxAdmissionController,
require_cloud_admission_policy,
)
from langbot.pkg.box.service import BoxService
from langbot.pkg.cloud.entitlements import (
EntitlementResolver,
EntitlementSnapshot,
EntitlementUnavailableError,
)
_UTC = dt.timezone.utc
_CONTEXT = ExecutionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=3,
entitlement_revision=7,
)
def _snapshot(
*,
revision: int = 7,
managed: bool = True,
sessions: int = 1,
expires_at: int = 2_000,
) -> EntitlementSnapshot:
return EntitlementSnapshot(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
entitlement_revision=revision,
status='active',
not_before=1,
expires_at=expires_at,
features={'managed_sandbox': managed},
limits={'managed_sandbox_sessions': sessions},
)
def _controller(snapshot: EntitlementSnapshot, *, now: float = 1_000.25):
provider = SimpleNamespace(get_workspace_entitlement=AsyncMock(return_value=snapshot))
resolver = EntitlementResolver('instance-a', provider)
client = SimpleNamespace(
upsert_sandbox_admission_grant=AsyncMock(
side_effect=lambda grant: {
'installed': True,
'workspace_uuid': grant.workspace_uuid,
'execution_generation': grant.execution_generation,
'entitlement_revision': grant.entitlement_revision,
'max_sessions': grant.max_sessions,
'max_managed_processes': grant.max_managed_processes,
}
),
revoke_sandbox_admission_grant=AsyncMock(
side_effect=lambda revocation: {
'revoked': True,
'workspace_uuid': revocation.workspace_uuid,
'entitlement_revision': revocation.entitlement_revision,
}
),
)
app = SimpleNamespace(entitlement_resolver=resolver, logger=Mock())
controller = SandboxAdmissionController(
app,
client,
policy=SandboxAdmissionPolicy(required=True, max_grant_ttl_sec=300),
wall_time=lambda: now,
)
return controller, client, provider
def test_cloud_admission_policy_requires_positive_workspace_quota():
with pytest.raises(BoxAdmissionError, match='workspace quota must be a positive integer'):
require_cloud_admission_policy(
{
'required': True,
'workspace_quota_mb': 0,
}
)
policy = require_cloud_admission_policy(
{
'required': True,
'workspace_quota_mb': 32,
}
)
assert policy.workspace_quota_mb == 32
@pytest.mark.asyncio
async def test_active_generic_entitlement_installs_short_lived_numeric_grant():
controller, client, provider = _controller(_snapshot())
grant = await controller.require(_CONTEXT)
assert grant.instance_uuid == _CONTEXT.instance_uuid
assert grant.workspace_uuid == _CONTEXT.workspace_uuid
assert grant.execution_generation == _CONTEXT.placement_generation
assert grant.entitlement_revision == 7
assert grant.max_sessions == 1
assert grant.max_managed_processes == 0
assert grant.expires_at == dt.datetime.fromtimestamp(1_300, tz=_UTC)
assert (grant.expires_at - dt.datetime.fromtimestamp(1_000.25, tz=_UTC)).total_seconds() < 300
provider.get_workspace_entitlement.assert_awaited_once_with('workspace-a')
client.upsert_sandbox_admission_grant.assert_awaited_once_with(grant)
client.revoke_sandbox_admission_grant.assert_not_awaited()
@pytest.mark.asyncio
@pytest.mark.parametrize(
'snapshot',
[
_snapshot(managed=False),
_snapshot(sessions=0),
_snapshot(sessions=2),
],
)
async def test_non_eligible_entitlement_revokes_and_fails_closed(snapshot):
controller, client, _provider = _controller(snapshot)
with pytest.raises(EntitlementUnavailableError):
await controller.require(_CONTEXT)
client.upsert_sandbox_admission_grant.assert_not_awaited()
revocation = client.revoke_sandbox_admission_grant.await_args.args[0]
assert revocation.entitlement_revision == snapshot.entitlement_revision
@pytest.mark.asyncio
async def test_transient_entitlement_failure_does_not_tombstone_valid_revision():
controller, client, provider = _controller(_snapshot())
await controller.require(_CONTEXT)
provider.get_workspace_entitlement.side_effect = RuntimeError('control plane unavailable')
with pytest.raises(RuntimeError, match='control plane unavailable'):
await controller.require(_CONTEXT)
client.revoke_sandbox_admission_grant.assert_not_awaited()
provider.get_workspace_entitlement.side_effect = None
provider.get_workspace_entitlement.return_value = _snapshot()
recovered = await controller.require(_CONTEXT)
assert recovered.entitlement_revision == 7
@pytest.mark.asyncio
async def test_runtime_receipt_mismatch_is_revoked_and_never_admitted():
controller, client, _provider = _controller(_snapshot())
client.upsert_sandbox_admission_grant.return_value = {'installed': True, 'workspace_uuid': 'other'}
client.upsert_sandbox_admission_grant.side_effect = None
with pytest.raises(Exception, match='invalid sandbox admission receipt'):
await controller.require(_CONTEXT)
client.revoke_sandbox_admission_grant.assert_not_awaited()
@pytest.mark.asyncio
async def test_authoritative_cancelled_revision_is_revoked():
cancelled = _snapshot(revision=8).model_copy(update={'status': 'cancelled'})
controller, client, _provider = _controller(cancelled)
with pytest.raises(EntitlementUnavailableError, match='not active'):
await controller.require(_CONTEXT)
revocation = client.revoke_sandbox_admission_grant.await_args.args[0]
assert revocation.entitlement_revision == 8
@pytest.mark.asyncio
async def test_cloud_box_readiness_failure_aborts_service_initialization(tmp_path):
workspace_root = tmp_path / 'box' / 'workspaces'
workspace_root.mkdir(parents=True)
box_config = {
'enabled': True,
'backend': 'nsjail',
'runtime': {'endpoint': 'ws://box:5410'},
'local': {
'host_root': str(tmp_path / 'box'),
'default_workspace': str(workspace_root),
'allowed_mount_roots': [str(tmp_path / 'box')],
},
'admission': {
'required': True,
'logical_session_id': 'global',
'required_backend': 'nsjail',
'max_sessions': 1,
'max_managed_processes': 0,
'max_grant_ttl_sec': 300,
'workspace_quota_mb': 32,
},
}
client = SimpleNamespace(
initialize=AsyncMock(),
verify_shared_workspace=AsyncMock(
side_effect=lambda marker_name: {
'marker_name': marker_name,
'size': (workspace_root / marker_name).stat().st_size,
'sha256': hashlib.sha256((workspace_root / marker_name).read_bytes()).hexdigest(),
}
),
get_backend_info=AsyncMock(return_value={'name': 'docker', 'available': True}),
)
app = SimpleNamespace(
logger=Mock(),
deployment=SimpleNamespace(multi_workspace_enabled=True),
entitlement_resolver=Mock(),
workspace_service=SimpleNamespace(instance_uuid='instance-a'),
instance_config=SimpleNamespace(data={'box': box_config}),
)
service = BoxService(app, client=client)
with pytest.raises(Exception, match='nsjail isolation readiness failed'):
await service.initialize()
assert service.available is False
tests/unit_tests/box/test_box_connector.py
+142 -1
View File
@@ -6,14 +6,28 @@ from unittest.mock import AsyncMock, Mock
import pytest
from langbot_plugin.box.client import ActionRPCBoxClient
from langbot.pkg.box import connector as connector_module
from langbot_plugin.box.client import ActionRPCBoxClient
from langbot_plugin.box.errors import BoxRuntimeUnavailableError
from langbot_plugin.box.security import (
BOX_CONTROL_TOKEN_ENV,
BOX_CONTROL_TOKEN_HEADER,
BOX_INSTANCE_HEADER,
BOX_PLACEMENT_GENERATION_HEADER,
BOX_TRUSTED_INSTANCE_ENV,
BOX_WORKSPACE_HEADER,
)
from langbot_plugin.entities.io.context import ActionContext
from langbot.pkg.box.connector import BoxRuntimeConnector
_CONTROL_TOKEN = 'box-control-token-that-is-longer-than-32-bytes'
def make_app(logger: Mock, runtime_endpoint: str = ''):
return SimpleNamespace(
logger=logger,
workspace_service=SimpleNamespace(instance_uuid='instance-a'),
instance_config=SimpleNamespace(
data={
'box': {
@@ -239,3 +253,130 @@ async def test_box_disconnect_notifies_once_and_clears_handler(
disconnect.assert_awaited_once_with(connector)
assert connector._handler is None
await connector.aclose()
def test_box_runtime_connector_builds_host_control_headers(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv(BOX_CONTROL_TOKEN_ENV, _CONTROL_TOKEN)
connector = BoxRuntimeConnector(make_app(Mock(), runtime_endpoint='http://box-runtime:5410'))
headers = connector.get_control_headers()
assert headers == {
BOX_CONTROL_TOKEN_HEADER: _CONTROL_TOKEN,
BOX_INSTANCE_HEADER: 'instance-a',
}
assert _CONTROL_TOKEN not in connector._resolve_rpc_ws_url()
def test_box_runtime_connector_builds_placement_scoped_relay_headers(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setenv(BOX_CONTROL_TOKEN_ENV, _CONTROL_TOKEN)
connector = BoxRuntimeConnector(make_app(Mock(), runtime_endpoint='http://box-runtime:5410'))
headers = connector.get_relay_headers(
ActionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=7,
)
)
assert headers == {
BOX_CONTROL_TOKEN_HEADER: _CONTROL_TOKEN,
BOX_INSTANCE_HEADER: 'instance-a',
BOX_WORKSPACE_HEADER: 'workspace-a',
BOX_PLACEMENT_GENERATION_HEADER: '7',
}
def test_box_runtime_connector_rejects_relay_context_from_other_instance(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setenv(BOX_CONTROL_TOKEN_ENV, _CONTROL_TOKEN)
connector = BoxRuntimeConnector(make_app(Mock(), runtime_endpoint='http://box-runtime:5410'))
with pytest.raises(BoxRuntimeUnavailableError, match='another LangBot instance'):
connector.get_relay_headers(
ActionContext(
instance_uuid='instance-b',
workspace_uuid='workspace-a',
placement_generation=1,
)
)
def test_external_box_runtime_fails_closed_without_control_token(monkeypatch: pytest.MonkeyPatch):
monkeypatch.delenv(BOX_CONTROL_TOKEN_ENV, raising=False)
connector = BoxRuntimeConnector(make_app(Mock(), runtime_endpoint='http://box-runtime:5410'))
with pytest.raises(BoxRuntimeUnavailableError, match=BOX_CONTROL_TOKEN_ENV):
connector.get_control_headers()
async def test_local_stdio_injects_generated_token_and_trusted_instance(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.delenv(BOX_CONTROL_TOKEN_ENV, raising=False)
captured = {}
class FakeStdioClientController:
def __init__(self, **kwargs):
captured.update(kwargs)
self.process = Mock()
async def run(self, callback):
await callback(None)
monkeypatch.setattr(
'langbot_plugin.runtime.io.controllers.stdio.client.StdioClientController',
FakeStdioClientController,
)
connector = BoxRuntimeConnector(make_app(Mock()))
def fake_callback(_transport_name, connected, _connect_error, _generation):
async def callback(_connection):
connected.set()
return callback
monkeypatch.setattr(connector, '_make_connection_callback', fake_callback)
await connector._start_local_stdio()
assert len(captured['env'][BOX_CONTROL_TOKEN_ENV]) >= 32
assert captured['env'][BOX_TRUSTED_INSTANCE_ENV] == 'instance-a'
async def test_websocket_controller_receives_control_headers(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv(BOX_CONTROL_TOKEN_ENV, _CONTROL_TOKEN)
captured = {}
class FakeWebSocketClientController:
def __init__(self, **kwargs):
captured.update(kwargs)
async def run(self, callback):
await callback(None)
monkeypatch.setattr(
'langbot_plugin.runtime.io.controllers.ws.client.WebSocketClientController',
FakeWebSocketClientController,
)
connector = BoxRuntimeConnector(make_app(Mock(), runtime_endpoint='http://box-runtime:5410'))
def fake_callback(_transport_name, connected, _connect_error, _generation):
async def callback(_connection):
connected.set()
return callback
monkeypatch.setattr(connector, '_make_connection_callback', fake_callback)
await connector._connect_ws('ws://box-runtime:5410/rpc/ws', 'WebSocket')
assert captured['additional_headers'] == {
BOX_CONTROL_TOKEN_HEADER: _CONTROL_TOKEN,
BOX_INSTANCE_HEADER: 'instance-a',
}
assert _CONTROL_TOKEN not in captured['ws_url']
tests/unit_tests/box/test_box_deployment_config.py
+65
View File
@@ -0,0 +1,65 @@
from pathlib import Path
_REPO_ROOT = Path(__file__).resolve().parents[3]
def test_compose_injects_the_same_box_control_token_into_host_and_runtime():
compose = (_REPO_ROOT / 'docker' / 'docker-compose.yaml').read_text(encoding='utf-8')
box_service = compose.split(' langbot_box:', 1)[1].split(' langbot:', 1)[0]
langbot_service = compose.split(' langbot:', 1)[1]
token_env = 'LANGBOT_BOX_CONTROL_TOKEN=${LANGBOT_BOX_CONTROL_TOKEN:-}'
assert token_env in box_service
assert token_env in langbot_service
def test_kubernetes_uses_one_secret_for_box_runtime_and_langbot():
manifest = (_REPO_ROOT / 'docker' / 'kubernetes.yaml').read_text(encoding='utf-8')
box_deployment = manifest.split('name: langbot-box', 1)[1].split('# Service for LangBot Box runtime', 1)[0]
langbot_deployment = manifest.split('# Deployment for LangBot\n', 1)[1]
secret_reference = '\n'.join(
[
'- name: LANGBOT_BOX_CONTROL_TOKEN',
' valueFrom:',
' secretKeyRef:',
' name: langbot-box-control',
' key: token',
]
)
assert secret_reference in box_deployment
assert secret_reference in langbot_deployment
assert '--from-literal=token="$(openssl rand -hex 32)"' in manifest
def test_compose_injects_same_plugin_runtime_control_token_into_both_services():
compose = (_REPO_ROOT / 'docker' / 'docker-compose.yaml').read_text(encoding='utf-8')
runtime_service = compose.split(' langbot_plugin_runtime:', 1)[1].split(' langbot_box:', 1)[0]
langbot_service = compose.split(' langbot:', 1)[1]
token_env = 'LANGBOT_PLUGIN_RUNTIME_CONTROL_TOKEN=${LANGBOT_PLUGIN_RUNTIME_CONTROL_TOKEN:-}'
assert token_env in runtime_service
assert token_env in langbot_service
def test_kubernetes_uses_one_secret_for_plugin_runtime_and_langbot():
manifest = (_REPO_ROOT / 'docker' / 'kubernetes.yaml').read_text(encoding='utf-8')
runtime_deployment = manifest.split('# Deployment for LangBot Plugin Runtime', 1)[1].split(
'# Service for LangBot Plugin Runtime',
1,
)[0]
langbot_deployment = manifest.split('# Deployment for LangBot\n', 1)[1]
secret_reference = '\n'.join(
[
'- name: LANGBOT_PLUGIN_RUNTIME_CONTROL_TOKEN',
' valueFrom:',
' secretKeyRef:',
' name: langbot-plugin-runtime-control',
' key: token',
]
)
assert secret_reference in runtime_deployment
assert secret_reference in langbot_deployment
assert 'create secret generic langbot-plugin-runtime-control' in manifest
tests/unit_tests/box/test_box_service.py
+552 -76
View File
@@ -3,6 +3,7 @@ from __future__ import annotations
import asyncio
import datetime as dt
import os
import pathlib
import tempfile
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
@@ -15,6 +16,7 @@ from langbot_plugin.box.backend import BaseSandboxBackend
from langbot_plugin.box.client import BoxRuntimeClient, ActionRPCBoxClient
from langbot_plugin.box.errors import (
BoxBackendUnavailableError,
BoxError,
BoxSessionConflictError,
BoxSessionNotFoundError,
BoxValidationError,
@@ -30,9 +32,27 @@ from langbot_plugin.box.models import (
BoxSpec,
)
from langbot_plugin.box.runtime import BoxRuntime
from langbot_plugin.box.security import (
BOX_CONTROL_TOKEN_HEADER,
BOX_INSTANCE_HEADER,
BOX_PLACEMENT_GENERATION_HEADER,
BOX_WORKSPACE_HEADER,
)
from langbot_plugin.entities.io.context import ActionContext
from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.box.service import BoxService
_UTC = dt.timezone.utc
_CONTEXT = ExecutionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=1,
)
_ACTION_CONTEXT = ActionContext(
instance_uuid=_CONTEXT.instance_uuid,
workspace_uuid=_CONTEXT.workspace_uuid,
placement_generation=_CONTEXT.placement_generation,
)
class _InProcessBoxRuntimeClient(BoxRuntimeClient):
@@ -44,42 +64,63 @@ class _InProcessBoxRuntimeClient(BoxRuntimeClient):
async def initialize(self):
await self._runtime.initialize()
async def execute(self, spec):
async def execute(self, spec, *, action_context=None):
return await self._runtime.execute(spec)
async def shutdown(self):
await self._runtime.shutdown()
async def get_status(self):
async def get_status(self, *, action_context=None):
return await self._runtime.get_status()
async def get_sessions(self):
async def get_sessions(self, *, action_context=None):
return self._runtime.get_sessions()
async def get_backend_info(self):
return await self._runtime.get_backend_info()
async def delete_session(self, session_id):
async def delete_session(self, session_id, *, action_context=None):
await self._runtime.delete_session(session_id)
async def create_session(self, spec):
async def create_session(self, spec, *, action_context=None):
return await self._runtime.create_session(spec)
async def start_managed_process(self, session_id: str, spec: BoxManagedProcessSpec):
async def start_managed_process(
self,
session_id: str,
spec: BoxManagedProcessSpec,
*,
action_context=None,
):
return await self._runtime.start_managed_process(session_id, spec)
async def get_managed_process(self, session_id: str, process_id: str = 'default'):
async def get_managed_process(
self,
session_id: str,
process_id: str = 'default',
*,
action_context=None,
):
return self._runtime.get_managed_process(session_id, process_id)
async def stop_managed_process(self, session_id: str, process_id: str = 'default'):
async def stop_managed_process(
self,
session_id: str,
process_id: str = 'default',
*,
action_context=None,
):
await self._runtime.stop_managed_process(session_id, process_id)
async def get_session(self, session_id: str):
async def get_session(self, session_id: str, *, action_context=None):
return self._runtime.get_session(session_id)
async def init(self, config: dict) -> None:
self._runtime.init(config)
async def verify_shared_workspace(self, marker_name: str) -> dict:
return self._runtime.verify_shared_workspace(marker_name)
class FakeBackend(BaseSandboxBackend):
def __init__(self, logger: Mock, available: bool = True):
@@ -134,6 +175,12 @@ class FakeBackend(BaseSandboxBackend):
def make_query(query_id: int = 42) -> pipeline_query.Query:
return pipeline_query.Query.model_construct(
query_id=query_id,
query_uuid=f'query-{query_id}',
instance_uuid=_CONTEXT.instance_uuid,
workspace_uuid=_CONTEXT.workspace_uuid,
placement_generation=_CONTEXT.placement_generation,
bot_uuid='bot-a',
pipeline_uuid='pipeline-a',
launcher_type='person',
launcher_id='test_user',
sender_id='test_user',
@@ -170,8 +217,19 @@ def make_app(
if workspace_quota_mb is not None:
box_config['local']['workspace_quota_mb'] = workspace_quota_mb
workspace_service = SimpleNamespace(
instance_uuid=_CONTEXT.instance_uuid,
get_execution_binding=AsyncMock(
return_value=SimpleNamespace(
instance_uuid=_CONTEXT.instance_uuid,
workspace_uuid=_CONTEXT.workspace_uuid,
placement_generation=_CONTEXT.placement_generation,
)
),
)
return SimpleNamespace(
logger=logger,
workspace_service=workspace_service,
instance_config=SimpleNamespace(
data={
'box': box_config,
@@ -196,6 +254,46 @@ async def test_box_service_without_explicit_client_initializes_internal_connecto
connector.initialize.assert_awaited_once()
@pytest.mark.asyncio
async def test_cloud_initialize_validation_failure_closes_connector_and_cancels_reconnect(
monkeypatch: pytest.MonkeyPatch,
):
logger = Mock()
app = make_app(logger)
app.deployment = SimpleNamespace(multi_workspace_enabled=True)
app.instance_config.data['box'].update(
{
'backend': 'nsjail',
'admission': {'required': True, 'workspace_quota_mb': 32},
}
)
connector = Mock()
connector.client = Mock(spec=BoxRuntimeClient)
connector.initialize = AsyncMock()
connector.aclose = AsyncMock()
connector.runtime_disconnect_callback = Mock()
monkeypatch.setattr('langbot.pkg.box.service.BoxRuntimeConnector', Mock(return_value=connector))
service = BoxService(app)
service._ensure_default_workspace = Mock()
readiness_error = BoxValidationError('Cloud Box nsjail isolation readiness failed')
service._verify_cloud_runtime = AsyncMock(side_effect=readiness_error)
reconnect_task = asyncio.create_task(asyncio.Event().wait())
service._reconnect_task = reconnect_task
service._reconnecting = True
with pytest.raises(BoxValidationError) as exc_info:
await service.initialize()
assert exc_info.value is readiness_error
assert service.available is False
assert service._closing is True
assert service._reconnecting is False
assert service._reconnect_task is None
assert reconnect_task.cancelled()
assert connector.runtime_disconnect_callback is None
connector.aclose.assert_awaited_once()
class TestSharesFilesystemWithBox:
"""``shares_filesystem_with_box`` must reflect the real LangBot<->Box
filesystem topology, which is derived from the connector transport:
@@ -268,6 +366,43 @@ def test_separated_box_runtime_does_not_create_default_workspace_in_langbot(tmp_
assert not (host_root / 'default').exists()
@pytest.mark.asyncio
async def test_cloud_initialize_fails_when_core_and_runtime_volumes_are_separated(tmp_path):
logger = Mock()
core_root = tmp_path / 'core-box'
runtime_root = tmp_path / 'runtime-box'
(core_root / 'default').mkdir(parents=True)
runtime = BoxRuntime(logger=logger, backends=[FakeBackend(logger)], session_ttl_sec=300)
runtime.init(
{
'local': {
'host_root': str(runtime_root),
'default_workspace': 'default',
'allowed_mount_roots': [str(runtime_root)],
}
}
)
app = make_app(logger, host_root=str(core_root))
app.deployment = SimpleNamespace(multi_workspace_enabled=True)
app.instance_config.data['box'].update(
{
'backend': 'nsjail',
'admission': {'required': True, 'workspace_quota_mb': 32},
}
)
service = BoxService(
app,
client=_InProcessBoxRuntimeClient(logger, runtime),
)
with pytest.raises(BoxValidationError, match='shared durable Workspace volume'):
await service.initialize()
assert service.available is False
assert list((core_root / 'default').glob('.langbot-box-volume-probe-*')) == []
await runtime.shutdown()
def test_separated_box_runtime_allows_box_owned_missing_host_path(tmp_path):
logger = Mock()
runtime = BoxRuntime(logger=logger, backends=[FakeBackend(logger)], session_ttl_sec=300)
@@ -289,12 +424,77 @@ async def test_box_service_get_sessions_delegates_to_client():
service = BoxService(make_app(Mock()), client=client)
service._available = True
sessions = await service.get_sessions()
sessions = await service.get_sessions(_CONTEXT)
assert sessions == [{'session_id': 'test-session'}]
client.get_sessions.assert_awaited_once()
@pytest.mark.asyncio
async def test_box_service_relay_connection_is_binding_checked_and_scoped():
app = make_app(Mock())
client = Mock()
client.get_managed_process_websocket_url = Mock(
return_value='ws://box/v1/sessions/physical/managed-process/server-a/ws'
)
connector = Mock()
connector.ws_relay_base_url = 'http://box:5410'
connector.get_relay_headers = Mock(
return_value={
BOX_CONTROL_TOKEN_HEADER: 'secret',
BOX_INSTANCE_HEADER: _CONTEXT.instance_uuid,
BOX_WORKSPACE_HEADER: _CONTEXT.workspace_uuid,
BOX_PLACEMENT_GENERATION_HEADER: '1',
}
)
service = BoxService(app, client=client)
service._runtime_connector = connector
url, headers = await service.get_managed_process_websocket_connection(
_CONTEXT,
'mcp-shared',
'server-a',
)
assert url == 'ws://box/v1/sessions/physical/managed-process/server-a/ws'
assert headers[BOX_WORKSPACE_HEADER] == _CONTEXT.workspace_uuid
assert headers[BOX_PLACEMENT_GENERATION_HEADER] == '1'
app.workspace_service.get_execution_binding.assert_awaited_once_with(
_CONTEXT.workspace_uuid,
expected_generation=_CONTEXT.placement_generation,
)
action_context = connector.get_relay_headers.call_args.args[0]
assert action_context == _ACTION_CONTEXT
client.get_managed_process_websocket_url.assert_called_once_with(
'mcp-shared',
'http://box:5410',
'server-a',
action_context=_ACTION_CONTEXT,
)
@pytest.mark.asyncio
async def test_box_service_relay_connection_rejects_stale_binding():
app = make_app(Mock())
app.workspace_service.get_execution_binding.return_value = SimpleNamespace(
instance_uuid=_CONTEXT.instance_uuid,
workspace_uuid=_CONTEXT.workspace_uuid,
placement_generation=2,
)
client = Mock()
service = BoxService(app, client=client)
service._runtime_connector = Mock()
with pytest.raises(BoxValidationError, match='stale Workspace placement'):
await service.get_managed_process_websocket_connection(
_CONTEXT,
'mcp-shared',
'server-a',
)
service._runtime_connector.get_relay_headers.assert_not_called()
def test_box_service_dispose_delegates_to_internal_connector(monkeypatch: pytest.MonkeyPatch):
connector = Mock()
connector.client = Mock()
@@ -369,6 +569,28 @@ async def test_box_service_reconnect_restores_workspace_and_runs_cleanup(
assert service.available is True
@pytest.mark.asyncio
async def test_cloud_box_service_reconnect_does_not_reload_unscoped_skills(
monkeypatch: pytest.MonkeyPatch,
):
app = make_app(Mock())
app.skill_mgr = SimpleNamespace(reload_skills=AsyncMock())
service = BoxService(app, client=Mock(spec=BoxRuntimeClient))
service._cloud_managed = True
connector = Mock()
connector.reconnect = AsyncMock()
service._ensure_default_workspace = Mock()
service._verify_cloud_runtime = AsyncMock()
monkeypatch.setattr('langbot.pkg.box.service.asyncio.sleep', AsyncMock())
await service._reconnect_loop(connector)
connector.reconnect.assert_awaited_once()
service._verify_cloud_runtime.assert_awaited_once()
app.skill_mgr.reload_skills.assert_not_awaited()
assert service.available is True
@pytest.mark.asyncio
async def test_box_runtime_reuses_request_session():
logger = Mock()
@@ -409,7 +631,14 @@ async def test_box_service_session_id_uses_query_attributes_without_variables():
service = BoxService(make_app(logger), client=_InProcessBoxRuntimeClient(logger, runtime))
await service.initialize()
query = pipeline_query.Query.model_construct(query_id=7, launcher_type='group', launcher_id='room-1')
query = pipeline_query.Query.model_construct(
query_id=7,
instance_uuid=_CONTEXT.instance_uuid,
workspace_uuid=_CONTEXT.workspace_uuid,
placement_generation=_CONTEXT.placement_generation,
launcher_type='group',
launcher_id='room-1',
)
result = await service.execute_tool({'command': 'pwd'}, query)
assert result['session_id'] == 'group_room-1'
@@ -425,7 +654,12 @@ async def test_box_service_session_id_falls_back_to_query_id_for_synthetic_queri
service = BoxService(make_app(logger), client=_InProcessBoxRuntimeClient(logger, runtime))
await service.initialize()
query = pipeline_query.Query.model_construct(query_id=7)
query = pipeline_query.Query.model_construct(
query_id=7,
instance_uuid=_CONTEXT.instance_uuid,
workspace_uuid=_CONTEXT.workspace_uuid,
placement_generation=_CONTEXT.placement_generation,
)
result = await service.execute_tool({'command': 'pwd'}, query)
assert result['session_id'] == 'query_7'
@@ -447,8 +681,22 @@ async def test_box_service_forced_global_scope_overrides_pipeline_template():
await service.initialize()
# Two distinct callers that would otherwise get separate sandboxes.
q1 = pipeline_query.Query.model_construct(query_id=1, launcher_type='group', launcher_id='room-1')
q2 = pipeline_query.Query.model_construct(query_id=2, launcher_type='person', launcher_id='alice')
q1 = pipeline_query.Query.model_construct(
query_id=1,
instance_uuid=_CONTEXT.instance_uuid,
workspace_uuid=_CONTEXT.workspace_uuid,
placement_generation=_CONTEXT.placement_generation,
launcher_type='group',
launcher_id='room-1',
)
q2 = pipeline_query.Query.model_construct(
query_id=2,
instance_uuid=_CONTEXT.instance_uuid,
workspace_uuid=_CONTEXT.workspace_uuid,
placement_generation=_CONTEXT.placement_generation,
launcher_type='person',
launcher_id='alice',
)
r1 = await service.execute_tool({'command': 'pwd'}, q1)
r2 = await service.execute_tool({'command': 'pwd'}, q2)
@@ -551,7 +799,7 @@ async def test_box_service_uses_default_workspace_when_host_path_omitted(tmp_pat
assert result['ok'] is True
assert backend.start_calls == ['person_test_user']
assert backend.exec_calls == [('person_test_user', 'pwd')]
assert backend.start_specs[0].host_path == os.path.realpath(host_dir)
assert backend.start_specs[0].host_path == service._tenant_workspace(_CONTEXT)
@pytest.mark.asyncio
@@ -994,11 +1242,16 @@ async def test_box_service_rejects_execution_when_workspace_already_exceeds_quot
runtime = BoxRuntime(logger=logger, backends=[backend], session_ttl_sec=300)
host_dir = tmp_path / 'quota-workspace'
host_dir.mkdir()
(host_dir / 'already-too-large.bin').write_bytes(b'x' * (2 * 1024 * 1024))
app = make_app(logger, [str(tmp_path)], workspace_quota_mb=1)
app.instance_config.data['box']['local']['default_workspace'] = str(host_dir)
service = BoxService(app, client=_InProcessBoxRuntimeClient(logger, runtime))
tenant_host_dir = service._tenant_workspace(_CONTEXT)
assert tenant_host_dir is not None
os.makedirs(tenant_host_dir, exist_ok=True)
with open(os.path.join(tenant_host_dir, 'already-too-large.bin'), 'wb') as handle:
handle.write(b'x' * (2 * 1024 * 1024))
await service.initialize()
with pytest.raises(BoxValidationError, match='workspace quota exceeded before execution'):
@@ -1027,6 +1280,45 @@ async def test_box_service_rejects_and_cleans_up_when_execution_exceeds_workspac
assert backend.stop_calls == ['person_test_user']
@pytest.mark.asyncio
async def test_box_service_rejects_workspace_inode_bomb_before_execution(tmp_path):
logger = Mock()
backend = FakeBackend(logger)
runtime = BoxRuntime(logger=logger, backends=[backend], session_ttl_sec=300)
host_dir = tmp_path / 'quota-workspace-entries'
host_dir.mkdir()
app = make_app(logger, [str(tmp_path)], workspace_quota_mb=1)
app.instance_config.data['box']['local']['default_workspace'] = str(host_dir)
app.instance_config.data['box']['limits'] = {'max_workspace_entries': 2}
service = BoxService(app, client=_InProcessBoxRuntimeClient(logger, runtime))
tenant_host_dir = service._tenant_workspace(_CONTEXT)
assert tenant_host_dir is not None
os.makedirs(tenant_host_dir, exist_ok=True)
for index in range(3):
pathlib.Path(tenant_host_dir, f'tiny-{index}').write_bytes(b'x')
await service.initialize()
with pytest.raises(BoxValidationError, match='workspace entry limit exceeded before execution'):
await service.execute_tool({'command': 'echo hi'}, make_query(46))
assert backend.start_calls == []
def test_box_service_workspace_entry_limit_is_hard_clamped():
app = make_app(Mock())
app.instance_config.data['box']['limits'] = {'max_workspace_entries': 10_000_000}
service = BoxService(app, client=Mock(spec=BoxRuntimeClient))
assert service._max_workspace_entries() == 1_000_000
app.instance_config.data['box']['limits']['max_workspace_entries'] = 0
assert service._max_workspace_entries() == 1
app.instance_config.data['box']['limits']['max_workspace_entries'] = 'invalid'
assert service._max_workspace_entries() == 100_000
@pytest.mark.asyncio
async def test_profile_offline_readonly_locks_read_only_rootfs():
"""offline_readonly locks read_only_rootfs so it cannot be overridden."""
@@ -1137,7 +1429,7 @@ async def test_service_records_errors_on_failure():
with pytest.raises(Exception):
await service.execute_tool({'command': 'echo hello'}, make_query(50))
errors = service.get_recent_errors()
errors = service.get_recent_errors(_CONTEXT)
assert len(errors) == 1
assert errors[0]['type'] == 'BoxBackendUnavailableError'
assert errors[0]['query_id'] == '50'
@@ -1156,7 +1448,7 @@ async def test_service_error_ring_buffer_capped():
with pytest.raises(Exception):
await service.execute_tool({'command': 'fail'}, make_query(100 + i))
errors = service.get_recent_errors()
errors = service.get_recent_errors(_CONTEXT)
assert len(errors) == 50
# Oldest should have been evicted, newest kept
assert errors[0]['query_id'] == '110'
@@ -1171,7 +1463,7 @@ async def test_service_get_status_aggregates_runtime_and_profile():
service = BoxService(make_app(logger), client=_InProcessBoxRuntimeClient(logger, runtime))
await service.initialize()
status = await service.get_status()
status = await service.get_status(_CONTEXT)
assert status['profile'] == 'default'
assert status['backend']['name'] == 'fake'
assert status['backend']['available'] is True
@@ -1215,7 +1507,12 @@ async def _make_rpc_pair(runtime: BoxRuntime):
client_conn, server_conn = _make_queue_connection_pair()
server_handler = BoxServerHandler(server_conn, runtime)
server_handler = BoxServerHandler(
server_conn,
runtime,
host_control_authenticated=True,
trusted_instance_uuid=_CONTEXT.instance_uuid,
)
server_task = asyncio.create_task(server_handler.run())
client_handler = Handler.__new__(Handler)
@@ -1239,7 +1536,7 @@ async def test_rpc_client_execute():
client, server_task, client_task = await _make_rpc_pair(runtime)
try:
spec = BoxSpec.model_validate({'cmd': 'echo remote', 'session_id': 'r-1'})
result = await client.execute(spec)
result = await client.execute(spec, action_context=_ACTION_CONTEXT)
assert result.session_id == 'r-1'
assert result.status == BoxExecutionStatus.COMPLETED
@@ -1261,9 +1558,9 @@ async def test_rpc_client_get_sessions():
client, server_task, client_task = await _make_rpc_pair(runtime)
try:
spec = BoxSpec.model_validate({'cmd': 'echo hi', 'session_id': 'r-2'})
await client.execute(spec)
await client.execute(spec, action_context=_ACTION_CONTEXT)
sessions = await client.get_sessions()
sessions = await client.get_sessions(action_context=_ACTION_CONTEXT)
assert len(sessions) == 1
assert sessions[0]['session_id'] == 'r-2'
finally:
@@ -1272,6 +1569,34 @@ async def test_rpc_client_get_sessions():
await runtime.shutdown()
@pytest.mark.asyncio
async def test_rpc_generation_advance_retires_old_session_and_rejects_old_context():
logger = Mock()
backend = FakeBackend(logger)
runtime = BoxRuntime(logger=logger, backends=[backend], session_ttl_sec=300)
await runtime.initialize()
second_context = _ACTION_CONTEXT.model_copy(update={'placement_generation': 2})
client, server_task, client_task = await _make_rpc_pair(runtime)
try:
spec = BoxSpec.model_validate({'cmd': 'echo generation', 'session_id': 'shared'})
await client.execute(spec, action_context=_ACTION_CONTEXT)
await client.execute(spec, action_context=second_context)
assert len(backend.start_calls) == 2
assert backend.start_calls[0] != backend.start_calls[1]
assert backend.stop_calls == [backend.start_calls[0]]
assert [session['session_id'] for session in await client.get_sessions(action_context=second_context)] == [
'shared'
]
with pytest.raises(BoxError, match='Stale Box placement generation'):
await client.execute(spec, action_context=_ACTION_CONTEXT)
finally:
server_task.cancel()
client_task.cancel()
await runtime.shutdown()
@pytest.mark.asyncio
async def test_rpc_client_get_status():
logger = Mock()
@@ -1281,7 +1606,7 @@ async def test_rpc_client_get_status():
client, server_task, client_task = await _make_rpc_pair(runtime)
try:
status = await client.get_status()
status = await client.get_status(action_context=_ACTION_CONTEXT)
assert 'backend' in status
assert 'active_sessions' in status
@@ -1323,11 +1648,11 @@ async def test_rpc_client_delete_session():
client, server_task, client_task = await _make_rpc_pair(runtime)
try:
spec = BoxSpec.model_validate({'cmd': 'echo hi', 'session_id': 'r-del-1'})
await client.execute(spec)
await client.execute(spec, action_context=_ACTION_CONTEXT)
await client.delete_session('r-del-1')
await client.delete_session('r-del-1', action_context=_ACTION_CONTEXT)
sessions = await client.get_sessions()
sessions = await client.get_sessions(action_context=_ACTION_CONTEXT)
assert len(sessions) == 0
finally:
server_task.cancel()
@@ -1345,7 +1670,7 @@ async def test_rpc_client_delete_session_raises_not_found():
client, server_task, client_task = await _make_rpc_pair(runtime)
try:
with pytest.raises(BoxSessionNotFoundError):
await client.delete_session('nonexistent')
await client.delete_session('nonexistent', action_context=_ACTION_CONTEXT)
finally:
server_task.cancel()
client_task.cancel()
@@ -1362,11 +1687,11 @@ async def test_rpc_client_create_session():
client, server_task, client_task = await _make_rpc_pair(runtime)
try:
spec = BoxSpec.model_validate({'cmd': 'placeholder', 'session_id': 'r-create-1'})
info = await client.create_session(spec)
info = await client.create_session(spec, action_context=_ACTION_CONTEXT)
assert info['session_id'] == 'r-create-1'
assert info['backend_name'] == 'fake'
sessions = await client.get_sessions()
sessions = await client.get_sessions(action_context=_ACTION_CONTEXT)
assert len(sessions) == 1
finally:
server_task.cancel()
@@ -1384,11 +1709,11 @@ async def test_rpc_client_exec_raises_conflict_error():
client, server_task, client_task = await _make_rpc_pair(runtime)
try:
spec1 = BoxSpec.model_validate({'cmd': 'echo first', 'session_id': 'r-conflict-1', 'network': 'off'})
await client.execute(spec1)
await client.execute(spec1, action_context=_ACTION_CONTEXT)
spec2 = BoxSpec.model_validate({'cmd': 'echo second', 'session_id': 'r-conflict-1', 'network': 'on'})
with pytest.raises(BoxSessionConflictError):
await client.execute(spec2)
await client.execute(spec2, action_context=_ACTION_CONTEXT)
finally:
server_task.cancel()
client_task.cancel()
@@ -1487,7 +1812,7 @@ class TestBoxDisabledByConfig:
service = BoxService(make_app(logger, enabled=False), client=Mock(spec=BoxRuntimeClient))
await service.initialize()
status = await service.get_status()
status = await service.get_status(_CONTEXT)
assert status['available'] is False
assert status['enabled'] is False
@@ -1502,7 +1827,7 @@ class TestBoxDisabledByConfig:
await service.initialize()
status = await service.get_status()
status = await service.get_status(_CONTEXT)
assert status['available'] is False
assert status['enabled'] is True
assert 'docker daemon' in status['connector_error']
@@ -1526,7 +1851,7 @@ class TestBoxDisabledByConfig:
service = BoxService(make_app(logger, enabled=True), client=client)
await service.initialize()
status = await service.get_status()
status = await service.get_status(_CONTEXT)
assert status['available'] is False
assert status['enabled'] is True
# The detailed backend object is preserved for the dialog
@@ -1547,7 +1872,7 @@ class TestBoxDisabledByConfig:
service = BoxService(make_app(logger, enabled=True), client=client)
await service.initialize()
status = await service.get_status()
status = await service.get_status(_CONTEXT)
assert status['available'] is True
assert status['backend'] == {'name': 'docker', 'available': True}
# No spurious connector_error overlay when everything is healthy
@@ -1565,6 +1890,57 @@ class TestBoxDisabledByConfig:
assert service._reconnecting is False
@pytest.mark.asyncio
async def test_disconnect_callback_does_not_schedule_on_closing_event_loop(
monkeypatch: pytest.MonkeyPatch,
):
service = BoxService(make_app(Mock()), client=Mock(spec=BoxRuntimeClient))
closed_loop = Mock()
closed_loop.is_closed.return_value = True
monkeypatch.setattr('langbot.pkg.box.service.asyncio.get_running_loop', Mock(return_value=closed_loop))
await service._on_runtime_disconnect(connector=Mock())
closed_loop.create_task.assert_not_called()
assert service._reconnect_task is None
assert service._reconnecting is False
@pytest.mark.asyncio
async def test_disconnect_callback_closes_reconnect_coroutine_when_task_creation_races_with_loop_close(
monkeypatch: pytest.MonkeyPatch,
):
service = BoxService(make_app(Mock()), client=Mock(spec=BoxRuntimeClient))
loop = Mock()
loop.is_closed.return_value = False
loop.create_task.side_effect = RuntimeError('event loop is closed')
monkeypatch.setattr('langbot.pkg.box.service.asyncio.get_running_loop', Mock(return_value=loop))
async def reconnect():
await asyncio.Event().wait()
reconnect_coroutine = reconnect()
service._reconnect_loop = Mock(return_value=reconnect_coroutine)
await service._on_runtime_disconnect(connector=Mock())
loop.create_task.assert_called_once_with(reconnect_coroutine)
assert reconnect_coroutine.cr_frame is None
assert service._reconnect_task is None
assert service._reconnecting is False
def test_disconnect_callback_does_not_schedule_without_running_event_loop():
service = BoxService(make_app(Mock()), client=Mock(spec=BoxRuntimeClient))
callback = service._on_runtime_disconnect(connector=Mock())
with pytest.raises(StopIteration):
callback.send(None)
assert service._reconnect_task is None
assert service._reconnecting is False
class TestBuildSkillExtraMounts:
"""Robustness of skill mount construction against a stale skill cache.
@@ -1577,7 +1953,7 @@ class TestBuildSkillExtraMounts:
def _make_service(self, logger, skills, *, shares_filesystem=True):
app = make_app(logger)
app.skill_mgr = SimpleNamespace(skills=skills)
app.skill_mgr = SimpleNamespace(skills=skills, get_skills=Mock(return_value=skills))
client = Mock(spec=BoxRuntimeClient)
service = BoxService(app, client=client)
# Tests construct BoxService with an injected client (no connector), so
@@ -1714,6 +2090,17 @@ class TestAttachmentHelpers:
component = SimpleNamespace(base64=None, url=None, path=None)
assert await BoxService._component_to_bytes(component) is None
@pytest.mark.asyncio
async def test_component_to_bytes_rejects_oversized_base64(self, monkeypatch):
monkeypatch.setattr(BoxService, '_ATTACHMENT_MAX_BYTES', 4)
component = SimpleNamespace(
base64='data:application/octet-stream;base64,' + ('A' * 12),
url=None,
path=None,
)
assert await BoxService._component_to_bytes(component) is None
class TestInboundOutboundRoundTrip:
def _service(self) -> BoxService:
@@ -1745,7 +2132,7 @@ class TestInboundOutboundRoundTrip:
assert '/workspace/inbox/' in parameters['command']
return {
'ok': True,
'stdout': '["/workspace/inbox/42/image_1.png"]',
'stdout': '["/workspace/inbox/query-42/image_1.png"]',
'stderr': '',
}
@@ -1755,7 +2142,7 @@ class TestInboundOutboundRoundTrip:
assert len(descriptors) == 1
d = descriptors[0]
assert d['type'] == 'Image'
assert d['path'] == '/workspace/inbox/42/image_1.png'
assert d['path'] == '/workspace/inbox/query-42/image_1.png'
assert d['size'] == len(img_bytes)
@pytest.mark.asyncio
@@ -1778,7 +2165,7 @@ class TestInboundOutboundRoundTrip:
async def fake_execute_tool(parameters, q):
calls.append(parameters['command'])
if 'os.walk' in parameters['command']:
if 'os.scandir' in parameters['command']:
return {
'ok': True,
'stdout': '[{"name": "out.png", "b64": "QUJD"}]',
@@ -1808,7 +2195,7 @@ class TestInboundOutboundRoundTrip:
async def fake_execute_tool(parameters, q):
calls.append(parameters['command'])
if 'os.walk' in parameters['command']:
if 'os.scandir' in parameters['command']:
return {'ok': True, 'stdout': '[]', 'stderr': ''}
return {'ok': True, 'stdout': '', 'stderr': ''}
@@ -1835,14 +2222,17 @@ class TestAttachmentHostPath:
"""
def _service_with_workspace(self, tmp_path):
ws = str(tmp_path / 'box' / 'default')
os.makedirs(ws, exist_ok=True)
default_workspace = str(tmp_path / 'box' / 'default')
os.makedirs(default_workspace, exist_ok=True)
app = make_app(Mock(), allowed_mount_roots=[str(tmp_path)], host_root=str(tmp_path / 'box'))
service = BoxService(app, client=Mock(spec=BoxRuntimeClient))
service._available = True
# Force the default_workspace to our tmp dir so _host_query_dir resolves.
service.default_workspace = ws
return service, ws
service.default_workspace = default_workspace
tenant_workspace = service._tenant_workspace(_CONTEXT)
assert tenant_workspace is not None
os.makedirs(tenant_workspace, exist_ok=True)
return service, tenant_workspace
@pytest.mark.asyncio
async def test_inbound_writes_to_host_no_exec(self, tmp_path):
@@ -1865,9 +2255,9 @@ class TestAttachmentHostPath:
assert d['type'] == 'Image'
assert d['size'] == len(big)
# File actually landed on the host workspace.
host_file = os.path.join(ws, 'inbox', str(query.query_id), d['name'])
host_file = os.path.join(ws, 'inbox', str(query.query_uuid), d['name'])
assert os.path.isfile(host_file)
assert open(host_file, 'rb').read() == big
assert pathlib.Path(host_file).read_bytes() == big
@pytest.mark.asyncio
async def test_inbound_host_clears_stale_query_dir(self, tmp_path):
@@ -1877,9 +2267,9 @@ class TestAttachmentHostPath:
service, ws = self._service_with_workspace(tmp_path)
# Seed a stale file under the same query_id (simulates webchat id reuse).
stale_dir = os.path.join(ws, 'inbox', '42')
stale_dir = os.path.join(ws, 'inbox', 'query-42')
os.makedirs(stale_dir, exist_ok=True)
open(os.path.join(stale_dir, 'image_1.png'), 'wb').write(b'STALE-OLD-IMAGE')
pathlib.Path(stale_dir, 'image_1.png').write_bytes(b'STALE-OLD-IMAGE')
new = b'\x89PNG\r\n\x1a\n NEW'
b64 = 'data:image/png;base64,' + base64.b64encode(new).decode()
@@ -1889,20 +2279,50 @@ class TestAttachmentHostPath:
descriptors = await service.materialize_inbound_attachments(query)
# The new write recreated the dir; the stale file is gone, new bytes present.
host_file = os.path.join(stale_dir, descriptors[0]['name'])
assert open(host_file, 'rb').read() == new
host_bytes = pathlib.Path(host_file).read_bytes()
assert host_bytes == new
# No leftover content from the stale image.
assert b'STALE-OLD-IMAGE' not in open(host_file, 'rb').read()
assert b'STALE-OLD-IMAGE' not in host_bytes
@pytest.mark.asyncio
async def test_inbound_host_replaces_query_symlink_without_touching_other_workspace(self, tmp_path):
import base64
import langbot_plugin.api.entities.builtin.platform.message as platform_message
service, ws = self._service_with_workspace(tmp_path)
other_workspace = tmp_path / 'other-workspace'
other_workspace.mkdir()
protected = other_workspace / 'protected.txt'
protected.write_bytes(b'workspace-b-secret')
inbox = os.path.join(ws, 'inbox')
os.makedirs(inbox, exist_ok=True)
os.symlink(other_workspace, os.path.join(inbox, 'query-42'))
query = make_query()
payload = b'workspace-a-input'
query.message_chain = platform_message.MessageChain(
[platform_message.File(name='input.bin', base64=base64.b64encode(payload).decode())]
)
service.execute_tool = AsyncMock(side_effect=AssertionError('exec must not be used on host path'))
descriptors = await service.materialize_inbound_attachments(query)
assert descriptors[0]['path'] == '/workspace/inbox/query-42/input.bin'
assert protected.read_bytes() == b'workspace-b-secret'
assert not os.path.islink(os.path.join(inbox, 'query-42'))
assert pathlib.Path(inbox, 'query-42', 'input.bin').read_bytes() == payload
@pytest.mark.asyncio
async def test_outbound_reads_host_and_clears(self, tmp_path):
service, ws = self._service_with_workspace(tmp_path)
query = make_query()
outbox = os.path.join(ws, 'outbox', str(query.query_id))
outbox = os.path.join(ws, 'outbox', str(query.query_uuid))
os.makedirs(outbox, exist_ok=True)
# A large file that would be truncated on the exec/stdout path:
big_png = b'\x89PNG\r\n\x1a\n' + b'y' * (400 * 1024)
open(os.path.join(outbox, 'result.png'), 'wb').write(big_png)
open(os.path.join(outbox, 'notes.txt'), 'wb').write(b'hello')
pathlib.Path(outbox, 'result.png').write_bytes(big_png)
pathlib.Path(outbox, 'notes.txt').write_bytes(b'hello')
service.execute_tool = AsyncMock(side_effect=AssertionError('exec must not be used on host path'))
attachments = await service.collect_outbound_attachments(query)
@@ -1917,20 +2337,76 @@ class TestAttachmentHostPath:
# Outbox cleared after collection.
assert os.listdir(outbox) == []
@pytest.mark.asyncio
async def test_outbound_host_never_follows_query_or_file_symlinks(self, tmp_path):
service, ws = self._service_with_workspace(tmp_path)
query = make_query()
other_workspace = tmp_path / 'other-workspace'
other_workspace.mkdir()
secret = other_workspace / 'secret.txt'
secret.write_bytes(b'workspace-b-secret')
outbox_root = os.path.join(ws, 'outbox')
os.makedirs(outbox_root, exist_ok=True)
# A hostile query-directory replacement is rejected rather than read.
query_dir = os.path.join(outbox_root, str(query.query_uuid))
os.symlink(other_workspace, query_dir)
with pytest.raises(BoxValidationError, match='symbolic link'):
await service.collect_outbound_attachments(query)
assert secret.read_bytes() == b'workspace-b-secret'
os.unlink(query_dir)
os.makedirs(query_dir)
os.symlink(secret, os.path.join(query_dir, 'leak.txt'))
service.execute_tool = AsyncMock(side_effect=AssertionError('exec must not be used on host path'))
assert await service.collect_outbound_attachments(query) == []
assert secret.read_bytes() == b'workspace-b-secret'
@pytest.mark.asyncio
async def test_outbound_host_fails_closed_on_inode_bomb(self, tmp_path):
service, ws = self._service_with_workspace(tmp_path)
query = make_query()
outbox = os.path.join(ws, 'outbox', str(query.query_uuid))
os.makedirs(outbox, exist_ok=True)
harmless_target = tmp_path / 'harmless-target'
harmless_target.write_bytes(b'x')
# Symlinks do not count toward the 20 returned files, so this proves
# traversal itself has a bounded entry budget.
for index in range(513):
os.symlink(harmless_target, os.path.join(outbox, f'entry-{index}'))
with pytest.raises(BoxValidationError, match='symbolic link'):
await service.collect_outbound_attachments(query)
assert harmless_target.read_bytes() == b'x'
def test_host_attachment_directories_use_query_uuid_not_process_local_id(self, tmp_path):
service, _ws = self._service_with_workspace(tmp_path)
first = make_query(query_id=7)
second = make_query(query_id=7)
object.__setattr__(first, 'query_uuid', 'replica-a-query')
object.__setattr__(second, 'query_uuid', 'replica-b-query')
first_path = service._host_query_dir(service.OUTBOX_SUBDIR, first)
second_path = service._host_query_dir(service.OUTBOX_SUBDIR, second)
assert first_path is not None and first_path.endswith('/outbox/replica-a-query')
assert second_path is not None and second_path.endswith('/outbox/replica-b-query')
assert first_path != second_path
@pytest.mark.asyncio
async def test_outbound_empty_clears_stale_host_dir(self, tmp_path):
# Reusing a query_id (counter resets on restart) must not re-send files
# a previous run left in the outbox: an empty collection still clears it.
service, ws = self._service_with_workspace(tmp_path)
query = make_query()
outbox = os.path.join(ws, 'outbox', str(query.query_id))
outbox = os.path.join(ws, 'outbox', str(query.query_uuid))
os.makedirs(outbox, exist_ok=True)
# Stale file from a prior turn; the agent produced nothing this turn —
# but _read_outbox_host would still pick it up, so collection must drop
# it and then wipe the dir. Simulate "nothing produced this turn" by
# treating any present file as stale and asserting it is not re-sent
# across a second, genuinely-empty collection.
open(os.path.join(outbox, 'stale.png'), 'wb').write(b'\x89PNG\r\n\x1a\n old')
pathlib.Path(outbox, 'stale.png').write_bytes(b'\x89PNG\r\n\x1a\n old')
service.execute_tool = AsyncMock(side_effect=AssertionError('exec must not be used on host path'))
# First collection drains + clears the dir.
@@ -1952,7 +2428,7 @@ class TestAttachmentHostPath:
for sub in ('inbox', 'outbox'):
d = os.path.join(ws, sub, '0')
os.makedirs(d, exist_ok=True)
open(os.path.join(d, 'leftover.bin'), 'wb').write(b'from a previous process')
pathlib.Path(d, 'leftover.bin').write_bytes(b'from a previous process')
service.execute_tool = AsyncMock(side_effect=AssertionError('exec must not be used for host-owned files'))
await service._purge_attachment_dirs()
@@ -1963,37 +2439,37 @@ class TestAttachmentHostPath:
assert os.path.isdir(ws)
@pytest.mark.asyncio
async def test_purge_attachment_dirs_falls_back_to_exec_for_root_owned(self, tmp_path, monkeypatch):
# When the host delete cannot remove a dir (root-owned container output),
# purge must fall back to deleting from inside the sandbox via exec.
async def test_purge_attachment_dirs_never_uses_unscoped_exec_for_root_owned(self, tmp_path, monkeypatch):
# Startup has no trusted Workspace context. If host deletion cannot
# remove root-owned output, cleanup must fail closed instead of issuing
# an unscoped Box exec that could cross a tenant boundary.
service, ws = self._service_with_workspace(tmp_path)
outbox = os.path.join(ws, 'outbox')
os.makedirs(os.path.join(outbox, '0'), exist_ok=True)
# Simulate a host delete that cannot remove the root-owned outbox.
import shutil as _shutil
from langbot.pkg.box import secure_fs
real_rmtree = _shutil.rmtree
real_purge = secure_fs.purge_subdirectory
def fake_rmtree(path, *a, **k):
if os.path.abspath(path) == os.path.abspath(outbox):
return # "permission denied" — silently leaves the dir
return real_rmtree(path, *a, **k)
def fake_purge(root, subdir):
if os.path.abspath(os.path.join(root, subdir)) == os.path.abspath(outbox):
raise PermissionError('root-owned')
return real_purge(root, subdir)
monkeypatch.setattr(_shutil, 'rmtree', fake_rmtree)
monkeypatch.setattr(secure_fs, 'purge_subdirectory', fake_purge)
executed = {}
spec_obj = object()
service.build_spec = Mock(return_value=spec_obj)
service.client.execute = AsyncMock(side_effect=lambda s: executed.setdefault('spec', s))
service.build_spec = Mock()
service.client.execute = AsyncMock()
await service._purge_attachment_dirs()
# build_spec was asked to rm the surviving outbox via exec.
cmd = service.build_spec.call_args.args[0]['cmd']
assert 'rm -rf' in cmd and '/workspace/outbox' in cmd
assert '/workspace/inbox' not in cmd # inbox was host-deletable
service.client.execute.assert_awaited_once_with(spec_obj)
assert os.path.isdir(outbox)
service.build_spec.assert_not_called()
service.client.execute.assert_not_awaited()
assert any(
'no trusted Workspace context' in str(call.args[0]) for call in service.ap.logger.warning.call_args_list
)
@pytest.mark.asyncio
async def test_purge_attachment_dirs_noop_without_workspace(self):
tests/unit_tests/box/test_workspace.py
+47 -2
View File
@@ -7,6 +7,7 @@ from unittest.mock import AsyncMock, Mock
import pytest
from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.box.workspace import (
BoxWorkspaceSession,
classify_python_workspace,
@@ -16,6 +17,13 @@ from langbot.pkg.box.workspace import (
)
_CONTEXT = ExecutionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=1,
)
def test_rewrite_mounted_path_translates_host_prefix():
result = rewrite_mounted_path('/tmp/demo/project/app.py', '/tmp/demo/project')
assert result == '/workspace/app.py'
@@ -57,6 +65,9 @@ def test_wrap_python_command_with_env_contains_bootstrap_and_command():
assert '_LB_SYSTEM_PYTHON="$(command -v python3 || command -v python || true)"' in command
assert '"$_LB_SYSTEM_PYTHON" -m venv "$_LB_VENV_DIR"' in command
assert 'kill -0 "$_LB_LOCK_OWNER"' in command
assert 'max_manifest_bytes = 10 * 1024 * 1024' in command
assert 'handle.read(1024 * 1024)' in command
assert 'digest.update(handle.read())' not in command
assert 'export VIRTUAL_ENV="$_LB_VENV_DIR"' in command
assert command.rstrip().endswith('python script.py')
@@ -66,6 +77,7 @@ async def test_workspace_session_execute_for_query_uses_session_payload():
box_service = SimpleNamespace(execute_spec_payload=AsyncMock(return_value={'ok': True}))
workspace = BoxWorkspaceSession(
box_service,
_CONTEXT,
'skill-person_123-demo',
host_path='/tmp/project',
host_path_mode='rw',
@@ -94,6 +106,7 @@ async def test_workspace_session_start_managed_process_rewrites_command_and_args
box_service = SimpleNamespace(start_managed_process=AsyncMock(return_value={'status': 'running'}))
workspace = BoxWorkspaceSession(
box_service,
_CONTEXT,
'mcp-u1',
host_path='/tmp/project',
host_path_mode='ro',
@@ -106,8 +119,10 @@ async def test_workspace_session_start_managed_process_rewrites_command_and_args
)
assert result == {'status': 'running'}
session_id = box_service.start_managed_process.await_args.args[0]
payload = box_service.start_managed_process.await_args.args[1]
execution_context = box_service.start_managed_process.await_args.args[0]
session_id = box_service.start_managed_process.await_args.args[1]
payload = box_service.start_managed_process.await_args.args[2]
assert execution_context == _CONTEXT
assert session_id == 'mcp-u1'
assert payload == {
'command': 'python',
@@ -118,9 +133,39 @@ async def test_workspace_session_start_managed_process_rewrites_command_and_args
}
@pytest.mark.asyncio
async def test_workspace_session_relay_connection_keeps_execution_context():
box_service = SimpleNamespace(
get_managed_process_websocket_connection=AsyncMock(
return_value=(
'ws://box/relay',
{'X-LangBot-Placement-Generation': '1'},
)
)
)
workspace = BoxWorkspaceSession(
box_service,
_CONTEXT,
'mcp-shared',
)
connection = await workspace.get_managed_process_websocket_connection('server-a')
assert connection == (
'ws://box/relay',
{'X-LangBot-Placement-Generation': '1'},
)
box_service.get_managed_process_websocket_connection.assert_awaited_once_with(
_CONTEXT,
'mcp-shared',
'server-a',
)
def test_workspace_session_build_session_payload_keeps_generic_workspace_shape():
workspace = BoxWorkspaceSession(
Mock(),
_CONTEXT,
'workspace-1',
host_path='/tmp/project',
host_path_mode='rw',
tests/unit_tests/cloud/test_bootstrap.py
+380
View File
@@ -0,0 +1,380 @@
from __future__ import annotations
import dataclasses
from types import SimpleNamespace
import pytest
from langbot.pkg.cloud.bootstrap import (
CloudBootstrapError,
CloudManifestRefreshService,
CloudRuntimeUnavailableError,
DeploymentAdmissionGuard,
OpenSourceDeployment,
VerifiedCloudDeployment,
resolve_deployment,
)
from langbot.pkg.cloud.entitlements import EntitlementSnapshot
pytestmark = pytest.mark.asyncio
class _Entitlements:
async def get_workspace_entitlement(self, workspace_uuid: str) -> EntitlementSnapshot:
return EntitlementSnapshot(
instance_uuid='instance-a',
workspace_uuid=workspace_uuid,
entitlement_revision=1,
status='active',
not_before=1,
expires_at=4_000_000_000,
features={'managed_sandbox': True},
limits={'managed_sandbox_sessions': 1},
)
class _Directory:
async def fetch_snapshot(self, instance_uuid: str):
del instance_uuid
raise AssertionError('not used by bootstrap contract tests')
async def fetch_events(self, instance_uuid: str, after_cursor: int, limit: int):
del instance_uuid, after_cursor, limit
raise AssertionError('not used by bootstrap contract tests')
async def fetch_workspaces(self, instance_uuid: str, workspace_uuids: tuple[str, ...]):
del instance_uuid, workspace_uuids
raise AssertionError('not used by bootstrap contract tests')
class _Manifest:
def __init__(self):
self.candidate = None
self.closed = False
async def refresh_manifest(self):
if self.candidate is None:
raise AssertionError('no refreshed Manifest was configured')
return self.candidate
async def aclose(self) -> None:
self.closed = True
class _Provider:
def __init__(self):
self.manifest_provider = _Manifest()
def bootstrap(self, *, instance_uuid: str, instance_config: dict):
del instance_config
return VerifiedCloudDeployment(
instance_uuid=instance_uuid,
manifest_jti='manifest-a',
manifest_generation=3,
expires_at=4_000_000_000,
release='cloud-v2',
capabilities=frozenset({'multi_workspace_v2'}),
tenant_isolation_version=2,
entitlement_provider=_Entitlements(),
directory_provider=_Directory(),
manifest_provider=self.manifest_provider,
verification_key_id='root-2026',
)
class _EntryPoint:
def __init__(self, value):
self.value = value
def load(self):
return self.value
class _EntryPoints(list):
def select(self, *, group: str):
return self if group == 'langbot.cloud_bootstrap' else []
def _cloud_config() -> dict:
return {
'database': {'use': 'postgresql'},
'vdb': {
'use': 'pgvector',
'pgvector': {
'use_business_database': True,
'allowed_dimensions': [384, 768, 1536],
},
},
'mcp': {'stdio': {'enabled': False}},
'plugin': {'worker': {'require_hard_limits': True}},
'box': {
'enabled': True,
'backend': 'nsjail',
'runtime': {'endpoint': 'ws://langbot-box:5410'},
'admission': {
'required': True,
'logical_session_id': 'global',
'required_backend': 'nsjail',
'max_sessions': 1,
'max_managed_processes': 0,
'max_grant_ttl_sec': 300,
'workspace_quota_mb': 32,
},
'local': {
'host_root': '/var/lib/langbot/box',
'default_workspace': '/var/lib/langbot/box/workspaces',
'allowed_mount_roots': ['/var/lib/langbot/box'],
},
},
# Proves mutable product metadata does not participate in selection.
'system': {'edition': 'community'},
}
async def test_no_closed_entry_point_selects_oss_singleton_even_if_edition_says_cloud():
deployment = await resolve_deployment(
instance_uuid='instance-a',
instance_config={'system': {'edition': 'cloud'}},
entry_points=lambda: _EntryPoints(),
)
assert isinstance(deployment, OpenSourceDeployment)
assert deployment.multi_workspace_enabled is False
async def test_verified_closed_entry_point_activates_cloud_policy():
deployment = await resolve_deployment(
instance_uuid='instance-a',
instance_config=_cloud_config(),
entry_points=lambda: _EntryPoints([_EntryPoint(_Provider)]),
now=1_000,
)
assert isinstance(deployment, VerifiedCloudDeployment)
assert deployment.multi_workspace_enabled is True
assert deployment.persistence_mode == 'cloud_runtime'
@pytest.mark.parametrize(
('field', 'value', 'message'),
[
('database', {'use': 'sqlite'}, 'database.use=postgresql'),
('vdb', {'use': 'chroma'}, 'vdb.use=pgvector'),
('mcp', {'stdio': {'enabled': True}}, 'mcp.stdio.enabled=false'),
('plugin', {'worker': {'require_hard_limits': False}}, 'plugin.worker.require_hard_limits=true'),
],
)
async def test_cloud_runtime_config_is_fail_closed(field, value, message):
config = _cloud_config()
config[field] = value
with pytest.raises(CloudBootstrapError, match=message):
await resolve_deployment(
instance_uuid='instance-a',
instance_config=config,
entry_points=lambda: _EntryPoints([_EntryPoint(_Provider())]),
now=1_000,
)
@pytest.mark.parametrize(
('directory_config', 'message'),
[
({'max_active_workspaces': 0}, 'greater than or equal to 1'),
({'max_active_workspaces': True}, 'must be an integer'),
(
{
'max_active_workspaces': 10,
'max_snapshot_workspaces': 9,
},
'max_snapshot_workspaces',
),
({'max_response_bytes': 64 * 1024 * 1024 + 1}, 'less than or equal to'),
],
)
async def test_cloud_directory_capacity_contract_is_fail_closed(directory_config, message):
config = _cloud_config()
config['cloud'] = {'directory': directory_config}
with pytest.raises(CloudBootstrapError, match=message):
await resolve_deployment(
instance_uuid='instance-a',
instance_config=config,
entry_points=lambda: _EntryPoints([_EntryPoint(_Provider())]),
now=1_000,
)
@pytest.mark.parametrize(
('pgvector_config', 'message'),
[
({'use_business_database': False, 'allowed_dimensions': [1536]}, 'use_business_database=true'),
({'use_business_database': True, 'allowed_dimensions': []}, 'allowed_dimensions'),
({'use_business_database': True, 'allowed_dimensions': [3072]}, 'allowed_dimensions'),
({'use_business_database': True, 'allowed_dimensions': [True]}, 'allowed_dimensions'),
],
)
async def test_cloud_pgvector_contract_is_fail_closed(pgvector_config, message):
config = _cloud_config()
config['vdb']['pgvector'] = pgvector_config
with pytest.raises(CloudBootstrapError, match=message):
await resolve_deployment(
instance_uuid='instance-a',
instance_config=config,
entry_points=lambda: _EntryPoints([_EntryPoint(_Provider())]),
now=1_000,
)
@pytest.mark.parametrize(
('mutate', 'message'),
[
(lambda config: config['box'].update(enabled=False), 'box.enabled=true'),
(lambda config: config['box'].update(backend='docker'), 'box.backend=nsjail'),
(lambda config: config['box']['runtime'].update(endpoint=''), 'box.runtime.endpoint'),
(
lambda config: config['box']['admission'].update(max_sessions=2),
'grant-enforced Box admission',
),
(
lambda config: config['box']['admission'].update(max_managed_processes=1),
'zero managed processes',
),
(
lambda config: config['box']['admission'].update(max_grant_ttl_sec=301),
'max_grant_ttl_sec',
),
(
lambda config: config['box']['admission'].update(workspace_quota_mb=0),
'workspace_quota_mb must be a positive integer',
),
(
lambda config: config['box']['admission'].update(workspace_quota_mb=True),
'workspace_quota_mb must be a positive integer',
),
(
lambda config: config['box']['local'].update(default_workspace='relative/workspaces'),
'default_workspace must be an absolute',
),
(
lambda config: config['box']['local'].update(
default_workspace='/other/workspaces',
),
'under allowed_mount_roots',
),
],
)
async def test_cloud_box_contract_is_fail_closed(mutate, message):
config = _cloud_config()
mutate(config)
with pytest.raises(CloudBootstrapError, match=message):
await resolve_deployment(
instance_uuid='instance-a',
instance_config=config,
entry_points=lambda: _EntryPoints([_EntryPoint(_Provider())]),
now=1_000,
)
async def test_invalid_provider_never_falls_back_to_oss():
provider = SimpleNamespace(bootstrap=lambda **_: object())
with pytest.raises(CloudBootstrapError, match='must return VerifiedCloudDeployment'):
await resolve_deployment(
instance_uuid='instance-a',
instance_config=_cloud_config(),
entry_points=lambda: _EntryPoints([_EntryPoint(provider)]),
now=1_000,
)
async def test_duplicate_closed_providers_fail_closed():
with pytest.raises(CloudBootstrapError, match='Exactly one'):
await resolve_deployment(
instance_uuid='instance-a',
instance_config=_cloud_config(),
entry_points=lambda: _EntryPoints([_EntryPoint(_Provider()), _EntryPoint(_Provider())]),
now=1_000,
)
async def test_deployment_admission_expires_even_after_wall_clock_rollback():
wall = [1_000.0]
monotonic = [50.0]
deployment = dataclasses.replace(
_Provider().bootstrap(instance_uuid='instance-a', instance_config={}),
expires_at=1_010,
)
guard = DeploymentAdmissionGuard(
'instance-a',
deployment,
wall_time=lambda: wall[0],
monotonic_time=lambda: monotonic[0],
)
assert guard.require_active() is deployment
wall[0] = 900.0
monotonic[0] = 60.0
with pytest.raises(CloudRuntimeUnavailableError, match='expired'):
guard.require_active()
async def test_deployment_admission_accepts_only_monotonic_non_conflicting_renewal():
wall = [1_000.0]
monotonic = [50.0]
current = dataclasses.replace(
_Provider().bootstrap(instance_uuid='instance-a', instance_config={}),
expires_at=1_010,
)
guard = DeploymentAdmissionGuard(
'instance-a',
current,
wall_time=lambda: wall[0],
monotonic_time=lambda: monotonic[0],
)
renewed = dataclasses.replace(
current,
manifest_jti='manifest-b',
manifest_generation=4,
expires_at=2_000,
)
guard.replace(renewed)
assert guard.require_active() is renewed
rollback = dataclasses.replace(current, manifest_generation=2)
with pytest.raises(CloudRuntimeUnavailableError, match='rolled back'):
guard.replace(rollback)
conflicting = dataclasses.replace(renewed, manifest_jti='different')
with pytest.raises(CloudRuntimeUnavailableError, match='conflicting'):
guard.replace(conflicting)
async def test_manifest_refresh_replaces_receipt_before_short_ttl_expires():
wall = [1_000.0]
provider = _Provider()
current = dataclasses.replace(
provider.bootstrap(instance_uuid='instance-a', instance_config={}),
expires_at=1_300,
)
guard = DeploymentAdmissionGuard('instance-a', current, wall_time=lambda: wall[0])
renewed = dataclasses.replace(
current,
manifest_jti='manifest-renewed',
manifest_generation=current.manifest_generation + 1,
expires_at=2_000,
)
provider.manifest_provider.candidate = renewed
service = CloudManifestRefreshService(
guard,
provider.manifest_provider,
SimpleNamespace(exception=lambda *_: None),
wall_time=lambda: wall[0],
)
assert service.next_refresh_delay() == 120
assert await service.refresh_once() is renewed
assert guard.deployment is renewed
tests/unit_tests/cloud/test_directory_projection.py
+1062
View File
File diff suppressed because it is too large Load Diff
tests/unit_tests/cloud/test_entitlements.py
+161
View File
@@ -0,0 +1,161 @@
from __future__ import annotations
import asyncio
import pytest
from unittest.mock import AsyncMock
from langbot.pkg.cloud.entitlements import EntitlementResolver, EntitlementSnapshot, EntitlementUnavailableError
def _snapshot(**overrides) -> EntitlementSnapshot:
values = {
'instance_uuid': 'instance-a',
'workspace_uuid': 'workspace-a',
'entitlement_revision': 7,
'status': 'active',
'not_before': 100,
'expires_at': 200,
'features': {'managed_sandbox': True, 'mcp_stdio': False},
'limits': {'managed_sandbox_sessions': 1},
}
values.update(overrides)
return EntitlementSnapshot(**values)
def test_active_snapshot_exposes_only_generic_features_and_limits():
snapshot = _snapshot().require_active(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
now=150,
)
snapshot.require_feature('managed_sandbox')
assert snapshot.limit('managed_sandbox_sessions') == 1
assert 'plan' not in snapshot.model_fields
@pytest.mark.parametrize(
'snapshot,now',
[
(_snapshot(status='suspended'), 150),
(_snapshot(), 99),
(_snapshot(), 200),
],
)
def test_inactive_or_expired_snapshot_fails_closed(snapshot, now):
with pytest.raises(EntitlementUnavailableError):
snapshot.require_active(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
now=now,
)
def test_scope_mismatch_fails_closed():
with pytest.raises(EntitlementUnavailableError, match='scope'):
_snapshot().require_active(
instance_uuid='instance-a',
workspace_uuid='workspace-b',
now=150,
)
@pytest.mark.asyncio
async def test_resolver_rejects_revision_rollback():
provider = AsyncMock()
provider.get_workspace_entitlement = AsyncMock(side_effect=[_snapshot(), _snapshot(entitlement_revision=6)])
resolver = EntitlementResolver('instance-a', provider)
await resolver.resolve('workspace-a', now=150)
with pytest.raises(EntitlementUnavailableError, match='rolled back'):
await resolver.resolve('workspace-a', now=150)
@pytest.mark.asyncio
async def test_resolver_rejects_same_revision_with_different_contents():
provider = AsyncMock()
provider.get_workspace_entitlement = AsyncMock(
side_effect=[
_snapshot(),
_snapshot(features={'managed_sandbox': False}),
]
)
resolver = EntitlementResolver('instance-a', provider)
await resolver.resolve('workspace-a', now=150)
with pytest.raises(EntitlementUnavailableError, match='conflicting contents'):
await resolver.resolve('workspace-a', now=150)
@pytest.mark.asyncio
async def test_resolver_checks_deployment_admission_before_and_after_provider_call():
checks = 0
def require_admission() -> None:
nonlocal checks
checks += 1
if checks == 2:
raise RuntimeError('manifest expired during provider call')
provider = AsyncMock()
provider.get_workspace_entitlement = AsyncMock(return_value=_snapshot())
resolver = EntitlementResolver(
'instance-a',
provider,
deployment_admission=require_admission,
)
with pytest.raises(RuntimeError, match='expired during provider call'):
await resolver.resolve('workspace-a', now=150)
assert checks == 2
@pytest.mark.asyncio
async def test_directory_activity_reconciliation_drops_historical_snapshots():
provider = AsyncMock()
provider.get_workspace_entitlement = AsyncMock(return_value=_snapshot())
resolver = EntitlementResolver('instance-a', provider)
await resolver.reconcile_active_workspaces({'workspace-a', 'workspace-b'})
await resolver.resolve('workspace-a', now=150)
await resolver.reconcile_active_workspaces({'workspace-b'})
assert resolver.snapshot_counts() == {
'active_workspaces': 1,
'cached_snapshots': 0,
}
with pytest.raises(EntitlementUnavailableError, match='directory projection'):
await resolver.resolve('workspace-a', now=150)
provider.get_workspace_entitlement.assert_awaited_once()
@pytest.mark.asyncio
async def test_directory_fence_wins_race_with_inflight_entitlement_fetch():
provider_started = asyncio.Event()
release_provider = asyncio.Event()
async def fetch(_workspace_uuid: str) -> EntitlementSnapshot:
provider_started.set()
await release_provider.wait()
return _snapshot()
provider = AsyncMock()
provider.get_workspace_entitlement = AsyncMock(side_effect=fetch)
resolver = EntitlementResolver('instance-a', provider)
await resolver.reconcile_active_workspaces({'workspace-a'})
resolve_task = asyncio.create_task(resolver.resolve('workspace-a', now=150))
await provider_started.wait()
await resolver.update_workspace_activity(
active_workspace_uuids=set(),
inactive_workspace_uuids={'workspace-a'},
)
release_provider.set()
with pytest.raises(EntitlementUnavailableError, match='directory projection'):
await resolve_task
assert resolver.snapshot_counts() == {
'active_workspaces': 0,
'cached_snapshots': 0,
}
tests/unit_tests/cloud/test_space_launch.py
+164
View File
@@ -0,0 +1,164 @@
from __future__ import annotations
import base64
import json
import time
import uuid
from types import SimpleNamespace
import pytest
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
from langbot.pkg.cloud.launch import SpaceLaunchError, SpaceLaunchService
pytestmark = pytest.mark.asyncio
INSTANCE_UUID = 'instance-test'
ACCOUNT_UUID = '11111111-1111-4111-8111-111111111111'
WORKSPACE_UUID = '22222222-2222-4222-8222-222222222222'
KEY_ID = 'space-key-1'
def _base64url(raw: bytes) -> str:
return base64.urlsafe_b64encode(raw).rstrip(b'=').decode('ascii')
def _sign(private_key: Ed25519PrivateKey, claims: dict, *, key_id: str = KEY_ID) -> str:
header = {'alg': 'EdDSA', 'kid': key_id, 'typ': 'langbot-control-plane+jwt'}
encoded_header = _base64url(json.dumps(header, separators=(',', ':')).encode('utf-8'))
encoded_claims = _base64url(json.dumps(claims, separators=(',', ':')).encode('utf-8'))
signing_input = f'{encoded_header}.{encoded_claims}'
return f'{signing_input}.{_base64url(private_key.sign(signing_input.encode("ascii")))}'
def _claims(*, now: int, jti: str | None = None, workspace_uuid: str = WORKSPACE_UUID) -> dict:
return {
'iss': 'langbot-space',
'aud': 'langbot-cloud-runtime',
'sub': f'langbot-instance:{INSTANCE_UUID}',
'jti': jti or str(uuid.uuid4()),
'iat': now,
'nbf': now - 5,
'exp': now + 90,
'instance_uuid': INSTANCE_UUID,
'kind': 'workspace.launch',
'payload': {
'account_uuid': ACCOUNT_UUID,
'workspace_uuid': workspace_uuid,
},
}
def _service(private_key: Ed25519PrivateKey, *, now: int) -> SpaceLaunchService:
public_key = private_key.public_key().public_bytes(
encoding=serialization.Encoding.Raw,
format=serialization.PublicFormat.Raw,
)
app = SimpleNamespace(
deployment=SimpleNamespace(multi_workspace_enabled=True, verification_key_id=KEY_ID),
workspace_service=SimpleNamespace(instance_uuid=INSTANCE_UUID),
instance_config=SimpleNamespace(
data={
'space': {
'launch': {
'control_plane_public_key': _base64url(public_key),
}
}
}
),
)
return SpaceLaunchService(app, wall_time=lambda: now)
async def test_consumes_valid_workspace_launch_assertion_once():
private_key = Ed25519PrivateKey.generate()
now = int(time.time())
service = _service(private_key, now=now)
token = _sign(private_key, _claims(now=now))
launch = await service.consume_assertion(token, expected_workspace_uuid=WORKSPACE_UUID)
assert launch == {'account_uuid': ACCOUNT_UUID, 'workspace_uuid': WORKSPACE_UUID}
with pytest.raises(SpaceLaunchError, match='already been consumed'):
await service.consume_assertion(token, expected_workspace_uuid=WORKSPACE_UUID)
async def test_replay_cache_does_not_scan_all_live_assertions(monkeypatch):
private_key = Ed25519PrivateKey.generate()
now = int(time.time())
service = _service(private_key, now=now)
for index in range(512):
await service._consume_jti(f'jti-{index}', now + 90)
class NoGlobalIterationDict(dict):
def __iter__(self):
raise AssertionError('replay admission scanned all live assertions')
def keys(self):
raise AssertionError('replay admission scanned all live assertions')
def items(self):
raise AssertionError('replay admission scanned all live assertions')
def values(self):
raise AssertionError('replay admission scanned all live assertions')
guarded_jtis = NoGlobalIterationDict(service._consumed_jtis)
monkeypatch.setattr(service, '_consumed_jtis', guarded_jtis)
await service._consume_jti('jti-new', now + 90)
assert len(guarded_jtis) == 513
async def test_replay_cache_fails_closed_at_capacity(monkeypatch):
from langbot.pkg.cloud import launch
private_key = Ed25519PrivateKey.generate()
now = int(time.time())
service = _service(private_key, now=now)
monkeypatch.setattr(launch, '_CONSUMED_JTI_MAX_ENTRIES', 2)
await service._consume_jti('jti-1', now + 90)
await service._consume_jti('jti-2', now + 90)
with pytest.raises(SpaceLaunchError, match='replay cache capacity'):
await service._consume_jti('jti-3', now + 90)
with pytest.raises(SpaceLaunchError, match='already been consumed'):
await service._consume_jti('jti-1', now + 90)
async def test_rejects_expired_wrong_workspace_and_wrong_instance_assertions():
private_key = Ed25519PrivateKey.generate()
now = int(time.time())
service = _service(private_key, now=now)
expired = _claims(now=now)
expired['exp'] = now - 60
with pytest.raises(SpaceLaunchError, match='expired'):
await service.consume_assertion(_sign(private_key, expired), expected_workspace_uuid=WORKSPACE_UUID)
wrong_workspace = _sign(private_key, _claims(now=now, workspace_uuid='33333333-3333-4333-8333-333333333333'))
with pytest.raises(SpaceLaunchError, match='another Workspace'):
await service.consume_assertion(wrong_workspace, expected_workspace_uuid=WORKSPACE_UUID)
wrong_instance = _claims(now=now)
wrong_instance['instance_uuid'] = 'other-instance'
with pytest.raises(SpaceLaunchError, match='instance UUID'):
await service.consume_assertion(_sign(private_key, wrong_instance), expected_workspace_uuid=WORKSPACE_UUID)
async def test_rejects_invalid_signature_and_non_cloud_mode():
private_key = Ed25519PrivateKey.generate()
now = int(time.time())
token = _sign(private_key, _claims(now=now))
service = _service(Ed25519PrivateKey.generate(), now=now)
with pytest.raises(SpaceLaunchError, match='signature'):
await service.consume_assertion(token, expected_workspace_uuid=WORKSPACE_UUID)
oss_service = _service(private_key, now=now)
oss_service.ap.deployment.multi_workspace_enabled = False
with pytest.raises(SpaceLaunchError, match='verified Cloud mode'):
await oss_service.consume_assertion(token, expected_workspace_uuid=WORKSPACE_UUID)
tests/unit_tests/core/test_app_resource_maintenance.py
+113
View File
@@ -0,0 +1,113 @@
from __future__ import annotations
import asyncio
from types import SimpleNamespace
import pytest
from langbot.pkg.core.app import Application
pytestmark = pytest.mark.asyncio
class _TaskManager:
def __init__(self, stop: asyncio.Event) -> None:
self.stop = stop
self.tasks: list[asyncio.Task] = []
def create_task(self, coro, *, name='', **_kwargs):
task = asyncio.create_task(coro, name=name)
self.tasks.append(task)
return SimpleNamespace(task=task)
async def wait_all(self) -> None:
await self.stop.wait()
for task in self.tasks:
task.cancel()
await asyncio.gather(*self.tasks, return_exceptions=True)
async def _wait_forever() -> None:
await asyncio.Event().wait()
async def test_resource_maintenance_waits_and_shares_workspace_discovery() -> None:
stop = asyncio.Event()
completed = asyncio.Event()
discovery_calls = 0
job_calls: list[str] = []
async def list_bindings():
nonlocal discovery_calls
discovery_calls += 1
return [
SimpleNamespace(
instance_uuid='instance',
workspace_uuid='workspace',
placement_generation=1,
)
]
async def cleanup_monitoring(_context, _retention_days, *, batch_size):
assert batch_size == 10
job_calls.append('monitoring')
return {}
async def cleanup_storage(_context):
job_calls.append('storage')
completed.set()
return {}
application = Application()
application.event_loop = asyncio.get_running_loop()
application.event_loop_monitor = SimpleNamespace(start=lambda: None)
application.task_mgr = _TaskManager(stop)
application.plugin_connector = SimpleNamespace(initialize_plugins=lambda: asyncio.sleep(0))
application.platform_mgr = SimpleNamespace(run=_wait_forever)
application.ctrl = SimpleNamespace(run=_wait_forever)
application.http_ctrl = SimpleNamespace(run=_wait_forever)
application.telemetry = None
application.workspace_collaboration_service = None
application.workspace_service = SimpleNamespace(list_active_execution_bindings=list_bindings)
application.monitoring_service = SimpleNamespace(cleanup_expired_records=cleanup_monitoring)
application.maintenance_service = SimpleNamespace(cleanup_expired_files=cleanup_storage)
application.instance_config = SimpleNamespace(
data={
'monitoring': {
'auto_cleanup': {
'enabled': True,
'retention_days': 30,
'delete_batch_size': 10,
'check_interval_hours': 0.00002,
}
},
'storage': {
'cleanup': {
'enabled': True,
'check_interval_hours': 0.00002,
}
},
}
)
application.logger = SimpleNamespace(
info=lambda *_args, **_kwargs: None,
warning=lambda *_args, **_kwargs: None,
error=lambda *_args, **_kwargs: None,
debug=lambda *_args, **_kwargs: None,
)
async def no_web_info() -> None:
return None
application.print_web_access_info = no_web_info
run_task = asyncio.create_task(application.run())
try:
await asyncio.sleep(0.01)
assert discovery_calls == 0
await asyncio.wait_for(completed.wait(), timeout=1)
assert discovery_calls == 1
assert job_calls == ['monitoring', 'storage']
finally:
stop.set()
await asyncio.wait_for(run_task, timeout=1)
tests/unit_tests/core/test_app_shutdown.py
+142
View File
@@ -0,0 +1,142 @@
from __future__ import annotations
import asyncio
from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
from langbot.pkg.core.app import Application
@pytest.mark.asyncio
async def test_shutdown_closes_mcp_session_manager_once() -> None:
app = Application()
stop_session_manager = AsyncMock()
app.platform_mgr = SimpleNamespace(shutdown=AsyncMock())
app.tool_mgr = SimpleNamespace(shutdown=AsyncMock())
app.model_mgr = SimpleNamespace(shutdown=AsyncMock())
app.box_service = SimpleNamespace(shutdown=AsyncMock())
app.plugin_connector = SimpleNamespace(aclose=AsyncMock())
app.telemetry = SimpleNamespace(shutdown=AsyncMock())
app.vector_db_mgr = SimpleNamespace(shutdown=AsyncMock())
app.storage_mgr = SimpleNamespace(shutdown=AsyncMock())
manifest_provider = SimpleNamespace(aclose=AsyncMock())
app.deployment = SimpleNamespace(manifest_provider=manifest_provider)
persistence_engine = SimpleNamespace(dispose=AsyncMock())
app.persistence_mgr = SimpleNamespace(db=SimpleNamespace(engine=persistence_engine))
app.http_ctrl = SimpleNamespace(mcp_mount=SimpleNamespace(stop_session_manager=stop_session_manager))
await app.shutdown()
await app.shutdown()
stop_session_manager.assert_awaited_once()
app.platform_mgr.shutdown.assert_awaited_once()
app.tool_mgr.shutdown.assert_awaited_once()
app.model_mgr.shutdown.assert_awaited_once()
app.box_service.shutdown.assert_awaited_once()
app.plugin_connector.aclose.assert_awaited_once()
app.telemetry.shutdown.assert_awaited_once()
app.vector_db_mgr.shutdown.assert_awaited_once()
app.storage_mgr.shutdown.assert_awaited_once()
manifest_provider.aclose.assert_awaited_once()
persistence_engine.dispose.assert_awaited_once()
@pytest.mark.asyncio
async def test_dispose_tracks_only_one_shutdown_task() -> None:
app = Application()
app.event_loop = asyncio.get_running_loop()
app.dispose()
shutdown_task = app._shutdown_task
app.dispose()
assert shutdown_task is not None
assert app._shutdown_task is shutdown_task
await shutdown_task
app.dispose()
assert app._shutdown_task is shutdown_task
@pytest.mark.asyncio
async def test_runtime_resource_stats_are_aggregate_and_constant_time() -> None:
app = Application()
app.event_loop = asyncio.get_running_loop()
app.blocking_executor = SimpleNamespace(
snapshot=lambda: {
'inflight': 3,
'running': 2,
'pending': 1,
'rejected_total': 4,
}
)
app.task_mgr = SimpleNamespace(get_stats=lambda: {'total': 5, 'completed': 2})
app.query_pool = SimpleNamespace(
queries=[object()],
cached_queries={},
active_query_count_by_workspace={'workspace-a': 1},
)
app.model_mgr = SimpleNamespace(
provider_dict={'provider': object()},
llm_model_dict={},
embedding_model_dict={},
rerank_model_dict={},
)
app.platform_mgr = SimpleNamespace(_bots_by_key={})
app.pipeline_mgr = SimpleNamespace(_pipelines_by_key={})
app.rag_mgr = SimpleNamespace(knowledge_bases={})
app.plugin_connector = SimpleNamespace(_known_desired_states={'installation': object()})
app.persistence_mgr = SimpleNamespace(
get_resource_stats=lambda: {
'configured_capacity': 20,
'checked_out': 3,
}
)
app.directory_projection_service = SimpleNamespace(
resource_snapshot=lambda: {
'active_workspaces': 10,
'max_active_workspaces': 1000,
}
)
app.tool_mgr = SimpleNamespace(
mcp_tool_loader=SimpleNamespace(
_sessions={},
_hosted_mcp_tasks=[],
_host_dispatch_tasks=set(),
)
)
app.telemetry = SimpleNamespace(send_tasks=[])
stats = app.get_runtime_resource_stats()
assert stats['asyncio_tasks'] >= 1
assert stats['event_loop'] == {
'running': False,
'samples_total': 0,
'last_lag_ms': 0,
'recent_p95_lag_ms': 0,
'recent_max_lag_ms': 0,
'max_lag_ms': 0,
}
assert stats['blocking_executor']['rejected_total'] == 4
assert stats['application_tasks'] == {
'total': 5,
'completed': 2,
}
assert stats['database_pool'] == {
'configured_capacity': 20,
'checked_out': 3,
}
assert stats['directory'] == {
'active_workspaces': 10,
'max_active_workspaces': 1000,
}
assert stats['query_pool'] == {
'queued': 1,
'cached': 0,
'active_workspaces': 1,
}
assert stats['models']['providers'] == 1
assert stats['runtimes']['plugin_installations'] == 1
tests/unit_tests/core/test_boot.py
+25 -1
View File
@@ -2,13 +2,37 @@ from __future__ import annotations
import signal
from types import SimpleNamespace
from unittest.mock import Mock
from unittest.mock import AsyncMock, Mock
import pytest
from langbot.pkg.core import boot
@pytest.mark.asyncio
async def test_make_app_shuts_down_partially_built_application(monkeypatch):
app_inst = SimpleNamespace(
event_loop=None,
shutdown=AsyncMock(),
initialize=AsyncMock(),
)
class FailingStage:
async def run(self, ap):
assert ap is app_inst
raise RuntimeError('startup failed')
monkeypatch.setattr(boot.app, 'Application', lambda: app_inst)
monkeypatch.setattr(boot, 'stage_order', ['FailingStage'])
monkeypatch.setitem(boot.stage.preregistered_stages, 'FailingStage', FailingStage)
with pytest.raises(RuntimeError, match='startup failed'):
await boot.make_app(SimpleNamespace())
app_inst.shutdown.assert_awaited_once()
app_inst.initialize.assert_not_awaited()
@pytest.mark.asyncio
async def test_main_signal_handler_handles_sigint_before_app_created(monkeypatch):
captured_handler = {}
tests/unit_tests/core/test_load_config.py
+121
View File
@@ -35,6 +35,22 @@ class TestApplyEnvOverridesToConfig:
assert result['system']['name'] == 'custom_name'
def test_override_log_never_prints_secret_value(self, capsys):
"""Environment-backed credentials must not be copied into logs."""
load_config = get_load_config_module()
secret = 'database-password-that-must-not-leak'
cfg = {'database': {'postgresql': {'password': ''}}}
env = {'DATABASE__POSTGRESQL__PASSWORD': secret}
with patch.dict(os.environ, env, clear=True):
result = load_config._apply_env_overrides_to_config(cfg)
captured = capsys.readouterr().out
assert result['database']['postgresql']['password'] == secret
assert 'DATABASE__POSTGRESQL__PASSWORD' in captured
assert secret not in captured
def test_override_int_value(self):
"""Test overriding an int value with proper conversion."""
load_config = get_load_config_module()
@@ -48,6 +64,20 @@ class TestApplyEnvOverridesToConfig:
assert result['concurrency']['pipeline'] == 10
assert isinstance(result['concurrency']['pipeline'], int)
def test_cloud_directory_limit_override_keeps_integer_type_on_upgraded_config(self):
load_config = get_load_config_module()
cfg = load_config._complete_runtime_policy_defaults({})
with patch.dict(
os.environ,
{'CLOUD__DIRECTORY__MAX_ACTIVE_WORKSPACES': '250'},
clear=True,
):
result = load_config._apply_env_overrides_to_config(cfg)
assert result['cloud']['directory']['max_active_workspaces'] == 250
assert isinstance(result['cloud']['directory']['max_active_workspaces'], int)
def test_override_int_value_invalid_conversion(self):
"""Test that invalid int conversion keeps string value."""
load_config = get_load_config_module()
@@ -196,6 +226,19 @@ class TestApplyEnvOverridesToConfig:
assert result['system']['name'] == 'default'
def test_skip_env_vars_with_empty_path_segments(self, capsys):
"""Platform variables such as __CF_USER_TEXT_ENCODING are not config."""
load_config = get_load_config_module()
cfg = {'system': {'name': 'default'}}
env = {'__CF_USER_TEXT_ENCODING': '0x1F5:0x0:0x64'}
with patch.dict(os.environ, env, clear=True):
result = load_config._apply_env_overrides_to_config(cfg)
assert result == cfg
assert capsys.readouterr().out == ''
def test_nested_config_path(self):
"""Test overriding deeply nested config."""
load_config = get_load_config_module()
@@ -259,6 +302,84 @@ class TestApplyEnvOverridesToConfig:
assert result['system']['enable'] is False
assert result['concurrency']['pipeline'] == 10
def test_plugin_worker_and_stdio_policy_native_env_overrides(self):
load_config = get_load_config_module()
cfg = {
'plugin': {
'worker': {
'max_cpus': 1.0,
'max_memory_mb': 512,
'max_pids': 128,
'max_open_files': 256,
'max_file_size_mb': 512,
'max_concurrent_restarts': 1,
'restart_failure_threshold': 8,
'restart_failure_window_seconds': 30.0,
'restart_circuit_open_seconds': 60.0,
}
},
'mcp': {'stdio': {'enabled': True}},
}
env = {
'PLUGIN__WORKER__MAX_CPUS': '2.5',
'PLUGIN__WORKER__MAX_MEMORY_MB': '1024',
'PLUGIN__WORKER__MAX_PIDS': '64',
'PLUGIN__WORKER__MAX_OPEN_FILES': '128',
'PLUGIN__WORKER__MAX_FILE_SIZE_MB': '256',
'PLUGIN__WORKER__MAX_CONCURRENT_RESTARTS': '2',
'PLUGIN__WORKER__RESTART_FAILURE_THRESHOLD': '12',
'PLUGIN__WORKER__RESTART_FAILURE_WINDOW_SECONDS': '45.5',
'PLUGIN__WORKER__RESTART_CIRCUIT_OPEN_SECONDS': '90.0',
'MCP__STDIO__ENABLED': 'false',
}
with patch.dict(os.environ, env, clear=True):
result = load_config._apply_env_overrides_to_config(cfg)
assert result['plugin']['worker'] == {
'max_cpus': 2.5,
'max_memory_mb': 1024,
'max_pids': 64,
'max_open_files': 128,
'max_file_size_mb': 256,
'max_concurrent_restarts': 2,
'restart_failure_threshold': 12,
'restart_failure_window_seconds': 45.5,
'restart_circuit_open_seconds': 90.0,
}
assert result['mcp']['stdio']['enabled'] is False
def test_runtime_policy_defaults_preserve_env_types_for_upgraded_config(self):
load_config = get_load_config_module()
cfg = {'plugin': {'enable': True}}
completed = load_config._complete_runtime_policy_defaults(cfg)
with patch.dict(
os.environ,
{
'PLUGIN__WORKER__MAX_MEMORY_MB': '768',
'MCP__STDIO__ENABLED': 'false',
'SYSTEM__BLOCKING_EXECUTOR__MAX_WORKERS': '12',
'SYSTEM__BLOCKING_EXECUTOR__MAX_PENDING': '256',
'SYSTEM__BLOCKING_EXECUTOR__MAX_INFLIGHT_PER_SCOPE': '3',
},
clear=True,
):
result = load_config._apply_env_overrides_to_config(completed)
assert result['system']['blocking_executor'] == {
'max_workers': 12,
'max_pending': 256,
'max_inflight_per_scope': 3,
}
assert isinstance(
result['system']['blocking_executor']['max_workers'],
int,
)
assert result['plugin']['worker']['max_memory_mb'] == 768
assert isinstance(result['plugin']['worker']['max_memory_mb'], int)
assert result['mcp']['stdio']['enabled'] is False
def test_webhook_prefix_override(self):
"""Test overriding webhook_prefix via environment variable."""
load_config = get_load_config_module()
tests/unit_tests/core/test_taskmgr.py
+121
View File
@@ -12,6 +12,8 @@ from __future__ import annotations
import pytest
import asyncio
import contextvars
import inspect
import sys
from unittest.mock import Mock, MagicMock
from contextlib import contextmanager
@@ -264,6 +266,28 @@ class TestTaskWrapper:
wrapper.cancel()
@pytest.mark.asyncio
async def test_workspace_task_sets_blocking_work_scope(self):
"""Detached tasks recover tenant fairness from durable ownership."""
_, TaskWrapper, _ = get_taskmgr_classes()
from langbot.pkg.utils.bounded_executor import (
current_blocking_work_scope,
)
mock_app = create_mock_app()
async def read_scope():
return current_blocking_work_scope()
wrapper = TaskWrapper(
mock_app,
read_scope(),
workspace_uuid='workspace-a',
)
assert await wrapper.task == 'workspace-a'
assert current_blocking_work_scope() is None
@pytest.mark.asyncio
async def test_to_dict_serialization(self):
"""Test TaskWrapper.to_dict serialization."""
@@ -360,6 +384,53 @@ class TestAsyncTaskManager:
wrapper.cancel()
@pytest.mark.asyncio
async def test_create_task_does_not_inherit_request_context(self):
"""Long-lived tasks must receive identity through explicit arguments."""
_, _, AsyncTaskManager = get_taskmgr_classes()
mock_app = create_mock_app()
manager = AsyncTaskManager(mock_app)
request_value = contextvars.ContextVar('request_value', default=None)
token = request_value.set('request-scoped-transaction')
observed = []
async def detached_task(captured_workspace: str) -> None:
observed.append((request_value.get(), captured_workspace))
try:
wrapper = manager.create_task(detached_task('workspace-a'))
await wrapper.task
finally:
request_value.reset(token)
assert observed == [(None, 'workspace-a')]
@pytest.mark.asyncio
async def test_create_task_waits_for_registered_transaction_commit(self):
_, _, AsyncTaskManager = get_taskmgr_classes()
mock_app = create_mock_app()
gate = asyncio.get_running_loop().create_future()
class PersistenceManagerStub:
def create_after_commit_gate(self):
return gate
mock_app.persistence_mgr = PersistenceManagerStub()
manager = AsyncTaskManager(mock_app)
observed = []
async def background_work() -> None:
observed.append('started')
wrapper = manager.create_task(background_work())
await asyncio.sleep(0)
assert observed == []
gate.set_result(None)
await wrapper.task
assert observed == ['started']
@pytest.mark.asyncio
async def test_get_stats_counts_correctly(self):
"""Test get_stats returns correct counts."""
@@ -482,6 +553,56 @@ class TestAsyncTaskManager:
wrapper.cancel()
@pytest.mark.asyncio
async def test_create_user_task_enforces_workspace_active_limit_and_closes_rejected_coroutine(self):
"""A noisy Workspace cannot accumulate unbounded background work."""
_, _, AsyncTaskManager = get_taskmgr_classes()
mock_app = create_mock_app()
mock_app.instance_config.data['system']['task_retention'].update(
{
'max_active_user_tasks': 10,
'max_active_user_tasks_per_workspace': 1,
}
)
manager = AsyncTaskManager(mock_app)
async def long_coro():
await asyncio.sleep(10)
first = manager.create_user_task(long_coro(), workspace_uuid='workspace-a')
rejected = long_coro()
with pytest.raises(RuntimeError, match='Workspace has too many active user operations'):
manager.create_user_task(rejected, workspace_uuid='workspace-a')
assert inspect.getcoroutinestate(rejected) == inspect.CORO_CLOSED
other_workspace = manager.create_user_task(long_coro(), workspace_uuid='workspace-b')
first.cancel()
other_workspace.cancel()
@pytest.mark.asyncio
async def test_create_user_task_enforces_instance_active_limit(self):
"""The shared process retains a hard cap even across Workspaces."""
_, _, AsyncTaskManager = get_taskmgr_classes()
mock_app = create_mock_app()
mock_app.instance_config.data['system']['task_retention'].update(
{
'max_active_user_tasks': 1,
'max_active_user_tasks_per_workspace': 10,
}
)
manager = AsyncTaskManager(mock_app)
async def long_coro():
await asyncio.sleep(10)
first = manager.create_user_task(long_coro(), workspace_uuid='workspace-a')
rejected = long_coro()
with pytest.raises(RuntimeError, match='instance has too many active user operations'):
manager.create_user_task(rejected, workspace_uuid='workspace-b')
assert inspect.getcoroutinestate(rejected) == inspect.CORO_CLOSED
first.cancel()
@pytest.mark.asyncio
async def test_get_task_by_id(self):
"""Test get_task_by_id returns correct task."""
tests/unit_tests/persistence/test_database_decorator.py
+11 -1
View File
@@ -11,9 +11,19 @@ Note: Uses import isolation to break circular import chains.
from __future__ import annotations
import sys
from unittest.mock import Mock, MagicMock
from contextlib import contextmanager
from typing import Generator
from unittest.mock import MagicMock, Mock
import pytest
@pytest.fixture(autouse=True)
def isolate_database_manager_registry(monkeypatch: pytest.MonkeyPatch) -> None:
"""Keep decorator tests from mutating the process-wide manager registry."""
from langbot.pkg.persistence import database
monkeypatch.setattr(database, 'preregistered_managers', list(database.preregistered_managers))
@contextmanager
tests/unit_tests/persistence/test_postgresql_database.py
+233
View File
@@ -0,0 +1,233 @@
from __future__ import annotations
from types import SimpleNamespace
import pytest
# Persistence manager performs the package's database-manager registration;
# importing a concrete manager first would enter the historical app/mgr cycle.
from langbot.pkg.persistence import mgr as _persistence_mgr # noqa: F401
from langbot.pkg.persistence.databases import postgresql
@pytest.mark.asyncio
async def test_postgresql_manager_parses_explicit_url_without_string_reassembly(monkeypatch) -> None:
captured = None
captured_options = None
sentinel_engine = object()
def create_engine(url, **options):
nonlocal captured, captured_options
captured = url
captured_options = options
return sentinel_engine
monkeypatch.setattr(postgresql.sqlalchemy_asyncio, 'create_async_engine', create_engine)
ap = SimpleNamespace(
instance_config=SimpleNamespace(
data={
'database': {
'postgresql': {
'url': 'postgresql://runtime:p%40ss@db.internal:5432/langbot?sslmode=require',
}
}
}
)
)
manager = postgresql.PostgreSQLDatabaseManager(ap)
await manager.initialize()
assert captured.drivername == 'postgresql+asyncpg'
assert captured.password == 'p@ss'
assert captured.query['ssl'] == 'require'
assert 'sslmode' not in captured.query
assert captured_options == {
'pool_size': 10,
'max_overflow': 10,
'pool_timeout': 30,
'pool_recycle': 1800,
'pool_pre_ping': True,
}
assert manager.engine is sentinel_engine
@pytest.mark.asyncio
async def test_postgresql_manager_builds_structured_url_with_special_password(monkeypatch) -> None:
captured = None
def create_engine(url, **_options):
nonlocal captured
captured = url
return object()
monkeypatch.setattr(postgresql.sqlalchemy_asyncio, 'create_async_engine', create_engine)
ap = SimpleNamespace(
instance_config=SimpleNamespace(
data={
'database': {
'postgresql': {
'host': 'db.internal',
'port': 5432,
'user': 'runtime',
'password': 'p@ss:/?#word',
'database': 'langbot',
}
}
}
)
)
await postgresql.PostgreSQLDatabaseManager(ap).initialize()
assert captured.password == 'p@ss:/?#word'
assert captured.host == 'db.internal'
assert captured.database == 'langbot'
@pytest.mark.asyncio
async def test_postgresql_manager_applies_explicit_bounded_pool_options(monkeypatch) -> None:
captured_options = None
def create_engine(_url, **options):
nonlocal captured_options
captured_options = options
return object()
monkeypatch.setattr(postgresql.sqlalchemy_asyncio, 'create_async_engine', create_engine)
ap = SimpleNamespace(
instance_config=SimpleNamespace(
data={
'database': {
'postgresql': {
'pool_size': 24,
'max_overflow': 0,
'pool_timeout_seconds': 7,
'pool_recycle_seconds': 600,
}
}
}
)
)
await postgresql.PostgreSQLDatabaseManager(ap).initialize()
assert captured_options == {
'pool_size': 24,
'max_overflow': 0,
'pool_timeout': 7,
'pool_recycle': 600,
'pool_pre_ping': True,
}
@pytest.mark.asyncio
async def test_cloud_postgresql_manager_applies_bounded_server_timeouts(monkeypatch) -> None:
captured_options = None
def create_engine(_url, **options):
nonlocal captured_options
captured_options = options
return object()
monkeypatch.setattr(postgresql.sqlalchemy_asyncio, 'create_async_engine', create_engine)
ap = SimpleNamespace(
instance_config=SimpleNamespace(
data={
'database': {
'postgresql': {
'statement_timeout_ms': 45_000,
'lock_timeout_ms': 4_000,
'idle_in_transaction_session_timeout_ms': 55_000,
}
}
}
)
)
manager = postgresql.PostgreSQLDatabaseManager(ap)
manager.persistence_mode = 'cloud_runtime'
await manager.initialize()
assert captured_options['connect_args'] == {
'server_settings': {
'statement_timeout': '45000',
'lock_timeout': '4000',
'idle_in_transaction_session_timeout': '55000',
}
}
@pytest.mark.asyncio
@pytest.mark.parametrize(
('name', 'value'),
[
('pool_size', 0),
('pool_size', True),
('max_overflow', -1),
('pool_size', 101),
('max_overflow', 101),
('pool_timeout_seconds', 0),
('pool_timeout_seconds', 301),
('pool_recycle_seconds', '1800'),
('pool_recycle_seconds', 86401),
],
)
async def test_postgresql_manager_rejects_invalid_pool_options(name, value) -> None:
ap = SimpleNamespace(instance_config=SimpleNamespace(data={'database': {'postgresql': {name: value}}}))
with pytest.raises(ValueError, match=rf'database\.postgresql\.{name}'):
await postgresql.PostgreSQLDatabaseManager(ap).initialize()
@pytest.mark.asyncio
async def test_postgresql_manager_rejects_combined_pool_capacity_above_hard_ceiling() -> None:
ap = SimpleNamespace(
instance_config=SimpleNamespace(
data={
'database': {
'postgresql': {
'pool_size': 60,
'max_overflow': 41,
}
}
}
)
)
with pytest.raises(ValueError, match=r'pool_size \+ max_overflow'):
await postgresql.PostgreSQLDatabaseManager(ap).initialize()
@pytest.mark.asyncio
@pytest.mark.parametrize(
('name', 'value'),
[
('statement_timeout_ms', 0),
('statement_timeout_ms', 300_001),
('lock_timeout_ms', 60_001),
('idle_in_transaction_session_timeout_ms', True),
('idle_in_transaction_session_timeout_ms', 300_001),
],
)
async def test_cloud_postgresql_manager_rejects_unsafe_server_timeouts(name, value) -> None:
ap = SimpleNamespace(instance_config=SimpleNamespace(data={'database': {'postgresql': {name: value}}}))
with pytest.raises(ValueError, match=rf'database\.postgresql\.{name}'):
manager = postgresql.PostgreSQLDatabaseManager(ap)
manager.persistence_mode = 'cloud_runtime'
await manager.initialize()
@pytest.mark.asyncio
async def test_postgresql_manager_rejects_non_postgresql_url_without_echoing_secret() -> None:
ap = SimpleNamespace(
instance_config=SimpleNamespace(
data={'database': {'postgresql': {'url': 'sqlite:///operator-super-secret.db'}}}
)
)
manager = postgresql.PostgreSQLDatabaseManager(ap)
with pytest.raises(ValueError, match='valid PostgreSQL') as exc_info:
await manager.initialize()
assert 'operator-super-secret' not in str(exc_info.value)
tests/unit_tests/persistence/test_release_migration.py
+155
View File
@@ -0,0 +1,155 @@
from __future__ import annotations
import logging
from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
import sqlalchemy
from langbot.__main__ import _build_parser
from langbot.pkg.persistence import release_migration
from langbot.pkg.persistence.mgr import PersistenceManager, PersistenceMode
def _cloud_config(*, database_use: str = 'postgresql', runtime_user: str = 'langbot_runtime') -> dict:
return {
'database': {
'use': database_use,
'postgresql': {
'host': 'runtime-db',
'port': 5432,
'user': runtime_user,
'password': 'runtime-secret',
'database': 'langbot',
},
'cloud_migration': {
'operator_dsn_env': 'TEST_LANGBOT_OPERATOR_DSN',
},
},
'vdb': {
'use': 'pgvector',
'pgvector': {
'use_business_database': True,
'allowed_dimensions': [384, 1536],
},
},
}
def _operator_environ(
*,
user: str = 'langbot_migrator',
database: str = 'langbot',
host: str = 'runtime-db',
port: int = 5432,
) -> dict[str, str]:
return {
'TEST_LANGBOT_OPERATOR_DSN': (f'postgresql://{user}:operator%40secret@{host}:{port}/{database}?sslmode=require')
}
def test_cloud_migration_cli_is_explicit() -> None:
args = _build_parser().parse_args(['migrate', '--cloud'])
assert args.command == 'migrate'
assert args.cloud is True
with pytest.raises(SystemExit) as exc_info:
_build_parser().parse_args(['migrate'])
assert exc_info.value.code == 2
def test_operator_url_is_separate_and_preserves_escaped_secret() -> None:
url = release_migration._operator_database_url(
_cloud_config(),
environ=_operator_environ(),
)
assert url.drivername == 'postgresql+asyncpg'
assert url.username == 'langbot_migrator'
assert url.password == 'operator@secret'
assert url.host == 'runtime-db'
assert url.port == 5432
assert url.database == 'langbot'
assert url.query['ssl'] == 'require'
assert 'sslmode' not in url.query
@pytest.mark.parametrize(
('config', 'environ', 'message'),
[
(_cloud_config(database_use='sqlite'), _operator_environ(), 'SQLite fallback is forbidden'),
(_cloud_config(), {}, 'requires the operator DSN'),
(_cloud_config(), {'TEST_LANGBOT_OPERATOR_DSN': 'not a secret://operator-password'}, 'DSN is invalid'),
(
_cloud_config(),
{'TEST_LANGBOT_OPERATOR_DSN': 'postgresql://operator:secret@runtime-db:not-a-port/langbot'},
'DSN is invalid',
),
(_cloud_config(), _operator_environ(user='langbot_runtime'), 'distinct operator role'),
(_cloud_config(), _operator_environ(database='another_database'), 'configured runtime database'),
(_cloud_config(), _operator_environ(host='other-cluster'), 'runtime PostgreSQL endpoint'),
(_cloud_config(), _operator_environ(port=6432), 'runtime PostgreSQL endpoint'),
],
)
def test_operator_url_rejects_unsafe_configuration(config: dict, environ: dict[str, str], message: str) -> None:
with pytest.raises(release_migration.CloudReleaseMigrationConfigurationError, match=message) as exc_info:
release_migration._operator_database_url(config, environ=environ)
assert 'operator-password' not in str(exc_info.value)
@pytest.mark.asyncio
async def test_release_migration_disposes_operator_engine_on_failure(monkeypatch) -> None:
engine = SimpleNamespace(dispose=AsyncMock())
manager = SimpleNamespace(
db=SimpleNamespace(engine=engine),
initialize=AsyncMock(side_effect=RuntimeError('migration failed')),
shutdown=AsyncMock(side_effect=engine.dispose),
)
def manager_factory(*args, **kwargs):
del args, kwargs
return manager
monkeypatch.setattr(release_migration, 'PersistenceManager', manager_factory)
ap = SimpleNamespace(
instance_config=SimpleNamespace(data=_cloud_config()),
logger=logging.getLogger('release-migration-disposal-test'),
persistence_mgr=None,
)
with pytest.raises(RuntimeError, match='migration failed'):
await release_migration.run_cloud_release_migration(ap, environ=_operator_environ())
assert ap.persistence_mgr is manager
manager.shutdown.assert_awaited_once()
engine.dispose.assert_awaited_once()
@pytest.mark.asyncio
async def test_release_mode_rejects_sqlite_before_schema_changes(tmp_path, monkeypatch) -> None:
from langbot.pkg.persistence import mgr as persistence_mgr_module
from langbot.pkg.persistence.databases.sqlite import SQLiteDatabaseManager
monkeypatch.setattr(persistence_mgr_module.database, 'preregistered_managers', [SQLiteDatabaseManager])
ap = SimpleNamespace(
instance_config=SimpleNamespace(
data={
'database': {
'use': 'sqlite',
'sqlite': {'path': str(tmp_path / 'must-not-migrate.db')},
}
}
),
logger=logging.getLogger('release-migration-sqlite-rejection-test'),
)
manager = PersistenceManager(ap, mode=PersistenceMode.RELEASE_MIGRATION)
with pytest.raises(RuntimeError, match='requires PostgreSQL'):
await manager.initialize()
await manager.get_db_engine().dispose()
engine = sqlalchemy.create_engine(f'sqlite:///{tmp_path / "must-not-migrate.db"}')
try:
assert sqlalchemy.inspect(engine).get_table_names() == []
finally:
engine.dispose()
tests/unit_tests/persistence/test_tenant_uow.py
+1170
View File
File diff suppressed because it is too large Load Diff
tests/unit_tests/pipeline/conftest.py
+24
View File
@@ -17,6 +17,7 @@ from unittest.mock import AsyncMock, Mock
# this, running a stage test in isolation triggers a circular-import error:
# stage.py → core.app → pipelinemgr → stage.stage_class (not yet bound).
import langbot.pkg.pipeline.pipelinemgr # noqa: F401
from langbot.pkg.api.http.context import ExecutionContext
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
import langbot_plugin.api.entities.builtin.platform.message as platform_message
@@ -40,6 +41,14 @@ class MockApplication:
self.query_pool = self._create_mock_query_pool()
self.instance_config = self._create_mock_instance_config()
self.task_mgr = self._create_mock_task_manager()
self.workspace_service = AsyncMock()
self.workspace_service.get_execution_binding = AsyncMock(
return_value=Mock(
instance_uuid='test-instance',
workspace_uuid='test-workspace',
placement_generation=1,
)
)
# Skill manager is optional; PreProcessor only touches it for the
# local-agent runner. None keeps the skill-binding branch inert.
self.skill_mgr = None
@@ -83,6 +92,7 @@ class MockApplication:
query_pool.cached_queries = {}
query_pool.queries = []
query_pool.condition = AsyncMock()
query_pool.remove_query = AsyncMock(return_value=True)
return query_pool
def _create_mock_instance_config(self):
@@ -191,6 +201,9 @@ def sample_query(sample_message_chain, sample_message_event, mock_adapter):
# Use model_construct to bypass Pydantic validation for test purposes
query = pipeline_query.Query.model_construct(
instance_uuid='test-instance',
workspace_uuid='test-workspace',
placement_generation=1,
query_id='test-query-id',
launcher_type=provider_session.LauncherTypes.PERSON,
launcher_id=12345,
@@ -219,6 +232,17 @@ def sample_query(sample_message_chain, sample_message_event, mock_adapter):
resp_message_chain=None,
current_stage_name=None,
)
object.__setattr__(
query,
'_execution_context',
ExecutionContext(
instance_uuid='test-instance',
workspace_uuid='test-workspace',
placement_generation=1,
bot_uuid='test-bot-uuid',
pipeline_uuid='test-pipeline-uuid',
),
)
return query
tests/unit_tests/pipeline/test_aggregator.py
+422 -24
View File
@@ -13,8 +13,11 @@ from __future__ import annotations
import pytest
import asyncio
import contextvars
from contextlib import asynccontextmanager
from unittest.mock import Mock, AsyncMock
from importlib import import_module
from types import SimpleNamespace
from tests.factories import (
FakeApp,
@@ -25,6 +28,49 @@ from tests.factories import (
import langbot_plugin.api.entities.builtin.provider.session as provider_session
from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.pipeline.pool import (
ExecutionContextMismatchError,
ExecutionContextRequiredError,
bind_execution_context,
)
from langbot.pkg.workspace.errors import WorkspaceGenerationMismatchError
def execution_context(
workspace_uuid='workspace-test',
*,
bot_uuid='test-bot',
pipeline_uuid=None,
placement_generation=1,
):
return ExecutionContext(
instance_uuid='instance-test',
workspace_uuid=workspace_uuid,
placement_generation=placement_generation,
bot_uuid=bot_uuid,
pipeline_uuid=pipeline_uuid,
)
def aggregation_key(
context,
*,
launcher_type=provider_session.LauncherTypes.PERSON,
launcher_id=12345,
bot_uuid='test-bot',
pipeline_uuid=None,
):
return (
context.instance_uuid,
context.workspace_uuid,
context.placement_generation,
bot_uuid,
pipeline_uuid,
launcher_type.value,
launcher_id,
)
def get_aggregator_module():
"""Lazy import to avoid circular import issues."""
@@ -36,12 +82,66 @@ def make_aggregator_app():
app = FakeApp()
# Ensure query_pool has add_query method
app.query_pool.add_query = AsyncMock()
async def resolve_context(
context,
*,
bot_uuid,
pipeline_uuid,
query_uuid=None,
):
if context is None:
raise ExecutionContextRequiredError('ExecutionContext required in test')
return bind_execution_context(
context,
bot_uuid=bot_uuid,
pipeline_uuid=pipeline_uuid,
query_uuid=query_uuid,
)
app.query_pool.resolve_execution_context = AsyncMock(side_effect=resolve_context)
# Add pipeline_mgr mock
app.pipeline_mgr = AsyncMock()
app.pipeline_mgr.get_pipeline_by_uuid = AsyncMock(return_value=None)
app.workspace_service = Mock()
app.workspace_service.get_execution_binding = AsyncMock(
return_value=Mock(
instance_uuid='instance-test',
workspace_uuid='workspace-test',
placement_generation=1,
)
)
return app
def enable_aggregation(app, *, delay=10.0):
pipeline = Mock()
pipeline.pipeline_entity.config = {
'trigger': {
'message-aggregation': {
'enabled': True,
'delay': delay,
}
}
}
app.pipeline_mgr.get_pipeline_by_uuid = AsyncMock(return_value=pipeline)
def scoped_message_kwargs(context, *, launcher_id=12345, text='hello'):
chain = text_chain(text)
return {
'execution_context': context,
'bot_uuid': context.bot_uuid,
'launcher_type': provider_session.LauncherTypes.PERSON,
'launcher_id': launcher_id,
'sender_id': launcher_id,
'message_event': friend_message_event(chain),
'message_chain': chain,
'adapter': mock_adapter(),
'pipeline_uuid': context.pipeline_uuid,
}
class TestPendingMessage:
"""Tests for PendingMessage dataclass."""
@@ -54,6 +154,7 @@ class TestPendingMessage:
adapter = mock_adapter()
pending = aggregator.PendingMessage(
execution_context=execution_context(pipeline_uuid='test-pipeline'),
bot_uuid='test-bot',
launcher_type=provider_session.LauncherTypes.PERSON,
launcher_id=12345,
@@ -77,9 +178,14 @@ class TestSessionBuffer:
"""SessionBuffer should be created with correct fields."""
aggregator = get_aggregator_module()
buffer = aggregator.SessionBuffer(session_id='test-session')
context = execution_context()
key = aggregation_key(context)
buffer = aggregator.SessionBuffer(
aggregation_key=key,
execution_context=context,
)
assert buffer.session_id == 'test-session'
assert buffer.aggregation_key == key
assert buffer.messages == []
assert buffer.timer_task is None
assert buffer.last_message_time is not None
@@ -93,6 +199,7 @@ class TestSessionBuffer:
adapter = mock_adapter()
pending = aggregator.PendingMessage(
execution_context=execution_context(),
bot_uuid='test-bot',
launcher_type=provider_session.LauncherTypes.PERSON,
launcher_id=12345,
@@ -103,8 +210,10 @@ class TestSessionBuffer:
pipeline_uuid=None,
)
context = execution_context()
buffer = aggregator.SessionBuffer(
session_id='test-session',
aggregation_key=aggregation_key(context),
execution_context=context,
messages=[pending],
)
@@ -127,7 +236,7 @@ class TestMessageAggregatorInit:
class TestMessageAggregatorSessionId:
"""Tests for session ID generation."""
"""Tests for scoped aggregation key generation."""
def test_session_id_format(self):
"""Session ID should be correctly formatted."""
@@ -136,13 +245,24 @@ class TestMessageAggregatorSessionId:
app = make_aggregator_app()
agg = aggregator.MessageAggregator(app)
session_id = agg._get_session_id(
context = execution_context()
session_id = agg._get_aggregation_key(
context,
bot_uuid='bot-123',
launcher_type=provider_session.LauncherTypes.PERSON,
launcher_id=45678,
pipeline_uuid=None,
)
assert session_id == 'bot-123:person:45678'
assert session_id == (
'instance-test',
'workspace-test',
1,
'bot-123',
None,
'person',
45678,
)
def test_session_id_different_launchers(self):
"""Different launcher types should produce different IDs."""
@@ -151,16 +271,21 @@ class TestMessageAggregatorSessionId:
app = make_aggregator_app()
agg = aggregator.MessageAggregator(app)
person_id = agg._get_session_id(
context = execution_context()
person_id = agg._get_aggregation_key(
context,
bot_uuid='bot',
launcher_type=provider_session.LauncherTypes.PERSON,
launcher_id=123,
pipeline_uuid=None,
)
group_id = agg._get_session_id(
group_id = agg._get_aggregation_key(
context,
bot_uuid='bot',
launcher_type=provider_session.LauncherTypes.GROUP,
launcher_id=123,
pipeline_uuid=None,
)
assert person_id != group_id
@@ -177,7 +302,7 @@ class TestMessageAggregatorConfig:
app = make_aggregator_app()
agg = aggregator.MessageAggregator(app)
enabled, delay = await agg._get_aggregation_config(None)
enabled, delay = await agg._get_aggregation_config(execution_context(), None)
assert enabled == False
assert delay == 1.5
@@ -191,7 +316,10 @@ class TestMessageAggregatorConfig:
app.pipeline_mgr.get_pipeline_by_uuid = AsyncMock(return_value=None)
agg = aggregator.MessageAggregator(app)
enabled, delay = await agg._get_aggregation_config('unknown-pipeline')
enabled, delay = await agg._get_aggregation_config(
execution_context(pipeline_uuid='unknown-pipeline'),
'unknown-pipeline',
)
assert enabled == False
assert delay == 1.5
@@ -217,7 +345,10 @@ class TestMessageAggregatorConfig:
agg = aggregator.MessageAggregator(app)
enabled, delay = await agg._get_aggregation_config('test-pipeline')
enabled, delay = await agg._get_aggregation_config(
execution_context(pipeline_uuid='test-pipeline'),
'test-pipeline',
)
assert enabled == True
assert delay == 2.0
@@ -243,7 +374,10 @@ class TestMessageAggregatorConfig:
agg = aggregator.MessageAggregator(app)
enabled, delay = await agg._get_aggregation_config('test-pipeline')
enabled, delay = await agg._get_aggregation_config(
execution_context(pipeline_uuid='test-pipeline'),
'test-pipeline',
)
assert delay == 1.0 # Clamped to minimum
@@ -268,7 +402,10 @@ class TestMessageAggregatorConfig:
agg = aggregator.MessageAggregator(app)
enabled, delay = await agg._get_aggregation_config('test-pipeline')
enabled, delay = await agg._get_aggregation_config(
execution_context(pipeline_uuid='test-pipeline'),
'test-pipeline',
)
assert delay == 10.0 # Clamped to maximum
@@ -293,7 +430,10 @@ class TestMessageAggregatorConfig:
agg = aggregator.MessageAggregator(app)
enabled, delay = await agg._get_aggregation_config('test-pipeline')
enabled, delay = await agg._get_aggregation_config(
execution_context(pipeline_uuid='test-pipeline'),
'test-pipeline',
)
assert delay == 1.5 # Default
@@ -314,6 +454,7 @@ class TestMessageAggregatorAddMessage:
adapter = mock_adapter()
await agg.add_message(
execution_context=execution_context(),
bot_uuid='test-bot',
launcher_type=provider_session.LauncherTypes.PERSON,
launcher_id=12345,
@@ -353,6 +494,7 @@ class TestMessageAggregatorAddMessage:
adapter = mock_adapter()
await agg.add_message(
execution_context=execution_context(pipeline_uuid='test-pipeline'),
bot_uuid='test-bot',
launcher_type=provider_session.LauncherTypes.PERSON,
launcher_id=12345,
@@ -394,6 +536,7 @@ class TestMessageAggregatorAddMessage:
# Add messages up to MAX_BUFFER_MESSAGES
for i in range(aggregator.MAX_BUFFER_MESSAGES):
await agg.add_message(
execution_context=execution_context(pipeline_uuid='test-pipeline'),
bot_uuid='test-bot',
launcher_type=provider_session.LauncherTypes.PERSON,
launcher_id=12345,
@@ -405,7 +548,14 @@ class TestMessageAggregatorAddMessage:
)
# Buffer should be flushed (empty or no buffer)
session_id = agg._get_session_id('test-bot', provider_session.LauncherTypes.PERSON, 12345)
context = execution_context(pipeline_uuid='test-pipeline')
session_id = agg._get_aggregation_key(
context,
'test-bot',
provider_session.LauncherTypes.PERSON,
12345,
'test-pipeline',
)
assert session_id not in agg.buffers or len(agg.buffers[session_id].messages) == 0
@@ -424,6 +574,7 @@ class TestMessageAggregatorMerge:
adapter = mock_adapter()
pending = aggregator.PendingMessage(
execution_context=execution_context(),
bot_uuid='test-bot',
launcher_type=provider_session.LauncherTypes.PERSON,
launcher_id=12345,
@@ -451,6 +602,7 @@ class TestMessageAggregatorMerge:
adapter = mock_adapter()
pending1 = aggregator.PendingMessage(
execution_context=execution_context(),
bot_uuid='test-bot',
launcher_type=provider_session.LauncherTypes.PERSON,
launcher_id=12345,
@@ -462,6 +614,7 @@ class TestMessageAggregatorMerge:
)
pending2 = aggregator.PendingMessage(
execution_context=execution_context(),
bot_uuid='test-bot',
launcher_type=provider_session.LauncherTypes.PERSON,
launcher_id=12345,
@@ -492,6 +645,7 @@ class TestMessageAggregatorMerge:
adapter = mock_adapter()
pending1 = aggregator.PendingMessage(
execution_context=execution_context(pipeline_uuid='test-pipeline-uuid'),
bot_uuid='test-bot',
launcher_type=provider_session.LauncherTypes.PERSON,
launcher_id=12345,
@@ -504,6 +658,7 @@ class TestMessageAggregatorMerge:
)
pending2 = aggregator.PendingMessage(
execution_context=execution_context(pipeline_uuid='test-pipeline-uuid'),
bot_uuid='test-bot',
launcher_type=provider_session.LauncherTypes.PERSON,
launcher_id=12345,
@@ -532,7 +687,8 @@ class TestMessageAggregatorFlush:
app = make_aggregator_app()
agg = aggregator.MessageAggregator(app)
await agg._flush_buffer('nonexistent-session')
context = execution_context()
await agg._flush_buffer(aggregation_key(context), context)
# Should not call query_pool
assert not app.query_pool.add_query.called
@@ -550,6 +706,7 @@ class TestMessageAggregatorFlush:
adapter = mock_adapter()
pending = aggregator.PendingMessage(
execution_context=execution_context(),
bot_uuid='test-bot',
launcher_type=provider_session.LauncherTypes.PERSON,
launcher_id=12345,
@@ -560,17 +717,57 @@ class TestMessageAggregatorFlush:
pipeline_uuid=None,
)
context = execution_context()
key = aggregation_key(context)
buffer = aggregator.SessionBuffer(
session_id='test-session',
aggregation_key=key,
execution_context=context,
messages=[pending],
)
agg.buffers['test-session'] = buffer
agg.buffers[key] = buffer
await agg._flush_buffer('test-session')
await agg._flush_buffer(key, context)
assert app.query_pool.add_query.called
assert 'test-session' not in agg.buffers
assert key not in agg.buffers
@pytest.mark.asyncio
async def test_flush_drops_buffer_when_placement_generation_is_stale(self):
"""A debounce timer cannot enqueue work after its placement is fenced."""
aggregator = get_aggregator_module()
app = make_aggregator_app()
app.workspace_service.get_execution_binding.side_effect = WorkspaceGenerationMismatchError('stale generation')
agg = aggregator.MessageAggregator(app)
context = execution_context(placement_generation=3)
pending = aggregator.PendingMessage(
execution_context=context,
bot_uuid='test-bot',
launcher_type=provider_session.LauncherTypes.PERSON,
launcher_id=12345,
sender_id=12345,
message_event=friend_message_event(text_chain('stale')),
message_chain=text_chain('stale'),
adapter=mock_adapter(),
pipeline_uuid=None,
)
key = aggregation_key(context)
agg.buffers[key] = aggregator.SessionBuffer(
aggregation_key=key,
execution_context=context,
messages=[pending],
)
with pytest.raises(WorkspaceGenerationMismatchError):
await agg._flush_buffer(key, context)
app.workspace_service.get_execution_binding.assert_awaited_once_with(
'workspace-test',
expected_generation=3,
)
app.query_pool.add_query.assert_not_awaited()
assert key not in agg.buffers
class TestMessageAggregatorFlushAll:
@@ -603,6 +800,7 @@ class TestMessageAggregatorFlushAll:
# Create two buffers
pending1 = aggregator.PendingMessage(
execution_context=execution_context(),
bot_uuid='test-bot',
launcher_type=provider_session.LauncherTypes.PERSON,
launcher_id=12345,
@@ -614,6 +812,7 @@ class TestMessageAggregatorFlushAll:
)
pending2 = aggregator.PendingMessage(
execution_context=execution_context(),
bot_uuid='test-bot',
launcher_type=provider_session.LauncherTypes.PERSON,
launcher_id=67890,
@@ -624,14 +823,213 @@ class TestMessageAggregatorFlushAll:
pipeline_uuid=None,
)
buffer1 = aggregator.SessionBuffer(session_id='session-1', messages=[pending1])
buffer2 = aggregator.SessionBuffer(session_id='session-2', messages=[pending2])
context = execution_context()
key1 = aggregation_key(context, launcher_id=12345)
key2 = aggregation_key(context, launcher_id=67890)
buffer1 = aggregator.SessionBuffer(
aggregation_key=key1,
execution_context=context,
messages=[pending1],
)
buffer2 = aggregator.SessionBuffer(
aggregation_key=key2,
execution_context=context,
messages=[pending2],
)
agg.buffers['session-1'] = buffer1
agg.buffers['session-2'] = buffer2
agg.buffers[key1] = buffer1
agg.buffers[key2] = buffer2
await agg.flush_all()
# Both buffers should be flushed
assert len(agg.buffers) == 0
assert app.query_pool.add_query.call_count == 2
class TestMessageAggregatorWorkspaceIsolation:
"""Regression coverage for fail-closed and cross-workspace behavior."""
@pytest.mark.asyncio
async def test_missing_execution_context_fails_closed(self):
app = make_aggregator_app()
agg = get_aggregator_module().MessageAggregator(app)
kwargs = scoped_message_kwargs(execution_context())
kwargs['execution_context'] = None
with pytest.raises(ExecutionContextRequiredError):
await agg.add_message(**kwargs)
app.query_pool.add_query.assert_not_awaited()
@pytest.mark.asyncio
async def test_same_launcher_in_two_workspaces_uses_separate_buffers(self):
app = make_aggregator_app()
enable_aggregation(app)
agg = get_aggregator_module().MessageAggregator(app)
await agg.add_message(**scoped_message_kwargs(execution_context('workspace-a', pipeline_uuid='test-pipeline')))
await agg.add_message(**scoped_message_kwargs(execution_context('workspace-b', pipeline_uuid='test-pipeline')))
assert len(agg.buffers) == 2
assert {key[1] for key in agg.buffers} == {'workspace-a', 'workspace-b'}
await agg.flush_all()
@pytest.mark.asyncio
async def test_new_buffer_uses_scope_counter_without_global_scan(self):
class NoGlobalIterationDict(dict):
def __iter__(self):
raise AssertionError('aggregation admission scanned all buffers')
def items(self):
raise AssertionError('aggregation admission scanned all buffers')
def values(self):
raise AssertionError('aggregation admission scanned all buffers')
app = make_aggregator_app()
enable_aggregation(app)
agg = get_aggregator_module().MessageAggregator(app)
agg.max_buffers = 2_000
agg.max_buffers_per_workspace = 2_000
existing = {
(
'instance-test',
f'workspace-{index}',
1,
'bot',
'pipeline',
'person',
index,
): object()
for index in range(1_000)
}
agg.buffers = NoGlobalIterationDict(existing)
agg._buffer_counts_by_scope = {key[:3]: 1 for key in existing}
context = execution_context(
'workspace-target',
pipeline_uuid='test-pipeline',
)
await agg.add_message(**scoped_message_kwargs(context))
key = aggregation_key(
context,
pipeline_uuid='test-pipeline',
)
assert key in agg.buffers
assert agg._buffer_counts_by_scope[key[:3]] == 1
timer_task = agg.buffers[key].timer_task
assert timer_task is not None
timer_task.cancel()
await asyncio.gather(timer_task, return_exceptions=True)
await agg._flush_buffer(key, context)
assert key[:3] not in agg._buffer_counts_by_scope
@pytest.mark.asyncio
async def test_same_launcher_in_two_bots_uses_separate_buffers(self):
app = make_aggregator_app()
enable_aggregation(app)
agg = get_aggregator_module().MessageAggregator(app)
await agg.add_message(
**scoped_message_kwargs(execution_context(bot_uuid='bot-a', pipeline_uuid='test-pipeline'))
)
await agg.add_message(
**scoped_message_kwargs(execution_context(bot_uuid='bot-b', pipeline_uuid='test-pipeline'))
)
assert len(agg.buffers) == 2
assert {key[3] for key in agg.buffers} == {'bot-a', 'bot-b'}
await agg.flush_all()
@pytest.mark.asyncio
async def test_timer_receives_exact_captured_execution_context(self, monkeypatch):
app = make_aggregator_app()
enable_aggregation(app)
agg = get_aggregator_module().MessageAggregator(app)
request_value = contextvars.ContextVar('aggregator_request_value', default=None)
token = request_value.set('request-scope')
observed = []
async def delayed_flush(*args):
observed.append((request_value.get(), args[2]))
monkeypatch.setattr(agg, '_delayed_flush', delayed_flush)
context = execution_context(pipeline_uuid='test-pipeline')
try:
await agg.add_message(**scoped_message_kwargs(context))
await asyncio.sleep(0)
finally:
request_value.reset(token)
assert observed == [(None, context)]
await agg.flush_all()
@pytest.mark.asyncio
async def test_delayed_flush_opens_explicit_workspace_uow(self, monkeypatch):
app = make_aggregator_app()
app.persistence_mgr.mode = SimpleNamespace(value='cloud_runtime')
scopes = []
@asynccontextmanager
async def tenant_uow(workspace_uuid):
scopes.append(workspace_uuid)
yield
app.persistence_mgr.tenant_uow = tenant_uow
agg = get_aggregator_module().MessageAggregator(app)
flush = AsyncMock()
monkeypatch.setattr(agg, '_flush_buffer', flush)
context = execution_context('workspace-a', pipeline_uuid='test-pipeline')
key = aggregation_key(context, pipeline_uuid='test-pipeline')
await agg._delayed_flush(key, 0, context)
assert scopes == ['workspace-a']
flush.assert_awaited_once_with(key, context)
@pytest.mark.asyncio
async def test_flush_rejects_context_from_another_workspace(self):
app = make_aggregator_app()
enable_aggregation(app)
agg = get_aggregator_module().MessageAggregator(app)
context_a = execution_context('workspace-a', pipeline_uuid='test-pipeline')
context_b = execution_context('workspace-b', pipeline_uuid='test-pipeline')
await agg.add_message(**scoped_message_kwargs(context_a))
key = next(iter(agg.buffers))
with pytest.raises(ExecutionContextMismatchError):
await agg._flush_buffer(key, context_b)
assert key in agg.buffers
await agg.flush_all()
def test_merge_rejects_messages_from_different_workspaces(self):
app = make_aggregator_app()
agg = get_aggregator_module().MessageAggregator(app)
aggregator = get_aggregator_module()
with pytest.raises(ExecutionContextMismatchError):
agg._merge_messages(
[
aggregator.PendingMessage(**scoped_message_kwargs(execution_context('workspace-a'))),
aggregator.PendingMessage(**scoped_message_kwargs(execution_context('workspace-b'))),
]
)
@pytest.mark.asyncio
async def test_flush_all_preserves_each_workspace_context(self):
app = make_aggregator_app()
enable_aggregation(app)
agg = get_aggregator_module().MessageAggregator(app)
await agg.add_message(**scoped_message_kwargs(execution_context('workspace-a', pipeline_uuid='test-pipeline')))
await agg.add_message(**scoped_message_kwargs(execution_context('workspace-b', pipeline_uuid='test-pipeline')))
await agg.flush_all()
forwarded_workspaces = {
call.kwargs['execution_context'].workspace_uuid for call in app.query_pool.add_query.await_args_list
}
assert forwarded_workspaces == {'workspace-a', 'workspace-b'}
tests/unit_tests/pipeline/test_chat_handler.py
+10
View File
@@ -464,3 +464,13 @@ class TestChatHandlerHelper:
handler = chat.ChatMessageHandler(fake_app)
result = handler.cut_str('first line\nsecond line')
assert '...' in result
def test_response_size_limit_uses_instance_config(self, fake_app):
from langbot_plugin.api.entities.builtin.provider.message import Message
fake_app.instance_config.data['system'] = {'response_limits': {'max_generated_chars': 4}}
chat = get_chat_handler()
handler = chat.ChatMessageHandler(fake_app)
with pytest.raises(RuntimeError, match='configured limit'):
handler._check_response_size(Message(role='assistant', content='12345'))
tests/unit_tests/pipeline/test_chat_session_limit.py
+7 -1
View File
@@ -8,6 +8,7 @@ from unittest.mock import AsyncMock, Mock
import pytest
import yaml
from langbot_plugin.api.entities.builtin.provider import session as provider_session
def _preproc_module():
@@ -43,7 +44,12 @@ def _prompt_preprocessing_context(default_prompt=None, prompt=None):
async def _run_preprocessor(mock_app, sample_query, conversation):
session = SimpleNamespace(launcher_type=sample_query.launcher_type, launcher_id=sample_query.launcher_id)
session = provider_session.Session(
launcher_type=sample_query.launcher_type,
launcher_id=sample_query.launcher_id,
sender_id=sample_query.sender_id,
bot_uuid=sample_query.bot_uuid,
)
mock_app.sess_mgr.get_session = AsyncMock(return_value=session)
mock_app.sess_mgr.get_conversation = AsyncMock(return_value=conversation)
mock_app.plugin_connector.emit_event = AsyncMock(return_value=_prompt_preprocessing_context())
tests/unit_tests/pipeline/test_controller_tenancy.py
+145
View File
@@ -0,0 +1,145 @@
from __future__ import annotations
import asyncio
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, Mock
import pytest
import sqlalchemy as sa
from sqlalchemy.ext.asyncio import create_async_engine
from langbot.pkg.persistence.mgr import PersistenceManager, PersistenceMode
from langbot.pkg.persistence.tenant_uow import PersistenceScopeKind
from langbot.pkg.pipeline.controller import Controller
from langbot.pkg.workspace.errors import WorkspaceGenerationMismatchError
def _prepare_scheduler(mock_app):
query_pool = MagicMock()
query_pool.remove_query = AsyncMock(return_value=True)
query_pool.__aenter__ = AsyncMock(return_value=query_pool)
query_pool.__aexit__ = AsyncMock(return_value=None)
query_pool.condition = SimpleNamespace(notify_all=Mock())
mock_app.query_pool = query_pool
session = SimpleNamespace(_semaphore=SimpleNamespace(release=Mock()))
mock_app.sess_mgr.get_session = AsyncMock(return_value=session)
mock_app.pipeline_mgr = SimpleNamespace(get_pipeline_by_uuid=AsyncMock())
return query_pool, session
@pytest.mark.asyncio
async def test_controller_drops_stale_query_before_pipeline_lookup(
mock_app,
sample_query,
):
query_pool, session = _prepare_scheduler(mock_app)
mock_app.workspace_service.get_execution_binding.side_effect = WorkspaceGenerationMismatchError('stale generation')
controller = Controller(mock_app)
initial_slots = controller.semaphore._value
await controller._process_query(sample_query)
mock_app.workspace_service.get_execution_binding.assert_awaited_once_with(
'test-workspace',
expected_generation=1,
)
mock_app.pipeline_mgr.get_pipeline_by_uuid.assert_not_awaited()
query_pool.remove_query.assert_awaited_once_with(sample_query)
session._semaphore.release.assert_called_once_with()
query_pool.condition.notify_all.assert_called_once_with()
assert controller.semaphore._value == initial_slots
@pytest.mark.asyncio
async def test_cloud_controller_releases_database_connection_during_pipeline_wait(
tmp_path,
mock_app,
sample_query,
):
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "pipeline-short-scope.db"}')
table = sa.Table('pipeline_scope_probe', sa.MetaData(), sa.Column('id', sa.Integer, primary_key=True))
manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
manager.db = SimpleNamespace(get_engine=lambda: engine)
checked_out = 0
def on_checkout(*_args):
nonlocal checked_out
checked_out += 1
def on_checkin(*_args):
nonlocal checked_out
checked_out -= 1
sa.event.listen(engine.sync_engine, 'checkout', on_checkout)
sa.event.listen(engine.sync_engine, 'checkin', on_checkin)
try:
async with engine.begin() as conn:
await conn.run_sync(table.metadata.create_all)
_prepare_scheduler(mock_app)
mock_app.persistence_mgr = manager
pipeline_waiting = asyncio.Event()
release_pipeline = asyncio.Event()
async def get_binding(*_args, **_kwargs):
assert manager.current_scope().kind is PersistenceScopeKind.WORKSPACE
assert manager.current_session() is None
await manager.execute_async(sa.select(table.c.id))
assert manager.current_session() is None
return SimpleNamespace(
instance_uuid='test-instance',
workspace_uuid='test-workspace',
placement_generation=1,
)
async def run_pipeline(_query):
await manager.execute_async(sa.select(table.c.id))
assert manager.current_session() is None
pipeline_waiting.set()
await release_pipeline.wait()
assert manager.current_scope().kind is PersistenceScopeKind.WORKSPACE
assert manager.current_session() is None
runtime_pipeline = SimpleNamespace(run=AsyncMock(side_effect=run_pipeline))
async def get_pipeline(*_args, **_kwargs):
await manager.execute_async(sa.select(table.c.id))
assert manager.current_session() is None
return runtime_pipeline
mock_app.workspace_service.get_execution_binding = AsyncMock(side_effect=get_binding)
mock_app.pipeline_mgr.get_pipeline_by_uuid = AsyncMock(side_effect=get_pipeline)
controller = Controller(mock_app)
task = asyncio.create_task(controller._process_query(sample_query))
await asyncio.wait_for(pipeline_waiting.wait(), timeout=2)
assert checked_out == 0
assert not task.done()
release_pipeline.set()
await asyncio.wait_for(task, timeout=2)
assert checked_out == 0
runtime_pipeline.run.assert_awaited_once_with(sample_query)
finally:
await engine.dispose()
@pytest.mark.asyncio
async def test_controller_revalidates_generation_before_running_pipeline(
mock_app,
sample_query,
):
query_pool, session = _prepare_scheduler(mock_app)
runtime_pipeline = SimpleNamespace(run=AsyncMock())
mock_app.pipeline_mgr.get_pipeline_by_uuid.return_value = runtime_pipeline
controller = Controller(mock_app)
await controller._process_query(sample_query)
mock_app.workspace_service.get_execution_binding.assert_awaited_once_with(
'test-workspace',
expected_generation=1,
)
runtime_pipeline.run.assert_awaited_once_with(sample_query)
query_pool.remove_query.assert_awaited_once_with(sample_query)
session._semaphore.release.assert_called_once_with()
tests/unit_tests/pipeline/test_longtext_image.py
+50
View File
@@ -0,0 +1,50 @@
from types import SimpleNamespace
from unittest.mock import Mock
import pytest
import langbot_plugin.api.entities.builtin.platform.message as platform_message
from langbot.pkg.pipeline.longtext.strategies.image import Text2ImageStrategy
from langbot.pkg.pipeline.longtext.strategies import image
class _WideFont:
def getlength(self, text: str) -> int:
return len(text) * 100
def test_image_strategy_line_split_always_consumes_input():
strategy = Text2ImageStrategy(Mock())
lines = strategy._split_text_lines('abc', 1, _WideFont())
assert lines == ['a', 'b', 'c']
assert ''.join(lines) == 'abc'
def test_image_strategy_numeric_boundaries_are_found_in_linear_order():
strategy = Text2ImageStrategy(Mock())
assert strategy.indexNumber('a12-b12-c345') == [['12', 1], ['12', 5], ['345', 9]]
def test_image_strategy_rejects_unbounded_line_count_before_allocating_canvas(monkeypatch):
strategy = Text2ImageStrategy(Mock())
monkeypatch.setattr(image, '_MAX_TEXT_TO_IMAGE_LINES', 2)
with pytest.raises(ValueError, match='2 lines'):
strategy._split_text_lines('one\ntwo\nthree', 1000, _WideFont())
@pytest.mark.asyncio
async def test_image_strategy_falls_back_to_forward_for_oversized_text(monkeypatch):
app = Mock()
strategy = Text2ImageStrategy(app)
monkeypatch.setattr(image, '_MAX_TEXT_TO_IMAGE_CHARS', 4)
query = SimpleNamespace(adapter=SimpleNamespace(bot_account_id='bot'))
components = await strategy.process('12345', query)
assert len(components) == 1
assert isinstance(components[0], platform_message.Forward)
app.logger.warning.assert_called_once()
tests/unit_tests/pipeline/test_n8nsvapi.py
+16 -1
View File
@@ -37,7 +37,11 @@ _saved_modules = {name: sys.modules.get(name) for name in _import_stubs}
for _name, _stub in _import_stubs.items():
sys.modules[_name] = _stub
try:
from langbot.pkg.provider.runners.n8nsvapi import N8nServiceAPIRunner
from langbot.pkg.provider.runners.n8nsvapi import (
N8nAPIError,
N8nServiceAPIRunner,
_MAX_N8N_RESPONSE_CHARS,
)
finally:
for _name, _original in _saved_modules.items():
if _original is None:
@@ -230,6 +234,17 @@ async def test_plain_json_non_dict_response():
assert chunks[0].content == '["a", "b"]'
@pytest.mark.asyncio
async def test_response_size_is_bounded():
runner = make_runner()
with pytest.raises(N8nAPIError, match='exceeds the runtime limit'):
await collect_chunks(
runner,
[b'x' * (_MAX_N8N_RESPONSE_CHARS + 1)],
)
@pytest.mark.asyncio
async def test_invalid_json_returns_raw_text():
"""Non-JSON response returns raw text as-is."""
tests/unit_tests/pipeline/test_pipeline_service.py
+7 -3
View File
@@ -5,6 +5,9 @@ import pytest
from langbot.pkg.api.http.service.pipeline import PipelineService
WORKSPACE_UUID = 'workspace-a'
@pytest.mark.asyncio
async def test_update_pipeline_filters_protected_fields_without_mutating_input(mock_app):
service = PipelineService(mock_app)
@@ -27,7 +30,7 @@ async def test_update_pipeline_filters_protected_fields_without_mutating_input(m
}
original_pipeline_data = pipeline_data.copy()
await service.update_pipeline('pipeline-uuid', pipeline_data)
await service.update_pipeline(WORKSPACE_UUID, 'pipeline-uuid', pipeline_data)
assert pipeline_data == original_pipeline_data
@@ -36,8 +39,9 @@ async def test_update_pipeline_filters_protected_fields_without_mutating_input(m
assert updated_fields == {'name'}
mock_app.bot_service.update_bot.assert_awaited_once_with(
WORKSPACE_UUID,
'bot-uuid',
{'use_pipeline_name': 'Updated pipeline'},
)
mock_app.pipeline_mgr.remove_pipeline.assert_awaited_once_with('pipeline-uuid')
mock_app.pipeline_mgr.load_pipeline.assert_awaited_once_with(loaded_pipeline)
mock_app.pipeline_mgr.remove_pipeline.assert_awaited_once_with('workspace-a', 'pipeline-uuid')
mock_app.pipeline_mgr.load_pipeline.assert_awaited_once_with('workspace-a', loaded_pipeline)
tests/unit_tests/pipeline/test_pipelinemgr.py
+201 -12
View File
@@ -3,9 +3,23 @@ PipelineManager unit tests
"""
import pytest
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
from importlib import import_module
from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.workspace.entities import WorkspaceExecutionBinding
from langbot.pkg.workspace.errors import WorkspaceGenerationMismatchError, WorkspaceInvariantError
def _context(pipeline_uuid: str = 'test-uuid') -> ExecutionContext:
return ExecutionContext(
instance_uuid='test-instance',
workspace_uuid='test-workspace',
placement_generation=1,
pipeline_uuid=pipeline_uuid,
)
def get_pipelinemgr_module():
return import_module('langbot.pkg.pipeline.pipelinemgr')
@@ -37,6 +51,95 @@ async def test_pipeline_manager_initialize(mock_app):
assert len(manager.pipelines) == 0
@pytest.mark.asyncio
async def test_cloud_startup_reuses_validated_pipeline_binding(mock_app):
class TenantUow:
async def __aenter__(self):
return self
async def __aexit__(self, *_args):
return False
binding = WorkspaceExecutionBinding(
instance_uuid='test-instance',
workspace_uuid='test-workspace',
placement_generation=1,
write_fenced=False,
state='active',
)
pipeline_entity = Mock(
uuid='test-uuid',
workspace_uuid='test-workspace',
stages=[],
config={},
extensions_preferences={},
)
mock_app.persistence_mgr.mode = SimpleNamespace(value='cloud_runtime')
mock_app.persistence_mgr.tenant_uow = lambda _workspace_uuid: TenantUow()
mock_app.persistence_mgr.execute_async = AsyncMock(return_value=Mock(all=Mock(return_value=[pipeline_entity])))
mock_app.workspace_service.list_active_execution_bindings = AsyncMock(return_value=[binding])
mock_app.workspace_service.get_execution_binding = AsyncMock(
side_effect=AssertionError('startup pipeline loader repeated a validated binding lookup')
)
manager = get_pipelinemgr_module().PipelineManager(mock_app)
manager.stage_dict = {}
await manager.load_pipelines_from_db()
assert len(manager.pipelines) == 1
mock_app.workspace_service.get_execution_binding.assert_not_awaited()
def test_generation_advance_prunes_superseded_workspace_pipelines(mock_app):
class NoGlobalIterationDict(dict):
def __iter__(self):
raise AssertionError('generation advance scanned every pipeline')
def items(self):
raise AssertionError('generation advance scanned every pipeline')
def values(self):
raise AssertionError('generation advance scanned every pipeline')
pipelinemgr = get_pipelinemgr_module()
manager = pipelinemgr.PipelineManager(mock_app)
old_context = _context()
next_context = ExecutionContext(
instance_uuid=old_context.instance_uuid,
workspace_uuid=old_context.workspace_uuid,
placement_generation=2,
pipeline_uuid=old_context.pipeline_uuid,
)
old_pipeline = SimpleNamespace(
execution_context=old_context,
workspace_uuid=old_context.workspace_uuid,
placement_generation=old_context.placement_generation,
)
other_pipelines = [
SimpleNamespace(
execution_context=ExecutionContext(
instance_uuid='test-instance',
workspace_uuid=f'workspace-{index}',
placement_generation=1,
pipeline_uuid=f'pipeline-{index}',
),
workspace_uuid=f'workspace-{index}',
placement_generation=1,
)
for index in range(1_000)
]
manager.pipelines = [old_pipeline, *other_pipelines]
manager._observe_execution_context(old_context)
manager._pipelines_by_key = NoGlobalIterationDict(manager._pipelines_by_key)
manager._observe_execution_context(next_context)
manager._pipelines_by_key = dict(manager._pipelines_by_key)
assert manager.pipelines == other_pipelines
with pytest.raises(WorkspaceInvariantError, match='rolled back'):
manager._observe_execution_context(old_context)
@pytest.mark.asyncio
async def test_load_pipeline(mock_app):
"""Test loading a single pipeline"""
@@ -51,11 +154,12 @@ async def test_load_pipeline(mock_app):
# Create test pipeline entity
pipeline_entity = Mock(spec=persistence_pipeline.LegacyPipeline)
pipeline_entity.uuid = 'test-uuid'
pipeline_entity.workspace_uuid = 'test-workspace'
pipeline_entity.stages = []
pipeline_entity.config = {'test': 'config'}
pipeline_entity.extensions_preferences = {'plugins': []}
await manager.load_pipeline(pipeline_entity)
await manager.load_pipeline(_context(), pipeline_entity)
assert len(manager.pipelines) == 1
assert manager.pipelines[0].pipeline_entity.uuid == 'test-uuid'
@@ -75,19 +179,20 @@ async def test_get_pipeline_by_uuid(mock_app):
# Create and add test pipeline
pipeline_entity = Mock(spec=persistence_pipeline.LegacyPipeline)
pipeline_entity.uuid = 'test-uuid'
pipeline_entity.workspace_uuid = 'test-workspace'
pipeline_entity.stages = []
pipeline_entity.config = {}
pipeline_entity.extensions_preferences = {'plugins': []}
await manager.load_pipeline(pipeline_entity)
await manager.load_pipeline(_context(), pipeline_entity)
# Test retrieval
result = await manager.get_pipeline_by_uuid('test-uuid')
result = await manager.get_pipeline_by_uuid(_context(), 'test-uuid')
assert result is not None
assert result.pipeline_entity.uuid == 'test-uuid'
# Test non-existent UUID
result = await manager.get_pipeline_by_uuid('non-existent')
result = await manager.get_pipeline_by_uuid(_context('non-existent'), 'non-existent')
assert result is None
@@ -105,15 +210,16 @@ async def test_remove_pipeline(mock_app):
# Create and add test pipeline
pipeline_entity = Mock(spec=persistence_pipeline.LegacyPipeline)
pipeline_entity.uuid = 'test-uuid'
pipeline_entity.workspace_uuid = 'test-workspace'
pipeline_entity.stages = []
pipeline_entity.config = {}
pipeline_entity.extensions_preferences = {'plugins': []}
await manager.load_pipeline(pipeline_entity)
await manager.load_pipeline(_context(), pipeline_entity)
assert len(manager.pipelines) == 1
# Remove pipeline
await manager.remove_pipeline('test-uuid')
await manager.remove_pipeline(_context(), 'test-uuid')
assert len(manager.pipelines) == 0
@@ -143,25 +249,104 @@ async def test_runtime_pipeline_execute(mock_app, sample_query):
# Create pipeline entity
pipeline_entity = Mock(spec=persistence_pipeline.LegacyPipeline)
pipeline_entity.uuid = 'test-pipeline-uuid'
pipeline_entity.workspace_uuid = 'test-workspace'
pipeline_entity.config = sample_query.pipeline_config
pipeline_entity.extensions_preferences = {'plugins': []}
# Create runtime pipeline
runtime_pipeline = pipelinemgr.RuntimePipeline(mock_app, pipeline_entity, [stage_container])
runtime_pipeline = pipelinemgr.RuntimePipeline(
mock_app,
pipeline_entity,
[stage_container],
_context('test-pipeline-uuid'),
)
# Mock plugin connector
event_ctx = Mock()
event_ctx.is_prevented_default = Mock(return_value=False)
mock_app.plugin_connector.emit_event = AsyncMock(return_value=event_ctx)
# Add query to cached_queries to prevent KeyError in finally block
mock_app.query_pool.cached_queries[sample_query.query_id] = sample_query
# Execute pipeline
await runtime_pipeline.run(sample_query)
# Verify stage was called
mock_stage.process.assert_called_once()
mock_app.query_pool.remove_query.assert_awaited_once_with(sample_query)
@pytest.mark.asyncio
async def test_runtime_pipeline_rejects_stale_generation_before_side_effects(
mock_app,
sample_query,
):
pipelinemgr = get_pipelinemgr_module()
persistence_pipeline = get_persistence_pipeline_module()
pipeline_entity = Mock(spec=persistence_pipeline.LegacyPipeline)
pipeline_entity.uuid = 'test-pipeline-uuid'
pipeline_entity.workspace_uuid = 'test-workspace'
pipeline_entity.config = sample_query.pipeline_config
pipeline_entity.extensions_preferences = {'plugins': []}
runtime_pipeline = pipelinemgr.RuntimePipeline(
mock_app,
pipeline_entity,
[],
_context('test-pipeline-uuid'),
)
mock_app.workspace_service.get_execution_binding.side_effect = WorkspaceGenerationMismatchError('stale generation')
with pytest.raises(WorkspaceGenerationMismatchError):
await runtime_pipeline.run(sample_query)
mock_app.plugin_connector.emit_event.assert_not_awaited()
sample_query.adapter.reply_message.assert_not_awaited()
sample_query.adapter.reply_message_chunk.assert_not_awaited()
@pytest.mark.asyncio
async def test_runtime_pipeline_revalidates_after_awaited_stage(
mock_app,
sample_query,
):
pipelinemgr = get_pipelinemgr_module()
stage = get_stage_module()
persistence_pipeline = get_persistence_pipeline_module()
entities = get_entities_module()
pipeline_entity = Mock(spec=persistence_pipeline.LegacyPipeline)
pipeline_entity.uuid = 'test-pipeline-uuid'
pipeline_entity.workspace_uuid = 'test-workspace'
pipeline_entity.config = sample_query.pipeline_config
pipeline_entity.extensions_preferences = {'plugins': []}
result = entities.StageProcessResult(
result_type=entities.ResultType.CONTINUE,
new_query=sample_query,
user_notice='must not be sent',
console_notice='',
debug_notice='',
error_notice='',
)
async def stage_process(*_args):
mock_app.workspace_service.get_execution_binding.side_effect = WorkspaceGenerationMismatchError(
'generation changed during stage'
)
return result
mock_stage = Mock(spec=stage.PipelineStage)
mock_stage.process = Mock(side_effect=stage_process)
runtime_pipeline = pipelinemgr.RuntimePipeline(
mock_app,
pipeline_entity,
[pipelinemgr.StageInstContainer(inst_name='TestStage', inst=mock_stage)],
_context('test-pipeline-uuid'),
)
with pytest.raises(WorkspaceGenerationMismatchError):
await runtime_pipeline._execute_from_stage(0, sample_query)
sample_query.adapter.reply_message.assert_not_awaited()
sample_query.adapter.reply_message_chunk.assert_not_awaited()
def test_runtime_pipeline_prefers_local_agent_mcp_resources(mock_app):
@@ -170,6 +355,8 @@ def test_runtime_pipeline_prefers_local_agent_mcp_resources(mock_app):
persistence_pipeline = get_persistence_pipeline_module()
pipeline_entity = Mock(spec=persistence_pipeline.LegacyPipeline)
pipeline_entity.uuid = 'test-uuid'
pipeline_entity.workspace_uuid = 'test-workspace'
pipeline_entity.config = {
'ai': {
'local-agent': {
@@ -183,7 +370,7 @@ def test_runtime_pipeline_prefers_local_agent_mcp_resources(mock_app):
'mcp_resource_agent_read_enabled': True,
}
runtime_pipeline = pipelinemgr.RuntimePipeline(mock_app, pipeline_entity, [])
runtime_pipeline = pipelinemgr.RuntimePipeline(mock_app, pipeline_entity, [], _context())
assert runtime_pipeline.mcp_resource_attachments == [{'server_uuid': 'srv-new', 'uri': 'file:///new.md'}]
assert runtime_pipeline.mcp_resource_agent_read_enabled is False
@@ -195,13 +382,15 @@ def test_runtime_pipeline_falls_back_to_extension_mcp_resources(mock_app):
persistence_pipeline = get_persistence_pipeline_module()
pipeline_entity = Mock(spec=persistence_pipeline.LegacyPipeline)
pipeline_entity.uuid = 'test-uuid'
pipeline_entity.workspace_uuid = 'test-workspace'
pipeline_entity.config = {'ai': {'local-agent': {}}}
pipeline_entity.extensions_preferences = {
'mcp_resources': [{'server_uuid': 'srv-old', 'uri': 'file:///old.md'}],
'mcp_resource_agent_read_enabled': False,
}
runtime_pipeline = pipelinemgr.RuntimePipeline(mock_app, pipeline_entity, [])
runtime_pipeline = pipelinemgr.RuntimePipeline(mock_app, pipeline_entity, [], _context())
assert runtime_pipeline.mcp_resource_attachments == [{'server_uuid': 'srv-old', 'uri': 'file:///old.md'}]
assert runtime_pipeline.mcp_resource_agent_read_enabled is False
tests/unit_tests/pipeline/test_pool.py
+212 -14
View File
@@ -6,10 +6,52 @@ Tests query management, ID generation, and async context handling.
from __future__ import annotations
import uuid
from types import SimpleNamespace
import pytest
from unittest.mock import Mock, patch
from langbot.pkg.pipeline.pool import QueryPool
from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.pipeline.pool import (
ExecutionContextMismatchError,
ExecutionContextRequiredError,
QueryNotFoundError,
QueryPool,
QueryPoolCapacityError,
get_query_execution_context,
)
TEST_CONTEXT = ExecutionContext(
instance_uuid='instance-test',
workspace_uuid='workspace-test',
placement_generation=1,
)
def oss_pool():
"""Build the explicit singleton resolver used by the OSS compatibility path."""
return QueryPool(singleton_context_resolver=lambda: TEST_CONTEXT)
async def add_scoped_mock_query(pool, context, *, bot_uuid='bot-a'):
"""Create a Query through the real pool while keeping SDK details mocked."""
query = Mock()
query.bot_uuid = bot_uuid
query.pipeline_uuid = None
query.query_id = pool.query_id_counter
with patch('langbot.pkg.pipeline.pool.pipeline_query.Query', return_value=query):
return await pool.add_query(
bot_uuid=bot_uuid,
launcher_type=Mock(),
launcher_id='launcher-1',
sender_id='sender-1',
message_event=Mock(),
message_chain=Mock(),
adapter=Mock(),
execution_context=context,
)
pytestmark = pytest.mark.asyncio
@@ -39,7 +81,7 @@ class TestQueryPoolAddQuery:
async def test_add_query_adds_query_with_id(self):
"""add_query creates, stores, and caches a Query with the correct ID."""
pool = QueryPool()
pool = oss_pool()
# Mock Query creation
mock_query = Mock()
@@ -62,12 +104,12 @@ class TestQueryPoolAddQuery:
# Query is added to list and cache
assert pool.queries[0] is mock_query
assert pool.cached_queries[0] is mock_query
assert pool.cached_queries[('workspace-test', mock_query.query_uuid)] is mock_query
assert mock_query.query_id == 0
async def test_add_query_increments_counter(self):
"""Each add_query increments the counter."""
pool = QueryPool()
pool = oss_pool()
mock_query1 = Mock()
mock_query1.query_id = 0
@@ -103,7 +145,7 @@ class TestQueryPoolAddQuery:
async def test_add_query_appends_to_list(self):
"""Query is appended to queries list."""
pool = QueryPool()
pool = oss_pool()
mock_query = Mock()
mock_query.query_id = 0
@@ -126,7 +168,7 @@ class TestQueryPoolAddQuery:
async def test_add_query_caches_query(self):
"""Query is cached by query_id."""
pool = QueryPool()
pool = oss_pool()
mock_query = Mock()
mock_query.query_id = 0
@@ -144,12 +186,13 @@ class TestQueryPoolAddQuery:
adapter=Mock(),
)
assert 0 in pool.cached_queries
assert pool.cached_queries[0] is mock_query
cache_key = ('workspace-test', mock_query.query_uuid)
assert cache_key in pool.cached_queries
assert pool.cached_queries[cache_key] is mock_query
async def test_add_query_with_pipeline_uuid(self):
"""Query can have pipeline_uuid set."""
pool = QueryPool()
pool = oss_pool()
mock_query = Mock()
mock_query.query_id = 0
@@ -175,7 +218,7 @@ class TestQueryPoolAddQuery:
async def test_add_query_sets_routed_by_rule_variable(self):
"""Query has _routed_by_rule variable."""
pool = QueryPool()
pool = oss_pool()
mock_query = Mock()
mock_query.query_id = 0
@@ -201,7 +244,7 @@ class TestQueryPoolAddQuery:
async def test_add_query_notifier_condition(self):
"""add_query notifies waiting consumers."""
pool = QueryPool()
pool = oss_pool()
mock_query = Mock()
mock_query.query_id = 0
@@ -237,7 +280,7 @@ class TestQueryPoolContext:
async def test_aenter_acquires_lock(self):
"""__aenter__ acquires the pool lock."""
pool = QueryPool()
pool = oss_pool()
async with pool as p:
# Lock is acquired
@@ -260,7 +303,7 @@ class TestQueryPoolEdgeCases:
async def test_multiple_queries_cached_correctly(self):
"""Multiple queries are cached separately."""
pool = QueryPool()
pool = oss_pool()
mock_queries = []
for i in range(5):
@@ -287,4 +330,159 @@ class TestQueryPoolEdgeCases:
# Each query is cached by its ID
for i in range(5):
assert pool.cached_queries[i] is mock_queries[i]
query = mock_queries[i]
assert pool.cached_queries[('workspace-test', query.query_uuid)] is query
class TestQueryPoolWorkspaceIsolation:
"""Regression coverage for trusted scope and scoped cache indexes."""
async def test_add_query_requires_execution_context_by_default(self):
with pytest.raises(ExecutionContextRequiredError):
await QueryPool().add_query(
bot_uuid='bot-a',
launcher_type=Mock(),
launcher_id='launcher-1',
sender_id='sender-1',
message_event=Mock(),
message_chain=Mock(),
adapter=Mock(),
)
async def test_serialized_scope_fields_are_not_trusted_context(self):
forged_query = SimpleNamespace(
instance_uuid='instance-test',
workspace_uuid='workspace-test',
placement_generation=1,
bot_uuid='bot-a',
pipeline_uuid=None,
query_uuid='forged-query',
)
with pytest.raises(ExecutionContextRequiredError):
get_query_execution_context(forged_query)
async def test_query_lookup_is_workspace_scoped(self):
pool = QueryPool()
query = await add_scoped_mock_query(pool, TEST_CONTEXT)
uuid.UUID(query.query_uuid)
assert await pool.get_query('workspace-test', query.query_uuid) is query
assert await pool.get_query('workspace-other', query.query_uuid) is None
assert await pool.get_query_by_legacy_id('workspace-test', 0) is query
assert await pool.get_query_by_legacy_id('workspace-other', 0) is None
with pytest.raises(QueryNotFoundError):
await pool.require_query('workspace-other', query.query_uuid)
async def test_cache_separates_same_opaque_id_between_workspaces(self, monkeypatch):
fixed_uuid = uuid.UUID('11111111-1111-4111-8111-111111111111')
monkeypatch.setattr('langbot.pkg.pipeline.pool.uuid.uuid4', lambda: fixed_uuid)
pool = QueryPool()
context_a = TEST_CONTEXT
context_b = ExecutionContext(
instance_uuid='instance-test',
workspace_uuid='workspace-other',
placement_generation=1,
)
query_a = await add_scoped_mock_query(pool, context_a)
query_b = await add_scoped_mock_query(pool, context_b)
assert query_a.query_uuid == query_b.query_uuid
assert await pool.get_query('workspace-test', query_a.query_uuid) is query_a
assert await pool.get_query('workspace-other', query_b.query_uuid) is query_b
async def test_remove_query_cleans_both_scoped_indexes(self):
pool = QueryPool()
query = await add_scoped_mock_query(pool, TEST_CONTEXT)
assert await pool.remove_query(query) is True
assert await pool.get_query('workspace-test', query.query_uuid) is None
assert await pool.get_query_by_legacy_id('workspace-test', query.query_id) is None
assert await pool.remove_query(query) is False
async def test_context_cannot_substitute_bot_identity(self):
context = ExecutionContext(
instance_uuid='instance-test',
workspace_uuid='workspace-test',
placement_generation=1,
bot_uuid='bot-b',
)
with pytest.raises(ExecutionContextMismatchError):
await add_scoped_mock_query(QueryPool(), context, bot_uuid='bot-a')
async def test_query_counter_is_scoped_by_workspace_and_generation(self):
pool = QueryPool()
workspace_a = TEST_CONTEXT
workspace_b = ExecutionContext(
instance_uuid='instance-test',
workspace_uuid='workspace-other',
placement_generation=1,
)
next_generation = ExecutionContext(
instance_uuid='instance-test',
workspace_uuid='workspace-test',
placement_generation=2,
)
await add_scoped_mock_query(pool, workspace_a)
await add_scoped_mock_query(pool, workspace_a)
await add_scoped_mock_query(pool, workspace_b)
assert pool.get_query_count(workspace_a) == 2
assert pool.get_query_count(workspace_b) == 1
assert pool.get_query_count(next_generation) == 0
assert pool.query_id_counter == 3
async def test_workspace_capacity_discards_oldest_queued_query(self):
pool = QueryPool(max_queries=3, max_queries_per_workspace=2)
first = await add_scoped_mock_query(pool, TEST_CONTEXT)
second = await add_scoped_mock_query(pool, TEST_CONTEXT)
third = await add_scoped_mock_query(pool, TEST_CONTEXT)
assert await pool.get_query(TEST_CONTEXT.workspace_uuid, first.query_uuid) is None
assert await pool.get_query(TEST_CONTEXT.workspace_uuid, second.query_uuid) is second
assert await pool.get_query(TEST_CONTEXT.workspace_uuid, third.query_uuid) is third
assert pool.active_query_count_by_workspace == {TEST_CONTEXT.workspace_uuid: 2}
assert pool.get_dropped_query_count(TEST_CONTEXT) == 1
async def test_capacity_rejects_when_every_query_is_already_running(self):
pool = QueryPool(max_queries=1, max_queries_per_workspace=1)
running = await add_scoped_mock_query(pool, TEST_CONTEXT)
async with pool:
pool.mark_query_running_locked(running)
with pytest.raises(QueryPoolCapacityError):
await add_scoped_mock_query(pool, TEST_CONTEXT)
assert pool.active_query_count_by_workspace == {TEST_CONTEXT.workspace_uuid: 1}
async def test_mark_query_running_keeps_active_indexes_but_removes_queue_entry(self):
pool = QueryPool(max_queries=1, max_queries_per_workspace=1)
running = await add_scoped_mock_query(pool, TEST_CONTEXT)
async with pool:
pool.mark_query_running_locked(running)
assert running not in pool.queries
assert await pool.get_query(TEST_CONTEXT.workspace_uuid, running.query_uuid) is running
assert pool.active_query_count_by_workspace == {TEST_CONTEXT.workspace_uuid: 1}
async def test_historical_workspace_counters_are_bounded(self):
pool = QueryPool(max_queries=2, max_queries_per_workspace=1)
contexts = [
ExecutionContext(
instance_uuid='instance-test',
workspace_uuid=f'workspace-{index}',
placement_generation=1,
)
for index in range(3)
]
for context in contexts:
query = await add_scoped_mock_query(pool, context)
await pool.remove_query(query)
assert len(pool.query_count_by_scope) == 2
assert (contexts[0].instance_uuid, contexts[0].workspace_uuid, 1) not in pool.query_count_by_scope
tests/unit_tests/pipeline/test_preproc.py
+30 -31
View File
@@ -16,6 +16,8 @@ from unittest.mock import AsyncMock, Mock
from importlib import import_module
from types import SimpleNamespace
from langbot_plugin.api.entities.builtin.provider import session as provider_session
from tests.factories import (
FakeApp,
text_query,
@@ -35,6 +37,20 @@ def get_entities_module():
return import_module('langbot.pkg.pipeline.entities')
def make_session(
launcher_type: provider_session.LauncherTypes = provider_session.LauncherTypes.PERSON,
launcher_id: int = 12345,
) -> provider_session.Session:
"""Build a scope-aware Session that matches the shared Query factory."""
return provider_session.Session(
launcher_type=launcher_type,
launcher_id=launcher_id,
sender_id=12345,
bot_uuid='test-bot-uuid',
)
class TestPreProcessorNormalText:
"""Tests for normal text message preprocessing."""
@@ -46,9 +62,7 @@ class TestPreProcessorNormalText:
app = FakeApp()
# Mock session manager to return a session
mock_session = Mock()
mock_session.launcher_type = Mock(value='person')
mock_session.launcher_id = 12345
mock_session = make_session()
app.sess_mgr.get_session = AsyncMock(return_value=mock_session)
# Mock conversation
@@ -92,9 +106,7 @@ class TestPreProcessorNormalText:
preproc = get_preproc_module()
app = FakeApp()
mock_session = Mock()
mock_session.launcher_type = Mock(value='person')
mock_session.launcher_id = 12345
mock_session = make_session()
app.sess_mgr.get_session = AsyncMock(return_value=mock_session)
mock_conversation = Mock()
@@ -132,9 +144,7 @@ class TestPreProcessorEmptyMessage:
entities = get_entities_module()
app = FakeApp()
mock_session = Mock()
mock_session.launcher_type = Mock(value='person')
mock_session.launcher_id = 12345
mock_session = make_session()
app.sess_mgr.get_session = AsyncMock(return_value=mock_session)
mock_conversation = Mock()
@@ -171,9 +181,7 @@ class TestPreProcessorImageSegment:
preproc = get_preproc_module()
app = FakeApp()
mock_session = Mock()
mock_session.launcher_type = Mock(value='person')
mock_session.launcher_id = 12345
mock_session = make_session()
app.sess_mgr.get_session = AsyncMock(return_value=mock_session)
mock_conversation = Mock()
@@ -219,9 +227,7 @@ class TestPreProcessorImageSegment:
preproc = get_preproc_module()
app = FakeApp()
mock_session = Mock()
mock_session.launcher_type = Mock(value='person')
mock_session.launcher_id = 12345
mock_session = make_session()
app.sess_mgr.get_session = AsyncMock(return_value=mock_session)
mock_conversation = Mock()
@@ -258,9 +264,7 @@ class TestPreProcessorModelSelection:
preproc = get_preproc_module()
app = FakeApp()
mock_session = Mock()
mock_session.launcher_type = Mock(value='person')
mock_session.launcher_id = 12345
mock_session = make_session()
app.sess_mgr.get_session = AsyncMock(return_value=mock_session)
mock_conversation = Mock()
@@ -305,9 +309,7 @@ class TestPreProcessorModelSelection:
preproc = get_preproc_module()
app = FakeApp()
mock_session = Mock()
mock_session.launcher_type = Mock(value='person')
mock_session.launcher_id = 12345
mock_session = make_session()
app.sess_mgr.get_session = AsyncMock(return_value=mock_session)
mock_conversation = Mock()
@@ -324,7 +326,7 @@ class TestPreProcessorModelSelection:
mock_fallback = Mock()
mock_fallback.model_entity = Mock(uuid='fallback-uuid', abilities=['func_call'])
async def mock_get_model(uuid):
async def mock_get_model(_context, uuid):
if uuid == 'primary-uuid':
return mock_primary
elif uuid == 'fallback-uuid':
@@ -368,9 +370,7 @@ class TestPreProcessorVariables:
preproc = get_preproc_module()
app = FakeApp()
mock_session = Mock()
mock_session.launcher_type = Mock(value='person')
mock_session.launcher_id = 12345
mock_session = make_session()
app.sess_mgr.get_session = AsyncMock(return_value=mock_session)
mock_conversation = Mock()
@@ -405,9 +405,10 @@ class TestPreProcessorVariables:
preproc = get_preproc_module()
app = FakeApp()
mock_session = Mock()
mock_session.launcher_type = Mock(value='group')
mock_session.launcher_id = 99999
mock_session = make_session(
provider_session.LauncherTypes.GROUP,
99999,
)
app.sess_mgr.get_session = AsyncMock(return_value=mock_session)
mock_conversation = Mock()
@@ -443,9 +444,7 @@ class TestPreProcessorToolSelection:
preproc = get_preproc_module()
app = FakeApp()
mock_session = Mock()
mock_session.launcher_type = Mock(value='person')
mock_session.launcher_id = 12345
mock_session = make_session()
app.sess_mgr.get_session = AsyncMock(return_value=mock_session)
mock_conversation = Mock()
tests/unit_tests/pipeline/test_query_pool.py
+8 -1
View File
@@ -9,6 +9,7 @@ import langbot_plugin.api.definition.abstract.platform.adapter as abstract_platf
import langbot_plugin.api.definition.abstract.platform.event_logger as abstract_platform_logger
from langbot.pkg.pipeline.pool import QueryPool
from langbot.pkg.api.http.context import ExecutionContext
class DummyEventLogger(abstract_platform_logger.AbstractEventLogger):
@@ -64,12 +65,18 @@ async def test_add_query_returns_created_query_and_preserves_side_effects(
adapter=adapter,
pipeline_uuid='test-pipeline-uuid',
routed_by_rule=True,
execution_context=ExecutionContext(
instance_uuid='test-instance-uuid',
workspace_uuid='test-workspace-uuid',
placement_generation=1,
),
)
assert query is query_pool.queries[0]
assert query_pool.cached_queries[0] is query
assert query_pool.cached_queries[('test-workspace-uuid', query.query_uuid)] is query
assert query_pool.query_id_counter == 1
assert query.query_id == 0
assert query.bot_uuid == 'test-bot-uuid'
assert query.pipeline_uuid == 'test-pipeline-uuid'
assert query.workspace_uuid == 'test-workspace-uuid'
assert query.variables == {'_routed_by_rule': True}
tests/unit_tests/pipeline/test_ratelimit.py
+24 -4
View File
@@ -152,8 +152,18 @@ class TestFixedWindowAlgo:
# First request creates container
await algo.require_access(sample_query_with_rate_limit, provider_session.LauncherTypes.PERSON, '12345')
# Key format: 'LauncherTypes.PERSON_12345' (enum string representation)
expected_key = 'LauncherTypes.PERSON_12345'
context = sample_query_with_rate_limit._execution_context
expected_key = ':'.join(
(
context.instance_uuid,
context.workspace_uuid,
str(context.placement_generation),
str(sample_query_with_rate_limit.bot_uuid),
str(sample_query_with_rate_limit.pipeline_uuid),
str(provider_session.LauncherTypes.PERSON),
'12345',
)
)
assert expected_key in algo.containers
container = algo.containers[expected_key]
@@ -191,8 +201,18 @@ class TestFixedWindowAlgo:
for i in range(5):
await algo.require_access(sample_query, provider_session.LauncherTypes.PERSON, 'test')
# Key format: 'LauncherTypes.PERSON_test'
expected_key = 'LauncherTypes.PERSON_test'
context = sample_query._execution_context
expected_key = ':'.join(
(
context.instance_uuid,
context.workspace_uuid,
str(context.placement_generation),
str(sample_query.bot_uuid),
str(sample_query.pipeline_uuid),
str(provider_session.LauncherTypes.PERSON),
'test',
)
)
container = algo.containers[expected_key]
assert window_start in container.records
assert container.records[window_start] == 5
tests/unit_tests/platform/test_aiocqhttp_message_converter.py
+19
View File
@@ -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()
tests/unit_tests/platform/test_botmgr_tenancy.py
+402
View File
@@ -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()
tests/unit_tests/platform/test_dingtalk_adapter.py
+12
View File
@@ -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(
{
tests/unit_tests/platform/test_discord_limits.py
+12
View File
@@ -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)
tests/unit_tests/platform/test_http_bot_tenancy.py
+150
View File
@@ -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'}
tests/unit_tests/platform/test_kook_limits.py
+24
View File
@@ -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)
tests/unit_tests/platform/test_lark_adapter.py
+27
View File
@@ -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'},
tests/unit_tests/platform/test_line_limits.py
+30
View File
@@ -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()
tests/unit_tests/platform/test_matrix_limits.py
+21
View File
@@ -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))
tests/unit_tests/platform/test_openclaw_weixin_client.py
+31
View File
@@ -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,
)
tests/unit_tests/platform/test_openclaw_weixin_tenancy.py
+95
View File
@@ -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)
tests/unit_tests/platform/test_qqofficial_api.py
+23
View File
@@ -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)]
tests/unit_tests/platform/test_routing_rules.py
+16
View File
@@ -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'
tests/unit_tests/platform/test_telegram_adapter.py
+22
View File
@@ -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()
tests/unit_tests/platform/test_webhook_pusher.py
+99
View File
@@ -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
tests/unit_tests/platform/test_websocket_adapter_attachments.py
+51 -24
View File
@@ -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()
tests/unit_tests/platform/test_websocket_session_isolation.py
+203 -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,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()
tests/unit_tests/platform/test_wechatpad_limits.py
+39
View File
@@ -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)
tests/unit_tests/platform/test_wecom_api_limits.py
+31
View File
@@ -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=')
tests/unit_tests/platform/test_wecombot_template_card.py
+23 -10
View File
@@ -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():
tests/unit_tests/plugin/test_connector_methods.py
+84 -27
View File
@@ -9,11 +9,32 @@ Tests cover:
from __future__ import annotations
from contextlib import nullcontext
from types import SimpleNamespace
import pytest
from unittest.mock import Mock, AsyncMock
from unittest.mock import AsyncMock, Mock
from importlib import import_module
from tests.factories import text_query
from langbot_plugin.entities.io.context import InstallationBinding
from langbot.pkg.api.http.context import ExecutionContext
TEST_EXECUTION_CONTEXT = ExecutionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=1,
)
TEST_INSTALLATION_BINDING = InstallationBinding(
instance_uuid=TEST_EXECUTION_CONTEXT.instance_uuid,
workspace_uuid=TEST_EXECUTION_CONTEXT.workspace_uuid,
placement_generation=TEST_EXECUTION_CONTEXT.placement_generation,
installation_uuid='00000000-0000-4000-8000-000000000001',
runtime_revision=1,
artifact_digest='a' * 64,
)
def get_connector_module():
@@ -29,6 +50,7 @@ def create_mock_app():
mock_app.instance_config.data = {'plugin': {'enable': True}}
mock_app.persistence_mgr = AsyncMock()
mock_app.persistence_mgr.execute_async = AsyncMock()
mock_app.persistence_mgr.tenant_uow = None
return mock_app
@@ -39,7 +61,19 @@ def create_mock_connector():
async def mock_disconnect_callback(conn):
pass
return connector.PluginRuntimeConnector(create_mock_app(), mock_disconnect_callback)
instance = connector.PluginRuntimeConnector(create_mock_app(), mock_disconnect_callback)
instance._execution_context.set(TEST_EXECUTION_CONTEXT)
instance._operation_bindings = AsyncMock(return_value=[TEST_INSTALLATION_BINDING])
instance._target_binding = AsyncMock(return_value=TEST_INSTALLATION_BINDING)
instance._load_workspace_settings = AsyncMock(return_value=[])
instance.require_workspace_context = AsyncMock(side_effect=lambda context: context)
return instance
def configure_handler(connector, runtime_handler):
runtime_handler.installation_scope = Mock(side_effect=lambda _binding: nullcontext())
connector.handler = runtime_handler
return runtime_handler
class TestListPlugins:
@@ -87,7 +121,7 @@ class TestListPlugins:
get_connector_module()
connector = create_mock_connector()
connector.handler = AsyncMock()
configure_handler(connector, AsyncMock())
connector.handler.list_plugins = AsyncMock(
return_value=[{'manifest': {'manifest': {'metadata': {'author': 'test', 'name': 'plugin'}}}}]
)
@@ -96,6 +130,7 @@ class TestListPlugins:
connector.handler.list_plugins.assert_called_once()
assert result == [{'manifest': {'manifest': {'metadata': {'author': 'test', 'name': 'plugin'}}}}]
connector._load_workspace_settings.assert_awaited_once_with(TEST_EXECUTION_CONTEXT)
@pytest.mark.asyncio
async def test_filters_by_component_kinds(self):
@@ -103,7 +138,7 @@ class TestListPlugins:
get_connector_module()
connector = create_mock_connector()
connector.handler = AsyncMock()
configure_handler(connector, AsyncMock())
connector.handler.list_plugins = AsyncMock(
return_value=[
{
@@ -130,7 +165,7 @@ class TestListPlugins:
get_connector_module()
connector = create_mock_connector()
connector.handler = AsyncMock()
configure_handler(connector, AsyncMock())
connector.handler.list_plugins = AsyncMock(
return_value=[
{
@@ -177,7 +212,7 @@ class TestPluginDiagnostics:
'response_sources': response_sources,
}
connector.handler = AsyncMock()
configure_handler(connector, AsyncMock())
connector.handler.emit_event = AsyncMock(side_effect=emit_event_response)
fake_event_ctx = Mock()
@@ -221,7 +256,7 @@ class TestPluginDiagnostics:
],
}
connector.handler = AsyncMock()
configure_handler(connector, AsyncMock())
connector.handler.emit_event = AsyncMock(side_effect=emit_event_response)
fake_event_ctx = Mock()
@@ -244,7 +279,7 @@ class TestPluginDiagnostics:
connector_module.context.EventContext.from_event = original_from_event
connector_module.context.EventContext.model_validate = original_model_validate
assert '_response_sources' not in vars(event_ctx)
assert event_ctx._response_sources == []
assert event_ctx._emitted_plugins == [
{'manifest': {'metadata': {'author': 'tester', 'name': 'demo'}}},
]
@@ -259,7 +294,7 @@ class TestPluginDiagnostics:
mock_app = create_mock_app()
mock_app.instance_config.data = {'plugin': {'enable': False}}
connector = connector_module.PluginRuntimeConnector(mock_app, mock_disconnect)
connector.handler = AsyncMock()
configure_handler(connector, AsyncMock())
await connector.notify_plugin_diagnostic({'code': 'response_delivery_failed'})
@@ -268,7 +303,7 @@ class TestPluginDiagnostics:
@pytest.mark.asyncio
async def test_notify_plugin_diagnostic_is_best_effort(self):
connector = create_mock_connector()
connector.handler = AsyncMock()
configure_handler(connector, AsyncMock())
connector.handler.notify_plugin_diagnostic = AsyncMock(side_effect=RuntimeError('action not found'))
await connector.notify_plugin_diagnostic({'code': 'response_delivery_failed'})
@@ -303,7 +338,7 @@ class TestListKnowledgeEngines:
get_connector_module()
connector = create_mock_connector()
connector.handler = AsyncMock()
configure_handler(connector, AsyncMock())
connector.handler.list_knowledge_engines = AsyncMock(
return_value=[{'plugin_id': 'author/engine', 'name': 'Engine'}]
)
@@ -346,7 +381,7 @@ class TestListParsers:
get_connector_module()
connector = create_mock_connector()
connector.handler = AsyncMock()
configure_handler(connector, AsyncMock())
connector.handler.list_parsers = AsyncMock(
return_value=[{'plugin_id': 'author/parser', 'supported_mime_types': ['text/plain']}]
)
@@ -372,7 +407,7 @@ class TestCallParser:
get_connector_module()
connector = create_mock_connector()
connector.handler = AsyncMock()
configure_handler(connector, AsyncMock())
connector.handler.parse_document = AsyncMock(return_value={'content': 'parsed'})
result = await connector.call_parser(
@@ -399,7 +434,7 @@ class TestRAGMethods:
get_connector_module()
connector = create_mock_connector()
connector.handler = AsyncMock()
configure_handler(connector, AsyncMock())
connector.handler.rag_ingest_document = AsyncMock(return_value={'status': 'success'})
result = await connector.call_rag_ingest('author/engine', {'file': 'test.pdf'})
@@ -413,7 +448,7 @@ class TestRAGMethods:
get_connector_module()
connector = create_mock_connector()
connector.handler = AsyncMock()
configure_handler(connector, AsyncMock())
connector.handler.retrieve_knowledge = AsyncMock(
return_value={
'results': [
@@ -442,7 +477,7 @@ class TestRAGMethods:
get_connector_module()
connector = create_mock_connector()
connector.handler = AsyncMock()
configure_handler(connector, AsyncMock())
connector.handler.get_rag_creation_schema = AsyncMock(return_value={'properties': {'name': {'type': 'string'}}})
result = await connector.get_rag_creation_schema('author/engine')
@@ -456,7 +491,7 @@ class TestRAGMethods:
get_connector_module()
connector = create_mock_connector()
connector.handler = AsyncMock()
configure_handler(connector, AsyncMock())
connector.handler.get_rag_retrieval_schema = AsyncMock(
return_value={'properties': {'top_k': {'type': 'integer'}}}
)
@@ -472,7 +507,7 @@ class TestRAGMethods:
get_connector_module()
connector = create_mock_connector()
connector.handler = AsyncMock()
configure_handler(connector, AsyncMock())
connector.handler.rag_on_kb_create = AsyncMock(return_value={'status': 'ok'})
await connector.rag_on_kb_create('author/engine', 'kb-uuid', {'model': 'test'})
@@ -485,7 +520,7 @@ class TestRAGMethods:
get_connector_module()
connector = create_mock_connector()
connector.handler = AsyncMock()
configure_handler(connector, AsyncMock())
connector.handler.rag_on_kb_delete = AsyncMock(return_value={'status': 'ok'})
await connector.rag_on_kb_delete('author/engine', 'kb-uuid')
@@ -498,7 +533,7 @@ class TestRAGMethods:
get_connector_module()
connector = create_mock_connector()
connector.handler = AsyncMock()
configure_handler(connector, AsyncMock())
connector.handler.rag_delete_document = AsyncMock(return_value=True)
result = await connector.call_rag_delete_document('author/engine', 'doc-uuid', 'kb-uuid')
@@ -592,7 +627,7 @@ class TestGetPluginInfo:
get_connector_module()
connector = create_mock_connector()
connector.handler = AsyncMock()
configure_handler(connector, AsyncMock())
connector.handler.get_plugin_info = AsyncMock(return_value={'manifest': {'metadata': {'name': 'plugin'}}})
result = await connector.get_plugin_info('author', 'plugin')
@@ -605,17 +640,39 @@ class TestSetPluginConfig:
"""Tests for set_plugin_config method."""
@pytest.mark.asyncio
async def test_calls_handler_set_plugin_config(self):
"""Test that handler.set_plugin_config is called."""
async def test_updates_revision_then_applies_desired_state(self):
"""Config changes are fenced by a new runtime revision."""
get_connector_module()
connector = create_mock_connector()
connector.handler = AsyncMock()
connector.handler.set_plugin_config = AsyncMock(return_value={'status': 'ok'})
configure_handler(connector, AsyncMock())
connector.handler.register_installation_binding = Mock()
connector.handler.apply_plugin_installation = AsyncMock(return_value={'state': 'running'})
setting = SimpleNamespace(
installation_uuid=TEST_INSTALLATION_BINDING.installation_uuid,
runtime_revision=1,
artifact_digest=TEST_INSTALLATION_BINDING.artifact_digest,
enabled=True,
install_info={'_artifact_storage': 'tenant_binary_storage_v1'},
)
connector._setting_for_plugin = AsyncMock(return_value=(TEST_EXECUTION_CONTEXT, setting))
connector.ap.persistence_mgr.execute_async = AsyncMock(return_value=SimpleNamespace(rowcount=1))
await connector.set_plugin_config('author', 'plugin', {'setting': 'value'})
connector.handler.set_plugin_config.assert_called_once_with('author', 'plugin', {'setting': 'value'})
applied_binding = connector.handler.apply_plugin_installation.await_args.args[0]
assert applied_binding.runtime_revision == 2
assert applied_binding.installation_uuid == TEST_INSTALLATION_BINDING.installation_uuid
connector.handler.register_installation_binding.assert_called_once_with(
applied_binding,
plugin_author='author',
plugin_name='plugin',
)
connector.handler.apply_plugin_installation.assert_awaited_once_with(
applied_binding,
artifact_package=None,
enabled=True,
)
class TestPingPluginRuntime:
@@ -639,7 +696,7 @@ class TestPingPluginRuntime:
get_connector_module()
connector = create_mock_connector()
connector.handler = AsyncMock()
configure_handler(connector, AsyncMock())
connector.handler.ping = AsyncMock(return_value={'status': 'ok'})
await connector.ping_plugin_runtime()
tests/unit_tests/plugin/test_connector_ping.py
+235 -1
View File
@@ -6,14 +6,34 @@ from unittest.mock import AsyncMock, Mock
import pytest
from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.plugin import connector as connector_module
from langbot.pkg.plugin.connector import PluginRuntimeConnector, PluginRuntimeNotConnectedError
from langbot_plugin.runtime.security import (
PLUGIN_RUNTIME_CONTROL_TOKEN_ENV,
PLUGIN_RUNTIME_CONTROL_TOKEN_HEADER,
)
def make_connector() -> PluginRuntimeConnector:
app = SimpleNamespace(
logger=Mock(),
instance_config=SimpleNamespace(data={'plugin': {'enable': True}, 'space': {'url': ''}}),
instance_config=SimpleNamespace(
data={
'plugin': {
'enable': True,
'worker': {
'max_cpus': 1.0,
'max_memory_mb': 512,
'max_pids': 128,
'max_open_files': 256,
'max_file_size_mb': 512,
'require_hard_limits': False,
},
},
'space': {'url': ''},
}
),
)
return PluginRuntimeConnector(app, AsyncMock())
@@ -57,7 +77,9 @@ async def test_stdio_runtime_connection_does_not_capture_unconsumed_stderr(
monkeypatch: pytest.MonkeyPatch,
):
connector = make_connector()
connector._prepare_connected_runtime = AsyncMock()
created = {}
monkeypatch.setattr(connector_module.constants, 'instance_id', 'instance-a')
class FakeRuntimeHandler:
def __init__(self, connection, disconnect_callback, ap):
@@ -106,6 +128,7 @@ async def test_stdio_runtime_connection_does_not_capture_unconsumed_stderr(
assert created['capture_stderr'] is False
assert connector._connected.is_set()
connector._prepare_connected_runtime.assert_awaited_once()
await connector.aclose()
@@ -115,6 +138,8 @@ async def test_runtime_disconnect_notifies_once_and_clears_handler(
):
disconnect = AsyncMock()
connector = PluginRuntimeConnector(make_connector().ap, disconnect)
connector._prepare_connected_runtime = AsyncMock()
monkeypatch.setattr(connector_module.constants, 'instance_id', 'instance-a')
class FakeRuntimeHandler:
def __init__(self, connection, disconnect_callback, ap):
@@ -165,3 +190,212 @@ async def test_runtime_disconnect_notifies_once_and_clears_handler(
disconnect.assert_awaited_once_with(connector)
assert not hasattr(connector, 'handler')
await connector.aclose()
@pytest.mark.asyncio
async def test_disabled_connector_validates_workspace_without_runtime_handler():
app = SimpleNamespace(
instance_config=SimpleNamespace(data={'plugin': {'enable': False}}),
workspace_service=SimpleNamespace(
get_execution_binding=AsyncMock(
return_value=SimpleNamespace(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=3,
)
)
),
)
connector = PluginRuntimeConnector(app, AsyncMock())
request_context = ExecutionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=3,
)
result = await connector.require_workspace_context(request_context)
assert result == request_context
app.workspace_service.get_execution_binding.assert_awaited_once_with(
'workspace-a',
expected_generation=3,
)
@pytest.mark.asyncio
async def test_enabled_connector_reports_not_connected_after_workspace_validation():
connector = make_connector()
connector.ap.workspace_service = SimpleNamespace(
get_execution_binding=AsyncMock(
return_value=SimpleNamespace(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=3,
)
)
)
request_context = ExecutionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=3,
)
with pytest.raises(PluginRuntimeNotConnectedError, match='Plugin runtime is not connected'):
await connector.require_workspace_context(request_context)
@pytest.mark.asyncio
async def test_oss_connector_resolves_singleton_only_for_legacy_callers():
connector = make_connector()
connector.ap.workspace_service = SimpleNamespace(
get_local_execution_binding=AsyncMock(
return_value=SimpleNamespace(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=3,
)
)
)
assert await connector._current_execution_context() == ExecutionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=3,
)
def test_edition_metadata_cannot_enable_shared_runtime_profile():
connector = make_connector()
connector.ap.instance_config.data['system'] = {'edition': 'cloud'}
assert connector.runtime_profile == 'oss_dev'
def test_closed_deployment_selects_instance_scoped_shared_profile():
app = SimpleNamespace(
instance_config=SimpleNamespace(data={'plugin': {'enable': True}}),
deployment=SimpleNamespace(mode='cloud'),
)
connector = PluginRuntimeConnector(app, AsyncMock())
assert connector.runtime_profile == 'shared'
def test_external_runtime_control_headers_require_strong_secret(monkeypatch):
monkeypatch.delenv(PLUGIN_RUNTIME_CONTROL_TOKEN_ENV, raising=False)
connector = make_connector()
with pytest.raises(PluginRuntimeNotConnectedError, match=PLUGIN_RUNTIME_CONTROL_TOKEN_ENV):
connector._control_headers(allow_generate=False)
def test_local_runtime_control_headers_generate_ephemeral_secret(monkeypatch):
monkeypatch.delenv(PLUGIN_RUNTIME_CONTROL_TOKEN_ENV, raising=False)
connector = make_connector()
headers = connector._control_headers(allow_generate=True)
assert len(headers[PLUGIN_RUNTIME_CONTROL_TOKEN_HEADER]) >= 32
@pytest.mark.asyncio
async def test_oss_legacy_fallback_fails_without_workspace_service():
connector = make_connector()
with pytest.raises(AttributeError):
await connector._current_execution_context()
@pytest.mark.asyncio
async def test_cloud_connector_never_falls_back_to_ghost_local_workspace():
app = SimpleNamespace(
instance_config=SimpleNamespace(data={'plugin': {'enable': True}}),
deployment=SimpleNamespace(mode='cloud'),
)
get_local_binding = AsyncMock()
app.workspace_service = SimpleNamespace(
get_local_execution_binding=get_local_binding,
)
connector = PluginRuntimeConnector(app, AsyncMock())
with pytest.raises(Exception, match='Plugin resource not found'):
await connector._current_execution_context()
get_local_binding.assert_not_awaited()
def test_worker_policy_is_loaded_only_from_instance_configuration():
app = SimpleNamespace(
instance_config=SimpleNamespace(
data={
'plugin': {
'enable': True,
'worker': {
'max_cpus': 1.5,
'max_memory_mb': 768,
'max_pids': 64,
'max_open_files': 128,
'max_file_size_mb': 32,
'max_concurrent_restarts': 2,
'restart_failure_threshold': 12,
'restart_failure_window_seconds': 45,
'restart_circuit_open_seconds': 90,
'require_hard_limits': True,
},
# A plugin-controlled value at any other path is ignored.
'manifest': {'max_memory_mb': 99999},
}
}
)
)
connector = PluginRuntimeConnector(app, AsyncMock())
policy = connector._load_worker_policy()
assert policy.max_cpus == 1.5
assert policy.max_memory_mb == 768
assert policy.max_pids == 64
assert policy.max_open_files == 128
assert policy.max_file_size_mb == 32
assert policy.max_concurrent_restarts == 2
assert policy.restart_failure_threshold == 12
assert policy.restart_failure_window_seconds == 45
assert policy.restart_circuit_open_seconds == 90
assert policy.require_hard_limits is True
@pytest.mark.asyncio
async def test_explicit_cloud_binding_is_revalidated_against_projection():
app = SimpleNamespace(
instance_config=SimpleNamespace(
data={
'plugin': {'enable': True},
}
),
deployment=SimpleNamespace(mode='cloud'),
)
app.workspace_service = SimpleNamespace(
get_execution_binding=AsyncMock(
return_value=SimpleNamespace(
instance_uuid='instance-a',
workspace_uuid='workspace-cloud-a',
placement_generation=7,
)
),
)
connector = PluginRuntimeConnector(app, AsyncMock())
connector.handler = SimpleNamespace()
connector._synchronize_workspace = AsyncMock()
configured = ExecutionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-cloud-a',
placement_generation=7,
)
assert await connector.require_workspace_context(configured) == configured
app.workspace_service.get_execution_binding.assert_awaited_once_with(
'workspace-cloud-a',
expected_generation=7,
)
connector._synchronize_workspace.assert_awaited_once_with(configured)
tests/unit_tests/plugin/test_connector_pure.py
+22
View File
@@ -12,6 +12,7 @@ import zipfile
from types import SimpleNamespace
from unittest.mock import MagicMock
import httpx
import pytest
@@ -123,6 +124,17 @@ class TestExtractDepsMetadata:
# Should find requirements.txt in subdirectory
assert task_context.metadata['deps_total'] == 2
def test_archive_preview_rejects_extreme_compression_ratio(self):
from langbot.pkg.plugin.connector import inspect_plugin_archive_metadata
zip_buffer = io.BytesIO()
with zipfile.ZipFile(zip_buffer, 'w', compression=zipfile.ZIP_DEFLATED) as zf:
zf.writestr('manifest.yaml', 'kind: Plugin\nmetadata: {}\n')
zf.writestr('bomb.py', b'A' * (1024 * 1024))
with pytest.raises(ValueError, match='compression-ratio limit'):
inspect_plugin_archive_metadata(zip_buffer.getvalue())
class TestParsePluginId:
"""Tests for _parse_plugin_id static method."""
@@ -141,3 +153,13 @@ class TestParsePluginId:
with pytest.raises(ValueError):
PluginRuntimeConnector._parse_plugin_id('')
@pytest.mark.asyncio
async def test_marketplace_response_reader_is_bounded():
from langbot.pkg.plugin.connector import _read_httpx_response_limited
response = httpx.Response(200, content=b'oversized')
with pytest.raises(ValueError, match='exceeds'):
await _read_httpx_response_limited(response, max_bytes=4)
tests/unit_tests/plugin/test_connector_reconcile.py
+499
View File
@@ -0,0 +1,499 @@
from __future__ import annotations
import datetime
import hashlib
from contextlib import nullcontext
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
import pytest
from langbot_plugin.entities.io.context import InstallationBinding
from langbot_plugin.runtime.plugin.mgr import PluginInstallSource
from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.plugin.connector import (
PluginInstallationFailedError,
PluginRuntimeConnector,
)
def connection_result_connector(execute_async: AsyncMock) -> PluginRuntimeConnector:
app = SimpleNamespace(
instance_config=SimpleNamespace(data={'plugin': {'enable': True}}),
deployment=SimpleNamespace(mode='cloud'),
persistence_mgr=SimpleNamespace(
tenant_uow=None,
execute_async=execute_async,
),
logger=Mock(),
)
return PluginRuntimeConnector(app, AsyncMock())
def execution_binding(workspace_uuid: str, generation: int = 1) -> SimpleNamespace:
return SimpleNamespace(
instance_uuid='instance-a',
workspace_uuid=workspace_uuid,
placement_generation=generation,
)
def plugin_setting(
workspace_suffix: str,
artifact_digest: str,
*,
durable: bool = True,
) -> SimpleNamespace:
return SimpleNamespace(
plugin_author='author',
plugin_name=f'plugin-{workspace_suffix}',
installation_uuid=f'00000000-0000-4000-8000-0000000000{workspace_suffix}',
runtime_revision=1,
artifact_digest=artifact_digest,
enabled=True,
priority=0,
created_at=datetime.datetime(2026, 1, 1),
install_source='local',
install_info={'_artifact_storage': 'tenant_binary_storage_v1'} if durable else {},
)
def runtime_handler(
*,
missing_artifacts: list[str] | None = None,
failed_installations: list[dict[str, str]] | None = None,
) -> SimpleNamespace:
return SimpleNamespace(
register_installation_binding=Mock(),
unregister_installation_binding=Mock(),
reconcile_plugin_installations=AsyncMock(
return_value={
'applied': [],
'removed': [],
'missing_artifacts': missing_artifacts or [],
'failed_installations': failed_installations or [],
}
),
apply_plugin_installation=AsyncMock(return_value={'state': 'starting'}),
installation_scope=Mock(side_effect=lambda _binding: nullcontext()),
list_plugins=AsyncMock(return_value=[]),
)
def shared_connector(
projected_bindings: list[list[SimpleNamespace]],
settings: dict[str, list[SimpleNamespace]],
) -> PluginRuntimeConnector:
async def get_execution_binding(workspace_uuid: str, *, expected_generation: int | None = None):
for binding_set in projected_bindings:
for binding in binding_set:
if binding.workspace_uuid == workspace_uuid:
assert expected_generation in (None, binding.placement_generation)
return binding
raise AssertionError(f'unexpected Workspace {workspace_uuid}')
app = SimpleNamespace(
instance_config=SimpleNamespace(data={'plugin': {'enable': True}}),
deployment=SimpleNamespace(mode='cloud'),
workspace_service=SimpleNamespace(
list_active_execution_bindings=AsyncMock(side_effect=projected_bindings),
get_execution_binding=AsyncMock(side_effect=get_execution_binding),
),
persistence_mgr=SimpleNamespace(tenant_uow=None),
logger=Mock(),
)
connector = PluginRuntimeConnector(app, AsyncMock())
connector._load_workspace_settings = AsyncMock(side_effect=lambda context: settings[context.workspace_uuid])
return connector
@pytest.mark.asyncio
async def test_shared_reconnect_replays_two_workspaces_and_removes_missing_projection():
binding_a = execution_binding('workspace-a')
binding_b = execution_binding('workspace-b')
setting_a = plugin_setting('01', 'a' * 64)
setting_b = plugin_setting('02', 'b' * 64)
connector = shared_connector(
[[binding_a, binding_b], [binding_a]],
{'workspace-a': [setting_a], 'workspace-b': [setting_b]},
)
first_handler = runtime_handler()
connector.handler = first_handler
await connector._prepare_connected_runtime()
first_desired = first_handler.reconcile_plugin_installations.await_args.args[0]
assert {state.binding.workspace_uuid for state in first_desired} == {'workspace-a', 'workspace-b'}
second_handler = runtime_handler()
connector.handler = second_handler
await connector._prepare_connected_runtime()
second_desired = second_handler.reconcile_plugin_installations.await_args.args[0]
assert [state.binding.workspace_uuid for state in second_desired] == ['workspace-a']
second_handler.unregister_installation_binding.assert_called_once_with(first_desired[1].binding)
assert set(connector._workspace_installations) == {'workspace-a'}
assert set(connector._known_desired_states) == {setting_a.installation_uuid}
@pytest.mark.asyncio
async def test_empty_projected_workspaces_do_not_retain_installation_sets():
binding_a = execution_binding('workspace-a')
binding_b = execution_binding('workspace-b')
connector = shared_connector(
[[binding_a, binding_b]],
{'workspace-a': [], 'workspace-b': []},
)
connector.handler = runtime_handler()
await connector._prepare_connected_runtime()
assert connector._workspace_installations == {}
assert connector._known_desired_states == {}
connector.handler.reconcile_plugin_installations.assert_awaited_once_with(())
@pytest.mark.asyncio
async def test_fresh_shared_runtime_cache_replays_persisted_local_package():
package = b'local-lbpkg-bytes'
digest = hashlib.sha256(package).hexdigest()
binding = execution_binding('workspace-a')
setting = plugin_setting('01', digest)
connector = shared_connector([[binding]], {'workspace-a': [setting]})
connector.handler = runtime_handler(missing_artifacts=[setting.installation_uuid])
connector._load_artifact_package = AsyncMock(return_value=package)
await connector._prepare_connected_runtime()
desired = connector.handler.reconcile_plugin_installations.await_args.args[0][0]
connector._load_artifact_package.assert_awaited_once()
connector.handler.apply_plugin_installation.assert_awaited_once_with(
desired.binding,
artifact_package=package,
enabled=True,
)
@pytest.mark.asyncio
async def test_oss_upgrade_keeps_legacy_data_plugins_when_no_lbpkg_was_backfilled():
binding = execution_binding('workspace-a')
setting = plugin_setting('01', hashlib.sha256(b'legacy-installation').hexdigest(), durable=False)
legacy_plugin = {
'debug': False,
'manifest': {'manifest': {'metadata': {'author': 'author', 'name': 'plugin-01'}}},
'components': [],
}
app = SimpleNamespace(
instance_config=SimpleNamespace(data={'plugin': {'enable': True}}),
deployment=SimpleNamespace(mode='oss'),
workspace_service=SimpleNamespace(get_local_execution_binding=AsyncMock(return_value=binding)),
persistence_mgr=SimpleNamespace(tenant_uow=None),
logger=Mock(),
)
connector = PluginRuntimeConnector(app, AsyncMock())
connector.handler = runtime_handler(missing_artifacts=[setting.installation_uuid])
connector.handler.list_plugins = AsyncMock(return_value=[legacy_plugin])
connector.handler.apply_plugin_installation = AsyncMock(return_value={'state': 'artifact_missing'})
connector._load_workspace_settings = AsyncMock(return_value=[setting])
connector._load_artifact_package = AsyncMock(return_value=None)
await connector._prepare_connected_runtime()
plugins = await connector.list_plugins()
assert plugins == [legacy_plugin]
assert connector.handler.list_plugins.await_count >= 2
connector._load_artifact_package.assert_awaited_once()
assert not hasattr(connector.handler, 'delete_plugin')
@pytest.mark.asyncio
async def test_local_install_persists_verified_package_before_runtime_apply():
package = b'local-lbpkg-bytes'
digest = hashlib.sha256(package).hexdigest()
execution_context = ExecutionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=1,
)
binding = InstallationBinding(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=1,
installation_uuid='00000000-0000-4000-8000-000000000001',
runtime_revision=1,
artifact_digest=digest,
)
app = SimpleNamespace(
instance_config=SimpleNamespace(data={'plugin': {'enable': True}}),
deployment=SimpleNamespace(mode='cloud'),
logger=Mock(),
)
connector = PluginRuntimeConnector(app, AsyncMock())
connector.handler = runtime_handler()
connector._current_execution_context = AsyncMock(return_value=execution_context)
connector._inspect_plugin_package = Mock(return_value=('author', 'plugin'))
connector._store_artifact_package = AsyncMock()
connector._persist_installation_package = AsyncMock(return_value=(binding, None, False))
connector._wait_for_installed_plugin_ready = AsyncMock()
await connector.install_plugin(
PluginInstallSource.LOCAL,
{'plugin_file': package},
)
connector._store_artifact_package.assert_awaited_once_with(execution_context, digest, package)
connector.handler.apply_plugin_installation.assert_awaited_once_with(
binding,
artifact_package=package,
enabled=True,
)
@pytest.mark.asyncio
@pytest.mark.parametrize(('remaining_references', 'statement_count'), [(1, 1), (0, 2)])
async def test_artifact_cleanup_is_reference_counted_within_workspace(
remaining_references: int,
statement_count: int,
):
app = SimpleNamespace(
instance_config=SimpleNamespace(data={'plugin': {'enable': True}}),
deployment=SimpleNamespace(mode='cloud'),
)
connector = PluginRuntimeConnector(app, AsyncMock())
execution_context = ExecutionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=1,
)
statements = []
async def execute(statement):
statements.append(statement)
return SimpleNamespace(scalar_one=lambda: remaining_references)
await connector._delete_artifact_if_unreferenced(
execution_context,
'a' * 64,
execute=execute,
)
assert len(statements) == statement_count
@pytest.mark.asyncio
async def test_artifact_store_and_load_accept_connection_scalar_results():
package = b'persisted-lbpkg'
digest = hashlib.sha256(package).hexdigest()
execution_context = ExecutionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=1,
)
execute_async = AsyncMock(return_value=SimpleNamespace(scalar_one_or_none=lambda: package))
connector = connection_result_connector(execute_async)
await connector._store_artifact_package(execution_context, digest, package)
loaded = await connector._load_artifact_package(execution_context, digest)
assert loaded == package
assert execute_async.await_count == 2
@pytest.mark.asyncio
async def test_workspace_settings_accept_connection_mapping_results():
created_at = datetime.datetime(2026, 1, 1)
row = {
'workspace_uuid': 'workspace-a',
'plugin_author': 'author',
'plugin_name': 'plugin',
'installation_uuid': '00000000-0000-4000-8000-000000000001',
'artifact_digest': 'a' * 64,
'runtime_revision': 2,
'enabled': True,
'priority': 3,
'config': {'key': 'value'},
'install_source': 'local',
'install_info': {'_artifact_storage': 'tenant_binary_storage_v1'},
'created_at': created_at,
'updated_at': created_at,
}
mapped_result = SimpleNamespace(mappings=lambda: SimpleNamespace(all=lambda: [row]))
connector = connection_result_connector(AsyncMock(return_value=mapped_result))
settings = await connector._load_workspace_settings(
ExecutionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=1,
)
)
assert len(settings) == 1
assert settings[0].installation_uuid == row['installation_uuid']
assert settings[0].runtime_revision == 2
assert settings[0].created_at == created_at
@pytest.mark.asyncio
async def test_installation_update_accepts_connection_row_result():
old_digest = 'a' * 64
new_digest = 'b' * 64
existing = SimpleNamespace(
installation_uuid='00000000-0000-4000-8000-000000000001',
runtime_revision=4,
artifact_digest=old_digest,
install_info={'_artifact_storage': 'tenant_binary_storage_v1'},
)
execute_async = AsyncMock(
side_effect=[
SimpleNamespace(first=lambda: existing),
SimpleNamespace(rowcount=1),
]
)
connector = connection_result_connector(execute_async)
execution_context = ExecutionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=7,
)
binding, previous_digest, previous_was_durable = await connector._persist_installation_package(
execution_context,
plugin_author='author',
plugin_name='plugin',
install_source=PluginInstallSource.LOCAL,
install_info={},
artifact_digest=new_digest,
)
assert binding.installation_uuid == existing.installation_uuid
assert binding.runtime_revision == 5
assert binding.artifact_digest == new_digest
assert previous_digest == old_digest
assert previous_was_durable is True
@pytest.mark.asyncio
async def test_apply_dependency_failure_raises_stable_observable_error():
binding = execution_binding('workspace-a')
setting = plugin_setting('01', 'a' * 64)
connector = shared_connector([[binding]], {'workspace-a': [setting]})
connector.handler = runtime_handler()
desired = (
await connector._load_workspace_desired_states(
ExecutionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=1,
)
)
)[0]
connector.handler.apply_plugin_installation.return_value = {
'installation_uuid': setting.installation_uuid,
'state': 'failed',
'error_code': 'dependency_prepare_failed',
'message': 'Plugin dependency installer exited with code 1',
}
with pytest.raises(PluginInstallationFailedError) as exc_info:
await connector._apply_desired_state(desired)
error = exc_info.value
assert error.installation_uuid == setting.installation_uuid
assert error.error_code == 'dependency_prepare_failed'
assert '[dependency_prepare_failed]' in str(error)
assert connector._installation_failures[setting.installation_uuid] == {
'installation_uuid': setting.installation_uuid,
'error_code': 'dependency_prepare_failed',
'message': 'Plugin dependency installer exited with code 1',
}
connector.ap.logger.error.assert_called_once()
@pytest.mark.asyncio
async def test_shared_reconcile_records_one_failure_without_blocking_other_state():
binding_a = execution_binding('workspace-a')
binding_b = execution_binding('workspace-b')
setting_a = plugin_setting('01', 'a' * 64)
setting_b = plugin_setting('02', 'b' * 64)
failure = {
'installation_uuid': setting_a.installation_uuid,
'error_code': 'dependency_prepare_failed',
'message': 'Plugin dependency installer exited with code 1',
}
connector = shared_connector(
[[binding_a, binding_b]],
{'workspace-a': [setting_a], 'workspace-b': [setting_b]},
)
connector.handler = runtime_handler(failed_installations=[failure])
await connector._prepare_connected_runtime()
desired = connector.handler.reconcile_plugin_installations.await_args.args[0]
assert {item.binding.installation_uuid for item in desired} == {
setting_a.installation_uuid,
setting_b.installation_uuid,
}
assert connector._installation_failures == {
setting_a.installation_uuid: failure,
}
assert set(connector._known_desired_states) == {
setting_a.installation_uuid,
setting_b.installation_uuid,
}
connector.ap.logger.error.assert_called_once_with(
'Plugin installation %s failed during reconcile [%s]: %s',
setting_a.installation_uuid,
'dependency_prepare_failed',
'Plugin dependency installer exited with code 1',
)
@pytest.mark.asyncio
async def test_missing_artifact_repair_adds_dependency_failure_and_continues():
package_a = b'package-a'
package_b = b'package-b'
binding = execution_binding('workspace-a')
setting_a = plugin_setting('01', hashlib.sha256(package_a).hexdigest())
setting_b = plugin_setting('02', hashlib.sha256(package_b).hexdigest())
connector = shared_connector(
[[binding]],
{'workspace-a': [setting_a, setting_b]},
)
connector.handler = runtime_handler(
missing_artifacts=[
setting_a.installation_uuid,
setting_b.installation_uuid,
]
)
connector._load_artifact_package = AsyncMock(side_effect=[package_a, package_b])
connector.handler.apply_plugin_installation.side_effect = [
{
'installation_uuid': setting_a.installation_uuid,
'state': 'failed',
'error_code': 'dependency_prepare_failed',
'message': 'Plugin dependency installer exited with code 1',
},
{'installation_uuid': setting_b.installation_uuid, 'state': 'starting'},
]
result = await connector.reconcile_projected_workspaces(
[
ExecutionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=1,
)
]
)
assert connector.handler.apply_plugin_installation.await_count == 2
assert result['failed_installations'] == [
{
'installation_uuid': setting_a.installation_uuid,
'error_code': 'dependency_prepare_failed',
'message': 'Plugin dependency installer exited with code 1',
}
]
assert setting_a.installation_uuid in connector._installation_failures
assert setting_b.installation_uuid not in connector._installation_failures
tests/unit_tests/plugin/test_connector_static.py
+7
View File
@@ -53,3 +53,10 @@ class TestParsePluginId:
author, name = connector.PluginRuntimeConnector._parse_plugin_id('lang-bot/my_rag_engine')
assert author == 'lang-bot'
assert name == 'my_rag_engine'
def test_runtime_id_is_stable_across_core_restarts(monkeypatch):
connector = get_connector_module()
monkeypatch.setattr(connector.constants, 'instance_id', 'instance-a')
assert connector.PluginRuntimeConnector._build_runtime_id() == 'instance-a:plugin-runtime'
tests/unit_tests/plugin/test_github_install_security.py
+174
View File
@@ -0,0 +1,174 @@
from __future__ import annotations
import httpx
import pytest
from langbot.pkg.plugin import connector as connector_module
from .test_connector_methods import create_mock_connector
TRUSTED_LEGACY_REQUEST = {
'owner': 'langbot-app',
'repo': 'demo-plugin',
'release_tag': 'v1.0.0',
'asset_url': 'https://github.com/langbot-app/demo-plugin/releases/download/v1.0.0/demo.lbpkg',
}
def _patch_client(monkeypatch, handler):
real_async_client = httpx.AsyncClient
observed: list[dict] = []
def client_factory(*args, **kwargs):
observed.append(dict(kwargs))
return real_async_client(
transport=httpx.MockTransport(handler),
follow_redirects=kwargs.get('follow_redirects', False),
trust_env=kwargs.get('trust_env', True),
timeout=kwargs.get('timeout'),
)
monkeypatch.setattr(connector_module.httpx, 'AsyncClient', client_factory)
return observed
@pytest.mark.parametrize(
'asset_url',
[
'http://127.0.0.1/internal.lbpkg',
'https://169.254.169.254/latest/meta-data',
'https://github.com@127.0.0.1/internal.lbpkg',
'https://evil.example/langbot-app/demo-plugin/releases/download/v1.0.0/demo.lbpkg',
],
)
@pytest.mark.asyncio
async def test_github_install_rejects_internal_or_untrusted_asset_url_before_network(
monkeypatch,
asset_url,
):
connector = create_mock_connector()
monkeypatch.setattr(
connector_module.httpx,
'AsyncClient',
lambda *_args, **_kwargs: pytest.fail('network client must not be created'),
)
with pytest.raises(ValueError, match='GitHub release asset URL'):
await connector._download_github_package(
{**TRUSTED_LEGACY_REQUEST, 'asset_url': asset_url},
None,
)
@pytest.mark.asyncio
async def test_github_asset_id_is_resolved_server_side_and_redirect_escape_is_rejected(
monkeypatch,
):
connector = create_mock_connector()
def handler(request: httpx.Request) -> httpx.Response:
if request.url.path.endswith('/releases/42'):
return httpx.Response(
200,
json={
'id': 42,
'tag_name': 'v1.0.0',
'assets': [{'id': 99, 'size': 128, 'state': 'uploaded'}],
},
)
if request.url.path.endswith('/releases/assets/99'):
return httpx.Response(
302,
headers={'location': 'http://169.254.169.254/latest/meta-data'},
)
raise AssertionError(f'unexpected request: {request.url}')
observed = _patch_client(monkeypatch, handler)
with pytest.raises(ValueError, match='untrusted host'):
await connector._download_github_package(
{
'owner': 'langbot-app',
'repo': 'demo-plugin',
'release_tag': 'v1.0.0',
'release_id': 42,
'asset_id': 99,
'asset_url': 'https://attacker.invalid/ignored',
},
None,
)
assert len(observed) == 1
assert observed[0]['trust_env'] is False
assert observed[0]['follow_redirects'] is False
@pytest.mark.asyncio
async def test_github_download_rejects_oversized_content_length(monkeypatch):
connector = create_mock_connector()
monkeypatch.setattr(connector_module, '_GITHUB_PLUGIN_DOWNLOAD_MAX_BYTES', 8)
def handler(_request: httpx.Request) -> httpx.Response:
return httpx.Response(200, headers={'content-length': '9'}, content=b'')
_patch_client(monkeypatch, handler)
with pytest.raises(ValueError, match='10 MiB download limit'):
await connector._download_github_package(TRUSTED_LEGACY_REQUEST, None)
@pytest.mark.asyncio
async def test_github_download_counts_chunked_stream_bytes(monkeypatch):
connector = create_mock_connector()
monkeypatch.setattr(connector_module, '_GITHUB_PLUGIN_DOWNLOAD_MAX_BYTES', 8)
class ChunkedBody(httpx.AsyncByteStream):
async def __aiter__(self):
yield b'1234'
yield b'56789'
def handler(_request: httpx.Request) -> httpx.Response:
return httpx.Response(200, stream=ChunkedBody())
_patch_client(monkeypatch, handler)
with pytest.raises(ValueError, match='10 MiB download limit'):
await connector._download_github_package(TRUSTED_LEGACY_REQUEST, None)
@pytest.mark.asyncio
async def test_github_asset_id_download_allows_github_object_redirect(monkeypatch):
connector = create_mock_connector()
def handler(request: httpx.Request) -> httpx.Response:
if request.url.path.endswith('/releases/42'):
return httpx.Response(
200,
json={
'id': 42,
'tag_name': 'v1.0.0',
'assets': [{'id': 99, 'size': 7, 'state': 'uploaded'}],
},
)
if request.url.path.endswith('/releases/assets/99'):
return httpx.Response(
302,
headers={
'location': 'https://release-assets.githubusercontent.com/github-production-release-asset/demo'
},
)
if request.url.host == 'release-assets.githubusercontent.com':
return httpx.Response(200, content=b'package')
raise AssertionError(f'unexpected request: {request.url}')
_patch_client(monkeypatch, handler)
package = await connector._download_github_package(
{
'owner': 'langbot-app',
'repo': 'demo-plugin',
'release_tag': 'v1.0.0',
'release_id': 42,
'asset_id': 99,
},
None,
)
assert package == b'package'
tests/unit_tests/plugin/test_handler.py
+49 -1
View File
@@ -10,13 +10,61 @@ from unittest.mock import AsyncMock, MagicMock, Mock
import pytest
from langbot_plugin.entities.io.actions.enums import PluginToRuntimeAction
from langbot_plugin.entities.io.context import ActionContext, InstallationBinding
def make_handler(app):
"""Create a RuntimeConnectionHandler with mocked external connection."""
from langbot.pkg.plugin.handler import RuntimeConnectionHandler
return RuntimeConnectionHandler(Mock(), AsyncMock(return_value=True), app)
workspace_context = ActionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=1,
)
app.workspace_service = SimpleNamespace(
get_execution_binding=AsyncMock(
return_value=SimpleNamespace(
instance_uuid=workspace_context.instance_uuid,
workspace_uuid=workspace_context.workspace_uuid,
placement_generation=workspace_context.placement_generation,
)
)
)
runtime_handler = RuntimeConnectionHandler(
Mock(),
AsyncMock(return_value=True),
app,
)
installation_binding = InstallationBinding(
**workspace_context.model_dump(exclude_none=True),
installation_uuid='00000000-0000-4000-8000-000000000001',
runtime_revision=1,
artifact_digest='a' * 64,
)
runtime_handler.register_installation_binding(
installation_binding,
plugin_author='test-author',
plugin_name='test-plugin',
)
runtime_handler._current_action_context.set(installation_binding)
query_pool = getattr(app, 'query_pool', None)
if query_pool is not None and hasattr(query_pool, 'cached_queries'):
def scoped_query(query):
if query is not None:
query.instance_uuid = workspace_context.instance_uuid
query.workspace_uuid = workspace_context.workspace_uuid
query.placement_generation = workspace_context.placement_generation
return query
query_pool.get_query = AsyncMock(
side_effect=lambda workspace_uuid, query_uuid: scoped_query(query_pool.cached_queries.get(query_uuid))
)
query_pool.get_query_by_legacy_id = AsyncMock(
side_effect=lambda workspace_uuid, query_id: scoped_query(query_pool.cached_queries.get(query_id))
)
return runtime_handler
class TestHandlerQueryVariables:
tests/unit_tests/plugin/test_handler_actions.py
+276 -59
View File
@@ -8,13 +8,80 @@ from unittest.mock import AsyncMock, Mock
import pytest
from langbot_plugin.entities.io.actions.enums import PluginToRuntimeAction, RuntimeToLangBotAction
from langbot_plugin.entities.io.context import ActionContext, InstallationBinding
from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.storage.mgr import StorageMgr
def make_handler(app):
TEST_EXECUTION_CONTEXT = ExecutionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=1,
)
def canonical_binary_key(owner_type: str, owner: str, key: str) -> str:
return StorageMgr.canonical_binary_storage_key(
TEST_EXECUTION_CONTEXT,
owner_type=owner_type,
owner=owner,
key=key,
)
def make_handler(app, workspace_context: ActionContext | None = None):
"""Create a RuntimeConnectionHandler with mocked external connection."""
from langbot.pkg.plugin.handler import RuntimeConnectionHandler
return RuntimeConnectionHandler(Mock(), AsyncMock(return_value=True), app)
workspace_context = workspace_context or ActionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=1,
)
app.workspace_service = SimpleNamespace(
get_execution_binding=AsyncMock(
return_value=SimpleNamespace(
instance_uuid=workspace_context.instance_uuid,
workspace_uuid=workspace_context.workspace_uuid,
placement_generation=workspace_context.placement_generation,
)
)
)
runtime_handler = RuntimeConnectionHandler(
Mock(),
AsyncMock(return_value=True),
app,
)
installation_binding = InstallationBinding(
**workspace_context.model_dump(exclude_none=True),
installation_uuid='00000000-0000-4000-8000-000000000001',
runtime_revision=1,
artifact_digest='a' * 64,
)
runtime_handler.register_installation_binding(
installation_binding,
plugin_author='test-author',
plugin_name='test-plugin',
)
runtime_handler._current_action_context.set(installation_binding)
query_pool = getattr(app, 'query_pool', None)
if query_pool is not None and hasattr(query_pool, 'cached_queries'):
def scoped_query(query):
if query is not None:
query.instance_uuid = workspace_context.instance_uuid
query.workspace_uuid = workspace_context.workspace_uuid
query.placement_generation = workspace_context.placement_generation
return query
query_pool.get_query = AsyncMock(
side_effect=lambda workspace_uuid, query_uuid: scoped_query(query_pool.cached_queries.get(query_uuid))
)
query_pool.get_query_by_legacy_id = AsyncMock(
side_effect=lambda workspace_uuid, query_id: scoped_query(query_pool.cached_queries.get(query_id))
)
return runtime_handler
def make_result(first_item=None):
@@ -34,6 +101,8 @@ class TestRagRerankAction:
def app(self):
mock_app = Mock()
mock_app.model_mgr = Mock()
mock_app.persistence_mgr = Mock()
mock_app.persistence_mgr.execute_async = AsyncMock(return_value=make_result(SimpleNamespace(uuid='rerank-1')))
mock_app.logger = Mock()
return mock_app
@@ -63,12 +132,16 @@ class TestRagRerankAction:
assert response.code == 0
assert response.data['results'] == [{'index': 1, 'relevance_score': 0.9}]
app.model_mgr.get_rerank_model_by_uuid.assert_awaited_once_with('rerank-1')
app.model_mgr.get_rerank_model_by_uuid.assert_awaited_once_with(
TEST_EXECUTION_CONTEXT,
'rerank-1',
)
provider.invoke_rerank.assert_awaited_once_with(
model=rerank_model,
query='hello',
documents=['a', 'b'],
extra_args={'return_documents': False},
execution_context=TEST_EXECUTION_CONTEXT,
)
@pytest.mark.asyncio
@@ -101,13 +174,10 @@ class TestInitializePluginSettings:
return mock_app
@pytest.mark.asyncio
async def test_creates_new_setting_when_not_exists(self, app):
"""New plugin settings use default enabled, priority and config values."""
async def test_rejects_desired_installation_when_setting_not_exists(self, app):
"""A desired-state worker cannot create an unowned Core setting row."""
runtime_handler = make_handler(app)
app.persistence_mgr.execute_async.side_effect = [
make_result(),
Mock(),
]
app.persistence_mgr.execute_async.return_value = make_result()
response = await runtime_handler.actions[RuntimeToLangBotAction.INITIALIZE_PLUGIN_SETTINGS.value](
{
@@ -118,33 +188,23 @@ class TestInitializePluginSettings:
}
)
assert response.code == 0
assert app.persistence_mgr.execute_async.await_count == 2
insert_params = compiled_params(app.persistence_mgr.execute_async.await_args_list[1].args[0])
assert insert_params == {
'plugin_author': 'test-author',
'plugin_name': 'test-plugin',
'install_source': 'local',
'install_info': {'path': '/test'},
'enabled': True,
'priority': 0,
'config': {},
}
assert response.code != 0
assert 'Plugin installation setting was not found' in response.message
app.persistence_mgr.execute_async.assert_awaited_once()
@pytest.mark.asyncio
async def test_inherits_values_from_existing_setting(self, app):
"""Existing settings are replaced while preserving user-controlled values."""
async def test_existing_desired_setting_remains_core_owned(self, app):
"""Runtime initialization validates identity without rewriting Core state."""
runtime_handler = make_handler(app)
existing_setting = SimpleNamespace(
enabled=False,
priority=5,
config={'key': 'value'},
installation_uuid='00000000-0000-4000-8000-000000000001',
runtime_revision=1,
artifact_digest='a' * 64,
)
app.persistence_mgr.execute_async.side_effect = [
make_result(existing_setting),
Mock(),
Mock(),
]
app.persistence_mgr.execute_async.return_value = make_result(existing_setting)
response = await runtime_handler.actions[RuntimeToLangBotAction.INITIALIZE_PLUGIN_SETTINGS.value](
{
@@ -156,13 +216,7 @@ class TestInitializePluginSettings:
)
assert response.code == 0
assert app.persistence_mgr.execute_async.await_count == 3
insert_params = compiled_params(app.persistence_mgr.execute_async.await_args_list[2].args[0])
assert insert_params['enabled'] is False
assert insert_params['priority'] == 5
assert insert_params['config'] == {'key': 'value'}
assert insert_params['install_source'] == 'github'
assert insert_params['install_info'] == {'repo': 'author/name'}
app.persistence_mgr.execute_async.assert_awaited_once()
class TestSetBinaryStorage:
@@ -203,7 +257,7 @@ class TestSetBinaryStorage:
)
assert response.code != 0
assert '2048 > 1024 bytes' in response.message
assert '1024-byte limit' in response.message
app.persistence_mgr.execute_async.assert_not_awaited()
@pytest.mark.asyncio
@@ -218,7 +272,12 @@ class TestSetBinaryStorage:
assert response.code == 0
assert app.persistence_mgr.execute_async.await_count == 2
insert_params = compiled_params(app.persistence_mgr.execute_async.await_args_list[1].args[0])
assert insert_params['unique_key'] == 'plugin:test-owner:test-key'
assert insert_params['workspace_uuid'] == 'workspace-a'
assert insert_params['unique_key'] == canonical_binary_key(
'plugin',
'test-author/test-plugin',
'test-key',
)
assert insert_params['value'] == b'x' * 512
@pytest.mark.asyncio
@@ -231,7 +290,15 @@ class TestSetBinaryStorage:
assert response.code == 0
assert app.persistence_mgr.execute_async.await_count == 2
select_params = compiled_params(app.persistence_mgr.execute_async.await_args_list[0].args[0])
update_params = compiled_params(app.persistence_mgr.execute_async.await_args_list[1].args[0])
expected_key = canonical_binary_key(
'plugin',
'test-author/test-plugin',
'test-key',
)
assert expected_key in select_params.values()
assert expected_key in update_params.values()
assert update_params['value'] == b'new'
@pytest.mark.asyncio
@@ -245,21 +312,25 @@ class TestSetBinaryStorage:
)
assert response.code != 0
assert '10485761 > 10485760 bytes' in response.message
assert '10485760' in response.message
app.persistence_mgr.execute_async.assert_not_awaited()
@pytest.mark.asyncio
async def test_negative_limit_disables_size_check(self, app):
"""Negative max_value_bytes allows values larger than the normal default."""
async def test_negative_limit_falls_back_to_bounded_default(self, app, monkeypatch):
"""Negative max_value_bytes cannot disable the process memory boundary."""
import langbot.pkg.plugin.handler as handler_module
runtime_handler = make_handler(app)
app.instance_config.data['plugin']['binary_storage']['max_value_bytes'] = -1
monkeypatch.setattr(handler_module, '_DEFAULT_BINARY_STORAGE_VALUE_BYTES', 1024)
response = await runtime_handler.actions[RuntimeToLangBotAction.SET_BINARY_STORAGE.value](
self.payload(b'x' * 2048)
)
assert response.code == 0
assert app.persistence_mgr.execute_async.await_count == 2
assert response.code != 0
assert '1024-byte limit' in response.message
app.persistence_mgr.execute_async.assert_not_awaited()
@pytest.mark.asyncio
async def test_zero_limit_rejects_non_empty_values(self, app):
@@ -285,26 +356,18 @@ class TestGetPluginSettings:
return mock_app
@pytest.mark.asyncio
async def test_returns_defaults_when_setting_not_found(self, app):
"""Default plugin settings are returned when no persisted row exists."""
async def test_rejects_desired_installation_when_setting_not_found(self, app):
"""A desired-state worker cannot synthesize settings for a missing row."""
runtime_handler = make_handler(app)
app.persistence_mgr.execute_async.return_value = make_result()
response = await runtime_handler.actions[RuntimeToLangBotAction.GET_PLUGIN_SETTINGS.value](
{
'plugin_author': 'test-author',
'plugin_name': 'test-plugin',
}
)
assert response.code == 0
assert response.data == {
'enabled': True,
'priority': 0,
'plugin_config': {},
'install_source': 'local',
'install_info': {},
}
with pytest.raises(ValueError, match='Plugin installation setting was not found'):
await runtime_handler.actions[RuntimeToLangBotAction.GET_PLUGIN_SETTINGS.value](
{
'plugin_author': 'test-author',
'plugin_name': 'test-plugin',
}
)
@pytest.mark.asyncio
async def test_returns_actual_values_when_setting_exists(self, app):
@@ -316,6 +379,9 @@ class TestGetPluginSettings:
config={'custom': 'config'},
install_source='github',
install_info={'repo': 'test/repo'},
installation_uuid='00000000-0000-4000-8000-000000000001',
runtime_revision=1,
artifact_digest='a' * 64,
)
app.persistence_mgr.execute_async.return_value = make_result(setting)
@@ -333,9 +399,93 @@ class TestGetPluginSettings:
'plugin_config': {'custom': 'config'},
'install_source': 'github',
'install_info': {'repo': 'test/repo'},
'installation_uuid': '00000000-0000-4000-8000-000000000001',
'runtime_revision': 1,
'artifact_digest': 'a' * 64,
}
class TestGetConfigFile:
"""Plugin config files remain bound to the trusted runtime placement."""
WORKSPACE_A = '11111111-1111-4111-8111-111111111111'
WORKSPACE_B = '22222222-2222-4222-8222-222222222222'
@pytest.fixture
def app(self):
mock_app = Mock()
mock_app.persistence_mgr = Mock()
mock_app.persistence_mgr.execute_async = AsyncMock()
mock_app.storage_mgr = StorageMgr(mock_app)
mock_app.storage_mgr.storage_provider = SimpleNamespace(
load_bounded=AsyncMock(return_value=b'plugin config bytes')
)
mock_app.logger = Mock()
return mock_app
@staticmethod
def object_key(
*,
workspace_uuid: str = WORKSPACE_A,
placement_generation: int = 1,
owner_type: str = 'plugin_config',
) -> str:
return StorageMgr.scoped_object_key(
ExecutionContext(
instance_uuid='instance-a',
workspace_uuid=workspace_uuid,
placement_generation=placement_generation,
),
owner_type=owner_type,
owner=workspace_uuid,
key='config.json',
)
async def invoke(self, app, file_key: str):
runtime_handler = make_handler(
app,
ActionContext(
instance_uuid='instance-a',
workspace_uuid=self.WORKSPACE_A,
placement_generation=1,
),
)
app.persistence_mgr.execute_async.return_value = make_result(
SimpleNamespace(config={'uploaded_file': file_key})
)
return await runtime_handler.actions[PluginToRuntimeAction.GET_CONFIG_FILE.value]({'file_key': file_key})
@pytest.mark.asyncio
async def test_loads_config_file_from_same_workspace_generation(self, app):
file_key = self.object_key()
response = await self.invoke(app, file_key)
assert response.code == 0
assert base64.b64decode(response.data['file_base64']) == b'plugin config bytes'
app.storage_mgr.storage_provider.load_bounded.assert_awaited_once_with(
file_key,
max_bytes=10 * 1024 * 1024,
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
'file_key',
[
object_key.__func__(workspace_uuid=WORKSPACE_B),
object_key.__func__(placement_generation=2),
object_key.__func__(owner_type='upload_document'),
],
ids=['other-workspace', 'stale-generation', 'wrong-owner-type'],
)
async def test_rejects_config_file_outside_trusted_scope(self, app, file_key):
response = await self.invoke(app, file_key)
assert response.code != 0
assert 'Failed to load config file' in response.message
app.storage_mgr.storage_provider.load_bounded.assert_not_awaited()
class TestGetBinaryStorage:
"""Tests for get_binary_storage action handler."""
@@ -364,6 +514,16 @@ class TestGetBinaryStorage:
assert response.data == {
'value_base64': base64.b64encode(b'test binary content').decode('utf-8'),
}
statement_params = compiled_params(app.persistence_mgr.execute_async.await_args.args[0])
assert 'workspace-a' in statement_params.values()
assert (
canonical_binary_key(
'plugin',
'test-author/test-plugin',
'test-key',
)
in statement_params.values()
)
@pytest.mark.asyncio
async def test_returns_error_when_not_found(self, app):
@@ -383,6 +543,63 @@ class TestGetBinaryStorage:
assert 'Storage with key test-key not found' in response.message
class TestDeleteAndListBinaryStorage:
"""Delete/list remain fenced to the trusted canonical owner scope."""
@pytest.fixture
def app(self):
mock_app = Mock()
mock_app.persistence_mgr = Mock()
mock_app.persistence_mgr.execute_async = AsyncMock()
return mock_app
@pytest.mark.asyncio
async def test_delete_uses_workspace_and_canonical_unique_key(self, app):
runtime_handler = make_handler(app)
response = await runtime_handler.actions[RuntimeToLangBotAction.DELETE_BINARY_STORAGE.value](
{
'key': 'test-key',
'owner_type': 'plugin',
'owner': 'forged-owner',
}
)
assert response.code == 0
statement_params = compiled_params(app.persistence_mgr.execute_async.await_args.args[0])
assert 'workspace-a' in statement_params.values()
assert (
canonical_binary_key(
'plugin',
'test-author/test-plugin',
'test-key',
)
in statement_params.values()
)
assert 'forged-owner' not in statement_params.values()
@pytest.mark.asyncio
async def test_list_keys_uses_trusted_plugin_owner(self, app):
result = Mock()
result.scalars.return_value.all.return_value = ['first', 'second']
app.persistence_mgr.execute_async.return_value = result
runtime_handler = make_handler(app)
response = await runtime_handler.actions[RuntimeToLangBotAction.GET_BINARY_STORAGE_KEYS.value](
{
'owner_type': 'plugin',
'owner': 'forged-owner',
}
)
assert response.code == 0
assert response.data == {'keys': ['first', 'second']}
statement_params = compiled_params(app.persistence_mgr.execute_async.await_args.args[0])
assert 'workspace-a' in statement_params.values()
assert 'test-author/test-plugin' in statement_params.values()
assert 'forged-owner' not in statement_params.values()
class TestHandlerQueryLookup:
"""Tests for query lookup in cached_queries."""
tests/unit_tests/plugin/test_handler_tenancy.py
+446
View File
@@ -0,0 +1,446 @@
from __future__ import annotations
import asyncio
import contextlib
import json
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
import pytest
from sqlalchemy.ext.asyncio import create_async_engine
from langbot_plugin.entities.io.actions.enums import PluginToRuntimeAction
from langbot_plugin.entities.io.context import ActionContext, InstallationBinding, PluginWorkerPolicy, RuntimeIdentity
from langbot_plugin.entities.io.resp import ActionResponse
from langbot_plugin.runtime.io.connection import Connection
from langbot.pkg.plugin.handler import RuntimeConnectionHandler
from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.persistence.mgr import PersistenceManager, PersistenceMode
class EmptyResult:
def first(self):
return None
def all(self):
return []
class RecordingConnection(Connection):
def __init__(self):
self.sent: list[str] = []
async def send(self, message: str) -> None:
self.sent.append(message)
async def receive(self) -> str:
raise NotImplementedError
async def close(self) -> None:
return None
def workspace_context(workspace_uuid: str = 'workspace-a') -> ActionContext:
return ActionContext(
instance_uuid='instance-a',
workspace_uuid=workspace_uuid,
placement_generation=7,
)
def make_handler(workspace_uuid: str = 'workspace-a'):
context = workspace_context(workspace_uuid)
app = SimpleNamespace(
deployment=SimpleNamespace(mode='cloud'),
persistence_mgr=SimpleNamespace(execute_async=AsyncMock(return_value=EmptyResult())),
logger=Mock(),
workspace_service=SimpleNamespace(
get_execution_binding=AsyncMock(
return_value=SimpleNamespace(
instance_uuid=context.instance_uuid,
workspace_uuid=context.workspace_uuid,
placement_generation=context.placement_generation,
)
)
),
)
runtime_handler = RuntimeConnectionHandler(
Mock(),
AsyncMock(return_value=True),
app,
)
installation_context = InstallationBinding(
**context.model_dump(exclude_none=True),
installation_uuid='00000000-0000-4000-8000-000000000001',
runtime_revision=1,
artifact_digest='a' * 64,
)
runtime_handler.register_installation_binding(
installation_context,
plugin_author='author-a',
plugin_name='plugin-a',
)
return runtime_handler, app, installation_context
async def invoke_with_context(
runtime_handler: RuntimeConnectionHandler,
action_context: ActionContext,
action: PluginToRuntimeAction,
data: dict,
):
token = runtime_handler._current_action_context.set(action_context)
try:
return await runtime_handler.actions[action.value](data)
finally:
runtime_handler._current_action_context.reset(token)
@pytest.mark.asyncio
async def test_plugin_action_requires_installation_capability():
runtime_handler, _app, _installation_context = make_handler()
with pytest.raises(ValueError, match='trusted Workspace context'):
await runtime_handler.actions[PluginToRuntimeAction.GET_LANGBOT_VERSION.value]({})
@pytest.mark.asyncio
async def test_runtime_action_enters_trusted_workspace_scope():
runtime_handler, app, installation_context = make_handler()
scope_events: list[tuple[str, str]] = []
@contextlib.asynccontextmanager
async def tenant_scope(workspace_uuid: str):
scope_events.append(('enter', workspace_uuid))
try:
yield SimpleNamespace()
finally:
scope_events.append(('exit', workspace_uuid))
class RecordingPersistenceManager:
def __init__(self, execute_async):
self.execute_async = execute_async
def tenant_scope(self, workspace_uuid: str):
return tenant_scope(workspace_uuid)
app.persistence_mgr = RecordingPersistenceManager(app.persistence_mgr.execute_async)
response = await invoke_with_context(
runtime_handler,
installation_context,
PluginToRuntimeAction.GET_LANGBOT_VERSION,
{},
)
assert response.code == 0
assert scope_events == [('enter', 'workspace-a'), ('exit', 'workspace-a')]
@pytest.mark.asyncio
async def test_blocked_llm_provider_does_not_hold_tenant_database_session():
entered = asyncio.Event()
release = asyncio.Event()
observations: list[bool] = []
engine = create_async_engine('sqlite+aiosqlite:///:memory:')
manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
manager.db = SimpleNamespace(get_engine=lambda: engine)
runtime_handler, app, installation_context = make_handler()
app.persistence_mgr = manager
async def invoke_llm(**_kwargs):
observations.append(manager.current_session() is None)
entered.set()
await release.wait()
observations.append(manager.current_session() is None)
return SimpleNamespace(model_dump=lambda: {'role': 'assistant', 'content': 'ok'})
app.model_mgr = SimpleNamespace(
get_model_by_uuid=AsyncMock(
return_value=SimpleNamespace(
model_entity=SimpleNamespace(workspace_uuid='workspace-a'),
provider=SimpleNamespace(invoke_llm=invoke_llm),
)
)
)
runtime_handler._require_plugin_action_context = AsyncMock(return_value=(installation_context, SimpleNamespace()))
runtime_handler._require_active_action_context = AsyncMock()
runtime_handler._resource_exists = AsyncMock(return_value=True)
action = asyncio.create_task(
invoke_with_context(
runtime_handler,
installation_context,
PluginToRuntimeAction.INVOKE_LLM,
{
'llm_model_uuid': 'model-a',
'messages': [],
},
)
)
try:
await entered.wait()
assert observations == [True]
release.set()
response = await action
assert response.code == 0
assert observations == [True, True]
finally:
release.set()
if not action.done():
await action
await engine.dispose()
@pytest.mark.asyncio
async def test_plugin_action_rejects_forged_installation_capability():
runtime_handler, _app, installation_context = make_handler()
forged = installation_context.model_copy(update={'installation_uuid': 'forged-installation'})
with pytest.raises(
ValueError,
match='installation is not registered in this Workspace',
):
await invoke_with_context(
runtime_handler,
forged,
PluginToRuntimeAction.GET_LANGBOT_VERSION,
{},
)
@pytest.mark.asyncio
async def test_plugin_action_rejects_stale_workspace_generation():
runtime_handler, app, installation_context = make_handler()
app.workspace_service.get_execution_binding.side_effect = ValueError('generation is fenced')
with pytest.raises(ValueError, match='generation is fenced'):
await invoke_with_context(
runtime_handler,
installation_context,
PluginToRuntimeAction.GET_LANGBOT_VERSION,
{},
)
app.workspace_service.get_execution_binding.assert_awaited_once_with(
'workspace-a',
expected_generation=7,
)
@pytest.mark.asyncio
async def test_query_uuid_and_forged_payload_workspace_cannot_cross_tenants():
runtime_handler, app, installation_context = make_handler()
query_a = SimpleNamespace(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=7,
bot_uuid='bot-a',
)
query_b = SimpleNamespace(
instance_uuid='instance-a',
workspace_uuid='workspace-b',
placement_generation=7,
bot_uuid='bot-b',
)
async def get_query(workspace_uuid, query_uuid):
return {
('workspace-a', 'query-a'): query_a,
('workspace-b', 'query-b'): query_b,
}.get((workspace_uuid, query_uuid))
app.query_pool = SimpleNamespace(
get_query=AsyncMock(side_effect=get_query),
get_query_by_legacy_id=AsyncMock(),
)
response = await invoke_with_context(
runtime_handler,
installation_context,
PluginToRuntimeAction.GET_BOT_UUID,
{
'query_id': 2,
'query_uuid': 'query-b',
'workspace_uuid': 'workspace-b',
},
)
assert response.code != 0
app.query_pool.get_query.assert_awaited_once_with('workspace-a', 'query-b')
@pytest.mark.asyncio
async def test_legacy_query_id_fallback_is_workspace_scoped():
runtime_handler, app, installation_context = make_handler()
query = SimpleNamespace(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=7,
bot_uuid='bot-a',
)
app.query_pool = SimpleNamespace(
get_query=AsyncMock(),
get_query_by_legacy_id=AsyncMock(return_value=query),
)
response = await invoke_with_context(
runtime_handler,
installation_context,
PluginToRuntimeAction.GET_BOT_UUID,
{'query_id': 19},
)
assert response.code == 0
assert response.data == {'bot_uuid': 'bot-a'}
app.query_pool.get_query_by_legacy_id.assert_awaited_once_with(
'workspace-a',
19,
)
def test_runtime_connection_is_instance_scoped_and_unbound():
runtime_handler, _app, _installation_context = make_handler()
assert runtime_handler.bound_action_context is None
def test_inbound_tenant_action_requires_complete_installation_envelope():
runtime_handler, _app, installation_context = make_handler()
assert (
runtime_handler.validate_inbound_action_context(
PluginToRuntimeAction.GET_BOTS.value,
installation_context,
)
== installation_context
)
with pytest.raises(ValueError, match='complete InstallationBinding'):
runtime_handler.validate_inbound_action_context(
PluginToRuntimeAction.GET_BOTS.value,
workspace_context('workspace-b').for_installation('installation-b'),
)
@pytest.mark.asyncio
async def test_legacy_oss_worker_capability_remains_usable_after_identity_migration():
runtime_handler, app, installation_context = make_handler()
app.deployment = SimpleNamespace(mode='oss')
setting = SimpleNamespace(
plugin_author='author-a',
plugin_name='plugin-a',
installation_uuid=installation_context.installation_uuid,
runtime_revision=installation_context.runtime_revision,
artifact_digest=installation_context.artifact_digest,
)
result = Mock()
result.first.return_value = setting
app.persistence_mgr.execute_async.return_value = result
legacy_context = workspace_context().for_installation(installation_context.installation_uuid)
assert (
runtime_handler.validate_inbound_action_context(
PluginToRuntimeAction.GET_LANGBOT_VERSION.value,
legacy_context,
)
== legacy_context
)
response = await invoke_with_context(
runtime_handler,
legacy_context,
PluginToRuntimeAction.GET_LANGBOT_VERSION,
{},
)
assert response.code == 0
def test_installation_uuid_cannot_move_between_workspaces():
runtime_handler, _app, binding = make_handler()
moved = binding.model_copy(update={'workspace_uuid': 'workspace-b', 'runtime_revision': 2})
with pytest.raises(ValueError, match='cannot move between Workspaces'):
runtime_handler.register_installation_binding(
moved,
plugin_author='author-a',
plugin_name='plugin-a',
)
@pytest.mark.asyncio
async def test_newer_revision_fences_old_binding():
runtime_handler, _app, binding = make_handler()
newer = binding.model_copy(update={'runtime_revision': 2, 'artifact_digest': 'b' * 64})
runtime_handler.register_installation_binding(
newer,
plugin_author='author-a',
plugin_name='plugin-a',
)
with pytest.raises(ValueError, match='revision or artifact is stale'):
await runtime_handler._resolve_installation_identity(binding)
@pytest.mark.asyncio
async def test_plugin_vector_action_forwards_trusted_context_and_logical_collection():
runtime_handler, app, installation_context = make_handler()
app.rag_runtime_service = SimpleNamespace(vector_upsert=AsyncMock())
response = await invoke_with_context(
runtime_handler,
installation_context,
PluginToRuntimeAction.VECTOR_UPSERT,
{
'workspace_uuid': 'workspace-forged',
'collection_id': 'plugin-supplied-name',
'vectors': [[0.1]],
'ids': ['point-a'],
},
)
assert response.code == 0
execution_context = app.rag_runtime_service.vector_upsert.await_args.args[0]
assert execution_context == ExecutionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=7,
)
assert app.rag_runtime_service.vector_upsert.await_args.args[1] == 'plugin-supplied-name'
@pytest.mark.asyncio
async def test_host_to_runtime_action_carries_trusted_connector_context():
app = SimpleNamespace(logger=Mock())
connection = RecordingConnection()
runtime_handler = RuntimeConnectionHandler(
connection,
AsyncMock(return_value=True),
app,
)
task = asyncio.create_task(
runtime_handler.set_runtime_config(
runtime_identity=RuntimeIdentity(instance_uuid='instance-a', runtime_id='runtime-a'),
worker_policy=PluginWorkerPolicy(
max_cpus=1.0,
max_memory_mb=256,
max_pids=32,
max_open_files=64,
max_file_size_mb=128,
),
runtime_profile='oss_dev',
cloud_service_url=None,
)
)
for _ in range(10):
if connection.sent:
break
await asyncio.sleep(0)
request = json.loads(connection.sent[0])
runtime_handler.resp_waiters[request['seq_id']].set_result(ActionResponse.success({}))
await task
assert request['data']['runtime_identity'] == {
'instance_uuid': 'instance-a',
'runtime_id': 'runtime-a',
}
assert request.get('context') is None
tests/unit_tests/plugin/test_plugin_component_filtering.py
+35 -4
View File
@@ -1,7 +1,34 @@
"""Test plugin list filtering by component kinds."""
from contextlib import nullcontext
from unittest.mock import AsyncMock, MagicMock
import pytest
from langbot_plugin.entities.io.context import InstallationBinding
from langbot.pkg.api.http.context import ExecutionContext
TEST_EXECUTION_CONTEXT = ExecutionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=1,
)
TEST_INSTALLATION_BINDING = InstallationBinding(
instance_uuid=TEST_EXECUTION_CONTEXT.instance_uuid,
workspace_uuid=TEST_EXECUTION_CONTEXT.workspace_uuid,
placement_generation=TEST_EXECUTION_CONTEXT.placement_generation,
installation_uuid='00000000-0000-4000-8000-000000000001',
runtime_revision=1,
artifact_digest='a' * 64,
)
def configure_connector(connector) -> None:
connector._execution_context.set(TEST_EXECUTION_CONTEXT)
connector._operation_bindings = AsyncMock(return_value=[TEST_INSTALLATION_BINDING])
connector._load_workspace_settings = AsyncMock(return_value=[])
connector.handler.installation_scope = MagicMock(side_effect=lambda _binding: nullcontext())
@pytest.mark.asyncio
@@ -17,6 +44,7 @@ async def test_plugin_list_filter_by_component_kinds():
# Create connector
connector = PluginRuntimeConnector(mock_app, AsyncMock())
connector.handler = MagicMock()
configure_connector(connector)
# Mock plugin data with different component kinds
mock_plugins = [
@@ -90,7 +118,7 @@ async def test_plugin_list_filter_by_component_kinds():
# Mock database query
async def mock_execute_async(query):
mock_result = MagicMock()
mock_result.__iter__ = lambda self: iter([])
mock_result.scalars.return_value.all.return_value = []
return mock_result
mock_app.persistence_mgr.execute_async = mock_execute_async
@@ -123,6 +151,7 @@ async def test_plugin_list_filter_no_filter():
# Create connector
connector = PluginRuntimeConnector(mock_app, AsyncMock())
connector.handler = MagicMock()
configure_connector(connector)
# Mock plugin data with different component kinds
mock_plugins = [
@@ -157,7 +186,7 @@ async def test_plugin_list_filter_no_filter():
# Mock database query
async def mock_execute_async(query):
mock_result = MagicMock()
mock_result.__iter__ = lambda self: iter([])
mock_result.scalars.return_value.all.return_value = []
return mock_result
mock_app.persistence_mgr.execute_async = mock_execute_async
@@ -184,6 +213,7 @@ async def test_plugin_list_filter_empty_result():
# Create connector
connector = PluginRuntimeConnector(mock_app, AsyncMock())
connector.handler = MagicMock()
configure_connector(connector)
# Mock plugin data - only KnowledgeEngine plugins
mock_plugins = [
@@ -206,7 +236,7 @@ async def test_plugin_list_filter_empty_result():
# Mock database query
async def mock_execute_async(query):
mock_result = MagicMock()
mock_result.__iter__ = lambda self: iter([])
mock_result.scalars.return_value.all.return_value = []
return mock_result
mock_app.persistence_mgr.execute_async = mock_execute_async
@@ -230,6 +260,7 @@ async def test_plugin_list_filter_plugin_without_components():
# Create connector
connector = PluginRuntimeConnector(mock_app, AsyncMock())
connector.handler = MagicMock()
configure_connector(connector)
# Mock plugin data - one with components, one without
mock_plugins = [
@@ -264,7 +295,7 @@ async def test_plugin_list_filter_plugin_without_components():
# Mock database query
async def mock_execute_async(query):
mock_result = MagicMock()
mock_result.__iter__ = lambda self: iter([])
mock_result.scalars.return_value.all.return_value = []
return mock_result
mock_app.persistence_mgr.execute_async = mock_execute_async
tests/unit_tests/plugin/test_plugin_list_sorting.py
+35 -16
View File
@@ -1,8 +1,34 @@
"""Test plugin list sorting functionality."""
from contextlib import nullcontext
from datetime import datetime, timedelta
from unittest.mock import AsyncMock, MagicMock
import pytest
from langbot_plugin.entities.io.context import InstallationBinding
from langbot.pkg.api.http.context import ExecutionContext
TEST_EXECUTION_CONTEXT = ExecutionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=1,
)
TEST_INSTALLATION_BINDING = InstallationBinding(
instance_uuid=TEST_EXECUTION_CONTEXT.instance_uuid,
workspace_uuid=TEST_EXECUTION_CONTEXT.workspace_uuid,
placement_generation=TEST_EXECUTION_CONTEXT.placement_generation,
installation_uuid='00000000-0000-4000-8000-000000000001',
runtime_revision=1,
artifact_digest='a' * 64,
)
def configure_connector(connector) -> None:
connector._execution_context.set(TEST_EXECUTION_CONTEXT)
connector._operation_bindings = AsyncMock(return_value=[TEST_INSTALLATION_BINDING])
connector.handler.installation_scope = MagicMock(side_effect=lambda _binding: nullcontext())
@pytest.mark.asyncio
@@ -18,6 +44,7 @@ async def test_plugin_list_sorting_debug_first():
# Create connector
connector = PluginRuntimeConnector(mock_app, AsyncMock())
connector.handler = MagicMock()
configure_connector(connector)
# Mock plugin data with different debug states and timestamps
now = datetime.now()
@@ -60,9 +87,7 @@ async def test_plugin_list_sorting_debug_first():
connector.handler.list_plugins = AsyncMock(return_value=mock_plugins)
# Mock database query to return all timestamps in a single batch
async def mock_execute_async(query):
mock_result = MagicMock()
async def mock_load_workspace_settings(_execution_context):
# Create mock rows for all plugins with timestamps
mock_rows = []
@@ -85,12 +110,9 @@ async def test_plugin_list_sorting_debug_first():
mock_row3.created_at = now
mock_rows.append(mock_row3)
# Make the result iterable
mock_result.__iter__ = lambda self: iter(mock_rows)
return mock_rows
return mock_result
mock_app.persistence_mgr.execute_async = mock_execute_async
connector._load_workspace_settings = AsyncMock(side_effect=mock_load_workspace_settings)
# Call list_plugins
result = await connector.list_plugins()
@@ -120,6 +142,7 @@ async def test_plugin_list_sorting_by_installation_time():
# Create connector
connector = PluginRuntimeConnector(mock_app, AsyncMock())
connector.handler = MagicMock()
configure_connector(connector)
# Mock plugin data - all non-debug with different installation times
now = datetime.now()
@@ -162,9 +185,7 @@ async def test_plugin_list_sorting_by_installation_time():
connector.handler.list_plugins = AsyncMock(return_value=mock_plugins)
# Mock database query to return all timestamps in a single batch
async def mock_execute_async(query):
mock_result = MagicMock()
async def mock_load_workspace_settings(_execution_context):
# Create mock rows for all plugins with timestamps
mock_rows = []
@@ -187,12 +208,9 @@ async def test_plugin_list_sorting_by_installation_time():
mock_row3.created_at = now
mock_rows.append(mock_row3)
# Make the result iterable
mock_result.__iter__ = lambda self: iter(mock_rows)
return mock_rows
return mock_result
mock_app.persistence_mgr.execute_async = mock_execute_async
connector._load_workspace_settings = AsyncMock(side_effect=mock_load_workspace_settings)
# Call list_plugins
result = await connector.list_plugins()
@@ -217,6 +235,7 @@ async def test_plugin_list_empty():
# Create connector
connector = PluginRuntimeConnector(mock_app, AsyncMock())
connector.handler = MagicMock()
configure_connector(connector)
# Mock empty plugin list
connector.handler.list_plugins = AsyncMock(return_value=[])
tests/unit_tests/provider/conftest.py
+43
View File
@@ -16,6 +16,18 @@ from langbot.pkg.provider.modelmgr import token
from langbot.pkg.provider.modelmgr.modelmgr import ModelManager
from langbot.pkg.entity.persistence import model as persistence_model
from langbot.pkg.discover import engine as discover_engine
from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.workspace.entities import WorkspaceExecutionBinding
TEST_INSTANCE_UUID = 'test-instance'
TEST_WORKSPACE_UUID = 'test-workspace'
TEST_GENERATION = 1
TEST_EXECUTION_CONTEXT = ExecutionContext(
instance_uuid=TEST_INSTANCE_UUID,
workspace_uuid=TEST_WORKSPACE_UUID,
placement_generation=TEST_GENERATION,
)
class FakeProviderAPIRequester(requester.ProviderAPIRequester):
@@ -157,6 +169,26 @@ def mock_app_for_modelmgr():
app.llm_model_service = AsyncMock()
app.embedding_models_service = AsyncMock()
app.monitoring_service = AsyncMock()
app.workspace_service = SimpleNamespace(
get_execution_binding=AsyncMock(
return_value=WorkspaceExecutionBinding(
instance_uuid=TEST_INSTANCE_UUID,
workspace_uuid=TEST_WORKSPACE_UUID,
placement_generation=TEST_GENERATION,
write_fenced=False,
state='active',
)
),
get_local_execution_binding=AsyncMock(
return_value=WorkspaceExecutionBinding(
instance_uuid=TEST_INSTANCE_UUID,
workspace_uuid=TEST_WORKSPACE_UUID,
placement_generation=TEST_GENERATION,
write_fenced=False,
state='active',
)
),
)
return app
@@ -184,6 +216,7 @@ def fake_persistence_data():
providers = [
persistence_model.ModelProvider(
workspace_uuid=TEST_WORKSPACE_UUID,
uuid=provider_uuid,
name='Test Provider',
requester='fake-requester',
@@ -191,6 +224,7 @@ def fake_persistence_data():
api_keys=['test-api-key-1', 'test-api-key-2'],
),
persistence_model.ModelProvider(
workspace_uuid=TEST_WORKSPACE_UUID,
uuid=provider_uuid2,
name='Test Provider 2',
requester='another-fake-requester',
@@ -201,6 +235,7 @@ def fake_persistence_data():
llm_models = [
persistence_model.LLMModel(
workspace_uuid=TEST_WORKSPACE_UUID,
uuid='test-llm-uuid-1',
name='TestLLM-1',
provider_uuid=provider_uuid,
@@ -208,6 +243,7 @@ def fake_persistence_data():
extra_args={'temperature': 0.7},
),
persistence_model.LLMModel(
workspace_uuid=TEST_WORKSPACE_UUID,
uuid='test-llm-uuid-2',
name='TestLLM-2',
provider_uuid=provider_uuid,
@@ -218,6 +254,7 @@ def fake_persistence_data():
embedding_models = [
persistence_model.EmbeddingModel(
workspace_uuid=TEST_WORKSPACE_UUID,
uuid='test-embedding-uuid-1',
name='TestEmbedding-1',
provider_uuid=provider_uuid,
@@ -227,6 +264,7 @@ def fake_persistence_data():
rerank_models = [
persistence_model.RerankModel(
workspace_uuid=TEST_WORKSPACE_UUID,
uuid='test-rerank-uuid-1',
name='TestRerank-1',
provider_uuid=provider_uuid2,
@@ -252,6 +290,7 @@ def runtime_provider(fake_persistence_data, mock_app_for_modelmgr):
requester_inst = FakeProviderAPIRequester(mock_app_for_modelmgr, {'base_url': provider_entity.base_url})
return requester.RuntimeProvider(
execution_context=TEST_EXECUTION_CONTEXT,
provider_entity=provider_entity,
token_mgr=token_mgr,
requester=requester_inst,
@@ -263,6 +302,7 @@ def runtime_llm_model(fake_persistence_data, runtime_provider):
"""Provides a RuntimeLLMModel instance for testing."""
model_entity = fake_persistence_data['llm_models'][0]
return requester.RuntimeLLMModel(
execution_context=TEST_EXECUTION_CONTEXT,
model_entity=model_entity,
provider=runtime_provider,
)
@@ -273,6 +313,7 @@ def runtime_embedding_model(fake_persistence_data, runtime_provider):
"""Provides a RuntimeEmbeddingModel instance for testing."""
model_entity = fake_persistence_data['embedding_models'][0]
return requester.RuntimeEmbeddingModel(
execution_context=TEST_EXECUTION_CONTEXT,
model_entity=model_entity,
provider=runtime_provider,
)
@@ -286,6 +327,7 @@ def runtime_rerank_model(fake_persistence_data, mock_app_for_modelmgr):
requester_inst = AnotherFakeRequester(mock_app_for_modelmgr, {'base_url': provider_entity.base_url})
provider = requester.RuntimeProvider(
execution_context=TEST_EXECUTION_CONTEXT,
provider_entity=provider_entity,
token_mgr=token_mgr,
requester=requester_inst,
@@ -293,6 +335,7 @@ def runtime_rerank_model(fake_persistence_data, mock_app_for_modelmgr):
model_entity = fake_persistence_data['rerank_models'][0]
return requester.RuntimeRerankModel(
execution_context=TEST_EXECUTION_CONTEXT,
model_entity=model_entity,
provider=provider,
)
tests/unit_tests/provider/runners/test_difysvapi_runner.py
+181 -7
View File
@@ -6,6 +6,7 @@ Tests the helper methods that don't require real Dify API calls.
from __future__ import annotations
import pytest
import time
from unittest.mock import AsyncMock, MagicMock
import langbot_plugin.api.entities.builtin.platform.message as platform_message
@@ -18,13 +19,12 @@ class TestDifyWorkflowSubmitClient:
class FakeResponse:
status_code = 503
headers = {}
async def aread(self):
return b''
async def aiter_lines(self):
raise AssertionError('error responses must not enter the SSE loop')
yield
async def aiter_bytes(self, chunk_size=None):
del chunk_size
if False:
yield b''
class FakeStreamContext:
async def __aenter__(self):
@@ -66,6 +66,33 @@ class TestDifyWorkflowSubmitClient:
)
)
@pytest.mark.asyncio
async def test_sse_parser_rejects_an_unbounded_line(self):
from langbot.libs.dify_service_api.v1 import client, errors
class FakeResponse:
async def aiter_bytes(self, chunk_size=None):
del chunk_size
for _ in range(129):
yield b'x' * 8192
with pytest.raises(errors.DifyAPIError, match='SSE event exceeds'):
await anext(client._iter_sse_json(FakeResponse()))
@pytest.mark.asyncio
async def test_upload_rejects_oversized_local_file(self, tmp_path):
from langbot.libs.dify_service_api.v1 import client
file_path = tmp_path / 'large.bin'
file_path.write_bytes(b'x' * (client._MAX_DIFY_UPLOAD_BYTES + 1))
dify_client = client.AsyncDifyServiceClient(
'test-key',
'https://dify.example/v1',
)
with pytest.raises(ValueError, match='exceeds the size limit'):
await dify_client.upload_file(file_path, 'person_user-1')
class TestDifyExtractTextOutput:
"""Tests for _extract_dify_text_output method."""
@@ -252,42 +279,65 @@ class TestDifyHumanInputForms:
runner.dify_client.upload_file = AsyncMock(return_value={'id': 'upload-1'})
return runner
def test_pending_forms_are_isolated_by_bot_and_pipeline(self):
def test_pending_forms_are_isolated_by_workspace_generation_bot_and_pipeline(self):
from langbot.pkg.provider.runners import difysvapi
query_a = MagicMock()
query_a.instance_uuid = 'instance-a'
query_a.workspace_uuid = 'workspace-a'
query_a.placement_generation = 1
query_a.bot_uuid = 'bot-a'
query_a.pipeline_uuid = 'pipeline-a'
query_a.session.launcher_type.value = 'person'
query_a.session.launcher_id = 'shared-user'
query_b = MagicMock()
query_b.instance_uuid = 'instance-a'
query_b.workspace_uuid = 'workspace-a'
query_b.placement_generation = 1
query_b.bot_uuid = 'bot-b'
query_b.pipeline_uuid = 'pipeline-a'
query_b.session.launcher_type.value = 'person'
query_b.session.launcher_id = 'shared-user'
query_c = MagicMock()
query_c.instance_uuid = 'instance-a'
query_c.workspace_uuid = 'workspace-a'
query_c.placement_generation = 1
query_c.bot_uuid = 'bot-a'
query_c.pipeline_uuid = 'pipeline-b'
query_c.session.launcher_type.value = 'person'
query_c.session.launcher_id = 'shared-user'
query_d = MagicMock()
query_d.instance_uuid = 'instance-a'
query_d.workspace_uuid = 'workspace-b'
query_d.placement_generation = 2
query_d.bot_uuid = 'bot-a'
query_d.pipeline_uuid = 'pipeline-a'
query_d.session.launcher_type.value = 'person'
query_d.session.launcher_id = 'shared-user'
key_a = difysvapi._session_key_from_query(query_a)
key_b = difysvapi._session_key_from_query(query_b)
key_c = difysvapi._session_key_from_query(query_c)
key_d = difysvapi._session_key_from_query(query_d)
difysvapi._PENDING_FORMS.clear()
difysvapi._set_pending_form(key_a, {'form_token': 'token-a', 'workflow_run_id': 'run-a'})
difysvapi._set_pending_form(key_b, {'form_token': 'token-b', 'workflow_run_id': 'run-b'})
difysvapi._set_pending_form(key_c, {'form_token': 'token-c', 'workflow_run_id': 'run-c'})
difysvapi._set_pending_form(key_d, {'form_token': 'token-d', 'workflow_run_id': 'run-d'})
assert key_a != key_b
assert key_a != key_c
assert key_a != key_d
assert difysvapi._get_pending_form_by_token(key_a, 'token-a') is not None
assert difysvapi._get_pending_form_by_token(key_a, 'token-b') is None
assert difysvapi._get_pending_form_by_token(key_a, 'token-c') is None
assert difysvapi._get_pending_form_by_token(key_a, 'token-d') is None
assert difysvapi._get_pending_form_by_token(key_b, 'token-b') is not None
assert difysvapi._get_pending_form_by_token(key_c, 'token-c') is not None
assert difysvapi._get_pending_form_by_token(key_d, 'token-d') is not None
assert difysvapi._get_latest_pending_form(key_a)['workflow_run_id'] == 'run-a'
assert difysvapi._get_latest_pending_form(key_b)['workflow_run_id'] == 'run-b'
assert difysvapi._get_latest_pending_form(key_c)['workflow_run_id'] == 'run-c'
@@ -297,6 +347,130 @@ class TestDifyHumanInputForms:
assert difysvapi._dify_user_from_query(query_a) == difysvapi._dify_user_from_query(query_c)
difysvapi._PENDING_FORMS.clear()
def test_pending_form_lookup_does_not_scan_unrelated_sessions(self, monkeypatch):
from langbot.pkg.provider.runners import difysvapi
def session_key(index: int):
return (
'instance',
f'workspace-{index}',
1,
'bot',
'pipeline',
'adapter',
'person',
f'user-{index}',
)
difysvapi._PENDING_FORMS.clear()
for index in range(512):
difysvapi._set_pending_form(
session_key(index),
{
'form_token': f'token-{index}',
'workflow_run_id': f'run-{index}',
},
)
class NoGlobalIterationDict(dict):
def __iter__(self):
raise AssertionError('pending form lookup scanned all sessions')
def keys(self):
raise AssertionError('pending form lookup scanned all sessions')
def items(self):
raise AssertionError('pending form lookup scanned all sessions')
def values(self):
raise AssertionError('pending form lookup scanned all sessions')
guarded_forms = NoGlobalIterationDict(difysvapi._PENDING_FORMS)
monkeypatch.setattr(difysvapi, '_PENDING_FORMS', guarded_forms)
assert difysvapi._get_pending_form_by_token(session_key(511), 'token-511')['workflow_run_id'] == 'run-511'
difysvapi._set_pending_form(
session_key(512),
{'form_token': 'token-512', 'workflow_run_id': 'run-512'},
)
assert len(guarded_forms) == 513
def test_pending_form_expiry_heap_ignores_stale_overwrite_and_stays_bounded(self):
from langbot.pkg.provider.runners import difysvapi
session_key = (
'instance',
'workspace',
1,
'bot',
'pipeline',
'adapter',
'person',
'user',
)
difysvapi._PENDING_FORMS.clear()
now = time.time()
difysvapi._set_pending_form(
session_key,
{
'form_token': 'token',
'workflow_run_id': 'stale',
'expiration_time': now + 1,
},
)
for revision in range(500):
difysvapi._set_pending_form(
session_key,
{
'form_token': 'token',
'workflow_run_id': f'current-{revision}',
'expiration_time': now + 3600 + revision,
},
)
difysvapi._prune_pending_forms(now + 2)
assert difysvapi._get_pending_form_by_token(session_key, 'token')['workflow_run_id'] == 'current-499'
assert difysvapi._PENDING_FORM_ACTIVE_COUNT == 1
assert len(difysvapi._PENDING_FORM_EXPIRY_HEAP) <= max(
difysvapi._PENDING_FORM_HEAP_COMPACT_FLOOR,
difysvapi._PENDING_FORM_ACTIVE_COUNT * difysvapi._PENDING_FORM_HEAP_MAX_MULTIPLIER,
)
def test_pending_form_capacity_evicts_earliest_session_without_full_scan(
self,
monkeypatch,
):
from langbot.pkg.provider.runners import difysvapi
def session_key(index: int):
return (
'instance',
f'workspace-{index}',
1,
'bot',
'pipeline',
'adapter',
'person',
f'user-{index}',
)
difysvapi._PENDING_FORMS.clear()
monkeypatch.setattr(difysvapi, '_PENDING_FORM_MAX_SESSIONS', 2)
now = time.time()
for index, expires_in in ((1, 300), (2, 100), (3, 200)):
difysvapi._set_pending_form(
session_key(index),
{
'form_token': f'token-{index}',
'expiration_time': now + expires_in,
},
)
assert session_key(1) in difysvapi._PENDING_FORMS
assert session_key(2) not in difysvapi._PENDING_FORMS
assert session_key(3) in difysvapi._PENDING_FORMS
def test_interactive_form_data_preserves_pipeline_uuid(self):
from langbot.pkg.provider.runners import difysvapi
tests/unit_tests/provider/runners/test_remote_stream_limits.py
+49
View File
@@ -0,0 +1,49 @@
from __future__ import annotations
import pytest
from langbot.libs.deerflow_api.client import (
ERROR_BODY_MAX_BYTES,
_read_error_body,
)
from langbot.libs.deerflow_api.errors import DeerFlowAPIError
from langbot.pkg.provider.runners.langflowapi import (
_MAX_LANGFLOW_LINE_CHARS,
_MAX_LANGFLOW_RESPONSE_BYTES,
_iter_limited_lines,
_read_limited_response,
)
class _ChunkedResponse:
def __init__(self, chunks: list[bytes]):
self._chunks = chunks
async def aiter_bytes(self, chunk_size=None):
del chunk_size
for chunk in self._chunks:
yield chunk
@pytest.mark.asyncio
async def test_langflow_rejects_oversized_stream_event():
response = _ChunkedResponse([b'x' * (_MAX_LANGFLOW_LINE_CHARS + 1)])
with pytest.raises(ValueError, match='event exceeds'):
await anext(_iter_limited_lines(response))
@pytest.mark.asyncio
async def test_langflow_rejects_oversized_blocking_response():
response = _ChunkedResponse([b'x' * (_MAX_LANGFLOW_RESPONSE_BYTES + 1)])
with pytest.raises(ValueError, match='response exceeds'):
await _read_limited_response(response)
@pytest.mark.asyncio
async def test_deerflow_rejects_oversized_error_body():
response = _ChunkedResponse([b'x' * (ERROR_BODY_MAX_BYTES + 1)])
with pytest.raises(DeerFlowAPIError, match='response exceeds'):
await _read_error_body(response)
tests/unit_tests/provider/runners/test_runner_resource_limits.py
+62
View File
@@ -0,0 +1,62 @@
from __future__ import annotations
import asyncio
import threading
from unittest.mock import AsyncMock
import pytest
from langbot.pkg.provider import runner
from langbot.pkg.provider.runners import (
cozeapi,
dashscopeapi,
tboxapi,
weknoraapi,
)
@pytest.mark.asyncio
async def test_blocking_provider_iterator_runs_outside_event_loop():
release = threading.Event()
def values():
release.wait(timeout=2)
yield 'ready'
task = asyncio.create_task(anext(runner.iterate_sync(values())))
await asyncio.sleep(0)
assert not task.done()
release.set()
assert await asyncio.wait_for(task, timeout=1) == 'ready'
@pytest.mark.asyncio
async def test_sync_provider_iterator_has_event_limit():
with pytest.raises(RuntimeError, match='event limit'):
async for _ in runner.iterate_sync(iter([1, 2]), max_items=1):
pass
@pytest.mark.asyncio
async def test_coze_runner_closes_request_scoped_client():
request_runner = object.__new__(cozeapi.CozeAPIRunner)
request_runner.coze = AsyncMock()
await request_runner.aclose()
request_runner.coze.close.assert_awaited_once()
@pytest.mark.parametrize(
('append', 'exception_type'),
[
(cozeapi._append_bounded, ValueError),
(dashscopeapi._append_bounded, dashscopeapi.DashscopeAPIError),
(tboxapi._append_bounded, tboxapi.TboxAPIError),
(weknoraapi._append_bounded, weknoraapi.errors.WeKnoraAPIError),
],
)
def test_provider_accumulators_reject_oversized_output(append, exception_type):
with pytest.raises(exception_type, match='exceeds the runtime limit'):
append('x' * (1024 * 1024), 'y')
tests/unit_tests/provider/test_localagent_no_duplicate.py
+13 -1
View File
@@ -10,6 +10,7 @@ import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
import langbot_plugin.api.entities.builtin.provider.message as provider_message
import langbot_plugin.api.entities.builtin.provider.session as provider_session
from langbot.pkg.api.http.context import ExecutionContext, PrincipalContext, PrincipalType
from langbot.pkg.provider.runners.localagent import LocalAgentRunner
@@ -97,7 +98,7 @@ def make_query() -> pipeline_query.Query:
adapter = AsyncMock()
adapter.is_stream_output_supported = AsyncMock(return_value=False)
return pipeline_query.Query.model_construct(
query = pipeline_query.Query.model_construct(
query_id='no-dup-query',
launcher_type=provider_session.LauncherTypes.PERSON,
launcher_id=12345,
@@ -124,6 +125,17 @@ def make_query() -> pipeline_query.Query:
use_llm_model_uuid='test-model-uuid',
variables={},
)
object.__setattr__(
query,
'_execution_context',
ExecutionContext(
instance_uuid='instance-test',
workspace_uuid='workspace-test',
placement_generation=1,
trigger_principal=PrincipalContext(PrincipalType.SYSTEM),
),
)
return query
def _make_app(provider) -> SimpleNamespace:

Some files were not shown because too many files have changed in this diff Show More