xinyin025/LangBot
1
0
Fork 0
mirror of https://github.com/langbot-app/LangBot.git synced 2026-08-08 20:30:59 +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
scripts/cloud_runtime_soak.py
+1241
View File
File diff suppressed because it is too large Load Diff
scripts/runtime_resource_probe.py
+466
View File
@@ -0,0 +1,466 @@
#!/usr/bin/env python3
"""Exercise long-lived Core registries and verify that they reach a plateau.
This probe is intentionally separate from the default test suite because the
audit profile creates tens of thousands of historical identities. It uses the
real admission, eviction, and cleanup code while replacing external platform
objects that are irrelevant to registry retention.
"""
from __future__ import annotations
import argparse
import asyncio
import gc
import json
import time
import tracemalloc
from dataclasses import asdict, dataclass
from types import SimpleNamespace
from unittest.mock import patch
import psutil
from langbot.pkg.api.http.context import ExecutionContext
# Import the Application graph before taskmgr. The production boot path has
# this same ordering; importing taskmgr first exposes its historical cycle
# through HTTP route annotations.
from langbot.pkg.core import app as _core_app # noqa: F401
from langbot.pkg.core.taskmgr import AsyncTaskManager
from langbot.pkg.pipeline.pool import QueryPool
from langbot.pkg.pipeline.ratelimit.algos.fixedwin import FixedWindowAlgo
from langbot.pkg.plugin.connector import PluginRuntimeConnector
from langbot.pkg.platform.sources.websocket_adapter import (
WebSocketMessage,
WebSocketSession,
)
from langbot.pkg.provider.modelmgr.modelmgr import ModelManager
from langbot.pkg.provider.session.sessionmgr import SessionManager
from langbot_plugin.api.entities.builtin.provider.session import LauncherTypes
@dataclass(frozen=True, slots=True)
class ProbeScale:
query_churn_per_phase: int
session_churn_per_phase: int
rate_limit_churn_per_phase: int
task_churn_per_phase: int
websocket_churn_per_phase: int
empty_workspace_churn_per_phase: int
SCALES = {
'quick': ProbeScale(
query_churn_per_phase=2_500,
session_churn_per_phase=500,
rate_limit_churn_per_phase=10_000,
task_churn_per_phase=1_000,
websocket_churn_per_phase=500,
empty_workspace_churn_per_phase=1_000,
),
'audit': ProbeScale(
query_churn_per_phase=25_000,
session_churn_per_phase=2_500,
rate_limit_churn_per_phase=10_000,
task_churn_per_phase=5_000,
websocket_churn_per_phase=2_500,
empty_workspace_churn_per_phase=10_000,
),
}
class _ProbeQuery:
"""Small weak-referenceable stand-in for SDK Query construction."""
def __init__(self, **values):
self.__dict__.update(values)
class _EmptyResult:
def all(self) -> list:
return []
class _EmptyPluginRuntimeHandler:
async def reconcile_plugin_installations(self, _states: tuple) -> dict:
return {
'applied': [],
'removed': [],
'missing_artifacts': [],
'failed_installations': [],
}
def unregister_installation_binding(self, _binding) -> None:
raise AssertionError('An empty Workspace exposed an installation binding')
@dataclass(frozen=True, slots=True)
class ProcessSample:
rss_bytes: int
traced_current_bytes: int
traced_peak_bytes: int
asyncio_tasks: int
threads: int
open_fds: int | None
def _sample_process() -> ProcessSample:
gc.collect()
process = psutil.Process()
try:
open_fds = process.num_fds()
except (AttributeError, psutil.Error):
open_fds = None
traced_current, traced_peak = tracemalloc.get_traced_memory()
return ProcessSample(
rss_bytes=process.memory_info().rss,
traced_current_bytes=traced_current,
traced_peak_bytes=traced_peak,
asyncio_tasks=len(asyncio.all_tasks()),
threads=process.num_threads(),
open_fds=open_fds,
)
def _execution_context(index: int, *, query_uuid: str | None = None) -> ExecutionContext:
return ExecutionContext(
instance_uuid='runtime-resource-probe',
workspace_uuid=f'workspace-{index}',
placement_generation=1,
bot_uuid='probe-bot',
pipeline_uuid='probe-pipeline',
query_uuid=query_uuid,
)
class CoreRuntimeProbe:
"""Own the same manager instances across two equal churn phases."""
def __init__(self) -> None:
self.query_pool = QueryPool(max_queries=100, max_queries_per_workspace=1)
app = SimpleNamespace(
event_loop=asyncio.get_running_loop(),
persistence_mgr=None,
instance_config=SimpleNamespace(
data={
'concurrency': {'session': 1},
'system': {
'session_retention': {
'idle_ttl_seconds': 86_400,
'max_entries': 200,
'max_entries_per_workspace': 200,
'max_conversations_per_session': 20,
'max_messages_per_conversation': 100,
},
'task_retention': {
'completed_limit': 200,
'max_log_chars': 4_096,
'max_active_user_tasks': 256,
'max_active_user_tasks_per_workspace': 8,
},
},
}
),
)
self.session_manager = SessionManager(app)
self.task_manager = AsyncTaskManager(app)
self.rate_limit = FixedWindowAlgo(SimpleNamespace())
self.websocket_session = WebSocketSession(
'resource-probe',
max_conversations=200,
max_messages=100,
)
logger = SimpleNamespace(
debug=lambda *_args, **_kwargs: None,
info=lambda *_args, **_kwargs: None,
warning=lambda *_args, **_kwargs: None,
error=lambda *_args, **_kwargs: None,
)
self.empty_model_queries = 0
async def execute_empty(_statement):
self.empty_model_queries += 1
return _EmptyResult()
model_app = SimpleNamespace(
logger=logger,
persistence_mgr=SimpleNamespace(execute_async=execute_empty),
)
self.empty_model_manager = ModelManager(model_app)
async def runtime_disconnect_callback(_connector) -> None:
return None
plugin_app = SimpleNamespace(
instance_config=SimpleNamespace(data={'plugin': {'enable': True}}),
deployment=SimpleNamespace(mode='cloud'),
logger=logger,
)
self.empty_plugin_connector = PluginRuntimeConnector(
plugin_app,
runtime_disconnect_callback,
)
self.empty_plugin_connector.handler = _EmptyPluginRuntimeHandler()
async def validate_context(context):
return context
async def load_desired_states(_context):
return []
self.empty_plugin_connector._validate_execution_context = validate_context
self.empty_plugin_connector._load_workspace_desired_states = load_desired_states
async def initialize(self) -> None:
await self.rate_limit.initialize()
async def run_phase(self, scale: ProbeScale, phase: int) -> None:
offsets = {
'query': (phase - 1) * scale.query_churn_per_phase,
'session': (phase - 1) * scale.session_churn_per_phase,
'rate': (phase - 1) * scale.rate_limit_churn_per_phase,
'task': (phase - 1) * scale.task_churn_per_phase,
'websocket': (phase - 1) * scale.websocket_churn_per_phase,
'empty_workspace': ((phase - 1) * scale.empty_workspace_churn_per_phase),
}
await self._churn_queries(offsets['query'], scale.query_churn_per_phase)
await self._churn_sessions(offsets['session'], scale.session_churn_per_phase)
await self._churn_rate_limits(offsets['rate'], scale.rate_limit_churn_per_phase)
await self._churn_tasks(offsets['task'], scale.task_churn_per_phase)
self._churn_websocket_history(
offsets['websocket'],
scale.websocket_churn_per_phase,
)
await self._churn_empty_workspaces(
offsets['empty_workspace'],
scale.empty_workspace_churn_per_phase,
)
await asyncio.sleep(0)
async def _churn_queries(self, start: int, count: int) -> None:
def make_query(**values):
return _ProbeQuery(**values)
with patch(
'langbot.pkg.pipeline.pool.pipeline_query.Query',
side_effect=make_query,
):
for index in range(start, start + count):
context = _execution_context(index)
query = await self.query_pool.add_query(
bot_uuid='probe-bot',
launcher_type=LauncherTypes.PERSON,
launcher_id=f'launcher-{index}',
sender_id=f'sender-{index}',
message_event=SimpleNamespace(),
message_chain=SimpleNamespace(),
adapter=None,
pipeline_uuid='probe-pipeline',
execution_context=context,
)
removed = await self.query_pool.remove_query(query)
if not removed:
raise AssertionError('Query cleanup failed')
async def _churn_sessions(self, start: int, count: int) -> None:
for index in range(start, start + count):
workspace_index = index % 100
context = _execution_context(
workspace_index,
query_uuid=f'session-query-{index}',
)
query = SimpleNamespace(
launcher_type=LauncherTypes.PERSON,
launcher_id=f'launcher-{index}',
sender_id=f'sender-{index}',
bot_uuid='probe-bot',
pipeline_uuid='probe-pipeline',
query_uuid=context.query_uuid,
_execution_context=context,
)
await self.session_manager.get_session(query)
async def _churn_rate_limits(self, start: int, count: int) -> None:
for index in range(start, start + count):
context = _execution_context(
index % 1_000,
query_uuid=f'rate-query-{index}',
)
query = SimpleNamespace(
bot_uuid='probe-bot',
pipeline_uuid='probe-pipeline',
_execution_context=context,
pipeline_config={
'safety': {
'rate-limit': {
'window-length': 60,
'limitation': 100_000,
'strategy': 'drop',
}
}
},
)
admitted = await self.rate_limit.require_access(
query,
LauncherTypes.PERSON,
f'rate-identity-{index}',
)
if not admitted:
raise AssertionError('Rate-limit registry rejected bounded churn')
async def _churn_tasks(self, start: int, count: int) -> None:
async def complete_immediately() -> None:
return None
for batch_start in range(start, start + count, 256):
batch_size = min(256, start + count - batch_start)
wrappers = [
self.task_manager.create_task(
complete_immediately(),
name=f'resource-probe-{batch_start + offset}',
)
for offset in range(batch_size)
]
await asyncio.gather(*(wrapper.task for wrapper in wrappers))
await asyncio.sleep(0)
def _churn_websocket_history(self, start: int, count: int) -> None:
for index in range(start, start + count):
conversation_key = f'conversation-{index}'
response_id = f'response-{index}'
indexes = self.websocket_session.get_stream_message_indexes(conversation_key)
indexes[response_id] = 0
self.websocket_session.append_message(
conversation_key,
WebSocketMessage(
id=self.websocket_session.next_message_id(conversation_key),
role='assistant',
content='probe',
message_chain=[],
timestamp='1970-01-01T00:00:00+00:00',
is_final=True,
),
)
async def _churn_empty_workspaces(self, start: int, count: int) -> None:
for index in range(start, start + count):
await self.empty_model_manager._load_workspace_models(_execution_context(index))
await self.empty_plugin_connector.reconcile_projected_workspaces(
_execution_context(index) for index in range(start, start + count)
)
def retained_state(self) -> dict[str, int]:
return {
'query_cached': len(self.query_pool.cached_queries),
'query_queued': len(self.query_pool.queries),
'query_active_workspaces': len(self.query_pool.active_query_count_by_workspace),
'query_scope_counters': len(self.query_pool.query_count_by_scope),
'sessions': len(self.session_manager.session_list),
'session_index': len(self.session_manager._session_index),
'rate_limit_containers': len(self.rate_limit.containers),
'task_records': len(self.task_manager.tasks),
'websocket_conversations': len(self.websocket_session.message_lists),
'websocket_stream_indexes': len(self.websocket_session.stream_message_indexes),
'empty_model_scopes': len(self.empty_model_manager._scope_generations),
'empty_model_providers': len(self.empty_model_manager.provider_dict),
'empty_model_llms': len(self.empty_model_manager.llm_model_dict),
'empty_plugin_workspace_sets': len(self.empty_plugin_connector._workspace_installations),
'empty_plugin_installations': len(self.empty_plugin_connector._known_desired_states),
}
def assert_bounded(self) -> None:
state = self.retained_state()
expected_maximums = {
'query_cached': 0,
'query_queued': 0,
'query_active_workspaces': 0,
'query_scope_counters': 100,
'sessions': 200,
'session_index': 200,
'rate_limit_containers': 10_000,
'task_records': 200,
'websocket_conversations': 200,
'websocket_stream_indexes': 200,
'empty_model_scopes': 0,
'empty_model_providers': 0,
'empty_model_llms': 0,
'empty_plugin_workspace_sets': 0,
'empty_plugin_installations': 0,
}
violations = {key: (state[key], maximum) for key, maximum in expected_maximums.items() if state[key] > maximum}
if violations:
raise AssertionError(f'Core retained-state limits failed: {violations}')
async def _run(args: argparse.Namespace) -> dict:
scale = SCALES[args.scale]
tracemalloc.start()
started_at = time.monotonic()
probe = CoreRuntimeProbe()
await probe.initialize()
baseline = _sample_process()
await probe.run_phase(scale, 1)
probe.assert_bounded()
phase_one = _sample_process()
state_one = probe.retained_state()
await probe.run_phase(scale, 2)
probe.assert_bounded()
phase_two = _sample_process()
state_two = probe.retained_state()
if state_two != state_one:
raise AssertionError(f'Core retained state did not plateau: phase_one={state_one}, phase_two={state_two}')
traced_growth = phase_two.traced_current_bytes - phase_one.traced_current_bytes
rss_growth = phase_two.rss_bytes - phase_one.rss_bytes
max_traced_growth = int(args.max_traced_growth_mib * 1024 * 1024)
max_rss_growth = int(args.max_rss_growth_mib * 1024 * 1024)
if traced_growth > max_traced_growth:
raise AssertionError(f'Second-phase traced memory grew by {traced_growth} bytes (limit {max_traced_growth})')
if rss_growth > max_rss_growth:
raise AssertionError(f'Second-phase RSS grew by {rss_growth} bytes (limit {max_rss_growth})')
return {
'component': 'langbot-core',
'scale': args.scale,
'work_per_phase': asdict(scale),
'elapsed_seconds': round(time.monotonic() - started_at, 3),
'samples': {
'baseline': asdict(baseline),
'phase_one': asdict(phase_one),
'phase_two': asdict(phase_two),
},
'second_phase_growth': {
'rss_bytes': rss_growth,
'traced_current_bytes': traced_growth,
},
'retained_state': {
'phase_one': state_one,
'phase_two': state_two,
},
'passed': True,
}
def _parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--scale', choices=tuple(SCALES), default='quick')
parser.add_argument('--max-traced-growth-mib', type=float, default=8.0)
parser.add_argument('--max-rss-growth-mib', type=float, default=64.0)
parser.add_argument('--json', action='store_true', help='Print compact JSON')
return parser.parse_args()
def main() -> None:
args = _parse_args()
result = asyncio.run(_run(args))
if args.json:
print(json.dumps(result, sort_keys=True))
else:
print(json.dumps(result, indent=2, sort_keys=True))
if __name__ == '__main__':
main()
scripts/workspace_runtime_capacity_probe.py
+561
View File
@@ -0,0 +1,561 @@
#!/usr/bin/env python3
"""Measure populated Workspace runtime replacement cost and retention.
Unlike ``runtime_resource_probe.py``, which stresses historical request keys
and empty tenants, this probe keeps one representative Provider, LLM,
Embedding model, Rerank model, Pipeline, Bot, and Knowledge Base per Workspace.
It then advances every Workspace to a new placement generation and verifies
that old runtime objects are closed and collectible while active registry
cardinality remains constant.
"""
from __future__ import annotations
import argparse
import asyncio
import gc
import json
import time
import tracemalloc
import weakref
from dataclasses import asdict, dataclass
from types import SimpleNamespace
import psutil
from langbot.pkg.api.http.context import ExecutionContext
# Match the production import order; importing a leaf manager first exposes a
# historical annotation cycle that the application graph resolves.
from langbot.pkg.core import app as _core_app # noqa: F401
from langbot.pkg.entity.persistence import bot as persistence_bot
from langbot.pkg.entity.persistence import model as persistence_model
from langbot.pkg.entity.persistence import pipeline as persistence_pipeline
from langbot.pkg.entity.persistence import rag as persistence_rag
from langbot.pkg.pipeline.pipelinemgr import PipelineManager
from langbot.pkg.platform.botmgr import PlatformManager
from langbot.pkg.provider.modelmgr import requester
from langbot.pkg.provider.modelmgr.modelmgr import ModelManager
from langbot.pkg.provider.tools.loaders.mcp import MCPLoader
from langbot.pkg.rag.knowledge.kbmgr import RAGManager
from langbot.pkg.workspace.entities import WorkspaceExecutionBinding
@dataclass(frozen=True, slots=True)
class ProbeScale:
workspaces: int
SCALES = {
'quick': ProbeScale(workspaces=250),
'audit': ProbeScale(workspaces=5_000),
}
@dataclass(frozen=True, slots=True)
class ProcessSample:
rss_bytes: int
traced_current_bytes: int
traced_peak_bytes: int
asyncio_tasks: int
threads: int
open_fds: int | None
class _ProbeLogger:
def debug(self, *_args, **_kwargs) -> None:
return None
def info(self, *_args, **_kwargs) -> None:
return None
def warning(self, *_args, **_kwargs) -> None:
return None
def error(self, *_args, **_kwargs) -> None:
return None
class _ProbeWorkspaceService:
instance_uuid = 'runtime-capacity-probe'
def __init__(self) -> None:
self.generations: dict[str, int] = {}
self.binding_lookups = 0
async def get_execution_binding(
self,
workspace_uuid: str,
*,
expected_generation: int | None = None,
) -> WorkspaceExecutionBinding:
self.binding_lookups += 1
generation = self.generations[workspace_uuid]
if expected_generation is not None and expected_generation != generation:
raise AssertionError(f'stale probe generation {expected_generation} != {generation}')
return WorkspaceExecutionBinding(
instance_uuid=self.instance_uuid,
workspace_uuid=workspace_uuid,
placement_generation=generation,
write_fenced=False,
state='active',
)
class _ProbeRequester(requester.ProviderAPIRequester):
name = 'capacity-probe'
closed = 0
async def invoke_llm(
self,
query,
model,
messages,
funcs=None,
extra_args=None,
remove_think=False,
):
return None
async def aclose(self) -> None:
type(self).closed += 1
class _ProbeAdapter:
killed = 0
def __init__(self, _config, _logger) -> None:
self.listeners = []
def register_listener(self, event_type, listener) -> None:
self.listeners.append((event_type, listener))
async def kill(self) -> None:
type(self).killed += 1
class _ProbeMCPSession:
closed = 0
def __init__(self, server_name: str) -> None:
self.server_name = server_name
async def shutdown(self) -> None:
type(self).closed += 1
def _sample_process() -> ProcessSample:
gc.collect()
process = psutil.Process()
try:
open_fds = process.num_fds()
except (AttributeError, psutil.Error):
open_fds = None
traced_current, traced_peak = tracemalloc.get_traced_memory()
return ProcessSample(
rss_bytes=process.memory_info().rss,
traced_current_bytes=traced_current,
traced_peak_bytes=traced_peak,
asyncio_tasks=len(asyncio.all_tasks()),
threads=process.num_threads(),
open_fds=open_fds,
)
class PopulatedWorkspaceProbe:
def __init__(self) -> None:
_ProbeRequester.closed = 0
_ProbeAdapter.killed = 0
_ProbeMCPSession.closed = 0
self.workspace_service = _ProbeWorkspaceService()
self.logger = _ProbeLogger()
self.app = SimpleNamespace(
logger=self.logger,
workspace_service=self.workspace_service,
persistence_mgr=SimpleNamespace(
mode=SimpleNamespace(value='cloud_runtime'),
),
pipeline_config_meta_trigger={'name': 'trigger', 'stages': []},
pipeline_config_meta_safety={'name': 'safety', 'stages': []},
pipeline_config_meta_ai={'name': 'ai', 'stages': []},
pipeline_config_meta_output={'name': 'output', 'stages': []},
task_mgr=SimpleNamespace(
cancel_by_scope=lambda *_args, **_kwargs: None,
cancel_task=lambda *_args, **_kwargs: None,
),
)
self.model_manager = ModelManager(self.app)
self.model_manager.requester_dict = {
_ProbeRequester.name: _ProbeRequester,
}
self.pipeline_manager = PipelineManager(self.app)
self.pipeline_manager.stage_dict = {}
self.rag_manager = RAGManager(self.app)
self.mcp_loader = MCPLoader(self.app)
self.platform_manager = PlatformManager(self.app)
self.platform_manager.adapter_dict = {
'capacity-probe': _ProbeAdapter,
}
self.generation_refs: dict[
int,
list[weakref.ReferenceType],
] = {}
def _context(
self,
workspace_uuid: str,
generation: int,
*,
bot_uuid: str | None = None,
pipeline_uuid: str | None = None,
) -> ExecutionContext:
return ExecutionContext(
instance_uuid=self.workspace_service.instance_uuid,
workspace_uuid=workspace_uuid,
placement_generation=generation,
bot_uuid=bot_uuid,
pipeline_uuid=pipeline_uuid,
)
async def load_generation(self, workspaces: int, generation: int) -> None:
for index in range(workspaces):
workspace_uuid = f'workspace-{index}'
provider_uuid = f'provider-{index}'
llm_uuid = f'llm-{index}'
embedding_uuid = f'embedding-{index}'
rerank_uuid = f'rerank-{index}'
pipeline_uuid = f'pipeline-{index}'
bot_uuid = f'bot-{index}'
kb_uuid = f'knowledge-{index}'
mcp_server_name = f'mcp-{index}'
self.workspace_service.generations[workspace_uuid] = generation
context = self._context(workspace_uuid, generation)
runtime_provider = await self.model_manager.load_provider(
context,
persistence_model.ModelProvider(
uuid=provider_uuid,
workspace_uuid=workspace_uuid,
name='Capacity Provider',
requester=_ProbeRequester.name,
base_url='https://capacity.invalid',
api_keys=['probe'],
),
)
await self.model_manager.cache_provider(context, runtime_provider)
runtime_llm = await self.model_manager.load_llm_model_with_provider(
context,
persistence_model.LLMModel(
uuid=llm_uuid,
workspace_uuid=workspace_uuid,
name='Capacity LLM',
provider_uuid=provider_uuid,
abilities=['func_call'],
extra_args={'temperature': 0.1},
),
runtime_provider,
)
await self.model_manager.cache_llm_model(context, runtime_llm)
runtime_embedding = await self.model_manager.load_embedding_model_with_provider(
context,
persistence_model.EmbeddingModel(
uuid=embedding_uuid,
workspace_uuid=workspace_uuid,
name='Capacity Embedding',
provider_uuid=provider_uuid,
extra_args={'dimensions': 1_024},
),
runtime_provider,
)
await self.model_manager.cache_embedding_model(
context,
runtime_embedding,
)
runtime_rerank = await self.model_manager.load_rerank_model_with_provider(
context,
persistence_model.RerankModel(
uuid=rerank_uuid,
workspace_uuid=workspace_uuid,
name='Capacity Rerank',
provider_uuid=provider_uuid,
extra_args={},
),
runtime_provider,
)
await self.model_manager.cache_rerank_model(
context,
runtime_rerank,
)
pipeline_context = self._context(
workspace_uuid,
generation,
pipeline_uuid=pipeline_uuid,
)
await self.pipeline_manager.load_pipeline(
pipeline_context,
persistence_pipeline.LegacyPipeline(
uuid=pipeline_uuid,
workspace_uuid=workspace_uuid,
name='Capacity Pipeline',
description='',
for_version='probe',
is_default=True,
stages=[],
config={},
extensions_preferences={},
),
_binding_validated=True,
)
runtime_pipeline = self.pipeline_manager._pipelines_by_key[
(
self.workspace_service.instance_uuid,
workspace_uuid,
pipeline_uuid,
)
]
runtime_kb = await self.rag_manager.load_knowledge_base(
context,
persistence_rag.KnowledgeBase(
uuid=kb_uuid,
workspace_uuid=workspace_uuid,
name='Capacity Knowledge',
description='',
knowledge_engine_plugin_id=None,
collection_id=kb_uuid,
creation_settings={},
retrieval_settings={},
),
_binding_validated=True,
)
await self.mcp_loader._assert_execution_active(context)
runtime_mcp = _ProbeMCPSession(mcp_server_name)
self.mcp_loader._register_session(
context,
mcp_server_name,
runtime_mcp,
)
bot_context = self._context(
workspace_uuid,
generation,
bot_uuid=bot_uuid,
)
runtime_bot = await self.platform_manager.load_bot(
bot_context,
persistence_bot.Bot(
uuid=bot_uuid,
workspace_uuid=workspace_uuid,
name='Capacity Bot',
description='',
adapter='capacity-probe',
adapter_config={},
enable=True,
use_pipeline_uuid=pipeline_uuid,
pipeline_routing_rules=[],
),
_binding_validated=True,
)
self.generation_refs.setdefault(generation, []).extend(
(
weakref.ref(runtime_provider),
weakref.ref(runtime_llm),
weakref.ref(runtime_embedding),
weakref.ref(runtime_rerank),
weakref.ref(runtime_pipeline),
weakref.ref(runtime_kb),
weakref.ref(runtime_mcp),
weakref.ref(runtime_bot),
)
)
await asyncio.sleep(0)
def retained_state(self) -> dict[str, int]:
return {
'model_providers': len(self.model_manager.provider_dict),
'llm_models': len(self.model_manager.llm_model_dict),
'embedding_models': len(self.model_manager.embedding_model_dict),
'rerank_models': len(self.model_manager.rerank_model_dict),
'model_scopes': len(self.model_manager._scope_generations),
'pipelines': len(self.pipeline_manager._pipelines_by_key),
'pipeline_scopes': len(self.pipeline_manager._scope_generations),
'knowledge_bases': len(self.rag_manager.knowledge_bases),
'knowledge_scopes': len(self.rag_manager._scope_generations),
'mcp_sessions': len(self.mcp_loader.sessions),
'mcp_scopes': len(self.mcp_loader._scope_generations),
'bots': len(self.platform_manager._bots_by_key),
'bot_scopes': len(self.platform_manager._scope_generations),
'requesters_closed': _ProbeRequester.closed,
'adapters_killed': _ProbeAdapter.killed,
'mcp_sessions_closed': _ProbeMCPSession.closed,
'binding_lookups': self.workspace_service.binding_lookups,
}
def assert_generation_state(
self,
workspaces: int,
generation: int,
) -> None:
state = self.retained_state()
cardinality_keys = (
'model_providers',
'llm_models',
'embedding_models',
'rerank_models',
'model_scopes',
'pipelines',
'pipeline_scopes',
'knowledge_bases',
'knowledge_scopes',
'mcp_sessions',
'mcp_scopes',
'bots',
'bot_scopes',
)
invalid = {key: value for key in cardinality_keys if (value := state[key]) != workspaces}
if invalid:
raise AssertionError(f'populated Workspace cardinality mismatch: {invalid}')
expected_retired = (generation - 1) * workspaces
if state['requesters_closed'] != expected_retired:
raise AssertionError(f'retired requester count {state["requesters_closed"]} != {expected_retired}')
if state['adapters_killed'] != expected_retired:
raise AssertionError(f'retired adapter count {state["adapters_killed"]} != {expected_retired}')
if state['mcp_sessions_closed'] != expected_retired:
raise AssertionError(f'retired MCP session count {state["mcp_sessions_closed"]} != {expected_retired}')
def assert_generation_collected(self, generation: int) -> None:
gc.collect()
references = self.generation_refs.pop(generation)
retained = sum(reference() is not None for reference in references)
if retained:
raise AssertionError(f'{retained} generation-{generation} runtime objects remain reachable')
async def _run(args: argparse.Namespace) -> dict:
scale = SCALES[args.scale]
tracemalloc.start()
probe = PopulatedWorkspaceProbe()
baseline = _sample_process()
phase_one_started = time.monotonic()
await probe.load_generation(scale.workspaces, 1)
phase_one_seconds = time.monotonic() - phase_one_started
probe.assert_generation_state(scale.workspaces, 1)
phase_one = _sample_process()
phase_one_state = probe.retained_state()
phase_two_started = time.monotonic()
await probe.load_generation(scale.workspaces, 2)
phase_two_seconds = time.monotonic() - phase_two_started
probe.assert_generation_state(scale.workspaces, 2)
probe.assert_generation_collected(1)
phase_two = _sample_process()
phase_two_state = probe.retained_state()
phase_three_started = time.monotonic()
await probe.load_generation(scale.workspaces, 3)
phase_three_seconds = time.monotonic() - phase_three_started
probe.assert_generation_state(scale.workspaces, 3)
probe.assert_generation_collected(2)
phase_three = _sample_process()
phase_three_state = probe.retained_state()
cardinality_keys = (
'model_providers',
'llm_models',
'embedding_models',
'rerank_models',
'model_scopes',
'pipelines',
'pipeline_scopes',
'knowledge_bases',
'knowledge_scopes',
'mcp_sessions',
'mcp_scopes',
'bots',
'bot_scopes',
)
if any(
phase_two_state[key] != phase_one_state[key] or phase_three_state[key] != phase_one_state[key]
for key in cardinality_keys
):
raise AssertionError(
'populated Workspace registries did not plateau: '
f'phase_one={phase_one_state}, phase_two={phase_two_state}, '
f'phase_three={phase_three_state}'
)
traced_growth = phase_three.traced_current_bytes - phase_two.traced_current_bytes
rss_growth = phase_three.rss_bytes - phase_two.rss_bytes
max_traced_growth = int(args.max_traced_growth_mib * 1024 * 1024)
max_rss_growth = int(args.max_rss_growth_mib * 1024 * 1024)
if traced_growth > max_traced_growth:
raise AssertionError(f'replacement traced memory grew by {traced_growth} bytes (limit {max_traced_growth})')
if rss_growth > max_rss_growth:
raise AssertionError(f'replacement RSS grew by {rss_growth} bytes (limit {max_rss_growth})')
phase_ratio = max(
phase_two_seconds,
phase_three_seconds,
) / max(phase_one_seconds, 0.000_001)
if phase_ratio > args.max_replacement_time_ratio:
raise AssertionError(f'replacement phase ratio {phase_ratio:.3f} exceeds {args.max_replacement_time_ratio:.3f}')
return {
'component': 'langbot-populated-workspaces',
'scale': args.scale,
'workspaces': scale.workspaces,
'passed': True,
'phase_seconds': {
'initial': round(phase_one_seconds, 3),
'replacement_one': round(phase_two_seconds, 3),
'replacement_two': round(phase_three_seconds, 3),
'maximum_replacement_ratio': round(phase_ratio, 3),
},
'samples': {
'baseline': asdict(baseline),
'phase_one': asdict(phase_one),
'phase_two': asdict(phase_two),
'phase_three': asdict(phase_three),
},
'replacement_growth': {
'rss_bytes': rss_growth,
'traced_current_bytes': traced_growth,
},
'retained_state': {
'phase_one': phase_one_state,
'phase_two': phase_two_state,
'phase_three': phase_three_state,
},
}
def _parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--scale', choices=tuple(SCALES), default='quick')
parser.add_argument('--max-traced-growth-mib', type=float, default=16.0)
parser.add_argument('--max-rss-growth-mib', type=float, default=64.0)
parser.add_argument(
'--max-replacement-time-ratio',
type=float,
default=3.0,
)
parser.add_argument('--json', action='store_true')
return parser.parse_args()
def main() -> None:
args = _parse_args()
result = asyncio.run(_run(args))
if args.json:
print(json.dumps(result, sort_keys=True))
else:
print(json.dumps(result, indent=2, sort_keys=True))
if __name__ == '__main__':
main()