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
src/langbot/pkg/cloud/bootstrap.py
+5 -1
View File
@@ -11,7 +11,7 @@ from collections.abc import Awaitable, Callable
from typing import Any, Protocol, runtime_checkable
from ..workspace.policy import CloudWorkspacePolicy, SingleWorkspacePolicy
from .directory import DirectoryProjectionProvider
from .directory import DirectoryProjectionProvider, directory_projection_limits_from_config
from .entitlements import EntitlementProvider, OpenSourceEntitlementProvider
@@ -112,6 +112,10 @@ class VerifiedCloudDeployment:
raise CloudBootstrapError('Verified Cloud bootstrap did not provide a Manifest renewal adapter')
def validate_instance_config(self, config: dict[str, Any]) -> None:
try:
directory_projection_limits_from_config(config)
except (TypeError, ValueError) as exc:
raise CloudBootstrapError(f'Cloud directory limits are invalid: {exc}') from exc
if config.get('database', {}).get('use') != 'postgresql':
raise CloudBootstrapError('Cloud runtime requires database.use=postgresql')
if config.get('vdb', {}).get('use') != self.required_vector_backend:
src/langbot/pkg/cloud/directory.py
+125 -29
View File
@@ -7,10 +7,86 @@ from typing import Any, Protocol, runtime_checkable
import pydantic
DEFAULT_MAX_ACTIVE_WORKSPACES = 1_000
HARD_MAX_ACTIVE_WORKSPACES = 5_000
DEFAULT_MAX_SNAPSHOT_WORKSPACES = 1_000
HARD_MAX_SNAPSHOT_WORKSPACES = 5_000
DEFAULT_MAX_SNAPSHOT_MEMBERSHIPS = 20_000
HARD_MAX_SNAPSHOT_MEMBERSHIPS = 100_000
DEFAULT_MAX_CONTROL_PLANE_RESPONSE_BYTES = 32 * 1024 * 1024
HARD_MAX_CONTROL_PLANE_RESPONSE_BYTES = 64 * 1024 * 1024
class DirectoryProjectionUnavailableError(RuntimeError):
"""Raised when the verified Cloud directory cannot safely admit work."""
class DirectoryProjectionLimits(pydantic.BaseModel):
"""Instance-owned cardinality limits for verified Cloud directory data.
These are operational safety limits, not subscription entitlements. Core
fails the complete projection transaction when a limit is exceeded instead
of truncating an authoritative directory and accidentally hiding tenants.
The closed adapter consumes the same limits before priming entitlement
caches, and additionally bounds the HTTP response buffered for signature
verification.
"""
model_config = pydantic.ConfigDict(frozen=True, extra='forbid')
max_active_workspaces: int = pydantic.Field(
default=DEFAULT_MAX_ACTIVE_WORKSPACES,
ge=1,
le=HARD_MAX_ACTIVE_WORKSPACES,
)
max_snapshot_workspaces: int = pydantic.Field(
default=DEFAULT_MAX_SNAPSHOT_WORKSPACES,
ge=1,
le=HARD_MAX_SNAPSHOT_WORKSPACES,
)
max_snapshot_memberships: int = pydantic.Field(
default=DEFAULT_MAX_SNAPSHOT_MEMBERSHIPS,
ge=1,
le=HARD_MAX_SNAPSHOT_MEMBERSHIPS,
)
max_response_bytes: int = pydantic.Field(
default=DEFAULT_MAX_CONTROL_PLANE_RESPONSE_BYTES,
ge=1024 * 1024,
le=HARD_MAX_CONTROL_PLANE_RESPONSE_BYTES,
)
@pydantic.field_validator(
'max_active_workspaces',
'max_snapshot_workspaces',
'max_snapshot_memberships',
'max_response_bytes',
mode='before',
)
@classmethod
def _reject_boolean_limits(cls, value: object) -> object:
if isinstance(value, bool):
raise ValueError('must be an integer')
return value
@pydantic.model_validator(mode='after')
def _validate_workspace_limits(self) -> DirectoryProjectionLimits:
if self.max_snapshot_workspaces < self.max_active_workspaces:
raise ValueError('max_snapshot_workspaces must be greater than or equal to max_active_workspaces')
return self
def directory_projection_limits_from_config(config: dict[str, Any]) -> DirectoryProjectionLimits:
"""Parse typed Cloud directory limits from the instance configuration."""
cloud_config = config.get('cloud', {})
if not isinstance(cloud_config, dict):
raise ValueError('cloud must be a mapping')
directory_config = cloud_config.get('directory', {})
if not isinstance(directory_config, dict):
raise ValueError('cloud.directory must be a mapping')
return DirectoryProjectionLimits.model_validate(directory_config)
class DirectoryMember(pydantic.BaseModel):
"""One account membership published by the SaaS control plane."""
@@ -48,7 +124,10 @@ class DirectoryWorkspace(pydantic.BaseModel):
created_by_account_uuid: str = pydantic.Field(min_length=1, max_length=36)
projection_revision: int = pydantic.Field(ge=1)
execution_generation: int = pydantic.Field(ge=1)
members: tuple[DirectoryMember, ...] = ()
members: tuple[DirectoryMember, ...] = pydantic.Field(
default=(),
max_length=HARD_MAX_SNAPSHOT_MEMBERSHIPS,
)
@pydantic.field_validator('members', mode='before')
@classmethod
@@ -57,13 +136,16 @@ class DirectoryWorkspace(pydantic.BaseModel):
@pydantic.model_validator(mode='after')
def _validate_members(self) -> DirectoryWorkspace:
membership_uuids = [member.membership_uuid for member in self.members]
account_uuids = [member.account_uuid for member in self.members]
if len(membership_uuids) != len(set(membership_uuids)):
raise ValueError('Directory Workspace contains duplicate membership UUIDs')
if len(account_uuids) != len(set(account_uuids)):
raise ValueError('Directory Workspace contains duplicate account UUIDs')
if self.created_by_account_uuid not in set(account_uuids):
membership_uuids: set[str] = set()
account_uuids: set[str] = set()
for member in self.members:
if member.membership_uuid in membership_uuids:
raise ValueError('Directory Workspace contains duplicate membership UUIDs')
if member.account_uuid in account_uuids:
raise ValueError('Directory Workspace contains duplicate account UUIDs')
membership_uuids.add(member.membership_uuid)
account_uuids.add(member.account_uuid)
if self.created_by_account_uuid not in account_uuids:
raise ValueError('Directory Workspace must include its creator')
return self
@@ -76,7 +158,10 @@ class DirectorySnapshot(pydantic.BaseModel):
instance_uuid: str = pydantic.Field(min_length=1, max_length=255)
cursor: int = pydantic.Field(ge=0)
generated_at: datetime.datetime
workspaces: tuple[DirectoryWorkspace, ...] = ()
workspaces: tuple[DirectoryWorkspace, ...] = pydantic.Field(
default=(),
max_length=HARD_MAX_SNAPSHOT_WORKSPACES,
)
@pydantic.field_validator('workspaces', mode='before')
@classmethod
@@ -85,15 +170,20 @@ class DirectorySnapshot(pydantic.BaseModel):
@pydantic.model_validator(mode='after')
def _validate_workspaces(self) -> DirectorySnapshot:
workspace_uuids = [workspace.uuid for workspace in self.workspaces]
slugs = [workspace.slug for workspace in self.workspaces]
membership_uuids = [member.membership_uuid for workspace in self.workspaces for member in workspace.members]
if len(workspace_uuids) != len(set(workspace_uuids)):
raise ValueError('Directory snapshot contains duplicate Workspace UUIDs')
if len(slugs) != len(set(slugs)):
raise ValueError('Directory snapshot contains duplicate Workspace slugs')
if len(membership_uuids) != len(set(membership_uuids)):
raise ValueError('Directory snapshot contains duplicate membership UUIDs')
workspace_uuids: set[str] = set()
slugs: set[str] = set()
membership_uuids: set[str] = set()
for workspace in self.workspaces:
if workspace.uuid in workspace_uuids:
raise ValueError('Directory snapshot contains duplicate Workspace UUIDs')
if workspace.slug in slugs:
raise ValueError('Directory snapshot contains duplicate Workspace slugs')
workspace_uuids.add(workspace.uuid)
slugs.add(workspace.slug)
for member in workspace.members:
if member.membership_uuid in membership_uuids:
raise ValueError('Directory snapshot contains duplicate membership UUIDs')
membership_uuids.add(member.membership_uuid)
return self
@@ -127,17 +217,23 @@ class DirectoryDelta(pydantic.BaseModel):
if len(requested) != len(set(requested)):
raise ValueError('Directory delta contains duplicate requested Workspace UUIDs')
workspace_uuids = [workspace.uuid for workspace in self.workspaces]
slugs = [workspace.slug for workspace in self.workspaces]
membership_uuids = [member.membership_uuid for workspace in self.workspaces for member in workspace.members]
if len(workspace_uuids) != len(set(workspace_uuids)):
raise ValueError('Directory delta contains duplicate Workspace UUIDs')
if not set(workspace_uuids).issubset(set(requested)):
raise ValueError('Directory delta returned an unrequested Workspace')
if len(slugs) != len(set(slugs)):
raise ValueError('Directory delta contains duplicate Workspace slugs')
if len(membership_uuids) != len(set(membership_uuids)):
raise ValueError('Directory delta contains duplicate membership UUIDs')
requested_set = set(requested)
workspace_uuids: set[str] = set()
slugs: set[str] = set()
membership_uuids: set[str] = set()
for workspace in self.workspaces:
if workspace.uuid in workspace_uuids:
raise ValueError('Directory delta contains duplicate Workspace UUIDs')
if workspace.uuid not in requested_set:
raise ValueError('Directory delta returned an unrequested Workspace')
if workspace.slug in slugs:
raise ValueError('Directory delta contains duplicate Workspace slugs')
workspace_uuids.add(workspace.uuid)
slugs.add(workspace.slug)
for member in workspace.members:
if member.membership_uuid in membership_uuids:
raise ValueError('Directory delta contains duplicate membership UUIDs')
membership_uuids.add(member.membership_uuid)
return self
src/langbot/pkg/cloud/directory_projection.py
+184 -15
View File
@@ -29,6 +29,7 @@ from .directory import (
DirectoryEvent,
DirectoryEventBatch,
DirectoryMember,
DirectoryProjectionLimits,
DirectoryProjectionProvider,
DirectoryProjectionUnavailableError,
DirectorySnapshot,
@@ -65,6 +66,7 @@ _MEMBERSHIP_STATUS_MAP = {
'removed': MembershipStatus.REMOVED.value,
}
_INCREMENTAL_PROJECTION_FINGERPRINT = hashlib.sha256(b'langbot-directory-incremental-v1').hexdigest()
_ACCOUNT_QUERY_CHUNK_SIZE = 500
class _DirectorySnapshotSuperseded(DirectoryProjectionUnavailableError):
@@ -89,6 +91,7 @@ class DirectoryProjectionService:
sync_interval_seconds: float = 5.0,
max_staleness_seconds: float = 60.0,
event_limit: int = 100,
limits: DirectoryProjectionLimits | None = None,
monotonic_time: Callable[[], float] = time.monotonic,
) -> None:
if not isinstance(provider, DirectoryProjectionProvider):
@@ -101,15 +104,21 @@ class DirectoryProjectionService:
raise ValueError('Directory max staleness must exceed the sync interval')
if event_limit <= 0 or event_limit > 100:
raise ValueError('Directory event limit must be between 1 and 100')
if limits is not None and not isinstance(limits, DirectoryProjectionLimits):
raise TypeError('Directory projection limits must be a DirectoryProjectionLimits value')
self.ap = ap
self.provider = provider
self.instance_uuid = instance_uuid.strip()
self.sync_interval_seconds = float(sync_interval_seconds)
self.max_staleness_seconds = float(max_staleness_seconds)
self.event_limit = event_limit
self.limits = limits or DirectoryProjectionLimits()
self._monotonic_time = monotonic_time
self._last_success_monotonic: float | None = None
self._ready = False
self._active_workspace_count = 0
self._last_batch_workspace_count = 0
self._last_batch_membership_count = 0
# Every runtime replica must consume the event stream independently:
# entitlement snapshots live in the closed adapter's process memory.
# The database cursor remains the shared projection high-water mark,
@@ -184,6 +193,90 @@ class DirectoryProjectionService:
if self._monotonic_time() - last_success >= self.max_staleness_seconds:
raise DirectoryProjectionUnavailableError('Cloud directory projection is stale')
def resource_snapshot(self) -> dict[str, int]:
"""Return aggregate, tenant-free cardinality gauges for health checks."""
return {
'active_workspaces': self._active_workspace_count,
'max_active_workspaces': self.limits.max_active_workspaces,
'last_batch_workspaces': self._last_batch_workspace_count,
'last_batch_memberships': self._last_batch_membership_count,
'max_snapshot_workspaces': self.limits.max_snapshot_workspaces,
'max_snapshot_memberships': self.limits.max_snapshot_memberships,
}
def _validate_batch_capacity(
self,
workspaces: tuple[DirectoryWorkspace, ...],
*,
full_snapshot: bool,
) -> tuple[int, int]:
workspace_count = len(workspaces)
if full_snapshot and workspace_count > self.limits.max_snapshot_workspaces:
raise DirectoryProjectionUnavailableError(
'Directory snapshot Workspace capacity exceeded '
f'({workspace_count} > {self.limits.max_snapshot_workspaces})'
)
active_count = 0
membership_count = 0
for workspace in workspaces:
if workspace.status == WorkspaceStatus.ACTIVE.value:
active_count += 1
membership_count += len(workspace.members)
if membership_count > self.limits.max_snapshot_memberships:
raise DirectoryProjectionUnavailableError(
'Directory membership capacity exceeded '
f'({membership_count} > {self.limits.max_snapshot_memberships})'
)
if active_count > self.limits.max_active_workspaces:
raise DirectoryProjectionUnavailableError(
f'Directory active Workspace capacity exceeded ({active_count} > {self.limits.max_active_workspaces})'
)
return workspace_count, membership_count
async def _enforce_active_workspace_capacity(self, session: Any) -> int:
"""Count the committed candidate state while holding the projection lock.
Full snapshots can validate their own active count before doing any
database work. Incremental deltas cannot know the instance total, so
every projection path also checks the database after applying fences.
The caller holds the per-instance DirectoryProjectionState row lock;
concurrent replicas therefore cannot race two individually-admitted
deltas above the instance ceiling.
"""
active_count = int(
(
await session.scalar(
sqlalchemy.select(sqlalchemy.func.count())
.select_from(Workspace)
.where(
Workspace.instance_uuid == self.instance_uuid,
Workspace.source == WorkspaceSource.CLOUD_PROJECTION.value,
Workspace.status == WorkspaceStatus.ACTIVE.value,
)
)
)
or 0
)
if active_count > self.limits.max_active_workspaces:
raise DirectoryProjectionUnavailableError(
f'Projected active Workspace capacity exceeded ({active_count} > {self.limits.max_active_workspaces})'
)
return active_count
def _record_batch_cardinality(
self,
*,
active_workspaces: int,
workspaces: int,
memberships: int,
) -> None:
self._active_workspace_count = active_workspaces
self._last_batch_workspace_count = workspaces
self._last_batch_membership_count = memberships
async def apply_snapshot(
self,
snapshot: DirectorySnapshot,
@@ -194,6 +287,10 @@ class DirectoryProjectionService:
if not isinstance(snapshot, DirectorySnapshot):
raise DirectoryProjectionUnavailableError('Directory provider returned an invalid snapshot')
workspace_count, membership_count = self._validate_batch_capacity(
snapshot.workspaces,
full_snapshot=True,
)
snapshot = DirectorySnapshot.model_validate(snapshot.model_dump())
if snapshot.instance_uuid != self.instance_uuid:
raise DirectoryProjectionUnavailableError('Directory snapshot targets another LangBot instance')
@@ -244,9 +341,10 @@ class DirectoryProjectionService:
raise DirectoryProjectionUnavailableError('Directory snapshot cursor has conflicting contents')
await self._record_events(session, event_list, now=now)
await self._apply_accounts(session, snapshot)
await self._apply_workspaces(session, snapshot)
accounts_by_uuid = await self._apply_accounts(session, snapshot)
await self._apply_workspaces(session, snapshot, accounts_by_uuid=accounts_by_uuid)
await self._fence_absent_workspaces(session, snapshot)
active_workspace_count = await self._enforce_active_workspace_capacity(session)
state.cursor = snapshot.cursor
state.snapshot_coverage_cursor = snapshot.cursor
@@ -260,6 +358,11 @@ class DirectoryProjectionService:
await self._reconcile_entitlement_snapshot_set(snapshot)
self._publish_runtime_execution_projection(snapshot.workspaces)
self._record_batch_cardinality(
active_workspaces=active_workspace_count,
workspaces=workspace_count,
memberships=membership_count,
)
self._record_success()
self._consumer_cursor = snapshot.cursor
@@ -270,6 +373,10 @@ class DirectoryProjectionService:
raise DirectoryProjectionUnavailableError('Directory provider returned an invalid delta')
if not isinstance(batch, DirectoryEventBatch):
raise DirectoryProjectionUnavailableError('Directory provider returned an invalid event batch')
workspace_count, membership_count = self._validate_batch_capacity(
delta.workspaces,
full_snapshot=False,
)
delta = DirectoryDelta.model_validate(delta.model_dump())
batch = DirectoryEventBatch.model_validate(batch.model_dump())
self._validate_batch(batch, expected_after_cursor=batch.after_cursor)
@@ -325,8 +432,12 @@ class DirectoryProjectionService:
generated_at=delta.generated_at,
workspaces=delta.workspaces,
)
await self._apply_accounts(session, projected_delta)
await self._apply_workspaces(session, projected_delta)
accounts_by_uuid = await self._apply_accounts(session, projected_delta)
await self._apply_workspaces(
session,
projected_delta,
accounts_by_uuid=accounts_by_uuid,
)
await self._fence_workspaces(
session,
{
@@ -340,6 +451,7 @@ class DirectoryProjectionService:
# incremental path; a later full snapshot replaces this marker.
state.snapshot_fingerprint = _INCREMENTAL_PROJECTION_FINGERPRINT
active_workspace_count = await self._enforce_active_workspace_capacity(session)
state.last_applied_at = now
state.lease_expires_at = lease_expires_at
await self._mark_events_applied(session, batch.events, now=now)
@@ -354,6 +466,11 @@ class DirectoryProjectionService:
returned.values(),
affected_workspace_uuids=requested,
)
self._record_batch_cardinality(
active_workspaces=active_workspace_count,
workspaces=workspace_count,
memberships=membership_count,
)
if projection_caught_up:
self._record_success()
self._consumer_cursor = batch.cursor
@@ -580,7 +697,7 @@ class DirectoryProjectionService:
for row in inbox_rows:
row.applied_at = now
async def _apply_accounts(self, session: Any, snapshot: DirectorySnapshot) -> None:
async def _apply_accounts(self, session: Any, snapshot: DirectorySnapshot) -> dict[str, User]:
selected: dict[str, DirectoryMember] = {}
emails: dict[str, str] = {}
for workspace in snapshot.workspaces:
@@ -596,11 +713,35 @@ class DirectoryProjectionService:
if previous is None:
selected[member.account_uuid] = member
# Fetch existing UUID and email owners in bounded batches. The previous
# two SELECTs per unique account made a large but valid directory
# snapshot produce tens of thousands of serial round trips during
# startup. The configured membership ceiling bounds the materialized
# maps, while batching stays below PostgreSQL parameter limits.
accounts_by_uuid: dict[str, User] = {}
accounts_by_email: dict[str, User] = {}
selected_items = list(selected.items())
for start in range(0, len(selected_items), _ACCOUNT_QUERY_CHUNK_SIZE):
chunk = selected_items[start : start + _ACCOUNT_QUERY_CHUNK_SIZE]
account_uuids = [account_uuid for account_uuid, _member in chunk]
normalized_emails = [member.normalized_email for _account_uuid, member in chunk]
rows = (
await session.scalars(
sqlalchemy.select(User).where(
sqlalchemy.or_(
User.uuid.in_(account_uuids),
User.normalized_email.in_(normalized_emails),
)
)
)
).all()
for account in rows:
accounts_by_uuid[account.uuid] = account
accounts_by_email[account.normalized_email] = account
for account_uuid, member in selected.items():
account = await session.scalar(sqlalchemy.select(User).where(User.uuid == account_uuid))
email_account = await session.scalar(
sqlalchemy.select(User).where(User.normalized_email == member.normalized_email)
)
account = accounts_by_uuid.get(account_uuid)
email_account = accounts_by_email.get(member.normalized_email)
if email_account is not None and email_account.uuid != account_uuid:
raise DirectoryProjectionUnavailableError('Directory account email collides with another Core account')
if account is None:
@@ -616,6 +757,8 @@ class DirectoryProjectionService:
space_account_uuid=account_uuid,
)
session.add(account)
accounts_by_uuid[account_uuid] = account
accounts_by_email[member.normalized_email] = account
continue
if account.source != AccountSource.CLOUD_PROJECTION.value:
raise DirectoryProjectionUnavailableError('Directory account UUID collides with a local Core account')
@@ -637,8 +780,15 @@ class DirectoryProjectionService:
account.account_type = 'space'
account.space_account_uuid = account_uuid
await session.flush()
return accounts_by_uuid
async def _apply_workspaces(self, session: Any, snapshot: DirectorySnapshot) -> None:
async def _apply_workspaces(
self,
session: Any,
snapshot: DirectorySnapshot,
*,
accounts_by_uuid: dict[str, User],
) -> None:
for candidate in snapshot.workspaces:
workspace = await session.get(Workspace, candidate.uuid)
if workspace is None:
@@ -649,7 +799,7 @@ class DirectoryProjectionService:
slug=candidate.slug,
type=candidate.type,
status=candidate.status,
created_by_account_uuid=await self._projected_creator_uuid(session, candidate),
created_by_account_uuid=self._projected_creator_uuid(candidate, accounts_by_uuid),
source=WorkspaceSource.CLOUD_PROJECTION.value,
projection_revision=candidate.projection_revision,
)
@@ -661,14 +811,18 @@ class DirectoryProjectionService:
workspace.slug = candidate.slug
workspace.type = candidate.type
workspace.status = candidate.status
workspace.created_by_account_uuid = await self._projected_creator_uuid(session, candidate)
workspace.created_by_account_uuid = self._projected_creator_uuid(candidate, accounts_by_uuid)
workspace.projection_revision = candidate.projection_revision
await self._apply_memberships(session, workspace, candidate)
await self._apply_execution_state(session, workspace, candidate)
async def _projected_creator_uuid(self, session: Any, candidate: DirectoryWorkspace) -> str | None:
creator = await session.scalar(sqlalchemy.select(User).where(User.uuid == candidate.created_by_account_uuid))
@staticmethod
def _projected_creator_uuid(
candidate: DirectoryWorkspace,
accounts_by_uuid: dict[str, User],
) -> str | None:
creator = accounts_by_uuid.get(candidate.created_by_account_uuid)
if creator is None:
if candidate.status == WorkspaceStatus.ACTIVE.value:
raise DirectoryProjectionUnavailableError('Active Directory Workspace creator is not projected')
@@ -801,9 +955,24 @@ class DirectoryProjectionService:
included = {workspace.uuid for workspace in snapshot.workspaces}
projected = (
await session.scalars(
sqlalchemy.select(Workspace).where(
sqlalchemy.select(Workspace)
.outerjoin(
WorkspaceExecutionState,
WorkspaceExecutionState.workspace_uuid == Workspace.uuid,
)
.where(
Workspace.instance_uuid == self.instance_uuid,
Workspace.source == WorkspaceSource.CLOUD_PROJECTION.value,
sqlalchemy.or_(
Workspace.status.not_in(
(
WorkspaceStatus.ARCHIVED.value,
WorkspaceStatus.DELETED.value,
)
),
WorkspaceExecutionState.state == WorkspaceExecutionStatus.ACTIVE.value,
WorkspaceExecutionState.write_fenced == sqlalchemy.false(),
),
)
)
).all()
src/langbot/pkg/core/app.py
+12
View File
@@ -268,11 +268,23 @@ class Application:
}
)
directory_stats = {}
directory_snapshot = getattr(self.directory_projection_service, 'resource_snapshot', None)
if callable(directory_snapshot):
directory_stats = directory_snapshot()
database_stats = {}
database_snapshot = getattr(self.persistence_mgr, 'get_resource_stats', None)
if callable(database_snapshot):
database_stats = database_snapshot()
return {
'asyncio_tasks': asyncio_tasks,
'event_loop': self.event_loop_monitor.snapshot(),
'blocking_executor': (self.blocking_executor.snapshot() if self.blocking_executor is not None else {}),
'application_tasks': task_stats,
'database_pool': database_stats,
'directory': directory_stats,
'query_pool': query_pool_stats,
'models': model_stats,
'runtimes': runtime_stats,
src/langbot/pkg/core/stages/build_app.py
+2
View File
@@ -42,6 +42,7 @@ from ...workspace import collaboration as workspace_collaboration_module
from ...workspace import invitation_delivery as invitation_delivery_module
from ...cloud import bootstrap as cloud_bootstrap
from ...cloud import launch as cloud_launch_module
from ...cloud.directory import directory_projection_limits_from_config
from ...cloud.directory_projection import DirectoryProjectionService
from ...cloud.entitlements import EntitlementResolver
from ...api.http.context import ExecutionContext, PrincipalContext, PrincipalType
@@ -154,6 +155,7 @@ class BuildAppStage(stage.BootingStage):
ap,
deployment.directory_provider,
constants.instance_id,
limits=directory_projection_limits_from_config(ap.instance_config.data),
)
await directory_projection_service.initialize()
ap.directory_projection_service = directory_projection_service
src/langbot/pkg/core/stages/load_config.py
+20 -1
View File
@@ -14,6 +14,25 @@ from ..bootutils import config
_RUNTIME_POLICY_DEFAULTS = {
'cloud': {
'directory': {
'max_active_workspaces': 1000,
'max_snapshot_workspaces': 1000,
'max_snapshot_memberships': 20000,
'max_response_bytes': 33554432,
}
},
'database': {
'postgresql': {
'pool_size': 10,
'max_overflow': 10,
'pool_timeout_seconds': 30,
'pool_recycle_seconds': 1800,
'statement_timeout_ms': 60000,
'lock_timeout_ms': 5000,
'idle_in_transaction_session_timeout_ms': 60000,
}
},
'system': {
'blocking_executor': {
'max_workers': bounded_executor.DEFAULT_MAX_WORKERS,
@@ -49,7 +68,7 @@ def _complete_runtime_policy_defaults(cfg: dict) -> dict:
The historic config loader intentionally does not deep-complete the whole
template. These fields are different: their native env overrides must
retain boolean/numeric types on upgraded instances, so their defaults must
exist before ``PLUGIN__...`` and ``MCP__...`` are parsed.
exist before ``CLOUD__...``, ``PLUGIN__...`` and ``MCP__...`` are parsed.
"""
def merge(target: dict, defaults: dict, path: tuple[str, ...] = ()) -> None:
src/langbot/pkg/persistence/database.py
+1
View File
@@ -39,6 +39,7 @@ class BaseDatabaseManager(abc.ABC):
) -> None:
self.ap = ap
self.url_override = url_override
self.persistence_mode: str | None = None
@abc.abstractmethod
async def initialize(self) -> None:
src/langbot/pkg/persistence/databases/postgresql.py
+110 -26
View File
@@ -7,6 +7,14 @@ from .. import database
from ..postgresql_url import normalize_asyncpg_url
MAX_POOL_CONNECTIONS = 100
MAX_POOL_TIMEOUT_SECONDS = 300
MAX_POOL_RECYCLE_SECONDS = 86_400
MAX_STATEMENT_TIMEOUT_MS = 300_000
MAX_LOCK_TIMEOUT_MS = 60_000
MAX_IDLE_TRANSACTION_TIMEOUT_MS = 300_000
@database.manager_class('postgresql')
class PostgreSQLDatabaseManager(database.BaseDatabaseManager):
"""PostgreSQL database manager"""
@@ -18,14 +26,16 @@ class PostgreSQLDatabaseManager(database.BaseDatabaseManager):
default: int,
*,
minimum: int,
maximum: int,
) -> int:
value = config.get(name, default)
if isinstance(value, bool) or not isinstance(value, int) or value < minimum:
if isinstance(value, bool) or not isinstance(value, int) or not minimum <= value <= maximum:
comparator = 'non-negative' if minimum == 0 else 'positive'
raise ValueError(f'database.postgresql.{name} must be a {comparator} integer')
raise ValueError(f'database.postgresql.{name} must be a {comparator} integer no greater than {maximum}')
return value
async def initialize(self) -> None:
self._pool_timeouts_total = 0
postgresql_config = self.ap.instance_config.data.get('database', {}).get('postgresql', {})
if not isinstance(postgresql_config, dict):
raise ValueError('database.postgresql must be an object')
@@ -53,31 +63,105 @@ class PostgreSQLDatabaseManager(database.BaseDatabaseManager):
port=postgresql_config.get('port', 5432),
database=postgresql_config.get('database', 'postgres'),
)
self.pool_size = self._pool_integer(
postgresql_config,
'pool_size',
10,
minimum=1,
maximum=MAX_POOL_CONNECTIONS,
)
self.max_overflow = self._pool_integer(
postgresql_config,
'max_overflow',
10,
minimum=0,
maximum=MAX_POOL_CONNECTIONS,
)
if self.pool_size + self.max_overflow > MAX_POOL_CONNECTIONS:
raise ValueError(f'database.postgresql pool_size + max_overflow must not exceed {MAX_POOL_CONNECTIONS}')
self.pool_timeout_seconds = self._pool_integer(
postgresql_config,
'pool_timeout_seconds',
30,
minimum=1,
maximum=MAX_POOL_TIMEOUT_SECONDS,
)
self.pool_recycle_seconds = self._pool_integer(
postgresql_config,
'pool_recycle_seconds',
1800,
minimum=1,
maximum=MAX_POOL_RECYCLE_SECONDS,
)
connect_args = {}
self.statement_timeout_ms = 0
self.lock_timeout_ms = 0
self.idle_transaction_timeout_ms = 0
if self.persistence_mode == 'cloud_runtime':
self.statement_timeout_ms = self._pool_integer(
postgresql_config,
'statement_timeout_ms',
60_000,
minimum=1,
maximum=MAX_STATEMENT_TIMEOUT_MS,
)
self.lock_timeout_ms = self._pool_integer(
postgresql_config,
'lock_timeout_ms',
5_000,
minimum=1,
maximum=MAX_LOCK_TIMEOUT_MS,
)
self.idle_transaction_timeout_ms = self._pool_integer(
postgresql_config,
'idle_in_transaction_session_timeout_ms',
60_000,
minimum=1,
maximum=MAX_IDLE_TRANSACTION_TIMEOUT_MS,
)
connect_args = {
'server_settings': {
'statement_timeout': str(self.statement_timeout_ms),
'lock_timeout': str(self.lock_timeout_ms),
'idle_in_transaction_session_timeout': str(self.idle_transaction_timeout_ms),
}
}
self.engine = sqlalchemy_asyncio.create_async_engine(
engine_url,
pool_size=self._pool_integer(
postgresql_config,
'pool_size',
10,
minimum=1,
),
max_overflow=self._pool_integer(
postgresql_config,
'max_overflow',
10,
minimum=0,
),
pool_timeout=self._pool_integer(
postgresql_config,
'pool_timeout_seconds',
30,
minimum=1,
),
pool_recycle=self._pool_integer(
postgresql_config,
'pool_recycle_seconds',
1800,
minimum=1,
),
pool_size=self.pool_size,
max_overflow=self.max_overflow,
pool_timeout=self.pool_timeout_seconds,
pool_recycle=self.pool_recycle_seconds,
pool_pre_ping=True,
**({'connect_args': connect_args} if connect_args else {}),
)
def resource_stats(self) -> dict[str, int]:
"""Return aggregate pool gauges without exposing connection details."""
pool = self.engine.pool
def read(name: str) -> int:
method = getattr(pool, name, None)
if not callable(method):
return 0
try:
return int(method())
except Exception:
return 0
return {
'configured_size': self.pool_size,
'configured_max_overflow': self.max_overflow,
'configured_capacity': self.pool_size + self.max_overflow,
'statement_timeout_ms': self.statement_timeout_ms,
'lock_timeout_ms': self.lock_timeout_ms,
'idle_in_transaction_session_timeout_ms': self.idle_transaction_timeout_ms,
'timeouts_total': self._pool_timeouts_total,
'checked_in': read('checkedin'),
'checked_out': read('checkedout'),
'overflow': max(read('overflow'), 0),
}
def record_pool_timeout(self) -> None:
self._pool_timeouts_total += 1
src/langbot/pkg/persistence/mgr.py
+14
View File
@@ -159,6 +159,7 @@ class PersistenceManager:
for manager in database.preregistered_managers:
if manager.name == database_type:
self.db = manager(self.ap, url_override=self._database_url_override)
self.db.persistence_mode = self.mode.value
await self.db.initialize()
selected_manager = self.db
break
@@ -195,6 +196,17 @@ class PersistenceManager:
if engine is not None:
await engine.dispose()
def get_resource_stats(self) -> dict[str, int]:
"""Return database-manager-owned aggregate resource gauges."""
resource_stats = getattr(getattr(self, 'db', None), 'resource_stats', None)
if not callable(resource_stats):
return {}
try:
return resource_stats()
except Exception:
return {}
@contextlib.asynccontextmanager
async def _release_migration_lock(self) -> typing.AsyncIterator[None]:
"""Serialize the complete PostgreSQL migration and validation window."""
@@ -1865,11 +1877,13 @@ class PersistenceManager:
return active.session
def _scoped_uow(self, scope: PersistenceScope) -> TenantUnitOfWork:
on_pool_timeout = getattr(getattr(self, 'db', None), 'record_pool_timeout', None)
return TenantUnitOfWork(
self.get_db_engine(),
scope=scope,
active_transaction=self._active_transaction,
active_scope=self._active_scope,
on_pool_timeout=(on_pool_timeout if callable(on_pool_timeout) else None),
)
def _get_active_transaction(self) -> ActiveScopedTransaction | None:
src/langbot/pkg/persistence/tenant_uow.py
+5 -1
View File
@@ -1220,6 +1220,7 @@ class TenantUnitOfWork:
scope: PersistenceScope | None = None,
active_transaction: ActiveTransactionVar | None = None,
active_scope: ActivePersistenceScopeVar | None = None,
on_pool_timeout: typing.Callable[[], None] | None = None,
) -> None:
if (workspace_uuid is None) == (scope is None):
raise ValueError('TenantUnitOfWork requires exactly one Workspace or persistence scope')
@@ -1229,6 +1230,7 @@ class TenantUnitOfWork:
self.workspace_uuid = self.scope.settings[0][1] if self.scope.kind == PersistenceScopeKind.WORKSPACE else None
self._active_transaction = active_transaction
self._active_scope = active_scope
self._on_pool_timeout = on_pool_timeout
self._session: sqlalchemy_asyncio.AsyncSession | None = None
self._transaction: sqlalchemy_asyncio.AsyncSessionTransaction | None = None
self._active_state: ActiveScopedTransaction | None = None
@@ -1322,7 +1324,9 @@ class TenantUnitOfWork:
self._context_token = self._active_transaction.set(state)
self._database_operation_token = _DATABASE_OPERATION_TRANSACTION.set(state)
self._owns_transaction = True
except BaseException:
except BaseException as exc:
if isinstance(exc, sqlalchemy.exc.TimeoutError) and self._on_pool_timeout is not None:
self._on_pool_timeout()
if self._database_operation_token is not None:
_DATABASE_OPERATION_TRANSACTION.reset(self._database_operation_token)
self._database_operation_token = None
src/langbot/templates/config.yaml
+21
View File
@@ -50,6 +50,22 @@ concurrency:
# Hard admission limits for queued + running pipeline queries.
pending_queries: 1000
pending_queries_per_workspace: 100
cloud:
# Operational safety ceilings for the one logical Cloud instance. These
# are not subscription entitlements. An authoritative directory update
# that would exceed them is rejected atomically rather than truncated.
directory:
# Tune downward from the measured production capacity curve. Core has
# an absolute safety ceiling of 5,000 active Workspaces.
max_active_workspaces: 1000
# Full snapshots contain current Workspaces only. Archived tombstones
# are delivered through bounded per-Workspace deltas.
max_snapshot_workspaces: 1000
# Aggregate memberships accepted in one signed snapshot or delta.
max_snapshot_memberships: 20000
# Signed control-plane envelope buffered by the closed adapter before
# JSON/JWS verification (32 MiB; absolute maximum 64 MiB).
max_response_bytes: 33554432
proxy:
http: ''
https: ''
@@ -135,6 +151,11 @@ database:
max_overflow: 10
pool_timeout_seconds: 30
pool_recycle_seconds: 1800
# Applied only to Cloud runtime connections. The one-shot release
# migration uses its operator connection without these short limits.
statement_timeout_ms: 60000
lock_timeout_ms: 5000
idle_in_transaction_session_timeout_ms: 60000
cloud_migration:
# `langbot migrate --cloud` reads an operator-only PostgreSQL DSN from
# this environment variable. The operator role must differ from the