mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-08-09 04:40:57 +00:00
chore(merge): sync master into dev/4.11.x
This commit is contained in:
huanghuoguoguo
@@ -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
|
||||
|
||||
@@ -49,7 +54,7 @@ async def test_update_bot_copies_input_before_filtering_legacy_routing_fields():
|
||||
'pipeline_routing_rules': [{'type': 'launcher_type'}],
|
||||
}
|
||||
|
||||
await service.update_bot('bot-1', payload)
|
||||
await service.update_bot(WORKSPACE_UUID, 'bot-1', payload)
|
||||
|
||||
# caller's dict must not be mutated
|
||||
assert payload == {
|
||||
|
||||
@@ -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'}
|
||||
@@ -0,0 +1,381 @@
|
||||
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(ValueError, match='Pipeline not found'):
|
||||
await application.bot_service.update_bot(
|
||||
WORKSPACE_A,
|
||||
'bot-a',
|
||||
{
|
||||
'event_bindings': [
|
||||
{
|
||||
'event_pattern': 'message.*',
|
||||
'target_type': 'pipeline',
|
||||
'target_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
|
||||
)
|
||||
@@ -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
|
||||
@@ -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']
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user