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

fix(cloud): enforce instance capacity ceilings

Browse Source
This commit is contained in:
Junyan Qin
2026-07-29 13:45:58 +08:00
parent e52d6880f5
commit c89e6f3bd2
25 changed files with 1062 additions and 87 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/unit_tests/cloud/test_bootstrap.py
+28
View File
@@ -178,6 +178,34 @@ async def test_cloud_runtime_config_is_fail_closed(field, value, message):
)
@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'),
[
tests/unit_tests/cloud/test_directory_projection.py
+182
View File
@@ -14,6 +14,7 @@ from langbot.pkg.cloud.directory import (
DirectoryEvent,
DirectoryEventBatch,
DirectoryMember,
DirectoryProjectionLimits,
DirectoryProjectionUnavailableError,
DirectorySnapshot,
DirectoryWorkspace,
@@ -274,6 +275,187 @@ async def test_event_limit_matches_single_delta_request_limit(projection_context
)
async def test_snapshot_capacity_rejects_atomically_before_projection(projection_context):
application, session_factory = projection_context
service = DirectoryProjectionService(
application,
_Provider([_snapshot(1)]),
INSTANCE_UUID,
limits=DirectoryProjectionLimits(
max_active_workspaces=1,
max_snapshot_workspaces=1,
max_snapshot_memberships=1,
),
)
second_workspace = DirectoryWorkspace(
uuid=SECOND_WORKSPACE_UUID,
name='Second Workspace',
slug='second-workspace',
type='personal',
status='active',
created_by_account_uuid='20000000-0000-0000-0000-000000000002',
projection_revision=1,
execution_generation=1,
members=[
DirectoryMember(
membership_uuid=SECOND_MEMBERSHIP_UUID,
account_uuid='20000000-0000-0000-0000-000000000002',
normalized_email='second@example.com',
display_name='Second Owner',
account_status='active',
role='owner',
membership_status='active',
projection_revision=1,
)
],
)
with pytest.raises(DirectoryProjectionUnavailableError, match='Workspace capacity exceeded'):
await service.apply_snapshot(
_snapshot(
1,
workspaces=[
_workspace(),
second_workspace,
],
)
)
async with session_factory() as session:
assert await session.scalar(sqlalchemy.select(sqlalchemy.func.count()).select_from(Workspace)) == 0
assert await session.get(DirectoryProjectionState, INSTANCE_UUID) is None
async def test_incremental_capacity_rolls_back_without_advancing_cursor(projection_context):
application, session_factory = projection_context
service = DirectoryProjectionService(
application,
_Provider([_snapshot(1)]),
INSTANCE_UUID,
limits=DirectoryProjectionLimits(
max_active_workspaces=1,
max_snapshot_workspaces=1,
max_snapshot_memberships=2,
),
)
await service.initialize()
second_account_uuid = '20000000-0000-0000-0000-000000000002'
second_workspace = DirectoryWorkspace(
uuid=SECOND_WORKSPACE_UUID,
name='Second Workspace',
slug='second-workspace',
type='personal',
status='active',
created_by_account_uuid=second_account_uuid,
projection_revision=2,
execution_generation=1,
members=[
DirectoryMember(
membership_uuid=SECOND_MEMBERSHIP_UUID,
account_uuid=second_account_uuid,
normalized_email='second@example.com',
display_name='Second Owner',
account_status='active',
role='owner',
membership_status='active',
projection_revision=2,
)
],
)
event = DirectoryEvent(
cursor=2,
uuid='40000000-0000-0000-0000-000000000002',
aggregate_uuid=SECOND_WORKSPACE_UUID,
event_type='directory.changed',
revision=2,
payload={
'workspace_uuid': SECOND_WORKSPACE_UUID,
'directory_revision': 2,
},
created_at=datetime.datetime(2026, 7, 24, 12, 2, tzinfo=datetime.UTC),
)
batch = DirectoryEventBatch(
instance_uuid=INSTANCE_UUID,
after_cursor=1,
cursor=2,
high_water_cursor=2,
events=[event],
)
delta = DirectoryDelta(
instance_uuid=INSTANCE_UUID,
requested_workspace_uuids=[SECOND_WORKSPACE_UUID],
generated_at=datetime.datetime(2026, 7, 24, 12, 2, tzinfo=datetime.UTC),
workspaces=[second_workspace],
)
with pytest.raises(DirectoryProjectionUnavailableError, match='Projected active Workspace capacity exceeded'):
await service.apply_delta(delta, batch)
async with session_factory() as session:
assert await session.get(Workspace, SECOND_WORKSPACE_UUID) is None
state = await session.get(DirectoryProjectionState, INSTANCE_UUID)
assert state is not None
assert state.cursor == 1
assert service.resource_snapshot()['active_workspaces'] == 1
async def test_account_projection_reads_large_directory_in_bounded_batches(projection_context):
application, session_factory = projection_context
service = DirectoryProjectionService(
application,
_Provider([_snapshot(1)]),
INSTANCE_UUID,
)
workspaces = []
for number in range(501):
account_uuid = f'20000000-0000-0000-0000-{number:012d}'
workspaces.append(
DirectoryWorkspace(
uuid=f'10000000-0000-0000-0000-{number:012d}',
name=f'Workspace {number}',
slug=f'workspace-{number}',
type='personal',
status='active',
created_by_account_uuid=account_uuid,
projection_revision=1,
execution_generation=1,
members=[
DirectoryMember(
membership_uuid=f'30000000-0000-0000-0000-{number:012d}',
account_uuid=account_uuid,
normalized_email=f'owner-{number}@example.com',
display_name=f'Owner {number}',
account_status='active',
role='owner',
membership_status='active',
projection_revision=1,
)
],
)
)
snapshot = _snapshot(1, workspaces=workspaces)
user_selects = 0
def count_user_selects(_connection, _cursor, statement, _parameters, _context, _executemany):
nonlocal user_selects
if statement.lstrip().upper().startswith('SELECT') and 'users' in statement:
user_selects += 1
engine = application.persistence_mgr.engine
sqlalchemy.event.listen(engine.sync_engine, 'before_cursor_execute', count_user_selects)
try:
async with application.persistence_mgr.directory_projection_uow(INSTANCE_UUID) as uow:
projected = await service._apply_accounts(uow.session, snapshot)
finally:
sqlalchemy.event.remove(engine.sync_engine, 'before_cursor_execute', count_user_selects)
assert len(projected) == 501
assert user_selects == 2
async with session_factory() as session:
assert await session.scalar(sqlalchemy.select(sqlalchemy.func.count()).select_from(User)) == 501
async def test_archived_and_absent_workspaces_are_execution_fenced(projection_context):
application, session_factory = projection_context
service = DirectoryProjectionService(
tests/unit_tests/core/test_app_shutdown.py
+20
View File
@@ -88,6 +88,18 @@ async def test_runtime_resource_stats_are_aggregate_and_constant_time() -> None:
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={},
@@ -113,6 +125,14 @@ async def test_runtime_resource_stats_are_aggregate_and_constant_time() -> None:
'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,
tests/unit_tests/core/test_load_config.py
+14
View File
@@ -64,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()
tests/unit_tests/persistence/test_postgresql_database.py
+80
View File
@@ -121,6 +121,43 @@ async def test_postgresql_manager_applies_explicit_bounded_pool_options(monkeypa
}
@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'),
@@ -128,8 +165,12 @@ async def test_postgresql_manager_applies_explicit_bounded_pool_options(monkeypa
('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:
@@ -139,6 +180,45 @@ async def test_postgresql_manager_rejects_invalid_pool_options(name, value) -> N
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(
tests/unit_tests/persistence/test_tenant_uow.py
+28
View File
@@ -26,6 +26,7 @@ from langbot.pkg.persistence.tenant_uow import (
CrossScopeTransactionError,
PersistenceScopeKind,
ScopedSessionTransactionError,
TenantScopedAsyncSession,
TenantScopedSyncSession,
TenantScopeRequiredError,
TenantUnitOfWork,
@@ -37,6 +38,33 @@ from langbot.pkg.persistence.tenant_uow import (
pytestmark = pytest.mark.asyncio
async def test_tenant_uow_reports_pool_timeout_during_transaction_admission(monkeypatch) -> None:
engine = create_async_engine('sqlite+aiosqlite:///:memory:')
pool_timeouts = 0
def record_pool_timeout() -> None:
nonlocal pool_timeouts
pool_timeouts += 1
async def fail_transaction_start(self, capability):
del self, capability
raise sa.exc.TimeoutError('pool exhausted')
monkeypatch.setattr(TenantScopedAsyncSession, '_start_owned_transaction', fail_transaction_start)
try:
with pytest.raises(sa.exc.TimeoutError, match='pool exhausted'):
async with TenantUnitOfWork(
engine,
'10000000-0000-0000-0000-000000000001',
on_pool_timeout=record_pool_timeout,
):
pass
finally:
await engine.dispose()
assert pool_timeouts == 1
def _on_conflict_statement(*, update_value, update_key='value', index_element=None):
table = sa.table('conflict_rows', sa.column('id'), sa.column('value'))
if index_element is None:
tests/unit_tests/utils/test_cloud_runtime_soak.py
+52
View File
@@ -461,6 +461,58 @@ def test_evaluate_gate_detects_stuck_mcp_projection_cleanup() -> None:
assert any('mcp_projection_reconcile_active above zero' in failure for failure in result.failures)
def test_evaluate_gate_detects_directory_and_database_capacity_violation() -> None:
prefix = 'body.resources'
state = _state(
'endpoint',
[
_sample(
0,
**{
f'{prefix}.directory.active_workspaces': 1000,
f'{prefix}.directory.max_active_workspaces': 1000,
f'{prefix}.database_pool.checked_out': 20,
f'{prefix}.database_pool.configured_capacity': 20,
},
),
_sample(
60,
**{
f'{prefix}.directory.active_workspaces': 1001,
f'{prefix}.directory.max_active_workspaces': 1000,
f'{prefix}.database_pool.checked_out': 21,
f'{prefix}.database_pool.configured_capacity': 20,
},
),
],
)
result = soak.evaluate_gate(
[state],
analysis_start_seconds=0,
thresholds=_thresholds(),
)
assert any('directory.active_workspaces' in failure for failure in result.failures)
assert any('database_pool.checked_out' in failure for failure in result.failures)
def test_evaluate_gate_requires_capacity_gauges_from_named_core_endpoint() -> None:
state = _state(
'endpoint',
[_sample(0, **{'http.ok': 1}), _sample(60, **{'http.ok': 1})],
)
state.target = soak.Target(name='core', kind='endpoint', location='/core')
result = soak.evaluate_gate(
[state],
analysis_start_seconds=0,
thresholds=_thresholds(require_event_loop_metrics=False),
)
assert any('required capacity gauge' in failure for failure in result.failures)
def test_evaluate_gate_detects_event_loop_stall_and_sustained_lag() -> None:
prefix = 'body.resources.event_loop'
samples = [