From b6c4e2c395bae2cdd4bdcc25c94b34b97b64baf9 Mon Sep 17 00:00:00 2001 From: RockChinQ Date: Sun, 11 Oct 2026 04:09:52 +0000 Subject: [PATCH] feat(cloud): enforce signed workspace subscription resource limits --- .../pkg/api/http/controller/groups/plugins.py | 4 +- .../api/http/controller/groups/workspaces.py | 6 +++ src/langbot/pkg/api/http/service/bot.py | 4 +- src/langbot/pkg/api/http/service/knowledge.py | 6 +-- src/langbot/pkg/api/http/service/mcp.py | 4 +- src/langbot/pkg/api/http/service/pipeline.py | 8 +++- .../pkg/api/http/service/workspace_limits.py | 20 +++++++++ tests/integration/api/test_workspaces.py | 2 + .../test_workspace_subscription_limits.py | 41 +++++++++++++++++++ .../useWorkspaceQuotaStatus.ts | 14 ++++--- web/src/app/infra/entities/workspace.ts | 1 + 11 files changed, 97 insertions(+), 13 deletions(-) create mode 100644 src/langbot/pkg/api/http/service/workspace_limits.py create mode 100644 tests/unit_tests/test_workspace_subscription_limits.py diff --git a/src/langbot/pkg/api/http/controller/groups/plugins.py b/src/langbot/pkg/api/http/controller/groups/plugins.py index 950a8954d..eb8b09592 100644 --- a/src/langbot/pkg/api/http/controller/groups/plugins.py +++ b/src/langbot/pkg/api/http/controller/groups/plugins.py @@ -296,7 +296,9 @@ class PluginsRouterGroup(group.RouterGroup): async def _check_extensions_limit(self, request_context: RequestContext) -> str | None: """Check if extensions limit is reached. Returns error response if limit exceeded, None otherwise.""" await self.ap.plugin_connector.require_workspace_context(request_context) - limitation = self.ap.instance_config.data.get('system', {}).get('limitation', {}) + from ...service.workspace_limits import workspace_limits + + limitation = await workspace_limits(self.ap, request_context.workspace_uuid) max_extensions = limitation.get('max_extensions', -1) if max_extensions >= 0: plugins = await self.ap.plugin_connector.list_plugins() diff --git a/src/langbot/pkg/api/http/controller/groups/workspaces.py b/src/langbot/pkg/api/http/controller/groups/workspaces.py index 7873e9f17..39dc9c577 100644 --- a/src/langbot/pkg/api/http/controller/groups/workspaces.py +++ b/src/langbot/pkg/api/http/controller/groups/workspaces.py @@ -129,6 +129,7 @@ class WorkspacesRouterGroup(group.RouterGroup): membership = quart.g.workspace_membership workspace = await self.ap.workspace_service.get_workspace(request_context.workspace_uuid) plan_name: str | None = None + limitation = None resolver = getattr(self.ap, 'entitlement_resolver', None) if workspace.source == WorkspaceSource.CLOUD_PROJECTION.value and resolver is not None: entitlement = await resolver.resolve( @@ -136,6 +137,9 @@ class WorkspacesRouterGroup(group.RouterGroup): minimum_revision=request_context.entitlement_revision, ) plan_name = entitlement.plan_name + from ...service.workspace_limits import LIMIT_KEYS + + limitation = {key: entitlement.limit(value) for key, value in LIMIT_KEYS.items()} if request_context.principal.principal_type == PrincipalType.SUPPORT_ADMIN: return self.success( data={ @@ -154,6 +158,7 @@ class WorkspacesRouterGroup(group.RouterGroup): 'permissions': sorted(request_context.workspace.permissions), 'placement_generation': request_context.placement_generation, 'plan_name': plan_name, + 'limitation': limitation, } ) account = await self.ap.user_service.get_user_by_uuid(request_context.account_uuid) @@ -170,6 +175,7 @@ class WorkspacesRouterGroup(group.RouterGroup): 'permissions': sorted(request_context.workspace.permissions), 'placement_generation': request_context.placement_generation, 'plan_name': plan_name, + 'limitation': limitation, } ) diff --git a/src/langbot/pkg/api/http/service/bot.py b/src/langbot/pkg/api/http/service/bot.py index 59ab6fe58..981c7cf50 100644 --- a/src/langbot/pkg/api/http/service/bot.py +++ b/src/langbot/pkg/api/http/service/bot.py @@ -685,7 +685,9 @@ class BotService: """Create bot""" workspace_uuid = require_workspace_uuid(context) # Check limitation - limitation = self.ap.instance_config.data.get('system', {}).get('limitation', {}) + from .workspace_limits import workspace_limits + + limitation = await workspace_limits(self.ap, workspace_uuid) max_bots = limitation.get('max_bots', -1) if max_bots >= 0: existing_bots = await self.get_bots(context) diff --git a/src/langbot/pkg/api/http/service/knowledge.py b/src/langbot/pkg/api/http/service/knowledge.py index 654bf2f17..5897aa578 100644 --- a/src/langbot/pkg/api/http/service/knowledge.py +++ b/src/langbot/pkg/api/http/service/knowledge.py @@ -56,9 +56,9 @@ class KnowledgeService: require_workspace_uuid(context) # In new architecture, we delegate entirely to RAGManager which uses plugins. # Legacy internal KB creation is removed. - limitation = ( - getattr(getattr(self.ap, 'instance_config', None), 'data', {}).get('system', {}).get('limitation', {}) - ) + from .workspace_limits import workspace_limits + + limitation = await workspace_limits(self.ap, require_workspace_uuid(context)) max_knowledge_bases = limitation.get('max_knowledge_bases', -1) if max_knowledge_bases >= 0: knowledge_bases = await self.ap.rag_mgr.get_all_knowledge_base_details(context) diff --git a/src/langbot/pkg/api/http/service/mcp.py b/src/langbot/pkg/api/http/service/mcp.py index 67e8e14ae..89a78f79a 100644 --- a/src/langbot/pkg/api/http/service/mcp.py +++ b/src/langbot/pkg/api/http/service/mcp.py @@ -212,7 +212,9 @@ class MCPService: # stdio MCP even though Box Runtime itself remains available. require_stdio_mcp_enabled(self.ap, server_data) - limitation = self.ap.instance_config.data.get('system', {}).get('limitation', {}) + from .workspace_limits import workspace_limits + + limitation = await workspace_limits(self.ap, workspace_uuid) max_extensions = limitation.get('max_extensions', -1) if max_extensions >= 0: mcp_count_result = await self.ap.persistence_mgr.execute_async( diff --git a/src/langbot/pkg/api/http/service/pipeline.py b/src/langbot/pkg/api/http/service/pipeline.py index bdd24f3d7..22a1c86c0 100644 --- a/src/langbot/pkg/api/http/service/pipeline.py +++ b/src/langbot/pkg/api/http/service/pipeline.py @@ -173,7 +173,9 @@ class PipelineService: if 'config' in pipeline_data: RunnerConfigResolver.validate_pipeline_config(pipeline_data['config']) # Check limitation - limitation = self.ap.instance_config.data.get('system', {}).get('limitation', {}) + from .workspace_limits import workspace_limits + + limitation = await workspace_limits(self.ap, workspace_uuid) max_pipelines = limitation.get('max_pipelines', -1) if max_pipelines >= 0: existing_pipelines = await self.get_pipelines(context) @@ -290,7 +292,9 @@ class PipelineService: """Copy a pipeline with all its configurations""" workspace_uuid = require_workspace_uuid(context) # Check limitation - limitation = self.ap.instance_config.data.get('system', {}).get('limitation', {}) + from .workspace_limits import workspace_limits + + limitation = await workspace_limits(self.ap, workspace_uuid) max_pipelines = limitation.get('max_pipelines', -1) if max_pipelines >= 0: existing_pipelines = await self.get_pipelines(context) diff --git a/src/langbot/pkg/api/http/service/workspace_limits.py b/src/langbot/pkg/api/http/service/workspace_limits.py new file mode 100644 index 000000000..dc41a2a2b --- /dev/null +++ b/src/langbot/pkg/api/http/service/workspace_limits.py @@ -0,0 +1,20 @@ +"""Resolve resource limits at the Workspace service boundary.""" + +from ....cloud.entitlements import EntitlementUnavailableError + +LIMIT_KEYS = { + 'max_bots': 'bots.max', + 'max_pipelines': 'pipelines.max', + 'max_extensions': 'plugins.max', + 'max_knowledge_bases': 'knowledge_bases.max', +} + + +async def workspace_limits(ap, workspace_uuid: str, *, minimum_revision: int = 0) -> dict[str, int]: + if getattr(getattr(ap, 'deployment', None), 'mode', None) != 'cloud': + return dict(ap.instance_config.data.get('system', {}).get('limitation', {})) + resolver = getattr(ap, 'entitlement_resolver', None) + if resolver is None: + raise EntitlementUnavailableError('Workspace entitlement resolver is unavailable') + snapshot = await resolver.resolve(workspace_uuid, minimum_revision=minimum_revision) + return {config_key: snapshot.limit(limit_key) for config_key, limit_key in LIMIT_KEYS.items()} diff --git a/tests/integration/api/test_workspaces.py b/tests/integration/api/test_workspaces.py index a561cd44c..7a6254d01 100644 --- a/tests/integration/api/test_workspaces.py +++ b/tests/integration/api/test_workspaces.py @@ -824,6 +824,7 @@ async def test_cloud_projection_is_selected_explicitly_and_collaboration_runs_in not_before=1, expires_at=4102444800, plan_name='free', + limits={'bots.max': 2, 'pipelines.max': 3, 'plugins.max': 4, 'knowledge_bases.max': 5}, ) application.entitlement_resolver = PlanResolver() @@ -843,6 +844,7 @@ async def test_cloud_projection_is_selected_explicitly_and_collaboration_runs_in assert current['workspace']['source'] == 'cloud_projection' assert current['placement_generation'] == 12 assert current['plan_name'] == 'free' + assert current['limitation']['max_bots'] == 2 create_workspace = await client.post( '/api/v1/workspaces', diff --git a/tests/unit_tests/test_workspace_subscription_limits.py b/tests/unit_tests/test_workspace_subscription_limits.py new file mode 100644 index 000000000..c16604ed9 --- /dev/null +++ b/tests/unit_tests/test_workspace_subscription_limits.py @@ -0,0 +1,41 @@ +import asyncio +from types import SimpleNamespace + +import pytest + +from langbot.pkg.cloud.entitlements import EntitlementSnapshot, EntitlementUnavailableError +from langbot.pkg.api.http.service.workspace_limits import workspace_limits + + +def test_oss_retains_instance_limits(): + ap = SimpleNamespace(instance_config=SimpleNamespace(data={'system': {'limitation': {'max_bots': 7}}})) + assert asyncio.run(workspace_limits(ap, 'one')) == {'max_bots': 7} + + +def test_cloud_uses_signed_workspace_limits(): + class Resolver: + async def resolve(self, workspace_uuid, **kwargs): + assert workspace_uuid == 'one' + return EntitlementSnapshot( + instance_uuid='i', + workspace_uuid='one', + entitlement_revision=1, + status='active', + not_before=0, + expires_at=9999999999, + limits={'bots.max': 2, 'pipelines.max': 3, 'plugins.max': 4, 'knowledge_bases.max': 5}, + ) + + ap = SimpleNamespace(deployment=SimpleNamespace(mode='cloud'), entitlement_resolver=Resolver()) + assert asyncio.run(workspace_limits(ap, 'one')) == { + 'max_bots': 2, + 'max_pipelines': 3, + 'max_extensions': 4, + 'max_knowledge_bases': 5, + } + + +def test_cloud_missing_resolver_fails_closed(): + ap = SimpleNamespace(deployment=SimpleNamespace(mode='cloud')) + with pytest.raises(EntitlementUnavailableError): + asyncio.run(workspace_limits(ap, 'one')) diff --git a/web/src/app/home/components/workspace-quota/useWorkspaceQuotaStatus.ts b/web/src/app/home/components/workspace-quota/useWorkspaceQuotaStatus.ts index 4788b7477..8786caf39 100644 --- a/web/src/app/home/components/workspace-quota/useWorkspaceQuotaStatus.ts +++ b/web/src/app/home/components/workspace-quota/useWorkspaceQuotaStatus.ts @@ -1,3 +1,4 @@ +import { useCurrentWorkspace } from '@/app/infra/http'; import { systemInfo } from '@/app/infra/http/HttpClient'; import { useSidebarData } from '@/app/home/components/home-sidebar/SidebarDataContext'; @@ -46,27 +47,30 @@ export function useWorkspaceQuotaStatus(): WorkspaceQuotaStatus { skills, quotaDataLoaded, } = useSidebarData(); - const limitation = systemInfo.limitation; + const currentWorkspace = useCurrentWorkspace(); + const cloud = currentWorkspace?.workspace.source === 'cloud_projection'; + const limitation = cloud ? currentWorkspace.limitation : systemInfo.limitation; + const loaded = quotaDataLoaded && (!cloud || limitation != null); const botQuota = quotaItem( bots.length, limitation?.max_bots, - quotaDataLoaded, + loaded, ); const pipelineQuota = quotaItem( pipelines.length, limitation?.max_pipelines, - quotaDataLoaded, + loaded, ); const knowledgeBaseQuota = quotaItem( knowledgeBases.length, limitation?.max_knowledge_bases, - quotaDataLoaded, + loaded, ); const extensionQuota = quotaItem( pluginCount + mcpServers.length + skills.length, limitation?.max_extensions, - quotaDataLoaded, + loaded, ); return { diff --git a/web/src/app/infra/entities/workspace.ts b/web/src/app/infra/entities/workspace.ts index c369961a2..8ce39485e 100644 --- a/web/src/app/infra/entities/workspace.ts +++ b/web/src/app/infra/entities/workspace.ts @@ -34,6 +34,7 @@ export interface CurrentWorkspace { placement_generation: number; /** Signed Cloud display metadata; never used for client-side authorization. */ plan_name?: string | null; + limitation?: Record | null; } export interface WorkspaceSpaceBilling {