diff --git a/.github/workflows/deploy-prod.yml b/.github/workflows/deploy-prod.yml new file mode 100644 index 000000000..acf33d0e5 --- /dev/null +++ b/.github/workflows/deploy-prod.yml @@ -0,0 +1,59 @@ +name: Build and deploy production + +on: + push: + branches: [deploy/prod] + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: langbot-production + cancel-in-progress: false + +env: + CORE_IMAGE: ${{ secrets.DOCKER_USERNAME }}/langbot + CLOUD_IMAGE: ${{ secrets.DOCKER_USERNAME }}/langbot-cloud-core + SPACE_REF: 178b1634c104e635c706e1fb4ca37cb3de073cf9 + +jobs: + build-and-deploy: + runs-on: ubuntu-latest + environment: production + steps: + - uses: actions/checkout@v4 + - uses: docker/setup-buildx-action@v3 + - uses: docker/login-action@v3 + with: + username: ${{ secrets.DOCKER_USERNAME }} + password: ${{ secrets.DOCKER_PASSWORD }} + - name: Build exact Core image + uses: docker/build-push-action@v6 + with: + context: . + push: true + tags: | + ${{ env.CORE_IMAGE }}:prod-${{ github.sha }} + ${{ env.CORE_IMAGE }}:deploy-prod + cache-from: type=gha,scope=core-prod + cache-to: type=gha,mode=max,scope=core-prod + - name: Checkout production Cloud adapter + uses: actions/checkout@v4 + with: + repository: langbot-app/langbot-space + ref: ${{ env.SPACE_REF }} + token: ${{ secrets.CLA_PAT }} + path: .space + - name: Build exact Cloud Core image + uses: docker/build-push-action@v6 + with: + context: .space + file: .space/Dockerfile.cloud + push: true + build-args: LANGBOT_CORE_IMAGE=${{ env.CORE_IMAGE }}:prod-${{ github.sha }} + tags: | + ${{ env.CLOUD_IMAGE }}:prod-${{ github.sha }} + ${{ env.CLOUD_IMAGE }}:deploy-prod + cache-from: type=gha,scope=cloud-core-prod + cache-to: type=gha,mode=max,scope=cloud-core-prod diff --git a/deploy/prod/deploy.sh b/deploy/prod/deploy.sh new file mode 100755 index 000000000..dc51f89a5 --- /dev/null +++ b/deploy/prod/deploy.sh @@ -0,0 +1,97 @@ +#!/usr/bin/env bash +set -Eeuo pipefail + +cd /opt/langbot-cloud-prod +TAG=${1:?usage: deploy.sh prod-<40-char-sha>} +[[ "$TAG" =~ ^prod-[0-9a-f]{40}$ ]] || { echo 'invalid immutable image tag' >&2; exit 2; } +[[ -s .env ]] || { echo '/opt/langbot-cloud-prod/.env is missing' >&2; exit 3; } + +rendered_compose=$(docker compose config) +grep -Fq 'LANGBOT_SPACE_CONTROL_PLANE_URL: https://space.langbot.app' <<<"$rendered_compose" || { + echo 'Cloud control-plane URL must be https://space.langbot.app' >&2 + exit 4 +} +grep -Fq 'SPACE__URL: https://space.langbot.app' <<<"$rendered_compose" || { + echo 'Cloud user-facing Space URL must be https://space.langbot.app' >&2 + exit 5 +} +grep -Eq 'LANGBOT_TELEMETRY_INGEST_TOKEN: .+' <<<"$rendered_compose" || { + echo 'Cloud telemetry ingest token must be configured' >&2 + exit 6 +} + +update_env() { + local key=$1 value=$2 + python3 - "$key" "$value" <<'PY' +from pathlib import Path +import os +import sys + +path = Path('.env') +key, value = sys.argv[1:] +lines = path.read_text().splitlines() +updated = False +for index, line in enumerate(lines): + if line.startswith(f'{key}='): + lines[index] = f'{key}={value}' + updated = True + break +if not updated: + lines.append(f'{key}={value}') +temporary = Path('.env.tmp') +temporary.write_text('\n'.join(lines) + '\n') +os.chmod(temporary, 0o600) +temporary.replace(path) +PY +} +update_env LANGBOT_IMAGE_TAG "$TAG" +set -a +. ./.env +set +a +: "${CLOUD_V2_CONTROL_PLANE_TOKEN:?CLOUD_V2_CONTROL_PLANE_TOKEN is required}" + +for attempt in 1 2 3 4 5; do + if docker compose pull postgres redis migrate plugin-runtime core; then + break + fi + if [ "$attempt" -eq 5 ]; then + echo "docker compose pull failed after $attempt attempts" >&2 + exit 1 + fi + delay=$((attempt * 10)) + echo "docker compose pull failed (attempt $attempt/5); retrying in ${delay}s" >&2 + sleep "$delay" +done +docker compose up -d postgres redis +for _ in $(seq 1 60); do + if docker compose exec -T postgres pg_isready -U langbot_operator -d langbot >/dev/null 2>&1; then break; fi + sleep 2 +done +docker compose exec -T postgres pg_isready -U langbot_operator -d langbot >/dev/null + +docker compose exec -T postgres psql -v ON_ERROR_STOP=1 -U langbot_operator -d langbot \ + -v runtime_password="$POSTGRES_RUNTIME_PASSWORD" <<'SQL' +SELECT format('CREATE ROLE langbot_runtime LOGIN PASSWORD %L', :'runtime_password') +WHERE NOT EXISTS (SELECT 1 FROM pg_roles WHERE rolname = 'langbot_runtime')\gexec +ALTER ROLE langbot_runtime PASSWORD :'runtime_password'; +GRANT CONNECT ON DATABASE langbot TO langbot_runtime; +REVOKE CREATE ON SCHEMA public FROM PUBLIC, langbot_runtime; +REVOKE ALL PRIVILEGES ON ALL TABLES IN SCHEMA public FROM langbot_runtime; +REVOKE ALL PRIVILEGES ON ALL SEQUENCES IN SCHEMA public FROM langbot_runtime; +ALTER DEFAULT PRIVILEGES FOR ROLE langbot_operator IN SCHEMA public REVOKE ALL ON TABLES FROM langbot_runtime; +ALTER DEFAULT PRIVILEGES FOR ROLE langbot_operator IN SCHEMA public REVOKE ALL ON SEQUENCES FROM langbot_runtime; +GRANT USAGE ON SCHEMA public TO langbot_runtime; +SQL + +docker compose --profile tools run --rm migrate + +docker compose up -d --remove-orphans plugin-runtime core +for _ in $(seq 1 90); do + if docker compose exec -T core python -c 'import urllib.request; urllib.request.urlopen("http://127.0.0.1:5300/healthz", timeout=3)' >/dev/null 2>&1; then + docker compose ps + exit 0 + fi + sleep 2 +done +docker compose logs --tail=200 core plugin-runtime >&2 +exit 1 diff --git a/deploy/prod/docker-compose.yml b/deploy/prod/docker-compose.yml new file mode 100644 index 000000000..db224b4f0 --- /dev/null +++ b/deploy/prod/docker-compose.yml @@ -0,0 +1,162 @@ +services: + postgres: + image: pgvector/pgvector:pg17 + container_name: langbot-cloud-postgres + restart: unless-stopped + environment: + POSTGRES_DB: langbot + POSTGRES_USER: langbot_operator + POSTGRES_PASSWORD: ${POSTGRES_OPERATOR_PASSWORD} + volumes: + - postgres-data:/var/lib/postgresql/data + healthcheck: + test: [CMD-SHELL, "pg_isready -U langbot_operator -d langbot"] + interval: 5s + timeout: 5s + retries: 30 + networks: [internal] + + redis: + image: redis:7.4-alpine + container_name: langbot-cloud-redis + restart: unless-stopped + command: [redis-server, --appendonly, "yes", --requirepass, "${REDIS_PASSWORD}"] + volumes: + - redis-data:/data + healthcheck: + test: [CMD-SHELL, "redis-cli -a \"$${REDIS_PASSWORD}\" ping | grep PONG"] + interval: 5s + timeout: 5s + retries: 20 + environment: + REDIS_PASSWORD: ${REDIS_PASSWORD} + networks: [internal] + + migrate: + image: rockchin/langbot-cloud-core:${LANGBOT_IMAGE_TAG} + profiles: [tools] + command: [uv, run, langbot, migrate, --cloud] + environment: &core-env + TZ: Asia/Shanghai + SYSTEM__INSTANCE_ID: ${CLOUD_V2_INSTANCE_UUID} + SYSTEM__EDITION: cloud + SYSTEM__RECOVERY_KEY: ${SYSTEM_RECOVERY_KEY} + SYSTEM__JWT__SECRET: ${JWT_SECRET} + SYSTEM__LIMITATION__MAX_BOTS: "2" + SYSTEM__LIMITATION__MAX_PIPELINES: "3" + SYSTEM__LIMITATION__MAX_EXTENSIONS: "3" + SYSTEM__LIMITATION__MAX_KNOWLEDGE_BASES: "2" + API__WEBHOOK_PREFIX: https://cloud.langbot.app + API__WEBUI_URL: https://cloud.langbot.app + WORKSPACE__INVITATIONS__PUBLIC_WEB_URL: https://cloud.langbot.app + DATABASE__USE: postgresql + DATABASE__POSTGRESQL__URL: postgresql+asyncpg://langbot_runtime:${POSTGRES_RUNTIME_PASSWORD}@postgres:5432/langbot + DATABASE__CLOUD_MIGRATION__OPERATOR_DSN_ENV: LANGBOT_CLOUD_MIGRATION_DSN + LANGBOT_CLOUD_MIGRATION_DSN: postgresql://langbot_operator:${POSTGRES_OPERATOR_PASSWORD}@postgres:5432/langbot + VDB__USE: pgvector + VDB__PGVECTOR__USE_BUSINESS_DATABASE: "true" + VDB__PGVECTOR__ALLOWED_DIMENSIONS: "384,512,768,1024,1536" + PLUGIN__ENABLE: "true" + PLUGIN__RUNTIME_WS_URL: ws://plugin-runtime:5400/control/ws + PLUGIN__DISPLAY_PLUGIN_DEBUG_URL: wss://cloud.langbot.app/plugin/debug/ws + PLUGIN__WORKER__MAX_CPUS: "0.25" + PLUGIN__WORKER__MAX_MEMORY_MB: "256" + PLUGIN__WORKER__MAX_PIDS: "128" + PLUGIN__WORKER__MAX_WORKERS: "16" + PLUGIN__WORKER__MAX_TOTAL_CPUS: "4.0" + PLUGIN__WORKER__MAX_TOTAL_MEMORY_MB: "4096" + PLUGIN__WORKER__REQUIRE_HARD_LIMITS: "true" + LANGBOT_PLUGIN_RUNTIME_CONTROL_TOKEN: ${PLUGIN_RUNTIME_CONTROL_TOKEN} + # Cloud v2 currently grants no managed Box capability. Keep the shared + # runtime deployed but disable Core integration until a hard-quota-capable + # backend can satisfy the fail-closed Cloud readiness contract. + BOX__ENABLED: "false" + BOX__BACKEND: nsjail + BOX__RUNTIME__ENDPOINT: ws://box:5410 + BOX__ADMISSION__REQUIRED: "true" + BOX__ADMISSION__LOGICAL_SESSION_ID: global + BOX__ADMISSION__REQUIRED_BACKEND: nsjail + BOX__ADMISSION__MAX_SESSIONS: "1" + BOX__ADMISSION__MAX_MANAGED_PROCESSES: "0" + BOX__ADMISSION__CPUS: "0.25" + BOX__ADMISSION__MEMORY_MB: "256" + BOX__ADMISSION__WORKSPACE_QUOTA_MB: "256" + BOX__LOCAL__HOST_ROOT: /app/data/box + BOX__LOCAL__DEFAULT_WORKSPACE: /app/data/box + BOX__LOCAL__ALLOWED_MOUNT_ROOTS: /app/data/box + LANGBOT_BOX_CONTROL_TOKEN: ${BOX_CONTROL_TOKEN} + MCP__STDIO__ENABLED: "false" + LANGBOT_SPACE_CONTROL_PLANE_URL: https://space.langbot.app + LANGBOT_SPACE_CONTROL_PLANE_TOKEN: ${CLOUD_V2_CONTROL_PLANE_TOKEN} + LANGBOT_TELEMETRY_INGEST_TOKEN: ${CLOUD_V2_CONTROL_PLANE_TOKEN} + LANGBOT_SPACE_CONTROL_PLANE_PUBLIC_KEY: ${CLOUD_V2_MANIFEST_PUBLIC_KEY} + LANGBOT_SPACE_CONTROL_PLANE_KEY_ID: ${CLOUD_V2_MANIFEST_KEY_ID} + SPACE__URL: https://space.langbot.app + depends_on: + postgres: {condition: service_healthy} + networks: [internal] + + plugin-runtime: + image: rockchin/langbot:${LANGBOT_IMAGE_TAG} + container_name: langbot-cloud-plugin-runtime + restart: unless-stopped + command: [uv, run, python, -m, langbot_plugin.cli.__init__, rt] + environment: + LANGBOT_PLUGIN_RUNTIME_CONTROL_TOKEN: ${PLUGIN_RUNTIME_CONTROL_TOKEN} + volumes: + - plugin-data:/app/data + - /sys/fs/cgroup:/sys/fs/cgroup:rw + cgroup: host + privileged: true + expose: ["5400"] + networks: [internal] + + box: + image: rockchin/langbot:${LANGBOT_IMAGE_TAG} + container_name: langbot-cloud-box + restart: unless-stopped + command: [uv, run, lbp, box, --host, 0.0.0.0, --ws-control-port, "5410"] + environment: + LANGBOT_BOX_CONTROL_TOKEN: ${BOX_CONTROL_TOKEN} + LANGBOT_BOX_ROOT: /app/data/box + volumes: + - box-data:/app/data/box + - /sys/fs/cgroup:/sys/fs/cgroup:rw + cgroup: host + privileged: true + expose: ["5410"] + networks: [internal] + + core: + image: rockchin/langbot-cloud-core:${LANGBOT_IMAGE_TAG} + container_name: langbot-cloud-core + restart: unless-stopped + environment: *core-env + volumes: + - core-data:/app/data + - box-data:/app/data/box + depends_on: + postgres: {condition: service_healthy} + redis: {condition: service_healthy} + plugin-runtime: {condition: service_started} + box: {condition: service_started} + expose: ["5300"] + healthcheck: + test: [CMD-SHELL, "python -c 'import urllib.request; urllib.request.urlopen(\"http://127.0.0.1:5300/healthz\", timeout=3)'" ] + interval: 10s + timeout: 5s + retries: 30 + start_period: 30s + networks: [internal, shared-network] + +networks: + internal: + shared-network: + external: true + +volumes: + postgres-data: + redis-data: + plugin-data: + box-data: + core-data: diff --git a/pyproject.toml b/pyproject.toml index d1e588fa8..63eb27a84 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "langbot" -version = "4.10.6" +version = "4.10.7" description = "Production-grade platform for building agentic IM bots" readme = "README.md" license-files = ["LICENSE"] @@ -71,7 +71,7 @@ dependencies = [ "chromadb>=1.0.0,<2.0.0", "qdrant-client (>=1.15.1,<2.0.0)", "pyseekdb==1.1.0.post3", - "langbot-plugin @ git+https://github.com/langbot-app/langbot-plugin-sdk.git@1d65ed301a6afc52150a998043f73cd6032c8162", + "langbot-plugin @ git+https://github.com/langbot-app/langbot-plugin-sdk.git@101e453e916b39465a6294d6471c9eaae8725d5c", "asyncpg>=0.30.0", "line-bot-sdk>=3.19.0", "matrix-nio>=0.25.2", diff --git a/src/langbot/pkg/api/http/controller/groups/user.py b/src/langbot/pkg/api/http/controller/groups/user.py index 8b86ba361..49b59906b 100644 --- a/src/langbot/pkg/api/http/controller/groups/user.py +++ b/src/langbot/pkg/api/http/controller/groups/user.py @@ -285,8 +285,17 @@ class UserRouterGroup(group.RouterGroup): request_context.workspace_uuid, ) owner = await self.ap.user_service.get_workspace_owner(access.workspace.uuid) - owner_space_bound = bool(owner and owner.space_account_uuid) - credits = await self.ap.space_service.get_credits(owner.user) if owner_space_bound else None + cloud_mode = getattr(getattr(self.ap, 'deployment', None), 'mode', 'oss') == 'cloud' + owner_has_local_space_credentials = bool(owner and owner.space_account_uuid) + # Cloud Accounts authenticate through LangBot Account, so every projected + # Workspace owner is already bound even when this Core has no local OAuth + # token row (model billing uses the owner's control-plane API key). + owner_space_bound = cloud_mode or owner_has_local_space_credentials + credits = ( + await self.ap.space_service.get_credits(owner.user) + if owner is not None and owner.space_account_uuid + else None + ) return self.success( data={ 'credits': credits, @@ -302,8 +311,10 @@ class UserRouterGroup(group.RouterGroup): return self.success(data={'initialized': False}) capabilities = await self.ap.user_service.get_login_capabilities() - if getattr(getattr(self.ap, 'deployment', None), 'mode', 'oss') == 'cloud': + cloud_mode = getattr(getattr(self.ap, 'deployment', None), 'mode', 'oss') == 'cloud' + if cloud_mode: capabilities['password_login_enabled'] = False + capabilities['authenticated_invitation_acceptance_enabled'] = cloud_mode return self.success(data={'initialized': True, **capabilities}) @self.route('/set-password', methods=['POST'], auth_type=group.AuthType.USER_TOKEN) diff --git a/src/langbot/pkg/api/http/service/model.py b/src/langbot/pkg/api/http/service/model.py index 88801c6b8..898a791e2 100644 --- a/src/langbot/pkg/api/http/service/model.py +++ b/src/langbot/pkg/api/http/service/model.py @@ -5,6 +5,7 @@ import uuid import sqlalchemy from langbot_plugin.api.entities.builtin.provider import message as provider_message +from ....cloud.model_catalog import LANGBOT_MODELS_PROVIDER_REQUESTER from ....core import app from ....entity.persistence import model as persistence_model from ....entity.persistence import pipeline as persistence_pipeline @@ -113,6 +114,23 @@ async def _require_workspace_provider( return provider +def _is_cloud_runtime(ap: app.Application) -> bool: + mode = getattr(ap.persistence_mgr, 'mode', None) + return getattr(mode, 'value', None) == 'cloud_runtime' + + +async def _assert_cloud_managed_provider_mutable( + ap: app.Application, + context: TenantContext, + provider_uuid: str, +) -> None: + if not _is_cloud_runtime(ap): + return + provider = await _require_workspace_provider(ap, context, provider_uuid) + if provider.get('requester') == LANGBOT_MODELS_PROVIDER_REQUESTER: + raise ValueError('LangBot Models is managed by Cloud and cannot be modified') + + async def _require_runtime_provider( ap: app.Application, context: TenantContext, @@ -213,6 +231,7 @@ class LLMModelsService: model_data['provider_uuid'] = provider_uuid await _require_workspace_provider(self.ap, context, model_data['provider_uuid']) + await _assert_cloud_managed_provider_mutable(self.ap, context, model_data['provider_uuid']) await _validate_provider_supports(self.ap, context, model_data['provider_uuid'], 'llm') await self.ap.persistence_mgr.execute_async(sqlalchemy.insert(persistence_model.LLMModel).values(**model_data)) @@ -291,11 +310,17 @@ class LLMModelsService: return model_dict - async def update_llm_model(self, context: TenantContext, model_uuid: str, model_data: dict) -> None: + async def update_llm_model( + self, + context: TenantContext, + model_uuid: str, + model_data: dict, + ) -> None: """Update an existing LLM model""" existing_model = await self.get_llm_model(context, model_uuid, include_secret=True) if existing_model is None: raise WorkspaceNotFoundError('Model not found') + await _assert_cloud_managed_provider_mutable(self.ap, context, existing_model['provider_uuid']) model_data = model_data.copy() model_data.pop('uuid', None) model_data.pop('workspace_uuid', None) @@ -321,6 +346,7 @@ class LLMModelsService: provider_uuid = model_data.get('provider_uuid', existing_model['provider_uuid']) await _require_workspace_provider(self.ap, context, provider_uuid) + await _assert_cloud_managed_provider_mutable(self.ap, context, provider_uuid) await _validate_provider_supports(self.ap, context, provider_uuid, 'llm') result = await self.ap.persistence_mgr.execute_async( @@ -355,6 +381,11 @@ class LLMModelsService: async def delete_llm_model(self, context: TenantContext, model_uuid: str) -> None: """Delete an LLM model""" + if _is_cloud_runtime(self.ap): + existing_model = await self.get_llm_model(context, model_uuid, include_secret=True) + if existing_model is None: + raise WorkspaceNotFoundError('Model not found') + await _assert_cloud_managed_provider_mutable(self.ap, context, existing_model['provider_uuid']) result = await self.ap.persistence_mgr.execute_async( scope_statement( sqlalchemy.delete(persistence_model.LLMModel).where(persistence_model.LLMModel.uuid == model_uuid), @@ -448,7 +479,10 @@ class EmbeddingModelsService: return serialized if include_secret else [_redact_model_secrets(model) for model in serialized] async def create_embedding_model( - self, context: TenantContext, model_data: dict, preserve_uuid: bool = False + self, + context: TenantContext, + model_data: dict, + preserve_uuid: bool = False, ) -> str: """Create a new embedding model""" model_data = model_data.copy() @@ -472,6 +506,7 @@ class EmbeddingModelsService: model_data['provider_uuid'] = provider_uuid await _require_workspace_provider(self.ap, context, model_data['provider_uuid']) + await _assert_cloud_managed_provider_mutable(self.ap, context, model_data['provider_uuid']) await _validate_provider_supports(self.ap, context, model_data['provider_uuid'], 'text-embedding') await self.ap.persistence_mgr.execute_async( @@ -530,11 +565,17 @@ class EmbeddingModelsService: return model_dict - async def update_embedding_model(self, context: TenantContext, model_uuid: str, model_data: dict) -> None: + async def update_embedding_model( + self, + context: TenantContext, + model_uuid: str, + model_data: dict, + ) -> None: """Update an existing embedding model""" existing_model = await self.get_embedding_model(context, model_uuid, include_secret=True) if existing_model is None: raise WorkspaceNotFoundError('Model not found') + await _assert_cloud_managed_provider_mutable(self.ap, context, existing_model['provider_uuid']) model_data = model_data.copy() model_data.pop('uuid', None) model_data.pop('workspace_uuid', None) @@ -559,6 +600,7 @@ class EmbeddingModelsService: provider_uuid = model_data.get('provider_uuid', existing_model['provider_uuid']) await _require_workspace_provider(self.ap, context, provider_uuid) + await _assert_cloud_managed_provider_mutable(self.ap, context, provider_uuid) await _validate_provider_supports(self.ap, context, provider_uuid, 'text-embedding') result = await self.ap.persistence_mgr.execute_async( @@ -593,6 +635,11 @@ class EmbeddingModelsService: async def delete_embedding_model(self, context: TenantContext, model_uuid: str) -> None: """Delete an embedding model""" + if _is_cloud_runtime(self.ap): + existing_model = await self.get_embedding_model(context, model_uuid, include_secret=True) + if existing_model is None: + raise WorkspaceNotFoundError('Model not found') + await _assert_cloud_managed_provider_mutable(self.ap, context, existing_model['provider_uuid']) result = await self.ap.persistence_mgr.execute_async( scope_statement( sqlalchemy.delete(persistence_model.EmbeddingModel).where( @@ -685,7 +732,12 @@ class RerankModelsService: serialized = [self.ap.persistence_mgr.serialize_model(persistence_model.RerankModel, m) for m in models] return serialized if include_secret else [_redact_model_secrets(model) for model in serialized] - async def create_rerank_model(self, context: TenantContext, model_data: dict, preserve_uuid: bool = False) -> str: + async def create_rerank_model( + self, + context: TenantContext, + model_data: dict, + preserve_uuid: bool = False, + ) -> str: """Create a new rerank model""" model_data = model_data.copy() if not preserve_uuid: @@ -708,6 +760,7 @@ class RerankModelsService: model_data['provider_uuid'] = provider_uuid await _require_workspace_provider(self.ap, context, model_data['provider_uuid']) + await _assert_cloud_managed_provider_mutable(self.ap, context, model_data['provider_uuid']) await _validate_provider_supports(self.ap, context, model_data['provider_uuid'], 'rerank') await self.ap.persistence_mgr.execute_async( @@ -766,11 +819,17 @@ class RerankModelsService: return model_dict - async def update_rerank_model(self, context: TenantContext, model_uuid: str, model_data: dict) -> None: + async def update_rerank_model( + self, + context: TenantContext, + model_uuid: str, + model_data: dict, + ) -> None: """Update an existing rerank model""" existing_model = await self.get_rerank_model(context, model_uuid, include_secret=True) if existing_model is None: raise WorkspaceNotFoundError('Model not found') + await _assert_cloud_managed_provider_mutable(self.ap, context, existing_model['provider_uuid']) model_data = model_data.copy() model_data.pop('uuid', None) model_data.pop('workspace_uuid', None) @@ -795,6 +854,7 @@ class RerankModelsService: provider_uuid = model_data.get('provider_uuid', existing_model['provider_uuid']) await _require_workspace_provider(self.ap, context, provider_uuid) + await _assert_cloud_managed_provider_mutable(self.ap, context, provider_uuid) await _validate_provider_supports(self.ap, context, provider_uuid, 'rerank') result = await self.ap.persistence_mgr.execute_async( @@ -829,6 +889,11 @@ class RerankModelsService: async def delete_rerank_model(self, context: TenantContext, model_uuid: str) -> None: """Delete a rerank model""" + if _is_cloud_runtime(self.ap): + existing_model = await self.get_rerank_model(context, model_uuid, include_secret=True) + if existing_model is None: + raise WorkspaceNotFoundError('Model not found') + await _assert_cloud_managed_provider_mutable(self.ap, context, existing_model['provider_uuid']) result = await self.ap.persistence_mgr.execute_async( scope_statement( sqlalchemy.delete(persistence_model.RerankModel).where( diff --git a/src/langbot/pkg/api/http/service/provider.py b/src/langbot/pkg/api/http/service/provider.py index 74647d7ab..dc29d3858 100644 --- a/src/langbot/pkg/api/http/service/provider.py +++ b/src/langbot/pkg/api/http/service/provider.py @@ -5,6 +5,7 @@ import traceback import sqlalchemy +from ....cloud.model_catalog import LANGBOT_MODELS_PROVIDER_REQUESTER from ....core import app from ....entity.persistence import model as persistence_model from ....workspace.errors import WorkspaceNotFoundError @@ -20,6 +21,20 @@ class ModelProviderService: def __init__(self, ap: app.Application) -> None: self.ap = ap + def _is_cloud_runtime(self) -> bool: + mode = getattr(self.ap.persistence_mgr, 'mode', None) + return getattr(mode, 'value', None) == 'cloud_runtime' + + def _system_requester_is_reserved(self, requester: object) -> bool: + return self._is_cloud_runtime() and requester == LANGBOT_MODELS_PROVIDER_REQUESTER + + async def _assert_provider_mutable(self, context: TenantContext, provider_uuid: str) -> None: + if not self._is_cloud_runtime(): + return + provider = await self.get_provider(context, provider_uuid) + if provider is not None and self._system_requester_is_reserved(provider.get('requester')): + raise ValueError('LangBot Models is managed by Cloud and cannot be modified') + @staticmethod def _normalize_api_keys(api_keys: str | list[str] | tuple[str, ...] | None) -> list[str]: if api_keys is None: @@ -99,6 +114,8 @@ class ModelProviderService: async def create_provider(self, context: TenantContext, provider_data: dict) -> str: """Create a new provider""" provider_data = provider_data.copy() + if self._system_requester_is_reserved(provider_data.get('requester')): + raise ValueError('space-chat-completions is reserved for the Cloud-managed LangBot Models provider') provider_data['uuid'] = str(uuid.uuid4()) provider_data['workspace_uuid'] = require_workspace_uuid(context) provider_data['api_keys'] = self._normalize_api_keys( @@ -115,7 +132,10 @@ class ModelProviderService: async def update_provider(self, context: TenantContext, provider_uuid: str, provider_data: dict) -> None: """Update an existing provider""" + await self._assert_provider_mutable(context, provider_uuid) provider_data = provider_data.copy() + if self._system_requester_is_reserved(provider_data.get('requester')): + raise ValueError('space-chat-completions is reserved for the Cloud-managed LangBot Models provider') provider_data.pop('uuid', None) provider_data.pop('workspace_uuid', None) if 'api_keys' in provider_data: @@ -145,6 +165,7 @@ class ModelProviderService: async def delete_provider(self, context: TenantContext, provider_uuid: str) -> None: """Delete a provider (only if no models reference it)""" + await self._assert_provider_mutable(context, provider_uuid) workspace_uuid = require_workspace_uuid(context) # Check if any models use this provider llm_result = await self.ap.persistence_mgr.execute_async( @@ -245,6 +266,8 @@ class ModelProviderService: api_keys: list, ) -> str: """Find existing provider or create new one""" + if self._system_requester_is_reserved(requester): + raise ValueError('space-chat-completions is reserved for the Cloud-managed LangBot Models provider') workspace_uuid = require_workspace_uuid(context) api_keys = self._normalize_api_keys(restore_secret_placeholders(api_keys, sensitive=True)) diff --git a/src/langbot/pkg/cloud/bootstrap.py b/src/langbot/pkg/cloud/bootstrap.py index 8e9ee1886..2c7ddafd2 100644 --- a/src/langbot/pkg/cloud/bootstrap.py +++ b/src/langbot/pkg/cloud/bootstrap.py @@ -13,6 +13,7 @@ from typing import Any, Protocol, runtime_checkable from ..workspace.policy import CloudWorkspacePolicy, SingleWorkspacePolicy from .directory import DirectoryProjectionProvider, directory_projection_limits_from_config from .entitlements import EntitlementProvider, OpenSourceEntitlementProvider +from .model_catalog import CloudModelCatalogProvider CLOUD_BOOTSTRAP_ENTRY_POINT = 'langbot.cloud_bootstrap' @@ -50,6 +51,7 @@ class OpenSourceDeployment: ) directory_provider: None = None manifest_provider: None = None + model_catalog_provider: None = None persistence_mode: str = 'oss_compat' required_vector_backend: str | None = None @@ -80,6 +82,7 @@ class VerifiedCloudDeployment: entitlement_provider: EntitlementProvider directory_provider: DirectoryProjectionProvider manifest_provider: CloudManifestProvider + model_catalog_provider: CloudModelCatalogProvider verification_key_id: str mode: str = dataclasses.field(default='cloud', init=False) workspace_policy: CloudWorkspacePolicy = dataclasses.field(default_factory=CloudWorkspacePolicy, init=False) @@ -110,6 +113,8 @@ class VerifiedCloudDeployment: raise CloudBootstrapError('Verified Cloud bootstrap did not provide a directory adapter') if not isinstance(self.manifest_provider, CloudManifestProvider): raise CloudBootstrapError('Verified Cloud bootstrap did not provide a Manifest renewal adapter') + if not isinstance(self.model_catalog_provider, CloudModelCatalogProvider): + raise CloudBootstrapError('Verified Cloud bootstrap did not provide a model catalog adapter') def validate_instance_config(self, config: dict[str, Any]) -> None: try: @@ -138,8 +143,14 @@ class VerifiedCloudDeployment: if plugin_worker.get('require_hard_limits') is not True: raise CloudBootstrapError('Cloud Runtime requires plugin.worker.require_hard_limits=true') box_config = config.get('box', {}) - if box_config.get('enabled') is not True: - raise CloudBootstrapError('Cloud runtime requires box.enabled=true') + box_enabled = box_config.get('enabled') + if box_enabled is False: + # Explicitly disabling Box removes the sandbox surface entirely and + # therefore does not weaken tenant isolation. Validate the strict + # runtime/admission contract only when the surface is enabled. + return + if box_enabled is not True: + raise CloudBootstrapError('Cloud runtime requires box.enabled to be an explicit boolean') if box_config.get('backend') != 'nsjail': raise CloudBootstrapError('Cloud runtime requires box.backend=nsjail') runtime_endpoint = str(box_config.get('runtime', {}).get('endpoint', '') or '').strip() diff --git a/src/langbot/pkg/cloud/model_catalog.py b/src/langbot/pkg/cloud/model_catalog.py new file mode 100644 index 000000000..a76376505 --- /dev/null +++ b/src/langbot/pkg/cloud/model_catalog.py @@ -0,0 +1,320 @@ +from __future__ import annotations + +import asyncio +import uuid +from datetime import datetime +from typing import Any, Literal, Protocol, runtime_checkable + +import sqlalchemy +from pydantic import BaseModel, ConfigDict, Field, SecretStr, field_validator + +from ..entity.persistence import model as persistence_model + + +LANGBOT_MODELS_PROVIDER_REQUESTER = 'space-chat-completions' +LANGBOT_MODELS_PROVIDER_NAME = 'LangBot Models' +_MODEL_RESOURCE_NAMESPACE = uuid.UUID('94c703ca-1df5-4e91-bcd3-74ac65cb7921') +_SUPPORTED_CATEGORIES = {'chat', 'embedding', 'rerank'} +_MODEL_TABLES = ( + persistence_model.LLMModel, + persistence_model.EmbeddingModel, + persistence_model.RerankModel, +) + + +class CloudModelCatalogItem(BaseModel): + model_config = ConfigDict(extra='forbid', frozen=True) + + uuid: str = Field(min_length=1, max_length=255) + model_id: str = Field(min_length=1, max_length=255) + category: Literal['chat', 'embedding', 'rerank'] + llm_abilities: tuple[str, ...] = () + is_featured: bool = False + featured_order: int = 0 + + @field_validator('llm_abilities', mode='before') + @classmethod + def normalize_missing_abilities(cls, value: Any) -> Any: + return () if value is None else value + + @field_validator('llm_abilities') + @classmethod + def validate_abilities(cls, value: tuple[str, ...]) -> tuple[str, ...]: + if any(not item.strip() or len(item) > 64 for item in value): + raise ValueError('Model abilities must be non-empty strings of at most 64 characters') + if len(set(value)) != len(value): + raise ValueError('Model abilities must be unique') + return value + + +class CloudWorkspaceModelBilling(BaseModel): + model_config = ConfigDict(extra='forbid', frozen=True) + + workspace_uuid: str = Field(min_length=36, max_length=36) + owner_account_uuid: str | None = Field(default=None, min_length=36, max_length=36) + api_key: SecretStr | None = None + + @field_validator('workspace_uuid') + @classmethod + def validate_uuid(cls, value: str) -> str: + return str(uuid.UUID(value)) + + @field_validator('owner_account_uuid') + @classmethod + def validate_optional_uuid(cls, value: str | None) -> str | None: + return None if value is None else str(uuid.UUID(value)) + + +class CloudModelCatalogSnapshot(BaseModel): + model_config = ConfigDict(extra='forbid', frozen=True) + + instance_uuid: str = Field(min_length=1, max_length=255) + generated_at: datetime + base_url: str = Field(min_length=1, max_length=512) + models: tuple[CloudModelCatalogItem, ...] + workspaces: tuple[CloudWorkspaceModelBilling, ...] + + @field_validator('base_url') + @classmethod + def validate_base_url(cls, value: str) -> str: + normalized = value.rstrip('/') + if not normalized.startswith('https://'): + raise ValueError('Cloud model gateway base URL must use HTTPS') + return normalized + + @field_validator('models') + @classmethod + def validate_models(cls, value: tuple[CloudModelCatalogItem, ...]) -> tuple[CloudModelCatalogItem, ...]: + if len(value) > 500: + raise ValueError('Cloud model catalog exceeds 500 models') + identities = {(item.category, item.uuid) for item in value} + if len(identities) != len(value): + raise ValueError('Cloud model catalog contains duplicate model identities') + return value + + @field_validator('workspaces') + @classmethod + def validate_workspaces( + cls, value: tuple[CloudWorkspaceModelBilling, ...] + ) -> tuple[CloudWorkspaceModelBilling, ...]: + if len(value) > 10_000: + raise ValueError('Cloud model catalog exceeds 10000 Workspaces') + identities = {item.workspace_uuid for item in value} + if len(identities) != len(value): + raise ValueError('Cloud model catalog contains duplicate Workspaces') + return value + + +@runtime_checkable +class CloudModelCatalogProvider(Protocol): + async def fetch_model_catalog(self, instance_uuid: str) -> CloudModelCatalogSnapshot: + """Fetch and verify the complete model catalog and Workspace billing projection.""" + ... + + +def system_provider_uuid(workspace_uuid: str) -> str: + workspace = str(uuid.UUID(workspace_uuid)) + return str(uuid.uuid5(_MODEL_RESOURCE_NAMESPACE, f'{workspace}:provider:{LANGBOT_MODELS_PROVIDER_REQUESTER}')) + + +def system_model_uuid(workspace_uuid: str, category: str, upstream_uuid: str) -> str: + workspace = str(uuid.UUID(workspace_uuid)) + if category not in _SUPPORTED_CATEGORIES: + raise ValueError(f'Unsupported model category: {category}') + if not upstream_uuid: + raise ValueError('Upstream model UUID is required') + return str(uuid.uuid5(_MODEL_RESOURCE_NAMESPACE, f'{workspace}:model:{category}:{upstream_uuid}')) + + +class CloudModelCatalogSyncService: + """Reconcile Space-owned model catalog and Owner billing tokens into every Cloud Workspace.""" + + def __init__( + self, + ap: Any, + provider: CloudModelCatalogProvider, + instance_uuid: str, + *, + sync_interval_seconds: float = 3600.0, + ) -> None: + if not isinstance(provider, CloudModelCatalogProvider): + raise TypeError('Cloud model catalog sync requires a CloudModelCatalogProvider') + if sync_interval_seconds < 10: + raise ValueError('Cloud model catalog sync interval must be at least 10 seconds') + self.ap = ap + self.provider = provider + self.instance_uuid = instance_uuid + self.sync_interval_seconds = float(sync_interval_seconds) + # A tenant UoW commits one Workspace at a time. Keep a durable in-memory + # convergence marker so a failed runtime reload is retried even when the + # following database reconciliation is a no-op. + self._runtime_reload_pending = False + + async def initialize(self) -> None: + await self.sync_once(reload_runtime=False) + + async def run(self) -> None: + while True: + await asyncio.sleep(self.sync_interval_seconds) + try: + await self.sync_once(reload_runtime=True) + except asyncio.CancelledError: + raise + except Exception as exc: + # Exception messages can contain rendered SQL bound values, + # including provider API keys. Log only the exception class. + self.ap.logger.warning(f'Cloud model catalog synchronization failed ({type(exc).__name__})') + + async def sync_once(self, *, reload_runtime: bool = True) -> dict[str, int]: + summary = {'workspaces': 0, 'created': 0, 'updated': 0, 'deleted': 0} + snapshot: CloudModelCatalogSnapshot | None = None + sync_error: Exception | None = None + reload_error: Exception | None = None + try: + snapshot = await self.provider.fetch_model_catalog(self.instance_uuid) + if snapshot.instance_uuid != self.instance_uuid: + raise ValueError('Cloud model catalog targets another LangBot instance') + + bindings = await self.ap.workspace_service.list_active_execution_bindings() + billing_by_workspace = {item.workspace_uuid: item for item in snapshot.workspaces} + missing = sorted( + binding.workspace_uuid for binding in bindings if binding.workspace_uuid not in billing_by_workspace + ) + if missing: + raise ValueError( + f'Cloud model catalog is missing billing projections for {len(missing)} active Workspaces' + ) + + for binding in bindings: + counts = await self._sync_workspace( + binding.workspace_uuid, + snapshot, + billing_by_workspace[binding.workspace_uuid], + ) + summary['workspaces'] += 1 + workspace_changed = any(counts[key] > 0 for key in ('created', 'updated', 'deleted')) + if workspace_changed: + # _sync_workspace returns only after its tenant UoW commits. + self._runtime_reload_pending = True + for key in ('created', 'updated', 'deleted'): + summary[key] += counts[key] + except Exception as exc: + sync_error = exc + finally: + model_mgr = getattr(self.ap, 'model_mgr', None) + if reload_runtime and self._runtime_reload_pending and model_mgr is not None: + try: + await model_mgr.load_models_from_db() + except Exception as exc: + reload_error = exc + else: + self._runtime_reload_pending = False + + if sync_error is not None: + if reload_error is not None: + raise sync_error from reload_error + raise sync_error + if reload_error is not None: + raise reload_error + + changed = any(summary[key] > 0 for key in ('created', 'updated', 'deleted')) + if changed and snapshot is not None: + self.ap.logger.info( + 'Cloud model catalog synchronized ' + f'({summary["workspaces"]} Workspaces, {len(snapshot.models)} models, ' + f'created={summary["created"]}, updated={summary["updated"]}, deleted={summary["deleted"]})' + ) + return summary + + async def _sync_workspace( + self, + workspace_uuid: str, + snapshot: CloudModelCatalogSnapshot, + billing: CloudWorkspaceModelBilling, + ) -> dict[str, int]: + counts = {'created': 0, 'updated': 0, 'deleted': 0} + provider_uuid = system_provider_uuid(workspace_uuid) + desired_keys = [billing.api_key.get_secret_value()] if billing.api_key is not None else [] + + async with self.ap.persistence_mgr.tenant_uow(workspace_uuid) as uow: + provider = await uow.session.scalar( + sqlalchemy.select(persistence_model.ModelProvider).where( + persistence_model.ModelProvider.uuid == provider_uuid + ) + ) + provider_values = { + 'workspace_uuid': workspace_uuid, + 'name': LANGBOT_MODELS_PROVIDER_NAME, + 'requester': LANGBOT_MODELS_PROVIDER_REQUESTER, + 'base_url': snapshot.base_url, + 'api_keys': desired_keys, + } + if provider is None: + provider = persistence_model.ModelProvider(uuid=provider_uuid, **provider_values) + uow.session.add(provider) + await uow.session.flush() + counts['created'] += 1 + elif self._update_entity(provider, provider_values): + counts['updated'] += 1 + + existing_by_table: dict[type, dict[str, Any]] = {} + for table in _MODEL_TABLES: + rows = ( + await uow.session.scalars(sqlalchemy.select(table).where(table.provider_uuid == provider_uuid)) + ).all() + existing_by_table[table] = {row.uuid: row for row in rows} + + desired_ids: dict[type, set[str]] = {table: set() for table in _MODEL_TABLES} + for item in snapshot.models: + table, values = self._model_values(workspace_uuid, provider_uuid, item) + model_uuid = system_model_uuid(workspace_uuid, item.category, item.uuid) + desired_ids[table].add(model_uuid) + existing = existing_by_table[table].get(model_uuid) + if existing is None: + uow.session.add(table(uuid=model_uuid, **values)) + counts['created'] += 1 + elif self._update_entity(existing, values): + counts['updated'] += 1 + + for table, entities in existing_by_table.items(): + for model_uuid, entity in entities.items(): + if model_uuid not in desired_ids[table]: + await uow.session.delete(entity) + counts['deleted'] += 1 + + return counts + + @staticmethod + def _update_entity(entity: Any, values: dict[str, Any]) -> bool: + changed = False + for key, value in values.items(): + if getattr(entity, key) != value: + setattr(entity, key, value) + changed = True + return changed + + @staticmethod + def _model_values( + workspace_uuid: str, + provider_uuid: str, + item: CloudModelCatalogItem, + ) -> tuple[type, dict[str, Any]]: + ranking = 100 - item.featured_order if item.is_featured else 0 + common = { + 'workspace_uuid': workspace_uuid, + 'name': item.model_id, + 'provider_uuid': provider_uuid, + 'extra_args': {}, + 'prefered_ranking': ranking, + } + if item.category == 'chat': + return persistence_model.LLMModel, { + **common, + 'abilities': list(item.llm_abilities), + 'context_length': None, + } + if item.category == 'embedding': + return persistence_model.EmbeddingModel, common + if item.category == 'rerank': + return persistence_model.RerankModel, common + raise ValueError(f'Unsupported model category: {item.category}') diff --git a/src/langbot/pkg/core/app.py b/src/langbot/pkg/core/app.py index bfd01387d..3f4a17329 100644 --- a/src/langbot/pkg/core/app.py +++ b/src/langbot/pkg/core/app.py @@ -54,6 +54,7 @@ from ..cloud import launch as cloud_launch_module from ..cloud import support_admin as cloud_support_admin_module from ..cloud import directory_projection as cloud_directory_projection_module from ..cloud import entitlements as cloud_entitlements_module +from ..cloud import model_catalog as cloud_model_catalog_module from ..api.http.context import ExecutionContext, PrincipalContext, PrincipalType @@ -142,13 +143,12 @@ class Application: deployment: cloud_bootstrap_module.OpenSourceDeployment | cloud_bootstrap_module.VerifiedCloudDeployment = None deployment_admission: cloud_bootstrap_module.DeploymentAdmissionGuard = None - + directory_projection_service: cloud_directory_projection_module.DirectoryProjectionService | None = None + cloud_model_catalog_service: cloud_model_catalog_module.CloudModelCatalogSyncService | None = None manifest_refresh_service: cloud_bootstrap_module.CloudManifestRefreshService | None = None entitlement_resolver: cloud_entitlements_module.EntitlementResolver | None = None - directory_projection_service: cloud_directory_projection_module.DirectoryProjectionService | None = None - vector_db_mgr: vectordb_mgr.VectorDBManager = None http_ctrl: http_controller.HTTPController = None @@ -306,6 +306,12 @@ class Application: name='cloud-directory-projection', scopes=[core_entities.LifecycleControlScope.APPLICATION], ) + if self.cloud_model_catalog_service is not None: + self.task_mgr.create_task( + self.cloud_model_catalog_service.run(), + name='cloud-model-catalog-sync', + scopes=[core_entities.LifecycleControlScope.APPLICATION], + ) if self.manifest_refresh_service is not None: self.task_mgr.create_task( self.manifest_refresh_service.run(), diff --git a/src/langbot/pkg/core/stages/build_app.py b/src/langbot/pkg/core/stages/build_app.py index b08fdfaf7..14d533dcc 100644 --- a/src/langbot/pkg/core/stages/build_app.py +++ b/src/langbot/pkg/core/stages/build_app.py @@ -46,6 +46,7 @@ from ...cloud import support_admin as cloud_support_admin_module from ...cloud.directory import directory_projection_limits_from_config from ...cloud.directory_projection import DirectoryProjectionService from ...cloud.entitlements import EntitlementResolver +from ...cloud.model_catalog import CloudModelCatalogSyncService from ...api.http.context import ExecutionContext, PrincipalContext, PrincipalType from ...api.http.authz import WorkspaceRequiredError @@ -176,6 +177,16 @@ class BuildAppStage(stage.BootingStage): # of repeating tenant validation for every manager. await workspace_service_inst.prime_startup_execution_bindings() + if not isinstance(deployment, cloud_bootstrap.VerifiedCloudDeployment): + raise RuntimeError('Multi-Workspace runtime requires a verified Cloud deployment') + cloud_model_catalog_service = CloudModelCatalogSyncService( + ap, + deployment.model_catalog_provider, + constants.instance_id, + ) + await cloud_model_catalog_service.initialize() + ap.cloud_model_catalog_service = cloud_model_catalog_service + ap.workspace_collaboration_service = workspace_collaboration_module.WorkspaceCollaborationService( ap, workspace_service_inst, diff --git a/src/langbot/pkg/persistence/tenant_uow.py b/src/langbot/pkg/persistence/tenant_uow.py index 53be402fa..0e460f335 100644 --- a/src/langbot/pkg/persistence/tenant_uow.py +++ b/src/langbot/pkg/persistence/tenant_uow.py @@ -209,7 +209,7 @@ _ALLOWED_SCOPED_BUILTIN_FUNCTION_TYPES = { 'now': sqlalchemy.sql.functions.now, 'sum': sqlalchemy.sql.functions.sum, } -_ALLOWED_SCOPED_GENERIC_FUNCTIONS = frozenset({'length', 'nullif'}) +_ALLOWED_SCOPED_GENERIC_FUNCTIONS = frozenset({'date_trunc', 'length', 'nullif'}) _ALLOWED_SCOPED_CUSTOM_OPERATORS = frozenset({'<=>'}) _ALLOWED_SCOPED_STATEMENT_TYPES = ( sqlalchemy.sql.dml.UpdateBase, diff --git a/src/langbot/pkg/pipeline/controller.py b/src/langbot/pkg/pipeline/controller.py index d14288c29..1502ae4de 100644 --- a/src/langbot/pkg/pipeline/controller.py +++ b/src/langbot/pkg/pipeline/controller.py @@ -132,9 +132,7 @@ class Controller: break - if selected_query: # 找到了 - queries.remove(selected_query) - else: # 没找到 说明:没有请求 或者 所有query对应的session都已达到并发上限 + if not selected_query: # 没找到 说明:没有请求 或者 所有query对应的session都已达到并发上限 await self.ap.query_pool.condition.wait() continue diff --git a/src/langbot/pkg/platform/sources/websocket_adapter.py b/src/langbot/pkg/platform/sources/websocket_adapter.py index 6176e2d93..0752225c6 100644 --- a/src/langbot/pkg/platform/sources/websocket_adapter.py +++ b/src/langbot/pkg/platform/sources/websocket_adapter.py @@ -707,28 +707,37 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter) if len(listener_tasks) >= 100: await self.logger.warning('WebSocket inbound listener capacity reached; dropping message') return - token = _current_pipeline_uuid.set(pipeline_uuid) - try: - task_manager = getattr(self.ap, 'task_mgr', None) - if task_manager is None or not isinstance(getattr(task_manager, 'tasks', None), list): - listener_task = asyncio.create_task(listeners[event.__class__](event, callback_adapter)) - else: - listener_task = task_manager.create_task( - listeners[event.__class__](event, callback_adapter), - kind='websocket-message', - name=f'websocket-message-{connection.connection_id}', - scopes=[ - core_entities.LifecycleControlScope.APPLICATION, - core_entities.LifecycleControlScope.PLATFORM, - ], - instance_uuid=connection.instance_uuid, - workspace_uuid=connection.workspace_uuid, - placement_generation=connection.placement_generation, - ).task - listener_tasks.add(listener_task) - listener_task.add_done_callback(self._listener_task_done) - finally: - _current_pipeline_uuid.reset(token) + listener = typing.cast( + typing.Callable[[typing.Any, typing.Any], typing.Awaitable[None]], + listeners[event.__class__], + ) + + async def run_listener(): + token = _current_pipeline_uuid.set(pipeline_uuid) + try: + await listener(event, callback_adapter) + finally: + _current_pipeline_uuid.reset(token) + + listener_coro = run_listener() + task_manager = getattr(self.ap, 'task_mgr', None) + if task_manager is None or not isinstance(getattr(task_manager, 'tasks', None), list): + listener_task = asyncio.create_task(listener_coro) + else: + listener_task = task_manager.create_task( + listener_coro, + kind='websocket-message', + name=f'websocket-message-{connection.connection_id}', + scopes=[ + core_entities.LifecycleControlScope.APPLICATION, + core_entities.LifecycleControlScope.PLATFORM, + ], + instance_uuid=connection.instance_uuid, + workspace_uuid=connection.workspace_uuid, + placement_generation=connection.placement_generation, + ).task + listener_tasks.add(listener_task) + listener_task.add_done_callback(self._listener_task_done) def get_websocket_messages( self, diff --git a/tests/integration/api/test_smoke.py b/tests/integration/api/test_smoke.py index 9c2927503..dfd5054f2 100644 --- a/tests/integration/api/test_smoke.py +++ b/tests/integration/api/test_smoke.py @@ -9,6 +9,8 @@ Run: uv run pytest tests/integration/api/test_smoke.py -q from __future__ import annotations +from types import SimpleNamespace + import pytest from unittest.mock import MagicMock, AsyncMock, Mock @@ -304,12 +306,34 @@ class TestUserInitEndpoint: data = await response.get_json() assert data['data'] == { 'initialized': True, + 'authenticated_invitation_acceptance_enabled': False, 'password_login_enabled': True, 'space_login_enabled': False, } fake_api_app.user_service.get_login_capabilities.assert_awaited_once_with() fake_api_app.user_service.get_first_user.assert_not_awaited() + @pytest.mark.asyncio + async def test_account_info_enables_authenticated_invitation_acceptance_in_cloud( + self, quart_test_client, fake_api_app + ): + fake_api_app.deployment = SimpleNamespace(mode='cloud') + fake_api_app.user_service.is_initialized.return_value = True + fake_api_app.user_service.get_login_capabilities = AsyncMock( + return_value={'password_login_enabled': True, 'space_login_enabled': True} + ) + + response = await quart_test_client.get('/api/v1/user/account-info') + + assert response.status_code == 200 + data = await response.get_json() + assert data['data'] == { + 'initialized': True, + 'authenticated_invitation_acceptance_enabled': True, + 'password_login_enabled': False, + 'space_login_enabled': True, + } + @pytest.mark.asyncio async def test_recovery_key_resets_any_existing_account(self, quart_test_client, fake_api_app, monkeypatch): fake_api_app.user_service.is_initialized.return_value = True diff --git a/tests/integration/api/test_user_space_oauth.py b/tests/integration/api/test_user_space_oauth.py index deb1e93e2..35e048053 100644 --- a/tests/integration/api/test_user_space_oauth.py +++ b/tests/integration/api/test_user_space_oauth.py @@ -272,9 +272,10 @@ async def test_space_credits_are_resolved_from_workspace_owner(space_oauth_api): '/api/v1/user/space-credits', headers={'Authorization': 'Bearer account-token', 'X-Workspace-Id': WORKSPACE_UUID}, ) + payload = await response.get_json() assert response.status_code == 200 - assert (await response.get_json())['data'] == { + assert payload['data'] == { 'credits': 25000, 'owner_space_bound': True, 'is_workspace_owner': True, @@ -282,6 +283,28 @@ async def test_space_credits_are_resolved_from_workspace_owner(space_oauth_api): application.space_service.get_credits.assert_awaited_once_with('owner@example.com') +@pytest.mark.asyncio +async def test_cloud_workspace_owner_is_always_space_bound_after_login(space_oauth_api): + application, client = space_oauth_api + application.deployment.mode = 'cloud' + application.user_service.get_workspace_owner = AsyncMock(return_value=None) + application.space_service.get_credits = AsyncMock() + + response = await client.get( + '/api/v1/user/space-credits', + headers={'Authorization': 'Bearer account-token', 'X-Workspace-Id': WORKSPACE_UUID}, + ) + payload = await response.get_json() + + assert response.status_code == 200 + assert payload['data'] == { + 'credits': None, + 'owner_space_bound': True, + 'is_workspace_owner': True, + } + application.space_service.get_credits.assert_not_awaited() + + @pytest.mark.asyncio async def test_bind_callback_uses_opaque_state_and_never_treats_it_as_jwt(space_oauth_api): application, client = space_oauth_api diff --git a/tests/unit_tests/api/service/test_cloud_managed_model_protection.py b/tests/unit_tests/api/service/test_cloud_managed_model_protection.py new file mode 100644 index 000000000..62ed06049 --- /dev/null +++ b/tests/unit_tests/api/service/test_cloud_managed_model_protection.py @@ -0,0 +1,112 @@ +"""Cloud Runtime write protection for the managed LangBot Models catalog.""" + +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from langbot.pkg.api.http.service import model as model_service_module +from langbot.pkg.api.http.service.model import ( + EmbeddingModelsService, + LLMModelsService, + RerankModelsService, + _assert_cloud_managed_provider_mutable, +) +from langbot.pkg.cloud.model_catalog import LANGBOT_MODELS_PROVIDER_REQUESTER + + +WORKSPACE = 'workspace-a' +PROVIDER = 'managed-provider' +MODEL = 'managed-model' + + +@pytest.mark.asyncio +async def test_managed_provider_guard_is_cloud_only(monkeypatch) -> None: + async def managed_provider(_ap, _context, provider_uuid): + assert provider_uuid == PROVIDER + return {'uuid': PROVIDER, 'requester': LANGBOT_MODELS_PROVIDER_REQUESTER} + + monkeypatch.setattr(model_service_module, '_require_workspace_provider', managed_provider) + application = SimpleNamespace(persistence_mgr=SimpleNamespace(mode=SimpleNamespace(value='cloud_runtime'))) + + with pytest.raises(ValueError, match='managed by Cloud'): + await _assert_cloud_managed_provider_mutable( + application, + WORKSPACE, + PROVIDER, + ) + + application.persistence_mgr.mode.value = 'normal' + await _assert_cloud_managed_provider_mutable( + application, + WORKSPACE, + PROVIDER, + ) + + +@pytest.mark.parametrize( + ('service_type', 'create_method', 'model_data'), + [ + (LLMModelsService, 'create_llm_model', {'provider_uuid': PROVIDER, 'name': 'chat', 'abilities': []}), + (EmbeddingModelsService, 'create_embedding_model', {'provider_uuid': PROVIDER, 'name': 'embedding'}), + (RerankModelsService, 'create_rerank_model', {'provider_uuid': PROVIDER, 'name': 'rerank'}), + ], +) +@pytest.mark.asyncio +async def test_all_model_types_reject_creation_under_managed_provider( + monkeypatch, + service_type, + create_method: str, + model_data: dict, +) -> None: + guard = AsyncMock(side_effect=ValueError('LangBot Models is managed by Cloud and cannot be modified')) + monkeypatch.setattr(model_service_module, '_assert_cloud_managed_provider_mutable', guard) + application = SimpleNamespace( + persistence_mgr=SimpleNamespace(), + provider_service=SimpleNamespace( + get_provider=AsyncMock(return_value={'uuid': PROVIDER, 'requester': LANGBOT_MODELS_PROVIDER_REQUESTER}) + ), + model_mgr=None, + ) + service = service_type(application) + + with pytest.raises(ValueError, match='managed by Cloud'): + await getattr(service, create_method)(WORKSPACE, model_data) + + guard.assert_awaited_once() + + +@pytest.mark.parametrize( + ('service_type', 'get_method', 'write_method', 'payload'), + [ + (LLMModelsService, 'get_llm_model', 'update_llm_model', {'name': 'changed'}), + (LLMModelsService, 'get_llm_model', 'delete_llm_model', None), + (EmbeddingModelsService, 'get_embedding_model', 'update_embedding_model', {'name': 'changed'}), + (EmbeddingModelsService, 'get_embedding_model', 'delete_embedding_model', None), + (RerankModelsService, 'get_rerank_model', 'update_rerank_model', {'name': 'changed'}), + (RerankModelsService, 'get_rerank_model', 'delete_rerank_model', None), + ], +) +@pytest.mark.asyncio +async def test_all_model_types_reject_update_and_delete_for_managed_provider( + monkeypatch, + service_type, + get_method: str, + write_method: str, + payload: dict | None, +) -> None: + guard = AsyncMock(side_effect=ValueError('LangBot Models is managed by Cloud and cannot be modified')) + monkeypatch.setattr(model_service_module, '_assert_cloud_managed_provider_mutable', guard) + application = SimpleNamespace(persistence_mgr=SimpleNamespace(mode=SimpleNamespace(value='cloud_runtime'))) + service = service_type(application) + monkeypatch.setattr( + service, + get_method, + AsyncMock(return_value={'uuid': MODEL, 'provider_uuid': PROVIDER, 'extra_args': {}}), + ) + + args = (WORKSPACE, MODEL) if payload is None else (WORKSPACE, MODEL, payload) + with pytest.raises(ValueError, match='managed by Cloud'): + await getattr(service, write_method)(*args) + + guard.assert_awaited_once() diff --git a/tests/unit_tests/api/service/test_provider_service.py b/tests/unit_tests/api/service/test_provider_service.py index 15b995895..fdf8b01bc 100644 --- a/tests/unit_tests/api/service/test_provider_service.py +++ b/tests/unit_tests/api/service/test_provider_service.py @@ -25,6 +25,7 @@ from langbot.pkg.workspace.errors import WorkspaceNotFoundError pytestmark = pytest.mark.asyncio WORKSPACE_UUID = 'workspace-a' +SYSTEM_REQUESTER = 'space-chat-completions' def _create_mock_provider( @@ -1005,3 +1006,56 @@ class TestProviderSecretRoundtrip: ) ap.persistence_mgr.execute_async.assert_not_awaited() + + +class TestCloudManagedProviderProtection: + @staticmethod + def _service() -> ModelProviderService: + ap = SimpleNamespace( + persistence_mgr=SimpleNamespace( + mode=SimpleNamespace(value='cloud_runtime'), + execute_async=AsyncMock(), + ), + model_mgr=SimpleNamespace(), + ) + return ModelProviderService(ap) + + async def test_cloud_rejects_user_created_system_requester(self): + service = self._service() + + with pytest.raises(ValueError, match='reserved'): + await service.create_provider( + WORKSPACE_UUID, + { + 'name': 'Fake LangBot Models', + 'requester': SYSTEM_REQUESTER, + 'base_url': 'https://example.invalid/v1', + 'api_keys': ['fake'], + }, + ) + + with pytest.raises(ValueError, match='reserved'): + await service.find_or_create_provider( + WORKSPACE_UUID, + SYSTEM_REQUESTER, + 'https://api.langbot.cloud/v1', + ['fake'], + ) + service.ap.persistence_mgr.execute_async.assert_not_awaited() + + async def test_cloud_rejects_update_and_delete_of_managed_provider(self): + service = self._service() + service.get_provider = AsyncMock( + return_value={'uuid': 'system-provider', 'requester': SYSTEM_REQUESTER} + ) + + with pytest.raises(ValueError, match='managed by Cloud'): + await service.update_provider(WORKSPACE_UUID, 'system-provider', {'name': 'Renamed'}) + with pytest.raises(ValueError, match='managed by Cloud'): + await service.delete_provider(WORKSPACE_UUID, 'system-provider') + service.ap.persistence_mgr.execute_async.assert_not_awaited() + + async def test_oss_does_not_reserve_space_requester(self): + ap = SimpleNamespace(persistence_mgr=SimpleNamespace(mode=SimpleNamespace(value='oss_compat'))) + service = ModelProviderService(ap) + assert service._system_requester_is_reserved(SYSTEM_REQUESTER) is False diff --git a/tests/unit_tests/cloud/test_bootstrap.py b/tests/unit_tests/cloud/test_bootstrap.py index 77ccefa5d..6225ba7b7 100644 --- a/tests/unit_tests/cloud/test_bootstrap.py +++ b/tests/unit_tests/cloud/test_bootstrap.py @@ -66,6 +66,10 @@ class _Provider: def __init__(self): self.manifest_provider = _Manifest() + async def fetch_model_catalog(self, instance_uuid: str): + del instance_uuid + raise AssertionError('not used by bootstrap contract tests') + def bootstrap(self, *, instance_uuid: str, instance_config: dict): del instance_config return VerifiedCloudDeployment( @@ -79,6 +83,7 @@ class _Provider: entitlement_provider=_Entitlements(), directory_provider=_Directory(), manifest_provider=self.manifest_provider, + model_catalog_provider=self, verification_key_id='root-2026', ) @@ -228,10 +233,23 @@ async def test_cloud_pgvector_contract_is_fail_closed(pgvector_config, message): ) +async def test_cloud_runtime_allows_explicitly_disabled_box(): + config = _cloud_config() + config['box']['enabled'] = False + + deployment = await resolve_deployment( + instance_uuid='instance-a', + instance_config=config, + entry_points=lambda: _EntryPoints([_EntryPoint(_Provider())]), + now=1_000, + ) + + assert isinstance(deployment, VerifiedCloudDeployment) + + @pytest.mark.parametrize( ('mutate', 'message'), [ - (lambda config: config['box'].update(enabled=False), 'box.enabled=true'), (lambda config: config['box'].update(backend='docker'), 'box.backend=nsjail'), (lambda config: config['box']['runtime'].update(endpoint=''), 'box.runtime.endpoint'), ( diff --git a/tests/unit_tests/cloud/test_model_catalog.py b/tests/unit_tests/cloud/test_model_catalog.py new file mode 100644 index 000000000..6f4b0a275 --- /dev/null +++ b/tests/unit_tests/cloud/test_model_catalog.py @@ -0,0 +1,372 @@ +from __future__ import annotations + +import asyncio +import logging +from datetime import UTC, datetime +from types import SimpleNamespace + +import pytest +import sqlalchemy +from sqlalchemy.ext.asyncio import create_async_engine + +from langbot.pkg.cloud.model_catalog import ( + CloudModelCatalogSnapshot, + CloudModelCatalogSyncService, + system_model_uuid, + system_provider_uuid, +) +from langbot.pkg.entity.persistence.base import Base +from langbot.pkg.entity.persistence.model import EmbeddingModel, LLMModel, ModelProvider +from langbot.pkg.entity.persistence.workspace import Workspace +from langbot.pkg.persistence.mgr import PersistenceManager, PersistenceMode + + +pytestmark = pytest.mark.asyncio +INSTANCE_UUID = 'instance-model-catalog' +WORKSPACE_A = '00000000-0000-4000-8000-000000000001' +WORKSPACE_B = '00000000-0000-4000-8000-000000000002' +OWNER_A = '10000000-0000-4000-8000-000000000001' +OWNER_B = '10000000-0000-4000-8000-000000000002' + + +class _CatalogProvider: + def __init__(self, snapshot: CloudModelCatalogSnapshot) -> None: + self.snapshot = snapshot + + async def fetch_model_catalog(self, instance_uuid: str) -> CloudModelCatalogSnapshot: + assert instance_uuid == INSTANCE_UUID + return self.snapshot + + +def _snapshot( + *, + key_a: str | None = 'owner-a-key', + model_id: str = 'gpt-test', + include_embedding: bool = True, +) -> CloudModelCatalogSnapshot: + models = [ + { + 'uuid': 'upstream-chat', + 'model_id': model_id, + 'category': 'chat', + 'llm_abilities': ['chat', 'vision'], + 'is_featured': True, + 'featured_order': 7, + } + ] + if include_embedding: + models.append( + { + 'uuid': 'upstream-embedding', + 'model_id': 'embedding-test', + 'category': 'embedding', + } + ) + return CloudModelCatalogSnapshot.model_validate( + { + 'instance_uuid': INSTANCE_UUID, + 'generated_at': datetime.now(UTC), + 'base_url': 'https://api.langbot.cloud/v1/', + 'models': models, + 'workspaces': [ + { + 'workspace_uuid': WORKSPACE_A, + 'owner_account_uuid': OWNER_A, + 'api_key': key_a, + }, + { + 'workspace_uuid': WORKSPACE_B, + 'owner_account_uuid': OWNER_B, + 'api_key': 'owner-b-key', + }, + ], + } + ) + + +async def test_catalog_snapshot_treats_null_model_abilities_as_empty() -> None: + payload = _snapshot().model_dump(mode='json') + payload['models'][0]['llm_abilities'] = None + + snapshot = CloudModelCatalogSnapshot.model_validate(payload) + + assert snapshot.models[0].llm_abilities == () + + +async def test_catalog_reconciles_every_workspace_idempotently_and_tracks_owner_and_downlisting(tmp_path) -> None: + engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "model-catalog.db"}') + manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME) + manager.db = SimpleNamespace(get_engine=lambda: engine) + bindings = [ + SimpleNamespace(instance_uuid=INSTANCE_UUID, workspace_uuid=WORKSPACE_A, placement_generation=1), + SimpleNamespace(instance_uuid=INSTANCE_UUID, workspace_uuid=WORKSPACE_B, placement_generation=1), + ] + workspace_service = SimpleNamespace(list_active_execution_bindings=lambda: _async_value(bindings)) + reload_counter = _AsyncCounter() + runtime_reload = SimpleNamespace(load_models_from_db=reload_counter) + app = SimpleNamespace( + persistence_mgr=manager, + workspace_service=workspace_service, + model_mgr=runtime_reload, + logger=logging.getLogger(__name__), + ) + provider = _CatalogProvider(_snapshot()) + service = CloudModelCatalogSyncService(app, provider, INSTANCE_UUID) + + try: + async with engine.begin() as connection: + await connection.run_sync(Base.metadata.create_all) + await connection.execute( + sqlalchemy.insert(Workspace), + [ + { + 'uuid': WORKSPACE_A, + 'instance_uuid': INSTANCE_UUID, + 'name': 'A', + 'slug': 'a', + 'source': 'cloud_projection', + }, + { + 'uuid': WORKSPACE_B, + 'instance_uuid': INSTANCE_UUID, + 'name': 'B', + 'slug': 'b', + 'source': 'cloud_projection', + }, + ], + ) + await connection.execute( + sqlalchemy.insert(ModelProvider).values( + uuid='custom-provider', + workspace_uuid=WORKSPACE_A, + name='Custom', + requester='openai-chat-completions', + base_url='https://custom.example/v1', + api_keys=['custom-key'], + ) + ) + await connection.execute( + sqlalchemy.insert(LLMModel).values( + uuid='custom-model', + workspace_uuid=WORKSPACE_A, + name='custom-model', + provider_uuid='custom-provider', + abilities=['chat'], + extra_args={}, + prefered_ranking=0, + ) + ) + + first = await service.sync_once() + assert first == {'workspaces': 2, 'created': 6, 'updated': 0, 'deleted': 0} + assert reload_counter.calls == 1 + + async with engine.connect() as connection: + providers = ( + await connection.execute( + sqlalchemy.select( + ModelProvider.uuid, + ModelProvider.workspace_uuid, + ModelProvider.api_keys, + ).where(ModelProvider.requester == 'space-chat-completions') + ) + ).all() + assert {item.workspace_uuid for item in providers} == {WORKSPACE_A, WORKSPACE_B} + assert {item.uuid for item in providers} == { + system_provider_uuid(WORKSPACE_A), + system_provider_uuid(WORKSPACE_B), + } + assert {item.workspace_uuid: item.api_keys for item in providers} == { + WORKSPACE_A: ['owner-a-key'], + WORKSPACE_B: ['owner-b-key'], + } + assert await connection.scalar(sqlalchemy.select(sqlalchemy.func.count()).select_from(LLMModel)) == 3 + assert await connection.scalar(sqlalchemy.select(sqlalchemy.func.count()).select_from(EmbeddingModel)) == 2 + + second = await service.sync_once() + assert second == {'workspaces': 2, 'created': 0, 'updated': 0, 'deleted': 0} + assert reload_counter.calls == 1 + + provider.snapshot = _snapshot( + key_a='new-owner-key', + model_id='gpt-renamed', + include_embedding=False, + ) + third = await service.sync_once() + assert third == {'workspaces': 2, 'created': 0, 'updated': 3, 'deleted': 2} + assert reload_counter.calls == 2 + + async with engine.connect() as connection: + provider_a_keys = await connection.scalar( + sqlalchemy.select(ModelProvider.api_keys).where(ModelProvider.uuid == system_provider_uuid(WORKSPACE_A)) + ) + assert provider_a_keys == ['new-owner-key'] + system_model_names = ( + ( + await connection.execute( + sqlalchemy.select(LLMModel.name).where( + LLMModel.provider_uuid.in_( + [system_provider_uuid(WORKSPACE_A), system_provider_uuid(WORKSPACE_B)] + ) + ) + ) + ) + .scalars() + .all() + ) + assert set(system_model_names) == {'gpt-renamed'} + assert await connection.scalar(sqlalchemy.select(sqlalchemy.func.count()).select_from(EmbeddingModel)) == 0 + assert ( + await connection.scalar( + sqlalchemy.select(sqlalchemy.func.count()) + .select_from(ModelProvider) + .where(ModelProvider.uuid == 'custom-provider') + ) + == 1 + ) + assert ( + await connection.scalar( + sqlalchemy.select(sqlalchemy.func.count()) + .select_from(LLMModel) + .where(LLMModel.uuid == 'custom-model') + ) + == 1 + ) + + provider.snapshot = _snapshot(key_a=None, model_id='gpt-renamed', include_embedding=False) + fourth = await service.sync_once() + assert fourth == {'workspaces': 2, 'created': 0, 'updated': 1, 'deleted': 0} + assert reload_counter.calls == 3 + async with engine.connect() as connection: + provider_a_keys = await connection.scalar( + sqlalchemy.select(ModelProvider.api_keys).where(ModelProvider.uuid == system_provider_uuid(WORKSPACE_A)) + ) + assert provider_a_keys == [] + finally: + await engine.dispose() + + +def test_workspace_scoped_ids_are_stable_and_secrets_are_redacted() -> None: + assert system_provider_uuid(WORKSPACE_A) == system_provider_uuid(WORKSPACE_A) + assert system_provider_uuid(WORKSPACE_A) != system_provider_uuid(WORKSPACE_B) + assert system_model_uuid(WORKSPACE_A, 'chat', 'upstream') != system_model_uuid(WORKSPACE_B, 'chat', 'upstream') + snapshot = _snapshot() + assert 'owner-a-key' not in repr(snapshot) + + +async def test_snapshot_must_cover_every_active_workspace() -> None: + snapshot = _snapshot().model_copy(update={'workspaces': _snapshot().workspaces[:1]}) + app = SimpleNamespace( + workspace_service=SimpleNamespace( + list_active_execution_bindings=lambda: _async_value( + [SimpleNamespace(workspace_uuid=WORKSPACE_A), SimpleNamespace(workspace_uuid=WORKSPACE_B)] + ) + ), + logger=logging.getLogger(__name__), + ) + service = CloudModelCatalogSyncService(app, _CatalogProvider(snapshot), INSTANCE_UUID) + with pytest.raises(ValueError, match='missing billing projections for 1 active Workspaces'): + await service.sync_once() + + +async def _async_value(value): + return value + + +class _AsyncCounter: + def __init__(self) -> None: + self.calls = 0 + + async def __call__(self) -> None: + self.calls += 1 + + +async def test_partial_workspace_failure_reloads_already_committed_changes() -> None: + bindings = [ + SimpleNamespace(workspace_uuid=WORKSPACE_A), + SimpleNamespace(workspace_uuid=WORKSPACE_B), + ] + reload_counter = _AsyncCounter() + app = SimpleNamespace( + workspace_service=SimpleNamespace(list_active_execution_bindings=lambda: _async_value(bindings)), + model_mgr=SimpleNamespace(load_models_from_db=reload_counter), + logger=logging.getLogger(__name__), + ) + service = CloudModelCatalogSyncService(app, _CatalogProvider(_snapshot()), INSTANCE_UUID) + calls = 0 + + async def sync_workspace(*_args): + nonlocal calls + calls += 1 + if calls == 1: + return {'created': 1, 'updated': 0, 'deleted': 0} + raise RuntimeError('second Workspace failed') + + service._sync_workspace = sync_workspace # type: ignore[method-assign] + + with pytest.raises(RuntimeError, match='second Workspace failed'): + await service.sync_once() + assert reload_counter.calls == 1 + + +async def test_failed_runtime_reload_is_retried_after_noop_sync() -> None: + bindings = [SimpleNamespace(workspace_uuid=WORKSPACE_A)] + + class _FlakyReload: + def __init__(self) -> None: + self.calls = 0 + + async def __call__(self) -> None: + self.calls += 1 + if self.calls == 1: + raise RuntimeError('reload failed') + + runtime_reload = _FlakyReload() + app = SimpleNamespace( + workspace_service=SimpleNamespace(list_active_execution_bindings=lambda: _async_value(bindings)), + model_mgr=SimpleNamespace(load_models_from_db=runtime_reload), + logger=logging.getLogger(__name__), + ) + service = CloudModelCatalogSyncService(app, _CatalogProvider(_snapshot()), INSTANCE_UUID) + calls = 0 + + async def sync_workspace(*_args): + nonlocal calls + calls += 1 + if calls == 1: + return {'created': 1, 'updated': 0, 'deleted': 0} + return {'created': 0, 'updated': 0, 'deleted': 0} + + service._sync_workspace = sync_workspace # type: ignore[method-assign] + + with pytest.raises(RuntimeError, match='reload failed'): + await service.sync_once() + summary = await service.sync_once() + assert summary == {'workspaces': 1, 'created': 0, 'updated': 0, 'deleted': 0} + assert runtime_reload.calls == 2 + + +async def test_background_sync_log_redacts_exception_message(caplog) -> None: + secret = 'owner-secret-api-key' + attempted = asyncio.Event() + + class _FailingProvider: + async def fetch_model_catalog(self, instance_uuid: str) -> CloudModelCatalogSnapshot: + del instance_uuid + attempted.set() + raise RuntimeError(f'database parameters include {secret}') + + app = SimpleNamespace(logger=logging.getLogger(__name__)) + service = CloudModelCatalogSyncService(app, _FailingProvider(), INSTANCE_UUID) + service.sync_interval_seconds = 0.001 + task = asyncio.create_task(service.run()) + try: + await asyncio.wait_for(attempted.wait(), timeout=1) + await asyncio.sleep(0.01) + finally: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert secret not in caplog.text + assert 'Cloud model catalog synchronization failed (RuntimeError)' in caplog.text diff --git a/tests/unit_tests/persistence/test_tenant_uow.py b/tests/unit_tests/persistence/test_tenant_uow.py index 7be62bf2e..bc8369f08 100644 --- a/tests/unit_tests/persistence/test_tenant_uow.py +++ b/tests/unit_tests/persistence/test_tenant_uow.py @@ -961,6 +961,7 @@ async def test_scoped_session_rejects_raw_or_unapproved_sql( sa.select(sa.func.coalesce(sa.func.sum(sa.literal(1)), sa.literal(0))), sa.select( sa.func.now(), + sa.func.date_trunc('hour', sa.column('timestamp')), sa.func.length(sa.literal('value')), sa.func.nullif(sa.literal('value'), sa.literal('')), ), diff --git a/tests/unit_tests/pipeline/test_controller_tenancy.py b/tests/unit_tests/pipeline/test_controller_tenancy.py index b99e54bb2..ac999e56b 100644 --- a/tests/unit_tests/pipeline/test_controller_tenancy.py +++ b/tests/unit_tests/pipeline/test_controller_tenancy.py @@ -11,6 +11,7 @@ from sqlalchemy.ext.asyncio import create_async_engine from langbot.pkg.persistence.mgr import PersistenceManager, PersistenceMode from langbot.pkg.persistence.tenant_uow import PersistenceScopeKind from langbot.pkg.pipeline.controller import Controller +from langbot.pkg.pipeline.pool import QueryPool from langbot.pkg.workspace.errors import WorkspaceGenerationMismatchError @@ -143,3 +144,31 @@ async def test_controller_revalidates_generation_before_running_pipeline( runtime_pipeline.run.assert_awaited_once_with(sample_query) query_pool.remove_query.assert_awaited_once_with(sample_query) session._semaphore.release.assert_called_once_with() + + +@pytest.mark.asyncio +async def test_controller_schedules_query_without_removing_it_twice(mock_app, sample_query): + query_pool = QueryPool() + query_pool.queries.append(sample_query) + mock_app.query_pool = query_pool + mock_app.sess_mgr.get_session = AsyncMock(return_value=SimpleNamespace(_semaphore=asyncio.Semaphore(1))) + + scheduler_errors: list[str] = [] + + def stop_on_scheduler_error(message): + scheduler_errors.append(str(message)) + raise asyncio.CancelledError + + def stop_after_scheduling(process_coro, **_kwargs): + process_coro.close() + raise asyncio.CancelledError + + mock_app.logger.error.side_effect = stop_on_scheduler_error + mock_app.task_mgr.create_task.side_effect = stop_after_scheduling + controller = Controller(mock_app) + + with pytest.raises(asyncio.CancelledError): + await controller.consumer() + + assert scheduler_errors == [] + assert query_pool.queries == [] diff --git a/tests/unit_tests/platform/test_websocket_session_isolation.py b/tests/unit_tests/platform/test_websocket_session_isolation.py index 682a42184..958ba5888 100644 --- a/tests/unit_tests/platform/test_websocket_session_isolation.py +++ b/tests/unit_tests/platform/test_websocket_session_isolation.py @@ -1,6 +1,7 @@ """Regression tests for isolated embed-widget conversations.""" import asyncio +import contextvars from pathlib import Path from unittest.mock import AsyncMock, Mock @@ -204,6 +205,48 @@ async def test_embed_event_uses_stable_session_launcher(monkeypatch): assert received[0].sender.id == f'websocket_pipeline-1:{session_id}' +@pytest.mark.asyncio +async def test_pipeline_override_survives_detached_listener_task(monkeypatch): + manager = WebSocketConnectionManager() + connection = await manager.add_connection( + websocket=Mock(), + scope=SCOPE_A, + pipeline_uuid='pipeline-1', + session_type='person', + ) + monkeypatch.setattr(websocket_adapter_module, 'ws_connection_manager', manager) + + class DetachedTaskManager: + def __init__(self): + self.tasks = [] + + def create_task(self, coro, **_kwargs): + task = asyncio.create_task(coro, context=contextvars.Context()) + self.tasks.append(task) + return Mock(task=task) + + task_manager = DetachedTaskManager() + adapter = WebSocketAdapter.model_construct( + ap=Mock(task_mgr=task_manager), + logger=_adapter_logger(), + ) + adapter.websocket_person_session = WebSocketSession(id='person') + adapter.websocket_group_session = WebSocketSession(id='group') + pipeline_overrides = [] + + async def listener(_event, callback_adapter): + pipeline_overrides.append(callback_adapter.get_pipeline_uuid_override()) + + adapter.listeners = {platform_events.FriendMessage: listener} + await adapter.handle_websocket_message( + connection, + {'message': [{'type': 'Plain', 'text': 'hello'}], 'stream': False}, + ) + await asyncio.gather(*task_manager.tasks) + + assert pipeline_overrides == ['pipeline-1'] + + @pytest.mark.asyncio async def test_embed_group_event_uses_stable_session_launcher(monkeypatch): manager = WebSocketConnectionManager() diff --git a/uv.lock b/uv.lock index a23277e25..c2f7bef76 100644 --- a/uv.lock +++ b/uv.lock @@ -1999,7 +1999,7 @@ wheels = [ [[package]] name = "langbot" -version = "4.10.6" +version = "4.10.7" source = { editable = "." } dependencies = [ { name = "aiocqhttp" }, @@ -2116,7 +2116,7 @@ requires-dist = [ { name = "ebooklib", specifier = ">=0.18" }, { name = "gewechat-client", specifier = ">=0.1.5" }, { name = "html2text", specifier = ">=2024.2.26" }, - { name = "langbot-plugin", git = "https://github.com/langbot-app/langbot-plugin-sdk.git?rev=1d65ed301a6afc52150a998043f73cd6032c8162" }, + { name = "langbot-plugin", git = "https://github.com/langbot-app/langbot-plugin-sdk.git?rev=101e453e916b39465a6294d6471c9eaae8725d5c" }, { name = "langchain", specifier = ">=1.3.9" }, { name = "langchain-core", specifier = ">=1.3.3" }, { name = "langchain-text-splitters", specifier = ">=1.1.2" }, @@ -2182,8 +2182,8 @@ dev = [ [[package]] name = "langbot-plugin" -version = "0.4.18" -source = { git = "https://github.com/langbot-app/langbot-plugin-sdk.git?rev=1d65ed301a6afc52150a998043f73cd6032c8162#1d65ed301a6afc52150a998043f73cd6032c8162" } +version = "0.5.0" +source = { git = "https://github.com/langbot-app/langbot-plugin-sdk.git?rev=101e453e916b39465a6294d6471c9eaae8725d5c#101e453e916b39465a6294d6471c9eaae8725d5c" } dependencies = [ { name = "aiofiles" }, { name = "aiohttp" }, diff --git a/web/src/app/auth/space/callback/page.tsx b/web/src/app/auth/space/callback/page.tsx index 1028de6ac..c2c142589 100644 --- a/web/src/app/auth/space/callback/page.tsx +++ b/web/src/app/auth/space/callback/page.tsx @@ -5,6 +5,7 @@ import { beginAuthenticatedSession, beginSupportAdminSession, bootstrapWorkspaceSession, + clearPendingInvitationToken, getPendingInvitationToken, } from '@/app/infra/http'; import { toast } from 'sonner'; @@ -66,6 +67,10 @@ function SpaceOAuthCallbackContent() { const [searchParams] = useSearchParams(); const { t } = useTranslation(); const isMountedRef = useRef(true); + const directLaunchFragmentRef = useRef<{ + workspaceUuid: string | null; + launchAssertion: string | null; + } | null>(null); const [status, setStatus] = useState< 'loading' | 'confirm' | 'success' | 'error' @@ -108,8 +113,31 @@ function SpaceOAuthCallbackContent() { } beginAuthenticatedSession(response.token, response.user); - if (getPendingInvitationToken()) { - navigate('/invitations/accept', { replace: true }); + const invitationToken = getPendingInvitationToken(); + if (invitationToken) { + let invitation; + try { + invitation = + await httpClient.acceptWorkspaceInvitation(invitationToken); + } catch (error) { + const code = (error as { code?: string }).code; + const path = code + ? `/invitations/accept?error=${encodeURIComponent(code)}` + : '/invitations/accept'; + navigate(path, { replace: true }); + return; + } + + beginAuthenticatedSession(invitation.token, response.user); + clearPendingInvitationToken(); + const workspaceResult = await bootstrapWorkspaceSession({ + preferredWorkspaceUuid: invitation.workspace_uuid, + }); + if (workspaceResult.status === 'unavailable') { + navigate('/workspace-unavailable', { replace: true }); + return; + } + navigate('/home', { replace: true }); return; } const workspaceResult = await bootstrapWorkspaceSession({ @@ -220,8 +248,28 @@ function SpaceOAuthCallbackContent() { const errorDescription = searchParams.get('error_description'); const mode = searchParams.get('mode'); const state = searchParams.get('state'); - const workspaceUuid = searchParams.get('workspace_uuid'); - const launchAssertion = searchParams.get('launch_assertion'); + if (directLaunchFragmentRef.current === null) { + const fragmentParams = new URLSearchParams( + window.location.hash.startsWith('#') + ? window.location.hash.slice(1) + : window.location.hash, + ); + directLaunchFragmentRef.current = { + workspaceUuid: fragmentParams.get('workspace_uuid'), + launchAssertion: fragmentParams.get('launch_assertion'), + }; + if (window.location.hash) { + window.history.replaceState( + null, + '', + `${window.location.pathname}${window.location.search}`, + ); + } + } + const workspaceUuid = + directLaunchFragmentRef.current.workspaceUuid ?? + searchParams.get('workspace_uuid'); + const launchAssertion = directLaunchFragmentRef.current.launchAssertion; if (error) { setStatus('error'); diff --git a/web/src/app/infra/http/BackendClient.ts b/web/src/app/infra/http/BackendClient.ts index 63d6f37c8..03e0899f6 100644 --- a/web/src/app/infra/http/BackendClient.ts +++ b/web/src/app/infra/http/BackendClient.ts @@ -1179,6 +1179,7 @@ export class BackendClient extends BaseHttpClient { public getAccountInfo(): Promise<{ initialized: boolean; + authenticated_invitation_acceptance_enabled?: boolean; password_login_enabled?: boolean; space_login_enabled?: boolean; }> { diff --git a/web/src/app/invitations/accept/page.tsx b/web/src/app/invitations/accept/page.tsx index e8d412ecc..0d83331b5 100644 --- a/web/src/app/invitations/accept/page.tsx +++ b/web/src/app/invitations/accept/page.tsx @@ -93,6 +93,10 @@ export default function AcceptInvitationPage() { const [confirmPassword, setConfirmPassword] = useState(''); const [passwordRegistrationEnabled, setPasswordRegistrationEnabled] = useState(false); + const [ + authenticatedInvitationAcceptanceEnabled, + setAuthenticatedInvitationAcceptanceEnabled, + ] = useState(false); useEffect(() => { const handleHashChange = () => setInvitationHash(window.location.hash); @@ -113,6 +117,9 @@ export default function AcceptInvitationPage() { .getAccountInfo() .then((info) => { setPasswordRegistrationEnabled(info.password_login_enabled !== false); + setAuthenticatedInvitationAcceptanceEnabled( + info.authenticated_invitation_acceptance_enabled === true, + ); }) .catch(() => setPasswordRegistrationEnabled(false)); if (!invitationToken) { @@ -304,7 +311,18 @@ export default function AcceptInvitationPage() { )} - {hasLoginToken ? ( + {hasLoginToken && authenticatedInvitationAcceptanceEnabled ? ( + + ) : hasLoginToken ? (
{t('workspace.authenticatedInvitationNotice')} diff --git a/web/tests/e2e/invitations.spec.ts b/web/tests/e2e/invitations.spec.ts index 5e8b217ec..b036838e2 100644 --- a/web/tests/e2e/invitations.spec.ts +++ b/web/tests/e2e/invitations.spec.ts @@ -165,3 +165,164 @@ test('an authenticated OSS invitation requires logout before registration', asyn invitation: 'logout-invitation', }); }); + +test('an authenticated Cloud Account can accept its invitation directly', async ({ + page, +}) => { + await installLangBotApiMocks(page, { + authenticated: true, + storage: { + token: 'invited-account-token', + userEmail: 'invited@example.com', + }, + }); + await page.route('**/api/v1/user/account-info', async (route) => { + await route.fulfill({ + status: 200, + contentType: 'application/json', + body: JSON.stringify({ + code: 0, + data: { + initialized: true, + authenticated_invitation_acceptance_enabled: true, + password_login_enabled: false, + space_login_enabled: true, + }, + msg: 'ok', + }), + }); + }); + await page.route('**/api/v1/invitations/inspect', async (route) => { + await route.fulfill({ + status: 200, + contentType: 'application/json', + body: JSON.stringify({ + code: 0, + data: { + invitation: { + uuid: 'cloud-invitation', + workspace_uuid: 'workspace-playwright', + normalized_email: 'invited@example.com', + role: 'viewer', + status: 'pending', + }, + workspace: { + uuid: 'workspace-playwright', + name: 'Playwright Workspace', + }, + }, + msg: 'ok', + }), + }); + }); + + let acceptanceAuthorization = ''; + await page.route('**/api/v1/invitations/accept', async (route) => { + acceptanceAuthorization = route.request().headers().authorization ?? ''; + await route.fulfill({ + status: 200, + contentType: 'application/json', + body: JSON.stringify({ + code: 0, + data: { + token: 'accepted-cloud-account-token', + workspace_uuid: 'workspace-playwright', + }, + msg: 'ok', + }), + }); + }); + + await page.goto('/invitations/accept#token=cloud-invitation'); + + await expect( + page.getByRole('button', { name: 'Accept Invitation' }), + ).toBeVisible(); + await page.getByRole('button', { name: 'Accept Invitation' }).click(); + await expect(page).toHaveURL(/\/home(?:\/monitoring)?$/); + expect(acceptanceAuthorization).toBe('Bearer invited-account-token'); +}); + +test('Space OAuth accepts a pending invitation with the freshly authenticated account', async ({ + page, +}) => { + await installLangBotApiMocks(page, { + authenticated: false, + storage: { + token: 'stale-other-account-token', + userEmail: 'other@example.com', + }, + }); + await page.addInitScript(() => { + sessionStorage.setItem( + 'langbot_pending_invitation_token', + 'matching-invitation', + ); + }); + + await page.route('**/api/v1/user/space/callback', async (route) => { + await route.fulfill({ + status: 200, + contentType: 'application/json', + body: JSON.stringify({ + code: 0, + data: { + token: 'fresh-invited-account-token', + user: 'invited@example.com', + }, + msg: 'ok', + }), + }); + }); + await page.route('**/api/v1/user/info', async (route) => { + await route.fulfill({ + status: 200, + contentType: 'application/json', + body: JSON.stringify({ + code: 0, + data: { + account_uuid: 'invited-account', + user: 'invited@example.com', + account_type: 'space', + has_password: false, + }, + msg: 'ok', + }), + }); + }); + + let acceptanceAuthorization = ''; + await page.route('**/api/v1/invitations/accept', async (route) => { + acceptanceAuthorization = route.request().headers().authorization ?? ''; + await route.fulfill({ + status: 200, + contentType: 'application/json', + body: JSON.stringify({ + code: 0, + data: { + token: 'accepted-invited-account-token', + workspace_uuid: 'workspace-playwright', + }, + msg: 'ok', + }), + }); + }); + + await page.goto('/auth/space/callback?code=oauth-code&state=oauth-state'); + + await expect(page).toHaveURL(/\/home(?:\/monitoring)?$/, { + timeout: 5_000, + }); + expect(acceptanceAuthorization).toBe('Bearer fresh-invited-account-token'); + expect( + await page.evaluate(() => ({ + token: localStorage.getItem('token'), + userEmail: localStorage.getItem('userEmail'), + invitation: sessionStorage.getItem('langbot_pending_invitation_token'), + })), + ).toEqual({ + token: 'accepted-invited-account-token', + userEmail: 'invited@example.com', + invitation: null, + }); +}); diff --git a/web/tests/unit/space-launch-callback.test.mjs b/web/tests/unit/space-launch-callback.test.mjs new file mode 100644 index 000000000..7debabf6c --- /dev/null +++ b/web/tests/unit/space-launch-callback.test.mjs @@ -0,0 +1,24 @@ +import assert from 'node:assert/strict'; +import fs from 'node:fs'; +import test from 'node:test'; + +const source = fs.readFileSync( + new URL('../../src/app/auth/space/callback/page.tsx', import.meta.url), + 'utf8', +); + +test('direct launch assertion is fragment-only and removed before exchange', () => { + assert.doesNotMatch(source, /searchParams\.get\(['"]launch_assertion['"]\)/); + const readIndex = source.indexOf("fragmentParams.get('launch_assertion')"); + const clearIndex = source.indexOf('window.history.replaceState'); + const exchangeIndex = source.indexOf('handleOAuthCallback(', clearIndex); + assert.ok(readIndex >= 0, 'fragment assertion read is missing'); + assert.ok( + clearIndex > readIndex, + 'URL fragment is not cleared after copying the assertion', + ); + assert.ok( + exchangeIndex > clearIndex, + 'assertion exchange starts before the fragment is cleared', + ); +});